diff --git a/CHANGELOG.md b/CHANGELOG.md index 16d18a1e..c23be273 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,12 +17,16 @@ for the providers are recognised by value, whatever their format. The console shows the rules' names and descriptions as thinkwatch-core writes them, in both languages. A value recognised by mistake can be made an exception without switching its rule off, and long conversations are -no longer scanned from the start every turn. The thinkwatch-core crates +no longer scanned from the start every turn. A provider can set the +slowest generation its streams may keep up, and a stream that falls +below it is cut so the client can try again. The thinkwatch-core crates move from v0.83.1 to v0.86.0. -No new setting or column. The redaction policy gains one optional field -(`exceptions`), with endpoints of its own under -`/api/admin/security/redact/exceptions`. +No new setting. Two new provider columns +(`providers.min_speed_tokens_per_sec`, `providers.min_speed_window_secs`, +added on startup) and one new metric (`gateway_provider_too_slow_total`). +The redaction policy gains one optional field (`exceptions`), with +endpoints of its own under `/api/admin/security/redact/exceptions`. ### Read before upgrading @@ -89,6 +93,24 @@ No new setting or column. The redaction policy gains one optional field it and scans only what is new. What is found is the same as without it. The cache is rebuilt whenever the rules, the providers' keys or the exceptions change. +- **Stop a stream that generates too slowly** (the provider's + "Stop when generation is slower than … tokens/s for … s" setting, + `min_speed_tokens_per_sec` 1–1000 and `min_speed_window_secs` 5–120, + default 20, on the provider endpoints; blank or `null` is off, the + default). From the first token on, a stream is judged once a second: + when the text it streamed over the last window (answer, reasoning and + tool arguments, four ASCII characters or one other character to a + token; Gemini's running count when higher) comes to fewer tokens per + second than the minimum, the stream ends with a 503 error in the + client's format (`overloaded_error` for Anthropic, `response.failed` + for Responses, `UNAVAILABLE` for Gemini), the upstream connection is + closed, and the route is charged a failure. Pauses of 5 s or more are + left out of the window, and a window that is more than half pauses is + not judged. The session's affinity to that route is dropped, the + answer so far is billed as for any stream cut short, and the log row + carries the error `UpstreamTooSlow` and the speed at the cut + (`tokens_per_sec`). Whole answers and Responses WebSocket turns are + not measured. ### Changed diff --git a/crates/common/src/dto/mod.rs b/crates/common/src/dto/mod.rs index 6e897009..e72c64a3 100644 --- a/crates/common/src/dto/mod.rs +++ b/crates/common/src/dto/mod.rs @@ -257,6 +257,15 @@ pub struct CreateProviderRequest { #[serde(default)] pub headers: Vec, pub config: Option, + /// Cut a stream that generates fewer tokens per second than this + /// (1–1000), averaged over `min_speed_window_secs`. Absent or `null` + /// ⇒ off. + #[serde(default)] + pub min_speed_tokens_per_sec: Option, + /// The window the speed is averaged over, in seconds (5–120). + /// Absent ⇒ 20. + #[serde(default)] + pub min_speed_window_secs: Option, } #[derive(Debug, Serialize, Deserialize, Clone, utoipa::ToSchema)] diff --git a/crates/common/src/models/provider.rs b/crates/common/src/models/provider.rs index cc02a0e7..7b087269 100644 --- a/crates/common/src/models/provider.rs +++ b/crates/common/src/models/provider.rs @@ -15,6 +15,12 @@ pub struct Provider { pub config_json: serde_json::Value, pub created_at: DateTime, pub deleted_at: Option>, + /// The slowest generation a stream from this upstream may keep up, + /// in tokens per second (1–1000), averaged over the last + /// `min_speed_window_secs` seconds (5–120). `None` ⇒ off. See + /// `think_watch_gateway::proxy::min_speed`. + pub min_speed_tokens_per_sec: Option, + pub min_speed_window_secs: i32, } #[derive(Debug, Clone, Serialize, Deserialize, FromRow)] diff --git a/crates/gateway/src/error.rs b/crates/gateway/src/error.rs index fa771ff5..c376f283 100644 --- a/crates/gateway/src/error.rs +++ b/crates/gateway/src/error.rs @@ -103,6 +103,21 @@ pub enum GatewayError { conversation." )] SealedHistoryUnavailable { provider: String, model: String }, + /// A stream generated slower than its provider's minimum (see + /// `proxy::min_speed`): `rate` tokens per second over the last `secs` + /// seconds, below `min`. The stream was cut; maps to 503, the + /// upstream failing, so the client can try again and the route's + /// breaker counts it. + #[error( + "Upstream `{provider}` generated {rate} tokens per second over the last {secs} seconds, \ + below the minimum of {min}." + )] + UpstreamTooSlow { + provider: String, + rate: u32, + min: u32, + secs: u32, + }, } impl GatewayError { @@ -126,6 +141,7 @@ impl GatewayError { GatewayError::PolicyBlocked(_) => 403, GatewayError::SealedHistoryRefused { status, .. } => i64::from(*status), GatewayError::SealedHistoryUnavailable { .. } => 409, + GatewayError::UpstreamTooSlow { .. } => 503, } } @@ -146,6 +162,7 @@ impl GatewayError { GatewayError::PolicyBlocked(_) => "PolicyBlocked", GatewayError::SealedHistoryRefused { .. } => "SealedHistoryRefused", GatewayError::SealedHistoryUnavailable { .. } => "SealedHistoryUnavailable", + GatewayError::UpstreamTooSlow { .. } => "UpstreamTooSlow", } } @@ -256,4 +273,25 @@ mod tests { assert_eq!(down.retry_after_secs(), Some(30)); assert_eq!(down.should_retry(), None); } + + #[test] + fn a_stream_too_slow_is_the_upstream_failing_and_says_how_slow() { + let e = GatewayError::UpstreamTooSlow { + provider: "local-vllm".into(), + rate: 3, + min: 15, + secs: 20, + }; + assert_eq!(e.status_code(), 503); + assert_eq!(e.error_tag(), "UpstreamTooSlow"); + assert!(crate::proxy::upstream_failed( + e.error_tag(), + e.status_code() + )); + assert_eq!( + e.to_string(), + "Upstream `local-vllm` generated 3 tokens per second over the last 20 seconds, \ + below the minimum of 15." + ); + } } diff --git a/crates/gateway/src/lifecycle/mod.rs b/crates/gateway/src/lifecycle/mod.rs index 6e857c46..542ad26f 100644 --- a/crates/gateway/src/lifecycle/mod.rs +++ b/crates/gateway/src/lifecycle/mod.rs @@ -34,6 +34,7 @@ use tw_dialect::ir::Dialect; use crate::guards::Guards; use crate::proxy::generate::{Wire, priced, tokens}; +use crate::proxy::min_speed::{CHECK_EVERY, MinSpeed, SpeedWatch}; use crate::proxy::retry::{KEEP_ALIVE, KEEP_ALIVE_FRAME}; use crate::proxy::shaper::{StreamShaper, rewrite_model}; use crate::proxy::{ @@ -141,6 +142,9 @@ pub(crate) struct ChatPickedRoute { /// The route's capacity caps: its token cap counts the answer's /// tokens in `record_usage`. pub caps: crate::route_caps::RouteCaps, + /// The slowest generation the provider's streams may keep up (see + /// `proxy::min_speed`). + pub min_speed: Option, } /// Everything the post-invoke hooks read. Built once before the upstream @@ -188,6 +192,9 @@ pub(crate) struct Sent { /// The attempt under way sent a different reasoning effort than the /// caller asked (see `proxy::effort`). pub effort: Option, + /// The stream was cut for generating too slowly, at this many tokens + /// per second (see `proxy::min_speed`). + pub too_slow: Option, } /// Build the streaming pump: forward the upstream's bytes to the caller — @@ -251,6 +258,7 @@ pub(crate) fn build_chat_pump( guards: &Guards, request: &ChatRequestSnapshot, provider: &str, + min_speed: Option, ) -> (axum::response::Response, StreamTail) { let readers = Arc::new(Mutex::new(Readers::default())); let readers_for_tail = Arc::clone(&readers); @@ -288,6 +296,7 @@ pub(crate) fn build_chat_pump( last_sent: std::time::Instant::now(), waiting: false, effort: None, + too_slow: None, })); // Gemini clients get no keep-alive comments: the google-genai SDK reads // every SSE line as JSON. @@ -318,7 +327,7 @@ pub(crate) fn build_chat_pump( } let sent_at = std::time::Instant::now(); if let Ok(mut s) = sent.lock() { - *s = Sent { retries: resend.budget.retries(), last_sent: sent_at, waiting: false, effort: None }; + *s = Sent { retries: resend.budget.retries(), last_sent: sent_at, waiting: false, effort: None, too_slow: None }; } // What an attempt that failed had read is not this one's: its // usage is not billed, its error not the request's. @@ -392,16 +401,47 @@ pub(crate) fn build_chat_pump( // While a failure could still be sent again, the upstream's // opening is held back until its first content (see `Opening`). let mut hold = resend.budget.open().then(|| Opening::new(wire.dialect)); + // The provider's slowest generation, judged once a second from + // the first token on (see `proxy::min_speed`). + let mut speed = min_speed.map(|m| SpeedWatch::new(m, wire.dialect)); + let mut judge = tokio::time::interval(CHECK_EVERY); + judge.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); let mut source = upstream.bytes_stream(); loop { let next = tokio::select! { - item = source.next() => Some(item), - _ = tokio::time::sleep_until(quiet_since + KEEP_ALIVE), if hold.is_some() && keep_alive => None, + item = source.next() => Pulled::Item(item), + _ = tokio::time::sleep_until(quiet_since + KEEP_ALIVE), if hold.is_some() && keep_alive => Pulled::KeepAlive, + _ = judge.tick(), if speed.as_ref().is_some_and(SpeedWatch::started) => Pulled::Judge, }; - let Some(item) = next else { - quiet_since = tokio::time::Instant::now(); - yield Ok(Bytes::from_static(KEEP_ALIVE_FRAME)); - continue; + let item = match next { + Pulled::Item(item) => item, + Pulled::KeepAlive => { + quiet_since = tokio::time::Instant::now(); + yield Ok(Bytes::from_static(KEEP_ALIVE_FRAME)); + continue; + } + Pulled::Judge => { + let Some((rate, min)) = speed.as_mut().and_then(|s| { + Some((s.too_slow(std::time::Instant::now())?, s.min())) + }) else { + continue; + }; + // Too slow: the client is told in its own format, and + // returning drops the upstream's connection. + let err = too_slow(&provider, rate, min); + let mut out = shaper.process(&error_frame(client, err.status_code(), &err.to_string())); + out.extend(shaper.finish()); + if let Ok(mut s) = sent.lock() { + s.too_slow = Some(rate); + } + tell(&mut done_tx, StreamOutcome::UpstreamError { + error_type: err.error_tag().to_string(), + message: err.to_string(), + status_code: err.status_code(), + }); + yield Ok(Bytes::from(out)); + return; + } }; let Some(item) = item else { // The upstream's stream ended. One whose opening is still @@ -448,6 +488,9 @@ pub(crate) fn build_chat_pump( }; match item { Ok(chunk) => { + if let Some(s) = speed.as_mut() { + s.feed(&chunk, std::time::Instant::now()); + } let (sealed, said) = match readers.lock() { Ok(mut r) => { r.feed(&chunk); @@ -691,6 +734,7 @@ pub(crate) fn build_chat_pump( last_sent: started_at, waiting: false, effort: None, + too_slow: None, }); let invoked = Invoked { client_ip: identity.ip_address.clone(), @@ -1108,6 +1152,7 @@ fn resending( last_sent: std::time::Instant::now(), waiting: true, effort: None, + too_slow: None, }; } tracing::info!( @@ -1224,6 +1269,37 @@ impl Opening { } } +/// What the pump waits on next: a piece of the upstream's stream (or its +/// end), a keep-alive due, or a second to judge the stream's speed. +enum Pulled { + Item(T), + KeepAlive, + Judge, +} + +/// A stream cut for generating at `rate` tokens per second, below `min`: +/// the error, and the log line saying so. +fn too_slow(provider: &str, rate: u32, min: MinSpeed) -> GatewayError { + tracing::warn!( + provider, + tokens_per_sec = rate, + min_tokens_per_sec = min.tokens_per_sec, + window_secs = min.window_secs, + "upstream generated too slowly, cutting the stream" + ); + metrics::counter!( + "gateway_provider_too_slow_total", + "provider" => crate::metrics_labels::normalize_provider_label(provider), + ) + .increment(1); + GatewayError::UpstreamTooSlow { + provider: provider.to_string(), + rate, + min: min.tokens_per_sec, + secs: min.window_secs, + } +} + /// Tell the tail how the stream ended. Only the first outcome counts. fn tell(done_tx: &mut Option>, outcome: StreamOutcome) { if let Some(tx) = done_tx.take() { @@ -1533,12 +1609,15 @@ impl Surface for ChatCompletionSurface { deps.request.request_started_at.elapsed().as_millis() as i64, logged_status, with_usage_detail( - effort_detail( - retried_detail( - pinned_detail(error_detail, deps.route.pinned), - deps.route.sel_record.retries, + speed_detail( + effort_detail( + retried_detail( + pinned_detail(error_detail, deps.route.pinned), + deps.route.sel_record.retries, + ), + deps.route.sel_record.effort, ), - deps.route.sel_record.effort, + deps.route.sel_record.too_slow, ), &usage, usage_estimated(&invoked.view), @@ -1623,6 +1702,24 @@ fn effort_detail( Some(serde_json::Value::Object(d)) } +/// The audit detail, with `"tokens_per_sec"` when the stream was cut for +/// generating too slowly (see `proxy::min_speed`): its speed over the +/// window at the cut. A failed request otherwise records no speed. +fn speed_detail( + detail: Option, + too_slow: Option, +) -> Option { + let Some(rate) = too_slow else { + return detail; + }; + let mut d = match detail { + Some(serde_json::Value::Object(d)) => d, + _ => serde_json::Map::new(), + }; + d.insert("tokens_per_sec".into(), rate.into()); + Some(serde_json::Value::Object(d)) +} + /// The usage a captured view bills — shared by `record_usage` and /// `emit_audit` so the budget and the audit row can never disagree. fn extract_usage(view: &CapturedView) -> tw_dialect::usage::Usage { @@ -1852,6 +1949,7 @@ mod tests { last_sent: std::time::Instant::now(), waiting: false, effort: None, + too_slow: None, }); resending( "up", diff --git a/crates/gateway/src/proxy/generate.rs b/crates/gateway/src/proxy/generate.rs index 0d351761..eb85bc3d 100644 --- a/crates/gateway/src/proxy/generate.rs +++ b/crates/gateway/src/proxy/generate.rs @@ -1233,6 +1233,7 @@ async fn run( upstream_model: route.upstream_model.clone(), sel_record, caps: crate::route_caps::RouteCaps::of(route), + min_speed: route.upstream.min_speed, }, cache_enabled: surface.caches, }; diff --git a/crates/gateway/src/proxy/min_speed.rs b/crates/gateway/src/proxy/min_speed.rs new file mode 100644 index 00000000..3cad0a21 --- /dev/null +++ b/crates/gateway/src/proxy/min_speed.rs @@ -0,0 +1,447 @@ +//! Cut a stream an upstream generates too slowly. +//! +//! A provider can set the slowest generation it may keep a stream at +//! (`providers.min_speed_tokens_per_sec`, averaged over +//! `providers.min_speed_window_secs`). Some upstreams — an overloaded +//! relay, a small self-hosted model — keep a stream open and dribble a +//! token every second or two: the client neither gets its answer nor an +//! error it could retry on. With a minimum set, such a stream is cut and +//! the client told why (503), so it can try again — elsewhere, since the +//! route has a failure against it and its affinity is dropped. +//! +//! # The rule +//! +//! - Only streams. A whole answer and a WebSocket turn are not measured. +//! - Counting starts at the first token. Until `window` has passed since +//! then, nothing is judged. +//! - Once a second, the speed over the last `window` is +//! `tokens in it / (window − long pauses in it)`; below the minimum, the +//! stream is cut. +//! - A **long pause** is [`LONG_PAUSE`] or more without content. It and +//! its (zero) tokens are left out of the window: an upstream that went +//! quiet is not generating slowly, it is not generating at all, and a +//! model thinking behind a hidden reasoning block looks the same. When +//! the pauses leave less than half the window, that second is not +//! judged — a few seconds of data are not enough to conclude on. +//! +//! # Counting tokens +//! +//! The text the upstream streams — answer text, reasoning text and tool +//! arguments, read in its own format with `tw_dialect::convert::Reader` — +//! estimated as the desktop gateway estimates its live speed: an ASCII +//! byte is a quarter of a token, any other character one token. Bytes +//! alone (four to a token, as [`crate::usage_estimate`] bills an answer +//! nobody reported) would count a Chinese character as three quarters of +//! a token where tokenizers give it one or more, and cut Chinese answers +//! that are not slow. Gemini reports its running output count in each +//! chunk; the larger of that and the estimate counts. +//! +//! What is billed is not this: a stream cut here is billed like any other +//! stream cut short (see `usage_estimate::complete`). + +use std::collections::VecDeque; +use std::time::{Duration, Instant}; + +use tw_dialect::convert::Reader; +use tw_dialect::ir::{Delta, Dialect, Event}; + +/// A pause this long or longer is left out of the window. +pub const LONG_PAUSE: Duration = Duration::from_secs(5); + +/// How often a stream's speed is judged. +pub const CHECK_EVERY: Duration = Duration::from_secs(1); + +/// The ranges the columns allow, checked again where the row is read. +pub const TOKENS_PER_SEC: std::ops::RangeInclusive = 1..=1000; +pub const WINDOW_SECS: std::ops::RangeInclusive = 5..=120; + +/// The slowest generation a provider's streams may keep up. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MinSpeed { + pub tokens_per_sec: u32, + pub window_secs: u32, +} + +impl MinSpeed { + /// From the provider row: `None` when off, or when a value is out of + /// range (the columns' CHECKs keep that from happening). + pub fn from_row(tokens_per_sec: Option, window_secs: i32) -> Option { + let tokens_per_sec = u32::try_from(tokens_per_sec?).ok()?; + let window_secs = u32::try_from(window_secs).ok()?; + (TOKENS_PER_SEC.contains(&tokens_per_sec) && WINDOW_SECS.contains(&window_secs)).then_some( + Self { + tokens_per_sec, + window_secs, + }, + ) + } + + fn window(&self) -> Duration { + Duration::from_secs(u64::from(self.window_secs)) + } +} + +/// Text tokens, estimated as the desktop gateway does: four ASCII bytes +/// to a token, one token for any other character. +#[derive(Debug, Default, Clone, Copy)] +struct Text { + ascii: u64, + other: u64, +} + +impl Text { + fn add(&mut self, s: &str) { + for &b in s.as_bytes() { + if b.is_ascii() { + self.ascii += 1; + } else if b & 0xC0 != 0x80 { + // A character's first byte; continuation bytes are 10xxxxxx. + self.other += 1; + } + } + } + + fn tokens(&self) -> u64 { + self.ascii.div_ceil(4) + self.other + } +} + +/// One stream's speed, fed the upstream's bytes as they arrive. +pub struct SpeedWatch { + min: MinSpeed, + reader: Reader, + text: Text, + /// Gemini's running output count. + reported: u64, + gemini: bool, + /// Tokens already put down in `arrivals`. + counted: u64, + /// When the first token arrived. + first: Option, + /// The upstream said the answer is done (its stop reason): what it + /// sends or keeps open after that is not generation. + done: bool, + /// Each piece of content: when, and how many tokens it brought. The + /// window's, plus the last one before it (the pause that runs into + /// the window starts there). + arrivals: VecDeque<(Instant, u64)>, +} + +impl SpeedWatch { + pub fn new(min: MinSpeed, upstream: Dialect) -> Self { + Self { + min, + reader: Reader::new(upstream), + text: Text::default(), + reported: 0, + gemini: upstream == Dialect::Gemini, + counted: 0, + first: None, + done: false, + arrivals: VecDeque::new(), + } + } + + pub fn min(&self) -> MinSpeed { + self.min + } + + /// Whether there is anything to judge: the first token is in, and the + /// upstream has not said the answer is done. + pub fn started(&self) -> bool { + self.first.is_some() && !self.done + } + + /// A piece of the upstream's stream (SSE, in its own format) arrived + /// at `now`. + pub fn feed(&mut self, chunk: &[u8], now: Instant) { + let mut content = false; + for e in self.reader.feed(chunk) { + match e { + Event::Delta { + delta: Delta::Text(t) | Delta::Thinking(t) | Delta::ToolInput(t), + .. + } if !t.is_empty() => { + self.text.add(&t); + content = true; + } + Event::Usage(u) if self.gemini => self.reported = self.reported.max(u.output), + Event::Stop(_) => self.done = true, + _ => {} + } + } + if content { + self.arrived(now); + } + } + + fn arrived(&mut self, now: Instant) { + let total = self.text.tokens().max(self.reported); + let tokens = total.saturating_sub(self.counted); + self.counted = self.counted.max(total); + self.first.get_or_insert(now); + self.arrivals.push_back((now, tokens)); + } + + /// The speed over the last window, in tokens per second, when it is + /// below the minimum at `now` — `None` when it is not, or when it is + /// too early to judge (see the module docs). + pub fn too_slow(&mut self, now: Instant) -> Option { + if self.done { + return None; + } + let rate = self.rate(now)?; + (rate < f64::from(self.min.tokens_per_sec)).then(|| rate.floor() as u32) + } + + /// The speed over the last window at `now`, when it can be judged. + fn rate(&mut self, now: Instant) -> Option { + let first = self.first?; + let window = self.min.window(); + if now.saturating_duration_since(first) < window { + return None; + } + let start = now - window; + // Keep one arrival at or before the window's start: a pause that + // runs into the window is measured from it. + while self.arrivals.len() >= 2 && self.arrivals[1].0 <= start { + self.arrivals.pop_front(); + } + let tokens: u64 = self + .arrivals + .iter() + .filter(|(at, _)| *at > start) + .map(|(_, n)| n) + .sum(); + // The gaps between arrivals, and from the last one to now. + let ends = self + .arrivals + .iter() + .map(|(at, _)| *at) + .skip(1) + .chain(std::iter::once(now)); + let paused: Duration = self + .arrivals + .iter() + .map(|(at, _)| *at) + .zip(ends) + .filter(|(from, to)| to.saturating_duration_since(*from) >= LONG_PAUSE) + .map(|(from, to)| to.min(now).saturating_duration_since(from.max(start))) + .sum(); + let measured = window.saturating_sub(paused); + if measured < window / 2 { + return None; + } + Some(tokens as f64 / measured.as_secs_f64()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn secs(n: f64) -> Duration { + Duration::from_secs_f64(n) + } + + fn min(tokens_per_sec: u32, window_secs: u32) -> MinSpeed { + MinSpeed { + tokens_per_sec, + window_secs, + } + } + + /// A Chat chunk carrying `text`. + fn chat(text: &str) -> Vec { + format!( + "data: {}\n\n", + json!({"id": "c", "object": "chat.completion.chunk", "model": "m", + "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": null}]}) + ) + .into_bytes() + } + + /// `tokens` ASCII tokens of text. + fn words(tokens: usize) -> String { + "abcd".repeat(tokens) + } + + #[test] + fn the_row_turns_it_on_only_with_a_speed_in_range() { + assert_eq!(MinSpeed::from_row(None, 20), None); + assert_eq!(MinSpeed::from_row(Some(15), 20), Some(min(15, 20))); + assert_eq!(MinSpeed::from_row(Some(0), 20), None); + assert_eq!(MinSpeed::from_row(Some(1001), 20), None); + assert_eq!(MinSpeed::from_row(Some(15), 4), None); + assert_eq!(MinSpeed::from_row(Some(15), 121), None); + } + + #[test] + fn ascii_counts_by_the_quarter_and_other_characters_whole() { + let mut t = Text::default(); + t.add("abcdefgh"); + assert_eq!(t.tokens(), 2); + t.add("你好"); + assert_eq!(t.tokens(), 4); + t.add("é"); + assert_eq!(t.tokens(), 5); + } + + #[test] + fn nothing_is_judged_before_the_first_token_or_within_a_window_of_it() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(10, 20), Dialect::Chat); + // The opening carries no content. + w.feed( + b"data: {\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"}}]}\n\n", + t, + ); + assert!(!w.started()); + assert_eq!(w.too_slow(t + secs(60.0)), None); + w.feed(&chat("a"), t + secs(1.0)); + assert!(w.started()); + // One token in 19.9 seconds is slow, but not a window yet. + for i in 0..19 { + w.feed(&chat("a"), t + secs(2.0 + i as f64 * 0.5)); + } + assert_eq!(w.too_slow(t + secs(20.9)), None); + assert!(w.too_slow(t + secs(21.0)).is_some()); + } + + #[test] + fn a_stream_under_the_minimum_is_cut_with_the_rate_it_had() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(10, 10), Dialect::Chat); + // Four tokens a second: two every half second. + for i in 0..20 { + w.feed(&chat(&words(2)), t + secs(i as f64 * 0.5)); + } + assert_eq!(w.too_slow(t + secs(9.9)), None, "not a window yet"); + for i in 20..24 { + w.feed(&chat(&words(2)), t + secs(i as f64 * 0.5)); + } + // At 11.5 s the window holds (1.5, 11.5]: 20 arrivals, 40 tokens. + assert_eq!(w.too_slow(t + secs(11.5)), Some(4)); + } + + #[test] + fn a_stream_at_the_minimum_runs_on() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(10, 10), Dialect::Chat); + for i in 0..60 { + let at = t + secs(i as f64 * 0.5); + w.feed(&chat(&words(5)), at); + assert_eq!(w.too_slow(at), None, "at {i}"); + } + } + + #[test] + fn a_long_pause_is_left_out_of_the_window() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(15, 20), Dialect::Chat); + // 20 tokens a second for 12 s, then 8 s of silence: over the whole + // 20 s window that is 12 a second, without the pause 20. + for i in 0..=24 { + w.feed(&chat(&words(10)), t + secs(i as f64 * 0.5)); + } + assert_eq!(w.too_slow(t + secs(20.0)), None); + assert_eq!(w.rate(t + secs(20.0)).map(f64::round), Some(20.0)); + } + + #[test] + fn a_shorter_silence_slows_the_stream() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(10, 10), Dialect::Chat); + // 16 tokens a second until 10 s, then silent. + for i in 0..=20 { + w.feed(&chat(&words(8)), t + secs(i as f64 * 0.5)); + } + // 4.9 s of silence is not a long pause: (4.9, 14.9] holds the + // arrivals at 5.0..=10.0, 88 tokens over 10 s. + assert_eq!(w.too_slow(t + secs(14.9)), Some(8)); + // At 5 s it is one: (5, 15] is 80 tokens over the 5 s before it. + assert_eq!(w.too_slow(t + secs(15.0)), None); + assert_eq!(w.rate(t + secs(15.0)).map(f64::round), Some(16.0)); + } + + /// An upstream can keep its stream open a while after the answer's + /// end, or send its usage later: none of that is generating slowly. + #[test] + fn nothing_is_judged_once_the_answer_is_done() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(10, 5), Dialect::Chat); + for i in 0..10 { + w.feed(&chat(&words(10)), t + secs(i as f64 * 0.5)); + } + let stop = json!({"id": "c", "object": "chat.completion.chunk", "model": "m", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}); + w.feed(format!("data: {stop}\n\n").as_bytes(), t + secs(5.0)); + assert!(!w.started()); + assert_eq!(w.too_slow(t + secs(9.0)), None); + } + + #[test] + fn a_window_mostly_paused_is_not_judged() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(10, 20), Dialect::Chat); + w.feed(&chat("a"), t); + w.feed(&chat("a"), t + secs(1.0)); + // From 1 s on the stream is silent: the pause covers 11 of the 20 + // seconds at 21 s, and only 9 are left. + assert_eq!(w.too_slow(t + secs(21.0)), None); + // A trickle after the pause: still mostly paused. + w.feed(&chat("a"), t + secs(22.0)); + assert_eq!(w.too_slow(t + secs(22.5)), None); + } + + #[test] + fn a_pause_running_into_the_window_is_cut_at_its_start() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(1, 10), Dialect::Chat); + // Steady until 8 s, silent until 14 s, steady after. + for i in 0..=16 { + w.feed(&chat(&words(4)), t + secs(i as f64 * 0.5)); + } + for i in 0..=12 { + w.feed(&chat(&words(4)), t + secs(14.0 + i as f64 * 0.5)); + } + // Window (10, 20]: 4 s of the pause 8 → 14 are in it, 6 s are + // measured, and the arrivals at 14..=20 bring 13 × 4 tokens. + assert_eq!( + w.rate(t + secs(20.0)).map(|r| (r * 100.0).round()), + Some(867.0) + ); + } + + #[test] + fn reasoning_and_tool_arguments_count_and_signatures_do_not() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(1, 5), Dialect::Anthropic); + let event = + |v: serde_json::Value| format!("event: {}\ndata: {v}\n\n", v["type"].as_str().unwrap()); + let stream = [ + event(json!({"type": "message_start", "message": {"id": "m", "type": "message", + "role": "assistant", "model": "m", "content": [], "usage": {"input_tokens": 1, "output_tokens": 1}}})), + event(json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}})), + event(json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "abcd"}})), + event(json!({"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "x".repeat(400)}})), + event(json!({"type": "content_block_stop", "index": 0})), + event(json!({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "id": "t", "name": "ls", "input": {}}})), + event(json!({"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "{\"a\":1}"}})), + ] + .concat(); + w.feed(stream.as_bytes(), t); + assert_eq!(w.counted, 3); + } + + #[test] + fn geminis_running_count_counts_when_it_is_higher() { + let t = Instant::now(); + let mut w = SpeedWatch::new(min(1, 5), Dialect::Gemini); + let chunk = json!({"candidates": [{"content": {"role": "model", "parts": [{"text": "ab"}]}}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 7}}); + w.feed(format!("data: {chunk}\n\n").as_bytes(), t); + assert_eq!(w.counted, 7); + } +} diff --git a/crates/gateway/src/proxy/mod.rs b/crates/gateway/src/proxy/mod.rs index a72563e1..3bd787b5 100644 --- a/crates/gateway/src/proxy/mod.rs +++ b/crates/gateway/src/proxy/mod.rs @@ -28,6 +28,7 @@ pub(crate) mod generate; mod headers; mod identity; mod log_ctx; +pub mod min_speed; mod models; mod pipeline; mod protocol_relearn; @@ -42,7 +43,9 @@ pub mod transport; pub(crate) use accounting::post_flight_account; pub(crate) use body_capture::prepare_body_capture; pub(crate) use log_ctx::emit_gateway_log_with_extra; -pub(crate) use routing::{Resend, SelectionRecord, fails as upstream_failed, finalize_health}; +pub(crate) use routing::{ + Resend, SelectionRecord, fails as upstream_failed, finalize_health, forget_affinity, +}; // pub re-exports — `server::app` mounts these as route handlers. pub use early_cancel::{EarlyCancel, EarlyCancelSlot}; diff --git a/crates/gateway/src/proxy/pipeline.rs b/crates/gateway/src/proxy/pipeline.rs index 77f9731d..e0deb24e 100644 --- a/crates/gateway/src/proxy/pipeline.rs +++ b/crates/gateway/src/proxy/pipeline.rs @@ -130,9 +130,24 @@ pub(super) fn launch_stream_pump( &deps.guards, &deps.request, &deps.route.provider_name, + deps.route.min_speed, ); tokio::spawn(async move { let (invoked, sent) = tail.await; + // Cut for generating too slowly: the session's next request is not + // sent back to the same upstream by its affinity, and the row says + // how fast it was going. + if let Some(rate) = sent.too_slow { + super::forget_affinity( + &deps.state.redis, + deps.request.identity.user_id.as_deref(), + &deps.request.mapped_model, + deps.route.provider_id, + deps.route.sel_record.picked_route_id, + ) + .await; + deps.route.sel_record.too_slow = Some(rate); + } // A stream that went out again: the route's time is the attempt // that ended it, and the row says how many retries there were. if sent.retries > 0 { diff --git a/crates/gateway/src/proxy/routing.rs b/crates/gateway/src/proxy/routing.rs index 1cb9639a..1d3db53c 100644 --- a/crates/gateway/src/proxy/routing.rs +++ b/crates/gateway/src/proxy/routing.rs @@ -91,6 +91,32 @@ pub(super) async fn set_affinity( .await; } +/// Drop a session's affinity to the route a stream was just cut on for +/// generating too slowly (see `super::min_speed`): the client's retry goes +/// where the strategy sends it, not back to the same upstream. Only an +/// affinity that still points at this route or its provider is dropped — +/// another request of the session may have moved it since. +pub(crate) async fn forget_affinity( + redis: &fred::clients::Client, + user_id: Option<&str>, + model: &str, + provider_id: Uuid, + route_id: Uuid, +) { + use fred::interfaces::KeysInterface; + let Some(uid) = user_id else { return }; + for (mode, id) in [ + (AffinityMode::Provider, provider_id), + (AffinityMode::Route, route_id), + ] { + let key = format!("affinity:{}:{uid}:{model}", mode.as_str()); + let held: Option = redis.get(&key).await.ok().flatten(); + if held.as_deref() == Some(id.to_string().as_str()) { + let _: Result = redis.del(&key).await; + } + } +} + /// Resolve `(strategy, affinity_mode, affinity_ttl)` for a model: /// per-model override falls through to gateway-wide defaults. async fn resolve_routing_config( @@ -305,6 +331,10 @@ pub(crate) struct SelectionRecord { /// The request's last attempt was already counted against the route: /// a stream whose client left while it waited to go out again. pub counted: bool, + /// The stream was cut for generating too slowly (see + /// `super::min_speed`), at this many tokens per second over its + /// window. Said on the log row as `"tokens_per_sec"`. + pub too_slow: Option, } /// Record health for the picked route once the request finishes. @@ -550,6 +580,7 @@ pub(super) async fn select_route_with_failover<'a>( retries: budget.retries(), effort, counted: false, + too_slow: None, }, )); } @@ -662,6 +693,7 @@ pub(super) async fn select_route_for_stream<'a>( retries: 0, effort: None, counted: false, + too_slow: None, }, )); } diff --git a/crates/gateway/src/proxy/transport.rs b/crates/gateway/src/proxy/transport.rs index 6da9652f..84431279 100644 --- a/crates/gateway/src/proxy/transport.rs +++ b/crates/gateway/src/proxy/transport.rs @@ -68,6 +68,9 @@ pub struct Upstream { /// The models this upstream refused the cache breakpoints a /// conversion adds for (see `proxy::cache_marks`). pub(crate) cache_marks: crate::proxy::cache_marks::Refused, + /// The slowest generation a stream from this upstream may keep up + /// (see `proxy::min_speed`). `None` ⇒ not measured. + pub min_speed: Option, } impl Upstream { @@ -79,9 +82,16 @@ impl Upstream { shape, label: label.to_string(), cache_marks: Default::default(), + min_speed: None, } } + /// The same upstream, with streams cut below `min_speed`. + pub fn with_min_speed(mut self, min_speed: Option) -> Self { + self.min_speed = min_speed; + self + } + /// Is this the vendor's own endpoint rather than a relay? /// /// The conversion layer needs to know: official endpoints are diff --git a/crates/server/src/gateway_adapters.rs b/crates/server/src/gateway_adapters.rs index b05d4be6..db4b3116 100644 --- a/crates/server/src/gateway_adapters.rs +++ b/crates/server/src/gateway_adapters.rs @@ -24,6 +24,8 @@ pub(crate) struct ProviderMaterials { /// Bedrock API key (an `Authorization` header) has no keys and is /// never signed. Only read by the Bedrock adapter. pub(crate) aws: Credential, + /// The slowest generation the provider's streams may keep up. + pub(crate) min_speed: Option, } impl ProviderMaterials { @@ -86,6 +88,10 @@ impl ProviderMaterials { .and_then(|v| v.as_str()) .map(|s| s.to_string()), aws, + min_speed: think_watch_gateway::proxy::min_speed::MinSpeed::from_row( + provider.min_speed_tokens_per_sec, + provider.min_speed_window_secs, + ), } } @@ -134,12 +140,9 @@ pub(crate) fn build_upstream(m: &ProviderMaterials) -> Arc { } _ => Shape::Standard, }; - Arc::new(Upstream::new( - &m.base_url, - m.headers.clone(), - shape, - &m.name, - )) + Arc::new( + Upstream::new(&m.base_url, m.headers.clone(), shape, &m.name).with_min_speed(m.min_speed), + ) } /// The API version an Azure deployment is addressed with when the @@ -164,6 +167,8 @@ mod tests { config_json: config, created_at: chrono::Utc::now(), deleted_at: None, + min_speed_tokens_per_sec: None, + min_speed_window_secs: 20, } } diff --git a/crates/server/src/handlers/providers.rs b/crates/server/src/handlers/providers.rs index f54a59be..72ad74e3 100644 --- a/crates/server/src/handlers/providers.rs +++ b/crates/server/src/handlers/providers.rs @@ -280,6 +280,12 @@ pub async fn create_provider( encrypt_aws_secret_in_config(&mut config, &state.config.encryption_key)?; config["headers"] = encrypt_headers_for_storage(&req.headers, &state.config.encryption_key)?; + let min_speed = min_speed( + req.min_speed_tokens_per_sec, + req.min_speed_window_secs + .unwrap_or(i64::from(DEFAULT_MIN_SPEED_WINDOW_SECS)), + )?; + let mut provider = repo::insert( &state.db, &req.name, @@ -287,6 +293,7 @@ pub async fn create_provider( &req.provider_type, &req.base_url, &config, + min_speed, ) .await?; @@ -311,6 +318,52 @@ pub struct UpdateProviderRequest { pub base_url: Option, /// Unified request headers (auth + custom + identity templates). pub headers: Option>, + /// PATCH-clearable slowest generation, in tokens per second. Absent + /// = unchanged, JSON `null` = off, number (1 to 1000) = cut a stream + /// that averages fewer over `min_speed_window_secs`. + #[serde( + default, + deserialize_with = "crate::handlers::serde_util::deserialize_some" + )] + #[schema(value_type = Option)] + pub min_speed_tokens_per_sec: Option>, + /// The window that speed is averaged over, in seconds (5 to 120). + /// Absent = unchanged. + pub min_speed_window_secs: Option, +} + +/// A stream's window when none is given, in seconds. +pub(crate) const DEFAULT_MIN_SPEED_WINDOW_SECS: i32 = 20; + +/// A provider's slowest generation as stored: `tokens_per_sec` 1 to 1000 +/// or none (off), `window_secs` 5 to 120 — the columns' CHECKs, as a +/// useful 400. +pub(crate) fn min_speed( + tokens_per_sec: Option, + window_secs: i64, +) -> Result { + let tokens_per_sec = tokens_per_sec + .map(|n| { + i32::try_from(n) + .ok() + .filter(|n| (1..=1000).contains(n)) + .ok_or_else(|| { + AppError::BadRequest( + "min_speed_tokens_per_sec must be between 1 and 1000".into(), + ) + }) + }) + .transpose()?; + let window_secs = i32::try_from(window_secs) + .ok() + .filter(|n| (5..=120).contains(n)) + .ok_or_else(|| { + AppError::BadRequest("min_speed_window_secs must be between 5 and 120".into()) + })?; + Ok(repo::MinSpeed { + tokens_per_sec, + window_secs, + }) } #[utoipa::path( @@ -372,7 +425,24 @@ pub async fn update_provider( config }; - let mut updated = repo::update(&state.db, id, display_name, base_url, &config_json).await?; + let min_speed = min_speed( + match req.min_speed_tokens_per_sec { + None => existing.min_speed_tokens_per_sec.map(i64::from), + Some(set) => set, + }, + req.min_speed_window_secs + .unwrap_or(i64::from(existing.min_speed_window_secs)), + )?; + + let mut updated = repo::update( + &state.db, + id, + display_name, + base_url, + &config_json, + min_speed, + ) + .await?; // A new base URL or credential can mean an entirely different // upstream, so every dialect we learned for this provider's routes @@ -620,6 +690,8 @@ fn test_materials( headers, api_version: stored.and_then(|s| s.api_version.clone()), aws, + // A connection test lists models; it streams nothing. + min_speed: None, } } @@ -788,6 +860,65 @@ fn probe_url( mod tests { use super::*; + #[test] + fn a_minimum_speed_is_in_the_columns_ranges_or_off() { + let ok = |t, w| min_speed(t, w).unwrap(); + assert_eq!( + ok(None, 20), + repo::MinSpeed { + tokens_per_sec: None, + window_secs: 20 + } + ); + assert_eq!( + ok(Some(1), 5), + repo::MinSpeed { + tokens_per_sec: Some(1), + window_secs: 5 + } + ); + assert_eq!(ok(Some(1000), 120).tokens_per_sec, Some(1000)); + for (t, w) in [ + (Some(0), 20), + (Some(1001), 20), + (Some(-1), 20), + (Some(i64::MAX), 20), + (None, 4), + (None, 121), + (None, 0), + ] { + assert!(min_speed(t, w).is_err(), "{t:?} / {w}"); + } + assert!( + min_speed(Some(0), 20) + .unwrap_err() + .to_string() + .contains("min_speed_tokens_per_sec") + ); + assert!( + min_speed(None, 4) + .unwrap_err() + .to_string() + .contains("min_speed_window_secs") + ); + } + + #[test] + fn a_patch_tells_an_absent_speed_from_a_cleared_one() { + let absent: UpdateProviderRequest = serde_json::from_value(serde_json::json!({})).unwrap(); + assert_eq!(absent.min_speed_tokens_per_sec, None); + assert_eq!(absent.min_speed_window_secs, None); + let cleared: UpdateProviderRequest = + serde_json::from_value(serde_json::json!({"min_speed_tokens_per_sec": null})).unwrap(); + assert_eq!(cleared.min_speed_tokens_per_sec, Some(None)); + let set: UpdateProviderRequest = serde_json::from_value( + serde_json::json!({"min_speed_tokens_per_sec": 15, "min_speed_window_secs": 30}), + ) + .unwrap(); + assert_eq!(set.min_speed_tokens_per_sec, Some(Some(15))); + assert_eq!(set.min_speed_window_secs, Some(30)); + } + fn test_hex_key() -> &'static str { // 32 bytes of zeroes → valid 64-char hex key for AES-256. "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" @@ -903,6 +1034,7 @@ mod tests { headers: vec![("Authorization".into(), "Bearer ABSK-saved".into())], api_version: None, aws: keys("AKIA-saved", "secret-saved"), + min_speed: None, } } diff --git a/crates/server/src/services/provider_repository.rs b/crates/server/src/services/provider_repository.rs index 03a07e96..a45cc254 100644 --- a/crates/server/src/services/provider_repository.rs +++ b/crates/server/src/services/provider_repository.rs @@ -37,6 +37,15 @@ pub async fn name_of(pool: &PgPool, id: Uuid) -> Result, AppError ) } +/// The slowest generation a provider's streams may keep up (see +/// `think_watch_gateway::proxy::min_speed`): `tokens_per_sec` `None` is +/// off. Checked against the columns' ranges by the handler. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MinSpeed { + pub tokens_per_sec: Option, + pub window_secs: i32, +} + pub async fn insert( pool: &PgPool, name: &str, @@ -44,16 +53,20 @@ pub async fn insert( provider_type: &str, base_url: &str, config_json: &serde_json::Value, + min_speed: MinSpeed, ) -> Result { Ok(sqlx::query_as::<_, Provider>( - r#"INSERT INTO providers (name, display_name, provider_type, base_url, config_json) - VALUES ($1, $2, $3, $4, $5) RETURNING *"#, + r#"INSERT INTO providers (name, display_name, provider_type, base_url, config_json, + min_speed_tokens_per_sec, min_speed_window_secs) + VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING *"#, ) .bind(name) .bind(display_name) .bind(provider_type) .bind(base_url) .bind(config_json) + .bind(min_speed.tokens_per_sec) + .bind(min_speed.window_secs) .fetch_one(pool) .await?) } @@ -64,15 +77,19 @@ pub async fn update( display_name: &str, base_url: &str, config_json: &serde_json::Value, + min_speed: MinSpeed, ) -> Result { Ok(sqlx::query_as::<_, Provider>( - r#"UPDATE providers SET display_name = $2, base_url = $3, config_json = $4 + r#"UPDATE providers SET display_name = $2, base_url = $3, config_json = $4, + min_speed_tokens_per_sec = $5, min_speed_window_secs = $6 WHERE id = $1 RETURNING *"#, ) .bind(id) .bind(display_name) .bind(base_url) .bind(config_json) + .bind(min_speed.tokens_per_sec) + .bind(min_speed.window_secs) .fetch_one(pool) .await?) } diff --git a/crates/test-support/src/mock_provider.rs b/crates/test-support/src/mock_provider.rs index 4f63a9e4..d6ce8a16 100644 --- a/crates/test-support/src/mock_provider.rs +++ b/crates/test-support/src/mock_provider.rs @@ -192,6 +192,61 @@ pub async fn sse_upstream_lingering( format!("http://{addr}") } +/// An upstream that answers every `POST {path}` with an event stream +/// sent piece by piece: each `(wait, piece)` goes out `wait` after the one +/// before it. Returns its base URL and a flag that is set when the +/// gateway hangs up before the last piece went out. +/// +/// wiremock sends a body in one piece, so this is a server of its own. +pub async fn sse_upstream_paced( + path: &str, + pieces: Vec<(std::time::Duration, String)>, +) -> (String, std::sync::Arc) { + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + /// Set on drop unless the stream ran to its end. + struct HungUp { + flag: Arc, + finished: bool, + } + impl Drop for HungUp { + fn drop(&mut self) { + if !self.finished { + self.flag.store(true, Ordering::SeqCst); + } + } + } + + let hung_up = Arc::new(AtomicBool::new(false)); + let flag = Arc::clone(&hung_up); + let answer = axum::routing::post(move || { + let pieces = pieces.clone(); + let flag = Arc::clone(&flag); + async move { + let body = async_stream::stream! { + let mut guard = HungUp { flag, finished: false }; + for (wait, piece) in pieces { + tokio::time::sleep(wait).await; + yield Ok::<_, std::convert::Infallible>(bytes::Bytes::from(piece)); + } + guard.finished = true; + }; + axum::response::Response::builder() + .header("content-type", "text/event-stream") + .body(axum::body::Body::from_stream(body)) + .unwrap() + } + }); + let app = axum::Router::new().route(path, answer); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + (format!("http://{addr}"), hung_up) +} + fn openai_sse_chunks(model: &str) -> Vec { let chunk = |delta: Value| { format!( diff --git a/crates/test-support/tests/gateway_min_speed.rs b/crates/test-support/tests/gateway_min_speed.rs new file mode 100644 index 00000000..e8995aa8 --- /dev/null +++ b/crates/test-support/tests/gateway_min_speed.rs @@ -0,0 +1,389 @@ +//! Cutting a stream its upstream generates too slowly. +//! +//! A provider with `min_speed_tokens_per_sec` set has its streams judged +//! once a second from the first token on: below the minimum over the last +//! `min_speed_window_secs` (pauses of 5 s or more left out), the stream is +//! cut with a 503 error in the client's format, the upstream's connection +//! dropped and the route charged a failure. The upstreams here are paced +//! servers (`sse_upstream_paced`), so each test takes as long as its +//! stream: the windows are the shortest allowed. + +use std::sync::atomic::Ordering; +use std::time::{Duration, Instant}; + +use serde_json::Value; +use think_watch_test_support::mock_provider::sse_upstream_paced; +use think_watch_test_support::prelude::*; + +/// A model routed to `uri` under a provider that cuts streams below +/// `tokens_per_sec` over `window_secs`; a key that may call it, and the +/// route's id. +async fn seed( + app: &TestApp, + uri: &str, + provider_type: &str, + model: &str, + tokens_per_sec: i32, + window_secs: i32, +) -> (String, uuid::Uuid) { + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let provider = + fixtures::create_provider(&app.db, &unique_name("slow"), provider_type, uri, None) + .await + .unwrap(); + sqlx::query( + "UPDATE providers SET min_speed_tokens_per_sec = $2, min_speed_window_secs = $3 \ + WHERE id = $1", + ) + .bind(provider.id) + .bind(tokens_per_sec) + .bind(window_secs) + .execute(&app.db) + .await + .unwrap(); + fixtures::create_model_and_route(&app.db, provider.id, model) + .await + .unwrap(); + let route_id: uuid::Uuid = + sqlx::query_scalar("SELECT id FROM model_routes WHERE provider_id = $1") + .bind(provider.id) + .fetch_one(&app.db) + .await + .unwrap(); + app.rebuild_gateway_router().await; + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("slow-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap() + .plaintext; + (key, route_id) +} + +fn chat_chunk(model: &str, delta: Value, finish: Option<&str>) -> String { + format!( + "data: {}\n\n", + json!({ + "id": "chatcmpl-slow", + "object": "chat.completion.chunk", + "created": 1_700_000_000_i64, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + }) + ) +} + +/// A Chat stream: the opening, then `tokens` ASCII tokens of text in each +/// `(wait, tokens)` piece, then the end. +fn chat_stream(model: &str, pieces: &[(f64, usize)]) -> Vec<(Duration, String)> { + let mut out = vec![( + Duration::ZERO, + chat_chunk(model, json!({"role": "assistant"}), None), + )]; + for (wait, tokens) in pieces { + out.push(( + Duration::from_secs_f64(*wait), + chat_chunk(model, json!({"content": "abcd".repeat(*tokens)}), None), + )); + } + out.push(( + Duration::ZERO, + [ + chat_chunk(model, json!({}), Some("stop")), + "data: [DONE]\n\n".to_string(), + ] + .concat(), + )); + out +} + +fn chat(model: &str) -> Value { + json!({"model": model, "stream": true, "messages": [{"role": "user", "content": "hi"}]}) +} + +/// The route's failures in the breaker's window, once the stream's tail +/// has recorded it. +async fn errors_on(app: &TestApp, route_id: uuid::Uuid) -> u32 { + use think_watch_gateway::health::{CircuitBreakerConfig, HealthTracker}; + let tracker = HealthTracker::new(app.state.redis.clone()); + let cfg = CircuitBreakerConfig::load(&app.state.dynamic_config).await; + for _ in 0..50 { + let h = tracker.snapshot(route_id, cfg).await; + if h.total > 0 { + return h.errors; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + panic!("the stream's outcome was never recorded against the route"); +} + +/// One token a second against a minimum of five: cut about five seconds +/// after the first token, with a Chat error the client can read, the +/// upstream hung up on, and a failure on the route. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_stream_slower_than_the_minimum_is_cut() { + let app = TestApp::spawn().await; + let model = "slow-chat"; + let pieces: Vec<(f64, usize)> = (0..30) + .map(|i| (if i == 0 { 0.0 } else { 1.0 }, 1)) + .collect(); + let (uri, hung_up) = + sse_upstream_paced("/v1/chat/completions", chat_stream(model, &pieces)).await; + let (key, route_id) = seed(&app, &uri, "openai", model, 5, 5).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let started = Instant::now(); + let resp = gw.post("/v1/chat/completions", chat(model)).await.unwrap(); + let took = started.elapsed(); + resp.assert_ok(); + let text = resp.text(); + assert!( + text.contains("abcd"), + "the answer so far reaches the client: {text}" + ); + assert!( + text.contains("tokens per second over the last 5 seconds, below the minimum of 5."), + "{text}" + ); + assert!(!text.contains("[DONE]"), "{text}"); + assert!( + took >= Duration::from_secs(5) && took < Duration::from_secs(12), + "cut after {took:?}" + ); + // The error frame is Chat's, with a 503's class. + let error: Value = text + .lines() + .filter_map(|l| l.strip_prefix("data: ")) + .filter_map(|d| serde_json::from_str::(d).ok()) + .find(|v| v.get("error").is_some()) + .expect("an error frame"); + assert!( + error["error"]["message"] + .as_str() + .is_some_and(|m| m.starts_with("Upstream `")), + "{error}" + ); + + assert_eq!(errors_on(&app, route_id).await, 1); + for _ in 0..20 { + if hung_up.load(Ordering::SeqCst) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!( + hung_up.load(Ordering::SeqCst), + "the upstream was not hung up on" + ); +} + +/// An Anthropic client is told with Anthropic's error for a 503, +/// `overloaded_error`, which its SDK retries. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn an_anthropic_client_is_told_the_upstream_is_overloaded() { + let app = TestApp::spawn().await; + let model = "slow-claude"; + let event = |v: Value| format!("event: {}\ndata: {v}\n\n", v["type"].as_str().unwrap()); + let mut pieces = vec![( + Duration::ZERO, + [ + event(json!({"type": "message_start", "message": { + "id": "msg_1", "type": "message", "role": "assistant", "model": model, + "content": [], "stop_reason": null, + "usage": {"input_tokens": 12, "output_tokens": 1} + }})), + event(json!({"type": "content_block_start", "index": 0, + "content_block": {"type": "text", "text": ""}})), + ] + .concat(), + )]; + for _ in 0..30 { + pieces.push(( + Duration::from_secs(1), + event(json!({"type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": "abcd"}})), + )); + } + let (uri, _) = sse_upstream_paced("/v1/messages", pieces).await; + let (key, _) = seed(&app, &uri, "anthropic", model, 5, 5).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw + .post( + "/v1/messages", + json!({"model": model, "max_tokens": 100, "stream": true, + "messages": [{"role": "user", "content": "hi"}]}), + ) + .await + .unwrap(); + resp.assert_ok(); + let text = resp.text(); + assert!(text.contains("event: error"), "{text}"); + assert!(text.contains("overloaded_error"), "{text}"); + assert!(text.contains("below the minimum of 5."), "{text}"); +} + +/// Eight tokens a second against a minimum of five runs to its end. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_stream_at_speed_runs_to_its_end() { + let app = TestApp::spawn().await; + let model = "fast-chat"; + let pieces: Vec<(f64, usize)> = (0..16) + .map(|i| (if i == 0 { 0.0 } else { 0.5 }, 4)) + .collect(); + let (uri, hung_up) = + sse_upstream_paced("/v1/chat/completions", chat_stream(model, &pieces)).await; + let (key, route_id) = seed(&app, &uri, "openai", model, 5, 5).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw.post("/v1/chat/completions", chat(model)).await.unwrap(); + resp.assert_ok(); + let text = resp.text(); + assert!(text.contains("[DONE]"), "{text}"); + assert!(!text.contains("tokens per second"), "{text}"); + assert_eq!(errors_on(&app, route_id).await, 0); + assert!(!hung_up.load(Ordering::SeqCst)); +} + +/// A long pause is the upstream not generating, not generating slowly: +/// eight tokens a second for 6 s, 8 s of silence, then on. Counting the +/// pause, the last 10 s would drop below three a second around 13 s. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_long_pause_does_not_cut_the_stream() { + let app = TestApp::spawn().await; + let model = "pausing-chat"; + let mut pieces: Vec<(f64, usize)> = (0..13) + .map(|i| (if i == 0 { 0.0 } else { 0.5 }, 4)) + .collect(); + pieces.push((8.0, 4)); + pieces.extend((0..4).map(|_| (0.5, 4))); + let (uri, _) = sse_upstream_paced("/v1/chat/completions", chat_stream(model, &pieces)).await; + let (key, route_id) = seed(&app, &uri, "openai", model, 3, 10).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw.post("/v1/chat/completions", chat(model)).await.unwrap(); + resp.assert_ok(); + let text = resp.text(); + assert!(text.contains("[DONE]"), "{text}"); + assert!(!text.contains("tokens per second"), "{text}"); + assert_eq!(errors_on(&app, route_id).await, 0); +} + +/// Without a minimum, a slow stream runs as it always did. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn without_a_minimum_a_slow_stream_is_left_alone() { + let app = TestApp::spawn().await; + let model = "slow-unlimited"; + let pieces: Vec<(f64, usize)> = (0..7) + .map(|i| (if i == 0 { 0.0 } else { 1.0 }, 1)) + .collect(); + let (uri, _) = sse_upstream_paced("/v1/chat/completions", chat_stream(model, &pieces)).await; + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let provider = fixtures::create_provider(&app.db, &unique_name("slow"), "openai", &uri, None) + .await + .unwrap(); + fixtures::create_model_and_route(&app.db, provider.id, model) + .await + .unwrap(); + app.rebuild_gateway_router().await; + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("slow-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap() + .plaintext; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw.post("/v1/chat/completions", chat(model)).await.unwrap(); + resp.assert_ok(); + assert!(resp.text().contains("[DONE]"), "{}", resp.text()); +} + +/// The console takes the minimum in range, turns it off with `null`, and +/// refuses one out of range. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_console_sets_and_clears_a_providers_minimum() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + let provider = fixtures::create_provider( + &app.db, + &unique_name("speed"), + "openai", + "https://api.openai.com", + None, + ) + .await + .unwrap(); + let at = format!("/api/admin/providers/{}", provider.id); + + let got = con.get(&at).await.unwrap(); + got.assert_ok(); + let v: Value = got.json().unwrap(); + assert_eq!(v["min_speed_tokens_per_sec"], Value::Null); + assert_eq!(v["min_speed_window_secs"], 20); + + let set = con + .patch( + &at, + json!({"min_speed_tokens_per_sec": 15, "min_speed_window_secs": 30}), + ) + .await + .unwrap(); + set.assert_ok(); + let v: Value = set.json().unwrap(); + assert_eq!(v["min_speed_tokens_per_sec"], 15); + assert_eq!(v["min_speed_window_secs"], 30); + + // Another edit leaves it as it is. + let renamed = con + .patch(&at, json!({"display_name": "Renamed"})) + .await + .unwrap(); + renamed.assert_ok(); + assert_eq!( + renamed.json::().unwrap()["min_speed_tokens_per_sec"], + 15 + ); + + for bad in [ + json!({"min_speed_tokens_per_sec": 0}), + json!({"min_speed_tokens_per_sec": 1001}), + json!({"min_speed_window_secs": 4}), + json!({"min_speed_window_secs": 121}), + ] { + con.patch(&at, bad.clone()) + .await + .unwrap() + .assert_status(400); + } + + let cleared = con + .patch(&at, json!({"min_speed_tokens_per_sec": null})) + .await + .unwrap(); + cleared.assert_ok(); + let v: Value = cleared.json().unwrap(); + assert_eq!(v["min_speed_tokens_per_sec"], Value::Null); + assert_eq!(v["min_speed_window_secs"], 30); +} diff --git a/db/schema.sql b/db/schema.sql index 00eb4a63..1295605e 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -329,9 +329,25 @@ CREATE TABLE IF NOT EXISTS providers ( is_active BOOLEAN NOT NULL DEFAULT TRUE, config_json JSONB DEFAULT '{}', created_at TIMESTAMPTZ NOT NULL DEFAULT now(), - deleted_at TIMESTAMPTZ + deleted_at TIMESTAMPTZ, + -- The slowest generation this upstream may keep a stream at: once + -- its first token is in, a stream that averages fewer than + -- `min_speed_tokens_per_sec` over the last `min_speed_window_secs` + -- seconds (pauses of 5 s or more left out) is cut with a 503 and + -- counted against the route. NULL ⇒ off. See + -- crates/gateway/src/proxy/min_speed.rs. + min_speed_tokens_per_sec INT + CHECK (min_speed_tokens_per_sec IS NULL OR min_speed_tokens_per_sec BETWEEN 1 AND 1000), + min_speed_window_secs INT NOT NULL DEFAULT 20 + CHECK (min_speed_window_secs BETWEEN 5 AND 120) ); +-- Added after first boot; see the note above `model_routes`'s ALTER. +ALTER TABLE providers ADD COLUMN IF NOT EXISTS min_speed_tokens_per_sec + INT CHECK (min_speed_tokens_per_sec IS NULL OR min_speed_tokens_per_sec BETWEEN 1 AND 1000); +ALTER TABLE providers ADD COLUMN IF NOT EXISTS min_speed_window_secs + INT NOT NULL DEFAULT 20 CHECK (min_speed_window_secs BETWEEN 5 AND 120); + CREATE INDEX IF NOT EXISTS idx_providers_not_deleted ON providers(created_at) WHERE deleted_at IS NULL; -- Exposed catalog of model IDs clients can call via `/v1/models`. diff --git a/web/src/i18n/en.json b/web/src/i18n/en.json index c0c4fc92..b7a4ef30 100644 --- a/web/src/i18n/en.json +++ b/web/src/i18n/en.json @@ -583,6 +583,14 @@ "headers": "Request Headers", "headersDesc": "Headers sent with every request to this provider. Use {{user_id}} or {{user_email}} in values to forward caller identity.", "headerValueStored": "Saved — leave blank to keep", + "minSpeed": "Minimum generation speed", + "minSpeedBefore": "Stop when generation is slower than", + "minSpeedMiddle": "tokens/s for", + "minSpeedAfter": "s", + "minSpeedRate": "Generation speed (tokens/s)", + "minSpeedWindow": "Window (seconds)", + "minSpeedHint": "Leave the speed blank for no limit. Measured from the first token; pauses of 5 s or more are not counted. After a stop, the client retries.", + "minSpeedInvalid": "The speed must be a whole number from 1 to 1000, and the window from 5 to 120 seconds.", "customHeaders": "Custom Headers", "customHeadersDesc": "Custom HTTP headers forwarded with requests. Use {{user_id}} or {{user_email}} in values to forward caller identity.", "addHeader": "Add Header", diff --git a/web/src/i18n/zh.json b/web/src/i18n/zh.json index 1aec6bbb..8bfd2cbb 100644 --- a/web/src/i18n/zh.json +++ b/web/src/i18n/zh.json @@ -583,6 +583,14 @@ "headers": "请求头", "headersDesc": "发送给此提供商的请求头。值中可使用 {{user_id}} 或 {{user_email}} 转发调用者身份。", "headerValueStored": "已保存,留空即保持不变", + "minSpeed": "生成速度下限", + "minSpeedBefore": "生成速度低于", + "minSpeedMiddle": "token/秒持续", + "minSpeedAfter": "秒时中断", + "minSpeedRate": "生成速度(token/秒)", + "minSpeedWindow": "判定时长(秒)", + "minSpeedHint": "速度留空即不限制。从第一个 token 起计,5 秒以上的停顿不计入。中断后由客户端重试。", + "minSpeedInvalid": "生成速度须为 1–1000 的整数,判定时长须为 5–120 秒。", "customHeaders": "自定义请求头", "customHeadersDesc": "随请求转发的自定义 HTTP 头。值中可使用 {{user_id}} 或 {{user_email}} 转发调用者身份。", "addHeader": "添加请求头", diff --git a/web/src/routes/gateway/provider-dialogs.test.tsx b/web/src/routes/gateway/provider-dialogs.test.tsx index 37c8fda1..8b255375 100644 --- a/web/src/routes/gateway/provider-dialogs.test.tsx +++ b/web/src/routes/gateway/provider-dialogs.test.tsx @@ -2,6 +2,7 @@ import { describe, it, expect, vi, beforeAll, beforeEach } from 'vitest' import { render, screen } from '@testing-library/react' import userEvent, { type UserEvent } from '@testing-library/user-event' import { CreateProviderDialog, EditProviderDialog } from './provider-dialogs' +import { minSpeedFields } from './provider-min-speed' import type { Provider } from './provider-types' vi.mock('@/lib/api', () => ({ @@ -44,6 +45,8 @@ function bedrock(headers: { key: string; value: string; encrypted?: boolean }[]) base_url: 'us-east-1', is_active: true, config_json: { headers }, + min_speed_tokens_per_sec: null, + min_speed_window_secs: 20, created_at: '2026-09-28T00:00:00Z', } } @@ -153,6 +156,8 @@ describe('EditProviderDialog', () => { display_name: 'Bedrock', base_url: 'us-east-1', headers: [{ key: 'Authorization', value: 'Bearer ABSK-new' }], + min_speed_tokens_per_sec: null, + min_speed_window_secs: 20, }) }) @@ -188,4 +193,58 @@ describe('EditProviderDialog', () => { provider_id: 'prov-1', }) }) + + it('sets the slowest generation a stream may keep up', async () => { + const user = userEvent.setup() + renderEdit(bedrock([])) + + const rate = screen.getByLabelText('Generation speed (tokens/s)') + const window = screen.getByLabelText('Window (seconds)') + expect(rate).toHaveValue('') + expect(window).toHaveValue('20') + await user.type(rate, '15') + await user.clear(window) + await user.type(window, '30') + await user.click(screen.getByRole('button', { name: 'Save' })) + + expect(apiPatch).toHaveBeenCalledWith( + '/api/admin/providers/prov-1', + expect.objectContaining({ min_speed_tokens_per_sec: 15, min_speed_window_secs: 30 }), + ) + }) + + it('turns the minimum off when the speed is cleared', async () => { + const user = userEvent.setup() + renderEdit({ ...bedrock([]), min_speed_tokens_per_sec: 15, min_speed_window_secs: 30 }) + + const rate = screen.getByLabelText('Generation speed (tokens/s)') + expect(rate).toHaveValue('15') + await user.clear(rate) + await user.click(screen.getByRole('button', { name: 'Save' })) + + expect(apiPatch).toHaveBeenCalledWith( + '/api/admin/providers/prov-1', + expect.objectContaining({ min_speed_tokens_per_sec: null, min_speed_window_secs: 30 }), + ) + }) + + it('will not save a speed or window out of range', async () => { + const user = userEvent.setup() + renderEdit(bedrock([])) + + await user.type(screen.getByLabelText('Generation speed (tokens/s)'), '1001') + expect(screen.getByText(/whole number from 1 to 1000/)).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Save' })).toBeDisabled() + }) +}) + +describe('minSpeedFields', () => { + it('takes whole numbers in range, and a blank speed as off', () => { + expect(minSpeedFields('', '20')).toEqual({ min_speed_tokens_per_sec: null, min_speed_window_secs: 20 }) + expect(minSpeedFields(' 1000 ', '5')).toEqual({ min_speed_tokens_per_sec: 1000, min_speed_window_secs: 5 }) + expect(minSpeedFields('1', '120')).toEqual({ min_speed_tokens_per_sec: 1, min_speed_window_secs: 120 }) + for (const [rate, window] of [['0', '20'], ['1001', '20'], ['1.5', '20'], ['x', '20'], ['', '4'], ['', '121'], ['', '']]) { + expect(minSpeedFields(rate, window), `${rate} / ${window}`).toBeNull() + } + }) }) diff --git a/web/src/routes/gateway/provider-dialogs.tsx b/web/src/routes/gateway/provider-dialogs.tsx index d6ec8444..64f32def 100644 --- a/web/src/routes/gateway/provider-dialogs.tsx +++ b/web/src/routes/gateway/provider-dialogs.tsx @@ -21,6 +21,7 @@ import { ScrollArea } from '@/components/ui/scroll-area'; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from '@/components/ui/collapsible'; import { toast } from 'sonner'; import type { Provider, TestResult } from './provider-types'; +import { DEFAULT_MIN_SPEED_WINDOW_SECS, minSpeedFields } from './provider-min-speed'; const defaultHeadersForType = (type: string): [string, string][] => { switch (type) { @@ -427,6 +428,8 @@ export function EditProviderDialog({ open, onOpenChange, provider, onSuccess }: const [editStoredKeys, setEditStoredKeys] = useState>(new Set()); const [editSaving, setEditSaving] = useState(false); const [editError, setEditError] = useState(''); + const [minSpeedRate, setMinSpeedRate] = useState(''); + const [minSpeedWindow, setMinSpeedWindow] = useState(String(DEFAULT_MIN_SPEED_WINDOW_SECS)); const [testing, setTesting] = useState(false); const [testResult, setTestResult] = useState(null); @@ -437,6 +440,8 @@ export function EditProviderDialog({ open, onOpenChange, provider, onSuccess }: if (provider) { setEditDisplayName(provider.display_name); setEditBaseUrl(provider.base_url); + setMinSpeedRate(provider.min_speed_tokens_per_sec?.toString() ?? ''); + setMinSpeedWindow(String(provider.min_speed_window_secs ?? DEFAULT_MIN_SPEED_WINDOW_SECS)); setEditError(''); setTestResult(null); const existing = provider.config_json?.headers ?? []; @@ -453,8 +458,10 @@ export function EditProviderDialog({ open, onOpenChange, provider, onSuccess }: } } + const minSpeed = minSpeedFields(minSpeedRate, minSpeedWindow); + const handleEdit = async () => { - if (!provider) return; + if (!provider || !minSpeed) return; setEditError(''); setEditSaving(true); try { @@ -462,6 +469,7 @@ export function EditProviderDialog({ open, onOpenChange, provider, onSuccess }: display_name: editDisplayName, base_url: editBaseUrl, headers: editHeaders.filter(([k]) => k.trim()).map(([k, v]) => ({ key: k, value: v })), + ...minSpeed, }); onOpenChange(false); await onSuccess(); @@ -531,6 +539,36 @@ export function EditProviderDialog({ open, onOpenChange, provider, onSuccess }: presets={headerPresets(t)} /> +
+ +
+ {t('providers.minSpeedBefore')} + setMinSpeedRate(e.target.value)} + /> + {t('providers.minSpeedMiddle')} + setMinSpeedWindow(e.target.value)} + /> + {t('providers.minSpeedAfter')} +
+ {minSpeed ? ( +

{t('providers.minSpeedHint')}

+ ) : ( +

{t('providers.minSpeedInvalid')}

+ )} +
{editError && ( @@ -557,7 +595,7 @@ export function EditProviderDialog({ open, onOpenChange, provider, onSuccess }: setTestResult, )} /> - diff --git a/web/src/routes/gateway/provider-min-speed.ts b/web/src/routes/gateway/provider-min-speed.ts new file mode 100644 index 00000000..da41e058 --- /dev/null +++ b/web/src/routes/gateway/provider-min-speed.ts @@ -0,0 +1,16 @@ +/** A stream's window when none is set, in seconds. */ +export const DEFAULT_MIN_SPEED_WINDOW_SECS = 20; + +/** + * The minimum generation speed as the PATCH sends it, or `null` when a + * field is out of range: the speed blank (off) or a whole number from 1 + * to 1000, the window a whole number from 5 to 120. + */ +export function minSpeedFields(rate: string, window: string) { + const whole = (v: string, lo: number, hi: number) => + /^\d+$/.test(v.trim()) && Number(v) >= lo && Number(v) <= hi ? Number(v) : null; + const tokens = rate.trim() === '' ? null : whole(rate, 1, 1000); + const secs = whole(window, 5, 120); + if ((rate.trim() !== '' && tokens === null) || secs === null) return null; + return { min_speed_tokens_per_sec: tokens, min_speed_window_secs: secs }; +} diff --git a/web/src/routes/gateway/provider-types.ts b/web/src/routes/gateway/provider-types.ts index c5182d23..c1048b1f 100644 --- a/web/src/routes/gateway/provider-types.ts +++ b/web/src/routes/gateway/provider-types.ts @@ -11,6 +11,12 @@ export interface Provider { * re-submitting them blank keeps it (see `merge_headers_for_storage`). */ config_json?: { headers?: { key: string; value: string; encrypted?: boolean }[] }; + /** + * The slowest generation a stream may keep up, in tokens per second + * over `min_speed_window_secs` seconds. `null` ⇒ off. + */ + min_speed_tokens_per_sec: number | null; + min_speed_window_secs: number; created_at: string; }