fix(routing): honor routing.retries on /v1/messages and /v1/responses

crates/aisix-proxy/src/lib.rs

@@ -4111,6 +4111,151 @@ data: [DONE]\n\n";
         assert_eq!(events[1].completion_tokens, 1);
     }
 
+    /// #641 parity: `/v1/messages` must honor `routing.retries` — re-hitting
+    /// the SAME target before failing over, like chat.rs. A single-target group
+    /// with `retries=1` and an always-502 target makes TWO attempts (initial +
+    /// one same-target retry), classified `initial` then `retry`. Before the fix
+    /// `/v1/messages` ignored `retries` and made only one attempt (so the
+    /// upstream `.expect(2)` and the second event would never arrive).
+    #[tokio::test]
+    async fn messages_routing_honors_same_target_retries() {
+        use aisix_obs::UsageSink;
+
+        let upstream = MockServer::start().await;
+        Mock::given(method("POST"))
+            .and(path("/chat/completions"))
+            .respond_with(ResponseTemplate::new(502).set_body_string("upstream down"))
+            .expect(2)
+            .mount(&upstream)
+            .await;
+
+        let hub = Arc::new(Hub::new());
+        hub.register_specialized("openai", Arc::new(openai_test_bridge()));
+
+        let snap = AisixSnapshot::new();
+        snap.provider_keys
+            .insert(pk_entry_with_id("pk-bad", &upstream.uri()));
+        snap.models
+            .insert(model_entry_with_id("m-bad", "primary", "pk-bad"));
+        snap.models.insert(routing_entry(
+            "smart",
+            "failover",
+            &["primary"],
+            Some(1),
+            None,
+            None,
+        ));
+        snap.apikeys.insert(apikey_entry("sk-caller", &["smart"]));
+
+        let (tx, mut rx) = tokio::sync::mpsc::channel(8);
+        let app = build_router(build_state(snap, hub).with_usage_sink(UsageSink::new(tx)));
+        let body = serde_json::json!({
+            "model": "smart",
+            "max_tokens": 100,
+            "messages": [{"role": "user", "content": "hi"}]
+        });
+        let req = Request::builder()
+            .method("POST")
+            .uri("/v1/messages")
+            .header("authorization", "Bearer sk-caller")
+            .header("content-type", "application/json")
+            .body(Body::from(body.to_string()))
+            .unwrap();
+
+        let resp = run(app, req).await;
+        assert_eq!(resp.status(), StatusCode::BAD_GATEWAY);
+
+        let mut events = Vec::new();
+        for _ in 0..2 {
+            let ev = tokio::time::timeout(std::time::Duration::from_millis(3000), rx.recv())
+                .await
+                .expect("two attempts (initial + retry) must each emit an event")
+                .expect("sender dropped");
+            events.push(ev);
+        }
+        events.sort_by_key(|e| e.attempt_index);
+        assert_eq!(events[0].attempt_kind, "initial");
+        assert_eq!(
+            events[1].attempt_kind, "retry",
+            "the same-target second attempt must be classified as a retry"
+        );
+        assert!(
+            events.iter().all(|e| e.attempt_model == "primary"),
+            "both attempts hit the SAME target (retry, not fallover)"
+        );
+        assert!(events.iter().all(|e| e.model_id == "m-bad"));
+        assert!(events.iter().all(|e| e.status_code == 502));
+        // upstream `.expect(2)` asserts exactly two upstream calls on Drop.
+    }
+
+    /// #641 parity for `/v1/responses` (Codex): same-target `routing.retries`
+    /// before fail-over. Single-target group, `retries=1`, always-502 →
+    /// initial + one retry, both hitting the same target.
+    #[tokio::test]
+    async fn responses_routing_honors_same_target_retries() {
+        use aisix_obs::UsageSink;
+
+        let upstream = MockServer::start().await;
+        Mock::given(method("POST"))
+            .and(path("/v1/responses"))
+            .respond_with(ResponseTemplate::new(502).set_body_string("upstream down"))
+            .expect(2)
+            .mount(&upstream)
+            .await;
+
+        let hub = Arc::new(Hub::new());
+        hub.register_specialized("openai", Arc::new(openai_test_bridge()));
+
+        let snap = AisixSnapshot::new();
+        snap.provider_keys
+            .insert(pk_entry_with_id("pk-bad", &upstream.uri()));
+        snap.models
+            .insert(model_entry_with_id("m-bad", "primary", "pk-bad"));
+        snap.models.insert(routing_entry(
+            "smart",
+            "failover",
+            &["primary"],
+            Some(1),
+            None,
+            None,
+        ));
+        snap.apikeys.insert(apikey_entry("sk-caller", &["smart"]));
+
+        let (tx, mut rx) = tokio::sync::mpsc::channel(8);
+        let app = build_router(build_state(snap, hub).with_usage_sink(UsageSink::new(tx)));
+        let body = serde_json::json!({"model": "smart", "input": "hi"});
+        let req = Request::builder()
+            .method("POST")
+            .uri("/v1/responses")
+            .header("authorization", "Bearer sk-caller")
+            .header("content-type", "application/json")
+            .body(Body::from(body.to_string()))
+            .unwrap();
+
+        let resp = run(app, req).await;
+        assert_eq!(resp.status(), StatusCode::BAD_GATEWAY);
+
+        let mut events = Vec::new();
+        for _ in 0..2 {
+            let ev = tokio::time::timeout(std::time::Duration::from_millis(3000), rx.recv())
+                .await
+                .expect("two attempts (initial + retry) must each emit an event")
+                .expect("sender dropped");
+            events.push(ev);
+        }
+        events.sort_by_key(|e| e.attempt_index);
+        assert_eq!(events[0].attempt_kind, "initial");
+        assert_eq!(
+            events[1].attempt_kind, "retry",
+            "the same-target second attempt must be classified as a retry"
+        );
+        assert!(
+            events.iter().all(|e| e.attempt_model == "primary"),
+            "both attempts hit the SAME target (retry, not fallover)"
+        );
+        assert!(events.iter().all(|e| e.status_code == 502));
+    }
+
     /// AISIX-Cloud#790: a plain direct-model request (no routing group)
     /// records the client-sent name in `requested_model` and keeps the
     /// direct model's own id in `model_id` — on both the OpenAI and the

