diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 6ecf5ed8..5f18db35 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/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 diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 53988258..08712f6e 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/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 = 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; + } } } } diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 955187f4..07e66cf9 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -367,6 +367,15 @@ async fn dispatch( .unwrap_or(false); let is_routing_request = model_entry.value.routing.is_some(); let mut routing = RoutingTelemetry::default(); + // `routing.retries` — same-target retries (with backoff) before failing + // over, honoured exactly like chat.rs / messages.rs (#641). 0 (default) + // keeps fail-over-only; /v1/responses previously ignored it entirely. + let retries = model_entry + .value + .routing + .as_ref() + .map(|r| r.retries_or_default()) + .unwrap_or(0); // Walk the targets, failing over on a retryable failure. Streaming and // non-streaming share this loop: the per-target dispatch branches @@ -374,99 +383,112 @@ async fn dispatch( // has arrived under `stream_timeout` (#554) — so the 200 is committed to // exactly one target and a slow first chunk fails over. let mut last_err: Option = None; - for target in &attempt_models { - 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(); + 'targets: for target in &attempt_models { // Resolved ProviderKey UUID for this target — feeds the per-PK // telemetry attribution tags on the emitted UsageEvent // (AISIX-Cloud#867). Recorded on the AttemptRecord (success + failure) // so both the winner and each failed-attempt event can attribute it. let pk_id = target.model.provider_key_id.clone().unwrap_or_default(); - // Winning-attempt classification (#655) for the streaming path's - // end-of-stream UsageEvent. The non-streaming / buffered paths emit - // from the handler and ignore it. - let attempt = AttemptInfo { - index: idx, - kind: kind.to_string(), - model: target_model.clone(), - ..Default::default() - }; - let result = if target.model.provider.as_deref() == Some("openai") { - responses_to_target( - state, - &snapshot, - body, - &target.model, - &target.id, - request_id, - resolved_chain.as_ref(), - started, - &model_name, - &auth.entry.id, - client, - attempt, - ) - .await - } else { - responses_cross_provider_to_target( - state, - &snapshot, - body, - &target.model, - &target.id, - request_id, - resolved_chain.clone(), - started, - &model_name, - &auth.entry.id, - client, - attempt, - ) - .await - }; - match result { - Ok(mut success) => { - routing.attempts.push(AttemptRecord { - index: idx, - kind, - target_model, - target_model_id: target.id.clone(), - provider_key_id: pk_id, - status: success.response.status().as_u16(), - success: true, - error_class: String::new(), - error_message: String::new(), - latency_ms: ms_since(attempt_started), - }); - success.routing = routing; - return Ok(success); + 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 { - 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 { - break; + 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(); + // Winning-attempt classification (#655) for the streaming path's + // end-of-stream UsageEvent. The non-streaming / buffered paths emit + // from the handler and ignore it. + let attempt = AttemptInfo { + index: idx, + kind: kind.to_string(), + model: target_model.clone(), + ..Default::default() + }; + let result = if target.model.provider.as_deref() == Some("openai") { + responses_to_target( + state, + &snapshot, + body, + &target.model, + &target.id, + request_id, + resolved_chain.as_ref(), + started, + &model_name, + &auth.entry.id, + client, + attempt, + ) + .await + } else { + responses_cross_provider_to_target( + state, + &snapshot, + body, + &target.model, + &target.id, + request_id, + resolved_chain.clone(), + started, + &model_name, + &auth.entry.id, + client, + attempt, + ) + .await + }; + match result { + Ok(mut success) => { + routing.attempts.push(AttemptRecord { + index: idx, + kind, + target_model, + target_model_id: target.id.clone(), + provider_key_id: pk_id.clone(), + status: success.response.status().as_u16(), + success: true, + error_class: String::new(), + error_message: String::new(), + latency_ms: ms_since(attempt_started), + }); + success.routing = routing; + return Ok(success); + } + 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. Retryable → re-hit the + // same target until `retries` is exhausted, then fall over + // to the next target (the outer loop advances). + if !retryable { + break 'targets; + } + if attempt_idx == retries { + break; + } } } }