turul_a2a_aws_lambda/

lib.rs

//! AWS Lambda adapter for [`turul-a2a`](https://docs.rs/turul-a2a): it feeds
//! Lambda events to the same axum router the binary server uses, so one
//! [`AgentExecutor`] runs unchanged
//! behind a Function URL, API Gateway, ALB or SQS.
//!
//! Use it to deploy an A2A agent as a Lambda function; executor, storage and
//! middleware come from `turul-a2a`.
//!
//! # Example
//!
//! ```toml
//! [dependencies]
//! turul-a2a = "0.2"
//! turul-a2a-proto = "0.2"
//! turul-a2a-aws-lambda = "0.2"
//! async-trait = "0.1"
//! lambda_runtime = "1"
//! tokio = { version = "1", features = ["full"] }
//! ```
//!
//! ```rust,no_run
//! use turul_a2a::storage::InMemoryA2aStorage;
//! use turul_a2a_aws_lambda::LambdaA2aServerBuilder;
//! # use turul_a2a::prelude::*;
//! # struct Echo;
//! # #[async_trait::async_trait]
//! # impl AgentExecutor for Echo {
//! #     async fn execute(&self, _: &mut Task, _: &Message, _: &ExecutionContext)
//! #         -> Result<(), A2aError> { Ok(()) }
//! #     fn agent_card(&self) -> turul_a2a_proto::AgentCard { Default::default() }
//! # }
//!
//! #[tokio::main]
//! async fn main() -> Result<(), lambda_runtime::Error> {
//!     // In-memory storage suits a demo; use a shared backend in production.
//!     LambdaA2aServerBuilder::new()
//!         .executor(Echo)
//!         .storage(InMemoryA2aStorage::new())
//!         .run()
//!         .await
//! }
//! ```
//!
//! # Choosing a runner
//!
//! `LambdaA2aServerBuilder::run` is the default entry point. Without the
//! `sqs` feature it is `run_http_only`; with it, `run_http_and_sqs`, which
//! accepts either event shape. The named runners are strict and fail on any
//! other shape: `run_http_only` (Function URL / API Gateway),
//! `run_sqs_only` and `run_http_and_sqs` (both `sqs`). For a third trigger,
//! match on [`classify_event`] yourself and pass HTTP events to
//! [`LambdaA2aHandler::handle_http_event_value`].
//! Behind a REST API Gateway stage, set
//! [`LambdaA2aServerBuilder::strip_path_prefix`] so the router sees root-rooted paths.
//!
//! # Lambda behaviour
//!
//! - SSE responses are buffered: the task runs, its events are collected, and
//!   one response is returned. `GET /tasks/{id}:subscribe` serves tasks that
//!   are not yet terminal and closes on a terminal state or `INPUT_REQUIRED`.
//! - Push-notification delivery cannot rely on work continuing after the
//!   response. Recovery is carried by [`LambdaScheduledRecoveryHandler`]
//!   (required on every backend) and, on DynamoDB, [`LambdaStreamRecoveryHandler`].
//! - No cross-instance cancel poller runs. `CancelTask` writes its marker to the
//!   shared backend, but a live executor on a different invocation does not
//!   observe it.
//!
//! # Features
//!
//! | Feature | Enables |
//! |---|---|
//! | `sqs` | SQS-backed durable executor and the SQS runners |
//! | `dynamodb` | `turul-a2a`'s DynamoDB backend |
//!
//! # Where next
//!
//! - [`LambdaAuthorizerMiddleware`] to map an API Gateway authorizer to a caller identity
//! - [`lambda-agent`](https://github.com/aussierobots/turul-a2a/tree/main/examples/lambda-agent)
//!   and the worker examples on GitHub

mod adapter;
mod auth;
mod event;
mod no_streaming;
mod scheduled_recovery;
mod stream_recovery;

#[cfg(feature = "sqs")]
mod durable;

#[cfg(test)]
mod builder_tests;
#[cfg(test)]
mod scheduled_recovery_tests;
#[cfg(test)]
mod stream_recovery_tests;

pub use adapter::{axum_to_lambda_response, lambda_to_axum_request};
pub use auth::{AuthorizerMapping, LambdaAuthorizerMiddleware};
pub use event::{LambdaEvent, classify_event};
pub use no_streaming::NoStreamingLayer;
pub use scheduled_recovery::{
    LambdaScheduledRecoveryConfig, LambdaScheduledRecoveryHandler, LambdaScheduledRecoveryResponse,
    handle_maintenance_trigger,
};
pub use stream_recovery::LambdaStreamRecoveryHandler;

#[cfg(feature = "sqs")]
pub use durable::{SqsDurableExecutorQueue, drive_sqs_batch};

use std::sync::Arc;

use turul_a2a::error::A2aError;
use turul_a2a::executor::AgentExecutor;
use turul_a2a::middleware::{A2aMiddleware, MiddlewareStack};
use turul_a2a::push::A2aPushDeliveryStore;
use turul_a2a::router::{AppState, build_router};
use turul_a2a::server::RuntimeConfig;
use turul_a2a::storage::{
    A2aAtomicStore, A2aCancellationSupervisor, A2aEventStore, A2aPushNotificationStorage,
    A2aTaskStorage,
};
use turul_a2a::streaming::TaskEventBroker;

