Skip to main content

turul_mcp_aws_lambda/
lib.rs

1//! AWS Lambda integration for turul-mcp-framework
2//!
3//! This crate provides seamless integration between the turul-mcp-framework and AWS Lambda,
4//! enabling serverless deployment of MCP servers with proper session management, CORS handling,
5//! and SSE streaming support.
6//!
7//! ## Architecture
8//!
9//! The crate bridges the gap between Lambda's HTTP execution model and the framework's
10//! hyper-based architecture through:
11//!
12//! - **Type Conversion**: Clean conversion between `lambda_http` and `hyper` types
13//! - **Handler Registration**: Direct tool registration with `JsonRpcDispatcher`
14//! - **Session Management**: DynamoDB-backed session persistence across invocations
15//! - **CORS Support**: Proper CORS header injection for browser clients
16//! - **SSE Streaming**: Server-Sent Events adaptation through Lambda's streaming response
17//!
18//! ## Quick Start
19//!
20//! ```rust,no_run
21//! use turul_mcp_aws_lambda::LambdaMcpServerBuilder;
22//! use turul_mcp_derive::McpTool;
23//! use turul_mcp_server::{McpResult, SessionContext};
24//!
25//! #[derive(McpTool, Clone, Default)]
26//! #[tool(name = "example", description = "Example tool")]
27//! struct ExampleTool {
28//!     #[param(description = "Example parameter")]
29//!     value: String,
30//! }
31//!
32//! impl ExampleTool {
33//!     async fn execute(&self, _session: Option<SessionContext>) -> McpResult<String> {
34//!         Ok(format!("Got: {}", self.value))
35//!     }
36//! }
37//!
38//! #[tokio::main]
39//! async fn main() -> Result<(), lambda_http::Error> {
40//!     let server = LambdaMcpServerBuilder::new()
41//!         .tool(ExampleTool::default())
42//!         .build()
43//!         .await?;
44//!
45//!     let handler = server.handler().await?;
46//!
47//!     // run_streaming handles API Gateway completion invocations gracefully
48//!     turul_mcp_aws_lambda::run_streaming(handler).await
49//! }
50//! ```
51//!
52//! ## Streaming Entry Points
53//!
54//! Two entry points replace `lambda_http::run_with_streaming_response()`:
55//!
56//! - [`run_streaming()`] — standard path: pass a [`LambdaMcpHandler`] directly
57//! - [`run_streaming_with()`] — custom dispatch: provide your own closure for
58//!   pre-dispatch logic (e.g., `.well-known` routing) while still getting
59//!   completion-invocation handling for free
60
61pub mod adapter;
62pub mod builder;
63pub mod error;
64pub mod handler;
65pub mod prelude;
66pub mod server;
67
68#[cfg(feature = "cors")]
69pub mod cors;
70
71#[cfg(feature = "sse")]
72pub mod streaming;
73
74// Re-exports for convenience
75/// Builder for creating Lambda MCP servers with fluent configuration API
76pub use builder::LambdaMcpServerBuilder;
77/// Lambda-specific error types and result aliases
78pub use error::{LambdaError, Result};
79/// Lambda request handler with session management and protocol conversion
80pub use handler::LambdaMcpHandler;
81/// Core Lambda MCP server implementation with DynamoDB integration
82pub use server::LambdaMcpServer;
83
84#[cfg(feature = "cors")]
85pub use cors::{CorsConfig, create_preflight_response, inject_cors_headers};
86
87/// Classification of a raw Lambda runtime event payload.
88///
89/// Used by [`run_streaming()`] to distinguish API Gateway requests from
90/// streaming completion invocations and unknown event shapes.
91#[derive(Debug)]
92enum RuntimeEventClassification {
93    /// Valid API Gateway / ALB / Function URL event.
94    ///
95    /// Stores `Box<LambdaRequest>` to avoid a large enum variant (clippy
96    /// `large_enum_variant`). Callers dereference with `(*lambda_request).into()`
97    /// to move the inner `LambdaRequest` into an `http::Request`.
98    ApiGatewayEvent(Box<lambda_http::request::LambdaRequest>),
99    /// AWS streaming completion invocation (contains `invokeCompletionStatus`)
100    StreamingCompletion,
101    /// Unrecognized payload — not API Gateway, not completion
102    UnrecognizedEvent,
103}
104
105/// Classify a raw JSON payload into one of three categories.
106///
107/// Order matters:
108/// 1. Try API Gateway/ALB/WebSocket deserialization first (most common path)
109/// 2. Check for streaming completion signature (`invokeCompletionStatus` at top level)
110/// 3. Everything else is unrecognized
111///
112/// # Completion Detection Heuristic
113///
114/// Streaming completion payloads are identified by the presence of an
115/// `invokeCompletionStatus` field at the top level. This is a compatibility
116/// heuristic based on observed AWS behavior as of 2026-03 — AWS does not
117/// officially document this payload shape.
118///
119/// As of this writing, no API Gateway v1/v2, ALB, or WebSocket event
120/// produced by `lambda_http` contains this field at the top level.
121/// The fixture corpus in `src/fixtures/` guards against drift.
122///
123/// **Precedence**: API Gateway deserialization is attempted first.
124/// If a payload is both a valid API Gateway event AND contains
125/// `invokeCompletionStatus`, it will be classified as `ApiGatewayEvent`.
126/// The completion heuristic only applies to payloads that fail API
127/// Gateway parsing. This means false-positive completion detection is
128/// preferred over retry storms — an intentional design choice.
129fn classify_runtime_event(payload: serde_json::Value) -> RuntimeEventClassification {
130    // Fast path: try API Gateway event deserialization
131    if let Ok(request) =
132        serde_json::from_value::<lambda_http::request::LambdaRequest>(payload.clone())
133    {
134        return RuntimeEventClassification::ApiGatewayEvent(Box::new(request));
135    }
136
137    // Check for streaming completion signature
138    if payload.get("invokeCompletionStatus").is_some() {
139        return RuntimeEventClassification::StreamingCompletion;
140    }
141
142    // Unknown payload shape
143    RuntimeEventClassification::UnrecognizedEvent
144}
145
146type StreamBody = http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>;
147type StreamResult =
148    lambda_runtime::StreamResponse<http_body_util::BodyDataStream<EnsureOneFrame<StreamBody>>>;
149
150/// Result of [`handle_runtime_payload()`], carrying both the Lambda response
151/// and a static string identifying the event type for logging/observability.
152struct HandleResult {
153    response: StreamResult,
154    /// One of `"api_gateway_event"`, `"streaming_completion"`, or
155    /// `"unrecognized_lambda_payload"`.
156    event_type: &'static str,
157}
158
159/// Process a raw Lambda runtime payload into a streaming response.
160///
161/// Classifies the payload via [`classify_runtime_event()`], dispatches API
162/// Gateway events through `dispatch`, and acknowledges non-API payloads with
163/// an empty 200 response. Returns a [`HandleResult`] so the caller can
164/// inspect `event_type` for logging decisions.
165async fn handle_runtime_payload<F, Fut>(
166    payload: serde_json::Value,
167    context: lambda_runtime::Context,
168    dispatch: F,
169) -> std::result::Result<HandleResult, lambda_http::Error>
170where
171    F: FnOnce(lambda_http::Request) -> Fut,
172    Fut: std::future::Future<
173            Output = std::result::Result<http::Response<StreamBody>, lambda_http::Error>,
174        >,
175{
176    match classify_runtime_event(payload) {
177        RuntimeEventClassification::ApiGatewayEvent(lambda_request) => {
178            use lambda_http::RequestExt;
179            let request: lambda_http::Request = (*lambda_request).into();
180            let request = request.with_lambda_context(context);
181            let response = dispatch(request).await?;
182            Ok(HandleResult {
183                response: into_lambda_stream_response(response),
184                event_type: "api_gateway_event",
185            })
186        }
187        RuntimeEventClassification::StreamingCompletion => Ok(HandleResult {
188            response: into_lambda_stream_response(empty_streaming_response()),
189            event_type: "streaming_completion",
190        }),
191        RuntimeEventClassification::UnrecognizedEvent => Ok(HandleResult {
192            response: into_lambda_stream_response(empty_streaming_response()),
193            event_type: "unrecognized_lambda_payload",
194        }),
195    }
196}
197
198/// Map an event type string to the appropriate tracing log level.
199///
200/// Returns `Some(Level::WARN)` for unrecognized payloads (surfaced in
201/// CloudWatch), `Some(Level::DEBUG)` for completion acks (normally silent),
202/// and `None` for API Gateway events (no extra logging needed).
203fn event_log_level(event_type: &str) -> Option<tracing::Level> {
204    match event_type {
205        "streaming_completion" => Some(tracing::Level::DEBUG),
206        "unrecognized_lambda_payload" => Some(tracing::Level::WARN),
207        _ => None,
208    }
209}
210
211/// Run the Lambda MCP handler with streaming response support.
212///
213/// This replaces `lambda_http::run_with_streaming_response(service_fn(...))` and
214/// gracefully handles API Gateway streaming completion invocations that would
215/// otherwise cause deserialization errors in the Lambda runtime.
216///
217/// ## Problem
218///
219/// When API Gateway uses `response-streaming-invocations`, it sends a secondary
220/// "completion" invocation after the streaming response finishes. This invocation
221/// is NOT an API Gateway proxy event — `lambda_http` cannot deserialize it, causing
222/// ERROR logs and CloudWatch Lambda Error metrics for every streaming response.
223///
224/// ## Solution
225///
226/// This function uses `lambda_runtime::run()` directly with `serde_json::Value`
227/// (which always deserializes), then classifies the payload three ways via
228/// `classify_runtime_event()`:
229///
230/// - **`ApiGatewayEvent`** — dispatched to the handler normally
231/// - **`StreamingCompletion`** — acknowledged silently (`debug` log)
232/// - **`UnrecognizedEvent`** — acknowledged with a `warn` log to surface
233///   unexpected payload shapes in CloudWatch without triggering Lambda retries
234///
235/// ## Usage
236///
237/// ```rust,ignore
238/// let handler = server.handler().await?;
239/// turul_mcp_aws_lambda::run_streaming(handler).await
240/// ```
241pub async fn run_streaming(
242    handler: LambdaMcpHandler,
243) -> std::result::Result<(), lambda_http::Error> {
244    use lambda_runtime::{LambdaEvent, service_fn};
245
246    lambda_runtime::run(service_fn(move |event: LambdaEvent<serde_json::Value>| {
247        let handler = handler.clone();
248        async move {
249            let result = handle_runtime_payload(event.payload, event.context, |req| {
250                handler.handle_streaming(req)
251            })
252            .await?;
253
254            match event_log_level(result.event_type) {
255                Some(level) if level == tracing::Level::WARN => {
256                    tracing::warn!(
257                        event_type = result.event_type,
258                        "Received unrecognized Lambda invocation payload"
259                    );
260                }
261                Some(_) => {
262                    tracing::debug!(
263                        event_type = result.event_type,
264                        "Acknowledging streaming completion"
265                    );
266                }
267                None => {}
268            }
269
270            Ok::<_, lambda_http::Error>(result.response)
271        }
272    }))
273    .await
274}
275
276/// Run a custom dispatch function with streaming response support.
277///
278/// Like [`run_streaming()`], but accepts a custom dispatch closure instead of
279/// a [`LambdaMcpHandler`]. Use this when you need pre-dispatch logic
280/// (e.g., `.well-known` routing) that runs before the MCP handler.
281///
282/// The dispatch closure is called once per API Gateway invocation. Streaming
283/// completion invocations and unrecognized payloads are acknowledged
284/// automatically without invoking the closure.
285///
286/// ## Usage
287///
288/// ```rust,ignore
289/// async fn lambda_handler(request: Request) -> Result<Response, Error> {
290///     // pre-dispatch logic here (e.g., .well-known short-circuit)
291///     let handler = HANDLER.get_or_try_init(|| async { ... }).await?;
292///     handler.handle_streaming(request).await
293/// }
294///
295/// turul_mcp_aws_lambda::run_streaming_with(lambda_handler).await
296/// ```
297pub async fn run_streaming_with<F, Fut>(dispatch: F) -> std::result::Result<(), lambda_http::Error>
298where
299    F: Fn(lambda_http::Request) -> Fut + Clone + Send + 'static,
300    Fut: std::future::Future<
301            Output = std::result::Result<http::Response<StreamBody>, lambda_http::Error>,
302        > + Send,
303{
304    use lambda_runtime::{LambdaEvent, service_fn};
305
306    lambda_runtime::run(service_fn(move |event: LambdaEvent<serde_json::Value>| {
307        let dispatch = dispatch.clone();
308        async move {
309            let result = handle_runtime_payload(event.payload, event.context, dispatch).await?;
310
311            match event_log_level(result.event_type) {
312                Some(level) if level == tracing::Level::WARN => {
313                    tracing::warn!(
314                        event_type = result.event_type,
315                        "Received unrecognized Lambda invocation payload"
316                    );
317                }
318                Some(_) => {
319                    tracing::debug!(
320                        event_type = result.event_type,
321                        "Acknowledging streaming completion"
322                    );
323                }
324                None => {}
325            }
326
327            Ok::<_, lambda_http::Error>(result.response)
328        }
329    }))
330    .await
331}
332
333/// Body adapter that guarantees at least one data frame is yielded.
334///
335/// Lambda Response Streaming framing requires the multipart body stream
336/// to produce at least one chunk before EOF. If the wrapped body yields
337/// zero data frames (e.g. `Empty::new()`, or `Full::new(Bytes::new())`
338/// after `Full` short-circuits, or any body whose first poll returns
339/// `None`), the Runtime API request body terminates without matching
340/// the framing contract, the connection closes with `hyper::Error`
341/// `IncompleteMessage`, and AWS reports the invocation as a 60-second
342/// timeout. API Gateway then emits 502 to the client.
343///
344/// This adapter is invisible for bodies that already produce ≥1 data
345/// frame. For zero-frame bodies, it emits one zero-length `Bytes` data
346/// frame, which satisfies the multipart framing contract without
347/// changing the visible response body.
348///
349/// See ADR-026.
350struct EnsureOneFrame<B> {
351    inner: B,
352    /// Tracks whether we have observed (or fabricated) the first data
353    /// frame, so we only inject the fallback once and only if needed.
354    first_frame_state: FirstFrameState,
355}
356
357#[derive(Debug, Clone, Copy, PartialEq, Eq)]
358enum FirstFrameState {
359    /// Haven't yet seen a frame from the underlying body. If the
360    /// underlying poll returns `Ready(None)`, we inject the fallback.
361    Initial,
362    /// At least one frame (real or fallback-injected) has been emitted
363    /// downstream. Subsequent polls forward unchanged.
364    Done,
365    /// We injected the fallback frame on the last poll; the next poll
366    /// must return `Ready(None)` to terminate the stream.
367    FallbackEmitted,
368}
369
370impl<B> http_body::Body for EnsureOneFrame<B>
371where
372    B: http_body::Body<Data = bytes::Bytes> + Unpin,
373{
374    type Data = bytes::Bytes;
375    type Error = B::Error;
376
377    fn poll_frame(
378        mut self: std::pin::Pin<&mut Self>,
379        cx: &mut std::task::Context<'_>,
380    ) -> std::task::Poll<Option<std::result::Result<http_body::Frame<Self::Data>, Self::Error>>>
381    {
382        use std::task::Poll;
383        match self.first_frame_state {
384            FirstFrameState::FallbackEmitted => Poll::Ready(None),
385            FirstFrameState::Done => std::pin::Pin::new(&mut self.inner).poll_frame(cx),
386            FirstFrameState::Initial => {
387                match std::pin::Pin::new(&mut self.inner).poll_frame(cx) {
388                    Poll::Pending => Poll::Pending,
389                    Poll::Ready(Some(Ok(frame))) => {
390                        self.first_frame_state = FirstFrameState::Done;
391                        Poll::Ready(Some(Ok(frame)))
392                    }
393                    Poll::Ready(Some(Err(e))) => {
394                        // An error before any frame is yielded still satisfies
395                        // the "we have something to send" contract from
396                        // lambda_runtime's perspective — let the error
397                        // propagate as-is. No fallback needed.
398                        self.first_frame_state = FirstFrameState::Done;
399                        Poll::Ready(Some(Err(e)))
400                    }
401                    Poll::Ready(None) => {
402                        self.first_frame_state = FirstFrameState::FallbackEmitted;
403                        Poll::Ready(Some(Ok(http_body::Frame::data(bytes::Bytes::new()))))
404                    }
405                }
406            }
407        }
408    }
409
410    fn is_end_stream(&self) -> bool {
411        match self.first_frame_state {
412            FirstFrameState::Initial => false,
413            FirstFrameState::Done => self.inner.is_end_stream(),
414            FirstFrameState::FallbackEmitted => true,
415        }
416    }
417
418    fn size_hint(&self) -> http_body::SizeHint {
419        // Defer to inner. The fallback frame is zero bytes so adds nothing
420        // to the upper/lower bound; consumers reading `size_hint` are
421        // typically only concerned with content-length, which is irrelevant
422        // to streaming framing.
423        self.inner.size_hint()
424    }
425}
426
427/// Convert an HTTP response into a Lambda `StreamResponse`.
428///
429/// Replicates `lambda_http::streaming::into_stream_response` (which is private)
430/// by extracting status/headers/cookies into `MetadataPrelude` and converting
431/// the body into a `Stream`.
432///
433/// The body is wrapped in [`EnsureOneFrame`] to guarantee the resulting
434/// stream yields at least one chunk — Lambda Response Streaming framing
435/// requires this. See ADR-026.
436fn into_lambda_stream_response<B>(
437    response: http::Response<B>,
438) -> lambda_runtime::StreamResponse<http_body_util::BodyDataStream<EnsureOneFrame<B>>>
439where
440    B: http_body::Body<Data = bytes::Bytes> + Unpin + Send + 'static,
441{
442    let (parts, body) = response.into_parts();
443    let mut headers = parts.headers;
444
445    // Extract Set-Cookie headers into the cookies vec (Lambda streaming protocol)
446    let cookies = headers
447        .get_all(http::header::SET_COOKIE)
448        .iter()
449        .map(|c| String::from_utf8_lossy(c.as_bytes()).to_string())
450        .collect::<Vec<_>>();
451    headers.remove(http::header::SET_COOKIE);
452
453    let body = EnsureOneFrame {
454        inner: body,
455        first_frame_state: FirstFrameState::Initial,
456    };
457
458    lambda_runtime::StreamResponse {
459        metadata_prelude: lambda_runtime::MetadataPrelude {
460            headers,
461            status_code: parts.status,
462            cookies,
463        },
464        stream: http_body_util::BodyDataStream::new(body),
465    }
466}
467
468/// Build an empty 200 response for acknowledging completion invocations.
469fn empty_streaming_response()
470-> http::Response<http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>> {
471    use http_body_util::{BodyExt, Full};
472    let body = Full::new(bytes::Bytes::new())
473        .map_err(|e: std::convert::Infallible| match e {})
474        .boxed_unsync();
475    http::Response::builder().status(200).body(body).unwrap()
476}
477
478#[cfg(test)]
479mod streaming_completion_tests {
480    use super::*;
481    use serde_json::json;
482
483    /// Load a test fixture from `src/fixtures/` via `include_str!`.
484    /// Compile-time verified — missing files cause a build error.
485    fn load_fixture(name: &str) -> serde_json::Value {
486        let json_str = match name {
487            "apigw_v1" => include_str!("fixtures/apigw_v1_proxy_event.json"),
488            "apigw_v2" => include_str!("fixtures/apigw_v2_http_api_event.json"),
489            "completion_success" => include_str!("fixtures/streaming_completion_success.json"),
490            "completion_failure" => include_str!("fixtures/streaming_completion_failure.json"),
491            "completion_extra" => include_str!("fixtures/streaming_completion_extra_fields.json"),
492            "completion_api_like" => {
493                include_str!("fixtures/completion_with_api_like_fields.json")
494            }
495            other => panic!("Unknown fixture: {other}"),
496        };
497        serde_json::from_str(json_str).unwrap_or_else(|e| panic!("Bad fixture {name}: {e}"))
498    }
499
500    // ── Fixture tests: API Gateway events → ApiGatewayEvent ──
501
502    #[test]
503    fn test_classify_api_gateway_v1_event() {
504        let payload = load_fixture("apigw_v1");
505        assert!(
506            matches!(
507                classify_runtime_event(payload),
508                RuntimeEventClassification::ApiGatewayEvent(_)
509            ),
510            "API Gateway v1 proxy event must classify as ApiGatewayEvent"
511        );
512    }
513
514    #[test]
515    fn test_classify_api_gateway_v2_event() {
516        let payload = load_fixture("apigw_v2");
517        assert!(
518            matches!(
519                classify_runtime_event(payload),
520                RuntimeEventClassification::ApiGatewayEvent(_)
521            ),
522            "API Gateway v2 HTTP API event must classify as ApiGatewayEvent"
523        );
524    }
525
526    // ── Fixture tests: Streaming completion → StreamingCompletion ──
527
528    #[test]
529    fn test_classify_streaming_completion() {
530        let payload = load_fixture("completion_success");
531        assert!(matches!(
532            classify_runtime_event(payload),
533            RuntimeEventClassification::StreamingCompletion
534        ));
535    }
536
537    #[test]
538    fn test_classify_completion_failure_status() {
539        let payload = load_fixture("completion_failure");
540        assert!(matches!(
541            classify_runtime_event(payload),
542            RuntimeEventClassification::StreamingCompletion
543        ));
544    }
545
546    #[test]
547    fn test_classify_completion_extra_fields() {
548        let payload = load_fixture("completion_extra");
549        assert!(matches!(
550            classify_runtime_event(payload),
551            RuntimeEventClassification::StreamingCompletion
552        ));
553    }
554
555    // ── R5: Precedence edge case ──
556
557    #[test]
558    fn test_classify_completion_with_api_like_fields() {
559        // Intentional: prefer false-positive ack over retries.
560        // A payload with invokeCompletionStatus + partial API Gateway fields
561        // is classified as StreamingCompletion (not UnrecognizedEvent),
562        // because invokeCompletionStatus is the discriminator.
563        //
564        // NOTE: This fixture is intentionally NOT a valid API Gateway event.
565        // Do not "fix" it into one — that would change the expected classification.
566        let payload = load_fixture("completion_api_like");
567        assert!(matches!(
568            classify_runtime_event(payload),
569            RuntimeEventClassification::StreamingCompletion
570        ));
571    }
572
573    // ── Inline tests: Unrecognized payloads → UnrecognizedEvent ──
574
575    #[test]
576    fn test_classify_empty_object() {
577        assert!(matches!(
578            classify_runtime_event(json!({})),
579            RuntimeEventClassification::UnrecognizedEvent
580        ));
581    }
582
583    #[test]
584    fn test_classify_random_object() {
585        assert!(matches!(
586            classify_runtime_event(json!({"foo": "bar", "baz": 123})),
587            RuntimeEventClassification::UnrecognizedEvent
588        ));
589    }
590
591    #[test]
592    fn test_classify_null_payload() {
593        assert!(matches!(
594            classify_runtime_event(json!(null)),
595            RuntimeEventClassification::UnrecognizedEvent
596        ));
597    }
598
599    #[test]
600    fn test_classify_string_payload() {
601        assert!(matches!(
602            classify_runtime_event(json!("hello")),
603            RuntimeEventClassification::UnrecognizedEvent
604        ));
605    }
606
607    #[test]
608    fn test_classify_array_payload() {
609        assert!(matches!(
610            classify_runtime_event(json!([1, 2, 3])),
611            RuntimeEventClassification::UnrecognizedEvent
612        ));
613    }
614
615    #[test]
616    fn test_classify_nested_invoke_status() {
617        // invokeCompletionStatus must be at top level to match
618        let payload = json!({
619            "data": {"invokeCompletionStatus": "Success"}
620        });
621        assert!(matches!(
622            classify_runtime_event(payload),
623            RuntimeEventClassification::UnrecognizedEvent
624        ));
625    }
626
627    // ── Property-style tests ──
628
629    #[test]
630    fn test_classify_never_panics_on_arbitrary_json() {
631        // R4: Only assert no panics — no brittle !ApiGatewayEvent assertion.
632        let payloads = vec![
633            json!(null),
634            json!(true),
635            json!(false),
636            json!(42),
637            json!(-1.5),
638            json!(""),
639            json!("some string"),
640            json!([]),
641            json!([1, "two", null, false]),
642            json!({}),
643            json!({"a": 1}),
644            json!({"requestContext": null}),
645            json!({"requestContext": "not-an-object"}),
646            json!({"httpMethod": "POST"}),
647            json!({"version": "2.0"}),
648            json!({"version": "2.0", "routeKey": "GET /"}),
649            json!({"resource": "/", "httpMethod": "GET"}),
650            json!({"deeply": {"nested": {"invokeCompletionStatus": "Success"}}}),
651            // Large payload
652            serde_json::Value::Object((0..100).map(|i| (format!("key_{i}"), json!(i))).collect()),
653        ];
654
655        for payload in payloads {
656            let _result = classify_runtime_event(payload);
657        }
658    }
659
660    #[test]
661    fn test_classify_invoke_completion_status_always_wins() {
662        // Any object with top-level invokeCompletionStatus that doesn't parse as API GW
663        // should be classified as StreamingCompletion
664        let payloads = vec![
665            json!({"invokeCompletionStatus": "Success"}),
666            json!({"invokeCompletionStatus": "Failure"}),
667            json!({"invokeCompletionStatus": "Unknown"}),
668            json!({"invokeCompletionStatus": 42}),
669            json!({"invokeCompletionStatus": null}),
670            json!({"invokeCompletionStatus": "Success", "requestId": "abc-123"}),
671            json!({"invokeCompletionStatus": "Success", "extra": "field", "nested": {"a": 1}}),
672        ];
673
674        for payload in payloads {
675            let result = classify_runtime_event(payload.clone());
676            assert!(
677                matches!(result, RuntimeEventClassification::StreamingCompletion),
678                "Payload with top-level invokeCompletionStatus must be StreamingCompletion: {payload}"
679            );
680        }
681    }
682
683    // ── event_log_level contract tests (R1) ──
684
685    #[test]
686    fn test_unrecognized_logs_at_warn_level() {
687        assert_eq!(
688            event_log_level("unrecognized_lambda_payload"),
689            Some(tracing::Level::WARN)
690        );
691    }
692
693    #[test]
694    fn test_completion_logs_at_debug_level() {
695        assert_eq!(
696            event_log_level("streaming_completion"),
697            Some(tracing::Level::DEBUG)
698        );
699    }
700
701    #[test]
702    fn test_api_gateway_has_no_extra_logging() {
703        assert_eq!(event_log_level("api_gateway_event"), None);
704    }
705
706    // ── handle_runtime_payload action-path tests (R3) ──
707
708    #[tokio::test]
709    async fn test_handle_completion_does_not_dispatch() {
710        let dispatched = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
711        let dispatched_clone = dispatched.clone();
712
713        let result = handle_runtime_payload(
714            load_fixture("completion_success"),
715            lambda_runtime::Context::default(),
716            |_req| {
717                let d = dispatched_clone.clone();
718                async move {
719                    d.store(true, std::sync::atomic::Ordering::SeqCst);
720                    Ok(empty_streaming_response())
721                }
722            },
723        )
724        .await
725        .expect("handle should succeed");
726
727        assert!(
728            !dispatched.load(std::sync::atomic::Ordering::SeqCst),
729            "Completion events must not dispatch to handler"
730        );
731        assert_eq!(result.event_type, "streaming_completion");
732        assert_eq!(result.response.metadata_prelude.status_code, 200);
733    }
734
735    #[tokio::test]
736    async fn test_handle_unrecognized_does_not_dispatch() {
737        let dispatched = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
738        let dispatched_clone = dispatched.clone();
739
740        let result = handle_runtime_payload(
741            json!({"foo": "bar"}),
742            lambda_runtime::Context::default(),
743            |_req| {
744                let d = dispatched_clone.clone();
745                async move {
746                    d.store(true, std::sync::atomic::Ordering::SeqCst);
747                    Ok(empty_streaming_response())
748                }
749            },
750        )
751        .await
752        .expect("handle should succeed");
753
754        assert!(
755            !dispatched.load(std::sync::atomic::Ordering::SeqCst),
756            "Unrecognized events must not dispatch to handler"
757        );
758        assert_eq!(result.event_type, "unrecognized_lambda_payload");
759        assert_eq!(result.response.metadata_prelude.status_code, 200);
760    }
761
762    #[tokio::test]
763    async fn test_handle_apigw_v1_dispatches() {
764        let dispatched = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
765        let dispatched_clone = dispatched.clone();
766
767        let result = handle_runtime_payload(
768            load_fixture("apigw_v1"),
769            lambda_runtime::Context::default(),
770            |_req| {
771                let d = dispatched_clone.clone();
772                async move {
773                    d.store(true, std::sync::atomic::Ordering::SeqCst);
774                    Ok(empty_streaming_response())
775                }
776            },
777        )
778        .await
779        .expect("handle should succeed");
780
781        assert!(
782            dispatched.load(std::sync::atomic::Ordering::SeqCst),
783            "API Gateway v1 events must dispatch to handler"
784        );
785        assert_eq!(result.event_type, "api_gateway_event");
786    }
787
788    #[tokio::test]
789    async fn test_handle_apigw_v2_dispatches() {
790        let dispatched = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
791        let dispatched_clone = dispatched.clone();
792
793        let result = handle_runtime_payload(
794            load_fixture("apigw_v2"),
795            lambda_runtime::Context::default(),
796            |_req| {
797                let d = dispatched_clone.clone();
798                async move {
799                    d.store(true, std::sync::atomic::Ordering::SeqCst);
800                    Ok(empty_streaming_response())
801                }
802            },
803        )
804        .await
805        .expect("handle should succeed");
806
807        assert!(
808            dispatched.load(std::sync::atomic::Ordering::SeqCst),
809            "API Gateway v2 events must dispatch to handler"
810        );
811        assert_eq!(result.event_type, "api_gateway_event");
812    }
813
814    #[tokio::test]
815    async fn test_handle_unrecognized_surfaces_distinct_event_type() {
816        let result = handle_runtime_payload(
817            json!({"unknown": true}),
818            lambda_runtime::Context::default(),
819            |_req| async { Ok(empty_streaming_response()) },
820        )
821        .await
822        .expect("handle should succeed");
823
824        assert_eq!(result.event_type, "unrecognized_lambda_payload");
825    }
826
827    // ── Existing response conversion tests ──
828
829    #[test]
830    fn test_empty_streaming_response() {
831        let resp = empty_streaming_response();
832        assert_eq!(resp.status(), 200);
833    }
834
835    #[test]
836    fn test_into_lambda_stream_response_preserves_metadata() {
837        use http_body_util::{BodyExt, Full};
838
839        let response = http::Response::builder()
840            .status(401)
841            .header("WWW-Authenticate", "Bearer realm=\"mcp\"")
842            .header("X-Custom", "test")
843            .body(
844                Full::new(bytes::Bytes::from("Unauthorized"))
845                    .map_err(|e: std::convert::Infallible| match e {})
846                    .boxed_unsync(),
847            )
848            .unwrap();
849
850        let stream_resp = into_lambda_stream_response(response);
851        assert_eq!(stream_resp.metadata_prelude.status_code, 401);
852        assert_eq!(
853            stream_resp
854                .metadata_prelude
855                .headers
856                .get("WWW-Authenticate")
857                .unwrap(),
858            "Bearer realm=\"mcp\""
859        );
860        assert_eq!(
861            stream_resp
862                .metadata_prelude
863                .headers
864                .get("X-Custom")
865                .unwrap(),
866            "test"
867        );
868    }
869
870    #[test]
871    fn test_into_lambda_stream_response_extracts_cookies() {
872        use http_body_util::{BodyExt, Full};
873
874        let response = http::Response::builder()
875            .status(200)
876            .header("Set-Cookie", "session=abc; Path=/")
877            .header("Set-Cookie", "theme=dark")
878            .body(
879                Full::new(bytes::Bytes::new())
880                    .map_err(|e: std::convert::Infallible| match e {})
881                    .boxed_unsync(),
882            )
883            .unwrap();
884
885        let stream_resp = into_lambda_stream_response(response);
886        assert_eq!(stream_resp.metadata_prelude.cookies.len(), 2);
887        assert!(
888            stream_resp
889                .metadata_prelude
890                .cookies
891                .contains(&"session=abc; Path=/".to_string())
892        );
893        assert!(
894            stream_resp
895                .metadata_prelude
896                .cookies
897                .contains(&"theme=dark".to_string())
898        );
899        assert!(
900            stream_resp
901                .metadata_prelude
902                .headers
903                .get("Set-Cookie")
904                .is_none()
905        );
906    }
907
908    // ── run_streaming_with dispatch tests ──
909
910    #[tokio::test]
911    async fn test_run_streaming_with_dispatches_apigw_events() {
912        let dispatched = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
913        let dispatched_clone = dispatched.clone();
914
915        let dispatch = move |_req: lambda_http::Request| {
916            let d = dispatched_clone.clone();
917            async move {
918                d.store(true, std::sync::atomic::Ordering::SeqCst);
919                Ok(empty_streaming_response())
920            }
921        };
922
923        let result = handle_runtime_payload(
924            load_fixture("apigw_v1"),
925            lambda_runtime::Context::default(),
926            dispatch,
927        )
928        .await
929        .expect("handle should succeed");
930
931        assert!(
932            dispatched.load(std::sync::atomic::Ordering::SeqCst),
933            "run_streaming_with dispatch must be called for API Gateway events"
934        );
935        assert_eq!(result.event_type, "api_gateway_event");
936    }
937
938    // ── Empty-body Lambda streaming envelope contract ──
939    //
940    // `run_streaming_with` accepts any `Response<B>` where `B: Body + Unpin
941    // + Send + 'static`. Consumers (notably custom OPTIONS short-circuits
942    // for `.well-known` routes) construct empty-body responses two ways:
943    //
944    //   - `Empty::new()`              — zero data frames, immediate end-of-stream
945    //   - `Full::new(Bytes::new())`   — at most one zero-byte data frame
946    //
947    // `BodyDataStream` adapts a `Body` into a `Stream<Item = Bytes>`,
948    // yielding only Data frames (Trailers frames are dropped). When the
949    // underlying body produces zero data frames, the resulting Lambda
950    // `StreamResponse` carries a stream that completes without ever
951    // yielding an item — the Runtime API multipart streaming framing
952    // never sees a body chunk, the request body terminates without
953    // matching the expected framing, and the Runtime API connection
954    // closes with `IncompleteMessage`. Production symptom: APIGW 502
955    // after the function timeout while the dispatch closure already
956    // returned `Ok(resp)`.
957    //
958    // Contract these tests enforce (ADR-026): the framework MUST produce
959    // a Lambda streaming response whose `BodyDataStream` yields at least
960    // one data frame for any input `Response<B>`, including bodies that
961    // would otherwise produce zero data frames natively.
962
963    async fn drain_stream<S, B>(
964        stream: lambda_runtime::StreamResponse<S>,
965    ) -> (u16, Vec<bytes::Bytes>)
966    where
967        S: futures::Stream<Item = std::result::Result<bytes::Bytes, B>> + Unpin,
968    {
969        use futures::StreamExt;
970        let status = stream.metadata_prelude.status_code.as_u16();
971        let mut frames = Vec::new();
972        let mut s = stream.stream;
973        while let Some(item) = s.next().await {
974            frames.push(item.unwrap_or_else(|_| bytes::Bytes::new()));
975        }
976        (status, frames)
977    }
978
979    /// Adapter-level: an `Empty::new()` body must still produce at least
980    /// one data frame after passing through `into_lambda_stream_response`.
981    /// This is the smallest possible repro of the production failure.
982    #[tokio::test]
983    async fn test_into_lambda_stream_response_empty_body_yields_data_frame() {
984        use http_body_util::{BodyExt, Empty};
985
986        let body = Empty::<bytes::Bytes>::new()
987            .map_err(|_: std::convert::Infallible| -> hyper::Error { unreachable!() })
988            .boxed_unsync();
989        let response = http::Response::builder()
990            .status(204)
991            .header("Access-Control-Allow-Origin", "*")
992            .header("Access-Control-Allow-Methods", "GET, OPTIONS")
993            .body(body)
994            .unwrap();
995
996        let stream = into_lambda_stream_response(response);
997        let (status, frames) = drain_stream(stream).await;
998
999        assert_eq!(status, 204);
1000        assert!(
1001            !frames.is_empty(),
1002            "empty-body streaming response must yield at least one data frame \
1003             (got zero frames — Runtime API will see IncompleteMessage); \
1004             frames: {:?}",
1005            frames,
1006        );
1007    }
1008
1009    /// Service-fn-level: mirror sd-mcp's exact dispatch shape — custom
1010    /// OPTIONS short-circuit returning `Response<UnsyncBoxBody>` with
1011    /// `Empty::new()` body — through `handle_runtime_payload` (which is
1012    /// what `run_streaming_with` wraps around the dispatch closure).
1013    /// Production failure repros here if (and only if) the adapter
1014    /// layer doesn't guarantee at least one data frame.
1015    #[tokio::test]
1016    async fn test_run_streaming_with_empty_body_dispatch_yields_data_frame() {
1017        let dispatch = |_req: lambda_http::Request| async move {
1018            use http_body_util::{BodyExt, Empty};
1019            let body = Empty::<bytes::Bytes>::new()
1020                .map_err(|_: std::convert::Infallible| -> hyper::Error { unreachable!() })
1021                .boxed_unsync();
1022            let resp = http::Response::builder()
1023                .status(204)
1024                .header("Access-Control-Allow-Origin", "https://client.example.test")
1025                .header("Access-Control-Allow-Methods", "GET, OPTIONS")
1026                .header("Access-Control-Allow-Headers", "Content-Type")
1027                .body(body)
1028                .unwrap();
1029            Ok::<_, lambda_http::Error>(resp)
1030        };
1031
1032        let result = handle_runtime_payload(
1033            load_fixture("apigw_v2"),
1034            lambda_runtime::Context::default(),
1035            dispatch,
1036        )
1037        .await
1038        .expect("handle_runtime_payload must succeed");
1039
1040        assert_eq!(result.event_type, "api_gateway_event");
1041        let (status, frames) = drain_stream(result.response).await;
1042        assert_eq!(status, 204);
1043        assert!(
1044            !frames.is_empty(),
1045            "run_streaming_with dispatch returning Empty::new() body must yield \
1046             at least one data frame through the envelope; got zero — this is \
1047             the production .well-known OPTIONS 502 / IncompleteMessage path; \
1048             frames: {:?}",
1049            frames,
1050        );
1051    }
1052
1053    /// Negative control: a non-empty body must of course pass through.
1054    /// Catches a faulty fix that suppresses all data frames.
1055    #[tokio::test]
1056    async fn test_into_lambda_stream_response_non_empty_body_preserves_bytes() {
1057        use http_body_util::{BodyExt, Full};
1058
1059        let payload = bytes::Bytes::from_static(b"hello");
1060        let body = Full::new(payload.clone())
1061            .map_err(|_: std::convert::Infallible| -> hyper::Error { unreachable!() })
1062            .boxed_unsync();
1063        let response = http::Response::builder().status(200).body(body).unwrap();
1064
1065        let stream = into_lambda_stream_response(response);
1066        let (status, frames) = drain_stream(stream).await;
1067        assert_eq!(status, 200);
1068        let joined: Vec<u8> = frames.iter().flat_map(|f| f.iter().copied()).collect();
1069        assert_eq!(
1070            joined,
1071            payload.to_vec(),
1072            "non-empty body must round-trip its bytes through the envelope",
1073        );
1074    }
1075
1076    #[tokio::test]
1077    async fn test_run_streaming_with_acks_completion_without_dispatch() {
1078        let dispatched = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
1079        let dispatched_clone = dispatched.clone();
1080
1081        let dispatch = move |_req: lambda_http::Request| {
1082            let d = dispatched_clone.clone();
1083            async move {
1084                d.store(true, std::sync::atomic::Ordering::SeqCst);
1085                Ok(empty_streaming_response())
1086            }
1087        };
1088
1089        let result = handle_runtime_payload(
1090            load_fixture("completion_success"),
1091            lambda_runtime::Context::default(),
1092            dispatch,
1093        )
1094        .await
1095        .expect("handle should succeed");
1096
1097        assert!(
1098            !dispatched.load(std::sync::atomic::Ordering::SeqCst),
1099            "run_streaming_with dispatch must NOT be called for completion events"
1100        );
1101        assert_eq!(result.event_type, "streaming_completion");
1102        assert_eq!(result.response.metadata_prelude.status_code, 200);
1103    }
1104}