diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 43e73e190..6e705cf56 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -1234,6 +1234,7 @@ dependencies = [ "serde", "serde_json", "thiserror", + "tracing", "uuid", ] diff --git a/backend/crates/gateway-admin/Cargo.toml b/backend/crates/gateway-admin/Cargo.toml index 1f0429024..f4c08e8b8 100644 --- a/backend/crates/gateway-admin/Cargo.toml +++ b/backend/crates/gateway-admin/Cargo.toml @@ -28,7 +28,7 @@ tracing = "0.1.44" uuid = { version = "1.23.3", features = ["v7"] } [dev-dependencies] -serde_json = { version = "1.0.150", features = ["preserve_order"] } +serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] } tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt"] } [lints] diff --git a/backend/crates/gateway-admin/src/model/observability.rs b/backend/crates/gateway-admin/src/model/observability.rs index b4b05a342..9cc27bce9 100644 --- a/backend/crates/gateway-admin/src/model/observability.rs +++ b/backend/crates/gateway-admin/src/model/observability.rs @@ -330,6 +330,7 @@ pub struct LatencyPercentiles { #[derive(Debug, Clone, PartialEq, Eq)] pub struct ProviderBillingInput { pub upstream_model_id: String, + pub service_tier: Option, pub input_tokens: Option, pub output_tokens: Option, pub cached_tokens: Option, @@ -345,6 +346,7 @@ pub struct UsageCalculatedBillingFact { pub bucket_start: DateTime, pub provider_kind: String, pub upstream_model_id: String, + pub service_tier: Option, pub input_tokens: Option, pub output_tokens: Option, pub cached_tokens: Option, @@ -574,6 +576,7 @@ pub struct UsageRecord { pub upstream_transport: Option, pub http_version: Option, pub websocket_pool: Option, + pub service_tier: Option, /// Provider 已筛选的专有观测 JSON；管理领域保持不透明。 pub provider_metadata_json: Option, pub attempt_count: u32, diff --git a/backend/crates/gateway-admin/src/ports/provider.rs b/backend/crates/gateway-admin/src/ports/provider.rs index 2621931b0..9eeb38623 100644 --- a/backend/crates/gateway-admin/src/ports/provider.rs +++ b/backend/crates/gateway-admin/src/ports/provider.rs @@ -64,6 +64,12 @@ pub trait ProviderAdmin: Send + Sync { /// 成功之后，不参与事务成败，也不得恢复或改写已经提交的账号状态。 async fn account_unavailable(&self, account_id: &ProviderAccountId); + /// 账号资格事实已经由控制面提交，失效 Provider 持有的可重建派生状态。 + /// + /// 该通知发生在 Store 事务成功之后、下一份 RuntimeSnapshot 编译之前；通知 + /// 不参与已提交事务成败。没有账号派生状态的 Provider 可使用默认空实现。 + async fn account_facts_changed(&self, _account_ids: &[ProviderAccountId]) {} + /// 生成一次连接测试所需的 Provider-owned operation；Core 负责实际执行与落账。 fn connection_test_operation( &self, diff --git a/backend/crates/gateway-admin/src/use_case/accounts.rs b/backend/crates/gateway-admin/src/use_case/accounts.rs index 3ab74b3a9..e0fae70ee 100644 --- a/backend/crates/gateway-admin/src/use_case/accounts.rs +++ b/backend/crates/gateway-admin/src/use_case/accounts.rs @@ -370,6 +370,9 @@ impl AccountsService for DefaultAccountsService { "provider credential refresh", ) .await?; + provider + .account_facts_changed(std::slice::from_ref(&result.account_id)) + .await; publish_committed(self.snapshot.as_ref(), result.config_revision).await?; let account = self.load_directory_item(&result.account_id, false).await?; Ok(AccountRefreshResult { diff --git a/backend/crates/gateway-admin/src/use_case/mod.rs b/backend/crates/gateway-admin/src/use_case/mod.rs index 96a2dd1fe..376541286 100644 --- a/backend/crates/gateway-admin/src/use_case/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/mod.rs @@ -307,6 +307,9 @@ async fn set_credential_enabled( if !enabled { provider.account_unavailable(&account_id).await; } + provider + .account_facts_changed(std::slice::from_ref(&account_id)) + .await; Ok(CredentialMutationResult { config_revision: revision, account_id, @@ -339,6 +342,7 @@ async fn delete_credentials( for account_id in &account_ids { provider.account_unavailable(account_id).await; } + provider.account_facts_changed(&account_ids).await; Ok(CredentialDeletionResult { config_revision: revision, account_ids, diff --git a/backend/crates/gateway-admin/src/use_case/observability.rs b/backend/crates/gateway-admin/src/use_case/observability.rs index 7fc9812f0..f3437ce15 100644 --- a/backend/crates/gateway-admin/src/use_case/observability.rs +++ b/backend/crates/gateway-admin/src/use_case/observability.rs @@ -428,6 +428,7 @@ fn recover_standard_costs( }; let input = ProviderBillingInput { upstream_model_id: fact.upstream_model_id, + service_tier: fact.service_tier, input_tokens: fact.input_tokens, output_tokens: fact.output_tokens, cached_tokens: fact.cached_tokens, @@ -512,6 +513,7 @@ impl DefaultObservabilityService { }; let input = ProviderBillingInput { upstream_model_id: upstream_model_id.clone(), + service_tier: record.service_tier.clone(), input_tokens: record.input_tokens, output_tokens: record.output_tokens, cached_tokens: record.cached_tokens, diff --git a/backend/crates/gateway-admin/src/use_case/openai.rs b/backend/crates/gateway-admin/src/use_case/openai.rs index 5124ba550..af6f59dac 100644 --- a/backend/crates/gateway-admin/src/use_case/openai.rs +++ b/backend/crates/gateway-admin/src/use_case/openai.rs @@ -151,6 +151,9 @@ impl OpenAiService for DefaultOpenAiService { .commit_credential_import(CredentialImportCommit { prepared }, &context) .await .map_err(|error| map_store_error(error, "OpenAI credential import"))?; + self.provider + .account_facts_changed(&result.credential_ids) + .await; publish_committed(self.snapshot.as_ref(), result.config_revision).await?; // 导入完成后立即做一次观察；失败不回滚已经提交的 credential。 self.observe_initial_quotas(&result.credential_ids, &context.request_id) @@ -203,6 +206,9 @@ impl OpenAiService for DefaultOpenAiService { "OpenAI authorization", ) .await?; + self.provider + .account_facts_changed(std::slice::from_ref(&result.account_id)) + .await; publish_committed(self.snapshot.as_ref(), result.config_revision).await?; if creates_account { self.observe_initial_quotas( @@ -244,6 +250,9 @@ impl OpenAiService for DefaultOpenAiService { "OpenAI credential rotation", ) .await?; + self.provider + .account_facts_changed(std::slice::from_ref(&result.account_id)) + .await; publish_committed(self.snapshot.as_ref(), result.config_revision).await?; Ok(result) } diff --git a/backend/crates/gateway-admin/tests/use_case/accounts.rs b/backend/crates/gateway-admin/tests/use_case/accounts.rs index 8eeffc8af..af1483a7a 100644 --- a/backend/crates/gateway-admin/tests/use_case/accounts.rs +++ b/backend/crates/gateway-admin/tests/use_case/accounts.rs @@ -180,6 +180,10 @@ impl ProviderAdmin for FakeProviderAdmin { self.record("provider.account_unavailable"); } + async fn account_facts_changed(&self, _: &[ProviderAccountId]) { + self.record("provider.account_facts_changed"); + } + fn connection_test_operation( &self, model: &gateway_core::routing::UpstreamModelId, @@ -852,6 +856,7 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() { "provider.prepare_refresh", "store.commit_refresh", "guard.finish", + "provider.account_facts_changed", "store.load_account", ] ); @@ -1157,8 +1162,7 @@ impl AccountProbe for FailingAccountProbe { "included usage exhausted", Some("usage_exhausted".to_owned()), Some("invalid_request_error".to_owned()), - ) - .expect("safe provider error"); + ); Err(GatewayError::new( GatewayErrorKind::RateLimited, "upstream capacity is temporarily unavailable", diff --git a/backend/crates/gateway-admin/tests/use_case/observability.rs b/backend/crates/gateway-admin/tests/use_case/observability.rs index 7333bf3d0..ecf5ffe54 100644 --- a/backend/crates/gateway-admin/tests/use_case/observability.rs +++ b/backend/crates/gateway-admin/tests/use_case/observability.rs @@ -249,6 +249,7 @@ async fn observability_services_should_calculate_usage_insights_and_diagnostic_s bucket_start: quarter_hour_start(now), provider_kind: "openai".to_owned(), upstream_model_id: "gpt-5.5".to_owned(), + service_tier: None, input_tokens: Some(800), output_tokens: Some(200), cached_tokens: Some(0), @@ -589,6 +590,7 @@ fn calculated_total_record( upstream_transport: None, http_version: None, websocket_pool: None, + service_tier: None, provider_metadata_json: None, attempt_count: 1, upstream_send_state: "sent".to_owned(), diff --git a/backend/crates/gateway-admin/tests/use_case/openai.rs b/backend/crates/gateway-admin/tests/use_case/openai.rs index 93ddba9ad..889e04c06 100644 --- a/backend/crates/gateway-admin/tests/use_case/openai.rs +++ b/backend/crates/gateway-admin/tests/use_case/openai.rs @@ -30,7 +30,11 @@ async fn openai_enable_should_use_scoped_store_without_provider_mutation() { assert_eq!( recorded(&events), - ["store.credential_details", "store.set_enabled"] + [ + "store.credential_details", + "store.set_enabled", + "provider.account_facts_changed", + ] ); assert_eq!(store.audit_requests(), ["request-openai"]); } @@ -54,6 +58,7 @@ async fn openai_disable_should_commit_then_release_provider_resources() { "store.credential_details", "store.set_enabled", "provider.account_unavailable", + "provider.account_facts_changed", ] ); assert_eq!(store.audit_requests(), ["request-openai"]); @@ -78,6 +83,7 @@ async fn openai_delete_should_commit_then_release_provider_resources() { "store.credential_details", "store.delete", "provider.account_unavailable", + "provider.account_facts_changed", ] ); assert_eq!(store.audit_requests(), ["request-openai"]); @@ -124,6 +130,7 @@ async fn openai_import_should_prepare_before_atomic_store_commit() { [ "provider.prepare_import", "store.commit_import", + "provider.account_facts_changed", "provider.quota", ] ); @@ -235,6 +242,7 @@ async fn openai_authorization_create_should_observe_initial_quota() { "provider.start_authorization", "provider.complete_authorization", "store.commit_authorization", + "provider.account_facts_changed", "provider.quota", ] ); @@ -324,6 +332,7 @@ async fn openai_authorization_store_failure_should_release_claim_for_retry() { "provider.complete_authorization", "store.commit_authorization", "authorization_guard.commit", + "provider.account_facts_changed", "provider.quota", ] ); @@ -401,6 +410,7 @@ async fn openai_reauthorization_should_commit_after_credential_revision_advances "provider.complete_authorization", "store.commit_authorization", "guard.finish", + "provider.account_facts_changed", ] ); assert_eq!(store.audit_requests(), ["oauth-complete-openai"]); diff --git a/backend/crates/gateway-admin/tests/use_case/xai.rs b/backend/crates/gateway-admin/tests/use_case/xai.rs index b05fdb68c..2b12d635e 100644 --- a/backend/crates/gateway-admin/tests/use_case/xai.rs +++ b/backend/crates/gateway-admin/tests/use_case/xai.rs @@ -29,6 +29,7 @@ async fn xai_disable_should_commit_then_notify_the_stateless_provider() { "store.credential_details", "store.set_enabled", "provider.account_unavailable", + "provider.account_facts_changed", ] ); assert_eq!(store.audit_requests(), ["request-xai"]); diff --git a/backend/crates/gateway-api/Cargo.toml b/backend/crates/gateway-api/Cargo.toml index 29ed50383..c7e5a30ee 100644 --- a/backend/crates/gateway-api/Cargo.toml +++ b/backend/crates/gateway-api/Cargo.toml @@ -21,7 +21,7 @@ gateway-admin = { path = "../gateway-admin" } gateway-core = { path = "../gateway-core" } gateway-protocol = { path = "../gateway-protocol" } serde = { version = "1.0.228", features = ["derive"] } -serde_json = { version = "1.0.150", features = ["preserve_order"] } +serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] } thiserror = "2.0.18" tokio = { version = "1.52.3", default-features = false, features = ["macros", "time"] } tower-http = { version = "0.7.0", features = ["cors", "fs", "request-id", "timeout", "trace"] } diff --git a/backend/crates/gateway-api/src/admin/observability.rs b/backend/crates/gateway-api/src/admin/observability.rs index 1bbc05ffc..a96ed713e 100644 --- a/backend/crates/gateway-api/src/admin/observability.rs +++ b/backend/crates/gateway-api/src/admin/observability.rs @@ -1800,7 +1800,7 @@ fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView { model, requested_model: Some(record.requested_model_id), upstream_model: record.upstream_model_id, - service_tier: None, + service_tier: record.service_tier, status_code, transport, attempt_index: None, diff --git a/backend/crates/gateway-api/src/openai/error.rs b/backend/crates/gateway-api/src/openai/error.rs index 17769d6db..8cc924f23 100644 --- a/backend/crates/gateway-api/src/openai/error.rs +++ b/backend/crates/gateway-api/src/openai/error.rs @@ -2,7 +2,8 @@ use axum::{ Json, - http::StatusCode, + body::Body, + http::{HeaderName, HeaderValue, StatusCode, header::CONTENT_TYPE}, response::{IntoResponse, Response}, }; use gateway_core::{ @@ -121,6 +122,33 @@ pub fn gateway_error_response(error: &GatewayError) -> Response { .into_response() } +/// 在下游尚未提交时优先交付 Provider 的原始 HTTP 失败响应。 +pub fn engine_error_response(error: &EngineError) -> Response { + if let EngineError::Provider(error) = error + && let Some(upstream) = error.client_visible_upstream_response() + && let Ok(status) = StatusCode::from_u16(upstream.status()) + { + let mut response = Response::new(Body::from(upstream.body().clone())); + *response.status_mut() = status; + if let Some(content_type) = upstream.content_type() + && let Ok(content_type) = HeaderValue::from_bytes(content_type) + { + response.headers_mut().insert(CONTENT_TYPE, content_type); + } + for header in upstream.headers() { + let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else { + continue; + }; + let Ok(value) = HeaderValue::from_bytes(header.value()) else { + continue; + }; + response.headers_mut().append(name, value); + } + return response; + } + gateway_error_response(&gateway_error_from_engine(error)) +} + /// Gateway 错误稳定映射，供 HTTP、SSE 和 WebSocket 共用。 #[must_use] pub const fn gateway_error_contract( diff --git a/backend/crates/gateway-api/src/openai/responses/error.rs b/backend/crates/gateway-api/src/openai/responses/error.rs index 48331fb34..3894beb34 100644 --- a/backend/crates/gateway-api/src/openai/responses/error.rs +++ b/backend/crates/gateway-api/src/openai/responses/error.rs @@ -152,9 +152,6 @@ pub enum ResponseEncodeError { /// 协议原生流缺少可返回的终态 response object。 #[error("responses wire stream has no terminal response")] MissingWireTerminal, - /// 协议原生事件携带了不一致的响应身份。 - #[error("responses wire stream changed its response identity")] - WireIdentityChanged, /// 终态 response object 无法序列化为完整 HTTP JSON 响应。 #[error("responses wire terminal serialization failed")] Serialization, @@ -167,7 +164,7 @@ impl ResponseEncodeError { ProtocolErrorBody { error: ProtocolError { kind: "server_error", - code: "invalid_canonical_response", + code: "invalid_upstream_response", message: "The gateway could not forward the upstream response.".to_owned(), param: None, }, diff --git a/backend/crates/gateway-api/src/openai/responses/http.rs b/backend/crates/gateway-api/src/openai/responses/http.rs index 48fa62719..66907c7fc 100644 --- a/backend/crates/gateway-api/src/openai/responses/http.rs +++ b/backend/crates/gateway-api/src/openai/responses/http.rs @@ -8,7 +8,7 @@ use axum::{ body::{Body, Bytes}, extract::{Extension, State, connect_info::ConnectInfo}, http::{ - HeaderMap, HeaderValue, StatusCode, + HeaderMap, HeaderName, HeaderValue, StatusCode, header::{CACHE_CONTROL, CONTENT_TYPE, USER_AGENT}, }, response::{IntoResponse, Response}, @@ -25,8 +25,8 @@ use crate::ApiState; use crate::openai::{ auth::{authenticate_client, authentication_error_response}, error::{ - gateway_error_contract, gateway_error_from_engine, gateway_error_response, - protocol_error_response, runtime_unavailable_response, + engine_error_response, gateway_error_contract, gateway_error_from_engine, + gateway_error_response, protocol_error_response, runtime_unavailable_response, }, }; @@ -180,7 +180,7 @@ pub async fn collect_execution_response(session: Box) -> R let events = match session.collect_uncommitted().await { Ok(events) => events, Err(error) => { - let response = gateway_error_response(&gateway_error_from_engine(&error)); + let response = engine_error_response(&error); return execution.record_response_status(response).await; } }; @@ -206,7 +206,7 @@ pub async fn collect_execution_response(session: Box) -> R .await { execution.cancel_and_finalize().await; - let response = gateway_error_response(&gateway_error_from_engine(&error)); + let response = engine_error_response(&error); return execution.record_response_status(response).await; } if !session.is_finalized() { @@ -237,9 +237,7 @@ fn encode_collected_events( ) -> Result, BufferedResponseEncodeError> { let mut encoder = OpenAiResponsesEncoder::new(); for event in events { - encoder - .push_sse(event) - .map_err(BufferedResponseEncodeError::Canonical)?; + encoder.push_sse(event); } let response = encoder .finish() @@ -253,23 +251,25 @@ fn json_body_response(encoded: Vec, response_headers: &[ProviderResponseHead response .headers_mut() .insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); - apply_safe_response_headers(response, response_headers) + apply_response_headers(response, response_headers) } -fn apply_safe_response_headers( +fn apply_response_headers( mut response: Response, response_headers: &[ProviderResponseHeader], ) -> Response { + let connection_options = super::response_connection_options(response_headers); for header in response_headers { - let Some(name) = super::safe_response_header_name(header.name()) else { + if !super::response_header_is_forwardable(header.name(), &connection_options) { + continue; + } + let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else { continue; }; - let Ok(value) = HeaderValue::from_str(header.value().as_str()) else { + let Ok(value) = HeaderValue::from_bytes(header.value()) else { continue; }; - response - .headers_mut() - .insert(axum::http::header::HeaderName::from_static(name), value); + response.headers_mut().append(name, value); } response } @@ -295,7 +295,7 @@ pub async fn stream_execution_response( return response; } Err(error) => { - let response = gateway_error_response(&gateway_error_from_engine(&error)); + let response = engine_error_response(&error); return execution.record_response_status(response).await; } }; @@ -310,18 +310,7 @@ pub async fn stream_execution_response( let mut encoder = OpenAiResponsesEncoder::new(); let mut frames = Vec::new(); for event in &first_events { - match encoder.push_sse(event) { - Ok(encoded) => frames.extend(encoded), - Err(error) => { - let response = protocol_error_response( - StatusCode::INTERNAL_SERVER_ERROR, - error.protocol_body(), - ); - let response = execution.record_response_status(response).await; - execution.cancel_and_finalize().await; - return response; - } - } + frames.extend(encoder.push_sse(event)); } if frames.is_empty() { let response = internal_gateway_response("gateway commit batch encoded no output"); @@ -338,7 +327,7 @@ pub async fn stream_execution_response( .await { execution.cancel_and_finalize().await; - let response = gateway_error_response(&gateway_error_from_engine(&error)); + let response = engine_error_response(&error); return execution.record_response_status(response).await; } let Some(session) = execution.into_session() else { @@ -371,7 +360,7 @@ fn event_stream_response(body: Body, response_headers: &[ProviderResponseHeader] response .headers_mut() .insert(CACHE_CONTROL, HeaderValue::from_static("no-cache")); - apply_safe_response_headers(response, response_headers) + apply_response_headers(response, response_headers) } fn internal_gateway_response(message: &'static str) -> Response { @@ -403,15 +392,10 @@ impl PendingExecution { } async fn record_response_status(&mut self, response: Response) -> Response { - let Some(session) = self.session.as_mut() else { - return internal_gateway_response("gateway response session is unavailable"); - }; - if session - .record_client_status(response.status().as_u16()) - .await - .is_err() - { - return internal_gateway_response("gateway client status could not be persisted"); + if let Some(session) = self.session.as_mut() { + let _ = session + .record_client_status(response.status().as_u16()) + .await; } response } @@ -499,10 +483,16 @@ impl ResponsesStreamState { .session .as_ref() .is_some_and(|session| session.is_finalized()); - self.finish_with_gateway_error(GatewayError::new( - GatewayErrorKind::Internal, - "gateway response ended without a terminal event", - )); + if self.execution_terminal { + self.pending + .push_back(Bytes::from_static(DONE_SSE_FRAME.as_bytes())); + self.output_finished = true; + } else { + self.finish_with_gateway_error(GatewayError::new( + GatewayErrorKind::Internal, + "gateway response ended without finalizing the execution", + )); + } } Err(error) => { self.execution_terminal = self @@ -521,18 +511,11 @@ impl ResponsesStreamState { } async fn push_event(&mut self, event: ProviderEvent) { - match self.encoder.push_sse(&event) { - Ok(frames) => { - if self.encoder.is_completed() { - self.finish_completed(frames).await; - } else { - self.pending.extend(frames); - } - } - Err(error) => { - self.cancel_execution().await; - self.finish_with_encode_error(error); - } + let frames = self.encoder.push_sse(&event); + if self.encoder.is_completed() { + self.finish_completed(frames).await; + } else { + self.pending.extend(frames); } } @@ -607,20 +590,6 @@ impl ResponsesStreamState { .push_back(Bytes::from_static(DONE_SSE_FRAME.as_bytes())); self.output_finished = true; } - - fn finish_with_encode_error(&mut self, error: ResponseEncodeError) { - let protocol = error.protocol_body().error; - self.pending - .push_back(Bytes::from(response_failed_sse_event_with_id( - self.encoder.response_id(), - protocol.kind, - protocol.code, - &protocol.message, - ))); - self.pending - .push_back(Bytes::from_static(DONE_SSE_FRAME.as_bytes())); - self.output_finished = true; - } } impl Drop for ResponsesStreamState { diff --git a/backend/crates/gateway-api/src/openai/responses/mod.rs b/backend/crates/gateway-api/src/openai/responses/mod.rs index 9b01abbbe..d98425561 100644 --- a/backend/crates/gateway-api/src/openai/responses/mod.rs +++ b/backend/crates/gateway-api/src/openai/responses/mod.rs @@ -17,13 +17,63 @@ pub use response::OpenAiResponsesEncoder; pub(crate) use websocket::responses_websocket; pub use websocket::{ResponseCreateFrameError, decode_response_create_with_context}; -pub(super) fn safe_response_header_name(name: &str) -> Option<&'static str> { - match name { - "x-request-id" => Some("x-request-id"), - "openai-model" => Some("openai-model"), - "x-models-etag" => Some("x-models-etag"), - "x-reasoning-included" => Some("x-reasoning-included"), - "openai-processing-ms" => Some("openai-processing-ms"), - _ => None, +use gateway_core::event::ProviderResponseHeader; + +pub(super) fn response_connection_options(headers: &[ProviderResponseHeader]) -> Vec { + headers + .iter() + .filter(|header| header.name().trim().eq_ignore_ascii_case("connection")) + .filter_map(|header| std::str::from_utf8(header.value()).ok()) + .flat_map(|value| value.split(',')) + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(str::to_ascii_lowercase) + .collect() +} + +pub(super) fn response_header_is_forwardable(name: &str, connection_options: &[String]) -> bool { + let name = name.trim().to_ascii_lowercase(); + if connection_options + .iter() + .any(|option| option.eq_ignore_ascii_case(&name)) + || name.starts_with("sec-websocket-") + { + return false; } + + !matches!( + name.as_str(), + "connection" + | "keep-alive" + | "proxy-connection" + | "te" + | "trailer" + | "transfer-encoding" + | "upgrade" + | "content-length" + | "content-type" + | "content-encoding" + | "authorization" + | "x-api-key" + | "www-authenticate" + | "authentication-info" + | "proxy-authenticate" + | "proxy-authorization" + | "proxy-authentication-info" + | "cookie" + | "cookie2" + | "set-cookie" + | "set-cookie2" + | "chatgpt-account-id" + | "chatgpt-organization-id" + | "chatgpt-org-id" + | "chatgpt-project-id" + | "openai-organization" + | "openai-project" + | "x-openai-organization" + | "x-openai-project" + | "x-codex-installation-id" + | "x-codex-turn-state" + | "x-codex-turn-metadata" + ) } diff --git a/backend/crates/gateway-api/src/openai/responses/request.rs b/backend/crates/gateway-api/src/openai/responses/request.rs index 6535b06c0..b9f6ed349 100644 --- a/backend/crates/gateway-api/src/openai/responses/request.rs +++ b/backend/crates/gateway-api/src/openai/responses/request.rs @@ -3,6 +3,7 @@ use std::{fmt, net::IpAddr}; use axum::http::HeaderMap; +use base64::{Engine as _, engine::general_purpose::STANDARD}; use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState}; use gateway_protocol::openai::{ X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER, @@ -14,6 +15,7 @@ use super::error::RequestDecodeError; const OPENAI_PROTOCOL: &str = "openai"; const REVIEW_SUBAGENT: &str = "review"; const OPENAI_SUBAGENT_KEY: &str = "x-openai-subagent"; +const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers"; #[derive(Clone, Default)] pub struct OpenAiRequestHeaders { @@ -34,6 +36,7 @@ pub struct OpenAiRequestHeaders { responses_lite: Option, memgen_request: Option, subagent: Option, + passthrough_headers: Vec, } impl OpenAiRequestHeaders { @@ -57,12 +60,8 @@ impl OpenAiRequestHeaders { turn_id: header_string(headers, "x-codex-turn-id"), responses_lite: header_string(headers, X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER), memgen_request: header_string(headers, X_OPENAI_MEMGEN_REQUEST_HEADER), - subagent: header_string(headers, OPENAI_SUBAGENT_KEY).filter(|value| { - matches!( - value.as_str(), - "review" | "compact" | "memory_consolidation" | "collab_spawn" - ) - }), + subagent: header_string(headers, OPENAI_SUBAGENT_KEY), + passthrough_headers: passthrough_headers(headers), } } @@ -108,6 +107,12 @@ impl OpenAiRequestHeaders { insert_protocol_context(&mut context, "turn_id", self.turn_id.as_ref()); insert_protocol_context(&mut context, "responses_lite", self.responses_lite.as_ref()); insert_protocol_context(&mut context, "memgen_request", self.memgen_request.as_ref()); + if !self.passthrough_headers.is_empty() { + context.insert( + PASSTHROUGH_HEADERS_CONTEXT_KEY.to_owned(), + Value::Array(self.passthrough_headers.clone()), + ); + } if let Some(use_websocket) = use_websocket { context.insert("use_websocket".to_owned(), Value::Bool(use_websocket)); } @@ -322,13 +327,14 @@ pub(super) fn decode_request_object( review: bool, request_headers: &OpenAiRequestHeaders, ) -> Result { - // `use_websocket` 只影响本地 transport。无论值是否可识别，都不得进入上游 body。 - let use_websocket = object - .remove("use_websocket") - .and_then(|value| value.as_bool()); + // 仅消费已识别的本地 transport 开关；未知同名值保留给未来上游协议。 + let use_websocket = object.get("use_websocket").and_then(Value::as_bool); + if use_websocket.is_some() { + object.remove("use_websocket"); + } - let model = required_non_empty_string(&object, "model", "model")?.trim(); - if model.is_empty() { + let model = required_non_empty_string(&object, "model", "model")?; + if model.trim().is_empty() { return Err(RequestDecodeError::EmptyField { field: "model".to_owned(), }); @@ -375,26 +381,98 @@ fn inject_subagent_metadata(body: &mut Map, subagent: &str) { let metadata = body .entry("client_metadata".to_owned()) .or_insert_with(|| Value::Object(Map::new())); - if !metadata.is_object() { - *metadata = Value::Object(Map::new()); - } - if let Some(metadata) = metadata.as_object_mut() { - metadata.insert( - OPENAI_SUBAGENT_KEY.to_owned(), - Value::String(subagent.to_owned()), - ); - } + let Some(metadata) = metadata.as_object_mut() else { + return; + }; + metadata.insert( + OPENAI_SUBAGENT_KEY.to_owned(), + Value::String(subagent.to_owned()), + ); } fn header_string(headers: &HeaderMap, name: &str) -> Option { headers .get(name) .and_then(|value| value.to_str().ok()) - .map(str::trim) - .filter(|value| !value.is_empty()) .map(ToOwned::to_owned) } +fn passthrough_headers(headers: &HeaderMap) -> Vec { + let connection_headers = headers + .get_all("connection") + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(str::to_ascii_lowercase) + .collect::>(); + + headers + .iter() + .filter(|(name, _)| passthrough_header_name(name.as_str(), &connection_headers)) + .map(|(name, value)| { + Value::Array(vec![ + Value::String(name.as_str().to_owned()), + Value::String(STANDARD.encode(value.as_bytes())), + ]) + }) + .collect() +} + +fn passthrough_header_name(name: &str, connection_headers: &[String]) -> bool { + if connection_headers + .iter() + .any(|connection_header| connection_header.eq_ignore_ascii_case(name)) + || name.starts_with("sec-websocket-") + || name.starts_with("x-grok-") + || name.starts_with("x-xai-") + { + return false; + } + + !matches!( + name, + // 逐跳头、上游 authority 和重序列化后的实体长度由 transport 重建。 + "connection" + | "keep-alive" + | "proxy-connection" + | "proxy-authenticate" + | "proxy-authorization" + | "te" + | "trailer" + | "transfer-encoding" + | "upgrade" + | "host" + | "content-length" + // 下游连接诊断事实只用于本地观测，不能伪装成 OpenAI 请求事实。 + | "forwarded" + | "x-forwarded-for" + | "x-forwarded-host" + | "x-forwarded-proto" + | "x-forwarded-port" + | "x-real-ip" + | "true-client-ip" + | "cf-connecting-ip" + | "x-request-id" + // 下游鉴权和账号 cookie 绝不能成为上游账号身份。 + | "authorization" + | "x-api-key" + | "cookie" + | "cookie2" + | "chatgpt-account-id" + | "chatgpt-organization-id" + | "chatgpt-org-id" + | "chatgpt-project-id" + | "openai-organization" + | "openai-project" + | "x-openai-organization" + | "x-openai-project" + // installation ID 始终由当前 lease 重建。 + | "x-codex-installation-id" + ) +} + fn required_non_empty_string<'a>( object: &'a Map, key: &str, diff --git a/backend/crates/gateway-api/src/openai/responses/response.rs b/backend/crates/gateway-api/src/openai/responses/response.rs index 424ad111c..67a4cf6ab 100644 --- a/backend/crates/gateway-api/src/openai/responses/response.rs +++ b/backend/crates/gateway-api/src/openai/responses/response.rs @@ -11,15 +11,14 @@ const OPENAI_PROTOCOL: &str = "openai"; /// OpenAI Responses wire 的转发器。 /// -/// OpenAI Provider 与 xAI adapter 都必须交付 OpenAI wire。这里仅保留同一 -/// 响应的身份和终态校验，不再从 canonical facts 重建客户端 JSON；因此 -/// OpenAI 上游的未知字段、字段顺序和事件语义不会被 API 层二次解释。 +/// OpenAI Provider 与 xAI adapter 都必须交付 OpenAI wire。wire 是客户端交付 +/// 与终态判断的事实来源；canonical facts 只在 wire 暂未携带 ID 时提供旁路观测， +/// 不得反过来拒绝或改写上游事件。 #[derive(Debug, Default)] pub struct OpenAiResponsesEncoder { response_id: Option, wire_terminal: Option, wire_failure: bool, - canonical_completed: bool, } impl OpenAiResponsesEncoder { @@ -30,53 +29,41 @@ impl OpenAiResponsesEncoder { response_id: None, wire_terminal: None, wire_failure: false, - canonical_completed: false, } } /// 消费一个 Provider event，并返回 SSE frames。 - /// - /// # Errors - /// - /// 同一响应身份变化时返回错误。 - pub fn push_sse(&mut self, event: &ProviderEvent) -> Result, ResponseEncodeError> { - self.observe_identity(event)?; + pub fn push_sse(&mut self, event: &ProviderEvent) -> Vec { + self.observe_canonical_identity(event); let Some(wire) = openai_wire(event) else { - return Ok(Vec::new()); + return Vec::new(); }; - self.observe_wire_terminal(wire); + self.observe_wire(wire); if let Some(raw_sse_frame) = wire.raw_sse_frame() { - return Ok(vec![raw_sse_frame.clone()]); + return vec![raw_sse_frame.clone()]; } - Ok(vec![Bytes::from(encode_sse_event_with_metadata( + vec![Bytes::from(encode_sse_event_with_metadata( wire.event_type().unwrap_or_default(), &wire.data().to_string(), wire.sse_id(), wire.sse_retry(), - ))]) + ))] } /// 消费一个 Provider event，并返回 WebSocket JSON messages。 - /// - /// # Errors - /// - /// 同一响应身份变化时返回错误。 - pub fn push_websocket( - &mut self, - event: &ProviderEvent, - ) -> Result, ResponseEncodeError> { - self.observe_identity(event)?; + pub fn push_websocket(&mut self, event: &ProviderEvent) -> Vec { + self.observe_canonical_identity(event); let Some(wire) = openai_wire(event).filter(|wire| wire.has_json_data()) else { - return Ok(Vec::new()); + return Vec::new(); }; - self.observe_wire_terminal(wire); - Ok(vec![wire.data().to_string()]) + self.observe_wire(wire); + vec![wire.data().to_string()] } - /// 返回是否已经看到同一响应的 canonical 与 wire 终态。 + /// 返回是否已经看到客户端可见的 wire 终态。 #[must_use] pub fn is_completed(&self) -> bool { - self.wire_terminal.is_some() && self.canonical_completed + self.wire_terminal.is_some() } /// 返回是否已把 Provider 原生失败 event 交付给客户端。 @@ -95,39 +82,31 @@ impl OpenAiResponsesEncoder { /// /// # Errors /// - /// 缺少 wire 或 canonical 终态时返回错误。 + /// 缺少可转换为非流式响应的 wire 终态时返回错误。 pub fn finish(self) -> Result { - if self.canonical_completed { - self.wire_terminal - .ok_or(ResponseEncodeError::MissingWireTerminal) - } else { - Err(ResponseEncodeError::MissingWireTerminal) - } + self.wire_terminal + .ok_or(ResponseEncodeError::MissingWireTerminal) } - fn observe_identity(&mut self, event: &ProviderEvent) -> Result<(), ResponseEncodeError> { + fn observe_canonical_identity(&mut self, event: &ProviderEvent) { for fact in event.canonical_facts() { let metadata = match fact { - GatewayEvent::Started(metadata) => metadata, - GatewayEvent::Completed(metadata) => { - self.canonical_completed = true; - metadata - } + GatewayEvent::Started(metadata) | GatewayEvent::Completed(metadata) => metadata, _ => continue, }; - if self - .response_id - .as_deref() - .is_some_and(|current| current != metadata.response_id()) - { - return Err(ResponseEncodeError::WireIdentityChanged); - } self.response_id = Some(metadata.response_id().to_owned()); } - Ok(()) } - fn observe_wire_terminal(&mut self, wire: &ProtocolWireEvent) { + fn observe_wire(&mut self, wire: &ProtocolWireEvent) { + if let Some(response_id) = wire + .data() + .pointer("/response/id") + .or_else(|| wire.data().get("response_id")) + .and_then(Value::as_str) + { + self.response_id = Some(response_id.to_owned()); + } let effective_type = wire .event_type() .or_else(|| wire.data().get("type").and_then(Value::as_str)); diff --git a/backend/crates/gateway-api/src/openai/responses/websocket.rs b/backend/crates/gateway-api/src/openai/responses/websocket.rs index 7d85824b7..c494e5991 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket.rs @@ -9,7 +9,7 @@ use axum::{ connect_info::ConnectInfo, ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade, close_code}, }, - http::{HeaderMap, StatusCode}, + http::{HeaderMap, HeaderName, StatusCode}, response::{IntoResponse, Response}, }; use futures::StreamExt; @@ -27,7 +27,6 @@ use crate::{ openai::{ auth::{authenticate_client, authentication_error_response}, error::{gateway_error_contract, gateway_error_from_engine, runtime_unavailable_response}, - router::MAX_CLIENT_REQUEST_BODY_BYTES, service::OpenAiService, }, }; @@ -102,8 +101,10 @@ impl ResponsesWebSocketAdapter { _connection_guard: connection_guard, }; websocket - .max_message_size(MAX_CLIENT_REQUEST_BODY_BYTES) - .max_frame_size(MAX_CLIENT_REQUEST_BODY_BYTES) + // 覆盖 axum/tungstenite 的私有 64 MiB message 与 16 MiB frame 默认值。 + // Responses JSON 的协议可接受性由上游决定，代理不另设 wire 长度上限。 + .max_message_size(usize::MAX) + .max_frame_size(usize::MAX) .on_upgrade(move |socket| async move { serve_responses_websocket(socket, session).await; }) @@ -303,18 +304,7 @@ async fn forward_execution( if let Some(update) = event.take_session_update() { provider_state = Some(update); } - match encoder.push_websocket(event) { - Ok(messages) => first_messages.extend(messages), - Err(error) => { - return send_protocol_error( - socket, - StatusCode::INTERNAL_SERVER_ERROR, - error.protocol_body(), - &request_id, - ) - .await; - } - } + first_messages.extend(encoder.push_websocket(event)); } if first_messages.is_empty() { let error = GatewayError::new( @@ -336,9 +326,7 @@ async fn forward_execution( { return outcome; } - if let Err(error) = commit_connection_replay(replay, &encoder, provider_state.take()) { - return send_gateway_error(socket, &error, &request_id).await; - } + commit_connection_replay(replay, &encoder, provider_state.take()); if !send_text( socket, response_metadata_event(&request_id, &response_headers), @@ -378,18 +366,7 @@ async fn forward_execution( if let Some(update) = event.take_session_update() { provider_state = Some(update); } - match encoder.push_websocket(event) { - Ok(encoded) => messages.extend(encoded), - Err(error) => { - return send_protocol_error( - socket, - StatusCode::INTERNAL_SERVER_ERROR, - error.protocol_body(), - &request_id, - ) - .await; - } - } + messages.extend(encoder.push_websocket(event)); } if encoder.is_completed() { if let Err(outcome) = @@ -397,11 +374,7 @@ async fn forward_execution( { return outcome; } - if let Err(error) = - commit_connection_replay(replay, &encoder, provider_state.take()) - { - return send_gateway_error(socket, &error, &request_id).await; - } + commit_connection_replay(replay, &encoder, provider_state.take()); if !send_messages(socket, messages).await { return ForwardOutcome::Disconnect; } @@ -413,9 +386,17 @@ async fn forward_execution( } } ActiveInput::Event(Ok(None)) => { + if execution + .session_mut() + .is_some_and(|session| session.is_finalized()) + { + commit_connection_replay(replay, &encoder, provider_state.take()); + execution.disarm(); + return ForwardOutcome::Continue; + } let error = GatewayError::new( GatewayErrorKind::Internal, - "gateway response ended without a terminal event", + "gateway response ended without finalizing the execution", ); return send_gateway_error(socket, &error, &request_id).await; } @@ -441,15 +422,10 @@ fn commit_connection_replay( replay: &mut ConnectionReplaySnapshot, encoder: &OpenAiResponsesEncoder, provider_state: Option, -) -> Result<(), GatewayError> { - let response_id = encoder.response_id().ok_or_else(|| { - GatewayError::new( - GatewayErrorKind::Internal, - "gateway completed a response without an identity", - ) - })?; - replay.commit(response_id.to_owned(), provider_state); - Ok(()) +) { + if let Some(response_id) = encoder.response_id() { + replay.commit(response_id.to_owned(), provider_state); + } } async fn confirm_completed_execution( @@ -595,21 +571,24 @@ fn response_metadata_event( request_id: &str, response_headers: &[ProviderResponseHeader], ) -> String { - let mut headers = response_headers - .iter() - .filter_map(|header| { - super::safe_response_header_name(header.name()).map(|name| { - ( - name.to_owned(), - Value::String(header.value().as_str().to_owned()), - ) - }) - }) - .collect::>(); - headers.insert( - "x-request-id".to_owned(), - Value::String(request_id.to_owned()), - ); + let connection_options = super::response_connection_options(response_headers); + let mut headers = Map::new(); + for header in response_headers { + if !super::response_header_is_forwardable(header.name(), &connection_options) { + continue; + } + let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else { + continue; + }; + let Ok(value) = std::str::from_utf8(header.value()) else { + continue; + }; + // 官方 response.metadata 使用 string map，无法表达同名多值；后值保持既有覆盖语义。 + headers.insert(name.as_str().to_owned(), Value::String(value.to_owned())); + } + headers + .entry("x-request-id".to_owned()) + .or_insert_with(|| Value::String(request_id.to_owned())); json!({ "type": "response.metadata", "headers": headers, diff --git a/backend/crates/gateway-api/src/openai/router.rs b/backend/crates/gateway-api/src/openai/router.rs index 9a8b53f60..b3f02cb78 100644 --- a/backend/crates/gateway-api/src/openai/router.rs +++ b/backend/crates/gateway-api/src/openai/router.rs @@ -13,9 +13,6 @@ use super::{ use crate::ApiState; -/// OpenAI 请求正文上限。 -pub const MAX_CLIENT_REQUEST_BODY_BYTES: usize = 16 * 1024 * 1024; - /// 构造 OpenAI 客户端协议路由。 pub(crate) fn router() -> Router { Router::new() @@ -26,5 +23,7 @@ pub(crate) fn router() -> Router { .route("/v1/models/{model_id}/info", get(model_info)) // 官方 OpenAI 模型详情合同使用 path ID；它不属于 Admin API 约束。 .route("/v1/models/{model_id}", get(model_detail)) - .layer(DefaultBodyLimit::max(MAX_CLIENT_REQUEST_BODY_BYTES)) + // Responses 正文属于 OpenAI/Codex 协议；代理不能用私有大小上限提前拒绝 + // 上游本可接受的未来 payload。 + .layer(DefaultBodyLimit::disable()) } diff --git a/backend/crates/gateway-api/src/openai/service.rs b/backend/crates/gateway-api/src/openai/service.rs index 2da1c6bcf..939cca27c 100644 --- a/backend/crates/gateway-api/src/openai/service.rs +++ b/backend/crates/gateway-api/src/openai/service.rs @@ -81,12 +81,6 @@ impl OpenAiService { "requested model was not found", ) })?; - if !self.execution.contains_public_model(&client, &public_model) { - return Err(GatewayError::new( - GatewayErrorKind::ModelNotFound, - "requested model was not found", - )); - } let previous_response_id = match metadata.continuation() { ContinuationIntent::None => None, ContinuationIntent::PreviousResponseId(value) => { diff --git a/backend/crates/gateway-api/tests/admin/observability.rs b/backend/crates/gateway-api/tests/admin/observability.rs index aed2f79ba..01dc4e950 100644 --- a/backend/crates/gateway-api/tests/admin/observability.rs +++ b/backend/crates/gateway-api/tests/admin/observability.rs @@ -326,6 +326,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() { upstream_transport: Some("http_sse".to_owned()), http_version: Some("h2".to_owned()), websocket_pool: Some("reuse".to_owned()), + service_tier: Some("priority".to_owned()), provider_metadata_json: Some( serde_json::json!({ "effectiveModel": "grok-4.5", @@ -402,6 +403,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() { assert_eq!( serde_json::json!({ "route": value["data"]["items"][0]["route"], + "serviceTier": value["data"]["items"][0]["serviceTier"], "authenticationKind": value["data"]["items"][0]["authenticationKind"], "imageInputTokens": value["data"]["items"][0]["tokenDetails"]["imageInputTokens"], "imageOutputTokens": value["data"]["items"][0]["tokenDetails"]["imageOutputTokens"], @@ -428,6 +430,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() { }), serde_json::json!({ "route": "/v1/responses", + "serviceTier": "priority", "authenticationKind": "oauth", "imageInputTokens": 31, "imageOutputTokens": 9, diff --git a/backend/crates/gateway-api/tests/openai/error.rs b/backend/crates/gateway-api/tests/openai/error.rs index 2abe8763a..7e59ab659 100644 --- a/backend/crates/gateway-api/tests/openai/error.rs +++ b/backend/crates/gateway-api/tests/openai/error.rs @@ -141,8 +141,7 @@ fn local_provider_capacity_exhaustion_should_preserve_safe_provider_detail() { "no account is eligible for the requested model", Some("no_eligible_account".to_owned()), Some("account_unavailable_error".to_owned()), - ) - .expect("safe provider detail"); + ); let error = EngineError::Provider( ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent) .with_client_visible_upstream_error(detail), @@ -327,14 +326,11 @@ fn openai_error_response_should_preserve_only_safe_contract_fields() { async fn gateway_error_response_should_expose_only_structured_client_visible_upstream_fields() { let error = GatewayError::from_provider( &ProviderError::new(ProviderErrorKind::QuotaExhausted, UpstreamSendState::Sent) - .with_client_visible_upstream_error( - ClientVisibleUpstreamError::new( - "Your Codex quota is exhausted", - Some("quota_exhausted".to_owned()), - Some("rate_limit_error".to_owned()), - ) - .expect("safe structured upstream error"), - ), + .with_client_visible_upstream_error(ClientVisibleUpstreamError::new( + "Your Codex quota is exhausted", + Some("quota_exhausted".to_owned()), + Some("rate_limit_error".to_owned()), + )), ); let response = gateway_error_response(&error); diff --git a/backend/crates/gateway-api/tests/openai/responses/http.rs b/backend/crates/gateway-api/tests/openai/responses/http.rs index 9aef9a031..0f64608fa 100644 --- a/backend/crates/gateway-api/tests/openai/responses/http.rs +++ b/backend/crates/gateway-api/tests/openai/responses/http.rs @@ -8,7 +8,10 @@ use std::sync::{ use axum::{ body::{Body, to_bytes}, extract::connect_info::ConnectInfo, - http::{HeaderMap, Request, StatusCode, header::AUTHORIZATION}, + http::{ + HeaderMap, HeaderValue, Request, StatusCode, + header::{AUTHORIZATION, CONTENT_TYPE}, + }, }; use bytes::Bytes; use futures::future::{BoxFuture, pending}; @@ -18,8 +21,8 @@ use gateway_core::engine::execution::{ }; use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError, UpstreamSendState}; use gateway_core::error::{ - ClientVisibleUpstreamError, GatewayError, GatewayErrorKind, ProviderError, ProviderErrorKind, - SafeUpstreamValue, + ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, GatewayError, GatewayErrorKind, + ProviderError, ProviderErrorKind, }; use gateway_core::event::{ ContentItem, ContentKind, GatewayEvent, ProtocolWireEvent, ProviderEvent, @@ -101,6 +104,8 @@ struct ContextCaptureExecution { #[derive(Debug, Clone, PartialEq)] struct CapturedClientContext { + public_model: String, + wire_model: Option, client_ip: Option, user_agent: Option, endpoint: String, @@ -125,11 +130,14 @@ impl ExecutionService for ContextCaptureExecution { } fn public_models(&self, _: &AuthenticatedClient) -> Vec { - Vec::new() + ["model-a", "model-b"] + .into_iter() + .map(|model| PublicModelId::new(model).expect("catalog model")) + .collect() } - fn contains_public_model(&self, _: &AuthenticatedClient, _: &PublicModelId) -> bool { - true + fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool { + matches!(model.as_str(), "model-a" | "model-b") } fn start( @@ -137,11 +145,16 @@ impl ExecutionService for ContextCaptureExecution { request: StartExecution, ) -> BoxFuture<'_, Result> { Box::pin(async move { + let public_model = request.public_model.as_str().to_owned(); let operation_kind = request.operation.kind(); let generation = match &request.operation { Operation::Generate(generation) => Some(generation), _ => None, }; + let wire_model = generation + .and_then(|generation| generation.protocol_payload().body().get("model")) + .and_then(Value::as_str) + .map(ToOwned::to_owned); let (client_metadata, protocol_context, prompt_cache_key, input) = generation.map_or((None, None, None, None), |generation| { let payload = generation.protocol_payload(); @@ -159,6 +172,8 @@ impl ExecutionService for ContextCaptureExecution { .as_ref() .map(|value| value.as_str().to_owned()); *self.observed.lock().expect("context capture lock") = Some(CapturedClientContext { + public_model, + wire_model, client_ip: request.metadata.client_ip, user_agent: request.metadata.user_agent, endpoint: request.metadata.endpoint, @@ -510,13 +525,17 @@ async fn request_context_should_resolve_forwarded_precedence_and_peer_fallback() assert_eq!( captured_client_context(headers, peer).await, CapturedClientContext { + public_model: "smart-code".to_owned(), + wire_model: Some("smart-code".to_owned()), client_ip: Some("198.51.100.1".parse().expect("expected IP")), user_agent: Some("Codex-CLI/1.0".to_owned()), endpoint: "/v1/responses".to_owned(), operation_kind: OperationKind::Generate, input: Some(json!("hello")), client_metadata: None, - protocol_context: None, + protocol_context: Some(json!({ + "opaque_request_headers": [["user-agent", "IENvZGV4LUNMSS8xLjAg"]] + })), prompt_cache_key: None, previous_response_id: None, } @@ -540,6 +559,21 @@ async fn request_context_should_resolve_forwarded_precedence_and_peer_fallback() ); } +#[tokio::test] +async fn stale_model_catalog_should_not_block_a_new_model_request() { + let model = " gpt-future-codex "; + let captured = captured_http_request( + "openai", + json!({"model": model, "input": "hello"}), + HeaderMap::new(), + None, + ) + .await; + + assert_eq!(captured.public_model, model); + assert_eq!(captured.wire_model.as_deref(), Some(model)); +} + #[tokio::test] async fn http_request_should_pass_opaque_previous_response_id_to_core() { for response_id in [ @@ -567,14 +601,17 @@ async fn http_request_should_pass_opaque_previous_response_id_to_core() { } #[tokio::test] -async fn http_request_should_forward_safe_codex_headers_without_projecting_xai_headers() { +async fn http_request_should_forward_codex_headers_without_projecting_xai_headers() { let peer = "192.0.2.10:443".parse().expect("peer address"); let mut headers = HeaderMap::new(); headers.insert("x-codex-turn-state", "turn-state".parse().expect("header")); headers.insert("conversation-id", "conversation-1".parse().expect("header")); headers.insert("session-id", "session-1".parse().expect("header")); headers.insert("x-grok-turn-idx", "7".parse().expect("header")); - headers.insert("x-openai-subagent", "compact".parse().expect("header")); + headers.insert( + "x-openai-subagent", + "future_codex_mode".parse().expect("header"), + ); headers.insert( "x-openai-internal-codex-responses-lite", "true".parse().expect("header"), @@ -600,10 +637,33 @@ async fn http_request_should_forward_safe_codex_headers_without_projecting_xai_h .client_metadata .as_ref() .and_then(|metadata| metadata.get("x-openai-subagent")), - Some(&json!("compact")) + Some(&json!("future_codex_mode")) ); } +#[tokio::test] +async fn subagent_header_should_not_replace_non_object_client_metadata() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-openai-subagent", + "future_codex_mode".parse().expect("header"), + ); + + let captured = captured_http_request( + "openai", + json!({ + "model": "smart-code", + "input": "preserve metadata", + "client_metadata": "opaque-client-value" + }), + headers, + None, + ) + .await; + + assert_eq!(captured.client_metadata, Some(json!("opaque-client-value"))); +} + #[tokio::test] async fn xai_private_headers_should_not_enter_openai_request_facts() { let peer = "192.0.2.10:443".parse().expect("peer address"); @@ -666,6 +726,43 @@ async fn review_route_should_inject_review_subagent_and_record_its_endpoint() { ); } +#[tokio::test] +async fn review_route_should_not_replace_non_object_client_metadata() { + let observed = Arc::new(Mutex::new(None)); + let execution = Arc::new(ContextCaptureExecution { + observed: Arc::clone(&observed), + client: authenticated_client("sk_context_test"), + }); + let request = Request::post("/v1/responses/review") + .header(AUTHORIZATION, "Bearer sk_context_test") + .body(Body::from( + json!({ + "model": "smart-code", + "input": "review this", + "client_metadata": ["opaque", "client", "value"] + }) + .to_string(), + )) + .expect("review request"); + + let response = api_router(execution) + .await + .oneshot(request) + .await + .expect("review response"); + let captured = observed + .lock() + .expect("review capture lock") + .clone() + .expect("captured review request"); + + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!( + captured.client_metadata, + Some(json!(["opaque", "client", "value"])) + ); +} + #[tokio::test] async fn streaming_encodes_first_frame_before_commit_and_http_delivery() { let trace = Arc::new(Trace::default()); @@ -687,19 +784,22 @@ async fn streaming_encodes_first_frame_before_commit_and_http_delivery() { } #[tokio::test] -async fn streaming_response_forwards_only_allowlisted_provider_headers() { +async fn streaming_response_should_append_ordinary_headers_and_skip_unrepresentable_or_sensitive_ones() + { let trace = Arc::new(Trace::default()); let headers = vec![ - ProviderResponseHeader::new( - "x-models-etag", - SafeUpstreamValue::new("models-v2").expect("etag"), - ) - .expect("safe header"), + ProviderResponseHeader::new("x-models-etag", Bytes::from_static(b"models-v2")), + ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"first")), + ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"second")), + ProviderResponseHeader::new("x-future-bytes", Bytes::from_static(b"\xffopaque")), ProviderResponseHeader::new( "authorization", - SafeUpstreamValue::new("should-not-cross-boundary").expect("header value"), - ) - .expect("syntactically safe header"), + Bytes::from_static(b"should-not-cross-boundary"), + ), + ProviderResponseHeader::new("connection", Bytes::from_static(b"x-hop-secret")), + ProviderResponseHeader::new("x-hop-secret", Bytes::from_static(b"hop-secret")), + ProviderResponseHeader::new("content-type", Bytes::from_static(b"application/private")), + ProviderResponseHeader::new("bad\0name", Bytes::from_static(b"unrepresentable")), ]; let session = FakeSession::streaming( Arc::clone(&trace), @@ -719,7 +819,29 @@ async fn streaming_response_forwards_only_allowlisted_provider_headers() { .and_then(|value| value.to_str().ok()), Some("models-v2") ); + assert_eq!( + response + .headers() + .get_all("x-future-multi") + .iter() + .map(HeaderValue::as_bytes) + .collect::>(), + vec![b"first".as_slice(), b"second".as_slice()] + ); + assert_eq!( + response + .headers() + .get("x-future-bytes") + .map(HeaderValue::as_bytes), + Some(b"\xffopaque".as_slice()) + ); assert!(response.headers().get("authorization").is_none()); + assert!(response.headers().get("connection").is_none()); + assert!(response.headers().get("x-hop-secret").is_none()); + assert_eq!( + response.headers().get(CONTENT_TYPE), + Some(&HeaderValue::from_static("text/event-stream")) + ); std::mem::forget(response); } @@ -835,7 +957,107 @@ async fn streaming_rate_limit_before_first_frame_should_persist_the_returned_429 } #[tokio::test] -async fn streaming_postcommit_encode_failure_cancels_without_second_commit() { +async fn streaming_upstream_http_failure_before_first_frame_should_preserve_raw_response() { + let trace = Arc::new(Trace::default()); + let raw_body = Bytes::from_static( + b"{\"error\":{\"message\":\"rate limited\",\"future_field\":{\"kept\":true}},\"top_level\":17}\x00", + ); + let error = ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent) + .with_client_visible_upstream_response( + ClientVisibleUpstreamResponse::new( + 429, + Some(b"application/problem+json; charset=utf-8".to_vec()), + raw_body.clone(), + ) + .with_headers(vec![ + ProviderResponseHeader::new("retry-after", Bytes::from_static(b"17")), + ProviderResponseHeader::new("x-request-id", Bytes::from_static(b"req-failure")), + ProviderResponseHeader::new("x-future-error", Bytes::from_static(b"first")), + ProviderResponseHeader::new("x-future-error", Bytes::from_static(b"second")), + ProviderResponseHeader::new("x-future-bytes", Bytes::from_static(b"\xffopaque")), + ]), + ); + let session = FakeSession::streaming( + Arc::clone(&trace), + vec![NextStep::Error(EngineError::Provider(error))], + ); + + let response = stream_execution_response(Box::new(session), None).await; + + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!( + response + .headers() + .get(CONTENT_TYPE) + .map(|value| value.as_bytes()), + Some(b"application/problem+json; charset=utf-8".as_slice()) + ); + assert_eq!( + response + .headers() + .get("retry-after") + .map(|value| value.as_bytes()), + Some(b"17".as_slice()) + ); + assert_eq!( + response + .headers() + .get("x-request-id") + .map(|value| value.as_bytes()), + Some(b"req-failure".as_slice()) + ); + assert_eq!( + response + .headers() + .get_all("x-future-error") + .iter() + .map(|value| value.as_bytes()) + .collect::>(), + vec![b"first".as_slice(), b"second".as_slice()] + ); + assert_eq!( + response + .headers() + .get("x-future-bytes") + .map(|value| value.as_bytes()), + Some(b"\xffopaque".as_slice()) + ); + assert_eq!( + to_bytes(response.into_body(), usize::MAX) + .await + .expect("read raw error body"), + raw_body + ); + assert_eq!(trace.client_statuses(), vec![429]); +} + +#[tokio::test] +async fn buffered_upstream_http_failure_should_preserve_raw_response() { + let trace = Arc::new(Trace::default()); + let raw_body = Bytes::from_static(b"not-json\xffstill-upstream"); + let error = ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::Sent) + .with_client_visible_upstream_response(ClientVisibleUpstreamResponse::new( + 503, + Some(b"application/octet-stream".to_vec()), + raw_body.clone(), + )); + let session = FakeSession::buffered(Arc::clone(&trace), vec![started(), completed()]) + .with_collect_error(EngineError::Provider(error)); + + let response = collect_execution_response(Box::new(session)).await; + + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + to_bytes(response.into_body(), usize::MAX) + .await + .expect("read raw error body"), + raw_body + ); + assert_eq!(trace.client_statuses(), vec![503]); +} + +#[tokio::test] +async fn streaming_canonical_identity_change_does_not_interrupt_wire_delivery() { let trace = Arc::new(Trace::default()); let session = FakeSession::streaming( Arc::clone(&trace), @@ -845,7 +1067,7 @@ async fn streaming_postcommit_encode_failure_cancels_without_second_commit() { mismatched_terminal_event(), CommitRequirement::AlreadyCommitted, )), - NextStep::FinalizeCancelled, + NextStep::FinalizeSuccess, ], ); @@ -854,12 +1076,16 @@ async fn streaming_postcommit_encode_failure_cancels_without_second_commit() { .await .expect("read SSE body"); - assert!(String::from_utf8_lossy(&body).contains("response.failed")); + let body = String::from_utf8_lossy(&body); + assert!(body.contains("response.completed")); + assert!(body.contains("resp_other")); + assert!(!body.contains("response.failed")); + assert!(body.ends_with("data: [DONE]\n\n")); assert_eq!(trace.client_statuses(), vec![200]); - assert!(trace.is_cancelled()); + assert!(!trace.is_cancelled()); assert_eq!( trace.snapshot(), - vec!["next_event", "commit", "next_event", "cancel_finalize"] + vec!["next_event", "commit", "next_event", "next_end"] ); } @@ -908,6 +1134,38 @@ async fn streaming_success_should_emit_terminal_event_and_done_marker() { ); } +#[tokio::test] +async fn streaming_finalized_unknown_wire_should_end_cleanly_without_synthetic_failure() { + let trace = Arc::new(Trace::default()); + let raw = Bytes::from_static( + b"event: response.future_terminal\ndata: {\"type\":\"response.future_terminal\",\"opaque\":true}\n\n", + ); + let event = ProviderEvent::wire( + ProtocolWireEvent::raw_sse("openai", raw.clone()).expect("raw future event"), + ); + let session = FakeSession::streaming( + Arc::clone(&trace), + vec![ + NextStep::Event(delivery_provider( + event, + CommitRequirement::CommitBeforeDelivery, + )), + NextStep::FinalizeSuccess, + ], + ); + + let response = stream_execution_response(Box::new(session), None).await; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("read transparent SSE body"); + + assert!(body.as_ref().starts_with(raw.as_ref())); + assert!(body.as_ref().ends_with(b"data: [DONE]\n\n")); + assert!(!String::from_utf8_lossy(&body).contains("response.failed")); + assert!(!trace.is_cancelled()); + assert_eq!(trace.snapshot(), vec!["next_event", "commit", "next_end"]); +} + #[tokio::test] async fn streaming_completed_event_without_finalized_execution_should_fail_closed() { let trace = Arc::new(Trace::default()); @@ -976,14 +1234,11 @@ async fn streaming_empty_terminal_should_emit_failure_done_and_cancel_execution( async fn streaming_error_should_emit_client_visible_upstream_details() { let trace = Arc::new(Trace::default()); let error = ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent) - .with_client_visible_upstream_error( - ClientVisibleUpstreamError::new( - "Your Codex quota is exhausted", - Some("quota_exhausted".to_owned()), - Some("rate_limit_error".to_owned()), - ) - .expect("safe structured upstream error"), - ); + .with_client_visible_upstream_error(ClientVisibleUpstreamError::new( + "Your Codex quota is exhausted", + Some("quota_exhausted".to_owned()), + Some("rate_limit_error".to_owned()), + )); let session = FakeSession::streaming( Arc::clone(&trace), vec![ @@ -1246,13 +1501,10 @@ async fn buffered_response_commits_only_after_complete_json_is_encoded() { } #[tokio::test] -async fn buffered_response_forwards_allowlisted_provider_headers() { +async fn buffered_response_forwards_long_ordinary_provider_header_values() { let trace = Arc::new(Trace::default()); - let header = ProviderResponseHeader::new( - "openai-model", - SafeUpstreamValue::new("gpt-5.5-effective").expect("model header"), - ) - .expect("safe header"); + let model = format!("gpt-{}", "x".repeat(512)); + let header = ProviderResponseHeader::new("openai-model", Bytes::from(model.clone())); let session = FakeSession::buffered(Arc::clone(&trace), vec![started(), completed()]) .with_response_headers(vec![header]); @@ -1263,7 +1515,7 @@ async fn buffered_response_forwards_allowlisted_provider_headers() { .headers() .get("openai-model") .and_then(|value| value.to_str().ok()), - Some("gpt-5.5-effective") + Some(model.as_str()) ); } diff --git a/backend/crates/gateway-api/tests/openai/responses/mod.rs b/backend/crates/gateway-api/tests/openai/responses/mod.rs index c0131dc50..6d4725cff 100644 --- a/backend/crates/gateway-api/tests/openai/responses/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/mod.rs @@ -2,6 +2,7 @@ mod http; mod websocket; use axum::http::{HeaderMap, HeaderValue}; +use base64::{Engine as _, engine::general_purpose::STANDARD}; use bytes::Bytes; use gateway_api::openai::responses::{ ContinuationIntent, DecodedResponsesRequest, OpenAiRequestHeaders, OpenAiResponsesEncoder, @@ -117,6 +118,28 @@ fn decoder_should_preserve_the_openai_body_and_only_derive_stable_routing_facts( assert!(decoded.metadata().store()); } +#[test] +fn decoder_should_preserve_opaque_client_model_values() { + for model in [ + format!("future-{}", "x".repeat(512)), + "future\0model".to_owned(), + " future-model ".to_owned(), + ] { + let decoded = generate_request(json!({ + "model": model.clone(), + "input": "hello" + })); + + assert_eq!(decoded.metadata().public_model(), model); + assert_eq!( + openai_wire_body(&decoded) + .get("model") + .and_then(Value::as_str), + Some(model.as_str()) + ); + } +} + #[test] fn decoder_should_preserve_compaction_trigger_for_the_openai_provider() { let decoded = generate_request(json!({ @@ -149,6 +172,21 @@ fn decoder_should_keep_transport_override_out_of_the_openai_wire_body() { ); } +#[test] +fn decoder_should_preserve_unrecognized_use_websocket_values() { + let decoded = generate_request(json!({ + "model": "smart-code", + "input": "hello", + "use_websocket": {"future": "transport-mode"} + })); + + assert_eq!( + openai_wire_body(&decoded).get("use_websocket"), + Some(&json!({"future": "transport-mode"})) + ); + assert!(!openai_protocol_context(&decoded).contains_key("use_websocket")); +} + #[test] fn decoder_should_preserve_connection_metadata_outside_the_openai_wire_body() { let mut headers = HeaderMap::new(); @@ -170,10 +208,142 @@ fn decoder_should_preserve_connection_metadata_outside_the_openai_wire_body() { ("turn_state".to_owned(), json!("turn-state")), ("turn_metadata".to_owned(), json!("{\"kind\":\"review\"}")), ("conversation_id".to_owned(), json!("conversation")), + ( + "opaque_request_headers".to_owned(), + json!([ + ["x-codex-turn-state", STANDARD.encode(b"turn-state")], + [ + "x-codex-turn-metadata", + STANDARD.encode(br#"{"kind":"review"}"#) + ], + ["conversation-id", STANDARD.encode(b"conversation")] + ]), + ), ]) ); } +#[test] +fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() { + let mut headers = HeaderMap::new(); + headers.append( + "x-openai-future-mode", + HeaderValue::from_static("future-ascii"), + ); + headers.append( + "x-openai-future-mode", + HeaderValue::from_bytes(b"\x80\xff").expect("opaque header bytes"), + ); + headers.append("version", HeaderValue::from_static("future-v1")); + headers.append( + "version", + HeaderValue::from_bytes(b"future-\x80").expect("opaque known header bytes"), + ); + headers.insert( + "accept", + HeaderValue::from_static("application/vnd.openai.responses+json"), + ); + headers.insert( + "content-type", + HeaderValue::from_static("application/vnd.openai.responses.request+json"), + ); + headers.insert( + "user-agent", + HeaderValue::from_static("Codex future-client"), + ); + headers.insert("originator", HeaderValue::from_static("future-codex")); + headers.append( + "openai-beta", + HeaderValue::from_static("future_responses=v2"), + ); + headers.append("openai-beta", HeaderValue::from_static("future_tools=v3")); + headers.insert( + "x-openai-internal-codex-residency", + HeaderValue::from_static("future-region"), + ); + headers.insert( + "x-codex-turn-state", + HeaderValue::from_static("same-account-turn"), + ); + + headers.insert( + "connection", + HeaderValue::from_static("keep-alive, x-private-hop"), + ); + headers.insert("x-private-hop", HeaderValue::from_static("drop-me")); + headers.insert("host", HeaderValue::from_static("downstream.invalid")); + headers.insert("content-length", HeaderValue::from_static("999")); + headers.insert("authorization", HeaderValue::from_static("Bearer client")); + headers.insert("x-api-key", HeaderValue::from_static("client-key")); + headers.insert("cookie", HeaderValue::from_static("client=cookie")); + headers.insert( + "chatgpt-account-id", + HeaderValue::from_static("client-account"), + ); + headers.insert( + "chatgpt-project-id", + HeaderValue::from_static("client-project"), + ); + headers.insert( + "x-codex-installation-id", + HeaderValue::from_static("client-installation"), + ); + + let decoded = + decode_request_with_headers(br#"{"model":"smart-code","input":"hello"}"#, &headers) + .expect("opaque request headers should not affect body decoding"); + let entries = openai_protocol_context(&decoded) + .get("opaque_request_headers") + .and_then(Value::as_array) + .expect("opaque request header context"); + let values = |target: &str| { + entries + .iter() + .filter_map(Value::as_array) + .filter(|entry| entry.first().and_then(Value::as_str) == Some(target)) + .filter_map(|entry| entry.get(1).and_then(Value::as_str)) + .map(|encoded| STANDARD.decode(encoded).expect("valid base64 header value")) + .collect::>() + }; + + assert_eq!( + values("x-openai-future-mode"), + vec![b"future-ascii".to_vec(), b"\x80\xff".to_vec()] + ); + assert_eq!( + values("version"), + vec![b"future-v1".to_vec(), b"future-\x80".to_vec()] + ); + assert_eq!( + values("openai-beta"), + vec![b"future_responses=v2".to_vec(), b"future_tools=v3".to_vec()] + ); + for preserved in [ + "accept", + "content-type", + "user-agent", + "originator", + "x-openai-internal-codex-residency", + "x-codex-turn-state", + ] { + assert_eq!(values(preserved).len(), 1, "missing {preserved}"); + } + for excluded in [ + "connection", + "x-private-hop", + "host", + "content-length", + "authorization", + "x-api-key", + "cookie", + "chatgpt-account-id", + "chatgpt-project-id", + "x-codex-installation-id", + ] { + assert!(values(excluded).is_empty(), "unexpected {excluded}"); + } +} + #[test] fn decoder_should_preserve_unknown_nested_values_without_debug_disclosure() { let secret = "nested-private-value"; @@ -190,6 +360,18 @@ fn decoder_should_preserve_unknown_nested_values_without_debug_disclosure() { assert!(!format!("{decoded:?}").contains(secret)); } +#[test] +fn decoder_should_preserve_unknown_arbitrary_precision_numbers() { + let number = "12345678901234567890123456789012345678901234567890"; + let body = format!(r#"{{"model":"smart-code","input":"hello","future_number":{number}}}"#); + + let decoded = decode_request_with_headers(body.as_bytes(), &HeaderMap::new()) + .expect("opaque numeric field should decode"); + let encoded = serde_json::to_string(openai_wire_body(&decoded)).expect("encode request body"); + + assert!(encoded.contains(&format!(r#""future_number":{number}"#))); +} + #[test] fn decoder_should_preserve_opaque_continuation_and_mark_the_routing_requirement() { for response_id in [ @@ -330,7 +512,7 @@ fn transparent_encoder_should_forward_raw_sse_frames_unknown_events_and_terminal let mut encoder = OpenAiResponsesEncoder::new(); let frames = events .iter() - .flat_map(|event| encoder.push_sse(event).expect("wire event should encode")) + .flat_map(|event| encoder.push_sse(event)) .collect::>(); let response = encoder.finish().expect("wire response should finish"); let body = frames @@ -384,15 +566,11 @@ fn transparent_encoder_should_use_identical_json_for_sse_and_websocket() { let sse = events .iter() - .flat_map(|event| sse_encoder.push_sse(event).expect("SSE wire event")) + .flat_map(|event| sse_encoder.push_sse(event)) .collect::>(); let websocket = events .iter() - .flat_map(|event| { - websocket_encoder - .push_websocket(event) - .expect("WebSocket wire event") - }) + .flat_map(|event| websocket_encoder.push_websocket(event)) .collect::>(); let sse_text = String::from_utf8( sse.iter() @@ -414,8 +592,40 @@ fn transparent_encoder_should_use_identical_json_for_sse_and_websocket() { } #[test] -fn transparent_encoder_should_require_matching_canonical_and_wire_terminals() { +fn transparent_encoder_should_follow_wire_when_canonical_identity_changes() { let mut encoder = OpenAiResponsesEncoder::new(); + let started = openai_wire_event( + vec![GatewayEvent::Started(ResponseMeta::new( + "resp_1", "gpt-test", + ))], + "response.created", + json!({"type": "response.created", "response": {"id": "wire_resp_1"}}), + ); + encoder.push_sse(&started); + let changed_identity = openai_wire_event( + vec![GatewayEvent::Completed(ResponseMeta::new( + "resp_2", "gpt-test", + ))], + "response.completed", + json!({"type": "response.completed", "response": {"id": "wire_resp_2"}}), + ); + + let completed = encoder.push_sse(&changed_identity); + + assert_eq!(completed.len(), 1); + assert!(encoder.is_completed()); + assert_eq!(encoder.response_id(), Some("wire_resp_2")); + assert_eq!( + encoder + .finish() + .expect("wire terminal remains authoritative"), + json!({"id": "wire_resp_2"}) + ); +} + +#[test] +fn transparent_encoder_should_only_require_wire_terminal_for_buffered_conversion() { + let mut missing_terminal = OpenAiResponsesEncoder::new(); let started = openai_wire_event( vec![GatewayEvent::Started(ResponseMeta::new( "resp_1", "gpt-test", @@ -423,24 +633,8 @@ fn transparent_encoder_should_require_matching_canonical_and_wire_terminals() { "response.created", json!({"type": "response.created", "response": {"id": "resp_1"}}), ); - encoder.push_sse(&started).expect("started event"); - let changed_identity = openai_wire_event( - vec![GatewayEvent::Completed(ResponseMeta::new( - "resp_2", "gpt-test", - ))], - "response.completed", - json!({"type": "response.completed", "response": {"id": "resp_2"}}), - ); - assert_eq!( - encoder - .push_sse(&changed_identity) - .expect_err("identity changes must fail"), - ResponseEncodeError::WireIdentityChanged - ); - - let mut missing_terminal = OpenAiResponsesEncoder::new(); - missing_terminal.push_sse(&started).expect("started event"); + missing_terminal.push_sse(&started); assert_eq!( missing_terminal .finish() diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket.rs b/backend/crates/gateway-api/tests/openai/responses/websocket.rs index 72143f217..13da2e569 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket.rs @@ -8,6 +8,7 @@ use axum::{ body::Body, http::{Request, StatusCode, header::AUTHORIZATION}, }; +use bytes::Bytes; use futures::{SinkExt, StreamExt, future::BoxFuture}; use gateway_api::openai::responses::ResponseCreateFrameError; use gateway_core::engine::execution::{ @@ -177,6 +178,7 @@ struct AtomicFailureTrace { struct AtomicFailureSession { trace: Arc, + response_headers: Vec, } impl ExecutionSession for AtomicFailureSession { @@ -201,7 +203,7 @@ impl ExecutionSession for AtomicFailureSession { } fn response_headers(&self) -> &[ProviderResponseHeader] { - &[] + &self.response_headers } fn commit_downstream(&mut self, _: Option) -> BoxFuture<'_, Result<(), EngineError>> { @@ -231,6 +233,7 @@ impl ExecutionSession for AtomicFailureSession { struct AtomicFailureExecution { client: AuthenticatedClient, trace: Arc, + response_headers: Vec, } impl ExecutionService for AtomicFailureExecution { @@ -263,6 +266,7 @@ impl ExecutionService for AtomicFailureExecution { stream: request.metadata.stream, session: Box::new(AtomicFailureSession { trace: Arc::clone(&self.trace), + response_headers: self.response_headers.clone(), }), }) }) @@ -320,6 +324,7 @@ async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() { let execution = Arc::new(AtomicFailureExecution { client: authenticated_client("sk_ws_atomic"), trace: Arc::clone(&trace), + response_headers: Vec::new(), }); let app = api_router(execution).await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0") @@ -396,6 +401,172 @@ async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() { server.abort(); } +#[tokio::test] +async fn websocket_response_metadata_should_preserve_ordinary_headers_without_blocking_events() { + let trace = Arc::new(AtomicFailureTrace::default()); + let execution = Arc::new(AtomicFailureExecution { + client: authenticated_client("sk_ws_atomic"), + trace, + response_headers: vec![ + ProviderResponseHeader::new("x-future-header", Bytes::from_static(b"future-value")), + ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"first")), + ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"second")), + ProviderResponseHeader::new("x-future-bytes", Bytes::from_static(b"\xffopaque")), + ProviderResponseHeader::new("bad\0name", Bytes::from_static(b"unrepresentable")), + ProviderResponseHeader::new( + "authorization", + Bytes::from_static(b"should-not-cross-boundary"), + ), + ProviderResponseHeader::new("connection", Bytes::from_static(b"x-private-hop")), + ProviderResponseHeader::new("x-private-hop", Bytes::from_static(b"private-hop")), + ProviderResponseHeader::new("content-type", Bytes::from_static(b"application/private")), + ProviderResponseHeader::new("x-request-id", Bytes::from_static(b"req_upstream")), + ], + }); + let app = api_router(execution).await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket test server"); + let address = listener.local_addr().expect("test server address"); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .await + .expect("serve WebSocket test router"); + }); + let mut request = format!("ws://{address}/v1/responses") + .into_client_request() + .expect("WebSocket request"); + request.headers_mut().insert( + AUTHORIZATION, + "Bearer sk_ws_atomic".parse().expect("authorization"), + ); + let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket"); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + socket + .send(ClientMessage::Text( + json!({ + "type": "response.create", + "model": "model-a", + "input": "hello" + }) + .to_string() + .into(), + )) + .await + .expect("send response.create"); + + let mut metadata = None; + let mut saw_failed = false; + while metadata.is_none() || !saw_failed { + let message = tokio::time::timeout(Duration::from_secs(2), socket.next()) + .await + .expect("WebSocket response timeout") + .expect("WebSocket remains open") + .expect("valid WebSocket frame"); + let ClientMessage::Text(text) = message else { + continue; + }; + let value: Value = serde_json::from_str(&text).expect("JSON WebSocket event"); + match value.get("type").and_then(Value::as_str) { + Some("response.metadata") => metadata = Some(value), + Some("response.failed") => saw_failed = true, + _ => {} + } + } + + let headers = metadata + .as_ref() + .and_then(|value| value.get("headers")) + .and_then(Value::as_object) + .expect("response metadata headers"); + assert_eq!(headers.get("x-future-header"), Some(&json!("future-value"))); + // 官方 metadata 是 string map，同名多值只能保持既有的后值覆盖语义。 + assert_eq!(headers.get("x-future-multi"), Some(&json!("second"))); + assert_eq!(headers.get("x-request-id"), Some(&json!("req_upstream"))); + for omitted in [ + "x-future-bytes", + "bad\0name", + "authorization", + "connection", + "x-private-hop", + "content-type", + ] { + assert!( + !headers.contains_key(omitted), + "unexpected metadata header: {omitted:?}" + ); + } + assert!(saw_failed); + + socket.close(None).await.expect("close WebSocket"); + server.abort(); +} + +#[tokio::test] +async fn websocket_response_create_should_not_have_a_private_16_mib_frame_limit() { + let trace = Arc::new(AtomicFailureTrace::default()); + let execution = Arc::new(AtomicFailureExecution { + client: authenticated_client("sk_ws_atomic"), + trace, + response_headers: Vec::new(), + }); + let app = api_router(execution).await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket test server"); + let address = listener.local_addr().expect("test server address"); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .await + .expect("serve WebSocket test router"); + }); + let mut request = format!("ws://{address}/v1/responses") + .into_client_request() + .expect("WebSocket request"); + request.headers_mut().insert( + AUTHORIZATION, + "Bearer sk_ws_atomic".parse().expect("authorization"), + ); + let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket"); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + let payload = json!({ + "type": "response.create", + "model": "model-a", + "input": "x".repeat(16 * 1024 * 1024 + 1), + }) + .to_string(); + assert!(payload.len() > 16 * 1024 * 1024); + + socket + .send(ClientMessage::Text(payload.into())) + .await + .expect("send response.create above the former frame limit"); + + let mut saw_upstream_event = false; + for _ in 0..3 { + let message = tokio::time::timeout(Duration::from_secs(10), socket.next()) + .await + .expect("WebSocket response timeout") + .expect("WebSocket remains open") + .expect("valid WebSocket frame"); + let ClientMessage::Text(text) = message else { + continue; + }; + let value: Value = serde_json::from_str(&text).expect("JSON WebSocket event"); + if matches!( + value.get("type").and_then(Value::as_str), + Some("response.created" | "response.failed") + ) { + saw_upstream_event = true; + break; + } + } + assert!(saw_upstream_event); + + socket.close(None).await.expect("close WebSocket"); + server.abort(); +} + const TEST_WEBSOCKET_KEY: &str = "AAAAAAAAAAAAAAAAAAAAAA=="; fn upgrade_request(authorization: &str) -> Request { diff --git a/backend/crates/gateway-api/tests/openai/router.rs b/backend/crates/gateway-api/tests/openai/router.rs index 92ed7d10e..79636230b 100644 --- a/backend/crates/gateway-api/tests/openai/router.rs +++ b/backend/crates/gateway-api/tests/openai/router.rs @@ -1,57 +1,29 @@ use axum::{ - Router, - body::{Body, Bytes}, - extract::DefaultBodyLimit, + body::Body, http::{Method, Request, StatusCode}, - routing::post, }; -use gateway_api::openai::router::MAX_CLIENT_REQUEST_BODY_BYTES; use tower::ServiceExt; use super::api_router_with_origins; use super::models::ModelsExecution; -fn body_limit_app() -> Router { - Router::new() - .route( - "/v1/responses", - post(|body: Bytes| async move { - if body.len() == MAX_CLIENT_REQUEST_BODY_BYTES { - StatusCode::NO_CONTENT - } else { - StatusCode::INTERNAL_SERVER_ERROR - } - }), - ) - .layer(DefaultBodyLimit::max(MAX_CLIENT_REQUEST_BODY_BYTES)) -} +const REMOVED_BODY_LIMIT_BYTES: usize = 16 * 1024 * 1024; #[tokio::test] -async fn responses_body_limit_should_accept_exactly_sixteen_mibibytes() { - let response = body_limit_app() +async fn responses_body_should_accept_payload_above_the_removed_private_limit() { + let router = api_router_with_origins(ModelsExecution::new(), Vec::new()).await; + let response = router .oneshot( Request::post("/v1/responses") - .body(Body::from(vec![b'a'; MAX_CLIENT_REQUEST_BODY_BYTES])) - .expect("build exact-limit request"), + .body(Body::from(vec![b'a'; REMOVED_BODY_LIMIT_BYTES + 1])) + .expect("build request above the removed limit"), ) .await - .expect("route exact-limit request"); + .expect("route request above the removed limit"); - assert_eq!(response.status(), StatusCode::NO_CONTENT); -} - -#[tokio::test] -async fn responses_body_limit_should_reject_one_byte_over_sixteen_mibibytes() { - let response = body_limit_app() - .oneshot( - Request::post("/v1/responses") - .body(Body::from(vec![b'a'; MAX_CLIENT_REQUEST_BODY_BYTES + 1])) - .expect("build over-limit request"), - ) - .await - .expect("route over-limit request"); - - assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE); + // The handler sees the body and rejects the missing credentials. A restored body limit + // would return 413 before authentication runs. + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); } #[tokio::test] diff --git a/backend/crates/gateway-core/Cargo.toml b/backend/crates/gateway-core/Cargo.toml index 4a4d8c7ac..8f7ab4fdd 100644 --- a/backend/crates/gateway-core/Cargo.toml +++ b/backend/crates/gateway-core/Cargo.toml @@ -17,8 +17,9 @@ bytes = "1.11.1" futures = "0.3.32" futures-timer = "3.0.3" serde = { version = "1.0.228", features = ["derive"] } -serde_json = { version = "1.0.150", features = ["preserve_order"] } +serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] } thiserror = "2.0.18" +tracing = "0.1.44" uuid = { version = "1.23.3", features = ["v7"] } [lints] diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index 287543ae6..7924bc7ec 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -1,6 +1,7 @@ //! 唯一的账号重试、发送与下游 commit barrier owner。 use std::collections::BTreeSet; +use std::future::Future; use std::num::NonZeroU32; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime}; @@ -17,7 +18,7 @@ use crate::engine::{ ModelRequestTimings, NewModelRequest, ProviderAccountStateOwner, ProviderAttemptOutcome, RequestAttemptContext, UpstreamSendState, }; -use crate::error::{GatewayError, GatewayErrorKind, ProviderError, ProviderErrorKind}; +use crate::error::{GatewayError, GatewayErrorKind, ProviderError, ProviderErrorKind, StoreError}; use crate::event::{ GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation, }; @@ -271,10 +272,9 @@ where return Err(EngineError::DownstreamCommitRequired); } if let Some(pending) = self.pending_terminal_failure.take() { - let engine_error = provider_engine_error(&pending.error); self.finish_provider_error_with_send_state(&pending.error, pending.send_state) .await?; - return Err(engine_error); + return Err(provider_engine_error(pending.error)); } if self.finalized { return Ok(None); @@ -341,10 +341,9 @@ where PullOutcome::TerminalFailure { error, send_state, .. } => { - let engine_error = provider_engine_error(&error); self.finish_provider_error_with_send_state(&error, send_state) .await?; - return Err(engine_error); + return Err(provider_engine_error(error)); } PullOutcome::End => { if events.is_empty() { @@ -357,11 +356,11 @@ where } } - /// 在协议 adapter 真正写出首字节前持久化下游不可撤回边界。 + /// 在协议 adapter 真正写出首字节前记录下游不可撤回边界。 /// /// # Errors /// - /// 没有待提交结果、重复提交或 Store 失败时返回错误。 + /// 没有待提交结果或重复提交时返回错误；观测 Store 失败只记服务端告警。 pub async fn commit_downstream( &mut self, client_status_code: Option, @@ -370,10 +369,14 @@ where return Err(EngineError::InvalidDeliveryState); } let committed_at = SystemTime::now(); - self.engine - .store() - .mark_downstream_committed(&self.request_id, committed_at, client_status_code) - .await?; + if self.request_persisted { + best_effort_store_write(self.engine.store().mark_downstream_committed( + &self.request_id, + committed_at, + client_status_code, + )) + .await; + } self.downstream_committed_at = Some(committed_at); self.client_status_code = client_status_code; self.delivery_pending = false; @@ -387,7 +390,7 @@ where /// /// # Errors /// - /// 状态已经写入或 Store 无法写回时返回错误。 + /// 状态已经写入时返回错误；观测 Store 失败只记服务端告警。 pub async fn record_client_status( &mut self, client_status_code: u16, @@ -396,10 +399,12 @@ where return Err(EngineError::InvalidDeliveryState); } if self.request_persisted { - self.engine - .store() - .record_client_status(&self.request_id, client_status_code) - .await?; + best_effort_store_write( + self.engine + .store() + .record_client_status(&self.request_id, client_status_code), + ) + .await; } self.client_status_code = Some(client_status_code); Ok(()) @@ -513,40 +518,17 @@ where return Err(EngineError::Deadline); } PollBoundary::Item(Some(Ok(mut event))) => { - if let Some(observation) = event.take_observation() - && let Err(error) = self.observe_response(observation) - { - match self.handle_stream_error(error).await? { - StreamErrorOutcome::AttemptDiscarded => { - return Ok(PullOutcome::AttemptDiscarded); - } - terminal @ StreamErrorOutcome::TerminalFailure { .. } => { - return Ok(terminal.into_pull_outcome()); - } - } - } - let mut identity_error = None; - for fact in event.canonical_facts() { - if let Err(error) = self.observe_response_identity(fact) { - identity_error = Some(error); - break; - } - } - if let Some(error) = identity_error { - match self.handle_stream_error(error).await? { - StreamErrorOutcome::AttemptDiscarded => { - return Ok(PullOutcome::AttemptDiscarded); - } - terminal @ StreamErrorOutcome::TerminalFailure { .. } => { - return Ok(terminal.into_pull_outcome()); - } - } + if let Some(observation) = event.take_observation() { + self.observe_response(observation); } for fact in event.canonical_facts() { - self.observe_event(fact).await?; + self.observe_response_identity(fact); + } + for fact in event.canonical_facts() { + self.observe_event(fact).await; } if event.wire_event().is_some() && !event.has_canonical_facts() { - self.observe_wire_event().await?; + self.observe_wire_event().await; } if !event.has_client_event() { continue; @@ -672,21 +654,20 @@ where Ok(stream) => stream, Err(error) => { if error.kind() == ProviderErrorKind::NoEligibleAccount - && let Some(exhaustion) = self.last_account_exhaustion.clone() + && let Some(exhaustion) = self.last_account_exhaustion.take() { let send_state = self.current_send_state(); let mut events = std::mem::take(&mut self.last_account_exhaustion_events); if !events.is_empty() { - self.observe_atomic_terminal_events(&mut events)?; + self.observe_atomic_terminal_events(&mut events); return Ok(Some(PullOutcome::TerminalFailure { events, error: exhaustion, send_state, })); } - let engine_error = provider_engine_error(&exhaustion); self.finish_provider_error(&exhaustion).await?; - return Err(engine_error); + return Err(provider_engine_error(exhaustion)); } if !(matches!( error.kind(), @@ -695,9 +676,8 @@ where { self.record_provider_failure(candidate.provider().clone(), error.kind()); } - let engine_error = provider_engine_error(&error); self.finish_provider_error(&error).await?; - return Err(engine_error); + return Err(provider_engine_error(error)); } }, }; @@ -786,17 +766,23 @@ where http_version: None, }; if self.request_persisted { - self.engine.store().record_attempt(attempt_record).await?; + best_effort_store_write(self.engine.store().record_attempt(attempt_record)).await; } else { let request = self .pending_request - .take() + .as_ref() .ok_or(EngineError::InvalidDeliveryState)?; - self.engine - .store() - .create_model_request_with_attempt(request, attempt_record) - .await?; - self.request_persisted = true; + if best_effort_store_write( + self.engine + .store() + .create_model_request_with_attempt(request.clone(), attempt_record), + ) + .await + .is_some() + { + self.pending_request = None; + self.request_persisted = true; + } } self.attempts = next_attempt.get(); self.current = Some(CurrentAttempt { @@ -811,18 +797,9 @@ where Ok(None) } - async fn observe_event(&mut self, event: &GatewayEvent) -> Result<(), EngineError> { + async fn observe_event(&mut self, event: &GatewayEvent) { self.observe_event_facts(event); - let current = self.current.as_mut().ok_or(EngineError::NoActiveAttempt)?; - if !current.send_observed { - self.engine - .store() - .mark_send_state(&self.request_id, UpstreamSendState::Sent) - .await?; - current.send_observed = true; - self.send_state_watermark = UpstreamSendState::Sent; - } - Ok(()) + self.mark_send_observed().await; } fn observe_event_facts(&mut self, event: &GatewayEvent) { @@ -841,34 +818,38 @@ where } } - async fn observe_wire_event(&mut self) -> Result<(), EngineError> { - let current = self.current.as_mut().ok_or(EngineError::NoActiveAttempt)?; + async fn observe_wire_event(&mut self) { + self.mark_send_observed().await; + } + + async fn mark_send_observed(&mut self) { + let Some(current) = self.current.as_mut() else { + return; + }; if !current.send_observed { - self.engine - .store() - .mark_send_state(&self.request_id, UpstreamSendState::Sent) - .await?; + if self.request_persisted { + best_effort_store_write( + self.engine + .store() + .mark_send_state(&self.request_id, UpstreamSendState::Sent), + ) + .await; + } current.send_observed = true; self.send_state_watermark = UpstreamSendState::Sent; } - Ok(()) } - fn observe_response( - &mut self, - observation: ProviderResponseObservation, - ) -> Result<(), ProviderError> { - let send_state = self.current_send_state(); - let current = self - .current - .as_mut() - .ok_or_else(|| ProviderError::new(ProviderErrorKind::Protocol, send_state))?; + fn observe_response(&mut self, observation: ProviderResponseObservation) { + let Some(current) = self.current.as_mut() else { + return; + }; if current .response_observation .as_ref() .is_some_and(|existing| existing.transport() != observation.transport()) { - return Err(ProviderError::new(ProviderErrorKind::Protocol, send_state)); + return; } let observed = observation.timings(); if let Some(value) = observed.transport_decision_wait_ms { @@ -896,32 +877,16 @@ where self.timings.provider_processing_ms = Some(value); } current.response_observation = Some(observation); - Ok(()) } - fn observe_response_identity(&mut self, event: &GatewayEvent) -> Result<(), ProviderError> { + fn observe_response_identity(&mut self, event: &GatewayEvent) { let metadata = match event { GatewayEvent::Started(metadata) | GatewayEvent::Completed(metadata) => metadata, - _ => return Ok(()), + _ => return, }; - let response_id = metadata.response_id(); - if response_id.is_empty() { - return Ok(()); - } - if self - .upstream_response_id - .as_deref() - .is_some_and(|expected| expected != response_id) - { - return Err(ProviderError::new( - ProviderErrorKind::Protocol, - UpstreamSendState::Sent, - )); - } - let response_id = response_id.to_owned(); + let response_id = metadata.response_id().to_owned(); self.client_response_id = Some(response_id.clone()); self.upstream_response_id = Some(response_id); - Ok(()) } /// 失败可重试时要求丢弃本 attempt；预算耗尽时可返回最后一批原始失败事件。 @@ -937,10 +902,6 @@ where error.kind(), ProviderErrorKind::RateLimited | ProviderErrorKind::QuotaExhausted ); - if account_exhausted { - self.last_account_exhaustion = Some(error.clone()); - self.last_account_exhaustion_events.clear(); - } self.record_provider_failure(current.metadata.provider().clone(), error.kind()); // attempt_send_state 是本 attempt 自身的发送事实，驱动重试门； // 持久化与终态用请求级水位，二者不可混用（水位会把早先 attempt 的 @@ -951,10 +912,14 @@ where error.send_state() }; let send_state = self.raise_send_watermark(attempt_send_state); - self.engine - .store() - .mark_send_state(&self.request_id, send_state) - .await?; + if self.request_persisted { + best_effort_store_write( + self.engine + .store() + .mark_send_state(&self.request_id, send_state), + ) + .await; + } let provider_proved_replay_safe = provider_proved_replay_safe(&error); let continuation_retry = self.prepare_continuation_retry( ¤t, @@ -962,6 +927,12 @@ where attempt_send_state, provider_proved_replay_safe, ); + let pre_delivery_retry = self.account_selection.required_account().is_none() + && self.continuation_attempt == ContinuationAttempt::None + && self.downstream_committed_at.is_none() + && !self.delivery_pending + && error.allows_pre_delivery_retry() + && self.attempts < self.plan.max_attempts().get(); let ordinary_retry = self.account_selection.required_account().is_none() && self.continuation_attempt == ContinuationAttempt::None && self.downstream_committed_at.is_none() @@ -984,12 +955,13 @@ where .credential_recovery_attempted_accounts .contains(account) }); - let retryable = continuation_retry || same_account_retry || ordinary_retry; + let retryable = + continuation_retry || same_account_retry || ordinary_retry || pre_delivery_retry; if retryable { - if account_exhausted { - self.last_account_exhaustion_events = atomic_client_events; - } + // 普通 clone 只保留稳定事实；原始失败响应只能随最后一次客户端可见 + // 的账号耗尽错误继续，不能进入中间失败记录。 + let persistence_error = self.request_persisted.then(|| error.clone()); if same_account_retry { if let Some(account) = current.metadata.provider_account_id() { self.credential_recovery_attempted_accounts @@ -1003,36 +975,43 @@ where { self.excluded_accounts.insert(account.clone()); } - self.engine - .store() - .record_intermediate_failure(IntermediateFailure { - request_id: self.request_id.clone(), - attempt_index: current.index, - trigger: current.trigger, - provider_kind: current.metadata.provider().clone(), - account_id: current.metadata.provider_account_id().cloned(), - upstream_model_id: current.metadata.upstream_model().clone(), - upstream_status_code: current - .response_observation - .as_ref() - .and_then(ProviderResponseObservation::status_code), - upstream_request_id: current - .response_observation - .as_ref() - .and_then(ProviderResponseObservation::request_id) - .map(|value| value.as_str().to_owned()), - latency: current.started_at.elapsed().unwrap_or_default(), - error, - }) - .await?; + if account_exhausted { + self.last_account_exhaustion_events = atomic_client_events; + self.last_account_exhaustion = Some(error); + } + if let Some(error) = persistence_error { + best_effort_store_write( + self.engine + .store() + .record_intermediate_failure(IntermediateFailure { + request_id: self.request_id.clone(), + attempt_index: current.index, + trigger: current.trigger, + provider_kind: current.metadata.provider().clone(), + account_id: current.metadata.provider_account_id().cloned(), + upstream_model_id: current.metadata.upstream_model().clone(), + upstream_status_code: current + .response_observation + .as_ref() + .and_then(ProviderResponseObservation::status_code), + upstream_request_id: current + .response_observation + .as_ref() + .and_then(ProviderResponseObservation::request_id) + .map(|value| value.as_str().to_owned()), + latency: current.started_at.elapsed().unwrap_or_default(), + error, + }), + ) + .await; + } self.reset_uncommitted_observations(); return Ok(StreamErrorOutcome::AttemptDiscarded); } - let engine_error = provider_engine_error(&error); self.current = Some(current); if !atomic_client_events.is_empty() { - self.observe_atomic_terminal_events(&mut atomic_client_events)?; + self.observe_atomic_terminal_events(&mut atomic_client_events); return Ok(StreamErrorOutcome::TerminalFailure { events: atomic_client_events, error, @@ -1041,25 +1020,19 @@ where } self.finish_provider_error_with_send_state(&error, send_state) .await?; - Err(engine_error) + Err(provider_engine_error(error)) } - fn observe_atomic_terminal_events( - &mut self, - events: &mut [ProviderEvent], - ) -> Result<(), EngineError> { + fn observe_atomic_terminal_events(&mut self, events: &mut [ProviderEvent]) { for event in events { if let Some(observation) = event.take_observation() { - self.observe_response(observation) - .map_err(EngineError::Provider)?; + self.observe_response(observation); } for fact in event.canonical_facts() { - self.observe_response_identity(fact) - .map_err(EngineError::Provider)?; + self.observe_response_identity(fact); self.observe_event_facts(fact); } } - Ok(()) } fn prepare_continuation_retry( @@ -1149,34 +1122,38 @@ where .and_then(ProviderResponseObservation::status_code); let (upstream_transport, http_version, websocket_pool) = self.current_transport_observation(); + let service_tier = self.current_service_tier(); let provider_metadata_json = self.current_provider_metadata_json(); - self.engine - .store() - .finalize_model_request(ModelRequestFinalization { - request_id: self.request_id.clone(), - outcome: ExecutionOutcome::Succeeded, - send_state: UpstreamSendState::Sent, - attempt_count: self.attempts, - downstream_committed_at: self.downstream_committed_at, - client_status_code: self.client_status_code, - upstream_status_code, - client_response_id: self.client_response_id.clone(), - upstream_request_id, - upstream_response_id: self.upstream_response_id.clone(), - upstream_transport, - http_version, - websocket_pool, - provider_metadata_json, - error: None, - provider_error_code: None, - retry_after_ms: None, - usage: self.usage.clone(), - image_generation_succeeded: self.image_generation_succeeded(), - cost: self.cost.clone(), - timings: self.timings.clone(), - completed_at, - }) - .await?; + if self.request_persisted { + best_effort_store_write(self.engine.store().finalize_model_request( + ModelRequestFinalization { + request_id: self.request_id.clone(), + outcome: ExecutionOutcome::Succeeded, + send_state: UpstreamSendState::Sent, + attempt_count: self.attempts, + downstream_committed_at: self.downstream_committed_at, + client_status_code: self.client_status_code, + upstream_status_code, + client_response_id: self.client_response_id.clone(), + upstream_request_id, + upstream_response_id: self.upstream_response_id.clone(), + upstream_transport, + http_version, + websocket_pool, + service_tier, + provider_metadata_json, + error: None, + provider_error_code: None, + retry_after_ms: None, + usage: self.usage.clone(), + image_generation_succeeded: self.image_generation_succeeded(), + cost: self.cost.clone(), + timings: self.timings.clone(), + completed_at, + }, + )) + .await; + } self.finalized = true; Ok(()) } @@ -1245,11 +1222,13 @@ where UpstreamSendState::Ambiguous }; let send_state = self.raise_send_watermark(send_state); - if self.attempts > 0 { - self.engine - .store() - .mark_send_state(&self.request_id, send_state) - .await?; + if self.attempts > 0 && self.request_persisted { + best_effort_store_write( + self.engine + .store() + .mark_send_state(&self.request_id, send_state), + ) + .await; } self.finish_failure(FailureFinalization { outcome, @@ -1299,36 +1278,42 @@ where .and_then(ProviderResponseObservation::status_code); let (upstream_transport, http_version, websocket_pool) = self.current_transport_observation(); + let service_tier = self.current_service_tier(); let provider_metadata_json = self.current_provider_metadata_json(); - self.engine - .store() - .finalize_model_request(ModelRequestFinalization { - request_id: self.request_id.clone(), - outcome: finalization.outcome, - send_state: finalization.send_state, - attempt_count: self.attempts, - downstream_committed_at: self.downstream_committed_at, - client_status_code: self.client_status_code, - upstream_status_code: finalization.upstream_status_code.or(observed_status_code), - client_response_id: self.client_response_id.clone(), - upstream_request_id, - upstream_response_id: finalization - .provider_response_id - .or_else(|| self.upstream_response_id.clone()), - upstream_transport, - http_version, - websocket_pool, - provider_metadata_json, - error: Some(finalization.error), - provider_error_code: finalization.provider_error_code, - retry_after_ms: finalization.retry_after_ms, - usage: self.usage.clone(), - image_generation_succeeded: self.image_generation_succeeded(), - cost: self.cost.clone(), - timings: self.timings.clone(), - completed_at, - }) - .await?; + best_effort_store_write( + self.engine + .store() + .finalize_model_request(ModelRequestFinalization { + request_id: self.request_id.clone(), + outcome: finalization.outcome, + send_state: finalization.send_state, + attempt_count: self.attempts, + downstream_committed_at: self.downstream_committed_at, + client_status_code: self.client_status_code, + upstream_status_code: finalization + .upstream_status_code + .or(observed_status_code), + client_response_id: self.client_response_id.clone(), + upstream_request_id, + upstream_response_id: finalization + .provider_response_id + .or_else(|| self.upstream_response_id.clone()), + upstream_transport, + http_version, + websocket_pool, + service_tier, + provider_metadata_json, + error: Some(finalization.error), + provider_error_code: finalization.provider_error_code, + retry_after_ms: finalization.retry_after_ms, + usage: self.usage.clone(), + image_generation_succeeded: self.image_generation_succeeded(), + cost: self.cost.clone(), + timings: self.timings.clone(), + completed_at, + }), + ) + .await; self.finalized = true; Ok(()) } @@ -1359,6 +1344,14 @@ where .map(|metadata| metadata.as_json().to_owned()) } + fn current_service_tier(&self) -> Option { + self.current + .as_ref() + .and_then(|current| current.response_observation.as_ref()) + .and_then(ProviderResponseObservation::service_tier) + .map(str::to_owned) + } + fn image_generation_succeeded(&self) -> Option { self.image_generation_requested .then(|| self.usage.image_output_tokens.unwrap_or_default() > 0) @@ -1408,6 +1401,12 @@ where } } +async fn best_effort_store_write( + write: impl Future>, +) -> Option { + write.await.ok() +} + enum PullOutcome { Event(ProviderEvent), AttemptDiscarded, @@ -1517,11 +1516,11 @@ fn initial_continuation_attempt( } } -fn provider_engine_error(error: &ProviderError) -> EngineError { +fn provider_engine_error(error: ProviderError) -> EngineError { if error.kind() == ProviderErrorKind::Cancelled { EngineError::Cancelled } else { - EngineError::Provider(error.clone()) + EngineError::Provider(error) } } diff --git a/backend/crates/gateway-core/src/engine/credential.rs b/backend/crates/gateway-core/src/engine/credential.rs index 8e1604c9b..ae574afbb 100644 --- a/backend/crates/gateway-core/src/engine/credential.rs +++ b/backend/crates/gateway-core/src/engine/credential.rs @@ -932,6 +932,43 @@ impl AccountAvailabilityPolicy { } } +/// 候选账号未进入调度池的约束。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AccountSchedulingBlocker { + LocalAvailability, + Excluded, + ConcurrencyLimit, + RequestInterval, +} + +/// 优先账号在本次选择中的处理结果。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PreferredAccountSelection { + NotRequested, + Hit, + Missing, + Blocked(AccountSchedulingBlocker), +} + +/// 一次完整账号选择；同时保留优先账号决策，供 Provider 记录调度遥测。 +#[derive(Debug, Clone, Copy)] +pub struct AccountSelection<'a> { + candidate: &'a AccountCandidate, + preferred: PreferredAccountSelection, +} + +impl<'a> AccountSelection<'a> { + #[must_use] + pub const fn candidate(self) -> &'a AccountCandidate { + self.candidate + } + + #[must_use] + pub const fn preferred(self) -> PreferredAccountSelection { + self.preferred + } +} + /// 同一 target 内唯一的账号排序器。 #[derive(Debug, Default, Clone, Copy)] pub struct AccountSelector; @@ -943,43 +980,55 @@ impl AccountSelector { &self, candidates: &'a [AccountCandidate], context: &AccountSelectionContext, - ) -> Option<&'a AccountCandidate> { + ) -> Option> { + let preferred = if let Some(preferred) = context.preferred_account.as_ref() { + match candidates + .iter() + .find(|candidate| candidate.account.id() == preferred) + { + Some(candidate) => match self.scheduling_blocker(candidate, context) { + None => { + return Some(AccountSelection { + candidate, + preferred: PreferredAccountSelection::Hit, + }); + } + Some(blocker) => PreferredAccountSelection::Blocked(blocker), + }, + None => PreferredAccountSelection::Missing, + } + } else { + PreferredAccountSelection::NotRequested + }; + let mut eligible = candidates .iter() - .filter(|candidate| eligible(candidate, context)) + .filter(|candidate| self.scheduling_blocker(candidate, context).is_none()) .collect::>(); if eligible.is_empty() { return None; } - if let Some(preferred) = context.preferred_account.as_ref() - && let Some(candidate) = eligible - .iter() - .find(|candidate| candidate.account.id() == preferred) - && (context.policy.strategy() == RotationStrategy::Sticky - || (context.policy.strategy() == RotationStrategy::Smart - && !smart_preference_should_escape(&candidate.signals))) - { - return Some(candidate); - } - - match context.policy.strategy() { - RotationStrategy::QuotaResetPriority => eligible.sort_by_key(|candidate| { - ( - candidate.signals.quota_reset_at.is_none(), - candidate.signals.quota_reset_at, - candidate.signals.in_flight, - candidate.signals.last_started_at, - candidate.account.id().clone(), - ) - }), + let candidate = match context.policy.strategy() { + RotationStrategy::QuotaResetPriority => { + eligible.sort_by_key(|candidate| { + ( + candidate.signals.quota_reset_at.is_none(), + candidate.signals.quota_reset_at, + candidate.signals.in_flight, + candidate.signals.last_started_at, + candidate.account.id().clone(), + ) + }); + eligible.first().copied()? + } RotationStrategy::RoundRobin => { eligible.sort_by_key(|candidate| candidate.account.id().clone()); let index = context.round_robin_cursor as usize % eligible.len(); - return Some(eligible[index]); + eligible.get(index).copied()? } RotationStrategy::Smart => { - return select_smart_candidate(&eligible, context.round_robin_cursor); + select_smart_candidate(&eligible, context.round_robin_cursor)? } RotationStrategy::Sticky => { eligible.sort_by_key(|candidate| { @@ -988,9 +1037,44 @@ impl AccountSelector { candidate.account.id().clone(), ) }); + eligible.first().copied()? } + }; + Some(AccountSelection { + candidate, + preferred, + }) + } + + fn scheduling_blocker( + &self, + candidate: &AccountCandidate, + context: &AccountSelectionContext, + ) -> Option { + if !context.availability.bypasses_local_availability() + && !candidate.account.is_schedulable(context.now) + { + return Some(AccountSchedulingBlocker::LocalAvailability); } - eligible.into_iter().next() + if context.excluded_accounts.contains(candidate.account.id()) { + return Some(AccountSchedulingBlocker::Excluded); + } + if candidate.signals.in_flight >= context.policy.max_concurrent_per_account().get() { + return Some(AccountSchedulingBlocker::ConcurrencyLimit); + } + if candidate + .signals + .last_started_at + .is_some_and(|last_started| { + !context + .now + .duration_since(last_started) + .is_ok_and(|elapsed| elapsed >= context.policy.request_interval()) + }) + { + return Some(AccountSchedulingBlocker::RequestInterval); + } + None } } @@ -998,8 +1082,6 @@ const SMART_LOAD_WEIGHT: f64 = 1.0; const SMART_QUOTA_WEIGHT: f64 = 0.8; const SMART_FAILURE_WEIGHT: f64 = 1.0; const SMART_LATENCY_WEIGHT: f64 = 0.5; -const SMART_PREFERENCE_FAILURE_ESCAPE_BASIS_POINTS: u16 = 5_000; -const SMART_PREFERENCE_LATENCY_ESCAPE_MS: u64 = 15_000; struct SmartNormalization { max_in_flight: u32, @@ -1120,35 +1202,6 @@ fn higher_is_better(value: u64, minimum: u64, maximum: u64) -> f64 { (value.saturating_sub(minimum) as f64 / (maximum - minimum) as f64).clamp(0.0, 1.0) } -fn smart_preference_should_escape(signals: &AccountRuntimeSignals) -> bool { - signals - .failure_rate_basis_points - .is_some_and(|rate| rate > SMART_PREFERENCE_FAILURE_ESCAPE_BASIS_POINTS) - || signals - .first_output_latency_ms - .is_some_and(|latency| latency > SMART_PREFERENCE_LATENCY_ESCAPE_MS) -} - -fn eligible(candidate: &AccountCandidate, context: &AccountSelectionContext) -> bool { - if (!context.availability.bypasses_local_availability() - && !candidate.account.is_schedulable(context.now)) - || context.excluded_accounts.contains(candidate.account.id()) - || candidate.signals.in_flight >= context.policy.max_concurrent_per_account().get() - { - return false; - } - - candidate - .signals - .last_started_at - .is_none_or(|last_started| { - context - .now - .duration_since(last_started) - .is_ok_and(|elapsed| elapsed >= context.policy.request_interval()) - }) -} - /// Credential 值对象构造错误。 #[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] pub enum CredentialError { diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index 69f0a87eb..180331141 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -36,7 +36,7 @@ use crate::routing::{ }; const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60); -const CONTINUATION_AFFINITY_TIMEOUT: Duration = Duration::from_millis(100); +const REDIS_COORDINATION_TIMEOUT: Duration = Duration::from_millis(100); #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ClientTransport { @@ -233,15 +233,6 @@ impl DefaultExecutionService { request.client.policy.authorize().map_err(|_| { GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled") })?; - if !request.client.snapshot.contains_public_model_for_provider( - &request.public_model, - request.client.policy.provider_kind(), - ) { - return Err(GatewayError::new( - GatewayErrorKind::ModelNotFound, - "requested model was not found", - )); - } let started_at = SystemTime::now(); let deadline_at = started_at .checked_add(MODEL_REQUEST_DEADLINE) @@ -260,7 +251,7 @@ impl DefaultExecutionService { let continuation = match request.metadata.previous_response_id.as_ref() { Some(previous) => { let resolve = self.continuation.resolve(previous).fuse(); - let timeout = Delay::new(CONTINUATION_AFFINITY_TIMEOUT).fuse(); + let timeout = Delay::new(REDIS_COORDINATION_TIMEOUT).fuse(); pin_mut!(resolve, timeout); let pin = select_biased! { result = resolve => result.ok().flatten(), @@ -398,14 +389,33 @@ impl DefaultExecutionService { &self, provider_kind: &ProviderKind, ) -> Result { - let blocked_providers = match self.circuits.decision(provider_kind).await.map_err(|_| { - GatewayError::new( - GatewayErrorKind::NoAvailableProvider, - "provider health state is temporarily unavailable", - ) - })? { - ProviderCircuitDecision::Allow => BTreeSet::new(), - ProviderCircuitDecision::BlockedUntil(_) => BTreeSet::from([provider_kind.clone()]), + let decision = self.circuits.decision(provider_kind).fuse(); + let timeout = Delay::new(REDIS_COORDINATION_TIMEOUT).fuse(); + pin_mut!(decision, timeout); + let decision = select_biased! { + result = decision => Some(result), + _ = timeout => None, + }; + let blocked_providers = match decision { + Some(Ok(ProviderCircuitDecision::Allow)) => BTreeSet::new(), + Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => { + BTreeSet::from([provider_kind.clone()]) + } + Some(Err(error)) => { + tracing::warn!( + provider = provider_kind.as_str(), + %error, + "Provider circuit 读取失败，按可重建协调状态 fail-open" + ); + BTreeSet::new() + } + None => { + tracing::warn!( + provider = provider_kind.as_str(), + "Provider circuit 读取超时，按可重建协调状态 fail-open" + ); + BTreeSet::new() + } }; Ok(RoutingContext { provider_kind: Some(provider_kind.clone()), @@ -665,10 +675,13 @@ struct AdmissionLease { impl AdmissionLease { async fn release(self) { - let _ = self + if let Err(error) = self .port .release(&self.client_api_key_id, &self.model_request_id) - .await; + .await + { + tracing::warn!(%error, "Client admission 释放失败，依赖租约 TTL 收敛"); + } } } @@ -743,13 +756,7 @@ impl DefaultExecutionSession { return; }; self.continuation_recorded = true; - let record = self.continuation.record(pin).fuse(); - let timeout = Delay::new(CONTINUATION_AFFINITY_TIMEOUT).fuse(); - pin_mut!(record, timeout); - select_biased! { - _ = record => {}, - _ = timeout => {}, - } + record_native_continuation(self.continuation.as_ref(), pin).await; } async fn finalize_detached(&mut self) { @@ -757,8 +764,10 @@ impl DefaultExecutionSession { return; }; core.cancel(); - if !core.is_finalized() { - let _ = core.cancel_and_finalize().await; + if !core.is_finalized() + && let Err(error) = core.cancel_and_finalize().await + { + tracing::warn!(%error, "Detached execution 终态收敛失败"); } let pending = core .provider_attempt_outcomes() @@ -879,7 +888,32 @@ async fn publish_provider_attempt_outcomes( } Some(_) => continue, }; - let _ = result; + if let Err(error) = result { + tracing::warn!( + provider = outcome.provider_kind().as_str(), + %error, + "Provider circuit feedback 写入失败，数据面不受影响" + ); + } + } +} + +async fn record_native_continuation( + continuation: &dyn NativeContinuationPort, + pin: NativeContinuationPin, +) { + let record = continuation.record(pin).fuse(); + let timeout = Delay::new(REDIS_COORDINATION_TIMEOUT).fuse(); + pin_mut!(record, timeout); + select_biased! { + result = record => { + if let Err(error) = result { + tracing::warn!(%error, "Continuation affinity 写入失败，后续请求将退化为外部续接"); + } + }, + _ = timeout => { + tracing::warn!("Continuation affinity 后台写入超时，已丢弃本次亲和记录"); + }, } } diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 8c290a884..a071804ad 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -615,6 +615,7 @@ pub struct ModelRequestFinalization { pub upstream_transport: Option, pub http_version: Option, pub websocket_pool: Option, + pub service_tier: Option, /// Provider 已筛选的专有观测 JSON；Core 不解释字段。 pub provider_metadata_json: Option, pub error: Option, diff --git a/backend/crates/gateway-core/src/engine/provider.rs b/backend/crates/gateway-core/src/engine/provider.rs index 42907d57c..6041f107c 100644 --- a/backend/crates/gateway-core/src/engine/provider.rs +++ b/backend/crates/gateway-core/src/engine/provider.rs @@ -16,7 +16,7 @@ use crate::engine::credential::{ }; use crate::engine::{AttemptContext, UpstreamSendState}; use crate::error::{ - IdentifierError, ProviderError, ProviderErrorKind, SafeUpstreamValue, validate_text, + IdentifierError, OpaqueUpstreamValue, ProviderError, ProviderErrorKind, validate_text, }; use crate::event::{EventSequenceValidator, ProviderEvent}; use crate::operation::Operation; @@ -125,7 +125,7 @@ pub struct ProviderCallMetadata { provider: ProviderKind, upstream_model: UpstreamModelId, resource: ProviderResource, - upstream_request_id: Option, + upstream_request_id: Option, transport: UpstreamTransport, } @@ -149,7 +149,7 @@ impl ProviderCallMetadata { /// 设置 adapter 已分类为非 bearer 的 request ID。 #[must_use] - pub fn with_upstream_request_id(mut self, request_id: SafeUpstreamValue) -> Self { + pub fn with_upstream_request_id(mut self, request_id: OpaqueUpstreamValue) -> Self { self.upstream_request_id = Some(request_id); self } @@ -186,7 +186,7 @@ impl ProviderCallMetadata { /// 返回安全上游 request ID。 #[must_use] - pub const fn upstream_request_id(&self) -> Option<&SafeUpstreamValue> { + pub const fn upstream_request_id(&self) -> Option<&OpaqueUpstreamValue> { self.upstream_request_id.as_ref() } diff --git a/backend/crates/gateway-core/src/error.rs b/backend/crates/gateway-core/src/error.rs index 4a1ea9fa9..f62e9c079 100644 --- a/backend/crates/gateway-core/src/error.rs +++ b/backend/crates/gateway-core/src/error.rs @@ -3,10 +3,11 @@ use std::fmt; use std::time::Duration; +use bytes::Bytes; use thiserror::Error; use crate::engine::UpstreamSendState; -use crate::event::ProviderEvent; +use crate::event::{ProviderEvent, ProviderResponseHeader}; /// 应用层标识不满足约束。 #[derive(Debug, Clone, PartialEq, Eq, Error)] @@ -175,22 +176,18 @@ pub enum ContinuationFailure { HistoryUnavailable, } -/// Adapter 已明确分类为非 bearer、可用于诊断的上游值。 +/// Adapter 已明确分类为非 bearer 的不透明上游值。 +/// +/// Core 不解释或校验内容，只通过自定义 [`Debug`] 避免它意外进入日志。协议 adapter +/// 可在原客户端响应中读取原值。 #[derive(Clone, PartialEq, Eq, Hash)] -pub struct SafeUpstreamValue(String); +pub struct OpaqueUpstreamValue(String); -impl SafeUpstreamValue { - const MAX_BYTES: usize = 256; - - /// 校验并创建诊断值。 - /// - /// # Errors - /// - /// 空值、超长值或控制字符会返回 [`IdentifierError`]。 - pub fn new(value: impl Into) -> Result { - let value = value.into(); - validate_text(&value, Self::MAX_BYTES, false, None)?; - Ok(Self(value)) +impl OpaqueUpstreamValue { + /// 原样保存上游值。 + #[must_use] + pub fn new(value: impl Into) -> Self { + Self(value.into()) } /// 返回已经由 adapter 安全分类的原值。 @@ -200,9 +197,9 @@ impl SafeUpstreamValue { } } -impl fmt::Debug for SafeUpstreamValue { +impl fmt::Debug for OpaqueUpstreamValue { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("SafeUpstreamValue()") + formatter.write_str("OpaqueUpstreamValue()") } } @@ -213,31 +210,23 @@ impl fmt::Debug for SafeUpstreamValue { #[derive(Clone, PartialEq, Eq)] pub struct ClientVisibleUpstreamError { message: String, - code: Option, - error_type: Option, + code: Option, + error_type: Option, } impl ClientVisibleUpstreamError { - const MAX_MESSAGE_BYTES: usize = 8_192; - - /// 从结构化上游 `message`、`code` 和 `type` 创建客户端可见错误。 - /// - /// # Errors - /// - /// message 为空、过长或含控制字符时返回错误；可选 code/type 不满足安全 - /// 约束时会被省略，不会把整段上游 body 重新暴露。 + /// 原样保存结构化上游 `message`、`code` 和 `type`。 + #[must_use] pub fn new( message: impl Into, code: Option, error_type: Option, - ) -> Result { - let message = message.into(); - validate_text(&message, Self::MAX_MESSAGE_BYTES, false, None)?; - Ok(Self { - message, - code: code.and_then(|value| SafeUpstreamValue::new(value).ok()), - error_type: error_type.and_then(|value| SafeUpstreamValue::new(value).ok()), - }) + ) -> Self { + Self { + message: message.into(), + code: code.map(OpaqueUpstreamValue::new), + error_type: error_type.map(OpaqueUpstreamValue::new), + } } /// 返回原上游的结构化 message。 @@ -249,13 +238,13 @@ impl ClientVisibleUpstreamError { /// 返回原上游的结构化 code。 #[must_use] pub fn code(&self) -> Option<&str> { - self.code.as_ref().map(SafeUpstreamValue::as_str) + self.code.as_ref().map(OpaqueUpstreamValue::as_str) } /// 返回原上游的结构化 type。 #[must_use] pub fn error_type(&self) -> Option<&str> { - self.error_type.as_ref().map(SafeUpstreamValue::as_str) + self.error_type.as_ref().map(OpaqueUpstreamValue::as_str) } } @@ -265,6 +254,74 @@ impl fmt::Debug for ClientVisibleUpstreamError { } } +/// 只供当前客户端请求使用的原始上游 HTTP 失败响应。 +/// +/// 该值不属于稳定诊断事实，不能进入日志、探针或持久化。它刻意不实现 +/// [`Clone`]；[`ProviderError`] 的普通 clone 会丢弃它，只有最终失败的原对象 +/// 才能把响应交给协议 adapter。 +#[derive(PartialEq, Eq)] +pub struct ClientVisibleUpstreamResponse { + status: u16, + content_type: Option>, + headers: Vec, + body: Bytes, +} + +impl ClientVisibleUpstreamResponse { + /// 保存 transport 实际收到的状态码、Content-Type 原值和正文。 + #[must_use] + pub fn new(status: u16, content_type: Option>, body: Bytes) -> Self { + Self { + status, + content_type, + headers: Vec::new(), + body, + } + } + + /// 附加 Provider 已剔除账号绑定与 framing 信息的普通响应头。 + #[must_use] + pub fn with_headers(mut self, headers: Vec) -> Self { + self.headers = headers; + self + } + + #[must_use] + pub const fn status(&self) -> u16 { + self.status + } + + #[must_use] + pub fn content_type(&self) -> Option<&[u8]> { + self.content_type.as_deref() + } + + #[must_use] + pub fn headers(&self) -> &[ProviderResponseHeader] { + &self.headers + } + + #[must_use] + pub const fn body(&self) -> &Bytes { + &self.body + } +} + +impl fmt::Debug for ClientVisibleUpstreamResponse { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ClientVisibleUpstreamResponse") + .field("status", &self.status) + .field( + "content_type", + &self.content_type.as_ref().map(|_| ""), + ) + .field("header_count", &self.headers.len()) + .field("body", &"") + .finish() + } +} + /// 单次 Provider 调用的稳定错误。 /// /// 稳定诊断字段不接收原始响应正文，也不会在 `Debug` 或 `Display` 中打印上游 @@ -274,24 +331,30 @@ impl fmt::Debug for ClientVisibleUpstreamError { /// 唯一例外是 [`ProviderError::with_atomic_client_events`]：它只承载尚未交付客户端、 /// 必须与本错误一起完成重试判断的协议事件。Core 会在任何持久化或日志记录前取走 /// 这些事件，不能把它们当作诊断上下文使用。 -#[derive(Clone)] pub struct ProviderError { kind: ProviderErrorKind, send_state: UpstreamSendState, upstream_status: Option, - upstream_code: Option, - upstream_request_id: Option, - upstream_response_id: Option, + upstream_values: Option>, retry_after: Option, continuation_failure: Option, replay_safe: bool, + pre_delivery_retry: bool, credential_recovery_required: bool, retry_same_account: bool, sensitive_context_redacted: bool, client_visible_upstream_error: Option>, + client_visible_upstream_response: Option>, atomic_client_events: Option>, } +#[derive(Clone, Default)] +struct ProviderErrorUpstreamValues { + code: Option, + request_id: Option, + response_id: Option, +} + #[derive(Clone)] struct AtomicClientEvents(Vec); @@ -303,16 +366,16 @@ impl ProviderError { kind, send_state, upstream_status: None, - upstream_code: None, - upstream_request_id: None, - upstream_response_id: None, + upstream_values: None, retry_after: None, continuation_failure: None, replay_safe: false, + pre_delivery_retry: false, credential_recovery_required: false, retry_same_account: false, sensitive_context_redacted: false, client_visible_upstream_error: None, + client_visible_upstream_response: None, atomic_client_events: None, } } @@ -328,25 +391,30 @@ impl ProviderError { /// 附加 adapter 已分类为安全的上游错误 code。 #[must_use] - pub fn with_upstream_code(mut self, code: SafeUpstreamValue) -> Self { - self.upstream_code = Some(code); + pub fn with_upstream_code(mut self, code: OpaqueUpstreamValue) -> Self { + self.upstream_values_mut().code = Some(code); self } /// 附加 adapter 已分类为非 bearer 的上游 request ID。 #[must_use] - pub fn with_upstream_request_id(mut self, request_id: SafeUpstreamValue) -> Self { - self.upstream_request_id = Some(request_id); + pub fn with_upstream_request_id(mut self, request_id: OpaqueUpstreamValue) -> Self { + self.upstream_values_mut().request_id = Some(request_id); self } /// 附加 adapter 已分类为非 bearer 的上游 response ID。 #[must_use] - pub fn with_upstream_response_id(mut self, response_id: SafeUpstreamValue) -> Self { - self.upstream_response_id = Some(response_id); + pub fn with_upstream_response_id(mut self, response_id: OpaqueUpstreamValue) -> Self { + self.upstream_values_mut().response_id = Some(response_id); self } + fn upstream_values_mut(&mut self) -> &mut ProviderErrorUpstreamValues { + self.upstream_values + .get_or_insert_with(|| Box::new(ProviderErrorUpstreamValues::default())) + } + /// 附加 Provider 建议的冷却时间。 #[must_use] pub const fn with_retry_after(mut self, retry_after: Duration) -> Self { @@ -368,6 +436,17 @@ impl ProviderError { self } + /// 允许 Core 仅在客户端尚未收到任何事件时执行一次受预算约束的换号恢复。 + /// + /// 该标记不证明上游未执行，也不等价于 [`Self::with_replay_safe`]；它只表达 + /// Provider 选择了“客户端无感恢复优先”的传输策略。Core 在 continuation、 + /// 指定账号或下游已经进入提交状态时必须忽略它。 + #[must_use] + pub const fn with_pre_delivery_retry(mut self) -> Self { + self.pre_delivery_retry = true; + self + } + /// 要求 Core 在账号凭据已恢复后仅对原账号重放一次。 #[must_use] pub const fn with_same_account_retry(mut self) -> Self { @@ -389,6 +468,16 @@ impl ProviderError { self } + /// 附加仅可返回给当前请求方的原始上游 HTTP 失败响应。 + #[must_use] + pub fn with_client_visible_upstream_response( + mut self, + response: ClientVisibleUpstreamResponse, + ) -> Self { + self.client_visible_upstream_response = Some(Box::new(response)); + self + } + /// 附加必须先由 Core 判断重试、再决定丢弃或交付的客户端协议事件。 /// /// Provider 只能放入当前 attempt 尚未交付的事件；调用方应设置明确的有界缓冲， @@ -427,20 +516,26 @@ impl ProviderError { /// 返回安全分类的上游错误 code。 #[must_use] - pub fn upstream_code(&self) -> Option<&SafeUpstreamValue> { - self.upstream_code.as_ref() + pub fn upstream_code(&self) -> Option<&OpaqueUpstreamValue> { + self.upstream_values + .as_deref() + .and_then(|values| values.code.as_ref()) } /// 返回安全分类的上游 request ID。 #[must_use] - pub fn upstream_request_id(&self) -> Option<&SafeUpstreamValue> { - self.upstream_request_id.as_ref() + pub fn upstream_request_id(&self) -> Option<&OpaqueUpstreamValue> { + self.upstream_values + .as_deref() + .and_then(|values| values.request_id.as_ref()) } /// 返回安全分类的上游 response ID。 #[must_use] - pub fn upstream_response_id(&self) -> Option<&SafeUpstreamValue> { - self.upstream_response_id.as_ref() + pub fn upstream_response_id(&self) -> Option<&OpaqueUpstreamValue> { + self.upstream_values + .as_deref() + .and_then(|values| values.response_id.as_ref()) } /// 返回 Provider 建议的冷却时间。 @@ -461,6 +556,12 @@ impl ProviderError { self.replay_safe } + /// 返回 Provider 是否允许在首个客户端事件前换号恢复传输失败。 + #[must_use] + pub const fn allows_pre_delivery_retry(&self) -> bool { + self.pre_delivery_retry + } + /// 返回 Provider 是否已完成凭据恢复并要求原账号重放。 #[must_use] pub const fn retries_same_account(&self) -> bool { @@ -484,6 +585,12 @@ impl ProviderError { self.client_visible_upstream_error.as_deref() } + /// 返回只供当前协议响应使用的原始上游 HTTP 失败响应。 + #[must_use] + pub fn client_visible_upstream_response(&self) -> Option<&ClientVisibleUpstreamResponse> { + self.client_visible_upstream_response.as_deref() + } + /// 取走只供本次重试/提交决策使用的原子客户端事件。 /// /// Core 必须在克隆错误、记录中间失败或构造终态前调用本方法，避免原始 wire @@ -502,6 +609,27 @@ impl ProviderError { } } +impl Clone for ProviderError { + fn clone(&self) -> Self { + Self { + kind: self.kind, + send_state: self.send_state, + upstream_status: self.upstream_status, + upstream_values: self.upstream_values.clone(), + retry_after: self.retry_after, + continuation_failure: self.continuation_failure, + replay_safe: self.replay_safe, + pre_delivery_retry: self.pre_delivery_retry, + credential_recovery_required: self.credential_recovery_required, + retry_same_account: self.retry_same_account, + sensitive_context_redacted: self.sensitive_context_redacted, + client_visible_upstream_error: self.client_visible_upstream_error.clone(), + client_visible_upstream_response: None, + atomic_client_events: None, + } + } +} + impl fmt::Debug for ProviderError { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter @@ -511,25 +639,20 @@ impl fmt::Debug for ProviderError { .field("upstream_status", &self.upstream_status) .field( "upstream_code", - &self.upstream_code.as_ref().map(|_| ""), + &self.upstream_code().map(|_| ""), ) .field( "upstream_request_id", - &self - .upstream_request_id - .as_ref() - .map(|_| ""), + &self.upstream_request_id().map(|_| ""), ) .field( "upstream_response_id", - &self - .upstream_response_id - .as_ref() - .map(|_| ""), + &self.upstream_response_id().map(|_| ""), ) .field("retry_after", &self.retry_after) .field("continuation_failure", &self.continuation_failure) .field("replay_safe", &self.replay_safe) + .field("pre_delivery_retry", &self.pre_delivery_retry) .field( "credential_recovery_required", &self.credential_recovery_required, @@ -546,6 +669,13 @@ impl fmt::Debug for ProviderError { .as_ref() .map(|_| ""), ) + .field( + "client_visible_upstream_response", + &self + .client_visible_upstream_response + .as_ref() + .map(|_| ""), + ) .field( "atomic_client_events", &self diff --git a/backend/crates/gateway-core/src/event.rs b/backend/crates/gateway-core/src/event.rs index a80ffc916..3b8af2246 100644 --- a/backend/crates/gateway-core/src/event.rs +++ b/backend/crates/gateway-core/src/event.rs @@ -9,7 +9,7 @@ use thiserror::Error; use crate::accounting::{CalculatedCost, ProviderReportedCost, Usage}; use crate::engine::provider::UpstreamTransport; -use crate::error::{IdentifierError, SafeUpstreamValue, validate_text}; +use crate::error::{IdentifierError, OpaqueUpstreamValue, validate_text}; use crate::operation::ProviderSessionState; /// 一次响应的稳定元数据。 @@ -199,7 +199,7 @@ impl ProtocolWireEvent { /// /// # Errors /// - /// 协议名或显式事件名为空、过长或含控制字符时返回错误。 + /// 协议名不满足内部路由标识约束时返回错误。上游事件名保持不透明。 pub fn json( protocol: impl Into, event_type: Option, @@ -212,7 +212,7 @@ impl ProtocolWireEvent { /// /// # Errors /// - /// 协议名、显式事件名或 SSE ID 不满足 wire 安全约束时返回错误。 + /// 协议名不满足内部路由标识约束时返回错误。SSE 元数据保持不透明。 pub fn json_with_sse_metadata( protocol: impl Into, event_type: Option, @@ -222,15 +222,6 @@ impl ProtocolWireEvent { ) -> Result { let protocol = protocol.into(); validate_text(&protocol, 64, true, None)?; - if let Some(event_type) = event_type.as_deref() { - validate_text(event_type, 256, true, None)?; - } - if sse_id - .as_deref() - .is_some_and(|id| id.contains(['\0', '\r', '\n'])) - { - return Err(IdentifierError::ControlCharacter); - } Ok(Self { protocol, event_type, @@ -440,24 +431,21 @@ impl fmt::Debug for ProviderResponseMetadata { } } -/// Provider 已筛选、可由协议 adapter 再次按白名单表达的响应头。 +/// Provider 已筛选、可由协议 adapter 尝试表达的不透明响应头。 #[derive(Clone, PartialEq, Eq)] pub struct ProviderResponseHeader { name: String, - value: SafeUpstreamValue, + value: Bytes, } impl ProviderResponseHeader { - /// 创建规范化的小写响应头；非 ASCII token 名称直接拒绝。 + /// 保存 Provider 交付的原始名称和值；HTTP 可表达性由最终 adapter 判断。 #[must_use] - pub fn new(name: impl Into, value: SafeUpstreamValue) -> Option { - let name = name.into().trim().to_ascii_lowercase(); - let valid = !name.is_empty() - && name.len() <= 64 - && name - .bytes() - .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-'); - valid.then_some(Self { name, value }) + pub fn new(name: impl Into, value: Bytes) -> Self { + Self { + name: name.into(), + value, + } } #[must_use] @@ -466,7 +454,7 @@ impl ProviderResponseHeader { } #[must_use] - pub const fn value(&self) -> &SafeUpstreamValue { + pub const fn value(&self) -> &Bytes { &self.value } } @@ -475,7 +463,7 @@ impl fmt::Debug for ProviderResponseHeader { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter .debug_struct("ProviderResponseHeader") - .field("name", &self.name) + .field("name", &"") .field("value", &"") .finish() } @@ -488,7 +476,8 @@ pub struct ProviderResponseObservation { http_version: Option, websocket_pool: Option, status_code: Option, - request_id: Option, + request_id: Option, + service_tier: Option, timings: ProviderResponseTimings, client_headers: Vec, provider_metadata: Option, @@ -503,6 +492,7 @@ impl ProviderResponseObservation { websocket_pool: None, status_code: None, request_id: None, + service_tier: None, timings: ProviderResponseTimings::default(), client_headers: Vec::new(), provider_metadata: None, @@ -528,11 +518,36 @@ impl ProviderResponseObservation { } #[must_use] - pub fn with_request_id(mut self, request_id: SafeUpstreamValue) -> Self { + pub fn with_request_id(mut self, request_id: OpaqueUpstreamValue) -> Self { self.request_id = Some(request_id); self } + /// 附加 Provider 在响应中确认的服务档位。 + /// + /// # Errors + /// + /// 档位为空、超过 64 字节或包含控制字符时返回错误。 + pub fn try_with_service_tier( + mut self, + service_tier: impl Into, + ) -> Result { + let service_tier = service_tier.into(); + validate_text(&service_tier, 64, false, None)?; + self.service_tier = Some(service_tier); + Ok(self) + } + + /// 附加合法服务档位；不适合持久化的上游观测值会被忽略。 + #[must_use] + pub fn with_service_tier_if_valid(mut self, service_tier: impl Into) -> Self { + let service_tier = service_tier.into(); + if validate_text(&service_tier, 64, false, None).is_ok() { + self.service_tier = Some(service_tier); + } + self + } + #[must_use] pub const fn with_timings(mut self, timings: ProviderResponseTimings) -> Self { self.timings = timings; @@ -573,10 +588,15 @@ impl ProviderResponseObservation { } #[must_use] - pub const fn request_id(&self) -> Option<&SafeUpstreamValue> { + pub const fn request_id(&self) -> Option<&OpaqueUpstreamValue> { self.request_id.as_ref() } + #[must_use] + pub fn service_tier(&self) -> Option<&str> { + self.service_tier.as_deref() + } + #[must_use] pub const fn timings(&self) -> ProviderResponseTimings { self.timings diff --git a/backend/crates/gateway-core/src/routing/mod.rs b/backend/crates/gateway-core/src/routing/mod.rs index 4eeeb16f1..683ff1e3c 100644 --- a/backend/crates/gateway-core/src/routing/mod.rs +++ b/backend/crates/gateway-core/src/routing/mod.rs @@ -89,6 +89,18 @@ impl UpstreamModelId { Ok(Self(value)) } + /// 从未命中显式映射的客户端模型名构造同名上游模型。 + /// + /// 目录和配置仍使用 [`Self::new`] 的内部标识约束；客户端 wire 只要求非空， + /// 具体模型字符串是否可用由绑定 Provider 决定。 + pub fn from_client_wire(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() { + return Err(IdentifierError::Empty); + } + Ok(Self(value)) + } + #[must_use] pub fn as_str(&self) -> &str { &self.0 diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index 5eed44a7e..1e2ad8bc1 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -488,15 +488,11 @@ impl RuntimeSnapshot { #[must_use] pub fn mapped_model(&self, requested: &str) -> String { - let original = requested.trim(); + let original = requested; let mut current = original.to_owned(); let mut seen = BTreeSet::new(); for _ in 0..20 { - let Some(target) = self - .model_mappings - .get(¤t) - .map(|target| target.trim()) - else { + let Some(target) = self.model_mappings.get(¤t).map(String::as_str) else { return current; }; if !seen.insert(current.clone()) || seen.contains(target) { @@ -529,8 +525,14 @@ impl RuntimeSnapshot { { continue; } - let upstream_model = UpstreamModelId::new(self.mapped_model(public_model.as_str())) - .map_err(|_| RoutingError::InvalidIdentifier)?; + let requested_model = public_model.as_str(); + let mapped_model = self.mapped_model(requested_model); + let upstream_model = if self.model_mappings.contains_key(requested_model) { + UpstreamModelId::new(mapped_model) + } else { + UpstreamModelId::from_client_wire(mapped_model) + } + .map_err(|_| RoutingError::InvalidIdentifier)?; let emulated_features = match self .provider_models .get(provider) diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index 85ed093c9..4e6bdd496 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -6,6 +6,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; use async_trait::async_trait; +use bytes::Bytes; use futures::StreamExt; use futures::executor::block_on; @@ -28,8 +29,8 @@ use gateway_core::engine::{ ProviderAttemptOutcome, RecoveryReport, UpstreamSendState, }; use gateway_core::error::{ - ContinuationFailure, GatewayErrorKind, ProviderError, ProviderErrorKind, SafeUpstreamValue, - StoreError, + ClientVisibleUpstreamResponse, ContinuationFailure, GatewayErrorKind, OpaqueUpstreamValue, + ProviderError, ProviderErrorKind, StoreError, StoreErrorKind, }; use gateway_core::event::{ ContentItem, ContentKind, GatewayEvent, ProtocolWireEvent, ProviderEvent, @@ -65,6 +66,7 @@ struct FinalState { upstream_transport: Option, http_version: Option, websocket_pool: Option, + service_tier: Option, upstream_request_id: Option, upstream_status_code: Option, transport_decision_wait_ms: Option, @@ -97,11 +99,37 @@ struct StoreState { #[derive(Default)] struct FakeStore { state: Mutex, + failures: BTreeSet, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum StoreWriteFailure { + Create, + Commit, + Finalize, +} + +impl FakeStore { + fn failing(failure: StoreWriteFailure) -> Self { + Self { + state: Mutex::new(StoreState::default()), + failures: BTreeSet::from([failure]), + } + } + + fn fails(&self, failure: StoreWriteFailure) -> Result<(), StoreError> { + if self.failures.contains(&failure) { + Err(StoreError::new(StoreErrorKind::Unavailable)) + } else { + Ok(()) + } + } } #[async_trait] impl ExecutionStore for FakeStore { async fn create_model_request(&self, _request: NewModelRequest) -> Result<(), StoreError> { + self.fails(StoreWriteFailure::Create)?; self.state.lock().expect("store lock").created += 1; Ok(()) } @@ -134,6 +162,7 @@ impl ExecutionStore for FakeStore { _committed_at: SystemTime, client_status_code: Option, ) -> Result<(), StoreError> { + self.fails(StoreWriteFailure::Commit)?; let mut state = self.state.lock().expect("store lock"); state.commits += 1; state.committed_statuses.push(client_status_code); @@ -172,6 +201,7 @@ impl ExecutionStore for FakeStore { &self, finalization: ModelRequestFinalization, ) -> Result<(), StoreError> { + self.fails(StoreWriteFailure::Finalize)?; self.state .lock() .expect("store lock") @@ -194,6 +224,7 @@ impl ExecutionStore for FakeStore { upstream_transport: finalization.upstream_transport, http_version: finalization.http_version, websocket_pool: finalization.websocket_pool, + service_tier: finalization.service_tier, upstream_request_id: finalization.upstream_request_id, upstream_status_code: finalization.upstream_status_code, transport_decision_wait_ms: finalization.timings.transport_decision_wait_ms, @@ -291,15 +322,11 @@ impl Provider for ScriptedProvider { }, UpstreamTransport::new("http_sse").expect("transport"), ) - .with_upstream_request_id( - SafeUpstreamValue::new("upstream-request").expect("safe request id"), - ); + .with_upstream_request_id(OpaqueUpstreamValue::new("upstream-request")); Ok(ProviderStream::new( metadata, Box::pin(futures::stream::iter( - items - .into_iter() - .map(|event| event.map(ProviderEvent::from)), + items.into_iter().map(canonical_provider_event), )), (), )) @@ -318,12 +345,8 @@ impl Provider for ScriptedProvider { Ok(ProviderStream::new( metadata, Box::pin( - futures::stream::iter( - items - .into_iter() - .map(|event| event.map(ProviderEvent::from)), - ) - .chain(futures::stream::pending()), + futures::stream::iter(items.into_iter().map(canonical_provider_event)) + .chain(futures::stream::pending()), ), (), )) @@ -399,6 +422,12 @@ fn complete_stream(total_tokens: Option) -> Vec, +) -> Result { + event.map(ProviderEvent::canonical) +} + fn atomic_response_failed(response_id: &str, marker: &str) -> ProviderError { let started = ProviderEvent::canonical_with_wire( vec![GatewayEvent::Started(ResponseMeta::new( @@ -547,6 +576,17 @@ fn coordinator( Arc, ) { let store = Arc::new(FakeStore::default()); + coordinator_with_store(scripts, store) +} + +fn coordinator_with_store( + scripts: Vec