/// Builder for Lambda A2A handler.
///
/// Use `.storage(my_storage)` to supply a single backend implementing all storage traits.
/// All four storage traits must share the same backend so streaming, push delivery,
/// and cross-instance cancellation coordinate correctly.
///
/// For push-notification delivery, wire `.storage()` with a backend that implements
/// [`A2aPushDeliveryStore`] (all four first-party backends do). The atomic store
/// must also opt in via `with_push_dispatch_enabled(true)`; the
/// builder rejects configurations where the flag and the delivery store disagree.
pub struct LambdaA2aServerBuilder {
    executor: Option<Arc<dyn AgentExecutor>>,
    task_storage: Option<Arc<dyn A2aTaskStorage>>,
    push_storage: Option<Arc<dyn A2aPushNotificationStorage>>,
    event_store: Option<Arc<dyn A2aEventStore>>,
    atomic_store: Option<Arc<dyn A2aAtomicStore>>,
    cancellation_supervisor: Option<Arc<dyn A2aCancellationSupervisor>>,
    push_delivery_store: Option<Arc<dyn A2aPushDeliveryStore>>,
    middleware: Vec<Arc<dyn A2aMiddleware>>,
    runtime_config: RuntimeConfig,
    durable_executor_queue: Option<Arc<dyn turul_a2a::durable_executor::DurableExecutorQueue>>,
    path_prefix: Option<String>,
    outbound_url_validator: Option<turul_a2a::push::ssrf::OutboundUrlValidator>,
}

impl LambdaA2aServerBuilder {
    pub fn new() -> Self {
        Self {
            executor: None,
            task_storage: None,
            push_storage: None,
            event_store: None,
            atomic_store: None,
            cancellation_supervisor: None,
            push_delivery_store: None,
            middleware: vec![],
            runtime_config: RuntimeConfig::default(),
            durable_executor_queue: None,
            path_prefix: None,
            outbound_url_validator: None,
        }
    }

