From 2bd8fb4caa1ae5bcfae9c31f6585cbdbb3e3a90f Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Fri, 9 Oct 2026 15:11:05 +0800 Subject: [PATCH 1/2] feat(failover): idle timeout with automatic failover, and manual abort of requests and sessions An upstream could answer 200 and then send nothing, or stop halfway, and nothing ever gave up on it: there is deliberately no overall timeout, and `failover.next_on_slow_start` (off by default) only covered the start of a stream and always let the last candidate wait forever. `failover.idle_timeout_secs` (default 300, 30 to 3600) replaces it. The timer starts when the request is sent upstream, so slot waits do not count, and restarts on every piece of real content; keep-alives do not count, so an upstream that only pings still runs out of time. What counts is defined per dialect in one place (`tw_gateway::pulse`) and is shared with the opening hold, so "first content" means the same thing in both. - Before any content has reached the client the upstream counts as a failure (cooldown rules apply), the attempt is recorded as `idle_timeout` with the input it may have billed, the conversation no longer stays on it for the turn, and the request moves on. Streams and whole answers are held until their first content on every candidate but the last, so the next upstream starts afresh. With no candidate left the client gets a 504 in its own format. - After content has reached the client (or on the last candidate, whose stream is passed on as it arrives) the answer ends with an error event in the client's format and the request is recorded as failed. `stream_start_wait_secs` and `next_on_slow_start` are removed: holding the opening only up to a shorter window would make the before-content failover impossible, so the hold now lasts until the idle timeout. Old configs that still name them fail as unknown fields and the safe-mode repair offers to delete them (tested). Manual abort: `POST /request/{id}/abort` and `POST /sessions/{id}/abort` throw a per-request switch registered when the request starts and dropped with its ending. The upstream call is dropped at once, the client gets an error in its format (499 before the answer started), the request ends as `RequestFailed` with the new source `aborted` and code `gw.request.aborted`, an in-flight hop is recorded as `aborted`, and the upstream is not set aside. Upstream health counts aborts with cancellations. Requests no longer running are a 404 with their own codes. WebSocket turns are not abortable. Protocol 43. The store schema goes to 26 because `slow_start` disappears from stored attempt chains. Co-Authored-By: Claude Opus 5.5 --- crates/tw-api/msg-codes.txt | 7 +- crates/tw-api/src/ep.rs | 5 + crates/tw-api/src/lib.rs | 80 +- crates/tw-config/src/failover.rs | 51 +- crates/tw-config/src/lib.rs | 3 +- crates/tw-config/src/repair.rs | 36 + crates/tw-config/src/validate.rs | 45 +- crates/tw-config/tests/manual/schema.rs | 17 +- crates/tw-control/src/lib.rs | 41 +- crates/tw-control/tests/resources.rs | 83 +- crates/tw-dialect/src/gemini/response.rs | 2 + crates/tw-gateway/src/abort.rs | 186 +++ crates/tw-gateway/src/affinity.rs | 43 + crates/tw-gateway/src/ending.rs | 17 + crates/tw-gateway/src/error.rs | 46 + crates/tw-gateway/src/latency.rs | 2 +- crates/tw-gateway/src/lib.rs | 2 + crates/tw-gateway/src/pulse.rs | 378 +++++++ crates/tw-gateway/src/server.rs | 12 +- crates/tw-gateway/src/server/pipeline.rs | 17 +- crates/tw-gateway/src/server/pipeline/hop.rs | 477 ++++---- crates/tw-gateway/src/server/pipeline/idle.rs | 160 +++ .../tw-gateway/src/server/pipeline/opening.rs | 224 ++-- .../tw-gateway/src/server/pipeline/relay.rs | 100 +- crates/tw-gateway/src/server/pipeline/slow.rs | 101 -- crates/tw-gateway/src/state.rs | 7 + crates/tw-gateway/src/wire.rs | 2 + crates/tw-gateway/tests/idle_timeout.rs | 1007 +++++++++++++++++ crates/tw-gateway/tests/latency_samples.rs | 29 +- crates/tw-gateway/tests/slow_start.rs | 730 ------------ crates/tw-store/src/db.rs | 5 +- crates/tw-store/src/health.rs | 26 +- docs/config.md | 53 +- docs/config.zh-CN.md | 29 +- scripts/smoke.sh | 7 + 35 files changed, 2698 insertions(+), 1332 deletions(-) create mode 100644 crates/tw-gateway/src/abort.rs create mode 100644 crates/tw-gateway/src/pulse.rs create mode 100644 crates/tw-gateway/src/server/pipeline/idle.rs delete mode 100644 crates/tw-gateway/src/server/pipeline/slow.rs create mode 100644 crates/tw-gateway/tests/idle_timeout.rs delete mode 100644 crates/tw-gateway/tests/slow_start.rs diff --git a/crates/tw-api/msg-codes.txt b/crates/tw-api/msg-codes.txt index f1960fac..52d372f6 100644 --- a/crates/tw-api/msg-codes.txt +++ b/crates/tw-api/msg-codes.txt @@ -107,7 +107,6 @@ config.schema_too_new config.secret.empty_name config.secret.env_missing config.secret.unterminated -config.slow_start_too_short config.store.conflict config.store.missing config.store.read_failed @@ -206,6 +205,7 @@ control.replay_bedrock control.request_body_gone control.request_body_truncated control.request_not_found +control.request_not_running control.request_rejected control.reset_card_result_unreadable control.reset_cards_unreadable @@ -224,6 +224,7 @@ control.rule.no_such_probe_class control.rule.no_such_upstream control.rule.unknown_dialect control.session_not_found +control.session_not_running control.sheet_in_use control.sheet_not_found control.shutdown @@ -372,6 +373,7 @@ gw.probe.connect gw.probe.key_rejected gw.probe.request_failed gw.probe.timeout +gw.request.aborted gw.request.body_declared_too_large gw.request.body_over_limit gw.request.body_unreadable @@ -383,7 +385,6 @@ gw.route.protocol_mismatch gw.route.rule_failed gw.route.selected_upstream_missing gw.route.upstream_missing -gw.slow_start gw.toolcall.connection_cut gw.toolcall.response_cut gw.toolcall.response_withheld @@ -393,6 +394,8 @@ gw.upstream.bedrock_refused gw.upstream.bedrock_refused_unnamed gw.upstream.eventstream_broken gw.upstream.forward_failed +gw.upstream.idle_timeout +gw.upstream.idle_timeout_mid_stream gw.upstream.rate_limited gw.upstream.sign_failed gw.upstream.status diff --git a/crates/tw-api/src/ep.rs b/crates/tw-api/src/ep.rs index 11543588..edbcbefc 100644 --- a/crates/tw-api/src/ep.rs +++ b/crates/tw-api/src/ep.rs @@ -60,6 +60,9 @@ endpoints! { /// 延迟与速度，各带样本数和别家的参照。不给时间窗是最近 7 天 UpstreamHealth: GET "/upstreams/health", api::Window => api::UpstreamHealth; RequestDetail: GET "/request/{id}" [id], () => api::RequestDetail; + /// 中止一个在跑的请求：和上游的连接断开，客户端收到错误，请求记成手动中止。已经 + /// 结束了的是 404 + AbortRequest: POST "/request/{id}/abort" [id], () => api::Aborted; /// 把一条记录变成回放用例（YAML） Fixture: GET "/request/{id}/fixture" [id], () => String, text; Sessions: GET "/sessions", api::ListQuery => Vec; @@ -67,6 +70,8 @@ endpoints! { /// 一次会话读成一段对话：每一轮新说的话、回答、工具调用和结果（已脱敏）。可以只要 /// 从某一轮起的那些 SessionTranscript: GET "/sessions/{id}/transcript" [id], api::TranscriptQuery => api::Transcript; + /// 中止一次会话里所有在跑的请求。一个都没有是 404 + AbortSession: POST "/sessions/{id}/abort" [id], () => api::Aborted; // ─────────────────────────────────────────────── 测速、回放、试路由 SpeedQuote: POST "/speed/quote", api::SpeedRunRequest => api::SpeedQuote; diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index ec1f5830..52bb265d 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -224,6 +224,11 @@ slug_enum! { Request = "request", RateLimited = "rate_limited", Denied = "denied", + /// 在界面上手动中止的（`POST /request/{id}/abort`、`POST /sessions/{id}/abort`）。 + /// **不是上游的错**：那一家不停用、不算失败；和客户端自己走掉（`RequestCancelled`） + /// 也不是一回事。`message` 是 [`ABORTED`] 那个码。还没开始回答的，客户端收到 499； + /// 回答到一半的，按它的格式以一条错误收尾 + Aborted = "aborted", /// 网关自己的代码崩掉了 Internal = "internal", } @@ -814,7 +819,22 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// `config.manual_model_blank`、`_padded`、`_wildcard`、`_duplicate`、`_too_long`（400）， /// 整份配置的校验也查这几条。`source` 是 `discovered` 时 `models` 里也有手动添加的那些。 /// 照 41 写的界面分不出哪些模型是手动添加的。 -pub const CONTROL_API_VERSION: u32 = 42; +/// +/// **43 起上游不出声有了上限，请求能手动中止**：[`FailoverView`] 的 `stream_start_wait_secs` +/// 和 `next_on_slow_start` 删了，换成 `idle_timeout_secs`（无响应超时，默认 300 秒，30 到 +/// 3600）：从请求发出去起算，每来一段内容重新计时，心跳不算。客户端还什么都没收到时，这一家 +/// 记一次失败、换下一家，没有下一家了回 504；已经收到一部分的，回答按客户端的格式以错误收尾。 +/// 配置里写 `stream_start_wait_secs`、`next_on_slow_start` 加载不了（不认识的字段，一键修复删掉 +/// 它们），消息码 `config.slow_start_too_short`、`gw.slow_start` 跟着删。尝试链的结果 +/// （[`AttemptOutcome`]）删了 `slow_start`，多了 `idle_timeout`（说等了多少秒的 +/// `gw.upstream.idle_timeout`）和 `aborted`。新端点 `POST /request/{id}/abort` 和 +/// `POST /sessions/{id}/abort`（→ [`Aborted`]）叫停一个在跑的请求、一次会话里所有在跑的请求： +/// 和上游的连接立刻断开，客户端收到它自己格式的错误（还没开始回答的是 499），请求照常报 +/// [`Event::RequestFailed`]，`source` 是新的 [`FailureSource::Aborted`]、`message` 是 +/// [`ABORTED`]，上游不停用。已经结束了的回 404（`control.request_not_running`、 +/// `control.session_not_running`）。流中途停了的另有 `gw.upstream.idle_timeout_mid_stream`。 +/// 照 42 写的界面读不到 `idle_timeout_secs`，不认 `aborted` 这个来源。 +pub const CONTROL_API_VERSION: u32 = 43; #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -1515,11 +1535,17 @@ slug_enum! { /// （`status` 是它回的那个）。尝试链到此为止，这一行记成网关自己答的 /// （[`HistoryRow::local`]），费用 0 Estimated = "estimated", - /// 流式回答等了 `failover.stream_start_wait_secs` 还没有内容，开着 - /// `failover.next_on_slow_start`，放弃这一家、换下一家（连接断开，上游不再生成）。 - /// 响应头到了的有 `status`，没到的没有。**这一家不停用、不算失败**。上游可能已经按 - /// 输入收了钱：知道多少的在 `usage` 里 - SlowStart = "slow_start", + /// 从请求发出去起 `failover.idle_timeout_secs` 秒没有内容（无响应超时），客户端也还 + /// 什么都没收到：放弃这一家（连接断开，上游不再生成），换下一家，没有下一家了就回 + /// 超时错误。响应头到了的有 `status`，没到的没有。`error` 是 `gw.upstream.idle_timeout`， + /// 说等了多少秒。**这一家记一次失败**（和 5xx 一样算进停用的账）。上游可能已经按输入 + /// 收了钱：知道多少的在 `usage` 里 + IdleTimeout = "idle_timeout", + /// 这一跳在等上游时被手动中止（`POST /request/{id}/abort`、`POST /sessions/{id}/abort`）： + /// 连接断开，上游不再生成，尝试链到此为止。**这一家不停用、不算失败**。已经接下、 + /// 在交回答的那一跳不改成它（还是 `served`），请求本身记成手动中止（见 + /// [`FailureSource::Aborted`]） + Aborted = "aborted", } } @@ -1545,11 +1571,12 @@ pub struct AttemptView { /// 上游返回的状态码。`error` 时没有 #[serde(default, skip_serializing_if = "Option::is_none")] pub status: Option, - /// `error` 时的说明。和这一跳报给客户端的那条错误是同一句。`slow_start` 时说等了多久 + /// `error` 时的说明。和这一跳报给客户端的那条错误是同一句。`idle_timeout` 时说等了多久， + /// `aborted` 时是 `gw.request.aborted` #[serde(default, skip_serializing_if = "Option::is_none")] pub error: Option, pub ms: u64, - /// 放弃了的这一跳（`slow_start`）上游可能已经收了钱的输入（见 [`AttemptUsage`]）。估不 + /// 放弃了的这一跳（`idle_timeout`）上游可能已经收了钱的输入（见 [`AttemptUsage`]）。估不 /// 出来的（请求解不开）没有。别的结果都没有：接下请求的那一跳的用量在结局里 #[serde(default, skip_serializing_if = "Option::is_none")] pub usage: Option, @@ -1575,12 +1602,19 @@ pub const FAILED_AFTER_SENDING: &[&str] = &[ "gw.upstream.stream_opening_error", ]; +/// 手动中止的请求报的码（[`FailureSource::Aborted`] 的 `message`，尝试链上 `aborted` 那一跳的 +/// `error`）。记录里只存码，上游体检按它把手动中止的和客户端取消的一样数，不算上游失败 +pub const ABORTED: &str = "gw.request.aborted"; + impl AttemptView { - /// 这一跳发到了上游：上游回了话（`served`、`status`、`slow_start`，`estimated` 里带着 + /// 这一跳发到了上游：上游回了话（`served`、`status`、`idle_timeout`、`aborted`，`estimated` 里带着 /// 状态码的），或者发出去之后才失败（[`FAILED_AFTER_SENDING`]）。 pub fn sent(&self) -> bool { match self.outcome { - AttemptOutcome::Served | AttemptOutcome::Status | AttemptOutcome::SlowStart => true, + AttemptOutcome::Served + | AttemptOutcome::Status + | AttemptOutcome::IdleTimeout + | AttemptOutcome::Aborted => true, AttemptOutcome::Estimated => self.status.is_some(), AttemptOutcome::Error => { self.skipped.is_none() @@ -1609,7 +1643,7 @@ impl RoutingView { } } -/// 放弃了的一跳（[`AttemptOutcome::SlowStart`]）上游可能已经收了钱的输入。 +/// 放弃了的一跳（[`AttemptOutcome::IdleTimeout`]）上游可能已经收了钱的输入。 /// /// 上游在流开头报了的（Anthropic 的 `message_start`）是它报的数；没报的只有 `input`，是网关 /// 估的（`estimated`，和 [`Event::RequestStarted`] 的 `input_estimate` 同一个数）。**输出不知道**： @@ -1841,6 +1875,17 @@ pub struct InFlightRequest { pub events: Vec, } +/// 手动中止（`POST /request/{id}/abort`、`POST /sessions/{id}/abort`）叫停了哪几个请求。 +/// +/// **叫停是立刻的，结局随后到**：每个请求照常报一条 [`Event::RequestFailed`]（`source` 是 +/// [`FailureSource::Aborted`]），通常在这条响应之前或紧跟着。叫停之后才跑完的不改。 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct Aborted { + /// 叫停的请求号，开始得早的在前 + pub requests: Vec, +} + /// 配置文件没通过校验的那一次。字段和 [`Event::ConfigRejected`] 一样。 /// /// **只记外部改动**（在编辑器里改的、命令行写的）：界面自己写坏的根本没落盘， @@ -1973,7 +2018,7 @@ pub struct RetentionView { pub body_bytes_now: u64, } -/// 上游失败之后停用多久、流开头最多等多久。和配置的 `failover` 一一对应， +/// 上游失败之后停用多久、多久没有内容就不再等。和配置的 `failover` 一一对应， /// 没写的是默认值 —— **界面显示的就是真在用的数**，不是「空 = 默认」。 #[derive(Debug, Clone, Default, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -1990,10 +2035,9 @@ pub struct FailoverView { pub quota_pause_secs: u64, /// 限流时按 `Retry-After` 停用，最多多少秒 pub rate_limit_max_pause_secs: u64, - /// 流式回答的开头最多等多少秒 - pub stream_start_wait_secs: u64, - /// 等过 `stream_start_wait_secs` 还没有内容就换下一家（最后一家照常等） - pub next_on_slow_start: bool, + /// 上游多少秒没有内容就不再等它（无响应超时）：从请求发出去起算，每来一段内容重新 + /// 计时。客户端还什么都没收到时换下一家，已经收到一部分的报错收尾 + pub idle_timeout_secs: u64, /// 一个请求合计最多等多少秒：等密钥的分钟、小时上限空出名额，和等满着（`max_concurrent`） /// 的上游空出位置，共用这一段。0 是不等 pub slot_wait_secs: u64, @@ -5833,7 +5877,9 @@ mod tests { for a in [ attempt(Served, Some(200), None), attempt(Status, Some(503), None), - attempt(SlowStart, None, Some("gw.slow_start")), + attempt(IdleTimeout, None, Some("gw.upstream.idle_timeout")), + attempt(IdleTimeout, Some(200), Some("gw.upstream.idle_timeout")), + attempt(Aborted, None, Some("gw.request.aborted")), attempt(Estimated, Some(404), None), attempt(Error, None, Some("gw.upstream.timeout")), attempt(Error, None, Some("gw.upstream.stream_opening_error")), diff --git a/crates/tw-config/src/failover.rs b/crates/tw-config/src/failover.rs index 9b29b7da..23dbfcd5 100644 --- a/crates/tw-config/src/failover.rs +++ b/crates/tw-config/src/failover.rs @@ -1,5 +1,4 @@ -//! 故障转移：一家上游失败之后停用多久、流开头最多等多久、等不到内容换不换下一家、上游 -//! 满着时最多等多久。 +//! 故障转移：一家上游失败之后停用多久、多久没有内容就不再等它、上游满着时最多等多久。 use serde::{Deserialize, Serialize}; @@ -34,15 +33,15 @@ pub struct Failover { /// 原因的失败 #[serde(default = "d_rate_limit_max_pause_secs")] pub rate_limit_max_pause_secs: u64, - /// 流式回答开头最多等多少秒。在第一段内容到达之前，上游报的错误照样换下一家； - /// 等过这么久还没有内容，就不再等，把已经收到的交给客户端 - #[serde(default = "d_stream_start_wait_secs")] - pub stream_start_wait_secs: u64, - /// 流式回答等过 [`Self::stream_start_wait_secs`] 还没有内容时，放弃这一家、换下一家。 - /// **最后一家不换**，照常等下去；这一家不停用，也不算一次失败。默认关：先想好再 - /// 输出的模型开头本来就慢，开着时要把等待调长 - #[serde(default)] - pub next_on_slow_start: bool, + /// 上游多少秒没有内容就不再等它（无响应超时）。从请求发给它的那一刻算起，**每来一段 + /// 内容重新计时**：正文、推理、工具调用都算，心跳不算（SSE 注释、Anthropic 的 `ping`、 + /// 只有角色的空块……）—— 否则一家只发心跳的上游会一直挂着。整包的回答没有「一段段」， + /// 从发出去到整份回来算一段。 + /// + /// 到点时还没有内容交给客户端的，这一家记一次失败，换下一家；没有下一家了就回超时 + /// 错误。已经交出去一部分的换不了（客户端会收到两遍开头），按客户端的格式报错收尾 + #[serde(default = "d_idle_timeout_secs")] + pub idle_timeout_secs: u64, /// 一个请求最多等多少秒，**整个请求合起来算**：准入时等密钥的分钟、小时上限空出名额， /// 之后上游的并发数满了（`providers[].max_concurrent`）时等空位，共用这一段。留在那一家 /// 的对话等它空出来，候选都满了时等先空出来的那一家；等不到的换下一家，或者回 429。 @@ -69,8 +68,8 @@ fn d_quota_pause_secs() -> u64 { fn d_rate_limit_max_pause_secs() -> u64 { 3600 } -fn d_stream_start_wait_secs() -> u64 { - 15 +fn d_idle_timeout_secs() -> u64 { + 300 } fn d_slot_wait_secs() -> u64 { 30 @@ -85,8 +84,7 @@ impl Default for Failover { no_balance_pause_secs: d_no_balance_pause_secs(), quota_pause_secs: d_quota_pause_secs(), rate_limit_max_pause_secs: d_rate_limit_max_pause_secs(), - stream_start_wait_secs: d_stream_start_wait_secs(), - next_on_slow_start: false, + idle_timeout_secs: d_idle_timeout_secs(), slot_wait_secs: d_slot_wait_secs(), } } @@ -96,12 +94,11 @@ impl Default for Failover { /// 说的时刻，周额度本来就可能在七天之后 pub const MAX_PAUSE_SECS: u64 = 7 * 24 * 3600; -/// 流开头最多等多少秒。再长的话，一家卡在半路的上游会让客户端先超时 -pub const MAX_STREAM_START_WAIT_SECS: u64 = 120; +/// 无响应超时最少写多少秒。再短的话，先想好再输出的模型还没开口就被放弃了 +pub const MIN_IDLE_TIMEOUT_SECS: u64 = 30; -/// 开着「开头慢就换下一家」时，流开头至少等多少秒。再短的话，平常的请求还没开口就被 -/// 切掉了 -pub const MIN_SLOW_START_WAIT_SECS: u64 = 5; +/// 无响应超时最多写多少秒：一小时。再长就等于不设 +pub const MAX_IDLE_TIMEOUT_SECS: u64 = 3600; /// 等空位最多写多少秒。等的时候客户端一个字节都收不到，再长的话它先超时了 pub const MAX_SLOT_WAIT_SECS: u64 = 300; @@ -137,10 +134,10 @@ impl Failover { MAX_PAUSE_SECS, ), ( - "stream_start_wait_secs", - self.stream_start_wait_secs, - 1, - MAX_STREAM_START_WAIT_SECS, + "idle_timeout_secs", + self.idle_timeout_secs, + MIN_IDLE_TIMEOUT_SECS, + MAX_IDLE_TIMEOUT_SECS, ), ("slot_wait_secs", self.slot_wait_secs, 0, MAX_SLOT_WAIT_SECS), ]; @@ -148,10 +145,4 @@ impl Failover { .into_iter() .find(|(_, v, min, max)| v < min || v > max) } - - /// 开着「开头慢就换下一家」、等待却短于 [`MIN_SLOW_START_WAIT_SECS`]：写的等待秒数。 - pub(crate) fn slow_start_too_short(&self) -> Option { - (self.next_on_slow_start && self.stream_start_wait_secs < MIN_SLOW_START_WAIT_SECS) - .then_some(self.stream_start_wait_secs) - } } diff --git a/crates/tw-config/src/lib.rs b/crates/tw-config/src/lib.rs index a4042271..32d48cf4 100644 --- a/crates/tw-config/src/lib.rs +++ b/crates/tw-config/src/lib.rs @@ -1201,8 +1201,7 @@ pub fn write(path: &Path, cfg: &Config) -> Result<(), WriteError> { } pub use failover::{ - Failover, MAX_PAUSE_SECS, MAX_SLOT_WAIT_SECS, MAX_STREAM_START_WAIT_SECS, - MIN_SLOW_START_WAIT_SECS, + Failover, MAX_IDLE_TIMEOUT_SECS, MAX_PAUSE_SECS, MAX_SLOT_WAIT_SECS, MIN_IDLE_TIMEOUT_SECS, }; pub use probes::{ClientProbes, ProbeAction}; pub use reload::{Rejected, Stage, stand_in, try_parse}; diff --git a/crates/tw-config/src/repair.rs b/crates/tw-config/src/repair.rs index b483e137..d7d9dd1b 100644 --- a/crates/tw-config/src/repair.rs +++ b/crates/tw-config/src/repair.rs @@ -243,6 +243,42 @@ mod tests { assert!(try_parse(&r.text).is_ok()); } + /// 0.67 删掉的两项故障转移设置（「开头慢就换下一家」和流开头的等待，并进了无响应超时）： + /// 写着它们的老配置读不进来，一键修复删掉它们，故障转移照默认值走 + #[test] + fn the_removed_slow_start_settings_are_offered_for_deletion() { + let bad = format!( + "{GOOD}failover:\n pause_secs: 90\n stream_start_wait_secs: 30\n next_on_slow_start: true\n" + ); + let r = repair(&bad).expect("该修得了"); + let got: Vec<_> = r + .fixes + .iter() + .map(|f| (f.kind, f.field.as_str(), f.value.as_deref())) + .collect(); + assert_eq!( + got, + vec![ + ( + FixKind::UnknownField, + "failover.stream_start_wait_secs", + Some("30") + ), + ( + FixKind::UnknownField, + "failover.next_on_slow_start", + Some("true") + ), + ] + ); + assert_eq!(r.text, format!("{GOOD}failover:\n pause_secs: 90\n")); + let cfg = try_parse(&r.text).unwrap(); + assert_eq!(cfg.failover.idle_timeout_secs, 300); + // 只写了这两项的，删完连 `failover:` 一起去掉 + let only = format!("{GOOD}failover:\n next_on_slow_start: false\n"); + assert_eq!(repair(&only).unwrap().text, GOOD); + } + #[test] fn a_misspelled_field_in_a_list_item_is_removed() { let bad = "version: 1\nlisten:\n control:\n key: c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00\nclients:\n - name: c\n key: tw-k\n colour: red\n"; diff --git a/crates/tw-config/src/validate.rs b/crates/tw-config/src/validate.rs index 622915a4..90707b2a 100644 --- a/crates/tw-config/src/validate.rs +++ b/crates/tw-config/src/validate.rs @@ -66,9 +66,6 @@ pub enum ValidationError { min: u64, max: u64, }, - /// 开着「开头慢就换下一家」，流开头的等待却短于 [`crate::MIN_SLOW_START_WAIT_SECS`] - #[error("{}", self.msg())] - SlowStartTooShort { secs: u64 }, #[error("{}", self.msg())] ControlKeyMissing, #[error("{}", self.msg())] @@ -273,12 +270,6 @@ impl ValidationError { "config.failover_range", field = field, value = value, min = min, max = max => "failover.{field} is {value}; it has to be between {min} and {max}" ), - SlowStartTooShort { secs } => msg!( - "config.slow_start_too_short", - secs = secs, min = crate::MIN_SLOW_START_WAIT_SECS => - "failover.stream_start_wait_secs is {secs} while failover.next_on_slow_start is on; \ - it has to be at least {min}, or ordinary answers are cut off before they start" - ), ControlKeyMissing => msg!( "config.control_key_missing" => "the configuration has no listen.control.key, the key the desktop app connects \ @@ -656,10 +647,6 @@ pub fn validate(cfg: &Config) -> Result<(), ValidationError> { max, }); } - // 开头慢就换下一家：等得太短的话，平常的回答还没开口就被切到下一家 - if let Some(secs) = cfg.failover.slow_start_too_short() { - return Err(ValidationError::SlowStartTooShort { secs }); - } // 控制面的钥匙。**缺了、短了、不是十六进制，整份配置都不收**，旧的继续 // 服务：一份没有钥匙的配置换进来，下一条连接谁都进不来 —— 包括要把它 // 改回去的那个界面；一把好猜的短钥匙和没有差不多 @@ -1498,15 +1485,18 @@ groups: let base = with_rules(&[], &[]); assert!(validate(&base).is_ok()); type Bend = fn(&mut crate::Failover); - let cases: [(&str, Bend); 5] = [ + let cases: [(&str, Bend); 6] = [ ("failures_to_pause", |f| f.failures_to_pause = 0), ("pause_secs", |f| f.pause_secs = 0), ("max_pause_secs", |f| { f.pause_secs = 120; f.max_pause_secs = 60; }), - ("stream_start_wait_secs", |f| { - f.stream_start_wait_secs = crate::MAX_STREAM_START_WAIT_SECS + 1 + ("idle_timeout_secs", |f| { + f.idle_timeout_secs = crate::MIN_IDLE_TIMEOUT_SECS - 1 + }), + ("idle_timeout_secs", |f| { + f.idle_timeout_secs = crate::MAX_IDLE_TIMEOUT_SECS + 1 }), ("slot_wait_secs", |f| { f.slot_wait_secs = crate::MAX_SLOT_WAIT_SECS + 1 @@ -1524,23 +1514,12 @@ groups: let mut x = base.clone(); x.failover.slot_wait_secs = 0; assert!(validate(&x).is_ok()); - } - - /// 开头慢就换下一家：开着时等待至少 5 秒，关着时 1 秒也照收（只是交得早） - #[test] - fn switching_on_a_slow_start_needs_a_long_enough_wait() { - let mut x = with_rules(&[], &[]); - x.failover.stream_start_wait_secs = 3; - assert!(validate(&x).is_ok(), "关着时不管"); - x.failover.next_on_slow_start = true; - let e = validate(&x).unwrap_err(); - assert!( - matches!(e, ValidationError::SlowStartTooShort { secs: 3 }), - "{e:?}" - ); - assert_eq!(e.msg().code, "config.slow_start_too_short"); - x.failover.stream_start_wait_secs = crate::MIN_SLOW_START_WAIT_SECS; - assert!(validate(&x).is_ok()); + // 无响应超时的两头都收 + for secs in [crate::MIN_IDLE_TIMEOUT_SECS, crate::MAX_IDLE_TIMEOUT_SECS] { + let mut x = base.clone(); + x.failover.idle_timeout_secs = secs; + assert!(validate(&x).is_ok(), "{secs}"); + } } /// 别名表：名字一个一个、不带通配、不撞内置前缀，每个别名列着别的名称， diff --git a/crates/tw-config/tests/manual/schema.rs b/crates/tw-config/tests/manual/schema.rs index ed18514e..dde9e529 100644 --- a/crates/tw-config/tests/manual/schema.rs +++ b/crates/tw-config/tests/manual/schema.rs @@ -1326,21 +1326,12 @@ pub fn sections() -> Vec
{ ), ), row( - "stream_start_wait_secs", + "idle_timeout_secs", Kind::Int, - Def::Is("15"), + Def::Is("300"), t( - "Seconds to hold a streamed answer until its first content arrives. An error before then moves the request to the next upstream; after this long, what has arrived is passed on. From 1 to 120.", - "流式回答在第一段内容到达前最多暂存的秒数。在此之前上游报错，请求换到下一家；超过这个时间，已收到的部分照常交给客户端。取值 1 到 120。", - ), - ), - row( - "next_on_slow_start", - Kind::Bool, - Def::Is("false"), - t( - "When a streamed answer still has no content `stream_start_wait_secs` after the request was sent, give up on that upstream and send the request to the next one. The last upstream always waits. The upstream given up on is not set aside. Needs `stream_start_wait_secs` of at least 5.", - "流式回答在请求发出 `stream_start_wait_secs` 秒后仍没有内容时，放弃这家上游，把请求交给下一家。最后一家总是等下去。被放弃的上游不会停用。开启时 `stream_start_wait_secs` 至少为 5。", + "Seconds an upstream may go without sending content before the gateway stops waiting for it. Counted from the moment the request is sent and started again by every piece of content: text, reasoning and tool calls count, keep-alives do not. A whole (non-streamed) answer counts from sending to the complete answer. Before any content has reached the client, the upstream counts as failed and the request moves to the next one; with none left, the client gets a timeout error. After content has reached the client, the answer ends with an error. From 30 to 3600.", + "上游多少秒没有发出内容，网关就不再等它。从请求发出的那一刻算起，每来一段内容重新计时：正文、推理、工具调用都算，心跳不算。整包（非流式）的回答从发出算到整份回来。还没有内容交给客户端时，这家上游记一次失败，请求换到下一家；没有下一家了，客户端收到超时错误。已经有内容交给客户端的，回答以一条错误收尾。取值 30 到 3600。", ), ), row( diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index 59af3e47..ba3e1352 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -126,6 +126,7 @@ pub fn router(state: ControlState) -> Router { .at(ep::SpeedQuote, speed_quote) .at(ep::SpeedRun, speed_run) .at(ep::RequestDetail, request_detail) + .at(ep::AbortRequest, abort_request) // 诊断包（脱敏纪律）。**只读，不写任何文件** .at(ep::Diagnostics, diagnostics::bundle) // 把一条真实请求变成回放用例。**录制不是新功能** —— @@ -137,6 +138,7 @@ pub fn router(state: ControlState) -> Router { .at(ep::Sessions, sessions) .at(ep::SessionDetail, session_detail) .at(ep::SessionTranscript, session_transcript) + .at(ep::AbortSession, abort_session) .at(ep::DryRun, dryrun::dry_run) // 为客户端发专用密钥。接管本身在桌面端做 .at(ep::ClientKey, clients::client_key) @@ -415,8 +417,7 @@ async fn overview(State(s): State) -> Json { no_balance_pause_secs: f.no_balance_pause_secs, quota_pause_secs: f.quota_pause_secs, rate_limit_max_pause_secs: f.rate_limit_max_pause_secs, - stream_start_wait_secs: f.stream_start_wait_secs, - next_on_slow_start: f.next_on_slow_start, + idle_timeout_secs: f.idle_timeout_secs, slot_wait_secs: f.slot_wait_secs, } }, @@ -1713,6 +1714,42 @@ async fn session_detail( })) } +/// 中止一个在跑的请求（见 `tw_gateway::abort`）。**叫停是立刻的**：和上游的连接在请求自己 +/// 那边断开，结局（`RequestFailed`，来源 `aborted`）随后从事件流上到。已经结束了的、从来 +/// 没有过的、跑在 WebSocket 连接上的，都是不在跑：404 +async fn abort_request( + State(s): State, + axum::extract::Path(id): axum::extract::Path, +) -> Result, Fail> { + s.gateway.aborts.request(id).map_err(|_| { + fail( + StatusCode::NOT_FOUND, + msg!("control.request_not_running", id = id => "Request {id} is not in progress."), + ) + })?; + tracing::info!(id, "aborted a request over the control plane"); + Ok(Json(tw_api::Aborted { requests: vec![id] })) +} + +/// 中止一次会话里所有在跑的请求。一个都没有是 404 +async fn abort_session( + State(s): State, + axum::extract::Path(id): axum::extract::Path, +) -> Result, Fail> { + let requests = s.gateway.aborts.session(&id); + if requests.is_empty() { + return Err(fail( + StatusCode::NOT_FOUND, + msg!( + "control.session_not_running", id = id => + "Session {id} has no request in progress." + ), + )); + } + tracing::info!(session = %id, count = requests.len(), "aborted a session's requests over the control plane"); + Ok(Json(tw_api::Aborted { requests })) +} + fn no_such_session(id: &str) -> Fail { fail( StatusCode::NOT_FOUND, diff --git a/crates/tw-control/tests/resources.rs b/crates/tw-control/tests/resources.rs index 8333eee7..01335248 100644 --- a/crates/tw-control/tests/resources.rs +++ b/crates/tw-control/tests/resources.rs @@ -37,9 +37,13 @@ routes: struct Bed { dir: tempfile::TempDir, app: axum::Router, + gateway: tw_gateway::AppState, } impl Bed { + fn gateway(&self) -> &tw_gateway::AppState { + &self.gateway + } fn file(&self) -> String { std::fs::read_to_string(self.dir.path().join("config.yaml")).unwrap() } @@ -59,7 +63,7 @@ fn bed(yaml: &str) -> Bed { shutdown: Default::default(), remote: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), - gateway: gw, + gateway: gw.clone(), store: None, started: std::time::Instant::now(), price_updater: Default::default(), @@ -69,6 +73,7 @@ fn bed(yaml: &str) -> Bed { Bed { app: tw_control::router(state), dir: d, + gateway: gw, } } @@ -179,8 +184,10 @@ async fn failover_settings_show_their_defaults_and_take_an_edit() { let f = &json(&body)["failover"]; assert_eq!(f["failures_to_pause"], 3, "{body}"); assert_eq!(f["pause_secs"], 60); - assert_eq!(f["stream_start_wait_secs"], 15); + assert_eq!(f["idle_timeout_secs"], 300); assert_eq!(f["slot_wait_secs"], 30); + assert!(f.get("stream_start_wait_secs").is_none(), "{body}"); + assert!(f.get("next_on_slow_start").is_none(), "{body}"); let (st, body) = call( &b.app, @@ -204,7 +211,7 @@ async fn failover_settings_show_their_defaults_and_take_an_edit() { "PATCH", "/config", serde_json::json!({ - "ops": [{ "op": "replace", "path": "/failover/stream_start_wait_secs", "value": 0 }], + "ops": [{ "op": "replace", "path": "/failover/idle_timeout_secs", "value": 10 }], }), ) .await; @@ -225,46 +232,80 @@ async fn failover_settings_show_their_defaults_and_take_an_edit() { assert_eq!(b.parsed().failover.slot_wait_secs, 0); } -/// 开头慢就换下一家：默认关；打开要等得够久，等得太短的被拒 +/// 无响应超时：概览给出真在用的数，改了就生效；删掉的两项写不进去 #[tokio::test] -async fn switching_on_a_slow_start_is_shown_and_needs_a_long_enough_wait() { +async fn the_idle_timeout_is_shown_and_takes_an_edit() { let b = bed(BASE); - let (_, body) = call(&b.app, "GET", "/overview", serde_json::Value::Null).await; - assert_eq!( - json(&body)["failover"]["next_on_slow_start"], - false, - "{body}" - ); - let (st, body) = call( &b.app, "PATCH", "/config", serde_json::json!({ - "ops": [{ "op": "replace", "path": "/failover/next_on_slow_start", "value": true }], + "ops": [{ "op": "replace", "path": "/failover/idle_timeout_secs", "value": 600 }], }), ) .await; assert_eq!(st, StatusCode::OK, "{body}"); - assert!(b.parsed().failover.next_on_slow_start); + assert_eq!(b.parsed().failover.idle_timeout_secs, 600); let (_, body) = call(&b.app, "GET", "/overview", serde_json::Value::Null).await; - assert_eq!( - json(&body)["failover"]["next_on_slow_start"], - true, - "{body}" - ); + assert_eq!(json(&body)["failover"]["idle_timeout_secs"], 600, "{body}"); let (st, body) = call( &b.app, "PATCH", "/config", serde_json::json!({ - "ops": [{ "op": "replace", "path": "/failover/stream_start_wait_secs", "value": 3 }], + "ops": [{ "op": "replace", "path": "/failover/next_on_slow_start", "value": true }], }), ) .await; assert_eq!(st, StatusCode::BAD_REQUEST, "{body}"); - assert!(body.contains("config.slow_start_too_short"), "{body}"); + assert!(body.contains("config.unknown_field"), "{body}"); +} + +/// 中止：没有在跑的请求、会话，是 404，各有各的码 +#[tokio::test] +async fn aborting_what_is_not_running_is_a_404() { + let b = bed(BASE); + let (st, body) = call(&b.app, "POST", "/request/42/abort", serde_json::Value::Null).await; + assert_eq!(st, StatusCode::NOT_FOUND, "{body}"); + assert!(body.contains("control.request_not_running"), "{body}"); + let (st, body) = call( + &b.app, + "POST", + "/sessions/abc-1/abort", + serde_json::Value::Null, + ) + .await; + assert_eq!(st, StatusCode::NOT_FOUND, "{body}"); + assert!(body.contains("control.session_not_running"), "{body}"); +} + +/// 中止在跑的：叫停，交回叫停了哪几个 +#[tokio::test] +async fn aborting_a_running_request_or_session_says_which_were_stopped() { + let b = bed(BASE); + let gateway = b.gateway(); + let (one, two) = ( + tw_gateway::abort::Switch::default(), + tw_gateway::abort::Switch::default(), + ); + let _a = gateway.aborts.enter(7, Some("s-1".into()), one.clone()); + let _b = gateway.aborts.enter(9, Some("s-1".into()), two.clone()); + let (st, body) = call(&b.app, "POST", "/request/7/abort", serde_json::Value::Null).await; + assert_eq!(st, StatusCode::OK, "{body}"); + assert_eq!(json(&body)["requests"], serde_json::json!([7])); + assert!(one.thrown() && !two.thrown()); + let (st, body) = call( + &b.app, + "POST", + "/sessions/s-1/abort", + serde_json::Value::Null, + ) + .await; + assert_eq!(st, StatusCode::OK, "{body}"); + assert_eq!(json(&body)["requests"], serde_json::json!([7, 9])); + assert!(two.thrown()); } // ─────────────────────────────────────────────────────────── 上游 diff --git a/crates/tw-dialect/src/gemini/response.rs b/crates/tw-dialect/src/gemini/response.rs index bd7f798a..9a6f5389 100644 --- a/crates/tw-dialect/src/gemini/response.rs +++ b/crates/tw-dialect/src/gemini/response.rs @@ -166,6 +166,8 @@ pub fn error_body(status: u16, message: &str) -> Value { 403 => "PERMISSION_DENIED", 404 => "NOT_FOUND", 429 => "RESOURCE_EXHAUSTED", + // Google 的接口给「操作被取消」用的就是 499 + 499 => "CANCELLED", 502 | 503 | 529 => "UNAVAILABLE", 504 => "DEADLINE_EXCEEDED", _ => "INTERNAL", diff --git a/crates/tw-gateway/src/abort.rs b/crates/tw-gateway/src/abort.rs new file mode 100644 index 00000000..f4026cc2 --- /dev/null +++ b/crates/tw-gateway/src/abort.rs @@ -0,0 +1,186 @@ +//! 手动中止在跑的请求（控制面的 `POST /request/{id}/abort`、`POST /sessions/{id}/abort`）。 +//! +//! 每个请求进门时带一个开关（[`Switch`]），发出开始事件、有了请求号时登记在这里（见 +//! [`Aborts::enter`]）。叫停就是扳开关：在等上游的那几处（发出去等响应头、等开头、等空位） +//! 和交回答的那个循环都看着它，扳了就**丢掉和上游的那条连接**（上游不再生成），给客户端 +//! 回一个它那种格式的错误，请求记成手动中止（[`tw_api::FailureSource::Aborted`]）。 +//! +//! **不是上游的错**：那一家不停用、不算失败，快慢样本也不记。 +//! +//! 登记跟着请求的结局走（挂在 [`crate::ending::Ending`] 上）：结局报了、或者被丢掉了，登记 +//! 就没了。所以「在不在这张表里」就是「这个请求还在不在跑」，叫停一个已经结束的请求得到 +//! 的是「没有在跑」，而不是一次什么都没发生的成功。 +//! +//! WebSocket 那条路不登记：一轮回答跑在一条长连接上，叫停一轮要连带处置整条连接，这一版 +//! 不做。 + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +use tokio::sync::Notify; + +/// 一个请求的中止开关。克隆出来的是同一个。 +#[derive(Clone, Default)] +pub struct Switch(Arc); + +#[derive(Default)] +struct Inner { + thrown: AtomicBool, + notify: Notify, +} + +impl Switch { + /// 扳下去：等着它的都醒过来，之后再等的立刻返回。扳过一次再扳什么都不变 + pub fn throw(&self) { + self.0.thrown.store(true, Ordering::SeqCst); + self.0.notify.notify_waiters(); + } + + /// 扳过了没有 + pub fn thrown(&self) -> bool { + self.0.thrown.load(Ordering::SeqCst) + } + + /// 等它被扳下去。**没人扳就一直等**：放在 `select!` 里和真正要等的东西并排 + pub async fn wait(&self) { + loop { + let notified = self.0.notify.notified(); + let mut notified = std::pin::pin!(notified); + // 先登记再看：看完和开始等之间扳下去的，`notify_waiters` 也叫得醒它 + notified.as_mut().enable(); + if self.thrown() { + return; + } + notified.await; + } + } +} + +/// 一个登记着的请求。 +struct Entry { + /// 它属于哪次会话（开始事件里的那个）。认不出会话的没有 + session: Option, + switch: Switch, +} + +/// 此刻在跑、可以叫停的请求。**跨重载存活**（挂在 `AppState` 上）：改一条规则不该让 +/// 在跑的请求叫不停。 +#[derive(Default)] +pub struct Aborts { + entries: Mutex>, +} + +/// 登记。**丢掉它就是注销**：跟着请求的结局走，结局报了、请求被丢掉了，它跟着没了。 +#[must_use = "dropping the registration takes the request off the list"] +pub struct Registered { + aborts: Arc, + id: u64, +} + +impl Drop for Registered { + fn drop(&mut self) { + self.aborts.lock().remove(&self.id); + } +} + +/// 要叫停的不在跑（从没有过、已经结束了）。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct NotRunning; + +impl Aborts { + fn lock(&self) -> std::sync::MutexGuard<'_, HashMap> { + // 锁中毒了照样拿里面的表：叫不停一个请求，好过一个不转发的网关 + self.entries.lock().unwrap_or_else(|p| p.into_inner()) + } + + /// 请求 `id`（会话 `session`）开始了，用 `switch` 叫停它。 + pub fn enter(self: &Arc, id: u64, session: Option, switch: Switch) -> Registered { + self.lock().insert(id, Entry { session, switch }); + Registered { + aborts: Arc::clone(self), + id, + } + } + + /// 叫停请求 `id`。 + pub fn request(&self, id: u64) -> Result<(), NotRunning> { + let map = self.lock(); + let e = map.get(&id).ok_or(NotRunning)?; + e.switch.throw(); + Ok(()) + } + + /// 叫停会话 `session` 里所有在跑的请求，交回它们的号，从小到大（也就是开始的先后）。 + /// 一个都没有是空的 + pub fn session(&self, session: &str) -> Vec { + let map = self.lock(); + let mut ids: Vec = map + .iter() + .filter(|(_, e)| e.session.as_deref() == Some(session)) + .map(|(id, e)| { + e.switch.throw(); + *id + }) + .collect(); + ids.sort_unstable(); + ids + } + + /// 此刻登记着几个 + pub fn len(&self) -> usize { + self.lock().len() + } + + /// 一个都没有 + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[tokio::test] + async fn a_thrown_switch_wakes_who_waits_and_whoever_comes_later() { + let s = Switch::default(); + let waiting = tokio::spawn({ + let s = s.clone(); + async move { s.wait().await } + }); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(!waiting.is_finished()); + s.throw(); + tokio::time::timeout(Duration::from_secs(1), waiting) + .await + .expect("扳下去要叫醒在等的") + .unwrap(); + // 扳过之后再等的立刻返回 + tokio::time::timeout(Duration::from_millis(50), s.wait()) + .await + .expect("扳过的开关不用再等"); + } + + #[test] + fn only_what_is_running_can_be_aborted_and_a_session_takes_all_of_its_own() { + let aborts = Arc::new(Aborts::default()); + let (a, b, c) = (Switch::default(), Switch::default(), Switch::default()); + let ra = aborts.enter(7, Some("s1".into()), a.clone()); + let _rb = aborts.enter(3, Some("s1".into()), b.clone()); + let _rc = aborts.enter(9, None, c.clone()); + assert_eq!(aborts.len(), 3); + + assert_eq!(aborts.request(42), Err(NotRunning)); + assert_eq!(aborts.session("s1"), vec![3, 7]); + assert!(a.thrown() && b.thrown() && !c.thrown()); + assert!(aborts.session("nope").is_empty()); + + // 结局报了、登记丢掉了：再叫停就是不在跑 + drop(ra); + assert_eq!(aborts.request(7), Err(NotRunning)); + assert_eq!(aborts.request(9), Ok(())); + assert!(c.thrown()); + } +} diff --git a/crates/tw-gateway/src/affinity.rs b/crates/tw-gateway/src/affinity.rs index 7a3384a4..b1698d0d 100644 --- a/crates/tw-gateway/src/affinity.rs +++ b/crates/tw-gateway/src/affinity.rs @@ -263,6 +263,22 @@ impl Affinity { self.upsert(&t.key, now_ms, |e| e.answer = Some(answer)); } + /// `provider` 在这段对话上没了声音（无响应超时，见 `server::pipeline::idle`）：上次回答它 + /// 的要是这一家，就不再记着 —— 这一轮接下来的请求照常排序，不会因为「上次回答的就是它」 + /// 又被送回去。这一轮的路由决定照旧沿用 + pub fn left(&self, c: &Conversation, provider: &str) { + self.forget(&c.key, provider); + } + + fn forget(&self, key: &str, provider: &str) { + let mut map = self.lock(); + if let Some(e) = map.get_mut(key) + && e.answer.as_ref().is_some_and(|a| a.provider == provider) + { + e.answer = None; + } + } + fn upsert(&self, key: &str, now_ms: u64, f: impl FnOnce(&mut Entry)) { let mut map = self.lock(); if !map.contains_key(key) && map.len() >= CAP { @@ -315,6 +331,11 @@ impl Ticket { pub fn answered(self, cache: u64) { self.store.answered(&self, cache, crate::server::now_ms()); } + + /// 这一家答到一半没了声音：上次回答这段对话的要是它，不再记着（见 [`Affinity::left`]） + pub fn left(self) { + self.store.forget(&self.key, &self.provider); + } } #[cfg(test)] @@ -328,6 +349,28 @@ mod tests { Conversation::new("default", "对话", Turn { number, within }) } + /// 没了声音的那一家：这一轮不再留在它那儿，别家的回答不受影响 + #[test] + fn a_provider_that_went_quiet_is_not_stayed_on() { + let a = Arc::new(Affinity::default()); + let c = conv(1, true); + a.ticket(&c, None, "甲").answered(5000); + let stay = |a: &Affinity| { + let mut cands = vec!["乙".to_string(), "甲".to_string()]; + a.stay(&c, None, &mut cands, |_| true, crate::server::now_ms()) + .map(|_| cands[0].clone()) + }; + assert_eq!(stay(&a).as_deref(), Some("甲")); + a.left(&c, "乙"); + assert_eq!(stay(&a).as_deref(), Some("甲"), "说的不是它，照旧"); + a.left(&c, "甲"); + assert_eq!(stay(&a), None); + // 通过回答的票据放开也一样 + a.ticket(&c, None, "甲").answered(5000); + a.ticket(&c, None, "甲").left(); + assert_eq!(stay(&a), None); + } + fn engine() -> Arc { Arc::new(Engine::with_default_rules( vec!["甲".into(), "乙".into()], diff --git a/crates/tw-gateway/src/ending.rs b/crates/tw-gateway/src/ending.rs index dc1e2245..6b66958b 100644 --- a/crates/tw-gateway/src/ending.rs +++ b/crates/tw-gateway/src/ending.rs @@ -79,6 +79,9 @@ pub struct Ending { /// 这段对话这一次由谁回答（见 [`crate::affinity`]）。**成功走完了才记**：失败的、 /// 半路断了的不算回答过，下一次照常排序 answer: Option, + /// 手动中止的登记（见 [`crate::abort`]）。**跟着结局走**：结局报了、或者被丢掉了，它跟着 + /// 没了，这个请求就不再算在跑 + abort: Option, /// 报过了。**只能报一次** told: bool, } @@ -222,6 +225,7 @@ impl Ending { upstream_error: None, refusal: None, answer: None, + abort: None, told: false, } } @@ -292,6 +296,19 @@ impl Ending { self.answer = Some(ticket); } + /// 这个请求可以手动中止了：登记交给结局保管（见 [`crate::abort`]）。 + pub fn abortable(&mut self, r: crate::abort::Registered) { + self.abort = Some(r); + } + + /// 回答的那一家答到一半没了声音（无响应超时，见 `server::pipeline::idle`）：这段对话这一轮 + /// 不再留在它那儿 + pub fn went_quiet(&mut self) { + if let Some(t) = self.answer.take() { + t.left(); + } + } + /// 上游的响应头到了。从这里起，客户端再走掉，报出去的取消带着状态码。 pub fn responded(&mut self, status: u16) { self.status = Some(status); diff --git a/crates/tw-gateway/src/error.rs b/crates/tw-gateway/src/error.rs index aa2af8d9..e1c88632 100644 --- a/crates/tw-gateway/src/error.rs +++ b/crates/tw-gateway/src/error.rs @@ -47,6 +47,14 @@ pub enum Source { /// 请求体超过网关的上限（见 `server::intake`）。**413，不是 400**：Anthropic 的格式里 /// 它是 `request_too_large`，和 Anthropic 自己嫌请求太大时一样。对外的词表里算 `request` TooLarge, + /// 上游在 `failover.idle_timeout_secs` 里一直没有内容，候选也用完了。**504，不是 502**： + /// 上游没坏，是没在时限里回话 —— Anthropic 的格式里它是 `timeout_error`，Gemini 是 + /// `DEADLINE_EXCEEDED`。对外的词表里算 `upstream` + Timeout, + /// 在界面上手动中止的（见 [`crate::abort`]）。**499**：客户端的 SDK 不重试 4xx + /// （408、409、429 除外），不会把用户叫停的请求自己再发一遍；Google 的接口给「操作被 + /// 取消」用的也是它。对外的词表里是 `aborted` + Aborted, } /// 上游都满着时告诉客户端过几秒再来（`Retry-After`）。 @@ -66,6 +74,8 @@ impl Source { Source::RateLimited | Source::Busy => "rate_limited", Source::Denied => "denied", Source::NotSupported | Source::TooLarge => "request", + Source::Timeout => "upstream", + Source::Aborted => "aborted", } } fn status(&self) -> StatusCode { @@ -79,6 +89,8 @@ impl Source { Source::Denied => StatusCode::FORBIDDEN, Source::NotSupported => StatusCode::NOT_IMPLEMENTED, Source::TooLarge => StatusCode::PAYLOAD_TOO_LARGE, + Source::Timeout => StatusCode::GATEWAY_TIMEOUT, + Source::Aborted => StatusCode::from_u16(499).expect("499 is a valid status code"), } } } @@ -180,6 +192,16 @@ impl GatewayError { pub fn too_large(detail: Msg) -> Self { Self::new(Source::TooLarge, detail) } + pub fn timeout(detail: Msg) -> Self { + Self::new(Source::Timeout, detail) + } + /// 手动中止（见 [`crate::abort`]）。句子只有一句，码是 [`tw_api::ABORTED`] + pub fn aborted() -> Self { + Self::new( + Source::Aborted, + tw_types::msg!("gw.request.aborted" => "The request was aborted by the user."), + ) + } } impl GatewayError { @@ -409,6 +431,30 @@ mod tests { assert_eq!(head(&r, "x-should-retry"), None); } + #[tokio::test] + async fn a_timeout_is_504_and_an_abort_is_499() { + let (status, slug, json) = + body_of(GatewayError::timeout(msg!("t.x" => "quiet")).in_dialect(Dialect::Gemini)) + .await; + assert_eq!(status, StatusCode::GATEWAY_TIMEOUT); + assert_eq!(slug, "upstream"); + assert_eq!(json["error"]["status"], "DEADLINE_EXCEEDED"); + let (_, _, json) = body_of(GatewayError::timeout(msg!("t.x" => "quiet"))).await; + assert_eq!(json["error"]["type"], "timeout_error"); + let (status, slug, json) = body_of(GatewayError::aborted()).await; + assert_eq!(status.as_u16(), 499); + assert_eq!(slug, "aborted"); + assert_eq!(GatewayError::aborted().detail.code, tw_api::ABORTED); + assert!( + json["error"]["message"] + .as_str() + .unwrap() + .starts_with("[ThinkWatch]") + ); + let (_, _, json) = body_of(GatewayError::aborted().in_dialect(Dialect::Gemini)).await; + assert_eq!(json["error"]["status"], "CANCELLED"); + } + #[test] fn a_responses_stream_is_told_with_response_failed() { // Chat 形状的 `{"error":…}` 在 Responses 的流里是一帧没人认的数据，客户端 diff --git a/crates/tw-gateway/src/latency.rs b/crates/tw-gateway/src/latency.rs index 2c8567ad..e35d6bc8 100644 --- a/crates/tw-gateway/src/latency.rs +++ b/crates/tw-gateway/src/latency.rs @@ -15,7 +15,7 @@ //! `server::pipeline::opening`），最后一家直接转发 —— 量的都是同一段。 //! - **只有流式回答有样本**：整包的回答只有「全到了」那一个时刻，分不出哪段在排队、哪段 //! 在说话。 -//! - **开头慢被放弃的那一家**（`failover.next_on_slow_start`）记它被给的那段时间：它至少 +//! - **没有内容超时被放弃的那一家**（`failover.idle_timeout_secs`）记它被给的那段时间：它至少 //! 这么慢，记成这个数，它就排到慢的那一头去。什么都不记的话，它留着的还是从前快的 //! 样本，下一个请求还先发给它。 //! diff --git a/crates/tw-gateway/src/lib.rs b/crates/tw-gateway/src/lib.rs index 344e80b7..128a93ed 100644 --- a/crates/tw-gateway/src/lib.rs +++ b/crates/tw-gateway/src/lib.rs @@ -5,6 +5,7 @@ //! 就把接缝留对**：`forward` 已经按「选中的 provider」取参数，M1 加路由 //! 时只需要换掉挑选逻辑。 +pub mod abort; pub mod access; pub mod affinity; pub mod answer_model; @@ -39,6 +40,7 @@ pub mod oauth; pub mod outbound; pub mod plugin; pub mod probe; +mod pulse; pub mod quota; pub mod quote; pub mod seal; diff --git a/crates/tw-gateway/src/pulse.rs b/crates/tw-gateway/src/pulse.rs new file mode 100644 index 00000000..02e6e263 --- /dev/null +++ b/crates/tw-gateway/src/pulse.rs @@ -0,0 +1,378 @@ +//! 回答在不在往前走：上游来的一段是真内容，还是心跳。 +//! +//! 无响应超时（`failover.idle_timeout_secs`）只认这个：每来一段真内容重新计时，心跳不算。 +//! 心跳也算的话，一家只发心跳、永远不出内容的上游会一直挂着，客户端跟着干等。流开头的等待 +//! （`server::pipeline::opening`）用的是同一个判据：开头的例行事件和心跳之后，第一段真内容 +//! 才算上游开口了。 +//! +//! **心跳**（不算）： +//! +//! - SSE 的注释行（`: keep-alive`、`: OPENROUTER PROCESSING`）和空的 `data:`； +//! - 叫心跳的事件，哪种格式都一样：`ping`、`keepalive`、`heartbeat`； +//! - Anthropic：`message_start`（收到请求就发的开场）、`ping`； +//! - Responses：`response.created`、`response.in_progress`、`response.queued`，和 Codex 后端 +//! 开头报额度的 `codex.*`； +//! - Chat Completions：只有角色、内容为空的块，`choices` 为空又没有用量的块； +//! - Gemini：候选里一段内容都没有、也没有结束原因的块（只带 `usageMetadata` 的那种）； +//! - Bedrock：`messageStart`。 +//! +//! **真内容**（重新计时）：正文、推理（字和摘要）、工具调用（开头和参数）、拒答、结束原因、 +//! 收尾的用量、`[DONE]`、上游在流里报的错，以及**认不出来的一切** —— 认错了往宽里错：把 +//! 心跳当内容，最多是一家坏了的上游多挂一阵；把内容当心跳，会把一个好好在答的回答掐断。 +//! +//! 不是流的回答（整包的 JSON）：有一个不是空白的字节就算 —— 有的中转站在整包生成完之前 +//! 隔一阵发一个空格保活。 + +use tw_dialect::ir::Dialect; + +/// 一个事件攒到这么大还没收齐，就当它是内容、不再攒：一个坏了的上游不能让这里无限长 +const MAX_EVENT: usize = 1024 * 1024; + +/// 哪种格式都算心跳的事件名 +const HEARTBEATS: &[&str] = &["ping", "keepalive", "keep-alive", "keep_alive", "heartbeat"]; + +/// 一个解析好的事件（`ty` 是 `event:` 行写的名字，没有就是 JSON 里的 `type`）是不是让回答 +/// 往前走了一步。上游报的错也算（流跟着就结束了） +pub(crate) fn advances(dialect: Dialect, ty: &str, v: &serde_json::Value) -> bool { + if HEARTBEATS.contains(&ty) { + return false; + } + match dialect { + Dialect::Anthropic => ty != "message_start", + Dialect::Responses => { + !matches!( + ty, + "response.created" | "response.in_progress" | "response.queued" + ) && !ty.starts_with("codex.") + } + Dialect::Chat => chat_advances(v), + Dialect::Gemini => gemini_advances(v), + Dialect::Bedrock => ty != "messageStart", + } +} + +/// Chat Completions 的一块：有错误、有内容、有结束原因，或者是收尾的那块用量 +fn chat_advances(v: &serde_json::Value) -> bool { + if v.get("error").is_some() { + return true; + } + let Some(choices) = v.get("choices").and_then(|c| c.as_array()) else { + // 没有 `choices`：只有用量的收尾块算，别的认不出来，也算 + return true; + }; + if choices.is_empty() { + return v.get("usage").is_some_and(|u| !u.is_null()); + } + choices.iter().any(|c| { + let finished = c.get("finish_reason").is_some_and(|f| !f.is_null()); + let said = c.get("delta").is_some_and(|d| { + [ + "content", + "reasoning_content", + "reasoning", + "tool_calls", + "function_call", + "refusal", + "audio", + ] + .iter() + .any(|k| d.get(*k).is_some_and(filled)) + }); + finished || said + }) +} + +/// Gemini 的一块：有错误、被拦下，或者某个候选有内容、有结束原因 +fn gemini_advances(v: &serde_json::Value) -> bool { + if v.get("error").is_some() || v.get("promptFeedback").is_some() { + return true; + } + let Some(candidates) = v.get("candidates").and_then(|c| c.as_array()) else { + // 只带用量、模型版本的块不是内容；别的认不出来的算 + return !v.as_object().is_some_and(|o| { + o.keys() + .all(|k| matches!(k.as_str(), "usageMetadata" | "modelVersion" | "responseId")) + }); + }; + candidates.iter().any(|c| { + let finished = c.get("finishReason").is_some_and(|f| !f.is_null()); + let parts = c + .pointer("/content/parts") + .and_then(|p| p.as_array()) + .is_some_and(|parts| { + parts.iter().any(|p| { + p.as_object().is_some_and(|o| { + o.iter().any(|(k, x)| k.as_str() != "thought" && filled(x)) + }) + }) + }); + finished || parts + }) +} + +/// 有东西：不是 null、空串、空数组、空对象 +fn filled(x: &serde_json::Value) -> bool { + match x { + serde_json::Value::Null => false, + serde_json::Value::String(s) => !s.is_empty(), + serde_json::Value::Array(a) => !a.is_empty(), + serde_json::Value::Object(o) => !o.is_empty(), + _ => true, + } +} + +/// 一个 SSE 事件（`event:` 行的名字、拼好的 `data:`）是不是真内容。 +pub(crate) fn event_advances(dialect: Dialect, event: Option<&str>, data: &str) -> bool { + // 带名字的格式先看名字：不用为每一个 token 解析一遍 JSON + if let Some(name) = event.map(str::trim).filter(|n| !n.is_empty()) { + if HEARTBEATS.contains(&name) { + return false; + } + if matches!( + dialect, + Dialect::Anthropic | Dialect::Responses | Dialect::Bedrock + ) { + return advances(dialect, name, &serde_json::Value::Null); + } + } + let data = data.trim(); + if data.is_empty() { + return false; + } + if data == "[DONE]" { + return true; + } + let Ok(v) = serde_json::from_str::(data) else { + return true; + }; + let ty = event + .map(str::trim) + .filter(|n| !n.is_empty()) + .or_else(|| v.get("type").and_then(|t| t.as_str())) + .unwrap_or_default() + .to_string(); + advances(dialect, &ty, &v) +} + +/// 跟着一条回答，一块一块地看有没有真内容。 +pub(crate) struct Pulse { + /// 上游说的是 SSE 时是它的格式；整包的回答是 None + dialect: Option, + /// 还没收到换行的那半行 + line: Vec, + /// 这个事件的名字，和攒着的 `data:` + event: Option, + data: String, +} + +impl Pulse { + /// `dialect`：上游的回答是这种格式的 SSE（Bedrock 的二进制帧先转成 SSE 再喂）。整包的 + /// 回答给 None + pub(crate) fn new(dialect: Option) -> Self { + Self { + dialect, + line: Vec::new(), + event: None, + data: String::new(), + } + } + + /// 这一块里有没有真内容。**一块里有一个就够了**：后面的事件只拆不判，下一块从头再判 + pub(crate) fn feed(&mut self, chunk: &[u8]) -> bool { + let Some(dialect) = self.dialect else { + return chunk.iter().any(|b| !b.is_ascii_whitespace()); + }; + let mut found = false; + let mut rest = chunk; + while let Some(nl) = memchr::memchr(b'\n', rest) { + let (head, tail) = rest.split_at(nl); + rest = &tail[1..]; + let line = if self.line.is_empty() { + head + } else { + self.line.extend_from_slice(head); + self.line.as_slice() + }; + let line = line.strip_suffix(b"\r").unwrap_or(line); + let ended = line.is_empty(); + if !ended { + Self::field(&mut self.event, &mut self.data, line); + } + self.line.clear(); + if ended { + let (event, data) = (self.event.take(), std::mem::take(&mut self.data)); + if !found && (event.is_some() || !data.is_empty()) { + found = event_advances(dialect, event.as_deref(), &data); + } + } + } + self.line.extend_from_slice(rest); + if self.line.len() + self.data.len() > MAX_EVENT { + self.line.clear(); + self.event = None; + self.data.clear(); + return true; + } + found + } + + /// 一行 SSE：`event:` 记名字，`data:` 攒起来，注释和别的字段不管 + fn field(event: &mut Option, data: &mut String, line: &[u8]) { + if let Some(v) = line.strip_prefix(b"event:") { + *event = Some(String::from_utf8_lossy(v).trim().to_string()); + } else if let Some(v) = line.strip_prefix(b"data:") { + let v = v.strip_prefix(b" ").unwrap_or(v); + if !data.is_empty() { + data.push('\n'); + } + data.push_str(&String::from_utf8_lossy(v)); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 一段 SSE 整个喂进去，有没有真内容 + fn sse(dialect: Dialect, text: &str) -> bool { + Pulse::new(Some(dialect)).feed(text.as_bytes()) + } + + #[test] + fn anthropic_pings_and_the_opening_are_not_content() { + let d = Dialect::Anthropic; + assert!(!sse(d, ": keep-alive\n\n")); + assert!(!sse(d, "event: ping\ndata: {\"type\":\"ping\"}\n\n")); + assert!(!sse( + d, + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{}}\n\n" + )); + // 不带 `event:` 行、只在 JSON 里写类型的也认得 + assert!(!sse(d, "data: {\"type\":\"ping\"}\n\n")); + for text in [ + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"Hm\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"a\"}}\n\n", + "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n", + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\n\n", + ] { + assert!(sse(d, text), "{text}"); + } + } + + #[test] + fn responses_progress_events_are_content_and_the_status_ones_are_not() { + let d = Dialect::Responses; + for quiet in [ + "event: response.created\ndata: {\"type\":\"response.created\"}\n\n", + "event: response.in_progress\ndata: {\"type\":\"response.in_progress\"}\n\n", + "event: response.queued\ndata: {\"type\":\"response.queued\"}\n\n", + "event: codex.rate_limits\ndata: {\"type\":\"codex.rate_limits\"}\n\n", + "data: {\"type\":\"keepalive\"}\n\n", + ] { + assert!(!sse(d, quiet), "{quiet}"); + } + for said in [ + "event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"item\":{\"type\":\"reasoning\"}}\n\n", + "event: response.reasoning_summary_text.delta\ndata: {\"type\":\"response.reasoning_summary_text.delta\",\"delta\":\"Thinking\"}\n\n", + "event: response.function_call_arguments.delta\ndata: {\"delta\":\"{\"}\n\n", + "event: response.output_text.delta\ndata: {\"delta\":\"Hi\"}\n\n", + "event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n", + "event: response.failed\ndata: {\"type\":\"response.failed\"}\n\n", + ] { + assert!(sse(d, said), "{said}"); + } + } + + #[test] + fn chat_role_only_and_empty_chunks_are_not_content() { + let d = Dialect::Chat; + for quiet in [ + "data: {\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\"}}]}\n\n", + "data: {\"choices\":[{\"index\":0,\"delta\":{}}]}\n\n", + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":null,\"tool_calls\":[]}}]}\n\n", + "data: {\"choices\":[]}\n\n", + "data: \n\n", + ": OPENROUTER PROCESSING\n\n", + ] { + assert!(!sse(d, quiet), "{quiet}"); + } + for said in [ + "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"Let me think\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"reasoning\":\"Let me think\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\"}}]}}]}\n\n", + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":3}}\n\n", + "data: {\"error\":{\"message\":\"busy\"}}\n\n", + "data: [DONE]\n\n", + "data: not json\n\n", + ] { + assert!(sse(d, said), "{said}"); + } + } + + #[test] + fn gemini_chunks_without_parts_are_not_content() { + let d = Dialect::Gemini; + for quiet in [ + "data: {\"usageMetadata\":{\"promptTokenCount\":10},\"modelVersion\":\"gemini-3-pro\"}\r\n\r\n", + "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"\"}]}}]}\r\n\r\n", + "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"\",\"thought\":true}]}}]}\r\n\r\n", + ] { + assert!(!sse(d, quiet), "{quiet}"); + } + for said in [ + "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hi\"}]}}]}\r\n\r\n", + "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hm\",\"thought\":true}]}}]}\r\n\r\n", + "data: {\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"ls\"}}]}}]}\r\n\r\n", + "data: {\"candidates\":[{\"finishReason\":\"STOP\"}]}\r\n\r\n", + "data: {\"promptFeedback\":{\"blockReason\":\"SAFETY\"}}\r\n\r\n", + ] { + assert!(sse(d, said), "{said}"); + } + } + + #[test] + fn bedrock_message_start_is_not_content() { + let d = Dialect::Bedrock; + assert!(!sse( + d, + "event: messageStart\ndata: {\"role\":\"assistant\"}\n\n" + )); + assert!(sse( + d, + "event: contentBlockDelta\ndata: {\"delta\":{\"text\":\"Hi\"}}\n\n" + )); + } + + #[test] + fn an_event_split_across_chunks_counts_when_it_is_complete() { + let mut p = Pulse::new(Some(Dialect::Anthropic)); + assert!(!p.feed(b"event: content_block_del")); + assert!(!p.feed(b"ta\r\ndata: {\"type\":\"content_block_delta\"}\r")); + assert!(p.feed(b"\n\r\n")); + // 一块里先是心跳、后是内容：算 + assert!(p.feed(b"event: ping\ndata: {}\n\nevent: content_block_delta\ndata: {}\n\n")); + // 只有心跳：不算 + assert!(!p.feed(b"event: ping\ndata: {}\n\n: keep-alive\n\n")); + } + + #[test] + fn a_whole_body_counts_any_byte_that_is_not_whitespace() { + let mut p = Pulse::new(None); + assert!(!p.feed(b" \n\t")); + assert!(p.feed(b" {\"id\":")); + } + + #[test] + fn an_event_that_never_ends_is_taken_as_content_and_dropped() { + let mut p = Pulse::new(Some(Dialect::Chat)); + let big = vec![b'x'; MAX_EVENT + 1]; + let mut first = b"data: ".to_vec(); + first.extend_from_slice(&big); + assert!(p.feed(&first)); + assert!(!p.feed(b"\n\n"), "攒的已经扔掉了"); + } +} diff --git a/crates/tw-gateway/src/server.rs b/crates/tw-gateway/src/server.rs index 645a658e..4a76389c 100644 --- a/crates/tw-gateway/src/server.rs +++ b/crates/tw-gateway/src/server.rs @@ -191,6 +191,8 @@ async fn passthrough( let body = intake::read(&headers, body, intake::MAX_BODY) .await .map_err(|e| e.in_dialect(dialect))?; + // 手动中止的开关（见 `crate::abort`）。管线发出开始事件时登记上它 + let abort = crate::abort::Switch::default(); let req = pipeline::Inbound { uri, query, @@ -201,8 +203,16 @@ async fn passthrough( dialect, started, from, + abort: abort.clone(), + }; + // 被手动中止时，**管线先自己收场**（`biased`）：在等上游的那几处它看着开关，丢掉那一跳、 + // 报完尝试链再返回。停在别处（等密钥的上限、取凭据）的，由这里整个丢掉，结局照样按 + // 手动中止报 + let result = tokio::select! { + biased; + r = pipeline::pipeline(state, rt, req, live, &mut ending) => r, + _ = abort.wait() => Err(GatewayError::aborted()), }; - let result = pipeline::pipeline(state, rt, req, live, &mut ending).await; if let Some(end) = ending.take() { match &result { Err(e) => end.failed(e.source.into(), e.detail.clone()), diff --git a/crates/tw-gateway/src/server/pipeline.rs b/crates/tw-gateway/src/server/pipeline.rs index f2e8addf..b8d36ff4 100644 --- a/crates/tw-gateway/src/server/pipeline.rs +++ b/crates/tw-gateway/src/server/pipeline.rs @@ -21,10 +21,10 @@ use tw_types::msg; mod admission; mod hop; +mod idle; mod opening; mod plug; mod relay; -mod slow; pub(crate) use hop::stream_fault; @@ -41,6 +41,8 @@ pub(super) struct Inbound { pub(super) dialect: tw_dialect::ir::Dialect, pub(super) started: std::time::Instant, pub(super) from: Sender, + /// 手动中止的开关（见 [`crate::abort`]）：开始之后登记上，等上游、交回答时看着它 + pub(super) abort: crate::abort::Switch, } /// 发出开始事件之后，后面几步都要用的。 @@ -54,6 +56,9 @@ struct Started { choice: Choice, /// 这是哪段对话（见 [`crate::affinity::identity`]）。认不出来是 None conversation: Option, + /// 这段对话的这一轮（见 [`crate::affinity`]）：一家没有内容超时了，这一轮就不再留在它 + /// 那儿。数不出轮次的没有 + turn: Option, /// 出站脱敏的账本：拦截档下按客户端原文编好了号，每一跳接着它换（见 /// [`crate::guard::look`]）。别的档位是空的 ledger: tw_guard::redact::replace::Ledger, @@ -160,7 +165,7 @@ pub(super) async fn pipeline( choice, wait_until, &decision, - fp.as_deref(), + (fp.as_deref(), conv.clone()), ending, ); (screening, started) @@ -814,7 +819,7 @@ fn start( choice: Choice, wait_until: tokio::time::Instant, decision: &tw_engine::Decision, - fp: Option<&str>, + (fp, turn): (Option<&str>, Option), ending: &mut Option, ) -> Started { // 熔断过滤。**只有一个候选时完全旁路**，全都熔断时 fail-open —— @@ -872,6 +877,7 @@ fn start( alive, choice, conversation: crate::affinity::identity(&req.headers, fp), + turn, ledger, found, hits: hits.filter(|_| redact_mode.acts()), @@ -922,13 +928,14 @@ fn open( let facts = &reading.facts; let id = state.bus.next_id(); let at_ms = now_ms(); + let session = fp.map(|fp| state.sessions.assign(fp, at_ms)); state.bus.emit(tw_api::Event::RequestStarted { id, client: req.client_name.clone(), // **旁证，不是身份。**只用来显示和判断「接管生效了吗」， // 不参与鉴权、路由、配额（见 crate::hint）。 client_hint: crate::hint::client_hint(&req.headers), - session: fp.map(|fp| state.sessions.assign(fp, at_ms)), + session: session.clone(), peer: req.from.peer.clone(), key_masked: req.from.key.clone(), route: choice.route.clone(), @@ -958,6 +965,8 @@ fn open( sink.clone(), ); end.redact_with(redaction.clone()); + // 从这一刻起可以手动中止：登记跟着结局走，结局报了就不在跑了（见 `crate::abort`） + end.abortable(state.aborts.enter(id, session, req.abort.clone())); *ending = Some(end); // 请求体交给观测层。**这时候它已经完整在内存里了**，所以这一步 diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index 2c29542d..5407b05e 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -10,7 +10,8 @@ //! 管这一跳的插件从客户端的原话起改，这一跳的转换、脱敏、发送用改过的那一份。换到下一 //! 家时从原话重来；同一家重发（OAuth 换 token、去封存）用这一跳定好的请求体，不重跑。 //! -//! 开着 `failover.next_on_slow_start` 时，开头迟迟没有内容的一家也换掉（见 [`super::slow`]）。 +//! 发出去之后 `failover.idle_timeout_secs` 还没有内容的一家也换掉（见 [`super::idle`]）；在等 +//! 上游的时候被手动中止的，丢掉这一跳、不再换（见 [`crate::abort`]）。 use bytes::Bytes; @@ -48,6 +49,9 @@ pub(super) struct Served<'a> { pub(super) slot: crate::slots::Slot, /// 这一跳发出去的那一刻。这一家的快慢样本从这里算起（见 [`crate::latency`]） pub(super) sent_at: std::time::Instant, + /// 这一家答没答上还不知道：最后一家的开头不压，成败的账等回程看到第一段内容才记（见 + /// `relay`）。没有内容就超时了记失败，和压着开头时一样 + pub(super) unconfirmed: bool, } /// 这个请求的着落。 @@ -146,8 +150,8 @@ pub(super) async fn try_upstreams<'a>( let catalog = state.catalog.load(); // 这把密钥的模型范围：每一跳发出的名字都要过它（见 `crate::sent::name`） let allow = crate::models::key_allow(&rt.config, &req.client_name); - // 开头慢就换下一家：开着、客户端要的是流时，等多久（见 `super::slow`） - let slow_wait = super::slow::wait(rt, reading); + // 无响应超时：每一跳从发出去算起，等多久（见 `super::idle`） + let quiet = super::idle::Quiet::of(state, rt); // 还没看的候选，按顺序 let mut queue: std::collections::VecDeque<&String> = started.alive.iter().collect(); // 满着没发的那几家（见 `crate::slots`），按候选的顺序。候选都看过一遍还没有着落时， @@ -161,7 +165,7 @@ pub(super) async fn try_upstreams<'a>( let mut queued: Option<(usize, u64)> = None; // 等到最后，剩下的候选还都满着 let mut stalled = false; - // 有一跳真的发到了上游、失败了（回了错误、流在内容之前断了、连不上、开头慢被放弃）。 + // 有一跳真的发到了上游、失败了（回了错误、流在内容之前断了、连不上、没有内容超时了）。 // 等满着的那几家等到最后的话，交出去的是这一家的错，不是「都满着」（见下面的 `stalled`） let mut reached = false; @@ -173,7 +177,16 @@ pub(super) async fn try_upstreams<'a>( None if busy.is_empty() => break, None => { let t = std::time::Instant::now(); - match state.slots.first_free(&busy, until).await { + let freed = tokio::select! { + biased; + freed = state.slots.first_free(&busy, until) => freed, + // 等空位的时候被手动中止：一个字节都还没发出去，尝试链上不多一跳 + _ = req.abort.wait() => { + halt = Some(GatewayError::aborted()); + break; + } + }; + match freed { Some((k, slot)) => (busy.remove(k), Some((slot, t.elapsed()))), None => { stalled = true; @@ -337,7 +350,14 @@ pub(super) async fn try_upstreams<'a>( // 这个请求能等的已经等完了（或者配置的是不等）：不再等 if tokio::time::Instant::now() < until { let t = std::time::Instant::now(); - slot = state.slots.take_by(&provider.name, until).await; + slot = tokio::select! { + biased; + slot = state.slots.take_by(&provider.name, until) => slot, + _ = req.abort.wait() => { + halt = Some(GatewayError::aborted()); + break; + } + }; waited = Some(t.elapsed().as_millis() as u64); hop_started = std::time::Instant::now(); } @@ -408,15 +428,7 @@ pub(super) async fn try_upstreams<'a>( let asked = Asked::of(req, reading, &plugged); let out = match super::heavy(asked.body, || { - prepare( - state, - req, - reading, - &asked, - provider, - &effective_set, - Some(id), - ) + prepare(state, req, reading, &asked, provider, &effective_set, id) }) { Ok(out) => out, Err(err) => { @@ -524,28 +536,14 @@ pub(super) async fn try_upstreams<'a>( .clone() .unwrap_or_else(|| reading.facts.model.clone()); let (attempt, bridge) = (chain.len(), plugged.bridge); - // 后面还有没有此刻接得下这个请求的（见 `successor`）。开头慢了才问。后面的是还没看的 - // 候选，加上满着、跳过了的那几家：**此刻有空位的才算**（见 `crate::slots`）—— 满着的 - // 那一家要等，而这个请求能等的多半已经等完了，放弃了这一家，换来的是一个 429 - let others = || { - let rest = queue - .iter() - .chain(busy.iter()) - .copied() - .filter(|n| !state.slots.is_full(n)); - successor(state, rt, req, reading, decision, &catalog, allow, rest) - }; - // 这一跳发出去的那一刻：这一家的快慢样本从这里算起（见 `crate::latency`） + // 这一跳发出去的那一刻：这一家的快慢样本、无响应超时都从这里算起（见 `crate::latency`、 + // `super::idle`）。之前等空位、跑插件的时间都不算 let sent_at = std::time::Instant::now(); - // 开头慢就换下一家：等到什么时候，从这一刻（请求发出去）算起。最后一家不换。到点时 - // 问过、后面没有接得下的，清掉它：这一跳从此和不开时一样 - let mut slow_deadline = slow_wait - .filter(|_| !last) - .map(|w| tokio::time::Instant::from_std(sent_at) + w); + let quiet_by = quiet.after(sent_at); - // 发出去、等响应头。**等着的这个 future 只活在这一块里**：放弃这一家时它跟着丢掉， - // 连接随之断开 - let sent = { + // 发出去、等响应头。**等着的这个 future 只活在这一块里**：放弃这一家、被中止时它跟着 + // 丢掉，连接随之断开 + let waited = { let sending = async { let sent = send( state, @@ -573,37 +571,44 @@ pub(super) async fn try_upstreams<'a>( Err(e) => Err(e), } }; - let mut sending = std::pin::pin!(sending); - match slow_deadline { - None => sending.await, - Some(deadline) => match tokio::time::timeout_at(deadline, sending.as_mut()).await { - Ok(sent) => sent, - // 响应头都还没来，后面又有接得下的：放弃这一家。不停用、不算失败（见 - // `super::slow`） - Err(_) if others() => { - reached = true; - let waited = slow_wait.unwrap_or_default(); - super::slow::timed_out(state, &provider.name, waited); - chain.push(super::slow::abandoned( - &provider.name, - model.clone(), - None, - None, - reading, - waited, - hop_started, - )); - last_err = Some(GatewayError::upstream(super::slow::said( - &provider.name, - waited, - ))); - continue; - } - Err(_) => { - slow_deadline = None; - sending.await - } - }, + tokio::select! { + biased; + sent = sending => Waited::Done(sent), + _ = req.abort.wait() => Waited::Aborted, + _ = tokio::time::sleep_until(quiet_by) => Waited::Quiet, + } + }; + let sent = match waited { + Waited::Done(sent) => sent, + // 被手动中止：这一跳到此为止，不再换下一家。不停用、不算失败 + Waited::Aborted => { + chain.push(super::idle::aborted_hop( + &provider.name, + model.clone(), + None, + hop_started, + )); + halt = Some(GatewayError::aborted()); + break; + } + // 响应头都没等到：放弃这一家，换下一家；没有下一家了，交出去的就是这个超时 + Waited::Quiet => { + reached = true; + last_err = Some(gave_up( + state, + started, + reading, + Abandoned { + provider, + model: model.clone(), + status: None, + usage: None, + quiet, + hop_started, + }, + &mut chain, + )); + continue; } }; match sent { @@ -718,100 +723,116 @@ pub(super) async fn try_upstreams<'a>( refusal, slot, sent_at, + unconfirmed: false, }); break; } } } Ok(r) => { - // 流式回答：**第一段内容到之前**上游在流里报的错，照样换下一家（见 - // `opening`）。最后一家不等 —— 没有下一家可换，等只会让客户端晚一点 - // 看到同一个错误 - let r = match opening_of(req, provider, &out, &r).filter(|_| !last) { + // 回答的开头：**第一段内容到之前**上游在流里报的错照样换下一家，一直没有内容 + // 也换（见 `opening`）。最后一家不压 —— 没有下一家可换，压着只会让客户端晚 + // 一点看到同样的东西；它不出声由回程按无响应超时收尾（见 `relay`） + let hold = hold_of(req, reading, provider, &out, &r); + // 最后一家不压开头：它答没答上（成败的账）要等回程看到第一段内容才知道 + let unconfirmed = last && hold.is_some(); + let r = match hold.filter(|_| !last) { None => r, - Some((dialect, eventstream)) => { - // 开头慢就换下一家的，等到发出请求之后的那一刻；别的从响应头到了算起 - let deadline = slow_deadline.unwrap_or_else(|| { - tokio::time::Instant::now() - + std::time::Duration::from_secs( - rt.config.failover.stream_start_wait_secs, - ) - }); - match super::opening::watch(r, dialect, eventstream, deadline).await { - super::opening::Opening::Go(r) => r, - // 到点了还没有内容，后面又有接得下的：放弃这一家。**响应跟着这一轮 - // 循环丢掉**，和上游的连接随之断开，它不再接着生成。不停用、不算失败 - super::opening::Opening::Slow { response, usage } - if slow_deadline.is_some() && others() => - { + Some(hold) => { + let status = r.status().as_u16(); + let watching = async move { + match hold { + Hold::Stream(dialect, eventstream) => { + super::opening::watch(r, dialect, eventstream, quiet_by).await + } + Hold::Body => super::opening::first_bytes(r, quiet_by).await, + } + }; + let opened = tokio::select! { + biased; + o = watching => Some(o), + // 压着开头的时候被中止：响应跟着丢掉，连接断开 + _ = req.abort.wait() => None, + }; + match opened { + None => { + chain.push(super::idle::aborted_hop( + &provider.name, + model.clone(), + Some(status), + hop_started, + )); + halt = Some(GatewayError::aborted()); + break; + } + Some(super::opening::Opening::Go(r)) => r, + // 到点了还没有内容：放弃这一家。**响应跟着这一轮循环丢掉**，和上游的连接 + // 随之断开，它不再接着生成 + Some(super::opening::Opening::Slow { response, usage }) => { let status = response.status().as_u16(); drop(response); reached = true; // 上游回了话，说明代理是通的 state.note_proxy_ok(&provider.proxy); - let waited = slow_wait.unwrap_or_default(); - super::slow::timed_out(state, &provider.name, waited); - chain.push(super::slow::abandoned( - &provider.name, - model.clone(), - Some(status), - usage, + last_err = Some(gave_up( + state, + started, reading, - waited, - hop_started, + Abandoned { + provider, + model: model.clone(), + status: Some(status), + usage, + quiet, + hop_started, + }, + &mut chain, )); - last_err = Some(GatewayError::upstream(super::slow::said( - &provider.name, - waited, - ))); continue; } - // 不换（没开、或者后面没有接得下的）：和内容来了一样交出去 - super::opening::Opening::Slow { response, .. } => response, - super::opening::Opening::Failed { + Some(super::opening::Opening::Failed { status, headers, body, kind, message, response, - } => { - match crate::failure::classify(status, &headers, &body, now_ms()) { - Verdict::ClientError => response, - Verdict::Failed(cause) => { - reached = true; - state.note_quota(id, &provider.name, &headers); - state.note_proxy_ok(&provider.proxy); - let cause = known_reset(state, &provider.name, cause); - note_health( - &state.bus, - &state.health, - &provider.name, - state.health.record_cause(&provider.name, cause), - ); - let said = msg!( - "gw.upstream.stream_opening_error", - upstream = provider.name.clone(), kind = kind, message = message => - "Upstream `{upstream}` started the answer and reported an error \ - before any content ({kind}): {message}" - ); - let err = if status == 429 { - GatewayError::rate_limited(said) - } else { - GatewayError::upstream(said) - }; - chain.push(hop_failed( - &provider.name, - model.clone(), - err.detail.clone(), - hop_started, - )); - last_err = Some(err); - continue; - } + }) => match crate::failure::classify(status, &headers, &body, now_ms()) + { + Verdict::ClientError => response, + Verdict::Failed(cause) => { + reached = true; + state.note_quota(id, &provider.name, &headers); + state.note_proxy_ok(&provider.proxy); + let cause = known_reset(state, &provider.name, cause); + note_health( + &state.bus, + &state.health, + &provider.name, + state.health.record_cause(&provider.name, cause), + ); + let said = msg!( + "gw.upstream.stream_opening_error", + upstream = provider.name.clone(), kind = kind, message = message => + "Upstream `{upstream}` started the answer and reported an error \ + before any content ({kind}): {message}" + ); + let err = if status == 429 { + GatewayError::rate_limited(said) + } else { + GatewayError::upstream(said) + }; + chain.push(hop_failed( + &provider.name, + model.clone(), + err.detail.clone(), + hop_started, + )); + last_err = Some(err); + continue; } - } - super::opening::Opening::Broken(err) => { + }, + Some(super::opening::Opening::Broken(err)) => { reached = true; note_health( &state.bus, @@ -831,12 +852,14 @@ pub(super) async fn try_upstreams<'a>( } } }; - note_health( - &state.bus, - &state.health, - &provider.name, - state.health.record_success(&provider.name), - ); + if !unconfirmed { + note_health( + &state.bus, + &state.health, + &provider.name, + state.health.record_success(&provider.name), + ); + } chain.push(hop( &provider.name, model.clone(), @@ -855,6 +878,7 @@ pub(super) async fn try_upstreams<'a>( refusal: None, slot, sent_at, + unconfirmed, }); break; } @@ -1120,73 +1144,9 @@ fn unsendable_tool( Some(GatewayError::new(crate::error::Source::Request, msg)) } -/// 慢了的那一家后面，`rest` 里还有没有接得下这个请求的（见 [`super::slow`]）。 -/// -/// 停用着的不算，这一跳发不出去的也不算：配置里没有了、格式对不上、阶段二的规则拒绝、 -/// 发给它的名字对不上（别名、清单、密钥范围）、转换不了、强制要用的工具发不过去。**一个都 -/// 没有的话，慢了的这一家就是最后一家**，照常等下去 —— 放弃了它，换来的是一个注定失败的 -/// 请求。 -/// -/// **满着的不算**（`max_concurrent`，见 [`crate::slots`]），调用方先把它们滤掉：到点的这一刻 -/// 有空位的（没设上限的、跳过时满着、此刻空出来了的）才接得下。满着的那一家要等，而一个请求 -/// 只有一段等待期限（`failover.slot_wait_secs`），开头慢的这一段多半已经把它用完了：放弃了 -/// 一家正在答的，换来的是去等一家满着的、等不到回 429。 -/// -/// **只看不跑**：看的是客户端的原话，不跑插件、不取密钥（那两样在真发的那一跳才知道拒 -/// 不拒），不发转换事件 -#[allow(clippy::too_many_arguments)] -fn successor<'r>( - state: &AppState, - rt: &Runtime, - req: &Inbound, - reading: &crate::client_api::Reading, - decision: &tw_engine::Decision, - catalog: &tw_engine::Catalog, - allow: Option<&[String]>, - mut rest: impl Iterator, -) -> bool { - let asked = Asked { - body: &req.body, - path: req.uri.path(), - decoded: reading.decoded.as_ref(), - json: reading.json, - }; - rest.any(|name| { - let Some(provider) = rt.config.providers.iter().find(|p| &p.name == name) else { - return false; - }; - if !state.health.is_available(name) - || protocol_mismatch(req, reading.generates, provider).is_some() - { - return false; - } - let Ok(tw_engine::Outcome2::Proceed { - mut set, - model: renamed, - .. - }) = rt - .engine - .phase_two(&reading.facts, &provider.name, &decision.set) - else { - return false; - }; - let asked_model = - rt.engine - .asked_of(&reading.facts, decision, &provider.name, renamed.as_deref()); - let Ok(sent) = - crate::sent::name(&rt.config, catalog, decision, provider, &asked_model, allow) - else { - return false; - }; - set.model = (sent != reading.facts.model).then_some(sent); - prepare(state, req, reading, &asked, provider, &set, None).is_ok() - }) -} - /// 把这一跳的请求（客户端那种格式，插件改过的话是改过的，见 [`Asked`]）改成要发的 /// 样子：同格式时只做参数改写，跨格式时转换。转换不了就换下一家：同格式的上游可能 -/// 还在后面。`id` 是这个请求的号，做了转换、丢了字段要报在它上面；只看发不发得出去时 -/// （见 [`successor`]）是 None，什么都不报。 +/// 还在后面。`id` 是这个请求的号，做了转换、丢了字段要报在它上面。 fn prepare( state: &AppState, req: &Inbound, @@ -1194,7 +1154,7 @@ fn prepare( asked: &Asked<'_>, provider: &tw_config::Provider, effective_set: &tw_engine::SetAction, - id: Option, + id: u64, ) -> Result { let generates = reading.generates; // 方言互转。**同格式时是 None，这一整段零成本** @@ -1257,7 +1217,7 @@ fn prepare( .and_then(|d| crate::egress::strip_body_identity(d, &out)) .unwrap_or(out) }; - if let Some((d, id)) = client_dialect.filter(|_| !dropped.is_empty()).zip(id) { + if let Some(d) = client_dialect.filter(|_| !dropped.is_empty()) { let same = crate::wire::dialect(d); state.bus.emit(tw_api::Event::Translated { id, @@ -1331,16 +1291,14 @@ fn prepare( // 却没生效」而完全不知道从哪儿查起 let mut dropped = p.dropped.clone(); dropped.extend(limit); - if let Some(id) = id { - state.bus.emit(tw_api::Event::Translated { - id, - provider: provider.name.clone(), - from: crate::wire::dialect(d.client), - to: crate::wire::dialect(dialect), - dropped, - at_ms: crate::server::now_ms(), - }); - } + state.bus.emit(tw_api::Event::Translated { + id, + provider: provider.name.clone(), + from: crate::wire::dialect(d.client), + to: crate::wire::dialect(dialect), + dropped, + at_ms: crate::server::now_ms(), + }); path = p.path.clone(); query = p.query.clone(); // Bedrock 上的 Claude：客户端 `anthropic-beta` 里 Bedrock 认的那几个放进请求体 @@ -1764,15 +1722,23 @@ pub(crate) fn stream_fault( } } -/// 这个回答要不要等开头：生成回答的流式响应才等。等的话，上游说的是哪种格式、 -/// 是不是 Bedrock 的二进制帧 -fn opening_of( +/// 一跳的响应头到了之后怎么等开头（见 [`super::opening`]）。 +enum Hold { + /// 流式：等第一段内容。上游说的是哪种格式、是不是 Bedrock 的二进制帧 + Stream(tw_dialect::ir::Dialect, bool), + /// 整包：等正文的第一个字节 + Body, +} + +/// 这个回答要不要等开头：生成回答的成功响应才等。流式的等第一段内容，整包的等正文开头 +fn hold_of( req: &Inbound, + reading: &crate::client_api::Reading, provider: &tw_config::Provider, out: &Outbound, r: &reqwest::Response, -) -> Option<(tw_dialect::ir::Dialect, bool)> { - if !r.status().is_success() { +) -> Option { + if !r.status().is_success() || !reading.generates { return None; } let ct = r @@ -1782,11 +1748,66 @@ fn opening_of( .unwrap_or_default(); let eventstream = provider.is_bedrock() && ct.starts_with(tw_bedrock::eventstream::CONTENT_TYPE); - if !(eventstream || ct.starts_with("text/event-stream")) { - return None; + // Codex 后端的流式响应没有 Content-Type（发给它的一律是流，见 `relay::Plan`） + let sse = ct.starts_with("text/event-stream") || (out.chatgpt && ct.is_empty()); + if !(eventstream || sse) { + return Some(Hold::Body); } let dialect = out.target.or_else(|| req.api.map(|a| a.dialect()))?; - Some((dialect, eventstream)) + Some(Hold::Stream(dialect, eventstream)) +} + +/// 发出去、等响应头的结局 +enum Waited { + Done(T), + /// 被手动中止了 + Aborted, + /// 无响应超时到了 + Quiet, +} + +/// 无响应超时之前一个字都没给的那一跳 +struct Abandoned<'p> { + provider: &'p tw_config::Provider, + model: Option, + /// 上游回的状态码，响应头没到的没有 + status: Option, + /// 流开头里上游报的用量 + usage: Option, + quiet: super::idle::Quiet, + hop_started: std::time::Instant, +} + +/// 放弃这一家（见 [`super::idle`]）：记一次失败、记它被给的那段时间、这段对话这一轮不再留在 +/// 它那儿，尝试链上记一跳 `idle_timeout`。交回候选用完时交给客户端的那条错误（504） +fn gave_up( + state: &AppState, + started: &Started, + reading: &crate::client_api::Reading, + a: Abandoned<'_>, + chain: &mut Vec, +) -> GatewayError { + let name = &a.provider.name; + note_health( + &state.bus, + &state.health, + name, + state.health.record_failure(name), + ); + super::idle::timed_out(state, name, a.quiet.window); + if let Some(c) = &started.turn { + state.affinity.left(c, name); + } + chain.push(super::idle::abandoned( + name, + a.model, + a.status, + a.usage, + reading, + a.quiet, + a.hop_started, + )); + GatewayError::timeout(super::idle::said(name, a.quiet.secs)) } fn now_ms() -> u64 { diff --git a/crates/tw-gateway/src/server/pipeline/idle.rs b/crates/tw-gateway/src/server/pipeline/idle.rs new file mode 100644 index 00000000..081f320e --- /dev/null +++ b/crates/tw-gateway/src/server/pipeline/idle.rs @@ -0,0 +1,160 @@ +//! 无响应超时（配置的 `failover.idle_timeout_secs`）：上游多久没有内容就不再等它。 +//! +//! 有的上游收下请求之后一直不出内容：中转站排着队，上游过载却不报错，响应头都迟迟不来； +//! 也有的答到一半停住，连接却不断。整体超时不能设（一个跑了六分钟的回答不该被掐断，见 +//! `crate::outbound`），所以看的是**多久没有内容**：从请求发给这一家的那一刻算起（之前等 +//! 空位的时间不算），每来一段真内容重新计时，心跳不算（判据见 [`crate::pulse`]）。整包的回答 +//! 没有「一段段」，从发出去到整份回来算一段。 +//! +//! 到点了怎么办，看客户端收到了什么： +//! +//! - **还什么都没收到**：放弃这一家（丢掉响应或者还在等的请求，连接跟着断，上游不再接着 +//! 生成），记一次失败（和 5xx 一样算进停用的账），换下一家（见 `hop`）。客户端无感。没有 +//! 下一家了，回一个超时错误（504，见 [`crate::error::Source::Timeout`]），尝试链带着。 +//! - **最后一家的响应已经交出去了**（最后一家不压开头）、或者**已经有内容交给了客户端**： +//! 换不了 —— 换一家会把开头再发一遍。回答按客户端的格式以一条错误收尾，请求记成失败 +//! （见 `relay`）。还没有内容的照样记这一家一次失败；内容出到一半才停的不记，和流在半路 +//! 断了一样。 +//! +//! 几条规矩： +//! +//! - **它的快慢样本记它被给的那段时间**（见 [`timed_out`]）：`url-test` 和按快慢分的 +//! `load-balance` 照这个把它往后排。 +//! - **这段对话这一轮不再留在它那儿**（见 [`crate::affinity::Affinity::left`]）：下一个请求 +//! 照常排序，不会因为「上次回答的就是它」又被送回去。 +//! - **尝试链上记一跳 `idle_timeout`**，带着上游可能已经收了钱的输入（见 +//! [`tw_api::AttemptUsage`]）。 +//! +//! **等的时候不给客户端发保活。**响应头要等选定了哪一家才发（见 `relay`），这期间客户端 +//! 那条连接上什么都没有；先发响应头再发 `: keepalive` 的话，状态码就定死成了 200 —— 之后 +//! 几家全都失败，429、5xx 和最后一家原样交出的 4xx 都给不出去，只能在流里报错，客户端按 +//! 状态码重试的逻辑就落空了；上游的响应头（请求号、额度）也带不过去。何况 Gemini 官方的 +//! Python SDK 会把注释行当成一段 JSON 去解析，直接报错。 + +use std::time::Duration; + +use crate::state::{AppState, Runtime}; +use tw_types::msg; + +/// 这个请求的无响应超时：配置的秒数，和它在这个进程里有多长（测试把一秒调短，见 +/// `AppState::idle_tick`）。 +#[derive(Debug, Clone, Copy)] +pub(super) struct Quiet { + /// 配置写的秒数。报给人看的句子用它 + pub(super) secs: u64, + /// 真正等多久 + pub(super) window: Duration, +} + +impl Quiet { + pub(super) fn of(state: &AppState, rt: &Runtime) -> Self { + let secs = rt.config.failover.idle_timeout_secs; + let window = state + .idle_tick + .checked_mul(u32::try_from(secs).unwrap_or(u32::MAX)) + .unwrap_or(Duration::MAX); + Self { secs, window } + } + + /// 从 `from`（这一跳发出去、或者上一段内容到的那一刻）算起，到什么时候为止 + pub(super) fn after(&self, from: std::time::Instant) -> tokio::time::Instant { + let from = tokio::time::Instant::from_std(from); + from.checked_add(self.window) + .unwrap_or_else(|| from + Duration::from_secs(86_400 * 365)) + } +} + +/// 放弃了这一家：给它记一个快慢样本，就是它被给的那段时间（见 [`crate::latency`]）。 +/// +/// **它至少这么慢**，这是个下限：记成这个数，`url-test` 和按快慢分的 `load-balance` 就把它 +/// 排到慢的那一头。什么都不记的话，它留着的还是从前快的样本，下一个请求照样先发给它， +/// 而等它的这段时间算到了接下来那一家头上。 +pub(super) fn timed_out(state: &AppState, provider: &str, waited: Duration) { + state.latency.record(provider, crate::latency::ms(waited)); +} + +/// 放弃了的那一跳：尝试链上的一行。`status` 是上游回的（响应头没到的没有），`seen` 是流 +/// 开头里上游报的用量。 +#[allow(clippy::too_many_arguments)] +pub(super) fn abandoned( + provider: &str, + model: Option, + status: Option, + seen: Option, + reading: &crate::client_api::Reading, + quiet: Quiet, + started: std::time::Instant, +) -> tw_api::AttemptView { + tw_api::AttemptView { + provider: provider.to_string(), + model, + outcome: tw_api::AttemptOutcome::IdleTimeout, + status, + error: Some(said(provider, quiet.secs)), + ms: started.elapsed().as_millis() as u64, + usage: usage(seen, reading), + // 等过空位的话，等了多久由尝试链补上（`stamp_queued`） + queued_ms: None, + skipped: None, + } +} + +/// 等上游的时候被手动中止的那一跳（见 [`crate::abort`]）。`status` 是上游回的（响应头没到的 +/// 没有） +pub(super) fn aborted_hop( + provider: &str, + model: Option, + status: Option, + started: std::time::Instant, +) -> tw_api::AttemptView { + tw_api::AttemptView { + provider: provider.to_string(), + model, + outcome: tw_api::AttemptOutcome::Aborted, + status, + error: Some(crate::error::GatewayError::aborted().detail), + ms: started.elapsed().as_millis() as u64, + usage: None, + queued_ms: None, + skipped: None, + } +} + +/// 放弃的那一家可能已经收了钱的输入：上游报了的用它报的，没报的用网关估的（和开始事件的 +/// `input_estimate` 同一个数），估不出来（请求解不开）就没有。 +fn usage( + seen: Option, + reading: &crate::client_api::Reading, +) -> Option { + match seen.filter(|u| u.prompt_total() > 0) { + Some(u) => Some(tw_api::AttemptUsage { + input: u.input, + cache_read: u.cache_read, + cache_write: u.cache_write, + estimated: false, + }), + None => matches!(reading.decoded, Some(Ok(_))).then_some(tw_api::AttemptUsage { + input: reading.facts.input_tokens, + cache_read: 0, + cache_write: 0, + estimated: true, + }), + } +} + +/// 还没有内容就超时了：尝试链上那一跳的说明，也是候选用完时交给客户端的那条错误 +pub(super) fn said(upstream: &str, secs: u64) -> tw_types::Msg { + msg!( + "gw.upstream.idle_timeout", upstream = upstream, secs = secs => + "Upstream `{upstream}` sent no content within {secs} seconds." + ) +} + +/// 答到一半停住了 +pub(super) fn stalled(upstream: &str, secs: u64) -> tw_types::Msg { + msg!( + "gw.upstream.idle_timeout_mid_stream", upstream = upstream, secs = secs => + "Upstream `{upstream}` stopped sending content partway through the answer and sent \ + nothing more for {secs} seconds." + ) +} diff --git a/crates/tw-gateway/src/server/pipeline/opening.rs b/crates/tw-gateway/src/server/pipeline/opening.rs index 6fdcacba..ee74a0bf 100644 --- a/crates/tw-gateway/src/server/pipeline/opening.rs +++ b/crates/tw-gateway/src/server/pipeline/opening.rs @@ -1,4 +1,4 @@ -//! 流式回答的开头：**第一段内容到达之前，上游在流里报的错误照样换下一家。** +//! 回答的开头：**第一段内容到达之前，上游在流里报的错误照样换下一家；一直没有内容，也换。** //! //! 上游回了 200、流也开了，之后第一个事件却是错误 —— Anthropic 过载时的 //! `overloaded_error`、Codex 额度用完时的 `response.failed`、Bedrock 的 @@ -7,12 +7,16 @@ //! //! 所以响应头到手之后先不转发：把流读到第一段内容为止，读到的字节原样留着。 //! 期间是错误就换下一家；是内容，就把留着的字节和后面的流接在一起交出去 —— -//! 客户端收到的和直接转发一个字节都不差。 +//! 客户端收到的和直接转发一个字节都不差。什么是内容、什么是开头的例行事件和心跳， +//! 和无响应超时是同一个判据（见 [`crate::pulse`]）。 //! -//! 等待有上限（配置的 `failover.stream_start_wait_secs`，以及 [`HOLD_LIMIT`]）： -//! 上游迟迟不出内容时不能一直压着，那样客户端看到的就是一个卡住的请求。等到点了是 -//! [`Opening::Slow`]：开着 `failover.next_on_slow_start` 时由调用方放弃这一家、换下一家， -//! 不开就和内容来了一样交出去。 +//! 等到无响应超时那一刻（`failover.idle_timeout_secs`，从请求发出去算起）还没有内容是 +//! [`Opening::Slow`]：调用方放弃这一家、换下一家。**最后一家不在这里等**（见 `hop`）：没有 +//! 下一家可换，压着它只会让客户端晚一点看到同样的东西。开头压得太多（[`HOLD_LIMIT`]）也不 +//! 再压，交出去由回程照常计时。 +//! +//! 整包的回答（不是流）也一样等：响应头先到、正文迟迟不来的上游，等到点了换下一家（见 +//! [`first_bytes`]）。 use bytes::Bytes; use futures::StreamExt; @@ -27,9 +31,9 @@ pub(super) const HOLD_LIMIT: usize = 1024 * 1024; pub(super) enum Opening { /// 内容来了（或者流结束了、开头压得太多了）：交给客户端。读过的字节已经接回去了 Go(reqwest::Response), - /// 等到点了还没有内容。`response` 和 [`Opening::Go`] 的一样，照旧交出去就是不换家； - /// **丢掉它就断开了和上游的连接**，上游不再接着生成。`usage` 是开头里上游报了的用量 - /// （Anthropic 的 `message_start` 带着输入），没报是 None + /// 等到无响应超时那一刻还没有内容。**丢掉 `response` 就断开了和上游的连接**，上游不再 + /// 接着生成。`usage` 是开头里上游报了的用量（Anthropic 的 `message_start` 带着输入）， + /// 没报是 None Slow { response: reqwest::Response, usage: Option, @@ -51,8 +55,8 @@ pub(super) enum Opening { Broken(GatewayError), } -/// 读到第一段内容为止，最多等到 `deadline`。`dialect` 是上游说的格式，`eventstream` -/// 表示流是 Bedrock 的二进制帧。 +/// 读到第一段内容为止，最多等到 `deadline`（无响应超时的那一刻）。`dialect` 是上游说的 +/// 格式，`eventstream` 表示流是 Bedrock 的二进制帧。 pub(super) async fn watch( r: reqwest::Response, dialect: Dialect, @@ -140,6 +144,48 @@ pub(super) async fn watch( } } +/// 整包的回答（不是流）：等到正文的第一个不是空白的字节，最多等到 `deadline`（无响应超时 +/// 的那一刻）。到了是 [`Opening::Go`]，读过的字节接回去；到点了还没有是 [`Opening::Slow`]。 +/// +/// 只发空格保活的中转站，空格不算（见 [`crate::pulse`]）。整包的回答不会在正文里报一个 +/// 「开头的错误」，所以这里没有 [`Opening::Failed`]。 +pub(super) async fn first_bytes(r: reqwest::Response, deadline: tokio::time::Instant) -> Opening { + let status = r.status(); + let headers = r.headers().clone(); + let mut stream = r.bytes_stream(); + let mut held: Vec = Vec::new(); + let mut size = 0usize; + let mut pulse = crate::pulse::Pulse::new(None); + let slow = loop { + match tokio::time::timeout_at(deadline, stream.next()).await { + Err(_) => break true, + Ok(None) => break false, + Ok(Some(Err(e))) => return Opening::Broken(crate::forward::map_reqwest_error(e)), + Ok(Some(Ok(chunk))) => { + size += chunk.len(); + let said = pulse.feed(&chunk); + held.push(chunk); + if said || size >= HOLD_LIMIT { + break false; + } + } + } + }; + let replay = futures::stream::iter(held.into_iter().map(Ok::<_, reqwest::Error>)); + let mut resp = http::Response::new(reqwest::Body::wrap_stream(replay.chain(stream))); + *resp.status_mut() = status; + *resp.headers_mut() = headers; + let response = reqwest::Response::from(resp); + if slow { + Opening::Slow { + response, + usage: None, + } + } else { + Opening::Go(response) + } +} + /// 一个事件是什么。 #[derive(Debug, PartialEq)] enum Judge { @@ -156,7 +202,8 @@ enum Judge { }, } -/// 看一个 SSE 事件。 +/// 看一个 SSE 事件：上游报的错，内容，还是开头的例行事件和心跳（后两样的判据在 +/// [`crate::pulse::advances`]，和无响应超时同一个）。 fn judge(dialect: Dialect, event: Option<&str>, data: &str) -> Judge { let data = data.trim(); if data.is_empty() { @@ -172,82 +219,70 @@ fn judge(dialect: Dialect, event: Option<&str>, data: &str) -> Judge { .map(str::to_string) .or_else(|| v.get("type").and_then(|t| t.as_str()).map(str::to_string)) .unwrap_or_default(); + if let Some(e) = error_in(dialect, &ty, &v, data) { + return e; + } + if crate::pulse::advances(dialect, &ty, &v) { + Judge::Content + } else { + Judge::Preamble + } +} + +/// 这个事件是不是上游报的错：是的话，那个错误对应的状态码和上游的原话 +fn error_in(dialect: Dialect, ty: &str, v: &serde_json::Value, data: &str) -> Option { match dialect { - Dialect::Anthropic => match ty.as_str() { - "message_start" | "ping" => Judge::Preamble, - "error" => { - let e = v.get("error").unwrap_or(&v); - let kind = text_of(e, "type"); - error(anthropic_status(&kind), data, kind, text_of(e, "message")) - } - _ => Judge::Content, - }, - Dialect::Responses => match ty.as_str() { - "response.created" | "response.in_progress" | "response.queued" => Judge::Preamble, - // Codex 后端在开头报一次额度，不是回答 - t if t.starts_with("codex.") => Judge::Preamble, + Dialect::Anthropic => (ty == "error").then(|| { + let e = v.get("error").unwrap_or(v); + let kind = text_of(e, "type"); + error(anthropic_status(&kind), data, kind, text_of(e, "message")) + }), + Dialect::Responses => match ty { "response.failed" => { - let e = v.pointer("/response/error").unwrap_or(&v); + let e = v.pointer("/response/error").unwrap_or(v); let kind = text_of(e, "code"); - error(openai_status(&kind), data, kind, text_of(e, "message")) + Some(error( + openai_status(&kind), + data, + kind, + text_of(e, "message"), + )) } "error" => { - let e = v.get("error").unwrap_or(&v); + let e = v.get("error").unwrap_or(v); let kind = Some(text_of(e, "code")) .filter(|c| !c.is_empty()) .unwrap_or_else(|| text_of(e, "type")); - error(openai_status(&kind), data, kind, text_of(e, "message")) + Some(error( + openai_status(&kind), + data, + kind, + text_of(e, "message"), + )) } - _ => Judge::Content, - }, - Dialect::Chat => { - if let Some(e) = v.get("error") { - let kind = Some(text_of(e, "code")) - .filter(|c| !c.is_empty()) - .unwrap_or_else(|| text_of(e, "type")); - let status = e - .get("code") - .and_then(|c| c.as_u64()) - .and_then(|c| u16::try_from(c).ok()) - .filter(|c| (400..600).contains(c)) - .unwrap_or_else(|| openai_status(&kind)); - return error(status, data, kind, text_of(e, "message")); - } - // 只有角色、没有内容的第一块（多数兼容接口都先发这么一块）还不算内容 - let delta = v.pointer("/choices/0/delta"); - let finished = v - .pointer("/choices/0/finish_reason") - .is_some_and(|f| !f.is_null()); - let said = delta.is_some_and(|d| { - ["content", "reasoning_content", "reasoning", "tool_calls"] - .iter() - .any(|k| d.get(*k).is_some_and(|x| !x.is_null() && x != "")) - }); - let usage_only = v.get("usage").is_some_and(|u| !u.is_null()) - && v.get("choices") - .and_then(|c| c.as_array()) - .is_none_or(|c| c.is_empty()); - if said || finished || usage_only { - Judge::Content - } else { - Judge::Preamble - } - } - Dialect::Gemini => match v.get("error") { - Some(e) => { - let status = e - .get("code") - .and_then(|c| c.as_u64()) - .and_then(|c| u16::try_from(c).ok()) - .unwrap_or(500); - error(status, data, text_of(e, "status"), text_of(e, "message")) - } - None => Judge::Content, - }, - Dialect::Bedrock => match ty.as_str() { - "messageStart" => Judge::Preamble, - _ => Judge::Content, + _ => None, }, + Dialect::Chat => v.get("error").map(|e| { + let kind = Some(text_of(e, "code")) + .filter(|c| !c.is_empty()) + .unwrap_or_else(|| text_of(e, "type")); + let status = e + .get("code") + .and_then(|c| c.as_u64()) + .and_then(|c| u16::try_from(c).ok()) + .filter(|c| (400..600).contains(c)) + .unwrap_or_else(|| openai_status(&kind)); + error(status, data, kind, text_of(e, "message")) + }), + Dialect::Gemini => v.get("error").map(|e| { + let status = e + .get("code") + .and_then(|c| c.as_u64()) + .and_then(|c| u16::try_from(c).ok()) + .unwrap_or(500); + error(status, data, text_of(e, "status"), text_of(e, "message")) + }), + Dialect::Bedrock => None, } } @@ -524,8 +559,39 @@ mod tests { } } + /// 整包的回答：先到的空格不算，到点了没有正文是 Slow；正文来了照原样交出去 + #[tokio::test] + async fn a_whole_answer_waits_for_its_first_real_byte() { + let deadline = tokio::time::Instant::now() + Duration::from_millis(200); + assert!(matches!( + first_bytes(stalled(" \n"), deadline).await, + Opening::Slow { usage: None, .. } + )); + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + match first_bytes(served(" {\"id\":1}").await, deadline).await { + Opening::Go(r) => assert_eq!(r.text().await.unwrap(), " {\"id\":1}"), + _ => panic!("该放行"), + } + // 空的正文也放行：交出去由回程照常收尾 + assert!(matches!( + first_bytes(served("").await, deadline).await, + Opening::Go(_) + )); + } + + /// Gemini 只带用量的块不算开口（和无响应超时同一个判据） + #[tokio::test] + async fn a_gemini_usage_only_chunk_is_not_the_first_content() { + let head = "data: {\"usageMetadata\":{\"promptTokenCount\":10}}\r\n\r\n"; + let deadline = tokio::time::Instant::now() + Duration::from_millis(200); + assert!(matches!( + watch(stalled(head), Dialect::Gemini, false, deadline).await, + Opening::Slow { .. } + )); + } + #[tokio::test] - async fn a_thinking_delta_is_content_not_a_slow_start() { + async fn a_thinking_delta_is_content_and_a_role_chunk_is_not() { // Chat 格式的推理字（DeepSeek、Qwen 的 `reasoning_content`）也是模型开口了 let head = "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":\"\"}}]}\n\n\ data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"Let me think\"}}]}\n\n"; diff --git a/crates/tw-gateway/src/server/pipeline/relay.rs b/crates/tw-gateway/src/server/pipeline/relay.rs index db683e98..f42f7544 100644 --- a/crates/tw-gateway/src/server/pipeline/relay.rs +++ b/crates/tw-gateway/src/server/pipeline/relay.rs @@ -7,6 +7,10 @@ //! //! 这些步骤要不要做、怎么做，在响应头到手的那一刻就全定了（[`Plan`]）； //! 流里每一块怎么处理在 [`Relay`] 上，`respond` 里的流只剩一个循环。 +//! +//! 那个循环还看着两件事：上游多久没有内容（无响应超时，见 [`super::idle`]），和这个请求有 +//! 没有被手动中止（见 [`crate::abort`]）。哪一件到了都丢掉上游的流，按客户端的格式补一条 +//! 错误收尾。 use axum::body::Body; use axum::http::{HeaderMap, StatusCode}; @@ -47,6 +51,7 @@ pub(super) fn respond( refusal, slot, sent_at, + unconfirmed, .. } = served; let status = @@ -159,6 +164,12 @@ pub(super) fn respond( && status.is_success()) .then_some(state.ping_every); let ping_for = state.ping_for; + // 无响应超时：从这一跳发出去算起，每来一段真内容重新计时（见 `super::idle`）。真内容按 + // 上游的格式认；不是流的回答，有一个不是空白的字节就算 + let quiet = super::idle::Quiet::of(state, rt); + let mut pulse = crate::pulse::Pulse::new(plan.is_sse.then_some(upstream_dialect)); + let abort = req.abort.clone(); + let (bus, health) = (state.bus.clone(), state.health.clone()); let stream = async_stream::stream! { // **通行证跟着响应体走。**这个流被丢掉的时候它才还回去：正常 // 发完是一种，客户端中途断开、hyper 丢掉响应体是另一种 —— 两种 @@ -180,24 +191,54 @@ pub(super) fn respond( // 半开的连接永远等不到下一个字节，一直补心跳的话客户端也永远不会放弃，这个 // 请求就挂在那儿了。停下之后由客户端自己的静默计时来断 let mut upstream_since = tokio::time::Instant::now(); + // 到这一刻还没有真内容就不再等（见 `super::idle`） + let mut quiet_by = quiet.after(sent_at); + // 上游给过真内容了没有：没给过的超时还是这一家的失败，给过的是答到一半停住了 + let mut said = false; + // 被手动中止了（结局报手动中止，不说「流断了」） + let mut aborted = false; + // 这一家的成败还没记（最后一家不压开头，见 `hop`）：第一段真内容到了记成功 + let mut unconfirmed = unconfirmed; loop { - let next = match ping_every { - None => chunks.next().await, - // `next()` 被超时丢掉不丢数据：它只是去问一次流，没拿走任何东西 - Some(every) => match tokio::time::timeout_at(quiet_since + every, chunks.next()).await { - Ok(next) => next, - Err(_) => { - // 停在一帧中间时这一轮不补，也要重新计时，否则会原地空转 - quiet_since = tokio::time::Instant::now(); - // **不经过留档、计量和审查**：心跳不是上游说的话，不进请求记录， - // 也不算输出。**只在帧的边界上插**，上游停在一帧中间时插进去 - // 会把那一帧拆坏 —— 那时宁可不补 - if relay.between_frames() && upstream_since.elapsed() < ping_for { - yield Ok::(Bytes::from_static(PING)); - } - continue; + // 补心跳的那一刻。不补的格式给一个不会用到的时刻，那一支不参与 + let ping_at = ping_every.map_or(quiet_by, |every| quiet_since + every); + let next = tokio::select! { + // 先看上游：已经到了的那一块不因为同时到点而丢掉 + biased; + next = chunks.next() => next, + _ = abort.wait() => { + aborted = true; + broke = Some(GatewayError::aborted().in_dialect(dialect)); + break; + } + _ = tokio::time::sleep_until(quiet_by) => { + let why = if said { + super::idle::stalled(&upstream_name, quiet.secs) + } else { + // 一个字都没给：和压着开头时超时一样，这一家记一次失败 + crate::server::note_health( + &bus, + &health, + &upstream_name, + health.record_failure(&upstream_name), + ); + super::idle::said(&upstream_name, quiet.secs) + }; + ending.went_quiet(); + broke = Some(GatewayError::timeout(why).in_dialect(dialect)); + break; + } + _ = tokio::time::sleep_until(ping_at), if ping_every.is_some() => { + // 停在一帧中间时这一轮不补，也要重新计时，否则会原地空转 + quiet_since = tokio::time::Instant::now(); + // **不经过留档、计量和审查**：心跳不是上游说的话，不进请求记录， + // 也不算输出。**只在帧的边界上插**，上游停在一帧中间时插进去 + // 会把那一帧拆坏 —— 那时宁可不补 + if relay.between_frames() && upstream_since.elapsed() < ping_for { + yield Ok::(Bytes::from_static(PING)); } - }, + continue; + } }; let Some(item) = next else { break }; // eventstream 拆成 SSE。一帧没收齐时这一块什么都转不出来，等下一块 @@ -218,6 +259,19 @@ pub(super) fn respond( match item { Ok(chunk) => { upstream_since = tokio::time::Instant::now(); + // 真内容才重新计时：心跳、开头的例行事件不算 + if pulse.feed(&chunk) { + said = true; + quiet_by = quiet.after(std::time::Instant::now()); + if std::mem::take(&mut unconfirmed) { + crate::server::note_health( + &bus, + &health, + &upstream_name, + health.record_success(&upstream_name), + ); + } + } // **旁路嗅探和留档，不缓冲**：字节照常流向客户端，同时 // 喂它一份。上游返回的 usage 是真相，而拿不到它就只能估。 // @@ -252,6 +306,15 @@ pub(super) fn respond( } // 响应体留档和结束事件都在 `ending` 里：三种结局要交出去的是同一份 // 东西，分开写就会有一种漏掉 + // 一段真内容都没有就好好收了尾：也是答上了 + if unconfirmed && broke.is_none() { + crate::server::note_health( + &bus, + &health, + &upstream_name, + health.record_success(&upstream_name), + ); + } match broke { None => ending.finished(status.as_u16()), Some(err) => { @@ -264,8 +327,11 @@ pub(super) fn respond( // `source` 用这个错误自己的：上游断了是 `upstream`，被 // 防火墙切断是 `denied` —— 后者不是上游坏了，是策略拦的。 // 码保持不变，只在句子前面点明它断在流里 —— 界面认的是码 + // 手动中止的不是流断了：原句照报，记录里是那个码 let mut why = err.detail.clone(); - why.text = format!("the response stream broke: {}", why.text); + if !aborted { + why.text = format!("the response stream broke: {}", why.text); + } ending.failed(err.source.into(), why); if let Some(frame) = relay.error_tail(&err) { yield Ok(Bytes::from(relay.plugins_tail(&frame))); diff --git a/crates/tw-gateway/src/server/pipeline/slow.rs b/crates/tw-gateway/src/server/pipeline/slow.rs deleted file mode 100644 index 85a88d9a..00000000 --- a/crates/tw-gateway/src/server/pipeline/slow.rs +++ /dev/null @@ -1,101 +0,0 @@ -//! 开头慢就换下一家（配置的 `failover.next_on_slow_start`）。 -//! -//! 有的上游收下请求之后很久不出内容：中转站排着队，上游过载却不报错，响应头都迟迟不来。 -//! 开着这一项时，从请求发出去算起等 `failover.stream_start_wait_secs`，还没有内容就**断开 -//! 这一家**（丢掉响应或者还在等的请求，连接跟着断，上游不再接着生成），换下一家。客户端 -//! 这时一个字节都还没收到，换一家它无感。 -//! -//! 几条规矩： -//! -//! - **最后一家不换**，照常等下去。「最后」按到点的那一刻后面还有没有接得下的算（见 -//! `hop::successor`）：停用着的、这一跳发不出去的、并发数满着的（`max_concurrent`）不算 —— -//! 否则放弃了一个慢的，换来的是一个注定失败的，或者一个等不到空位的 429。 -//! - **这一家不停用、不算失败**：慢不是坏，下一个请求它可能就快了。 -//! - **它的快慢样本记它被给的那段时间**（见 [`timed_out`]）：`url-test` 和按快慢分的 -//! `load-balance` 照这个把它往后排。 -//! - **尝试链上记一跳 `slow_start`**，带着上游可能已经收了钱的输入（见 -//! [`tw_api::AttemptUsage`]）。 -//! - 只管客户端要流式的请求：整包的请求本来就要等全部生成完，开头慢说明不了什么。 -//! -//! **等的时候不给客户端发保活。**响应头要等选定了哪一家才发（见 `relay`），这期间客户端 -//! 那条连接上什么都没有；先发响应头再发 `: keepalive` 的话，状态码就定死成了 200 —— 之后 -//! 几家全都失败，429、5xx 和最后一家原样交出的 4xx 都给不出去，只能在流里报错，客户端按 -//! 状态码重试的逻辑就落空了；上游的响应头（请求号、额度）也带不过去。何况 Gemini 官方的 -//! Python SDK 会把注释行当成一段 JSON 去解析，直接报错。 - -use std::time::Duration; - -use crate::state::Runtime; -use tw_types::msg; - -/// 这个请求开头慢了换不换、换的话等多久：开着这一项、客户端要的是流时才有。 -pub(super) fn wait(rt: &Runtime, reading: &crate::client_api::Reading) -> Option { - let f = &rt.config.failover; - let streams = matches!(&reading.decoded, Some(Ok(d)) if d.request.stream); - (f.next_on_slow_start && streams).then(|| Duration::from_secs(f.stream_start_wait_secs)) -} - -/// 放弃了这一家：给它记一个快慢样本，就是它被给的那段时间（见 [`crate::latency`]）。 -/// -/// **它至少这么慢**，这是个下限：记成这个数，`url-test` 和按快慢分的 `load-balance` 就把它 -/// 排到慢的那一头。什么都不记的话，它留着的还是从前快的样本，下一个请求照样先发给它， -/// 而等它的这段时间算到了接下来那一家头上。 -pub(super) fn timed_out(state: &crate::state::AppState, provider: &str, waited: Duration) { - state.latency.record(provider, crate::latency::ms(waited)); -} - -/// 放弃了的那一跳：尝试链上的一行。`status` 是上游回的（响应头没到的没有），`seen` 是流 -/// 开头里上游报的用量。 -pub(super) fn abandoned( - provider: &str, - model: Option, - status: Option, - seen: Option, - reading: &crate::client_api::Reading, - waited: Duration, - started: std::time::Instant, -) -> tw_api::AttemptView { - tw_api::AttemptView { - provider: provider.to_string(), - model, - outcome: tw_api::AttemptOutcome::SlowStart, - status, - error: Some(said(provider, waited)), - ms: started.elapsed().as_millis() as u64, - usage: usage(seen, reading), - // 等过空位的话，等了多久由尝试链补上（`stamp_queued`） - queued_ms: None, - skipped: None, - } -} - -/// 放弃的那一家可能已经收了钱的输入：上游报了的用它报的，没报的用网关估的（和开始事件的 -/// `input_estimate` 同一个数），估不出来（请求解不开）就没有。 -fn usage( - seen: Option, - reading: &crate::client_api::Reading, -) -> Option { - match seen.filter(|u| u.prompt_total() > 0) { - Some(u) => Some(tw_api::AttemptUsage { - input: u.input, - cache_read: u.cache_read, - cache_write: u.cache_write, - estimated: false, - }), - None => matches!(reading.decoded, Some(Ok(_))).then_some(tw_api::AttemptUsage { - input: reading.facts.input_tokens, - cache_read: 0, - cache_write: 0, - estimated: true, - }), - } -} - -/// 尝试链上那一跳的说明。后面几家也都不行时，它也是交给客户端的那条错误的退路 -pub(super) fn said(upstream: &str, waited: Duration) -> tw_types::Msg { - msg!( - "gw.slow_start", upstream = upstream, secs = waited.as_secs() => - "Upstream `{upstream}` sent no content within {secs} seconds, so the request moved on \ - to the next upstream." - ) -} diff --git a/crates/tw-gateway/src/state.rs b/crates/tw-gateway/src/state.rs index 74a624a1..43deda09 100644 --- a/crates/tw-gateway/src/state.rs +++ b/crates/tw-gateway/src/state.rs @@ -248,6 +248,11 @@ pub struct AppState { pub ping_for: std::time::Duration, /// 跑插件的线程池（见 [`crate::plugin::pool`]）。**跨重载存活**；线程第一次用到时才起 pub plugin_pool: Arc, + /// 此刻在跑、可以手动中止的请求（见 [`crate::abort`]）。**跨重载存活** + pub aborts: Arc, + /// `failover.idle_timeout_secs` 的一秒有多长（见 `server::pipeline::idle`）。**测试会把它 + /// 调短**，否则一条无响应超时的测试至少要干等三十秒 + pub idle_tick: std::time::Duration, } impl AppState { @@ -315,6 +320,8 @@ impl AppState { ping_every: crate::PING_EVERY, ping_for: crate::PING_FOR, plugin_pool: Arc::new(crate::plugin::pool::Pool::default_size()), + aborts: Default::default(), + idle_tick: std::time::Duration::from_secs(1), }; // 手写的清单马上可用；向上游问是后台的事，不挡启动 state.publish_catalog(); diff --git a/crates/tw-gateway/src/wire.rs b/crates/tw-gateway/src/wire.rs index 51c2fb9c..e122e8a2 100644 --- a/crates/tw-gateway/src/wire.rs +++ b/crates/tw-gateway/src/wire.rs @@ -14,6 +14,8 @@ impl From for FailureSource { Source::RateLimited | Source::Busy => Self::RateLimited, Source::Denied => Self::Denied, Source::NotSupported | Source::TooLarge => Self::Request, + Source::Timeout => Self::Upstream, + Source::Aborted => Self::Aborted, } } } diff --git a/crates/tw-gateway/tests/idle_timeout.rs b/crates/tw-gateway/tests/idle_timeout.rs new file mode 100644 index 00000000..d10215e7 --- /dev/null +++ b/crates/tw-gateway/tests/idle_timeout.rs @@ -0,0 +1,1007 @@ +//! 无响应超时（`failover.idle_timeout_secs`）和手动中止（`tw_gateway::abort`），端到端。 +//! +//! 假上游先回响应头（或者迟迟不回），之后按脚本隔一阵发一段，发完了挂着、隔一阵发一次心跳； +//! 它记下自己的响应被丢掉了没有 —— 网关放弃它、中止它时，连接要真的断开，上游才会停下。 +//! +//! 超时写最短的 30 秒，测试把一秒调成三十分之一秒（`AppState::idle_tick`）：等的是 1 秒。 + +use std::net::SocketAddr; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use axum::Router; +use bytes::Bytes; +use futures::StreamExt; +use serde_json::{Value, json}; +use tw_api::AttemptOutcome::{Aborted, Error, IdleTimeout, Served}; +use tw_api::Event; +use tw_config::{Client, Config, Failover, Protocol, Provider}; + +/// 无响应超时在测试里有多长 +const WINDOW: Duration = Duration::from_secs(1); + +/// 一个假上游的样子：响应头之前等多久，之后按顺序隔多少毫秒发哪一段，发完了挂不挂着 +#[derive(Clone)] +struct Script { + header_delay_ms: u64, + steps: Vec<(u64, &'static str)>, + /// 发完了不收尾，每 100 毫秒发一次 `beat`，直到连接断开 + hang: bool, + /// 挂着时发的心跳 + beat: &'static str, + content_type: &'static str, +} + +impl Default for Script { + fn default() -> Self { + Self { + header_delay_ms: 0, + steps: Vec::new(), + hang: false, + beat: ": keep-alive\n\n", + content_type: "text/event-stream", + } + } +} + +/// 假上游被打了几次，以及它的响应（或者还没回的那个请求）被丢掉了没有 +struct Upstream { + addr: SocketAddr, + hits: Arc, + dropped: Arc, +} + +/// 被丢掉时举旗 +struct Flag(Arc); + +impl Drop for Flag { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } +} + +async fn upstream(script: Script) -> Upstream { + let hits = Arc::new(AtomicUsize::new(0)); + let dropped = Arc::new(AtomicBool::new(false)); + let (h, d) = (hits.clone(), dropped.clone()); + let app = Router::new().fallback(axum::routing::any(move || { + let (script, h, d) = (script.clone(), h.clone(), d.clone()); + async move { + h.fetch_add(1, Ordering::SeqCst); + // 响应头还没回时请求就被丢掉的，也要看得见 + let waiting = Flag(d.clone()); + tokio::time::sleep(Duration::from_millis(script.header_delay_ms)).await; + std::mem::forget(waiting); + let flag = Flag(d); + let steps = futures::stream::iter(script.steps).then(|(ms, chunk)| async move { + tokio::time::sleep(Duration::from_millis(ms)).await; + Ok::<_, std::convert::Infallible>(Bytes::from_static(chunk.as_bytes())) + }); + let (hang, beat) = (script.hang, script.beat); + let tail = futures::stream::unfold((), move |_| async move { + if !hang { + return None; + } + tokio::time::sleep(Duration::from_millis(100)).await; + Some((Ok(Bytes::from_static(beat.as_bytes())), ())) + }); + let body = steps.chain(tail).map(move |x| { + let _ = &flag; + x + }); + axum::response::Response::builder() + .header("content-type", script.content_type) + .body(axum::body::Body::from_stream(body)) + .unwrap() + } + })); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + Upstream { + addr, + hits, + dropped, + } +} + +const MESSAGE_START: &str = "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1200,\"cache_read_input_tokens\":800,\"output_tokens\":1}}}\n\n"; +const PING: &str = "event: ping\ndata: {\"type\":\"ping\"}\n\n"; +const THINKING: &str = "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\n\ + event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"Let me think\"}}\n\n"; +const WORD: &str = "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"word \"}}\n\n"; +const STOP: &str = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + +/// 一段完整的回答，`word` 是正文 +fn answer(word: &'static str) -> &'static str { + let text = format!( + "event: content_block_delta\ndata: {{\"type\":\"content_block_delta\",\"index\":0,\"delta\":{{\"type\":\"text_delta\",\"text\":\"{word}\"}}}}\n\n\ + event: message_stop\ndata: {{\"type\":\"message_stop\"}}\n\n" + ); + Box::leak(text.into_boxed_str()) +} + +/// 开了流、报了输入，之后只有心跳 +fn stalled() -> Script { + Script { + steps: vec![(0, MESSAGE_START), (0, PING)], + hang: true, + ..Default::default() + } +} + +/// 马上就答 +fn prompt(word: &'static str) -> Script { + Script { + steps: vec![(0, MESSAGE_START), (0, answer(word))], + ..Default::default() + } +} + +/// 整包的回答 +fn whole(word: &'static str) -> &'static str { + let text = format!( + "{{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-5\",\"content\":[{{\"type\":\"text\",\"text\":\"{word}\"}}],\"stop_reason\":\"end_turn\",\"usage\":{{\"input_tokens\":5,\"output_tokens\":1}}}}" + ); + Box::leak(text.into_boxed_str()) +} + +fn provider(name: &str, up: &Upstream, protocol: Protocol) -> Provider { + Provider { + name: name.into(), + base_url: format!("http://{}", up.addr), + key: Some("sk-upstream".into()), + protocol: Some(protocol), + ..Default::default() + } +} + +/// 起网关，交回数据面的状态。`slot_wait_secs` 是等空位的期限 +async fn gateway_with( + providers: Vec, + slot_wait_secs: u64, +) -> ( + SocketAddr, + tokio::sync::broadcast::Receiver, + tw_gateway::AppState, +) { + let cfg = Config { + version: 1, + clients: vec![Client { + name: "c".into(), + key: "tw-k".into(), + ..Default::default() + }], + providers, + failover: Failover { + idle_timeout_secs: 30, + slot_wait_secs, + ..Default::default() + }, + ..Default::default() + }; + let mut state = tw_gateway::AppState::new(cfg).unwrap(); + state.idle_tick = WINDOW / 30; + let rx = state.bus.subscribe(); + let addr = tw_gateway::serve(state.clone(), ([127, 0, 0, 1], 0).into()) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + (addr, rx, state) +} + +async fn gateway( + providers: Vec, +) -> ( + SocketAddr, + tokio::sync::broadcast::Receiver, + tw_gateway::AppState, +) { + gateway_with(providers, 30).await +} + +fn messages(stream: bool) -> Value { + json!({"model": "claude-sonnet-5", "max_tokens": 1024, "stream": stream, + "messages": [{"role": "user", "content": "Say hello."}]}) +} + +/// 发出去，交回状态码、`x-thinkwatch-error` 和整个正文 +async fn send(gw: SocketAddr, path: &str, body: &Value) -> (u16, Option, String) { + let resp = reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{gw}{path}")) + .header("content-type", "application/json") + .header("x-api-key", "tw-k") + .header("x-goog-api-key", "tw-k") + .header("authorization", "Bearer tw-k") + .body(body.to_string()) + .send() + .await + .unwrap(); + let status = resp.status().as_u16(); + let source = resp + .headers() + .get("x-thinkwatch-error") + .map(|v| v.to_str().unwrap().to_string()); + (status, source, resp.text().await.unwrap()) +} + +async fn post(gw: SocketAddr, path: &str, body: &Value) -> (u16, String) { + let (status, _, text) = send(gw, path, body).await; + (status, text) +} + +/// 下一个满足 `f` 的事件，最多等十秒 +async fn next( + rx: &mut tokio::sync::broadcast::Receiver, + mut f: impl FnMut(&Event) -> bool, +) -> Event { + loop { + let e = tokio::time::timeout(Duration::from_secs(10), rx.recv()) + .await + .expect("没等到那个事件") + .unwrap(); + if f(&e) { + return e; + } + } +} + +/// 这个请求的估算输入（开始事件带着）和尝试链 +async fn estimate_and_attempts( + rx: &mut tokio::sync::broadcast::Receiver, +) -> (Option, Vec) { + let mut estimate = None; + loop { + match next(rx, |_| true).await { + Event::RequestStarted { input_estimate, .. } => estimate = input_estimate, + Event::RequestRouted { attempts, .. } => return (estimate, attempts), + _ => {} + } + } +} + +/// 请求的结局：结束是 Ok(状态码)，失败是 Err((来源, 码))，取消是 panic +async fn outcome( + rx: &mut tokio::sync::broadcast::Receiver, +) -> Result { + match next(rx, |e| { + matches!( + e, + Event::RequestFinished { .. } + | Event::RequestFailed { .. } + | Event::RequestCancelled { .. } + ) + }) + .await + { + Event::RequestFinished { status, .. } => Ok(status), + Event::RequestFailed { + source, message, .. + } => Err((source, message.code)), + e => panic!("不该是取消：{e:?}"), + } +} + +/// 等旗举起来，最多两秒 +async fn eventually(flag: &AtomicBool) -> bool { + for _ in 0..40 { + if flag.load(Ordering::SeqCst) { + return true; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + false +} + +fn outcomes(attempts: &[tw_api::AttemptView]) -> Vec<(&str, tw_api::AttemptOutcome)> { + attempts + .iter() + .map(|a| (a.provider.as_str(), a.outcome)) + .collect() +} + +// ───────────────────────────────────────────── 还没有内容：换下一家 + +#[tokio::test] +async fn a_quiet_first_upstream_is_dropped_after_the_window_and_the_next_one_answers() { + let slow = upstream(stalled()).await; + let good = upstream(prompt("hello")).await; + let (gw, mut rx, state) = gateway(vec![ + provider("slow", &slow, Protocol::Anthropic), + provider("good", &good, Protocol::Anthropic), + ]) + .await; + + let t = Instant::now(); + let (status, text) = post(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "{text}"); + assert!(text.contains("hello"), "{text}"); + assert!( + !text.contains("ping") && text.matches("event: message_start").count() == 1, + "放弃的那一家的开头不该到客户端：{text}" + ); + assert!(t.elapsed() >= WINDOW, "等够了才换"); + assert!( + eventually(&slow.dropped).await, + "放弃的那一家连接要断开，它才会停下" + ); + + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [("slow", IdleTimeout), ("good", Served)] + ); + let gave_up = &attempts[0]; + assert_eq!(gave_up.status, Some(200)); + assert_eq!( + gave_up.usage, + Some(tw_api::AttemptUsage { + input: 1200, + cache_read: 800, + cache_write: 0, + estimated: false, + }), + "上游在开头报了输入，记它报的" + ); + let said = gave_up.error.as_ref().expect("说等了多久"); + assert_eq!(said.code, "gw.upstream.idle_timeout"); + assert_eq!(said.args["secs"], "30", "说的是配置写的秒数"); + assert_eq!(outcome(&mut rx).await, Ok(200)); + + // **记一次失败**：和 5xx 一样算进停用的账，连着三次就停用 + for _ in 0..2 { + let (status, _) = post(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200); + } + assert!(!state.health.is_available("slow"), "三次超时之后停用"); + let rates = state + .health + .success_rates(&["slow".to_string(), "good".to_string()]); + assert_eq!(rates.get("slow"), None, "不到五个样本不给成功率：{rates:?}"); + let (status, _) = post(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200); + assert_eq!(slow.hits.load(Ordering::SeqCst), 3, "停用之后不再先试它"); +} + +#[tokio::test] +async fn response_headers_that_never_come_time_out_the_same_way() { + // 响应头都迟迟不来（中转站排着队）：一样从发出请求算起，到点放弃 + let queued = upstream(Script { + header_delay_ms: 5_000, + ..prompt("too late") + }) + .await; + let good = upstream(prompt("hello")).await; + let (gw, mut rx, _) = gateway(vec![ + provider("queued", &queued, Protocol::Anthropic), + provider("good", &good, Protocol::Anthropic), + ]) + .await; + let t = Instant::now(); + let (status, text) = post(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "{text}"); + assert!(text.contains("hello"), "{text}"); + assert!(t.elapsed() < Duration::from_secs(4), "没有等到响应头"); + assert!(eventually(&queued.dropped).await, "还在等的请求要断开"); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [("queued", IdleTimeout), ("good", Served)] + ); + assert_eq!(attempts[0].status, None, "响应头没到，没有状态码"); + assert!(attempts[0].usage.is_some_and(|u| u.estimated)); +} + +#[tokio::test] +async fn when_no_upstream_is_left_the_client_gets_a_timeout_and_nothing_waits_forever() { + // 两家的响应头都不来：第二家也是到点就不等了，客户端收到它那种格式的 504 + let a = upstream(Script { + header_delay_ms: 60_000, + ..prompt("never") + }) + .await; + let b = upstream(Script { + header_delay_ms: 60_000, + ..prompt("never") + }) + .await; + let (gw, mut rx, _) = gateway(vec![ + provider("a", &a, Protocol::Anthropic), + provider("b", &b, Protocol::Anthropic), + ]) + .await; + let t = Instant::now(); + let (status, source, text) = send(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 504, "{text}"); + assert_eq!(source.as_deref(), Some("upstream")); + let v: Value = serde_json::from_str(&text).unwrap(); + assert_eq!(v["error"]["type"], "timeout_error", "{text}"); + let message = v["error"]["message"].as_str().unwrap(); + assert!(message.starts_with("[ThinkWatch]"), "{message}"); + assert!(message.contains("tried: a → b"), "{message}"); + let waited = t.elapsed(); + assert!( + waited >= WINDOW * 2 && waited < Duration::from_secs(5), + "{waited:?}" + ); + assert!(eventually(&b.dropped).await); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [("a", IdleTimeout), ("b", IdleTimeout)] + ); + assert_eq!( + outcome(&mut rx).await, + Err(( + tw_api::FailureSource::Upstream, + "gw.upstream.idle_timeout".into() + )) + ); +} + +#[tokio::test] +async fn the_last_upstream_that_says_nothing_ends_the_stream_with_an_error() { + // 唯一的一家开了流、只发心跳：响应头已经交出去了（最后一家不压开头），到点由流里的 + // 一条错误收尾，记成没有内容的超时 + let only = upstream(stalled()).await; + let (gw, mut rx, state) = gateway(vec![provider("only", &only, Protocol::Anthropic)]).await; + let t = Instant::now(); + let (status, text) = post(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "{text}"); + assert!(t.elapsed() < Duration::from_secs(4)); + assert!(text.contains("event: error"), "{text}"); + assert!(text.contains("timeout_error"), "{text}"); + assert!(eventually(&only.dropped).await); + assert_eq!( + outcome(&mut rx).await, + Err(( + tw_api::FailureSource::Upstream, + "gw.upstream.idle_timeout".into() + )) + ); + // 一个字都没给：这一家也记一次失败，三次之后停用（只有它一家时请求照样发给它） + assert!(state.health.is_available("only")); + for _ in 0..2 { + let (status, _) = post(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200); + } + assert!( + !state.health.is_available("only"), + "三次没有内容的超时之后停用" + ); +} + +#[tokio::test] +async fn pings_alone_do_not_keep_an_upstream_alive() { + // 只发 Anthropic 的 `ping` 事件、Responses 的 `response.in_progress`：心跳不算，照样超时 + for (protocol, path, body, start, beat) in [ + ( + Protocol::Anthropic, + "/v1/messages", + messages(true), + MESSAGE_START, + PING, + ), + ( + Protocol::OpenaiResponses, + "/v1/responses", + json!({"model": "gpt-5", "stream": true, "input": "Say hello."}), + "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\n", + "event: response.in_progress\ndata: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_1\"}}\n\n", + ), + ( + Protocol::OpenaiChat, + "/v1/chat/completions", + json!({"model": "gpt-5", "stream": true, "messages": [{"role": "user", "content": "hi"}]}), + "data: {\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\"}}]}\n\n", + "data: {\"choices\":[{\"index\":0,\"delta\":{}}]}\n\n", + ), + ] { + let only = upstream(Script { + steps: vec![(0, start)], + hang: true, + beat, + ..Default::default() + }) + .await; + let (gw, mut rx, _) = gateway(vec![provider("only", &only, protocol)]).await; + let t = Instant::now(); + let (status, text) = post(gw, path, &body).await; + assert_eq!(status, 200, "{protocol:?}: {text}"); + assert!(t.elapsed() < Duration::from_secs(4), "{protocol:?}"); + assert!(text.contains("[ThinkWatch]"), "{protocol:?}: {text}"); + assert_eq!( + outcome(&mut rx).await, + Err(( + tw_api::FailureSource::Upstream, + "gw.upstream.idle_timeout".into() + )), + "{protocol:?}" + ); + } +} + +#[tokio::test] +async fn a_whole_answer_that_does_not_come_in_time_moves_on() { + // 整包的请求：从发出去到整份回来算一段。响应头迟迟不来的，换下一家 + let slow = upstream(Script { + header_delay_ms: 5_000, + steps: vec![(0, whole("late"))], + content_type: "application/json", + ..Default::default() + }) + .await; + // 响应头马上到、正文一直不来（只有空格）的，也换 + let mute = upstream(Script { + steps: vec![(0, " ")], + hang: true, + beat: " ", + content_type: "application/json", + ..Default::default() + }) + .await; + let good = upstream(Script { + steps: vec![(0, whole("whole"))], + content_type: "application/json", + ..Default::default() + }) + .await; + let (gw, mut rx, _) = gateway(vec![ + provider("slow", &slow, Protocol::Anthropic), + provider("mute", &mute, Protocol::Anthropic), + provider("good", &good, Protocol::Anthropic), + ]) + .await; + let (status, text) = post(gw, "/v1/messages", &messages(false)).await; + assert_eq!(status, 200, "{text}"); + assert!(text.contains("whole"), "{text}"); + assert!(eventually(&slow.dropped).await && eventually(&mute.dropped).await); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [ + ("slow", IdleTimeout), + ("mute", IdleTimeout), + ("good", Served) + ] + ); + assert_eq!(attempts[1].status, Some(200)); + + // 只有一家、整份不来：504，不是一直等 + let only = upstream(Script { + header_delay_ms: 60_000, + steps: vec![(0, whole("never"))], + content_type: "application/json", + ..Default::default() + }) + .await; + let (gw, _, _) = gateway(vec![provider("only", &only, Protocol::Anthropic)]).await; + let (status, source, text) = send(gw, "/v1/messages", &messages(false)).await; + assert_eq!(status, 504, "{text}"); + assert_eq!(source.as_deref(), Some("upstream")); +} + +// ───────────────────────────────────────────── 内容在来：不算超时 + +#[tokio::test] +async fn a_slow_but_steady_stream_is_not_cut_off() { + // 每 400 毫秒一段，一共三秒：总时长远超 1 秒，但没有哪一段间隔超过 + let mut steps = vec![(0, MESSAGE_START)]; + steps.extend(std::iter::repeat_n((400, WORD), 7)); + steps.push((400, STOP)); + let steady = upstream(Script { + steps, + ..Default::default() + }) + .await; + let other = upstream(prompt("other")).await; + let (gw, mut rx, _) = gateway(vec![ + provider("steady", &steady, Protocol::Anthropic), + provider("other", &other, Protocol::Anthropic), + ]) + .await; + let t = Instant::now(); + let (status, text) = post(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "{text}"); + assert!(t.elapsed() >= Duration::from_secs(3)); + assert_eq!(text.matches("word ").count(), 7, "{text}"); + assert!(!text.contains("event: error"), "{text}"); + assert_eq!(other.hits.load(Ordering::SeqCst), 0); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!(outcomes(&attempts), [("steady", Served)]); + assert_eq!(outcome(&mut rx).await, Ok(200)); +} + +#[tokio::test] +async fn thinking_counts_as_content() { + // 开头在想（推理的字在流），正文要过一阵才来：每一段间隔都不到 1 秒，加起来超过 + let thinker = upstream(Script { + steps: vec![ + (0, MESSAGE_START), + (600, THINKING), + (600, answer("thought")), + ], + ..Default::default() + }) + .await; + let other = upstream(prompt("other")).await; + let (gw, mut rx, _) = gateway(vec![ + provider("thinker", &thinker, Protocol::Anthropic), + provider("other", &other, Protocol::Anthropic), + ]) + .await; + let (status, text) = post(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "{text}"); + assert!( + text.contains("Let me think") && text.contains("thought"), + "{text}" + ); + assert_eq!(other.hits.load(Ordering::SeqCst), 0); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!(outcomes(&attempts), [("thinker", Served)]); + assert_eq!(outcome(&mut rx).await, Ok(200)); +} + +// ───────────────────────────────────────────── 已经有内容：报错收尾 + +#[tokio::test] +async fn content_then_silence_ends_with_an_error_the_client_understands() { + for (protocol, path, body, said, end) in [ + ( + Protocol::Anthropic, + "/v1/messages", + messages(true), + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"model\":\"claude-sonnet-5\",\"usage\":{\"input_tokens\":5,\"output_tokens\":1}}}\n\n\ + event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n\ + event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"half an ans\"}}\n\n", + "event: error\n", + ), + ( + Protocol::OpenaiChat, + "/v1/chat/completions", + json!({"model": "gpt-5", "stream": true, "messages": [{"role": "user", "content": "hi"}]}), + "data: {\"id\":\"c1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"half an ans\"}}]}\n\n", + "\"error\"", + ), + ( + Protocol::OpenaiResponses, + "/v1/responses", + json!({"model": "gpt-5", "stream": true, "input": "Say hello."}), + "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\n\ + event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"half an ans\"}\n\n", + "event: response.failed\n", + ), + ( + Protocol::Gemini, + "/v1beta/models/gemini-3-pro:streamGenerateContent?alt=sse", + json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}), + "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"half an ans\"}]}}]}\r\n\r\n", + "DEADLINE_EXCEEDED", + ), + ] { + let up = upstream(Script { + steps: vec![(0, said)], + hang: true, + ..Default::default() + }) + .await; + // 两家：第一家的开头压到了内容，交出去之后才停住 —— 那时已经换不了 + let spare = upstream(prompt("spare")).await; + let (gw, mut rx, state) = gateway(vec![ + provider("half", &up, protocol), + provider("spare", &spare, protocol), + ]) + .await; + let (status, text) = post(gw, path, &body).await; + assert_eq!(status, 200, "{protocol:?}: {text}"); + assert!(text.contains("half an ans"), "{protocol:?}: {text}"); + let tail = &text[text.find("half an ans").unwrap()..]; + assert!(tail.contains(end), "{protocol:?}: {text}"); + assert!(tail.contains("[ThinkWatch]"), "{protocol:?}: {text}"); + assert_eq!(spare.hits.load(Ordering::SeqCst), 0, "{protocol:?}"); + assert!(eventually(&up.dropped).await, "{protocol:?}"); + assert_eq!( + outcome(&mut rx).await, + Err(( + tw_api::FailureSource::Upstream, + "gw.upstream.idle_timeout_mid_stream".into() + )), + "{protocol:?}" + ); + // 答到一半停住的不记失败：和流在半路断了一样 + assert!(state.health.is_available("half")); + } +} + +// ───────────────────────────────────────────── 等空位不算在里面 + +#[tokio::test] +async fn waiting_for_a_slot_is_not_counted_as_silence() { + // 唯一的一家满着：等了 1.5 秒才空出来，之后 600 毫秒出内容。等的那段不算，照常答上 + let late = upstream(Script { + steps: vec![(0, MESSAGE_START), (600, answer("patience"))], + ..Default::default() + }) + .await; + let (gw, mut rx, state) = gateway_with( + vec![Provider { + max_concurrent: Some(1), + ..provider("late", &late, Protocol::Anthropic) + }], + 5, + ) + .await; + let held = state.slots.try_take("late").expect("一个空位"); + let asking = tokio::spawn(async move { post(gw, "/v1/messages", &messages(true)).await }); + tokio::time::sleep(Duration::from_millis(1_500)).await; + drop(held); + let (status, text) = asking.await.unwrap(); + assert_eq!(status, 200, "{text}"); + assert!(text.contains("patience"), "{text}"); + assert!(!text.contains("event: error"), "{text}"); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + // 先是满着跳过的那一行，等到空位之后发出去的是第二行 + assert_eq!(outcomes(&attempts), [("late", Error), ("late", Served)]); + assert!(attempts[1].queued_ms.is_some_and(|ms| ms >= 1_000)); +} + +/// 放弃的那一家占着的位置（`max_concurrent`，见 `tw_gateway::slots`）**当场**还回去：接下 +/// 请求的那一家还在答，慢的那一家已经空出来了，不等这个请求结束 +#[tokio::test] +async fn the_slot_of_an_upstream_given_up_on_is_free_at_once() { + const CONTENT: &str = "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n"; + let slow = upstream(stalled()).await; + // 先答上，隔一会儿才说完：这段时间里请求还没结束 + let good = upstream(Script { + steps: vec![(0, MESSAGE_START), (0, CONTENT), (800, STOP)], + ..Default::default() + }) + .await; + let (gw, _, state) = gateway(vec![ + Provider { + max_concurrent: Some(1), + ..provider("slow", &slow, Protocol::Anthropic) + }, + provider("good", &good, Protocol::Anthropic), + ]) + .await; + + let asking = tokio::spawn(async move { post(gw, "/v1/messages", &messages(true)).await }); + // 还在等慢的那一家：它的位置占着 + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!(slow.hits.load(Ordering::SeqCst), 1); + assert!(state.slots.try_take("slow").is_none(), "等着的时候占着位置"); + // 换到了好的那一家：慢的那一家的位置已经还回来了，这个请求还没结束 + for _ in 0..60 { + if good.hits.load(Ordering::SeqCst) == 1 { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert_eq!(good.hits.load(Ordering::SeqCst), 1); + assert!(!asking.is_finished(), "测的是请求还在进行时"); + assert!( + state.slots.try_take("slow").is_some(), + "放弃的那一家的位置要当场还回去" + ); + let (status, text) = asking.await.unwrap(); + assert_eq!(status, 200, "{text}"); + assert!(text.contains("hello"), "{text}"); +} + +/// 并发数满着的那一家（`max_concurrent`） +fn full(name: &str, up: &Upstream) -> Provider { + Provider { + max_concurrent: Some(1), + ..provider(name, up, Protocol::Anthropic) + } +} + +/// 跳过时满着的那一家，在放弃不出声的那一家之后空出来了：换到它 +#[tokio::test] +async fn an_upstream_that_frees_meanwhile_takes_over() { + let busy = upstream(prompt("freed")).await; + let slow = upstream(stalled()).await; + let (gw, mut rx, state) = gateway_with( + vec![ + full("busy", &busy), + provider("slow", &slow, Protocol::Anthropic), + ], + 3, + ) + .await; + let held = state.slots.try_take("busy").expect("一个空位"); + let asking = tokio::spawn(async move { post(gw, "/v1/messages", &messages(true)).await }); + tokio::time::sleep(Duration::from_millis(1_400)).await; + drop(held); + let (status, text) = asking.await.unwrap(); + assert_eq!(status, 200, "{text}"); + assert!(text.contains("freed"), "{text}"); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [("busy", Error), ("slow", IdleTimeout), ("busy", Served)] + ); +} + +/// 剩下的那一家一直满着：不出声的那一家放弃之后等它到期限，等不到就交出超时，不是「都满着」 +/// 的 429 —— 有一家真的收到过请求 +#[tokio::test] +async fn a_next_upstream_that_never_frees_leaves_the_timeout_to_tell() { + let busy = upstream(prompt("busy")).await; + let slow = upstream(stalled()).await; + let (gw, mut rx, state) = gateway_with( + vec![ + full("busy", &busy), + provider("slow", &slow, Protocol::Anthropic), + ], + 2, + ) + .await; + let _held = state.slots.try_take("busy").expect("一个空位"); + let (status, source, text) = send(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 504, "{text}"); + assert_eq!(source.as_deref(), Some("upstream")); + assert_eq!(busy.hits.load(Ordering::SeqCst), 0); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [("busy", Error), ("slow", IdleTimeout)] + ); +} + +// ───────────────────────────────────────────── 手动中止 + +/// 下一个开始的请求：号和会话 +async fn started(rx: &mut tokio::sync::broadcast::Receiver) -> (u64, Option) { + match next(rx, |e| matches!(e, Event::RequestStarted { .. })).await { + Event::RequestStarted { id, session, .. } => (id, session), + _ => unreachable!(), + } +} + +#[tokio::test] +async fn aborting_before_the_answer_starts_answers_499_and_leaves_the_upstream_alone() { + let queued = upstream(Script { + header_delay_ms: 60_000, + ..prompt("never") + }) + .await; + let other = upstream(prompt("other")).await; + let (gw, mut rx, state) = gateway(vec![ + provider("queued", &queued, Protocol::Anthropic), + provider("other", &other, Protocol::Anthropic), + ]) + .await; + let asking = tokio::spawn(async move { send(gw, "/v1/messages", &messages(true)).await }); + let (id, _) = started(&mut rx).await; + tokio::time::sleep(Duration::from_millis(200)).await; + state.aborts.request(id).expect("还在跑"); + + let (status, source, text) = asking.await.unwrap(); + assert_eq!(status, 499, "{text}"); + assert_eq!(source.as_deref(), Some("aborted")); + let v: Value = serde_json::from_str(&text).unwrap(); + assert!( + v["error"]["message"] + .as_str() + .unwrap() + .starts_with("[ThinkWatch]"), + "{text}" + ); + assert!(eventually(&queued.dropped).await, "和上游的连接要断开"); + // 不换下一家 + assert_eq!(other.hits.load(Ordering::SeqCst), 0); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!(outcomes(&attempts), [("queued", Aborted)]); + assert_eq!( + attempts[0].error.as_ref().map(|m| m.code.as_str()), + Some(tw_api::ABORTED) + ); + assert_eq!( + outcome(&mut rx).await, + Err((tw_api::FailureSource::Aborted, tw_api::ABORTED.into())) + ); + // 不是上游的错 + assert!(state.health.is_available("queued")); + for _ in 0..5 { + state.health.record_success("queued"); + } + assert_eq!( + state + .health + .success_rates(&["queued".to_string()]) + .get("queued"), + Some(&1.0), + "中止不算一次失败" + ); + // 结束了的再叫停：不在跑 + assert!(state.aborts.request(id).is_err()); + assert!(state.aborts.is_empty()); +} + +#[tokio::test] +async fn aborting_mid_stream_ends_the_answer_with_an_error_in_its_format() { + let mut steps = vec![(0, MESSAGE_START)]; + steps.extend(std::iter::repeat_n((200, WORD), 50)); + let talker = upstream(Script { + steps, + ..Default::default() + }) + .await; + let (gw, mut rx, state) = gateway(vec![provider("talker", &talker, Protocol::Anthropic)]).await; + let asking = tokio::spawn(async move { post(gw, "/v1/messages", &messages(true)).await }); + let (id, _) = started(&mut rx).await; + // 等它答上几个字 + tokio::time::sleep(Duration::from_millis(700)).await; + state.aborts.request(id).expect("还在跑"); + let (status, text) = asking.await.unwrap(); + assert_eq!(status, 200, "{text}"); + assert!(text.contains("word "), "{text}"); + assert!(text.matches("word ").count() < 50, "没等到说完：{text}"); + let tail = &text[text.rfind("word ").unwrap()..]; + assert!(tail.contains("event: error"), "{text}"); + assert!(tail.contains("[ThinkWatch]"), "{text}"); + assert!(eventually(&talker.dropped).await, "上游要停下"); + match next(&mut rx, |e| matches!(e, Event::RequestFailed { .. })).await { + Event::RequestFailed { + source, + message, + usage, + .. + } => { + assert_eq!(source, tw_api::FailureSource::Aborted); + assert_eq!(message.code, tw_api::ABORTED); + assert!( + !message.text.contains("stream broke"), + "中止不是流断了：{}", + message.text + ); + assert!(usage.is_some(), "上游已经计了费，用量照记"); + } + _ => unreachable!(), + } + assert!(state.health.is_available("talker")); +} + +#[tokio::test] +async fn aborting_a_session_stops_every_request_of_it_and_nothing_else() { + let mut steps = vec![(0, MESSAGE_START)]; + steps.extend(std::iter::repeat_n((200, WORD), 50)); + let talker = upstream(Script { + steps, + ..Default::default() + }) + .await; + let (gw, mut rx, state) = gateway(vec![provider("talker", &talker, Protocol::Anthropic)]).await; + // 同一段对话的两个请求（指纹一样，归同一次会话），和另一段对话的一个 + let other = json!({"model": "claude-sonnet-5", "max_tokens": 1024, "stream": true, + "system": "Another conversation entirely.", + "messages": [{"role": "user", "content": "Something else."}]}); + let a = tokio::spawn(async move { post(gw, "/v1/messages", &messages(true)).await }); + let (first, session) = started(&mut rx).await; + let b = tokio::spawn(async move { post(gw, "/v1/messages", &messages(true)).await }); + let (second, again) = started(&mut rx).await; + let c = tokio::spawn(async move { post(gw, "/v1/messages", &other).await }); + let (third, elsewhere) = started(&mut rx).await; + let session = session.expect("认得出会话"); + assert_eq!(again.as_deref(), Some(session.as_str())); + assert_ne!(elsewhere.as_deref(), Some(session.as_str())); + tokio::time::sleep(Duration::from_millis(400)).await; + + assert_eq!(state.aborts.session(&session), vec![first, second]); + for asking in [a, b] { + let (status, text) = asking.await.unwrap(); + assert_eq!(status, 200); + assert!(text.contains("event: error"), "{text}"); + } + assert!(state.aborts.session(&session).is_empty(), "都结束了"); + // 另一段对话照常 + state.aborts.request(third).expect("另一段对话还在跑"); + let (_, text) = c.await.unwrap(); + assert!(text.contains("event: error"), "{text}"); +} diff --git a/crates/tw-gateway/tests/latency_samples.rs b/crates/tw-gateway/tests/latency_samples.rs index 3f451ef8..88a5b171 100644 --- a/crates/tw-gateway/tests/latency_samples.rs +++ b/crates/tw-gateway/tests/latency_samples.rs @@ -1,6 +1,6 @@ //! 快慢样本（`url-test` 和按快慢分的 `load-balance` 看的那个数）量的是哪一段，端到端：从 //! 这一跳发出去到回答的第一段内容 —— 不含之前失败了的几跳，不是响应头，和这一跳排在第几家 -//! 无关；整包的回答不记；开头慢被放弃的那一家记它被给的那段时间。 +//! 无关；整包的回答不记；没有内容超时被放弃的那一家记它被给的那段时间。 use std::net::SocketAddr; use std::sync::Arc; @@ -93,9 +93,9 @@ fn provider(name: &str, at: SocketAddr) -> Provider { } } -/// 起网关，按声明的顺序故障转移。开头最多等 1 秒（测试里图快；配置校验要求换家时至少 -/// 5 秒，网关自己不查） -async fn gateway(providers: Vec, switch: bool) -> (SocketAddr, Arc) { +/// 起网关，按声明的顺序故障转移。无响应超时 1 秒：配置写最短的 30 秒，测试把一秒调成 +/// 三十分之一秒 +async fn gateway(providers: Vec) -> (SocketAddr, Arc) { let cfg = Config { version: 1, clients: vec![Client { @@ -105,13 +105,13 @@ async fn gateway(providers: Vec, switch: bool) -> (SocketAddr, Arc, - /// 发完了不收尾，每 100 毫秒发一行注释（SSE 的心跳），直到连接断开 - hang: bool, -} - -/// 假上游被打了几次，以及它的响应（或者还没回的那个请求）被丢掉了没有 -struct Upstream { - addr: SocketAddr, - hits: Arc, - dropped: Arc, -} - -/// 被丢掉时举旗 -struct Flag(Arc); - -impl Drop for Flag { - fn drop(&mut self) { - self.0.store(true, Ordering::SeqCst); - } -} - -async fn upstream(script: Script) -> Upstream { - let hits = Arc::new(AtomicUsize::new(0)); - let dropped = Arc::new(AtomicBool::new(false)); - let (h, d) = (hits.clone(), dropped.clone()); - let app = Router::new().fallback(axum::routing::any(move || { - let (script, h, d) = (script.clone(), h.clone(), d.clone()); - async move { - h.fetch_add(1, Ordering::SeqCst); - // 响应头还没回时请求就被丢掉的，也要看得见 - let waiting = Flag(d.clone()); - tokio::time::sleep(Duration::from_millis(script.header_delay_ms)).await; - std::mem::forget(waiting); - let flag = Flag(d); - let steps = futures::stream::iter(script.steps).then(|(ms, chunk)| async move { - tokio::time::sleep(Duration::from_millis(ms)).await; - Ok::<_, std::convert::Infallible>(Bytes::from_static(chunk.as_bytes())) - }); - let hang = script.hang; - let tail = futures::stream::unfold((), move |_| async move { - if !hang { - return None; - } - tokio::time::sleep(Duration::from_millis(100)).await; - Some((Ok(Bytes::from_static(b": ping\n\n")), ())) - }); - let body = steps.chain(tail).map(move |x| { - let _ = &flag; - x - }); - axum::response::Response::builder() - .header("content-type", "text/event-stream") - .body(axum::body::Body::from_stream(body)) - .unwrap() - } - })); - let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = l.local_addr().unwrap(); - tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); - Upstream { - addr, - hits, - dropped, - } -} - -const MESSAGE_START: &str = "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1200,\"cache_read_input_tokens\":800,\"output_tokens\":1}}}\n\n"; -const PING: &str = "event: ping\ndata: {\"type\":\"ping\"}\n\n"; -const THINKING: &str = "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\n\ - event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"Let me think\"}}\n\n"; - -/// 一段完整的回答，`word` 是正文 -fn answer(word: &'static str) -> &'static str { - let text = format!( - "event: content_block_delta\ndata: {{\"type\":\"content_block_delta\",\"index\":0,\"delta\":{{\"type\":\"text_delta\",\"text\":\"{word}\"}}}}\n\n\ - event: message_stop\ndata: {{\"type\":\"message_stop\"}}\n\n" - ); - Box::leak(text.into_boxed_str()) -} - -/// 开了流、报了输入，之后只有心跳 -fn stalled() -> Script { - Script { - header_delay_ms: 0, - steps: vec![(0, MESSAGE_START), (0, PING)], - hang: true, - } -} - -/// 马上就答 -fn prompt(word: &'static str) -> Script { - Script { - header_delay_ms: 0, - steps: vec![(0, MESSAGE_START), (0, answer(word))], - hang: false, - } -} - -/// 开了流，隔 `ms` 毫秒才答 -fn late(ms: u64, word: &'static str) -> Script { - Script { - header_delay_ms: 0, - steps: vec![(0, MESSAGE_START), (ms, answer(word))], - hang: false, - } -} - -fn provider(name: &str, up: &Upstream, protocol: Protocol) -> Provider { - Provider { - name: name.into(), - base_url: format!("http://{}", up.addr), - key: Some("sk-upstream".into()), - protocol: Some(protocol), - ..Default::default() - } -} - -/// 起网关。等 1 秒（测试里图快；配置校验要求开着时至少 5 秒，网关自己不查） -async fn gateway( - providers: Vec, - switch: bool, -) -> ( - SocketAddr, - tokio::sync::broadcast::Receiver, - Arc, -) { - let cfg = Config { - version: 1, - clients: vec![Client { - name: "c".into(), - key: "tw-k".into(), - ..Default::default() - }], - providers, - failover: Failover { - stream_start_wait_secs: 1, - next_on_slow_start: switch, - ..Default::default() - }, - ..Default::default() - }; - let state = tw_gateway::AppState::new(cfg).unwrap(); - let rx = state.bus.subscribe(); - let health = state.health.clone(); - let addr = tw_gateway::serve(state, ([127, 0, 0, 1], 0).into()) - .await - .unwrap(); - tokio::time::sleep(Duration::from_millis(50)).await; - (addr, rx, health) -} - -fn messages(stream: bool) -> Value { - json!({"model": "claude-sonnet-5", "max_tokens": 1024, "stream": stream, - "messages": [{"role": "user", "content": "Say hello."}]}) -} - -async fn post(gw: SocketAddr, path: &str, body: &Value) -> (u16, String) { - let resp = reqwest::Client::new() - .post(format!("http://{gw}{path}")) - .header("content-type", "application/json") - .header("x-api-key", "tw-k") - .header("authorization", "Bearer tw-k") - .body(body.to_string()) - .send() - .await - .unwrap(); - let status = resp.status().as_u16(); - (status, resp.text().await.unwrap()) -} - -/// 这个请求的估算输入（开始事件带着）和尝试链 -async fn estimate_and_attempts( - rx: &mut tokio::sync::broadcast::Receiver, -) -> (Option, Vec) { - let mut estimate = None; - loop { - let e = tokio::time::timeout(Duration::from_secs(10), rx.recv()) - .await - .expect("no routing event") - .unwrap(); - match e { - tw_api::Event::RequestStarted { input_estimate, .. } => estimate = input_estimate, - tw_api::Event::RequestRouted { attempts, .. } => return (estimate, attempts), - _ => {} - } - } -} - -/// 等旗举起来，最多两秒 -async fn eventually(flag: &AtomicBool) -> bool { - for _ in 0..40 { - if flag.load(Ordering::SeqCst) { - return true; - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - false -} - -fn outcomes(attempts: &[tw_api::AttemptView]) -> Vec<(&str, tw_api::AttemptOutcome)> { - attempts - .iter() - .map(|a| (a.provider.as_str(), a.outcome)) - .collect() -} - -#[tokio::test] -async fn a_slow_first_upstream_is_dropped_after_the_wait_and_the_next_one_answers() { - let slow = upstream(stalled()).await; - let good = upstream(prompt("hello")).await; - let (gw, mut rx, health) = gateway( - vec![ - provider("slow", &slow, Protocol::Anthropic), - provider("good", &good, Protocol::Anthropic), - ], - true, - ) - .await; - - let t = Instant::now(); - let (status, text) = post(gw, "/v1/messages", &messages(true)).await; - assert_eq!(status, 200, "{text}"); - assert!(text.contains("hello"), "{text}"); - assert!( - !text.contains("ping"), - "慢的那一家的开头不该到客户端：{text}" - ); - assert!(t.elapsed() >= Duration::from_secs(1), "等够了才换"); - assert!( - eventually(&slow.dropped).await, - "放弃的那一家连接要断开，它才会停下" - ); - - let (_, attempts) = estimate_and_attempts(&mut rx).await; - use tw_api::AttemptOutcome::{Served, SlowStart}; - assert_eq!(outcomes(&attempts), [("slow", SlowStart), ("good", Served)]); - let gave_up = &attempts[0]; - assert_eq!(gave_up.status, Some(200)); - assert_eq!( - gave_up.usage, - Some(tw_api::AttemptUsage { - input: 1200, - cache_read: 800, - cache_write: 0, - estimated: false, - }), - "上游在开头报了输入，记它报的" - ); - assert_eq!( - gave_up.error.as_ref().map(|m| m.code.as_str()), - Some("gw.slow_start") - ); - - // **慢不是坏**：不停用，几次之后也照样先试它 - for _ in 0..4 { - let (status, _) = post(gw, "/v1/messages", &messages(true)).await; - assert_eq!(status, 200); - } - assert!(health.is_available("slow"), "慢的那一家不该被停用"); - assert_eq!(slow.hits.load(Ordering::SeqCst), 5, "每次都还先试它"); - // 按成败分的负载均衡看的成功率也不记它：换走了五次，它一次都没失败过，也没答上过 - let rates = health.success_rates(&["slow".to_string(), "good".to_string()]); - assert_eq!(rates.get("slow"), None, "{rates:?}"); - assert_eq!(rates.get("good"), Some(&1.0), "{rates:?}"); -} - -#[tokio::test] -async fn the_last_upstream_is_not_given_up_on() { - // 第一家慢、被放弃；第二家也慢，但它是最后一家：照常等它 - let slow = upstream(stalled()).await; - let last = upstream(late(1_600, "finally")).await; - let (gw, mut rx, _) = gateway( - vec![ - provider("slow", &slow, Protocol::Anthropic), - provider("last", &last, Protocol::Anthropic), - ], - true, - ) - .await; - let (status, text) = post(gw, "/v1/messages", &messages(true)).await; - assert_eq!(status, 200, "{text}"); - assert!(text.contains("finally"), "{text}"); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - use tw_api::AttemptOutcome::{Served, SlowStart}; - assert_eq!(outcomes(&attempts), [("slow", SlowStart), ("last", Served)]); -} - -#[tokio::test] -async fn a_next_upstream_that_is_paused_by_then_does_not_count() { - // 第一家慢，第二家在等的时候被停用了：后面没有接得下的，第一家就是最后一家，照常等它 - let slow = upstream(late(1_600, "patience")).await; - let other = upstream(prompt("other")).await; - let (gw, mut rx, health) = gateway( - vec![ - provider("slow", &slow, Protocol::Anthropic), - provider("other", &other, Protocol::Anthropic), - ], - true, - ) - .await; - let pause = tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(300)).await; - for _ in 0..3 { - health.record_failure("other"); - } - assert!(!health.is_available("other")); - }); - let (status, text) = post(gw, "/v1/messages", &messages(true)).await; - pause.await.unwrap(); - assert_eq!(status, 200, "{text}"); - assert!(text.contains("patience"), "{text}"); - assert_eq!(other.hits.load(Ordering::SeqCst), 0); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - assert_eq!( - outcomes(&attempts), - [("slow", tw_api::AttemptOutcome::Served)] - ); -} - -#[tokio::test] -async fn a_next_upstream_that_cannot_take_the_request_does_not_count() { - // Claude Code 搜网页的请求强制要用服务端工具 `web_search`，发不到 Chat 格式的上游： - // 那一家接不下，第一家就是最后一家 - let slow = upstream(late(1_600, "searched")).await; - let chat = upstream(prompt("other")).await; - let (gw, mut rx, _) = gateway( - vec![ - provider("slow", &slow, Protocol::Anthropic), - provider("chat", &chat, Protocol::OpenaiChat), - ], - true, - ) - .await; - let body = json!({ - "model": "claude-sonnet-5", "max_tokens": 1024, "stream": true, - "messages": [{"role": "user", "content": "Perform a web search for the query: bedrock pricing"}], - "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}], - "tool_choice": {"type": "tool", "name": "web_search"} - }); - let (status, text) = post(gw, "/v1/messages", &body).await; - assert_eq!(status, 200, "{text}"); - assert!(text.contains("searched"), "{text}"); - assert_eq!(chat.hits.load(Ordering::SeqCst), 0); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - assert_eq!( - outcomes(&attempts), - [("slow", tw_api::AttemptOutcome::Served)] - ); -} - -#[tokio::test] -async fn switched_off_the_slow_start_is_handed_on_as_before() { - let slow = upstream(late(1_600, "eventually")).await; - let other = upstream(prompt("other")).await; - let (gw, mut rx, _) = gateway( - vec![ - provider("slow", &slow, Protocol::Anthropic), - provider("other", &other, Protocol::Anthropic), - ], - false, - ) - .await; - let (status, text) = post(gw, "/v1/messages", &messages(true)).await; - assert_eq!(status, 200, "{text}"); - assert!(text.contains("eventually"), "{text}"); - assert_eq!(other.hits.load(Ordering::SeqCst), 0); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - assert_eq!( - outcomes(&attempts), - [("slow", tw_api::AttemptOutcome::Served)] - ); -} - -#[tokio::test] -async fn thinking_counts_as_content() { - // 开头就在想（推理的字在流），正文要过一阵才来：这不是开头慢 - let thinker = upstream(Script { - header_delay_ms: 0, - steps: vec![ - (0, MESSAGE_START), - (200, THINKING), - (1_400, answer("thought")), - ], - hang: false, - }) - .await; - let other = upstream(prompt("other")).await; - let (gw, mut rx, _) = gateway( - vec![ - provider("thinker", &thinker, Protocol::Anthropic), - provider("other", &other, Protocol::Anthropic), - ], - true, - ) - .await; - let (status, text) = post(gw, "/v1/messages", &messages(true)).await; - assert_eq!(status, 200, "{text}"); - assert!( - text.contains("Let me think") && text.contains("thought"), - "{text}" - ); - assert_eq!(other.hits.load(Ordering::SeqCst), 0); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - assert_eq!( - outcomes(&attempts), - [("thinker", tw_api::AttemptOutcome::Served)] - ); -} - -#[tokio::test] -async fn without_reported_usage_the_estimate_is_recorded_as_possibly_billed() { - // Chat 格式的上游开头只发一块角色、不报用量：记网关估的输入 - let silent = upstream(Script { - header_delay_ms: 0, - steps: vec![( - 0, - "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":\"\"}}]}\n\n", - )], - hang: true, - }) - .await; - let good = upstream(Script { - header_delay_ms: 0, - steps: vec![( - 0, - "data: {\"choices\":[{\"delta\":{\"content\":\"hi there\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n", - )], - hang: false, - }) - .await; - let (gw, mut rx, _) = gateway( - vec![ - provider("silent", &silent, Protocol::OpenaiChat), - provider("good", &good, Protocol::OpenaiChat), - ], - true, - ) - .await; - let body = json!({"model": "gpt-5", "stream": true, - "messages": [{"role": "user", "content": "Say hello."}]}); - let (status, text) = post(gw, "/v1/chat/completions", &body).await; - assert_eq!(status, 200, "{text}"); - assert!(text.contains("hi there"), "{text}"); - assert!(eventually(&silent.dropped).await); - let (estimate, attempts) = estimate_and_attempts(&mut rx).await; - use tw_api::AttemptOutcome::{Served, SlowStart}; - assert_eq!( - outcomes(&attempts), - [("silent", SlowStart), ("good", Served)] - ); - let estimate = estimate.expect("解得开的请求有估算"); - assert!(estimate > 0); - assert_eq!( - attempts[0].usage, - Some(tw_api::AttemptUsage { - input: estimate, - cache_read: 0, - cache_write: 0, - estimated: true, - }) - ); -} - -#[tokio::test] -async fn response_headers_that_never_come_count_as_a_slow_start() { - // 响应头都迟迟不来（中转站排着队）：一样从发出请求算起，到点放弃 - let queued = upstream(Script { - header_delay_ms: 5_000, - ..prompt("too late") - }) - .await; - let good = upstream(prompt("hello")).await; - let (gw, mut rx, _) = gateway( - vec![ - provider("queued", &queued, Protocol::Anthropic), - provider("good", &good, Protocol::Anthropic), - ], - true, - ) - .await; - let t = Instant::now(); - let (status, text) = post(gw, "/v1/messages", &messages(true)).await; - assert_eq!(status, 200, "{text}"); - assert!(text.contains("hello"), "{text}"); - assert!(t.elapsed() < Duration::from_secs(4), "没有等到响应头"); - assert!(eventually(&queued.dropped).await, "还在等的请求要断开"); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - use tw_api::AttemptOutcome::{Served, SlowStart}; - assert_eq!( - outcomes(&attempts), - [("queued", SlowStart), ("good", Served)] - ); - assert_eq!(attempts[0].status, None, "响应头没到，没有状态码"); - assert!(attempts[0].usage.is_some_and(|u| u.estimated)); -} - -#[tokio::test] -async fn a_request_that_does_not_stream_is_not_switched() { - // 整包的请求本来就要等全部生成完：响应头来得晚也照常等 - let slow = upstream(Script { - header_delay_ms: 1_600, - ..prompt("whole") - }) - .await; - let other = upstream(prompt("other")).await; - let (gw, mut rx, _) = gateway( - vec![ - provider("slow", &slow, Protocol::Anthropic), - provider("other", &other, Protocol::Anthropic), - ], - true, - ) - .await; - let (status, text) = post(gw, "/v1/messages", &messages(false)).await; - assert_eq!(status, 200, "{text}"); - assert_eq!(other.hits.load(Ordering::SeqCst), 0); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - assert_eq!( - outcomes(&attempts), - [("slow", tw_api::AttemptOutcome::Served)] - ); -} - -/// 放弃的那一家占着的位置（`max_concurrent`，见 `tw_gateway::slots`）**当场**还回去：接下 -/// 请求的那一家还在答，慢的那一家已经空出来了，不等这个请求结束 -#[tokio::test] -async fn the_slot_of_an_upstream_given_up_on_is_free_at_once() { - const CONTENT: &str = "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n"; - const STOP: &str = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; - let slow = upstream(stalled()).await; - // 先答上，隔两秒才说完：这两秒里请求还没结束 - let good = upstream(Script { - header_delay_ms: 0, - steps: vec![(0, MESSAGE_START), (0, CONTENT), (2_000, STOP)], - hang: false, - }) - .await; - let cfg = Config { - version: 1, - clients: vec![Client { - name: "c".into(), - key: "tw-k".into(), - ..Default::default() - }], - providers: vec![ - Provider { - max_concurrent: Some(1), - ..provider("slow", &slow, Protocol::Anthropic) - }, - provider("good", &good, Protocol::Anthropic), - ], - failover: Failover { - stream_start_wait_secs: 1, - next_on_slow_start: true, - ..Default::default() - }, - ..Default::default() - }; - let state = tw_gateway::AppState::new(cfg).unwrap(); - let gw = tw_gateway::serve(state.clone(), ([127, 0, 0, 1], 0).into()) - .await - .unwrap(); - tokio::time::sleep(Duration::from_millis(50)).await; - - let asking = tokio::spawn(async move { post(gw, "/v1/messages", &messages(true)).await }); - // 还在等慢的那一家：它的位置占着 - tokio::time::sleep(Duration::from_millis(300)).await; - assert_eq!(slow.hits.load(Ordering::SeqCst), 1); - assert!(state.slots.try_take("slow").is_none(), "等着的时候占着位置"); - // 换到了好的那一家：慢的那一家的位置已经还回来了，这个请求还没结束 - for _ in 0..60 { - if good.hits.load(Ordering::SeqCst) == 1 { - break; - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - assert_eq!(good.hits.load(Ordering::SeqCst), 1); - assert!(!asking.is_finished(), "测的是请求还在进行时"); - assert!( - state.slots.try_take("slow").is_some(), - "放弃的那一家的位置要当场还回去" - ); - let (status, text) = asking.await.unwrap(); - assert_eq!(status, 200, "{text}"); - assert!(text.contains("hello"), "{text}"); -} - -/// 起网关，交回数据面的状态（要占住某一家的位置）。`slot_wait_secs` 是等空位的期限 -async fn gateway_with_state( - providers: Vec, - slot_wait_secs: u64, -) -> ( - SocketAddr, - tokio::sync::broadcast::Receiver, - tw_gateway::AppState, -) { - let cfg = Config { - version: 1, - clients: vec![Client { - name: "c".into(), - key: "tw-k".into(), - ..Default::default() - }], - providers, - failover: Failover { - stream_start_wait_secs: 1, - next_on_slow_start: true, - slot_wait_secs, - ..Default::default() - }, - ..Default::default() - }; - let state = tw_gateway::AppState::new(cfg).unwrap(); - let rx = state.bus.subscribe(); - let addr = tw_gateway::serve(state.clone(), ([127, 0, 0, 1], 0).into()) - .await - .unwrap(); - tokio::time::sleep(Duration::from_millis(50)).await; - (addr, rx, state) -} - -/// 并发数满着的那一家（`max_concurrent`） -fn full(name: &str, up: &Upstream) -> Provider { - Provider { - max_concurrent: Some(1), - ..provider(name, up, Protocol::Anthropic) - } -} - -/// 后面只剩一家满着的（`max_concurrent`）：到点时它接不下，慢的这一家就是最后一家，照常等它。 -/// 放弃了它，换来的是去等那一家空出来、等不到回 429 —— 一个本来答得上的请求就这样丢了 -#[tokio::test] -async fn a_next_upstream_that_is_full_at_the_deadline_does_not_count() { - let busy = upstream(prompt("busy")).await; - let slow = upstream(late(2_500, "patience")).await; - let (gw, mut rx, state) = gateway_with_state( - vec![ - full("busy", &busy), - provider("slow", &slow, Protocol::Anthropic), - ], - 3, - ) - .await; - // 满着的那一家排在前面：它被当场跳过，慢的那一家发出去时后面还「有」它 - let _held = state.slots.try_take("busy").expect("一个空位"); - let (status, text) = post(gw, "/v1/messages", &messages(true)).await; - assert_eq!(status, 200, "{text}"); - assert!(text.contains("patience"), "{text}"); - assert_eq!(busy.hits.load(Ordering::SeqCst), 0); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - use tw_api::AttemptOutcome::{Error, Served}; - assert_eq!(outcomes(&attempts), [("busy", Error), ("slow", Served)]); - assert_eq!(attempts[0].skipped, Some(tw_api::ServeSkip::Busy)); -} - -/// 后面有一家此刻接得下：照常换过去，排在前面满着的那一家不挡它 -#[tokio::test] -async fn a_next_upstream_with_a_free_slot_still_takes_over() { - let busy = upstream(prompt("busy")).await; - let slow = upstream(stalled()).await; - let good = upstream(prompt("good")).await; - let (gw, mut rx, state) = gateway_with_state( - vec![ - full("busy", &busy), - provider("slow", &slow, Protocol::Anthropic), - provider("good", &good, Protocol::Anthropic), - ], - 3, - ) - .await; - let _held = state.slots.try_take("busy").expect("一个空位"); - let (status, text) = post(gw, "/v1/messages", &messages(true)).await; - assert_eq!(status, 200, "{text}"); - assert!(text.contains("good"), "{text}"); - assert_eq!(busy.hits.load(Ordering::SeqCst), 0); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - use tw_api::AttemptOutcome::{Error, Served, SlowStart}; - assert_eq!( - outcomes(&attempts), - [("busy", Error), ("slow", SlowStart), ("good", Served)] - ); -} - -/// 跳过时满着的那一家，到点之前空出来了：它此刻接得下，换到它 -#[tokio::test] -async fn an_upstream_that_frees_before_the_deadline_takes_over() { - let busy = upstream(prompt("freed")).await; - let slow = upstream(stalled()).await; - let (gw, mut rx, state) = gateway_with_state( - vec![ - full("busy", &busy), - provider("slow", &slow, Protocol::Anthropic), - ], - 3, - ) - .await; - let held = state.slots.try_take("busy").expect("一个空位"); - let asking = tokio::spawn(async move { post(gw, "/v1/messages", &messages(true)).await }); - tokio::time::sleep(Duration::from_millis(400)).await; - drop(held); - let (status, text) = asking.await.unwrap(); - assert_eq!(status, 200, "{text}"); - assert!(text.contains("freed"), "{text}"); - let (_, attempts) = estimate_and_attempts(&mut rx).await; - use tw_api::AttemptOutcome::{Error, Served, SlowStart}; - assert_eq!( - outcomes(&attempts), - [("busy", Error), ("slow", SlowStart), ("busy", Served)] - ); -} diff --git a/crates/tw-store/src/db.rs b/crates/tw-store/src/db.rs index b55c36db..85cdb888 100644 --- a/crates/tw-store/src/db.rs +++ b/crates/tw-store/src/db.rs @@ -27,7 +27,10 @@ use tw_api::Msg; /// `matching`、`revealed`），结局多了「已删除」。 /// /// 25：插件在每个请求上的运行记录（`plugin_runs`）。 -pub(crate) const SCHEMA: i64 = 25; +/// +/// 26：尝试链的结果（`routing` 里的 `outcome`）没有了 `slow_start`，多了 `idle_timeout` 和 +/// `aborted`。 +pub(crate) const SCHEMA: i64 = 26; /// 这一行算不出钱，**因为价目表里没有这个模型**：用量是有的，缺的是单价。 /// diff --git a/crates/tw-store/src/health.rs b/crates/tw-store/src/health.rs index 992b686e..1d42f5f1 100644 --- a/crates/tw-store/src/health.rs +++ b/crates/tw-store/src/health.rs @@ -3,7 +3,8 @@ //! **只摆事实和参照，不下结论**（见 [`tw_api::UpstreamHealth`]）。这里的每一个数都要 //! 说得清是从哪些请求里数出来的、数了几条： //! -//! - 失败：口径和概览一样（`error IS NOT NULL`），客户端取消的不算失败，单独数； +//! - 失败：口径和概览一样（`error IS NOT NULL`），客户端取消的不算失败，单独数；在界面上 +//! 手动中止的（[`tw_api::ABORTED`]）也不是上游的错，和取消的一起数； //! - 模型名：回答里写了的，归一之后和发出去的比（[`crate::model_name`]）； //! - 输入之比：上游报的输入（三项加起来）÷ 本地估算，取中位数，再拿别家服务同一个模型 //! 时的中位数作参照 —— 估算只准到两三成，参照才让这个数有意义； @@ -58,15 +59,21 @@ impl Db { to_ms: i64, ) -> Result { let mut st = self.conn.prepare( - "SELECT at_ms, provider, sent_model, answered_model, session, cancelled, - error IS NOT NULL, input_tokens, COALESCE(cache_read_tokens, 0), + "SELECT at_ms, provider, sent_model, answered_model, session, + cancelled OR COALESCE(error_code = ?3, 0), + error IS NOT NULL AND NOT COALESCE(error_code = ?3, 0), + input_tokens, COALESCE(cache_read_tokens, 0), COALESCE(cache_write_tokens, 0), input_estimate, ttft_ms, tokens_per_sec FROM requests WHERE at_ms >= ?1 AND at_ms < ?2 AND local = 0 ORDER BY at_ms, id", )?; let mut tally = Tally::default(); - let mut rows = st.query(params![from_ms.saturating_sub(CACHE_WARM_MS), to_ms])?; + let mut rows = st.query(params![ + from_ms.saturating_sub(CACHE_WARM_MS), + to_ms, + tw_api::ABORTED + ])?; while let Some(r) = rows.next()? { tally.add(&Line::read(r)?, from_ms); } @@ -525,13 +532,19 @@ mod tests { ); } - /// **本地应答、规则拒绝了的不归哪一家**；取消的不算失败，也不在请求数里 + /// **本地应答、规则拒绝了的不归哪一家**；取消的、手动中止的不算失败，也不在请求数里 #[test] fn requests_failures_and_cancellations_are_counted_apart() { let mut failed = on(2, T0 + 2, "中转", "claude-sonnet-4-5"); failed.error = Some(upstream_failed("502")); let mut cancelled = on(3, T0 + 3, "中转", "claude-sonnet-4-5"); cancelled.cancelled = true; + let mut aborted = on(9, T0 + 9, "中转", "claude-sonnet-4-5"); + aborted.error = Some(tw_api::Msg { + code: tw_api::ABORTED.into(), + args: Default::default(), + text: "The request was aborted by the user.".into(), + }); let mut local = on(4, T0 + 4, "中转", "claude-sonnet-4-5"); local.local = true; let mut refused = on(5, T0 + 5, "", "claude-sonnet-4-5"); @@ -545,6 +558,7 @@ mod tests { on(6, T0 + 6, "官方", "claude-sonnet-4-5"), on(7, T0 + 7, "官方", "claude-sonnet-4-5"), on(8, T0 + 8, "官方", "claude-sonnet-4-5"), + aborted, ]); assert_eq!( h.upstreams @@ -557,7 +571,7 @@ mod tests { let relay = checkup(&h, "中转"); assert_eq!( (relay.requests, relay.failed, relay.cancelled), - (2, 1, 1), + (2, 1, 2), "{relay:?}" ); let official = checkup(&h, "官方"); diff --git a/docs/config.md b/docs/config.md index fae4f75f..2709bf29 100644 --- a/docs/config.md +++ b/docs/config.md @@ -1016,19 +1016,35 @@ Before the first content of a streamed answer reaches the client, an error the upstream sends in the stream moves the request to the next candidate, the same as an error status would. -An upstream can also be slow to start: it accepts the request and then sends -nothing for a long time. With `next_on_slow_start`, the request moves on to the -next candidate when no content has arrived `stream_start_wait_secs` after it -was sent. It is off by default, because models that think before they write -can take long to start; with it on, wait 30 seconds or more. The last -candidate always waits, and the upstream given up on is not set aside. A -candidate that is at its `max_concurrent` at that moment does not count as a -next one: the slow upstream keeps the request. +An upstream can also go quiet: it accepts the request and then sends nothing, +or stops partway through. After `idle_timeout_secs` (300 by default) without +content, the gateway stops waiting for it. The time counts from the moment the +request is sent to that upstream, so waiting for a free slot does not count, +and starts again with every piece of content: text, reasoning and tool calls +count; keep-alives (SSE comments, Anthropic's `ping`, empty chunks, Responses' +`response.in_progress`) do not, so an upstream that only keeps the connection +alive still runs out of time. A whole, non-streamed answer counts from sending +to the complete answer. + +- When no content has reached the client yet, the upstream counts as failed + (towards `failures_to_pause`, like a 5xx), the attempt appears with the + outcome `idle_timeout`, and the request moves to the next candidate. Until + then a streamed answer is held back from the client, so the next upstream + starts it afresh. With no candidate left, the client gets a timeout error + (504) in its own format. The last candidate's stream is passed on as it + arrives, so once its response has started, a timeout there ends it with an + error event instead. +- When content has already reached the client, the request cannot move on + without repeating it: the answer ends with an error event in the client's + format, and the request is recorded as failed. + +The conversation then no longer stays on that upstream for the rest of its turn. +Models that think for a long time before they write anything need a longer +timeout. ```yaml failover: - stream_start_wait_secs: 30 - next_on_slow_start: true + idle_timeout_secs: 600 ``` When upstreams are at their `max_concurrent`, a request waits for a free slot @@ -1053,8 +1069,7 @@ failure instead. | `no_balance_pause_secs` | integer | `1800` | Seconds to set aside an upstream that reports an insufficient balance. | | `quota_pause_secs` | integer | `3600` | Seconds to set aside an upstream whose quota is used up when it does not say when the quota resets. When it does, the upstream is set aside until then. | | `rate_limit_max_pause_secs` | integer | `3600` | A rate-limited upstream is set aside for the time its `Retry-After` gives, at most this many seconds. Without `Retry-After` it counts as a failure without a stated reason. | -| `stream_start_wait_secs` | integer | `15` | Seconds to hold a streamed answer until its first content arrives. An error before then moves the request to the next upstream; after this long, what has arrived is passed on. From 1 to 120. | -| `next_on_slow_start` | bool | `false` | When a streamed answer still has no content `stream_start_wait_secs` after the request was sent, give up on that upstream and send the request to the next one. The last upstream always waits. The upstream given up on is not set aside. Needs `stream_start_wait_secs` of at least 5. | +| `idle_timeout_secs` | integer | `300` | Seconds an upstream may go without sending content before the gateway stops waiting for it. Counted from the moment the request is sent and started again by every piece of content: text, reasoning and tool calls count, keep-alives do not. A whole (non-streamed) answer counts from sending to the complete answer. Before any content has reached the client, the upstream counts as failed and the request moves to the next one; with none left, the client gets a timeout error. After content has reached the client, the answer ends with an error. From 30 to 3600. | | `slot_wait_secs` | integer | `30` | Seconds a request waits in all, counted once the key's own `max_concurrent` lets it in: for a key's `minute` or `hour` limit to free up, and for a free slot on upstreams at their `max_concurrent`. A key limit that does not free up in time refuses the request; without an upstream slot in time it goes to the next upstream, or, when every candidate is full, is answered with 429. `0`: never wait. From 0 to 300. | @@ -1168,11 +1183,11 @@ group shares out requests by the result in the same way as above. tenth. - `health`: upstreams that fail less get a larger share. It looks at the last 50 requests within the past 30 minutes. Server errors, rate limits, - used-up quota or balance, rejected credentials, timeouts and connection - errors count as failures; errors caused by the request itself do not, and - neither does a client that cancels, a switch away from a stream that is - slow to start, or an upstream skipped because it is at its - `max_concurrent`. An upstream that keeps failing keeps a twentieth of its + used-up quota or balance, rejected credentials, timeouts (including + `failover.idle_timeout_secs` before any content) and connection errors + count as failures; errors caused by the request itself do not, and neither + does a client that cancels, a request aborted by hand, or an upstream + skipped because it is at its `max_concurrent`. An upstream that keeps failing keeps a twentieth of its weight, so it still gets the occasional request and its recovery is noticed; one that fails outright is set aside by [`failover`](#cfg-failover) as before. @@ -1180,8 +1195,8 @@ group shares out requests by the result in the same way as above. Speed is measured on streamed answers only, from the moment the request is sent to that upstream, so waiting and upstreams that failed before it do not -count. An upstream given up on because its stream was slow to start -(`failover.next_on_slow_start`) counts as taking the whole wait. On a +count. An upstream given up on because it sent no content within +`failover.idle_timeout_secs` counts as taking the whole wait. On a Responses WebSocket connection, each `response.create` counts as one request for both speed and failures, its speed measured from the moment the upstream starts answering it. diff --git a/docs/config.zh-CN.md b/docs/config.zh-CN.md index 6e5f3e33..9747da32 100644 --- a/docs/config.zh-CN.md +++ b/docs/config.zh-CN.md @@ -785,16 +785,24 @@ security: 流式回答在第一段内容交给客户端之前，上游在流里报的错误和错误状态码一样，会把 请求换到下一个候选。 -上游也可能开头很慢：收下请求之后很久都不发内容。开启 `next_on_slow_start` 后，请求 -发出 `stream_start_wait_secs` 秒仍没有内容，就换到下一个候选。默认关闭，因为先思考 -再输出的模型本来就可能很久才开始；开启时建议等 30 秒以上。最后一个候选总是等下去， -被放弃的上游不会停用。到点那一刻并发数已满（`max_concurrent`）的候选不算下一个：请求 -留在慢的那一家。 +上游也可能不出声：收下请求之后一直不发内容，或者发到一半停住。连续 +`idle_timeout_secs` 秒（默认 300）没有内容，网关就不再等它。计时从请求发给这家上游的 +那一刻起，等空位的时间不算在内；每来一段内容重新计时：正文、推理、工具调用都算，心跳 +（SSE 注释、Anthropic 的 `ping`、空块、Responses 的 `response.in_progress`）不算，所以 +只保持连接、不出内容的上游同样会超时。整包（非流式）的回答从发出算到整份回来。 + +- 还没有内容交给客户端时，这家上游记一次失败（和 5xx 一样计入 `failures_to_pause`）， + 尝试链上这一跳的结果是 `idle_timeout`，请求换到下一个候选。在此之前流式回答不交给 + 客户端，下一家从头开始回答。没有候选了，客户端收到它自己格式的超时错误（504）。 + 最后一个候选的流是边收边交给客户端的，它的响应开始之后再超时，回答以一条错误事件收尾。 +- 已经有内容交给客户端时，换一家会重复已经发出的内容，所以不换：回答按客户端的格式 + 以一条错误事件收尾，请求记为失败。 + +之后这段对话在这一轮里不再留在这家上游。先思考很久才输出的模型，需要把超时调长。 ```yaml failover: - stream_start_wait_secs: 30 - next_on_slow_start: true + idle_timeout_secs: 600 ``` 上游的并发数满了（`max_concurrent`）时，一个请求等空位合计最多 `slot_wait_secs` @@ -815,8 +823,7 @@ failover: | `no_balance_pause_secs` | 整数 | `1800` | 上游报告余额不足时停用的秒数。 | | `quota_pause_secs` | 整数 | `3600` | 上游报告额度用完、但没有给出重置时间时停用的秒数。给出了重置时间的，停用到那一刻。 | | `rate_limit_max_pause_secs` | 整数 | `3600` | 被限流的上游按它给的 `Retry-After` 停用，最多这么多秒。没有 `Retry-After` 的按没有说明原因的失败计。 | -| `stream_start_wait_secs` | 整数 | `15` | 流式回答在第一段内容到达前最多暂存的秒数。在此之前上游报错，请求换到下一家；超过这个时间，已收到的部分照常交给客户端。取值 1 到 120。 | -| `next_on_slow_start` | 布尔 | `false` | 流式回答在请求发出 `stream_start_wait_secs` 秒后仍没有内容时，放弃这家上游，把请求交给下一家。最后一家总是等下去。被放弃的上游不会停用。开启时 `stream_start_wait_secs` 至少为 5。 | +| `idle_timeout_secs` | 整数 | `300` | 上游多少秒没有发出内容，网关就不再等它。从请求发出的那一刻算起，每来一段内容重新计时：正文、推理、工具调用都算，心跳不算。整包（非流式）的回答从发出算到整份回来。还没有内容交给客户端时，这家上游记一次失败，请求换到下一家；没有下一家了，客户端收到超时错误。已经有内容交给客户端的，回答以一条错误收尾。取值 30 到 3600。 | | `slot_wait_secs` | 整数 | `30` | 一个请求合计最多等的秒数，从过了密钥自己的 `max_concurrent` 时算起：等密钥的 `minute`、`hour` 上限空出名额，和等并发数满了（`max_concurrent`）的上游空出位置，都算在里面。密钥的上限到时空不出来就拒绝；等不到上游的空位就换下一家，候选全满时回 429。`0`：不等。取值 0 到 300。 | @@ -885,10 +892,10 @@ groups: - `weights`（默认）：只按权重。 - `latency`：越快的上游分得越多。快慢看典型的从发出请求到回答第一段内容的时间，与 `url-test` 使用同一份测量。比组内居中者快一倍的上游，权重乘以四；最多乘以十，最少乘以十分之一。 -- `health`：越少失败的上游分得越多。依据是最近 30 分钟内的最近 50 次请求：服务器错误、限流、额度或余额用尽、凭据被拒、超时和连接失败算作失败；请求本身导致的错误不算，客户端取消、因开头太慢而换走、因并发数满了（`max_concurrent`）而跳过也不算。经常失败的上游至少保留权重的二十分之一，仍会偶尔分到请求，以便发现它已经恢复；完全失败的上游照旧由 [`failover`](#cfg-failover) 暂停。 +- `health`：越少失败的上游分得越多。依据是最近 30 分钟内的最近 50 次请求：服务器错误、限流、额度或余额用尽、凭据被拒、超时和连接失败算作失败；请求本身导致的错误不算，客户端取消、手动中止、因并发数满了（`max_concurrent`）而跳过也不算；超过 `failover.idle_timeout_secs` 仍没有内容算作超时。经常失败的上游至少保留权重的二十分之一，仍会偶尔分到请求，以便发现它已经恢复；完全失败的上游照旧由 [`failover`](#cfg-failover) 暂停。 - `latency-health`：两个系数相乘。 -快慢只在流式回答上测，从请求发给这家上游的那一刻算起：之前的等待、之前失败的上游都不算在内。因开头太慢而被放弃的上游（`failover.next_on_slow_start`），按等满的那段时间计。Responses 的 WebSocket 连接上，每个 `response.create` 在快慢和成败上都算一个请求，快慢从上游开始回答它的那一刻算起。 +快慢只在流式回答上测，从请求发给这家上游的那一刻算起：之前的等待、之前失败的上游都不算在内。因超过 `failover.idle_timeout_secs` 没有内容而被放弃的上游，按等满的那段时间计。Responses 的 WebSocket 连接上，每个 `response.create` 在快慢和成败上都算一个请求，快慢从上游开始回答它的那一刻算起。 测量还不够的上游按中等对待。与只按权重时一样，进行中的对话留在原来的上游，差额由新对话补齐。 diff --git a/scripts/smoke.sh b/scripts/smoke.sh index 825b9599..69cff6cb 100755 --- a/scripts/smoke.sh +++ b/scripts/smoke.sh @@ -510,6 +510,13 @@ if [ "$ID" != "0" ]; then C=$(get "/request/$ID"); [ "$C" = "200" ] && ok "GET /request/{id}" || bad "返回 $C" C=$(post /replay/quote "{\"id\":$ID,\"provider\":\"official\"}") [ "$C" = "200" ] && ok "POST /replay/quote" || bad "返回 $C" + # 中止：注册着、认得出跑完了的请求（404 带着自己的码），不是 405 或者没有这个端点 + C=$(post "/request/$ID/abort" '{}') + [ "$C" = "404" ] && grep -q 'control.request_not_running' "$TMP/out" \ + && ok "POST /request/{id}/abort 说跑完了的请求不在跑" || bad "返回 $C" "$(head -c 200 "$TMP/out")" + C=$(post "/sessions/no-such-session/abort" '{}') + [ "$C" = "404" ] && grep -q 'control.session_not_running' "$TMP/out" \ + && ok "POST /sessions/{id}/abort 说没有在跑的" || bad "返回 $C" "$(head -c 200 "$TMP/out")" else bad "历史里一条记录都没有" fi From 696a12fbf9afd20b5e25128edf7a1f0fd986bbe3 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:13:26 +0800 Subject: [PATCH 2/2] feat(failover): commit 200 and keep-alives after a 15 s hold, keep failing over in the body Holding a streamed answer's headers until its first content (up to the idle timeout, 300 s by default) let clients' own header timeouts fire before the gateway's failover ever got a chance. A streamed answer is now held for at most `OPENING_HOLD` (15 s, a constant, not configurable): quick upstream errors still fail over and can be answered with a proper status. Past that, the client gets `200` and the streaming headers, then an SSE comment (`: keep-alive`) every `KEEPALIVE_EVERY` (15 s); Gemini clients get none, because Google's Python SDK parses comment lines as JSON. The rest of the pipeline (trying candidates and relaying the answer) owns everything it needs, so it simply moves into the response body and carries on: the upstream's own events stay held until its first content, an idle timeout, an in-stream error or an early close still moves the request to the next candidate, whose stream starts cleanly under the same `200`. When every candidate fails after the headers went out, the stream ends with the client-format error event used mid-stream instead of a 504; a last upstream's error answer is told the same way. The comments are ours and never touch the idle timer. Whole answers and Gemini's JSON array streams are held as before. An ending dropped after its abort switch was thrown now reports the abort rather than a client cancel, so a request aborted while the pipeline is not watching the switch is still recorded as aborted. Co-Authored-By: Claude Opus 5.5 --- crates/tw-api/src/lib.rs | 4 +- crates/tw-gateway/src/abort.rs | 20 +- crates/tw-gateway/src/ending.rs | 14 + crates/tw-gateway/src/lib.rs | 10 + crates/tw-gateway/src/server/pipeline.rs | 82 ++++- .../tw-gateway/src/server/pipeline/commit.rs | 182 ++++++++++ crates/tw-gateway/src/server/pipeline/idle.rs | 18 +- .../tw-gateway/src/server/pipeline/opening.rs | 4 +- crates/tw-gateway/src/state.rs | 6 + crates/tw-gateway/tests/idle_timeout.rs | 310 +++++++++++++++++- docs/config.md | 10 +- docs/config.zh-CN.md | 6 +- scripts/smoke.sh | 47 +++ 13 files changed, 691 insertions(+), 22 deletions(-) create mode 100644 crates/tw-gateway/src/server/pipeline/commit.rs diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 52bb265d..30b511dd 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -822,8 +822,10 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// /// **43 起上游不出声有了上限，请求能手动中止**：[`FailoverView`] 的 `stream_start_wait_secs` /// 和 `next_on_slow_start` 删了，换成 `idle_timeout_secs`（无响应超时，默认 300 秒，30 到 -/// 3600）：从请求发出去起算，每来一段内容重新计时，心跳不算。客户端还什么都没收到时，这一家 +/// 3600）：从请求发出去起算，每来一段内容重新计时，心跳不算。还没有内容交给客户端时，这一家 /// 记一次失败、换下一家，没有下一家了回 504；已经收到一部分的，回答按客户端的格式以错误收尾。 +/// 流式的回答压着第一段内容最多 15 秒，过了就先交出 `200` 和流的响应头、隔一阵发一行 SSE +/// 注释（`: keep-alive`，Gemini 的客户端不发），故障转移照旧，候选用完了在流里报错，不再是 504。 /// 配置里写 `stream_start_wait_secs`、`next_on_slow_start` 加载不了（不认识的字段，一键修复删掉 /// 它们），消息码 `config.slow_start_too_short`、`gw.slow_start` 跟着删。尝试链的结果 /// （[`AttemptOutcome`]）删了 `slow_start`，多了 `idle_timeout`（说等了多少秒的 diff --git a/crates/tw-gateway/src/abort.rs b/crates/tw-gateway/src/abort.rs index f4026cc2..06539dca 100644 --- a/crates/tw-gateway/src/abort.rs +++ b/crates/tw-gateway/src/abort.rs @@ -9,7 +9,8 @@ //! //! 登记跟着请求的结局走（挂在 [`crate::ending::Ending`] 上）：结局报了、或者被丢掉了，登记 //! 就没了。所以「在不在这张表里」就是「这个请求还在不在跑」，叫停一个已经结束的请求得到 -//! 的是「没有在跑」，而不是一次什么都没发生的成功。 +//! 的是「没有在跑」，而不是一次什么都没发生的成功。**扳过开关之后被丢掉的结局报手动中止**， +//! 不报客户端取消：停在看不着开关的地方（取凭据、跑插件）时，是外面那一层把整个请求丢掉的。 //! //! WebSocket 那条路不登记：一轮回答跑在一条长连接上，叫停一轮要连带处置整条连接，这一版 //! 不做。 @@ -76,6 +77,14 @@ pub struct Aborts { pub struct Registered { aborts: Arc, id: u64, + switch: Switch, +} + +impl Registered { + /// 这个请求被叫停了没有 + pub fn thrown(&self) -> bool { + self.switch.thrown() + } } impl Drop for Registered { @@ -96,10 +105,17 @@ impl Aborts { /// 请求 `id`（会话 `session`）开始了，用 `switch` 叫停它。 pub fn enter(self: &Arc, id: u64, session: Option, switch: Switch) -> Registered { - self.lock().insert(id, Entry { session, switch }); + self.lock().insert( + id, + Entry { + session, + switch: switch.clone(), + }, + ); Registered { aborts: Arc::clone(self), id, + switch, } } diff --git a/crates/tw-gateway/src/ending.rs b/crates/tw-gateway/src/ending.rs index 6b66958b..68561bd4 100644 --- a/crates/tw-gateway/src/ending.rs +++ b/crates/tw-gateway/src/ending.rs @@ -523,6 +523,20 @@ impl Drop for Ending { // 往通道里 try_send、往广播里 send、读一下时钟。 let (usage, answered_model) = self.settle(); let usage = usage.map(view); + // 在界面上叫停之后被丢掉的（见 `crate::abort`）：是手动中止，不是客户端走了 + if self.abort.as_ref().is_some_and(|r| r.thrown()) { + self.bus.emit(tw_api::Event::RequestFailed { + id: self.id, + model: std::mem::take(&mut self.model), + source: tw_api::FailureSource::Aborted, + message: crate::error::GatewayError::aborted().detail, + bytes: self.received(), + duration_ms: Some(self.duration_ms()), + usage, + answered_model, + }); + return; + } // 是网关自己的代码崩掉了。**记成取消会冤枉客户端** —— 排查的人 // 会去问一个根本没做过这件事的客户端。 if std::thread::panicking() { diff --git a/crates/tw-gateway/src/lib.rs b/crates/tw-gateway/src/lib.rs index 128a93ed..0539e735 100644 --- a/crates/tw-gateway/src/lib.rs +++ b/crates/tw-gateway/src/lib.rs @@ -87,6 +87,16 @@ pub const PING_EVERY: std::time::Duration = std::time::Duration::from_secs(15); /// 自己的静默计时会断开它。正常在排队的上游会发 `: keep-alive` 这样的注释，不受影响 pub const PING_FOR: std::time::Duration = std::time::Duration::from_secs(600); +/// 流式回答第一段内容到达之前，最多压多久不给客户端响应头（见 `server::pipeline::commit`）。 +/// +/// 压着的时候还能换一家上游、上游开头报的错也还能换成一个像样的状态码；压得太久，客户端 +/// 自己等响应头的计时先到了。过了这么久就先交出 `200` 和流的响应头，故障转移在响应体里 +/// 接着做。**不可配置**：它不是给用户调的取舍，是客户端那一侧的限制 +pub const OPENING_HOLD: std::time::Duration = std::time::Duration::from_secs(15); + +/// 响应头先交出去之后，等回答的时候隔多久给客户端发一行 SSE 注释（`: keep-alive`） +pub const KEEPALIVE_EVERY: std::time::Duration = std::time::Duration::from_secs(15); + /// 请求来自谁。**如实写 ThinkWatch** —— 我们从不把自己报成别的客户端。 pub const ORIGINATOR: &str = "thinkwatch"; diff --git a/crates/tw-gateway/src/server/pipeline.rs b/crates/tw-gateway/src/server/pipeline.rs index b8d36ff4..73cb9683 100644 --- a/crates/tw-gateway/src/server/pipeline.rs +++ b/crates/tw-gateway/src/server/pipeline.rs @@ -20,6 +20,7 @@ use crate::state::{AppState, Runtime}; use tw_types::msg; mod admission; +mod commit; mod hop; mod idle; mod opening; @@ -178,6 +179,73 @@ pub(super) async fn pipeline( if let Some(why) = crate::guard::report(&state.bus, started.id, provider, &screening) { return Err(GatewayError::denied(why)); } + // 管线第 5 步：试上游、交回答（见 `answer`）。它拿走这个请求要的一切：流式的回答等不到 + // 内容时，响应头先交给客户端，它在响应体里接着跑（见 `commit`） + let early = commit::early(&state, &req, &reading); + let abort = req.abort.clone(); + let dialect = req.dialect; + let run = answer( + state, + rt, + req, + Tail { + reading, + decision, + started, + conv, + passes: (live, pass), + }, + ending.take(), + ); + match early { + None => run.await, + Some(e) => commit::hold(run, dialect, e, abort).await, + } +} + +/// 发出开始事件之后，回答那一段要的：路由事实、决定、开始时定下的、这段对话这一轮，和两张 +/// 通行证（在服务中、这把密钥的并发）。 +struct Tail { + reading: crate::client_api::Reading, + decision: tw_engine::Decision, + started: Started, + conv: Option, + passes: (crate::live::Pass, crate::limits::Pass), +} + +/// 管线第 5 步：依次试候选上游（见 [`hop`]），把接下的那一家的回答交给客户端（见 [`relay`]）。 +/// +/// **拥有它要的一切**（不借管线的东西）：流式的回答等过一阵还没有内容时，响应头先交出去， +/// 它挪进响应体里接着跑（见 [`commit`]）。结局也在它手上：返回错误之前自己报失败，被丢掉 +/// 由结局的 Drop 报（客户端走了是取消，叫停了是手动中止）。 +async fn answer( + state: AppState, + rt: Arc, + req: Inbound, + tail: Tail, + mut ending: Option, +) -> Result { + let r = answer_with(&state, &rt, &req, tail, &mut ending).await; + if let (Err(e), Some(end)) = (&r, ending.take()) { + end.failed(e.source.into(), e.detail.clone()); + } + r +} + +async fn answer_with( + state: &AppState, + rt: &Arc, + req: &Inbound, + tail: Tail, + ending: &mut Option, +) -> Result { + let Tail { + reading, + decision, + started, + conv, + passes: (live, pass), + } = tail; // 插件的请求钩子在每一跳里跑（见 `plug`）：从客户端的原话起改 —— 内容过滤删过的话是 // 删过的那一份 —— 几跳共用原文的解析和密钥的编号。插件表跟着运行时走：**整个请求是 // 同一份**，回答钩子用的也是它 @@ -191,7 +259,7 @@ pub(super) async fn pipeline( &req.body, ); let answer = - hop::try_upstreams(&state, &rt, &req, &reading, &decision, &started, &mut hook).await?; + hop::try_upstreams(state, rt, req, &reading, &decision, &started, &mut hook).await?; // 网关估的数不是哪一家回答的：不记这段对话留在哪一家 let mut served = match answer { hop::Answer::Served(served) => *served, @@ -199,7 +267,7 @@ pub(super) async fn pipeline( let ending = ending .take() .expect("written when the start event was emitted"); - return Ok(estimated(&state, &req, started.id, body, ending)); + return Ok(estimated(state, req, started.id, body, ending)); } }; // 这一跳的账比开头那本多了号（插件往请求里写了新的值，拦截档下接着编了号）：回答 @@ -207,7 +275,7 @@ pub(super) async fn pipeline( if served.ledger.len() != started.ledger.len() && let Some(e) = ending.as_mut() { - e.redact_with(redaction(&rt, served.ledger.clone())); + e.redact_with(redaction(rt, served.ledger.clone())); } // 回答钩子：上游回了成功的回答才有。**在交出结局之前起实例**：起不来而策略是拒绝时， // 这个请求按返回的错误收场，客户端还一个字节都没收到 @@ -228,7 +296,7 @@ pub(super) async fn pipeline( request_id: started.id, attempt: served.attempt, }; - crate::plugin::reply::Chain::start(&state, &rt.plugins, bridge, &ctx).await? + crate::plugin::reply::Chain::start(state, &rt.plugins, bridge, &ctx).await? } _ => None, }; @@ -244,9 +312,9 @@ pub(super) async fn pipeline( )); } Ok(relay::respond( - &state, - &rt, - &req, + state, + rt, + req, reading.generates, &reading.facts.model, served, diff --git a/crates/tw-gateway/src/server/pipeline/commit.rs b/crates/tw-gateway/src/server/pipeline/commit.rs new file mode 100644 index 00000000..cb9c9c47 --- /dev/null +++ b/crates/tw-gateway/src/server/pipeline/commit.rs @@ -0,0 +1,182 @@ +//! 流式回答迟迟没有内容时，**响应头先交给客户端**，故障转移在响应体里接着做。 +//! +//! 第一段内容到达之前，回答是压着的（见 [`super::opening`]）：那时还能换一家上游，上游开头 +//! 报的错也还能换成一个像样的状态码。可压得太久，客户端自己的计时先到了 —— 等响应头等 +//! 不到，它就放弃了，网关的故障转移根本轮不到出手（无响应超时默认 300 秒）。 +//! +//! 所以压着的时间有上限：[`crate::OPENING_HOLD`]（15 秒，不可配置）。到点了还没有回答，就 +//! 先交出 `200` 和流的响应头，之后每隔 [`crate::KEEPALIVE_EVERY`] 发一行 SSE 注释 +//! （`: keep-alive`）—— 每一种 SSE 客户端都跳过注释，它不是内容。上游的事件照旧压着，直到 +//! 第一段真内容：在那之前这个请求照样能因为无响应超时、开头报错、流断了换下一家，被放弃的 +//! 那一家一个字节都没交出去，下一家的流在同一个 `200` 底下从头开始。 +//! +//! - **保活不算内容**：它是网关自己发的，不经过上游那条计时（见 [`super::idle`]）。 +//! - **候选都失败了**：状态码已经交出去了，按客户端的格式在流里报那个错 —— 和回答到一半 +//! 断了是同一种错误帧。 +//! - **Gemini 的客户端不发保活**：它官方的 Python SDK 把注释行当成一段 JSON 去解析，直接 +//! 报错。响应头照样先交，流里等到回答为止。 +//! - 只管客户端要 SSE 流的生成请求。整包的请求、Gemini 不带 `alt=sse` 的那种（一个 JSON +//! 数组）照旧压到回答为止。 +//! +//! 在等的时候被手动中止（见 [`crate::abort`]）：丢掉还在试的那一跳，流里报手动中止。 + +use std::future::Future; +use std::pin::Pin; +use std::time::Duration; + +use axum::body::Body; +use axum::response::Response; +use bytes::Bytes; +use futures::StreamExt; + +use super::Inbound; +use crate::error::GatewayError; +use crate::state::AppState; +use tw_dialect::ir::Dialect; + +/// 保活那一行：一行 SSE 注释 +const KEEPALIVE: &[u8] = b": keep-alive\n\n"; + +/// 先交响应头的话怎么交。 +#[derive(Debug, Clone, Copy)] +pub(super) struct Early { + /// 压着最多多久（[`crate::OPENING_HOLD`]，测试调短） + hold: Duration, + /// 隔多久发一行保活。Gemini 的客户端不发：None + keepalive: Option, +} + +/// 这个请求压不住时先不先交响应头：客户端要的是生成回答的 SSE 流才交。 +pub(super) fn early( + state: &AppState, + req: &Inbound, + reading: &crate::client_api::Reading, +) -> Option { + let streams = reading.generates && matches!(&reading.decoded, Some(Ok(d)) if d.request.stream); + if !streams { + return None; + } + let gemini = req.dialect == Dialect::Gemini; + // Gemini 的客户端不带 `alt=sse` 时收的是一个 JSON 数组，不是 SSE + let sse = !gemini + || req + .query + .as_deref() + .is_some_and(|q| q.split('&').any(|kv| kv == "alt=sse")); + sse.then_some(Early { + hold: state.opening_hold, + keepalive: (!gemini).then_some(state.keepalive_every), + }) +} + +type Running = Pin> + Send>>; + +/// 等 `run`（试上游、交回答的那一段，见 `super::answer`）最多 `early.hold`：等到了就照常交出去； +/// 等不到就先交 `200` 和流的响应头，`run` 挪进响应体里接着跑（见 [`committed`]）。`dialect` 是 +/// 客户端的格式：之后的错误按它写。 +pub(super) async fn hold( + run: impl Future> + Send + 'static, + dialect: Dialect, + early: Early, + abort: crate::abort::Switch, +) -> Result { + let mut run: Running = Box::pin(run); + let by = tokio::time::Instant::now() + early.hold; + tokio::select! { + biased; + r = &mut run => return r, + _ = tokio::time::sleep_until(by) => {} + } + Ok(committed(run, dialect, early.keepalive, abort)) +} + +/// 已经交出 `200` 的响应：先发保活，`run` 有了结论再接上 —— 一家接下了，就是它的回答； +/// 都失败了，按客户端的格式报错收尾。 +fn committed( + run: Running, + dialect: Dialect, + keepalive: Option, + abort: crate::abort::Switch, +) -> Response { + let stream = async_stream::stream! { + let mut run = run; + // 第一行保活跟着响应头一起走 + let mut next = tokio::time::Instant::now(); + let every = keepalive.unwrap_or(Duration::MAX); + let got = loop { + tokio::select! { + biased; + r = &mut run => break Some(r), + // 停在看不着开关的地方时由这里丢掉它：结局在它手上，按手动中止报（见 `Ending`） + _ = abort.wait() => break None, + _ = tokio::time::sleep_until(next), if keepalive.is_some() => { + next = tokio::time::Instant::now() + every; + yield Ok::(Bytes::from_static(KEEPALIVE)); + } + } + }; + match got { + None => { + drop(run); + yield Ok(Bytes::from(GatewayError::aborted().in_dialect(dialect).sse_frame())); + } + // 结局已经在 `answer` 里报过了 + Some(Err(e)) => yield Ok(Bytes::from(e.in_dialect(dialect).sse_frame())), + Some(Ok(resp)) if resp.status().is_success() => { + let mut body = resp.into_body().into_data_stream(); + while let Some(chunk) = body.next().await { + match chunk { + Ok(b) => yield Ok(b), + Err(e) => { + yield Err(std::io::Error::other(e)); + break; + } + } + } + } + // 最后一家回了错误、原样交出的那种：状态码交不出去了，把它说的话写成流里的错误 + Some(Ok(resp)) => { + let status = resp.status().as_u16(); + let body = axum::body::to_bytes(resp.into_body(), crate::failure::BODY_PEEK * 4) + .await + .unwrap_or_default(); + yield Ok(Bytes::from(said(dialect, status, &body))); + } + } + }; + let mut resp = Response::new(Body::from_stream(stream)); + let h = resp.headers_mut(); + h.insert( + axum::http::header::CONTENT_TYPE, + axum::http::HeaderValue::from_static("text/event-stream"), + ); + h.insert( + axum::http::header::CACHE_CONTROL, + axum::http::HeaderValue::from_static("no-cache"), + ); + resp +} + +/// 一个错误响应（`status`，正文 `body`，客户端那种格式）写成流里的一帧错误。上游的原话照说， +/// 不加前缀：那是它说的 +fn said(dialect: Dialect, status: u16, body: &[u8]) -> String { + let text = tw_dialect::convert::error_message(dialect, body) + .unwrap_or_else(|| String::from_utf8_lossy(body).trim().to_string()); + tw_dialect::convert::error_frame(dialect, status, &text) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn an_error_answer_becomes_an_error_event_in_the_clients_format() { + let body = br#"{"type":"error","error":{"type":"invalid_request_error","message":"prompt is too long"}}"#; + let f = said(Dialect::Anthropic, 400, body); + assert!(f.starts_with("event: error\n"), "{f}"); + assert!(f.contains("prompt is too long"), "{f}"); + assert!(f.contains("invalid_request_error"), "{f}"); + let f = said(Dialect::Responses, 400, br#"{"error":{"message":"bad"}}"#); + assert!(f.starts_with("event: response.failed\n"), "{f}"); + } +} diff --git a/crates/tw-gateway/src/server/pipeline/idle.rs b/crates/tw-gateway/src/server/pipeline/idle.rs index 081f320e..32c07201 100644 --- a/crates/tw-gateway/src/server/pipeline/idle.rs +++ b/crates/tw-gateway/src/server/pipeline/idle.rs @@ -8,9 +8,11 @@ //! //! 到点了怎么办，看客户端收到了什么： //! -//! - **还什么都没收到**：放弃这一家（丢掉响应或者还在等的请求，连接跟着断，上游不再接着 -//! 生成），记一次失败（和 5xx 一样算进停用的账），换下一家（见 `hop`）。客户端无感。没有 -//! 下一家了，回一个超时错误（504，见 [`crate::error::Source::Timeout`]），尝试链带着。 +//! - **还没有内容交给客户端**：放弃这一家（丢掉响应或者还在等的请求，连接跟着断，上游不再 +//! 接着生成），记一次失败（和 5xx 一样算进停用的账），换下一家（见 `hop`）。客户端无感 —— +//! 流式的回答压过 [`crate::OPENING_HOLD`] 之后响应头已经交出去了，它收到的只是保活注释（见 +//! [`super::commit`]）。没有下一家了，回一个超时错误（504，见 +//! [`crate::error::Source::Timeout`]），尝试链带着；响应头已经交出去的，在流里报这个错。 //! - **最后一家的响应已经交出去了**（最后一家不压开头）、或者**已经有内容交给了客户端**： //! 换不了 —— 换一家会把开头再发一遍。回答按客户端的格式以一条错误收尾，请求记成失败 //! （见 `relay`）。还没有内容的照样记这一家一次失败；内容出到一半才停的不记，和流在半路 @@ -25,11 +27,11 @@ //! - **尝试链上记一跳 `idle_timeout`**，带着上游可能已经收了钱的输入（见 //! [`tw_api::AttemptUsage`]）。 //! -//! **等的时候不给客户端发保活。**响应头要等选定了哪一家才发（见 `relay`），这期间客户端 -//! 那条连接上什么都没有；先发响应头再发 `: keepalive` 的话，状态码就定死成了 200 —— 之后 -//! 几家全都失败，429、5xx 和最后一家原样交出的 4xx 都给不出去，只能在流里报错，客户端按 -//! 状态码重试的逻辑就落空了；上游的响应头（请求号、额度）也带不过去。何况 Gemini 官方的 -//! Python SDK 会把注释行当成一段 JSON 去解析，直接报错。 +//! **头 15 秒不给客户端发保活**（[`crate::OPENING_HOLD`]）：响应头要等选定了哪一家才发（见 +//! `relay`），先发响应头再发 `: keep-alive` 的话，状态码就定死成了 200 —— 之后几家全都失败， +//! 429、5xx 和最后一家原样交出的 4xx 都给不出去，只能在流里报错；上游的响应头（请求号、额度） +//! 也带不过去。可一直压到无响应超时，客户端等响应头的计时又先到了，所以压着的时间有上限， +//! 过了才先交响应头、发保活（见 [`super::commit`]）。 use std::time::Duration; diff --git a/crates/tw-gateway/src/server/pipeline/opening.rs b/crates/tw-gateway/src/server/pipeline/opening.rs index ee74a0bf..d4a1be32 100644 --- a/crates/tw-gateway/src/server/pipeline/opening.rs +++ b/crates/tw-gateway/src/server/pipeline/opening.rs @@ -13,7 +13,9 @@ //! 等到无响应超时那一刻（`failover.idle_timeout_secs`，从请求发出去算起）还没有内容是 //! [`Opening::Slow`]：调用方放弃这一家、换下一家。**最后一家不在这里等**（见 `hop`）：没有 //! 下一家可换，压着它只会让客户端晚一点看到同样的东西。开头压得太多（[`HOLD_LIMIT`]）也不 -//! 再压，交出去由回程照常计时。 +//! 再压，交出去由回程照常计时。压着的这段时间里客户端等的是响应头：流式的回答过了 +//! [`crate::OPENING_HOLD`] 先拿到 `200` 和保活注释，上游的事件照旧压在这里（见 +//! [`super::commit`]）。 //! //! 整包的回答（不是流）也一样等：响应头先到、正文迟迟不来的上游，等到点了换下一家（见 //! [`first_bytes`]）。 diff --git a/crates/tw-gateway/src/state.rs b/crates/tw-gateway/src/state.rs index 43deda09..d2921cc7 100644 --- a/crates/tw-gateway/src/state.rs +++ b/crates/tw-gateway/src/state.rs @@ -253,6 +253,10 @@ pub struct AppState { /// `failover.idle_timeout_secs` 的一秒有多长（见 `server::pipeline::idle`）。**测试会把它 /// 调短**，否则一条无响应超时的测试至少要干等三十秒 pub idle_tick: std::time::Duration, + /// 流式回答最多压多久不给响应头（[`crate::OPENING_HOLD`]）。**测试会把它调短** + pub opening_hold: std::time::Duration, + /// 响应头先交出去之后隔多久发一行保活（[`crate::KEEPALIVE_EVERY`]）。**测试会把它调短** + pub keepalive_every: std::time::Duration, } impl AppState { @@ -322,6 +326,8 @@ impl AppState { plugin_pool: Arc::new(crate::plugin::pool::Pool::default_size()), aborts: Default::default(), idle_tick: std::time::Duration::from_secs(1), + opening_hold: crate::OPENING_HOLD, + keepalive_every: crate::KEEPALIVE_EVERY, }; // 手写的清单马上可用；向上游问是后台的事，不挡启动 state.publish_catalog(); diff --git a/crates/tw-gateway/tests/idle_timeout.rs b/crates/tw-gateway/tests/idle_timeout.rs index d10215e7..29ebdac8 100644 --- a/crates/tw-gateway/tests/idle_timeout.rs +++ b/crates/tw-gateway/tests/idle_timeout.rs @@ -4,6 +4,8 @@ //! 它记下自己的响应被丢掉了没有 —— 网关放弃它、中止它时，连接要真的断开，上游才会停下。 //! //! 超时写最短的 30 秒，测试把一秒调成三十分之一秒（`AppState::idle_tick`）：等的是 1 秒。 +//! 流式回答压着不给响应头的时限（`AppState::opening_hold`）默认调到 10 秒，比这里的场景都长 —— +//! 响应头先交出去的那几条（见 `tw_gateway::OPENING_HOLD`）另调成 300 毫秒。 use std::net::SocketAddr; use std::sync::Arc; @@ -31,6 +33,7 @@ struct Script { /// 挂着时发的心跳 beat: &'static str, content_type: &'static str, + status: u16, } impl Default for Script { @@ -39,8 +42,9 @@ impl Default for Script { header_delay_ms: 0, steps: Vec::new(), hang: false, - beat: ": keep-alive\n\n", + beat: ": ping\n\n", content_type: "text/event-stream", + status: 200, } } } @@ -91,6 +95,7 @@ async fn upstream(script: Script) -> Upstream { x }); axum::response::Response::builder() + .status(script.status) .header("content-type", script.content_type) .body(axum::body::Body::from_stream(body)) .unwrap() @@ -157,6 +162,10 @@ fn provider(name: &str, up: &Upstream, protocol: Protocol) -> Provider { } } +/// 响应头先交出去的那几条测试里，流式回答最多压多久、隔多久发一行保活 +const HOLD: Duration = Duration::from_millis(300); +const KEEPALIVE_EVERY: Duration = Duration::from_millis(150); + /// 起网关，交回数据面的状态。`slot_wait_secs` 是等空位的期限 async fn gateway_with( providers: Vec, @@ -165,6 +174,19 @@ async fn gateway_with( SocketAddr, tokio::sync::broadcast::Receiver, tw_gateway::AppState, +) { + gateway_held(providers, slot_wait_secs, Duration::from_secs(10)).await +} + +/// 起网关，流式回答最多压 `hold` 不给响应头 +async fn gateway_held( + providers: Vec, + slot_wait_secs: u64, + hold: Duration, +) -> ( + SocketAddr, + tokio::sync::broadcast::Receiver, + tw_gateway::AppState, ) { let cfg = Config { version: 1, @@ -183,6 +205,8 @@ async fn gateway_with( }; let mut state = tw_gateway::AppState::new(cfg).unwrap(); state.idle_tick = WINDOW / 30; + state.opening_hold = hold; + state.keepalive_every = KEEPALIVE_EVERY; let rx = state.bus.subscribe(); let addr = tw_gateway::serve(state.clone(), ([127, 0, 0, 1], 0).into()) .await @@ -229,6 +253,27 @@ async fn send(gw: SocketAddr, path: &str, body: &Value) -> (u16, Option, (status, source, resp.text().await.unwrap()) } +/// 发出去，交回状态码、响应头到的那一刻（从发出算起）和整个正文 +async fn timed(gw: SocketAddr, path: &str, body: &Value) -> (u16, Duration, String) { + let t = Instant::now(); + let resp = reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{gw}{path}")) + .header("content-type", "application/json") + .header("x-api-key", "tw-k") + .header("x-goog-api-key", "tw-k") + .header("authorization", "Bearer tw-k") + .body(body.to_string()) + .send() + .await + .unwrap(); + let headers_at = t.elapsed(); + let status = resp.status().as_u16(); + (status, headers_at, resp.text().await.unwrap()) +} + async fn post(gw: SocketAddr, path: &str, body: &Value) -> (u16, String) { let (status, _, text) = send(gw, path, body).await; (status, text) @@ -1005,3 +1050,266 @@ async fn aborting_a_session_stops_every_request_of_it_and_nothing_else() { let (_, text) = c.await.unwrap(); assert!(text.contains("event: error"), "{text}"); } + +// ───────────────────────────────────────────── 压不住了：响应头先交出去 + +/// 一行保活 +const KA: &str = ": keep-alive\n\n"; + +#[tokio::test] +async fn past_the_hold_the_headers_go_out_with_keepalives_and_failover_carries_on() { + let slow = upstream(stalled()).await; + let good = upstream(prompt("hello")).await; + let (gw, mut rx, _) = gateway_held( + vec![ + provider("slow", &slow, Protocol::Anthropic), + provider("good", &good, Protocol::Anthropic), + ], + 30, + HOLD, + ) + .await; + let (status, headers_at, text) = timed(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "{text}"); + assert!( + headers_at >= HOLD && headers_at < WINDOW, + "响应头在压满之后、超时之前交出：{headers_at:?}" + ); + // 先是几行保活（压满之后一行，之后每隔一阵一行），然后是下一家从头开始的回答 + let answer = text.find("event: message_start").expect("下一家的回答"); + assert!(text[..answer].matches(KA).count() >= 3, "{text}"); + assert!( + text[..answer].trim_start().starts_with(": keep-alive"), + "{text}" + ); + assert_eq!(text.matches("event: message_start").count(), 1, "{text}"); + assert!( + !text.contains("event: ping"), + "放弃的那一家一个字节都不交出去：{text}" + ); + assert!(text[answer..].contains("hello"), "{text}"); + assert!( + !text[answer..].contains("keep-alive"), + "回答开始之后不再插保活：{text}" + ); + assert!(eventually(&slow.dropped).await); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [("slow", IdleTimeout), ("good", Served)] + ); + assert_eq!(outcome(&mut rx).await, Ok(200)); +} + +#[tokio::test] +async fn an_error_before_content_after_the_headers_still_fails_over() { + // 开了流，压满之后才报过载：响应头已经交出去了，照样换下一家，错误不到客户端 + let overloaded = upstream(Script { + steps: vec![ + (0, MESSAGE_START), + ( + 600, + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n", + ), + ], + ..Default::default() + }) + .await; + let good = upstream(prompt("hello")).await; + let (gw, mut rx, _) = gateway_held( + vec![ + provider("busy", &overloaded, Protocol::Anthropic), + provider("good", &good, Protocol::Anthropic), + ], + 30, + HOLD, + ) + .await; + let (status, headers_at, text) = timed(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "{text}"); + assert!(headers_at < Duration::from_millis(550), "{headers_at:?}"); + assert!(text.starts_with(": keep-alive"), "{text}"); + assert!(!text.contains("Overloaded"), "{text}"); + assert!(text.contains("hello"), "{text}"); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!(outcomes(&attempts), [("busy", Error), ("good", Served)]); + assert_eq!( + attempts[0].error.as_ref().map(|m| m.code.as_str()), + Some("gw.upstream.stream_opening_error") + ); +} + +#[tokio::test] +async fn when_every_upstream_stays_silent_after_the_headers_the_stream_ends_with_an_error() { + let a = upstream(stalled()).await; + let b = upstream(Script { + header_delay_ms: 60_000, + ..prompt("never") + }) + .await; + let (gw, mut rx, _) = gateway_held( + vec![ + provider("a", &a, Protocol::Anthropic), + provider("b", &b, Protocol::Anthropic), + ], + 30, + HOLD, + ) + .await; + let (status, headers_at, text) = timed(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "状态码已经交出去了：{text}"); + assert!(headers_at < WINDOW, "{headers_at:?}"); + let error = text.find("event: error").expect("流里的错误"); + assert!(text[..error].matches(KA).count() >= 3, "{text}"); + assert!(!text.contains("event: ping"), "{text}"); + assert!(text[error..].contains("timeout_error"), "{text}"); + assert!(text[error..].contains("[ThinkWatch]"), "{text}"); + assert!(text[error..].contains("tried: a → b"), "{text}"); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [("a", IdleTimeout), ("b", IdleTimeout)] + ); + assert_eq!( + outcome(&mut rx).await, + Err(( + tw_api::FailureSource::Upstream, + "gw.upstream.idle_timeout".into() + )) + ); +} + +#[tokio::test] +async fn a_refusal_from_the_last_upstream_after_the_headers_is_told_in_the_stream() { + // 第一家不出声、压满之后放弃；最后一家回 400：状态码交不出去了，它说的话写成流里的错误 + let slow = upstream(stalled()).await; + let refusing = upstream(Script { + status: 400, + content_type: "application/json", + steps: vec![( + 0, + "{\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"prompt is too long\"}}", + )], + ..Default::default() + }) + .await; + let (gw, mut rx, _) = gateway_held( + vec![ + provider("slow", &slow, Protocol::Anthropic), + provider("strict", &refusing, Protocol::Anthropic), + ], + 30, + HOLD, + ) + .await; + let (status, _, text) = timed(gw, "/v1/messages", &messages(true)).await; + assert_eq!(status, 200, "{text}"); + assert!(text.starts_with(": keep-alive"), "{text}"); + let error = text.find("event: error").expect("流里的错误"); + assert!(text[error..].contains("prompt is too long"), "{text}"); + assert!(text[error..].contains("invalid_request_error"), "{text}"); + assert_eq!( + outcome(&mut rx).await, + Err(( + tw_api::FailureSource::Upstream, + "gw.upstream.status_message".into() + )) + ); +} + +#[tokio::test] +async fn aborting_after_the_headers_went_out_ends_the_stream_with_the_abort() { + let slow = upstream(stalled()).await; + let other = upstream(prompt("other")).await; + let (gw, mut rx, state) = gateway_held( + vec![ + provider("slow", &slow, Protocol::Anthropic), + provider("other", &other, Protocol::Anthropic), + ], + 30, + HOLD, + ) + .await; + let asking = tokio::spawn(async move { timed(gw, "/v1/messages", &messages(true)).await }); + let (id, _) = started(&mut rx).await; + tokio::time::sleep(HOLD + Duration::from_millis(250)).await; + state.aborts.request(id).expect("还在跑"); + let (status, _, text) = asking.await.unwrap(); + assert_eq!(status, 200, "{text}"); + assert!(text.starts_with(": keep-alive"), "{text}"); + let error = text.find("event: error").expect("流里的错误"); + assert!(text[error..].contains("aborted by the user"), "{text}"); + assert!(eventually(&slow.dropped).await, "和上游的连接要断开"); + assert_eq!(other.hits.load(Ordering::SeqCst), 0, "不换下一家"); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!(outcomes(&attempts), [("slow", Aborted)]); + assert_eq!( + outcome(&mut rx).await, + Err((tw_api::FailureSource::Aborted, tw_api::ABORTED.into())) + ); +} + +#[tokio::test] +async fn gemini_clients_get_the_headers_early_but_no_comments() { + // Gemini 官方的 Python SDK 把注释行当成 JSON 去解析：响应头照样先交，流里不插注释 + let silent = upstream(Script { + steps: vec![( + 0, + "data: {\"usageMetadata\":{\"promptTokenCount\":10}}\r\n\r\n", + )], + hang: true, + ..Default::default() + }) + .await; + let good = upstream(Script { + steps: vec![( + 0, + "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"hello\"}]},\"finishReason\":\"STOP\"}]}\r\n\r\n", + )], + ..Default::default() + }) + .await; + let (gw, mut rx, _) = gateway_held( + vec![ + provider("silent", &silent, Protocol::Gemini), + provider("good", &good, Protocol::Gemini), + ], + 30, + HOLD, + ) + .await; + let path = "/v1beta/models/gemini-3-pro:streamGenerateContent?alt=sse"; + let body = json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}); + let (status, headers_at, text) = timed(gw, path, &body).await; + assert_eq!(status, 200, "{text}"); + assert!(headers_at < WINDOW, "{headers_at:?}"); + assert!( + !text.lines().any(|l| l.starts_with(':')), + "没有注释行：{text}" + ); + assert!(text.trim_start().starts_with("data: "), "{text}"); + assert!(text.contains("hello"), "{text}"); + assert!(!text.contains("promptTokenCount\":10}}"), "{text}"); + let (_, attempts) = estimate_and_attempts(&mut rx).await; + assert_eq!( + outcomes(&attempts), + [("silent", IdleTimeout), ("good", Served)] + ); +} + +#[tokio::test] +async fn a_whole_answer_is_never_answered_early() { + // 整包的请求照旧压到回答为止：压满之后也不先交响应头，候选用完了是 504 + let only = upstream(Script { + header_delay_ms: 60_000, + steps: vec![(0, whole("never"))], + content_type: "application/json", + ..Default::default() + }) + .await; + let (gw, _, _) = + gateway_held(vec![provider("only", &only, Protocol::Anthropic)], 30, HOLD).await; + let (status, headers_at, text) = timed(gw, "/v1/messages", &messages(false)).await; + assert_eq!(status, 504, "{text}"); + assert!(headers_at >= WINDOW, "{headers_at:?}"); +} diff --git a/docs/config.md b/docs/config.md index 2709bf29..00ce27c7 100644 --- a/docs/config.md +++ b/docs/config.md @@ -1029,11 +1029,19 @@ to the complete answer. - When no content has reached the client yet, the upstream counts as failed (towards `failures_to_pause`, like a 5xx), the attempt appears with the outcome `idle_timeout`, and the request moves to the next candidate. Until - then a streamed answer is held back from the client, so the next upstream + then an upstream's answer is held back from the client, so the next upstream starts it afresh. With no candidate left, the client gets a timeout error (504) in its own format. The last candidate's stream is passed on as it arrives, so once its response has started, a timeout there ends it with an error event instead. +- A streamed answer is held for at most 15 seconds. If no content has + arrived by then, the client receives `200` and the streaming headers, so + that its own wait for headers does not run out, followed by an SSE comment + (`: keep-alive`) every 15 seconds; Gemini clients get no comments. The + upstream's events stay held until its first content, and failover continues + as before under the same `200`: the next upstream's stream starts cleanly, + and when no candidate is left, the stream ends with an error event in the + client's format instead of a 504. The comments do not count as content. - When content has already reached the client, the request cannot move on without repeating it: the answer ends with an error event in the client's format, and the request is recorded as failed. diff --git a/docs/config.zh-CN.md b/docs/config.zh-CN.md index 9747da32..d2efbc5d 100644 --- a/docs/config.zh-CN.md +++ b/docs/config.zh-CN.md @@ -792,9 +792,13 @@ security: 只保持连接、不出内容的上游同样会超时。整包（非流式）的回答从发出算到整份回来。 - 还没有内容交给客户端时，这家上游记一次失败（和 5xx 一样计入 `failures_to_pause`）， - 尝试链上这一跳的结果是 `idle_timeout`，请求换到下一个候选。在此之前流式回答不交给 + 尝试链上这一跳的结果是 `idle_timeout`，请求换到下一个候选。在此之前上游的回答不交给 客户端，下一家从头开始回答。没有候选了，客户端收到它自己格式的超时错误（504）。 最后一个候选的流是边收边交给客户端的，它的响应开始之后再超时，回答以一条错误事件收尾。 +- 流式回答最多暂存 15 秒。到时还没有内容，客户端先收到 `200` 和流式响应头，以免它自己 + 等响应头的时限先到，之后每 15 秒收到一行 SSE 注释（`: keep-alive`）；Gemini 的客户端 + 不发注释。上游的事件照旧暂存到第一段内容，故障转移在同一个 `200` 之下照常进行：下一家 + 的流从头开始；没有候选了，流以客户端格式的错误事件结束，不再是 504。注释不算内容。 - 已经有内容交给客户端时，换一家会重复已经发出的内容，所以不换：回答按客户端的格式 以一条错误事件收尾，请求记为失败。 diff --git a/scripts/smoke.sh b/scripts/smoke.sh index 69cff6cb..1bb43647 100755 --- a/scripts/smoke.sh +++ b/scripts/smoke.sh @@ -116,6 +116,22 @@ class H(http.server.BaseHTTPRequestHandler): "stopReason": "end_turn", "usage": {"inputTokens": 30, "outputTokens": 2, "totalTokens": 32}}).encode()) saw = "yes" if b"sk-ant-api03-SMOKEKEY" in body else "no" + if b"QUIETSTART" in body: + # 开了流、报了输入，17 秒不出内容，然后答完：比网关压着响应头的 15 秒长 + self.send_response(200); self.send_header('content-type','text/event-stream') + self.send_header('connection','close'); self.end_headers() + self.wfile.write(b'event: message_start\ndata: {"type":"message_start",' + b'"message":{"usage":{"input_tokens":12,"output_tokens":1}}}\n\n') + self.wfile.flush() + time.sleep(17) + self.wfile.write(b'event: content_block_start\ndata: {"type":"content_block_start",' + b'"index":0,"content_block":{"type":"text","text":""}}\n\n' + b'event: content_block_delta\ndata: {"type":"content_block_delta",' + b'"index":0,"delta":{"type":"text_delta","text":"quiet-done"}}\n\n' + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n') + self.wfile.flush() + return if b"SLOWSTREAM" in body: # 先吐开头（输入用量在 message_start 里），然后长时间「思考」—— 客户端 # 会在这期间走掉。不给 content-length，读到连接关闭为止。 @@ -199,6 +215,10 @@ listen: clients: - name: claude-code key: tw-smoketestkey0123456789 + # 走两家候选的那把：流开头压不住时响应头先交出去，只在有下一家可换时才压 + - name: quiet + key: tw-smokequietkey0123456789 + route: 两家 providers: - name: relay base_url: http://127.0.0.1:{upport} @@ -225,6 +245,9 @@ groups: - name: 只走中转 type: fallback providers: [relay] + - name: 两家 + type: fallback + providers: [relay, official] routes: - name: 默认 rules: @@ -234,6 +257,10 @@ routes: to: bedrock - name: 冒烟：这条必须走 relay，不许转移 to: 只走中转 + - name: 两家 + rules: + - name: 先中转再官方 + to: 两家 security: redact: mode: enforce @@ -400,6 +427,26 @@ print("ok" if good else json.dumps(r, ensure_ascii=False, sort_keys=True))' 2>/d || bad "客户端中途走掉的请求没有按取消落库" "$GOT" fi +# ---------------------------------------------------------------- 流开头压不住 +step "流开头压不住：响应头先交出去" +# 两家候选，第一家开了流、17 秒不出内容。网关最多压 15 秒（tw_gateway::OPENING_HOLD）： +# 到点先交 200 和流的响应头、发一行保活注释，内容到了接在后面交出去。单元测试里的时限是 +# 调短的，这里看真二进制上那个 15 秒和 hyper 真的把响应头先送出去 +T=$(curl -s -N -m 40 -o "$TMP/quiet.out" -w '%{http_code} %{time_starttransfer} %{time_total}' \ + -XPOST "http://127.0.0.1:$PORT/v1/messages" -H 'x-api-key: tw-smokequietkey0123456789' \ + -H 'content-type: application/json' \ + -d '{"model":"claude-sonnet-4-5","max_tokens":64,"stream":true,"messages":[{"role":"user","content":"QUIETSTART"}]}' 2>/dev/null) +C=$(python3 -c 'import sys +code, first, total = sys.argv[1].split() +print("ok" if code == "200" and 14.0 <= float(first) < 16.9 and float(total) >= 16.9 else sys.argv[1])' "$T") +[ "$C" = "ok" ] && ok "压了 15 秒先交出响应头（${T}）" || bad "响应头不是在 15 秒时先交出去的" "$C" +if head -c 13 "$TMP/quiet.out" | grep -q '^: keep-alive' && grep -q 'quiet-done' "$TMP/quiet.out" \ + && [ "$(grep -c 'event: message_start' "$TMP/quiet.out")" = 1 ]; then + ok "先是保活注释，然后是那一家的回答，开头只有一份" +else + bad "响应体不对" "$(head -c 400 "$TMP/quiet.out")" +fi + # ---------------------------------------------------------------- Bedrock step "Bedrock" # 转成 Converse、访问密钥签名（假 Bedrock 核对签名覆盖的请求体哈希）、