crates/aisix-proxy/src/messages.rs

@@ -460,86 +460,115 @@ async fn dispatch(
     let is_routing_request = model_entry.value.routing.is_some();
     let mut routing = RoutingTelemetry::default();
 
+    // `routing.retries` — how many times to re-hit the SAME target (with
+    // backoff) on a retryable failure before failing over to the next target.
+    // Honoured here exactly like chat.rs (#641); 0 (the default) keeps the
+    // fail-over-only behaviour. /v1/messages previously ignored it entirely.
+    let retries = model_entry
+        .value
+        .routing
+        .as_ref()
+        .map(|r| r.retries_or_default())
+        .unwrap_or(0);
+
     // Walk targets, failing over to the next only on a retryable upstream
     // failure. A 4xx / config error is returned as-is — retrying other
     // targets won't help. Streaming and non-streaming share this loop:
     // `dispatch_to_target` branches internally and, for streaming, only
     // returns Ok once the first chunk has arrived under `stream_timeout`
     // (#554) — so the 200 is committed to exactly one target and a slow
-    // first chunk fails over like any other retryable error. Each target
-    // attempt becomes its own per-attempt record (#655).
+    // first chunk fails over like any other retryable error. Each attempt
+    // (initial / same-target retry / fallover) becomes its own per-attempt
+    // record (#655).
     let n = attempt_models.len();
     let mut last_err: Option<ProxyError> = None;
-    for (i, target) in attempt_models.iter().enumerate() {
-        let (idx, kind) = routing.begin_attempt(&target.model.display_name);
-        let target_model = if is_routing_request {
-            target.model.display_name.clone()
-        } else {
-            String::new()
-        };
+    'targets: for (i, target) in attempt_models.iter().enumerate() {
         let pk_id = crate::dispatch::resolve_provider_key(&snapshot, &target.model)
             .map(|e| e.id.clone())
             .unwrap_or_default();
-        let attempt_started = Instant::now();
-        match dispatch_to_target(
-            state,
-            &snapshot,
-            body,
-            target,
-            &model_name,
-            request_id,
-            started,
-            &auth.entry.id,
-            auth.key().team_id.clone(),
-            auth.key().user_id.clone(),
-            resolved_chain.clone(),
-            client,
-            AttemptInfo {
-                index: idx,
-                kind: kind.to_string(),
-                model: target_model.clone(),
-                ..Default::default()
-            },
-        )
-        .await
-        {
-            Ok(mut outcome) => {
-                routing.attempts.push(AttemptRecord {
-                    index: idx,
-                    kind,
-                    target_model,
-                    target_model_id: target.id.clone(),
-                    provider_key_id: outcome.provider_key_id.clone(),
-                    status: 200,
-                    success: true,
-                    error_class: String::new(),
-                    error_message: String::new(),
-                    latency_ms: ms_since(attempt_started),
-                });
-                outcome.routing = routing;
-                return Ok(outcome);
+        for attempt_idx in 0..=retries {
+            // Exponential backoff + jitter before re-hitting the SAME target
+            // (#641); cross-target fall-over (the outer loop) stays immediate.
+            if attempt_idx > 0 {
+                tokio::time::sleep(crate::routing::retry_backoff(attempt_idx as u32)).await;
             }
-            Err(e) => {
-                let retryable = matches!(
-                    &e,
-                    ProxyError::Bridge(be) if crate::routing::is_retryable(be, retry_on_429)
-                );
-                let (error_class, error_message) = attempt_error_from_proxy(&e);
-                routing.attempts.push(AttemptRecord {
+            let (idx, kind) = routing.begin_attempt(&target.model.display_name);
+            let target_model = if is_routing_request {
+                target.model.display_name.clone()
+            } else {
+                String::new()
+            };
+            let attempt_started = Instant::now();
+            match dispatch_to_target(
+                state,
+                &snapshot,
+                body,
+                target,
+                &model_name,
+                request_id,
+                started,
+                &auth.entry.id,
+                auth.key().team_id.clone(),
+                auth.key().user_id.clone(),
+                resolved_chain.clone(),
+                client,
+                AttemptInfo {
                     index: idx,
-                    kind,
-                    target_model,
-                    target_model_id: target.id.clone(),
-                    provider_key_id: pk_id,
-                    status: e.status().as_u16(),
-                    success: false,
-                    error_class,
-                    error_message,
-                    latency_ms: ms_since(attempt_started),
-                });
-                last_err = Some(e);
-                if !(retryable && i + 1 < n) {
-                    break;
+                    kind: kind.to_string(),
+                    model: target_model.clone(),
+                    ..Default::default()
+                },
+            )
+            .await
+            {
+                Ok(mut outcome) => {
+                    routing.attempts.push(AttemptRecord {
+                        index: idx,
+                        kind,
+                        target_model,
+                        target_model_id: target.id.clone(),
+                        provider_key_id: outcome.provider_key_id.clone(),
+                        status: 200,
+                        success: true,
+                        error_class: String::new(),
+                        error_message: String::new(),
+                        latency_ms: ms_since(attempt_started),
+                    });
+                    outcome.routing = routing;
+                    return Ok(outcome);
+                }
+                Err(e) => {
+                    let retryable = matches!(
+                        &e,
+                        ProxyError::Bridge(be) if crate::routing::is_retryable(be, retry_on_429)
+                    );
+                    let (error_class, error_message) = attempt_error_from_proxy(&e);
+                    routing.attempts.push(AttemptRecord {
+                        index: idx,
+                        kind,
+                        target_model,
+                        target_model_id: target.id.clone(),
+                        provider_key_id: pk_id.clone(),
+                        status: e.status().as_u16(),
+                        success: false,
+                        error_class,
+                        error_message,
+                        latency_ms: ms_since(attempt_started),
+                    });
+                    last_err = Some(e);
+                    // Non-retryable → stop entirely (retrying or failing over
+                    // won't help). Retryable → re-hit the same target until
+                    // `retries` is exhausted, then fall over to the next target
+                    // if there is one.
+                    if !retryable {
+                        break 'targets;
+                    }
+                    if attempt_idx == retries {
+                        if i + 1 >= n {
+                            break 'targets;
+                        }
+                        break;
+                    }
                 }
             }
         }