    /// Configure an HTTP path prefix to strip before the A2A router
    /// sees the request. Needed when Lambda sits behind an API Gateway
    /// whose stage + resource tree is forwarded into the function
    /// (e.g. REST API `AWS_PROXY` integrations surface
    /// `/stage/agent/message:send` where the A2A
    /// router expects `/message:send`).
    ///
    /// The prefix must start with `/` and must not end with `/`
    /// (unless it is exactly `/`, which is a no-op). Segment boundaries
    /// are respected — a prefix of `/dev` will not strip `/devs/...`.
    /// Paths that do not start with the prefix pass through unchanged
    /// (the router will 404 on a genuinely unknown path, which is the
    /// correct failure mode). Applies to both
    /// [`LambdaA2aHandler::run_http_only`] and
    /// `LambdaA2aHandler::run_http_and_sqs`. SQS events are
    /// unaffected.
    pub fn strip_path_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.path_prefix = Some(prefix.into());
        self
    }

    /// Wire a durable executor queue. When set, the
    /// `return_immediately = true` path in `core_send_message` enqueues
    /// a [`turul_a2a::durable_executor::QueuedExecutorJob`] on this
    /// queue instead of spawning the executor locally, and the
    /// capability flag `RuntimeConfig::supports_return_immediately` is
    /// turned **back on** as an implementation detail — the capability
    /// cannot be claimed without supplying the queue.
    ///
    /// Adopters should typically reach for the SQS-specific helper
    /// `Self::with_sqs_return_immediately` instead; this generic
    /// method exists so adopters can inject in-memory fakes for tests
    /// or provide non-SQS transports (Kinesis, Step Functions task
    /// token, self-invoke, etc.).
    pub fn with_durable_executor(
        mut self,
        queue: Arc<dyn turul_a2a::durable_executor::DurableExecutorQueue>,
    ) -> Self {
        self.durable_executor_queue = Some(queue);
        self.runtime_config.supports_return_immediately = true;
        self
    }

    /// Wire the AWS SQS durable executor queue. Ergonomic
    /// helper over [`Self::with_durable_executor`] that constructs a
    /// [`SqsDurableExecutorQueue`] from a `queue_url` + shared SQS
    /// client. Requires the `sqs` feature.
    #[cfg(feature = "sqs")]
    pub fn with_sqs_return_immediately(
        self,
        queue_url: impl Into<String>,
        sqs_client: Arc<aws_sdk_sqs::Client>,
    ) -> Self {
        let queue = Arc::new(SqsDurableExecutorQueue::new(queue_url, sqs_client));
        self.with_durable_executor(queue)
    }

    pub fn executor(mut self, exec: impl AgentExecutor + 'static) -> Self {
        self.executor = Some(Arc::new(exec));
        self
    }

    /// Set all storage from a single backend instance.
    ///
    /// This is the **preferred** method — a single struct implementing all
    /// storage traits guarantees the same-backend requirement
    /// AND ensures the cancellation supervisor reads the same backend
    /// the cancel marker is written to. A mismatch here (e.g., DynamoDB
    /// task storage + in-memory cancellation supervisor) silently
    /// breaks cross-instance cancellation.
    ///
    /// errata: `.storage()` wires storage traits only. It does
    /// **not** auto-register the storage as a push-delivery store, even if
    /// the backend happens to implement [`A2aPushDeliveryStore`]. To opt in
    /// to push delivery, call [`Self::push_delivery_store`] explicitly and
    /// call `with_push_dispatch_enabled(true)` on the storage instance
    /// before passing it here. Non-push deployments need neither.
    pub fn storage<S>(mut self, storage: S) -> Self
    where
        S: A2aTaskStorage
            + A2aPushNotificationStorage
            + A2aEventStore
            + A2aAtomicStore
            + A2aCancellationSupervisor
            + Clone
            + 'static,
    {
        self.task_storage = Some(Arc::new(storage.clone()));
        self.push_storage = Some(Arc::new(storage.clone()));
        self.event_store = Some(Arc::new(storage.clone()));
        self.atomic_store = Some(Arc::new(storage.clone()));
        self.cancellation_supervisor = Some(Arc::new(storage));
        self
    }

    /// Set task storage individually. Must be paired with event_store/atomic_store
    /// AND `cancellation_supervisor()` for same-backend compliance.
    pub fn task_storage(mut self, s: impl A2aTaskStorage + 'static) -> Self {
        self.task_storage = Some(Arc::new(s));
        self
    }

    /// Set push notification storage individually.
    pub fn push_storage(mut self, s: impl A2aPushNotificationStorage + 'static) -> Self {
        self.push_storage = Some(Arc::new(s));
        self
    }

    /// Set event store individually. Must match task_storage backend.
    pub fn event_store(mut self, s: impl A2aEventStore + 'static) -> Self {
        self.event_store = Some(Arc::new(s));
        self
    }

    /// Set atomic store individually. Must match task_storage backend.
    pub fn atomic_store(mut self, s: impl A2aAtomicStore + 'static) -> Self {
        self.atomic_store = Some(Arc::new(s));
        self
    }

    /// Set the cancellation supervisor individually. Must match
    /// task_storage backend — otherwise `:cancel` writes the marker to
    /// one backend and the supervisor reads from another, silently
    /// breaking cross-instance cancellation. Prefer `.storage()` for
    /// unified wiring. Consumed by `core_cancel_task`.
    pub fn cancellation_supervisor(mut self, s: impl A2aCancellationSupervisor + 'static) -> Self {
        self.cancellation_supervisor = Some(Arc::new(s));
        self
    }

    /// Set the push delivery store individually.
    ///
    /// Prefer `.storage()` for unified wiring. The delivery store MUST
    /// be on the same backend as `.task_storage(...)` — the builder
    /// rejects mismatches. Passing a delivery store also requires the
    /// atomic store to have opted in via `with_push_dispatch_enabled(true)`
    ///.
    pub fn push_delivery_store(mut self, store: impl A2aPushDeliveryStore + 'static) -> Self {
        self.push_delivery_store = Some(Arc::new(store));
        self
    }

    /// Override the runtime configuration (timeouts, retry budgets,
    /// push tuning). Defaults to [`RuntimeConfig::default()`]. The
    /// builder validates push settings (claim expiry vs retry horizon,
    /// ) when `.push_delivery_store(...)` is wired.
    pub fn runtime_config(mut self, cfg: RuntimeConfig) -> Self {
        self.runtime_config = cfg;
        self
    }

    pub fn middleware(mut self, mw: Arc<dyn A2aMiddleware>) -> Self {
        self.middleware.push(mw);
        self
    }

    /// Operator allowlist for push webhook URLs.
    ///
    /// Called with each webhook URL on every delivery attempt, before
    /// DNS resolution. Returning `Err` refuses the delivery: nothing is
    /// resolved or POSTed, and the delivery is recorded as a failed
    /// delivery with error class `SSRFBlocked`. The reason string is
    /// not persisted. Runs in addition to the built-in scheme and
    /// private-address checks, not instead of them. Takes effect only
    /// when push delivery is wired (`push_delivery_store`).
    pub fn outbound_url_validator(
        mut self,
        validator: turul_a2a::push::ssrf::OutboundUrlValidator,
    ) -> Self {
        self.outbound_url_validator = Some(validator);
        self
    }

    pub fn build(self) -> Result<LambdaA2aHandler, A2aError> {
        let executor = self
            .executor
            .ok_or(A2aError::Internal("executor is required".into()))?;
        let task_storage = self.task_storage.ok_or(A2aError::Internal(
            "task_storage is required for Lambda".into(),
        ))?;
        let push_storage = self.push_storage.ok_or(A2aError::Internal(
            "push_storage is required for Lambda".into(),
        ))?;
        let event_store: Arc<dyn A2aEventStore> = self.event_store.ok_or(A2aError::Internal(
            "event_store is required for Lambda (use .storage() for unified backend)".into(),
        ))?;
        let atomic_store: Arc<dyn A2aAtomicStore> = self.atomic_store.ok_or(A2aError::Internal(
            "atomic_store is required for Lambda (use .storage() for unified backend)".into(),
        ))?;
        // the cancellation supervisor MUST be supplied. No
        // InMemoryA2aStorage fallback: that would silently break
        // cross-instance cancellation by reading markers from a
        // different backend than they were written to.
        let cancellation_supervisor: Arc<dyn A2aCancellationSupervisor> =
            self.cancellation_supervisor.ok_or(A2aError::Internal(
                "cancellation_supervisor is required for Lambda. Use .storage() for unified \
                 backend wiring, or .cancellation_supervisor(...) alongside the individual \
                 storage setters. Omitting it silently breaks cross-instance cancellation."
                    .into(),
            ))?;

        // Same-backend enforcement: all storage traits + the
        // cancellation supervisor must report the same `backend_name`.
        let task_backend = task_storage.backend_name();
        let push_backend = push_storage.backend_name();
        let event_backend = event_store.backend_name();
        let atomic_backend = atomic_store.backend_name();
        let supervisor_backend = cancellation_supervisor.backend_name();
        if task_backend != push_backend
            || task_backend != event_backend
            || task_backend != atomic_backend
            || task_backend != supervisor_backend
        {
            return Err(A2aError::Internal(format!(
                "Storage backend mismatch: task={task_backend}, push={push_backend}, \
                 event={event_backend}, atomic={atomic_backend}, \
                 cancellation_supervisor={supervisor_backend}. \
                 All storage traits, the cancellation supervisor included, \
                 must share one backend so cross-instance cancel markers \
                 are observable. Use .storage() for unified backend."
            )));
        }

        let push_delivery_store = self.push_delivery_store;

        // Push-dispatch consistency: the atomic store's
        // opt-in flag and the presence of `push_delivery_store` MUST
        // agree. Both-true = push fully wired; both-false = non-push
        // deployment. Mixed cases are build errors — this mirrors the
        // main server builder (server/mod.rs).
        match (
            push_delivery_store.is_some(),
            atomic_store.push_dispatch_enabled(),
        ) {
            (true, true) | (false, false) => {}
            (true, false) => {
                return Err(A2aError::Internal(
                    "push_delivery_store wired but atomic_store.push_dispatch_enabled() \
                     is false. Call .with_push_dispatch_enabled(true) on the backend \
                     storage before passing it to .storage()."
                        .into(),
                ));
            }
            (false, true) => {
                return Err(A2aError::Internal(
                    "atomic_store.push_dispatch_enabled() is true but no \
                     push_delivery_store is wired. Pending-dispatch markers would be \
                     written with no consumer, imposing load-bearing infra for no \
                     benefit. If you need to populate markers for an external \
                     consumer, open an issue for a distinctly-named opt-in — for now, \
                     this configuration is rejected."
                        .into(),
                ));
            }
        }

        // The Lambda execution environment can freeze once the HTTP
        // response is flushed, so a `tokio::spawn`'d executor
        // continuation is opportunistic only. Refuse
        // `SendMessageConfiguration.return_immediately = true` at the
        // `core_send_message` entry point (UnsupportedOperationError —
        // an aspect of SendMessage this runtime cannot honour, A2A
        // v1.0.1 §3.3.2) via this capability flag instead of silently
        // orphaning the executor.
        //
        // if a durable executor queue is wired via
        // [`Self::with_durable_executor`] (or the SQS helper), the
        // capability is re-enabled because enqueueing is durable
        // across the Lambda freeze. The explicit re-enable lives in
        // the builder method so the capability cannot be claimed
        // without the mechanism.
        let mut runtime_config = self.runtime_config;
        if self.durable_executor_queue.is_none() {
            runtime_config.supports_return_immediately = false;
        }
        // No background loops run on Lambda, so a lease written here
        // would never be renewed and peers would fail the live task.
        runtime_config.executor_lease_ttl = None;
        runtime_config.executor_stall_timeout = None;
        // The idle-task sweep needs no loop: the scheduled-recovery
        // handler runs it.
        if runtime_config
            .idle_task_timeout
            .is_some_and(|d| d.is_zero())
        {
            return Err(A2aError::Internal(
                "idle_task_timeout must be greater than zero".into(),
            ));
        }

        // Push dispatcher wiring. When
        // `push_delivery_store` is present, validate its backend matches,
        // validate the retry-horizon invariant, and construct the
        // `PushDispatcher`. Under Lambda the dispatcher runs
        // opportunistically post-return; correctness
        // comes from the atomic marker + stream/scheduler recovery.
        let push_dispatcher: Option<Arc<turul_a2a::push::PushDispatcher>> = if let Some(delivery) =
            push_delivery_store.as_ref()
        {
            let push_delivery_backend = delivery.backend_name();
            if task_backend != push_delivery_backend {
                return Err(A2aError::Internal(format!(
                    "Storage backend mismatch: task={task_backend}, \
                         push_delivery={push_delivery_backend}. \
                         All storage traits must share the same backend."
                )));
            }

            // claim expiry must exceed the worst-case
            // retry horizon so a live claim is not mis-classified as
            // stale mid-retry.
            let retry_horizon = runtime_config
                .push_backoff_cap
                .saturating_mul(runtime_config.push_max_attempts as u32);
            if runtime_config.push_claim_expiry <= retry_horizon {
                return Err(A2aError::Internal(format!(
                    "push_claim_expiry ({:?}) must be greater than retry horizon \
                         (push_max_attempts={} * push_backoff_cap={:?} = {:?}). \
                         Raise push_claim_expiry or lower push_max_attempts/push_backoff_cap.",
                    runtime_config.push_claim_expiry,
                    runtime_config.push_max_attempts,
                    runtime_config.push_backoff_cap,
                    retry_horizon
                )));
            }

            let delivery_cfg =
                turul_a2a::push::delivery::PushDeliveryConfig::from_runtime_config(&runtime_config);

            let instance_id = format!("a2a-lambda-{}", uuid::Uuid::now_v7());
            let worker = turul_a2a::push::delivery::PushDeliveryWorker::new(
                delivery.clone(),
                delivery_cfg,
                self.outbound_url_validator.clone(),
                instance_id,
            )
            .map_err(|e| A2aError::Internal(format!("push worker build failed: {e}")))?;

            Some(Arc::new(turul_a2a::push::PushDispatcher::new(
                Arc::new(worker),
                push_storage.clone(),
                task_storage.clone(),
            )))
        } else {
            None
        };

        let state = AppState {
            executor,
            task_storage,
            push_storage,
            event_store,
            atomic_store,
            event_broker: TaskEventBroker::new(),
            middleware_stack: Arc::new(MiddlewareStack::new(self.middleware)),
            runtime_config,
            in_flight: Arc::new(turul_a2a::server::in_flight::InFlightRegistry::new()),
            cancellation_supervisor,
            push_delivery_store,
            push_dispatcher,
            durable_executor_queue: self.durable_executor_queue,
        };

        let router = build_router(state.clone());

        let path_prefix = match self.path_prefix {
            None => None,
            Some(p) if p.is_empty() || p == "/" => None,
            Some(p) => {
                if !p.starts_with('/') {
                    return Err(A2aError::InvalidRequest {
                        message: format!(
                            "LambdaA2aServerBuilder::strip_path_prefix: prefix must start with '/'; got {p:?}"
                        ),
                    });
                }
                if p.ends_with('/') {
                    return Err(A2aError::InvalidRequest {
                        message: format!(
                            "LambdaA2aServerBuilder::strip_path_prefix: prefix must not end with '/' (except '/'); got {p:?}"
                        ),
                    });
                }
                Some(Arc::from(p))
            }
        };

        Ok(LambdaA2aHandler {
            router,
            state,
            path_prefix,
        })
    }
}

