fix: gRPC error handling — proper status codes and circuit breaker accuracy

model_gateway/src/routers/grpc/client.rs

@@ -224,22 +224,23 @@ impl GrpcClient {
 
     /// Generate streaming response from request
     ///
-    /// Dispatches to the appropriate backend client and wraps the result in ProtoStream
+    /// Dispatches to the appropriate backend client and wraps the result in ProtoStream.
+    /// Returns `tonic::Status` on error so callers can inspect the gRPC status code directly.
     pub async fn generate(
         &mut self,
         req: ProtoGenerateRequest,
-    ) -> Result<ProtoStream, Box<dyn std::error::Error + Send + Sync>> {
+    ) -> Result<ProtoStream, tonic::Status> {
         match (self, req) {
             (Self::Sglang(client), ProtoGenerateRequest::Sglang(boxed_req)) => {
-                let stream = client.generate(*boxed_req).await?;
+                let stream = client.generate(*boxed_req).await.map_err(into_status)?;
                 Ok(ProtoStream::Sglang(stream))
             }
             (Self::Vllm(client), ProtoGenerateRequest::Vllm(boxed_req)) => {
-                let stream = client.generate(*boxed_req).await?;
+                let stream = client.generate(*boxed_req).await.map_err(into_status)?;
                 Ok(ProtoStream::Vllm(stream))
             }
             (Self::Trtllm(client), ProtoGenerateRequest::Trtllm(boxed_req)) => {
-                let stream = client.generate(*boxed_req).await?;
+                let stream = client.generate(*boxed_req).await.map_err(into_status)?;
                 Ok(ProtoStream::Trtllm(stream))
             }
             #[expect(
@@ -253,10 +254,10 @@ impl GrpcClient {
     pub async fn embed(
         &mut self,
         req: ProtoEmbedRequest,
-    ) -> Result<ProtoEmbedResponse, Box<dyn std::error::Error + Send + Sync>> {
+    ) -> Result<ProtoEmbedResponse, tonic::Status> {
         match (self, req) {
             (Self::Sglang(client), ProtoEmbedRequest::Sglang(boxed_req)) => {
-                let resp = client.embed(*boxed_req).await?;
+                let resp = client.embed(*boxed_req).await.map_err(into_status)?;
                 Ok(ProtoEmbedResponse::Sglang(resp))
             }
             #[expect(
@@ -368,6 +369,15 @@ impl GrpcClient {
     }
 }
 
+/// Downcast a boxed error to `tonic::Status`, falling back to `Status::internal`.
+/// The grpc_client crate returns `Box<dyn Error>` but the concrete type is always `tonic::Status`.
+fn into_status(err: Box<dyn std::error::Error + Send + Sync>) -> tonic::Status {
+    match err.downcast::<tonic::Status>() {
+        Ok(status) => *status,
+        Err(other) => tonic::Status::internal(other.to_string()),
+    }
+}
+
 // ---------------------------------------------------------------------------
 // Metadata wrappers
 // ---------------------------------------------------------------------------

model_gateway/src/routers/grpc/common/stages/request_execution.rs

@@ -17,11 +17,12 @@ use crate::{
                 ProtoEmbedRequest, ProtoEmbedResponseVariant, ProtoGenerateRequest, ProtoRequest,
                 ProtoStream,
             },
+            tonic_ext::{TonicResultExt, TonicStatusExt},
         },
     },
 };
 
-type StreamResult = Result<ProtoStream, Box<dyn std::error::Error + Send + Sync>>;
+type StreamResult = Result<ProtoStream, tonic::Status>;
 
 /// Request execution stage: Execute gRPC requests (single or dual dispatch)
 pub(crate) struct RequestExecutionStage {
@@ -80,18 +81,9 @@ impl PipelineStage for RequestExecutionStage {
         ctx.state.load_guards = Some(LoadGuards::new(workers, ctx.input.headers.as_ref()));
 
         // Extract dispatch metadata for tracing span
-        let request_id = ctx
-            .state
-            .dispatch
-            .as_ref()
-            .map(|d| d.request_id.as_str())
-            .unwrap_or("unknown");
-        let model = ctx
-            .state
-            .dispatch
-            .as_ref()
-            .map(|d| d.model.as_str())
-            .unwrap_or("unknown");
+        let dispatch = ctx.state.dispatch.as_ref();
+        let request_id = dispatch.map(|d| d.request_id.as_str()).unwrap_or("unknown");
+        let model = dispatch.map(|d| d.model.as_str()).unwrap_or("unknown");
 
         // Create OTEL span for gRPC request execution
         let span = info_span!(
@@ -110,7 +102,8 @@ impl PipelineStage for RequestExecutionStage {
                         // Dispatch based on runtime type:
                         // - SGLang: parallel dual dispatch with bootstrap metadata
                         // - vLLM: sequential prefill-then-decode (NIXL handles KV transfer)
-                        match workers.pd_runtime_type() {
+                        let runtime_type = workers.pd_runtime_type();
+                        match runtime_type {
                             Some(RuntimeType::Vllm) => {
                                 self.execute_sequential_pd(req, clients, workers).await
                             }
@@ -120,7 +113,7 @@ impl PipelineStage for RequestExecutionStage {
                             Some(RuntimeType::Trtllm) | Some(RuntimeType::External) => {
                                 error!(
                                     function = "RequestExecutionStage::execute",
-                                    runtime_type = ?workers.pd_runtime_type(),
+                                    runtime_type = ?runtime_type,
                                     "Runtime does not support PD disaggregated mode"
                                 );
                                 Err(error::bad_request(
@@ -176,17 +169,11 @@ impl RequestExecutionStage {
         })?;
 
         let result = client.generate(proto_request).await;
-
-        // Record circuit breaker outcome
-        workers.record_outcome(result.is_ok());
+        workers.record_outcome(result.is_healthy());
 
         let stream = result.map_err(|e| {
-            error!(
-                function = "execute_single",
-                error = %e,
-                "Failed to start generation"
-            );
-            error::internal_error(
+            error!(function = "execute_single", error = %e, "Failed to start generation");
+            e.to_http_error(
                 "start_generation_failed",
                 format!("Failed to start generation: {e}"),
             )
@@ -212,12 +199,8 @@ impl RequestExecutionStage {
         })?;
 
         let response = client.embed(proto_request).await.map_err(|e| {
-            error!(
-                function = "execute_single_embed",
-                error = %e,
-                "Failed to start embedding"
-            );
-            error::internal_error(
+            error!(function = "execute_single_embed", error = %e, "Failed to start embedding");
+            e.to_http_error(
                 "start_embedding_failed",
                 format!("Failed to start embedding: {e}"),
             )
@@ -276,40 +259,23 @@ impl RequestExecutionStage {
             decode_client.generate(decode_request)
         );
 
-        // Record circuit breaker outcomes for each worker individually
-        workers.record_dual_outcomes(prefill_result.is_ok(), decode_result.is_ok());
+        // Record circuit breaker outcomes (client errors don't count as failures)
+        workers.record_dual_outcomes(prefill_result.is_healthy(), decode_result.is_healthy());
 
         // Handle prefill result
-        let prefill_stream = match prefill_result {
-            Ok(s) => s,
-            Err(e) => {
-                error!(
-                    function = "execute_dual_dispatch",
-                    error = %e,
-                    "Prefill worker failed to start"
-                );
-                return Err(error::internal_error(
-                    "prefill_worker_failed_to_start",
-                    format!("Prefill worker failed to start: {e}"),
-                ));
-            }
-        };
+        let prefill_stream = prefill_result.map_err(|e| {
+            error!(function = "execute_dual_dispatch", error = %e, "Prefill worker failed to start");
+            e.to_http_error("prefill_worker_failed_to_start", format!("Prefill worker failed to start: {e}"))
+        })?;
 
         // Handle decode result
-        let decode_stream = match decode_result {
-            Ok(s) => s,
-            Err(e) => {
-                error!(
-                    function = "execute_dual_dispatch",
-                    error = %e,
-                    "Decode worker failed to start"
-                );
-                return Err(error::internal_error(
-                    "decode_worker_failed_to_start",
-                    format!("Decode worker failed to start: {e}"),
-                ));
-            }
-        };
+        let decode_stream = decode_result.map_err(|e| {
+            error!(function = "execute_dual_dispatch", error = %e, "Decode worker failed to start");
+            e.to_http_error(
+                "decode_worker_failed_to_start",
+                format!("Decode worker failed to start: {e}"),
+            )
+        })?;
 
         Ok(ExecutionResult::Dual {
             prefill: prefill_stream,
@@ -391,16 +357,9 @@ impl RequestExecutionStage {
             .generate(prefill_request)
             .await
             .map_err(|e| {
-                workers.record_outcome_prefill(false);
-                error!(
-                    function = "execute_sequential_pd",
-                    error = %e,
-                    "Prefill worker failed to start"
-                );
-                error::internal_error(
-                    "prefill_worker_failed_to_start",
-                    format!("Prefill worker failed to start: {e}"),
-                )
+                workers.record_outcome_prefill(!e.is_cb_failure());
+                error!(function = "execute_sequential_pd", error = %e, "Prefill worker failed to start");
+                e.to_http_error("prefill_worker_failed_to_start", format!("Prefill worker failed to start: {e}"))
             })?;
 
         // Drain prefill response (we just need to wait for completion)
@@ -410,13 +369,9 @@ impl RequestExecutionStage {
                     // Just consume the response, we use bootstrap info from worker metadata
                 }
                 Err(e) => {
-                    workers.record_outcome_prefill(false);
-                    error!(
-                        function = "execute_sequential_pd",
-                        error = %e,
-                        "Prefill stream error"
-                    );
-                    return Err(error::internal_error(
+                    workers.record_outcome_prefill(!e.is_cb_failure());
+                    error!(function = "execute_sequential_pd", error = %e, "Prefill stream error");
+                    return Err(e.to_http_error(
                         "prefill_stream_error",
                         format!("Prefill stream error: {e}"),
                     ));
@@ -441,13 +396,9 @@ impl RequestExecutionStage {
 
         // Send request to decode
         let decode_stream = decode_client.generate(decode_request).await.map_err(|e| {
-            workers.record_outcome_decode(false);
-            error!(
-                function = "execute_sequential_pd",
-                error = %e,
-                "Decode worker failed to start"
-            );
-            error::internal_error(
+            workers.record_outcome_decode(!e.is_cb_failure());
+            error!(function = "execute_sequential_pd", error = %e, "Decode worker failed to start");
+            e.to_http_error(
                 "decode_worker_failed_to_start",
                 format!("Decode worker failed to start: {e}"),
             )

model_gateway/src/routers/grpc/mod.rs

@@ -12,6 +12,7 @@ pub(crate) mod pipeline;
 pub(crate) mod proto_wrapper;
 pub(crate) mod regular;
 pub(crate) mod router; // Used by routers/factory
+pub(crate) mod tonic_ext;
 pub mod utils; // Used by routers/http and bindings/golang
 
 // Re-export for convenience

model_gateway/src/routers/grpc/tonic_ext.rs

@@ -0,0 +1,67 @@
+//! Extension traits for tonic gRPC types.
+
+use axum::response::Response;
+use http::StatusCode;
+use tonic::Code;
+
+use crate::{core::is_retryable_status, routers::error};
+
+/// Extension methods for `tonic::Status`.
+pub(crate) trait TonicStatusExt {
+    /// Map gRPC status code to the corresponding HTTP status code.
+    fn http_status(&self) -> StatusCode;
+
+    /// Returns `true` if this error should trip the circuit breaker.
+    ///
+    /// Delegates to `is_retryable_status(http_status())` so the circuit breaker
+    /// uses the same predicate for both HTTP and gRPC paths. This covers
+    /// 408, 429, 500, 502, 503, 504 (i.e. ResourceExhausted, DeadlineExceeded,
+    /// Internal, Unavailable, Unknown, DataLoss, etc.).
+    fn is_cb_failure(&self) -> bool;
+
+    /// Convert this gRPC error into an HTTP error response with the appropriate status code.
+    fn to_http_error(&self, code: &str, msg: String) -> Response;
+}
+
+impl TonicStatusExt for tonic::Status {
+    fn http_status(&self) -> StatusCode {
+        match self.code() {
+            Code::Ok => StatusCode::OK,
+            Code::InvalidArgument
+            | Code::FailedPrecondition
+            | Code::OutOfRange
+            | Code::Cancelled => StatusCode::BAD_REQUEST,
+            Code::Unauthenticated => StatusCode::UNAUTHORIZED,
+            Code::PermissionDenied => StatusCode::FORBIDDEN,
+            Code::NotFound => StatusCode::NOT_FOUND,
+            Code::AlreadyExists | Code::Aborted => StatusCode::CONFLICT,
+            Code::ResourceExhausted => StatusCode::TOO_MANY_REQUESTS,
+            Code::Unavailable => StatusCode::SERVICE_UNAVAILABLE,
+            Code::DeadlineExceeded => StatusCode::GATEWAY_TIMEOUT,
+            Code::Unimplemented => StatusCode::NOT_IMPLEMENTED,
+            // Internal, Unknown, DataLoss
+            _ => StatusCode::INTERNAL_SERVER_ERROR,
+        }
+    }
+
+    fn is_cb_failure(&self) -> bool {
+        is_retryable_status(self.http_status())
+    }
+
+    fn to_http_error(&self, code: &str, msg: String) -> Response {
+        error::create_error(self.http_status(), code, msg)
+    }
+}
+
+/// Extension for `Result<T, tonic::Status>` to check circuit breaker health.
+pub(crate) trait TonicResultExt {
+    /// Returns `true` if the result is healthy for the circuit breaker.
+    /// `Ok` and client-error results are healthy; only retryable errors are failures.
+    fn is_healthy(&self) -> bool;
+}
+
+impl<T> TonicResultExt for Result<T, tonic::Status> {
+    fn is_healthy(&self) -> bool {
+        self.as_ref().map_or_else(|e| !e.is_cb_failure(), |_| true)
+    }
+}

model_gateway/src/routers/grpc/utils.rs

@@ -39,7 +39,10 @@ use super::{
 use crate::{
     core::Worker,
     observability::metrics::metrics_labels,
-    routers::{error, grpc::proto_wrapper::ProtoResponseVariant},
+    routers::{
+        error,
+        grpc::{proto_wrapper::ProtoResponseVariant, tonic_ext::TonicStatusExt},
+    },
 };
 
 /// Type alias for the SSE channel sender used across streaming endpoints.
@@ -662,7 +665,9 @@ pub(crate) async fn collect_stream_responses(
                         all_responses.push(complete);
                     }
                     ProtoResponseVariant::Error(err) => {
-                        error!(function = "collect_stream_responses", worker = %worker_name, error = %err.message(), "Worker generation error");
+                        // In-band error (legacy): backends should use context.abort() instead.
+                        // Kept for backward compatibility during the transition.
+                        warn!(function = "collect_stream_responses", worker = %worker_name, error = %err.message(), "Worker sent in-band error (legacy path, backend should use context.abort)");
                         // Don't mark as completed - let Drop send abort for error cases
                         return Err(error::internal_error(
                             "worker_generation_failed",
@@ -678,9 +683,9 @@ pub(crate) async fn collect_stream_responses(
                 }
             }
             Err(e) => {
-                error!(function = "collect_stream_responses", worker = %worker_name, error = ?e, "Worker stream error");
+                error!(function = "collect_stream_responses", worker = %worker_name, grpc_code = ?e.code(), error = ?e, "Worker stream error");
                 // Don't mark as completed - let Drop send abort for error cases
-                return Err(error::internal_error(
+                return Err(e.to_http_error(
                     "worker_stream_failed",
                     format!("{worker_name} stream failed: {e}"),
                 ));

model_gateway/src/routers/http/router.rs

@@ -317,7 +317,7 @@ impl Router {
         events::RequestReceivedEvent {}.emit();
 
         let status = response.status();
-        worker.record_outcome(status.is_success());
+        worker.record_outcome(!is_retryable_status(status));
 
         // Record worker errors for server errors (5xx)
         if status.is_server_error() {