[model-gateway] add retry support to OpenAI router chat endpoint (#15589)

sgl-model-gateway/src/routers/openai/router.rs

@@ -34,7 +34,11 @@ use super::{
 };
 use crate::{
     app_context::AppContext,
-    core::{model_type::Endpoint, ModelCard, ProviderType, RuntimeType, Worker, WorkerRegistry},
+    config::types::RetryConfig,
+    core::{
+        is_retryable_status, model_type::Endpoint, ModelCard, ProviderType, RetryExecutor,
+        RuntimeType, Worker, WorkerRegistry,
+    },
     data_connector::{ConversationId, ListParams, ResponseId, SortOrder},
     observability::metrics::{bool_to_static_str, metrics_labels, Metrics},
     protocols::{
@@ -53,6 +57,7 @@ pub struct OpenAIRouter {
     healthy: AtomicBool,
     shared_components: Arc<SharedComponents>,
     responses_components: Arc<ResponsesComponents>,
+    retry_config: RetryConfig,
 }
 
 impl std::fmt::Debug for OpenAIRouter {
@@ -176,6 +181,7 @@ impl OpenAIRouter {
             healthy: AtomicBool::new(true),
             shared_components,
             responses_components,
+            retry_config: ctx.router_config.effective_retry_config(),
         })
     }
 
@@ -659,77 +665,129 @@ impl crate::routers::RouterTrait for OpenAIRouter {
             previous_response_id: None,
         });
 
+        // Wrap values in Arc to avoid cloning large objects on each retry attempt
         let payload_ref = ctx.payload().expect("Payload not prepared");
-        let mut req = ctx.components.client().post(&url).json(&payload_ref.json);
+        let payload_json = Arc::new(payload_ref.json.clone());
-        let auth_header = extract_auth_header(ctx.headers(), worker.api_key());
+        let client = ctx.components.client().clone();
-        req = apply_provider_headers(req, &url, auth_header.as_ref());
+        let headers_cloned = Arc::new(ctx.headers().cloned());
+        let worker_api_key = Arc::new(worker.api_key().clone());
+        let is_streaming = ctx.is_streaming();
 
-        if ctx.is_streaming() {
+        let response = RetryExecutor::execute_response_with_retry(
-            req = req.header("Accept", "text/event-stream");
+            &self.retry_config,
-        }
+            |_attempt| {
+                let client = client.clone();
+                let url = url.clone();
+                let payload = Arc::clone(&payload_json);
+                let headers = Arc::clone(&headers_cloned);
+                let worker_api_key = Arc::clone(&worker_api_key);
+                let worker = Arc::clone(&worker);
 
-        let resp = match req.send().await {
+                async move {
-            Ok(r) => r,
+                    let mut req = client.post(&url).json(&*payload);
-            Err(e) => {
+                    let auth_header = extract_auth_header((*headers).as_ref(), &worker_api_key);
-                worker.circuit_breaker().record_failure();
+                    req = apply_provider_headers(req, &url, auth_header.as_ref());
-                Metrics::record_router_error(
-                    metrics_labels::ROUTER_OPENAI,
-                    metrics_labels::BACKEND_EXTERNAL,
-                    metrics_labels::CONNECTION_HTTP,
-                    model,
-                    metrics_labels::ENDPOINT_CHAT,
-                    metrics_labels::ERROR_BACKEND,
-                );
-                return (
-                    StatusCode::SERVICE_UNAVAILABLE,
-                    format!("Failed to contact upstream: {}", e),
-                )
-                    .into_response();
-            }
-        };
 
-        let status = StatusCode::from_u16(resp.status().as_u16())
+                    if is_streaming {
-            .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
+                        req = req.header("Accept", "text/event-stream");
+                    }
 
-        if !ctx.is_streaming() {
+                    let resp = match req.send().await {
-            let content_type = resp.headers().get(CONTENT_TYPE).cloned();
+                        Ok(r) => r,
-            match resp.bytes().await {
+                        Err(e) => {
-                Ok(body) => {
+                            worker.circuit_breaker().record_failure();
-                    worker.circuit_breaker().record_success();
+                            return (
-                    Metrics::record_router_duration(
+                                StatusCode::SERVICE_UNAVAILABLE,
-                        metrics_labels::ROUTER_OPENAI,
+                                format!("Failed to contact upstream: {}", e),
-                        metrics_labels::BACKEND_EXTERNAL,
+                            )
-                        metrics_labels::CONNECTION_HTTP,
+                                .into_response();
-                        model,
+                        }
-                        metrics_labels::ENDPOINT_CHAT,
+                    };
-                        start.elapsed(),
+
-                    );
+                    let status = StatusCode::from_u16(resp.status().as_u16())
-                    let mut response = Response::new(Body::from(body));
+                        .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
-                    *response.status_mut() = status;
+
-                    if let Some(ct) = content_type {
+                    // Record circuit breaker failure for error status codes
-                        response.headers_mut().insert(CONTENT_TYPE, ct);
+                    if !status.is_success() {
+                        worker.circuit_breaker().record_failure();
+                    }
+
+                    if !is_streaming {
+                        let content_type = resp.headers().get(CONTENT_TYPE).cloned();
+                        match resp.bytes().await {
+                            Ok(body) => {
+                                // Only record success after body is fully read
+                                if status.is_success() {
+                                    worker.circuit_breaker().record_success();
+                                }
+                                let mut response = Response::new(Body::from(body));
+                                *response.status_mut() = status;
+                                if let Some(ct) = content_type {
+                                    response.headers_mut().insert(CONTENT_TYPE, ct);
+                                }
+                                response
+                            }
+                            Err(e) => {
+                                worker.circuit_breaker().record_failure();
+                                (
+                                    StatusCode::INTERNAL_SERVER_ERROR,
+                                    format!("Failed to read response: {}", e),
+                                )
+                                    .into_response()
+                            }
+                        }
+                    } else {
+                        // Streaming response - record success when stream starts
+                        if status.is_success() {
+                            worker.circuit_breaker().record_success();
+                        }
+                        let stream = resp.bytes_stream();
+                        let (tx, rx) = mpsc::unbounded_channel();
+                        tokio::spawn(async move {
+                            let mut s = stream;
+                            while let Some(chunk) = s.next().await {
+                                match chunk {
+                                    Ok(bytes) => {
+                                        if tx.send(Ok(bytes)).is_err() {
+                                            break;
+                                        }
+                                    }
+                                    Err(e) => {
+                                        let _ = tx.send(Err(format!("Stream error: {}", e)));
+                                        break;
+                                    }
+                                }
+                            }
+                        });
+                        let mut response =
+                            Response::new(Body::from_stream(UnboundedReceiverStream::new(rx)));
+                        *response.status_mut() = status;
+                        response
+                            .headers_mut()
+                            .insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));
+                        response
                     }
-                    response
                 }
-                Err(e) => {
+            },
-                    worker.circuit_breaker().record_failure();
+            |res, _attempt| is_retryable_status(res.status()),
-                    Metrics::record_router_error(
+            |delay, attempt| {
-                        metrics_labels::ROUTER_OPENAI,
+                Metrics::record_worker_retry(
-                        metrics_labels::BACKEND_EXTERNAL,
+                    metrics_labels::BACKEND_EXTERNAL,
-                        metrics_labels::CONNECTION_HTTP,
+                    metrics_labels::ENDPOINT_CHAT,
-                        model,
+                );
-                        metrics_labels::ENDPOINT_CHAT,
+                Metrics::record_worker_retry_backoff(attempt, delay);
-                        metrics_labels::ERROR_BACKEND,
+            },
-                    );
+            || {
-                    (
+                Metrics::record_worker_retries_exhausted(
-                        StatusCode::INTERNAL_SERVER_ERROR,
+                    metrics_labels::BACKEND_EXTERNAL,
-                        format!("Failed to read response: {}", e),
+                    metrics_labels::ENDPOINT_CHAT,
-                    )
+                );
-                        .into_response()
+            },
-                }
+        )
-            }
+        .await;
-        } else {
+
-            // For streaming, record duration at start since we can't track completion
+        // Record duration/error metrics after retry completes
+        if response.status().is_success() {
             Metrics::record_router_duration(
                 metrics_labels::ROUTER_OPENAI,
                 metrics_labels::BACKEND_EXTERNAL,
@@ -738,31 +796,18 @@ impl crate::routers::RouterTrait for OpenAIRouter {
                 metrics_labels::ENDPOINT_CHAT,
                 start.elapsed(),
             );
-            let stream = resp.bytes_stream();
+        } else {
-            let (tx, rx) = mpsc::unbounded_channel();
+            Metrics::record_router_error(
-            tokio::spawn(async move {
+                metrics_labels::ROUTER_OPENAI,
-                let mut s = stream;
+                metrics_labels::BACKEND_EXTERNAL,
-                while let Some(chunk) = s.next().await {
+                metrics_labels::CONNECTION_HTTP,
-                    match chunk {
+                model,
-                        Ok(bytes) => {
+                metrics_labels::ENDPOINT_CHAT,
-                            if tx.send(Ok(bytes)).is_err() {
+                metrics_labels::ERROR_BACKEND,
-                                break;
+            );
-                            }
-                        }
-                        Err(e) => {
-                            let _ = tx.send(Err(format!("Stream error: {}", e)));
-                            break;
-                        }
-                    }
-                }
-            });
-            let mut response = Response::new(Body::from_stream(UnboundedReceiverStream::new(rx)));
-            *response.status_mut() = status;
-            response
-                .headers_mut()
-                .insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));
-            response
         }
+
+        response
     }
 
     async fn route_responses(