impl Default for LambdaA2aServerBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Lambda handler wrapping the axum Router.
#[derive(Clone)]
pub struct LambdaA2aHandler {
    router: axum::Router,
    /// Held for SQS-event dispatch (`handle_sqs`) and the idle-task
    /// sweep (`cancel_idle_tasks`).
    state: AppState,
    /// Leading path segment to strip from each incoming HTTP request
    /// before the A2A router sees it. `None` = no stripping. Set via
    /// [`LambdaA2aServerBuilder::strip_path_prefix`].
    path_prefix: Option<Arc<str>>,
}

impl LambdaA2aHandler {
    pub fn builder() -> LambdaA2aServerBuilder {
        LambdaA2aServerBuilder::new()
    }

    /// Cancel tasks left waiting on their client for longer than the
    /// runtime configuration's `idle_task_timeout`; see
    /// [`turul_a2a::server::A2aServerBuilder::idle_task_timeout`]. With no
    /// timeout set this does nothing. Run by
    /// [`LambdaScheduledRecoveryHandler::with_idle_task_sweep`].
    pub async fn cancel_idle_tasks(
        &self,
    ) -> Result<turul_a2a::server::IdleSweep, turul_a2a::storage::A2aStorageError> {
        match self.state.runtime_config.idle_task_timeout {
            Some(idle_for) => {
                turul_a2a::server::liveness::cancel_idle_tasks(
                    &self.state,
                    Some(&self.state.in_flight),
                    idle_for,
                )
                .await
            }
            None => Ok(Default::default()),
        }
    }

