chore: added middleware layer to catch json validation errors (#3182)

Signed-off-by: ayushag <ayushag@nvidia.com>

chore: added middleware layer to catch json validation errors (#3182)
Signed-off-by: ayushag <ayushag@nvidia.com>
4c56c8ae · Ayush Agarwal · GitHub · 37bc8444 · 4c56c8ae · 4c56c8ae
Unverified Commit 4c56c8ae authored Sep 23, 2025 by Ayush Agarwal Committed by GitHub Sep 23, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 34 additions and 7 deletions

lib/llm/src/http/service/openai.rs lib/llm/src/http/service/openai.rs +33 -1

lib/llm/tests/http-service.rs lib/llm/tests/http-service.rs +1 -6

No files found.
--- a/lib/llm/src/http/service/openai.rs
+++ b/lib/llm/src/http/service/openai.rs
@@ -9,8 +9,11 @@ use std::{

 use axum::{
    Json, Router,
+    body::Body,
    extract::State,
+    http::Request,
    http::{HeaderMap, StatusCode},
+    middleware::{self, Next},
    response::{
        IntoResponse, Response,
        sse::{KeepAlive, Sse},
@@ -65,7 +68,7 @@ fn get_body_limit() -> usize {

 pub type ErrorResponse = (StatusCode, Json<ErrorMessage>);

-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Debug)]
 pub(crate) struct ErrorMessage {
    error: String,
 }
@@ -165,6 +168,31 @@ impl From<HttpError> for ErrorMessage {
    }
 }

+// Problem: Currently we are using JSON from axum as the request validator. Whenever there is an invalid JSON, it will return a 422.
+// But all the downstream apps that relies on openai based APIs, expects to get 400 for all these cases otherwise they fail badly
+// Solution: Intercept the response from handlers and convert ANY 422 status codes to 400 with the actual error message.
+pub async fn smart_json_error_middleware(request: Request<Body>, next: Next) -> Response {
+    let response = next.run(request).await;
+
+    if response.status() == StatusCode::UNPROCESSABLE_ENTITY {
+        let (_parts, body) = response.into_parts();
+        let body_bytes = axum::body::to_bytes(body, usize::MAX)
+            .await
+            .unwrap_or_default();
+        let error_message = String::from_utf8_lossy(&body_bytes).to_string();
+        (
+            StatusCode::BAD_REQUEST,
+            Json(ErrorMessage {
+                error: error_message,
+            }),
+        )
+            .into_response()
+    } else {
+        // Pass through if it is not a 422
+        response
+    }
+}
+
 /// Get the request ID from a primary source, or next from the headers, or lastly create a new one if not present
 // TODO: Similar function exists in lib/llm/src/grpc/service/openai.rs but with different signature and simpler logic
 fn get_or_create_request_id(primary: Option<&str>, headers: &HeaderMap) -> String {
@@ -1054,6 +1082,7 @@ pub fn completions_router(
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
        .route(&path, post(handler_completions))
+        .layer(middleware::from_fn(smart_json_error_middleware))
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
        .with_state(state);
    (vec![doc], router)
@@ -1070,6 +1099,7 @@ pub fn chat_completions_router(
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
        .route(&path, post(handler_chat_completions))
+        .layer(middleware::from_fn(smart_json_error_middleware))
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
        .with_state((state, template));
    (vec![doc], router)
@@ -1085,6 +1115,7 @@ pub fn embeddings_router(
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
        .route(&path, post(embeddings))
+        .layer(middleware::from_fn(smart_json_error_middleware))
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
        .with_state(state);
    (vec![doc], router)
@@ -1117,6 +1148,7 @@ pub fn responses_router(
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
        .route(&path, post(handler_responses))
+        .layer(middleware::from_fn(smart_json_error_middleware))
        .with_state((state, template));
    (vec![doc], router)
 }

--- a/lib/llm/tests/http-service.rs
+++ b/lib/llm/tests/http-service.rs
@@ -531,12 +531,7 @@ async fn test_http_service() {
        .await
        .unwrap();

-    assert_eq!(
-        response.status(),
-        StatusCode::UNPROCESSABLE_ENTITY,
-        "{:?}",
-        response
-    );
+    assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{:?}", response);

    // =========== Query /metrics endpoint ===========
    let response = client