[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,12 +665,30 @@ 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();
+        let headers_cloned = Arc::new(ctx.headers().cloned());
+        let worker_api_key = Arc::new(worker.api_key().clone());
+        let is_streaming = ctx.is_streaming();
+
+        let response = RetryExecutor::execute_response_with_retry(
+            &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);
+
+                async move {
+                    let mut req = client.post(&url).json(&*payload);
+                    let auth_header = extract_auth_header((*headers).as_ref(), &worker_api_key);
                     req = apply_provider_headers(req, &url, auth_header.as_ref());
 
-        if ctx.is_streaming() {
+                    if is_streaming {
                         req = req.header("Accept", "text/event-stream");
                     }
 
@@ -672,14 +696,6 @@ impl crate::routers::RouterTrait for OpenAIRouter {
                         Ok(r) => r,
                         Err(e) => {
                             worker.circuit_breaker().record_failure();
-                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),
@@ -691,19 +707,19 @@ impl crate::routers::RouterTrait for OpenAIRouter {
                     let status = StatusCode::from_u16(resp.status().as_u16())
                         .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
 
-        if !ctx.is_streaming() {
+                    // Record circuit breaker failure for error status codes
+                    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();
-                    Metrics::record_router_duration(
+                                }
-                        metrics_labels::ROUTER_OPENAI,
-                        metrics_labels::BACKEND_EXTERNAL,
-                        metrics_labels::CONNECTION_HTTP,
-                        model,
-                        metrics_labels::ENDPOINT_CHAT,
-                        start.elapsed(),
-                    );
                                 let mut response = Response::new(Body::from(body));
                                 *response.status_mut() = status;
                                 if let Some(ct) = content_type {
@@ -713,14 +729,6 @@ impl crate::routers::RouterTrait for OpenAIRouter {
                             }
                             Err(e) => {
                                 worker.circuit_breaker().record_failure();
-                    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,
-                    );
                                 (
                                     StatusCode::INTERNAL_SERVER_ERROR,
                                     format!("Failed to read response: {}", e),
@@ -729,15 +737,10 @@ impl crate::routers::RouterTrait for OpenAIRouter {
                             }
                         }
                     } else {
-            // For streaming, record duration at start since we can't track completion
+                        // Streaming response - record success when stream starts
-            Metrics::record_router_duration(
+                        if status.is_success() {
-                metrics_labels::ROUTER_OPENAI,
+                            worker.circuit_breaker().record_success();
-                metrics_labels::BACKEND_EXTERNAL,
+                        }
-                metrics_labels::CONNECTION_HTTP,
-                model,
-                metrics_labels::ENDPOINT_CHAT,
-                start.elapsed(),
-            );
                         let stream = resp.bytes_stream();
                         let (tx, rx) = mpsc::unbounded_channel();
                         tokio::spawn(async move {
@@ -756,7 +759,8 @@ impl crate::routers::RouterTrait for OpenAIRouter {
                                 }
                             }
                         });
-            let mut response = Response::new(Body::from_stream(UnboundedReceiverStream::new(rx)));
+                        let mut response =
+                            Response::new(Body::from_stream(UnboundedReceiverStream::new(rx)));
                         *response.status_mut() = status;
                         response
                             .headers_mut()
@@ -764,6 +768,47 @@ impl crate::routers::RouterTrait for OpenAIRouter {
                         response
                     }
                 }
+            },
+            |res, _attempt| is_retryable_status(res.status()),
+            |delay, attempt| {
+                Metrics::record_worker_retry(
+                    metrics_labels::BACKEND_EXTERNAL,
+                    metrics_labels::ENDPOINT_CHAT,
+                );
+                Metrics::record_worker_retry_backoff(attempt, delay);
+            },
+            || {
+                Metrics::record_worker_retries_exhausted(
+                    metrics_labels::BACKEND_EXTERNAL,
+                    metrics_labels::ENDPOINT_CHAT,
+                );
+            },
+        )
+        .await;
+
+        // Record duration/error metrics after retry completes
+        if response.status().is_success() {
+            Metrics::record_router_duration(
+                metrics_labels::ROUTER_OPENAI,
+                metrics_labels::BACKEND_EXTERNAL,
+                metrics_labels::CONNECTION_HTTP,
+                model,
+                metrics_labels::ENDPOINT_CHAT,
+                start.elapsed(),
+            );
+        } else {
+            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,
+            );
+        }
+
+        response
+    }
 
     async fn route_responses(
         &self,