From d6200d62fb4ba76d6a661b923717f0bd8edbeca9 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 31 Oct 2026 21:00:49 +0800 Subject: [PATCH] fix(gateway): classify stream faults before recording health A terminal SSE error can arrive over HTTP 200 and a clean EOF. Pulse correctly treats errors as progress for idle timing, but the relay used that signal to record success before inspecting the error. Its EOF path also only applied rate limits, so balance and quota failures stayed healthy. Inspect errors before recording an opening success, reuse the shared stream-fault classifier and known quota reset, and flush undelimited error frames before final health accounting. Apply the same EOF tolerance while buffering a stream opening so the original error is retained there too. Keep rate limits out of the health tally and retain the existing success after content has started. Cover the terminal health tally, pause and reset behavior, original error forwarding, subsequent routing, all four SSE dialects and no replay after content. --- crates/tw-gateway/src/ending.rs | 32 +-- crates/tw-gateway/src/server/pipeline/hop.rs | 4 +- .../tw-gateway/src/server/pipeline/opening.rs | 24 +- .../tw-gateway/src/server/pipeline/relay.rs | 57 +++-- crates/tw-gateway/tests/passthrough.rs | 237 ++++++++++++++++++ 5 files changed, 303 insertions(+), 51 deletions(-) diff --git a/crates/tw-gateway/src/ending.rs b/crates/tw-gateway/src/ending.rs index a9f96ea5..f8523f94 100644 --- a/crates/tw-gateway/src/ending.rs +++ b/crates/tw-gateway/src/ending.rs @@ -76,7 +76,7 @@ pub struct Ending { /// 上游在流里报的错：哪一家、原话。**有它就不是成功** —— 响应头是 200，回答却断在了 /// 半路 upstream_error: Option<(String, String)>, - /// 上游报错的那一帧的 `data`：回程照它认是不是限流（见 [`Ending::upstream_fault`]） + /// 上游报错的那一帧的 `data`：回程照它认失败原因（见 [`Ending::upstream_fault`]） upstream_fault: Option, /// 上游报的错是解不开对话里别家封存的内容（见 [`tw_dialect::sealed`]）：失败的说明讲清楚 /// 怎么回事，不只转述原话 @@ -495,12 +495,26 @@ impl Ending { } /// 上游在流里报了错：那一帧的 `data`（上游的格式，[`Ending::streaming`] 时给的）。回程照它 - /// 认是不是限流 —— 答到一半才报的限流，回答照样以错误收尾，这一家照样暂停（见 - /// [`crate::throttle`]） + /// 认失败原因；流结束后先调 [`Self::flush_stream_error`]，没有空行的尾帧也算。 pub fn upstream_fault(&self) -> Option<&str> { self.upstream_fault.as_deref() } + /// 上游不再来字节了：看最后一帧有没有报错。只收错误嗅探器，不报结局、不结算用量， + /// 回程先按同一个错误记这一家的成败，之后再交给 [`Self::finished`] 或 [`Self::failed`]。 + pub fn flush_stream_error(&mut self) { + if let Some(mut w) = self.watch.take() { + let frames = w.frames.flush(); + if self.upstream_error.is_none() + && let Some((found, data)) = w.error_in(frames) + { + self.sealed_refused = sealed_refusal(w.dialect, &data); + self.upstream_error = Some(found); + self.upstream_fault = Some(data); + } + } + } + /// 上游 `provider` 回的不是 2xx，原样交给了客户端（4xx 是请求本身的问题，或者没有 /// 下一家可换了；3xx 交还客户端，由它决定跟不跟）。 /// @@ -749,17 +763,7 @@ impl Ending { self.fail(tw_api::FailureSource::Upstream, why); return; } - // 最后一帧后面不带空行的上游：收尾时再看一眼 - if let Some(mut w) = self.watch.take() { - let frames = w.frames.flush(); - if self.upstream_error.is_none() - && let Some((found, data)) = w.error_in(frames) - { - self.sealed_refused = sealed_refusal(w.dialect, &data); - self.upstream_error = Some(found); - self.upstream_fault = Some(data); - } - } + self.flush_stream_error(); if let Some((upstream, message)) = self.upstream_error.take() { let message: String = message.chars().take(500).collect(); let why = if self.sealed_refused { diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index d672cd2a..4e0bb8bf 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -2881,8 +2881,8 @@ fn known_reset(state: &AppState, provider: &str, cause: Cause) -> Cause { /// 一个流式回答在第一段内容之前就收了尾，收尾的那个事件是 `data`：这算不算 `provider` 的错、 /// 按什么原因停用。**和这里开头报错的一跳同一个判据**（[`super::opening`] 读出状态码， /// [`crate::failure::classify`] 判断）：5xx、限流、额度用完、凭据被拒是 `Some`；请求本身的 -/// 问题、不是错误的是 `None` —— 那是这一家答上了。WebSocket 上 Responses 的一轮用它（见 -/// `crate::ws::turn`） +/// 问题、不是错误的是 `None` —— 那是这一家答上了。HTTP 回程和 WebSocket 上 Responses 的一轮 +/// 都用它（见 [`super::relay`]、`crate::ws::turn`） pub(crate) fn stream_fault( state: &AppState, provider: &str, diff --git a/crates/tw-gateway/src/server/pipeline/opening.rs b/crates/tw-gateway/src/server/pipeline/opening.rs index 3a64e98e..b91a23cc 100644 --- a/crates/tw-gateway/src/server/pipeline/opening.rs +++ b/crates/tw-gateway/src/server/pipeline/opening.rs @@ -91,8 +91,14 @@ pub(super) async fn watch( } }; let chunk = match next { - // 只有开头的例行事件就收了尾：一段内容都没有，和断了一样 - None => return Opening::Broken(ended()), + None => { + // 和回程一样容忍尾帧没有空行。只补给解析器，转发的仍是上游原字节； + // 只有开头的例行事件就收了尾的，照旧算没有内容。 + let Some(judged) = sse.judge(dialect, b"\n\n") else { + return Opening::Broken(ended()); + }; + break judged; + } Some(Ok(c)) => c, Some(Err(e)) => { return Opening::Broken(crate::forward::map_reqwest_error(e)); @@ -112,12 +118,7 @@ pub(super) async fn watch( }, }; usage.feed(&text); - let judged = sse - .feed(&text) - .into_iter() - .map(|(event, data)| judge(dialect, event.as_deref(), &data)) - .find(|j| !matches!(j, Judge::Preamble)); - if let Some(j) = judged { + if let Some(j) = sse.judge(dialect, &text) { break j; } if size >= HOLD_LIMIT { @@ -422,6 +423,13 @@ struct Sse { } impl Sse { + fn judge(&mut self, dialect: Dialect, bytes: &[u8]) -> Option { + self.feed(bytes) + .into_iter() + .map(|(event, data)| judge(dialect, event.as_deref(), &data)) + .find(|j| !matches!(j, Judge::Preamble)) + } + fn feed(&mut self, bytes: &[u8]) -> Vec<(Option, String)> { self.buf.push_str(&String::from_utf8_lossy(bytes)); let mut out = Vec::new(); diff --git a/crates/tw-gateway/src/server/pipeline/relay.rs b/crates/tw-gateway/src/server/pipeline/relay.rs index 8b3044c4..dd141d0b 100644 --- a/crates/tw-gateway/src/server/pipeline/relay.rs +++ b/crates/tw-gateway/src/server/pipeline/relay.rs @@ -21,6 +21,7 @@ use futures::StreamExt; use super::Inbound; use super::hop::Served; use crate::error::GatewayError; +use crate::failure::Cause; use crate::forward; use crate::server::{because, flagged}; use crate::state::{AppState, Runtime}; @@ -197,6 +198,7 @@ pub(super) fn respond( let throttles = state.throttles.clone(); // 这个请求结束时记一笔，到了时候重读这一家的余额（见 `crate::balances`） let passing = state.balance_passing(&upstream_name); + let state = state.clone(); let stream = async_stream::stream! { // **通行证跟着响应体走。**这个流被丢掉的时候它才还回去：正常 // 发完是一种，客户端中途断开、hyper 丢掉响应体是另一种 —— 两种 @@ -327,11 +329,19 @@ pub(super) fn respond( match item { Ok(chunk) => { upstream_since = tokio::time::Instant::now(); + // **旁路嗅探和留档，不缓冲**：看的是上游原话（带占位符的那一版）， + // usage 和请求记录不受后面的还原、转换影响。先认错误，再记这一家的成败。 + ending.feed(&chunk); // 真内容才重新计时：心跳、开头的例行事件不算 if pulse.feed(&chunk) { said = true; quiet_by = quiet.after(std::time::Instant::now()); - if std::mem::take(&mut unconfirmed) { + // 错误帧也算上游有动静（不能报静默超时），却不是答上了。 + // 同一块里已经有内容再报错的，保留开头成功的那一笔。 + if unconfirmed + && (ending.has_content() || ending.upstream_fault().is_none()) + { + unconfirmed = false; crate::server::note_health( &bus, &health, @@ -340,13 +350,6 @@ pub(super) fn respond( ); } } - // **旁路嗅探和留档，不缓冲**：字节照常流向客户端，同时 - // 喂它一份。上游返回的 usage 是真相，而拿不到它就只能估。 - // - // 看的是上游原话（带占位符的那一版）：usage 数字不受 - // 影响，而请求详情里存的正是「我们发出去的和收回来的」， - // 把还原后的存进去会让那一页说谎。 - ending.feed(&chunk); let (out, cut) = relay.chunk(&chunk).await; relay.sent(&out); if !out.is_empty() { @@ -400,31 +403,31 @@ pub(super) fn respond( ); ending.too_slow(rate); } - // 上游在流里报的错是限流：回答照样以错误收尾（客户端已经收到了东西，换不了），这一家暂停 - if let Some((code, body)) = ending + // 尾帧没带空行的也先认出来：健康记录和请求结局看的是同一个错误。 + ending.flush_stream_error(); + let fault = ending .upstream_fault() - .and_then(|data| super::opening::stream_error(upstream_dialect, data)) - && let crate::failure::Verdict::Failed(crate::failure::Cause::RateLimited { - retry_after, - }) = crate::failure::classify( - code, - &http::HeaderMap::new(), - &body, - crate::server::now_ms(), - ) - { + .and_then(|data| super::hop::stream_fault(&state, &upstream_name, upstream_dialect, data)); + // 答到一半才报的限流也暂停；客户端已经收到了东西，这一轮不重放。 + if let Some(Cause::RateLimited { retry_after }) = fault { throttles.limited(&upstream_name, sent_at, retry_after); } // 响应体留档和结束事件都在 `ending` 里：三种结局要交出去的是同一份 // 东西，分开写就会有一种漏掉 - // 一段真内容都没有就好好收了尾：也是答上了 + // 还没有记过成败：正常 EOF 不代表成功，流里报的错和压着开头时按同一个原因记。 if unconfirmed && broke.is_none() { - crate::server::note_health( - &bus, - &health, - &upstream_name, - health.record_success(&upstream_name), - ); + let changed = match fault { + // 限流由上面暂停，不进健康的成功率和熔断。 + Some(Cause::RateLimited { .. }) => None, + Some(cause) => { + if matches!(cause, Cause::NoBalance | Cause::QuotaUsedUp { .. }) { + throttles.released(&upstream_name, sent_at); + } + health.record_cause(&upstream_name, cause) + } + None => health.record_success(&upstream_name), + }; + crate::server::note_health(&bus, &health, &upstream_name, changed); } match broke { None => { diff --git a/crates/tw-gateway/tests/passthrough.rs b/crates/tw-gateway/tests/passthrough.rs index 0998401b..301f1fde 100644 --- a/crates/tw-gateway/tests/passthrough.rs +++ b/crates/tw-gateway/tests/passthrough.rs @@ -2496,3 +2496,240 @@ async fn a_quota_reset_is_kept_as_the_moment_it_happens() { Some(resets_at) ); } + +/// 第一家暂时不可用，第二家是最后的候选；四笔成功让下一笔成败能从成功率上看出来。 +async fn last_stream_gateway( + body: &'static str, + wait_secs: u64, +) -> (SocketAddr, tw_gateway::AppState, Arc) { + let (first, _) = start_upstream(false).await; + let (last, hits) = start_answering_upstream(200, "text/event-stream", body).await; + let mut cfg = two(first, last); + cfg.groups = vec![ + serde_yaml_ng::from_str("name: ordered\ntype: fallback\nproviders: [second, first]\n") + .unwrap(), + ]; + cfg.routes = vec![tw_engine::RouteSet::default_with(vec![ + serde_yaml_ng::from_str("name: catch-all\nto: ordered\n").unwrap(), + ])]; + cfg.failover.wait_secs = wait_secs; + cfg.failover.failures_to_pause = 1; + let state = tw_gateway::AppState::new(cfg).unwrap(); + state + .health + .record_cause("first", tw_gateway::failure::Cause::NoBalance); + for _ in 0..4 { + state.health.record_success("second"); + } + let gw = tw_gateway::serve(state.clone(), ([127, 0, 0, 1], 0).into()) + .await + .unwrap(); + (gw, state, hits) +} + +async fn last_stream(gw: SocketAddr) -> String { + let response = reqwest::Client::new() + .post(format!("http://{gw}/v1/messages")) + .header("x-api-key", "tw-k") + .body(r#"{"model":"claude-sonnet-4-5","stream":true}"#) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + response.text().await.unwrap() +} + +#[tokio::test] +async fn a_last_stream_balance_error_pauses_the_provider_and_counts_once() { + let error = "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"Insufficient Balance\"}}\n\n"; + for (wait_secs, body) in [ + (0, error), + (0, error.trim_end()), + (2, error), + (2, error.trim_end()), + ] { + let (gw, state, hits) = last_stream_gateway(body, wait_secs).await; + let mut events = state.bus.subscribe(); + assert_eq!(last_stream(gw).await, body); + assert!( + !state.health.is_available("second"), + "wait={wait_secs}, {body}" + ); + assert_eq!( + state.health.success_rates(&["second".into()])["second"], + 0.8, + "一次失败不能既记成功又记失败" + ); + loop { + match tokio::time::timeout(Duration::from_secs(2), events.recv()) + .await + .unwrap() + .unwrap() + { + tw_api::Event::RequestFailed { message, .. } => { + assert_eq!(message.code, "gw.upstream.stream_error"); + assert!(message.text.contains("Insufficient Balance")); + break; + } + tw_api::Event::RequestFinished { .. } => panic!("错误流被记成了成功"), + _ => {} + } + } + // 第一家恢复后，下一次请求跳过刚刚报余额不足的第二家。 + state.health.record_success("first"); + assert_eq!(send_to(gw).await.status(), 200); + assert_eq!(hits.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn last_stream_errors_use_the_same_health_causes_as_opening_errors() { + for (body, available, rate) in [ + ( + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n", + false, + Some(0.8), + ), + ( + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"authentication_error\",\"message\":\"Bad key\"}}\n\n", + false, + Some(0.8), + ), + ( + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"not_found_error\",\"message\":\"model not found\"}}\n\n", + true, + None, + ), + ( + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"max_tokens must be positive\"}}\n\n", + true, + Some(1.0), + ), + ] { + let (gw, state, hits) = last_stream_gateway(body, 0).await; + assert_eq!(last_stream(gw).await, body); + assert_eq!(state.health.is_available("second"), available, "{body}"); + assert_eq!( + state + .health + .success_rates(&["second".into()]) + .get("second") + .copied(), + rate, + "{body}" + ); + assert_eq!(hits.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn a_last_stream_quota_error_uses_the_reported_reset() { + for known in [false, true] { + let body = if known { + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"usage_limit_reached\"}}" + } else { + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"usage_limit_reached\",\"resets_in_seconds\":120}}\n\n" + }; + let (gw, state, _) = last_stream_gateway(body, 0).await; + if known { + let headers = HeaderMap::from_iter([ + ( + "anthropic-ratelimit-unified-5h-utilization" + .parse() + .unwrap(), + "1.0".parse().unwrap(), + ), + ( + "anthropic-ratelimit-unified-5h-reset".parse().unwrap(), + "120".parse().unwrap(), + ), + ]); + let now = chrono::Utc::now().timestamp_millis() as u64; + state.record_quota(0, "second", tw_gateway::quota::from_headers(&headers, now)); + } + assert_eq!(last_stream(gw).await, body); + assert!(!state.health.is_available("second")); + let left = state.health.cooldown_left("second").unwrap(); + assert!((115..=120).contains(&left.as_secs()), "{left:?}"); + assert_eq!( + state.health.success_rates(&["second".into()])["second"], + 0.8 + ); + } +} + +#[tokio::test] +async fn a_last_stream_rate_limit_pauses_without_counting_as_success_or_failure() { + for body in [ + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"Too many requests\"}}\n\n", + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"Too many requests\"}}", + ] { + let (gw, state, hits) = last_stream_gateway(body, 0).await; + assert_eq!(last_stream(gw).await, body); + assert!(state.health.is_available("second")); + assert!(state.health.success_rates(&["second".into()]).is_empty()); + assert!(state.throttles.view("second").is_some()); + assert_eq!(hits.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn a_last_stream_with_content_keeps_its_success_and_is_not_replayed() { + for body in [ + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"half\"}}\n\nevent: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"Insufficient Balance\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"half\"}}\n\nevent: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"Too many requests\"}}", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ] { + let (gw, state, hits) = last_stream_gateway(body, 0).await; + assert_eq!(last_stream(gw).await, body); + assert!(state.health.is_available("second")); + assert_eq!( + state.health.success_rates(&["second".into()])["second"], + 1.0 + ); + assert_eq!( + state.throttles.view("second").is_some(), + body.contains("rate_limit_error") + ); + assert_eq!(hits.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn last_stream_errors_are_classified_in_the_upstream_dialect() { + for (path, request, body) in [ + ( + "/v1/chat/completions", + r#"{"model":"m","messages":[],"stream":true}"#, + "data: {\"error\":{\"type\":\"invalid_request_error\",\"message\":\"Insufficient Balance\"}}\n\n", + ), + ( + "/v1/responses", + r#"{"model":"m","input":"hi","stream":true}"#, + "event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"insufficient_quota\",\"message\":\"Insufficient Balance\"}}}\n\n", + ), + ( + "/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse", + r#"{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}"#, + "data: {\"error\":{\"code\":400,\"status\":\"INVALID_ARGUMENT\",\"message\":\"Insufficient Balance\"}}\n\n", + ), + ] { + let (gw, state, hits) = last_stream_gateway(body, 0).await; + let response = reqwest::Client::new() + .post(format!("http://{gw}{path}")) + .header("x-api-key", "tw-k") + .body(request) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200, "{path}"); + assert_eq!(response.text().await.unwrap(), body, "{path}"); + assert!(!state.health.is_available("second"), "{path}"); + assert_eq!( + state.health.success_rates(&["second".into()])["second"], + 0.8, + "{path}" + ); + assert_eq!(hits.load(Ordering::SeqCst), 1); + } +}