    /// Handle a Lambda HTTP request.
    pub async fn handle(
        &self,
        event: lambda_http::Request,
    ) -> Result<lambda_http::Response<lambda_http::Body>, lambda_http::Error> {
        let axum_req = lambda_to_axum_request(event)?;
        let axum_req = match &self.path_prefix {
            Some(prefix) => adapter::strip_request_path_prefix(axum_req, prefix),
            None => axum_req,
        };

        let axum_resp = tower::ServiceExt::oneshot(self.router.clone(), axum_req)
            .await
            .map_err(|e| lambda_http::Error::from(format!("Router error: {e}")))?;

        axum_to_lambda_response(axum_resp).await
    }

    /// Handle a raw Lambda HTTP event delivered as `serde_json::Value`
    /// and return a response envelope **matching the inbound event
    /// shape**:
    ///
    /// | Inbound shape detected by `lambda_http::LambdaRequest`   | Response envelope                  |
    /// |----------------------------------------------------------|------------------------------------|
    /// | API Gateway v1 (REST) — top-level `httpMethod` + `path`  | `ApiGatewayProxyResponse`          |
    /// | API Gateway v2 / Function URL — `requestContext.http`    | `ApiGatewayV2httpResponse`         |
    /// | ALB target group — `requestContext.elb`                  | `AlbTargetGroupResponse`           |
    /// | WebSocket / unrecognised                                  | `Err(...)`                         |
    ///
    /// Mismatching envelopes — most commonly a v2 response sent in
    /// response to a v1 invocation — cause API Gateway to return
    /// `502 Bad Gateway` regardless of the body, because the response
    /// fails the gateway's contract validation. This method ensures
    /// symmetry: response shape is always picked from the inbound
    /// shape, never hardcoded.
    ///
    /// This is the framework's canonical HTTP envelope conversion
    /// primitive — useful when a single Lambda function routes among
    /// HTTP and one or more non-framework triggers (EventBridge,
    /// DynamoDB streams, custom invoke) and the adopter drives
    /// [`lambda_runtime::run`] with `serde_json::Value` directly.
    /// Handles the `LambdaRequest` deserialization, the text-vs-base64
    /// content-type detection, and the response envelope rebuild.
    ///
    /// Body encoding: text for `text/*`, `application/json`,
    /// `application/xml`, `application/javascript`, any `+json` / `+xml`
    /// structured-syntax suffix (RFC 6839, e.g. `application/a2a+json`),
    /// or anything with `charset=`; base64 otherwise.
    ///
    /// Path-prefix stripping (see
    /// [`LambdaA2aServerBuilder::strip_path_prefix`]) is applied
    /// because dispatch flows through [`Self::handle`]. Adopters who
    /// only serve HTTP should continue to use [`Self::run_http_only`]
    /// or [`LambdaA2aServerBuilder::run`]; this entry point exists for
    /// composition with adopter-owned third triggers.
    pub async fn handle_http_event_value(
        &self,
        value: serde_json::Value,
    ) -> Result<serde_json::Value, lambda_runtime::Error> {
        use base64::Engine;
        use http_body_util::BodyExt;

        let lambda_req: lambda_http::request::LambdaRequest = serde_json::from_value(value)
            .map_err(|e| lambda_runtime::Error::from(format!("invalid HTTP event: {e}")))?;

        // Capture the inbound shape so we can emit the matching
        // response envelope. Lost once we convert to `lambda_http::Request`.
        let shape = match &lambda_req {
            lambda_http::request::LambdaRequest::ApiGatewayV1(_) => HttpEventShape::ApiGatewayV1,
            lambda_http::request::LambdaRequest::ApiGatewayV2(_) => HttpEventShape::ApiGatewayV2,
            lambda_http::request::LambdaRequest::Alb(_) => HttpEventShape::Alb,
            lambda_http::request::LambdaRequest::WebSocket(_) => {
                return Err(lambda_runtime::Error::from(
                    "WebSocket Lambda events are not supported by this adapter",
                ));
            }
            _ => {
                return Err(lambda_runtime::Error::from(
                    "unsupported Lambda HTTP event variant",
                ));
            }
        };

        let req: lambda_http::Request = lambda_req.into();

        let resp = self
            .handle(req)
            .await
            .map_err(|e| lambda_runtime::Error::from(format!("handler error: {e}")))?;
        let (parts, body) = resp.into_parts();

        let bytes = body
            .collect()
            .await
            .map_err(|e| lambda_runtime::Error::from(format!("body collect: {e}")))?
            .to_bytes();

        let ct = parts
            .headers
            .get(http::header::CONTENT_TYPE)
            .and_then(|v| v.to_str().ok())
            .unwrap_or("application/octet-stream")
            .to_ascii_lowercase();
        let essence = ct.split(';').next().unwrap_or("").trim();
        let is_text = essence.starts_with("text/")
            || matches!(
                essence,
                "application/json" | "application/xml" | "application/javascript"
            )
            || essence.ends_with("+json")
            || essence.ends_with("+xml")
            || ct.contains("charset=");

        let (body_str, is_base64) = if bytes.is_empty() {
            (None, false)
        } else if is_text {
            match std::str::from_utf8(&bytes) {
                Ok(s) => (Some(s.to_string()), false),
                Err(_) => (
                    Some(base64::engine::general_purpose::STANDARD.encode(&bytes)),
                    true,
                ),
            }
        } else {
            (
                Some(base64::engine::general_purpose::STANDARD.encode(&bytes)),
                true,
            )
        };

        let status = parts.status.as_u16() as i64;
        let headers = parts.headers;
        let body_lambda = body_str.map(aws_lambda_events::encodings::Body::Text);

        match shape {
            HttpEventShape::ApiGatewayV1 => {
                let mut api_resp = aws_lambda_events::apigw::ApiGatewayProxyResponse::default();
                api_resp.status_code = status;
                api_resp.headers = headers;
                api_resp.body = body_lambda;
                api_resp.is_base64_encoded = is_base64;
                serde_json::to_value(api_resp).map_err(|e| {
                    lambda_runtime::Error::from(format!("serialise APIGW v1 response: {e}"))
                })
            }
            HttpEventShape::ApiGatewayV2 => {
                let mut api_resp = aws_lambda_events::apigw::ApiGatewayV2httpResponse::default();
                api_resp.status_code = status;
                api_resp.headers = headers;
                api_resp.body = body_lambda;
                api_resp.is_base64_encoded = is_base64;
                serde_json::to_value(api_resp).map_err(|e| {
                    lambda_runtime::Error::from(format!("serialise APIGW v2 response: {e}"))
                })
            }
            HttpEventShape::Alb => {
                let mut api_resp = aws_lambda_events::alb::AlbTargetGroupResponse::default();
                api_resp.status_code = status;
                api_resp.headers = headers;
                api_resp.body = body_lambda;
                api_resp.is_base64_encoded = is_base64;
                serde_json::to_value(api_resp).map_err(|e| {
                    lambda_runtime::Error::from(format!("serialise ALB response: {e}"))
                })
            }
        }
    }
}

