diff --git a/CHANGELOG.md b/CHANGELOG.md index 7b364779..cebd2d12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,119 @@ target. ## [Unreleased] +## [3.10.0] — 2026-10-25 + +Upstreams that fail for the moment before any of the answer reached the +client are now tried again instead of the error going straight back, and +a model can set the reasoning effort its requests carry, with the +gateway adapting when an upstream refuses a level. The thinkwatch-core +crates move from v0.81.0 to v0.83.1. + +One new setting (`gateway.retry_wait_secs`), one new column +(`models.reasoning_effort`, added on startup) and one new metric; no +route or setting is removed. + +### Read before upgrading + +- **`models` gains a nullable column, `reasoning_effort`.** The schema + applied at startup adds it (`ADD COLUMN IF NOT EXISTS`); existing + models keep it empty and their requests go out unchanged. +- **A reasoning level an upstream refuses is changed and sent again, for + every model.** This applies whether or not the model has a reasoning + setting and cannot be turned off: when an upstream answers 400 because + of the reasoning level a request carries, the request goes out once + more with a level the upstream accepts, or without one, and the change + is remembered for that provider and model for 7 days. The level sent + can differ from the one the caller asked for, in either direction. +- **A request whose upstreams fail for the moment can take up to two + minutes to fail.** When every route of a model failed with an overload + or a server error, or could not be reached, before any of the answer + reached the client, the gateway now waits and sends the request again + instead of returning the error at once, for up to + `gateway.retry_wait_secs` (120 seconds by default) after the request + arrived. Clients and load balancers with a shorter timeout than that may + give up first. Set `gateway.retry_wait_secs` to 0 to keep the earlier + behaviour. +- **A whole answer that runs into the request timeout is now an error + with a log row.** A non-streaming request's upstream attempt is cut + 2 seconds before the gateway's 120-second request timeout (counted from + when the request reached the gateway port), and the client gets the + last upstream error, or a 504 `ProviderTimeout` when there was none, + with its `gateway_logs` row, instead of an empty 408. A retry starts + only when its wait and another attempt as long as the last one fit + before then. +- **A stream's opening events arrive with its first content.** While a + stream can still be sent again, what the upstream sends before any + content (Anthropic `message_start` and `ping`, Responses + `response.created` and `response.in_progress`, a Chat chunk with only + the role) is held back and goes out together with the first content. + While it is held, and while the gateway waits to send again, the client + receives an SSE comment (`: keep-alive`) after every 15 seconds without + data — except Gemini clients, which never receive one. The status is + 200 as before. + +### Added + +- **A model can set its reasoning** (`models.reasoning_effort`, the + Reasoning control in the model editor): `off` takes reasoning out of + requests that have it on, and a level (`low`, `medium`, `high`, + `xhigh`, `max`) changes the effort of requests whose reasoning is on. + Requests without reasoning, such as title generation, go out + unchanged. The setting is written the way each hop's upstream format + writes it: with `off`, OpenAI Chat and Responses upstreams receive + `none`, Gemini 2.5 Flash a thinking budget of 0, and Anthropic and + Bedrock upstreams no thinking; Gemini models that cannot stop thinking + (2.5 Pro, Gemini 3) get no thinking configuration. The response cache + keys on the request together with the setting. +- **A reasoning level the upstream refuses is sent again once with one it + accepts.** On a 400 about the reasoning level, the request goes once + more to the same route with the level the error names (the lowest one + it lists when `none` was refused), or without one when it names none. + This is not counted as a retry, though its time falls within the retry + window. When that is answered, the choice is remembered in Redis per + provider, upstream model and request format for 7 days, so later + requests send what worked straight away; a provider's memory is + cleared when its URL or credentials change, a route relearns its + protocol, or the provider is deleted. A hop whose reasoning level was + changed records `"effort": {asked, sent, reason}` on its + `gateway_logs` row. +- **Transient upstream failures are retried.** When every candidate route + has been tried and the last one failed for the moment — 500, 502, 503, + 504 or 529, a connection that could not be made or broke before the + answer's head, a stream that broke off or reported an overload or a + server error (`overloaded_error`, `api_error`, `server_error`, + `server_is_overloaded`, `slow_down`, Gemini `UNAVAILABLE`) before any + content — and none of the answer reached the client, the request waits + 2, 4, 8, then 16 seconds (each give or take 20%), or as long as the + upstream said in `retry-after-ms` or `Retry-After` (a zero wait is + ignored, and a shorter one than 250 ms is taken as 250 ms), and goes out again + to the routes that failed that way. A model with one route, and a + stream, goes back to the same route. Each attempt is picked as before: + the strategy, the circuit breaker, the route caps, session affinity + and the provider holding a conversation's encrypted history all apply, + and each failed attempt counts against its route's breaker as it did. + Retries stop when the next wait would end past the deadline; an attempt + already under way is not cut short. Rate limits and spent quota (429), + refused credentials (401, 403), any other 4xx, timeouts, the gateway's + own limits, budgets and policies, and a client that left are never + retried, and neither is a stream after part of its answer went out. +- **`gateway.retry_wait_secs`** (Settings → Gateway → Routing strategy): + how long after a request arrives it may still be sent again, 0 to 600 + seconds, default 120; 0 turns retries off. +- **`gateway_provider_retry_total{provider}`**: attempts sent again after + a transient failure, by the provider they went to, next to + `gateway_provider_fallback_total`. +- The `gateway_logs` row of a request that was sent again records + `"retries": n` in its detail, whether it then succeeded or failed. + +### Changed + +- A provider whose base URL cannot be made into a request fails as + `ProviderError` instead of `NetworkError`, and is not retried. +- An error a stream reports as OpenAI's `server_error` is logged as the + upstream's 500 (`ProviderHttpError`), as the same error answered with + HTTP 500 is, instead of 502 (`ProviderError`). + ## [3.9.1] — 2026-10-23 A conversation whose history a provider encrypted, such as a Codex @@ -2581,7 +2694,9 @@ unreleased builds should: stop the gateway, run `db/schema.sql` against PostgreSQL, restart against this tag. The schema is idempotent end-to-end, so the apply is safe to repeat. -[Unreleased]: https://github.com/ThinkWatchProject/ThinkWatch/compare/v3.9.0...HEAD +[Unreleased]: https://github.com/ThinkWatchProject/ThinkWatch/compare/v3.10.0...HEAD +[3.10.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.10.0 +[3.9.1]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.9.1 [3.9.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.9.0 [3.8.3]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.8.3 [3.8.2]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.8.2 diff --git a/Cargo.lock b/Cargo.lock index cad1926f..12d5252b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4189,7 +4189,7 @@ dependencies = [ [[package]] name = "think-watch-auth" -version = "3.9.1" +version = "3.10.0" dependencies = [ "anyhow", "argon2", @@ -4219,7 +4219,7 @@ dependencies = [ [[package]] name = "think-watch-common" -version = "3.9.1" +version = "3.10.0" dependencies = [ "aes-gcm", "anyhow", @@ -4262,7 +4262,7 @@ dependencies = [ [[package]] name = "think-watch-gateway" -version = "3.9.1" +version = "3.10.0" dependencies = [ "anyhow", "arc-swap", @@ -4301,7 +4301,7 @@ dependencies = [ [[package]] name = "think-watch-mcp-gateway" -version = "3.9.1" +version = "3.10.0" dependencies = [ "anyhow", "arc-swap", @@ -4330,7 +4330,7 @@ dependencies = [ [[package]] name = "think-watch-server" -version = "3.9.1" +version = "3.10.0" dependencies = [ "anyhow", "arc-swap", @@ -4381,7 +4381,7 @@ dependencies = [ [[package]] name = "think-watch-test-support" -version = "3.9.1" +version = "3.10.0" dependencies = [ "anyhow", "async-stream", @@ -4837,8 +4837,8 @@ dependencies = [ [[package]] name = "tw-bedrock" -version = "0.81.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.81.0#aae0e486da837c94189046a7bcf5095d4eb1948f" +version = "0.83.1" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.83.1#0a34ac1ad9d1a6691cb98e38b3a9b7e4a1803bb9" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -4852,16 +4852,16 @@ dependencies = [ [[package]] name = "tw-breaker" -version = "0.81.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.81.0#aae0e486da837c94189046a7bcf5095d4eb1948f" +version = "0.83.1" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.83.1#0a34ac1ad9d1a6691cb98e38b3a9b7e4a1803bb9" dependencies = [ "serde", ] [[package]] name = "tw-dialect" -version = "0.81.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.81.0#aae0e486da837c94189046a7bcf5095d4eb1948f" +version = "0.83.1" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.83.1#0a34ac1ad9d1a6691cb98e38b3a9b7e4a1803bb9" dependencies = [ "memchr", "serde", @@ -4870,8 +4870,8 @@ dependencies = [ [[package]] name = "tw-guard" -version = "0.81.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.81.0#aae0e486da837c94189046a7bcf5095d4eb1948f" +version = "0.83.1" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.83.1#0a34ac1ad9d1a6691cb98e38b3a9b7e4a1803bb9" dependencies = [ "aho-corasick", "base64 0.22.1", diff --git a/Cargo.toml b/Cargo.toml index 58592396..ea2d34fc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ members = [ ] [workspace.package] -version = "3.9.1" +version = "3.10.0" edition = "2024" # The MSRV: the newest `rust-version` among the locked dependencies # (sqlx 0.9 and the aws-smithy crates under aws-sigv4). Without it, @@ -70,10 +70,10 @@ opt-level = 3 # never re-exported through a local shim. And the reverse: something only # this side uses (the at-rest crypto, IMDSv2 credentials, the gateway error) # lives here, not in core. -tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.81.0" } -tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.81.0" } -tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.81.0" } -tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.81.0" } +tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.83.1" } +tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.83.1" } +tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.83.1" } +tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.83.1" } # Web framework axum = { version = "0.8", features = ["macros", "ws"] } diff --git a/crates/common/src/dynamic_config.rs b/crates/common/src/dynamic_config.rs index ed044a87..3a777b7d 100644 --- a/crates/common/src/dynamic_config.rs +++ b/crates/common/src/dynamic_config.rs @@ -216,6 +216,12 @@ macro_rules! dc_getters_u64_from_i64 { }; } +/// `gateway.retry_wait_secs` when the setting is missing. +pub const RETRY_WAIT_DEFAULT_SECS: i64 = 120; +/// Upper bound on `gateway.retry_wait_secs`, enforced both at write time +/// (the settings handler) and at read time. +pub const RETRY_WAIT_MAX_SECS: i64 = 600; + /// Upper bound on `security.signature_drift_secs`, enforced both at /// write time (admin settings validation) and at read time (the /// getter below). Anything beyond five minutes weakens replay @@ -440,6 +446,18 @@ impl DynamicConfig { raw.clamp(0, 86400) } + /// How long after it arrived a request may still be sent again to + /// upstreams that failed for the moment (overloaded, a server error, + /// a broken connection) before any of the answer reached the client. + /// Clamped to [0, 600]; 0 never sends a request again. + pub async fn retry_wait_secs(&self) -> u64 { + let raw = self + .get_i64("gateway.retry_wait_secs") + .await + .unwrap_or(RETRY_WAIT_DEFAULT_SECS); + raw.clamp(0, RETRY_WAIT_MAX_SECS) as u64 + } + /// Exponent for the `latency` strategy: higher = more aggressive /// (heavier penalty on slow routes). Default 2.0 (aggressive). /// Clamped to [0.5, 8.0] so a typo can't make weights collapse diff --git a/crates/common/src/models/provider.rs b/crates/common/src/models/provider.rs index 8c8a481f..cc02a0e7 100644 --- a/crates/common/src/models/provider.rs +++ b/crates/common/src/models/provider.rs @@ -65,6 +65,10 @@ pub struct Model { /// limit. #[serde(default)] pub max_output_tokens: Option, + /// The reasoning setting: `off`, or a level (`low` … `max`) for + /// requests whose reasoning is on. `None` ⇒ requests go out as sent. + #[serde(default)] + pub reasoning_effort: Option, } fn default_enabled() -> bool { diff --git a/crates/gateway/Cargo.toml b/crates/gateway/Cargo.toml index 9d138611..da2b5e87 100644 --- a/crates/gateway/Cargo.toml +++ b/crates/gateway/Cargo.toml @@ -40,3 +40,4 @@ rust_decimal = { workspace = true } [dev-dependencies] http_1x = { package = "http", version = "1" } +tokio = { workspace = true, features = ["test-util"] } diff --git a/crates/gateway/src/error.rs b/crates/gateway/src/error.rs index 03b3b44b..fa771ff5 100644 --- a/crates/gateway/src/error.rs +++ b/crates/gateway/src/error.rs @@ -12,9 +12,17 @@ pub enum GatewayError { /// Upstream returned a non-2xx, non-429, non-401 status. The /// status is kept structured so error-classifier metrics stay /// readable and the gateway can classify retry-eligible 5xx - /// versus poison 4xx without parsing the message. + /// versus poison 4xx without parsing the message. `retry_after` is + /// when the upstream said to try again (`retry-after-ms` or + /// `Retry-After`), which an overloaded upstream can send with a 503: + /// the gateway waits that long before sending to it again (see + /// `proxy::retry`). #[error("Provider HTTP {status}: {message}")] - ProviderHttpError { status: u16, message: String }, + ProviderHttpError { + status: u16, + message: String, + retry_after: Option, + }, /// Upstream took longer than the configured timeout. Distinct /// from a network drop because the request reached the upstream /// — only the response was missing in time. diff --git a/crates/gateway/src/lifecycle/mod.rs b/crates/gateway/src/lifecycle/mod.rs index 98e73f59..f57cd9b3 100644 --- a/crates/gateway/src/lifecycle/mod.rs +++ b/crates/gateway/src/lifecycle/mod.rs @@ -34,9 +34,10 @@ use tw_dialect::ir::Dialect; use crate::guards::Guards; use crate::proxy::generate::{Wire, priced, tokens}; +use crate::proxy::retry::{KEEP_ALIVE, KEEP_ALIVE_FRAME}; use crate::proxy::shaper::{StreamShaper, rewrite_model}; use crate::proxy::{ - GatewayRequestIdentity, GatewayState, SelectionRecord, emit_gateway_log_with_extra, + GatewayRequestIdentity, GatewayState, Resend, SelectionRecord, emit_gateway_log_with_extra, finalize_health, post_flight_account, prepare_body_capture, }; @@ -158,11 +159,37 @@ pub(crate) struct ChatPostInvokeDeps { pub cache_enabled: bool, } -/// The upstream call a stream makes, not yet started. -pub(crate) type OpenUpstream = Pin< - Box> + Send>, +/// The upstream call a stream makes, not yet started: each call starts +/// one attempt. +pub(crate) type OpenUpstream = Box< + dyn FnMut() -> Pin< + Box< + dyn std::future::Future> + + Send, + >, + > + Send, >; +/// What a stream's tail resolves to once the stream ends: what the hooks +/// read, and how the stream reached its upstream. +pub(crate) type StreamTail = + Pin, Sent)> + Send>>; + +/// How a stream reached its upstream, for the tail: how many times it went +/// out again after a transient failure, and when the attempt that ended it +/// went out. +#[derive(Debug, Clone, Copy)] +pub(crate) struct Sent { + pub retries: u32, + pub last_sent: std::time::Instant, + /// Waiting to go out again: the attempt that failed was counted + /// against the route already, and nothing is under way. + pub waiting: bool, + /// The attempt under way sent a different reasoning effort than the + /// caller asked (see `proxy::effort`). + pub effort: Option, +} + /// Build the streaming pump: forward the upstream's bytes to the caller — /// converted if the route speaks another format, shaped either way — and /// return a tail future that resolves once the stream ends. @@ -176,8 +203,19 @@ pub(crate) type OpenUpstream = Pin< /// rejected dialect is still retried inside `open`, before any byte /// reaches the caller. /// -/// Nothing is buffered. Usage is sniffed and the whole answer assembled -/// alongside the bytes, not by holding them back. +/// Nothing of the answer is buffered. Usage is sniffed and the whole +/// answer assembled alongside the bytes, not by holding them back. +/// +/// **A transient failure before any of the answer went out goes out +/// again** to the same route (see `proxy::retry`; `resend` decides): the +/// upstream could not be reached, answered 5xx, broke off, or reported an +/// overload or a server error before any content. While a failure could +/// still be sent again, what the upstream sends ahead of its content is +/// held back (see [`Opening`]), so that a retry's opening is not the +/// client's second one. While it is held, and while the gateway waits to +/// send again, the client gets an SSE keep-alive comment after every +/// [`KEEP_ALIVE`] of silence. The status is 200 from the start, as it +/// always was: a failure that is the answer is said in the stream. /// /// **A dropped stream is a cancelled request.** When the client goes, /// hyper drops the body, and with it the sender the tail is waiting on; @@ -205,6 +243,7 @@ pub(crate) type OpenUpstream = Pin< #[allow(clippy::too_many_arguments)] pub(crate) fn build_chat_pump( open: OpenUpstream, + resend: Resend, mut shaper: StreamShaper, client: Dialect, client_sse: bool, @@ -212,10 +251,7 @@ pub(crate) fn build_chat_pump( guards: &Guards, request: &ChatRequestSnapshot, provider: &str, -) -> ( - axum::response::Response, - Pin> + Send>>, -) { +) -> (axum::response::Response, StreamTail) { let readers = Arc::new(Mutex::new(Readers::default())); let readers_for_tail = Arc::clone(&readers); @@ -247,149 +283,305 @@ pub(crate) fn build_chat_pump( let mut last_frame = LastFrame::new(client, client_sse); let provider_for_tail = provider.to_string(); let provider = provider.to_string(); + let sent = Arc::new(Mutex::new(Sent { + retries: 0, + last_sent: std::time::Instant::now(), + waiting: false, + effort: None, + })); + // Gemini clients get no keep-alive comments: the google-genai SDK reads + // every SSE line as JSON. + let keep_alive = keeps_alive(client); + let sent_for_tail = Arc::clone(&sent); let body = async_stream::stream! { let mut done_tx = Some(done_tx); - - let (upstream, wire) = match open.await { - Ok(opened) => opened, - Err(e) => { - // Headers already went out as 200, so the refusal is said - // in the stream — and logged with the upstream's own status, - // so a throttled upstream stays 429 on the audit row. - let mut out = shaper.process(&error_frame(client, e.status_code(), &e.to_string())); - out.extend(shaper.finish()); - tell(&mut done_tx, StreamOutcome::UpstreamError { - error_type: e.error_tag().to_string(), - message: e.to_string(), - status_code: e.status_code(), - }); - yield Ok::(Bytes::from(out)); - return; + let mut open = open; + let mut resend = resend; + // When the client last received something: a keep-alive is due + // `KEEP_ALIVE` after it while the gateway holds the stream back. + let mut quiet_since = tokio::time::Instant::now(); + // The wait before the next attempt, after a transient failure. + let mut backoff: Option = None; + + 'attempt: loop { + if let Some(wait) = backoff.take() { + let until = tokio::time::Instant::now() + wait; + if keep_alive { + let mut ticks = std::pin::pin!(keep_alives(quiet_since, until)); + while ticks.next().await.is_some() { + quiet_since = tokio::time::Instant::now(); + yield Ok::(Bytes::from_static(KEEP_ALIVE_FRAME)); + } + } else { + tokio::time::sleep_until(until).await; + } } - }; - if let Ok(mut r) = readers.lock() { - r.sniffer = Some(tw_dialect::usage::Sniffer::new()); - r.collector = Some(wire.collect.collector()); - r.errors = Some(ErrorWatch::new(wire.dialect)); - r.ledger = Some(wire.ledger.clone()); - } - // The hop that answered can have numbered a value of its own. - shaper.restore_with(&wire.ledger); - let mut convert = wire.convert.as_ref().map(|s| s.stream()); - // Bedrock streams AWS eventstream frames, not SSE. Unframe them at - // the door, so the sniffer, the collector and the converter all - // read the same SSE they read from every other upstream. - let mut unframe = (wire.dialect == Dialect::Bedrock) - .then(tw_bedrock::eventstream::Transcoder::new); - let mut source = upstream.bytes_stream(); - while let Some(item) = source.next().await { - let item = match item { - Ok(raw) => match unframe.as_mut() { - None => Ok(raw), - Some(t) => t - .feed(&raw) - .map(Bytes::from) - .map_err(|e| bedrock_ended(&provider, e)), - }, - Err(e) => Err(broke_off(format!("The upstream stream broke off: {e}"))), + 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 }; + } + // What an attempt that failed had read is not this one's: its + // usage is not billed, its error not the request's. + if let Ok(mut r) = readers.lock() { + *r = Readers::default(); + } + // A retry's upstream says nothing until it answers: the client + // hears keep-alives meanwhile. + let resent = resend.budget.retries() > 0; + let mut opening = open(); + let opened = loop { + let tick = tokio::time::sleep_until(quiet_since + KEEP_ALIVE); + let next = tokio::select! { + opened = &mut opening => Some(opened), + _ = tick, if resent && keep_alive => None, + }; + match next { + Some(opened) => break opened, + None => { + quiet_since = tokio::time::Instant::now(); + yield Ok(Bytes::from_static(KEEP_ALIVE_FRAME)); + } + } }; - match item { - Ok(chunk) => { - let sealed = match readers.lock() { - Ok(mut r) => { - r.feed(&chunk); - r.sealed_refusal() + let (upstream, wire) = match opened { + Ok(opened) => opened, + Err(e) => { + // Nothing of the answer went out: one that failed for + // the moment goes out again after a wait. + if let Some(wait) = resend.budget.wait_after(&e) + && resend.admit(elapsed_ms(sent_at)).await + { + resending(&provider, resend.budget.retries(), wait, e.error_tag(), &readers, &sent); + backoff = Some(wait); + continue 'attempt; + } + // Headers already went out as 200, so the refusal is said + // in the stream — and logged with the upstream's own status, + // so a throttled upstream stays 429 on the audit row. + let mut out = shaper.process(&error_frame(client, e.status_code(), &e.to_string())); + out.extend(shaper.finish()); + tell(&mut done_tx, StreamOutcome::UpstreamError { + error_type: e.error_tag().to_string(), + message: e.to_string(), + status_code: e.status_code(), + }); + yield Ok(Bytes::from(out)); + return; + } + }; + if let Ok(mut s) = sent.lock() { + s.effort = wire.effort; + } + if let Ok(mut r) = readers.lock() { + *r = Readers { + sniffer: Some(tw_dialect::usage::Sniffer::new()), + collector: Some(wire.collect.collector()), + errors: Some(ErrorWatch::new(wire.dialect)), + ledger: Some(wire.ledger.clone()), + said: None, + }; + } + // The hop that answered can have numbered a value of its own. + shaper.restore_with(&wire.ledger); + let mut convert = wire.convert.as_ref().map(|s| s.stream()); + // Bedrock streams AWS eventstream frames, not SSE. Unframe them at + // the door, so the sniffer, the collector and the converter all + // read the same SSE they read from every other upstream. + let mut unframe = (wire.dialect == Dialect::Bedrock) + .then(tw_bedrock::eventstream::Transcoder::new); + // 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)); + 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, + }; + let Some(item) = next else { + quiet_since = tokio::time::Instant::now(); + yield Ok(Bytes::from_static(KEEP_ALIVE_FRAME)); + continue; + }; + let Some(item) = item else { + // The upstream's stream ended. One whose opening is still + // held ended without content: an error it reported that + // is transient goes out again; anything else goes to the + // client as it would have. + if let Some(h) = hold.take() { + let said = readers.lock().ok().and_then(|mut r| r.reported(true)); + if let Some(said) = said + && let Some(wait) = retry_after_said(&resend, &provider, &said) + && resend.admit(elapsed_ms(sent_at)).await + { + resending(&provider, resend.budget.retries(), wait, "in-stream error", &readers, &sent); + backoff = Some(wait); + continue 'attempt; + } + match deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client) { + Ok(out) => { + if !out.is_empty() { + if last_frame.is_in(&out) { + delivered.store(true, Ordering::Release); + } + yield Ok(Bytes::from(out)); + } + } + Err((out, outcome)) => { + tell(&mut done_tx, outcome); + yield Ok(Bytes::from(out)); + return; + } } - Err(_) => None, - }; - // The upstream cannot read the conversation's encrypted - // history: the client is told why in plain words, in - // place of the upstream's own error. - if let Some(said) = sealed { - let err = sealed_refused(&provider, &said); - tracing::warn!( - provider = %provider, - message = %said.message, - "upstream refused the conversation's encrypted history" - ); - let out = cut(&shaper, convert.as_mut(), client, &[], &err); - 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 client_bytes = match convert.as_mut() { - Some(c) => c.process(&chunk), - None => chunk.to_vec(), - }; - let out = shaper.process(&client_bytes); - // A tool call the inspection stops: what came before - // still goes out, then the refusal. - let stop = inspector.as_mut().and_then(|i| i.check(&out)); - if let Some((err, safe)) = stop { - let out = cut(&shaper, convert.as_mut(), client, &out[..safe], &err); - 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; + break; + }; + let item = match item { + Ok(raw) => match unframe.as_mut() { + None => Ok(raw), + Some(t) => t + .feed(&raw) + .map(Bytes::from) + .map_err(|e| bedrock_ended(&provider, e)), + }, + Err(e) => Err(broke_off(format!("The upstream stream broke off: {e}"))), + }; + match item { + Ok(chunk) => { + let (sealed, said) = match readers.lock() { + Ok(mut r) => { + r.feed(&chunk); + (r.sealed_refusal(), r.reported(false)) + } + Err(_) => (None, None), + }; + // The upstream cannot read the conversation's encrypted + // history: the client is told why in plain words, in + // place of the upstream's own error. + if let Some(said) = sealed { + let err = sealed_refused(&provider, &said); + tracing::warn!( + provider = %provider, + message = %said.message, + "upstream refused the conversation's encrypted history" + ); + let out = cut(&shaper, convert.as_mut(), client, &[], &err); + 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 chunk = match hold.as_mut() { + None => chunk.to_vec(), + Some(h) => { + // An error before any content: one that is + // transient goes out again. Anything else is + // the answer, and goes out with the opening. + if !h.feed(&chunk) { + let Some(said) = &said else { continue }; + if let Some(wait) = retry_after_said(&resend, &provider, said) + && resend.admit(elapsed_ms(sent_at)).await + { + resending(&provider, resend.budget.retries(), wait, "in-stream error", &readers, &sent); + backoff = Some(wait); + continue 'attempt; + } + } + hold.take().map(Opening::release).unwrap_or_default() + } + }; + match deliver(&chunk, &mut convert, &mut shaper, inspector.as_mut(), client) { + Ok(out) => { + if !out.is_empty() { + if last_frame.is_in(&out) { + delivered.store(true, Ordering::Release); + } + yield Ok(Bytes::from(out)); + } + } + // A tool call the inspection stops: what came + // before still goes out, then the refusal. + Err((out, outcome)) => { + tell(&mut done_tx, outcome); + yield Ok(Bytes::from(out)); + return; + } + } } - if !out.is_empty() { - if last_frame.is_in(&out) { - delivered.store(true, Ordering::Release); + Err((message, outcome)) => { + // Broken off before any content: the upstream's + // failure, which goes out again. + if let Some(h) = hold.take() { + if let StreamOutcome::UpstreamError { error_type, status_code, .. } = &outcome + && let Some(wait) = resend.budget.wait_after_outcome(error_type, *status_code) + && resend.admit(elapsed_ms(sent_at)).await + { + resending(&provider, resend.budget.retries(), wait, error_type, &readers, &sent); + backoff = Some(wait); + continue 'attempt; + } + // The opening it sent goes out before the error. + match deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client) { + Ok(out) => { + if !out.is_empty() { + yield Ok(Bytes::from(out)); + } + } + Err((out, outcome)) => { + tell(&mut done_tx, outcome); + yield Ok(Bytes::from(out)); + return; + } + } } + // Headers are gone; the only way left to say it is in + // the stream, in the caller's own format. + let tail = match convert.as_mut() { + Some(c) => c.fail(&message), + None => error_frame(client, 502, &message), + }; + let mut out = shaper.process(&tail); + out.extend(shaper.finish()); + tell(&mut done_tx, outcome); yield Ok(Bytes::from(out)); + return; } } - Err((message, outcome)) => { - // Headers are gone; the only way left to say it is in - // the stream, in the caller's own format. - let tail = match convert.as_mut() { - Some(c) => c.fail(&message), - None => error_frame(client, 502, &message), - }; - let mut out = shaper.process(&tail); - out.extend(shaper.finish()); - tell(&mut done_tx, outcome); - yield Ok(Bytes::from(out)); - return; - } } - } - let tail = convert.as_mut().map(|c| c.finish()).unwrap_or_default(); - let mut out = shaper.process(&tail); - out.extend(shaper.finish()); - // The converter's last bytes can complete a tool call (the block's - // stop), so they are inspected too. - // A Chat call with no `finish_reason` or `[DONE]` completes only - // as the stream closes; one that stops it keeps the tail back. - let stop = inspector.as_mut().and_then(|i| { - i.check(&out) - .or_else(|| i.finish().map(|err| (err, 0))) - }); - if let Some((err, safe)) = stop { - let out = cut(&shaper, None, client, &out[..safe], &err); - tell(&mut done_tx, StreamOutcome::UpstreamError { - error_type: err.error_tag().to_string(), - message: err.to_string(), - status_code: err.status_code(), + let tail = convert.as_mut().map(|c| c.finish()).unwrap_or_default(); + let mut out = shaper.process(&tail); + out.extend(shaper.finish()); + // The converter's last bytes can complete a tool call (the block's + // stop), so they are inspected too. + // A Chat call with no `finish_reason` or `[DONE]` completes only + // as the stream closes; one that stops it keeps the tail back. + let stop = inspector.as_mut().and_then(|i| { + i.check(&out) + .or_else(|| i.finish().map(|err| (err, 0))) }); - yield Ok(Bytes::from(out)); + if let Some((err, safe)) = stop { + let out = cut(&shaper, None, client, &out[..safe], &err); + 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 outcome = match readers.lock().ok().and_then(|mut r| r.upstream_error()) { + Some(said) => failed_partway(&provider, &said), + None => StreamOutcome::Natural, + }; + tell(&mut done_tx, outcome); + if !out.is_empty() { + yield Ok(Bytes::from(out)); + } return; } - let outcome = match readers.lock().ok().and_then(|mut r| r.upstream_error()) { - Some(said) => failed_partway(&provider, &said), - None => StreamOutcome::Natural, - }; - tell(&mut done_tx, outcome); - if !out.is_empty() { - yield Ok(Bytes::from(out)); - } }; // A Gemini caller that did not ask for SSE reads one JSON array. @@ -494,7 +686,13 @@ pub(crate) fn build_chat_pump( assembled, logged, }; - Invoked { + let sent = sent_for_tail.lock().map(|s| *s).unwrap_or(Sent { + retries: 0, + last_sent: started_at, + waiting: false, + effort: None, + }); + let invoked = Invoked { client_ip: identity.ip_address.clone(), identity, trace_id, @@ -504,7 +702,8 @@ pub(crate) fn build_chat_pump( }, access_candidate: mapped_model, view: CapturedView::Streaming { outcome, captured }, - } + }; + (invoked, sent) }); (response, tail) } @@ -566,6 +765,20 @@ impl Readers { self.said.as_ref().filter(|s| s.sealed).cloned() } + /// The error the upstream reported in its stream so far, left in place + /// for the tail. `ended`: the stream ended, so a last frame without a + /// blank line after it is read too. + fn reported(&mut self, ended: bool) -> Option { + if ended + && self.said.is_none() + && let Some(said) = self.errors.as_mut().and_then(ErrorWatch::flush) + { + self.said = Some(said); + self.errors = None; + } + self.said.clone() + } + /// The error the upstream reported in its stream, if it did — read to /// the end, for an upstream whose last frame has no blank line after /// it. @@ -723,7 +936,7 @@ fn status_named(name: &str) -> Option { | "insufficient_quota" | "usage_limit_reached" | "RESOURCE_EXHAUSTED" => 429, - "api_error" => 500, + "api_error" | "server_error" => 500, "server_is_overloaded" | "slow_down" | "UNAVAILABLE" => 503, "timeout_error" | "DEADLINE_EXCEEDED" => 504, "overloaded_error" => 529, @@ -749,6 +962,11 @@ fn failed_partway(provider: &str, said: &Said) -> StreamOutcome { message = %said.message, "upstream reported an error partway through a stream" ); + outcome_of(provider, said) +} + +/// [`failed_partway`], without the log line. +fn outcome_of(provider: &str, said: &Said) -> StreamOutcome { let words: String = said.message.chars().take(500).collect(); if said.sealed { let e = sealed_refused(provider, said); @@ -847,6 +1065,165 @@ fn bedrock_status(kind: &str) -> u16 { .map_or(500, |(_, status)| *status) } +/// The wait before sending again after an error the upstream reported +/// in its stream before any content, when it is transient (see +/// `proxy::retry`) and the wait fits. +fn retry_after_said(resend: &Resend, provider: &str, said: &Said) -> Option { + match outcome_of(provider, said) { + StreamOutcome::UpstreamError { + error_type, + status_code, + .. + } => resend.budget.wait_after_outcome(&error_type, status_code), + _ => None, + } +} + +/// Whether the client may be sent keep-alive comments while the gateway +/// holds its stream back. Not a Gemini client: the google-genai SDK reads +/// every SSE line as JSON, and a JSON-array stream has no place for them. +fn keeps_alive(client: Dialect) -> bool { + client != Dialect::Gemini +} + +/// A stream goes out again after `wait`: what the failed attempt left for +/// the tail is cleared — a client that leaves during the wait left a +/// request whose failed attempt is already counted against the route +/// (see `routing::Resend::admit`), with nothing to bill — and the log +/// line written. +fn resending( + provider: &str, + retry: u32, + wait: std::time::Duration, + error: &str, + readers: &Mutex, + sent: &Mutex, +) { + if let Ok(mut r) = readers.lock() { + *r = Readers::default(); + } + if let Ok(mut s) = sent.lock() { + *s = Sent { + retries: retry, + last_sent: std::time::Instant::now(), + waiting: true, + effort: None, + }; + } + tracing::info!( + provider, + retry, + wait_ms = wait.as_millis() as u64, + error, + "Upstream failed for the moment before any of the answer went out, sending again after a wait" + ); +} + +fn elapsed_ms(since: std::time::Instant) -> u32 { + since.elapsed().as_millis().min(u32::MAX as u128) as u32 +} + +/// A tick each time the client is due a keep-alive while the gateway waits +/// until `until`: [`KEEP_ALIVE`] after `quiet_since`, then every +/// [`KEEP_ALIVE`]. Ends at `until`. +fn keep_alives( + quiet_since: tokio::time::Instant, + until: tokio::time::Instant, +) -> impl futures::Stream { + async_stream::stream! { + let mut due = quiet_since + KEEP_ALIVE; + while due < until { + tokio::time::sleep_until(due).await; + yield (); + due += KEEP_ALIVE; + } + tokio::time::sleep_until(until).await; + } +} + +/// What the client gets for `chunk` of the upstream's bytes: converted if +/// the route speaks another format, shaped, and inspected for tool calls. +/// `Err`: a tool call the inspection stops — what came before it, then +/// the refusal — and the stream's outcome. +fn deliver( + chunk: &[u8], + convert: &mut Option, + shaper: &mut StreamShaper, + inspector: Option<&mut crate::tool_inspection::StreamInspector>, + client: Dialect, +) -> Result, (Vec, StreamOutcome)> { + let client_bytes = match convert.as_mut() { + Some(c) => c.process(chunk), + None => chunk.to_vec(), + }; + let out = shaper.process(&client_bytes); + match inspector.and_then(|i| i.check(&out)) { + Some((err, safe)) => { + let out = cut(shaper, convert.as_mut(), client, &out[..safe], &err); + Err(( + out, + StreamOutcome::UpstreamError { + error_type: err.error_tag().to_string(), + message: err.to_string(), + status_code: err.status_code(), + }, + )) + } + None => Ok(out), + } +} + +/// A stream's opening, held back from the client while a failure could +/// still be sent again (see `proxy::retry`): what an upstream sends as +/// soon as it takes the request — Anthropic's `message_start` and `ping`, +/// Responses' `response.created` and `response.in_progress`, a Chat chunk +/// with only the role, Bedrock's `messageStart` — is not yet the answer. +/// Sent ahead of a failure, it could not be taken back, and a retry would +/// send it twice. It goes out with the first content, or with an error +/// that is the answer. +/// +/// Read on the upstream's bytes in its own format, as +/// `tw_dialect::convert::Reader` reads them: anything past the start and +/// the usage is content. An opening that grows past [`Opening::MAX`] +/// without any is let through: a broken upstream must not make it grow +/// without end. +struct Opening { + reader: tw_dialect::convert::Reader, + held: Vec, +} + +impl Opening { + const MAX: usize = 1024 * 1024; + + fn new(upstream: Dialect) -> Self { + Self { + reader: tw_dialect::convert::Reader::new(upstream), + held: Vec::new(), + } + } + + /// Hold `chunk` back. True once it carries content ahead of any error, + /// and with it what was held can go out: an error after content is + /// one after the answer began, whatever piece of the upstream's bytes + /// it came in. + fn feed(&mut self, chunk: &[u8]) -> bool { + use tw_dialect::ir::Event; + self.held.extend_from_slice(chunk); + let content = self + .reader + .feed(chunk) + .iter() + .take_while(|e| !matches!(e, Event::Error { .. })) + .any(|e| !matches!(e, Event::Start { .. } | Event::Usage(_))); + content || self.held.len() > Self::MAX + } + + /// What was held back, to go out. + fn release(self) -> Vec { + self.held + } +} + /// 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() { @@ -1156,7 +1533,13 @@ impl Surface for ChatCompletionSurface { deps.request.request_started_at.elapsed().as_millis() as i64, logged_status, with_usage_detail( - pinned_detail(error_detail, deps.route.pinned), + effort_detail( + retried_detail( + pinned_detail(error_detail, deps.route.pinned), + deps.route.sel_record.retries, + ), + deps.route.sel_record.effort, + ), &usage, usage_estimated(&invoked.view), ), @@ -1209,6 +1592,37 @@ fn pinned_detail(detail: Option, pinned: bool) -> Option, retries: u32) -> Option { + if retries == 0 { + return detail; + } + let mut d = match detail { + Some(serde_json::Value::Object(d)) => d, + _ => serde_json::Map::new(), + }; + d.insert("retries".into(), retries.into()); + Some(serde_json::Value::Object(d)) +} + +/// The audit detail, with `"effort"` when the hop that answered sent a +/// different reasoning effort than the caller asked (see `proxy::effort`). +fn effort_detail( + detail: Option, + effort: Option, +) -> Option { + let Some(effort) = effort else { + return detail; + }; + let mut d = match detail { + Some(serde_json::Value::Object(d)) => d, + _ => serde_json::Map::new(), + }; + d.insert("effort".into(), effort.detail()); + 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 { @@ -1233,6 +1647,243 @@ fn usage_estimated(view: &CapturedView) -> bool { mod tests { use super::*; + fn frame(kind: &str, v: serde_json::Value) -> String { + format!("event: {kind}\ndata: {v}\n\n") + } + + /// What an upstream sends as soon as it takes the request is held; + /// the first content lets it all through. + #[test] + fn a_streams_opening_is_held_until_its_first_content() { + use serde_json::json; + + let start = frame( + "message_start", + json!({"type": "message_start", "message": {"id": "msg_1", "model": "m", + "usage": {"input_tokens": 3, "output_tokens": 1}}}), + ); + let ping = frame("ping", json!({"type": "ping"})); + let block = frame( + "content_block_start", + json!({"type": "content_block_start", "index": 0, + "content_block": {"type": "text", "text": ""}}), + ); + let mut o = Opening::new(Dialect::Anthropic); + assert!(!o.feed(start.as_bytes())); + assert!(!o.feed(ping.as_bytes())); + assert!(o.feed(block.as_bytes())); + assert_eq!(o.release(), [start, ping, block].concat().into_bytes()); + + let mut o = Opening::new(Dialect::Responses); + let created = json!({"response": {"id": "resp_1", "model": "m"}}); + assert!(!o.feed(frame("response.created", created.clone()).as_bytes())); + assert!(!o.feed(frame("response.in_progress", created).as_bytes())); + assert!( + o.feed( + frame( + "response.output_text.delta", + json!({"type": "response.output_text.delta", "output_index": 0, + "content_index": 0, "delta": "Hi"}), + ) + .as_bytes() + ) + ); + + let chunk = |delta: serde_json::Value| { + format!( + "data: {}\n\n", + json!({"id": "c", "model": "m", "choices": [{"index": 0, "delta": delta}]}) + ) + }; + let mut o = Opening::new(Dialect::Chat); + assert!(!o.feed(chunk(json!({"role": "assistant", "content": ""})).as_bytes())); + assert!(o.feed(chunk(json!({"content": "Hi"})).as_bytes())); + } + + /// An error is not content: whether it goes out again is decided + /// with the opening still held. Content ahead of it in the same piece + /// of bytes began the answer. + #[test] + fn an_error_in_the_opening_is_not_content() { + use serde_json::json; + + let start = frame( + "message_start", + json!({"type": "message_start", "message": {"id": "msg_1", "model": "m"}}), + ); + let overload = frame( + "error", + json!({"type": "error", + "error": {"type": "overloaded_error", "message": "Overloaded"}}), + ); + let text = frame( + "content_block_start", + json!({"type": "content_block_start", "index": 0, + "content_block": {"type": "text", "text": ""}}), + ); + let mut o = Opening::new(Dialect::Anthropic); + assert!(!o.feed([start.as_str(), overload.as_str()].concat().as_bytes())); + + let mut o = Opening::new(Dialect::Anthropic); + assert!(o.feed([start, text, overload].concat().as_bytes())); + } + + /// An opening that grows without any content is let through. + #[test] + fn an_opening_does_not_grow_without_end() { + let mut o = Opening::new(Dialect::Chat); + let comment = format!(": {}\n\n", "x".repeat(64 * 1024)); + let fed = (0..32) + .map(|_| o.feed(comment.as_bytes())) + .collect::>(); + assert!(!fed[0]); + assert!(fed.last().copied().unwrap_or(false)); + } + + /// Which errors an upstream reports in its stream are worth sending + /// again for: an overload or a server error, not a refusal. + #[test] + fn an_overload_or_a_server_error_in_a_stream_is_transient() { + let cases = [ + ( + Dialect::Anthropic, + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n", + true, + ), + ( + Dialect::Anthropic, + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"Internal\"}}\n\n", + true, + ), + ( + Dialect::Chat, + "data: {\"error\":{\"message\":\"boom\",\"type\":\"server_error\"}}\n\n", + true, + ), + ( + Dialect::Responses, + "event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"error\":{\"code\":\"server_is_overloaded\",\"message\":\"busy\"}}}\n\n", + true, + ), + ( + Dialect::Gemini, + "data: {\"error\":{\"code\":503,\"message\":\"busy\",\"status\":\"UNAVAILABLE\"}}\n\n", + true, + ), + ( + Dialect::Anthropic, + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"too long\"}}\n\n", + false, + ), + ( + Dialect::Anthropic, + "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down\"}}\n\n", + false, + ), + // Nothing recognisable: not sent again on a guess. + ( + Dialect::Chat, + "data: {\"error\":{\"message\":\"odd\",\"code\":\"upstream_wobble\"}}\n\n", + false, + ), + ]; + for (dialect, frames, transient) in cases { + let mut r = Readers { + errors: Some(ErrorWatch::new(dialect)), + ..Default::default() + }; + r.feed(frames.as_bytes()); + let said = r.reported(false).unwrap_or_else(|| panic!("{frames}")); + let StreamOutcome::UpstreamError { + error_type, + status_code, + .. + } = outcome_of("up", &said) + else { + panic!("not a failure"); + }; + assert_eq!( + crate::proxy::retry::transient(&error_type, status_code), + transient, + "{frames}" + ); + // Left in place for the tail. + assert_eq!(r.upstream_error(), Some(said)); + } + } + + /// A last error frame without a blank line after it is read once the + /// stream ends. + #[test] + fn an_error_in_the_last_frame_is_reported_at_the_end() { + let mut r = Readers { + errors: Some(ErrorWatch::new(Dialect::Chat)), + ..Default::default() + }; + r.feed(b"data: {\"error\":{\"message\":\"boom\",\"type\":\"server_error\"}}"); + assert_eq!(r.reported(false), None); + let said = r.reported(true).unwrap(); + assert_eq!(said.status, Some(500)); + } + + /// A Gemini client, SSE or JSON array, never gets a keep-alive comment. + #[test] + fn only_non_gemini_clients_get_keep_alives() { + assert!(!keeps_alive(Dialect::Gemini)); + for client in [Dialect::Chat, Dialect::Anthropic, Dialect::Responses] { + assert!(keeps_alive(client)); + } + } + + /// A retry clears what the failed attempt left for the tail, and says + /// the stream is waiting. + #[test] + fn deciding_a_retry_clears_the_failed_attempt() { + let readers = Mutex::new(Readers { + errors: Some(ErrorWatch::new(Dialect::Anthropic)), + sniffer: Some(tw_dialect::usage::Sniffer::new()), + ..Default::default() + }); + readers.lock().unwrap().feed( + b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"x\"}}\n\n", + ); + let sent = Mutex::new(Sent { + retries: 0, + last_sent: std::time::Instant::now(), + waiting: false, + effort: None, + }); + resending( + "up", + 1, + std::time::Duration::from_secs(2), + "x", + &readers, + &sent, + ); + let mut r = readers.lock().unwrap(); + assert!(r.sniffer.is_none() && r.upstream_error().is_none()); + let s = *sent.lock().unwrap(); + assert_eq!((s.retries, s.waiting), (1, true)); + } + + /// While the gateway waits, the client hears a keep-alive every 15 + /// seconds of silence, and the wait ends on time. + #[tokio::test(start_paused = true)] + async fn a_wait_sends_a_keep_alive_every_fifteen_seconds() { + let start = tokio::time::Instant::now(); + let quiet_since = start - std::time::Duration::from_secs(5); + let until = start + std::time::Duration::from_secs(40); + let ticks = keep_alives(quiet_since, until); + let mut ticks = std::pin::pin!(ticks); + let mut at = Vec::new(); + while ticks.next().await.is_some() { + at.push((tokio::time::Instant::now() - start).as_secs()); + } + assert_eq!(at, [10, 25]); + assert_eq!((tokio::time::Instant::now() - start).as_secs(), 40); + } + #[test] fn the_last_frame_is_seen_in_the_chunk_that_completes_it() { let mut w = LastFrame::new(Dialect::Responses, true); @@ -1366,6 +2017,24 @@ mod tests { assert!(r.sealed_refusal().is_none()); } + #[test] + fn a_changed_reasoning_effort_is_said_on_the_log_row() { + use crate::proxy::effort::{EffortChange, Reason}; + assert_eq!(effort_detail(None, None), None); + let change = EffortChange { + asked: "xhigh", + sent: Some("high"), + reason: Reason::Rejected, + }; + assert_eq!( + effort_detail(Some(serde_json::json!({"retries": 1})), Some(change)), + Some(serde_json::json!({ + "retries": 1, + "effort": {"asked": "xhigh", "sent": "high", "reason": "rejected"}, + })) + ); + } + #[test] fn a_pinned_route_is_said_on_the_log_row() { assert_eq!(pinned_detail(None, false), None); @@ -1451,7 +2120,7 @@ mod tests { /// the chunk's text, which can be the model's answer. #[test] fn an_error_that_names_no_status_is_the_upstreams_failure() { - let chat = "data: {\"choices\":[{\"delta\":{\"content\":\"the secret plan\"}}],\"error\":{\"code\":\"server_error\"}}\n\n"; + let chat = "data: {\"choices\":[{\"delta\":{\"content\":\"the secret plan\"}}],\"error\":{\"code\":\"upstream_wobble\"}}\n\n"; let (error_type, status, message, counted) = recorded(Dialect::Chat, chat).unwrap(); assert_eq!((error_type.as_str(), status), ("ProviderError", 502)); assert!(counted); @@ -1483,11 +2152,12 @@ mod tests { let gemini = "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hi\"}]}}],\"error\":null}\n\n"; assert_eq!(recorded(Dialect::Gemini, gemini), None); - // And an error after such chunks is still seen. + // And an error after such chunks is still seen: OpenAI's + // `server_error` is its 500. let then = format!( "{chat}data: {{\"error\":{{\"message\":\"boom\",\"type\":\"server_error\"}}}}\n\n" ); - assert_eq!(recorded(Dialect::Chat, &then).unwrap().1, 502); + assert_eq!(recorded(Dialect::Chat, &then).unwrap().1, 500); } /// Cut to 500 characters, as a refusal's body is. diff --git a/crates/gateway/src/proxy/cache_marks.rs b/crates/gateway/src/proxy/cache_marks.rs index e07e7b50..3ac5c265 100644 --- a/crates/gateway/src/proxy/cache_marks.rs +++ b/crates/gateway/src/proxy/cache_marks.rs @@ -73,6 +73,7 @@ pub fn is_refusal(err: &GatewayError, label: &str) -> bool { let GatewayError::ProviderHttpError { status: 400, message, + .. } = err else { return false; @@ -102,6 +103,7 @@ mod tests { GatewayError::ProviderHttpError { status, message: format!("relay: {said}"), + retry_after: None, } } @@ -134,6 +136,7 @@ mod tests { &GatewayError::ProviderHttpError { status: 400, message: "prompt caching proxy: max_tokens: Field required".into(), + retry_after: None, }, "prompt caching proxy" )); diff --git a/crates/gateway/src/proxy/effort.rs b/crates/gateway/src/proxy/effort.rs new file mode 100644 index 00000000..8721e75c --- /dev/null +++ b/crates/gateway/src/proxy/effort.rs @@ -0,0 +1,739 @@ +//! Reasoning effort: the model's setting, and what an upstream refuses. +//! +//! # The setting (`models.reasoning_effort`) +//! +//! `off` turns reasoning off in a request that has it on (OpenAI's formats +//! write `none`; the others take the switch out, with the effort). A level +//! (`low` … `max`) changes the effort of a request **whose reasoning is +//! on**; a request without it — a title, a summary, a client that never +//! asked — goes out untouched: reasoning is never turned on for a caller. +//! Unset: requests go out as sent. It applies to each hop as it is +//! addressed, for the model named upstream (Gemini's budgets are capped +//! per model). A level is applied to the caller's request in its own +//! format; a request converted for another format carries the changed +//! effort the way that format writes it. `off` on a converted hop is +//! applied to the conversion instead ([`turn_off_converted`]): taking the +//! caller's reasoning out first would leave the converted request with no +//! effort at all, which OpenAI's models read as their default level. What +//! reasoning is on, and where each format writes the effort, is +//! `tw_dialect::params`, shared with the desktop gateway. +//! +//! # What an upstream refuses +//! +//! Not every model takes every level, and an upstream that does not says +//! so with a 400 (`tw_dialect::think::effort_rejection`; for a `none` it +//! does not take, `none_fix`). Both are read against the effort the +//! hop **wrote** (`tw_dialect::params::written_effort`), whether or not +//! reasoning is on — Claude with thinking off still refuses `xhigh` — and a +//! request that only carries a budget has none to refuse. The request then +//! goes to the **same upstream** once more, with the level it named (for a +//! refused `none`, the lowest it named) or, when it named none, without the +//! effort (the model's default). That is +//! not a retry of a failure of the moment: it does not wait and does not +//! count against `gateway.retry_wait_secs` (`super::retry`). A second +//! refusal is the answer. +//! +//! A resend that is answered is remembered in Redis, per provider, format +//! sent and model named upstream, for [`MEMORY_TTL_SECS`]: later requests +//! that would write the refused level write the one that worked straight +//! away. A remembered fix that is refused itself is forgotten (and replaced +//! if the resend that follows is answered); so is one whose resend was +//! answered at the level first refused. After the week the refused level +//! is tried again — a model can gain one. A provider's entries go when its +//! URL or credential changes, a route of it relearns its protocol, or it +//! is deleted ([`forget_provider`]). +//! +//! What a hop sent differently from what it would have is said on the +//! request's `gateway_logs` row, as `"effort"` ([`EffortChange::detail`]). + +use std::collections::BTreeMap; + +use fred::interfaces::{HashesInterface, KeysInterface}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use tw_dialect::ir::{Dialect, Effort}; +use tw_dialect::params::WrittenEffort; +use tw_dialect::think::EffortFix; +use uuid::Uuid; + +use crate::error::GatewayError; + +/// A model's reasoning setting (`models.reasoning_effort`). Unset is no +/// setting at all: `Option`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ReasoningSetting { + /// Turn reasoning off in requests that have it on. + Off, + /// Change the effort of requests whose reasoning is on. + Level(Effort), +} + +/// The values `models.reasoning_effort` may hold, in display order. +pub const SETTINGS: [&str; 6] = ["off", "low", "medium", "high", "xhigh", "max"]; + +impl ReasoningSetting { + /// One of [`SETTINGS`]. `minimal` is not one: few models take it, and + /// the ones that do not refuse it outright. + pub fn parse(s: &str) -> Option { + Some(match s { + "off" => Self::Off, + "low" => Self::Level(Effort::Low), + "medium" => Self::Level(Effort::Medium), + "high" => Self::Level(Effort::High), + "xhigh" => Self::Level(Effort::XHigh), + "max" => Self::Level(Effort::Max), + _ => return None, + }) + } + + /// As stored: one of [`SETTINGS`]. + pub fn slug(self) -> &'static str { + match self { + Self::Off => "off", + Self::Level(e) => tw_dialect::think::openai(e), + } + } + + /// Apply to `body`, a request in `dialect` for `model` (the name sent + /// upstream). Whether anything changed. + pub(crate) fn apply(self, dialect: Dialect, body: &mut Value, model: &str) -> bool { + match self { + Self::Off => { + if !tw_dialect::params::reasoning_on(dialect, body) { + return false; + } + let before = body.clone(); + tw_dialect::params::disable_thinking(dialect, body); + *body != before + } + Self::Level(e) => tw_dialect::params::set_reasoning_effort(dialect, body, e, model), + } + } +} + +/// `off` on a hop converted for `target`, applied to the caller's decoded +/// request `r` before it is encoded. Whether reasoning was on and is +/// turned off; the caller's reasoning is read as for the same-format hop. +/// +/// The reasoning configuration is taken out, except for Gemini, where it is +/// marked off: Gemini 2.5 thinks without one, and the encoder writes a +/// zero budget for the models that can stop (2.5 Flash) and reports the +/// setting dropped for those that cannot. For OpenAI's formats the +/// converted body then needs `none` written in it +/// (`tw_dialect::params::write_reasoning_none`): without an effort their +/// models reason at their default level. +pub(crate) fn turn_off_converted(r: &mut tw_dialect::ir::Request, target: Dialect) -> bool { + if !r.reasoning.as_ref().is_some_and(|re| re.enabled) { + return false; + } + r.reasoning = (target == Dialect::Gemini).then_some(tw_dialect::ir::Reasoning { + enabled: false, + effort: None, + budget: None, + summary: false, + }); + true +} + +/// What the caller's request asks for reasoning, read once. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(crate) struct Asked { + /// The model's setting, applied to each hop. + pub setting: Option, + /// The effort the caller asked for, with reasoning on: what a hop the + /// setting changed is said against. + pub effort: Option, + /// A hop can carry an effort an upstream could refuse: the caller has + /// reasoning on or writes one. Otherwise no hop asks the memory. + pub any: bool, +} + +impl Asked { + pub(crate) fn of(setting: Option, dialect: Dialect, body: &Value) -> Self { + Self { + setting, + effort: tw_dialect::params::reasoning_effort(dialect, body), + any: tw_dialect::params::reasoning_on(dialect, body) + || tw_dialect::params::written_effort(dialect, body).is_some(), + } + } + + /// What applying the setting to a hop changed, said against what the + /// caller asked: `after` is the hop's body in `dialect` once applied. + /// Nothing to say when the caller named no effort or the hop sends the + /// same. + pub(crate) fn change(&self, dialect: Dialect, after: &Value) -> Option { + let asked = slug(self.effort?); + let sent = shown(dialect, after); + (sent != Some(asked)).then_some(EffortChange { + asked, + sent, + reason: Reason::Setting, + }) + } +} + +/// Why a hop sent another effort. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Reason { + /// The model's setting. + Setting, + /// The upstream refused what went out first, and the request went + /// out again. + Rejected, + /// The upstream refused that level before ([`Memory`]). + Remembered, +} + +/// A hop that sent another reasoning effort. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct EffortChange { + /// What the hop would have sent before this change: for the setting, + /// what the caller asked; for a refusal or the memory, the level it + /// replaced. `none` is OpenAI's reasoning turned off. + pub asked: &'static str, + /// What it sent; `None`: the effort left out. + pub sent: Option<&'static str>, + pub reason: Reason, +} + +impl EffortChange { + /// As said on the log row: `{"asked": "xhigh", "sent": "high", + /// "reason": "rejected"}`, `sent` `null` when the effort was left out. + pub(crate) fn detail(&self) -> Value { + serde_json::json!({ + "asked": self.asked, + "sent": self.sent, + "reason": match self.reason { + Reason::Setting => "setting", + Reason::Rejected => "rejected", + Reason::Remembered => "remembered", + }, + }) + } +} + +/// `minimal` … `max`, as OpenAI writes them. +fn slug(e: Effort) -> &'static str { + tw_dialect::think::openai(e) +} + +/// The effort `v`, in `dialect`, sends: the field written, else the one +/// its reasoning implies (a budget read as a level). +pub(crate) fn shown(dialect: Dialect, v: &Value) -> Option<&'static str> { + tw_dialect::params::written_effort(dialect, v) + .map(WrittenEffort::slug) + .or_else(|| tw_dialect::params::reasoning_effort(dialect, v).map(slug)) +} + +/// The effort written in `body`, a hop in `dialect`. +pub(crate) fn written(dialect: Dialect, body: &[u8]) -> Option { + let v: Value = serde_json::from_slice(body).ok()?; + tw_dialect::params::written_effort(dialect, &v) +} + +/// Rewrite the effort written in `body` as `fix` says. What the hop sends +/// then ([`shown`]) when anything changed; `None` leaves `body` as it was. +pub(crate) fn apply_fix( + dialect: Dialect, + body: &mut Vec, + fix: EffortFix, +) -> Option> { + let mut v: Value = serde_json::from_slice(body).ok()?; + let changed = match fix { + EffortFix::Use(e) => tw_dialect::params::replace_reasoning_effort(dialect, &mut v, e), + EffortFix::Drop => tw_dialect::params::remove_reasoning_effort(dialect, &mut v), + }; + if !changed { + return None; + } + let bytes = serde_json::to_vec(&v).ok()?; + *body = bytes; + Some(shown(dialect, &v)) +} + +/// Did the upstream refuse the effort `sent`? How to send it again. +/// `label` is the upstream's name, which the error's message starts with +/// and which is not what it said. +pub(crate) fn rejection(err: &GatewayError, label: &str, sent: WrittenEffort) -> Option { + let GatewayError::ProviderHttpError { + status, message, .. + } = err + else { + return None; + }; + let said = message + .strip_prefix(label) + .and_then(|m| m.strip_prefix(": ")) + .unwrap_or(message); + match sent { + WrittenEffort::Level(e) => tw_dialect::think::effort_rejection(*status, said, e), + WrittenEffort::Off => tw_dialect::think::none_fix(*status, said), + } +} + +// ───────────────────────────────────────────── memory + +/// How long a refused level is remembered: 7 days. +pub const MEMORY_TTL_SECS: i64 = 7 * 24 * 60 * 60; + +/// One provider's memory: a hash, a field per format sent and model named +/// upstream ([`field`]). Hash-tagged on the provider, as the route keys +/// are on the route. +fn key(provider_id: Uuid) -> String { + format!("effort_memory:{{provider:{provider_id}}}") +} + +/// `openai-chat:gpt-5`: the format first, which has no `:` in it. +fn field(dialect: Dialect, model: &str) -> String { + format!("{}:{model}", dialect.slug()) +} + +/// What one model, asked in one format, refused: the written effort refused +/// (`none` too) → what worked instead. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct Memory(BTreeMap); + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct Entry { + /// The level that was answered; `None`: the effort left out. + #[serde(rename = "use")] + instead: Option, + /// Unix seconds. + at: i64, +} + +impl Memory { + fn parse(value: &str) -> Self { + serde_json::from_str(value).unwrap_or_default() + } + + /// How to rewrite a hop that writes `written`, as of `now`. Entries + /// older than [`MEMORY_TTL_SECS`] are no longer taken. + pub(crate) fn instead_of(&self, written: WrittenEffort, now: i64) -> Option { + let entry = self.0.get(written.slug())?; + if now - entry.at >= MEMORY_TTL_SECS { + return None; + } + Some(match &entry.instead { + Some(s) => EffortFix::Use(tw_dialect::think::parse_openai(s).flatten()?), + None => EffortFix::Drop, + }) + } + + /// `rejected` was refused and the resend `fix` answered, at `now`. A + /// fix that is `rejected` itself (the upstream took it after all) is + /// forgotten instead. Expired entries go. + pub(crate) fn note(&mut self, rejected: WrittenEffort, fix: EffortFix, now: i64) { + self.0.retain(|_, e| now - e.at < MEMORY_TTL_SECS); + let instead = match fix { + EffortFix::Use(e) if WrittenEffort::Level(e) == rejected => { + self.forget(rejected); + return; + } + EffortFix::Use(e) => Some(slug(e).to_string()), + EffortFix::Drop => None, + }; + self.0 + .insert(rejected.slug().to_string(), Entry { instead, at: now }); + } + + /// The fix remembered for `rejected` was refused itself. + pub(crate) fn forget(&mut self, rejected: WrittenEffort) { + self.0.remove(rejected.slug()); + } + + pub(crate) fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +pub(crate) fn now() -> i64 { + chrono::Utc::now().timestamp() +} + +/// What `provider_id` refused for `model` asked in `dialect`. A Redis that +/// cannot be read remembers nothing: the request goes out as it would have. +pub(crate) async fn recall( + redis: &fred::clients::Client, + provider_id: Uuid, + dialect: Dialect, + model: &str, +) -> Memory { + match redis + .hget::, _, _>(key(provider_id), field(dialect, model)) + .await + { + Ok(Some(v)) => Memory::parse(&v), + Ok(None) => Memory::default(), + Err(e) => { + tracing::debug!(error = %e, "could not read the reasoning effort memory"); + Memory::default() + } + } +} + +/// Change what `provider_id` remembers for `model` in `dialect`, and write +/// it back: emptied, the field goes. Best-effort. +async fn update( + redis: &fred::clients::Client, + provider_id: Uuid, + dialect: Dialect, + model: &str, + change: impl FnOnce(&mut Memory), +) { + let mut memory = recall(redis, provider_id, dialect, model).await; + let before = memory.clone(); + change(&mut memory); + if memory == before { + return; + } + let (key, field) = (key(provider_id), field(dialect, model)); + let written = async { + if memory.is_empty() { + return redis.hdel::<(), _, _>(&key, &field).await; + } + let value = serde_json::to_string(&memory).unwrap_or_default(); + redis + .hset::<(), _, _>(&key, (field.as_str(), value)) + .await?; + redis.expire::<(), _>(&key, MEMORY_TTL_SECS, None).await + } + .await; + if let Err(e) = written { + tracing::warn!(error = %e, "could not write the reasoning effort memory"); + } +} + +/// `provider_id` refused `rejected` for `model` in `dialect`, and the +/// resend `fix` was answered. +pub(crate) async fn remember( + redis: &fred::clients::Client, + provider_id: Uuid, + dialect: Dialect, + model: &str, + rejected: WrittenEffort, + fix: EffortFix, +) { + update(redis, provider_id, dialect, model, |m| { + m.note(rejected, fix, now()) + }) + .await; +} + +/// The fix remembered for `rejected` was refused itself. +pub(crate) async fn forget( + redis: &fred::clients::Client, + provider_id: Uuid, + dialect: Dialect, + model: &str, + rejected: WrittenEffort, +) { + update(redis, provider_id, dialect, model, |m| m.forget(rejected)).await; +} + +/// Forget what `provider_id` refused: it may be another upstream now. +/// Best-effort. +pub async fn forget_provider(redis: &fred::clients::Client, provider_id: Uuid) { + if let Err(e) = redis.del::<(), _>(key(provider_id)).await { + tracing::warn!(error = %e, %provider_id, "could not clear the reasoning effort memory"); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + const L: fn(Effort) -> WrittenEffort = WrittenEffort::Level; + + #[test] + fn the_settings_are_off_and_five_levels() { + for s in SETTINGS { + let setting = ReasoningSetting::parse(s).unwrap(); + assert_eq!(setting.slug(), s); + } + assert_eq!(ReasoningSetting::parse("off"), Some(ReasoningSetting::Off)); + assert_eq!( + ReasoningSetting::parse("xhigh"), + Some(ReasoningSetting::Level(Effort::XHigh)) + ); + for bad in ["minimal", "none", "on", "", "MAX"] { + assert_eq!(ReasoningSetting::parse(bad), None, "{bad}"); + } + } + + /// The setting applied to `body` (`dialect`), as one hop sends it, and + /// what is said about it. + fn applied(setting: &str, dialect: Dialect, mut body: Value) -> (Value, Option) { + let setting = ReasoningSetting::parse(setting); + let asked = Asked::of(setting, dialect, &body); + if let Some(s) = setting { + s.apply(dialect, &mut body, "gpt-5"); + } + let change = asked.change(dialect, &body); + (body, change) + } + + #[test] + fn a_level_changes_only_a_request_with_reasoning_on() { + let (body, change) = applied( + "max", + Dialect::Chat, + json!({"model": "m", "reasoning_effort": "xhigh", "messages": []}), + ); + assert_eq!(body["reasoning_effort"], "max"); + assert_eq!( + change.unwrap().detail(), + json!({"asked": "xhigh", "sent": "max", "reason": "setting"}) + ); + + // No reasoning: a title request goes out as sent. + let title = json!({"model": "m", "messages": []}); + assert_eq!(applied("max", Dialect::Chat, title.clone()), (title, None)); + + // Reasoning turned off by the caller stays off. + let none = json!({"model": "m", "reasoning_effort": "none"}); + assert_eq!(applied("max", Dialect::Chat, none.clone()), (none, None)); + + // Already at that level: nothing to say. + let same = json!({"model": "m", "reasoning": {"effort": "max"}, "input": "hi"}); + assert_eq!( + applied("max", Dialect::Responses, same.clone()), + (same, None) + ); + } + + #[test] + fn a_level_is_written_where_the_callers_format_puts_it() { + // Anthropic adaptive thinking, no effort named: one is added, and + // there was nothing asked to say it against. + let (body, change) = applied( + "high", + Dialect::Anthropic, + json!({"model": "m", "max_tokens": 64000, + "thinking": {"type": "adaptive"}, "messages": []}), + ); + assert_eq!(body["output_config"]["effort"], "high"); + assert_eq!(change, None); + + let (body, _) = applied( + "high", + Dialect::Responses, + json!({"model": "m", "reasoning": {"effort": "low"}, "input": "hi"}), + ); + assert_eq!(body["reasoning"]["effort"], "high"); + } + + #[test] + fn off_turns_reasoning_off_as_each_format_says_it() { + let (claude, change) = applied( + "off", + Dialect::Anthropic, + json!({"model": "m", "max_tokens": 64000, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "max"}, "messages": []}), + ); + assert!(claude.get("thinking").is_none(), "{claude}"); + assert!(claude.get("output_config").is_none(), "{claude}"); + assert_eq!( + change.unwrap().detail(), + json!({"asked": "max", "sent": null, "reason": "setting"}) + ); + + // OpenAI's formats say `none`: left out, the model's default applies. + let (chat, change) = applied( + "off", + Dialect::Chat, + json!({"model": "m", "reasoning_effort": "high"}), + ); + assert_eq!(chat["reasoning_effort"], "none"); + assert_eq!( + change.unwrap().detail(), + json!({"asked": "high", "sent": "none", "reason": "setting"}) + ); + + // Nothing to turn off. + let plain = json!({"model": "m", "messages": []}); + assert_eq!(applied("off", Dialect::Chat, plain.clone()), (plain, None)); + } + + #[test] + fn a_fix_rewrites_the_effort_as_written() { + let mut body = + serde_json::to_vec(&json!({"model": "m", "reasoning_effort": "xhigh"})).unwrap(); + assert_eq!(written(Dialect::Chat, &body), Some(L(Effort::XHigh))); + assert_eq!( + apply_fix(Dialect::Chat, &mut body, EffortFix::Use(Effort::High)), + Some(Some("high")) + ); + assert_eq!(written(Dialect::Chat, &body), Some(L(Effort::High))); + assert_eq!( + apply_fix(Dialect::Chat, &mut body, EffortFix::Drop), + Some(None) + ); + assert_eq!(written(Dialect::Chat, &body), None); + let before = body.clone(); + assert_eq!(apply_fix(Dialect::Chat, &mut body, EffortFix::Drop), None); + assert_eq!(body, before); + + // Written while thinking is off: still the field the upstream reads. + let mut off = serde_json::to_vec(&json!({"model": "m", "max_tokens": 9, + "thinking": {"type": "disabled"}, "output_config": {"effort": "xhigh"}})) + .unwrap(); + assert_eq!(written(Dialect::Anthropic, &off), Some(L(Effort::XHigh))); + assert_eq!( + apply_fix(Dialect::Anthropic, &mut off, EffortFix::Use(Effort::High)), + Some(Some("high")) + ); + + // A budget only: no written effort. + let budget = serde_json::to_vec(&json!({"model": "m", "max_tokens": 64000, + "thinking": {"type": "enabled", "budget_tokens": 32000}})) + .unwrap(); + assert_eq!(written(Dialect::Anthropic, &budget), None); + } + + fn http(status: u16, said: &str) -> GatewayError { + GatewayError::ProviderHttpError { + status, + message: format!("relay: {said}"), + retry_after: None, + } + } + + #[test] + fn a_refusal_of_the_effort_is_read_off_the_upstreams_words() { + let said = r#"{"error":{"message":"Unsupported value: 'xhigh' is not supported with this model. Supported values are: 'low', 'medium', and 'high'.","param":"reasoning_effort"}}"#; + assert_eq!( + rejection(&http(400, said), "relay", L(Effort::XHigh)), + Some(EffortFix::Use(Effort::High)) + ); + // Claude with thinking off: a ceiling. + let claude = r#"{"type":"error","error":{"type":"invalid_request_error","message":"effort 'xhigh' is not supported when thinking is disabled. Use effort 'high' or below."}}"#; + assert_eq!( + rejection(&http(400, claude), "relay", L(Effort::XHigh)), + Some(EffortFix::Use(Effort::High)) + ); + // `none` not taken, no level named: left out. + let none = r#"{"error":{"message":"Invalid value for reasoning_effort: 'none'"}}"#; + assert_eq!( + rejection(&http(400, none), "relay", WrittenEffort::Off), + Some(EffortFix::Drop) + ); + // `none` not taken, levels named: the lowest, not the model's + // default, which would reason more. + let listed = r#"{"error":{"message":"Unsupported value: 'none' is not supported with this model. Supported values are: 'minimal', 'low', 'medium', and 'high'.","param":"reasoning_effort"}}"#; + assert_eq!( + rejection(&http(400, listed), "relay", WrittenEffort::Off), + Some(EffortFix::Use(Effort::Minimal)) + ); + let listed = r#"{"error":{"message":"Unsupported value: 'none' is not supported with this model. Supported values are: 'low', 'medium', and 'high'.","param":"reasoning.effort"}}"#; + assert_eq!( + rejection(&http(400, listed), "relay", WrittenEffort::Off), + Some(EffortFix::Use(Effort::Low)) + ); + // Not a 400: an upstream failing is not a refusal. + assert_eq!(rejection(&http(500, said), "relay", L(Effort::XHigh)), None); + // A refusal of something else. + assert_eq!( + rejection( + &http( + 400, + r#"{"message":"text.verbosity must be one of low, medium, high"}"# + ), + "relay", + L(Effort::XHigh) + ), + None + ); + // The upstream's own name is not what it said. + assert_eq!( + rejection( + &GatewayError::ProviderHttpError { + status: 400, + message: "high effort reasoning relay: messages: Field required".into(), + retry_after: None, + }, + "high effort reasoning relay", + L(Effort::High) + ), + None + ); + assert_eq!( + rejection( + &GatewayError::UpstreamAuthError { + status: 403, + message: "reasoning effort".into() + }, + "relay", + L(Effort::High) + ), + None + ); + } + + #[test] + fn the_memory_keeps_each_refused_level_for_a_week() { + let t0 = 1_800_000_000; + let mut m = Memory::default(); + m.note(L(Effort::XHigh), EffortFix::Use(Effort::High), t0); + m.note(WrittenEffort::Off, EffortFix::Drop, t0 + 10); + assert_eq!( + m.instead_of(L(Effort::XHigh), t0 + 60), + Some(EffortFix::Use(Effort::High)) + ); + assert_eq!( + m.instead_of(WrittenEffort::Off, t0 + 60), + Some(EffortFix::Drop) + ); + assert_eq!(m.instead_of(L(Effort::High), t0 + 60), None); + // A week on, the refused level is tried again. + assert_eq!(m.instead_of(L(Effort::XHigh), t0 + MEMORY_TTL_SECS), None); + + // As stored, and read back. + let stored = serde_json::to_string(&m).unwrap(); + assert_eq!( + serde_json::from_str::(&stored).unwrap(), + json!({"xhigh": {"use": "high", "at": t0}, "none": {"use": null, "at": t0 + 10}}) + ); + assert_eq!(Memory::parse(&stored), m); + // Anything else reads as nothing remembered. + assert_eq!(Memory::parse("not json"), Memory::default()); + + // A fix refused itself is forgotten; replaced when another works. + m.forget(L(Effort::XHigh)); + assert_eq!(m.instead_of(L(Effort::XHigh), t0 + 60), None); + m.note(L(Effort::XHigh), EffortFix::Use(Effort::Medium), t0 + 20); + assert_eq!( + m.instead_of(L(Effort::XHigh), t0 + 60), + Some(EffortFix::Use(Effort::Medium)) + ); + // Answered at the level first refused: nothing to remember. + m.note(L(Effort::XHigh), EffortFix::Use(Effort::XHigh), t0 + 30); + assert_eq!(m.instead_of(L(Effort::XHigh), t0 + 60), None); + + // Noting drops what expired. + m.note( + L(Effort::Low), + EffortFix::Use(Effort::Medium), + t0 + MEMORY_TTL_SECS + 20, + ); + assert_eq!(m.0.len(), 1, "{m:?}"); + } + + #[test] + fn the_memory_is_keyed_on_the_provider_the_format_and_the_model() { + let id = Uuid::nil(); + assert_eq!( + key(id), + "effort_memory:{provider:00000000-0000-0000-0000-000000000000}" + ); + assert_eq!(field(Dialect::Chat, "gpt-5"), "openai-chat:gpt-5"); + assert_ne!( + field(Dialect::Anthropic, "m"), + field(Dialect::Responses, "m") + ); + } +} diff --git a/crates/gateway/src/proxy/generate.rs b/crates/gateway/src/proxy/generate.rs index b679b3ad..2c1c75b0 100644 --- a/crates/gateway/src/proxy/generate.rs +++ b/crates/gateway/src/proxy/generate.rs @@ -46,18 +46,22 @@ use tw_dialect::ir::{Dialect, Target}; use super::body_capture::prepare_body_capture; use super::early_cancel::EarlyCancelSlot; +use super::effort::{self, EffortChange, Reason}; use super::headers::{ request_id_header, resolve_conversation_id, resolve_session_id, resolve_trace_id, }; use super::log_ctx::{LogCtx, emit_gateway_error_log, emit_gateway_log}; use super::pipeline::{launch_stream_pump, run_buffered_post_invoke, run_preflight_stages}; +use super::retry::Budget; use super::routing::{ - build_selection_ctx, finalize_health, select_route_for_stream, select_route_with_failover, - set_affinity, + Resend, build_selection_ctx, finalize_health, select_route_for_stream, + select_route_with_failover, set_affinity, }; use super::sealed::Conversation; use super::shaper::{StreamShaper, rewrite_model}; use super::{GatewayErrorResponse, GatewayRequestIdentity, GatewayState}; +use tw_dialect::params::WrittenEffort; +use tw_dialect::think::EffortFix; use crate::cache::ResponseCache; use crate::guards::Caller; @@ -113,6 +117,7 @@ pub async fn proxy_chat_completion( headers: HeaderMap, axum::Extension(identity): axum::Extension, cancel: Option>, + arrived: Option>, body: Bytes, ) -> Result { generate( @@ -120,6 +125,7 @@ pub async fn proxy_chat_completion( headers, identity, cancel.map(|c| c.0), + arrived.map(|a| a.0.0), body, CHAT, "/v1/chat/completions", @@ -134,6 +140,7 @@ pub async fn proxy_anthropic_messages( headers: HeaderMap, axum::Extension(identity): axum::Extension, cancel: Option>, + arrived: Option>, body: Bytes, ) -> Result { generate( @@ -141,6 +148,7 @@ pub async fn proxy_anthropic_messages( headers, identity, cancel.map(|c| c.0), + arrived.map(|a| a.0.0), body, MESSAGES, "/v1/messages", @@ -155,6 +163,7 @@ pub async fn proxy_responses( headers: HeaderMap, axum::Extension(identity): axum::Extension, cancel: Option>, + arrived: Option>, body: Bytes, ) -> Result { generate( @@ -162,6 +171,7 @@ pub async fn proxy_responses( headers, identity, cancel.map(|c| c.0), + arrived.map(|a| a.0.0), body, RESPONSES, "/v1/responses", @@ -180,6 +190,7 @@ pub async fn proxy_gemini( headers: HeaderMap, axum::Extension(identity): axum::Extension, cancel: Option>, + arrived: Option>, body: Bytes, ) -> Result { generate( @@ -187,6 +198,7 @@ pub async fn proxy_gemini( headers, identity, cancel.map(|c| c.0), + arrived.map(|a| a.0.0), body, GEMINI, uri.path(), @@ -240,6 +252,9 @@ pub(crate) struct Outbound { /// The model's output cap (`models.max_output_tokens`), applied to each /// hop as it is addressed (see [`Outbound::cap_output`]). pub max_output_tokens: Option, + /// The request's reasoning: the model's setting, applied to each hop + /// as it is addressed, and what the caller asked (see `super::effort`). + pub reasoning: effort::Asked, } /// The request as it goes out to one upstream, and what it takes to read @@ -266,6 +281,14 @@ pub(crate) struct Wire { /// the caller marked none. An upstream that refuses them gets the /// request again without them (see `proxy::cache_marks`). pub auto_cache: bool, + /// This hop sends another reasoning effort than it would have, and why + /// (see `super::effort`). + pub effort: Option, + /// The effort written before the memory rewrote it: what a refusal of + /// the rewritten one is remembered under. `None`: not rewritten. + pub pre: Option, + /// The memory rewrote the effort. + pub remembered: bool, } impl Wire { @@ -284,6 +307,41 @@ impl Wire { } } + /// The reasoning effort this hop writes, as its own format writes it. + fn written_effort(&self) -> Option { + effort::written(self.dialect, &self.body) + } + + /// Rewrite the written effort as `fix` says, in place of `replaced`. + /// False when nothing changed. + fn fix_effort(&mut self, fix: EffortFix, replaced: WrittenEffort, reason: Reason) -> bool { + let Some(sent) = effort::apply_fix(self.dialect, &mut self.body, fix) else { + return false; + }; + self.effort = Some(EffortChange { + asked: replaced.slug(), + sent, + reason, + }); + true + } + + /// This hop with the effort `memory` says the upstream takes instead + /// of the one it writes. + fn remembered(mut self, memory: &effort::Memory) -> Wire { + if memory.is_empty() { + return self; + } + if let Some(written) = self.written_effort() + && let Some(fix) = memory.instead_of(written, effort::now()) + && self.fix_effort(fix, written, Reason::Remembered) + { + self.pre = Some(written); + self.remembered = true; + } + self + } + async fn send_to( &self, upstream: &super::transport::Upstream, @@ -405,10 +463,29 @@ impl Outbound { .map_err(|r| GatewayError::TransformError(r.0)) }; - if protocol.dialect() == client { + // The model's reasoning setting, for the model named upstream, on + // the caller's request in its own format — except `off` on a hop + // converted for another format, which is applied to the conversion + // (see `effort::turn_off_converted`). + let converted = protocol.dialect() != client; + let off_converted = + converted && self.reasoning.setting == Some(effort::ReasoningSetting::Off); + let mut source = std::borrow::Cow::Borrowed(&self.body); + let mut setting_change = None; + let mut level_applied = false; + if let Some(setting) = self.reasoning.setting.filter(|_| !off_converted) { + let mut body = self.body.clone(); + if setting.apply(client, &mut body, model) { + setting_change = self.reasoning.change(client, &body); + level_applied = true; + source = std::borrow::Cow::Owned(body); + } + } + + if !converted { // Forwarded as sent. Only the model changes, and a Chat // stream always asks for its usage (see `hides_usage`). - let mut body = self.body.clone(); + let mut body = source.into_owned(); let (path, query) = if client == Dialect::Gemini { // Gemini names the model in the path, and is always asked // for SSE (see `GEMINI_SSE`). @@ -464,11 +541,16 @@ impl Outbound { collect, ledger, auto_cache: false, + effort: setting_change, + pre: None, + remembered: false, }); } - let mut decoded = decode(&self.body)?; + let mut decoded = decode(&source)?; decoded.request.model = model.to_string(); + let turned_off = + off_converted && effort::turn_off_converted(&mut decoded.request, protocol.dialect()); let prepared = decoded.encode(&target(protocol.dialect())); if !prepared.dropped.is_empty() { tracing::info!( @@ -479,10 +561,17 @@ impl Outbound { ); } let mut body = prepared.body; - if self.max_output_tokens.is_some() + if (self.max_output_tokens.is_some() || turned_off || level_applied) && let Ok(mut v) = serde_json::from_slice::(&body) { self.cap_output(protocol.dialect(), &mut v, model, official); + if turned_off { + tw_dialect::params::write_reasoning_none(protocol.dialect(), &mut v); + } + // Said against what the converted request really carries. + if turned_off || level_applied { + setting_change = self.reasoning.change(protocol.dialect(), &v); + } body = serde_json::to_vec(&v).unwrap_or(body); } // The conversion moved the placeholders along with the text; one it @@ -505,6 +594,9 @@ impl Outbound { collect: prepared.session, ledger, auto_cache, + effort: setting_change, + pre: None, + remembered: false, }) } } @@ -516,15 +608,22 @@ impl Outbound { /// added gets the request once more without them, and is remembered as /// refusing them for this model (see `proxy::cache_marks`). /// -/// A rejected dialect or breakpoint is known before a single byte of body -/// arrives — the status check happens inside `send` — so a stream needs -/// no special handling: nothing has reached the client yet when the retry -/// happens. +/// An upstream that refuses the reasoning effort gets the request once +/// more with the one it named, or without one, and is remembered as +/// refusing it for this model when that is answered (see `proxy::effort`). +/// A level it refused before goes out as what worked then. +/// +/// A rejected dialect, breakpoint or effort is known before a single byte +/// of body arrives — the status check happens inside `send` — so a stream +/// needs no special handling: nothing has reached the client yet when the +/// retry happens. None of these count as a retry (`proxy::retry`): the +/// upstream did not fail, it refused something that can be left out. pub(crate) async fn send( entry: &RouteEntry, outbound: &Outbound, call_ctx: &CallCtx, db: &sqlx::PgPool, + redis: &fred::clients::Client, model: &str, ) -> Result<(reqwest::Response, Wire), GatewayError> { let upstream = &entry.upstream; @@ -532,25 +631,86 @@ pub(crate) async fn send( let first = outbound .address(entry.protocol, model, official)? .for_upstream(upstream, model); + // Only a request that can carry an effort asks the memory. + let memory = if outbound.reasoning.any { + effort::recall(redis, entry.provider_id, first.dialect, model).await + } else { + effort::Memory::default() + }; + let first = first.remembered(&memory); + if first.remembered { + tracing::info!( + provider = %entry.provider_name, + model, + effort = %first.effort.map(|c| c.detail()).unwrap_or_default(), + "Sending the reasoning effort the upstream took before instead of one it refused" + ); + } // The format of the last request sent, to read its refusal in. let mut wire_dialect = first.dialect; let mut last = match first.send_to(upstream, call_ctx).await { Ok(resp) => return Ok((resp, first)), Err(e) => e, }; + let mut wire = first; - if first.auto_cache && super::cache_marks::is_refusal(&last, &upstream.label) { - let mut wire = first; - if wire.drop_added_marks() { - upstream.cache_marks.note(model); + if wire.auto_cache + && super::cache_marks::is_refusal(&last, &upstream.label) + && wire.drop_added_marks() + { + upstream.cache_marks.note(model); + tracing::info!( + provider = %entry.provider_name, + model, + "Upstream refused the cache breakpoints added on conversion — sending again without them" + ); + match wire.send_to(upstream, call_ctx).await { + Ok(resp) => return Ok((resp, wire)), + Err(e) => last = e, + } + } + + // Read only off a 400: the body is parsed for it. A budget is no + // written effort, and never refused as one. + if matches!(last, GatewayError::ProviderHttpError { status: 400, .. }) + && let Some(sent) = wire.written_effort() + && let Some(fix) = effort::rejection(&last, &upstream.label, sent) + { + // What is remembered is keyed on what was written before the + // memory rewrote it. + let pre = wire.pre.unwrap_or(sent); + let remembered = wire.remembered; + if wire.fix_effort(fix, pre, Reason::Rejected) { tracing::info!( provider = %entry.provider_name, model, - "Upstream refused the cache breakpoints added on conversion — sending again without them" + effort = %wire.effort.map(|c| c.detail()).unwrap_or_default(), + "Upstream refused the reasoning effort — sending again with one it takes" ); - match wire.send_to(upstream, call_ctx).await { - Ok(resp) => return Ok((resp, wire)), - Err(e) => last = e, + let resent = wire.send_to(upstream, call_ctx).await; + // Answered: remembered, in place of a remembered fix that was + // refused itself. Refused again: that fix is forgotten. + let (redis, provider_id, dialect, model) = ( + redis.clone(), + entry.provider_id, + wire.dialect, + model.to_string(), + ); + match resent { + Ok(resp) => { + tokio::spawn(async move { + effort::remember(&redis, provider_id, dialect, &model, pre, fix).await; + }); + return Ok((resp, wire)); + } + Err(e) => { + if remembered { + tokio::spawn(async move { + effort::forget(&redis, provider_id, dialect, &model, pre).await; + }); + } + last = e; + } } } } @@ -567,10 +727,18 @@ pub(crate) async fn send( let wire = outbound .address(*protocol, model, official)? .for_upstream(upstream, model); + let wire = if outbound.reasoning.any { + let memory = effort::recall(redis, entry.provider_id, wire.dialect, model).await; + wire.remembered(&memory) + } else { + wire + }; wire_dialect = wire.dialect; match wire.send_to(upstream, call_ctx).await { Ok(resp) => { super::protocol_relearn::persist(db, entry.route_id, *protocol).await; + // What it refused was said in another format. + effort::forget_provider(redis, entry.provider_id).await; return Ok((resp, wire)); } Err(e) => { @@ -640,21 +808,27 @@ pub(crate) async fn read_whole( /// /// `cancel` is the middleware's record of a request whose client leaves /// before the response exists (see `early_cancel`); `None` on a -/// WebSocket turn, whose connection records its own end. +/// WebSocket turn, whose connection records its own end. `arrived` is +/// when the gateway port's request timeout started counting (see +/// [`super::Arrived`]); `None` on a WebSocket turn, which it does not +/// bound. #[allow(clippy::too_many_arguments)] pub(crate) async fn generate( state: GatewayState, headers: HeaderMap, identity: GatewayRequestIdentity, cancel: Option, + arrived: Option, body: Bytes, surface: ClientSurface, path: &str, query: Option<&str>, ) -> Result { - run(state, headers, identity, cancel, body, surface, path, query) - .await - .map_err(|e| e.in_dialect(surface.dialect)) + run( + state, headers, identity, cancel, arrived, body, surface, path, query, + ) + .await + .map_err(|e| e.in_dialect(surface.dialect)) } #[allow(clippy::too_many_arguments)] @@ -663,6 +837,7 @@ async fn run( headers: HeaderMap, identity: GatewayRequestIdentity, cancel: Option, + arrived: Option, body: Bytes, surface: ClientSurface, path: &str, @@ -807,15 +982,18 @@ async fn run( })? }; - // 6. The model's output cap. Applied to each hop as it is addressed - // (`Outbound::cap_output`): whether to fill one in depends on the - // upstream model, and the field on the upstream's format. The - // upstream stops there by itself; the answer is not measured. - let max_output_tokens = state - .router - .load() - .config_for(&mapped_model) - .max_output_tokens; + // 6. The model's settings. The output cap is applied to each hop as it + // is addressed (`Outbound::cap_output`): whether to fill one in + // depends on the upstream model, and the field on the upstream's + // format. The upstream stops there by itself; the answer is not + // measured. The reasoning setting is applied to each hop too, for + // the model named upstream, on the caller's request in its own format + // (see `super::effort`): a hop converted for another format carries + // the changed effort the way that one writes it. `off` on such a hop + // is applied to the conversion. + let model_config = state.router.load().config_for(&mapped_model); + let max_output_tokens = model_config.max_output_tokens; + let reasoning = effort::Asked::of(model_config.reasoning, surface.dialect, &outbound_body); let call_ctx = CallCtx::new( Some(trace_id.clone()), @@ -829,11 +1007,15 @@ async fn run( // limit. let cache_fingerprint = if surface.caches { ResponseCache::fingerprint(&outbound_body).map(|mut fp| { - // The cap is applied after this, per hop: an answer made under - // one cap must not be served under another. + // The cap and the reasoning setting are applied after this, per + // hop: an answer made under one must not be served under + // another. if let Some(cap) = max_output_tokens { fp.extend_from_slice(format!("\nmax_output_tokens={cap}").as_bytes()); } + if let Some(setting) = reasoning.setting { + fp.extend_from_slice(format!("\nreasoning_effort={}", setting.slug()).as_bytes()); + } fp }) } else { @@ -1025,6 +1207,7 @@ async fn run( redaction: guards.redaction.clone(), ledger: ledger.clone(), max_output_tokens, + reasoning, }; let snapshot = |route: &RouteEntry, sel_record| crate::lifecycle::ChatPostInvokeDeps { state: state.clone(), @@ -1054,14 +1237,21 @@ async fn run( cache_enabled: surface.caches, }; + // How long after it arrived the request may still go out again to + // upstreams that failed for the moment (see `super::retry`). + let retry_wait = std::time::Duration::from_secs(state.dynamic_config.retry_wait_secs().await); + if outbound.stream { - // One pick, no retry once bytes have gone to the client. A dialect - // rejection still gets its retry: it arrives before any body. + // One pick, no failover once bytes have gone to the client. A + // dialect rejection still gets its retry: it arrives before any + // body. A transient failure before any of the answer went out + // sends to the same route again (see `build_chat_pump`). let sel_ctx = build_selection_ctx(&state, &mapped_model, identity.user_id.as_deref(), pin).await; let (entry, sel_record) = select_route_for_stream(routes, &sel_ctx) .await .map_err(|e| GatewayErrorResponse::from(ctx.emit(e)))?; + let resend = Resend::new(entry, &sel_ctx, Budget::new(request_started_at, retry_wait)); set_affinity( &state.redis, @@ -1075,16 +1265,30 @@ async fn run( let hide_usage = outbound.hides_usage(); // Started on the stream's first poll — see `build_chat_pump` for - // why it must not be awaited here. + // why it must not be awaited here — and again for each retry. let open: crate::lifecycle::OpenUpstream = { let entry = entry.clone(); + let outbound = std::sync::Arc::new(outbound); let call_ctx = call_ctx.clone(); let db = state.db.clone(); + let redis = state.redis.clone(); let model = entry .upstream_model .clone() .unwrap_or_else(|| mapped_model.clone()); - Box::pin(async move { send(&entry, &outbound, &call_ctx, &db, &model).await }) + Box::new(move || { + let (entry, outbound, call_ctx, db, redis, model) = ( + entry.clone(), + outbound.clone(), + call_ctx.clone(), + db.clone(), + redis.clone(), + model.clone(), + ); + Box::pin( + async move { send(&entry, &outbound, &call_ctx, &db, &redis, &model).await }, + ) + }) }; let deps = snapshot(entry, sel_record); @@ -1094,6 +1298,7 @@ async fn run( return Ok(launch_stream_pump( deps, open, + resend, shaper, surface.dialect, client_sse, @@ -1114,28 +1319,44 @@ async fn run( None, ) .await; - let (entry, completed, answer_ledger, sel_record) = - select_route_with_failover(routes, &outbound, &call_ctx, &sel_ctx, &mapped_model) - .await - .map_err(|e| { - // Every candidate failed — no winning provider to name. - emit_gateway_error_log( - &state.audit, - &metadata.request_id, - session_id.as_deref(), - identity.user_id.as_deref(), - identity.user_email.as_deref(), - identity.api_key_id.as_deref(), - identity.api_key_lineage_id.as_deref(), - identity.ip_address.as_deref(), - &mapped_model, - None, - request_started_at.elapsed().as_millis() as i64, - &e, - error_capture, - ); - GatewayErrorResponse::from(e) - })?; + // A whole answer must be back within the gateway port's request + // timeout, which started counting before this handler did (auth, the + // body): past it the client gets an empty 408 and the request no log + // row. Each attempt is cut a little before it, and a retry starts only + // when it fits (see `super::retry`). + let mut budget = Budget::new(request_started_at, retry_wait); + if let Some(arrived) = arrived { + budget = budget.within(super::retry::hard_deadline(arrived)); + } + let (entry, completed, answer_ledger, sel_record) = select_route_with_failover( + routes, + &outbound, + &call_ctx, + &sel_ctx, + &mapped_model, + &mut budget, + ) + .await + .map_err(|e| { + // Every candidate failed — no winning provider to name. + emit_gateway_error_log( + &state.audit, + &metadata.request_id, + session_id.as_deref(), + identity.user_id.as_deref(), + identity.user_email.as_deref(), + identity.api_key_id.as_deref(), + identity.api_key_lineage_id.as_deref(), + identity.ip_address.as_deref(), + &mapped_model, + None, + request_started_at.elapsed().as_millis() as i64, + &e, + budget.retries(), + error_capture, + ); + GatewayErrorResponse::from(e) + })?; // The answer as the caller will receive it: their values restored, // then its text through the content rules that check answers. @@ -1267,6 +1488,7 @@ fn json_response(body: Vec) -> axum::response::Response { #[cfg(test)] mod tests { use super::*; + use tw_dialect::ir::Effort; #[test] fn a_gemini_path_names_the_model_and_whether_to_stream() { @@ -1299,9 +1521,296 @@ mod tests { redaction, ledger, max_output_tokens: cap, + reasoning: effort::Asked::default(), + } + } + + /// `body` from a `surface` caller, under the reasoning `setting`. + fn with_setting(surface: ClientSurface, path: &str, body: Value, setting: &str) -> Outbound { + let setting = effort::ReasoningSetting::parse(setting); + let reasoning = effort::Asked::of(setting, surface.dialect, &body); + Outbound { + reasoning, + ..outbound(surface, path, body, None) } } + fn sent(wire: &Wire) -> Value { + serde_json::from_slice(&wire.body).unwrap() + } + + #[test] + fn the_reasoning_setting_goes_out_on_every_hop_in_its_format() { + let ask = serde_json::json!({ + "model": "gpt-5", "reasoning_effort": "xhigh", + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(CHAT, "/v1/chat/completions", ask.clone(), "max"); + // The caller's request stays as sent: the setting is per hop. + assert_eq!(outbound.body, ask); + // As sent, to a Chat upstream. + let wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap(); + assert_eq!(sent(&wire)["reasoning_effort"], "max"); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "xhigh", "sent": "max", "reason": "setting"})) + ); + // Converted, to a Responses upstream. + let wire = outbound + .address(UpstreamProtocol::OpenAiResponses, "gpt-5", false) + .unwrap(); + assert_eq!(sent(&wire)["reasoning"]["effort"], "max"); + + // Off, for OpenAI's formats: `none`. + let ask = serde_json::json!({ + "model": "gpt-5", "reasoning": {"effort": "high"}, "input": "ping" + }); + let outbound = with_setting(RESPONSES, "/v1/responses", ask, "off"); + let wire = outbound + .address(UpstreamProtocol::OpenAiResponses, "gpt-5", false) + .unwrap(); + assert_eq!(sent(&wire)["reasoning"]["effort"], "none"); + + // A request without reasoning goes out as without the setting. + let ask = serde_json::json!({"model": "gpt-5", "messages": []}); + let outbound = with_setting(CHAT, "/v1/chat/completions", ask, "max"); + assert!(!outbound.reasoning.any); + let wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap(); + assert_eq!(wire.effort, None); + assert!(sent(&wire).get("reasoning_effort").is_none()); + } + + /// `off` on a hop converted for another format: the converted request + /// says it the way the upstream's format does, and the log says what + /// it carries. + #[test] + fn off_reaches_a_hop_converted_for_another_format() { + let claude = serde_json::json!({ + "model": "gpt-5", "max_tokens": 2000, + "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}, + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(MESSAGES, "/v1/messages", claude.clone(), "off"); + let wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap(); + assert_eq!(sent(&wire)["reasoning_effort"], "none"); + assert_eq!(wire.written_effort(), Some(WrittenEffort::Off)); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "high", "sent": "none", "reason": "setting"})) + ); + let wire = outbound + .address(UpstreamProtocol::OpenAiResponses, "gpt-5", false) + .unwrap(); + assert_eq!(sent(&wire)["reasoning"]["effort"], "none"); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "high", "sent": "none", "reason": "setting"})) + ); + // The same format, as before: the switch and the effort taken out. + let wire = outbound + .address(UpstreamProtocol::AnthropicMessages, "gpt-5", false) + .unwrap(); + assert!(sent(&wire).get("thinking").is_none(), "{}", sent(&wire)); + assert!(sent(&wire).get("output_config").is_none()); + + // Adaptive with no effort named converts to no effort at all: off + // still says `none`, with no asked level to say it against. + let mut adaptive = claude.clone(); + adaptive.as_object_mut().unwrap().remove("output_config"); + let outbound = with_setting(MESSAGES, "/v1/messages", adaptive, "off"); + let wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap(); + assert_eq!(sent(&wire)["reasoning_effort"], "none"); + assert_eq!(wire.effort, None); + + // Gemini 2.5 Flash thinks without a thinking config: a zero budget. + let outbound = with_setting(MESSAGES, "/v1/messages", claude, "off"); + let wire = outbound + .address(UpstreamProtocol::GoogleGenerate, "gemini-2.5-flash", false) + .unwrap(); + assert_eq!( + sent(&wire)["generationConfig"]["thinkingConfig"], + serde_json::json!({"thinkingBudget": 0}) + ); + + // A Chat caller to Anthropic: no thinking. + let chat = serde_json::json!({ + "model": "claude-opus-4-6", "reasoning_effort": "high", + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(CHAT, "/v1/chat/completions", chat, "off"); + let wire = outbound + .address( + UpstreamProtocol::AnthropicMessages, + "claude-opus-4-6", + false, + ) + .unwrap(); + assert!(sent(&wire).get("thinking").is_none(), "{}", sent(&wire)); + assert!(sent(&wire).get("output_config").is_none()); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "high", "sent": null, "reason": "setting"})) + ); + + // No reasoning: nothing turned off, nothing written. + let plain = serde_json::json!({ + "model": "gpt-5", "max_tokens": 50, + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(MESSAGES, "/v1/messages", plain, "off"); + let wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap(); + assert!(sent(&wire).get("reasoning_effort").is_none()); + assert_eq!(wire.effort, None); + } + + /// A level on a converted hop is written the way the upstream's format + /// does, and the log says what went out. + #[test] + fn a_level_on_a_converted_hop_is_said_as_sent() { + let claude = serde_json::json!({ + "model": "gpt-5", "max_tokens": 2000, + "thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}, + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(MESSAGES, "/v1/messages", claude, "high"); + let wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap(); + assert_eq!(sent(&wire)["reasoning_effort"], "high"); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "low", "sent": "high", "reason": "setting"})) + ); + let wire = outbound + .address(UpstreamProtocol::OpenAiResponses, "gpt-5", false) + .unwrap(); + assert_eq!(sent(&wire)["reasoning"]["effort"], "high"); + let chat = serde_json::json!({ + "model": "claude-opus-4-6", "reasoning_effort": "low", + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(CHAT, "/v1/chat/completions", chat, "high"); + let wire = outbound + .address( + UpstreamProtocol::AnthropicMessages, + "claude-opus-4-6", + false, + ) + .unwrap(); + assert_eq!(sent(&wire)["output_config"]["effort"], "high"); + assert_eq!(sent(&wire)["thinking"]["type"], "adaptive"); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "low", "sent": "high", "reason": "setting"})) + ); + } + + /// A converted `none` the upstream refuses goes out again with the + /// lowest level it named. + #[test] + fn a_refused_none_on_a_converted_hop_takes_the_lowest_level() { + let claude = serde_json::json!({ + "model": "gpt-5", "max_tokens": 2000, + "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}, + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(MESSAGES, "/v1/messages", claude, "off"); + let mut wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap(); + let refused = GatewayError::ProviderHttpError { + status: 400, + message: "relay: Unsupported value: 'none' is not supported with this model for reasoning_effort. Supported values are: 'minimal', 'low', 'medium', and 'high'.".into(), + retry_after: None, + }; + let sent_effort = wire.written_effort().unwrap(); + let fix = effort::rejection(&refused, "relay", sent_effort).unwrap(); + assert_eq!(fix, EffortFix::Use(Effort::Minimal)); + assert!(wire.fix_effort(fix, sent_effort, Reason::Rejected)); + assert_eq!(sent(&wire)["reasoning_effort"], "minimal"); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "none", "sent": "minimal", "reason": "rejected"})) + ); + } + + #[test] + fn a_refused_effort_is_sent_again_as_the_upstream_says() { + let ask = serde_json::json!({ + "model": "gpt-5", "reasoning_effort": "xhigh", + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(CHAT, "/v1/chat/completions", ask, ""); + let mut wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap(); + assert_eq!(wire.effort, None); + let xhigh = WrittenEffort::Level(Effort::XHigh); + assert_eq!(wire.written_effort(), Some(xhigh)); + + assert!(wire.fix_effort(EffortFix::Use(Effort::High), xhigh, Reason::Rejected)); + assert_eq!(sent(&wire)["reasoning_effort"], "high"); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "xhigh", "sent": "high", "reason": "rejected"})) + ); + // Left out: the model's default. + assert!(wire.fix_effort(EffortFix::Drop, xhigh, Reason::Rejected)); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "xhigh", "sent": null, "reason": "rejected"})) + ); + // Nothing left to change. + assert!(!wire.fix_effort(EffortFix::Drop, xhigh, Reason::Rejected)); + } + + #[test] + fn a_level_refused_before_goes_out_as_what_worked() { + let ask = serde_json::json!({ + "model": "gpt-5", "reasoning_effort": "xhigh", + "messages": [{"role": "user", "content": "ping"}] + }); + let outbound = with_setting(CHAT, "/v1/chat/completions", ask, ""); + let xhigh = WrittenEffort::Level(Effort::XHigh); + let mut memory = effort::Memory::default(); + memory.note(xhigh, EffortFix::Use(Effort::High), effort::now()); + let wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap() + .remembered(&memory); + assert_eq!(sent(&wire)["reasoning_effort"], "high"); + assert!(wire.remembered); + assert_eq!(wire.pre, Some(xhigh)); + assert_eq!( + wire.effort.map(|c| c.detail()), + Some(serde_json::json!({"asked": "xhigh", "sent": "high", "reason": "remembered"})) + ); + // Another level: as sent. + let mut other = effort::Memory::default(); + other.note( + WrittenEffort::Level(Effort::Max), + EffortFix::Drop, + effort::now(), + ); + let wire = outbound + .address(UpstreamProtocol::OpenAiChat, "gpt-5", false) + .unwrap() + .remembered(&other); + assert_eq!(wire.written_effort(), Some(xhigh)); + assert!(!wire.remembered); + assert_eq!(wire.effort, None); + } + /// The body a Chat caller's request goes out with to a Chat upstream, /// under a model cap of `cap`. fn sent_to_chat(ask: Value, cap: u32, official: bool) -> Value { diff --git a/crates/gateway/src/proxy/log_ctx.rs b/crates/gateway/src/proxy/log_ctx.rs index 6efb362c..8e3d4890 100644 --- a/crates/gateway/src/proxy/log_ctx.rs +++ b/crates/gateway/src/proxy/log_ctx.rs @@ -82,6 +82,7 @@ impl<'a> LogCtx<'a> { None, self.started.elapsed().as_millis() as i64, &err, + 0, // Pre-route-selection failures don't have a parsed request // we can safely serialize (transform errors, malformed JSON, // etc. land here). Body capture for these paths is a follow- @@ -96,6 +97,10 @@ impl<'a> LogCtx<'a> { /// blob mirrors the success-path shape but also carries `error_type` /// and `error_message` so operators can drill down without joining /// against a separate error-log stream. +/// +/// `retries`: how many times the request went out again after a +/// transient failure (see `super::retry`), said as `"retries"` when any +/// did. #[allow(clippy::too_many_arguments)] pub(super) fn emit_gateway_error_log( audit: &think_watch_common::audit::AuditLogger, @@ -110,9 +115,13 @@ pub(super) fn emit_gateway_error_log( provider: Option<&str>, latency_ms: i64, err: &GatewayError, + retries: u32, bodies: BodyCapture, ) { - let detail = error_detail(model_id, provider, latency_ms, err); + let mut detail = error_detail(model_id, provider, latency_ms, err); + if retries > 0 { + detail["retries"] = retries.into(); + } // Same `chat.completion` action as the success path — flush_gateway // drops the action when it writes ChGatewayRow, so the trace // endpoint distinguishes errors via `status_code` (>= 400) instead. @@ -272,6 +281,7 @@ mod tests { GatewayError::ProviderHttpError { status: 500, message: "bedrock: arn:aws:iam::123456789012:user/gateway".into(), + retry_after: None, }, GatewayError::UpstreamRateLimited { retry_after_secs: Some(12), diff --git a/crates/gateway/src/proxy/mod.rs b/crates/gateway/src/proxy/mod.rs index de83c925..a72563e1 100644 --- a/crates/gateway/src/proxy/mod.rs +++ b/crates/gateway/src/proxy/mod.rs @@ -23,6 +23,7 @@ mod accounting; mod body_capture; pub(crate) mod cache_marks; mod early_cancel; +pub mod effort; pub(crate) mod generate; mod headers; mod identity; @@ -31,6 +32,7 @@ mod models; mod pipeline; mod protocol_relearn; mod responses_ws; +pub(crate) mod retry; mod routing; pub(crate) mod sealed; pub mod shaper; @@ -40,7 +42,7 @@ 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::{SelectionRecord, fails as upstream_failed, finalize_health}; +pub(crate) use routing::{Resend, SelectionRecord, fails as upstream_failed, finalize_health}; // pub re-exports — `server::app` mounts these as route handlers. pub use early_cancel::{EarlyCancel, EarlyCancelSlot}; @@ -51,6 +53,30 @@ pub use identity::limits_for_ai_gateway; pub use models::{list_gemini_models_handler, list_models_handler}; pub use responses_ws::proxy_responses_ws; +/// How long the gateway port gives a handler to produce its response +/// before answering 408 (the server's `TimeoutLayer`). A stream's body is +/// not bound by it — only the time to its head — but a whole answer is, +/// failover and retries included, so its retries stop by then too. +pub const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120); + +/// When a request reached the gateway port, stamped just inside the +/// server's `TimeoutLayer` (see [`stamp_arrival`]): the moment its +/// [`REQUEST_TIMEOUT`] started counting, before authentication and the +/// body were read. +#[derive(Debug, Clone, Copy)] +pub struct Arrived(pub std::time::Instant); + +/// Stamp [`Arrived`] on the request. Layered directly inside the gateway +/// port's `TimeoutLayer`. +pub async fn stamp_arrival( + mut req: axum::extract::Request, + next: axum::middleware::Next, +) -> axum::response::Response { + req.extensions_mut() + .insert(Arrived(std::time::Instant::now())); + next.run(req).await +} + /// Shared application state for the gateway proxy handlers. #[derive(Clone)] pub struct GatewayState { @@ -259,6 +285,7 @@ mod helper_tests { GatewayError::ProviderHttpError { status: 418, message: "teapot".into(), + retry_after: None, }, 418, ), @@ -329,6 +356,7 @@ mod helper_tests { GatewayError::ProviderHttpError { status: 503, message: "down".into(), + retry_after: None, }, ] { let expected = err.status_code(); diff --git a/crates/gateway/src/proxy/pipeline.rs b/crates/gateway/src/proxy/pipeline.rs index eea8871c..77f9731d 100644 --- a/crates/gateway/src/proxy/pipeline.rs +++ b/crates/gateway/src/proxy/pipeline.rs @@ -15,6 +15,7 @@ use super::identity::limits_for_ai_gateway; use super::{GatewayErrorResponse, GatewayRequestIdentity, GatewayState}; +use super::routing::Resend; use super::shaper::StreamShaper; use crate::lifecycle::{ ChatCompletionOutcome, ChatCompletionSurface, ChatPostInvokeDeps, Completed, OpenUpstream, @@ -112,14 +113,16 @@ fn short_circuit_to_response(outcome: ChatCompletionOutcome) -> GatewayErrorResp /// awaits the tail future, then drives `run_post_invoke` through /// record_outcome → write_cache → record_usage → emit_audit. pub(super) fn launch_stream_pump( - deps: ChatPostInvokeDeps, + mut deps: ChatPostInvokeDeps, open: OpenUpstream, + resend: Resend, shaper: StreamShaper, client: Dialect, client_sse: bool, ) -> axum::response::Response { let (response, tail) = build_chat_pump( open, + resend, shaper, client, client_sse, @@ -129,7 +132,18 @@ pub(super) fn launch_stream_pump( &deps.route.provider_name, ); tokio::spawn(async move { - let invoked = tail.await; + let (invoked, sent) = tail.await; + // 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 { + deps.route.sel_record.started_at = sent.last_sent; + deps.route.sel_record.retries = sent.retries; + } + deps.route.sel_record.effort = sent.effort; + // The client left while the stream waited to go out again: the + // attempt that failed was counted against the route when the retry + // was decided, and nothing went out since. + deps.route.sel_record.counted = sent.waiting; run_post_invoke::(invoked, &deps).await; }); response diff --git a/crates/gateway/src/proxy/protocol_relearn.rs b/crates/gateway/src/proxy/protocol_relearn.rs index 42156a0c..c81c397d 100644 --- a/crates/gateway/src/proxy/protocol_relearn.rs +++ b/crates/gateway/src/proxy/protocol_relearn.rs @@ -27,7 +27,10 @@ use crate::protocol::UpstreamProtocol; /// once. Only 4xx bodies that name an API surface qualify — the shape /// upstreams actually use to report this. pub(super) fn is_protocol_mismatch(err: &GatewayError) -> bool { - let GatewayError::ProviderHttpError { status, message } = err else { + let GatewayError::ProviderHttpError { + status, message, .. + } = err + else { return false; }; // A 5xx is an upstream incident, not a dialect problem, and retrying @@ -104,10 +107,12 @@ mod tests { message: "The model 'anthropic.claude-opus-4' does not support the \ '/v1/chat/completions' API" .into(), + retry_after: None, })); assert!(is_protocol_mismatch(&GatewayError::ProviderHttpError { status: 404, message: "unsupported endpoint /v1/responses for this model".into(), + retry_after: None, })); } @@ -119,11 +124,13 @@ mod tests { assert!(!is_protocol_mismatch(&GatewayError::ProviderHttpError { status: 400, message: "messages: at least one message is required".into(), + retry_after: None, })); // Upstream incident, not a dialect problem. assert!(!is_protocol_mismatch(&GatewayError::ProviderHttpError { status: 503, message: "/v1/chat/completions is not available right now".into(), + retry_after: None, })); assert!(!is_protocol_mismatch(&GatewayError::UpstreamAuthError { status: 401, diff --git a/crates/gateway/src/proxy/responses_ws.rs b/crates/gateway/src/proxy/responses_ws.rs index 2e2ee8cf..9d67f6a9 100644 --- a/crates/gateway/src/proxy/responses_ws.rs +++ b/crates/gateway/src/proxy/responses_ws.rs @@ -110,6 +110,7 @@ async fn serve( headers.clone(), identity.clone(), None, + None, Bytes::from(Value::Object(body).to_string()), RESPONSES, "/v1/responses", diff --git a/crates/gateway/src/proxy/retry.rs b/crates/gateway/src/proxy/retry.rs new file mode 100644 index 00000000..50923680 --- /dev/null +++ b/crates/gateway/src/proxy/retry.rs @@ -0,0 +1,430 @@ +//! Sending a request again after the upstream failed for the moment. +//! +//! When every candidate route has been tried and the last one failed for +//! the moment, and none of the answer has reached the client, the request +//! waits and goes again to the routes that failed that way, instead of +//! handing the failure to the client. A model with one route — and a +//! stream, which is sent to the one route it picked — sends to it again. +//! +//! **For the moment** ([`transient`]): overloaded (503, 529, and the same +//! errors reported inside a stream), a server error (500, 502, 504), a +//! connection that could not be made or broke before the answer's head, a +//! stream that broke off before any of the answer. **Not** sent again: an +//! upstream's rate limit or spent quota (429), a refused credential (401, +//! 403), any other 4xx (a bad request, a model the upstream does not +//! have), a timeout, the gateway's own limits, budgets and policies, and a +//! client that left. Neither is anything after part of the answer went +//! out: the client would get its beginning twice. +//! +//! **How long it waits** ([`Budget`]): 2, 4, 8, then 16 seconds each +//! time, give or take a fifth so that many requests failing together do +//! not come back together; or as long as the upstream said to wait +//! (`retry-after-ms`, `Retry-After`) — but never less than +//! [`MIN_SAID_WAIT`]: an upstream that says "now" would otherwise be +//! asked again at once, hundreds of times over. It stops when the wait +//! would end past the request's deadline, `gateway.retry_wait_secs` after +//! the request arrived (0: never sent again). There is no count of its +//! own. An attempt already under way is not cut short by that deadline: +//! it only decides whether another one starts. +//! +//! A whole (non-streamed) answer has a second, hard deadline: the +//! gateway's request timeout (`super::REQUEST_TIMEOUT`), after which the +//! client gets an empty 408 and the request no log row. Each attempt is +//! cut at it, and a retry starts only when the wait and another attempt +//! as long as the last one fit before it (see [`Budget::within`]). +//! +//! Each attempt is a pick like any other: the strategy, the breaker, the +//! route caps, the session affinity and the provider holding a +//! conversation's encrypted history (`super::sealed`) all apply, and the +//! failed attempt counts against its route as before. + +use std::time::{Duration, Instant}; + +use rand::RngExt; + +use crate::error::GatewayError; + +/// What the client receives while the gateway holds back the opening of +/// a stream or waits to send again: an SSE comment, which every SSE +/// reader skips. Sent once nothing has gone out for [`KEEP_ALIVE`]. +pub(crate) const KEEP_ALIVE_FRAME: &[u8] = b": keep-alive\n\n"; + +/// How long the client may go without a byte while the gateway holds a +/// stream back. +pub(crate) const KEEP_ALIVE: Duration = Duration::from_secs(15); + +/// Did the upstream fail for the moment, so that the same request may +/// well be answered a little later? From what a stream records of its +/// error, as [`super::routing::fails`] reads it: the error's tag and +/// status. +/// +/// `transport` is a stream that broke off in transit (see +/// `lifecycle::broke_off`). `NetworkError` is a connection that could not +/// be made, a TLS failure, or one that broke before or while the answer +/// was read. A timeout (`ProviderTimeout`) is not: the upstream already +/// had all that time. +pub(crate) fn transient(error_tag: &str, status: i64) -> bool { + match error_tag { + "ProviderHttpError" => matches!(status, 500 | 502 | 503 | 504 | 529), + "NetworkError" | "transport" => true, + _ => false, + } +} + +/// [`transient`] for an error. +pub(crate) fn is_transient(err: &GatewayError) -> bool { + transient(err.error_tag(), err.status_code()) +} + +/// How long before the gateway's request timeout a whole answer's last +/// attempt is cut: time to write its error and its log row. +pub(crate) const TIMEOUT_MARGIN: Duration = Duration::from_secs(2); + +/// The hard deadline of a whole answer to a request that `arrived` at the +/// gateway port (see [`super::Arrived`]). +pub(crate) fn hard_deadline(arrived: Instant) -> Instant { + arrived + super::REQUEST_TIMEOUT.saturating_sub(TIMEOUT_MARGIN) +} + +/// The shortest wait an upstream's `retry-after-ms` / `Retry-After` is +/// taken for. Zero is taken as no word at all: the backoff applies. +pub(crate) const MIN_SAID_WAIT: Duration = Duration::from_millis(250); + +/// How long the upstream said to wait before trying again, held to +/// [`MIN_SAID_WAIT`]. +fn said_to_wait(err: &GatewayError) -> Option { + match err { + GatewayError::ProviderHttpError { retry_after, .. } => retry_after + .filter(|d| !d.is_zero()) + .map(|d| d.max(MIN_SAID_WAIT)), + _ => None, + } +} + +/// The wait before the `n`-th retry (1-based), before jitter: 2, 4, 8, +/// then 16 seconds. +fn backoff(n: u32) -> Duration { + Duration::from_secs(2u64 << n.clamp(1, 4).saturating_sub(1)) +} + +/// A request's allowance for sending again: its deadline, and how many +/// retries went out. +#[derive(Debug, Clone)] +pub(crate) struct Budget { + /// `None`: never sent again. + deadline: Option, + /// When the answer must be back by, whatever the retries: a whole + /// answer's request timeout. `None` for a stream, whose body is not + /// bound by it. + hard: Option, + retries: u32, +} + +impl Budget { + /// The allowance of a request that arrived at `arrived`: `wait` from + /// then. Zero: none. + pub(crate) fn new(arrived: Instant, wait: Duration) -> Self { + Self { + deadline: (!wait.is_zero()) + .then(|| arrived.checked_add(wait)) + .flatten(), + hard: None, + retries: 0, + } + } + + /// The same allowance for a request that must be answered by `hard`: + /// each attempt is cut there ([`Budget::hard_deadline`]), and a retry + /// starts only when its wait and another attempt as long as the last + /// one fit before it. + pub(crate) fn within(mut self, hard: Instant) -> Self { + self.hard = Some(hard); + self + } + + /// When the attempt under way must be given up, if ever. + pub(crate) fn hard_deadline(&self) -> Option { + self.hard + } + + /// Whether a failure now could still be followed by another attempt. + pub(crate) fn open(&self) -> bool { + self.deadline.is_some_and(|d| Instant::now() < d) + } + + /// How long to wait before sending again after `err`. `None`: `err` + /// is the answer — it is not transient, or the wait would end past + /// the deadline. + pub(crate) fn wait_after(&self, err: &GatewayError) -> Option { + self.wait_after_taking(err, Duration::ZERO) + } + + /// [`Budget::wait_after`] for an attempt that took `took`: under a + /// hard deadline, the next one is expected to take as long. + pub(crate) fn wait_after_taking(&self, err: &GatewayError, took: Duration) -> Option { + if !is_transient(err) { + return None; + } + self.wait(said_to_wait(err), took) + } + + /// [`Budget::wait_after`] for what a stream recorded of its error. + pub(crate) fn wait_after_outcome(&self, error_tag: &str, status: i64) -> Option { + if !transient(error_tag, status) { + return None; + } + self.wait(None, Duration::ZERO) + } + + fn wait(&self, said: Option, took: Duration) -> Option { + let jitter = rand::rng().random_range(0.8..=1.2); + self.wait_at(Instant::now(), said, jitter, took) + } + + /// The wait before the next retry, `jitter` the factor the backoff is + /// spread by: the upstream's own word if it gave one, else the + /// backoff. `None` when it would end past the deadline, or when it and + /// another attempt as long as the last (`took`) would not fit before + /// the hard deadline. + fn wait_at( + &self, + now: Instant, + said: Option, + jitter: f64, + took: Duration, + ) -> Option { + let deadline = self.deadline?; + let wait = said.unwrap_or_else(|| backoff(self.retries + 1).mul_f64(jitter)); + let resumes = now.checked_add(wait)?; + if resumes > deadline { + return None; + } + if let Some(hard) = self.hard + && resumes.checked_add(took)? > hard + { + return None; + } + Some(wait) + } + + /// The next attempt went out. + pub(crate) fn sent(&mut self) { + self.retries += 1; + } + + /// How many retries went out. + pub(crate) fn retries(&self) -> u32 { + self.retries + } +} + +/// `gateway_provider_retry_total{provider}`: an attempt that went out +/// again after a transient failure, by the provider it went to. Next to +/// `gateway_provider_fallback_total`, which counts moving on to another +/// route. +pub(crate) fn count(provider: &str) { + metrics::counter!( + "gateway_provider_retry_total", + "provider" => crate::metrics_labels::normalize_provider_label(provider), + ) + .increment(1); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn http(status: u16, retry_after: Option) -> GatewayError { + GatewayError::ProviderHttpError { + status, + message: "p: down".into(), + retry_after, + } + } + + #[test] + fn only_a_failure_of_the_moment_is_transient() { + for status in [500, 502, 503, 504, 529] { + assert!(is_transient(&http(status, None)), "{status}"); + } + for status in [400, 404, 408, 413, 422, 501] { + assert!(!is_transient(&http(status, None)), "{status}"); + } + assert!(is_transient(&GatewayError::NetworkError("reset".into()))); + // A stream broken off before any of the answer. + assert!(transient("transport", 502)); + + for e in [ + GatewayError::UpstreamRateLimited { + retry_after_secs: Some(1), + }, + GatewayError::UpstreamAuthError { + status: 401, + message: String::new(), + }, + GatewayError::ProviderTimeout("slow".into()), + GatewayError::ProviderInvalidResponse("not json".into()), + GatewayError::ProviderError("signing failed".into()), + GatewayError::TransformError("bad".into()), + GatewayError::PolicyBlocked("rule".into()), + GatewayError::rate_limited("user:requests/1m", 5), + GatewayError::budget_exhausted("user:budget/monthly", 5), + GatewayError::SealedHistoryRefused { + status: 400, + provider: "p".into(), + said: "invalid_encrypted_content".into(), + }, + ] { + assert!(!is_transient(&e), "{e:?}"); + } + // A stream the client left, or one that ended normally. + assert!(!transient("ClientCancelled", 499)); + } + + #[test] + fn the_backoff_doubles_up_to_sixteen_seconds() { + let secs: Vec = (1..=6).map(|n| backoff(n).as_secs()).collect(); + assert_eq!(secs, [2, 4, 8, 16, 16, 16]); + } + + #[test] + fn the_wait_is_spread_by_a_fifth_either_way() { + let now = Instant::now(); + let budget = Budget::new(now, Duration::from_secs(120)); + assert_eq!( + budget.wait_at(now, None, 0.8, Duration::ZERO), + Some(Duration::from_millis(1600)) + ); + assert_eq!( + budget.wait_at(now, None, 1.2, Duration::ZERO), + Some(Duration::from_millis(2400)) + ); + for _ in 0..100 { + let w = budget.wait(None, Duration::ZERO).unwrap(); + assert!( + w >= Duration::from_millis(1600) && w <= Duration::from_millis(2400), + "{w:?}" + ); + } + } + + #[test] + fn each_retry_waits_longer() { + let now = Instant::now(); + let mut budget = Budget::new(now, Duration::from_secs(600)); + let mut waits = Vec::new(); + for _ in 0..5 { + waits.push( + budget + .wait_at(now, None, 1.0, Duration::ZERO) + .unwrap() + .as_secs(), + ); + budget.sent(); + } + assert_eq!(waits, [2, 4, 8, 16, 16]); + assert_eq!(budget.retries(), 5); + } + + #[test] + fn the_upstreams_own_wait_wins_over_the_backoff() { + let now = Instant::now(); + let budget = Budget::new(now, Duration::from_secs(120)); + let said = Some(Duration::from_millis(300)); + assert_eq!(budget.wait_at(now, said, 1.2, Duration::ZERO), said); + assert_eq!( + budget.wait_at(now, Some(Duration::from_secs(30)), 1.0, Duration::ZERO), + Some(Duration::from_secs(30)) + ); + // Longer than what is left: the failure is the answer. + assert_eq!( + budget.wait_at(now, Some(Duration::from_secs(121)), 1.0, Duration::ZERO), + None + ); + } + + /// "Now" is no word at all, and "almost now" is held to a floor: an + /// upstream that says so must not be asked again hundreds of times. + #[test] + fn a_wait_of_nothing_from_the_upstream_is_not_taken_at_its_word() { + assert_eq!(said_to_wait(&http(503, Some(Duration::ZERO))), None); + assert_eq!( + said_to_wait(&http(503, Some(Duration::from_millis(1)))), + Some(MIN_SAID_WAIT) + ); + assert_eq!( + said_to_wait(&http(503, Some(Duration::from_secs(3)))), + Some(Duration::from_secs(3)) + ); + // Zero falls back to the backoff. + let budget = Budget::new(Instant::now(), Duration::from_secs(120)); + let w = budget.wait_after(&http(503, Some(Duration::ZERO))).unwrap(); + assert!(w >= Duration::from_millis(1600), "{w:?}"); + } + + /// A whole answer's hard deadline: a retry starts only when its wait + /// and another attempt as long as the last fit before it. + #[test] + fn a_retry_must_fit_before_the_hard_deadline() { + let arrived = Instant::now(); + let budget = Budget::new(arrived, Duration::from_secs(600)) + .within(arrived + Duration::from_secs(118)); + assert_eq!( + budget.hard_deadline(), + Some(arrived + Duration::from_secs(118)) + ); + // A 504 after 60 s, 62 s in: a 2 s wait and another 60 s do not fit. + let at = arrived + Duration::from_secs(62); + let minute = Duration::from_secs(60); + assert_eq!(budget.wait_at(at, None, 1.0, minute), None); + // A quick failure does. + assert_eq!( + budget.wait_at(at, None, 1.0, Duration::from_millis(50)), + Some(Duration::from_secs(2)) + ); + // Without a hard deadline (a stream), only the soft one counts. + let stream = Budget::new(arrived, Duration::from_secs(600)); + assert_eq!(stream.hard_deadline(), None); + assert!(stream.wait_at(at, None, 1.0, minute).is_some()); + } + + #[test] + fn no_wait_ends_past_the_deadline() { + let arrived = Instant::now(); + let budget = Budget::new(arrived, Duration::from_secs(10)); + assert!(budget.open()); + // 9 s in, 1 s left: a 2 s backoff does not fit. + assert_eq!( + budget.wait_at(arrived + Duration::from_secs(9), None, 1.0, Duration::ZERO), + None + ); + assert_eq!( + budget.wait_at(arrived + Duration::from_secs(8), None, 1.0, Duration::ZERO), + Some(Duration::from_secs(2)) + ); + } + + #[test] + fn a_zero_wait_never_sends_again() { + let now = Instant::now(); + let budget = Budget::new(now, Duration::ZERO); + assert!(!budget.open()); + assert_eq!(budget.wait_at(now, None, 1.0, Duration::ZERO), None); + assert_eq!(budget.wait_after(&http(503, None)), None); + } + + #[test] + fn a_failure_that_is_not_transient_is_the_answer() { + let budget = Budget::new(Instant::now(), Duration::from_secs(120)); + assert!(budget.wait_after(&http(503, None)).is_some()); + assert_eq!(budget.wait_after(&http(400, None)), None); + assert_eq!( + budget.wait_after(&GatewayError::UpstreamRateLimited { + retry_after_secs: None + }), + None + ); + assert!(budget.wait_after_outcome("transport", 502).is_some()); + assert_eq!(budget.wait_after_outcome("ProviderHttpError", 400), None); + } +} diff --git a/crates/gateway/src/proxy/routing.rs b/crates/gateway/src/proxy/routing.rs index c9d29a99..1cb9639a 100644 --- a/crates/gateway/src/proxy/routing.rs +++ b/crates/gateway/src/proxy/routing.rs @@ -16,6 +16,7 @@ use crate::route_caps::{self, Admission, CappedRoutes, RouteCaps}; use crate::router::{AffinityMode, RouteEntry}; use crate::strategy::{self, RoutingStrategy}; +use super::retry::Budget; use super::sealed::Pin; /// What the affinity layer can pin a session to. @@ -293,17 +294,31 @@ pub(crate) struct SelectionRecord { /// at finalize time — which is correct for streams since there's /// no failover loop inflating the e2e timer. pub picked_latency_ms: Option, + /// How many times the request went out again after a transient + /// failure before this route answered (see `super::retry`). Said on + /// the log row as `"retries"`. + pub retries: u32, + /// The hop that answered sent a different reasoning effort than the + /// caller asked (see `super::effort`). Said on the log row as + /// `"effort"`. + pub effort: Option, + /// 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, } /// Record health for the picked route once the request finishes. /// Best-effort; failures are logged. pub(crate) async fn finalize_health(state: &GatewayState, sel: &SelectionRecord, success: bool) { + if sel.counted { + return; + } let total_latency_ms = sel.started_at.elapsed().as_millis().min(u32::MAX as u128) as u32; // For health, we want the picked route's *own* time, not the // cumulative including prior failed attempts. Non-stream sets // this explicitly when it picks a winner; stream falls back to - // total elapsed (which is the picked route's time anyway since - // streams don't loop). + // total elapsed, which is the picked route's time: a stream that + // went out again starts `started_at` at the attempt that answered. let health_latency_ms = sel.picked_latency_ms.unwrap_or(total_latency_ms); let breaker = resolve_breaker_config(state).await; let _ = state @@ -405,12 +420,24 @@ async fn pick_admitted<'a>( /// fails (see [`is_upstream_failure`]) tries another candidate from the /// remaining set until exhausted. When every candidate left is at its /// cap and none failed, the answer is the caps' 429. +/// +/// Under `budget`'s hard deadline (a whole answer's request timeout), +/// each attempt is cut there and the last failure, or a timeout, is the +/// answer. +/// +/// When every candidate is exhausted and the last one failed for the +/// moment (see `super::retry`), the request waits as `budget` allows and +/// goes round again, to the providers whose last failure in this request +/// was transient — picked as before, so the strategy, the breaker, the +/// caps, the affinity and a sealed conversation's pin all apply. A model +/// with one route sends to it again. pub(super) async fn select_route_with_failover<'a>( routes: &'a [RouteEntry], outbound: &super::generate::Outbound, call_ctx: &CallCtx, ctx: &SelectionCtx<'_>, caller_model: &str, + budget: &mut Budget, ) -> Result, GatewayError> { let started_at = std::time::Instant::now(); let candidates: Vec<&RouteEntry> = routes.iter().collect(); @@ -418,102 +445,187 @@ pub(super) async fn select_route_with_failover<'a>( let mut last_error: Option = None; let mut tried: Vec = Vec::new(); let mut capped = CappedRoutes::default(); + // The providers whose last failure in this request was transient, in + // the order they failed: the ones a retry goes back to. + let mut transient: Vec = Vec::new(); + // How long the last attempt took: under a hard deadline, a retry + // starts only when one as long fits after its wait. + let mut last_took = std::time::Duration::ZERO; + + loop { + let retrying = budget.retries() > 0; + let mut attempted = false; + for _ in 0..candidates.len() { + let Some(entry) = pick_admitted(&candidates, &tried, &mut capped, ctx).await? else { + break; + }; + tried.push(entry.provider_id); + attempted = true; + if retrying { + super::retry::count(&entry.provider_name); + } - for _ in 0..candidates.len() { - let Some(entry) = pick_admitted(&candidates, &tried, &mut capped, ctx).await? else { - break; - }; - tried.push(entry.provider_id); - - let upstream_model = entry.upstream_model.as_deref().unwrap_or(caller_model); - - // Per-attempt clock: record latency against this route as - // *its* time, not "everything since the request started" - // (which would double-count earlier failed attempts in - // a failover chain and skew the latency strategy). - let attempt_started_at = std::time::Instant::now(); - // A rejected dialect is retried inside `send`, on this same route. - let result = - match super::generate::send(entry, outbound, call_ctx, &ctx.state.db, upstream_model) + let upstream_model = entry.upstream_model.as_deref().unwrap_or(caller_model); + + // Per-attempt clock: record latency against this route as + // *its* time, not "everything since the request started" + // (which would double-count earlier failed attempts in + // a failover chain and skew the latency strategy). + let attempt_started_at = std::time::Instant::now(); + // A rejected dialect is retried inside `send`, on this same route. + let attempt = async { + match super::generate::send( + entry, + outbound, + call_ctx, + &ctx.state.db, + &ctx.state.redis, + upstream_model, + ) .await - { - Ok((resp, wire)) => { - super::generate::read_whole(resp, &wire, caller_model, outbound.input_estimate) - .await - .map(|answer| (answer, wire.ledger)) + { + Ok((resp, wire)) => super::generate::read_whole( + resp, + &wire, + caller_model, + outbound.input_estimate, + ) + .await + .map(|answer| (answer, wire.ledger, wire.effort)), + Err(e) => Err(e), } - Err(e) => Err(e), + }; + // Cut at the hard deadline, so that the answer is an error the + // client and the log are told, not the request timeout's 408. + let result = match budget.hard_deadline() { + Some(hard) => { + tokio::time::timeout_at(tokio::time::Instant::from_std(hard), attempt).await + } + None => Ok(attempt.await), }; - let attempt_latency_ms = attempt_started_at - .elapsed() - .as_millis() - .min(u32::MAX as u128) as u32; + last_took = attempt_started_at.elapsed(); + let attempt_latency_ms = last_took.as_millis().min(u32::MAX as u128) as u32; - match result { - Ok((response, ledger)) => { - set_affinity( - &ctx.state.redis, - ctx.user_id, - ctx.model_id, - ctx.affinity_mode, - entry, - ctx.affinity_ttl_secs, - ) - .await; - return Ok(( - entry, - response, - ledger, - SelectionRecord { - picked_route_id: entry.route_id, - started_at, - picked_latency_ms: Some(attempt_latency_ms), - }, - )); - } - Err(e) if is_upstream_failure(&e) => { + let Ok(result) = result else { tracing::warn!( provider = %entry.provider_name, provider_id = %entry.provider_id, - error = %e, - "Route failed, trying next" + "Route still had not answered at the gateway's request timeout" ); - metrics::counter!( - "gateway_provider_fallback_total", - "from" => crate::metrics_labels::normalize_provider_label(&entry.provider_name), - ) - .increment(1); - // Record the failed attempt in health so the - // breaker can trip mid-failover. let _ = ctx .state .health .record(entry.route_id, attempt_latency_ms, true, ctx.breaker) .await; - last_error = Some(e); - continue; - } - Err(e) => { - // The request itself was refused. Sibling routes would - // refuse it the same way, so it goes back to the caller. - // An upstream that answered with a refusal is working: - // it counts as a success for the route, as it does on - // the desktop. - if matches!( - e, - GatewayError::ProviderHttpError { .. } - | GatewayError::SealedHistoryRefused { .. } - ) { + // An earlier failure says more than the cut; with none, + // the upstream simply took too long. + return Err(last_error.unwrap_or_else(|| { + GatewayError::ProviderTimeout(format!( + "{} did not answer within the gateway's request timeout", + entry.provider_name + )) + })); + }; + + match result { + Ok((response, ledger, effort)) => { + set_affinity( + &ctx.state.redis, + ctx.user_id, + ctx.model_id, + ctx.affinity_mode, + entry, + ctx.affinity_ttl_secs, + ) + .await; + return Ok(( + entry, + response, + ledger, + SelectionRecord { + picked_route_id: entry.route_id, + started_at, + picked_latency_ms: Some(attempt_latency_ms), + retries: budget.retries(), + effort, + counted: false, + }, + )); + } + Err(e) if is_upstream_failure(&e) => { + tracing::warn!( + provider = %entry.provider_name, + provider_id = %entry.provider_id, + error = %e, + "Route failed, trying next" + ); + metrics::counter!( + "gateway_provider_fallback_total", + "from" => crate::metrics_labels::normalize_provider_label(&entry.provider_name), + ) + .increment(1); + // Record the failed attempt in health so the + // breaker can trip mid-failover. let _ = ctx .state .health - .record(entry.route_id, attempt_latency_ms, false, ctx.breaker) + .record(entry.route_id, attempt_latency_ms, true, ctx.breaker) .await; + transient.retain(|p| *p != entry.provider_id); + if super::retry::is_transient(&e) { + transient.push(entry.provider_id); + } + last_error = Some(e); + continue; + } + Err(e) => { + // The request itself was refused. Sibling routes would + // refuse it the same way, so it goes back to the caller. + // An upstream that answered with a refusal is working: + // it counts as a success for the route, as it does on + // the desktop. + if matches!( + e, + GatewayError::ProviderHttpError { .. } + | GatewayError::SealedHistoryRefused { .. } + ) { + let _ = ctx + .state + .health + .record(entry.route_id, attempt_latency_ms, false, ctx.breaker) + .await; + } + return Err(e); } - return Err(e); } } + + // Every candidate is exhausted. Go round again only when this + // round sent something, its last failure was transient, and the + // wait fits before the deadline. + let wait = match &last_error { + Some(e) if attempted && !transient.is_empty() => budget.wait_after_taking(e, last_took), + _ => None, + }; + let Some(wait) = wait else { break }; + tracing::info!( + model = ctx.model_id, + retry = budget.retries() + 1, + wait_ms = wait.as_millis() as u64, + error = ?last_error.as_ref().map(GatewayError::error_tag), + "Every route failed for the moment, sending again after a wait" + ); + tokio::time::sleep(wait).await; + budget.sent(); + // Only the providers that failed for the moment are asked again. + // Caps are counted afresh: a window full then may have room now. + tried = candidates + .iter() + .map(|e| e.provider_id) + .filter(|p| !transient.contains(p)) + .collect(); + capped = CappedRoutes::default(); } // An upstream that failed is the answer's cause even when the rest @@ -547,6 +659,9 @@ pub(super) async fn select_route_for_stream<'a>( // back to elapsed-from-started_at, which IS the // picked route's wall time for the streaming case. picked_latency_ms: None, + retries: 0, + effort: None, + counted: false, }, )); } @@ -556,6 +671,67 @@ pub(super) async fn select_route_for_stream<'a>( })) } +/// What a stream needs to go out again to its route after a transient +/// failure before any of the answer reached the client (see +/// `super::retry`). A stream is sent to the one route it picked, so that +/// is where it goes again, under the pin and affinity it was picked with. +pub(crate) struct Resend { + state: GatewayState, + route_id: Uuid, + provider: String, + caps: RouteCaps, + breaker: CircuitBreakerConfig, + fail_closed: bool, + /// The route's provider alone can read the conversation's encrypted + /// history: asked again with its breaker open, as it was picked. + pinned_only: bool, + pub(crate) budget: Budget, +} + +impl Resend { + pub(super) fn new(entry: &RouteEntry, ctx: &SelectionCtx<'_>, budget: Budget) -> Self { + Self { + state: ctx.state.clone(), + route_id: entry.route_id, + provider: entry.provider_name.clone(), + caps: RouteCaps::of(entry), + breaker: ctx.breaker, + fail_closed: ctx.fail_closed, + pinned_only: matches!(ctx.pin, Some(Pin::Only(p)) if p == entry.provider_id), + budget, + } + } + + /// Before the wait: may the route be sent to again? Not while its + /// breaker is open, nor at one of its caps — the pick a failover makes; + /// the retry is counted against the caps now. When it may, the attempt + /// that failed (`failed_ms` long) counts against the route, as a failed + /// attempt in a failover does, and the retry is counted. When it may + /// not, the failure is the answer and the stream's tail records it. + pub(crate) async fn admit(&mut self, failed_ms: u32) -> bool { + let health = self + .state + .health + .snapshot(self.route_id, self.breaker) + .await; + if health.state == tw_breaker::State::Open && !self.pinned_only { + return false; + } + match route_caps::admit(&self.state.redis, &self.caps, self.fail_closed).await { + Ok(Admission::Admitted) => {} + Ok(Admission::Capped { .. }) | Err(_) => return false, + } + let _ = self + .state + .health + .record(self.route_id, failed_ms, true, self.breaker) + .await; + self.budget.sent(); + super::retry::count(&self.provider); + true + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/gateway/src/proxy/sealed.rs b/crates/gateway/src/proxy/sealed.rs index 28acdeef..b8b98d2b 100644 --- a/crates/gateway/src/proxy/sealed.rs +++ b/crates/gateway/src/proxy/sealed.rs @@ -179,9 +179,9 @@ pub(crate) fn explain_refusal( e: GatewayError, ) -> GatewayError { match e { - GatewayError::ProviderHttpError { status, message } - if (400..500).contains(&status) && refused(wire, &message) => - { + GatewayError::ProviderHttpError { + status, message, .. + } if (400..500).contains(&status) && refused(wire, &message) => { let said = message .strip_prefix(label) .and_then(|m| m.strip_prefix(": ")) @@ -294,6 +294,7 @@ mod tests { let e = GatewayError::ProviderHttpError { status: 400, message: format!("relay-b: {body}"), + retry_after: None, }; let e = explain_refusal(Dialect::Responses, "relay-b", "relay-b", e); let GatewayError::SealedHistoryRefused { @@ -323,6 +324,7 @@ mod tests { let other = || GatewayError::ProviderHttpError { status: 400, message: "p: max_output_tokens is too large".into(), + retry_after: None, }; assert!(matches!( explain_refusal(Dialect::Responses, "p", "p", other()), @@ -331,6 +333,7 @@ mod tests { let sealed = || GatewayError::ProviderHttpError { status: 400, message: "p: invalid_encrypted_content".into(), + retry_after: None, }; // Not sent as Responses: nothing of OpenAI's was sealed in it. assert!(matches!( @@ -341,6 +344,7 @@ mod tests { let e = GatewayError::ProviderHttpError { status: 500, message: "p: invalid_encrypted_content".into(), + retry_after: None, }; assert!(matches!( explain_refusal(Dialect::Responses, "p", "p", e), diff --git a/crates/gateway/src/proxy/transport.rs b/crates/gateway/src/proxy/transport.rs index 3e862b06..6da9652f 100644 --- a/crates/gateway/src/proxy/transport.rs +++ b/crates/gateway/src/proxy/transport.rs @@ -209,8 +209,14 @@ impl Upstream { /// A request that never got an answer: the upstream timed out, or the /// connection could not be made or broke. Either way it says nothing /// about the request, and another route may well answer it. +/// +/// A request that could not even be built — a provider whose base URL is +/// no URL — is the route's configuration, not a failure of the moment: +/// sending it again would fail the same way (see `retry::transient`). pub(crate) fn transport_error(e: reqwest::Error) -> GatewayError { - if e.is_timeout() { + if e.is_builder() { + GatewayError::ProviderError(format!("The upstream request could not be built: {e}")) + } else if e.is_timeout() { GatewayError::ProviderTimeout(e.to_string()) } else { GatewayError::NetworkError(e.to_string()) @@ -229,19 +235,41 @@ pub(crate) fn transport_error(e: reqwest::Error) -> GatewayError { /// Every other status is kept as it is, in `ProviderHttpError`. Whether /// it is the upstream failing (5xx, 408) or the upstream refusing this /// request (any other 4xx) decides failover and the circuit breaker — -/// see `routing::is_upstream_failure`. +/// see `routing::is_upstream_failure` — and whether it is worth sending +/// again after a wait, `proxy::retry`. The error keeps how long the +/// upstream said to wait. pub(crate) fn status_error( status: u16, - retry_after_secs: Option, + retry_after: Option, message: String, ) -> GatewayError { match status { - 429 => GatewayError::UpstreamRateLimited { retry_after_secs }, + 429 => GatewayError::UpstreamRateLimited { + retry_after_secs: retry_after.map(|d| u32::try_from(d.as_secs()).unwrap_or(u32::MAX)), + }, 401 | 403 => GatewayError::UpstreamAuthError { status, message }, - _ => GatewayError::ProviderHttpError { status, message }, + _ => GatewayError::ProviderHttpError { + status, + message, + retry_after, + }, } } +/// When an answer says to try again: OpenAI's `retry-after-ms`, or +/// `Retry-After` in seconds (see [`crate::error::parse_retry_after_seconds`]). +fn retry_after(headers: &reqwest::header::HeaderMap) -> Option { + let header = |name: &str| headers.get(name).and_then(|v| v.to_str().ok()); + header("retry-after-ms") + .and_then(|v| v.trim().parse::().ok()) + .map(Duration::from_millis) + .or_else(|| { + header(reqwest::header::RETRY_AFTER.as_str()) + .and_then(crate::error::parse_retry_after_seconds) + .map(|s| Duration::from_secs(u64::from(s))) + }) +} + /// Turn a non-2xx upstream answer into the error the caller sees (see /// [`status_error`]). /// @@ -256,14 +284,16 @@ async fn check_status( ) -> Result { let status = resp.status(); if status == reqwest::StatusCode::TOO_MANY_REQUESTS { + // Echoed to the caller in whole seconds, as the upstream wrote it. let retry_after_secs = resp .headers() .get(reqwest::header::RETRY_AFTER) .and_then(|v| v.to_str().ok()) .and_then(crate::error::parse_retry_after_seconds); - return Err(status_error(429, retry_after_secs, String::new())); + return Err(GatewayError::UpstreamRateLimited { retry_after_secs }); } if !status.is_success() { + let retry_after = retry_after(resp.headers()); let body = resp.text().await.unwrap_or_default(); tracing::warn!(provider = label, status = %status, body = %body, "upstream returned non-2xx"); const CLIENT_MAX: usize = 512; @@ -279,7 +309,7 @@ async fn check_status( }; return Err(status_error( status.as_u16(), - None, + retry_after, format!("{label}: {shown}"), )); } @@ -393,6 +423,51 @@ mod tests { .into() } + /// A base URL that is no URL fails every time: not worth sending + /// again, unlike a connection that could not be made. + #[tokio::test] + async fn a_request_that_cannot_be_built_is_not_a_failure_of_the_moment() { + let built = client().post("http://exa mple.invalid/v1").send().await; + let e = transport_error(built.unwrap_err()); + assert!(matches!(e, GatewayError::ProviderError(_)), "{e:?}"); + assert!(!crate::proxy::retry::is_transient(&e)); + + let refused = client().post("http://127.0.0.1:1/v1").send().await; + let e = transport_error(refused.unwrap_err()); + assert!(matches!(e, GatewayError::NetworkError(_)), "{e:?}"); + assert!(crate::proxy::retry::is_transient(&e)); + } + + #[tokio::test] + async fn an_overloaded_answer_keeps_how_long_to_wait() { + let with = |headers: &[(&str, &str)]| { + let mut b = http_1x::Response::builder().status(503); + for (k, v) in headers { + b = b.header(*k, *v); + } + reqwest::Response::from(b.body("busy".to_string()).unwrap()) + }; + let wait = |e: GatewayError| match e { + GatewayError::ProviderHttpError { retry_after, .. } => retry_after, + e => panic!("{e:?}"), + }; + let e = check_status( + with(&[("retry-after-ms", "250"), ("retry-after", "3")]), + "p", + ) + .await + .unwrap_err(); + assert_eq!(wait(e), Some(Duration::from_millis(250))); + let e = check_status(with(&[("retry-after", "3")]), "p") + .await + .unwrap_err(); + assert_eq!(wait(e), Some(Duration::from_secs(3))); + let e = check_status(with(&[("retry-after", "soon")]), "p") + .await + .unwrap_err(); + assert_eq!(wait(e), None); + } + #[tokio::test] async fn a_refused_credential_keeps_the_upstreams_reason_out_of_the_callers_sight() { let body = diff --git a/crates/gateway/src/router.rs b/crates/gateway/src/router.rs index 56c12b3b..b733bdd9 100644 --- a/crates/gateway/src/router.rs +++ b/crates/gateway/src/router.rs @@ -78,6 +78,9 @@ pub struct ModelRoutingConfig { /// limit. Carried alongside the routing overrides so the proxy reads /// everything model-scoped in one HashMap lookup. pub max_output_tokens: Option, + /// The model's reasoning setting (`models.reasoning_effort`): see + /// `proxy::effort`. `None` ⇒ requests go out as sent. + pub reasoning: Option, } /// Affinity scope — see `proxy.rs` for the runtime semantics. @@ -334,6 +337,7 @@ mod tests { affinity_mode: Some(AffinityMode::None), affinity_ttl_secs: Some(60), max_output_tokens: None, + reasoning: None, }, ); let cfg = router.config_for("gpt-4o-mini"); diff --git a/crates/server/src/app.rs b/crates/server/src/app.rs index 4b9d02ed..6c449749 100644 --- a/crates/server/src/app.rs +++ b/crates/server/src/app.rs @@ -409,9 +409,12 @@ pub async fn create_gateway_app(_config: &AppConfig, state: AppState) -> anyhow: .merge(key_routes) .merge(mcp_routes) .layer(RequestBodyLimitLayer::new(10 * 1024 * 1024)) // 10MB for large prompts + // When the timeout below starts counting: a whole answer's retries + // stop in time for its error to be written (see `gateway_proxy::retry`). + .layer(axum::middleware::from_fn(gateway_proxy::stamp_arrival)) .layer(TimeoutLayer::with_status_code( axum::http::StatusCode::REQUEST_TIMEOUT, - std::time::Duration::from_secs(120), // longer timeout for streaming + gateway_proxy::REQUEST_TIMEOUT, )) .layer(gateway_cors) .layer(SetResponseHeaderLayer::overriding( @@ -1249,10 +1252,11 @@ pub(crate) async fn load_providers_into_router( affinity_mode: Option, affinity_ttl_secs: Option, max_output_tokens: Option, + reasoning_effort: Option, } let model_rows = sqlx::query_as::<_, ModelRow>( r#"SELECT model_id, routing_strategy, affinity_mode, affinity_ttl_secs, - max_output_tokens + max_output_tokens, reasoning_effort FROM models"#, ) .fetch_all(&state.db) @@ -1278,6 +1282,10 @@ pub(crate) async fn load_providers_into_router( .max_output_tokens .and_then(|v| u32::try_from(v).ok()) .filter(|v| *v > 0), + reasoning: m + .reasoning_effort + .as_deref() + .and_then(think_watch_gateway::proxy::effort::ReasoningSetting::parse), }; // Skip storing the all-default config (saves a HashMap entry // per model that's just inheriting global defaults). @@ -1285,6 +1293,7 @@ pub(crate) async fn load_providers_into_router( || cfg.affinity_mode.is_some() || cfg.affinity_ttl_secs.is_some() || cfg.max_output_tokens.is_some() + || cfg.reasoning.is_some() { router.set_model_config(&m.model_id, cfg); } diff --git a/crates/server/src/handlers/admin/settings.rs b/crates/server/src/handlers/admin/settings.rs index b373948d..7e7a3cc9 100644 --- a/crates/server/src/handlers/admin/settings.rs +++ b/crates/server/src/handlers/admin/settings.rs @@ -656,6 +656,17 @@ fn validate_setting(key: &str, value: &serde_json::Value) -> Result<(), AppError )); } } + "gateway.retry_wait_secs" => { + let v = value + .as_i64() + .ok_or_else(|| AppError::BadRequest(format!("{key} must be an integer")))?; + let max = think_watch_common::dynamic_config::RETRY_WAIT_MAX_SECS; + if !(0..=max).contains(&v) { + return Err(AppError::BadRequest(format!( + "retry_wait_secs must be between 0 and {max}" + ))); + } + } "gateway.latency_strategy_k" => { let v = value .as_f64() @@ -751,6 +762,16 @@ mod tests { assert!(validate_setting("api_keys.default_expiry_days", &json!(-1)).is_err()); } + #[test] + fn the_retry_wait_takes_zero_up_to_ten_minutes() { + assert!(validate_setting("gateway.retry_wait_secs", &json!(0)).is_ok()); + assert!(validate_setting("gateway.retry_wait_secs", &json!(120)).is_ok()); + assert!(validate_setting("gateway.retry_wait_secs", &json!(600)).is_ok()); + assert!(validate_setting("gateway.retry_wait_secs", &json!(601)).is_err()); + assert!(validate_setting("gateway.retry_wait_secs", &json!(-1)).is_err()); + assert!(validate_setting("gateway.retry_wait_secs", &json!("120")).is_err()); + } + #[test] fn validates_the_public_host_with_the_gateway_address_rule() { for ok in [ diff --git a/crates/server/src/handlers/models.rs b/crates/server/src/handlers/models.rs index 5bad2cfe..0efe9083 100644 --- a/crates/server/src/handlers/models.rs +++ b/crates/server/src/handlers/models.rs @@ -124,6 +124,12 @@ pub struct CreateModelRequest { /// for nothing gets it. NULL/missing = no limit. #[serde(default)] pub max_output_tokens: Option, + /// The reasoning setting: `off` takes reasoning out of requests that + /// have it on; `low` / `medium` / `high` / `xhigh` / `max` change the + /// effort of requests whose reasoning is on. NULL/missing = requests + /// go out as sent. + #[serde(default)] + pub reasoning_effort: Option, /// Removed: see [`refuse_output_guardrails`]. #[serde(default)] #[schema(value_type = Option>, deprecated)] @@ -175,6 +181,7 @@ pub async fn create_model( )?; refuse_output_guardrails(req.output_guardrails.as_ref())?; let max_output_tokens = max_output_tokens(req.max_output_tokens)?; + let reasoning_effort = reasoning_effort(req.reasoning_effort)?; let model = repo::insert( &state.db, @@ -188,6 +195,7 @@ pub async fn create_model( affinity_ttl_secs: req.affinity_ttl_secs, tags: req.tags.as_deref(), max_output_tokens, + reasoning_effort: reasoning_effort.as_deref(), cache_weights: cache, }, ) @@ -207,9 +215,9 @@ pub async fn create_model( // the 5-min TTL elapses. state.weight_cache.invalidate_all().await; - // The output cap is read from the router's per-model config; a route - // naming this model id may already exist. - if max_output_tokens.is_some() { + // The output cap and the reasoning setting are read from the router's + // per-model config; a route naming this model id may already exist. + if max_output_tokens.is_some() || reasoning_effort.is_some() { crate::app::rebuild_gateway_router(&state).await; } @@ -254,6 +262,11 @@ pub struct UpdateModelRequest { #[serde(default, deserialize_with = "deserialize_some")] #[schema(value_type = Option)] pub max_output_tokens: Option>, + /// PATCH-clearable reasoning setting. Absent = unchanged, JSON `null` + /// = clear (requests go out as sent), string = one of off / low / + /// medium / high / xhigh / max. + #[serde(default, deserialize_with = "deserialize_some")] + pub reasoning_effort: Option>, /// Removed: see [`refuse_output_guardrails`]. #[serde(default)] #[schema(value_type = Option>, deprecated)] @@ -292,6 +305,20 @@ pub(crate) fn max_output_tokens(n: Option) -> Result, AppError> .transpose() } +/// A model's reasoning setting as stored: one of +/// `think_watch_gateway::proxy::effort::SETTINGS`, or none — the column's +/// CHECK, as a useful 400. +pub(crate) fn reasoning_effort(value: Option) -> Result, AppError> { + use think_watch_gateway::proxy::effort::SETTINGS; + match value { + Some(v) if !SETTINGS.contains(&v.as_str()) => Err(AppError::BadRequest(format!( + "reasoning_effort must be one of: {}", + SETTINGS.join(", ") + ))), + v => Ok(v), + } +} + /// Cache weights may be zero (an upstream that does not bill cache /// reads) but not negative — the column's CHECK, as a useful 400. fn validate_cache_weights(weights: &[Option]) -> Result<(), AppError> { @@ -402,6 +429,10 @@ pub async fn update_model( None => existing.max_output_tokens, Some(inner) => max_output_tokens(inner)?, }; + let new_reasoning_effort: Option = match req.reasoning_effort { + None => existing.reasoning_effort.clone(), + Some(inner) => reasoning_effort(inner)?, + }; validate_routing_overrides( new_strategy.as_deref(), new_affinity_mode.as_deref(), @@ -423,6 +454,7 @@ pub async fn update_model( affinity_ttl_secs: new_affinity_ttl, tags: new_tags.as_deref(), max_output_tokens: new_max_output_tokens, + reasoning_effort: new_reasoning_effort.as_deref(), cache_weights: cache, }, req.enabled.unwrap_or(existing.enabled), @@ -1648,5 +1680,28 @@ mod tests { let set: UpdateModelRequest = serde_json::from_value(serde_json::json!({"max_output_tokens": 4096})).unwrap(); assert_eq!(set.max_output_tokens, Some(Some(4096))); + + assert_eq!(absent.reasoning_effort, None); + let cleared: UpdateModelRequest = + serde_json::from_value(serde_json::json!({"reasoning_effort": null})).unwrap(); + assert_eq!(cleared.reasoning_effort, Some(None)); + let set: UpdateModelRequest = + serde_json::from_value(serde_json::json!({"reasoning_effort": "max"})).unwrap(); + assert_eq!(set.reasoning_effort, Some(Some("max".into()))); + } + + #[test] + fn the_reasoning_setting_is_off_or_a_level() { + assert_eq!(reasoning_effort(None).unwrap(), None); + for ok in ["off", "low", "medium", "high", "xhigh", "max"] { + assert_eq!( + reasoning_effort(Some(ok.into())).unwrap().as_deref(), + Some(ok) + ); + } + for bad in ["minimal", "none", "on", "", "High"] { + let e = reasoning_effort(Some(bad.into())).unwrap_err(); + assert!(e.to_string().contains("reasoning_effort"), "{bad}: {e}"); + } } } diff --git a/crates/server/src/handlers/providers.rs b/crates/server/src/handlers/providers.rs index 2211274b..f54a59be 100644 --- a/crates/server/src/handlers/providers.rs +++ b/crates/server/src/handlers/providers.rs @@ -386,6 +386,9 @@ pub async fn update_provider( // path back for an operator who fixed access upstream and // re-saved their credentials. let forgotten = crate::protocol_probe::clear_for_provider(&state.db, id).await; + // And for the reasoning efforts it refused: another upstream may + // take them. + think_watch_gateway::proxy::effort::forget_provider(&state.redis, id).await; if cleared > 0 || forgotten > 0 { tracing::info!( provider = %existing.name, @@ -464,6 +467,7 @@ pub async fn delete_provider( .await?; let name = repo::name_of(&state.db, id).await?; let routes_deleted = repo::soft_delete(&state.db, id).await?; + think_watch_gateway::proxy::effort::forget_provider(&state.redis, id).await; state.audit.log( auth_user diff --git a/crates/server/src/protocol_probe.rs b/crates/server/src/protocol_probe.rs index 6b0171ca..1d42d3ea 100644 --- a/crates/server/src/protocol_probe.rs +++ b/crates/server/src/protocol_probe.rs @@ -456,6 +456,7 @@ mod tests { assert!(is_inconclusive(&GatewayError::ProviderHttpError { status: 503, message: "OpenAI: overloaded".into(), + retry_after: None, })); } @@ -491,10 +492,12 @@ mod tests { message: "OpenAI: The model 'openai.gpt-5.5' does not support \ the '/v1/chat/completions' API" .into(), + retry_after: None, })); assert!(!is_inconclusive(&GatewayError::ProviderHttpError { status: 404, message: "no such model".into(), + retry_after: None, })); } diff --git a/crates/server/src/services/model_repository.rs b/crates/server/src/services/model_repository.rs index 79ef440e..15147ef6 100644 --- a/crates/server/src/services/model_repository.rs +++ b/crates/server/src/services/model_repository.rs @@ -52,6 +52,8 @@ pub struct ModelRow { /// The most output tokens a request to this model may ask for. /// `None` ⇒ no limit. pub max_output_tokens: Option, + /// The reasoning setting (`off`, `low` … `max`). `None` ⇒ unchanged. + pub reasoning_effort: Option, } #[derive(Debug, Serialize, sqlx::FromRow, utoipa::ToSchema)] @@ -92,7 +94,7 @@ pub struct ModelRouteRow { const MODEL_COLUMNS: &str = "id, model_id, display_name, input_weight, output_weight, \ cache_read_weight, cache_write_weight, cache_write_1h_weight, \ routing_strategy, affinity_mode, affinity_ttl_secs, tags, enabled, \ - max_output_tokens"; + max_output_tokens, reasoning_effort"; /// One page of the catalog, and the total matching `search` / `status`. /// @@ -144,7 +146,7 @@ pub async fn list( m.enabled, COALESCE(rc.providers, '{{}}'::text[]) AS providers, m.routing_strategy, m.affinity_mode, m.affinity_ttl_secs, - m.max_output_tokens + m.max_output_tokens, m.reasoning_effort FROM models m LEFT JOIN LATERAL ( SELECT COUNT(*) AS route_count, @@ -186,6 +188,7 @@ pub struct ModelFields<'a> { pub affinity_ttl_secs: Option, pub tags: Option<&'a [String]>, pub max_output_tokens: Option, + pub reasoning_effort: Option<&'a str>, /// Read, 5-minute write, 1-hour write. pub cache_weights: [Option; 3], } @@ -196,8 +199,9 @@ pub async fn insert(pool: &PgPool, model_id: &str, f: &ModelFields<'_>) -> Resul (model_id, display_name, input_weight, output_weight, routing_strategy, affinity_mode, affinity_ttl_secs, tags, max_output_tokens, - cache_read_weight, cache_write_weight, cache_write_1h_weight) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + cache_read_weight, cache_write_weight, cache_write_1h_weight, + reasoning_effort) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) RETURNING {MODEL_COLUMNS}"# ); Ok(sqlx::query_as::<_, Model>(AssertSqlSafe(sql)) @@ -213,6 +217,7 @@ pub async fn insert(pool: &PgPool, model_id: &str, f: &ModelFields<'_>) -> Resul .bind(f.cache_weights[0]) .bind(f.cache_weights[1]) .bind(f.cache_weights[2]) + .bind(f.reasoning_effort) .fetch_one(pool) .await?) } @@ -245,7 +250,8 @@ pub async fn update( max_output_tokens = $10, cache_read_weight = $11, cache_write_weight = $12, - cache_write_1h_weight = $13 + cache_write_1h_weight = $13, + reasoning_effort = $14 WHERE id = $1 RETURNING {MODEL_COLUMNS}"# ); @@ -263,6 +269,7 @@ pub async fn update( .bind(f.cache_weights[0]) .bind(f.cache_weights[1]) .bind(f.cache_weights[2]) + .bind(f.reasoning_effort) .fetch_one(pool) .await?) } diff --git a/crates/test-support/src/lib.rs b/crates/test-support/src/lib.rs index fc0791f9..ba3ed974 100644 --- a/crates/test-support/src/lib.rs +++ b/crates/test-support/src/lib.rs @@ -237,6 +237,12 @@ impl TestApp { config.validate().map_err(anyhow::Error::msg)?; let mut state = init::init_state(config.clone(), db.clone(), redis, ch_client).await?; + // An upstream that fails in a test is not sent to again: waiting + // out two minutes of retries on every failing upstream would only + // slow the suite. Tests of the retries turn them on with + // `set_setting("gateway.retry_wait_secs", …)`. + fixtures::set_setting(&db, "gateway.retry_wait_secs", serde_json::json!(0)).await?; + state.dynamic_config.reload().await?; if let Some(v) = opts.url_validator { // Every outbound fetch goes through it — webhook and Kafka // deliveries included, not only the admin handlers. diff --git a/crates/test-support/tests/analytics_clickhouse.rs b/crates/test-support/tests/analytics_clickhouse.rs index ba1e0296..5e9b849b 100644 --- a/crates/test-support/tests/analytics_clickhouse.rs +++ b/crates/test-support/tests/analytics_clickhouse.rs @@ -648,7 +648,7 @@ async fn an_error_the_upstream_reports_in_its_stream_fails_the_request() { assert_failed_partway( status, &detail, - (502, "ProviderError"), + (500, "ProviderHttpError"), "The model crashed.", ); } @@ -790,7 +790,7 @@ async fn a_caller_that_leaves_after_response_failed_saw_the_request_fail() { assert_failed_partway( rows[0].0, &rows[0].1, - (502, "ProviderError"), + (500, "ProviderHttpError"), "The model crashed.", ); } @@ -957,3 +957,91 @@ async fn v1_usage_reports_the_cost_this_month_from_the_request_log() { "b={b} one call={one_call}" ); } + +/// A request that went out again after a transient failure says how many +/// times on its log row — whole answer and stream alike. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_retried_request_says_so_on_its_log_row() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let app = TestApp::spawn_with_clickhouse().await; + app.set_setting("gateway.retry_wait_secs", json!(30)).await; + let model = "gpt-retried"; + let upstream = MockServer::start().await; + // Every request is overloaded once, then answered. + let overloaded = || { + ResponseTemplate::new(503) + .insert_header("retry-after-ms", "250") + .set_body_json(json!({"error": {"message": "overloaded"}})) + }; + let chunk = |delta: Value, finish: Value| { + format!( + "data: {}\n\n", + json!({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + "usage": null}) + ) + }; + let stream = [ + chunk(json!({"content": "hi"}), Value::Null), + chunk(json!({}), json!("stop")), + "data: [DONE]\n\n".to_string(), + ] + .concat(); + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(overloaded()) + .up_to_n_times(1) + .mount(&upstream) + .await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-retried", "object": "chat.completion", "created": 1, "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10} + }))) + .up_to_n_times(1) + .mount(&upstream) + .await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(overloaded()) + .up_to_n_times(1) + .mount(&upstream) + .await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_raw(stream, "text/event-stream")) + .mount(&upstream) + .await; + let (key, user_id) = seed_upstream(&app, &upstream.uri(), "openai", model).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + // Two prompts: the second must not be the first one's cache hit. + for (stream, prompt) in [(false, "hi"), (true, "hello")] { + let resp = gw + .post( + "/v1/chat/completions", + json!({"model": model, "stream": stream, + "messages": [{"role": "user", "content": prompt}]}), + ) + .await + .unwrap(); + resp.assert_ok(); + } + assert_eq!( + upstream.received_requests().await.unwrap_or_default().len(), + 4 + ); + + let ch = app.state.clickhouse.as_ref().expect("clickhouse client"); + for (status, detail) in gateway_rows(ch, user_id, 2).await { + assert_eq!(status, 200, "{detail}"); + assert_eq!(detail["retries"], 1, "{detail}"); + } +} diff --git a/crates/test-support/tests/gateway_proxy.rs b/crates/test-support/tests/gateway_proxy.rs index f8207697..ae4a92e1 100644 --- a/crates/test-support/tests/gateway_proxy.rs +++ b/crates/test-support/tests/gateway_proxy.rs @@ -665,3 +665,325 @@ async fn an_upstream_that_refuses_added_cache_breakpoints_gets_the_request_witho assert_eq!(sent.len(), 3); assert!(!String::from_utf8_lossy(&sent[2].body).contains("cache_control")); } + +/// A Chat upstream that takes reasoning efforts up to `high`: a request +/// asking for more is refused with a 400 naming the ones it takes, the +/// way OpenAI words it. It does not take `none` either, and names the same +/// ones. +async fn chat_upstream_up_to_high() -> MockProvider { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, Request, ResponseTemplate}; + let upstream = MockProvider { + server: wiremock::MockServer::start().await, + }; + upstream + .mount( + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(|req: &Request| { + let body: Value = serde_json::from_slice(&req.body).unwrap_or_default(); + let effort = body["reasoning_effort"].as_str().unwrap_or(""); + if matches!(effort, "none" | "xhigh" | "max") { + return ResponseTemplate::new(400).set_body_json(json!({"error": { + "message": format!("Unsupported value: '{effort}' is not supported with this model. Supported values are: 'low', 'medium', and 'high'."), + "type": "invalid_request_error", + "param": "reasoning_effort", + "code": "unsupported_value", + }})); + } + let usage = json!({"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}); + if body["stream"] == true { + let chunk = |choices: Value, usage: Value| { + format!( + "data: {}\n\n", + json!({"id": "c", "object": "chat.completion.chunk", "created": 0, + "model": "m", "choices": choices, "usage": usage}) + ) + }; + let sse = [ + chunk( + json!([{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": null}]), + Value::Null, + ), + chunk(json!([{"index": 0, "delta": {}, "finish_reason": "stop"}]), Value::Null), + chunk(json!([]), usage), + "data: [DONE]\n\n".to_string(), + ] + .concat(); + ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream") + } else { + ResponseTemplate::new(200).set_body_json(json!({ + "id": "c", "object": "chat.completion", "created": 0, "model": "m", + "choices": [{"index": 0, "finish_reason": "stop", + "message": {"role": "assistant", "content": "hi"}}], + "usage": usage, + })) + } + }), + ) + .await; + upstream +} + +/// The reasoning efforts `upstream` received, in order. +async fn efforts_sent(upstream: &MockProvider) -> Vec> { + upstream + .received_requests() + .await + .iter() + .map(|r| { + let body: Value = r.body_json().unwrap(); + body["reasoning_effort"].as_str().map(str::to_string) + }) + .collect() +} + +/// `model` on `upstream` (OpenAI Chat), with the reasoning `setting`, and a +/// key: the provider's id and the key. +async fn seed_reasoning( + app: &TestApp, + upstream: &str, + model: &str, + setting: Option<&str>, +) -> (Uuid, String) { + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let provider = + fixtures::create_provider(&app.db, &unique_name("effort"), "openai", upstream, None) + .await + .unwrap(); + fixtures::create_model_and_route(&app.db, provider.id, model) + .await + .unwrap(); + sqlx::query("UPDATE models SET reasoning_effort = $1 WHERE model_id = $2") + .bind(setting) + .bind(model) + .execute(&app.db) + .await + .unwrap(); + app.rebuild_gateway_router().await; + let key = + fixtures::create_api_key(&app.db, user.user.id, "effort", &["ai_gateway"], None, None) + .await + .unwrap() + .plaintext; + (provider.id, key) +} + +/// A refused reasoning effort goes to the same upstream once more with the +/// level it named, the client gets the answer, and later requests send +/// that level straight away — streamed or not. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_refused_reasoning_effort_is_resent_once_and_remembered() { + use fred::interfaces::HashesInterface; + // The provider's new base URL is the loopback mock again. + let app = TestApp::spawn_reaching_loopback().await; + let upstream = chat_upstream_up_to_high().await; + let model = "effort-refused"; + let (provider_id, key) = seed_reasoning(&app, &upstream.uri(), model, None).await; + let gw = app.gateway_client(); + gw.set_bearer(&key); + // Sampled at a temperature: a repeated request is not a cache hit. + let ask = |stream: bool| { + json!({"model": model, "stream": stream, "temperature": 0.5, + "reasoning_effort": "xhigh", + "messages": [{"role": "user", "content": "ping"}]}) + }; + + let resp = gw.post("/v1/chat/completions", ask(false)).await.unwrap(); + resp.assert_ok(); + assert_eq!( + efforts_sent(&upstream).await, + [Some("xhigh".into()), Some("high".into())] + ); + + // Remembered once the resend was answered (written in the background). + let memory_key = format!("effort_memory:{{provider:{provider_id}}}"); + let field = format!("openai-chat:{model}"); + let mut remembered: Option = None; + for _ in 0..50 { + remembered = app.state.redis.hget(&memory_key, &field).await.unwrap(); + if remembered.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + let remembered: Value = serde_json::from_str(&remembered.expect("remembered")).unwrap(); + assert_eq!(remembered["xhigh"]["use"], "high", "{remembered}"); + + // Now straight to `high`, for a whole answer and a stream alike. + for stream in [false, true] { + let before = upstream.received_requests().await.len(); + let resp = gw.post("/v1/chat/completions", ask(stream)).await.unwrap(); + resp.assert_ok(); + assert_eq!( + efforts_sent(&upstream).await[before..], + [Some("high".to_string())], + "stream={stream}" + ); + } + + // A changed endpoint forgets: another upstream may take `xhigh`. + let (console, _) = admin_session_with_user(&app).await; + console + .patch( + &format!("/api/admin/providers/{provider_id}"), + json!({"base_url": upstream.uri()}), + ) + .await + .unwrap() + .assert_ok(); + let gone: Option = app.state.redis.hget(&memory_key, &field).await.unwrap(); + assert_eq!(gone, None); + + // A remembered fix the upstream refuses itself is replaced by the one + // that is answered next. + let now = chrono::Utc::now().timestamp(); + let _: () = app + .state + .redis + .hset( + &memory_key, + ( + field.as_str(), + json!({"xhigh": {"use": "max", "at": now}}).to_string(), + ), + ) + .await + .unwrap(); + let before = upstream.received_requests().await.len(); + gw.post("/v1/chat/completions", ask(false)) + .await + .unwrap() + .assert_ok(); + assert_eq!( + efforts_sent(&upstream).await[before..], + [Some("max".to_string()), Some("high".to_string())] + ); + let mut replaced = Value::Null; + for _ in 0..50 { + let v: Option = app.state.redis.hget(&memory_key, &field).await.unwrap(); + replaced = v + .map(|v| serde_json::from_str(&v).unwrap()) + .unwrap_or_default(); + if replaced["xhigh"]["use"] == "high" { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + assert_eq!(replaced["xhigh"]["use"], "high", "{replaced}"); +} + +/// A refused effort refused again is the answer: one resend, no more. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn an_effort_refused_twice_is_the_answer() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, ResponseTemplate}; + let app = TestApp::spawn().await; + let upstream = MockProvider { + server: wiremock::MockServer::start().await, + }; + upstream + .mount( + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({"error": { + "message": "reasoning_effort: only 'low' or 'medium' are allowed", + "type": "invalid_request_error"}}))), + ) + .await; + let (_, key) = seed_reasoning(&app, &upstream.uri(), "effort-twice", None).await; + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw + .post( + "/v1/chat/completions", + json!({"model": "effort-twice", "reasoning_effort": "high", + "messages": [{"role": "user", "content": "ping"}]}), + ) + .await + .unwrap(); + resp.assert_status(400); + assert_eq!( + efforts_sent(&upstream).await, + [Some("high".into()), Some("medium".into())] + ); +} + +/// The model's setting: a level changes requests that have reasoning on +/// and leaves the rest alone; `off` writes OpenAI's `none`, and an upstream +/// that does not take it gets the request with the lowest level it names — +/// remembered. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_models_reasoning_setting_rewrites_requests_with_reasoning_on() { + let app = TestApp::spawn().await; + let upstream = chat_upstream_up_to_high().await; + let (_, key) = seed_reasoning(&app, &upstream.uri(), "effort-medium", Some("medium")).await; + let (_, off_key) = seed_reasoning(&app, &upstream.uri(), "effort-off", Some("off")).await; + let gw = app.gateway_client(); + + gw.set_bearer(&key); + for (n, effort) in [(1, Some("low")), (2, None)] { + let mut body = json!({"model": "effort-medium", "temperature": 0.1 * f64::from(n), + "messages": [{"role": "user", "content": "ping"}]}); + if let Some(e) = effort { + body["reasoning_effort"] = json!(e); + } + gw.post("/v1/chat/completions", body) + .await + .unwrap() + .assert_ok(); + } + // `low` became `medium`; the request without reasoning stayed without. + assert_eq!(efforts_sent(&upstream).await, [Some("medium".into()), None]); + + gw.set_bearer(&off_key); + for n in [1, 2] { + gw.post( + "/v1/chat/completions", + json!({"model": "effort-off", "reasoning_effort": "high", + "temperature": 0.1 * f64::from(n), + "messages": [{"role": "user", "content": "ping"}]}), + ) + .await + .unwrap() + .assert_ok(); + if n == 1 { + // Remembered in the background. + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + } + } + assert_eq!( + efforts_sent(&upstream).await[2..], + [Some("none".into()), Some("low".into()), Some("low".into())] + ); +} + +/// `off` on a request converted for an upstream of another format: an +/// Anthropic caller's thinking reaches an OpenAI Chat upstream as `none`, +/// not as no effort at all (the model's default), and a refused `none` +/// goes out again with the lowest level the upstream names. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn off_reaches_an_upstream_of_another_format() { + let app = TestApp::spawn().await; + let upstream = chat_upstream_up_to_high().await; + let (_, key) = seed_reasoning(&app, &upstream.uri(), "effort-off-converted", Some("off")).await; + let gw = app.gateway_client(); + gw.set_bearer(&key); + gw.post( + "/v1/messages", + json!({"model": "effort-off-converted", "max_tokens": 2000, + "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}, + "messages": [{"role": "user", "content": "ping"}]}), + ) + .await + .unwrap() + .assert_ok(); + assert_eq!( + efforts_sent(&upstream).await, + [Some("none".into()), Some("low".into())] + ); +} diff --git a/crates/test-support/tests/gateway_retry.rs b/crates/test-support/tests/gateway_retry.rs new file mode 100644 index 00000000..06a421da --- /dev/null +++ b/crates/test-support/tests/gateway_retry.rs @@ -0,0 +1,566 @@ +//! Sending a request again after the upstream failed for the moment. +//! +//! Every candidate route tried and the last one overloaded, failing with a +//! server error or its connection broken, before any of the answer went +//! to the client: the gateway waits and sends again, within +//! `gateway.retry_wait_secs` of the request's arrival. A model with one +//! route — and a stream — goes back to the same route. Refusals are the +//! answer as before, and nothing is sent again once part of a stream went +//! out. +//! +//! The harness turns retries off (see `TestApp::try_spawn_with`); each +//! test here turns them on. Upstreams that say how long to wait +//! (`retry-after-ms`) keep the tests quick; one test waits out the +//! backoff itself. + +use serde_json::Value; +use std::time::{Duration, Instant}; +use think_watch_test_support::prelude::*; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +/// A model routed to `uri`, and a key that may call it. +async fn seed(app: &TestApp, uri: &str, provider_type: &str, model: &str) -> String { + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let provider = + fixtures::create_provider(&app.db, &unique_name("retry"), provider_type, uri, None) + .await + .unwrap(); + fixtures::create_model_and_route(&app.db, provider.id, model) + .await + .unwrap(); + app.rebuild_gateway_router().await; + fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("retry-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap() + .plaintext +} + +async fn with_retries(secs: i64) -> TestApp { + let app = TestApp::spawn().await; + app.set_setting("gateway.retry_wait_secs", json!(secs)) + .await; + app +} + +/// An overloaded upstream that says to come back in 250 ms. +fn overloaded() -> ResponseTemplate { + ResponseTemplate::new(503) + .insert_header("retry-after-ms", "250") + .set_body_json(json!({"error": {"message": "overloaded"}})) +} + +fn chat_ok(model: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-retry", + "object": "chat.completion", + "created": 1_700_000_000_i64, + "model": model, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello world"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10} + })) +} + +/// `first` for the first `times` requests, `then` after. +async fn upstream( + at: &str, + first: ResponseTemplate, + times: u64, + then: ResponseTemplate, +) -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(at)) + .respond_with(first) + .up_to_n_times(times) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path(at)) + .respond_with(then) + .mount(&server) + .await; + server +} + +async fn sent(server: &MockServer) -> usize { + server.received_requests().await.unwrap_or_default().len() +} + +fn chat(model: &str, stream: bool) -> Value { + json!({"model": model, "stream": stream, "messages": [{"role": "user", "content": "hi"}]}) +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_whole_answer_goes_out_again_after_an_overload() { + let app = with_retries(30).await; + let server = upstream("/v1/chat/completions", overloaded(), 2, chat_ok("retry-1")).await; + let key = seed(&app, &server.uri(), "openai", "retry-1").await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw + .post("/v1/chat/completions", chat("retry-1", false)) + .await + .unwrap(); + resp.assert_ok(); + let body: Value = resp.json().unwrap(); + assert_eq!(body["choices"][0]["message"]["content"], "hello world"); + assert_eq!(sent(&server).await, 3, "two overloads, then the answer"); +} + +/// Without a word from the upstream, the first retry waits two seconds, +/// give or take a fifth. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn without_a_word_from_the_upstream_the_retry_waits_out_the_backoff() { + let app = with_retries(30).await; + let bad_gateway = ResponseTemplate::new(502).set_body_string("bad gateway"); + let server = upstream("/v1/chat/completions", bad_gateway, 1, chat_ok("retry-2")).await; + let key = seed(&app, &server.uri(), "openai", "retry-2").await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let started = Instant::now(); + gw.post("/v1/chat/completions", chat("retry-2", false)) + .await + .unwrap() + .assert_ok(); + let took = started.elapsed(); + assert!(took >= Duration::from_millis(1600), "{took:?}"); + assert_eq!(sent(&server).await, 2); +} + +/// An upstream that says to come back at once is not asked again at +/// once: a zero wait is no word at all, and the backoff applies. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_wait_of_zero_from_the_upstream_still_waits_out_the_backoff() { + let app = with_retries(30).await; + let now = ResponseTemplate::new(503) + .insert_header("retry-after-ms", "0") + .insert_header("retry-after", "0") + .set_body_json(json!({"error": {"message": "overloaded"}})); + let server = upstream("/v1/chat/completions", now, 1, chat_ok("retry-zero")).await; + let key = seed(&app, &server.uri(), "openai", "retry-zero").await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let started = Instant::now(); + gw.post("/v1/chat/completions", chat("retry-zero", false)) + .await + .unwrap() + .assert_ok(); + let took = started.elapsed(); + assert!(took >= Duration::from_millis(1600), "{took:?}"); + assert_eq!(sent(&server).await, 2); +} + +/// A refusal — a bad request, the upstream's rate limit — is the answer: +/// the upstream is asked once. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_refusal_is_not_sent_again() { + let app = with_retries(30).await; + for (status, model) in [ + (400u16, "retry-400"), + (429, "retry-429"), + (404, "retry-404"), + ] { + let refusal = ResponseTemplate::new(status) + .insert_header("retry-after-ms", "250") + .set_body_json(json!({"error": {"message": "no"}})); + let server = upstream("/v1/chat/completions", refusal, 1, chat_ok(model)).await; + let key = seed(&app, &server.uri(), "openai", model).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw + .post("/v1/chat/completions", chat(model, false)) + .await + .unwrap(); + resp.assert_status(status); + assert_eq!(sent(&server).await, 1, "{status}"); + } +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_zero_wait_never_sends_again() { + let app = with_retries(0).await; + let server = upstream( + "/v1/chat/completions", + overloaded(), + 1, + chat_ok("retry-off"), + ) + .await; + let key = seed(&app, &server.uri(), "openai", "retry-off").await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw + .post("/v1/chat/completions", chat("retry-off", false)) + .await + .unwrap(); + resp.assert_status(503); + assert_eq!(sent(&server).await, 1); +} + +/// The retries stop when the next wait would end past the deadline, and +/// the last failure is the answer. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_retries_stop_at_the_deadline() { + let app = with_retries(1).await; + let slow_down = || { + ResponseTemplate::new(503) + .insert_header("retry-after-ms", "400") + .set_body_json(json!({"error": {"message": "overloaded"}})) + }; + let server = upstream("/v1/chat/completions", slow_down(), 100, slow_down()).await; + let key = seed(&app, &server.uri(), "openai", "retry-deadline").await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let started = Instant::now(); + let resp = gw + .post("/v1/chat/completions", chat("retry-deadline", false)) + .await + .unwrap(); + resp.assert_status(503); + assert!(started.elapsed() < Duration::from_secs(3)); + // 0 ms, ~400 ms, ~800 ms; a wait to ~1200 ms ends past the deadline. + let n = sent(&server).await; + assert!((2..=3).contains(&n), "{n}"); +} + +/// Two routes failing for the moment: both are asked again, and the one +/// that has recovered answers. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn every_route_that_failed_for_the_moment_is_asked_again() { + let app = with_retries(30).await; + let down = upstream("/v1/chat/completions", overloaded(), 100, overloaded()).await; + let back = upstream( + "/v1/chat/completions", + overloaded(), + 1, + chat_ok("retry-two"), + ) + .await; + let user = fixtures::create_random_user(&app.db).await.unwrap(); + for server in [&down, &back] { + let p = fixtures::create_provider( + &app.db, + &unique_name("retry-two"), + "openai", + &server.uri(), + None, + ) + .await + .unwrap(); + fixtures::create_model_route(&app.db, p.id, "retry-two", 100) + .await + .unwrap(); + } + app.rebuild_gateway_router().await; + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("retry-two-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap(); + + let gw = app.gateway_client(); + gw.set_bearer(&key.plaintext); + gw.post("/v1/chat/completions", chat("retry-two", false)) + .await + .unwrap() + .assert_ok(); + assert_eq!(sent(&back).await, 2, "failed once, answered the retry"); + assert!(sent(&down).await >= 1); +} + +/// A stream whose upstream cannot open goes out again: the client gets +/// one whole answer. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_stream_that_cannot_open_goes_out_again() { + let app = with_retries(30).await; + let chunk = |delta: Value, finish: Value| { + format!( + "data: {}\n\n", + json!({"id": "c", "object": "chat.completion.chunk", "created": 1, + "model": "retry-stream", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}) + ) + }; + let answer = [ + chunk(json!({"role": "assistant"}), Value::Null), + chunk(json!({"content": "hello"}), Value::Null), + chunk(json!({}), json!("stop")), + "data: [DONE]\n\n".to_string(), + ] + .concat(); + let ok = ResponseTemplate::new(200).set_body_raw(answer, "text/event-stream"); + let server = upstream("/v1/chat/completions", overloaded(), 1, ok).await; + let key = seed(&app, &server.uri(), "openai", "retry-stream").await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw + .post("/v1/chat/completions", chat("retry-stream", true)) + .await + .unwrap(); + resp.assert_ok(); + let text = resp.text(); + assert!(text.contains("hello"), "{text}"); + assert!(!text.contains("overloaded"), "{text}"); + assert_eq!(text.matches("\"role\":\"assistant\"").count(), 1, "{text}"); + assert_eq!(sent(&server).await, 2); +} + +fn anthropic_event(kind: &str, v: Value) -> String { + format!("event: {kind}\ndata: {v}\n\n") +} + +fn message_start(model: &str) -> String { + anthropic_event( + "message_start", + 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} + }}), + ) +} + +fn overload_event() -> String { + anthropic_event( + "error", + json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}), + ) +} + +fn text_block(text: &str) -> String { + [ + anthropic_event( + "content_block_start", + json!({"type": "content_block_start", "index": 0, + "content_block": {"type": "text", "text": ""}}), + ), + anthropic_event( + "content_block_delta", + json!({"type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": text}}), + ), + ] + .concat() +} + +fn messages(model: &str) -> Value { + json!({"model": model, "max_tokens": 100, "stream": true, + "messages": [{"role": "user", "content": "hi"}]}) +} + +/// Anthropic opens its stream and then reports an overload, before any +/// content: the opening was held back, the stream goes out again, and +/// the client sees one answer. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_stream_overloaded_before_any_content_goes_out_again() { + let app = with_retries(30).await; + let model = "retry-claude"; + let failed = [message_start(model), overload_event()].concat(); + let answer = [ + message_start(model), + text_block("Hello"), + anthropic_event( + "content_block_stop", + json!({"type": "content_block_stop", "index": 0}), + ), + anthropic_event( + "message_delta", + json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 2}}), + ), + anthropic_event("message_stop", json!({"type": "message_stop"})), + ] + .concat(); + let server = upstream( + "/v1/messages", + ResponseTemplate::new(200).set_body_raw(failed, "text/event-stream"), + 1, + ResponseTemplate::new(200).set_body_raw(answer, "text/event-stream"), + ) + .await; + let key = seed(&app, &server.uri(), "anthropic", model).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw.post("/v1/messages", messages(model)).await.unwrap(); + resp.assert_ok(); + let text = resp.text(); + assert_eq!(text.matches("event: message_start").count(), 1, "{text}"); + assert!(text.contains("Hello"), "{text}"); + assert!(text.contains("event: message_stop"), "{text}"); + assert!(!text.contains("overloaded_error"), "{text}"); + assert_eq!(sent(&server).await, 2); +} + +/// Part of the answer went out before the overload: the client gets it +/// and the error, and the upstream is not asked again. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_stream_that_failed_after_content_is_not_sent_again() { + let app = with_retries(30).await; + let model = "retry-claude-late"; + let failed = [message_start(model), text_block("Hel"), overload_event()].concat(); + let server = upstream( + "/v1/messages", + ResponseTemplate::new(200).set_body_raw(failed.clone(), "text/event-stream"), + 1, + ResponseTemplate::new(200).set_body_raw(failed, "text/event-stream"), + ) + .await; + let key = seed(&app, &server.uri(), "anthropic", model).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw.post("/v1/messages", messages(model)).await.unwrap(); + resp.assert_ok(); + let text = resp.text(); + assert!(text.contains("Hel"), "{text}"); + assert!(text.contains("overloaded_error"), "{text}"); + assert_eq!(sent(&server).await, 1); +} + +/// A stream the upstream refuses before any content — a bad request — +/// is the answer: the opening and the refusal reach the client, once. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_stream_refused_before_any_content_is_not_sent_again() { + let app = with_retries(30).await; + let model = "retry-claude-refused"; + let refused = [ + message_start(model), + anthropic_event( + "error", + json!({"type": "error", "error": {"type": "invalid_request_error", + "message": "prompt is too long"}}), + ), + ] + .concat(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200).set_body_raw(refused, "text/event-stream")) + .mount(&server) + .await; + let key = seed(&app, &server.uri(), "anthropic", model).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key); + let resp = gw.post("/v1/messages", messages(model)).await.unwrap(); + resp.assert_ok(); + let text = resp.text(); + assert_eq!(text.matches("event: message_start").count(), 1, "{text}"); + assert!(text.contains("prompt is too long"), "{text}"); + assert_eq!(sent(&server).await, 1); +} + +/// While a stream waits to go out again, a Chat client hears a keep-alive +/// comment after 15 seconds of silence; a Gemini client reading SSE does +/// not — the google-genai SDK reads every line as JSON. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn only_non_gemini_clients_get_keep_alives_while_a_stream_waits() { + let app = with_retries(60).await; + let chunk = |model: &str, delta: Value, finish: Value| { + format!( + "data: {}\n\n", + json!({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}) + ) + }; + let stream_of = |model: &str| { + [ + chunk(model, json!({"content": "hello"}), Value::Null), + chunk(model, json!({}), json!("stop")), + "data: [DONE]\n\n".to_string(), + ] + .concat() + }; + // Overloaded once, saying to come back in 16 s: longer than the + // keep-alive interval. + let busy = || { + ResponseTemplate::new(503) + .insert_header("retry-after-ms", "16000") + .set_body_json(json!({"error": {"message": "overloaded"}})) + }; + let mut keys = Vec::new(); + for model in ["retry-ka-chat", "retry-ka-gemini"] { + let ok = ResponseTemplate::new(200).set_body_raw(stream_of(model), "text/event-stream"); + let server = upstream("/v1/chat/completions", busy(), 1, ok).await; + keys.push((seed(&app, &server.uri(), "openai", model).await, server)); + } + + let chat = async { + reqwest::Client::new() + .post(format!("{}/v1/chat/completions", app.gateway_url)) + .bearer_auth(&keys[0].0) + .json(&chat("retry-ka-chat", true)) + .send() + .await + .unwrap() + .text() + .await + .unwrap() + }; + let gemini = async { + reqwest::Client::new() + .post(format!( + "{}/v1beta/models/retry-ka-gemini:streamGenerateContent?alt=sse", + app.gateway_url + )) + .header("x-goog-api-key", &keys[1].0) + .json(&json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}]})) + .send() + .await + .unwrap() + .text() + .await + .unwrap() + }; + let (chat, gemini) = tokio::join!(chat, gemini); + + assert!(chat.contains(": keep-alive\n\n"), "{chat}"); + assert!(chat.contains("hello"), "{chat}"); + assert!(gemini.contains("hello"), "{gemini}"); + for line in gemini.lines().filter(|l| !l.is_empty()) { + let data = line + .strip_prefix("data: ") + .unwrap_or_else(|| panic!("not a data line: {line:?}\n{gemini}")); + serde_json::from_str::(data).unwrap_or_else(|e| panic!("{e}: {data}")); + } + for (_, server) in &keys { + assert_eq!(sent(server).await, 2); + } +} diff --git a/db/schema.sql b/db/schema.sql index 877519f4..00eb4a63 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -385,6 +385,13 @@ CREATE TABLE IF NOT EXISTS models ( -- is lowered to it, and a request without one gets it. NULL ⇒ no -- limit. See crates/gateway/src/proxy/generate.rs. max_output_tokens INTEGER CHECK (max_output_tokens IS NULL OR max_output_tokens > 0), + -- The reasoning setting: 'off' takes reasoning out of requests that + -- have it on; a level changes the effort of requests whose reasoning + -- is on. NULL ⇒ requests go out as sent. See + -- crates/gateway/src/proxy/effort.rs. + reasoning_effort TEXT + CHECK (reasoning_effort IS NULL + OR reasoning_effort IN ('off', 'low', 'medium', 'high', 'xhigh', 'max')), created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); @@ -401,6 +408,9 @@ ALTER TABLE models ADD COLUMN IF NOT EXISTS cache_write_1h_weight -- after this file (crates/common/src/guard_policy/legacy.rs). ALTER TABLE models ADD COLUMN IF NOT EXISTS max_output_tokens INTEGER CHECK (max_output_tokens IS NULL OR max_output_tokens > 0); +ALTER TABLE models ADD COLUMN IF NOT EXISTS reasoning_effort + TEXT CHECK (reasoning_effort IS NULL + OR reasoning_effort IN ('off', 'low', 'medium', 'high', 'xhigh', 'max')); -- Platform-wide per-token pricing baseline. Single-row singleton -- (PK pinned to 1 via CHECK). `cost($) = tokens × weight × baseline`. diff --git a/db/seeds.sql b/db/seeds.sql index 596ef90b..d670ba24 100644 --- a/db/seeds.sql +++ b/db/seeds.sql @@ -354,6 +354,8 @@ INSERT INTO system_settings (key, value, category, description) VALUES 'Default session affinity mode (none/provider/route)'), ('gateway.default_affinity_ttl_secs','300', 'gateway', 'Default affinity key TTL in seconds (0-86400)'), + ('gateway.retry_wait_secs', '120', 'gateway', + 'How long after a request arrives it may still be sent again to upstreams that failed for the moment (overloaded, server error, broken connection) before any of the answer reached the client, in seconds (0-600; 0 = never)'), ('gateway.latency_strategy_k', '2.0', 'gateway', 'Exponent for the latency-strategy weighting (higher = more aggressive). Default 2.0 (aggressive).'), ('gateway.cb_enabled', 'true', 'gateway', diff --git a/deploy/helm/think-watch/Chart.yaml b/deploy/helm/think-watch/Chart.yaml index 4f4bfa1b..8ffa1c3d 100644 --- a/deploy/helm/think-watch/Chart.yaml +++ b/deploy/helm/think-watch/Chart.yaml @@ -2,8 +2,8 @@ apiVersion: v2 name: think-watch description: Enterprise AI API Gateway & MCP Management Platform type: application -version: 3.9.1 -appVersion: "3.9.1" +version: 3.10.0 +appVersion: "3.10.0" keywords: - ai - gateway diff --git a/web/package.json b/web/package.json index dd575fe4..266747ef 100644 --- a/web/package.json +++ b/web/package.json @@ -1,7 +1,7 @@ { "name": "web", "private": true, - "version": "3.9.1", + "version": "3.10.0", "type": "module", "packageManager": "pnpm@12.14.0", "scripts": { diff --git a/web/src/components/segmented.tsx b/web/src/components/segmented.tsx index 4136ee05..3697176a 100644 --- a/web/src/components/segmented.tsx +++ b/web/src/components/segmented.tsx @@ -22,6 +22,7 @@ export function Segmented({ label, disabled = false, size = 'default', + fill = false, className, }: { value: T; @@ -31,6 +32,8 @@ export function Segmented({ label: string; disabled?: boolean; size?: 'sm' | 'default'; + /** Take the whole row, the options sharing it equally. */ + fill?: boolean; className?: string; }) { const refs = useRef<(HTMLButtonElement | null)[]>([]); @@ -57,7 +60,8 @@ export function Segmented({ aria-disabled={disabled || undefined} onKeyDown={onKeyDown} className={cn( - 'inline-flex items-center gap-0.5 rounded-md border bg-muted/30 p-0.5', + fill ? 'flex w-full' : 'inline-flex', + 'items-center gap-0.5 rounded-md border bg-muted/30 p-0.5', size === 'sm' ? 'text-xs' : 'text-sm', className, )} @@ -77,7 +81,8 @@ export function Segmented({ disabled={disabled || o.disabled} onClick={() => !checked && onChange(o.value)} className={cn( - 'rounded px-2.5 font-medium whitespace-nowrap transition-colors outline-none focus-visible:ring-2 focus-visible:ring-ring/50 disabled:cursor-not-allowed disabled:opacity-50', + fill ? 'flex-auto px-1.5' : 'px-2.5', + 'rounded font-medium whitespace-nowrap transition-colors outline-none focus-visible:ring-2 focus-visible:ring-ring/50 disabled:cursor-not-allowed disabled:opacity-50', size === 'sm' ? 'py-0.5' : 'py-1', checked ? 'bg-background text-foreground shadow-sm' diff --git a/web/src/i18n/en.json b/web/src/i18n/en.json index b49de707..d290e064 100644 --- a/web/src/i18n/en.json +++ b/web/src/i18n/en.json @@ -864,13 +864,19 @@ "cacheReadWeight": "Cache read weight", "cacheWriteWeight": "Cache write weight", "cacheWrite1hWeight": "1-hour cache write weight", - "maxOutputTokens": "Max output tokens" + "maxOutputTokens": "Max output tokens", + "reasoningEffort": "Reasoning" }, "useGlobalDefault": "Use global default", "unlimited": "Unlimited", "routingOverrideTitle": "Routing overrides", "routingOverrideHint": "Override the gateway-wide defaults for this model. Leave blank to inherit from Settings → Gateway.", "maxOutputTokensHint": "The most output tokens a request to this model can ask for. A request that sets a limit gets the smaller of the two; one that sets none gets this value if it is within the model family's default (32,000 for Claude, 8,192 for others), and the model's own default otherwise. Leave blank for no limit.", + "reasoningEffortHint": "A level applies only to requests that already have reasoning on; requests without reasoning (such as title generation) are left unchanged.", + "reasoning": { + "keep": "Unchanged", + "off": "Off" + }, "strategy": { "weighted": "Weighted random", "latency": "Latency-first", @@ -1816,6 +1822,8 @@ "defaultAffinityTtlHint": "How long an affinity assignment lasts. 0 disables affinity entirely.", "latencyStrategyK": "Latency aggressiveness (k)", "latencyStrategyKHint": "Exponent in 1/latency^k. Higher = more aggressive bias toward fast routes. Default 2.0 (aggressive). Range 0.5–8.0.", + "retryWaitSecs": "Retry window (seconds)", + "retryWaitSecsHint": "When every route failed for the moment (overloaded, a server error, a broken connection) before any of the answer reached the client, the request is sent again after a short wait, for up to this long after it arrived. 0 disables retries. Range 0–600.", "strategy": { "weighted": "Weighted random", "latency": "Latency-first", diff --git a/web/src/i18n/zh.json b/web/src/i18n/zh.json index 272d5000..2ace1a87 100644 --- a/web/src/i18n/zh.json +++ b/web/src/i18n/zh.json @@ -864,13 +864,19 @@ "cacheReadWeight": "缓存读取权重", "cacheWriteWeight": "缓存写入权重", "cacheWrite1hWeight": "1 小时缓存写入权重", - "maxOutputTokens": "最大输出 token" + "maxOutputTokens": "最大输出 token", + "reasoningEffort": "推理" }, "useGlobalDefault": "使用全局默认", "unlimited": "不限制", "routingOverrideTitle": "路由覆盖", "routingOverrideHint": "针对此模型覆盖网关全局默认。留空则继承「设置 → 网关」中的配置。", "maxOutputTokensHint": "发往此模型的请求最多输出这么多 token。请求写了上限时取两者中较小的；未写时，此值不超过该模型系列的默认上限（Claude 32,000，其他 8,192）则按此值发送，否则按模型默认。留空则不限制。", + "reasoningEffortHint": "指定档位时，仅改写已开启推理的请求；未开启推理的请求（如生成标题）保持不变。", + "reasoning": { + "keep": "不改变", + "off": "关闭" + }, "strategy": { "weighted": "加权随机", "latency": "延迟优先", @@ -1816,6 +1822,8 @@ "defaultAffinityTtlHint": "黏性键的存活时间。0 = 完全关闭黏性。", "latencyStrategyK": "延迟激进度 (k)", "latencyStrategyKHint": "权重 ∝ 1/延迟^k。k 越大对慢路由的惩罚越激进。默认 2.0(激进)。范围 0.5–8.0。", + "retryWaitSecs": "重试时限 (秒)", + "retryWaitSecsHint": "所有路由都暂时失败（过载、服务器错误、连接中断）且回答尚未有任何内容发给客户端时，稍等后重发请求，直到请求到达后的这段时间为止。0 = 不重试。范围 0–600。", "strategy": { "weighted": "加权随机", "latency": "延迟优先", diff --git a/web/src/routes/admin/settings.tsx b/web/src/routes/admin/settings.tsx index d5b63fc1..a7e7f614 100644 --- a/web/src/routes/admin/settings.tsx +++ b/web/src/routes/admin/settings.tsx @@ -105,6 +105,7 @@ export function SettingsPage() { const [defaultAffinityMode, setDefaultAffinityMode] = useState('provider'); const [defaultAffinityTtlSecs, setDefaultAffinityTtlSecs] = useState(300); const [latencyStrategyK, setLatencyStrategyK] = useState(2.0); + const [retryWaitSecs, setRetryWaitSecs] = useState(120); const [cbEnabled, setCbEnabled] = useState(true); const [cbErrorPct, setCbErrorPct] = useState(50); const [cbMinSamples, setCbMinSamples] = useState(10); @@ -162,6 +163,7 @@ export function SettingsPage() { setDefaultAffinityMode(str(getSettingValue(data, 'gateway', 'default_affinity_mode'), 'provider')); setDefaultAffinityTtlSecs(num(getSettingValue(data, 'gateway', 'default_affinity_ttl_secs'), 300)); setLatencyStrategyK(Number(getSettingValue(data, 'gateway', 'latency_strategy_k')) || 2.0); + setRetryWaitSecs(num(getSettingValue(data, 'gateway', 'retry_wait_secs'), 120)); setCbEnabled(getSettingValue(data, 'gateway', 'cb_enabled') !== false); setCbErrorPct(num(getSettingValue(data, 'gateway', 'cb_error_pct'), 50)); setCbMinSamples(num(getSettingValue(data, 'gateway', 'cb_min_samples'), 10)); @@ -258,6 +260,7 @@ export function SettingsPage() { const defaultAffinityModeSave = useFieldAutosave({ value: defaultAffinityMode, isLoaded, persist: (v) => patchOne('gateway.default_affinity_mode', v), debounceMs: 0 }); const defaultAffinityTtlSave = useFieldAutosave({ value: defaultAffinityTtlSecs, isLoaded, persist: (v) => patchOne('gateway.default_affinity_ttl_secs', v) }); const latencyStrategyKSave = useFieldAutosave({ value: latencyStrategyK, isLoaded, persist: (v) => patchOne('gateway.latency_strategy_k', v) }); + const retryWaitSecsSave = useFieldAutosave({ value: retryWaitSecs, isLoaded, persist: (v) => patchOne('gateway.retry_wait_secs', v) }); const cbEnabledSave = useFieldAutosave({ value: cbEnabled, isLoaded, persist: (v) => patchOne('gateway.cb_enabled', v), debounceMs: 0 }); const cbErrorPctSave = useFieldAutosave({ value: cbErrorPct, isLoaded, persist: (v) => patchOne('gateway.cb_error_pct', v) }); const cbMinSamplesSave = useFieldAutosave({ value: cbMinSamples, isLoaded, persist: (v) => patchOne('gateway.cb_min_samples', v) }); @@ -662,6 +665,15 @@ export function SettingsPage() { />

