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_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#[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(feature = "dynamic-tools")]
50 tool_registry: Option<Arc<turul_mcp_server::ToolRegistry>>,
51
52 #[cfg(feature = "cors")]
54 cors_config: Option<CorsConfig>,
55}
56
57impl LambdaMcpHandler {
58 #[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 let middleware_stack = Arc::new(turul_http_mcp_server::middleware::MiddlewareStack::new());
75
76 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 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, );
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 #[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 let middleware_stack = Arc::new(turul_http_mcp_server::middleware::MiddlewareStack::new());
123
124 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 let streamable_handler = StreamableHttpHandler::new(
136 Arc::new(config),
137 dispatcher,
138 session_storage,
139 stream_manager,
140 capabilities,
141 middleware_stack,
142 None, );
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 #[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 #[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 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 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 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(¬ifier));
240 self.streamable_handler = self.streamable_handler.with_tool_notifier(notifier);
241 self
242 }
243
244 #[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 #[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 pub fn get_stream_manager(&self) -> &Arc<StreamManager> {
260 self.session_handler.get_stream_manager()
261 }
262
263 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 #[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 #[cfg(feature = "dynamic-tools")]
297 if let Some(ref registry) = self.tool_registry {
298 if let Err(e) = registry.check_for_changes().await {
299 tracing::warn!(error = %e, "Failed to check for tool changes");
300 }
301 }
302
303 let hyper_req = crate::adapter::lambda_to_hyper_request(req)?;
305
306 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) => {} 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 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 let mut lambda_resp = crate::adapter::hyper_to_lambda_response(hyper_resp).await?;
356
357 #[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 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 #[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 return Ok(self.convert_lambda_response_to_streaming(preflight_response));
404 }
405
406 #[cfg(feature = "dynamic-tools")]
408 if let Some(ref registry) = self.tool_registry {
409 if let Err(e) = registry.check_for_changes().await {
410 tracing::warn!(error = %e, "Failed to check for tool changes (streaming)");
411 }
412 }
413
414 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 let path = hyper_req.uri().path().to_string();
420 if !self.route_registry.is_empty() {
421 match self.route_registry.match_route(&path) {
422 Ok(Some(route_handler)) => {
423 debug!("Custom route matched (streaming): {}", path);
424 use http_body_util::BodyExt;
425 let (parts, body) = hyper_req.into_parts();
426 let boxed_req = hyper::Request::from_parts(parts, body.boxed_unsync());
427 return Ok(route_handler.handle(boxed_req).await);
428 }
429 Ok(None) => {} Err(e) => {
431 debug!("Route validation error (streaming): {}", e);
432 return Ok(e.into_response());
433 }
434 }
435 }
436
437 use turul_http_mcp_server::protocol::McpProtocolVersion;
439 let protocol_version = hyper_req
440 .headers()
441 .get("MCP-Protocol-Version")
442 .and_then(|h| h.to_str().ok())
443 .and_then(McpProtocolVersion::parse_version)
444 .unwrap_or(McpProtocolVersion::V2025_06_18);
445
446 let hyper_resp = if protocol_version.supports_streamable_http() {
448 debug!(
450 "Using StreamableHttpHandler for protocol {}",
451 protocol_version.to_string()
452 );
453 self.streamable_handler.handle_request(hyper_req).await
454 } else {
455 debug!(
457 "Using SessionMcpHandler for legacy protocol {}",
458 protocol_version.to_string()
459 );
460 self.session_handler
461 .handle_mcp_request(hyper_req)
462 .await
463 .map_err(|e| {
464 Box::new(crate::error::LambdaError::McpFramework(e.to_string()))
465 as Box<dyn std::error::Error + Send + Sync>
466 })?
467 };
468
469 let mut lambda_resp = crate::adapter::hyper_to_lambda_streaming(hyper_resp);
471
472 #[cfg(feature = "cors")]
474 if let Some(ref cors_config) = self.cors_config {
475 inject_cors_headers(&mut lambda_resp, cors_config, request_origin.as_deref())
476 .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
477 }
478
479 Ok(lambda_resp)
480 }
481
482 fn convert_lambda_response_to_streaming(
484 &self,
485 lambda_response: LambdaResponse<LambdaBody>,
486 ) -> lambda_http::Response<http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>>
487 {
488 use bytes::Bytes;
489 use http_body_util::{BodyExt, Full};
490
491 let (parts, body) = lambda_response.into_parts();
492 let body_bytes = match body {
493 LambdaBody::Empty => Bytes::new(),
494 LambdaBody::Text(text) => Bytes::from(text),
495 LambdaBody::Binary(bytes) => Bytes::from(bytes),
496 _ => Bytes::new(),
497 };
498
499 let streaming_body = Full::new(body_bytes)
501 .map_err(|e: std::convert::Infallible| match e {})
502 .boxed_unsync();
503
504 lambda_http::Response::from_parts(parts, streaming_body)
505 }
506}
507
508#[cfg(test)]
509mod tests {
510 use super::*;
511 use http::Request;
512 use turul_mcp_session_storage::InMemorySessionStorage;
513
514 #[tokio::test]
515 async fn test_handler_creation() {
516 let session_storage = Arc::new(InMemorySessionStorage::new());
517 let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
518 let dispatcher = JsonRpcDispatcher::new();
519 let config = ServerConfig::default();
520 let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
521 let capabilities = ServerCapabilities::default();
522
523 let handler = LambdaMcpHandler::new(
524 dispatcher,
525 session_storage,
526 stream_manager,
527 config,
528 StreamConfig::default(),
529 implementation,
530 capabilities,
531 false, #[cfg(feature = "cors")]
533 None,
534 );
535
536 assert!(!handler.sse_enabled);
538 }
539
540 #[tokio::test]
541 async fn test_sse_enabled_with_handle_works() {
542 let session_storage = Arc::new(InMemorySessionStorage::new());
543 let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
544 let dispatcher = JsonRpcDispatcher::new();
545 let config = ServerConfig::default();
546 let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
547 let capabilities = ServerCapabilities::default();
548
549 let handler = LambdaMcpHandler::new(
551 dispatcher,
552 session_storage,
553 stream_manager,
554 config,
555 StreamConfig::default(),
556 implementation,
557 capabilities,
558 true, #[cfg(feature = "cors")]
560 None,
561 );
562
563 let lambda_req = Request::builder()
565 .method("POST")
566 .uri("/mcp")
567 .body(LambdaBody::Text(
568 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
569 ))
570 .unwrap();
571
572 let result = handler.handle(lambda_req).await;
574 assert!(
575 result.is_ok(),
576 "handle() should work with SSE enabled for snapshot-based responses"
577 );
578 }
579
580 #[tokio::test]
582 async fn test_stream_config_preservation() {
583 let session_storage = Arc::new(InMemorySessionStorage::new());
584 let dispatcher = JsonRpcDispatcher::new();
585 let config = ServerConfig::default();
586 let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
587 let capabilities = ServerCapabilities::default();
588
589 let custom_stream_config = StreamConfig {
591 channel_buffer_size: 1024, max_replay_events: 200, keepalive_interval_seconds: 10, cors_origin: "https://custom-test.example.com".to_string(), };
596
597 let stream_manager = Arc::new(StreamManager::with_config(
599 session_storage.clone(),
600 custom_stream_config.clone(),
601 ));
602
603 let handler = LambdaMcpHandler::new(
604 dispatcher,
605 session_storage,
606 stream_manager,
607 config,
608 custom_stream_config.clone(),
609 implementation,
610 capabilities,
611 false, #[cfg(feature = "cors")]
613 None,
614 );
615
616 assert!(!handler.sse_enabled);
618
619 let stream_manager = handler.get_stream_manager();
621
622 let actual_config = stream_manager.get_config();
624
625 assert_eq!(
626 actual_config.channel_buffer_size, custom_stream_config.channel_buffer_size,
627 "Custom channel_buffer_size was not propagated correctly"
628 );
629 assert_eq!(
630 actual_config.max_replay_events, custom_stream_config.max_replay_events,
631 "Custom max_replay_events was not propagated correctly"
632 );
633 assert_eq!(
634 actual_config.keepalive_interval_seconds,
635 custom_stream_config.keepalive_interval_seconds,
636 "Custom keepalive_interval_seconds was not propagated correctly"
637 );
638 assert_eq!(
639 actual_config.cors_origin, custom_stream_config.cors_origin,
640 "Custom cors_origin was not propagated correctly"
641 );
642
643 assert!(Arc::strong_count(stream_manager) >= 1);
645 }
646
647 #[tokio::test]
649 async fn test_full_builder_chain_stream_config() {
650 use crate::LambdaMcpServerBuilder;
651 use turul_mcp_session_storage::InMemorySessionStorage;
652
653 let custom_stream_config = turul_http_mcp_server::StreamConfig {
655 channel_buffer_size: 2048, max_replay_events: 500, keepalive_interval_seconds: 15, cors_origin: "https://full-chain-test.example.com".to_string(),
659 };
660
661 let server = LambdaMcpServerBuilder::new()
663 .name("full-chain-test")
664 .version("1.0.0")
665 .storage(Arc::new(InMemorySessionStorage::new()))
666 .sse(true) .stream_config(custom_stream_config.clone())
668 .build()
669 .await
670 .expect("Server should build successfully");
671
672 let handler = server
674 .handler()
675 .await
676 .expect("Handler should be created from server");
677
678 assert!(handler.sse_enabled, "SSE should be enabled");
680
681 let stream_manager = handler.get_stream_manager();
683 let actual_config = stream_manager.get_config();
684
685 assert_eq!(
686 actual_config.channel_buffer_size, custom_stream_config.channel_buffer_size,
687 "Custom channel_buffer_size should be preserved through builder → server → handler chain"
688 );
689 assert_eq!(
690 actual_config.max_replay_events, custom_stream_config.max_replay_events,
691 "Custom max_replay_events should be preserved through builder → server → handler chain"
692 );
693 assert_eq!(
694 actual_config.keepalive_interval_seconds,
695 custom_stream_config.keepalive_interval_seconds,
696 "Custom keepalive_interval_seconds should be preserved through builder → server → handler chain"
697 );
698 assert_eq!(
699 actual_config.cors_origin, custom_stream_config.cors_origin,
700 "Custom cors_origin should be preserved through builder → server → handler chain"
701 );
702
703 assert!(
705 Arc::strong_count(stream_manager) >= 1,
706 "Stream manager should be properly initialized"
707 );
708
709 let test_session_id = uuid::Uuid::now_v7().as_simple().to_string();
712
713 let subscriptions = stream_manager.get_subscriptions(&test_session_id).await;
716 assert!(
717 subscriptions.is_empty(),
718 "New session should have no subscriptions initially"
719 );
720
721 assert_eq!(
724 stream_manager.get_config().channel_buffer_size,
725 2048,
726 "Stream manager should be using the custom buffer size functionally"
727 );
728 }
729
730 #[tokio::test]
735 async fn test_non_streaming_runtime_sse_false() {
736 use crate::LambdaMcpServerBuilder;
737 use turul_mcp_session_storage::InMemorySessionStorage;
738
739 let server = LambdaMcpServerBuilder::new()
740 .name("test-non-streaming-sse-false")
741 .version("1.0.0")
742 .storage(Arc::new(InMemorySessionStorage::new()))
743 .sse(false) .build()
745 .await
746 .expect("Server should build successfully");
747
748 let handler = server
749 .handler()
750 .await
751 .expect("Handler should be created from server");
752
753 assert!(!handler.sse_enabled, "SSE should be disabled");
755
756 let lambda_req = Request::builder()
758 .method("POST")
759 .uri("/mcp")
760 .body(LambdaBody::Text(
761 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
762 ))
763 .unwrap();
764
765 let result = handler.handle(lambda_req).await;
767 assert!(
768 result.is_ok(),
769 "POST /mcp should work with non-streaming + sse(false)"
770 );
771 }
772
773 #[tokio::test]
775 async fn test_non_streaming_runtime_sse_true() {
776 use crate::LambdaMcpServerBuilder;
777 use turul_mcp_session_storage::InMemorySessionStorage;
778
779 let server = LambdaMcpServerBuilder::new()
780 .name("test-non-streaming-sse-true")
781 .version("1.0.0")
782 .storage(Arc::new(InMemorySessionStorage::new()))
783 .sse(true) .build()
785 .await
786 .expect("Server should build successfully");
787
788 let handler = server
789 .handler()
790 .await
791 .expect("Handler should be created from server");
792
793 assert!(handler.sse_enabled, "SSE should be enabled");
795
796 let lambda_req = Request::builder()
798 .method("POST")
799 .uri("/mcp")
800 .body(LambdaBody::Text(
801 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
802 ))
803 .unwrap();
804
805 let result = handler.handle(lambda_req).await;
807 assert!(
808 result.is_ok(),
809 "POST /mcp should work with non-streaming + sse(true)"
810 );
811
812 }
815
816 #[tokio::test]
818 async fn test_streaming_runtime_sse_false() {
819 use crate::LambdaMcpServerBuilder;
820 use turul_mcp_session_storage::InMemorySessionStorage;
821
822 let server = LambdaMcpServerBuilder::new()
823 .name("test-streaming-sse-false")
824 .version("1.0.0")
825 .storage(Arc::new(InMemorySessionStorage::new()))
826 .sse(false) .build()
828 .await
829 .expect("Server should build successfully");
830
831 let handler = server
832 .handler()
833 .await
834 .expect("Handler should be created from server");
835
836 assert!(!handler.sse_enabled, "SSE should be disabled");
838
839 let lambda_req = Request::builder()
841 .method("POST")
842 .uri("/mcp")
843 .body(LambdaBody::Text(
844 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
845 ))
846 .unwrap();
847
848 let result = handler.handle_streaming(lambda_req).await;
850 assert!(
851 result.is_ok(),
852 "Streaming runtime should work with sse(false)"
853 );
854 }
855
856 #[tokio::test]
858 async fn test_streaming_runtime_sse_true() {
859 use crate::LambdaMcpServerBuilder;
860 use turul_mcp_session_storage::InMemorySessionStorage;
861
862 let server = LambdaMcpServerBuilder::new()
863 .name("test-streaming-sse-true")
864 .version("1.0.0")
865 .storage(Arc::new(InMemorySessionStorage::new()))
866 .sse(true) .build()
868 .await
869 .expect("Server should build successfully");
870
871 let handler = server
872 .handler()
873 .await
874 .expect("Handler should be created from server");
875
876 assert!(handler.sse_enabled, "SSE should be enabled");
878
879 let lambda_req = Request::builder()
881 .method("POST")
882 .uri("/mcp")
883 .body(LambdaBody::Text(
884 r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
885 ))
886 .unwrap();
887
888 let result = handler.handle_streaming(lambda_req).await;
890 assert!(
891 result.is_ok(),
892 "Streaming runtime should work with sse(true) for real-time streaming"
893 );
894
895 }
898
899 async fn build_strict_streaming_handler() -> LambdaMcpHandler {
903 use crate::LambdaMcpServerBuilder;
904 use turul_mcp_session_storage::InMemorySessionStorage;
905
906 let server = LambdaMcpServerBuilder::new()
907 .name("lifecycle-test")
908 .version("1.0.0")
909 .tool(LifecycleTestTool)
910 .storage(Arc::new(InMemorySessionStorage::new()))
911 .strict_lifecycle(true) .sse(true)
913 .build()
914 .await
915 .expect("build should succeed");
916
917 server.handler().await.expect("handler should succeed")
918 }
919
920 #[derive(Clone, Default)]
922 struct LifecycleTestTool;
923
924 impl turul_mcp_builders::traits::HasBaseMetadata for LifecycleTestTool {
925 fn name(&self) -> &str {
926 "ping_tool"
927 }
928 }
929 impl turul_mcp_builders::traits::HasDescription for LifecycleTestTool {
930 fn description(&self) -> Option<&str> {
931 Some("test tool")
932 }
933 }
934 impl turul_mcp_builders::traits::HasInputSchema for LifecycleTestTool {
935 fn input_schema(&self) -> &turul_mcp_protocol::ToolSchema {
936 static SCHEMA: std::sync::OnceLock<turul_mcp_protocol::ToolSchema> =
937 std::sync::OnceLock::new();
938 SCHEMA.get_or_init(turul_mcp_protocol::ToolSchema::object)
939 }
940 }
941 impl turul_mcp_builders::traits::HasOutputSchema for LifecycleTestTool {
942 fn output_schema(&self) -> Option<&turul_mcp_protocol::ToolSchema> {
943 None
944 }
945 }
946 impl turul_mcp_builders::traits::HasAnnotations for LifecycleTestTool {
947 fn annotations(&self) -> Option<&turul_mcp_protocol::tools::ToolAnnotations> {
948 None
949 }
950 }
951 impl turul_mcp_builders::traits::HasToolMeta for LifecycleTestTool {
952 fn tool_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
953 None
954 }
955 }
956 impl turul_mcp_builders::traits::HasIcons for LifecycleTestTool {}
957 impl turul_mcp_builders::traits::HasExecution for LifecycleTestTool {}
958
959 #[async_trait::async_trait]
960 impl turul_mcp_server::McpTool for LifecycleTestTool {
961 async fn call(
962 &self,
963 _args: serde_json::Value,
964 _session: Option<turul_mcp_server::SessionContext>,
965 ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
966 Ok(turul_mcp_protocol::tools::CallToolResult::success(vec![
967 turul_mcp_protocol::tools::ToolResult::text("pong"),
968 ]))
969 }
970 }
971
972 fn streaming_mcp_request(body: &str, session_id: Option<&str>) -> LambdaRequest {
974 let mut builder = Request::builder()
975 .method("POST")
976 .uri("/mcp")
977 .header("Content-Type", "application/json")
978 .header("Accept", "application/json, text/event-stream")
979 .header("MCP-Protocol-Version", "2025-11-25");
980
981 if let Some(sid) = session_id {
982 builder = builder.header("Mcp-Session-Id", sid);
983 }
984
985 builder.body(LambdaBody::Text(body.to_string())).unwrap()
986 }
987
988 async fn collect_streaming_body(
990 response: lambda_http::Response<
991 http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
992 >,
993 ) -> (http::StatusCode, String) {
994 use http_body_util::BodyExt;
995 let status = response.status();
996 let session_id = response
997 .headers()
998 .get("Mcp-Session-Id")
999 .and_then(|v| v.to_str().ok())
1000 .map(String::from);
1001 let body_bytes = response
1002 .into_body()
1003 .collect()
1004 .await
1005 .map(|c| c.to_bytes())
1006 .unwrap_or_default();
1007 let body_str = String::from_utf8_lossy(&body_bytes).to_string();
1008 let _ = session_id; (status, body_str)
1010 }
1011
1012 fn extract_session_id(
1014 response: &lambda_http::Response<
1015 http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
1016 >,
1017 ) -> Option<String> {
1018 response
1019 .headers()
1020 .get("Mcp-Session-Id")
1021 .and_then(|v| v.to_str().ok())
1022 .map(String::from)
1023 }
1024
1025 fn parse_response_json(body: &str) -> serde_json::Value {
1027 let json_str = body
1029 .lines()
1030 .find(|line| line.starts_with("data: "))
1031 .map(|line| &line[6..])
1032 .unwrap_or(body.trim());
1033 serde_json::from_str(json_str)
1034 .unwrap_or_else(|e| panic!("Failed to parse JSON from body: {e}\nBody: {body}"))
1035 }
1036
1037 #[tokio::test]
1039 async fn test_lambda_streaming_strict_handshake_succeeds() {
1040 let handler = build_strict_streaming_handler().await;
1041
1042 let init_req = streaming_mcp_request(
1044 &serde_json::json!({
1045 "jsonrpc": "2.0", "method": "initialize", "id": 1,
1046 "params": {
1047 "protocolVersion": "2025-11-25",
1048 "capabilities": {},
1049 "clientInfo": { "name": "test", "version": "1.0.0" }
1050 }
1051 })
1052 .to_string(),
1053 None,
1054 );
1055 let init_resp = handler
1056 .handle_streaming(init_req)
1057 .await
1058 .expect("initialize should succeed");
1059 let session_id = extract_session_id(&init_resp).expect("must return session ID");
1060 let (status, _body) = collect_streaming_body(init_resp).await;
1061 assert_eq!(status, 200, "initialize should return 200");
1062
1063 let notif_req = streaming_mcp_request(
1065 &serde_json::json!({
1066 "jsonrpc": "2.0",
1067 "method": "notifications/initialized",
1068 "params": {}
1069 })
1070 .to_string(),
1071 Some(&session_id),
1072 );
1073 let notif_resp = handler
1074 .handle_streaming(notif_req)
1075 .await
1076 .expect("notification should succeed");
1077 let (status, _) = collect_streaming_body(notif_resp).await;
1078 assert_eq!(status, 202, "notifications/initialized should return 202");
1079
1080 let list_req = streaming_mcp_request(
1082 &serde_json::json!({
1083 "jsonrpc": "2.0", "method": "tools/list", "id": 2
1084 })
1085 .to_string(),
1086 Some(&session_id),
1087 );
1088 let list_resp = handler
1089 .handle_streaming(list_req)
1090 .await
1091 .expect("tools/list should succeed");
1092 let (status, body) = collect_streaming_body(list_resp).await;
1093 assert_eq!(status, 200, "tools/list should return 200");
1094 let json = parse_response_json(&body);
1095 assert!(
1096 json["result"]["tools"].is_array(),
1097 "tools/list should return tools array: {json}"
1098 );
1099
1100 let call_req = streaming_mcp_request(
1102 &serde_json::json!({
1103 "jsonrpc": "2.0", "method": "tools/call", "id": 3,
1104 "params": { "name": "ping_tool", "arguments": {} }
1105 })
1106 .to_string(),
1107 Some(&session_id),
1108 );
1109 let call_resp = handler
1110 .handle_streaming(call_req)
1111 .await
1112 .expect("tools/call should succeed");
1113 let (status, body) = collect_streaming_body(call_resp).await;
1114 assert_eq!(status, 200, "tools/call should return 200");
1115 let json = parse_response_json(&body);
1116 assert!(
1117 json["result"].is_object(),
1118 "tools/call should return result: {json}"
1119 );
1120 }
1121
1122 #[tokio::test]
1124 async fn test_lambda_streaming_strict_rejects_before_initialized() {
1125 let handler = build_strict_streaming_handler().await;
1126
1127 let init_req = streaming_mcp_request(
1129 &serde_json::json!({
1130 "jsonrpc": "2.0", "method": "initialize", "id": 1,
1131 "params": {
1132 "protocolVersion": "2025-11-25",
1133 "capabilities": {},
1134 "clientInfo": { "name": "test", "version": "1.0.0" }
1135 }
1136 })
1137 .to_string(),
1138 None,
1139 );
1140 let init_resp = handler.handle_streaming(init_req).await.unwrap();
1141 let session_id = extract_session_id(&init_resp).unwrap();
1142 let _ = collect_streaming_body(init_resp).await;
1143
1144 let list_req = streaming_mcp_request(
1146 &serde_json::json!({
1147 "jsonrpc": "2.0", "method": "tools/list", "id": 2
1148 })
1149 .to_string(),
1150 Some(&session_id),
1151 );
1152 let list_resp = handler.handle_streaming(list_req).await.unwrap();
1153 let (_, body) = collect_streaming_body(list_resp).await;
1154 let json = parse_response_json(&body);
1155 assert!(
1156 json["error"].is_object(),
1157 "tools/list should return JSON-RPC error: {json}"
1158 );
1159 assert_eq!(
1160 json["error"]["code"].as_i64().unwrap(),
1161 -32031,
1162 "tools/list must return SessionError code -32031, got: {json}"
1163 );
1164 assert!(
1165 json["error"]["message"]
1166 .as_str()
1167 .unwrap()
1168 .contains("notifications/initialized"),
1169 "Error must mention notifications/initialized: {}",
1170 json["error"]["message"]
1171 );
1172
1173 let call_req = streaming_mcp_request(
1175 &serde_json::json!({
1176 "jsonrpc": "2.0", "method": "tools/call", "id": 3,
1177 "params": { "name": "ping_tool", "arguments": {} }
1178 })
1179 .to_string(),
1180 Some(&session_id),
1181 );
1182 let call_resp = handler.handle_streaming(call_req).await.unwrap();
1183 let (_, body) = collect_streaming_body(call_resp).await;
1184 let json = parse_response_json(&body);
1185 assert!(
1186 json["error"].is_object(),
1187 "tools/call should return JSON-RPC error: {json}"
1188 );
1189 assert_eq!(
1190 json["error"]["code"].as_i64().unwrap(),
1191 -32031,
1192 "tools/call must return SessionError code -32031, got: {json}"
1193 );
1194 assert!(
1195 json["error"]["message"]
1196 .as_str()
1197 .unwrap()
1198 .contains("notifications/initialized"),
1199 "Error must mention notifications/initialized: {}",
1200 json["error"]["message"]
1201 );
1202 }
1203
1204 #[tokio::test]
1206 async fn test_lambda_streaming_initialized_is_effective_immediately() {
1207 let handler = build_strict_streaming_handler().await;
1208
1209 let init_req = streaming_mcp_request(
1211 &serde_json::json!({
1212 "jsonrpc": "2.0", "method": "initialize", "id": 1,
1213 "params": {
1214 "protocolVersion": "2025-11-25",
1215 "capabilities": {},
1216 "clientInfo": { "name": "test", "version": "1.0.0" }
1217 }
1218 })
1219 .to_string(),
1220 None,
1221 );
1222 let init_resp = handler.handle_streaming(init_req).await.unwrap();
1223 let session_id = extract_session_id(&init_resp).unwrap();
1224 let _ = collect_streaming_body(init_resp).await;
1225
1226 let notif_req = streaming_mcp_request(
1228 &serde_json::json!({
1229 "jsonrpc": "2.0",
1230 "method": "notifications/initialized",
1231 "params": {}
1232 })
1233 .to_string(),
1234 Some(&session_id),
1235 );
1236 let notif_resp = handler.handle_streaming(notif_req).await.unwrap();
1237 let (status, _) = collect_streaming_body(notif_resp).await;
1238 assert_eq!(status, 202);
1239
1240 let list_req = streaming_mcp_request(
1242 &serde_json::json!({
1243 "jsonrpc": "2.0", "method": "tools/list", "id": 2
1244 })
1245 .to_string(),
1246 Some(&session_id),
1247 );
1248 let list_resp = handler.handle_streaming(list_req).await.unwrap();
1249 let (status, body) = collect_streaming_body(list_resp).await;
1250 assert_eq!(
1251 status, 200,
1252 "tools/list must succeed immediately after initialized"
1253 );
1254 let json = parse_response_json(&body);
1255 assert!(
1256 json["result"]["tools"].is_array(),
1257 "Must return tools list, not error: {json}"
1258 );
1259 }
1260
1261 #[tokio::test]
1263 async fn test_lambda_streaming_lenient_mode_allows_without_initialized() {
1264 use crate::LambdaMcpServerBuilder;
1265 use turul_mcp_session_storage::InMemorySessionStorage;
1266
1267 let server = LambdaMcpServerBuilder::new()
1268 .name("lenient-test")
1269 .version("1.0.0")
1270 .tool(LifecycleTestTool)
1271 .storage(Arc::new(InMemorySessionStorage::new()))
1272 .strict_lifecycle(false) .sse(true)
1274 .build()
1275 .await
1276 .unwrap();
1277
1278 let handler = server.handler().await.unwrap();
1279
1280 let init_req = streaming_mcp_request(
1282 &serde_json::json!({
1283 "jsonrpc": "2.0", "method": "initialize", "id": 1,
1284 "params": {
1285 "protocolVersion": "2025-11-25",
1286 "capabilities": {},
1287 "clientInfo": { "name": "test", "version": "1.0.0" }
1288 }
1289 })
1290 .to_string(),
1291 None,
1292 );
1293 let init_resp = handler.handle_streaming(init_req).await.unwrap();
1294 let session_id = extract_session_id(&init_resp).unwrap();
1295 let _ = collect_streaming_body(init_resp).await;
1296
1297 let list_req = streaming_mcp_request(
1299 &serde_json::json!({
1300 "jsonrpc": "2.0", "method": "tools/list", "id": 2
1301 })
1302 .to_string(),
1303 Some(&session_id),
1304 );
1305 let list_resp = handler.handle_streaming(list_req).await.unwrap();
1306 let (status, body) = collect_streaming_body(list_resp).await;
1307 assert_eq!(
1308 status, 200,
1309 "Lenient mode should allow tools/list without initialized"
1310 );
1311 let json = parse_response_json(&body);
1312 assert!(
1313 json["result"]["tools"].is_array(),
1314 "Must return tools list in lenient mode: {json}"
1315 );
1316 }
1317}