The Complete Production SOP: How to Deploy and Run Tower Composable Asynchronous Services & Middleware Layers in Autonomous AI Agents

use std::task::{Context, Poll};
use std::time::Duration;
use tower::{Layer, Service};
use futures_util::future::BoxFuture;

#[derive(Clone)]
pub struct TimeoutLayer {
    timeout: Duration,
}

impl TimeoutLayer {
    pub fn new(timeout: Duration) -> Self {
        Self { timeout }
    }
}

impl<S> Layer<S> for TimeoutLayer {
    type Service = TimeoutMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        TimeoutMiddleware {
            inner,
            timeout: self.timeout,
        }
    }
}

#[derive(Clone)]
pub struct TimeoutMiddleware<S> {
    inner: S,
    timeout: Duration,
}

impl<S, Req> Service<Req> for TimeoutMiddleware<S>
where
    S: Service<Req> + Clone + Send + 'static,
    S::Future: Send + 'static,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>> + 'static,
    Req: Send + 'static,
{
    type Response = S::Response;
    type Error = Box<dyn std::error::Error + Send + Sync>;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx).map_err(Into::into)
    }

    fn call(&mut self, req: Req) -> Self::Future {
        let fut = self.inner.call(req);
        let sleep_dur = self.timeout;

        Box::pin(async move {
            tokio::select! {
                res = fut => res.map_err(Into::into),
                _ = tokio::time::sleep(sleep_dur) => {
                    Err("request timed out".into())
                }
            }
        })
    }
}

[dependencies]
tower = { version = "0.4", features = ["util", "buffer", "limit", "timeout"] }
futures-util = "0.3"
tokio = { version = "1", features = ["full"] }

Implement tower service trait for an API Key validation layer that inspects an incoming http::Request, checking the 'x-api-key' header. Use ServiceBuilder to compose layers with a 5-second timeout and 100 req/sec rate limit.

#[tokio::test]
async fn test_service_composition() {
    use tower::{Service, ServiceExt, ServiceBuilder};
    use std::time::Duration;

    let service = tower::service_fn(|req: String| async move {
        Ok::<_, std::io::Error>(format!("echo: {}", req))
    });

    let mut composed = ServiceBuilder::new()
        .timeout(Duration::from_secs(1))
        .service(service);

    let response = composed.ready().await.unwrap().call("hello".to_string()).await;
    assert!(response.is_ok());
}