/// Inbound HTTP event shape detected by [`LambdaA2aHandler::handle_http_event_value`].
///
/// The response envelope must match the inbound shape — see method
/// docs for the contract.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum HttpEventShape {
    ApiGatewayV1,
    ApiGatewayV2,
    Alb,
}

impl LambdaA2aHandler {
    /// Handle an SQS batch from the durable executor queue.
    ///
    /// Returns `SqsBatchResponse` whose `batch_item_failures` list
    /// tells the event source mapping which records to retry.
    /// Requires `FunctionResponseTypes: ["ReportBatchItemFailures"]`
    /// on the mapping.
    ///
    /// Each record: deserialize envelope → load task → terminal-no-op
    /// check → cancel-marker direct-CANCELED commit → run executor.
    /// See `durable::drive_sqs_batch` for the full semantics.
    #[cfg(feature = "sqs")]
    pub async fn handle_sqs(
        &self,
        event: aws_lambda_events::event::sqs::SqsEvent,
    ) -> aws_lambda_events::event::sqs::SqsBatchResponse {
        durable::drive_sqs_batch(&self.state, event).await
    }

    /// Run this handler as a pure HTTP Lambda. Strict: any non-HTTP
    /// event shape (SQS, DynamoDB stream, scheduler, …) causes
    /// `lambda_http::run` to return a deserialization error on
    /// invocation.
    ///
    /// Appropriate for Lambdas whose only trigger is a Function URL,
    /// API Gateway, or ALB. The request Lambda in the two-Lambda
    /// durable-executor topology is the canonical caller.
    pub async fn run_http_only(self) -> Result<(), lambda_runtime::Error> {
        let handler = Arc::new(self);
        lambda_http::run(lambda_http::service_fn(move |event| {
            let h = Arc::clone(&handler);
            async move { h.handle(event).await }
        }))
        .await
    }
}

// Default entrypoint when the `sqs` feature is off — `.run()` aliases
// the HTTP-only runner. With the `sqs` feature enabled, `.run()` lives
// in `durable.rs` and routes HTTP+SQS via the classifier.
#[cfg(not(feature = "sqs"))]
impl LambdaA2aHandler {
    /// Default Lambda runner. Dispatches based on the event shape the
    /// binary receives. Without the `sqs` feature this is HTTP-only
    /// (equivalent to [`LambdaA2aHandler::run_http_only`]).
    pub async fn run(self) -> Result<(), lambda_runtime::Error> {
        self.run_http_only().await
    }
}

impl LambdaA2aServerBuilder {
    /// Build and run in one call. Sugar for
    /// `self.build()?.run().await`. Hides the handler from the adopter —
    /// use `.build()` + `handler.run_*()` instead when the handler is
    /// needed for unit tests, custom middleware, or an explicit
    /// topology runner.
    pub async fn run(self) -> Result<(), lambda_runtime::Error> {
        let handler = self.build().map_err(|e| {
            lambda_runtime::Error::from(format!("LambdaA2aServerBuilder::run: {e}"))
        })?;
        handler.run().await
    }
}

#[cfg(test)]
mod handle_http_event_value_shape_tests {
    //! Regression tests for the event-shape symmetry contract on
    //! [`LambdaA2aHandler::handle_http_event_value`]. The wire-shape
    //! mismatch (v2 response to a v1 invocation) is what causes API
    //! Gateway REST integrations to return `502 Bad Gateway`. These
    //! tests pin the contract: response envelope must always match
    //! the inbound request envelope.

