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