{t('settings.latencyStrategyKHint')}

+ } + /> diff --git a/web/src/routes/gateway/models/ModelDetailSheet.tsx b/web/src/routes/gateway/models/ModelDetailSheet.tsx index 19b816e0..53fe83f5 100644 --- a/web/src/routes/gateway/models/ModelDetailSheet.tsx +++ b/web/src/routes/gateway/models/ModelDetailSheet.tsx @@ -164,6 +164,16 @@ export function ModelDetailSheet({ +
+
{t('models.field.reasoningEffort')}
+
+ {model.reasoning_effort == null + ? t('models.reasoning.keep') + : model.reasoning_effort === 'off' + ? t('models.reasoning.off') + : model.reasoning_effort} +
+
{CACHE_WEIGHTS.map((k) => (
diff --git a/web/src/routes/gateway/models/ModelEditorDialog.test.tsx b/web/src/routes/gateway/models/ModelEditorDialog.test.tsx index 0513142c..9da31cb8 100644 --- a/web/src/routes/gateway/models/ModelEditorDialog.test.tsx +++ b/web/src/routes/gateway/models/ModelEditorDialog.test.tsx @@ -86,3 +86,39 @@ describe('ModelEditorDialog max output tokens', () => { expect(apiPatch).not.toHaveBeenCalled() }) }) + +describe('ModelEditorDialog reasoning', () => { + it('leaves requests unchanged by default and sends a chosen level', async () => { + const user = userEvent.setup() + renderEditor(model) + + const group = screen.getByRole('radiogroup', { name: 'Reasoning' }) + expect(screen.getByRole('radio', { name: 'Unchanged' })).toHaveAttribute('aria-checked', 'true') + expect(group).toBeInTheDocument() + await user.click(screen.getByRole('radio', { name: 'max' })) + await user.click(screen.getByRole('button', { name: 'Save' })) + + await waitFor(() => + expect(apiPatch).toHaveBeenCalledWith( + '/api/admin/models/m-1', + expect.objectContaining({ reasoning_effort: 'max' }), + ), + ) + }) + + it('shows the stored setting and clears it back to unchanged', async () => { + const user = userEvent.setup() + renderEditor({ ...model, reasoning_effort: 'off' }) + + expect(screen.getByRole('radio', { name: 'Off' })).toHaveAttribute('aria-checked', 'true') + await user.click(screen.getByRole('radio', { name: 'Unchanged' })) + await user.click(screen.getByRole('button', { name: 'Save' })) + + await waitFor(() => + expect(apiPatch).toHaveBeenCalledWith( + '/api/admin/models/m-1', + expect.objectContaining({ reasoning_effort: null }), + ), + ) + }) +}) diff --git a/web/src/routes/gateway/models/ModelEditorDialog.tsx b/web/src/routes/gateway/models/ModelEditorDialog.tsx index 2d8d0f6f..debc86f4 100644 --- a/web/src/routes/gateway/models/ModelEditorDialog.tsx +++ b/web/src/routes/gateway/models/ModelEditorDialog.tsx @@ -13,6 +13,7 @@ import { } from '@/components/ui/dialog'; import { Input } from '@/components/ui/input'; import { Label } from '@/components/ui/label'; +import { Segmented } from '@/components/segmented'; import { Select, SelectContent, @@ -28,6 +29,7 @@ import { AFFINITY_MODES, CACHE_WEIGHTS, MAX_OUTPUT_TOKENS_CEILING, + REASONING_LEVELS, ROUTING_STRATEGIES, derivedCacheWeight, emptyModelForm, @@ -37,6 +39,7 @@ import { type ModelFormState, type ModelRow, type PlatformPricing, + type ReasoningEffort, type RoutingStrategy, } from './types'; @@ -89,6 +92,7 @@ export function ModelEditorDialog({ affinity_mode: (model.affinity_mode ?? '') as ModelFormState['affinity_mode'], affinity_ttl_secs: model.affinity_ttl_secs == null ? '' : String(model.affinity_ttl_secs), max_output_tokens: model.max_output_tokens == null ? '' : String(model.max_output_tokens), + reasoning_effort: model.reasoning_effort ?? '', }); } else { setForm(emptyModelForm); @@ -143,6 +147,7 @@ export function ModelEditorDialog({ affinity_mode: form.affinity_mode === '' ? null : form.affinity_mode, affinity_ttl_secs: ttlNum, max_output_tokens: maxOutputTokens, + reasoning_effort: form.reasoning_effort === '' ? null : form.reasoning_effort, }; setSaving(true); try { @@ -336,6 +341,24 @@ export function ModelEditorDialog({ />

{t('models.maxOutputTokensHint')}

+
+ +
+ + size="sm" + fill + label={t('models.field.reasoningEffort')} + value={form.reasoning_effort === '' ? 'keep' : form.reasoning_effort} + options={[ + { value: 'keep', label: t('models.reasoning.keep') }, + { value: 'off', label: t('models.reasoning.off') }, + ...REASONING_LEVELS.map((l) => ({ value: l, label: l })), + ]} + onChange={(v) => setForm({ ...form, reasoning_effort: v === 'keep' ? '' : v })} + /> +
+

{t('models.reasoningEffortHint')}

+
{error && ( diff --git a/web/src/routes/gateway/models/types.ts b/web/src/routes/gateway/models/types.ts index af32b694..e354107c 100644 --- a/web/src/routes/gateway/models/types.ts +++ b/web/src/routes/gateway/models/types.ts @@ -37,8 +37,21 @@ export interface ModelRow { /// gets it when it is within the model family's default (32,000 for /// Claude, 8,192 for others). null ⇒ no limit. max_output_tokens?: number | null; + /// Reasoning setting; null ⇒ requests go out as sent. + reasoning_effort?: ReasoningEffort | null; } +/// `off` takes reasoning out of requests that have it on; a level +/// changes the effort of requests whose reasoning is on. Mirrors +/// `SETTINGS` in `crates/gateway/src/proxy/effort.rs`. +export type ReasoningEffort = 'off' | 'low' | 'medium' | 'high' | 'xhigh' | 'max'; + +export const REASONING_EFFORTS: ReasoningEffort[] = ['off', 'low', 'medium', 'high', 'xhigh', 'max']; + +/// The five levels are shown as the APIs write them; only "off" is +/// translated. +export const REASONING_LEVELS: ReasoningEffort[] = ['low', 'medium', 'high', 'xhigh', 'max']; + /// Upper bound the form accepts for `max_output_tokens` — the largest /// value a Postgres `INTEGER` column holds. export const MAX_OUTPUT_TOKENS_CEILING = 2_147_483_647; @@ -156,6 +169,8 @@ export interface ModelFormState { affinity_ttl_secs: string; /// Empty string ⇒ no limit (stored as NULL). max_output_tokens: string; + /// Empty string ⇒ unchanged (stored as NULL). + reasoning_effort: '' | ReasoningEffort; } export interface RouteFormState { @@ -183,6 +198,7 @@ export const emptyModelForm: ModelFormState = { affinity_mode: '', affinity_ttl_secs: '', max_output_tokens: '', + reasoning_effort: '', }; export const emptyRouteForm: RouteFormState = {