    use super::*;
    use std::sync::Arc;
    use turul_a2a::executor::{AgentExecutor, ExecutionContext};
    use turul_a2a::server::RuntimeConfig;
    use turul_a2a::server::in_flight::InFlightRegistry;
    use turul_a2a::storage::InMemoryA2aStorage;
    use turul_a2a::streaming::TaskEventBroker;
    use turul_a2a_types::{Message, Task};

    struct NoOpExecutor;

    #[async_trait::async_trait]
    impl AgentExecutor for NoOpExecutor {
        async fn execute(
            &self,
            task: &mut Task,
            _msg: &Message,
            _ctx: &ExecutionContext,
        ) -> Result<(), turul_a2a::error::A2aError> {
            let mut proto = task.as_proto().clone();
            proto.status = Some(turul_a2a_proto::TaskStatus {
                state: turul_a2a_proto::TaskState::Completed.into(),
                message: None,
                timestamp: None,
            });
            *task = Task::try_from(proto).unwrap();
            Ok(())
        }
        fn agent_card(&self) -> turul_a2a_proto::AgentCard {
            turul_a2a_proto::AgentCard {
                name: "test-agent".to_string(),
                ..Default::default()
            }
        }
    }

    fn build_handler() -> LambdaA2aHandler {
        let s = Arc::new(InMemoryA2aStorage::new());
        let state = AppState {
            executor: Arc::new(NoOpExecutor),
            task_storage: s.clone(),
            push_storage: s.clone(),
            event_store: s.clone(),
            atomic_store: s.clone(),
            event_broker: TaskEventBroker::new(),
            middleware_stack: Arc::new(MiddlewareStack::new(vec![])),
            runtime_config: RuntimeConfig::default(),
            in_flight: Arc::new(InFlightRegistry::new()),
            cancellation_supervisor: s.clone(),
            push_delivery_store: None,
            push_dispatcher: None,
            durable_executor_queue: None,
        };
        let router = build_router(state.clone());
        LambdaA2aHandler {
            router,
            state,
            path_prefix: None,
        }
    }

    /// API Gateway v2 / Function URL event. `requestContext.http.method`
    /// is what `LambdaRequest::ApiGatewayV2` keys off.
    fn apigw_v2_agent_card_get() -> serde_json::Value {
        serde_json::json!({
            "version": "2.0",
            "routeKey": "GET /.well-known/agent-card.json",
            "rawPath": "/.well-known/agent-card.json",
            "rawQueryString": "",
            "headers": {"accept": "application/json", "a2a-version": "1.0"},
            "requestContext": {
                "accountId": "000000000000",
                "apiId": "api",
                "domainName": "fake.execute-api",
                "domainPrefix": "fake",
                "http": {
                    "method": "GET",
                    "path": "/.well-known/agent-card.json",
                    "protocol": "HTTP/1.1",
                    "sourceIp": "127.0.0.1",
                    "userAgent": "test"
                },
                "requestId": "rid",
                "routeKey": "GET /.well-known/agent-card.json",
                "stage": "$default",
                "time": "01/Jan/2026:00:00:00 +0000",
                "timeEpoch": 1_735_689_600_000i64
            },
            "isBase64Encoded": false
        })
    }

    /// API Gateway v1 (REST) event. Top-level `httpMethod` + `path` and
    /// `requestContext.identity` (no `requestContext.http`) is what
    /// `LambdaRequest::ApiGatewayV1` keys off. This is the shape REST
    /// API + AWS_PROXY integrations send into the Lambda — the live
    /// fleet topology that 502'd against a hardcoded v2 response.
    ///
    /// `stage` is `$default` here so `lambda_http` does not prepend a
    /// stage prefix — this fixture exists to pin the response-shape
    /// contract, not the path-rewriting behaviour. See the
    /// `strip_path_prefix_*` tests for stage-prefix coverage.
    fn apigw_v1_agent_card_get_with_stage(stage: &str, path: &str) -> serde_json::Value {
        let request_context_path = if stage == "$default" {
            path.to_string()
        } else {
            format!("/{stage}{path}")
        };
        serde_json::json!({
            "resource": "/{proxy+}",
            "path": path,
            "httpMethod": "GET",
            "headers": {"accept": "application/json", "a2a-version": "1.0"},
            "multiValueHeaders": {},
            "queryStringParameters": null,
            "multiValueQueryStringParameters": null,
            "pathParameters": {"proxy": path.trim_start_matches('/').to_string()},
            "stageVariables": null,
            "requestContext": {
                "accountId": "000000000000",
                "apiId": "abc123",
                "domainName": "fake.execute-api",
                "domainPrefix": "fake",
                "extendedRequestId": "xrid",
                "httpMethod": "GET",
                "identity": {
                    "sourceIp": "127.0.0.1",
                    "userAgent": "test"
                },
                "path": request_context_path,
                "protocol": "HTTP/1.1",
                "requestId": "rid",
                "requestTime": "01/Jan/2026:00:00:00 +0000",
                "requestTimeEpoch": 1_735_689_600_000i64,
                "resourceId": "rsrc",
                "resourcePath": "/{proxy+}",
                "stage": stage
            },
            "body": null,
            "isBase64Encoded": false
        })
    }

    fn apigw_v1_agent_card_get() -> serde_json::Value {
        apigw_v1_agent_card_get_with_stage("$default", "/.well-known/agent-card.json")
    }

    /// ALB target-group event. `requestContext.elb` is what
    /// `LambdaRequest::Alb` keys off.
    fn alb_agent_card_get() -> serde_json::Value {
        serde_json::json!({
            "requestContext": {
                "elb": {
                    "targetGroupArn": "arn:aws:elasticloadbalancing:ap-southeast-2:000:targetgroup/fake/00"
                }
            },
            "httpMethod": "GET",
            "path": "/.well-known/agent-card.json",
            "queryStringParameters": {},
            "headers": {"accept": "application/json", "a2a-version": "1.0"},
            "body": "",
            "isBase64Encoded": false
        })
    }

