Skip to main content

turul_a2a_aws_lambda/
lib.rs

1//! AWS Lambda adapter for turul-a2a.
2//!
3//! # Choosing a runner
4//!
5//! A Lambda binary receives raw event JSON; the adapter owns
6//! classifying that JSON and dispatching to the right A2A handler
7//! method. Adopters pick the runner whose name matches the Lambda's
8//! AWS trigger topology and end `main.rs` with one `.await?` call —
9//! `lambda_runtime` / `lambda_http` / event parsing never appear in
10//! adopter code.
11//!
12//! | Lambda triggers                   | Runner                                         |
13//! |-----------------------------------|------------------------------------------------|
14//! | HTTP only (Function URL / APIGW)  | [`LambdaA2aHandler::run_http_only`]            |
15//! | SQS only (event source mapping)   | `LambdaA2aHandler::run_sqs_only` *(`sqs`)*   |
16//! | HTTP **and** SQS on one function  | `LambdaA2aHandler::run_http_and_sqs` *(`sqs`)* |
17//! | not sure / just run it            | [`LambdaA2aHandler::run`]                      |
18//!
19//! [`LambdaA2aHandler::run`] is the default "it just works" entry
20//! point: without the `sqs` feature it aliases `run_http_only`; with
21//! `sqs` it aliases `run_http_and_sqs` (permissive, handles either
22//! shape). The explicit runners are **strict** — a non-matching
23//! event shape fails loudly, which is what hardened deployments
24//! and tests want.
25//!
26//! The one-call shortcut is [`LambdaA2aServerBuilder::run`] — it
27//! build-then-runs in a single fluent call. Use `.build()` + a
28//! handler method when the handler is needed for unit tests or
29//! custom middleware.
30//!
31//! ## Composing with non-framework triggers
32//!
33//! Lambdas wired to HTTP **plus** a third trigger the framework does
34//! not know about (EventBridge scheduler, DynamoDB stream, custom
35//! invoke) can't use a named runner directly — the Unknown branch
36//! is adopter-owned. Reach for the public classifier + the HTTP
37//! envelope primitive instead:
38//!
39//! ```ignore
40//! lambda_runtime::run(lambda_runtime::service_fn(
41//!     move |event: lambda_runtime::LambdaEvent<serde_json::Value>| {
42//!         let handler = handler.clone();
43//!         async move {
44//!             let (value, _ctx) = event.into_parts();
45//!             match turul_a2a_aws_lambda::classify_event(&value) {
46//!                 turul_a2a_aws_lambda::LambdaEvent::Http => {
47//!                     handler.handle_http_event_value(value).await
48//!                 }
49//!                 turul_a2a_aws_lambda::LambdaEvent::Sqs => {
50//!                     let sqs = serde_json::from_value(value)?;
51//!                     Ok(serde_json::to_value(handler.handle_sqs(sqs).await)?)
52//!                 }
53//!                 turul_a2a_aws_lambda::LambdaEvent::Unknown => {
54//!                     // Adopter-owned: run your scheduled sweep,
55//!                     // DynamoDB stream consumer, etc.
56//!                     run_adopter_specific_path(value).await
57//!                 }
58//!             }
59//!         }
60//!     }
61//! )).await
62//! ```
63//!
64//! [`LambdaA2aHandler::handle_http_event_value`] is available without
65//! the `sqs` feature; `handle_sqs` is gated on `sqs`. `classify_event`
66//! + `LambdaEvent` are always available.
67//!
68//! ## API Gateway path prefix
69//!
70//! When Lambda sits behind a REST API Gateway with `AWS_PROXY`
71//! integration, the event carries the full stage + resource tree in
72//! the path (e.g. `/stage/agent/message:send`),
73//! which the root-rooted A2A router would 404. Configure the prefix
74//! to strip via [`LambdaA2aServerBuilder::strip_path_prefix`]:
75//!
76//! ```ignore
77//! LambdaA2aServerBuilder::new()
78//!     .executor(my_executor)
79//!     .storage(my_storage)
80//!     .strip_path_prefix("/stage/agent")  // API GW prefix
81//!     .run()
82//!     .await
83//! ```
84//!
85//! The strip applies to `run_http_only` and `run_http_and_sqs`.
86//! SQS events are unaffected. Non-matching paths pass through
87//! unchanged (router still 404s on genuinely unknown paths —
88//! that's the correct failure mode).
89//!
90//! ```ignore
91//! // HTTP-only Lambda (request Lambda in two-Lambda topology):
92//! LambdaA2aServerBuilder::new()
93//!     .executor(my_executor)
94//!     .storage(my_storage)
95//!     .with_sqs_return_immediately(queue_url, sqs)  // enqueues durable jobs
96//!     .build()?
97//!     .run_http_only()
98//!     .await
99//!
100//! // Single-function HTTP+SQS demo:
101//! LambdaA2aServerBuilder::new()
102//!     .executor(my_executor)
103//!     .storage(my_storage)
104//!     .with_sqs_return_immediately(queue_url, sqs)
105//!     .run()            // builder shortcut; dispatches HTTP + SQS
106//!     .await
107//!
108//! // Pure SQS worker — no with_sqs_return_immediately(...) needed
109//! // (it consumes the queue, never enqueues):
110//! LambdaA2aServerBuilder::new()
111//!     .executor(my_executor)
112//!     .storage(my_storage)
113//!     .build()?
114//!     .run_sqs_only()
115//!     .await
116//! ```
117//!
118//! # Adapter internals
119//!
120//! Thin wrapper: converts Lambda events to axum requests, delegates to
121//! the same Router, converts responses back. Adapter contract:
122//! - Authorizer context mapped via synthetic headers with anti-spoofing.
123//! - SSE responses are buffered: task executes, events are collected,
124//!   returned as one response.
125//! - `POST /message:stream` executes the task within the Lambda
126//!   invocation and returns all events.
127//! - `GET /tasks/{id}:subscribe` is for tasks that are **not** in a
128//!   terminal state. Within one invocation it emits the initial `Task`
129//!   snapshot, delivers events committed after it, and closes when the
130//!   task reaches a terminal state or INPUT_REQUIRED. Subscribing to an
131//!   already-terminal task returns `UnsupportedOperationError` per
132//!   A2A v1.0 §3.1.6. For retrieving a terminal task's final state
133//!   use `GetTask`.
134//!
135//! # Push-notification delivery — external triggers are mandatory
136//!
137//! Push delivery on Lambda is architecturally different from the
138//! binary server. The request Lambda installed via
139//! [`LambdaA2aServerBuilder`] still constructs a `PushDispatcher`
140//! when `push_delivery_store` is wired, but any `tokio::spawn`
141//! continuation it emits post-return is **opportunistic only** — the
142//! Lambda execution environment may be frozen indefinitely between
143//! invocations, so nothing can depend on that continuation completing.
144//!
145//! Correctness for push delivery on Lambda is carried by:
146//!
147//! 1. The atomic pending-dispatch marker written inside the request
148//!    Lambda's commit transaction (opt in via
149//!    `StorageImpl::with_push_dispatch_enabled(true)`).
150//! 2. [`LambdaStreamRecoveryHandler`] — DynamoDB Streams trigger on
151//!    `a2a_push_pending_dispatches`. DynamoDB backends only.
152//! 3. [`LambdaScheduledRecoveryHandler`] — EventBridge Scheduler
153//!    backstop. **Required for all backends**; it is the sole
154//!    recovery path for SQLite / PostgreSQL / in-memory deployments.
155//!
156//! Without at least the scheduled worker, push delivery on Lambda is
157//! not durable — a marker written on a cold invocation may never be
158//! consumed. The example wiring lives in
159//! `examples/lambda-stream-worker` and
160//! `examples/lambda-scheduled-worker`.
161//!
162//! Lambda streaming is request-scoped (not persistent SSE connections). The durable
163//! event store ensures events survive across invocations. A client that loses
164//! its stream subscribes again and receives a fresh `Task` snapshot.
165//!
166//! # Cross-instance cancellation
167//!
168//! Lambda invocations are stateless and short-lived, so the Lambda
169//! adapter does **not** run the persistent cross-instance cancel
170//! poller that `A2aServer::run()` spawns. Cancellation behaviour on
171//! the Lambda adapter:
172//!
173//! - **Marker writes** — `CancelTask` on a Lambda invocation writes
174//!   the cancel marker to the shared backend (DynamoDB / PostgreSQL).
175//!   This works.
176//! - **Propagation to a live executor on the SAME Lambda invocation** —
177//!   works via the same-instance token-trip path in `core_cancel_task`.
178//! - **Propagation to a live executor on a DIFFERENT Lambda invocation
179//!   (warm container)** — **not currently supported**. There is no
180//!   persistent poller to observe markers written by other
181//!   invocations. A subsequent invocation whose handler reads the
182//!   marker directly may act on it, but that is not a substitute for
183//!   the server runtime's live propagation.
184//!
185//! The builder still **requires** an `A2aCancellationSupervisor`
186//! implementation on the same backend so that marker writes reach
187//! the correct backend and a future polling-adapter variant can
188//! consume them. If a deployment passes a non-matching supervisor,
189//! `build()` rejects the configuration.
190
191mod adapter;
192mod auth;
193mod event;
194mod no_streaming;
195mod scheduled_recovery;
196mod stream_recovery;
197
198#[cfg(feature = "sqs")]
199mod durable;
200
201#[cfg(test)]
202mod builder_tests;
203#[cfg(test)]
204mod scheduled_recovery_tests;
205#[cfg(test)]
206mod stream_recovery_tests;
207
208pub use adapter::{axum_to_lambda_response, lambda_to_axum_request};
209pub use auth::{AuthorizerMapping, LambdaAuthorizerMiddleware};
210pub use event::{LambdaEvent, classify_event};
211pub use no_streaming::NoStreamingLayer;
212pub use scheduled_recovery::{
213    LambdaScheduledRecoveryConfig, LambdaScheduledRecoveryHandler, LambdaScheduledRecoveryResponse,
214    handle_maintenance_trigger,
215};
216pub use stream_recovery::LambdaStreamRecoveryHandler;
217
218#[cfg(feature = "sqs")]
219pub use durable::{SqsDurableExecutorQueue, drive_sqs_batch};
220
221use std::sync::Arc;
222
223use turul_a2a::error::A2aError;
224use turul_a2a::executor::AgentExecutor;
225use turul_a2a::middleware::{A2aMiddleware, MiddlewareStack};
226use turul_a2a::push::A2aPushDeliveryStore;
227use turul_a2a::router::{AppState, build_router};
228use turul_a2a::server::RuntimeConfig;
229use turul_a2a::storage::{
230    A2aAtomicStore, A2aCancellationSupervisor, A2aEventStore, A2aPushNotificationStorage,
231    A2aTaskStorage,
232};
233use turul_a2a::streaming::TaskEventBroker;
234
235/// Builder for Lambda A2A handler.
236///
237/// Use `.storage(my_storage)` to supply a single backend implementing all storage traits.
238/// All four storage traits must share the same backend so streaming, push delivery,
239/// and cross-instance cancellation coordinate correctly.
240///
241/// For push-notification delivery, wire `.storage()` with a backend that implements
242/// [`A2aPushDeliveryStore`] (all four first-party backends do). The atomic store
243/// must also opt in via `with_push_dispatch_enabled(true)`; the
244/// builder rejects configurations where the flag and the delivery store disagree.
245pub struct LambdaA2aServerBuilder {
246    executor: Option<Arc<dyn AgentExecutor>>,
247    task_storage: Option<Arc<dyn A2aTaskStorage>>,
248    push_storage: Option<Arc<dyn A2aPushNotificationStorage>>,
249    event_store: Option<Arc<dyn A2aEventStore>>,
250    atomic_store: Option<Arc<dyn A2aAtomicStore>>,
251    cancellation_supervisor: Option<Arc<dyn A2aCancellationSupervisor>>,
252    push_delivery_store: Option<Arc<dyn A2aPushDeliveryStore>>,
253    middleware: Vec<Arc<dyn A2aMiddleware>>,
254    runtime_config: RuntimeConfig,
255    durable_executor_queue: Option<Arc<dyn turul_a2a::durable_executor::DurableExecutorQueue>>,
256    path_prefix: Option<String>,
257    outbound_url_validator: Option<turul_a2a::push::ssrf::OutboundUrlValidator>,
258}
259
260impl LambdaA2aServerBuilder {
261    pub fn new() -> Self {
262        Self {
263            executor: None,
264            task_storage: None,
265            push_storage: None,
266            event_store: None,
267            atomic_store: None,
268            cancellation_supervisor: None,
269            push_delivery_store: None,
270            middleware: vec![],
271            runtime_config: RuntimeConfig::default(),
272            durable_executor_queue: None,
273            path_prefix: None,
274            outbound_url_validator: None,
275        }
276    }
277
278    /// Configure an HTTP path prefix to strip before the A2A router
279    /// sees the request. Needed when Lambda sits behind an API Gateway
280    /// whose stage + resource tree is forwarded into the function
281    /// (e.g. REST API `AWS_PROXY` integrations surface
282    /// `/stage/agent/message:send` where the A2A
283    /// router expects `/message:send`).
284    ///
285    /// The prefix must start with `/` and must not end with `/`
286    /// (unless it is exactly `/`, which is a no-op). Segment boundaries
287    /// are respected — a prefix of `/dev` will not strip `/devs/...`.
288    /// Paths that do not start with the prefix pass through unchanged
289    /// (the router will 404 on a genuinely unknown path, which is the
290    /// correct failure mode). Applies to both
291    /// [`LambdaA2aHandler::run_http_only`] and
292    /// `LambdaA2aHandler::run_http_and_sqs`. SQS events are
293    /// unaffected.
294    pub fn strip_path_prefix(mut self, prefix: impl Into<String>) -> Self {
295        self.path_prefix = Some(prefix.into());
296        self
297    }
298
299    /// Wire a durable executor queue. When set, the
300    /// `return_immediately = true` path in `core_send_message` enqueues
301    /// a [`turul_a2a::durable_executor::QueuedExecutorJob`] on this
302    /// queue instead of spawning the executor locally, and the
303    /// capability flag `RuntimeConfig::supports_return_immediately` is
304    /// turned **back on** as an implementation detail — the capability
305    /// cannot be claimed without supplying the queue.
306    ///
307    /// Adopters should typically reach for the SQS-specific helper
308    /// `Self::with_sqs_return_immediately` instead; this generic
309    /// method exists so adopters can inject in-memory fakes for tests
310    /// or provide non-SQS transports (Kinesis, Step Functions task
311    /// token, self-invoke, etc.).
312    pub fn with_durable_executor(
313        mut self,
314        queue: Arc<dyn turul_a2a::durable_executor::DurableExecutorQueue>,
315    ) -> Self {
316        self.durable_executor_queue = Some(queue);
317        self.runtime_config.supports_return_immediately = true;
318        self
319    }
320
321    /// Wire the AWS SQS durable executor queue. Ergonomic
322    /// helper over [`Self::with_durable_executor`] that constructs a
323    /// [`SqsDurableExecutorQueue`] from a `queue_url` + shared SQS
324    /// client. Requires the `sqs` feature.
325    #[cfg(feature = "sqs")]
326    pub fn with_sqs_return_immediately(
327        self,
328        queue_url: impl Into<String>,
329        sqs_client: Arc<aws_sdk_sqs::Client>,
330    ) -> Self {
331        let queue = Arc::new(SqsDurableExecutorQueue::new(queue_url, sqs_client));
332        self.with_durable_executor(queue)
333    }
334
335    pub fn executor(mut self, exec: impl AgentExecutor + 'static) -> Self {
336        self.executor = Some(Arc::new(exec));
337        self
338    }
339
340    /// Set all storage from a single backend instance.
341    ///
342    /// This is the **preferred** method — a single struct implementing all
343    /// storage traits guarantees the same-backend requirement
344    /// AND ensures the cancellation supervisor reads the same backend
345    /// the cancel marker is written to. A mismatch here (e.g., DynamoDB
346    /// task storage + in-memory cancellation supervisor) silently
347    /// breaks cross-instance cancellation.
348    ///
349    /// errata: `.storage()` wires storage traits only. It does
350    /// **not** auto-register the storage as a push-delivery store, even if
351    /// the backend happens to implement [`A2aPushDeliveryStore`]. To opt in
352    /// to push delivery, call [`Self::push_delivery_store`] explicitly and
353    /// call `with_push_dispatch_enabled(true)` on the storage instance
354    /// before passing it here. Non-push deployments need neither.
355    pub fn storage<S>(mut self, storage: S) -> Self
356    where
357        S: A2aTaskStorage
358            + A2aPushNotificationStorage
359            + A2aEventStore
360            + A2aAtomicStore
361            + A2aCancellationSupervisor
362            + Clone
363            + 'static,
364    {
365        self.task_storage = Some(Arc::new(storage.clone()));
366        self.push_storage = Some(Arc::new(storage.clone()));
367        self.event_store = Some(Arc::new(storage.clone()));
368        self.atomic_store = Some(Arc::new(storage.clone()));
369        self.cancellation_supervisor = Some(Arc::new(storage));
370        self
371    }
372
373    /// Set task storage individually. Must be paired with event_store/atomic_store
374    /// AND `cancellation_supervisor()` for same-backend compliance.
375    pub fn task_storage(mut self, s: impl A2aTaskStorage + 'static) -> Self {
376        self.task_storage = Some(Arc::new(s));
377        self
378    }
379
380    /// Set push notification storage individually.
381    pub fn push_storage(mut self, s: impl A2aPushNotificationStorage + 'static) -> Self {
382        self.push_storage = Some(Arc::new(s));
383        self
384    }
385
386    /// Set event store individually. Must match task_storage backend.
387    pub fn event_store(mut self, s: impl A2aEventStore + 'static) -> Self {
388        self.event_store = Some(Arc::new(s));
389        self
390    }
391
392    /// Set atomic store individually. Must match task_storage backend.
393    pub fn atomic_store(mut self, s: impl A2aAtomicStore + 'static) -> Self {
394        self.atomic_store = Some(Arc::new(s));
395        self
396    }
397
398    /// Set the cancellation supervisor individually. Must match
399    /// task_storage backend — otherwise `:cancel` writes the marker to
400    /// one backend and the supervisor reads from another, silently
401    /// breaking cross-instance cancellation. Prefer `.storage()` for
402    /// unified wiring. Consumed by `core_cancel_task`.
403    pub fn cancellation_supervisor(mut self, s: impl A2aCancellationSupervisor + 'static) -> Self {
404        self.cancellation_supervisor = Some(Arc::new(s));
405        self
406    }
407
408    /// Set the push delivery store individually.
409    ///
410    /// Prefer `.storage()` for unified wiring. The delivery store MUST
411    /// be on the same backend as `.task_storage(...)` — the builder
412    /// rejects mismatches. Passing a delivery store also requires the
413    /// atomic store to have opted in via `with_push_dispatch_enabled(true)`
414    ///.
415    pub fn push_delivery_store(mut self, store: impl A2aPushDeliveryStore + 'static) -> Self {
416        self.push_delivery_store = Some(Arc::new(store));
417        self
418    }
419
420    /// Override the runtime configuration (timeouts, retry budgets,
421    /// push tuning). Defaults to [`RuntimeConfig::default()`]. The
422    /// builder validates push settings (claim expiry vs retry horizon,
423    /// ) when `.push_delivery_store(...)` is wired.
424    pub fn runtime_config(mut self, cfg: RuntimeConfig) -> Self {
425        self.runtime_config = cfg;
426        self
427    }
428
429    pub fn middleware(mut self, mw: Arc<dyn A2aMiddleware>) -> Self {
430        self.middleware.push(mw);
431        self
432    }
433
434    /// Operator allowlist for push webhook URLs.
435    ///
436    /// Called with each webhook URL on every delivery attempt, before
437    /// DNS resolution. Returning `Err` refuses the delivery: nothing is
438    /// resolved or POSTed, and the delivery is recorded as a failed
439    /// delivery with error class `SSRFBlocked`. The reason string is
440    /// not persisted. Runs in addition to the built-in scheme and
441    /// private-address checks, not instead of them. Takes effect only
442    /// when push delivery is wired (`push_delivery_store`).
443    pub fn outbound_url_validator(
444        mut self,
445        validator: turul_a2a::push::ssrf::OutboundUrlValidator,
446    ) -> Self {
447        self.outbound_url_validator = Some(validator);
448        self
449    }
450
451    pub fn build(self) -> Result<LambdaA2aHandler, A2aError> {
452        let executor = self
453            .executor
454            .ok_or(A2aError::Internal("executor is required".into()))?;
455        let task_storage = self.task_storage.ok_or(A2aError::Internal(
456            "task_storage is required for Lambda".into(),
457        ))?;
458        let push_storage = self.push_storage.ok_or(A2aError::Internal(
459            "push_storage is required for Lambda".into(),
460        ))?;
461        let event_store: Arc<dyn A2aEventStore> = self.event_store.ok_or(A2aError::Internal(
462            "event_store is required for Lambda (use .storage() for unified backend)".into(),
463        ))?;
464        let atomic_store: Arc<dyn A2aAtomicStore> = self.atomic_store.ok_or(A2aError::Internal(
465            "atomic_store is required for Lambda (use .storage() for unified backend)".into(),
466        ))?;
467        // the cancellation supervisor MUST be supplied. No
468        // InMemoryA2aStorage fallback: that would silently break
469        // cross-instance cancellation by reading markers from a
470        // different backend than they were written to.
471        let cancellation_supervisor: Arc<dyn A2aCancellationSupervisor> =
472            self.cancellation_supervisor.ok_or(A2aError::Internal(
473                "cancellation_supervisor is required for Lambda. Use .storage() for unified \
474                 backend wiring, or .cancellation_supervisor(...) alongside the individual \
475                 storage setters. Omitting it silently breaks cross-instance cancellation."
476                    .into(),
477            ))?;
478
479        // Same-backend enforcement: all storage traits + the
480        // cancellation supervisor must report the same `backend_name`.
481        let task_backend = task_storage.backend_name();
482        let push_backend = push_storage.backend_name();
483        let event_backend = event_store.backend_name();
484        let atomic_backend = atomic_store.backend_name();
485        let supervisor_backend = cancellation_supervisor.backend_name();
486        if task_backend != push_backend
487            || task_backend != event_backend
488            || task_backend != atomic_backend
489            || task_backend != supervisor_backend
490        {
491            return Err(A2aError::Internal(format!(
492                "Storage backend mismatch: task={task_backend}, push={push_backend}, \
493                 event={event_backend}, atomic={atomic_backend}, \
494                 cancellation_supervisor={supervisor_backend}. \
495                 All storage traits, the cancellation supervisor included, \
496                 must share one backend so cross-instance cancel markers \
497                 are observable. Use .storage() for unified backend."
498            )));
499        }
500
501        let push_delivery_store = self.push_delivery_store;
502
503        // Push-dispatch consistency: the atomic store's
504        // opt-in flag and the presence of `push_delivery_store` MUST
505        // agree. Both-true = push fully wired; both-false = non-push
506        // deployment. Mixed cases are build errors — this mirrors the
507        // main server builder (server/mod.rs).
508        match (
509            push_delivery_store.is_some(),
510            atomic_store.push_dispatch_enabled(),
511        ) {
512            (true, true) | (false, false) => {}
513            (true, false) => {
514                return Err(A2aError::Internal(
515                    "push_delivery_store wired but atomic_store.push_dispatch_enabled() \
516                     is false. Call .with_push_dispatch_enabled(true) on the backend \
517                     storage before passing it to .storage()."
518                        .into(),
519                ));
520            }
521            (false, true) => {
522                return Err(A2aError::Internal(
523                    "atomic_store.push_dispatch_enabled() is true but no \
524                     push_delivery_store is wired. Pending-dispatch markers would be \
525                     written with no consumer, imposing load-bearing infra for no \
526                     benefit. If you need to populate markers for an external \
527                     consumer, open an issue for a distinctly-named opt-in — for now, \
528                     this configuration is rejected."
529                        .into(),
530                ));
531            }
532        }
533
534        // The Lambda execution environment can freeze once the HTTP
535        // response is flushed, so a `tokio::spawn`'d executor
536        // continuation is opportunistic only. Refuse
537        // `SendMessageConfiguration.return_immediately = true` at the
538        // `core_send_message` entry point (UnsupportedOperationError —
539        // an aspect of SendMessage this runtime cannot honour, A2A
540        // v1.0.1 §3.3.2) via this capability flag instead of silently
541        // orphaning the executor.
542        //
543        // if a durable executor queue is wired via
544        // [`Self::with_durable_executor`] (or the SQS helper), the
545        // capability is re-enabled because enqueueing is durable
546        // across the Lambda freeze. The explicit re-enable lives in
547        // the builder method so the capability cannot be claimed
548        // without the mechanism.
549        let mut runtime_config = self.runtime_config;
550        if self.durable_executor_queue.is_none() {
551            runtime_config.supports_return_immediately = false;
552        }
553        // No background loops run on Lambda, so a lease written here
554        // would never be renewed and peers would fail the live task.
555        runtime_config.executor_lease_ttl = None;
556        runtime_config.executor_stall_timeout = None;
557        // The idle-task sweep needs no loop: the scheduled-recovery
558        // handler runs it.
559        if runtime_config
560            .idle_task_timeout
561            .is_some_and(|d| d.is_zero())
562        {
563            return Err(A2aError::Internal(
564                "idle_task_timeout must be greater than zero".into(),
565            ));
566        }
567
568        // Push dispatcher wiring. When
569        // `push_delivery_store` is present, validate its backend matches,
570        // validate the retry-horizon invariant, and construct the
571        // `PushDispatcher`. Under Lambda the dispatcher runs
572        // opportunistically post-return; correctness
573        // comes from the atomic marker + stream/scheduler recovery.
574        let push_dispatcher: Option<Arc<turul_a2a::push::PushDispatcher>> = if let Some(delivery) =
575            push_delivery_store.as_ref()
576        {
577            let push_delivery_backend = delivery.backend_name();
578            if task_backend != push_delivery_backend {
579                return Err(A2aError::Internal(format!(
580                    "Storage backend mismatch: task={task_backend}, \
581                         push_delivery={push_delivery_backend}. \
582                         All storage traits must share the same backend."
583                )));
584            }
585
586            // claim expiry must exceed the worst-case
587            // retry horizon so a live claim is not mis-classified as
588            // stale mid-retry.
589            let retry_horizon = runtime_config
590                .push_backoff_cap
591                .saturating_mul(runtime_config.push_max_attempts as u32);
592            if runtime_config.push_claim_expiry <= retry_horizon {
593                return Err(A2aError::Internal(format!(
594                    "push_claim_expiry ({:?}) must be greater than retry horizon \
595                         (push_max_attempts={} * push_backoff_cap={:?} = {:?}). \
596                         Raise push_claim_expiry or lower push_max_attempts/push_backoff_cap.",
597                    runtime_config.push_claim_expiry,
598                    runtime_config.push_max_attempts,
599                    runtime_config.push_backoff_cap,
600                    retry_horizon
601                )));
602            }
603
604            let delivery_cfg =
605                turul_a2a::push::delivery::PushDeliveryConfig::from_runtime_config(&runtime_config);
606
607            let instance_id = format!("a2a-lambda-{}", uuid::Uuid::now_v7());
608            let worker = turul_a2a::push::delivery::PushDeliveryWorker::new(
609                delivery.clone(),
610                delivery_cfg,
611                self.outbound_url_validator.clone(),
612                instance_id,
613            )
614            .map_err(|e| A2aError::Internal(format!("push worker build failed: {e}")))?;
615
616            Some(Arc::new(turul_a2a::push::PushDispatcher::new(
617                Arc::new(worker),
618                push_storage.clone(),
619                task_storage.clone(),
620            )))
621        } else {
622            None
623        };
624
625        let state = AppState {
626            executor,
627            task_storage,
628            push_storage,
629            event_store,
630            atomic_store,
631            event_broker: TaskEventBroker::new(),
632            middleware_stack: Arc::new(MiddlewareStack::new(self.middleware)),
633            runtime_config,
634            in_flight: Arc::new(turul_a2a::server::in_flight::InFlightRegistry::new()),
635            cancellation_supervisor,
636            push_delivery_store,
637            push_dispatcher,
638            durable_executor_queue: self.durable_executor_queue,
639        };
640
641        let router = build_router(state.clone());
642
643        let path_prefix = match self.path_prefix {
644            None => None,
645            Some(p) if p.is_empty() || p == "/" => None,
646            Some(p) => {
647                if !p.starts_with('/') {
648                    return Err(A2aError::InvalidRequest {
649                        message: format!(
650                            "LambdaA2aServerBuilder::strip_path_prefix: prefix must start with '/'; got {p:?}"
651                        ),
652                    });
653                }
654                if p.ends_with('/') {
655                    return Err(A2aError::InvalidRequest {
656                        message: format!(
657                            "LambdaA2aServerBuilder::strip_path_prefix: prefix must not end with '/' (except '/'); got {p:?}"
658                        ),
659                    });
660                }
661                Some(Arc::from(p))
662            }
663        };
664
665        Ok(LambdaA2aHandler {
666            router,
667            state,
668            path_prefix,
669        })
670    }
671}
672
673impl Default for LambdaA2aServerBuilder {
674    fn default() -> Self {
675        Self::new()
676    }
677}
678
679/// Lambda handler wrapping the axum Router.
680#[derive(Clone)]
681pub struct LambdaA2aHandler {
682    router: axum::Router,
683    /// Held for SQS-event dispatch (`handle_sqs`) and the idle-task
684    /// sweep (`cancel_idle_tasks`).
685    state: AppState,
686    /// Leading path segment to strip from each incoming HTTP request
687    /// before the A2A router sees it. `None` = no stripping. Set via
688    /// [`LambdaA2aServerBuilder::strip_path_prefix`].
689    path_prefix: Option<Arc<str>>,
690}
691
692impl LambdaA2aHandler {
693    pub fn builder() -> LambdaA2aServerBuilder {
694        LambdaA2aServerBuilder::new()
695    }
696
697    /// Cancel tasks left waiting on their client for longer than the
698    /// runtime configuration's `idle_task_timeout`; see
699    /// [`turul_a2a::server::A2aServerBuilder::idle_task_timeout`]. With no
700    /// timeout set this does nothing. Run by
701    /// [`LambdaScheduledRecoveryHandler::with_idle_task_sweep`].
702    pub async fn cancel_idle_tasks(
703        &self,
704    ) -> Result<turul_a2a::server::IdleSweep, turul_a2a::storage::A2aStorageError> {
705        match self.state.runtime_config.idle_task_timeout {
706            Some(idle_for) => {
707                turul_a2a::server::liveness::cancel_idle_tasks(
708                    &self.state,
709                    Some(&self.state.in_flight),
710                    idle_for,
711                )
712                .await
713            }
714            None => Ok(Default::default()),
715        }
716    }
717
718    /// Handle a Lambda HTTP request.
719    pub async fn handle(
720        &self,
721        event: lambda_http::Request,
722    ) -> Result<lambda_http::Response<lambda_http::Body>, lambda_http::Error> {
723        let axum_req = lambda_to_axum_request(event)?;
724        let axum_req = match &self.path_prefix {
725            Some(prefix) => adapter::strip_request_path_prefix(axum_req, prefix),
726            None => axum_req,
727        };
728
729        let axum_resp = tower::ServiceExt::oneshot(self.router.clone(), axum_req)
730            .await
731            .map_err(|e| lambda_http::Error::from(format!("Router error: {e}")))?;
732
733        axum_to_lambda_response(axum_resp).await
734    }
735
736    /// Handle a raw Lambda HTTP event delivered as `serde_json::Value`
737    /// and return a response envelope **matching the inbound event
738    /// shape**:
739    ///
740    /// | Inbound shape detected by `lambda_http::LambdaRequest`   | Response envelope                  |
741    /// |----------------------------------------------------------|------------------------------------|
742    /// | API Gateway v1 (REST) — top-level `httpMethod` + `path`  | `ApiGatewayProxyResponse`          |
743    /// | API Gateway v2 / Function URL — `requestContext.http`    | `ApiGatewayV2httpResponse`         |
744    /// | ALB target group — `requestContext.elb`                  | `AlbTargetGroupResponse`           |
745    /// | WebSocket / unrecognised                                  | `Err(...)`                         |
746    ///
747    /// Mismatching envelopes — most commonly a v2 response sent in
748    /// response to a v1 invocation — cause API Gateway to return
749    /// `502 Bad Gateway` regardless of the body, because the response
750    /// fails the gateway's contract validation. This method ensures
751    /// symmetry: response shape is always picked from the inbound
752    /// shape, never hardcoded.
753    ///
754    /// This is the framework's canonical HTTP envelope conversion
755    /// primitive — useful when a single Lambda function routes among
756    /// HTTP and one or more non-framework triggers (EventBridge,
757    /// DynamoDB streams, custom invoke) and the adopter drives
758    /// [`lambda_runtime::run`] with `serde_json::Value` directly.
759    /// Handles the `LambdaRequest` deserialization, the text-vs-base64
760    /// content-type detection, and the response envelope rebuild.
761    ///
762    /// Body encoding: text for `text/*`, `application/json`,
763    /// `application/xml`, `application/javascript`, any `+json` / `+xml`
764    /// structured-syntax suffix (RFC 6839, e.g. `application/a2a+json`),
765    /// or anything with `charset=`; base64 otherwise.
766    ///
767    /// Path-prefix stripping (see
768    /// [`LambdaA2aServerBuilder::strip_path_prefix`]) is applied
769    /// because dispatch flows through [`Self::handle`]. Adopters who
770    /// only serve HTTP should continue to use [`Self::run_http_only`]
771    /// or [`LambdaA2aServerBuilder::run`]; this entry point exists for
772    /// composition with adopter-owned third triggers.
773    pub async fn handle_http_event_value(
774        &self,
775        value: serde_json::Value,
776    ) -> Result<serde_json::Value, lambda_runtime::Error> {
777        use base64::Engine;
778        use http_body_util::BodyExt;
779
780        let lambda_req: lambda_http::request::LambdaRequest = serde_json::from_value(value)
781            .map_err(|e| lambda_runtime::Error::from(format!("invalid HTTP event: {e}")))?;
782
783        // Capture the inbound shape so we can emit the matching
784        // response envelope. Lost once we convert to `lambda_http::Request`.
785        let shape = match &lambda_req {
786            lambda_http::request::LambdaRequest::ApiGatewayV1(_) => HttpEventShape::ApiGatewayV1,
787            lambda_http::request::LambdaRequest::ApiGatewayV2(_) => HttpEventShape::ApiGatewayV2,
788            lambda_http::request::LambdaRequest::Alb(_) => HttpEventShape::Alb,
789            lambda_http::request::LambdaRequest::WebSocket(_) => {
790                return Err(lambda_runtime::Error::from(
791                    "WebSocket Lambda events are not supported by this adapter",
792                ));
793            }
794            _ => {
795                return Err(lambda_runtime::Error::from(
796                    "unsupported Lambda HTTP event variant",
797                ));
798            }
799        };
800
801        let req: lambda_http::Request = lambda_req.into();
802
803        let resp = self
804            .handle(req)
805            .await
806            .map_err(|e| lambda_runtime::Error::from(format!("handler error: {e}")))?;
807        let (parts, body) = resp.into_parts();
808
809        let bytes = body
810            .collect()
811            .await
812            .map_err(|e| lambda_runtime::Error::from(format!("body collect: {e}")))?
813            .to_bytes();
814
815        let ct = parts
816            .headers
817            .get(http::header::CONTENT_TYPE)
818            .and_then(|v| v.to_str().ok())
819            .unwrap_or("application/octet-stream")
820            .to_ascii_lowercase();
821        let essence = ct.split(';').next().unwrap_or("").trim();
822        let is_text = essence.starts_with("text/")
823            || matches!(
824                essence,
825                "application/json" | "application/xml" | "application/javascript"
826            )
827            || essence.ends_with("+json")
828            || essence.ends_with("+xml")
829            || ct.contains("charset=");
830
831        let (body_str, is_base64) = if bytes.is_empty() {
832            (None, false)
833        } else if is_text {
834            match std::str::from_utf8(&bytes) {
835                Ok(s) => (Some(s.to_string()), false),
836                Err(_) => (
837                    Some(base64::engine::general_purpose::STANDARD.encode(&bytes)),
838                    true,
839                ),
840            }
841        } else {
842            (
843                Some(base64::engine::general_purpose::STANDARD.encode(&bytes)),
844                true,
845            )
846        };
847
848        let status = parts.status.as_u16() as i64;
849        let headers = parts.headers;
850        let body_lambda = body_str.map(aws_lambda_events::encodings::Body::Text);
851
852        match shape {
853            HttpEventShape::ApiGatewayV1 => {
854                let mut api_resp = aws_lambda_events::apigw::ApiGatewayProxyResponse::default();
855                api_resp.status_code = status;
856                api_resp.headers = headers;
857                api_resp.body = body_lambda;
858                api_resp.is_base64_encoded = is_base64;
859                serde_json::to_value(api_resp).map_err(|e| {
860                    lambda_runtime::Error::from(format!("serialise APIGW v1 response: {e}"))
861                })
862            }
863            HttpEventShape::ApiGatewayV2 => {
864                let mut api_resp = aws_lambda_events::apigw::ApiGatewayV2httpResponse::default();
865                api_resp.status_code = status;
866                api_resp.headers = headers;
867                api_resp.body = body_lambda;
868                api_resp.is_base64_encoded = is_base64;
869                serde_json::to_value(api_resp).map_err(|e| {
870                    lambda_runtime::Error::from(format!("serialise APIGW v2 response: {e}"))
871                })
872            }
873            HttpEventShape::Alb => {
874                let mut api_resp = aws_lambda_events::alb::AlbTargetGroupResponse::default();
875                api_resp.status_code = status;
876                api_resp.headers = headers;
877                api_resp.body = body_lambda;
878                api_resp.is_base64_encoded = is_base64;
879                serde_json::to_value(api_resp).map_err(|e| {
880                    lambda_runtime::Error::from(format!("serialise ALB response: {e}"))
881                })
882            }
883        }
884    }
885}
886
887/// Inbound HTTP event shape detected by [`LambdaA2aHandler::handle_http_event_value`].
888///
889/// The response envelope must match the inbound shape — see method
890/// docs for the contract.
891#[derive(Debug, Clone, Copy, PartialEq, Eq)]
892enum HttpEventShape {
893    ApiGatewayV1,
894    ApiGatewayV2,
895    Alb,
896}
897
898impl LambdaA2aHandler {
899    /// Handle an SQS batch from the durable executor queue.
900    ///
901    /// Returns `SqsBatchResponse` whose `batch_item_failures` list
902    /// tells the event source mapping which records to retry.
903    /// Requires `FunctionResponseTypes: ["ReportBatchItemFailures"]`
904    /// on the mapping.
905    ///
906    /// Each record: deserialize envelope → load task → terminal-no-op
907    /// check → cancel-marker direct-CANCELED commit → run executor.
908    /// See `durable::drive_sqs_batch` for the full semantics.
909    #[cfg(feature = "sqs")]
910    pub async fn handle_sqs(
911        &self,
912        event: aws_lambda_events::event::sqs::SqsEvent,
913    ) -> aws_lambda_events::event::sqs::SqsBatchResponse {
914        durable::drive_sqs_batch(&self.state, event).await
915    }
916
917    /// Run this handler as a pure HTTP Lambda. Strict: any non-HTTP
918    /// event shape (SQS, DynamoDB stream, scheduler, …) causes
919    /// `lambda_http::run` to return a deserialization error on
920    /// invocation.
921    ///
922    /// Appropriate for Lambdas whose only trigger is a Function URL,
923    /// API Gateway, or ALB. The request Lambda in the two-Lambda
924    /// durable-executor topology is the canonical caller.
925    pub async fn run_http_only(self) -> Result<(), lambda_runtime::Error> {
926        let handler = Arc::new(self);
927        lambda_http::run(lambda_http::service_fn(move |event| {
928            let h = Arc::clone(&handler);
929            async move { h.handle(event).await }
930        }))
931        .await
932    }
933}
934
935// Default entrypoint when the `sqs` feature is off — `.run()` aliases
936// the HTTP-only runner. With the `sqs` feature enabled, `.run()` lives
937// in `durable.rs` and routes HTTP+SQS via the classifier.
938#[cfg(not(feature = "sqs"))]
939impl LambdaA2aHandler {
940    /// Default Lambda runner. Dispatches based on the event shape the
941    /// binary receives. Without the `sqs` feature this is HTTP-only
942    /// (equivalent to [`LambdaA2aHandler::run_http_only`]).
943    pub async fn run(self) -> Result<(), lambda_runtime::Error> {
944        self.run_http_only().await
945    }
946}
947
948impl LambdaA2aServerBuilder {
949    /// Build and run in one call. Sugar for
950    /// `self.build()?.run().await`. Hides the handler from the adopter —
951    /// use `.build()` + `handler.run_*()` instead when the handler is
952    /// needed for unit tests, custom middleware, or an explicit
953    /// topology runner.
954    pub async fn run(self) -> Result<(), lambda_runtime::Error> {
955        let handler = self.build().map_err(|e| {
956            lambda_runtime::Error::from(format!("LambdaA2aServerBuilder::run: {e}"))
957        })?;
958        handler.run().await
959    }
960}
961
962#[cfg(test)]
963mod handle_http_event_value_shape_tests {
964    //! Regression tests for the event-shape symmetry contract on
965    //! [`LambdaA2aHandler::handle_http_event_value`]. The wire-shape
966    //! mismatch (v2 response to a v1 invocation) is what causes API
967    //! Gateway REST integrations to return `502 Bad Gateway`. These
968    //! tests pin the contract: response envelope must always match
969    //! the inbound request envelope.
970
971    use super::*;
972    use std::sync::Arc;
973    use turul_a2a::executor::{AgentExecutor, ExecutionContext};
974    use turul_a2a::server::RuntimeConfig;
975    use turul_a2a::server::in_flight::InFlightRegistry;
976    use turul_a2a::storage::InMemoryA2aStorage;
977    use turul_a2a::streaming::TaskEventBroker;
978    use turul_a2a_types::{Message, Task};
979
980    struct NoOpExecutor;
981
982    #[async_trait::async_trait]
983    impl AgentExecutor for NoOpExecutor {
984        async fn execute(
985            &self,
986            task: &mut Task,
987            _msg: &Message,
988            _ctx: &ExecutionContext,
989        ) -> Result<(), turul_a2a::error::A2aError> {
990            let mut proto = task.as_proto().clone();
991            proto.status = Some(turul_a2a_proto::TaskStatus {
992                state: turul_a2a_proto::TaskState::Completed.into(),
993                message: None,
994                timestamp: None,
995            });
996            *task = Task::try_from(proto).unwrap();
997            Ok(())
998        }
999        fn agent_card(&self) -> turul_a2a_proto::AgentCard {
1000            turul_a2a_proto::AgentCard {
1001                name: "test-agent".to_string(),
1002                ..Default::default()
1003            }
1004        }
1005    }
1006
1007    fn build_handler() -> LambdaA2aHandler {
1008        let s = Arc::new(InMemoryA2aStorage::new());
1009        let state = AppState {
1010            executor: Arc::new(NoOpExecutor),
1011            task_storage: s.clone(),
1012            push_storage: s.clone(),
1013            event_store: s.clone(),
1014            atomic_store: s.clone(),
1015            event_broker: TaskEventBroker::new(),
1016            middleware_stack: Arc::new(MiddlewareStack::new(vec![])),
1017            runtime_config: RuntimeConfig::default(),
1018            in_flight: Arc::new(InFlightRegistry::new()),
1019            cancellation_supervisor: s.clone(),
1020            push_delivery_store: None,
1021            push_dispatcher: None,
1022            durable_executor_queue: None,
1023        };
1024        let router = build_router(state.clone());
1025        LambdaA2aHandler {
1026            router,
1027            state,
1028            path_prefix: None,
1029        }
1030    }
1031
1032    /// API Gateway v2 / Function URL event. `requestContext.http.method`
1033    /// is what `LambdaRequest::ApiGatewayV2` keys off.
1034    fn apigw_v2_agent_card_get() -> serde_json::Value {
1035        serde_json::json!({
1036            "version": "2.0",
1037            "routeKey": "GET /.well-known/agent-card.json",
1038            "rawPath": "/.well-known/agent-card.json",
1039            "rawQueryString": "",
1040            "headers": {"accept": "application/json", "a2a-version": "1.0"},
1041            "requestContext": {
1042                "accountId": "000000000000",
1043                "apiId": "api",
1044                "domainName": "fake.execute-api",
1045                "domainPrefix": "fake",
1046                "http": {
1047                    "method": "GET",
1048                    "path": "/.well-known/agent-card.json",
1049                    "protocol": "HTTP/1.1",
1050                    "sourceIp": "127.0.0.1",
1051                    "userAgent": "test"
1052                },
1053                "requestId": "rid",
1054                "routeKey": "GET /.well-known/agent-card.json",
1055                "stage": "$default",
1056                "time": "01/Jan/2026:00:00:00 +0000",
1057                "timeEpoch": 1_735_689_600_000i64
1058            },
1059            "isBase64Encoded": false
1060        })
1061    }
1062
1063    /// API Gateway v1 (REST) event. Top-level `httpMethod` + `path` and
1064    /// `requestContext.identity` (no `requestContext.http`) is what
1065    /// `LambdaRequest::ApiGatewayV1` keys off. This is the shape REST
1066    /// API + AWS_PROXY integrations send into the Lambda — the live
1067    /// fleet topology that 502'd against a hardcoded v2 response.
1068    ///
1069    /// `stage` is `$default` here so `lambda_http` does not prepend a
1070    /// stage prefix — this fixture exists to pin the response-shape
1071    /// contract, not the path-rewriting behaviour. See the
1072    /// `strip_path_prefix_*` tests for stage-prefix coverage.
1073    fn apigw_v1_agent_card_get_with_stage(stage: &str, path: &str) -> serde_json::Value {
1074        let request_context_path = if stage == "$default" {
1075            path.to_string()
1076        } else {
1077            format!("/{stage}{path}")
1078        };
1079        serde_json::json!({
1080            "resource": "/{proxy+}",
1081            "path": path,
1082            "httpMethod": "GET",
1083            "headers": {"accept": "application/json", "a2a-version": "1.0"},
1084            "multiValueHeaders": {},
1085            "queryStringParameters": null,
1086            "multiValueQueryStringParameters": null,
1087            "pathParameters": {"proxy": path.trim_start_matches('/').to_string()},
1088            "stageVariables": null,
1089            "requestContext": {
1090                "accountId": "000000000000",
1091                "apiId": "abc123",
1092                "domainName": "fake.execute-api",
1093                "domainPrefix": "fake",
1094                "extendedRequestId": "xrid",
1095                "httpMethod": "GET",
1096                "identity": {
1097                    "sourceIp": "127.0.0.1",
1098                    "userAgent": "test"
1099                },
1100                "path": request_context_path,
1101                "protocol": "HTTP/1.1",
1102                "requestId": "rid",
1103                "requestTime": "01/Jan/2026:00:00:00 +0000",
1104                "requestTimeEpoch": 1_735_689_600_000i64,
1105                "resourceId": "rsrc",
1106                "resourcePath": "/{proxy+}",
1107                "stage": stage
1108            },
1109            "body": null,
1110            "isBase64Encoded": false
1111        })
1112    }
1113
1114    fn apigw_v1_agent_card_get() -> serde_json::Value {
1115        apigw_v1_agent_card_get_with_stage("$default", "/.well-known/agent-card.json")
1116    }
1117
1118    /// ALB target-group event. `requestContext.elb` is what
1119    /// `LambdaRequest::Alb` keys off.
1120    fn alb_agent_card_get() -> serde_json::Value {
1121        serde_json::json!({
1122            "requestContext": {
1123                "elb": {
1124                    "targetGroupArn": "arn:aws:elasticloadbalancing:ap-southeast-2:000:targetgroup/fake/00"
1125                }
1126            },
1127            "httpMethod": "GET",
1128            "path": "/.well-known/agent-card.json",
1129            "queryStringParameters": {},
1130            "headers": {"accept": "application/json", "a2a-version": "1.0"},
1131            "body": "",
1132            "isBase64Encoded": false
1133        })
1134    }
1135
1136    /// REST v1 inbound MUST produce a v1 response shape. This is the
1137    /// regression for the live-fleet 502.
1138    #[tokio::test]
1139    async fn v1_inbound_event_emits_v1_response_envelope() {
1140        let handler = build_handler();
1141        let resp_value = handler
1142            .handle_http_event_value(apigw_v1_agent_card_get())
1143            .await
1144            .expect("handle_http_event_value(v1) must succeed");
1145
1146        // Parses cleanly as v1.
1147        let v1: aws_lambda_events::apigw::ApiGatewayProxyResponse =
1148            serde_json::from_value(resp_value.clone())
1149                .expect("response must deserialize as ApiGatewayProxyResponse (v1)");
1150        assert_eq!(v1.status_code, 200);
1151        match v1.body {
1152            Some(aws_lambda_events::encodings::Body::Text(s)) => {
1153                assert!(s.contains("\"name\""), "v1 body shape: {s}");
1154            }
1155            other => panic!("expected text body, got {other:?}"),
1156        }
1157
1158        // Negative: v1 response JSON must NOT contain the v2-only
1159        // `cookies` field. Presence of `cookies` is API Gateway REST's
1160        // 502 trigger.
1161        assert!(
1162            resp_value.get("cookies").is_none(),
1163            "v1 response must not carry v2-only `cookies` field; payload: {resp_value}"
1164        );
1165    }
1166
1167    /// API Gateway v2 / Function URL inbound MUST produce a v2 response
1168    /// shape (current behaviour, locked in to prevent the inverse
1169    /// regression).
1170    #[tokio::test]
1171    async fn v2_inbound_event_emits_v2_response_envelope() {
1172        let handler = build_handler();
1173        let resp_value = handler
1174            .handle_http_event_value(apigw_v2_agent_card_get())
1175            .await
1176            .expect("handle_http_event_value(v2) must succeed");
1177
1178        let v2: aws_lambda_events::apigw::ApiGatewayV2httpResponse =
1179            serde_json::from_value(resp_value.clone())
1180                .expect("response must deserialize as ApiGatewayV2httpResponse (v2)");
1181        assert_eq!(v2.status_code, 200);
1182        match v2.body {
1183            Some(aws_lambda_events::encodings::Body::Text(s)) => {
1184                assert!(s.contains("\"name\""), "v2 body shape: {s}");
1185            }
1186            other => panic!("expected text body, got {other:?}"),
1187        }
1188    }
1189
1190    /// ALB inbound MUST produce an ALB response shape.
1191    #[tokio::test]
1192    async fn alb_inbound_event_emits_alb_response_envelope() {
1193        let handler = build_handler();
1194        let resp_value = handler
1195            .handle_http_event_value(alb_agent_card_get())
1196            .await
1197            .expect("handle_http_event_value(alb) must succeed");
1198
1199        let alb: aws_lambda_events::alb::AlbTargetGroupResponse =
1200            serde_json::from_value(resp_value)
1201                .expect("response must deserialize as AlbTargetGroupResponse");
1202        assert_eq!(alb.status_code, 200);
1203    }
1204
1205    /// Live-fleet end-to-end regression: REST v1 stage-prefixed event
1206    /// + `strip_path_prefix("/dev")` on the handler routes through
1207    /// the agent-card endpoint AND emits a v1 response envelope. This
1208    /// is the exact scenario the live fleet was 502'ing on.
1209    #[tokio::test]
1210    async fn v1_stage_prefix_inbound_with_strip_emits_v1_response() {
1211        let s = Arc::new(InMemoryA2aStorage::new());
1212        let state = AppState {
1213            executor: Arc::new(NoOpExecutor),
1214            task_storage: s.clone(),
1215            push_storage: s.clone(),
1216            event_store: s.clone(),
1217            atomic_store: s.clone(),
1218            event_broker: TaskEventBroker::new(),
1219            middleware_stack: Arc::new(MiddlewareStack::new(vec![])),
1220            runtime_config: RuntimeConfig::default(),
1221            in_flight: Arc::new(InFlightRegistry::new()),
1222            cancellation_supervisor: s.clone(),
1223            push_delivery_store: None,
1224            push_dispatcher: None,
1225            durable_executor_queue: None,
1226        };
1227        let router = build_router(state.clone());
1228        let handler = LambdaA2aHandler {
1229            router,
1230            state,
1231            path_prefix: Some(Arc::from("/dev")),
1232        };
1233
1234        let event = apigw_v1_agent_card_get_with_stage("dev", "/.well-known/agent-card.json");
1235        let resp_value = handler
1236            .handle_http_event_value(event)
1237            .await
1238            .expect("v1 stage-prefixed event must succeed under strip_path_prefix");
1239
1240        let v1: aws_lambda_events::apigw::ApiGatewayProxyResponse =
1241            serde_json::from_value(resp_value)
1242                .expect("response must deserialize as ApiGatewayProxyResponse");
1243        assert_eq!(v1.status_code, 200);
1244        match v1.body {
1245            Some(aws_lambda_events::encodings::Body::Text(s)) => {
1246                assert!(s.contains("\"name\""), "v1 stage-prefix body: {s}");
1247            }
1248            other => panic!("expected text body, got {other:?}"),
1249        }
1250    }
1251
1252    /// Cross-shape decoding regression: a v1 response JSON carries
1253    /// `multiValueHeaders` (REST shape) but no `cookies`; a v2 response
1254    /// carries `cookies` but no v1-only `multiValueHeaders` semantics.
1255    /// Confirms the two shapes are wire-distinguishable, which is why
1256    /// the gateway 502s on a mismatch.
1257    #[tokio::test]
1258    async fn v1_response_does_not_match_v2_envelope_strictly() {
1259        let handler = build_handler();
1260        let v1_value = handler
1261            .handle_http_event_value(apigw_v1_agent_card_get())
1262            .await
1263            .unwrap();
1264
1265        // Sanity: v1 response JSON has no v2-only `cookies` field.
1266        assert!(v1_value.get("cookies").is_none());
1267
1268        let v2_value = handler
1269            .handle_http_event_value(apigw_v2_agent_card_get())
1270            .await
1271            .unwrap();
1272        // Sanity: v2 response JSON carries the v2 `cookies` field
1273        // (default empty array — still present on the wire).
1274        assert!(
1275            v2_value.get("cookies").is_some(),
1276            "v2 response should carry the v2-only cookies field"
1277        );
1278    }
1279}