litellm_rs/server/middleware/
request_id.rs

1//! Request ID middleware
2
3use crate::utils::error::gateway_error::{GatewayError, with_error_response_request_id};
4use actix_web::body::{BoxBody, MessageBody};
5use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform, forward_ready};
6use actix_web::http::header::HeaderValue;
7use futures::future::{Ready, ready};
8use std::future::Future;
9use std::pin::Pin;
10use tracing::debug;
11use uuid::Uuid;
12
13/// Request ID middleware for Actix-web
14pub struct RequestIdMiddleware;
15
16impl<S, B> Transform<S, ServiceRequest> for RequestIdMiddleware
17where
18    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
19    S::Future: 'static,
20    B: MessageBody + 'static,
21{
22    type Response = ServiceResponse<BoxBody>;
23    type Error = actix_web::Error;
24    type InitError = ();
25    type Transform = RequestIdMiddlewareService<S>;
26    type Future = Ready<Result<Self::Transform, Self::InitError>>;
27
28    fn new_transform(&self, service: S) -> Self::Future {
29        ready(Ok(RequestIdMiddlewareService { service }))
30    }
31}
32
33/// Service implementation for request ID middleware
34pub struct RequestIdMiddlewareService<S> {
35    service: S,
36}
37
38impl<S, B> Service<ServiceRequest> for RequestIdMiddlewareService<S>
39where
40    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
41    S::Future: 'static,
42    B: MessageBody + 'static,
43{
44    type Response = ServiceResponse<BoxBody>;
45    type Error = actix_web::Error;
46    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
47
48    forward_ready!(service);
49
50    fn call(&self, mut req: ServiceRequest) -> Self::Future {
51        let existing = req
52            .headers()
53            .get("x-request-id")
54            .and_then(|value| value.to_str().ok())
55            .filter(|s| !s.is_empty())
56            .map(str::to_string);
57
58        let request_id = if let Some(id) = existing {
59            id
60        } else {
61            let id = Uuid::new_v4().to_string();
62            req.headers_mut().insert(
63                actix_web::http::header::HeaderName::from_static("x-request-id"),
64                HeaderValue::from_str(&id).unwrap_or_else(|_| HeaderValue::from_static("invalid")),
65            );
66            id
67        };
68
69        debug!("Processing request: {}", request_id);
70
71        let fut = self.service.call(req);
72        Box::pin(async move {
73            let scoped_request_id = request_id.clone();
74            with_error_response_request_id(scoped_request_id, async move {
75                let header_name = actix_web::http::header::HeaderName::from_static("x-request-id");
76                let header_value = HeaderValue::from_str(&request_id)
77                    .unwrap_or_else(|_| HeaderValue::from_static("invalid"));
78
79                match fut.await {
80                    Ok(mut res) => {
81                        res.headers_mut().insert(header_name, header_value);
82                        Ok(res.map_into_boxed_body())
83                    }
84                    Err(err) => {
85                        let mut response = if let Some(gateway_error) =
86                            err.as_error::<GatewayError>()
87                        {
88                            gateway_error.error_response_with_request_id(Some(request_id.clone()))
89                        } else {
90                            err.error_response()
91                        };
92                        response.headers_mut().insert(header_name, header_value);
93                        Err(actix_web::error::InternalError::from_response(err, response).into())
94                    }
95                }
96            })
97            .await
98        })
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use actix_web::{App, HttpResponse, http::StatusCode, test, web};
106    use serde_json::Value;
107
108    #[actix_web::test]
109    async fn middleware_does_not_clone_request_before_routing() {
110        let app = test::init_service(App::new().wrap(RequestIdMiddleware).route(
111            "/health",
112            web::get().to(|| async { HttpResponse::Ok().finish() }),
113        ))
114        .await;
115
116        let req = test::TestRequest::get().uri("/health").to_request();
117        let res = test::call_service(&app, req).await;
118
119        assert_eq!(res.status(), StatusCode::OK);
120        assert!(res.headers().contains_key("x-request-id"));
121    }
122
123    #[actix_web::test]
124    async fn middleware_adds_request_id_to_error_responses() {
125        let app = test::init_service(App::new().wrap(RequestIdMiddleware).route(
126            "/error",
127            web::get().to(|| async {
128                Err::<HttpResponse, _>(actix_web::error::ErrorInternalServerError("test error"))
129            }),
130        ))
131        .await;
132
133        let req = test::TestRequest::get().uri("/error").to_request();
134        let res = test::call_service(&app, req).await;
135
136        assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR);
137        assert!(res.headers().contains_key("x-request-id"));
138    }
139
140    #[actix_web::test]
141    async fn middleware_adds_request_id_to_gateway_error_body() {
142        let app = test::init_service(App::new().wrap(RequestIdMiddleware).route(
143            "/error",
144            web::get().to(|| async {
145                Err::<HttpResponse, actix_web::Error>(
146                    GatewayError::BadRequest("bad request".to_string()).into(),
147                )
148            }),
149        ))
150        .await;
151
152        let req = test::TestRequest::get()
153            .uri("/error")
154            .insert_header(("x-request-id", "req-test-123"))
155            .to_request();
156        let res = test::call_service(&app, req).await;
157
158        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
159        assert_eq!(
160            res.headers()
161                .get("x-request-id")
162                .and_then(|value| value.to_str().ok()),
163            Some("req-test-123")
164        );
165        let body = test::read_body(res).await;
166        let json: Value = serde_json::from_slice(&body).unwrap();
167        assert_eq!(json["error"]["request_id"], "req-test-123");
168    }
169}
litellm_rs/server/middleware/request_id.rs

litellm_rs/server/middleware/
request_id.rs