    /// REST v1 inbound MUST produce a v1 response shape. This is the
    /// regression for the live-fleet 502.
    #[tokio::test]
    async fn v1_inbound_event_emits_v1_response_envelope() {
        let handler = build_handler();
        let resp_value = handler
            .handle_http_event_value(apigw_v1_agent_card_get())
            .await
            .expect("handle_http_event_value(v1) must succeed");

        // Parses cleanly as v1.
        let v1: aws_lambda_events::apigw::ApiGatewayProxyResponse =
            serde_json::from_value(resp_value.clone())
                .expect("response must deserialize as ApiGatewayProxyResponse (v1)");
        assert_eq!(v1.status_code, 200);
        match v1.body {
            Some(aws_lambda_events::encodings::Body::Text(s)) => {
                assert!(s.contains("\"name\""), "v1 body shape: {s}");
            }
            other => panic!("expected text body, got {other:?}"),
        }

        // Negative: v1 response JSON must NOT contain the v2-only
        // `cookies` field. Presence of `cookies` is API Gateway REST's
        // 502 trigger.
        assert!(
            resp_value.get("cookies").is_none(),
            "v1 response must not carry v2-only `cookies` field; payload: {resp_value}"
        );
    }

    /// API Gateway v2 / Function URL inbound MUST produce a v2 response
    /// shape (current behaviour, locked in to prevent the inverse
    /// regression).
    #[tokio::test]
    async fn v2_inbound_event_emits_v2_response_envelope() {
        let handler = build_handler();
        let resp_value = handler
            .handle_http_event_value(apigw_v2_agent_card_get())
            .await
            .expect("handle_http_event_value(v2) must succeed");

        let v2: aws_lambda_events::apigw::ApiGatewayV2httpResponse =
            serde_json::from_value(resp_value.clone())
                .expect("response must deserialize as ApiGatewayV2httpResponse (v2)");
        assert_eq!(v2.status_code, 200);
        match v2.body {
            Some(aws_lambda_events::encodings::Body::Text(s)) => {
                assert!(s.contains("\"name\""), "v2 body shape: {s}");
            }
            other => panic!("expected text body, got {other:?}"),
        }
    }

    /// ALB inbound MUST produce an ALB response shape.
    #[tokio::test]
    async fn alb_inbound_event_emits_alb_response_envelope() {
        let handler = build_handler();
        let resp_value = handler
            .handle_http_event_value(alb_agent_card_get())
            .await
            .expect("handle_http_event_value(alb) must succeed");

        let alb: aws_lambda_events::alb::AlbTargetGroupResponse =
            serde_json::from_value(resp_value)
                .expect("response must deserialize as AlbTargetGroupResponse");
        assert_eq!(alb.status_code, 200);
    }

    /// Live-fleet end-to-end regression: REST v1 stage-prefixed event
    /// + `strip_path_prefix("/dev")` on the handler routes through
    /// the agent-card endpoint AND emits a v1 response envelope. This
    /// is the exact scenario the live fleet was 502'ing on.
    #[tokio::test]
    async fn v1_stage_prefix_inbound_with_strip_emits_v1_response() {
        let s = Arc::new(InMemoryA2aStorage::new());
        let state = AppState {
            executor: Arc::new(NoOpExecutor),
            task_storage: s.clone(),
            push_storage: s.clone(),
            event_store: s.clone(),
            atomic_store: s.clone(),
            event_broker: TaskEventBroker::new(),
            middleware_stack: Arc::new(MiddlewareStack::new(vec![])),
            runtime_config: RuntimeConfig::default(),
            in_flight: Arc::new(InFlightRegistry::new()),
            cancellation_supervisor: s.clone(),
            push_delivery_store: None,
            push_dispatcher: None,
            durable_executor_queue: None,
        };
        let router = build_router(state.clone());
        let handler = LambdaA2aHandler {
            router,
            state,
            path_prefix: Some(Arc::from("/dev")),
        };

        let event = apigw_v1_agent_card_get_with_stage("dev", "/.well-known/agent-card.json");
        let resp_value = handler
            .handle_http_event_value(event)
            .await
            .expect("v1 stage-prefixed event must succeed under strip_path_prefix");

        let v1: aws_lambda_events::apigw::ApiGatewayProxyResponse =
            serde_json::from_value(resp_value)
                .expect("response must deserialize as ApiGatewayProxyResponse");
        assert_eq!(v1.status_code, 200);
        match v1.body {
            Some(aws_lambda_events::encodings::Body::Text(s)) => {
                assert!(s.contains("\"name\""), "v1 stage-prefix body: {s}");
            }
            other => panic!("expected text body, got {other:?}"),
        }
    }

    /// Cross-shape decoding regression: a v1 response JSON carries
    /// `multiValueHeaders` (REST shape) but no `cookies`; a v2 response
    /// carries `cookies` but no v1-only `multiValueHeaders` semantics.
    /// Confirms the two shapes are wire-distinguishable, which is why
    /// the gateway 502s on a mismatch.
    #[tokio::test]
    async fn v1_response_does_not_match_v2_envelope_strictly() {
        let handler = build_handler();
        let v1_value = handler
            .handle_http_event_value(apigw_v1_agent_card_get())
            .await
            .unwrap();

        // Sanity: v1 response JSON has no v2-only `cookies` field.
        assert!(v1_value.get("cookies").is_none());

        let v2_value = handler
            .handle_http_event_value(apigw_v2_agent_card_get())
            .await
            .unwrap();
        // Sanity: v2 response JSON carries the v2 `cookies` field
        // (default empty array — still present on the wire).
        assert!(
            v2_value.get("cookies").is_some(),
            "v2 response should carry the v2-only cookies field"
        );
    }
}