1use 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_protocol::{McpError, ServerCapabilities};
15use turul_mcp_session_storage::BoxedSessionStorage;
16use turul_rpc::JsonRpcDispatcher;
17
18use crate::error::Result;
19
20#[cfg(feature = "cors")]
21use crate::cors::{CorsConfig, create_preflight_response, inject_cors_headers};
22
23#[derive(Clone)]
34pub struct LambdaMcpHandler {
35 session_handler: SessionMcpHandler,
37
38 streamable_handler: StreamableHttpHandler,
40
41 #[allow(dead_code)]
43 sse_enabled: bool,
44
45 route_registry: Arc<turul_http_mcp_server::RouteRegistry>,
47
48 #[cfg_attr(not(test), allow(dead_code))]
51 dispatcher: Arc<JsonRpcDispatcher<McpError>>,
52
53 #[cfg(feature = "dynamic-tools")]
55 tool_registry: Option<Arc<turul_mcp_server::ToolRegistry>>,
56
57 #[cfg(feature = "cors")]
59 cors_config: Option<CorsConfig>,
60}
61
62impl LambdaMcpHandler {
63 pub fn with_server_info(mut self, info: turul_mcp_protocol::Implementation) -> Self {
68 self.streamable_handler = self.streamable_handler.with_server_info(info);
69 self
70 }
71
72 #[allow(clippy::too_many_arguments)]
74 pub fn new(
75 dispatcher: JsonRpcDispatcher<McpError>,
76 session_storage: Arc<BoxedSessionStorage>,
77 stream_manager: Arc<StreamManager>,
78 config: ServerConfig,
79 stream_config: StreamConfig,
80 implementation: turul_mcp_protocol::Implementation,
81 capabilities: ServerCapabilities,
82 sse_enabled: bool,
83 #[cfg(feature = "cors")] cors_config: Option<CorsConfig>,
84 ) -> Self {
85 let dispatcher = Arc::new(dispatcher);
86
87 let middleware_stack = Arc::new(turul_http_mcp_server::middleware::MiddlewareStack::new());
89
90 let session_handler = SessionMcpHandler::with_shared_stream_manager(
92 config.clone(),
93 dispatcher.clone(),
94 session_storage.clone(),
95 stream_config.clone(),
96 stream_manager.clone(),
97 middleware_stack.clone(),
98 );
99
100 let streamable_handler = StreamableHttpHandler::new(
102 Arc::new(config.clone()),
103 dispatcher.clone(),
104 session_storage.clone(),
105 stream_manager.clone(),
106 capabilities.clone(),
107 middleware_stack,
108 None, )
110 .with_server_info(implementation);
111
112 Self {
113 session_handler,
114 streamable_handler,
115 sse_enabled,
116 route_registry: Arc::new(turul_http_mcp_server::RouteRegistry::new()),
117 dispatcher,
118 #[cfg(feature = "dynamic-tools")]
119 tool_registry: None,
120 #[cfg(feature = "cors")]
121 cors_config,
122 }
123 }
124
125 #[allow(clippy::too_many_arguments)]
127 pub fn with_shared_stream_manager(
128 config: ServerConfig,
129 dispatcher: Arc<JsonRpcDispatcher<McpError>>,
130 session_storage: Arc<BoxedSessionStorage>,
131 stream_manager: Arc<StreamManager>,
132 stream_config: StreamConfig,
133 implementation: turul_mcp_protocol::Implementation,
134 capabilities: ServerCapabilities,
135 sse_enabled: bool,
136 ) -> Self {
137 let middleware_stack = Arc::new(turul_http_mcp_server::middleware::MiddlewareStack::new());
139
140 let session_handler = SessionMcpHandler::with_shared_stream_manager(
142 config.clone(),
143 dispatcher.clone(),
144 session_storage.clone(),
145 stream_config.clone(),
146 stream_manager.clone(),
147 middleware_stack.clone(),
148 );
149
150 let dispatcher_for_introspection = dispatcher.clone();
151
152 let streamable_handler = StreamableHttpHandler::new(
154 Arc::new(config),
155 dispatcher,
156 session_storage,
157 stream_manager,
158 capabilities,
159 middleware_stack,
160 None, )
162 .with_server_info(implementation);
163
164 Self {
165 session_handler,
166 streamable_handler,
167 sse_enabled,
168 route_registry: Arc::new(turul_http_mcp_server::RouteRegistry::new()),
169 dispatcher: dispatcher_for_introspection,
170 #[cfg(feature = "dynamic-tools")]
171 tool_registry: None,
172 #[cfg(feature = "cors")]
173 cors_config: None,
174 }
175 }
176
177 #[allow(clippy::too_many_arguments)]
179 pub fn with_middleware(
180 config: ServerConfig,
181 dispatcher: Arc<JsonRpcDispatcher<McpError>>,
182 session_storage: Arc<BoxedSessionStorage>,
183 stream_manager: Arc<StreamManager>,
184 stream_config: StreamConfig,
185 capabilities: ServerCapabilities,
186 middleware_stack: Arc<turul_http_mcp_server::middleware::MiddlewareStack>,
187 sse_enabled: bool,
188 route_registry: Arc<turul_http_mcp_server::RouteRegistry>,
189 ) -> Self {
190 Self::with_middleware_and_fingerprint(
191 config,
192 dispatcher,
193 session_storage,
194 stream_manager,
195 stream_config,
196 capabilities,
197 middleware_stack,
198 sse_enabled,
199 route_registry,
200 None,
201 )
202 }
203
204 #[allow(clippy::too_many_arguments)]
206 pub fn with_middleware_and_fingerprint(
207 config: ServerConfig,
208 dispatcher: Arc<JsonRpcDispatcher<McpError>>,
209 session_storage: Arc<BoxedSessionStorage>,
210 stream_manager: Arc<StreamManager>,
211 stream_config: StreamConfig,
212 capabilities: ServerCapabilities,
213 middleware_stack: Arc<turul_http_mcp_server::middleware::MiddlewareStack>,
214 sse_enabled: bool,
215 route_registry: Arc<turul_http_mcp_server::RouteRegistry>,
216 tool_fingerprint: Option<String>,
217 ) -> Self {
218 let session_handler = SessionMcpHandler::with_shared_stream_manager(
220 config.clone(),
221 dispatcher.clone(),
222 session_storage.clone(),
223 stream_config.clone(),
224 stream_manager.clone(),
225 middleware_stack.clone(),
226 )
227 .with_tool_fingerprint(tool_fingerprint.clone());
228
229 let dispatcher_for_introspection = dispatcher.clone();
230
231 let streamable_handler = StreamableHttpHandler::new(
233 Arc::new(config),
234 dispatcher,
235 session_storage,
236 stream_manager,
237 capabilities,
238 middleware_stack,
239 tool_fingerprint,
240 );
241
242 Self {
243 session_handler,
244 streamable_handler,
245 sse_enabled,
246 route_registry,
247 dispatcher: dispatcher_for_introspection,
248 #[cfg(feature = "dynamic-tools")]
249 tool_registry: None,
250 #[cfg(feature = "cors")]
251 cors_config: None,
252 }
253 }
254
255 pub fn with_tool_notifier(
257 mut self,
258 notifier: Arc<dyn turul_http_mcp_server::ToolChangeNotifier>,
259 ) -> Self {
260 self.session_handler = self
261 .session_handler
262 .with_tool_notifier(Arc::clone(¬ifier));
263 self.streamable_handler = self.streamable_handler.with_tool_notifier(notifier);
264 self
265 }
266
267 #[cfg(feature = "dynamic-tools")]
269 pub fn with_tool_registry(mut self, registry: Arc<turul_mcp_server::ToolRegistry>) -> Self {
270 self.tool_registry = Some(registry);
271 self
272 }
273
274 #[cfg(feature = "cors")]
276 pub fn with_cors(mut self, cors_config: CorsConfig) -> Self {
277 self.cors_config = Some(cors_config);
278 self
279 }
280
281 pub fn registered_methods(&self) -> Vec<String> {
285 self.dispatcher.registered_methods()
286 }
287
288 pub fn get_stream_manager(&self) -> &Arc<StreamManager> {
290 self.session_handler.get_stream_manager()
291 }
292
293 pub async fn handle(&self, req: LambdaRequest) -> Result<LambdaResponse<LambdaBody>> {
302 let method = req.method().clone();
303 let uri = req.uri().clone();
304
305 let request_origin = req
306 .headers()
307 .get("origin")
308 .and_then(|v| v.to_str().ok())
309 .map(|s| s.to_string());
310
311 info!(
312 "π Lambda MCP request: {} {} (origin: {:?})",
313 method, uri, request_origin
314 );
315
316 #[cfg(feature = "cors")]
318 if method == http::Method::OPTIONS
319 && let Some(ref cors_config) = self.cors_config
320 {
321 debug!("Handling CORS preflight request");
322 return create_preflight_response(cors_config, request_origin.as_deref());
323 }
324
325 #[cfg(feature = "dynamic-tools")]
327 if let Some(ref registry) = self.tool_registry
328 && let Err(e) = registry.check_for_changes().await
329 {
330 tracing::warn!(error = %e, "Failed to check for tool changes");
331 }
332
333 let hyper_req = crate::adapter::lambda_to_hyper_request(req)?;
335
336 let path = hyper_req.uri().path().to_string();
338 if !self.route_registry.is_empty() {
339 match self.route_registry.match_route(&path) {
340 Ok(Some(route_handler)) => {
341 debug!("Custom route matched: {}", path);
342 use http_body_util::BodyExt;
343 let (parts, body) = hyper_req.into_parts();
344 let boxed_req = hyper::Request::from_parts(parts, body.boxed_unsync());
345 let route_resp = route_handler.handle(boxed_req).await;
346 #[cfg_attr(not(feature = "cors"), allow(unused_mut))]
348 let mut lambda_resp =
349 crate::adapter::hyper_to_lambda_response(route_resp).await?;
350 #[cfg(feature = "cors")]
351 if let Some(ref cors_config) = self.cors_config {
352 inject_cors_headers(
353 &mut lambda_resp,
354 cors_config,
355 request_origin.as_deref(),
356 )?;
357 }
358 return Ok(lambda_resp);
359 }
360 Ok(None) => {} Err(e) => {
362 debug!("Route validation error: {}", e);
363 let route_resp = e.into_response();
364 #[cfg_attr(not(feature = "cors"), allow(unused_mut))]
366 let mut lambda_resp =
367 crate::adapter::hyper_to_lambda_response(route_resp).await?;
368 #[cfg(feature = "cors")]
369 if let Some(ref cors_config) = self.cors_config {
370 inject_cors_headers(
371 &mut lambda_resp,
372 cors_config,
373 request_origin.as_deref(),
374 )?;
375 }
376 return Ok(lambda_resp);
377 }
378 }
379 }
380
381 #[cfg(feature = "protocol-2026-07-28")]
387 let hyper_resp = self.streamable_handler.handle_request(hyper_req).await;
388
389 #[cfg(feature = "protocol-2025-11-25")]
393 let hyper_resp = self
394 .session_handler
395 .handle_mcp_request(hyper_req)
396 .await
397 .map_err(|e| crate::error::LambdaError::McpFramework(e.to_string()))?;
398
399 #[cfg_attr(not(feature = "cors"), allow(unused_mut))]
402 let mut lambda_resp = crate::adapter::hyper_to_lambda_response(hyper_resp).await?;
403
404 #[cfg(feature = "cors")]
406 if let Some(ref cors_config) = self.cors_config {
407 inject_cors_headers(&mut lambda_resp, cors_config, request_origin.as_deref())?;
408 }
409
410 Ok(lambda_resp)
411 }
412
413 pub async fn handle_streaming(
418 &self,
419 req: LambdaRequest,
420 ) -> std::result::Result<
421 lambda_http::Response<
422 http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
423 >,
424 Box<dyn std::error::Error + Send + Sync>,
425 > {
426 let method = req.method().clone();
427 let uri = req.uri().clone();
428 let request_origin = req
429 .headers()
430 .get("origin")
431 .and_then(|v| v.to_str().ok())
432 .map(|s| s.to_string());
433
434 debug!(
435 "π Lambda streaming MCP request: {} {} (origin: {:?})",
436 method, uri, request_origin
437 );
438
439 #[cfg(feature = "cors")]
441 if method == http::Method::OPTIONS
442 && let Some(ref cors_config) = self.cors_config
443 {
444 debug!("Handling CORS preflight request (streaming)");
445 let preflight_response =
446 create_preflight_response(cors_config, request_origin.as_deref())
447 .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
448
449 return Ok(self.convert_lambda_response_to_streaming(preflight_response));
451 }
452
453 #[cfg(feature = "dynamic-tools")]
455 if let Some(ref registry) = self.tool_registry
456 && let Err(e) = registry.check_for_changes().await
457 {
458 tracing::warn!(error = %e, "Failed to check for tool changes (streaming)");
459 }
460
461 let hyper_req = crate::adapter::lambda_to_hyper_request(req)
463 .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
464
465 let path = hyper_req.uri().path().to_string();
467 if !self.route_registry.is_empty() {
468 match self.route_registry.match_route(&path) {
469 Ok(Some(route_handler)) => {
470 debug!("Custom route matched (streaming): {}", path);
471 use http_body_util::BodyExt;
472 let (parts, body) = hyper_req.into_parts();
473 let boxed_req = hyper::Request::from_parts(parts, body.boxed_unsync());
474 #[cfg_attr(not(feature = "cors"), allow(unused_mut))]
476 let mut route_resp = route_handler.handle(boxed_req).await;
477 #[cfg(feature = "cors")]
478 if let Some(ref cors_config) = self.cors_config {
479 inject_cors_headers(
480 &mut route_resp,
481 cors_config,
482 request_origin.as_deref(),
483 )
484 .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
485 }
486 return Ok(route_resp);
487 }
488 Ok(None) => {} Err(e) => {
490 debug!("Route validation error (streaming): {}", e);
491 #[cfg_attr(not(feature = "cors"), allow(unused_mut))]
493 let mut err_resp = e.into_response();
494 #[cfg(feature = "cors")]
495 if let Some(ref cors_config) = self.cors_config {
496 inject_cors_headers(&mut err_resp, cors_config, request_origin.as_deref())
497 .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
498 }
499 return Ok(err_resp);
500 }
501 }
502 }
503
504 use turul_http_mcp_server::protocol::McpProtocolVersion;
506 let raw_version = hyper_req
511 .headers()
512 .get("MCP-Protocol-Version")
513 .and_then(|h| h.to_str().ok());
514 let parsed_version = raw_version.and_then(McpProtocolVersion::parse_version);
515 let unrecognised_version = raw_version.is_some() && parsed_version.is_none();
516 let protocol_version = parsed_version.unwrap_or(McpProtocolVersion::LATEST);
517
518 #[cfg(feature = "protocol-2026-07-28")]
524 let route_streamable = true;
525 #[cfg(feature = "protocol-2025-11-25")]
526 let route_streamable = unrecognised_version || protocol_version.supports_streamable_http();
527
528 let hyper_resp = if route_streamable {
530 debug!(
531 "Using StreamableHttpHandler for protocol {} (header={:?}, unrecognised={})",
532 protocol_version.as_str(),
533 raw_version,
534 unrecognised_version
535 );
536 self.streamable_handler.handle_request(hyper_req).await
537 } else {
538 debug!(
540 "Using SessionMcpHandler for legacy protocol {}",
541 protocol_version.as_str()
542 );
543 self.session_handler
544 .handle_mcp_request(hyper_req)
545 .await
546 .map_err(|e| {
547 Box::new(crate::error::LambdaError::McpFramework(e.to_string()))
548 as Box<dyn std::error::Error + Send + Sync>
549 })?
550 };
551
552 #[cfg_attr(not(feature = "cors"), allow(unused_mut))]
555 let mut lambda_resp = crate::adapter::hyper_to_lambda_streaming(hyper_resp);
556
557 #[cfg(feature = "cors")]
559 if let Some(ref cors_config) = self.cors_config {
560 inject_cors_headers(&mut lambda_resp, cors_config, request_origin.as_deref())
561 .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
562 }
563
564 Ok(lambda_resp)
565 }
566
567 #[cfg(feature = "cors")]
570 fn convert_lambda_response_to_streaming(
571 &self,
572 lambda_response: LambdaResponse<LambdaBody>,
573 ) -> lambda_http::Response<http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>>
574 {
575 use bytes::Bytes;
576 use http_body_util::{BodyExt, Full};
577
578 let (parts, body) = lambda_response.into_parts();
579 let body_bytes = match body {
580 LambdaBody::Empty => Bytes::new(),
581 LambdaBody::Text(text) => Bytes::from(text),
582 LambdaBody::Binary(bytes) => Bytes::from(bytes),
583 _ => Bytes::new(),
584 };
585
586 let streaming_body = Full::new(body_bytes)
588 .map_err(|e: std::convert::Infallible| match e {})
589 .boxed_unsync();
590
591 lambda_http::Response::from_parts(parts, streaming_body)
592 }
593}
594
595#[cfg(test)]
596mod tests {
597 use super::*;
598 use http::Request;
599 use turul_mcp_session_storage::InMemorySessionStorage;
600
601 #[tokio::test]
602 async fn test_handler_creation() {
603 let session_storage = Arc::new(InMemorySessionStorage::new());
604 let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
605 let dispatcher = JsonRpcDispatcher::new();
606 let config = ServerConfig::default();
607 let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
608 let capabilities = ServerCapabilities::default();
609
610 let handler = LambdaMcpHandler::new(
611 dispatcher,
612 session_storage,
613 stream_manager,
614 config,
615 StreamConfig::default(),
616 implementation,
617 capabilities,
618 false, #[cfg(feature = "cors")]
620 None,
621 );
622
623 assert!(!handler.sse_enabled);
625 }
626
627 #[tokio::test]
628 async fn test_sse_enabled_with_handle_works() {
629 let session_storage = Arc::new(InMemorySessionStorage::new());
630 let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
631 let dispatcher = JsonRpcDispatcher::new();
632 let config = ServerConfig::default();
633 let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
634 let capabilities = ServerCapabilities::default();
635
636 let handler = LambdaMcpHandler::new(
638 dispatcher,
639 session_storage,
640 stream_manager,
641 config,
642 StreamConfig::default(),
643 implementation,
644 capabilities,
645 true, #[cfg(feature = "cors")]
647 None,
648 );
649
650 let lambda_req = Request::builder()
652 .method("POST")
653 .uri("/mcp")
654 .body(LambdaBody::Text(
655 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
656 ))
657 .unwrap();
658
659 let result = handler.handle(lambda_req).await;
661 assert!(
662 result.is_ok(),
663 "handle() should work with SSE enabled for snapshot-based responses"
664 );
665 }
666
667 #[tokio::test]
669 async fn test_stream_config_preservation() {
670 let session_storage = Arc::new(InMemorySessionStorage::new());
671 let dispatcher = JsonRpcDispatcher::new();
672 let config = ServerConfig::default();
673 let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
674 let capabilities = ServerCapabilities::default();
675
676 let custom_stream_config = StreamConfig {
678 channel_buffer_size: 1024, max_replay_events: 200, keepalive_interval_seconds: 10, cors_origin: "https://custom-test.example.com".to_string(), };
683
684 let stream_manager = Arc::new(StreamManager::with_config(
686 session_storage.clone(),
687 custom_stream_config.clone(),
688 ));
689
690 let handler = LambdaMcpHandler::new(
691 dispatcher,
692 session_storage,
693 stream_manager,
694 config,
695 custom_stream_config.clone(),
696 implementation,
697 capabilities,
698 false, #[cfg(feature = "cors")]
700 None,
701 );
702
703 assert!(!handler.sse_enabled);
705
706 let stream_manager = handler.get_stream_manager();
708
709 let actual_config = stream_manager.get_config();
711
712 assert_eq!(
713 actual_config.channel_buffer_size, custom_stream_config.channel_buffer_size,
714 "Custom channel_buffer_size was not propagated correctly"
715 );
716 assert_eq!(
717 actual_config.max_replay_events, custom_stream_config.max_replay_events,
718 "Custom max_replay_events was not propagated correctly"
719 );
720 assert_eq!(
721 actual_config.keepalive_interval_seconds,
722 custom_stream_config.keepalive_interval_seconds,
723 "Custom keepalive_interval_seconds was not propagated correctly"
724 );
725 assert_eq!(
726 actual_config.cors_origin, custom_stream_config.cors_origin,
727 "Custom cors_origin was not propagated correctly"
728 );
729
730 assert!(Arc::strong_count(stream_manager) >= 1);
732 }
733
734 #[tokio::test]
736 async fn test_full_builder_chain_stream_config() {
737 use crate::LambdaMcpServerBuilder;
738 use turul_mcp_session_storage::InMemorySessionStorage;
739
740 let custom_stream_config = turul_http_mcp_server::StreamConfig {
742 channel_buffer_size: 2048, max_replay_events: 500, keepalive_interval_seconds: 15, cors_origin: "https://full-chain-test.example.com".to_string(),
746 };
747
748 let server = LambdaMcpServerBuilder::new()
750 .name("full-chain-test")
751 .version("1.0.0")
752 .storage(Arc::new(InMemorySessionStorage::new()))
753 .sse(true) .stream_config(custom_stream_config.clone())
755 .build()
756 .await
757 .expect("Server should build successfully");
758
759 let handler = server
761 .handler()
762 .await
763 .expect("Handler should be created from server");
764
765 assert!(handler.sse_enabled, "SSE should be enabled");
767
768 let stream_manager = handler.get_stream_manager();
770 let actual_config = stream_manager.get_config();
771
772 assert_eq!(
773 actual_config.channel_buffer_size, custom_stream_config.channel_buffer_size,
774 "Custom channel_buffer_size should be preserved through builder β server β handler chain"
775 );
776 assert_eq!(
777 actual_config.max_replay_events, custom_stream_config.max_replay_events,
778 "Custom max_replay_events should be preserved through builder β server β handler chain"
779 );
780 assert_eq!(
781 actual_config.keepalive_interval_seconds,
782 custom_stream_config.keepalive_interval_seconds,
783 "Custom keepalive_interval_seconds should be preserved through builder β server β handler chain"
784 );
785 assert_eq!(
786 actual_config.cors_origin, custom_stream_config.cors_origin,
787 "Custom cors_origin should be preserved through builder β server β handler chain"
788 );
789
790 assert!(
792 Arc::strong_count(stream_manager) >= 1,
793 "Stream manager should be properly initialized"
794 );
795
796 let test_session_id = uuid::Uuid::now_v7().as_simple().to_string();
799
800 let subscriptions = stream_manager.get_subscriptions(&test_session_id).await;
803 assert!(
804 subscriptions.is_empty(),
805 "New session should have no subscriptions initially"
806 );
807
808 assert_eq!(
811 stream_manager.get_config().channel_buffer_size,
812 2048,
813 "Stream manager should be using the custom buffer size functionally"
814 );
815 }
816
817 #[tokio::test]
822 async fn test_non_streaming_runtime_sse_false() {
823 use crate::LambdaMcpServerBuilder;
824 use turul_mcp_session_storage::InMemorySessionStorage;
825
826 let server = LambdaMcpServerBuilder::new()
827 .name("test-non-streaming-sse-false")
828 .version("1.0.0")
829 .storage(Arc::new(InMemorySessionStorage::new()))
830 .sse(false) .build()
832 .await
833 .expect("Server should build successfully");
834
835 let handler = server
836 .handler()
837 .await
838 .expect("Handler should be created from server");
839
840 assert!(!handler.sse_enabled, "SSE should be disabled");
842
843 let lambda_req = Request::builder()
845 .method("POST")
846 .uri("/mcp")
847 .body(LambdaBody::Text(
848 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
849 ))
850 .unwrap();
851
852 let result = handler.handle(lambda_req).await;
854 assert!(
855 result.is_ok(),
856 "POST /mcp should work with non-streaming + sse(false)"
857 );
858 }
859
860 #[tokio::test]
862 async fn test_non_streaming_runtime_sse_true() {
863 use crate::LambdaMcpServerBuilder;
864 use turul_mcp_session_storage::InMemorySessionStorage;
865
866 let server = LambdaMcpServerBuilder::new()
867 .name("test-non-streaming-sse-true")
868 .version("1.0.0")
869 .storage(Arc::new(InMemorySessionStorage::new()))
870 .sse(true) .build()
872 .await
873 .expect("Server should build successfully");
874
875 let handler = server
876 .handler()
877 .await
878 .expect("Handler should be created from server");
879
880 assert!(handler.sse_enabled, "SSE should be enabled");
882
883 let lambda_req = Request::builder()
885 .method("POST")
886 .uri("/mcp")
887 .body(LambdaBody::Text(
888 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
889 ))
890 .unwrap();
891
892 let result = handler.handle(lambda_req).await;
894 assert!(
895 result.is_ok(),
896 "POST /mcp should work with non-streaming + sse(true)"
897 );
898
899 }
902
903 #[tokio::test]
905 async fn test_streaming_runtime_sse_false() {
906 use crate::LambdaMcpServerBuilder;
907 use turul_mcp_session_storage::InMemorySessionStorage;
908
909 let server = LambdaMcpServerBuilder::new()
910 .name("test-streaming-sse-false")
911 .version("1.0.0")
912 .storage(Arc::new(InMemorySessionStorage::new()))
913 .sse(false) .build()
915 .await
916 .expect("Server should build successfully");
917
918 let handler = server
919 .handler()
920 .await
921 .expect("Handler should be created from server");
922
923 assert!(!handler.sse_enabled, "SSE should be disabled");
925
926 let lambda_req = Request::builder()
928 .method("POST")
929 .uri("/mcp")
930 .body(LambdaBody::Text(
931 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
932 ))
933 .unwrap();
934
935 let result = handler.handle_streaming(lambda_req).await;
937 assert!(
938 result.is_ok(),
939 "Streaming runtime should work with sse(false)"
940 );
941 }
942
943 #[tokio::test]
945 async fn test_streaming_runtime_sse_true() {
946 use crate::LambdaMcpServerBuilder;
947 use turul_mcp_session_storage::InMemorySessionStorage;
948
949 let server = LambdaMcpServerBuilder::new()
950 .name("test-streaming-sse-true")
951 .version("1.0.0")
952 .storage(Arc::new(InMemorySessionStorage::new()))
953 .sse(true) .build()
955 .await
956 .expect("Server should build successfully");
957
958 let handler = server
959 .handler()
960 .await
961 .expect("Handler should be created from server");
962
963 assert!(handler.sse_enabled, "SSE should be enabled");
965
966 let lambda_req = Request::builder()
968 .method("POST")
969 .uri("/mcp")
970 .body(LambdaBody::Text(
971 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
972 ))
973 .unwrap();
974
975 let result = handler.handle_streaming(lambda_req).await;
977 assert!(
978 result.is_ok(),
979 "Streaming runtime should work with sse(true) for real-time streaming"
980 );
981
982 }
985
986 #[cfg(feature = "protocol-2025-11-25")]
990 async fn build_strict_streaming_handler() -> LambdaMcpHandler {
991 use crate::LambdaMcpServerBuilder;
992 use turul_mcp_session_storage::InMemorySessionStorage;
993
994 let server = LambdaMcpServerBuilder::new()
995 .name("lifecycle-test")
996 .version("1.0.0")
997 .tool(LifecycleTestTool)
998 .storage(Arc::new(InMemorySessionStorage::new()))
999 .strict_lifecycle(true) .sse(true)
1001 .build()
1002 .await
1003 .expect("build should succeed");
1004
1005 server.handler().await.expect("handler should succeed")
1006 }
1007
1008 #[cfg(feature = "protocol-2025-11-25")]
1010 #[derive(Clone, Default)]
1011 struct LifecycleTestTool;
1012
1013 #[cfg(feature = "protocol-2025-11-25")]
1014 impl turul_mcp_builders::traits::HasBaseMetadata for LifecycleTestTool {
1015 fn name(&self) -> &str {
1016 "ping_tool"
1017 }
1018 }
1019 #[cfg(feature = "protocol-2025-11-25")]
1020 impl turul_mcp_builders::traits::HasDescription for LifecycleTestTool {
1021 fn description(&self) -> Option<&str> {
1022 Some("test tool")
1023 }
1024 }
1025 #[cfg(feature = "protocol-2025-11-25")]
1026 impl turul_mcp_builders::traits::HasInputSchema for LifecycleTestTool {
1027 fn input_schema(&self) -> &turul_mcp_protocol::ToolSchema {
1028 static SCHEMA: std::sync::OnceLock<turul_mcp_protocol::ToolSchema> =
1029 std::sync::OnceLock::new();
1030 SCHEMA.get_or_init(turul_mcp_protocol::ToolSchema::object)
1031 }
1032 }
1033 #[cfg(feature = "protocol-2025-11-25")]
1034 impl turul_mcp_builders::traits::HasOutputSchema for LifecycleTestTool {
1035 fn output_schema(&self) -> Option<&turul_mcp_protocol::ToolSchema> {
1036 None
1037 }
1038 }
1039 #[cfg(feature = "protocol-2025-11-25")]
1040 impl turul_mcp_builders::traits::HasAnnotations for LifecycleTestTool {
1041 fn annotations(&self) -> Option<&turul_mcp_protocol::tools::ToolAnnotations> {
1042 None
1043 }
1044 }
1045 #[cfg(feature = "protocol-2025-11-25")]
1046 impl turul_mcp_builders::traits::HasToolMeta for LifecycleTestTool {
1047 fn tool_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1048 None
1049 }
1050 }
1051 #[cfg(feature = "protocol-2025-11-25")]
1052 impl turul_mcp_builders::traits::HasIcons for LifecycleTestTool {}
1053 #[cfg(feature = "protocol-2025-11-25")]
1054 impl turul_mcp_builders::traits::HasExecution for LifecycleTestTool {}
1055
1056 #[async_trait::async_trait]
1057 #[cfg(feature = "protocol-2025-11-25")]
1058 impl turul_mcp_server::McpTool for LifecycleTestTool {
1059 async fn call(
1060 &self,
1061 _args: serde_json::Value,
1062 _session: Option<turul_mcp_server::SessionContext>,
1063 ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
1064 Ok(turul_mcp_protocol::tools::CallToolResult::success(vec![
1065 turul_mcp_protocol::tools::ToolResult::text("pong"),
1066 ]))
1067 }
1068 }
1069
1070 #[cfg(feature = "protocol-2025-11-25")]
1072 fn streaming_mcp_request(body: &str, session_id: Option<&str>) -> LambdaRequest {
1073 let mut builder = Request::builder()
1074 .method("POST")
1075 .uri("/mcp")
1076 .header("Content-Type", "application/json")
1077 .header("Accept", "application/json, text/event-stream")
1078 .header("MCP-Protocol-Version", "2025-11-25");
1079
1080 if let Some(sid) = session_id {
1081 builder = builder.header("Mcp-Session-Id", sid);
1082 }
1083
1084 builder.body(LambdaBody::Text(body.to_string())).unwrap()
1085 }
1086
1087 #[cfg(feature = "protocol-2025-11-25")]
1089 async fn collect_streaming_body(
1090 response: lambda_http::Response<
1091 http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
1092 >,
1093 ) -> (http::StatusCode, String) {
1094 use http_body_util::BodyExt;
1095 let status = response.status();
1096 let session_id = response
1097 .headers()
1098 .get("Mcp-Session-Id")
1099 .and_then(|v| v.to_str().ok())
1100 .map(String::from);
1101 let body_bytes = response
1102 .into_body()
1103 .collect()
1104 .await
1105 .map(|c| c.to_bytes())
1106 .unwrap_or_default();
1107 let body_str = String::from_utf8_lossy(&body_bytes).to_string();
1108 let _ = session_id; (status, body_str)
1110 }
1111
1112 #[cfg(feature = "protocol-2025-11-25")]
1114 fn extract_session_id(
1115 response: &lambda_http::Response<
1116 http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
1117 >,
1118 ) -> Option<String> {
1119 response
1120 .headers()
1121 .get("Mcp-Session-Id")
1122 .and_then(|v| v.to_str().ok())
1123 .map(String::from)
1124 }
1125
1126 #[cfg(feature = "protocol-2025-11-25")]
1128 fn parse_response_json(body: &str) -> serde_json::Value {
1129 let json_str = body
1131 .lines()
1132 .find(|line| line.starts_with("data: "))
1133 .map(|line| &line[6..])
1134 .unwrap_or(body.trim());
1135 serde_json::from_str(json_str)
1136 .unwrap_or_else(|e| panic!("Failed to parse JSON from body: {e}\nBody: {body}"))
1137 }
1138
1139 #[cfg(feature = "protocol-2025-11-25")]
1143 #[tokio::test]
1144 async fn test_lambda_streaming_strict_handshake_succeeds() {
1145 let handler = build_strict_streaming_handler().await;
1146
1147 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
1161 .handle_streaming(init_req)
1162 .await
1163 .expect("initialize should succeed");
1164 let session_id = extract_session_id(&init_resp).expect("must return session ID");
1165 let (status, _body) = collect_streaming_body(init_resp).await;
1166 assert_eq!(status, 200, "initialize should return 200");
1167
1168 let notif_req = streaming_mcp_request(
1170 &serde_json::json!({
1171 "jsonrpc": "2.0",
1172 "method": "notifications/initialized",
1173 "params": {}
1174 })
1175 .to_string(),
1176 Some(&session_id),
1177 );
1178 let notif_resp = handler
1179 .handle_streaming(notif_req)
1180 .await
1181 .expect("notification should succeed");
1182 let (status, _) = collect_streaming_body(notif_resp).await;
1183 assert_eq!(status, 202, "notifications/initialized should return 202");
1184
1185 let list_req = streaming_mcp_request(
1187 &serde_json::json!({
1188 "jsonrpc": "2.0", "method": "tools/list", "id": 2
1189 })
1190 .to_string(),
1191 Some(&session_id),
1192 );
1193 let list_resp = handler
1194 .handle_streaming(list_req)
1195 .await
1196 .expect("tools/list should succeed");
1197 let (status, body) = collect_streaming_body(list_resp).await;
1198 assert_eq!(status, 200, "tools/list should return 200");
1199 let json = parse_response_json(&body);
1200 assert!(
1201 json["result"]["tools"].is_array(),
1202 "tools/list should return tools array: {json}"
1203 );
1204
1205 let call_req = streaming_mcp_request(
1207 &serde_json::json!({
1208 "jsonrpc": "2.0", "method": "tools/call", "id": 3,
1209 "params": { "name": "ping_tool", "arguments": {} }
1210 })
1211 .to_string(),
1212 Some(&session_id),
1213 );
1214 let call_resp = handler
1215 .handle_streaming(call_req)
1216 .await
1217 .expect("tools/call should succeed");
1218 let (status, body) = collect_streaming_body(call_resp).await;
1219 assert_eq!(status, 200, "tools/call should return 200");
1220 let json = parse_response_json(&body);
1221 assert!(
1222 json["result"].is_object(),
1223 "tools/call should return result: {json}"
1224 );
1225 }
1226
1227 #[cfg(feature = "protocol-2025-11-25")]
1229 #[tokio::test]
1230 async fn test_lambda_streaming_strict_rejects_before_initialized() {
1231 let handler = build_strict_streaming_handler().await;
1232
1233 let init_req = streaming_mcp_request(
1235 &serde_json::json!({
1236 "jsonrpc": "2.0", "method": "initialize", "id": 1,
1237 "params": {
1238 "protocolVersion": "2025-11-25",
1239 "capabilities": {},
1240 "clientInfo": { "name": "test", "version": "1.0.0" }
1241 }
1242 })
1243 .to_string(),
1244 None,
1245 );
1246 let init_resp = handler.handle_streaming(init_req).await.unwrap();
1247 let session_id = extract_session_id(&init_resp).unwrap();
1248 let _ = collect_streaming_body(init_resp).await;
1249
1250 let list_req = streaming_mcp_request(
1252 &serde_json::json!({
1253 "jsonrpc": "2.0", "method": "tools/list", "id": 2
1254 })
1255 .to_string(),
1256 Some(&session_id),
1257 );
1258 let list_resp = handler.handle_streaming(list_req).await.unwrap();
1259 let (_, body) = collect_streaming_body(list_resp).await;
1260 let json = parse_response_json(&body);
1261 assert!(
1262 json["error"].is_object(),
1263 "tools/list should return JSON-RPC error: {json}"
1264 );
1265 assert_eq!(
1266 json["error"]["code"].as_i64().unwrap(),
1267 -32031,
1268 "tools/list must return SessionError code -32031, got: {json}"
1269 );
1270 assert!(
1271 json["error"]["message"]
1272 .as_str()
1273 .unwrap()
1274 .contains("notifications/initialized"),
1275 "Error must mention notifications/initialized: {}",
1276 json["error"]["message"]
1277 );
1278
1279 let call_req = streaming_mcp_request(
1281 &serde_json::json!({
1282 "jsonrpc": "2.0", "method": "tools/call", "id": 3,
1283 "params": { "name": "ping_tool", "arguments": {} }
1284 })
1285 .to_string(),
1286 Some(&session_id),
1287 );
1288 let call_resp = handler.handle_streaming(call_req).await.unwrap();
1289 let (_, body) = collect_streaming_body(call_resp).await;
1290 let json = parse_response_json(&body);
1291 assert!(
1292 json["error"].is_object(),
1293 "tools/call should return JSON-RPC error: {json}"
1294 );
1295 assert_eq!(
1296 json["error"]["code"].as_i64().unwrap(),
1297 -32031,
1298 "tools/call must return SessionError code -32031, got: {json}"
1299 );
1300 assert!(
1301 json["error"]["message"]
1302 .as_str()
1303 .unwrap()
1304 .contains("notifications/initialized"),
1305 "Error must mention notifications/initialized: {}",
1306 json["error"]["message"]
1307 );
1308 }
1309
1310 #[cfg(feature = "protocol-2025-11-25")]
1314 #[tokio::test]
1315 async fn test_lambda_streaming_initialized_is_effective_immediately() {
1316 let handler = build_strict_streaming_handler().await;
1317
1318 let init_req = streaming_mcp_request(
1320 &serde_json::json!({
1321 "jsonrpc": "2.0", "method": "initialize", "id": 1,
1322 "params": {
1323 "protocolVersion": "2025-11-25",
1324 "capabilities": {},
1325 "clientInfo": { "name": "test", "version": "1.0.0" }
1326 }
1327 })
1328 .to_string(),
1329 None,
1330 );
1331 let init_resp = handler.handle_streaming(init_req).await.unwrap();
1332 let session_id = extract_session_id(&init_resp).unwrap();
1333 let _ = collect_streaming_body(init_resp).await;
1334
1335 let notif_req = streaming_mcp_request(
1337 &serde_json::json!({
1338 "jsonrpc": "2.0",
1339 "method": "notifications/initialized",
1340 "params": {}
1341 })
1342 .to_string(),
1343 Some(&session_id),
1344 );
1345 let notif_resp = handler.handle_streaming(notif_req).await.unwrap();
1346 let (status, _) = collect_streaming_body(notif_resp).await;
1347 assert_eq!(status, 202);
1348
1349 let list_req = streaming_mcp_request(
1351 &serde_json::json!({
1352 "jsonrpc": "2.0", "method": "tools/list", "id": 2
1353 })
1354 .to_string(),
1355 Some(&session_id),
1356 );
1357 let list_resp = handler.handle_streaming(list_req).await.unwrap();
1358 let (status, body) = collect_streaming_body(list_resp).await;
1359 assert_eq!(
1360 status, 200,
1361 "tools/list must succeed immediately after initialized"
1362 );
1363 let json = parse_response_json(&body);
1364 assert!(
1365 json["result"]["tools"].is_array(),
1366 "Must return tools list, not error: {json}"
1367 );
1368 }
1369
1370 #[cfg(feature = "protocol-2025-11-25")]
1372 #[tokio::test]
1373 async fn test_lambda_streaming_lenient_mode_allows_without_initialized() {
1374 use crate::LambdaMcpServerBuilder;
1375 use turul_mcp_session_storage::InMemorySessionStorage;
1376
1377 let server = LambdaMcpServerBuilder::new()
1378 .name("lenient-test")
1379 .version("1.0.0")
1380 .tool(LifecycleTestTool)
1381 .storage(Arc::new(InMemorySessionStorage::new()))
1382 .strict_lifecycle(false) .sse(true)
1384 .build()
1385 .await
1386 .unwrap();
1387
1388 let handler = server.handler().await.unwrap();
1389
1390 let init_req = streaming_mcp_request(
1392 &serde_json::json!({
1393 "jsonrpc": "2.0", "method": "initialize", "id": 1,
1394 "params": {
1395 "protocolVersion": "2025-11-25",
1396 "capabilities": {},
1397 "clientInfo": { "name": "test", "version": "1.0.0" }
1398 }
1399 })
1400 .to_string(),
1401 None,
1402 );
1403 let init_resp = handler.handle_streaming(init_req).await.unwrap();
1404 let session_id = extract_session_id(&init_resp).unwrap();
1405 let _ = collect_streaming_body(init_resp).await;
1406
1407 let list_req = streaming_mcp_request(
1409 &serde_json::json!({
1410 "jsonrpc": "2.0", "method": "tools/list", "id": 2
1411 })
1412 .to_string(),
1413 Some(&session_id),
1414 );
1415 let list_resp = handler.handle_streaming(list_req).await.unwrap();
1416 let (status, body) = collect_streaming_body(list_resp).await;
1417 assert_eq!(
1418 status, 200,
1419 "Lenient mode should allow tools/list without initialized"
1420 );
1421 let json = parse_response_json(&body);
1422 assert!(
1423 json["result"]["tools"].is_array(),
1424 "Must return tools list in lenient mode: {json}"
1425 );
1426 }
1427
1428 #[cfg(feature = "cors")]
1436 mod cors_streaming_routes {
1437 use super::*;
1438 use async_trait::async_trait;
1439 use bytes::Bytes;
1440 use http_body_util::Full;
1441 use hyper::{Request as HyperRequest, Response as HyperResponse, StatusCode};
1442 use turul_http_mcp_server::middleware::MiddlewareStack;
1443 use turul_http_mcp_server::{
1444 RouteBody, RouteHandler, RouteRegistry, StreamConfig, StreamManager,
1445 };
1446
1447 struct StubRoute {
1448 status: StatusCode,
1449 body: &'static str,
1450 }
1451
1452 #[async_trait]
1453 impl RouteHandler for StubRoute {
1454 async fn handle(&self, _req: HyperRequest<RouteBody>) -> HyperResponse<RouteBody> {
1455 use http_body_util::BodyExt;
1456 HyperResponse::builder()
1457 .status(self.status)
1458 .header("Content-Type", "application/json")
1459 .body(
1460 Full::new(Bytes::from(self.body))
1461 .map_err(|never| match never {})
1462 .boxed_unsync(),
1463 )
1464 .unwrap()
1465 }
1466 }
1467
1468 fn handler_with_route_and_cors(
1469 registry: Arc<RouteRegistry>,
1470 cors: Option<CorsConfig>,
1471 ) -> LambdaMcpHandler {
1472 let session_storage = Arc::new(InMemorySessionStorage::new());
1473 let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
1474 let dispatcher = Arc::new(JsonRpcDispatcher::new());
1475 let config = ServerConfig::default();
1476 let capabilities = ServerCapabilities::default();
1477 let middleware_stack = Arc::new(MiddlewareStack::new());
1478
1479 let handler = LambdaMcpHandler::with_middleware(
1480 config,
1481 dispatcher,
1482 session_storage,
1483 stream_manager,
1484 StreamConfig::default(),
1485 capabilities,
1486 middleware_stack,
1487 false,
1488 registry,
1489 );
1490 match cors {
1491 Some(cfg) => handler.with_cors(cfg),
1492 None => handler,
1493 }
1494 }
1495
1496 fn get_request(path: &str, origin: &str) -> LambdaRequest {
1497 Request::builder()
1498 .method("GET")
1499 .uri(path)
1500 .header("Origin", origin)
1501 .body(LambdaBody::Empty)
1502 .unwrap()
1503 }
1504
1505 #[tokio::test]
1506 async fn streaming_custom_route_match_injects_cors() {
1507 let mut registry = RouteRegistry::new();
1508 registry.add_route(
1509 "/.well-known/oauth-protected-resource",
1510 Arc::new(StubRoute {
1511 status: StatusCode::OK,
1512 body: r#"{"resource":"https://example.test/mcp"}"#,
1513 }),
1514 );
1515 let handler =
1516 handler_with_route_and_cors(Arc::new(registry), Some(CorsConfig::default()));
1517
1518 let req = get_request(
1519 "/.well-known/oauth-protected-resource",
1520 "https://client.example.test",
1521 );
1522 let resp = handler.handle_streaming(req).await.unwrap();
1523
1524 assert_eq!(resp.status(), StatusCode::OK);
1525 assert!(
1526 resp.headers().contains_key("access-control-allow-origin"),
1527 "matched streaming route must carry CORS headers",
1528 );
1529 assert!(
1530 resp.headers().contains_key("access-control-expose-headers"),
1531 "matched streaming route must expose configured headers",
1532 );
1533 }
1534
1535 #[tokio::test]
1536 async fn streaming_route_validation_error_injects_cors() {
1537 let registry = Arc::new({
1539 let mut r = RouteRegistry::new();
1540 r.add_route(
1541 "/.well-known/oauth-protected-resource",
1542 Arc::new(StubRoute {
1543 status: StatusCode::OK,
1544 body: "{}",
1545 }),
1546 );
1547 r
1548 });
1549 let handler = handler_with_route_and_cors(registry, Some(CorsConfig::default()));
1550
1551 let req = get_request("/../etc/passwd", "https://client.example.test");
1552 let resp = handler.handle_streaming(req).await.unwrap();
1553
1554 assert!(
1555 resp.status().is_client_error(),
1556 "path-traversal must be a 4xx, got {}",
1557 resp.status(),
1558 );
1559 assert!(
1560 resp.headers().contains_key("access-control-allow-origin"),
1561 "validation-error streaming route must carry CORS headers",
1562 );
1563 }
1564
1565 #[tokio::test]
1566 async fn streaming_custom_route_without_cors_config_returns_untouched() {
1567 let mut registry = RouteRegistry::new();
1571 registry.add_route(
1572 "/.well-known/oauth-protected-resource",
1573 Arc::new(StubRoute {
1574 status: StatusCode::OK,
1575 body: "{}",
1576 }),
1577 );
1578 let handler = handler_with_route_and_cors(Arc::new(registry), None);
1579
1580 let req = get_request(
1581 "/.well-known/oauth-protected-resource",
1582 "https://client.example.test",
1583 );
1584 let resp = handler.handle_streaming(req).await.unwrap();
1585
1586 assert_eq!(resp.status(), StatusCode::OK);
1587 assert!(
1588 !resp.headers().contains_key("access-control-allow-origin"),
1589 "no CORS config β no CORS headers (got {:?})",
1590 resp.headers(),
1591 );
1592 }
1593 }
1594
1595 #[cfg(feature = "cors")]
1604 mod cors_streaming_oauth {
1605 use super::*;
1606 use async_trait::async_trait;
1607 use turul_http_mcp_server::middleware::{
1608 DispatcherResult, McpMiddleware, MiddlewareError, MiddlewareStack, RequestContext,
1609 SessionInjection,
1610 };
1611 use turul_http_mcp_server::{StreamConfig, StreamManager};
1612 use turul_mcp_session_storage::SessionView;
1613
1614 struct ForceChallenge;
1615
1616 #[async_trait]
1617 impl McpMiddleware for ForceChallenge {
1618 fn runs_before_session(&self) -> bool {
1619 true
1620 }
1621
1622 async fn before_dispatch(
1623 &self,
1624 _ctx: &mut RequestContext<'_>,
1625 _session: Option<&dyn SessionView>,
1626 _injection: &mut SessionInjection,
1627 ) -> std::result::Result<(), MiddlewareError> {
1628 Err(MiddlewareError::http_challenge(
1629 401,
1630 "Bearer realm=\"mcp\", resource_metadata=\"https://example.test/.well-known/oauth-protected-resource\"",
1631 ))
1632 }
1633
1634 async fn after_dispatch(
1635 &self,
1636 _ctx: &RequestContext<'_>,
1637 _result: &mut DispatcherResult,
1638 ) -> std::result::Result<(), MiddlewareError> {
1639 Ok(())
1640 }
1641 }
1642
1643 #[tokio::test]
1644 async fn streaming_401_challenge_has_cors_and_exposes_www_authenticate() {
1645 let session_storage = Arc::new(InMemorySessionStorage::new());
1646 let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
1647 let dispatcher = Arc::new(JsonRpcDispatcher::new());
1648 let config = ServerConfig {
1651 origin_policy: turul_http_mcp_server::OriginPolicy::Disabled,
1652 ..Default::default()
1653 };
1654 let capabilities = ServerCapabilities::default();
1655
1656 let mut middleware = MiddlewareStack::new();
1657 middleware.push(Arc::new(ForceChallenge));
1658 let middleware = Arc::new(middleware);
1659
1660 let route_registry = Arc::new(turul_http_mcp_server::RouteRegistry::new());
1661
1662 let handler = LambdaMcpHandler::with_middleware(
1663 config,
1664 dispatcher,
1665 session_storage,
1666 stream_manager,
1667 StreamConfig::default(),
1668 capabilities,
1669 middleware,
1670 false,
1671 route_registry,
1672 )
1673 .with_cors(CorsConfig::default());
1674
1675 let req = Request::builder()
1676 .method("POST")
1677 .uri("/mcp")
1678 .header("Content-Type", "application/json")
1679 .header("Accept", "application/json, text/event-stream")
1680 .header("MCP-Protocol-Version", "2025-11-25")
1681 .header("Origin", "https://client.example.test")
1682 .body(LambdaBody::Text(
1683 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
1684 ))
1685 .unwrap();
1686
1687 let resp = handler.handle_streaming(req).await.unwrap();
1688 let headers = resp.headers();
1689
1690 assert_eq!(resp.status(), 401, "challenge must be 401");
1691 assert!(
1692 headers.contains_key("www-authenticate"),
1693 "WWW-Authenticate must be preserved through streaming transport",
1694 );
1695 assert!(
1696 headers.contains_key("access-control-allow-origin"),
1697 "401 response must carry Access-Control-Allow-Origin",
1698 );
1699 let expose = headers
1700 .get("access-control-expose-headers")
1701 .and_then(|v| v.to_str().ok())
1702 .unwrap_or("");
1703 assert!(
1704 expose
1705 .split(',')
1706 .map(str::trim)
1707 .any(|h| h.eq_ignore_ascii_case("WWW-Authenticate")),
1708 "expose-headers must include WWW-Authenticate; got {expose:?}",
1709 );
1710 }
1711 }
1712}