Repository: tensorzero/tensorzero (language: rust)

Issue:
Add a simple 404 handler to the gateway
Currently, the gateway returns 404 with no body. It also doesn't log anything.

To simplify debugging, we could add a simple 404 handler that generates a `crate::errors::Error` (i.e. log + message).


Discussion / hints:
@GabrielBianconi Can I take up this one?
Sounds good, thanks @ChetanXpro!

Fix patch:
diff --git a/gateway/src/endpoints/fallback.rs b/gateway/src/endpoints/fallback.rs
new file mode 100644
index 000000000..ed955ca7a
--- /dev/null
+++ b/gateway/src/endpoints/fallback.rs
@@ -0,0 +1,70 @@
+use crate::error::{Error, ErrorDetails};
+use axum::{
+    body::Body,
+    http::Request,
+    response::{IntoResponse, Response},
+};
+
+pub async fn handle_404(req: Request<Body>) -> Response {
+    let path = req.uri().path().to_string();
+    let method = req.method().to_string();
+
+    Error::new(ErrorDetails::RouteNotFound { path, method }).into_response()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use axum::{
+        body::to_bytes,
+        http::{Method, StatusCode, Uri},
+    };
+    use serde_json::Value;
+
+    #[tokio::test]
+    async fn test_handle_404_get() {
+        let req = Request::builder()
+            .method(Method::GET)
+            .uri(Uri::from_static("/unknown/path"))
+            .body(Body::empty())
+            .unwrap();
+
+        let response = handle_404(req).await;
+
+        assert_eq!(response.status(), StatusCode::NOT_FOUND);
+
+        let body_bytes = to_bytes(response.into_body(), 1024).await.unwrap();
+        let body: Value = serde_json::from_slice(&body_bytes).unwrap();
+
+        let error_msg = body.get("error").and_then(Value::as_str).unwrap();
+        assert!(error_msg.contains("GET"));
+        assert!(error_msg.contains("/unknown/path"));
+    }
+
+    #[tokio::test]
+    async fn test_handle_404_post() {
+        let json_body = serde_json::json!({
+            "message": "Hello world",
+            "number": 42,
+            "active": true
+        });
+
+        let req = Request::builder()
+            .method(Method::POST)
+            .uri(Uri::from_static("/unknown/path"))
+            .header("content-type", "application/json")
+            .body(Body::from(serde_json::to_string(&json_body).unwrap()))
+            .unwrap();
+
+        let response = handle_404(req).await;
+
+        assert_eq!(response.status(), StatusCode::NOT_FOUND);
+
+        let body_bytes = to_bytes(response.into_body(), 1024).await.unwrap();
+        let body: Value = serde_json::from_slice(&body_bytes).unwrap();
+
+        let error_msg = body.get("error").and_then(Value::as_str).unwrap();
+        assert!(error_msg.contains("POST"));
+        assert!(error_msg.contains("/unknown/path"));
+    }
+}
diff --git a/gateway/src/endpoints/mod.rs b/gateway/src/endpoints/mod.rs
index c049fb811..fbeb5af78 100644
--- a/gateway/src/endpoints/mod.rs
+++ b/gateway/src/endpoints/mod.rs
@@ -1,4 +1,5 @@
 pub mod batch_inference;
+pub mod fallback;
 pub mod feedback;
 pub mod inference;
 pub mod openai_compatible;
diff --git a/gateway/src/error.rs b/gateway/src/error.rs
index a75d884a2..bc0a29247 100644
--- a/gateway/src/error.rs
+++ b/gateway/src/error.rs
@@ -219,6 +219,10 @@ pub enum ErrorDetails {
     UuidInFuture {
         raw_uuid: String,
     },
+    RouteNotFound {
+        path: String,
+        method: String,
+    },
 }
 
 impl ErrorDetails {
@@ -280,6 +284,7 @@ impl ErrorDetails {
             ErrorDetails::UnsupportedModelProviderForBatchInference { .. } => tracing::Level::WARN,
             ErrorDetails::UnsupportedVariantForBatchInference { .. } => tracing::Level::WARN,
             ErrorDetails::UuidInFuture { .. } => tracing::Level::WARN,
+            ErrorDetails::RouteNotFound { .. } => tracing::Level::WARN,
         }
     }
 
@@ -345,6 +350,7 @@ impl ErrorDetails {
             }
             ErrorDetails::UnsupportedVariantForBatchInference { .. } => StatusCode::BAD_REQUEST,
             ErrorDetails::UuidInFuture { .. } => StatusCode::BAD_REQUEST,
+            ErrorDetails::RouteNotFound { .. } => StatusCode::NOT_FOUND,
         }
     }
 
@@ -570,6 +576,9 @@ impl std::fmt::Display for ErrorDetails {
             ErrorDetails::UuidInFuture { raw_uuid } => {
                 write!(f, "UUID is in the future: {}", raw_uuid)
             }
+            ErrorDetails::RouteNotFound { path, method } => {
+                write!(f, "Route not found: {} {}", method, path)
+            }
         }
     }
 }
diff --git a/gateway/src/main.rs b/gateway/src/main.rs
index 7da8b680a..e3a0a2328 100644
--- a/gateway/src/main.rs
+++ b/gateway/src/main.rs
@@ -53,6 +53,7 @@ async fn main() {
             "/metrics",
             get(move || std::future::ready(metrics_handle.render())),
         )
+        .fallback(endpoints::fallback::handle_404)
         .with_state(app_state);
 
     // Bind to the socket address specified in the config, or default to 0.0.0.0:3000


Test patch:
diff --git a/gateway/tests/e2e/fallback.rs b/gateway/tests/e2e/fallback.rs
new file mode 100644
index 000000000..448a0a928
--- /dev/null
+++ b/gateway/tests/e2e/fallback.rs
@@ -0,0 +1,51 @@
+use reqwest::Client;
+use serde_json::Value;
+
+use crate::common::get_gateway_endpoint;
+
+#[tokio::test]
+async fn test_404_get() {
+    let client = Client::new();
+
+    // Make GET request to non-existent endpoint
+    let response = client
+        .get(get_gateway_endpoint("/non/existent/path"))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(response.status().as_u16(), 404);
+
+    let response_json = response.json::<Value>().await.unwrap();
+    let error_msg = response_json.get("error").and_then(Value::as_str).unwrap();
+
+    assert!(error_msg.contains("GET"));
+    assert!(error_msg.contains("/non/existent/path"));
+}
+
+#[tokio::test]
+async fn test_404_post() {
+    let client = Client::new();
+
+    // Make POST request to non-existent endpoint with JSON body
+    let json_body = serde_json::json!({
+        "message": "Hello world",
+        "number": 42,
+        "active": true
+    });
+
+    let response = client
+        .post(get_gateway_endpoint("/non/existent/path"))
+        .json(&json_body)
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(response.status().as_u16(), 404);
+
+    let response_json = response.json::<Value>().await.unwrap();
+    let error_msg = response_json.get("error").and_then(Value::as_str).unwrap();
+
+    assert!(error_msg.contains("POST"));
+    assert!(error_msg.contains("/non/existent/path"));
+}
diff --git a/gateway/tests/e2e/tests.rs b/gateway/tests/e2e/tests.rs
index aeda5f9e6..60eac55f4 100644
--- a/gateway/tests/e2e/tests.rs
+++ b/gateway/tests/e2e/tests.rs
@@ -3,6 +3,7 @@ mod best_of_n;
 mod clickhouse;
 mod common;
 mod dicl;
+mod fallback;
 mod feedback;
 mod health;
 mod inference;