Add error code in prometheus metrics and add X-SMG-Error-Code header (#15036)

sgl-model-gateway/src/middleware.rs

@@ -27,6 +27,7 @@ use tracing::{debug, error, field::Empty, info, info_span, warn, Span};
 pub use crate::core::token_bucket::TokenBucket;
 use crate::{
     observability::metrics::RouterMetrics,
+    routers::error::extract_error_code_from_response,
     server::AppState,
     wasm::{
         module::{MiddlewareAttachPoint, WasmModuleAttachPoint},
@@ -336,8 +337,10 @@ impl<B> OnResponse<B> for ResponseLogger {
         let status = response.status();
         let status_code = status.as_u16();
 
+        let error_code = extract_error_code_from_response(response);
+
         // TODO support `route` information
-        RouterMetrics::record_http_status_code(status_code);
+        RouterMetrics::record_http_status_code(status_code, error_code);
         RouterMetrics::record_request_duration(latency);
 
         // Record these in the span for structured logging/observability tools

sgl-model-gateway/src/observability/metrics.rs

@@ -275,7 +275,7 @@ pub fn init_metrics() {
     );
     describe_counter!(
         "sgl_router_http_responses_total",
-        "Total number of HTTP responses by status code"
+        "Total number of HTTP responses by status code and error code"
     );
 }
 
@@ -593,9 +593,10 @@ impl RouterMetrics {
         counter!("sgl_router_http_requests_total").increment(1);
     }
 
-    pub fn record_http_status_code(status_code: u16) {
+    pub fn record_http_status_code(status_code: u16, error_code: &str) {
         counter!("sgl_router_http_responses_total",
-            "status_code" => status_code.to_string()
+            "status_code" => status_code.to_string(),
+            "error_code" => error_code.to_string()
         )
         .increment(1);
     }

sgl-model-gateway/src/routers/error.rs

@@ -1,10 +1,12 @@
 use axum::{
-    http::StatusCode,
+    http::{HeaderMap, HeaderValue, StatusCode},
     response::{IntoResponse, Response},
     Json,
 };
 use serde_json::json;
 
+pub const HEADER_X_SMG_ERROR_CODE: &str = "X-SMG-Error-Code";
+
 pub fn internal_error(code: impl Into<String>, message: impl Into<String>) -> Response {
     create_error(StatusCode::INTERNAL_SERVER_ERROR, code, message)
 }
@@ -42,13 +44,23 @@ pub fn create_error(
     code: impl Into<String>,
     message: impl Into<String>,
 ) -> Response {
+    let code_str = code.into();
+    let message_str = message.into();
+
+    let mut headers = HeaderMap::new();
+    headers.insert(
+        HEADER_X_SMG_ERROR_CODE,
+        HeaderValue::from_str(&code_str).unwrap(),
+    );
+
     (
         status,
+        headers,
         Json(json!({
             "error": {
-                "message": message.into(),
                 "type": status_code_to_str(status),
-                "code": code.into(),
+                "code": code_str,
+                "message": message_str,
             }
         })),
     )
@@ -132,6 +144,14 @@ fn status_code_to_str(status_code: StatusCode) -> &'static str {
     }
 }
 
+pub fn extract_error_code_from_response<B>(response: &Response<B>) -> &str {
+    response
+        .headers()
+        .get(HEADER_X_SMG_ERROR_CODE)
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or_default()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;