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            && let Err(e) = registry.check_for_changes().await
299        {
300            tracing::warn!(error = %e, "Failed to check for tool changes");
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            && let Err(e) = registry.check_for_changes().await
410        {
411            tracing::warn!(error = %e, "Failed to check for tool changes (streaming)");
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                    let mut route_resp = route_handler.handle(boxed_req).await;
428                    #[cfg(feature = "cors")]
429                    if let Some(ref cors_config) = self.cors_config {
430                        inject_cors_headers(
431                            &mut route_resp,
432                            cors_config,
433                            request_origin.as_deref(),
434                        )
435                        .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
436                    }
437                    return Ok(route_resp);
438                }
439                Ok(None) => {} // No match, continue to MCP handler
440                Err(e) => {
441                    debug!("Route validation error (streaming): {}", e);
442                    let mut err_resp = e.into_response();
443                    #[cfg(feature = "cors")]
444                    if let Some(ref cors_config) = self.cors_config {
445                        inject_cors_headers(
446                            &mut err_resp,
447                            cors_config,
448                            request_origin.as_deref(),
449                        )
450                        .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
451                    }
452                    return Ok(err_resp);
453                }
454            }
455        }
456
457        // 🚀 PROTOCOL ROUTING: Check protocol version and route to appropriate handler
458        use turul_http_mcp_server::protocol::McpProtocolVersion;
459        let protocol_version = hyper_req
460            .headers()
461            .get("MCP-Protocol-Version")
462            .and_then(|h| h.to_str().ok())
463            .and_then(McpProtocolVersion::parse_version)
464            .unwrap_or(McpProtocolVersion::V2025_06_18);
465
466        // Route based on protocol version
467        let hyper_resp = if protocol_version.supports_streamable_http() {
468            // Use StreamableHttpHandler for MCP 2025-11-25 (proper headers, chunked SSE)
469            debug!(
470                "Using StreamableHttpHandler for protocol {}",
471                protocol_version.to_string()
472            );
473            self.streamable_handler.handle_request(hyper_req).await
474        } else {
475            // Legacy protocol: use SessionMcpHandler
476            debug!(
477                "Using SessionMcpHandler for legacy protocol {}",
478                protocol_version.to_string()
479            );
480            self.session_handler
481                .handle_mcp_request(hyper_req)
482                .await
483                .map_err(|e| {
484                    Box::new(crate::error::LambdaError::McpFramework(e.to_string()))
485                        as Box<dyn std::error::Error + Send + Sync>
486                })?
487        };
488
489        // 🚀 DELEGATION: Convert hyper response to Lambda streaming response (preserves streaming!)
490        let mut lambda_resp = crate::adapter::hyper_to_lambda_streaming(hyper_resp);
491
492        // Apply CORS headers if configured (Lambda-specific logic)
493        #[cfg(feature = "cors")]
494        if let Some(ref cors_config) = self.cors_config {
495            inject_cors_headers(&mut lambda_resp, cors_config, request_origin.as_deref())
496                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
497        }
498
499        Ok(lambda_resp)
500    }
501
502    /// Convert Lambda response to streaming format (helper for CORS preflight)
503    fn convert_lambda_response_to_streaming(
504        &self,
505        lambda_response: LambdaResponse<LambdaBody>,
506    ) -> lambda_http::Response<http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>>
507    {
508        use bytes::Bytes;
509        use http_body_util::{BodyExt, Full};
510
511        let (parts, body) = lambda_response.into_parts();
512        let body_bytes = match body {
513            LambdaBody::Empty => Bytes::new(),
514            LambdaBody::Text(text) => Bytes::from(text),
515            LambdaBody::Binary(bytes) => Bytes::from(bytes),
516            _ => Bytes::new(),
517        };
518
519        // Map error type from Infallible to hyper::Error
520        let streaming_body = Full::new(body_bytes)
521            .map_err(|e: std::convert::Infallible| match e {})
522            .boxed_unsync();
523
524        lambda_http::Response::from_parts(parts, streaming_body)
525    }
526}
527
528#[cfg(test)]
529mod tests {
530    use super::*;
531    use http::Request;
532    use turul_mcp_session_storage::InMemorySessionStorage;
533
534    #[tokio::test]
535    async fn test_handler_creation() {
536        let session_storage = Arc::new(InMemorySessionStorage::new());
537        let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
538        let dispatcher = JsonRpcDispatcher::new();
539        let config = ServerConfig::default();
540        let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
541        let capabilities = ServerCapabilities::default();
542
543        let handler = LambdaMcpHandler::new(
544            dispatcher,
545            session_storage,
546            stream_manager,
547            config,
548            StreamConfig::default(),
549            implementation,
550            capabilities,
551            false, // SSE disabled for test
552            #[cfg(feature = "cors")]
553            None,
554        );
555
556        // Test that handler was created successfully
557        assert!(!handler.sse_enabled);
558    }
559
560    #[tokio::test]
561    async fn test_sse_enabled_with_handle_works() {
562        let session_storage = Arc::new(InMemorySessionStorage::new());
563        let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
564        let dispatcher = JsonRpcDispatcher::new();
565        let config = ServerConfig::default();
566        let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
567        let capabilities = ServerCapabilities::default();
568
569        // Create handler with SSE enabled
570        let handler = LambdaMcpHandler::new(
571            dispatcher,
572            session_storage,
573            stream_manager,
574            config,
575            StreamConfig::default(),
576            implementation,
577            capabilities,
578            true, // SSE enabled - should work with handle() for snapshot-based SSE
579            #[cfg(feature = "cors")]
580            None,
581        );
582
583        // Create a test Lambda request
584        let lambda_req = Request::builder()
585            .method("POST")
586            .uri("/mcp")
587            .body(LambdaBody::Text(
588                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
589            ))
590            .unwrap();
591
592        // handle() should work (provides snapshot-based SSE rather than real-time streaming)
593        let result = handler.handle(lambda_req).await;
594        assert!(
595            result.is_ok(),
596            "handle() should work with SSE enabled for snapshot-based responses"
597        );
598    }
599
600    /// Test that verifies StreamConfig is properly threaded through the delegation
601    #[tokio::test]
602    async fn test_stream_config_preservation() {
603        let session_storage = Arc::new(InMemorySessionStorage::new());
604        let dispatcher = JsonRpcDispatcher::new();
605        let config = ServerConfig::default();
606        let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
607        let capabilities = ServerCapabilities::default();
608
609        // Create a custom StreamConfig with non-default values
610        let custom_stream_config = StreamConfig {
611            channel_buffer_size: 1024,      // Non-default value (default is 1000)
612            max_replay_events: 200,         // Non-default value (default is 100)
613            keepalive_interval_seconds: 10, // Non-default value (default is 30)
614            cors_origin: "https://custom-test.example.com".to_string(), // Non-default value
615        };
616
617        // Create stream manager with the custom config
618        let stream_manager = Arc::new(StreamManager::with_config(
619            session_storage.clone(),
620            custom_stream_config.clone(),
621        ));
622
623        let handler = LambdaMcpHandler::new(
624            dispatcher,
625            session_storage,
626            stream_manager,
627            config,
628            custom_stream_config.clone(),
629            implementation,
630            capabilities,
631            false, // SSE disabled for test
632            #[cfg(feature = "cors")]
633            None,
634        );
635
636        // The handler should be created successfully, proving the StreamConfig was accepted
637        assert!(!handler.sse_enabled);
638
639        // Verify that the stream manager has the custom configuration
640        let stream_manager = handler.get_stream_manager();
641
642        // Verify the StreamConfig values were propagated correctly
643        let actual_config = stream_manager.get_config();
644
645        assert_eq!(
646            actual_config.channel_buffer_size, custom_stream_config.channel_buffer_size,
647            "Custom channel_buffer_size was not propagated correctly"
648        );
649        assert_eq!(
650            actual_config.max_replay_events, custom_stream_config.max_replay_events,
651            "Custom max_replay_events was not propagated correctly"
652        );
653        assert_eq!(
654            actual_config.keepalive_interval_seconds,
655            custom_stream_config.keepalive_interval_seconds,
656            "Custom keepalive_interval_seconds was not propagated correctly"
657        );
658        assert_eq!(
659            actual_config.cors_origin, custom_stream_config.cors_origin,
660            "Custom cors_origin was not propagated correctly"
661        );
662
663        // Verify the stream manager is accessible (proves delegation worked)
664        assert!(Arc::strong_count(stream_manager) >= 1);
665    }
666
667    /// Test the full builder → server → handler chain with StreamConfig
668    #[tokio::test]
669    async fn test_full_builder_chain_stream_config() {
670        use crate::LambdaMcpServerBuilder;
671        use turul_mcp_session_storage::InMemorySessionStorage;
672
673        // Create a custom StreamConfig with non-default values
674        let custom_stream_config = turul_http_mcp_server::StreamConfig {
675            channel_buffer_size: 2048,      // Non-default value
676            max_replay_events: 500,         // Non-default value
677            keepalive_interval_seconds: 15, // Non-default value
678            cors_origin: "https://full-chain-test.example.com".to_string(),
679        };
680
681        // Test the complete builder → server → handler chain
682        let server = LambdaMcpServerBuilder::new()
683            .name("full-chain-test")
684            .version("1.0.0")
685            .storage(Arc::new(InMemorySessionStorage::new()))
686            .sse(true) // Enable SSE to test streaming functionality
687            .stream_config(custom_stream_config.clone())
688            .build()
689            .await
690            .expect("Server should build successfully");
691
692        // Create handler from server (this is the critical chain step)
693        let handler = server
694            .handler()
695            .await
696            .expect("Handler should be created from server");
697
698        // Verify the handler was created successfully
699        assert!(handler.sse_enabled, "SSE should be enabled");
700
701        // Verify that the custom StreamConfig was preserved through the entire chain
702        let stream_manager = handler.get_stream_manager();
703        let actual_config = stream_manager.get_config();
704
705        assert_eq!(
706            actual_config.channel_buffer_size, custom_stream_config.channel_buffer_size,
707            "Custom channel_buffer_size should be preserved through builder → server → handler chain"
708        );
709        assert_eq!(
710            actual_config.max_replay_events, custom_stream_config.max_replay_events,
711            "Custom max_replay_events should be preserved through builder → server → handler chain"
712        );
713        assert_eq!(
714            actual_config.keepalive_interval_seconds,
715            custom_stream_config.keepalive_interval_seconds,
716            "Custom keepalive_interval_seconds should be preserved through builder → server → handler chain"
717        );
718        assert_eq!(
719            actual_config.cors_origin, custom_stream_config.cors_origin,
720            "Custom cors_origin should be preserved through builder → server → handler chain"
721        );
722
723        // Verify the stream manager is functional
724        assert!(
725            Arc::strong_count(stream_manager) >= 1,
726            "Stream manager should be properly initialized"
727        );
728
729        // Additional verification: Test that the configuration is actually used functionally
730        // by verifying the stream manager can be used with the custom configuration
731        let test_session_id = uuid::Uuid::now_v7().as_simple().to_string();
732
733        // The stream manager should be able to handle session operations with the custom config
734        // This verifies the config isn't just preserved but actually used
735        let subscriptions = stream_manager.get_subscriptions(&test_session_id).await;
736        assert!(
737            subscriptions.is_empty(),
738            "New session should have no subscriptions initially"
739        );
740
741        // Verify the stream manager was constructed with our custom config values
742        // This confirms the config propagated through the entire builder → server → handler chain
743        assert_eq!(
744            stream_manager.get_config().channel_buffer_size,
745            2048,
746            "Stream manager should be using the custom buffer size functionally"
747        );
748    }
749
750    /// Test matrix: 4 combinations of streaming runtime vs SSE configuration
751    /// This ensures we don't have runtime hangs or configuration conflicts
752    ///
753    /// Test 1: Non-streaming runtime + sse(false) - This should work (snapshot mode)
754    #[tokio::test]
755    async fn test_non_streaming_runtime_sse_false() {
756        use crate::LambdaMcpServerBuilder;
757        use turul_mcp_session_storage::InMemorySessionStorage;
758
759        let server = LambdaMcpServerBuilder::new()
760            .name("test-non-streaming-sse-false")
761            .version("1.0.0")
762            .storage(Arc::new(InMemorySessionStorage::new()))
763            .sse(false) // Disable SSE for non-streaming runtime
764            .build()
765            .await
766            .expect("Server should build successfully");
767
768        let handler = server
769            .handler()
770            .await
771            .expect("Handler should be created from server");
772
773        // Verify configuration
774        assert!(!handler.sse_enabled, "SSE should be disabled");
775
776        // Create a test request (POST /mcp works in all configs)
777        let lambda_req = Request::builder()
778            .method("POST")
779            .uri("/mcp")
780            .body(LambdaBody::Text(
781                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
782            ))
783            .unwrap();
784
785        // This should work without hanging
786        let result = handler.handle(lambda_req).await;
787        assert!(
788            result.is_ok(),
789            "POST /mcp should work with non-streaming + sse(false)"
790        );
791    }
792
793    /// Test 2: Non-streaming runtime + sse(true) - This should work (snapshot-based SSE)
794    #[tokio::test]
795    async fn test_non_streaming_runtime_sse_true() {
796        use crate::LambdaMcpServerBuilder;
797        use turul_mcp_session_storage::InMemorySessionStorage;
798
799        let server = LambdaMcpServerBuilder::new()
800            .name("test-non-streaming-sse-true")
801            .version("1.0.0")
802            .storage(Arc::new(InMemorySessionStorage::new()))
803            .sse(true) // Enable SSE for snapshot-based responses
804            .build()
805            .await
806            .expect("Server should build successfully");
807
808        let handler = server
809            .handler()
810            .await
811            .expect("Handler should be created from server");
812
813        // Verify configuration
814        assert!(handler.sse_enabled, "SSE should be enabled");
815
816        // Create a test request (POST /mcp works in all configs)
817        let lambda_req = Request::builder()
818            .method("POST")
819            .uri("/mcp")
820            .body(LambdaBody::Text(
821                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
822            ))
823            .unwrap();
824
825        // This should work without hanging (provides snapshot-based SSE)
826        let result = handler.handle(lambda_req).await;
827        assert!(
828            result.is_ok(),
829            "POST /mcp should work with non-streaming + sse(true)"
830        );
831
832        // Note: GET /mcp would provide snapshot events, not real-time streaming
833        // This is the key difference from handle_streaming()
834    }
835
836    /// Test 3: Streaming runtime + sse(false) - This should work (SSE disabled)
837    #[tokio::test]
838    async fn test_streaming_runtime_sse_false() {
839        use crate::LambdaMcpServerBuilder;
840        use turul_mcp_session_storage::InMemorySessionStorage;
841
842        let server = LambdaMcpServerBuilder::new()
843            .name("test-streaming-sse-false")
844            .version("1.0.0")
845            .storage(Arc::new(InMemorySessionStorage::new()))
846            .sse(false) // Disable SSE even with streaming runtime
847            .build()
848            .await
849            .expect("Server should build successfully");
850
851        let handler = server
852            .handler()
853            .await
854            .expect("Handler should be created from server");
855
856        // Verify configuration
857        assert!(!handler.sse_enabled, "SSE should be disabled");
858
859        // Create a test request for streaming handler
860        let lambda_req = Request::builder()
861            .method("POST")
862            .uri("/mcp")
863            .body(LambdaBody::Text(
864                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
865            ))
866            .unwrap();
867
868        // This should work with streaming runtime even when SSE is disabled
869        let result = handler.handle_streaming(lambda_req).await;
870        assert!(
871            result.is_ok(),
872            "Streaming runtime should work with sse(false)"
873        );
874    }
875
876    /// Test 4: Streaming runtime + sse(true) - This should work (real-time SSE streaming)
877    #[tokio::test]
878    async fn test_streaming_runtime_sse_true() {
879        use crate::LambdaMcpServerBuilder;
880        use turul_mcp_session_storage::InMemorySessionStorage;
881
882        let server = LambdaMcpServerBuilder::new()
883            .name("test-streaming-sse-true")
884            .version("1.0.0")
885            .storage(Arc::new(InMemorySessionStorage::new()))
886            .sse(true) // Enable SSE with streaming runtime for real-time streaming
887            .build()
888            .await
889            .expect("Server should build successfully");
890
891        let handler = server
892            .handler()
893            .await
894            .expect("Handler should be created from server");
895
896        // Verify configuration
897        assert!(handler.sse_enabled, "SSE should be enabled");
898
899        // Create a test request for streaming handler
900        let lambda_req = Request::builder()
901            .method("POST")
902            .uri("/mcp")
903            .body(LambdaBody::Text(
904                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
905            ))
906            .unwrap();
907
908        // This should work and provide real-time SSE streaming
909        let result = handler.handle_streaming(lambda_req).await;
910        assert!(
911            result.is_ok(),
912            "Streaming runtime should work with sse(true) for real-time streaming"
913        );
914
915        // Note: GET /mcp would provide real-time streaming events
916        // This is the optimal configuration for real-time notifications
917    }
918
919    // ── Strict lifecycle tests over handle_streaming() ────────────────
920
921    /// Helper: build a Lambda handler with strict lifecycle and a test tool via the builder.
922    async fn build_strict_streaming_handler() -> LambdaMcpHandler {
923        use crate::LambdaMcpServerBuilder;
924        use turul_mcp_session_storage::InMemorySessionStorage;
925
926        let server = LambdaMcpServerBuilder::new()
927            .name("lifecycle-test")
928            .version("1.0.0")
929            .tool(LifecycleTestTool)
930            .storage(Arc::new(InMemorySessionStorage::new()))
931            .strict_lifecycle(true) // explicit — survives default changes
932            .sse(true)
933            .build()
934            .await
935            .expect("build should succeed");
936
937        server.handler().await.expect("handler should succeed")
938    }
939
940    // Test tool for lifecycle tests — satisfies all required traits
941    #[derive(Clone, Default)]
942    struct LifecycleTestTool;
943
944    impl turul_mcp_builders::traits::HasBaseMetadata for LifecycleTestTool {
945        fn name(&self) -> &str {
946            "ping_tool"
947        }
948    }
949    impl turul_mcp_builders::traits::HasDescription for LifecycleTestTool {
950        fn description(&self) -> Option<&str> {
951            Some("test tool")
952        }
953    }
954    impl turul_mcp_builders::traits::HasInputSchema for LifecycleTestTool {
955        fn input_schema(&self) -> &turul_mcp_protocol::ToolSchema {
956            static SCHEMA: std::sync::OnceLock<turul_mcp_protocol::ToolSchema> =
957                std::sync::OnceLock::new();
958            SCHEMA.get_or_init(turul_mcp_protocol::ToolSchema::object)
959        }
960    }
961    impl turul_mcp_builders::traits::HasOutputSchema for LifecycleTestTool {
962        fn output_schema(&self) -> Option<&turul_mcp_protocol::ToolSchema> {
963            None
964        }
965    }
966    impl turul_mcp_builders::traits::HasAnnotations for LifecycleTestTool {
967        fn annotations(&self) -> Option<&turul_mcp_protocol::tools::ToolAnnotations> {
968            None
969        }
970    }
971    impl turul_mcp_builders::traits::HasToolMeta for LifecycleTestTool {
972        fn tool_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
973            None
974        }
975    }
976    impl turul_mcp_builders::traits::HasIcons for LifecycleTestTool {}
977    impl turul_mcp_builders::traits::HasExecution for LifecycleTestTool {}
978
979    #[async_trait::async_trait]
980    impl turul_mcp_server::McpTool for LifecycleTestTool {
981        async fn call(
982            &self,
983            _args: serde_json::Value,
984            _session: Option<turul_mcp_server::SessionContext>,
985        ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
986            Ok(turul_mcp_protocol::tools::CallToolResult::success(vec![
987                turul_mcp_protocol::tools::ToolResult::text("pong"),
988            ]))
989        }
990    }
991
992    /// Helper: create a Lambda POST request for handle_streaming()
993    fn streaming_mcp_request(body: &str, session_id: Option<&str>) -> LambdaRequest {
994        let mut builder = Request::builder()
995            .method("POST")
996            .uri("/mcp")
997            .header("Content-Type", "application/json")
998            .header("Accept", "application/json, text/event-stream")
999            .header("MCP-Protocol-Version", "2025-11-25");
1000
1001        if let Some(sid) = session_id {
1002            builder = builder.header("Mcp-Session-Id", sid);
1003        }
1004
1005        builder.body(LambdaBody::Text(body.to_string())).unwrap()
1006    }
1007
1008    /// Helper: collect streaming response body into a string
1009    async fn collect_streaming_body(
1010        response: lambda_http::Response<
1011            http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
1012        >,
1013    ) -> (http::StatusCode, String) {
1014        use http_body_util::BodyExt;
1015        let status = response.status();
1016        let session_id = response
1017            .headers()
1018            .get("Mcp-Session-Id")
1019            .and_then(|v| v.to_str().ok())
1020            .map(String::from);
1021        let body_bytes = response
1022            .into_body()
1023            .collect()
1024            .await
1025            .map(|c| c.to_bytes())
1026            .unwrap_or_default();
1027        let body_str = String::from_utf8_lossy(&body_bytes).to_string();
1028        let _ = session_id; // available if needed
1029        (status, body_str)
1030    }
1031
1032    /// Helper: extract session ID from a streaming response
1033    fn extract_session_id(
1034        response: &lambda_http::Response<
1035            http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
1036        >,
1037    ) -> Option<String> {
1038        response
1039            .headers()
1040            .get("Mcp-Session-Id")
1041            .and_then(|v| v.to_str().ok())
1042            .map(String::from)
1043    }
1044
1045    /// Helper: parse JSON from a response body (handles SSE "data: " prefix)
1046    fn parse_response_json(body: &str) -> serde_json::Value {
1047        // Strip SSE framing if present
1048        let json_str = body
1049            .lines()
1050            .find(|line| line.starts_with("data: "))
1051            .map(|line| &line[6..])
1052            .unwrap_or(body.trim());
1053        serde_json::from_str(json_str)
1054            .unwrap_or_else(|e| panic!("Failed to parse JSON from body: {e}\nBody: {body}"))
1055    }
1056
1057    /// P0: Full strict lifecycle handshake succeeds on handle_streaming()
1058    #[tokio::test]
1059    async fn test_lambda_streaming_strict_handshake_succeeds() {
1060        let handler = build_strict_streaming_handler().await;
1061
1062        // Step 1: initialize
1063        let init_req = streaming_mcp_request(
1064            &serde_json::json!({
1065                "jsonrpc": "2.0", "method": "initialize", "id": 1,
1066                "params": {
1067                    "protocolVersion": "2025-11-25",
1068                    "capabilities": {},
1069                    "clientInfo": { "name": "test", "version": "1.0.0" }
1070                }
1071            })
1072            .to_string(),
1073            None,
1074        );
1075        let init_resp = handler
1076            .handle_streaming(init_req)
1077            .await
1078            .expect("initialize should succeed");
1079        let session_id = extract_session_id(&init_resp).expect("must return session ID");
1080        let (status, _body) = collect_streaming_body(init_resp).await;
1081        assert_eq!(status, 200, "initialize should return 200");
1082
1083        // Step 2: notifications/initialized
1084        let notif_req = streaming_mcp_request(
1085            &serde_json::json!({
1086                "jsonrpc": "2.0",
1087                "method": "notifications/initialized",
1088                "params": {}
1089            })
1090            .to_string(),
1091            Some(&session_id),
1092        );
1093        let notif_resp = handler
1094            .handle_streaming(notif_req)
1095            .await
1096            .expect("notification should succeed");
1097        let (status, _) = collect_streaming_body(notif_resp).await;
1098        assert_eq!(status, 202, "notifications/initialized should return 202");
1099
1100        // Step 3: tools/list
1101        let list_req = streaming_mcp_request(
1102            &serde_json::json!({
1103                "jsonrpc": "2.0", "method": "tools/list", "id": 2
1104            })
1105            .to_string(),
1106            Some(&session_id),
1107        );
1108        let list_resp = handler
1109            .handle_streaming(list_req)
1110            .await
1111            .expect("tools/list should succeed");
1112        let (status, body) = collect_streaming_body(list_resp).await;
1113        assert_eq!(status, 200, "tools/list should return 200");
1114        let json = parse_response_json(&body);
1115        assert!(
1116            json["result"]["tools"].is_array(),
1117            "tools/list should return tools array: {json}"
1118        );
1119
1120        // Step 4: tools/call
1121        let call_req = streaming_mcp_request(
1122            &serde_json::json!({
1123                "jsonrpc": "2.0", "method": "tools/call", "id": 3,
1124                "params": { "name": "ping_tool", "arguments": {} }
1125            })
1126            .to_string(),
1127            Some(&session_id),
1128        );
1129        let call_resp = handler
1130            .handle_streaming(call_req)
1131            .await
1132            .expect("tools/call should succeed");
1133        let (status, body) = collect_streaming_body(call_resp).await;
1134        assert_eq!(status, 200, "tools/call should return 200");
1135        let json = parse_response_json(&body);
1136        assert!(
1137            json["result"].is_object(),
1138            "tools/call should return result: {json}"
1139        );
1140    }
1141
1142    /// P0: Strict lifecycle rejects both tools/list and tools/call before notifications/initialized
1143    #[tokio::test]
1144    async fn test_lambda_streaming_strict_rejects_before_initialized() {
1145        let handler = build_strict_streaming_handler().await;
1146
1147        // Initialize to get session
1148        let init_req = streaming_mcp_request(
1149            &serde_json::json!({
1150                "jsonrpc": "2.0", "method": "initialize", "id": 1,
1151                "params": {
1152                    "protocolVersion": "2025-11-25",
1153                    "capabilities": {},
1154                    "clientInfo": { "name": "test", "version": "1.0.0" }
1155                }
1156            })
1157            .to_string(),
1158            None,
1159        );
1160        let init_resp = handler.handle_streaming(init_req).await.unwrap();
1161        let session_id = extract_session_id(&init_resp).unwrap();
1162        let _ = collect_streaming_body(init_resp).await;
1163
1164        // tools/list without notifications/initialized — must fail
1165        let list_req = streaming_mcp_request(
1166            &serde_json::json!({
1167                "jsonrpc": "2.0", "method": "tools/list", "id": 2
1168            })
1169            .to_string(),
1170            Some(&session_id),
1171        );
1172        let list_resp = handler.handle_streaming(list_req).await.unwrap();
1173        let (_, body) = collect_streaming_body(list_resp).await;
1174        let json = parse_response_json(&body);
1175        assert!(
1176            json["error"].is_object(),
1177            "tools/list should return JSON-RPC error: {json}"
1178        );
1179        assert_eq!(
1180            json["error"]["code"].as_i64().unwrap(),
1181            -32031,
1182            "tools/list must return SessionError code -32031, got: {json}"
1183        );
1184        assert!(
1185            json["error"]["message"]
1186                .as_str()
1187                .unwrap()
1188                .contains("notifications/initialized"),
1189            "Error must mention notifications/initialized: {}",
1190            json["error"]["message"]
1191        );
1192
1193        // tools/call without notifications/initialized — must also fail
1194        let call_req = streaming_mcp_request(
1195            &serde_json::json!({
1196                "jsonrpc": "2.0", "method": "tools/call", "id": 3,
1197                "params": { "name": "ping_tool", "arguments": {} }
1198            })
1199            .to_string(),
1200            Some(&session_id),
1201        );
1202        let call_resp = handler.handle_streaming(call_req).await.unwrap();
1203        let (_, body) = collect_streaming_body(call_resp).await;
1204        let json = parse_response_json(&body);
1205        assert!(
1206            json["error"].is_object(),
1207            "tools/call should return JSON-RPC error: {json}"
1208        );
1209        assert_eq!(
1210            json["error"]["code"].as_i64().unwrap(),
1211            -32031,
1212            "tools/call must return SessionError code -32031, got: {json}"
1213        );
1214        assert!(
1215            json["error"]["message"]
1216                .as_str()
1217                .unwrap()
1218                .contains("notifications/initialized"),
1219            "Error must mention notifications/initialized: {}",
1220            json["error"]["message"]
1221        );
1222    }
1223
1224    /// P0: tools/list succeeds immediately after notifications/initialized (race fix proof)
1225    #[tokio::test]
1226    async fn test_lambda_streaming_initialized_is_effective_immediately() {
1227        let handler = build_strict_streaming_handler().await;
1228
1229        // Initialize
1230        let init_req = streaming_mcp_request(
1231            &serde_json::json!({
1232                "jsonrpc": "2.0", "method": "initialize", "id": 1,
1233                "params": {
1234                    "protocolVersion": "2025-11-25",
1235                    "capabilities": {},
1236                    "clientInfo": { "name": "test", "version": "1.0.0" }
1237                }
1238            })
1239            .to_string(),
1240            None,
1241        );
1242        let init_resp = handler.handle_streaming(init_req).await.unwrap();
1243        let session_id = extract_session_id(&init_resp).unwrap();
1244        let _ = collect_streaming_body(init_resp).await;
1245
1246        // notifications/initialized
1247        let notif_req = streaming_mcp_request(
1248            &serde_json::json!({
1249                "jsonrpc": "2.0",
1250                "method": "notifications/initialized",
1251                "params": {}
1252            })
1253            .to_string(),
1254            Some(&session_id),
1255        );
1256        let notif_resp = handler.handle_streaming(notif_req).await.unwrap();
1257        let (status, _) = collect_streaming_body(notif_resp).await;
1258        assert_eq!(status, 202);
1259
1260        // Immediately — no delay — send tools/list
1261        let list_req = streaming_mcp_request(
1262            &serde_json::json!({
1263                "jsonrpc": "2.0", "method": "tools/list", "id": 2
1264            })
1265            .to_string(),
1266            Some(&session_id),
1267        );
1268        let list_resp = handler.handle_streaming(list_req).await.unwrap();
1269        let (status, body) = collect_streaming_body(list_resp).await;
1270        assert_eq!(
1271            status, 200,
1272            "tools/list must succeed immediately after initialized"
1273        );
1274        let json = parse_response_json(&body);
1275        assert!(
1276            json["result"]["tools"].is_array(),
1277            "Must return tools list, not error: {json}"
1278        );
1279    }
1280
1281    /// P1: Lenient mode allows operations without notifications/initialized
1282    #[tokio::test]
1283    async fn test_lambda_streaming_lenient_mode_allows_without_initialized() {
1284        use crate::LambdaMcpServerBuilder;
1285        use turul_mcp_session_storage::InMemorySessionStorage;
1286
1287        let server = LambdaMcpServerBuilder::new()
1288            .name("lenient-test")
1289            .version("1.0.0")
1290            .tool(LifecycleTestTool)
1291            .storage(Arc::new(InMemorySessionStorage::new()))
1292            .strict_lifecycle(false) // lenient mode
1293            .sse(true)
1294            .build()
1295            .await
1296            .unwrap();
1297
1298        let handler = server.handler().await.unwrap();
1299
1300        // Initialize (no notifications/initialized)
1301        let init_req = streaming_mcp_request(
1302            &serde_json::json!({
1303                "jsonrpc": "2.0", "method": "initialize", "id": 1,
1304                "params": {
1305                    "protocolVersion": "2025-11-25",
1306                    "capabilities": {},
1307                    "clientInfo": { "name": "test", "version": "1.0.0" }
1308                }
1309            })
1310            .to_string(),
1311            None,
1312        );
1313        let init_resp = handler.handle_streaming(init_req).await.unwrap();
1314        let session_id = extract_session_id(&init_resp).unwrap();
1315        let _ = collect_streaming_body(init_resp).await;
1316
1317        // Skip notifications/initialized — go straight to tools/list
1318        let list_req = streaming_mcp_request(
1319            &serde_json::json!({
1320                "jsonrpc": "2.0", "method": "tools/list", "id": 2
1321            })
1322            .to_string(),
1323            Some(&session_id),
1324        );
1325        let list_resp = handler.handle_streaming(list_req).await.unwrap();
1326        let (status, body) = collect_streaming_body(list_resp).await;
1327        assert_eq!(
1328            status, 200,
1329            "Lenient mode should allow tools/list without initialized"
1330        );
1331        let json = parse_response_json(&body);
1332        assert!(
1333            json["result"]["tools"].is_array(),
1334            "Must return tools list in lenient mode: {json}"
1335        );
1336    }
1337
1338    // ── Streaming custom-route CORS regression tests ──
1339    //
1340    // Guards the parity between the buffered `handle()` path and the
1341    // streaming `handle_streaming()` path: both must apply configured
1342    // CORS to custom-route responses (matched and validation-error)
1343    // before returning.
1344
1345    #[cfg(feature = "cors")]
1346    mod cors_streaming_routes {
1347        use super::*;
1348        use async_trait::async_trait;
1349        use bytes::Bytes;
1350        use http_body_util::Full;
1351        use hyper::{Request as HyperRequest, Response as HyperResponse, StatusCode};
1352        use turul_http_mcp_server::middleware::MiddlewareStack;
1353        use turul_http_mcp_server::{
1354            RouteBody, RouteHandler, RouteRegistry, StreamConfig, StreamManager,
1355        };
1356
1357        struct StubRoute {
1358            status: StatusCode,
1359            body: &'static str,
1360        }
1361
1362        #[async_trait]
1363        impl RouteHandler for StubRoute {
1364            async fn handle(&self, _req: HyperRequest<RouteBody>) -> HyperResponse<RouteBody> {
1365                use http_body_util::BodyExt;
1366                HyperResponse::builder()
1367                    .status(self.status)
1368                    .header("Content-Type", "application/json")
1369                    .body(
1370                        Full::new(Bytes::from(self.body))
1371                            .map_err(|never| match never {})
1372                            .boxed_unsync(),
1373                    )
1374                    .unwrap()
1375            }
1376        }
1377
1378        fn handler_with_route_and_cors(
1379            registry: Arc<RouteRegistry>,
1380            cors: Option<CorsConfig>,
1381        ) -> LambdaMcpHandler {
1382            let session_storage = Arc::new(InMemorySessionStorage::new());
1383            let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
1384            let dispatcher = Arc::new(JsonRpcDispatcher::new());
1385            let config = ServerConfig::default();
1386            let capabilities = ServerCapabilities::default();
1387            let middleware_stack = Arc::new(MiddlewareStack::new());
1388
1389            let handler = LambdaMcpHandler::with_middleware(
1390                config,
1391                dispatcher,
1392                session_storage,
1393                stream_manager,
1394                StreamConfig::default(),
1395                capabilities,
1396                middleware_stack,
1397                false,
1398                registry,
1399            );
1400            match cors {
1401                Some(cfg) => handler.with_cors(cfg),
1402                None => handler,
1403            }
1404        }
1405
1406        fn get_request(path: &str, origin: &str) -> LambdaRequest {
1407            Request::builder()
1408                .method("GET")
1409                .uri(path)
1410                .header("Origin", origin)
1411                .body(LambdaBody::Empty)
1412                .unwrap()
1413        }
1414
1415        #[tokio::test]
1416        async fn streaming_custom_route_match_injects_cors() {
1417            let mut registry = RouteRegistry::new();
1418            registry.add_route(
1419                "/.well-known/oauth-protected-resource",
1420                Arc::new(StubRoute {
1421                    status: StatusCode::OK,
1422                    body: r#"{"resource":"https://example.test/mcp"}"#,
1423                }),
1424            );
1425            let handler = handler_with_route_and_cors(
1426                Arc::new(registry),
1427                Some(CorsConfig::default()),
1428            );
1429
1430            let req = get_request(
1431                "/.well-known/oauth-protected-resource",
1432                "https://client.example.test",
1433            );
1434            let resp = handler.handle_streaming(req).await.unwrap();
1435
1436            assert_eq!(resp.status(), StatusCode::OK);
1437            assert!(
1438                resp.headers().contains_key("access-control-allow-origin"),
1439                "matched streaming route must carry CORS headers",
1440            );
1441            assert!(
1442                resp.headers().contains_key("access-control-expose-headers"),
1443                "matched streaming route must expose configured headers",
1444            );
1445        }
1446
1447        #[tokio::test]
1448        async fn streaming_route_validation_error_injects_cors() {
1449            // Empty registry + path-traversal path → validation error branch.
1450            let registry = Arc::new({
1451                let mut r = RouteRegistry::new();
1452                r.add_route(
1453                    "/.well-known/oauth-protected-resource",
1454                    Arc::new(StubRoute {
1455                        status: StatusCode::OK,
1456                        body: "{}",
1457                    }),
1458                );
1459                r
1460            });
1461            let handler = handler_with_route_and_cors(registry, Some(CorsConfig::default()));
1462
1463            let req = get_request("/../etc/passwd", "https://client.example.test");
1464            let resp = handler.handle_streaming(req).await.unwrap();
1465
1466            assert!(
1467                resp.status().is_client_error(),
1468                "path-traversal must be a 4xx, got {}",
1469                resp.status(),
1470            );
1471            assert!(
1472                resp.headers().contains_key("access-control-allow-origin"),
1473                "validation-error streaming route must carry CORS headers",
1474            );
1475        }
1476
1477        #[tokio::test]
1478        async fn streaming_custom_route_without_cors_config_returns_untouched() {
1479            // Sanity: without `.with_cors()`, the route response must NOT
1480            // gain CORS headers (regression guard so we never inject
1481            // default CORS for consumers who deliberately opted out).
1482            let mut registry = RouteRegistry::new();
1483            registry.add_route(
1484                "/.well-known/oauth-protected-resource",
1485                Arc::new(StubRoute {
1486                    status: StatusCode::OK,
1487                    body: "{}",
1488                }),
1489            );
1490            let handler = handler_with_route_and_cors(Arc::new(registry), None);
1491
1492            let req = get_request(
1493                "/.well-known/oauth-protected-resource",
1494                "https://client.example.test",
1495            );
1496            let resp = handler.handle_streaming(req).await.unwrap();
1497
1498            assert_eq!(resp.status(), StatusCode::OK);
1499            assert!(
1500                !resp.headers().contains_key("access-control-allow-origin"),
1501                "no CORS config → no CORS headers (got {:?})",
1502                resp.headers(),
1503            );
1504        }
1505    }
1506
1507    // ── OAuth-style 401 challenge through streaming + CORS ──
1508    //
1509    // Verifies the transport contract: a middleware that returns
1510    // `MiddlewareError::http_challenge(401, ...)` produces a response
1511    // that (a) keeps the WWW-Authenticate header, (b) carries
1512    // configured CORS, and (c) exposes WWW-Authenticate so browser
1513    // OAuth clients can read it for RFC 9728 discovery.
1514
1515    #[cfg(feature = "cors")]
1516    mod cors_streaming_oauth {
1517        use super::*;
1518        use async_trait::async_trait;
1519        use turul_http_mcp_server::middleware::{
1520            DispatcherResult, McpMiddleware, MiddlewareError, MiddlewareStack, RequestContext,
1521            SessionInjection,
1522        };
1523        use turul_http_mcp_server::{StreamConfig, StreamManager};
1524        use turul_mcp_session_storage::SessionView;
1525
1526        struct ForceChallenge;
1527
1528        #[async_trait]
1529        impl McpMiddleware for ForceChallenge {
1530            fn runs_before_session(&self) -> bool {
1531                true
1532            }
1533
1534            async fn before_dispatch(
1535                &self,
1536                _ctx: &mut RequestContext<'_>,
1537                _session: Option<&dyn SessionView>,
1538                _injection: &mut SessionInjection,
1539            ) -> std::result::Result<(), MiddlewareError> {
1540                Err(MiddlewareError::http_challenge(
1541                    401,
1542                    "Bearer realm=\"mcp\", resource_metadata=\"https://example.test/.well-known/oauth-protected-resource\"",
1543                ))
1544            }
1545
1546            async fn after_dispatch(
1547                &self,
1548                _ctx: &RequestContext<'_>,
1549                _result: &mut DispatcherResult,
1550            ) -> std::result::Result<(), MiddlewareError> {
1551                Ok(())
1552            }
1553        }
1554
1555        #[tokio::test]
1556        async fn streaming_401_challenge_has_cors_and_exposes_www_authenticate() {
1557            let session_storage = Arc::new(InMemorySessionStorage::new());
1558            let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
1559            let dispatcher = Arc::new(JsonRpcDispatcher::new());
1560            let config = ServerConfig::default();
1561            let capabilities = ServerCapabilities::default();
1562
1563            let mut middleware = MiddlewareStack::new();
1564            middleware.push(Arc::new(ForceChallenge));
1565            let middleware = Arc::new(middleware);
1566
1567            let route_registry =
1568                Arc::new(turul_http_mcp_server::RouteRegistry::new());
1569
1570            let handler = LambdaMcpHandler::with_middleware(
1571                config,
1572                dispatcher,
1573                session_storage,
1574                stream_manager,
1575                StreamConfig::default(),
1576                capabilities,
1577                middleware,
1578                false,
1579                route_registry,
1580            )
1581            .with_cors(CorsConfig::default());
1582
1583            let req = Request::builder()
1584                .method("POST")
1585                .uri("/mcp")
1586                .header("Content-Type", "application/json")
1587                .header("Accept", "application/json, text/event-stream")
1588                .header("MCP-Protocol-Version", "2025-11-25")
1589                .header("Origin", "https://client.example.test")
1590                .body(LambdaBody::Text(
1591                    r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
1592                ))
1593                .unwrap();
1594
1595            let resp = handler.handle_streaming(req).await.unwrap();
1596            let headers = resp.headers();
1597
1598            assert_eq!(resp.status(), 401, "challenge must be 401");
1599            assert!(
1600                headers.contains_key("www-authenticate"),
1601                "WWW-Authenticate must be preserved through streaming transport",
1602            );
1603            assert!(
1604                headers.contains_key("access-control-allow-origin"),
1605                "401 response must carry Access-Control-Allow-Origin",
1606            );
1607            let expose = headers
1608                .get("access-control-expose-headers")
1609                .and_then(|v| v.to_str().ok())
1610                .unwrap_or("");
1611            assert!(
1612                expose
1613                    .split(',')
1614                    .map(str::trim)
1615                    .any(|h| h.eq_ignore_ascii_case("WWW-Authenticate")),
1616                "expose-headers must include WWW-Authenticate; got {expose:?}",
1617            );
1618        }
1619    }
1620}