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); + } +}