diff --git a/CHANGELOG.md b/CHANGELOG.md index 44c100b1..67b31f38 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,40 @@ target. ## [Unreleased] +### Fixed + +- **A conversation whose history a provider encrypted goes back to that + provider.** OpenAI's Responses API with `store: false` (Codex always + sends it) hands the client encrypted state that comes back every turn: + reasoning, remote compactions, Codex's messages between agents, + encrypted tool output. Only the account that encrypted it can read it, + so a conversation that moved to another provider — after five idle + minutes, a tripped breaker or a failover — was refused on every turn + from then on with `invalid_encrypted_content`. The gateway now + remembers, for 30 days and whatever the affinity settings say, the + provider that last answered each conversation (the authenticated user + and the request's `x-session-id`, or else the session header the client + sends by itself, such as Codex's `session_id`; over HTTP and over the + WebSocket). The log row's `session_id` still comes from `x-session-id` + only. A + Responses request carrying encrypted content that cannot be left out + goes only to that provider: no strategy pick, no failover, and its open + breaker does not stop it. When that provider has no route for the model + any more, the request is refused up front with 409, naming the provider + and saying to start a new conversation. A request carrying only + encrypted reasoning goes to that provider first and still fails over. + Requests that name no conversation are routed as before. The request's log + row records `"pinned_by": "encrypted_history"` when the provider was + picked this way. +- **A provider that cannot read a conversation's encrypted history says + so.** Its refusal, as a 4xx answer or as the error event of a stream + that opened with 200, now reaches the caller and the log as + `SealedHistoryRefused`: the provider cannot read this conversation's + encrypted history because another provider encrypted it; route the + conversation back to that provider or start a new conversation. The + upstream's own words follow. It no longer counts against the route's + breaker. + ## [3.9.0] — 2026-10-20 The Add to ThinkWatch Lite link from a rotated key now carries the key's diff --git a/Cargo.lock b/Cargo.lock index f34d5d3c..2166287a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2554,7 +2554,7 @@ version = "5.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" dependencies = [ - "base64 0.22.1", + "base64 0.21.7", "chrono", "getrandom 0.2.17", "http", @@ -4837,8 +4837,8 @@ dependencies = [ [[package]] name = "tw-bedrock" -version = "0.77.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.77.0#00c9a892675c75c97e29609d762e63accb37fa34" +version = "0.80.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?rev=366833934f58c490fa5eceb8e7b6d3a3d979c719#366833934f58c490fa5eceb8e7b6d3a3d979c719" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -4852,16 +4852,16 @@ dependencies = [ [[package]] name = "tw-breaker" -version = "0.77.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.77.0#00c9a892675c75c97e29609d762e63accb37fa34" +version = "0.80.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?rev=366833934f58c490fa5eceb8e7b6d3a3d979c719#366833934f58c490fa5eceb8e7b6d3a3d979c719" dependencies = [ "serde", ] [[package]] name = "tw-dialect" -version = "0.77.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.77.0#00c9a892675c75c97e29609d762e63accb37fa34" +version = "0.80.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?rev=366833934f58c490fa5eceb8e7b6d3a3d979c719#366833934f58c490fa5eceb8e7b6d3a3d979c719" dependencies = [ "memchr", "serde", @@ -4870,8 +4870,8 @@ dependencies = [ [[package]] name = "tw-guard" -version = "0.77.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.77.0#00c9a892675c75c97e29609d762e63accb37fa34" +version = "0.80.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?rev=366833934f58c490fa5eceb8e7b6d3a3d979c719#366833934f58c490fa5eceb8e7b6d3a3d979c719" dependencies = [ "aho-corasick", "base64 0.22.1", diff --git a/Cargo.toml b/Cargo.toml index e582ef98..b013c84c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -70,10 +70,16 @@ 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.77.0" } -tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.77.0" } -tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.77.0" } -tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.77.0" } +# +# TEMPORARY: on core main at the merge of core PR #368 (`tw_dialect::sealed`) +# until core tags it — replace with the next core tag before this merges. +# Set on the dependencies themselves rather than in a `[patch]` table: core's +# CI appends its own `[patch]` table for this URL, and a second one would not +# parse. +tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", rev = "366833934f58c490fa5eceb8e7b6d3a3d979c719" } +tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", rev = "366833934f58c490fa5eceb8e7b6d3a3d979c719" } +tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", rev = "366833934f58c490fa5eceb8e7b6d3a3d979c719" } +tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", rev = "366833934f58c490fa5eceb8e7b6d3a3d979c719" } # Web framework axum = { version = "0.8", features = ["macros", "ws"] } diff --git a/crates/gateway/src/error.rs b/crates/gateway/src/error.rs index 56d88a1a..03b3b44b 100644 --- a/crates/gateway/src/error.rs +++ b/crates/gateway/src/error.rs @@ -69,6 +69,32 @@ pub enum GatewayError { /// pass it on. Maps to 403. #[error("Blocked by policy: {0}")] PolicyBlocked(String), + /// The upstream refused the conversation's encrypted history (OpenAI + /// Responses `encrypted_content`, see `tw_dialect::sealed`): another + /// provider's account encrypted it, and only that one can read it. + /// `status` is the upstream's (as an answer, or what its error in a + /// stream would have been), `said` its own words, truncated. Not the + /// upstream failing: no failover, nothing against its breaker. + #[error( + "Provider `{provider}` cannot read this conversation's encrypted history: another \ + provider encrypted it, and only that one can read it. Route the conversation back to \ + that provider, or start a new conversation. The provider said: {said}" + )] + SealedHistoryRefused { + status: u16, + provider: String, + said: String, + }, + /// The conversation's history is encrypted by `provider`, which has no + /// route for the model any more, and no other provider can read it. + /// Refused before anything is sent. Maps to 409: the request is well + /// formed, but the conversation cannot go on here. + #[error( + "This conversation's history is encrypted by provider `{provider}`, which can no \ + longer serve model `{model}`, and no other provider can read it. Start a new \ + conversation." + )] + SealedHistoryUnavailable { provider: String, model: String }, } impl GatewayError { @@ -90,6 +116,8 @@ impl GatewayError { GatewayError::UpstreamRateLimited { .. } | GatewayError::LocalRateLimited { .. } => 429, GatewayError::UpstreamAuthError { .. } => 401, GatewayError::PolicyBlocked(_) => 403, + GatewayError::SealedHistoryRefused { status, .. } => i64::from(*status), + GatewayError::SealedHistoryUnavailable { .. } => 409, } } @@ -108,6 +136,8 @@ impl GatewayError { GatewayError::LocalRateLimited { .. } => "LocalRateLimited", GatewayError::UpstreamAuthError { .. } => "UpstreamAuthError", GatewayError::PolicyBlocked(_) => "PolicyBlocked", + GatewayError::SealedHistoryRefused { .. } => "SealedHistoryRefused", + GatewayError::SealedHistoryUnavailable { .. } => "SealedHistoryUnavailable", } } diff --git a/crates/gateway/src/lifecycle/mod.rs b/crates/gateway/src/lifecycle/mod.rs index 3013224e..98e73f59 100644 --- a/crates/gateway/src/lifecycle/mod.rs +++ b/crates/gateway/src/lifecycle/mod.rs @@ -114,6 +114,10 @@ pub(crate) struct ChatRequestSnapshot { /// The request's input in tokens, estimated — billed only when the /// upstream reports no usage. pub input_estimate: u64, + /// The conversation a Responses request with a session id belongs to. + /// The provider that answers it is remembered for it (see + /// `proxy::sealed`). + pub conversation: Option, } /// Pre-flight rule + cap lists, reused by the post-flight debit. @@ -123,7 +127,12 @@ pub(crate) struct ChatPreflightLists { /// The route that actually served the request. pub(crate) struct ChatPickedRoute { + pub provider_id: uuid::Uuid, pub provider_name: String, + /// Picked because it holds the conversation's encrypted history, not by + /// the strategy or the affinity (see `proxy::sealed`). Said on the log + /// row as `"pinned_by": "encrypted_history"`. + pub pinned: bool, /// Upstream-side model id when the route remapped it. pub upstream_model: Option, /// Used by `finalize_health` inside `record_outcome`. @@ -286,8 +295,31 @@ pub(crate) fn build_chat_pump( }; match item { Ok(chunk) => { - if let Ok(mut r) = readers.lock() { - r.feed(&chunk); + let sealed = match readers.lock() { + Ok(mut r) => { + r.feed(&chunk); + r.sealed_refusal() + } + 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), @@ -528,6 +560,12 @@ impl Readers { } } + /// The upstream's refusal of the conversation's encrypted history, once + /// its stream has said it. Left in place for the tail. + fn sealed_refusal(&self) -> Option { + self.said.as_ref().filter(|s| s.sealed).cloned() + } + /// 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. @@ -575,13 +613,16 @@ impl ErrorWatch { } /// An error the upstream reported in its stream. -#[derive(Debug, PartialEq)] +#[derive(Debug, Clone, PartialEq)] struct Said { /// In its words, or [`NO_MESSAGE`] when the error carried none. message: String, /// The status the same error has as an answer, when the error names /// one (see [`status_of`]). status: Option, + /// The error refuses the conversation's encrypted history: another + /// provider sealed it (see `proxy::sealed`). + sealed: bool, } /// The message of an error the upstream reported without one. @@ -638,6 +679,7 @@ fn said_in(upstream: Dialect, f: &tw_dialect::frame::Frame) -> Option { Some(Said { message: message.to_string(), status: error.and_then(status_of), + sealed: crate::proxy::sealed::refused(upstream, &f.data), }) } @@ -708,6 +750,14 @@ fn failed_partway(provider: &str, said: &Said) -> StreamOutcome { "upstream reported an error partway through a stream" ); let words: String = said.message.chars().take(500).collect(); + if said.sealed { + let e = sealed_refused(provider, said); + return StreamOutcome::UpstreamError { + error_type: e.error_tag().to_string(), + message: e.to_string(), + status_code: e.status_code(), + }; + } let what = format!("{provider} reported an error partway through the answer: {words}"); let e = match said.status { Some(status) => crate::proxy::transport::status_error(status, None, what.clone()), @@ -727,6 +777,17 @@ fn failed_partway(provider: &str, said: &Said) -> StreamOutcome { } } +/// The upstream refused the conversation's encrypted history in its +/// stream: the error said plainly, with its words (500 characters at most). +/// A refusal that names no status is a bad request's. +fn sealed_refused(provider: &str, said: &Said) -> GatewayError { + GatewayError::SealedHistoryRefused { + status: said.status.unwrap_or(400), + provider: provider.to_string(), + said: said.message.chars().take(500).collect(), + } +} + /// A stream that broke off in transit: what the client is told, and the /// outcome. No usable answer arrived — the upstream's failure, 502. fn broke_off(message: String) -> (String, StreamOutcome) { @@ -759,6 +820,7 @@ fn bedrock_ended( let said = Said { message: format!("{kind}: {message}"), status: Some(bedrock_status(&kind)), + sealed: false, }; (told, failed_partway(provider, &said)) } @@ -957,6 +1019,25 @@ impl Surface for ChatCompletionSurface { CapturedView::Buffered(ChatCompletionOutcome::ShortCircuit(_)) => false, }; finalize_health(&deps.state, &deps.route.sel_record, success).await; + + // The provider that answered holds whatever it sealed in the answer + // (see `proxy::sealed`). Only a whole answer: one cut short never + // reached the client to be sent back. + let answered = matches!( + &invoked.view, + CapturedView::Streaming { + outcome: StreamOutcome::Natural, + .. + } | CapturedView::Buffered(ChatCompletionOutcome::Success(_)) + ); + if answered && let Some(c) = &deps.request.conversation { + c.answered_by( + &deps.state.redis, + deps.route.provider_id, + &deps.route.provider_name, + ) + .await; + } } async fn write_cache(deps: &Self::PostInvokeDeps, invoked: &Invoked) { @@ -1074,7 +1155,11 @@ impl Surface for ChatCompletionSurface { cost, deps.request.request_started_at.elapsed().as_millis() as i64, logged_status, - with_usage_detail(error_detail, &usage, usage_estimated(&invoked.view)), + with_usage_detail( + pinned_detail(error_detail, deps.route.pinned), + &usage, + usage_estimated(&invoked.view), + ), body_capture, ); } @@ -1110,6 +1195,20 @@ fn with_usage_detail( Some(serde_json::Value::Object(extra)) } +/// The audit detail, with `"pinned_by": "encrypted_history"` when the +/// route was picked because it holds the conversation's encrypted history. +fn pinned_detail(detail: Option, pinned: bool) -> Option { + if !pinned { + return detail; + } + let mut d = match detail { + Some(serde_json::Value::Object(d)) => d, + _ => serde_json::Map::new(), + }; + d.insert("pinned_by".into(), "encrypted_history".into()); + Some(serde_json::Value::Object(d)) +} + /// The usage a captured view bills — shared by `record_usage` and /// `emit_audit` so the budget and the audit row can never disagree. fn extract_usage(view: &CapturedView) -> tw_dialect::usage::Usage { @@ -1221,6 +1320,65 @@ mod tests { format!("event: response.failed\ndata: {v}\n\n") } + /// A refusal of the conversation's encrypted history in the stream — + /// as `response.failed` or as the Codex backend's bare `error` event — + /// is said plainly, keeps the upstream's words, and is a refusal, not + /// the route failing. + #[test] + fn a_refusal_of_encrypted_history_in_a_stream_is_said_plainly() { + let failed = responses_failed(serde_json::json!({ + "code": "invalid_encrypted_content", + "message": "The encrypted content for item rs_68 could not be verified.", + })); + let bare = "event: error\ndata: {\"type\":\"error\",\"message\":\"Encrypted function output content could not be decrypted or decoded.\"}\n\n"; + for frames in [failed.as_str(), bare] { + let mut r = Readers { + errors: Some(ErrorWatch::new(Dialect::Responses)), + ..Default::default() + }; + r.feed(frames.as_bytes()); + assert!(r.sealed_refusal().is_some(), "{frames}"); + let (error_type, status, message, counted) = + recorded(Dialect::Responses, frames).unwrap(); + assert_eq!( + (error_type.as_str(), status), + ("SealedHistoryRefused", 400), + "{frames}" + ); + assert!(!counted); + assert!( + message + .starts_with("Provider `up` cannot read this conversation's encrypted history"), + "{message}" + ); + assert!(message.contains("could not be"), "{message}"); + } + + // Another error is not one. + let other = responses_failed(serde_json::json!({ + "code": "invalid_prompt", "message": "Invalid prompt", + })); + let mut r = Readers { + errors: Some(ErrorWatch::new(Dialect::Responses)), + ..Default::default() + }; + r.feed(other.as_bytes()); + assert!(r.sealed_refusal().is_none()); + } + + #[test] + fn a_pinned_route_is_said_on_the_log_row() { + assert_eq!(pinned_detail(None, false), None); + assert_eq!( + pinned_detail(None, true), + Some(serde_json::json!({"pinned_by": "encrypted_history"})) + ); + assert_eq!( + pinned_detail(Some(serde_json::json!({"error_type": "x"})), true), + Some(serde_json::json!({"error_type": "x", "pinned_by": "encrypted_history"})) + ); + } + /// Throttling stays a 429, and counts against the route like a 429 /// answer does: that upstream's quota, which another route does not /// share. @@ -1338,6 +1496,7 @@ mod tests { let said = Said { message: "x".repeat(2_000), status: None, + sealed: false, }; let StreamOutcome::UpstreamError { message, .. } = failed_partway("anthropic-main", &said) else { @@ -1356,6 +1515,7 @@ mod tests { let said = Said { message: "arn:aws:iam::123456789012:user/gateway may not".into(), status: Some(403), + sealed: false, }; let StreamOutcome::UpstreamError { error_type, diff --git a/crates/gateway/src/proxy/generate.rs b/crates/gateway/src/proxy/generate.rs index 8245c9c6..b679b3ad 100644 --- a/crates/gateway/src/proxy/generate.rs +++ b/crates/gateway/src/proxy/generate.rs @@ -46,13 +46,16 @@ use tw_dialect::ir::{Dialect, Target}; use super::body_capture::prepare_body_capture; use super::early_cancel::EarlyCancelSlot; -use super::headers::{request_id_header, resolve_session_id, resolve_trace_id}; +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::routing::{ 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}; @@ -529,6 +532,8 @@ pub(crate) async fn send( let first = outbound .address(entry.protocol, model, official)? .for_upstream(upstream, model); + // 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, @@ -562,6 +567,7 @@ pub(crate) async fn send( let wire = outbound .address(*protocol, model, official)? .for_upstream(upstream, model); + wire_dialect = wire.dialect; match wire.send_to(upstream, call_ctx).await { Ok(resp) => { super::protocol_relearn::persist(db, entry.route_id, *protocol).await; @@ -577,7 +583,13 @@ pub(crate) async fn send( } } } - Err(last) + // A refusal of the conversation's encrypted history, said plainly. + Err(super::sealed::explain_refusal( + wire_dialect, + &entry.provider_name, + &upstream.label, + last, + )) } /// Read a whole answer and put it in the caller's format. @@ -969,6 +981,40 @@ async fn run( ) } }; + // A conversation's encrypted history goes back to the provider that + // can read it (see `super::sealed`). Only Responses seals history, and + // only a request that names its conversation belongs to one. + let conversation = (surface.dialect == Dialect::Responses) + .then(|| { + Conversation::of( + identity.user_id.as_deref(), + resolve_conversation_id(&headers).as_deref(), + ) + }) + .flatten(); + let pin = match (&conversation, tw_dialect::sealed::sealed(&outbound_body)) { + (Some(c), Some(sealed)) => { + let holder = c.holder(&state.redis).await; + let pin = super::sealed::decide( + Some(sealed), + holder.as_ref(), + routes.iter().map(|r| r.provider_id), + &mapped_model, + ) + .map_err(|e| ctx.emit(e))?; + if let (Some(pin), Some(h)) = (pin, &holder) { + tracing::info!( + request_id = %metadata.request_id, + provider = %h.provider_name, + ?sealed, + ?pin, + "Routing to the provider holding the conversation's encrypted history" + ); + } + pin + } + _ => None, + }; let outbound = Outbound { surface, path: path.to_string(), @@ -992,12 +1038,15 @@ async fn run( cache_fingerprint: cache_fingerprint.clone(), request_started_at, input_estimate, + conversation: conversation.clone(), }, preflight: crate::lifecycle::ChatPreflightLists { limits: preflight.limits.clone(), }, route: crate::lifecycle::ChatPickedRoute { + provider_id: route.provider_id, provider_name: route.provider_name.clone(), + pinned: pin.is_some_and(|p| p.provider_id() == route.provider_id), upstream_model: route.upstream_model.clone(), sel_record, caps: crate::route_caps::RouteCaps::of(route), @@ -1008,7 +1057,8 @@ async fn run( 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. - let sel_ctx = build_selection_ctx(&state, &mapped_model, identity.user_id.as_deref()).await; + 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)))?; @@ -1051,7 +1101,8 @@ async fn run( } // Buffered: full failover across healthy candidates. - let sel_ctx = build_selection_ctx(&state, &mapped_model, identity.user_id.as_deref()).await; + let sel_ctx = + build_selection_ctx(&state, &mapped_model, identity.user_id.as_deref(), pin).await; // Built before failover so the synchronous error closure can move it // in. Error paths capture the request only — nothing succeeded. let error_capture = prepare_body_capture( diff --git a/crates/gateway/src/proxy/headers.rs b/crates/gateway/src/proxy/headers.rs index 70e4ae84..e3354e5b 100644 --- a/crates/gateway/src/proxy/headers.rs +++ b/crates/gateway/src/proxy/headers.rs @@ -24,8 +24,38 @@ pub(super) fn resolve_trace_id(headers: &HeaderMap) -> String { /// store NULL so the ClickHouse `idx_session` bloom filter stays /// selective for rows that do carry an id. pub(super) fn resolve_session_id(headers: &HeaderMap) -> Option { + header_id(headers, "x-session-id") +} + +/// The headers coding agents name their conversation in by themselves: +/// Codex's `session_id` (and its older `conversation_id`), Claude Code's, +/// OpenCode's. The same list the desktop gateway reads. +const CLIENT_SESSION_HEADERS: &[&str] = &[ + "session_id", + "session-id", + "conversation_id", + "x-claude-code-session-id", + "x-session-affinity", + "x-opencode-session", +]; + +/// The conversation a request belongs to, for remembering which provider +/// holds its encrypted history (`super::sealed`): `x-session-id` when the +/// caller sent one, else the id the client names its conversation with on +/// its own — Codex sends `session_id` and no `x-session-id`. Validated as +/// `x-session-id` is. Only for routing; the log row keeps +/// [`resolve_session_id`]. +pub(super) fn resolve_conversation_id(headers: &HeaderMap) -> Option { + resolve_session_id(headers).or_else(|| { + CLIENT_SESSION_HEADERS + .iter() + .find_map(|name| header_id(headers, name)) + }) +} + +fn header_id(headers: &HeaderMap, name: &str) -> Option { headers - .get("x-session-id") + .get(name) .and_then(|v| v.to_str().ok()) .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty() && s.len() <= 128 && s.chars().all(|c| !c.is_control())) @@ -84,6 +114,21 @@ mod tests { assert_eq!(id.len(), 36, "blank header should fall back to UUID"); } + #[test] + fn a_conversation_is_named_by_x_session_id_or_by_the_client_itself() { + let mut h = HeaderMap::new(); + assert_eq!(resolve_conversation_id(&h), None); + h.insert("session_id", "codex-1".parse().unwrap()); + assert_eq!(resolve_conversation_id(&h).as_deref(), Some("codex-1")); + // Not a session id for the log row. + assert_eq!(resolve_session_id(&h), None); + h.insert("x-session-id", "mine".parse().unwrap()); + assert_eq!(resolve_conversation_id(&h).as_deref(), Some("mine")); + let mut blank = HeaderMap::new(); + blank.insert("session_id", " ".parse().unwrap()); + assert_eq!(resolve_conversation_id(&blank), None); + } + /// 128-char boundary: exactly 128 should pass, 129 should reject. #[test] fn resolve_trace_id_boundary() { diff --git a/crates/gateway/src/proxy/mod.rs b/crates/gateway/src/proxy/mod.rs index 76ef7f97..de83c925 100644 --- a/crates/gateway/src/proxy/mod.rs +++ b/crates/gateway/src/proxy/mod.rs @@ -32,6 +32,7 @@ mod pipeline; mod protocol_relearn; mod responses_ws; mod routing; +pub(crate) mod sealed; pub mod shaper; pub mod transport; diff --git a/crates/gateway/src/proxy/routing.rs b/crates/gateway/src/proxy/routing.rs index 183ab6e6..c9d29a99 100644 --- a/crates/gateway/src/proxy/routing.rs +++ b/crates/gateway/src/proxy/routing.rs @@ -1,6 +1,8 @@ //! Multi-route selection: strategy-driven weights + circuit-breaker //! filter + per-route capacity caps (`crate::route_caps`) + per-model -//! session affinity (none / provider / route). +//! session affinity (none / provider / route) — and, above all of them, +//! the provider that holds a conversation's encrypted history +//! (`super::sealed`). use rand::RngExt; use std::str::FromStr; @@ -14,6 +16,8 @@ use crate::route_caps::{self, Admission, CappedRoutes, RouteCaps}; use crate::router::{AffinityMode, RouteEntry}; use crate::strategy::{self, RoutingStrategy}; +use super::sealed::Pin; + /// What the affinity layer can pin a session to. #[derive(Debug, Clone, Copy)] enum AffinityHit { @@ -125,6 +129,10 @@ pub(super) struct SelectionCtx<'a> { pub(super) breaker: CircuitBreakerConfig, /// `security.rate_limit_fail_closed`, for the route caps. pub(super) fail_closed: bool, + /// The provider holding the conversation's encrypted history, when + /// the request carries some (see `super::sealed::decide`). Overrides + /// the strategy and the affinity. + pub(super) pin: Option, pub(super) state: &'a GatewayState, } @@ -132,6 +140,7 @@ pub(super) async fn build_selection_ctx<'a>( state: &'a GatewayState, model_id: &'a str, user_id: Option<&'a str>, + pin: Option, ) -> SelectionCtx<'a> { let (strategy, affinity_mode, affinity_ttl_secs) = resolve_routing_config(state, model_id).await; @@ -147,6 +156,7 @@ pub(super) async fn build_selection_ctx<'a>( latency_k, breaker, fail_closed, + pin, state, } } @@ -186,9 +196,10 @@ async fn pick_with_strategy<'a>( let mut excluded: Vec = Vec::with_capacity(group.len()); for (i, entry) in group.iter().enumerate() { let h = &healths[i]; - let excl = h.state == tw_breaker::State::Open - || tried.contains(&entry.provider_id) - || capped.contains(&entry.route_id); + // A provider that alone can read the history is asked even with + // its breaker open: there is nowhere else to go. + let broken = h.state == tw_breaker::State::Open && !matches!(ctx.pin, Some(Pin::Only(_))); + let excl = broken || tried.contains(&entry.provider_id) || capped.contains(&entry.route_id); excluded.push(excl); let success_rate = if h.total > 0 { Some((1.0 - h.error_pct / 100.0).clamp(0.0, 1.0)) @@ -203,6 +214,7 @@ async fn pick_with_strategy<'a>( } let weights = strategy::compute_weights(ctx.strategy, &signals, ctx.latency_k); + pin_to_holder(group, &mut excluded, ctx.pin); // If affinity points to a still-eligible candidate, use it. if let Some(hit) = affinity { @@ -249,6 +261,26 @@ async fn pick_with_strategy<'a>( picked_idx.map(|i| group[i]) } +/// Narrow the candidates to the provider holding the conversation's +/// encrypted history: always for [`Pin::Only`], for [`Pin::First`] while +/// one of its routes is still eligible. +fn pin_to_holder(group: &[&RouteEntry], excluded: &mut [bool], pin: Option) { + let Some(pin) = pin else { return }; + let holder = pin.provider_id(); + let eligible = group + .iter() + .zip(excluded.iter()) + .any(|(e, x)| !*x && e.provider_id == holder); + if matches!(pin, Pin::First(_)) && !eligible { + return; + } + for (e, x) in group.iter().zip(excluded.iter_mut()) { + if e.provider_id != holder { + *x = true; + } + } +} + /// What the proxy handler needs back from a selection in order to /// record health for the picked route after the request completes. #[derive(Clone)] @@ -300,7 +332,9 @@ pub(crate) async fn finalize_health(state: &GatewayState, sel: &SelectionRecord, /// gateway draws the same line. /// /// Refusals by the gateway itself (`TransformError`, `PolicyBlocked`, -/// local limits) are not the upstream's doing either. +/// local limits) are not the upstream's doing either. Nor is an upstream +/// refusing a conversation's encrypted history (`SealedHistoryRefused`): +/// the history belongs to another provider, not to a broken route. pub(crate) fn is_upstream_failure(err: &GatewayError) -> bool { fails(err.error_tag(), err.status_code()) } @@ -311,7 +345,11 @@ pub(crate) fn is_upstream_failure(err: &GatewayError) -> bool { pub(crate) fn fails(error_tag: &str, status: i64) -> bool { match error_tag { "ProviderHttpError" => status >= 500 || status == 408, - "TransformError" | "PolicyBlocked" | "LocalRateLimited" => false, + "TransformError" + | "PolicyBlocked" + | "LocalRateLimited" + | "SealedHistoryRefused" + | "SealedHistoryUnavailable" => false, _ => true, } } @@ -462,7 +500,11 @@ pub(super) async fn select_route_with_failover<'a>( // 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 { .. }) { + if matches!( + e, + GatewayError::ProviderHttpError { .. } + | GatewayError::SealedHistoryRefused { .. } + ) { let _ = ctx .state .health @@ -513,3 +555,65 @@ pub(super) async fn select_route_for_stream<'a>( GatewayError::ProviderError(format!("No provider found for model: {}", ctx.model_id)) })) } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + fn entry(provider_id: Uuid) -> RouteEntry { + RouteEntry { + upstream: Arc::new(crate::proxy::transport::Upstream::new( + "https://example.invalid", + Vec::new(), + crate::proxy::transport::Shape::Standard, + "p", + )), + provider_id, + route_id: Uuid::new_v4(), + provider_name: "p".into(), + upstream_model: None, + weight: 100, + label: None, + rpm_cap: None, + tpm_cap: None, + protocol: crate::protocol::UpstreamProtocol::OpenAiResponses, + alternates: Vec::new(), + } + } + + /// Which of `a`, `b`, `a` (a second route of the holder) are left. + fn left(pin: Option, a: Uuid, excluded: [bool; 3]) -> [bool; 3] { + let routes = [entry(a), entry(Uuid::new_v4()), entry(a)]; + let group: Vec<&RouteEntry> = routes.iter().collect(); + let mut excluded = excluded; + pin_to_holder(&group, &mut excluded, pin); + excluded.map(|x| !x) + } + + #[test] + fn a_pin_narrows_the_candidates_to_the_holders_routes() { + let a = Uuid::new_v4(); + let none = [false; 3]; + assert_eq!(left(None, a, none), [true, true, true]); + assert_eq!(left(Some(Pin::Only(a)), a, none), [true, false, true]); + assert_eq!(left(Some(Pin::First(a)), a, none), [true, false, true]); + } + + #[test] + fn only_sealed_reasoning_falls_over_once_the_holder_is_out() { + let a = Uuid::new_v4(); + // The holder tried (or capped, or broken): reasoning moves on… + let out = [true, false, true]; + assert_eq!(left(Some(Pin::First(a)), a, out), [false, true, false]); + // …anything else sealed has nowhere to go. + assert_eq!(left(Some(Pin::Only(a)), a, out), [false, false, false]); + } + + #[test] + fn sealed_refusals_and_refusing_up_front_do_not_count_against_the_route() { + assert!(!fails("SealedHistoryRefused", 400)); + assert!(!fails("SealedHistoryUnavailable", 409)); + assert!(fails("ProviderHttpError", 500)); + } +} diff --git a/crates/gateway/src/proxy/sealed.rs b/crates/gateway/src/proxy/sealed.rs new file mode 100644 index 00000000..28acdeef --- /dev/null +++ b/crates/gateway/src/proxy/sealed.rs @@ -0,0 +1,350 @@ +//! Conversations whose history an upstream sealed. +//! +//! OpenAI's Responses API, asked not to store a conversation (`store: +//! false` — Codex always asks that), hands the client what the next turn +//! needs encrypted, and the client sends it back every turn: reasoning, +//! remote compactions, what Codex's agents said to each other, encrypted +//! tool output. The key is OpenAI's, per account. Sent to another provider +//! — another account, a relay — the whole request is refused +//! (`invalid_encrypted_content`), and since the client resends the whole +//! history, so is every turn after it. Which parts are sealed, and how an +//! upstream says it cannot read them, is `tw_dialect::sealed`, shared with +//! the desktop gateway. +//! +//! So each conversation — the authenticated user and the caller's session +//! id — remembers the provider that last answered it ([`Conversation`]), +//! for [`BINDING_TTL_SECS`], whatever the affinity settings say: this is +//! not a preference for a warm cache, it is the one provider that can read +//! the history. A request carrying sealed history is then routed by +//! [`decide`]: +//! +//! * only sealed reasoning ([`Sealed::Reasoning`]): the provider that holds +//! it goes first; the others stay for failover — reasoning is the +//! model's own notes, and the conversation goes on without it; +//! * anything else ([`Sealed::Content`]): only that provider. Another +//! cannot read it, so there is no failover; one with no route for the +//! model any more is refused up front, with what to do about it. +//! +//! A conversation with no session id, or none remembered, is routed as +//! any other request; an upstream that then cannot read the history says +//! so in plain words (see [`explain_refusal`]). + +use fred::interfaces::KeysInterface; +use tw_dialect::ir::Dialect; +use tw_dialect::sealed::Sealed; +use uuid::Uuid; + +use crate::error::GatewayError; + +/// How long the provider that answered a conversation is remembered: 30 +/// days. A conversation picked up again days later (Codex's `resume`) still +/// carries the history only that provider can read. +pub(crate) const BINDING_TTL_SECS: i64 = 30 * 24 * 60 * 60; + +/// One conversation: an authenticated user's session. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct Conversation { + key: String, +} + +impl Conversation { + /// The conversation a request belongs to. `None` without a user or a + /// session id: nothing ties its turns together. + pub(crate) fn of(user_id: Option<&str>, session_id: Option<&str>) -> Option { + Some(Self { + key: format!("sealed_holder:{}:{}", user_id?, session_id?), + }) + } + + /// The provider that last answered this conversation, if remembered. + /// A Redis that cannot be read remembers nothing: the request is routed + /// as any other. + pub(crate) async fn holder(&self, redis: &fred::clients::Client) -> Option { + let value: Option = redis.get(&self.key).await.ok().flatten(); + Holder::parse(&value?) + } + + /// `provider_id` (`provider_name`) answered this conversation. + /// Best-effort. + pub(crate) async fn answered_by( + &self, + redis: &fred::clients::Client, + provider_id: Uuid, + provider_name: &str, + ) { + let value = Holder { + provider_id, + provider_name: provider_name.to_string(), + } + .encode(); + if let Err(e) = redis + .set::<(), _, _>( + &self.key, + value, + Some(fred::types::Expiration::EX(BINDING_TTL_SECS)), + None, + false, + ) + .await + { + tracing::warn!(error = %e, "could not remember which provider answered a conversation"); + } + } +} + +/// The provider that last answered a conversation. Its name is kept for +/// the refusal when it has no route left to name it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct Holder { + pub provider_id: Uuid, + pub provider_name: String, +} + +impl Holder { + /// ` `: the id has no space in it. + fn encode(&self) -> String { + format!("{} {}", self.provider_id, self.provider_name) + } + + fn parse(value: &str) -> Option { + let (id, name) = value.split_once(' ')?; + Some(Self { + provider_id: Uuid::parse_str(id).ok()?, + provider_name: name.to_string(), + }) + } +} + +/// Where a request carrying sealed history may go. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Pin { + /// Only this provider: no other can read the history. Its open breaker + /// does not keep the request from it — there is nowhere else to go, and + /// its own answer says more than "no route" would. + Only(Uuid), + /// This provider first; the others for failover. + First(Uuid), +} + +impl Pin { + pub(crate) fn provider_id(self) -> Uuid { + match self { + Pin::Only(id) | Pin::First(id) => id, + } + } +} + +/// Where a request goes, given what it carries sealed (`sealed`), the +/// provider that last answered its conversation (`holder`), and the +/// providers of the routes it may take (`providers`). +/// +/// `Ok(None)`: routed as any other request — nothing sealed, or nobody +/// known to hold it. Sealed reasoning whose holder has no route left is +/// routed as any other too: the conversation goes on without it. Anything +/// else sealed by a provider with no route left is refused: no provider can +/// read it. +pub(crate) fn decide( + sealed: Option, + holder: Option<&Holder>, + mut providers: impl Iterator, + model: &str, +) -> Result, GatewayError> { + let (Some(sealed), Some(holder)) = (sealed, holder) else { + return Ok(None); + }; + let routed = providers.any(|p| p == holder.provider_id); + match (sealed, routed) { + (Sealed::Content, true) => Ok(Some(Pin::Only(holder.provider_id))), + (Sealed::Content, false) => Err(GatewayError::SealedHistoryUnavailable { + provider: holder.provider_name.clone(), + model: model.to_string(), + }), + (Sealed::Reasoning, true) => Ok(Some(Pin::First(holder.provider_id))), + (Sealed::Reasoning, false) => Ok(None), + } +} + +/// An upstream's refusal of sealed history, said plainly: the upstream's +/// own words ("Encrypted function output content could not be decrypted") +/// do not tell anyone that the conversation came from another provider. +/// +/// `wire` is the format the request went out in. Only Responses is read, +/// as on the desktop: that is where the history is sealed per account. +/// `label` is the prefix the transport puts before the upstream's body +/// (see `transport::check_status`). Any other error is returned as it is. +pub(crate) fn explain_refusal( + wire: Dialect, + provider: &str, + label: &str, + e: GatewayError, +) -> GatewayError { + match e { + GatewayError::ProviderHttpError { status, message } + if (400..500).contains(&status) && refused(wire, &message) => + { + let said = message + .strip_prefix(label) + .and_then(|m| m.strip_prefix(": ")) + .unwrap_or(&message) + .to_string(); + GatewayError::SealedHistoryRefused { + status, + provider: provider.to_string(), + said, + } + } + e => e, + } +} + +/// Whether `said`, an error an upstream sent in `wire`, refuses sealed +/// history. +pub(crate) fn refused(wire: Dialect, said: &str) -> bool { + wire == Dialect::Responses && tw_dialect::sealed::refused(wire, said.as_bytes()) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn holder(id: Uuid) -> Holder { + Holder { + provider_id: id, + provider_name: "openai-team a".into(), + } + } + + #[test] + fn a_conversation_needs_a_user_and_a_session() { + assert!(Conversation::of(Some("u"), Some("s")).is_some()); + assert!(Conversation::of(None, Some("s")).is_none()); + assert!(Conversation::of(Some("u"), None).is_none()); + assert_ne!( + Conversation::of(Some("u1"), Some("s")), + Conversation::of(Some("u2"), Some("s")), + "two users' sessions are two conversations" + ); + } + + #[test] + fn the_holder_round_trips_with_a_name_that_has_spaces() { + let h = holder(Uuid::new_v4()); + assert_eq!(Holder::parse(&h.encode()), Some(h)); + assert_eq!(Holder::parse("not-a-uuid name"), None); + assert_eq!(Holder::parse(""), None); + } + + #[test] + fn content_only_another_provider_cannot_read_goes_only_to_its_holder() { + let (a, b) = (Uuid::new_v4(), Uuid::new_v4()); + let h = holder(a); + assert_eq!( + decide(Some(Sealed::Content), Some(&h), [b, a].into_iter(), "m").unwrap(), + Some(Pin::Only(a)) + ); + // The holder has no route left: nobody can read it. + match decide(Some(Sealed::Content), Some(&h), [b].into_iter(), "gpt-5") { + Err(e @ GatewayError::SealedHistoryUnavailable { .. }) => { + assert_eq!(e.status_code(), 409); + let said = e.to_string(); + assert!( + said.contains("`openai-team a`") && said.contains("`gpt-5`"), + "{said}" + ); + assert!(said.contains("Start a new conversation"), "{said}"); + } + other => panic!("{other:?}"), + } + } + + #[test] + fn sealed_reasoning_prefers_its_holder_and_can_do_without_it() { + let (a, b) = (Uuid::new_v4(), Uuid::new_v4()); + let h = holder(a); + assert_eq!( + decide(Some(Sealed::Reasoning), Some(&h), [b, a].into_iter(), "m").unwrap(), + Some(Pin::First(a)) + ); + assert_eq!( + decide(Some(Sealed::Reasoning), Some(&h), [b].into_iter(), "m").unwrap(), + None + ); + } + + #[test] + fn nothing_sealed_or_nobody_known_routes_as_usual() { + let a = Uuid::new_v4(); + assert_eq!( + decide(None, Some(&holder(a)), [a].into_iter(), "m").unwrap(), + None + ); + assert_eq!( + decide(Some(Sealed::Content), None, [a].into_iter(), "m").unwrap(), + None + ); + } + + #[test] + fn a_refusal_of_sealed_history_is_said_plainly() { + let body = json!({"error": { + "message": "The encrypted content for item rs_68 could not be verified.", + "type": "invalid_request_error", "code": "invalid_encrypted_content"}}) + .to_string(); + let e = GatewayError::ProviderHttpError { + status: 400, + message: format!("relay-b: {body}"), + }; + let e = explain_refusal(Dialect::Responses, "relay-b", "relay-b", e); + let GatewayError::SealedHistoryRefused { + status, + ref provider, + ref said, + } = e + else { + panic!("{e:?}"); + }; + assert_eq!((status, provider.as_str()), (400, "relay-b")); + assert_eq!(said, &body, "the upstream's own words, without the label"); + assert_eq!(e.status_code(), 400); + assert_eq!(e.error_tag(), "SealedHistoryRefused"); + let told = e.to_string(); + assert!(told.starts_with("Provider `relay-b` cannot read"), "{told}"); + assert!(told.contains("another provider encrypted it"), "{told}"); + assert!(told.contains("could not be verified"), "{told}"); + assert!( + !crate::proxy::upstream_failed(e.error_tag(), e.status_code()), + "a refusal, not the upstream failing" + ); + } + + #[test] + fn other_errors_and_other_formats_are_left_as_they_are() { + let other = || GatewayError::ProviderHttpError { + status: 400, + message: "p: max_output_tokens is too large".into(), + }; + assert!(matches!( + explain_refusal(Dialect::Responses, "p", "p", other()), + GatewayError::ProviderHttpError { .. } + )); + let sealed = || GatewayError::ProviderHttpError { + status: 400, + message: "p: invalid_encrypted_content".into(), + }; + // Not sent as Responses: nothing of OpenAI's was sealed in it. + assert!(matches!( + explain_refusal(Dialect::Chat, "p", "p", sealed()), + GatewayError::ProviderHttpError { .. } + )); + // A server error is the upstream failing, whatever it says. + let e = GatewayError::ProviderHttpError { + status: 500, + message: "p: invalid_encrypted_content".into(), + }; + assert!(matches!( + explain_refusal(Dialect::Responses, "p", "p", e), + GatewayError::ProviderHttpError { .. } + )); + } +} diff --git a/crates/test-support/tests/gateway_responses_ws.rs b/crates/test-support/tests/gateway_responses_ws.rs index 4373e722..fd7fc773 100644 --- a/crates/test-support/tests/gateway_responses_ws.rs +++ b/crates/test-support/tests/gateway_responses_ws.rs @@ -357,3 +357,46 @@ async fn an_answer_rule_filters_each_turn_across_frames() { } socket.close(None).await.unwrap(); } + +/// A turn on a socket opened with `x-session-id` belongs to that +/// conversation like an HTTP request does: the provider that answered it +/// is remembered for the encrypted history it may have handed back, and +/// the next connection's turns go back to it (see `gateway_sealed_history`). +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_turn_remembers_which_provider_answered_its_conversation() { + use fred::interfaces::KeysInterface; + + let app = TestApp::spawn().await; + let upstream = MockProvider::openai_chat_stream_ok("ws-sealed").await; + let (user_id, key) = seed(&app, &upstream.uri(), "ws-sealed").await; + let provider_id: Uuid = + sqlx::query_scalar("SELECT provider_id FROM model_routes WHERE model_id = $1") + .bind("ws-sealed") + .fetch_one(&app.db) + .await + .unwrap(); + + let mut req = ws_url(&app).into_client_request().unwrap(); + req.headers_mut() + .insert("authorization", format!("Bearer {key}").parse().unwrap()); + req.headers_mut() + .insert("x-session-id", "ws-conv".parse().unwrap()); + let (mut socket, _) = tokio_tungstenite::connect_async(req).await.unwrap(); + socket.send(create("ws-sealed", "hi")).await.unwrap(); + let events = turn(&mut socket).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + socket.close(None).await.unwrap(); + + let binding = format!("sealed_holder:{user_id}:ws-conv"); + let mut held: Option = None; + for _ in 0..50 { + held = app.state.redis.get(&binding).await.unwrap(); + if held.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + let held = held.expect("the answering provider is remembered"); + assert!(held.starts_with(&provider_id.to_string()), "{held}"); +} diff --git a/crates/test-support/tests/gateway_sealed_history.rs b/crates/test-support/tests/gateway_sealed_history.rs new file mode 100644 index 00000000..cc7ef0db --- /dev/null +++ b/crates/test-support/tests/gateway_sealed_history.rs @@ -0,0 +1,434 @@ +//! A conversation whose history an upstream encrypted (OpenAI Responses +//! `encrypted_content`, sent with `store: false` as Codex does) goes back +//! to the provider that answered it: only that account can read it. +//! +//! Two providers serve the same model over Responses. The first turn of a +//! conversation (`x-session-id`) is answered by `a`; later turns carry what +//! `a` encrypted, while the routing would send them to `b` — by weight, or +//! by an affinity pointing there. + +use fred::interfaces::KeysInterface; +use think_watch_test_support::prelude::*; +use wiremock::matchers::{body_partial_json, method, path}; +use wiremock::{Mock, ResponseTemplate}; + +const MODEL: &str = "sealed-model"; + +/// A Responses upstream that answers with `tag` — buffered, or streamed +/// when asked to stream. +async fn responses_upstream(tag: &str) -> MockProvider { + let mock = MockProvider { + server: wiremock::MockServer::start().await, + }; + let answer = json!({ + "id": format!("resp_{tag}"), "object": "response", "status": "completed", + "model": MODEL, + "output": [{"type": "message", "id": "msg_1", "role": "assistant", + "content": [{"type": "output_text", "text": tag}]}], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + }); + mock.mount( + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(body_partial_json(json!({"stream": true}))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + think_watch_test_support::mock_provider::responses_answer(MODEL, false), + "text/event-stream", + )) + .with_priority(3), + ) + .await; + mock.mount( + Mock::given(method("POST")) + .and(path("/v1/responses")) + .respond_with(ResponseTemplate::new(200).set_body_json(answer)) + .with_priority(5), + ) + .await; + mock +} + +/// From now on `upstream` fails every request with a 500. +async fn breaks(upstream: &MockProvider) { + upstream + .mount( + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(500).set_body_json(json!({"error": {"message": "down"}})), + ) + .with_priority(1), + ) + .await; +} + +/// OpenAI's refusal of content another account encrypted: as an answer, +/// and as the error event of a stream the Codex backend opens with 200. +async fn refuses_sealed_history(upstream: &MockProvider) { + let failed = json!({"type": "response.failed", "response": {"id": "resp_r", "status": "failed", + "error": {"code": "invalid_encrypted_content", + "message": "Encrypted function output content could not be decrypted or decoded."}}}); + upstream + .mount( + Mock::given(method("POST")) + .and(body_partial_json(json!({"stream": true}))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + format!("event: response.failed\ndata: {failed}\n\n"), + "text/event-stream", + )) + .with_priority(1), + ) + .await; + upstream + .mount( + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({"error": { + "message": "The encrypted content for item rs_68 could not be verified.", + "type": "invalid_request_error", "param": null, + "code": "invalid_encrypted_content"}}))) + .with_priority(2), + ) + .await; +} + +async fn hits(upstream: &MockProvider) -> usize { + upstream.received_requests().await.len() +} + +struct Setup { + app: TestApp, + a: MockProvider, + b: MockProvider, + a_id: Uuid, + a_name: String, + b_id: Uuid, + user_id: Uuid, + key: String, +} + +/// `a` serves [`MODEL`] over Responses; `b` joins once the first turn has +/// gone to `a` (see [`Setup::add_b`]). +async fn setup() -> Setup { + let app = TestApp::spawn().await; + let a = responses_upstream("from-a").await; + let b = responses_upstream("from-b").await; + let a_name = unique_name("acct-a"); + let pa = fixtures::create_provider(&app.db, &a_name, "openai", &a.uri(), None) + .await + .unwrap(); + let pb = fixtures::create_provider(&app.db, &unique_name("acct-b"), "openai", &b.uri(), None) + .await + .unwrap(); + fixtures::create_model_route(&app.db, pa.id, MODEL, 1) + .await + .unwrap(); + responses_routes(&app).await; + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let key = fixtures::create_api_key(&app.db, user.user.id, "s", &["ai_gateway"], None, None) + .await + .unwrap(); + Setup { + app, + a, + b, + a_id: pa.id, + a_name, + b_id: pb.id, + user_id: user.user.id, + key: key.plaintext, + } +} + +async fn responses_routes(app: &TestApp) { + sqlx::query( + "UPDATE model_routes SET upstream_protocol = 'openai_responses' WHERE model_id = $1", + ) + .bind(MODEL) + .execute(&app.db) + .await + .unwrap(); + app.rebuild_gateway_router().await; +} + +impl Setup { + /// `b` joins with a weight a thousand times `a`'s. + async fn add_b(&self) { + fixtures::create_model_route(&self.app.db, self.b_id, MODEL, 1000) + .await + .unwrap(); + responses_routes(&self.app).await; + } + + fn client(&self, session: Option<&str>) -> TestClient { + let gw = self.app.gateway_client(); + gw.set_bearer(&self.key); + if let Some(s) = session { + gw.set_header("x-session-id", s); + } + gw + } + + /// The provider-mode affinity of this user for the model points at `b`. + async fn affinity_to_b(&self) { + let key = format!("affinity:provider:{}:{MODEL}", self.user_id); + let _: () = self + .app + .state + .redis + .set(&key, self.b_id.to_string(), None, None, false) + .await + .unwrap(); + } + + fn binding_key(&self, session: &str) -> String { + format!("sealed_holder:{}:{session}", self.user_id) + } + + /// The provider remembered for `session`, once it is (a stream's is + /// written by its tail, after the response). + async fn holder(&self, session: &str) -> Option { + for _ in 0..50 { + let v: Option = self + .app + .state + .redis + .get(self.binding_key(session)) + .await + .unwrap(); + if v.is_some() { + return v; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + None + } +} + +fn first_turn(stream: bool) -> Json { + json!({"model": MODEL, "store": false, "stream": stream, "input": "hi"}) +} + +/// A turn carrying a remote compaction `a` encrypted: no other provider +/// can read it, and it cannot be left out. +fn content_turn() -> Json { + json!({"model": MODEL, "store": false, "input": [ + {"type": "compaction", "encrypted_content": "gAAAAABcompacted"}, + {"type": "message", "role": "user", "content": "go on"} + ]}) +} + +/// A turn carrying only reasoning `a` encrypted. +fn reasoning_turn() -> Json { + json!({"model": MODEL, "store": false, "input": [ + {"role": "user", "content": "hi"}, + {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAABthought"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "x"}]}, + {"role": "user", "content": "go on"} + ]}) +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn encrypted_content_stays_with_its_provider_whatever_the_routing_prefers() { + let s = setup().await; + let gw = s.client(Some("conv-1")); + + // The first turn, streamed, goes to `a` — the only route so far. + let resp = gw.post("/v1/responses", first_turn(true)).await.unwrap(); + resp.assert_ok(); + assert!(resp.text().contains("response.completed")); + let holder = s + .holder("conv-1") + .await + .expect("the answering provider is remembered"); + assert!(holder.starts_with(&s.a_id.to_string()), "{holder}"); + assert!(holder.ends_with(&s.a_name), "{holder}"); + + // `b` is preferred by weight and by the affinity; the conversation's + // history is `a`'s. + s.add_b().await; + s.affinity_to_b().await; + for _ in 0..3 { + gw.post("/v1/responses", content_turn()) + .await + .unwrap() + .assert_ok(); + } + assert_eq!(hits(&s.a).await, 4); + assert_eq!(hits(&s.b).await, 0, "sent where it cannot be read"); + + // A conversation without a session id, or one nobody remembers, is + // routed as before: the affinity sends it to `b`. + s.affinity_to_b().await; + s.client(None) + .post("/v1/responses", content_turn()) + .await + .unwrap() + .assert_ok(); + let _: i64 = s + .app + .state + .redis + .del(s.binding_key("conv-1")) + .await + .unwrap(); + s.affinity_to_b().await; + gw.post("/v1/responses", content_turn()) + .await + .unwrap() + .assert_ok(); + assert_eq!(hits(&s.a).await, 4); + assert_eq!(hits(&s.b).await, 2); + // …and is remembered again by whoever answered it. + let holder = s.holder("conv-1").await.unwrap(); + assert!(holder.starts_with(&s.b_id.to_string()), "{holder}"); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn encrypted_content_does_not_fail_over_to_a_provider_that_cannot_read_it() { + let s = setup().await; + let gw = s.client(Some("conv-2")); + gw.post("/v1/responses", first_turn(false)) + .await + .unwrap() + .assert_ok(); + s.add_b().await; + breaks(&s.a).await; + + // `a` answers with its own error, and nothing goes to `b` — buffered + // or streamed. + let resp = gw.post("/v1/responses", content_turn()).await.unwrap(); + resp.assert_status(500); + let mut streamed = content_turn(); + streamed["stream"] = json!(true); + let resp = gw.post("/v1/responses", streamed).await.unwrap(); + assert!(resp.text().contains("down"), "{}", resp.text()); + assert_eq!(hits(&s.a).await, 3); + assert_eq!(hits(&s.b).await, 0); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn encrypted_content_whose_provider_is_gone_is_refused_up_front() { + let s = setup().await; + let gw = s.client(Some("conv-3")); + gw.post("/v1/responses", first_turn(false)) + .await + .unwrap() + .assert_ok(); + s.add_b().await; + sqlx::query("UPDATE model_routes SET enabled = false WHERE provider_id = $1") + .bind(s.a_id) + .execute(&s.app.db) + .await + .unwrap(); + s.app.rebuild_gateway_router().await; + + let resp = gw.post("/v1/responses", content_turn()).await.unwrap(); + resp.assert_status(409); + let body: Json = resp.json().unwrap(); + let message = body["error"]["message"].as_str().unwrap(); + assert!( + message.contains(&format!("encrypted by provider `{}`", s.a_name)), + "{message}" + ); + assert!( + message.contains("no other provider can read it"), + "{message}" + ); + assert!(message.contains("Start a new conversation"), "{message}"); + assert_eq!(hits(&s.b).await, 0); + + // Reasoning alone does without it. + gw.post("/v1/responses", reasoning_turn()) + .await + .unwrap() + .assert_ok(); + assert_eq!(hits(&s.b).await, 1); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn encrypted_reasoning_prefers_its_provider_and_fails_over_without_it() { + let s = setup().await; + // Named the way Codex names its conversation, without `x-session-id`. + let gw = s.client(None); + gw.set_header("session_id", "conv-4"); + gw.post("/v1/responses", first_turn(false)) + .await + .unwrap() + .assert_ok(); + s.add_b().await; + s.affinity_to_b().await; + for _ in 0..3 { + gw.post("/v1/responses", reasoning_turn()) + .await + .unwrap() + .assert_ok(); + } + assert_eq!(hits(&s.a).await, 4); + assert_eq!(hits(&s.b).await, 0); + + // `a` down: `b` takes it. + breaks(&s.a).await; + let resp = gw.post("/v1/responses", reasoning_turn()).await.unwrap(); + resp.assert_ok(); + let body: Json = resp.json().unwrap(); + assert_eq!(body["output"][0]["content"][0]["text"], "from-b", "{body}"); + assert_eq!(hits(&s.a).await, 5); + assert_eq!(hits(&s.b).await, 1); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_provider_that_cannot_read_the_history_is_said_to_plainly() { + let s = setup().await; + // Nobody remembers this conversation, and `a` cannot read it. + refuses_sealed_history(&s.a).await; + let gw = s.client(Some("conv-5")); + + let resp = gw.post("/v1/responses", content_turn()).await.unwrap(); + resp.assert_status(400); + let body: Json = resp.json().unwrap(); + let message = body["error"]["message"].as_str().unwrap(); + assert!( + message.starts_with(&format!( + "Provider `{}` cannot read this conversation's encrypted history", + s.a_name + )), + "{message}" + ); + assert!( + message.contains("another provider encrypted it"), + "{message}" + ); + assert!(message.contains("start a new conversation"), "{message}"); + // The upstream's own words are kept. + assert!(message.contains("could not be verified"), "{message}"); + + // The same refusal partway through a stream. + let mut streamed = content_turn(); + streamed["stream"] = json!(true); + let resp = gw.post("/v1/responses", streamed).await.unwrap(); + resp.assert_ok(); + let text = resp.text(); + assert!( + text.contains("cannot read this conversation's encrypted history"), + "{text}" + ); + assert!(text.contains("could not be decrypted"), "{text}"); + assert!( + !text.contains("\"code\":\"invalid_encrypted_content\""), + "the upstream's own frame went out: {text}" + ); + + // Neither answer is remembered as `a` holding the conversation. + let held: Option = s + .app + .state + .redis + .get(s.binding_key("conv-5")) + .await + .unwrap(); + assert_eq!(held, None); +}