fix(proxy): restore transparent wire semantics

Preserve raw request and response data across HTTP, SSE, and WebSocket adapters, including opaque headers and upstream error envelopes. Remove private payload and catalog gates, keep account switching behind the pre-delivery barrier, and make observation failures unable to replace client-visible protocol data.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-29 00:10:35 +08:00
co-authored by Codex
parent 5a81856510
commit 0c90ed30fb
100 changed files with 5230 additions and 1465 deletions
+1
View File
@@ -1234,6 +1234,7 @@ dependencies = [
"serde",
"serde_json",
"thiserror",
"tracing",
"uuid",
]
+1 -1
View File
@@ -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]
@@ -330,6 +330,7 @@ pub struct LatencyPercentiles {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderBillingInput {
pub upstream_model_id: String,
pub service_tier: Option<String>,
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub cached_tokens: Option<u64>,
@@ -345,6 +346,7 @@ pub struct UsageCalculatedBillingFact {
pub bucket_start: DateTime<Utc>,
pub provider_kind: String,
pub upstream_model_id: String,
pub service_tier: Option<String>,
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub cached_tokens: Option<u64>,
@@ -574,6 +576,7 @@ pub struct UsageRecord {
pub upstream_transport: Option<String>,
pub http_version: Option<String>,
pub websocket_pool: Option<String>,
pub service_tier: Option<String>,
/// Provider 已筛选的专有观测 JSON;管理领域保持不透明。
pub provider_metadata_json: Option<String>,
pub attempt_count: u32,
@@ -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,
@@ -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 {
@@ -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,
@@ -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,
@@ -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)
}
@@ -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",
@@ -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(),
@@ -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"]);
@@ -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"]);
+1 -1
View File
@@ -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"] }
@@ -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,
+29 -1
View File
@@ -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(
@@ -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,
},
@@ -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<dyn ExecutionSession>) -> 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<dyn ExecutionSession>) -> 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<Vec<u8>, 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<u8>, 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 {
@@ -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<String> {
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"
)
}
@@ -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<String>,
memgen_request: Option<String>,
subagent: Option<String>,
passthrough_headers: Vec<Value>,
}
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<DecodedResponsesRequest, RequestDecodeError> {
// `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<String, Value>, 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<String> {
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<Value> {
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::<Vec<_>>();
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<String, Value>,
key: &str,
@@ -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<String>,
wire_terminal: Option<Value>,
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<Vec<Bytes>, ResponseEncodeError> {
self.observe_identity(event)?;
pub fn push_sse(&mut self, event: &ProviderEvent) -> Vec<Bytes> {
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<Vec<String>, ResponseEncodeError> {
self.observe_identity(event)?;
pub fn push_websocket(&mut self, event: &ProviderEvent) -> Vec<String> {
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<Value, ResponseEncodeError> {
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));
@@ -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<ProviderSessionState>,
) -> 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::<Map<String, Value>>();
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,
@@ -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<ApiState> {
Router::new()
@@ -26,5 +23,7 @@ pub(crate) fn router() -> Router<ApiState> {
.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())
}
@@ -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) => {
@@ -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,
@@ -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);
@@ -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<String>,
client_ip: Option<IpAddr>,
user_agent: Option<String>,
endpoint: String,
@@ -125,11 +130,14 @@ impl ExecutionService for ContextCaptureExecution {
}
fn public_models(&self, _: &AuthenticatedClient) -> Vec<PublicModelId> {
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<StartedExecution, GatewayError>> {
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<_>>(),
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<_>>(),
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())
);
}
@@ -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::<Vec<_>>()
};
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::<Vec<_>>();
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::<Vec<_>>();
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::<Vec<_>>();
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()
@@ -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<AtomicFailureTrace>,
response_headers: Vec<ProviderResponseHeader>,
}
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<u16>) -> BoxFuture<'_, Result<(), EngineError>> {
@@ -231,6 +233,7 @@ impl ExecutionSession for AtomicFailureSession {
struct AtomicFailureExecution {
client: AuthenticatedClient,
trace: Arc<AtomicFailureTrace>,
response_headers: Vec<ProviderResponseHeader>,
}
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<Body> {
@@ -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]
+2 -1
View File
@@ -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]
@@ -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<u16>,
@@ -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(
&current,
@@ -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<String> {
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<bool> {
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<T>(
write: impl Future<Output = Result<T, StoreError>>,
) -> Option<T> {
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)
}
}
@@ -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<AccountSelection<'a>> {
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::<Vec<_>>();
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<AccountSchedulingBlocker> {
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 {
@@ -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<RoutingContext, GatewayError> {
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 后台写入超时,已丢弃本次亲和记录");
},
}
}
@@ -615,6 +615,7 @@ pub struct ModelRequestFinalization {
pub upstream_transport: Option<String>,
pub http_version: Option<String>,
pub websocket_pool: Option<String>,
pub service_tier: Option<String>,
/// Provider 已筛选的专有观测 JSON;Core 不解释字段。
pub provider_metadata_json: Option<String>,
pub error: Option<GatewayError>,
@@ -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<SafeUpstreamValue>,
upstream_request_id: Option<OpaqueUpstreamValue>,
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()
}
+195 -65
View File
@@ -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<String>) -> Result<Self, IdentifierError> {
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<String>) -> 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(<redacted-from-Debug>)")
formatter.write_str("OpaqueUpstreamValue(<redacted-from-Debug>)")
}
}
@@ -213,31 +210,23 @@ impl fmt::Debug for SafeUpstreamValue {
#[derive(Clone, PartialEq, Eq)]
pub struct ClientVisibleUpstreamError {
message: String,
code: Option<SafeUpstreamValue>,
error_type: Option<SafeUpstreamValue>,
code: Option<OpaqueUpstreamValue>,
error_type: Option<OpaqueUpstreamValue>,
}
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<String>,
code: Option<String>,
error_type: Option<String>,
) -> Result<Self, IdentifierError> {
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<Vec<u8>>,
headers: Vec<ProviderResponseHeader>,
body: Bytes,
}
impl ClientVisibleUpstreamResponse {
/// 保存 transport 实际收到的状态码、Content-Type 原值和正文。
#[must_use]
pub fn new(status: u16, content_type: Option<Vec<u8>>, body: Bytes) -> Self {
Self {
status,
content_type,
headers: Vec::new(),
body,
}
}
/// 附加 Provider 已剔除账号绑定与 framing 信息的普通响应头。
#[must_use]
pub fn with_headers(mut self, headers: Vec<ProviderResponseHeader>) -> 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(|_| "<present>"),
)
.field("header_count", &self.headers.len())
.field("body", &"<redacted>")
.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<u16>,
upstream_code: Option<SafeUpstreamValue>,
upstream_request_id: Option<SafeUpstreamValue>,
upstream_response_id: Option<SafeUpstreamValue>,
upstream_values: Option<Box<ProviderErrorUpstreamValues>>,
retry_after: Option<Duration>,
continuation_failure: Option<ContinuationFailure>,
replay_safe: bool,
pre_delivery_retry: bool,
credential_recovery_required: bool,
retry_same_account: bool,
sensitive_context_redacted: bool,
client_visible_upstream_error: Option<Box<ClientVisibleUpstreamError>>,
client_visible_upstream_response: Option<Box<ClientVisibleUpstreamResponse>>,
atomic_client_events: Option<Box<AtomicClientEvents>>,
}
#[derive(Clone, Default)]
struct ProviderErrorUpstreamValues {
code: Option<OpaqueUpstreamValue>,
request_id: Option<OpaqueUpstreamValue>,
response_id: Option<OpaqueUpstreamValue>,
}
#[derive(Clone)]
struct AtomicClientEvents(Vec<ProviderEvent>);
@@ -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(|_| "<classified-safe>"),
&self.upstream_code().map(|_| "<classified-safe>"),
)
.field(
"upstream_request_id",
&self
.upstream_request_id
.as_ref()
.map(|_| "<classified-safe>"),
&self.upstream_request_id().map(|_| "<classified-safe>"),
)
.field(
"upstream_response_id",
&self
.upstream_response_id
.as_ref()
.map(|_| "<classified-safe>"),
&self.upstream_response_id().map(|_| "<classified-safe>"),
)
.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(|_| "<present>"),
)
.field(
"client_visible_upstream_response",
&self
.client_visible_upstream_response
.as_ref()
.map(|_| "<present>"),
)
.field(
"atomic_client_events",
&self
+48 -28
View File
@@ -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<String>,
event_type: Option<String>,
@@ -212,7 +212,7 @@ impl ProtocolWireEvent {
///
/// # Errors
///
/// 协议名、显式事件名或 SSE ID 不满足 wire 安全约束时返回错误。
/// 协议名不满足内部路由标识约束时返回错误。SSE 元数据保持不透明。
pub fn json_with_sse_metadata(
protocol: impl Into<String>,
event_type: Option<String>,
@@ -222,15 +222,6 @@ impl ProtocolWireEvent {
) -> Result<Self, IdentifierError> {
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<String>, value: SafeUpstreamValue) -> Option<Self> {
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<String>, 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", &"<redacted>")
.field("value", &"<redacted>")
.finish()
}
@@ -488,7 +476,8 @@ pub struct ProviderResponseObservation {
http_version: Option<UpstreamHttpVersion>,
websocket_pool: Option<WebSocketPoolKind>,
status_code: Option<u16>,
request_id: Option<SafeUpstreamValue>,
request_id: Option<OpaqueUpstreamValue>,
service_tier: Option<String>,
timings: ProviderResponseTimings,
client_headers: Vec<ProviderResponseHeader>,
provider_metadata: Option<ProviderResponseMetadata>,
@@ -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<String>,
) -> Result<Self, IdentifierError> {
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<String>) -> 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
@@ -89,6 +89,18 @@ impl UpstreamModelId {
Ok(Self(value))
}
/// 从未命中显式映射的客户端模型名构造同名上游模型。
///
/// 目录和配置仍使用 [`Self::new`] 的内部标识约束;客户端 wire 只要求非空,
/// 具体模型字符串是否可用由绑定 Provider 决定。
pub fn from_client_wire(value: impl Into<String>) -> Result<Self, IdentifierError> {
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
@@ -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(&current)
.map(|target| target.trim())
else {
let Some(target) = self.model_mappings.get(&current).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)
@@ -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<String>,
http_version: Option<String>,
websocket_pool: Option<String>,
service_tier: Option<String>,
upstream_request_id: Option<String>,
upstream_status_code: Option<u16>,
transport_decision_wait_ms: Option<u64>,
@@ -97,11 +99,37 @@ struct StoreState {
#[derive(Default)]
struct FakeStore {
state: Mutex<StoreState>,
failures: BTreeSet<StoreWriteFailure>,
}
#[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<u16>,
) -> 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<u64>) -> Vec<Result<GatewayEvent, Provid
events
}
fn canonical_provider_event(
event: Result<GatewayEvent, ProviderError>,
) -> Result<ProviderEvent, ProviderError> {
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<ScriptedProvider>,
) {
let store = Arc::new(FakeStore::default());
coordinator_with_store(scripts, store)
}
fn coordinator_with_store(
scripts: Vec<Script>,
store: Arc<FakeStore>,
) -> (
AttemptCoordinator<FakeStore>,
Arc<FakeStore>,
Arc<ScriptedProvider>,
) {
let provider = Arc::new(ScriptedProvider::new(scripts));
let mut registry = ProviderRegistry::builder();
registry
@@ -650,6 +690,78 @@ fn success_updates_one_model_request_and_persists_usage() {
);
}
#[test]
fn observation_store_failures_never_block_successful_delivery() {
for failure in [
StoreWriteFailure::Create,
StoreWriteFailure::Commit,
StoreWriteFailure::Finalize,
] {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation);
let store = Arc::new(FakeStore::failing(failure));
let (coordinator, _, provider) = coordinator_with_store(
vec![Script::Stream {
account_id: "acct_one",
items: complete_stream(None),
}],
store,
);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let events = block_on(session.collect_uncommitted())
.unwrap_or_else(|error| panic!("{failure:?} blocked collection: {error}"));
assert_eq!(events.len(), 2);
block_on(session.commit_downstream(Some(200)))
.unwrap_or_else(|error| panic!("{failure:?} blocked commit: {error}"));
assert!(session.is_finalized());
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1);
}
}
#[test]
fn terminal_provider_error_preserves_request_local_upstream_response() {
let operation = generate_operation();
let route_plan = plan(&operation);
let raw_body = Bytes::from_static(b"{\"upstream\":\"unchanged\"}\x00");
let provider_error =
ProviderError::new(ProviderErrorKind::InvalidRequest, UpstreamSendState::Sent)
.with_client_visible_upstream_response(ClientVisibleUpstreamResponse::new(
422,
Some(b"application/problem+json".to_vec()),
raw_body.clone(),
));
let (coordinator, _, _) = coordinator(vec![Script::Error(provider_error)]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let error = block_on(session.collect_uncommitted()).expect_err("provider rejection");
let EngineError::Provider(error) = error else {
panic!("expected provider error");
};
let response = error
.client_visible_upstream_response()
.expect("request-local upstream response");
assert_eq!(response.status(), 422);
assert_eq!(response.body(), &raw_body);
}
#[test]
fn success_preserves_opaque_upstream_response_id_and_builds_native_pin() {
let response_id = format!("resp_{}\0opaque", "x".repeat(4_096));
@@ -698,6 +810,50 @@ fn success_preserves_opaque_upstream_response_id_and_builds_native_pin() {
);
}
#[test]
fn canonical_identity_change_and_empty_terminal_id_are_observational() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Stream {
account_id: "acct_one",
items: vec![
Ok(GatewayEvent::Started(ResponseMeta::new(
"response-created",
"gpt-5",
))),
Ok(GatewayEvent::Completed(ResponseMeta::new("", "gpt-5"))),
],
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("identity observation cannot fail delivery");
let provider_state = ProviderSessionState::new("openai", Map::new()).expect("provider state");
let pin = session
.native_continuation_pin(&provider_state)
.expect("empty opaque response ID remains recordable");
assert_eq!(pin.previous_response_id().as_str(), "");
assert_eq!(pin.upstream_response_id().as_str(), "");
block_on(session.commit_downstream(Some(200))).expect("commit response");
let state = store.state.lock().expect("store lock");
assert_eq!(
state.finalizations[0].client_response_id.as_deref(),
Some("")
);
assert_eq!(
state.finalizations[0].upstream_response_id.as_deref(),
Some("")
);
}
fn finalized_image_request(image_output_tokens: Option<u64>) -> FinalState {
let operation = image_generate_operation();
let route_plan = plan(&operation);
@@ -780,7 +936,9 @@ fn response_observation_is_persisted_but_never_delivered() {
.with_http_version(UpstreamHttpVersion::Http2)
.with_websocket_pool(WebSocketPoolKind::New)
.with_status_code(200)
.with_request_id(SafeUpstreamValue::new("upstream-observed").expect("request id"))
.with_request_id(OpaqueUpstreamValue::new("upstream-observed"))
.try_with_service_tier("priority")
.expect("service tier")
.with_timings(ProviderResponseTimings {
transport_decision_wait_ms: Some(7),
connect_ms: Some(11),
@@ -792,7 +950,7 @@ fn response_observation_is_persisted_but_never_delivered() {
items.extend(
complete_stream(Some(21))
.into_iter()
.map(|event| event.map(ProviderEvent::canonical)),
.map(canonical_provider_event),
);
let (coordinator, store, _) = coordinator(vec![Script::ObservedStream {
account_id: "acct_observed",
@@ -817,6 +975,7 @@ fn response_observation_is_persisted_but_never_delivered() {
assert_eq!(finalization.upstream_transport.as_deref(), Some("http_sse"));
assert_eq!(finalization.http_version.as_deref(), Some("HTTP/2"));
assert_eq!(finalization.websocket_pool.as_deref(), Some("new"));
assert_eq!(finalization.service_tier.as_deref(), Some("priority"));
assert_eq!(
finalization.upstream_request_id.as_deref(),
Some("upstream-observed")
@@ -828,6 +987,48 @@ fn response_observation_is_persisted_but_never_delivered() {
assert_eq!(finalization.first_event_ms, Some(17));
}
#[test]
fn mismatched_response_observation_should_not_interrupt_client_events() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation);
let first = ProviderEvent::observation(ProviderResponseObservation::new(
UpstreamTransport::new("websocket").expect("first transport"),
));
let mismatched = ProviderEvent::observation(ProviderResponseObservation::new(
UpstreamTransport::new("http_sse").expect("mismatched transport"),
));
let mut items = vec![Ok(first), Ok(mismatched)];
items.extend(
complete_stream(None)
.into_iter()
.map(canonical_provider_event),
);
let (coordinator, store, _) = coordinator(vec![Script::ObservedStream {
account_id: "acct_observation_mismatch",
items,
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let events = block_on(session.collect_uncommitted()).expect("client events survive");
block_on(session.commit_downstream(Some(200))).expect("commit response");
assert_eq!(events.len(), 2);
assert_eq!(
store.state.lock().expect("store lock").finalizations[0]
.upstream_transport
.as_deref(),
Some("websocket")
);
}
#[test]
fn response_observation_should_persist_provider_metadata_and_output_timings() {
let operation = operation(RetrySafety::Idempotent);
@@ -852,7 +1053,7 @@ fn response_observation_should_persist_provider_metadata_and_output_timings() {
items.extend(
complete_stream(Some(21))
.into_iter()
.map(|event| event.map(ProviderEvent::canonical)),
.map(canonical_provider_event),
);
let (coordinator, store, _) = coordinator(vec![Script::ObservedStream {
account_id: "acct_observed",
@@ -992,7 +1193,7 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
)
.with_http_version(UpstreamHttpVersion::Http11)
.with_status_code(503)
.with_request_id(SafeUpstreamValue::new("discarded-request").expect("request id"))
.with_request_id(OpaqueUpstreamValue::new("discarded-request"))
.with_timings(ProviderResponseTimings {
transport_decision_wait_ms: Some(7),
connect_ms: Some(11),
@@ -1006,12 +1207,12 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
)
.with_http_version(UpstreamHttpVersion::Http2)
.with_status_code(200)
.with_request_id(SafeUpstreamValue::new("winning-request").expect("request id"));
.with_request_id(OpaqueUpstreamValue::new("winning-request"));
let mut second_items = vec![Ok(ProviderEvent::observation(second_observation))];
second_items.extend(
complete_stream(None)
.into_iter()
.map(|event| event.map(ProviderEvent::canonical)),
.map(canonical_provider_event),
);
let (coordinator, store, _) = coordinator(vec![
Script::ObservedStream {
@@ -1821,9 +2022,7 @@ fn rate_limited_account_exhaustion_survives_a_later_empty_selection() {
UpstreamSendState::Sent,
)
.with_status(429)
.with_upstream_code(
SafeUpstreamValue::new("rate_limit_exceeded").expect("provider code"),
)
.with_upstream_code(OpaqueUpstreamValue::new("rate_limit_exceeded"))
.with_retry_after(Duration::from_secs(30))
.with_replay_safe())],
},
@@ -1851,7 +2050,7 @@ fn rate_limited_account_exhaustion_survives_a_later_empty_selection() {
assert_eq!(error.kind(), ProviderErrorKind::RateLimited);
assert_eq!(error.upstream_status(), Some(429));
assert_eq!(
error.upstream_code().map(SafeUpstreamValue::as_str),
error.upstream_code().map(OpaqueUpstreamValue::as_str),
Some("rate_limit_exceeded")
);
assert_eq!(error.retry_after(), Some(Duration::from_secs(30)));
@@ -2312,6 +2511,100 @@ fn ambiguous_send_state_stops_retry() {
assert!(finalization.latency_ms.is_some());
}
#[test]
fn pre_delivery_transport_failure_rotates_account_for_non_idempotent_generation() {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::Ambiguous,
)
.with_pre_delivery_retry())],
},
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("second account succeeds before delivery");
block_on(session.commit_downstream(Some(200))).expect("commit winning response");
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 2);
assert!(
contexts[1]
.excluded_accounts()
.contains(&ProviderAccountId::new("acct_first").expect("account id"))
);
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn pre_delivery_retry_marker_is_ignored_after_downstream_commit() {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: vec![
Ok(GatewayEvent::Started(ResponseMeta::new(
"response-visible",
"gpt-5",
))),
Err(
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
.with_pre_delivery_retry(),
),
],
},
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let first = block_on(session.next_event())
.expect("first event")
.expect("visible event");
assert_eq!(
first.commit_requirement(),
CommitRequirement::CommitBeforeDelivery
);
block_on(session.commit_downstream(Some(200))).expect("commit first event");
let error = block_on(session.next_event()).expect_err("committed output cannot be replayed");
assert!(matches!(error, EngineError::Provider(_)));
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1);
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 0);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Incomplete);
assert!(state.finalizations[0].committed);
}
#[test]
fn structural_event_before_replay_safe_failure_should_switch_account_before_commit() {
let operation = operation(RetrySafety::Idempotent);
@@ -6,10 +6,10 @@ use serde_json::{Map, Value};
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountAvailabilityPolicy, AccountCandidate,
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSelectionContext,
AccountSelectionPolicy, AccountSelector, CredentialCasUpdate, CredentialRevision,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountUpdate, RotationStrategy,
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker,
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, CredentialCasUpdate,
CredentialRevision, OpaqueProviderData, PlaintextCredential, PreferredAccountSelection,
ProviderAccount, ProviderAccountId, ProviderAccountUpdate, RotationStrategy,
};
use gateway_core::routing::ProviderKind;
@@ -119,13 +119,13 @@ fn diagnostic_selection_bypasses_all_local_account_availability() {
assert_eq!(
AccountSelector
.select(std::slice::from_ref(&exhausted), &context)
.map(|candidate| candidate.account.id()),
.map(|selection| selection.candidate().account.id()),
Some(exhausted.account.id())
);
assert_eq!(
AccountSelector
.select(std::slice::from_ref(&disabled), &context)
.map(|candidate| candidate.account.id()),
.map(|selection| selection.candidate().account.id()),
Some(disabled.account.id())
);
}
@@ -169,7 +169,7 @@ fn smart_selector_should_prefer_lower_inflight_count() {
.select(&candidates, &context(RotationStrategy::Smart))
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_idle");
assert_eq!(selected.candidate().account.id().as_str(), "acct_idle");
}
#[test]
@@ -186,15 +186,13 @@ fn smart_selector_should_balance_signals_instead_of_using_lexicographic_load() {
.select(&candidates, &context(RotationStrategy::Smart))
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_healthy");
assert_eq!(selected.candidate().account.id().as_str(), "acct_healthy");
}
#[test]
fn smart_selector_should_keep_healthy_preferred_account_at_escape_boundary() {
fn smart_selector_should_report_a_schedulable_preferred_account_hit() {
let fallback = candidate("acct_fallback", 0, Some(100));
let mut preferred = candidate("acct_preferred", 0, Some(1));
preferred.signals.failure_rate_basis_points = Some(5_000);
preferred.signals.first_output_latency_ms = Some(15_000);
let preferred = candidate("acct_preferred", 0, Some(1));
let preferred_id = preferred.account.id().clone();
let mut selection = context(RotationStrategy::Smart);
selection.preferred_account = Some(preferred_id);
@@ -204,16 +202,22 @@ fn smart_selector_should_keep_healthy_preferred_account_at_escape_boundary() {
.select(&candidates, &selection)
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_preferred");
assert_eq!(
(
selected.candidate().account.id().as_str(),
selected.preferred(),
),
("acct_preferred", PreferredAccountSelection::Hit)
);
}
#[test]
fn smart_selector_should_escape_preferred_account_after_failure_threshold() {
fn smart_selector_should_keep_preferred_account_despite_failure_signals() {
let mut fallback = candidate("acct_fallback", 0, Some(100));
fallback.signals.failure_rate_basis_points = Some(0);
fallback.signals.first_output_latency_ms = Some(100);
let mut preferred = candidate("acct_preferred", 0, Some(100));
preferred.signals.failure_rate_basis_points = Some(5_001);
preferred.signals.failure_rate_basis_points = Some(10_000);
preferred.signals.first_output_latency_ms = Some(100);
let preferred_id = preferred.account.id().clone();
let mut selection = context(RotationStrategy::Smart);
@@ -222,19 +226,19 @@ fn smart_selector_should_escape_preferred_account_after_failure_threshold() {
let selected = AccountSelector
.select(&candidates, &selection)
.expect("fallback candidate available");
.expect("preferred candidate available");
assert_eq!(selected.account.id().as_str(), "acct_fallback");
assert_eq!(selected.candidate().account.id().as_str(), "acct_preferred");
}
#[test]
fn smart_selector_should_escape_preferred_account_after_latency_threshold() {
fn smart_selector_should_keep_preferred_account_despite_latency_signals() {
let mut fallback = candidate("acct_fallback", 0, Some(100));
fallback.signals.failure_rate_basis_points = Some(0);
fallback.signals.first_output_latency_ms = Some(100);
let mut preferred = candidate("acct_preferred", 0, Some(100));
preferred.signals.failure_rate_basis_points = Some(0);
preferred.signals.first_output_latency_ms = Some(15_001);
preferred.signals.first_output_latency_ms = Some(120_000);
let preferred_id = preferred.account.id().clone();
let mut selection = context(RotationStrategy::Smart);
selection.preferred_account = Some(preferred_id);
@@ -242,9 +246,9 @@ fn smart_selector_should_escape_preferred_account_after_latency_threshold() {
let selected = AccountSelector
.select(&candidates, &selection)
.expect("fallback candidate available");
.expect("preferred candidate available");
assert_eq!(selected.account.id().as_str(), "acct_fallback");
assert_eq!(selected.candidate().account.id().as_str(), "acct_preferred");
}
#[test]
@@ -260,7 +264,16 @@ fn smart_selector_should_escape_preferred_account_at_concurrency_limit() {
.select(&candidates, &selection)
.expect("fallback candidate available");
assert_eq!(selected.account.id().as_str(), "acct_fallback");
assert_eq!(
(
selected.candidate().account.id().as_str(),
selected.preferred(),
),
(
"acct_fallback",
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::ConcurrencyLimit),
)
);
}
#[test]
@@ -346,7 +359,10 @@ fn smart_selector_should_use_provider_quota_rank_after_load_is_equal() {
.select(&candidates, &context(RotationStrategy::Smart))
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_high_quota");
assert_eq!(
selected.candidate().account.id().as_str(),
"acct_high_quota"
);
}
#[test]
@@ -362,7 +378,7 @@ fn quota_reset_selector_should_prefer_known_earliest_window() {
.select(&candidates, &context(RotationStrategy::QuotaResetPriority))
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_earlier");
assert_eq!(selected.candidate().account.id().as_str(), "acct_earlier");
}
#[test]
@@ -375,7 +391,7 @@ fn smart_selector_should_prefer_known_quota_over_unknown_after_load_is_equal() {
.select(&candidates, &context(RotationStrategy::Smart))
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_known");
assert_eq!(selected.candidate().account.id().as_str(), "acct_known");
}
#[test]
@@ -388,7 +404,7 @@ fn selector_should_reject_account_at_concurrency_limit() {
.select(&candidates, &context(RotationStrategy::Smart))
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_available");
assert_eq!(selected.candidate().account.id().as_str(), "acct_available");
}
#[test]
@@ -410,7 +426,7 @@ fn selector_should_enforce_request_interval_until_boundary() {
.select(&candidates, &context)
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_ready");
assert_eq!(selected.candidate().account.id().as_str(), "acct_ready");
}
#[test]
@@ -427,7 +443,7 @@ fn sticky_selector_should_fall_back_when_requested_account_is_excluded() {
.select(&candidates, &context)
.expect("fallback candidate available");
assert_eq!(selected.account.id().as_str(), "acct_fallback");
assert_eq!(selected.candidate().account.id().as_str(), "acct_fallback");
}
#[test]
@@ -447,11 +463,38 @@ fn round_robin_selector_should_be_stable_for_unsorted_candidates() {
.expect("second candidate available");
assert_eq!(
(first.account.id().as_str(), second.account.id().as_str()),
(
first.candidate().account.id().as_str(),
second.candidate().account.id().as_str(),
),
("acct_first", "acct_second")
);
}
#[test]
fn round_robin_selector_should_honor_a_schedulable_preferred_account_before_the_cursor() {
let candidates = vec![
candidate("acct_first", 0, None),
candidate("acct_second", 0, None),
];
let mut context = context(RotationStrategy::RoundRobin);
context.round_robin_cursor = 0;
context.preferred_account =
Some(ProviderAccountId::new("acct_second").expect("preferred account"));
let selected = AccountSelector
.select(&candidates, &context)
.expect("preferred account available");
assert_eq!(
(
selected.candidate().account.id().as_str(),
selected.preferred(),
),
("acct_second", PreferredAccountSelection::Hit)
);
}
#[test]
fn selector_should_honor_request_local_exclusion() {
let candidates = vec![
@@ -466,7 +509,7 @@ fn selector_should_honor_request_local_exclusion() {
.select(&candidates, &context)
.expect("second account available");
assert_eq!(selected.account.id().as_str(), "acct_second");
assert_eq!(selected.candidate().account.id().as_str(), "acct_second");
}
#[test]
@@ -482,7 +525,7 @@ fn sticky_selector_should_prefer_requested_account() {
.select(&candidates, &context)
.expect("sticky account available");
assert_eq!(selected.account.id().as_str(), "acct_second");
assert_eq!(selected.candidate().account.id().as_str(), "acct_second");
}
#[test]
@@ -497,7 +540,7 @@ fn sticky_selector_without_request_binding_should_reuse_latest_account() {
.select(&candidates, &context(RotationStrategy::Sticky))
.expect("sticky candidate available");
assert_eq!(selected.account.id().as_str(), "acct_latest");
assert_eq!(selected.candidate().account.id().as_str(), "acct_latest");
}
#[test]
@@ -512,7 +555,7 @@ fn round_robin_selector_should_use_frozen_cursor() {
.select(&candidates, &context)
.expect("candidate available");
assert_eq!(selected.account.id().as_str(), "acct_second");
assert_eq!(selected.candidate().account.id().as_str(), "acct_second");
}
#[test]
@@ -4,7 +4,7 @@ use std::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
},
time::{Duration, SystemTime},
time::{Duration, Instant, SystemTime},
};
use async_trait::async_trait;
@@ -20,8 +20,9 @@ use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::engine::execution::{
ClientApiKeyUsageSink, DefaultExecutionService, ExecutionService, ProviderCircuitDecision,
ProviderCircuitError, ProviderCircuitPort, provider_failure_affects_circuit,
ClientApiKeyUsageSink, ClientTransport, DefaultExecutionService, ExecutionRequestMetadata,
ExecutionService, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
StartExecution, provider_failure_affects_circuit,
};
use gateway_core::engine::probe::{AccountProbe, AccountProbeRequest};
use gateway_core::engine::provider::{
@@ -37,7 +38,7 @@ use gateway_core::operation::{GenerateRequest, Operation, OperationKind, Protoco
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, RuntimeSnapshot,
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RuntimeSnapshot,
UpstreamModelId,
};
use serde_json::json;
@@ -159,6 +160,111 @@ fn successful_authentication_should_record_client_key_usage() {
assert_eq!(usage.recorded(), vec!["key_usage_test".to_owned()]);
}
#[test]
fn circuit_store_failure_should_fail_open_during_request_start() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(FailingDecisionCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let started = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::new("gpt-start").expect("public model"),
operation: start_operation(),
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}))
.expect("recoverable circuit state must not reject the request");
assert!(!started.session.is_finalized());
}
#[test]
fn slow_circuit_store_should_time_out_and_fail_open_during_request_start() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(PendingDecisionCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let started_at = Instant::now();
let started = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::new("gpt-start").expect("public model"),
operation: start_operation(),
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}))
.expect("slow recoverable circuit state must not reject the request");
assert!(started_at.elapsed() < Duration::from_secs(2));
assert!(!started.session.is_finalized());
}
#[test]
fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let model = "gpt-future-not-in-catalog";
let started = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::from_client_wire(model).expect("client model"),
operation: start_operation_for_model(model),
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}))
.expect("stale model catalog must not reject the request");
assert!(!started.session.is_finalized());
}
#[derive(Default)]
struct RecordingClientApiKeyUsage {
ids: Mutex<Vec<String>>,
@@ -309,6 +415,56 @@ impl ProviderCircuitPort for UnusedCircuits {
}
}
struct FailingDecisionCircuits;
impl ProviderCircuitPort for FailingDecisionCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { Err(ProviderCircuitError) })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
}
struct PendingDecisionCircuits;
impl ProviderCircuitPort for PendingDecisionCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(futures::future::pending())
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
}
struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
@@ -371,6 +527,34 @@ fn client_snapshot() -> RuntimeSnapshot {
.expect("client snapshot")
}
fn start_snapshot() -> RuntimeSnapshot {
let provider = ProviderKind::new("openai").expect("provider kind");
let capabilities =
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
AccountSelectionPolicy::new(
RotationStrategy::Smart,
std::num::NonZeroU32::new(1).expect("concurrency"),
Duration::from_millis(1),
),
vec![provider.clone()],
vec![ProviderModel::new(
provider.clone(),
UpstreamModelId::new("gpt-start").expect("model ID"),
capabilities,
)],
vec![ClientPolicy::new(
ClientApiKeyId::new("key_start_test").expect("client API key ID"),
PlaintextClientApiKey::new("sk_start_test").expect("plaintext client API key"),
provider,
true,
RateLimits::unlimited(),
)],
)
.expect("start snapshot")
}
fn probe_operation() -> Operation {
let body = json!({
"model": "gpt-probe",
@@ -381,3 +565,18 @@ fn probe_operation() -> Operation {
.expect("OpenAI payload"),
))
}
fn start_operation() -> Operation {
start_operation_for_model("gpt-start")
}
fn start_operation_for_model(model: &str) -> Operation {
let body = json!({
"model": model,
"input": [{"type": "message", "role": "user", "content": "ping"}],
});
Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object("openai", body.as_object().expect("request object").clone())
.expect("OpenAI payload"),
))
}
+83 -11
View File
@@ -1,8 +1,10 @@
use bytes::Bytes;
use gateway_core::engine::UpstreamSendState;
use gateway_core::error::{
ClientVisibleUpstreamError, GatewayError, ProviderError, ProviderErrorKind, SafeUpstreamValue,
ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, GatewayError, OpaqueUpstreamValue,
ProviderError, ProviderErrorKind,
};
use gateway_core::event::{ProtocolWireEvent, ProviderEvent};
use gateway_core::event::{ProtocolWireEvent, ProviderEvent, ProviderResponseHeader};
use serde_json::json;
#[test]
@@ -17,13 +19,22 @@ fn provider_error_debug_should_not_expose_sensitive_context() {
#[test]
fn provider_error_debug_should_not_print_classified_upstream_values() {
let diagnostic = "request-visible-only-through-explicit-accessor";
let value = SafeUpstreamValue::new(diagnostic).expect("test diagnostic is valid");
let value = OpaqueUpstreamValue::new(diagnostic);
let error = ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::Sent)
.with_upstream_request_id(value);
assert!(!format!("{error:?}").contains(diagnostic));
}
#[test]
fn opaque_upstream_value_should_preserve_arbitrary_text_without_logging_it() {
let original = format!("\0{}\n", "x".repeat(9_000));
let value = OpaqueUpstreamValue::new(original.clone());
assert_eq!(value.as_str(), original);
assert!(!format!("{value:?}").contains(&original));
}
#[test]
fn provider_error_replay_proof_should_default_to_false() {
let error = ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent);
@@ -59,18 +70,62 @@ fn provider_error_atomic_client_events_should_be_take_only_and_debug_redacted()
assert!(!error.has_atomic_client_events());
}
#[test]
fn provider_error_clone_should_drop_request_local_upstream_response() {
let marker = Bytes::from_static(b"raw-client-response-only");
let error = ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
.with_client_visible_upstream_response(
ClientVisibleUpstreamResponse::new(
429,
Some(b"application/problem+json".to_vec()),
marker.clone(),
)
.with_headers(vec![ProviderResponseHeader::new(
"x-future-error",
Bytes::from_static(b"opaque-header-value"),
)]),
);
let cloned = error.clone();
assert_eq!(
error
.client_visible_upstream_response()
.map(ClientVisibleUpstreamResponse::body),
Some(&marker)
);
assert!(cloned.client_visible_upstream_response().is_none());
assert!(!format!("{error:?}").contains("raw-client-response-only"));
assert!(!format!("{error:?}").contains("opaque-header-value"));
}
#[test]
fn provider_error_clone_should_drop_atomic_client_events() {
let event = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({"type": "response.failed"}),
)
.expect("atomic wire"),
);
let error = ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
.with_atomic_client_events(vec![event]);
let cloned = error.clone();
assert!(error.has_atomic_client_events());
assert!(!cloned.has_atomic_client_events());
}
#[test]
fn gateway_error_should_keep_client_visible_upstream_fields_out_of_safe_diagnostics() {
let message = "Your Codex quota is exhausted";
let error = ProviderError::new(ProviderErrorKind::QuotaExhausted, UpstreamSendState::Sent)
.with_client_visible_upstream_error(
ClientVisibleUpstreamError::new(
message,
Some("quota_exhausted".to_owned()),
Some("rate_limit_error".to_owned()),
)
.expect("safe structured upstream error"),
);
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
message,
Some("quota_exhausted".to_owned()),
Some("rate_limit_error".to_owned()),
));
let gateway = GatewayError::from_provider(&error);
assert_eq!(
@@ -82,3 +137,20 @@ fn gateway_error_should_keep_client_visible_upstream_fields_out_of_safe_diagnost
assert_eq!(gateway.client_error_type(), Some("rate_limit_error"));
assert!(!format!("{gateway:?}").contains(message));
}
#[test]
fn client_visible_upstream_error_should_preserve_opaque_structured_fields() {
let message = format!("\0{}\n", "m".repeat(9_000));
let code = format!("\0{}", "c".repeat(300));
let error_type = String::new();
let detail = ClientVisibleUpstreamError::new(
message.clone(),
Some(code.clone()),
Some(error_type.clone()),
);
assert_eq!(detail.message(), message);
assert_eq!(detail.code(), Some(code.as_str()));
assert_eq!(detail.error_type(), Some(error_type.as_str()));
assert!(!format!("{detail:?}").contains(&message));
}
+42 -2
View File
@@ -1,8 +1,9 @@
use bytes::Bytes;
use gateway_core::engine::provider::UpstreamTransport;
use gateway_core::event::{
ContentItem, ContentKind, EventSequenceError, EventSequenceValidator, GatewayEvent,
ProtocolWireEvent, ProviderEvent, ProviderResponseObservation, ReasoningDelta, ResponseMeta,
TextDelta, ToolCallDelta,
ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation,
ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta,
};
use serde_json::json;
@@ -21,6 +22,28 @@ fn provider_event_keeps_large_wire_and_observation_payloads_behind_one_indirecti
assert!(!observation.has_client_event());
}
#[test]
fn provider_response_header_should_preserve_opaque_name_and_bytes_without_debug_disclosure() {
let name = format!("x-future-{}\0", "n".repeat(512));
let value = Bytes::from_static(b"\xffopaque-response-header-secret");
let header = ProviderResponseHeader::new(name.clone(), value.clone());
assert_eq!(header.name(), name);
assert_eq!(header.value(), &value);
let rendered = format!("{header:?}");
assert!(!rendered.contains(&name));
assert!(!rendered.contains("opaque-response-header-secret"));
}
#[test]
fn invalid_observed_service_tier_should_be_ignored() {
let observation =
ProviderResponseObservation::new(UpstreamTransport::new("http_sse").expect("transport"))
.with_service_tier_if_valid(format!("priority\0{}", "x".repeat(128)));
assert_eq!(observation.service_tier(), None);
}
#[test]
fn protocol_wire_event_should_preserve_sse_metadata_without_exposing_id_in_debug() {
let wire = ProtocolWireEvent::json_with_sse_metadata(
@@ -37,6 +60,23 @@ fn protocol_wire_event_should_preserve_sse_metadata_without_exposing_id_in_debug
assert!(!format!("{wire:?}").contains("upstream-event-id"));
}
#[test]
fn protocol_wire_event_should_preserve_opaque_sse_metadata() {
let event_type = format!("\0{}", "event".repeat(80));
let sse_id = "id\0with\r\ncontrols".to_owned();
let wire = ProtocolWireEvent::json_with_sse_metadata(
"openai",
Some(event_type.clone()),
json!({"type": "future.event"}),
Some(sse_id.clone()),
None,
)
.expect("internal protocol name is valid");
assert_eq!(wire.event_type(), Some(event_type.as_str()));
assert_eq!(wire.sse_id(), Some(sse_id.as_str()));
}
#[test]
fn validator_should_accept_started_content_delta_completed_sequence() {
let mut validator = EventSequenceValidator::new();
@@ -190,21 +190,36 @@ fn global_mapping_should_apply_to_every_provider() {
#[test]
fn unmapped_model_should_pass_through_unchanged() {
let snapshot = snapshot();
let plan = snapshot
.plan(
&PublicModelId::new("future-openai-model").expect("model"),
&operation(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
.expect("unknown model remains transparent");
for requested in [
"future-openai-model".to_owned(),
format!("future-{}", "x".repeat(512)),
"future\0model".to_owned(),
" future-openai-model ".to_owned(),
] {
let plan = snapshot
.plan(
&PublicModelId::from_client_wire(requested.clone()).expect("client model"),
&operation(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
.expect("unknown model remains transparent");
assert_eq!(
plan.candidates()[0].upstream_model().as_str(),
"future-openai-model"
);
assert_eq!(plan.candidates()[0].upstream_model().as_str(), requested);
}
}
#[test]
fn model_mapping_should_use_exact_client_keys() {
let snapshot = snapshot().with_model_mappings(BTreeMap::from([(
" exact-alias ".to_owned(),
"gpt-5.5".to_owned(),
)]));
assert_eq!(snapshot.mapped_model(" exact-alias "), "gpt-5.5");
assert_eq!(snapshot.mapped_model("exact-alias"), "exact-alias");
}
#[test]
+1 -1
View File
@@ -13,7 +13,7 @@ path = "tests/main.rs"
[dependencies]
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"
uuid = { version = "1.23.3", features = ["v4"] }
@@ -317,11 +317,7 @@ pub fn is_rate_limit_header_name(name: &str) -> bool {
/// 从内部 `codex.rate_limits` 事件中解析限流信息。
pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
if value
.get("type")
.and_then(Value::as_str)
.is_some_and(|event_type| event_type != "codex.rate_limits")
{
if value.get("type").and_then(Value::as_str) != Some("codex.rate_limits") {
return None;
}
@@ -18,10 +18,11 @@ pub struct SseEvent {
pub retry: Option<u64>,
}
/// 保留完整原始字节与旁路解析结果的单个 SSE 帧。
/// 保留原始字节与旁路解析结果的 SSE wire 片段。
///
/// 这用于需要原样转发的 transport;`events` 只供旁路观测,不能反过来重建
/// `raw`。
/// 通常一个片段就是完整 SSE 帧。单帧超过旁路观测预算时会拆成多个原始片段,
/// 此时 `events` 为空;透明 transport 必须按顺序交付 `raw`,不能从 `events`
/// 反向重建或否决原始数据。
#[derive(Clone, PartialEq, Eq)]
pub struct SseFrame {
raw: Vec<u8>,
@@ -62,7 +63,7 @@ impl fmt::Debug for SseFrame {
}
}
/// 单事件缓冲上限。
/// 单事件旁路解析缓冲上限。
pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024;
/// SSE 流结束标记帧。
@@ -75,6 +76,9 @@ pub struct SseEventDecoder {
/// `pending` 中已确认不含帧分隔符的前缀长度。新数据到达时从该边界
/// 回退分隔符最大跨界宽度继续扫描,大帧跨多个 chunk 不会被反复重扫。
scanned: usize,
/// 超大未完成帧已开始直接交付;在遇到帧分隔符前只保留跨 chunk 扫描尾部,
/// 不再尝试解析该帧。
opaque_frame: bool,
}
impl SseEventDecoder {
@@ -129,28 +133,35 @@ impl SseEventDecoder {
/// 与 [`Self::push`] 相比,此方法仅应由需要原样转发的 transport 使用;普通
/// 解析路径无需为原始帧额外分配内存。旁路解析失败时仍保留原始帧,并以空
/// event 列表表示不可观测,避免观测失败截断客户端字节流。
pub fn push_frames(&mut self, chunk: &[u8]) -> Result<Vec<SseFrame>, SseError> {
pub fn push_frames(&mut self, chunk: &[u8]) -> Vec<SseFrame> {
self.pending.extend_from_slice(chunk);
let mut frames = Vec::new();
if self.opaque_frame && !self.finish_opaque_frame(&mut frames) {
self.flush_opaque_prefix(&mut frames);
return frames;
}
let mut consumed = 0usize;
while let Some(end) = self.next_frame_end(consumed) {
let raw = self.pending[consumed..end].to_vec();
let events = std::str::from_utf8(&raw)
.ok()
.and_then(|frame| parse_sse_events(frame).ok())
.unwrap_or_default();
let events = if raw.len() > MAX_SSE_EVENT_BUFFER_BYTES {
Vec::new()
} else {
std::str::from_utf8(&raw)
.ok()
.and_then(|frame| parse_sse_events(frame).ok())
.unwrap_or_default()
};
frames.push(SseFrame::new(raw, events));
consumed = end;
}
self.commit_consumed(consumed);
if self.pending.len() > MAX_SSE_EVENT_BUFFER_BYTES {
return Err(SseError::BufferExceeded {
max_bytes: MAX_SSE_EVENT_BUFFER_BYTES,
});
self.opaque_frame = true;
self.flush_opaque_prefix(&mut frames);
}
Ok(frames)
frames
}
/// 流结束时解析尚未带空行分隔符的最后一帧。
@@ -168,17 +179,53 @@ impl SseEventDecoder {
/// 在流结束时返回未带空行分隔符的最后一帧及其旁路解析结果。
///
/// 与 [`Self::push_frames`] 一致,旁路解析失败不会丢弃原始帧。
pub fn finish_frames(&mut self) -> Result<Vec<SseFrame>, SseError> {
pub fn finish_frames(&mut self) -> Vec<SseFrame> {
self.scanned = 0;
if self.pending.is_empty() {
return Ok(Vec::new());
self.opaque_frame = false;
return Vec::new();
}
let raw = std::mem::take(&mut self.pending);
let events = std::str::from_utf8(&raw)
.ok()
.and_then(|frame| parse_sse_events(frame).ok())
.unwrap_or_default();
Ok(vec![SseFrame::new(raw, events)])
let events = if self.opaque_frame || raw.len() > MAX_SSE_EVENT_BUFFER_BYTES {
Vec::new()
} else {
std::str::from_utf8(&raw)
.ok()
.and_then(|frame| parse_sse_events(frame).ok())
.unwrap_or_default()
};
self.opaque_frame = false;
vec![SseFrame::new(raw, events)]
}
/// 继续透明交付一个已超出观测预算的帧,并在分隔符后恢复普通解析。
fn finish_opaque_frame(&mut self, frames: &mut Vec<SseFrame>) -> bool {
let Some((position, separator_len)) = sse_frame_separator_bytes(&self.pending) else {
return false;
};
let end = position + separator_len;
let remaining = self.pending.split_off(end);
let raw = std::mem::replace(&mut self.pending, remaining);
frames.push(SseFrame::new(raw, Vec::new()));
self.opaque_frame = false;
self.scanned = 0;
true
}
/// 仅保留分隔符可能跨 chunk 的最大前缀宽度,其余原始字节立即释放给下游。
fn flush_opaque_prefix(&mut self, frames: &mut Vec<SseFrame>) {
const SEPARATOR_CROSS_CHUNK_BYTES: usize = 3;
let emit_len = self
.pending
.len()
.saturating_sub(SEPARATOR_CROSS_CHUNK_BYTES);
if emit_len == 0 {
return;
}
let tail = self.pending.split_off(emit_len);
let raw = std::mem::replace(&mut self.pending, tail);
self.scanned = 0;
frames.push(SseFrame::new(raw, Vec::new()));
}
}
@@ -366,12 +413,12 @@ pub fn encode_sse_event_with_metadata(
retry: Option<u64>,
) -> String {
let mut frame = String::new();
if let Some(id) = id {
if let Some(id) = id.filter(|id| sse_id_is_representable(id)) {
frame.push_str("id: ");
frame.push_str(id);
frame.push('\n');
}
if !event.is_empty() {
if !event.is_empty() && sse_field_value_is_single_line(event) {
frame.push_str("event: ");
frame.push_str(event);
frame.push('\n');
@@ -390,6 +437,14 @@ pub fn encode_sse_event_with_metadata(
frame
}
fn sse_field_value_is_single_line(value: &str) -> bool {
!value.contains(['\r', '\n'])
}
fn sse_id_is_representable(value: &str) -> bool {
sse_field_value_is_single_line(value) && !value.contains('\0')
}
fn track_event_buffer_bytes(current_bytes: &mut usize, line: &str) -> Result<(), SseError> {
let line_separator_bytes = usize::from(*current_bytes != 0);
*current_bytes = current_bytes
@@ -450,13 +450,17 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
#[test]
fn parse_rate_limits_event_should_reject_unrelated_event_type() {
assert_eq!(
parse_rate_limits_event(&json!({
for event in [
json!({
"type": "response.completed",
"rate_limits": {"primary": {"used_percent": 100}}
})),
None
);
}),
json!({
"rate_limits": {"primary": {"used_percent": 100}}
}),
] {
assert_eq!(parse_rate_limits_event(&event), None);
}
}
#[test]
@@ -135,13 +135,8 @@ fn raw_frame_decoder_should_preserve_exact_frame_across_chunk_boundaries() {
let raw = b": keep-alive\r\nid: evt_1\r\nevent: response.future\r\nretry: 250\r\ndata: { \"type\": \"response.future\", \"opaque\": true }\r\n\r\n";
let mut decoder = SseEventDecoder::default();
assert!(
decoder
.push_frames(&raw[..37])
.expect("partial raw frame")
.is_empty()
);
let frames = decoder.push_frames(&raw[37..]).expect("complete raw frame");
assert!(decoder.push_frames(&raw[..37]).is_empty());
let frames = decoder.push_frames(&raw[37..]);
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].raw(), raw);
@@ -161,15 +156,11 @@ fn raw_frame_decoder_should_preserve_unparseable_frames_for_transparent_transpor
let non_utf8 = b"data: \xff\n\n";
let mut decoder = SseEventDecoder::default();
let frames = decoder
.push_frames(invalid_retry)
.expect("raw frame remains deliverable");
let frames = decoder.push_frames(invalid_retry);
assert_eq!(frames[0].raw(), invalid_retry);
assert!(frames[0].events().is_empty());
let frames = decoder
.push_frames(non_utf8)
.expect("non-UTF-8 raw frame remains deliverable");
let frames = decoder.push_frames(non_utf8);
assert_eq!(frames[0].raw(), non_utf8);
assert!(frames[0].events().is_empty());
}
@@ -179,15 +170,8 @@ fn raw_frame_decoder_finish_should_preserve_unparseable_trailing_frame() {
let raw = b"retry: later\ndata: opaque";
let mut decoder = SseEventDecoder::default();
assert!(
decoder
.push_frames(raw)
.expect("unfinished raw frame remains buffered")
.is_empty()
);
let frames = decoder
.finish_frames()
.expect("trailing raw frame remains deliverable");
assert!(decoder.push_frames(raw).is_empty());
let frames = decoder.finish_frames();
assert_eq!(frames[0].raw(), raw);
assert!(frames[0].events().is_empty());
@@ -232,6 +216,32 @@ fn incremental_decoder_should_accept_exact_limit_and_reject_one_byte_over() {
);
}
#[test]
fn raw_frame_decoder_should_stream_oversized_unfinished_frame_and_resume_observation() {
let oversized = vec![b'x'; MAX_SSE_EVENT_BUFFER_BYTES + 1];
let mut decoder = SseEventDecoder::default();
let mut raw = decoder
.push_frames(&oversized)
.into_iter()
.flat_map(|frame| frame.into_parts().0)
.collect::<Vec<_>>();
assert!(!raw.is_empty());
let tail = decoder.push_frames(b"\n\n");
assert!(tail.iter().all(|frame| frame.events().is_empty()));
raw.extend(tail.into_iter().flat_map(|frame| frame.into_parts().0));
assert_eq!(raw.len(), oversized.len() + 2);
assert_eq!(&raw[..oversized.len()], oversized.as_slice());
assert_eq!(&raw[oversized.len()..], b"\n\n");
let observed = decoder.push_frames(b"event: response.future\ndata: {\"ok\":true}\n\n");
assert_eq!(observed.len(), 1);
assert_eq!(
observed[0].events()[0].event.as_deref(),
Some("response.future")
);
}
#[test]
fn encoder_should_round_trip_multiline_data() {
let frame = encode_sse_event("response.output_text.delta", "first\nsecond");
@@ -267,6 +277,20 @@ fn encoder_should_round_trip_sse_id_and_retry() {
);
}
#[test]
fn encoder_should_omit_metadata_that_sse_cannot_represent_as_one_field() {
let frame = encode_sse_event_with_metadata(
"future\nevent",
r#"{"type":"future\nevent"}"#,
Some("id\0with\nline"),
None,
);
assert!(!frame.contains("event:"));
assert!(!frame.contains("id:"));
assert!(frame.contains(r#"data: {"type":"future\nevent"}"#));
}
#[test]
fn response_failed_event_should_match_openai_failure_shape() {
let frame = response_failed_sse_event("server_error", "stream_disconnected", "closed early");
+1 -1
View File
@@ -20,7 +20,7 @@ gateway-core = { path = "../gateway-core" }
hex = "0.4.3"
redis = { version = "1.3.0", default-features = false, features = ["tokio-comp", "connection-manager", "script"] }
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"] }
sha2 = "0.11.0"
sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres", "chrono", "json", "macros", "migrate"] }
thiserror = "2.0.18"
@@ -224,6 +224,7 @@ pub struct ModelRequestFinalization {
pub upstream_transport: Option<String>,
pub http_version: Option<String>,
pub websocket_pool: Option<String>,
pub service_tier: Option<String>,
pub provider_metadata_json: Option<Value>,
pub error_kind: Option<String>,
pub provider_error_code: Option<String>,
@@ -276,6 +277,12 @@ impl ModelRequestFinalization {
{
return Err(invalid("websocket pool must be new or reuse"));
}
if let Some(service_tier) = self.service_tier.as_deref() {
require_nonempty(ENTITY, "service_tier", service_tier)?;
if service_tier.len() > 64 || service_tier.chars().any(char::is_control) {
return Err(invalid("service tier is invalid"));
}
}
if self
.provider_metadata_json
.as_ref()
@@ -573,7 +580,7 @@ impl ModelRequestRepository for PgExecutionStore {
latency_ms = $35, completed_at = $36,
upstream_transport = coalesce($37, upstream_transport),
http_version = coalesce($38, http_version), websocket_pool = $39,
provider_observation_json = $40
service_tier = $40, provider_observation_json = $41
where id = $1 and outcome = 'running'",
)
.bind(&finalization.model_request_id)
@@ -660,6 +667,7 @@ impl ModelRequestRepository for PgExecutionStore {
.bind(finalization.upstream_transport)
.bind(finalization.http_version)
.bind(finalization.websocket_pool)
.bind(finalization.service_tier)
.bind(finalization.provider_metadata_json.map(sqlx::types::Json))
.execute(&self.pool)
.await
@@ -927,6 +935,7 @@ impl ExecutionStore for PgExecutionStore {
upstream_transport: finalization.upstream_transport,
http_version: finalization.http_version,
websocket_pool: finalization.websocket_pool,
service_tier: finalization.service_tier,
provider_metadata_json,
error_kind,
provider_error_code: finalization.provider_error_code,
@@ -369,6 +369,7 @@ pub struct CalculatedUsageBillingFact {
pub bucket_start: DateTime<Utc>,
pub provider_kind: String,
pub upstream_model_id: String,
pub service_tier: Option<String>,
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub cached_tokens: Option<u64>,
@@ -521,6 +522,7 @@ pub struct UsageRecord {
pub upstream_transport: Option<String>,
pub http_version: Option<String>,
pub websocket_pool: Option<String>,
pub service_tier: Option<String>,
pub provider_metadata_json: Option<String>,
pub attempt_count: u32,
pub upstream_send_state: String,
@@ -1378,6 +1380,7 @@ fn admin_calculated_usage_billing_fact(
bucket_start: fact.bucket_start,
provider_kind: fact.provider_kind,
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,
@@ -1453,6 +1456,7 @@ fn admin_usage_record(record: UsageRecord) -> AdminStoreResult<admin_observabili
upstream_transport: record.upstream_transport,
http_version: record.http_version,
websocket_pool: record.websocket_pool,
service_tier: record.service_tier,
provider_metadata_json: record.provider_metadata_json,
attempt_count: record.attempt_count,
upstream_send_state: record.upstream_send_state,
@@ -1812,7 +1816,7 @@ async fn calculated_usage_billing_facts(
query.push_bind(granularity.sql_interval());
query.push(
"::interval, mr.started_at, timestamptz '1970-01-01 00:00:00+00') as bucket_start,
mr.provider_kind, mr.upstream_model_id,
mr.provider_kind, mr.upstream_model_id, mr.service_tier,
mr.input_tokens, mr.output_tokens, mr.cached_tokens, mr.cache_write_tokens,
mr.cost_currency, mr.cost_amount::text as amount
from model_requests mr where mr.started_at >= ",
@@ -2609,7 +2613,7 @@ const USAGE_RECORD_SELECT: &str =
pa.name as provider_account_name, pa.email as provider_account_email,
pa.authentication_kind as provider_account_authentication_kind,
mr.upstream_model_id, mr.upstream_transport, mr.http_version, mr.websocket_pool,
mr.provider_observation_json,
mr.service_tier, mr.provider_observation_json,
mr.attempt_count, mr.upstream_send_state, mr.downstream_committed_at,
mr.outcome, mr.client_status_code, mr.upstream_status_code,
mr.client_response_id, mr.upstream_request_id, mr.upstream_response_id,
@@ -3107,6 +3111,7 @@ fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<UsageRecord
upstream_transport: get(row, "upstream_transport")?,
http_version: get(row, "http_version")?,
websocket_pool: get(row, "websocket_pool")?,
service_tier: get(row, "service_tier")?,
provider_metadata_json: get::<Option<serde_json::Value>>(row, "provider_observation_json")?
.map(|value| serde_json::to_string(&value))
.transpose()
@@ -3202,6 +3207,7 @@ fn calculated_usage_billing_fact_from_row(
bucket_start: get(row, "bucket_start")?,
provider_kind: get(row, "provider_kind")?,
upstream_model_id: get(row, "upstream_model_id")?,
service_tier: get(row, "service_tier")?,
input_tokens: optional_unsigned(row, "input_tokens")?,
output_tokens: optional_unsigned(row, "output_tokens")?,
cached_tokens: optional_unsigned(row, "cached_tokens")?,
@@ -250,15 +250,24 @@ async fn core_adapter_should_persist_calculated_cost_exactly() {
let mut finalization = successful_core_finalization("req_calculated_cost");
finalization.websocket_pool = Some("reuse".to_owned());
finalization.service_tier = Some("priority".to_owned());
finalization.cost = CalculatedCost::from_usd_ticks(12_345)
.expect("calculated cost")
.into_estimate();
ExecutionStore::finalize_model_request(&repository, finalization)
.await
.expect("persist calculated cost");
let persisted: (String, String, String, String, String, String) = sqlx::query_as(
let persisted: (
String,
String,
String,
String,
String,
String,
Option<String>,
) = sqlx::query_as(
"select cost_source, cost_amount::text, cost_currency, upstream_transport, http_version,
websocket_pool
websocket_pool, service_tier
from model_requests where id = 'req_calculated_cost'",
)
.fetch_one(&database.pool)
@@ -274,6 +283,7 @@ async fn core_adapter_should_persist_calculated_cost_exactly() {
"websocket".to_owned(),
"HTTP/2".to_owned(),
"reuse".to_owned(),
Some("priority".to_owned()),
)
);
database.close().await;
@@ -376,6 +386,7 @@ fn successful_core_finalization(id: &str) -> CoreModelRequestFinalization {
upstream_transport: Some("websocket".to_owned()),
http_version: Some("HTTP/2".to_owned()),
websocket_pool: None,
service_tier: None,
provider_metadata_json: None,
error: None,
provider_error_code: None,
@@ -70,6 +70,12 @@ impl TestDatabase {
.execute(&mut *migration)
.await
.expect("apply opaque response storage migration");
sqlx::raw_sql(include_str!(
"../../../../migrations/0004_add_model_request_service_tier.sql"
))
.execute(&mut *migration)
.await
.expect("apply model request service tier migration");
migration.commit().await.expect("commit terminal migration");
Some(Self {
admin,
@@ -165,7 +171,7 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
.expect("drop migration test database");
admin.close().await;
assert_eq!((first_table_count, migration_count), (8, 3));
assert_eq!((first_table_count, migration_count), (8, 4));
assert_eq!(response_id_types, ["bytea", "bytea"]);
assert!(!raw_response_id_index_exists);
}
@@ -176,12 +182,13 @@ fn migrations_should_leave_transaction_ownership_to_sqlx() {
include_str!("../../../../migrations/0001_initial.sql"),
include_str!("../../../../migrations/0002_drop_opaque_response_id_index.sql"),
include_str!("../../../../migrations/0003_store_opaque_response_ids_as_bytes.sql"),
include_str!("../../../../migrations/0004_add_model_request_service_tier.sql"),
]
.into_iter()
.flat_map(str::lines)
.map(str::trim)
.filter(|line| matches!(*line, "begin;" | "commit;"))
.count();
.map(str::trim)
.filter(|line| matches!(*line, "begin;" | "commit;"))
.count();
assert_eq!(transaction_statements, 0);
}
@@ -99,8 +99,14 @@ async fn observability_preserves_and_filters_opaque_response_ids() {
.await
.expect("filter opaque response ID");
assert_eq!(records.total, 1);
assert_eq!(records.items[0].client_response_id.as_deref(), Some(response_id.as_str()));
assert_eq!(records.items[0].upstream_response_id.as_deref(), Some(response_id.as_str()));
assert_eq!(
records.items[0].client_response_id.as_deref(),
Some(response_id.as_str())
);
assert_eq!(
records.items[0].upstream_response_id.as_deref(),
Some(response_id.as_str())
);
database.close().await;
}
@@ -181,6 +187,7 @@ async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
assert_eq!(facts[0].upstream_model_id, "gpt-5.5");
assert_eq!(facts[0].input_tokens, Some(800));
assert_eq!(facts[0].output_tokens, Some(200));
assert_eq!(facts[0].service_tier.as_deref(), Some("priority"));
assert_eq!(facts[0].total.amount.as_str(), "1.25");
let store = PgAdminObservabilityStore::new(database.pool.clone(), None);
@@ -193,6 +200,7 @@ async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
.await
.expect("admin calculated usage billing facts");
assert_eq!(facts.len(), 1);
assert_eq!(facts[0].service_tier.as_deref(), Some("priority"));
assert_eq!(facts[0].total.amount.as_str(), "1.25");
database.close().await;
@@ -240,6 +248,10 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
1,
);
assert_eq!(dashboard.recent_requests.len(), 1);
assert_eq!(
dashboard.recent_requests[0].service_tier.as_deref(),
Some("priority")
);
assert_eq!(
dashboard.recent_requests[0].outcome,
admin_observability::RequestOutcome::Succeeded,
@@ -313,6 +325,10 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
assert_eq!(first_page.total, 1);
assert_eq!(first_page.items.len(), 1);
assert!(first_page.next_cursor.is_none());
assert_eq!(
first_page.items[0].service_tier.as_deref(),
Some("priority")
);
let filtered = store
.list_usage_records(admin_observability::UsageQuery {
@@ -339,6 +355,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
.await
.expect("fully filtered usage page");
assert_eq!(filtered.total, 1);
assert_eq!(filtered.items[0].service_tier.as_deref(), Some("priority"));
assert_eq!(
filtered.items[0].upstream_request_id.as_deref(),
Some("upstream_req_success")
@@ -371,6 +388,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
detail.request.outcome,
admin_observability::RequestOutcome::Succeeded
);
assert_eq!(detail.request.service_tier.as_deref(), Some("priority"));
assert_eq!(detail.attempts.len(), 1);
assert_eq!(
detail.attempts[0].outcome,
@@ -550,6 +568,10 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
);
assert_eq!(dashboard.recent_requests.len(), 1);
assert_eq!(dashboard.recent_requests[0].id, "req_observe_success");
assert_eq!(
dashboard.recent_requests[0].service_tier.as_deref(),
Some("priority")
);
assert_eq!(
dashboard
.trend
@@ -624,6 +646,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
assert_eq!(
(
successful_image.websocket_pool.as_deref(),
successful_image.service_tier.as_deref(),
successful_image
.provider_account_authentication_kind
.as_deref(),
@@ -634,6 +657,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
),
(
Some("reuse"),
Some("priority"),
Some("oauth"),
Some(31),
Some(9),
@@ -769,6 +793,7 @@ async fn seed_observability_facts(
client_transport, requested_model_id,
provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, websocket_pool,
service_tier,
attempt_count,
upstream_send_state, downstream_committed_at, outcome, client_status_code,
upstream_status_code, client_response_id, upstream_request_id, upstream_response_id,
@@ -781,7 +806,7 @@ async fn seed_observability_facts(
'req_observe_success', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe',
'acct_observe', 'upstream-model',
'http_sse', 'reuse', 1, 'sent', $1 - interval '19 minutes', 'succeeded', 200, 200,
'http_sse', 'reuse', 'priority', 1, 'sent', $1 - interval '19 minutes', 'succeeded', 200, 200,
'resp_observe_success', 'upstream_req_success', 'upstream_resp_success',
100, 20, 40, 3, 5, 31, 9, 120, true, true,
'provider_reported', 1.25, 'USD', 120, 900,
@@ -843,13 +868,13 @@ async fn seed_calculated_billing_facts(
client_transport, requested_model_id, provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
upstream_send_state, downstream_committed_at, outcome, client_status_code,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, total_tokens,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, total_tokens, service_tier,
cost_source, cost_amount, cost_currency, started_at, deadline_at, completed_at
) values (
'req_observe_calculated', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe', 'acct_observe', 'gpt-5.5',
'http_sse', 1, 'sent', $1 - interval '29 minutes', 'succeeded', 200,
800, 200, 0, 0, 1000, 'calculated', 1.25, 'USD',
800, 200, 0, 0, 1000, 'priority', 'calculated', 1.25, 'USD',
$1 - interval '30 minutes', $1 + interval '10 minutes', $1 - interval '29 minutes'
)",
)
+1 -1
View File
@@ -31,7 +31,7 @@ rustls-native-certs = "0.8.4"
rustls-pki-types = "1.14.1"
secrecy = "0.10.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"] }
sha2 = "0.11.0"
subtle = "2.6.1"
thiserror = "2.0.18"
+14 -1
View File
@@ -168,6 +168,19 @@ impl ProviderAdmin for OpenAiAdminProvider {
self.websocket_pool.evict_account(account_id.as_str()).await;
}
async fn account_facts_changed(&self, account_ids: &[ProviderAccountId]) {
if account_ids.is_empty() {
return;
}
if let Err(error) = self.catalog.invalidate() {
tracing::warn!(
account_count = account_ids.len(),
error = %error,
"OpenAI model catalog invalidation failed after account commit"
);
}
}
fn connection_test_operation(
&self,
upstream_model: &UpstreamModelId,
@@ -216,7 +229,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
output_tokens,
input.cached_tokens.unwrap_or_default(),
input.cache_write_tokens.unwrap_or_default(),
None,
input.service_tier.as_deref(),
) else {
return Ok(None);
};
@@ -7,7 +7,7 @@ use std::time::{Duration, SystemTime};
use chrono::{DateTime, Utc};
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId,
AccountAvailability, OpaqueProviderData, ProviderAccount, ProviderAccountId,
};
use gateway_core::engine::provider::ProviderCatalogGeneration;
use gateway_core::provider_ports::{
@@ -96,13 +96,7 @@ impl CodexPlanCatalog {
pub struct CodexCredentialCatalogSnapshot {
observed_at: SystemTime,
models: Vec<CodexCatalogModel>,
account_models: BTreeMap<ProviderAccountId, CodexAccountEntitlement>,
}
#[derive(Clone, PartialEq, Eq)]
struct CodexAccountEntitlement {
revision: CredentialRevision,
models: Vec<String>,
scope_models: BTreeMap<CodexCatalogScope, Vec<String>>,
}
impl CodexCredentialCatalogSnapshot {
@@ -116,16 +110,12 @@ impl CodexCredentialCatalogSnapshot {
&self.models
}
#[must_use]
pub fn account_models(
&self,
account_id: &ProviderAccountId,
revision: CredentialRevision,
) -> Option<&[String]> {
self.account_models
.get(account_id)
.filter(|entry| entry.revision == revision)
.map(|entry| entry.models.as_slice())
account: &ProviderAccount,
) -> Result<Option<&[String]>, CodexCredentialCatalogError> {
let scope = CodexCatalogScope::for_account(account)?;
Ok(self.scope_models.get(&scope).map(Vec::as_slice))
}
}
@@ -135,7 +125,7 @@ impl fmt::Debug for CodexCredentialCatalogSnapshot {
.debug_struct("CodexCredentialCatalogSnapshot")
.field("observed_at", &self.observed_at)
.field("model_count", &self.models.len())
.field("account_count", &self.account_models.len())
.field("scope_count", &self.scope_models.len())
.finish()
}
}
@@ -173,8 +163,6 @@ struct CatalogCacheState {
}
struct FetchedAccountModels {
account_id: ProviderAccountId,
revision: CredentialRevision,
models: Vec<CodexCatalogModel>,
etag: Option<String>,
}
@@ -246,30 +234,30 @@ impl CodexCredentialCatalogService {
.cloned())
}
/// 读取单账号当前仍新鲜的模型 entitlement;不触发网络。
/// 读取账号当前已验证套餐的模型 entitlement;不触发网络。
pub fn cached_account_models(
&self,
account_id: &ProviderAccountId,
revision: CredentialRevision,
account: &ProviderAccount,
) -> Result<Option<Vec<String>>, CodexCredentialCatalogError> {
Ok(self.cached()?.and_then(|snapshot| {
snapshot
.account_models(account_id, revision)
.map(<[String]>::to_vec)
}))
let Some(snapshot) = self.cached()? else {
return Ok(None);
};
snapshot
.account_models(account)
.map(|models| models.map(<[String]>::to_vec))
}
pub fn observed_model_support(
&self,
account_id: &ProviderAccountId,
revision: CredentialRevision,
account: &ProviderAccount,
model: &str,
) -> Result<Option<bool>, CodexCredentialCatalogError> {
Ok(self.cached()?.and_then(|snapshot| {
snapshot
.account_models(account_id, revision)
.map(|models| models.iter().any(|candidate| candidate == model))
}))
let Some(snapshot) = self.cached()? else {
return Ok(None);
};
Ok(Some(snapshot.account_models(account)?.is_some_and(
|models| models.iter().any(|candidate| candidate == model),
)))
}
/// 优先读取套餐目录 cache;缺失时才用当前账号所属套餐的有限候选集实时填充。
@@ -353,13 +341,13 @@ impl CodexCredentialCatalogService {
);
let mut union = BTreeMap::<String, CodexCatalogModel>::new();
let mut union_order = Vec::new();
let mut account_models = BTreeMap::new();
let mut scope_models = BTreeMap::new();
let mut etags = Vec::new();
for (scope, candidates) in groups {
let fetched = self.fetch_scope_models(&client, candidates).await?;
let entitlement = model_ids(&fetched.models);
self.replace_plan_catalog(&CodexPlanCatalog::new(
scope,
scope.clone(),
Utc::now(),
entitlement.clone(),
))
@@ -379,13 +367,7 @@ impl CodexCredentialCatalogService {
Entry::Occupied(_) => {}
}
}
account_models.insert(
fetched.account_id,
CodexAccountEntitlement {
revision: fetched.revision,
models: entitlement,
},
);
scope_models.insert(scope, entitlement);
etags.extend(fetched.etag);
}
let observed_at = SystemTime::now();
@@ -395,7 +377,7 @@ impl CodexCredentialCatalogService {
.into_iter()
.filter_map(|id| union.remove(&id))
.collect(),
account_models,
scope_models,
};
Ok(FetchedCatalog { snapshot, etags })
}
@@ -502,8 +484,6 @@ impl CodexCredentialCatalogService {
}
let snapshot = result.map_err(|_| CodexCredentialCatalogError::Upstream)?;
Ok(FetchedAccountModels {
account_id: prepared.account.id().clone(),
revision: prepared.account.revision(),
models: snapshot.models().to_vec(),
etag: snapshot.etag().map(str::to_owned),
})
@@ -782,7 +762,7 @@ fn same_catalog(
left: &CodexCredentialCatalogSnapshot,
right: &CodexCredentialCatalogSnapshot,
) -> bool {
left.models == right.models && left.account_models == right.account_models
left.models == right.models && left.scope_models == right.scope_models
}
fn validate_response_etag(etag: &str) -> Result<(), CodexCredentialCatalogError> {
@@ -7,8 +7,8 @@ use std::time::{Duration, SystemTime};
use gateway_core::engine::credential::{
AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats,
AccountRuntimeSignals, AccountSelectionContext, AccountSelector, ProviderAccount,
ProviderAccountId,
AccountRuntimeSignals, AccountSchedulingBlocker, AccountSelectionContext, AccountSelector,
PreferredAccountSelection, ProviderAccount, ProviderAccountId,
};
use gateway_core::engine::{AttemptContext, ContinuationAttempt};
use gateway_core::provider_ports::{
@@ -118,6 +118,118 @@ enum SessionAffinityLookup {
Unavailable,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AffinityEscapeReason {
HardUnavailable,
QuotaExhausted,
Cooldown,
LeaseSaturated,
PinnedAccount,
SelectionInvariant,
}
impl AffinityEscapeReason {
const fn as_str(self) -> &'static str {
match self {
Self::HardUnavailable => "hard_unavailable",
Self::QuotaExhausted => "quota_exhausted",
Self::Cooldown => "cooldown",
Self::LeaseSaturated => "lease_saturated",
Self::PinnedAccount => "pinned_account",
Self::SelectionInvariant => "selection_invariant",
}
}
}
#[derive(Debug, Default)]
struct AffinitySelection {
bound_account: Option<ProviderAccountId>,
preferred_account: Option<ProviderAccountId>,
escape_reason: Option<AffinityEscapeReason>,
}
impl AffinitySelection {
fn preferred(account_id: ProviderAccountId) -> Self {
Self {
bound_account: Some(account_id.clone()),
preferred_account: Some(account_id),
escape_reason: None,
}
}
fn escaped(account_id: ProviderAccountId, reason: AffinityEscapeReason) -> Self {
Self {
bound_account: Some(account_id),
preferred_account: None,
escape_reason: Some(reason),
}
}
fn bound_account(&self) -> Option<&ProviderAccountId> {
self.bound_account.as_ref()
}
fn preferred_account(&self) -> Option<&ProviderAccountId> {
self.preferred_account.as_ref()
}
fn escape(&mut self, reason: AffinityEscapeReason) {
if self.bound_account.is_some() && self.escape_reason.is_none() {
self.escape_reason = Some(reason);
self.preferred_account = None;
}
}
fn observe_preferred_selection(&mut self, selection: PreferredAccountSelection) {
if self.bound_account.is_none() || self.escape_reason.is_some() {
return;
}
match selection {
PreferredAccountSelection::Hit => {}
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::ConcurrencyLimit) => {
self.escape(AffinityEscapeReason::LeaseSaturated);
}
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::RequestInterval) => {
self.escape(AffinityEscapeReason::Cooldown);
}
PreferredAccountSelection::Blocked(
AccountSchedulingBlocker::LocalAvailability | AccountSchedulingBlocker::Excluded,
)
| PreferredAccountSelection::Missing => {
self.escape(AffinityEscapeReason::HardUnavailable);
}
PreferredAccountSelection::NotRequested => {
self.escape(AffinityEscapeReason::SelectionInvariant);
}
}
}
fn observe_lease_busy(&mut self, account_id: &ProviderAccountId) {
if self.bound_account.as_ref() == Some(account_id) {
self.escape(AffinityEscapeReason::LeaseSaturated);
}
}
fn telemetry(&self, selected_account: &ProviderAccountId) -> AffinityTelemetry {
AffinityTelemetry {
affinity_hit: self.bound_account.as_ref() == Some(selected_account)
&& self.escape_reason.is_none(),
escape_reason: self.escape_reason,
account_switch: self
.bound_account
.as_ref()
.is_some_and(|bound| bound != selected_account),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
struct AffinityTelemetry {
affinity_hit: bool,
escape_reason: Option<AffinityEscapeReason>,
account_switch: bool,
}
impl CodexCredentialSelector {
#[must_use]
// 选择器显式持有各能力边界,避免把 Provider 私有服务重新包装成通用容器。
@@ -169,15 +281,12 @@ impl CodexCredentialSelector {
.into_iter()
.filter(|account| {
account.provider() == &self.provider_kind
&& (diagnostic
|| !matches!(
self.catalog.observed_model_support(
account.id(),
account.revision(),
request.upstream_model,
),
Ok(Some(false))
))
&& (diagnostic || {
let observed_support = self
.catalog
.observed_model_support(account, request.upstream_model);
matches!(observed_support, Ok(None | Some(true)))
})
})
.collect::<Vec<_>>();
if !diagnostic {
@@ -228,27 +337,15 @@ impl CodexCredentialSelector {
.collect::<Vec<_>>()
});
let scheduling_candidates = quota_exhausted_candidates.as_deref().unwrap_or(&candidates);
let affinity_account = if diagnostic {
None
let mut affinity = if diagnostic {
AffinitySelection::default()
} else {
match request.session_affinity_key {
Some(key) => match self.lookup_session_affinity(key).await {
SessionAffinityLookup::Bound(account_id)
if scheduling_candidates.iter().any(|candidate| {
candidate.account.id() == &account_id
&& candidate.account.is_schedulable(SystemTime::now())
}) =>
{
Some(account_id)
}
SessionAffinityLookup::Bound(_) => {
self.clear_session_affinity(key).await;
None
}
SessionAffinityLookup::Missing | SessionAffinityLookup::Unavailable => None,
},
None => None,
}
self.resolve_session_affinity(
request.session_affinity_key,
&candidates,
SystemTime::now(),
)
.await
};
let cyber_policy_scope = self
.prepare_cyber_policy_scope(cyber_policy_session_key)
@@ -289,7 +386,16 @@ impl CodexCredentialSelector {
}
}
}
let preferred = pinned_account.or(affinity_account);
if pinned_account
.as_ref()
.zip(affinity.bound_account())
.is_some_and(|(pinned, bound)| pinned != bound)
{
affinity.escape(AffinityEscapeReason::PinnedAccount);
}
let preferred = pinned_account
.clone()
.or_else(|| affinity.preferred_account().cloned());
let mut shortest_retry = None;
loop {
@@ -305,7 +411,7 @@ impl CodexCredentialSelector {
AccountAvailabilityPolicy::Enforce
},
};
let Some(selected) = AccountSelector.select(scheduling_candidates, &context) else {
let Some(selection) = AccountSelector.select(scheduling_candidates, &context) else {
return match shortest_retry {
Some(retry_after) => Err(CredentialSelectionError::CapacityUnavailable {
retry_after: Some(retry_after),
@@ -313,6 +419,8 @@ impl CodexCredentialSelector {
None => Err(CredentialSelectionError::NoEligibleCredential),
};
};
affinity.observe_preferred_selection(selection.preferred());
let selected = selection.candidate();
let account = candidates
.iter()
.find(|candidate| candidate.account.id() == selected.account.id())
@@ -335,10 +443,24 @@ impl CodexCredentialSelector {
.await?
{
ProviderLeaseAcquisition::Busy { retry_after } => {
affinity.observe_lease_busy(account.id());
shortest_retry = minimum_duration(shortest_retry, retry_after);
excluded.insert(account.id().clone());
}
ProviderLeaseAcquisition::Acquired(guard) => {
let affinity_telemetry = affinity.telemetry(account.id());
tracing::info!(
request_id = %request.attempt.request_id(),
attempt_index = request.attempt.attempt_index().get(),
rotation_strategy = policy.strategy().as_str(),
account_id = %account.id(),
affinity_hit = affinity_telemetry.affinity_hit,
escape_reason = affinity_telemetry
.escape_reason
.map_or("", AffinityEscapeReason::as_str),
account_switch = affinity_telemetry.account_switch,
"OpenAI account selected"
);
let (account, runtime) = if account.authentication_kind()
== CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY
{
@@ -376,6 +498,7 @@ impl CodexCredentialSelector {
cookies,
cyber_policy_scope,
allows_account_state_mutation,
affinity_telemetry,
_guard: guard,
});
}
@@ -383,6 +506,34 @@ impl CodexCredentialSelector {
}
}
async fn resolve_session_affinity(
&self,
key: Option<&ProviderSessionAffinityKey>,
candidates: &[AccountCandidate],
now: SystemTime,
) -> AffinitySelection {
let Some(key) = key else {
return AffinitySelection::default();
};
let SessionAffinityLookup::Bound(account_id) = self.lookup_session_affinity(key).await
else {
return AffinitySelection::default();
};
let Some(candidate) = candidates
.iter()
.find(|candidate| candidate.account.id() == &account_id)
else {
self.clear_session_affinity(key).await;
return AffinitySelection::escaped(account_id, AffinityEscapeReason::HardUnavailable);
};
if candidate.account.is_schedulable(now) {
return AffinitySelection::preferred(account_id);
}
let reason = affinity_unavailable_reason(&candidate.account, now);
self.clear_session_affinity(key).await;
AffinitySelection::escaped(account_id, reason)
}
async fn lookup_session_affinity(
&self,
key: &ProviderSessionAffinityKey,
@@ -687,6 +838,7 @@ impl CodexCredentialSelector {
let Ok(current) = self.current_account(account.id()).await else {
return;
};
// QuotaExhausted 只能由结构化额度事实解除;一次成功响应本身不能证明额度已经恢复。
if current.provider() != &self.provider_kind
|| current.revision() != account.revision()
|| !current.enabled()
@@ -694,7 +846,6 @@ impl CodexCredentialSelector {
current.availability(),
AccountAvailability::Unknown
| AccountAvailability::Cooldown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
)
{
@@ -857,6 +1008,25 @@ fn quota_exhausted_scheduling_projection(account: &ProviderAccount) -> ProviderA
account.clone()
}
fn affinity_unavailable_reason(account: &ProviderAccount, now: SystemTime) -> AffinityEscapeReason {
if !account.enabled()
|| account
.access_token_expires_at()
.is_some_and(|expires_at| expires_at <= now)
{
return AffinityEscapeReason::HardUnavailable;
}
match account.availability() {
AccountAvailability::QuotaExhausted => AffinityEscapeReason::QuotaExhausted,
AccountAvailability::Cooldown => AffinityEscapeReason::Cooldown,
AccountAvailability::Unknown
| AccountAvailability::Ready
| AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::Invalid => AffinityEscapeReason::HardUnavailable,
}
}
fn active_risk_recovery<'a>(
recovery: &'a mut HashMap<String, RiskRecoveryState>,
account_id: &str,
@@ -895,6 +1065,7 @@ pub struct CodexCredentialLease {
installation_id: String,
cyber_policy_scope: Option<CodexCyberPolicyScope>,
allows_account_state_mutation: bool,
affinity_telemetry: AffinityTelemetry,
_guard: Box<dyn ProviderLeaseGuard>,
}
@@ -934,6 +1105,24 @@ impl CodexCredentialLease {
pub(crate) const fn allows_account_state_mutation(&self) -> bool {
self.allows_account_state_mutation
}
#[must_use]
pub const fn affinity_hit(&self) -> bool {
self.affinity_telemetry.affinity_hit
}
#[must_use]
pub const fn escape_reason(&self) -> Option<&'static str> {
match self.affinity_telemetry.escape_reason {
Some(reason) => Some(reason.as_str()),
None => None,
}
}
#[must_use]
pub const fn account_switch(&self) -> bool {
self.affinity_telemetry.account_switch
}
}
impl fmt::Debug for CodexCredentialLease {
+160 -143
View File
@@ -18,8 +18,8 @@ use gateway_core::engine::{
AttemptContext, CancellationToken, ContinuationAttempt, UpstreamSendState,
};
use gateway_core::error::{
ClientVisibleUpstreamError, ContinuationFailure, ProviderError, ProviderErrorKind,
SafeUpstreamValue,
ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ContinuationFailure,
OpaqueUpstreamValue, ProviderError, ProviderErrorKind,
};
use gateway_core::event::{
GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseMetadata,
@@ -70,16 +70,18 @@ use crate::transport::request::{
use crate::transport::session::CodexSessionIdentity;
use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason};
use crate::transport::{
CODEX_RESPONSES_PATH, CodexBackendClient, CodexBackendStreamingResponse, CodexBackendTransport,
CodexClientError, CodexRateLimitHeaderUpdates, CodexRequestContext, CodexResponseMetadata,
CodexTransportMetrics, CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url,
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError,
CodexRateLimitHeaderUpdates, CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics,
CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url,
};
const PROVIDER_NAME: &str = "openai";
const HTTP_SSE_TRANSPORT: &str = "http_sse";
const WEBSOCKET_TRANSPORT: &str = "websocket";
const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024;
/// 与 2.x 预取器一致:提交边界前最多保留 64 KiB 原始上游 chunk。
/// 提交边界前最多保留 64 KiB 原始上游 chunk;达到阈值后结束无感换号窗口,
/// 但不会把上游数据改写成协议失败。
const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api";
pub const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api";
@@ -169,7 +171,7 @@ impl Provider for CodexProvider {
};
let (semantics, reasoning_effort) = encode_generate_request(request, "observability")
.map(|mut encoded| {
if let Ok(Some(previous)) = decode_openai_session_state(request) {
if let Some(previous) = decode_openai_session_state(request) {
let mut input = previous
.transcript
.iter()
@@ -236,7 +238,7 @@ impl Provider for CodexProvider {
UpstreamSendState::NotSent,
));
};
let previous_session = decode_openai_session_state(generate)?;
let previous_session = decode_openai_session_state(generate);
let continuation_requested = generate.native_continuation_requested();
let mut upstream_request =
encode_generate_request(generate, candidate.upstream_model().as_str())
@@ -250,11 +252,7 @@ impl Provider for CodexProvider {
if let Some(identity) = &self.session_identity {
identity.prepare_local_conversation(&mut upstream_request);
}
if upstream_request
.turn_state
.as_deref()
.is_none_or(str::is_empty)
{
if upstream_request.turn_state.is_none() {
upstream_request.turn_state = previous_session
.as_ref()
.and_then(|state| state.turn_state.clone());
@@ -294,13 +292,7 @@ impl Provider for CodexProvider {
) || previous_session.as_ref().is_some_and(|state| {
state.continuation_scope == OpenAiContinuationScope::ReplayRequired
});
if replay_previous_response {
let state = previous_session.as_ref().ok_or_else(|| {
provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
)
})?;
if replay_previous_response && let Some(state) = previous_session.as_ref() {
let mut input = replay_input_for_account(state, lease.account_id().as_str());
input.reserve(request_input.len());
input.extend(request_input.iter().cloned());
@@ -465,17 +457,10 @@ struct OpenAiSessionCapture {
continuation_scope: Option<OpenAiContinuationScope>,
}
fn decode_openai_session_state(
request: &GenerateRequest,
) -> Result<Option<OpenAiSessionState>, ProviderError> {
fn decode_openai_session_state(request: &GenerateRequest) -> Option<OpenAiSessionState> {
request
.provider_session_state(PROVIDER_NAME)
.map(|state| {
serde_json::from_value(Value::Object(state.payload().clone())).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent)
})
})
.transpose()
.and_then(|state| serde_json::from_value(Value::Object(state.payload().clone())).ok())
}
fn encode_openai_session_state(
@@ -496,14 +481,14 @@ fn encode_openai_session_state(
fn attach_openai_session_update(
events: &mut [ProviderEvent],
capture: &mut Option<OpenAiSessionCapture>,
) -> Result<(), ProviderError> {
) {
let Some((terminal_index, output)) = events.iter().enumerate().find_map(|(index, event)| {
terminal_response_output(event).map(|output| (index, output.to_vec()))
}) else {
return Ok(());
return;
};
let Some(capture) = capture.take() else {
return Ok(());
return;
};
let mut transcript = capture
.previous
@@ -524,17 +509,19 @@ fn attach_openai_session_update(
item,
}),
);
let update = encode_openai_session_state(OpenAiSessionState {
let Some(continuation_scope) = capture.continuation_scope else {
return;
};
let Ok(update) = encode_openai_session_state(OpenAiSessionState {
account_id: capture.account_id,
conversation_id: capture.conversation_id,
turn_state: capture.turn_state,
continuation_scope: capture
.continuation_scope
.ok_or_else(|| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent))?,
continuation_scope,
transcript,
})?;
}) else {
return;
};
events[terminal_index].attach_session_update(update);
Ok(())
}
fn terminal_response_output(event: &ProviderEvent) -> Option<&[Value]> {
@@ -686,6 +673,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
})?;
let request_id = context.request_id().as_str().to_owned();
let cancellation = context.cancellation().clone();
let account_selection = CodexAccountSelectionTelemetry::new(
lease.affinity_hit(),
lease.escape_reason(),
lease.account_switch(),
);
let stream_commit_policy =
StreamCommitPolicy::for_continuation(context.continuation_attempt());
let mut response = create_response_attempt(
@@ -698,6 +690,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
lease.installation_id(),
&authorization,
cookie_header.as_ref(),
account_selection,
),
active_account.id().as_str(),
context.deadline(),
@@ -737,6 +730,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
&recovered.credential.installation_id,
&authorization,
cookie_header.as_ref(),
account_selection,
),
active_account.id().as_str(),
context.deadline(),
@@ -789,7 +783,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
active_account.id().as_str(),
context.attempt_index().get(),
);
yield ProviderEvent::observation(observation_state.observation()?);
if let Some(observation) = observation_state.observation() {
yield ProviderEvent::observation(observation);
}
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &response.rate_limit_headers).await;
}
@@ -858,8 +854,10 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
if observation_state.merge_rate_limit_headers(&updates) {
yield ProviderEvent::observation(observation_state.observation()?);
if observation_state.merge_rate_limit_headers(&updates)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
}
}
apply_failure(&failure_context, &active_account, &failure)
@@ -896,6 +894,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
}
};
pre_commit_events.observe_chunk(chunk_len);
let service_tier_changed = observation_state
.observe_service_tier(decoder.response_service_tier());
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
(
@@ -925,17 +925,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
apply_failure(&failure_context, &active_account, failure)
.await;
}
if let Err(error) = attach_openai_session_update(&mut events, &mut session_capture) {
if terminal_failure.is_none() {
let failure = MappedProviderFailure::plain(error);
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
} else {
Err(error)?;
}
return;
}
attach_openai_session_update(&mut events, &mut session_capture);
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 完成事件一旦交给下游,Core 可以立刻停止轮询 Provider stream;
// 在此之前持久化亲和关系,保证成功请求不会因流被提前 drop 而丢失绑定。
@@ -946,8 +936,13 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if rate_limits_changed || timing_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
if (rate_limits_changed
|| service_tier_changed
|| timing_changed
|| terminal_changed)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
}
if let Some((mut failure, atomic_upstream_failure)) = terminal_failure {
let failure_after_commit =
@@ -964,7 +959,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(failure.error)?;
return;
}
let events = pre_commit_events.stage(events, timing_signals, completed)?;
let events = pre_commit_events.stage(events, timing_signals, completed);
for event in events {
yield event;
}
@@ -995,6 +990,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
)
});
let timing_signals = decoder.take_timing_signals();
let service_tier_changed =
observation_state.observe_service_tier(decoder.response_service_tier());
let timing_changed = observation_state
.observe_timing_signals(timing_signals, context.timing_started_at());
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
@@ -1016,17 +1013,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
{
capture.turn_state = Some(turn_state);
}
if let Err(error) = attach_openai_session_update(&mut events, &mut session_capture) {
if terminal_failure.is_none() {
let failure = MappedProviderFailure::plain(error);
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
} else {
Err(error)?;
}
return;
}
attach_openai_session_update(&mut events, &mut session_capture);
let completed = events
.iter()
.flat_map(ProviderEvent::canonical_facts)
@@ -1042,8 +1029,10 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if timing_changed || rate_limits_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
if (service_tier_changed || timing_changed || rate_limits_changed || terminal_changed)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
}
if let Some((mut failure, atomic_upstream_failure)) = terminal_failure {
let failure_after_commit =
@@ -1060,7 +1049,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(failure.error)?;
return;
}
let events = pre_commit_events.finish(events, timing_signals, completed)?;
let events = pre_commit_events.finish(events, timing_signals, completed);
for event in events {
yield event;
}
@@ -1113,8 +1102,9 @@ impl OpenAiResponseObservationState {
requested_model: request.model().to_owned(),
stream: request.stream(),
compact: semantics.compact,
service_tier: request.service_tier().map(str::to_owned),
rate_limit_headers: safe_observation_headers(&response.rate_limit_headers),
// 请求档位保留在 requestSummary;一等观测只接受响应确认的实际档位。
service_tier: None,
rate_limit_headers: selected_observation_headers(&response.rate_limit_headers),
account_id: account_id.to_owned(),
attempt_index,
timings: openai_response_timings(
@@ -1125,7 +1115,7 @@ impl OpenAiResponseObservationState {
}
}
fn observation(&self) -> Result<ProviderResponseObservation, ProviderError> {
fn observation(&self) -> Option<ProviderResponseObservation> {
let mut observation = codex_response_observation(
self.transport,
&self.diagnostics,
@@ -1137,7 +1127,10 @@ impl OpenAiResponseObservationState {
if let Some(metadata) = self.provider_metadata() {
observation = observation.with_provider_metadata(metadata);
}
Ok(observation)
if let Some(service_tier) = self.service_tier.as_deref() {
observation = observation.with_service_tier_if_valid(service_tier.to_owned());
}
Some(observation)
}
fn observe_stream_chunk(&mut self, chunk: &[u8], started_at: Instant) -> bool {
@@ -1165,9 +1158,20 @@ impl OpenAiResponseObservationState {
changed
}
fn observe_service_tier(&mut self, service_tier: Option<&str>) -> bool {
let Some(service_tier) = service_tier else {
return false;
};
if self.service_tier.as_deref() == Some(service_tier) {
return false;
}
self.service_tier = Some(service_tier.to_owned());
true
}
fn merge_rate_limit_headers(&mut self, updates: &[(String, String)]) -> bool {
let mut changed = false;
for (name, value) in safe_observation_headers(updates) {
for (name, value) in selected_observation_headers(updates) {
if let Some(existing) = self
.rate_limit_headers
.iter_mut()
@@ -1236,7 +1240,9 @@ impl OpenAiResponseObservationState {
);
metadata.insert(
"upstreamTraceHeaders".to_owned(),
json!(safe_observation_headers(&self.diagnostics.trace_headers)),
json!(selected_observation_headers(
&self.diagnostics.trace_headers
)),
);
metadata.insert("requestSummary".to_owned(), self.request_summary.clone());
if let Some(service_tier) = &self.service_tier {
@@ -1268,7 +1274,7 @@ impl OpenAiResponseObservationState {
.trace_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("cf-ray"))
&& let Some((_, cf_ray)) = safe_observation_header("cf-ray", cf_ray)
&& let Some((_, cf_ray)) = selected_observation_header("cf-ray", cf_ray)
{
metadata.insert("cfRay".to_owned(), Value::String(cf_ray));
}
@@ -1313,10 +1319,9 @@ fn codex_response_observation(
metrics: &CodexTransportMetrics,
websocket_pool_decision: Option<crate::transport::WebSocketPoolDecision>,
timings: ProviderResponseTimings,
) -> Result<ProviderResponseObservation, ProviderError> {
) -> Option<ProviderResponseObservation> {
let mut observation = ProviderResponseObservation::new(
UpstreamTransport::new(actual_transport_name(transport))
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent))?,
UpstreamTransport::new(actual_transport_name(transport)).ok()?,
)
.with_timings(timings);
if let Some(version) = metrics
@@ -1336,23 +1341,16 @@ fn codex_response_observation(
if let Some(status_code) = diagnostics.status_code {
observation = observation.with_status_code(status_code);
}
if let Some(request_id) = diagnostics
.request_id
.as_deref()
.and_then(|request_id| SafeUpstreamValue::new(request_id.to_owned()).ok())
{
observation = observation.with_request_id(request_id);
if let Some(request_id) = diagnostics.request_id.as_deref() {
observation = observation.with_request_id(OpaqueUpstreamValue::new(request_id.to_owned()));
}
let client_headers = response_metadata
.client_headers
.iter()
.filter_map(|(name, value)| {
let value = SafeUpstreamValue::new(value.to_owned()).ok()?;
ProviderResponseHeader::new(name.to_owned(), value)
})
.map(|(name, value)| ProviderResponseHeader::new(name.to_owned(), value.clone()))
.collect();
observation = observation.with_client_headers(client_headers);
Ok(observation)
Some(observation)
}
fn codex_error_observation(error: &CodexClientError) -> Option<ProviderResponseObservation> {
@@ -1375,19 +1373,14 @@ fn codex_error_observation(error: &CodexClientError) -> Option<ProviderResponseO
transport_metrics,
None,
openai_response_timings(transport_metrics, &CodexResponseMetadata::default()),
)
.ok()?
)?
.with_status_code(status.as_u16());
}
CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => {
observation = observation.with_status_code(upstream.status_code);
if let Some(request_id) = upstream
.diagnostics
.request_id
.as_deref()
.and_then(|value| SafeUpstreamValue::new(value.to_owned()).ok())
{
observation = observation.with_request_id(request_id);
if let Some(request_id) = upstream.diagnostics.request_id.as_deref() {
observation =
observation.with_request_id(OpaqueUpstreamValue::new(request_id.to_owned()));
}
}
_ => {}
@@ -1414,7 +1407,8 @@ fn openai_processing_ms(response_metadata: &CodexResponseMetadata) -> Option<u64
.client_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("openai-processing-ms"))
.and_then(|(_, value)| value.parse::<u64>().ok())
.and_then(|(_, value)| std::str::from_utf8(value).ok())
.and_then(|value| value.trim().parse::<u64>().ok())
}
fn openai_response_request_summary(
@@ -1460,19 +1454,16 @@ const fn json_value_kind(value: Option<&Value>) -> &'static str {
}
}
fn safe_observation_headers(headers: &[(String, String)]) -> Vec<(String, String)> {
fn selected_observation_headers(headers: &[(String, String)]) -> Vec<(String, String)> {
headers
.iter()
.filter_map(|(name, value)| safe_observation_header(name, value))
.filter_map(|(name, value)| selected_observation_header(name, value))
.collect()
}
fn safe_observation_header(name: &str, value: &str) -> Option<(String, String)> {
fn selected_observation_header(name: &str, value: &str) -> Option<(String, String)> {
let name = name.trim().to_ascii_lowercase();
let value = value.trim();
let safe_value = SafeUpstreamValue::new(value.to_owned()).ok()?;
ProviderResponseHeader::new(name.clone(), safe_value)?;
Some((name, value.to_owned()))
(!name.is_empty()).then(|| (name, value.to_owned()))
}
fn insert_first_timing(target: &mut Option<u64>, started_at: Instant) -> bool {
@@ -1650,6 +1641,7 @@ fn codex_request_context<'a>(
installation_id: &'a str,
authorization: &'a SecretString,
cookie_header: Option<&'a SecretString>,
account_selection: CodexAccountSelectionTelemetry<'a>,
) -> CodexRequestContext<'a> {
CodexRequestContext {
authorization: authorization.expose_secret(),
@@ -1668,6 +1660,7 @@ fn codex_request_context<'a>(
thread_id: request.client_thread_id.as_deref(),
client_request_id: request.client_request_id.as_deref(),
turn_id: request.client_turn_id.as_deref(),
account_selection,
}
}
@@ -1816,8 +1809,9 @@ impl StreamCommitPolicy {
/// 提交边界前的上游事件预取。
///
/// 原始 chunk 计数而不是重编码后的 event 大小,保持 2.x 的 64 KiB 资源边界;一旦
/// 提交,后续事件不再具备无痕重放资格。
/// 原始 chunk 计数而不是重编码后的 event 大小。64 KiB 只限定无感换号窗口;达到
/// 阈值会提交已缓存 wire,绝不能因为网关私有资源规则伪造上游协议失败。一旦提交,
/// 后续事件不再具备无痕重放资格。
struct PreCommitClientEvents {
policy: StreamCommitPolicy,
pending: Vec<ProviderEvent>,
@@ -1846,21 +1840,18 @@ impl PreCommitClientEvents {
incoming: Vec<ProviderEvent>,
timing_signals: ResponseEventSignals,
completed: bool,
) -> Result<Vec<ProviderEvent>, ProviderError> {
) -> Vec<ProviderEvent> {
if self.committed {
return Ok(incoming);
return incoming;
}
let boundary_reached = self
.policy
.boundary_reached(&incoming, timing_signals, completed);
self.pending.extend(incoming);
if boundary_reached {
return Ok(self.commit_pending());
if boundary_reached || self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES {
return self.commit_pending();
}
if self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES {
return Err(prefetch_protocol_error());
}
Ok(Vec::new())
Vec::new()
}
fn finish(
@@ -1868,12 +1859,12 @@ impl PreCommitClientEvents {
incoming: Vec<ProviderEvent>,
timing_signals: ResponseEventSignals,
completed: bool,
) -> Result<Vec<ProviderEvent>, ProviderError> {
let events = self.stage(incoming, timing_signals, completed)?;
) -> Vec<ProviderEvent> {
let events = self.stage(incoming, timing_signals, completed);
if self.committed {
Ok(events)
events
} else {
Err(prefetch_protocol_error())
self.commit_pending()
}
}
@@ -1902,10 +1893,6 @@ impl PreCommitClientEvents {
}
}
fn prefetch_protocol_error() -> ProviderError {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
}
#[derive(Clone, Copy)]
enum ReplayBoundary {
BeforeSemanticOutput,
@@ -1991,7 +1978,29 @@ fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure {
}
fn map_stream_error(error: CodexClientError) -> MappedProviderFailure {
map_client_error(error, UpstreamSendState::Sent, false)
let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error);
let mut failure = map_client_error(error, UpstreamSendState::Sent, false);
if allows_pre_delivery_retry {
failure.error = failure.error.with_pre_delivery_retry();
}
failure
}
fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool {
match error {
CodexClientError::Http(_) | CodexClientError::StreamIdleTimeout { .. } => true,
CodexClientError::WebSocket(error) => matches!(
error,
CodexWebSocketExchangeError::Transport(_)
| CodexWebSocketExchangeError::PostSendAmbiguous { .. }
| CodexWebSocketExchangeError::SendTimeout { .. }
| CodexWebSocketExchangeError::ClosedBeforeTerminal
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
| CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. }
| CodexWebSocketExchangeError::InitialEventTimeout { .. }
),
_ => false,
}
}
fn map_canonical_error(
@@ -2088,7 +2097,7 @@ fn map_client_error(
}
fn map_upstream_failure(
failure: CodexUpstreamFailure,
mut failure: CodexUpstreamFailure,
observation: Option<ProviderResponseObservation>,
replay_boundary: ReplayBoundary,
) -> MappedProviderFailure {
@@ -2103,14 +2112,29 @@ fn map_upstream_failure(
.map(|_| ContinuationFailure::HistoryUnavailable);
let send_state = upstream_send_state(failure.send_phase);
let mut error = provider_error(provider_error_kind(category), send_state);
if let Some(message) = failure.client_message.as_ref()
&& let Ok(client_error) = ClientVisibleUpstreamError::new(
if let Some(message) = failure.client_message.as_ref() {
error = error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
message.clone(),
failure.client_code.clone(),
failure.client_error_type.clone(),
)
{
error = error.with_client_visible_upstream_error(client_error);
));
}
if let Some(response) = failure.client_response.take() {
let response = (*response).into_parts();
error = error.with_client_visible_upstream_response(
ClientVisibleUpstreamResponse::new(
response.status,
response.content_type,
response.body,
)
.with_headers(
response
.client_headers
.into_iter()
.map(|(name, value)| ProviderResponseHeader::new(name, value))
.collect(),
),
);
}
if let Some(status) = failure.status {
error = error.with_status(status.as_u16());
@@ -2126,18 +2150,11 @@ fn map_upstream_failure(
if let Some(retry_after) = failure.retry_after_seconds.map(Duration::from_secs) {
error = error.with_retry_after(retry_after);
}
if let Some(code) = failure
.persistable_code()
.and_then(|code| SafeUpstreamValue::new(code.to_owned()).ok())
{
error = error.with_upstream_code(code);
if let Some(code) = failure.persistable_code() {
error = error.with_upstream_code(OpaqueUpstreamValue::new(code.to_owned()));
}
if let Some(request_id) = failure
.request_id
.as_deref()
.and_then(|request_id| SafeUpstreamValue::new(request_id.to_owned()).ok())
{
error = error.with_upstream_request_id(request_id);
if let Some(request_id) = failure.request_id.as_deref() {
error = error.with_upstream_request_id(OpaqueUpstreamValue::new(request_id.to_owned()));
}
MappedProviderFailure {
error,
@@ -24,6 +24,7 @@ use super::protocol::responses::{
use super::usage::openai_billing_breakdown;
const CONTENTS_PER_OUTPUT: u32 = 1_024;
const MAX_SERVICE_TIER_BYTES: usize = 64;
/// 单 attempt 的增量 Responses decoder。
///
@@ -41,6 +42,7 @@ pub struct CodexCanonicalDecoder {
reasoning_output_seen: BTreeSet<u32>,
usage_emitted: bool,
semantic_output_seen: bool,
response_service_tier: Option<String>,
timing_signals: ResponseEventSignals,
raw_sse_passthrough: bool,
}
@@ -135,6 +137,7 @@ impl CodexCanonicalDecoder {
reasoning_output_seen: BTreeSet::new(),
usage_emitted: false,
semantic_output_seen: false,
response_service_tier: None,
timing_signals: ResponseEventSignals::default(),
raw_sse_passthrough: false,
}
@@ -152,15 +155,7 @@ impl CodexCanonicalDecoder {
pub fn push(&mut self, chunk: &[u8]) -> CodexCanonicalOutcome {
self.timing_signals = ResponseEventSignals::default();
if self.raw_sse_passthrough {
let frames = match self.decoder.push_frames(chunk) {
Ok(frames) => frames,
Err(error) => {
return self.failure(
Vec::new(),
CodexCanonicalError::Protocol(protocol_error(error)),
);
}
};
let frames = self.decoder.push_frames(chunk);
return self.decode_frames(frames);
}
let events = match self.decoder.push(chunk) {
@@ -178,15 +173,7 @@ impl CodexCanonicalDecoder {
pub fn finish(&mut self) -> CodexCanonicalOutcome {
self.timing_signals = ResponseEventSignals::default();
if self.raw_sse_passthrough {
let frames = match self.decoder.finish_frames() {
Ok(frames) => frames,
Err(error) => {
return self.failure(
Vec::new(),
CodexCanonicalError::Protocol(protocol_error(error)),
);
}
};
let frames = self.decoder.finish_frames();
return self.decode_frames(frames);
}
let events = match self.decoder.finish() {
@@ -210,6 +197,12 @@ impl CodexCanonicalDecoder {
std::mem::take(&mut self.timing_signals)
}
/// 返回本 attempt 已从上游响应生命周期帧观察到的实际服务档位。
#[must_use]
pub fn response_service_tier(&self) -> Option<&str> {
self.response_service_tier.as_deref()
}
fn decode(&mut self, events: Vec<SseEvent>) -> CodexCanonicalOutcome {
let mut output = Vec::new();
for event in events {
@@ -269,6 +262,7 @@ impl CodexCanonicalDecoder {
.event
.as_deref()
.or_else(|| value.get("type").and_then(Value::as_str));
self.observe_response_service_tier(&value);
let signals = response_event_signals(event_type, &value);
self.merge_timing_signals(signals);
if matches!(event_type, Some("response.failed" | "error")) {
@@ -313,29 +307,33 @@ impl CodexCanonicalDecoder {
self.timing_signals.text_output |= signals.text_output;
}
fn observe_response_service_tier(&mut self, value: &Value) {
let Some(service_tier) = response_object(value)
.and_then(|response| response.get("service_tier"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| {
!value.is_empty()
&& value.len() <= MAX_SERVICE_TIER_BYTES
&& !value.chars().any(char::is_control)
})
else {
return;
};
self.response_service_tier = Some(service_tier.to_ascii_lowercase());
}
/// 构造下发的 wire event。
///
/// 上游事件名或 SSE id 不满足 wire 安全约束时剥离该元数据继续下发,
/// 而不是把整帧静默丢弃;raw 帧存在时原始字节仍原样透传,JSON 正文
/// 继续供旁路观测使用。
/// raw 帧存在时原始字节原样透传;解析出的 event/id 只作旁路元数据,不能
/// 反过来决定客户端事件是否可交付。
fn wire_for_event(
event: SseEvent,
value: Value,
raw_sse_frame: Option<Bytes>,
) -> Option<ProtocolWireEvent> {
let event_type = event.event.filter(|name| {
ProtocolWireEvent::json("openai", Some(name.clone()), Value::Null).is_ok()
});
let sse_id = event.id.filter(|id| {
ProtocolWireEvent::json_with_sse_metadata(
"openai",
None,
Value::Null,
Some(id.clone()),
None,
)
.is_ok()
});
let event_type = event.event;
let sse_id = event.id;
match raw_sse_frame {
Some(raw_sse_frame) => ProtocolWireEvent::json_with_raw_sse_metadata(
"openai",
@@ -424,7 +422,7 @@ impl CodexCanonicalDecoder {
return Ok(());
}
let response = response_object(value).ok_or_else(protocol_error_marker)?;
let response_id = required_text(response, "id")?;
let response_id = required_string(response, "id")?;
let model = response
.get("model")
.and_then(Value::as_str)
@@ -810,7 +808,7 @@ impl CodexCanonicalDecoder {
self.completed = true;
return Ok(());
};
let Ok(response_id) = required_text(response, "id") else {
let Ok(response_id) = required_string(response, "id") else {
self.completed = true;
return Ok(());
};
@@ -838,6 +836,10 @@ impl CodexCanonicalDecoder {
.filter(|model| !model.is_empty())
.unwrap_or(&self.fallback_model)
.to_owned();
let service_tier = response
.get("service_tier")
.and_then(Value::as_str)
.or(self.response_service_tier.as_deref());
if let Some(breakdown) = usage
.filter(|usage| billable_usage_is_complete(response, *usage))
.and_then(|usage| {
@@ -847,7 +849,7 @@ impl CodexCanonicalDecoder {
usage.output_tokens,
usage.cached_tokens,
usage.cache_write_tokens,
response.get("service_tier").and_then(Value::as_str),
service_tier,
)
})
{
@@ -917,6 +919,14 @@ fn required_text(value: &Value, field: &str) -> Result<String, ProviderError> {
.ok_or_else(protocol_error_marker)
}
fn required_string(value: &Value, field: &str) -> Result<String, ProviderError> {
value
.get(field)
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.ok_or_else(protocol_error_marker)
}
fn optional_event_index(value: &Value, field: &str) -> Result<Option<u32>, ProviderError> {
value
.get(field)
@@ -38,7 +38,6 @@ use super::websocket::{
// 常量
// ---------------------------------------------------------------------------
const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 1024 * 1024;
const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
pub(super) const UPSTREAM_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
@@ -85,6 +84,78 @@ pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
// 错误类型
// ---------------------------------------------------------------------------
/// 当前 Responses 请求最终失败时可原样返回给客户端的上游 HTTP 响应。
#[derive(Clone, PartialEq, Eq)]
pub struct CodexClientVisibleUpstreamResponse {
status: u16,
content_type: Option<Vec<u8>>,
client_headers: Vec<(String, Bytes)>,
body: Bytes,
}
pub(crate) struct CodexClientVisibleUpstreamResponseParts {
pub(crate) status: u16,
pub(crate) content_type: Option<Vec<u8>>,
pub(crate) client_headers: Vec<(String, Bytes)>,
pub(crate) body: Bytes,
}
impl CodexClientVisibleUpstreamResponse {
pub(super) fn new(
status: StatusCode,
content_type: Option<Vec<u8>>,
client_headers: Vec<(String, Bytes)>,
body: Bytes,
) -> Self {
Self {
status: status.as_u16(),
content_type,
client_headers,
body,
}
}
pub const fn status(&self) -> u16 {
self.status
}
pub fn content_type(&self) -> Option<&[u8]> {
self.content_type.as_deref()
}
pub fn client_headers(&self) -> &[(String, Bytes)] {
&self.client_headers
}
pub const fn body(&self) -> &Bytes {
&self.body
}
pub(crate) fn into_parts(self) -> CodexClientVisibleUpstreamResponseParts {
CodexClientVisibleUpstreamResponseParts {
status: self.status,
content_type: self.content_type,
client_headers: self.client_headers,
body: self.body,
}
}
}
impl fmt::Debug for CodexClientVisibleUpstreamResponse {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CodexClientVisibleUpstreamResponse")
.field("status", &self.status)
.field(
"content_type",
&self.content_type.as_ref().map(|_| "<present>"),
)
.field("client_header_count", &self.client_headers.len())
.field("body", &"<redacted>")
.finish()
}
}
/// Codex 上游 HTTP 客户端错误。
#[derive(Error)]
pub enum CodexClientError {
@@ -125,6 +196,8 @@ pub enum CodexClientError {
status: StatusCode,
/// 上游错误体。
body: String,
/// 仅供当前 Responses 请求最终失败时原样返回的响应。
client_response: Option<Box<CodexClientVisibleUpstreamResponse>>,
/// 推导出的重试秒数。
retry_after_seconds: Option<u64>,
/// 上游诊断元数据。
@@ -208,6 +281,7 @@ impl CodexClientError {
body,
retry_after_seconds,
diagnostics,
client_response,
set_cookie_headers,
rate_limit_headers,
send_phase,
@@ -217,6 +291,7 @@ impl CodexClientError {
body,
*retry_after_seconds,
diagnostics,
client_response.as_deref(),
set_cookie_headers,
rate_limit_headers,
*send_phase,
@@ -227,6 +302,7 @@ impl CodexClientError {
&upstream.body,
upstream.retry_after_seconds,
&upstream.diagnostics,
upstream.client_response.as_deref(),
&upstream.set_cookie_headers,
&[],
upstream.send_phase,
@@ -248,6 +324,50 @@ pub type CodexBackendSseStream =
// 请求上下文与响应类型
// ---------------------------------------------------------------------------
/// 账号亲和决策传给 transport 的不可变遥测快照。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CodexAccountSelectionTelemetry<'a> {
affinity_hit: bool,
escape_reason: Option<&'a str>,
account_switch: bool,
}
impl<'a> CodexAccountSelectionTelemetry<'a> {
pub const NONE: Self = Self {
affinity_hit: false,
escape_reason: None,
account_switch: false,
};
#[must_use]
pub const fn new(
affinity_hit: bool,
escape_reason: Option<&'a str>,
account_switch: bool,
) -> Self {
Self {
affinity_hit,
escape_reason,
account_switch,
}
}
#[must_use]
pub const fn affinity_hit(self) -> bool {
self.affinity_hit
}
#[must_use]
pub const fn escape_reason(self) -> Option<&'a str> {
self.escape_reason
}
#[must_use]
pub const fn account_switch(self) -> bool {
self.account_switch
}
}
/// 单次 Codex 上游请求的上下文。
#[derive(Clone, Copy)]
pub struct CodexRequestContext<'a> {
@@ -283,6 +403,8 @@ pub struct CodexRequestContext<'a> {
pub client_request_id: Option<&'a str>,
/// 客户端 turn ID。
pub turn_id: Option<&'a str>,
/// 账号亲和选择结果;仅用于结构化遥测,不参与 transport 决策。
pub account_selection: CodexAccountSelectionTelemetry<'a>,
}
impl<'a> CodexRequestContext<'a> {
@@ -310,6 +432,7 @@ impl<'a> CodexRequestContext<'a> {
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: CodexAccountSelectionTelemetry::NONE,
}
}
}
@@ -328,6 +451,7 @@ impl fmt::Debug for CodexRequestContext<'_> {
"installation_id",
&self.installation_id.map(|_| "[PSEUDONYM]"),
)
.field("account_selection", &self.account_selection)
.finish_non_exhaustive()
}
}
@@ -453,20 +577,23 @@ pub(super) struct PreparedWebSocketRoute {
// ---------------------------------------------------------------------------
pub(super) fn log_websocket_pool_decision(
request_id: &str,
account_id: Option<&str>,
context: CodexRequestContext<'_>,
pool_account_id: Option<&str>,
pool_context: Option<&WebSocketPoolLogContext>,
decision: Option<WebSocketPoolDecision>,
) {
let Some(decision) = decision else {
return;
};
let rid_short = request_id.chars().take(8).collect::<String>();
let rid_short = context.request_id.chars().take(8).collect::<String>();
tracing::info!(
request_id = %request_id,
request_id = %context.request_id,
rid = %rid_short,
account_id = account_id.unwrap_or_default(),
account_id = pool_account_id.or(context.account_id).unwrap_or_default(),
ws_pool = decision.kind(),
affinity_hit = context.account_selection.affinity_hit(),
escape_reason = context.account_selection.escape_reason().unwrap_or_default(),
account_switch = context.account_selection.account_switch(),
conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()),
ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()),
"WebSocket pool decision"
@@ -576,14 +703,10 @@ pub(super) async fn read_capped_response_body(
})
}
pub(super) async fn read_capped_error_body(
pub(super) async fn read_error_response_body(
response: ReqwestResponse,
) -> Result<String, reqwest::Error> {
let body = read_capped_response_body(response, MAX_UPSTREAM_ERROR_BODY_BYTES).await?;
if body.limit_exceeded() {
return Ok("upstream error response exceeded the body limit".to_owned());
}
Ok(body.into_string())
) -> Result<Bytes, reqwest::Error> {
response.bytes().await
}
// ---------------------------------------------------------------------------
@@ -633,19 +756,13 @@ fn now_unix_timestamp_millis() -> u128 {
}
pub(super) fn openai_subagent_from_metadata(client_metadata: Option<&Value>) -> Option<String> {
let value = client_metadata?
.as_object()?
.get("x-openai-subagent")?
.as_str()?
.trim();
if matches!(
value,
"review" | "compact" | "memory_consolidation" | "collab_spawn"
) {
Some(value.to_string())
} else {
None
}
Some(
client_metadata?
.as_object()?
.get("x-openai-subagent")?
.as_str()?
.to_owned(),
)
}
// ---------------------------------------------------------------------------
@@ -662,6 +779,7 @@ pub(super) fn websocket_exchange_error_to_client_error(
status: StatusCode::from_u16(upstream.status_code)
.unwrap_or(StatusCode::BAD_GATEWAY),
body: upstream.body,
client_response: upstream.client_response,
retry_after_seconds: upstream.retry_after_seconds,
diagnostics: Box::new(upstream.diagnostics),
set_cookie_headers: upstream.set_cookie_headers,
@@ -20,7 +20,10 @@ use crate::transport::{
},
diagnostics::CodexUpstreamSendPhase,
endpoints::{CODEX_RESPONSES_PATH, endpoint_url},
headers::{build_codex_base_headers, insert_optional_header, websocket_header_pairs},
headers::{
build_codex_base_headers, insert_optional_header, insert_optional_protocol_header,
websocket_header_pairs,
},
profile::{CodexWireProfile, CodexWireProfileState},
protocol::{
responses::{CodexResponsesRequest, TransportRequirement, transport_requirement},
@@ -95,12 +98,25 @@ impl CodexBackendClient {
let retry_after_seconds = retry_after_seconds(response.headers(), None);
if !status.is_success() {
let body = read_capped_error_body(response).await?;
let content_type = response
.headers()
.get(CONTENT_TYPE)
.map(|value| value.as_bytes().to_vec());
let client_headers = response_meta::client_headers(response.headers());
let raw_body = read_error_response_body(response).await?;
let body = String::from_utf8_lossy(&raw_body).into_owned();
let retry_after_seconds =
retry_after_seconds.or_else(|| retry_after_seconds_from_body(&body));
return Err(CodexClientError::Upstream {
status,
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
client_response: Some(Box::new(CodexClientVisibleUpstreamResponse::new(
status,
content_type,
client_headers,
raw_body,
))),
retry_after_seconds,
diagnostics: Box::new(diagnostics),
set_cookie_headers,
rate_limit_headers,
@@ -254,8 +270,8 @@ impl CodexBackendClient {
http_version: Some("HTTP/1.1".to_string()),
};
log_websocket_pool_decision(
context.request_id,
pool_account_id.or(context.account_id),
context,
pool_account_id,
pool_log_context.as_ref(),
prepared.pool_decision(),
);
@@ -393,6 +409,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
client_response: None,
diagnostics: Box::new(diagnostics),
set_cookie_headers,
rate_limit_headers: Vec::new(),
@@ -412,21 +429,25 @@ impl CodexBackendClient {
let profile = self.profile.snapshot();
let mut headers = self.request_headers(&profile, context)?;
if let Some(subagent) = openai_subagent_from_metadata(request.client_metadata()) {
headers.insert(
HeaderName::from_static("x-openai-subagent"),
HeaderValue::from_str(&subagent)?,
);
insert_optional_protocol_header(&mut headers, "x-openai-subagent", Some(&subagent));
}
insert_optional_header(
insert_optional_protocol_header(
&mut headers,
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER,
request.responses_lite.as_deref(),
)?;
insert_optional_header(
);
insert_optional_protocol_header(
&mut headers,
X_OPENAI_MEMGEN_REQUEST_HEADER,
request.memgen_request.as_deref(),
)?;
);
// 客户端确实携带的普通协议头优先;没有携带时才使用上面的 Codex profile 默认值。
for name in request.passthrough_headers.keys() {
headers.remove(name);
for value in request.passthrough_headers.get_all(name) {
headers.append(name.clone(), value.clone());
}
}
Ok(headers)
}
@@ -445,7 +466,6 @@ impl CodexBackendClient {
profile: &CodexWireProfile,
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let request_id = context.client_request_id.unwrap_or(context.request_id);
let mut headers =
build_codex_base_headers(profile, context.authorization, context.account_id)?;
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
@@ -461,31 +481,44 @@ impl CodexBackendClient {
);
headers.insert(
HeaderName::from_static("x-client-request-id"),
HeaderValue::from_str(request_id)?,
HeaderValue::from_str(context.request_id)?,
);
insert_optional_protocol_header(
&mut headers,
"x-client-request-id",
context.client_request_id,
);
insert_optional_header(
&mut headers,
"x-codex-installation-id",
context.installation_id,
)?;
insert_optional_header(&mut headers, "session-id", context.session_id)?;
insert_optional_header(&mut headers, "thread-id", context.thread_id)?;
insert_optional_header(&mut headers, "x-codex-turn-id", context.turn_id)?;
insert_optional_header(&mut headers, "x-codex-window-id", context.codex_window_id)?;
insert_optional_header(&mut headers, "x-codex-turn-state", context.turn_state)?;
insert_optional_header(&mut headers, "x-codex-turn-metadata", context.turn_metadata)?;
insert_optional_header(&mut headers, "x-codex-beta-features", context.beta_features)?;
insert_optional_header(
insert_optional_protocol_header(&mut headers, "session-id", context.session_id);
insert_optional_protocol_header(&mut headers, "thread-id", context.thread_id);
insert_optional_protocol_header(&mut headers, "x-codex-turn-id", context.turn_id);
insert_optional_protocol_header(&mut headers, "x-codex-window-id", context.codex_window_id);
insert_optional_protocol_header(&mut headers, "x-codex-turn-state", context.turn_state);
insert_optional_protocol_header(
&mut headers,
"x-codex-turn-metadata",
context.turn_metadata,
);
insert_optional_protocol_header(
&mut headers,
"x-codex-beta-features",
context.beta_features,
);
insert_optional_protocol_header(
&mut headers,
"x-responsesapi-include-timing-metrics",
context.include_timing_metrics,
)?;
insert_optional_header(&mut headers, "version", context.version)?;
insert_optional_header(
);
insert_optional_protocol_header(&mut headers, "version", context.version);
insert_optional_protocol_header(
&mut headers,
"x-codex-parent-thread-id",
context.parent_thread_id,
)?;
);
Ok(headers)
}
@@ -1,5 +1,6 @@
//! 在 transport 边界采集的上游响应诊断信息。
use super::client::CodexClientVisibleUpstreamResponse;
use super::protocol::responses::ResponsesSseFailure;
use base64::{Engine as _, engine::general_purpose::STANDARD};
use reqwest::{StatusCode, header::HeaderMap};
@@ -107,6 +108,7 @@ pub struct CodexUpstreamFailure {
pub(crate) client_message: Option<String>,
pub(crate) client_code: Option<String>,
pub(crate) client_error_type: Option<String>,
pub(crate) client_response: Option<Box<CodexClientVisibleUpstreamResponse>>,
pub(crate) identity_error_code: Option<String>,
pub(crate) retry_after_seconds: Option<u64>,
pub(crate) request_id: Option<String>,
@@ -118,11 +120,16 @@ pub struct CodexUpstreamFailure {
impl CodexUpstreamFailure {
/// 用完整 HTTP 响应事实构造并分类一次上游拒绝。
#[expect(
clippy::too_many_arguments,
reason = "the classifier consumes one flat set of transport response facts"
)]
pub fn from_response(
status: StatusCode,
body: &str,
retry_after_seconds: Option<u64>,
diagnostics: &CodexUpstreamDiagnostics,
client_response: Option<&CodexClientVisibleUpstreamResponse>,
set_cookie_headers: &[String],
rate_limit_headers: &[(String, String)],
send_phase: CodexUpstreamSendPhase,
@@ -141,6 +148,7 @@ impl CodexUpstreamFailure {
client_message: fields.client_message,
client_code: fields.code,
client_error_type: fields.error_type,
client_response: client_response.cloned().map(Box::new),
identity_error_code: diagnostics.identity_error_code.clone(),
retry_after_seconds,
request_id: diagnostics.request_id.clone(),
@@ -180,6 +188,7 @@ impl CodexUpstreamFailure {
client_message: Some(fields.message),
client_code: fields.code,
client_error_type: fields.error_type,
client_response: None,
identity_error_code: diagnostics.identity_error_code.clone(),
retry_after_seconds: failure.retry_after_seconds,
request_id: diagnostics.request_id.clone(),
@@ -32,23 +32,32 @@ pub(super) fn insert_optional_header(
Ok(())
}
/// 尽力投影客户端协议值;无法表示为 HTTP header 时保留正文并跳过投影。
pub(super) fn insert_optional_protocol_header(
headers: &mut HeaderMap,
name: &'static str,
value: Option<&str>,
) {
let Some(value) = value.and_then(|value| HeaderValue::from_str(value).ok()) else {
return;
};
headers.insert(HeaderName::from_static(name), value);
}
fn header_pairs(headers: &HeaderMap) -> Vec<(String, String)> {
headers
.iter()
// tungstenite 的 opening serializer 只接受可 `to_str()` 的值;逐条跳过
// 无法构造的头,不能让一个扩展头中断业务 payload。
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_string(), value.to_string()))
.map(|value| (name.as_str().to_owned(), value.to_owned()))
})
.collect()
}
pub(super) fn websocket_header_pairs(headers: &HeaderMap) -> Vec<(String, String)> {
header_pairs(headers)
.into_iter()
.filter(|(name, _)| {
!name.eq_ignore_ascii_case("content-type") && !name.eq_ignore_ascii_case("accept")
})
.collect::<Vec<_>>()
}
@@ -25,10 +25,10 @@ pub use self::{
CodexModelCatalogSnapshot, MAX_CODEX_MODEL_CATALOG_BYTES, parse_codex_model_catalog,
},
client::{
CodexBackendClient, CodexBackendSseStream, CodexBackendStreamingResponse,
CodexBackendTransport, CodexClientError, CodexClientResult, CodexRateLimitHeaderUpdates,
CodexRequestContext, CodexTransportDecision, CodexTransportMetrics, CodexTurnStateUpdate,
build_reqwest_client,
CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendSseStream,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexClientResult,
CodexRateLimitHeaderUpdates, CodexRequestContext, CodexTransportDecision,
CodexTransportMetrics, CodexTurnStateUpdate, build_reqwest_client,
},
diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamSendPhase},
endpoints::{
@@ -4,6 +4,7 @@ use gateway_protocol::openai::{
CodexResponsesRequestSemantics as CodexRequestSemantics,
codex_responses_request_semantics_with_turn_metadata, events,
};
use reqwest::header::HeaderMap;
use serde::Serialize;
use serde_json::{Map, Value};
@@ -19,6 +20,8 @@ use serde_json::{Map, Value};
pub struct CodexResponsesRequest {
/// 上游请求体(唯一真相源)。
body: Map<String, Value>,
/// API 边界保存、Provider 逐条恢复的普通客户端请求头。
pub(crate) passthrough_headers: HeaderMap,
/// 是否由客户端显式提供了 prompt cache key。
pub explicit_prompt_cache_key: bool,
/// 客户端会话 ID。
@@ -402,7 +405,6 @@ fn failure_message(value: &Value) -> Option<String> {
.or_else(|| value.pointer("/error/message"))
.or_else(|| value.get("message"))
.and_then(Value::as_str)
.filter(|message| !message.trim().is_empty())
.map(ToString::to_string)
}
@@ -412,7 +414,6 @@ fn failure_code(value: &Value) -> Option<String> {
.or_else(|| value.pointer("/error/code"))
.or_else(|| value.get("code"))
.and_then(Value::as_str)
.filter(|code| !code.trim().is_empty())
.map(ToString::to_string)
}
@@ -420,7 +421,6 @@ fn failure_type(value: &Value) -> Option<String> {
failure_error(value)
.and_then(|error| error.get("type"))
.and_then(Value::as_str)
.filter(|error_type| !error_type.trim().is_empty())
.map(ToString::to_string)
}
@@ -446,6 +446,7 @@ impl CodexResponsesRequest {
pub fn from_body(body: Map<String, Value>) -> Self {
Self {
body,
passthrough_headers: HeaderMap::new(),
explicit_prompt_cache_key: false,
client_conversation_id: None,
client_session_id: None,
@@ -617,7 +618,7 @@ impl CodexResponsesRequest {
if !self.body.contains_key(key) {
return;
}
match value.filter(|value| !value.trim().is_empty()) {
match value {
Some(value) => {
self.body
.insert(key.to_string(), Value::String(value.to_string()));
@@ -90,6 +90,9 @@ pub fn websocket_metadata_headers(value: &Value) -> Vec<(String, String)> {
/// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。
pub fn websocket_metadata_turn_state(value: &Value) -> Option<String> {
if value.get("type").and_then(Value::as_str) != Some("response.metadata") {
return None;
}
websocket_metadata_headers(value)
.into_iter()
.find_map(|(name, value)| {
@@ -123,7 +126,6 @@ pub fn websocket_response_completed_id(value: &Value) -> Option<String> {
value
.pointer("/response/id")
.and_then(Value::as_str)
.filter(|id| !id.trim().is_empty())
.map(ToOwned::to_owned)
}
@@ -1,12 +1,16 @@
//! 核心 Generate operation 到 Codex Responses wire request 的严格编码。
use base64::{Engine as _, engine::general_purpose::STANDARD};
use gateway_core::operation::GenerateRequest;
use gateway_protocol::openai::WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::transport::protocol::responses::CodexResponsesRequest;
const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers";
const CROSS_ACCOUNT_IDENTITY_KEYS: &[&str] = &[
"authorization",
"Authorization",
@@ -40,24 +44,14 @@ const ACCOUNT_BOUND_STATE_KEYS: &[&str] = &[
"turnState",
"turn_state",
"x-codex-turn-state",
"turnMetadata",
"turn_metadata",
"x-codex-turn-metadata",
"previous_response_id",
"previousResponseId",
"response_id",
"responseId",
"conversation",
"conversation_id",
"conversationId",
];
const UNTRUSTED_CONTINUATION_KEYS: &[&str] = &[
"previous_response_id",
"previousResponseId",
"response_id",
"responseId",
];
const TURN_METADATA_KEYS: &[&str] = &["turnMetadata", "turn_metadata", "x-codex-turn-metadata"];
const INSTALLATION_ID_KEYS: &[&str] = &[
"installation_id",
@@ -146,9 +140,7 @@ impl ExtractedRequestContext {
.get("client_metadata")
.and_then(Value::as_object)
.and_then(|metadata| {
non_empty_string(
metadata.get(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY),
)
string_value(metadata.get(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY))
}),
// Memory consolidation 只由官方请求头提供;body 中同名字段不参与
// transport 事实提取。
@@ -158,15 +150,11 @@ impl ExtractedRequestContext {
}
fn body_string(body: &Map<String, Value>, key: &str) -> Option<String> {
non_empty_string(body.get(key))
string_value(body.get(key))
}
fn non_empty_string(value: Option<&Value>) -> Option<String> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
fn string_value(value: Option<&Value>) -> Option<String> {
value.and_then(Value::as_str).map(ToOwned::to_owned)
}
pub(crate) fn derive_conversation_anchor(
@@ -281,10 +269,6 @@ pub(crate) fn scope_request_to_account(
installation_id: &str,
cross_account: bool,
) {
for key in CROSS_ACCOUNT_IDENTITY_KEYS {
request.body_mut().remove(*key);
}
let client_metadata_turn_state = metadata_string(request, "x-codex-turn-state");
let preserve_turn_state =
!cross_account && (request.turn_state.is_some() || client_metadata_turn_state.is_some());
@@ -304,11 +288,13 @@ pub(crate) fn scope_request_to_account(
.and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account));
if cross_account {
request.passthrough_headers.remove("x-codex-turn-state");
request.passthrough_headers.remove("x-codex-turn-metadata");
sanitize_cross_account_input(request);
for key in ACCOUNT_BOUND_STATE_KEYS {
request.body_mut().remove(*key);
}
for key in UNTRUSTED_CONTINUATION_KEYS {
for key in CROSS_ACCOUNT_IDENTITY_KEYS
.iter()
.chain(ACCOUNT_BOUND_STATE_KEYS)
{
request.body_mut().remove(*key);
}
}
@@ -316,13 +302,13 @@ pub(crate) fn scope_request_to_account(
for key in INSTALLATION_ID_KEYS {
request.replace_existing_identity_field(key, Some(installation_id));
}
for (key, value) in [
("turnState", turn_state.as_deref()),
("x-codex-turn-state", turn_state.as_deref()),
("turnMetadata", turn_metadata.as_deref()),
("x-codex-turn-metadata", turn_metadata.as_deref()),
] {
request.replace_existing_identity_field(key, value);
for key in TURN_METADATA_KEYS {
let scoped = request
.body()
.get(*key)
.and_then(Value::as_str)
.and_then(|value| scope_turn_metadata(value, installation_id, cross_account));
replace_existing_body_string(request, key, scoped.as_deref());
}
if let Some(client_metadata) = request.client_metadata().cloned() {
@@ -337,11 +323,11 @@ pub(crate) fn scope_request_to_account(
}),
)
});
for key in CROSS_ACCOUNT_IDENTITY_KEYS {
metadata.remove(*key);
}
if cross_account {
for key in ACCOUNT_BOUND_STATE_KEYS {
for key in CROSS_ACCOUNT_IDENTITY_KEYS
.iter()
.chain(ACCOUNT_BOUND_STATE_KEYS)
{
metadata.remove(*key);
}
}
@@ -369,16 +355,12 @@ pub(crate) fn scope_request_to_account(
"x-codex-turn-metadata",
client_metadata_turn_metadata.as_deref(),
);
if !cross_account {
for (key, value) in scoped_turn_metadata {
replace_existing_metadata_field(&mut metadata, key, value.as_deref());
}
for (key, value) in scoped_turn_metadata {
replace_existing_metadata_field(&mut metadata, key, value.as_deref());
}
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
value if !cross_account => Some(value),
Value::Null => None,
_ => None,
value => Some(value),
};
request.set_client_metadata(scoped);
}
@@ -428,11 +410,12 @@ fn scope_turn_metadata(raw: &str, installation_id: &str, cross_account: bool) ->
return (!cross_account).then(|| raw.to_owned());
};
let mut changed = false;
for key in CROSS_ACCOUNT_IDENTITY_KEYS {
changed |= metadata.remove(*key).is_some();
}
if cross_account {
for key in ACCOUNT_BOUND_STATE_KEYS {
for key in CROSS_ACCOUNT_IDENTITY_KEYS
.iter()
.chain(ACCOUNT_BOUND_STATE_KEYS)
.chain(TURN_METADATA_KEYS)
{
changed |= metadata.remove(*key).is_some();
}
}
@@ -452,24 +435,37 @@ fn scope_turn_metadata(raw: &str, installation_id: &str, cross_account: bool) ->
serde_json::to_string(&metadata).ok()
}
fn replace_existing_body_string(
request: &mut CodexResponsesRequest,
key: &str,
value: Option<&str>,
) {
if request.body().get(key).is_some_and(Value::is_string)
&& let Some(value) = value
{
request
.body_mut()
.insert(key.to_owned(), Value::String(value.to_owned()));
}
}
fn replace_existing_metadata_field(
metadata: &mut Map<String, Value>,
key: &str,
value: Option<&str>,
) {
if metadata.contains_key(key) {
replace_metadata_field(metadata, key, value);
if metadata.contains_key(key)
&& let Some(value) = value
{
metadata.insert(key.to_owned(), Value::String(value.to_owned()));
}
}
fn replace_metadata_field(metadata: &mut Map<String, Value>, key: &str, value: Option<&str>) {
match value.filter(|value| !value.trim().is_empty()) {
Some(value) => {
metadata.insert(key.to_owned(), Value::String(value.to_owned()));
}
None => {
metadata.remove(key);
}
if let Some(value) = value
&& metadata.get(key).is_none_or(Value::is_string)
{
metadata.insert(key.to_owned(), Value::String(value.to_owned()));
}
}
@@ -477,6 +473,7 @@ fn replace_metadata_field(metadata: &mut Map<String, Value>, key: &str, value: O
// 优先;body 顶层别名只在 header 缺失时兜底(例如 WebSocket 帧无法逐轮携带
// 请求头的场景)。
fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map<String, Value>) {
request.passthrough_headers = decode_passthrough_headers(context);
request.turn_state =
context_string(context, "turn_state").or_else(|| request.turn_state.take());
request.turn_metadata =
@@ -519,10 +516,87 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map<Str
}
}
fn decode_passthrough_headers(context: &Map<String, Value>) -> HeaderMap {
let mut headers = HeaderMap::new();
let Some(entries) = context
.get(PASSTHROUGH_HEADERS_CONTEXT_KEY)
.and_then(Value::as_array)
else {
return headers;
};
for entry in entries {
let Some(entry) = entry.as_array().filter(|entry| entry.len() == 2) else {
continue;
};
let Some(name) = entry.first().and_then(Value::as_str) else {
continue;
};
let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else {
continue;
};
if provider_managed_header(name.as_str()) {
continue;
}
let Some(encoded) = entry.get(1).and_then(Value::as_str) else {
continue;
};
let Ok(bytes) = STANDARD.decode(encoded) else {
continue;
};
let Ok(value) = HeaderValue::from_bytes(&bytes) else {
continue;
};
headers.append(name, value);
}
headers
}
fn provider_managed_header(name: &str) -> bool {
name.starts_with("sec-websocket-")
|| name.starts_with("x-grok-")
|| name.starts_with("x-xai-")
|| matches!(
name,
"connection"
| "keep-alive"
| "proxy-connection"
| "proxy-authenticate"
| "proxy-authorization"
| "te"
| "trailer"
| "transfer-encoding"
| "upgrade"
| "host"
| "content-length"
| "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"
| "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"
| "x-codex-installation-id"
)
}
fn context_string(context: &Map<String, Value>, field: &str) -> Option<String> {
context
.get(field)
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.map(ToOwned::to_owned)
}
@@ -1,12 +1,15 @@
//! 上游响应元数据提取辅助。
use std::fmt;
use bytes::Bytes;
use gateway_protocol::openai::events::is_rate_limit_header_name;
use reqwest::header::{HeaderMap, SET_COOKIE};
use super::diagnostics::CodexUpstreamDiagnostics;
/// Codex Responses 上游响应元数据。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[derive(Clone, Default, PartialEq, Eq)]
pub struct CodexResponseMetadata {
/// 上游实际选用的模型。
pub effective_model: Option<String>,
@@ -14,17 +17,21 @@ pub struct CodexResponseMetadata {
pub models_etag: Option<String>,
/// 上游是否声明响应包含 reasoning。
pub reasoning_included: bool,
/// 允许透传给客户端的安全响应头。
pub client_headers: Vec<(String, String)>,
/// 允许交给 Core 的普通响应头;名称和值保持 transport 观察到的顺序与字节。
pub client_headers: Vec<(String, Bytes)>,
}
const CLIENT_RESPONSE_HEADERS: [&str; 5] = [
"x-request-id",
"openai-model",
"x-models-etag",
"x-reasoning-included",
"openai-processing-ms",
];
impl fmt::Debug for CodexResponseMetadata {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CodexResponseMetadata")
.field("effective_model", &self.effective_model)
.field("models_etag", &self.models_etag)
.field("reasoning_included", &self.reasoning_included)
.field("client_header_count", &self.client_headers.len())
.finish()
}
}
pub(super) fn diagnostics(
status_code: Option<u16>,
@@ -62,54 +69,128 @@ pub(super) fn rate_limit_headers(headers: &HeaderMap) -> Vec<(String, String)> {
}
pub(super) fn response_metadata(headers: &HeaderMap) -> CodexResponseMetadata {
response_metadata_from_pairs(
headers
.iter()
.filter_map(|(name, value)| value.to_str().ok().map(|value| (name.as_str(), value))),
)
response_metadata_from_client_headers(client_headers(headers))
}
/// 提取响应中可交给客户端 adapter 的普通头。
///
/// 账号身份、凭据、cookie、逐跳头和已经由 adapter 重建的 framing 在进入 Core 前剔除;
/// 其余名称和值保持 HeaderMap 的多值顺序和原始字节。
pub(super) fn client_headers(headers: &HeaderMap) -> Vec<(String, Bytes)> {
filter_client_headers(headers.iter().map(|(name, value)| {
(
name.as_str().to_owned(),
Bytes::copy_from_slice(value.as_bytes()),
)
}))
}
pub(super) fn merge_response_metadata(
metadata: &mut CodexResponseMetadata,
headers: impl IntoIterator<Item = (String, String)>,
) {
for (name, value) in headers {
apply_response_header(metadata, &name, &value);
for (name, value) in filter_client_headers(
headers
.into_iter()
.map(|(name, value)| (name, Bytes::from(value))),
) {
observe_typed_response_header(metadata, &name, &value);
metadata.client_headers.push((name, value));
}
}
fn response_metadata_from_pairs<'a>(
headers: impl IntoIterator<Item = (&'a str, &'a str)>,
fn response_metadata_from_client_headers(
client_headers: Vec<(String, Bytes)>,
) -> CodexResponseMetadata {
let mut metadata = CodexResponseMetadata::default();
for (name, value) in headers {
apply_response_header(&mut metadata, name, value);
for (name, value) in &client_headers {
observe_typed_response_header(&mut metadata, name, value);
}
metadata.client_headers = client_headers;
metadata
}
fn apply_response_header(metadata: &mut CodexResponseMetadata, name: &str, value: &str) {
let name = name.to_ascii_lowercase();
fn observe_typed_response_header(metadata: &mut CodexResponseMetadata, name: &str, value: &[u8]) {
let Ok(value) = std::str::from_utf8(value) else {
return;
};
let value = value.trim();
if value.is_empty() || !CLIENT_RESPONSE_HEADERS.contains(&name.as_str()) {
if value.is_empty() {
return;
}
match name.as_str() {
match name.trim().to_ascii_lowercase().as_str() {
"openai-model" => metadata.effective_model = Some(value.to_string()),
"x-models-etag" => metadata.models_etag = Some(value.to_string()),
"x-reasoning-included" => metadata.reasoning_included = true,
_ => {}
}
upsert_client_header(&mut metadata.client_headers, (name, value.to_string()));
}
fn upsert_client_header(headers: &mut Vec<(String, String)>, header: (String, String)) {
if let Some(existing) = headers
.iter_mut()
.find(|(name, _)| name.eq_ignore_ascii_case(&header.0))
fn filter_client_headers(
headers: impl IntoIterator<Item = (String, Bytes)>,
) -> Vec<(String, Bytes)> {
let headers = headers.into_iter().collect::<Vec<_>>();
let connection_options = headers
.iter()
.filter(|(name, _)| name.trim().eq_ignore_ascii_case("connection"))
.filter_map(|(_, value)| std::str::from_utf8(value).ok())
.flat_map(|value| value.split(','))
.map(str::trim)
.filter(|name| !name.is_empty())
.map(str::to_ascii_lowercase)
.collect::<Vec<_>>();
headers
.into_iter()
.filter(|(name, _)| client_response_header_is_forwardable(name, &connection_options))
.collect()
}
fn client_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-")
{
*existing = header;
} else {
headers.push(header);
return false;
}
!matches!(
name.as_str(),
// 逐跳和 framing 由下游 adapter 针对实际 JSON/SSE/WebSocket 载体重建。
"connection"
| "keep-alive"
| "proxy-connection"
| "te"
| "trailer"
| "transfer-encoding"
| "upgrade"
| "content-length"
| "content-type"
| "content-encoding"
// 上游账号、凭据和 cookie 不能跨越换号边界。
| "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"
)
}
@@ -347,6 +347,7 @@ impl CodexBackendClient {
},
retry_after_seconds,
body: "upstream usage response exceeded the body limit".to_owned(),
client_response: None,
diagnostics: Box::new(diagnostics),
set_cookie_headers: Vec::new(),
rate_limit_headers: Vec::new(),
@@ -367,6 +368,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
client_response: None,
diagnostics: Box::new(diagnostics),
set_cookie_headers: Vec::new(),
rate_limit_headers: Vec::new(),
@@ -389,6 +391,7 @@ impl CodexBackendClient {
|| "usage endpoint is unavailable".to_string(),
|body| format!("invalid usage response: {}", truncate_for_error(&body)),
),
client_response: None,
diagnostics: Box::new(CodexUpstreamDiagnostics::with_status(
StatusCode::BAD_GATEWAY.as_u16(),
)),
@@ -5,6 +5,7 @@ use std::{fmt, time::Duration};
use gateway_protocol::openai::sse::SseError;
use thiserror::Error;
use crate::transport::client::CodexClientVisibleUpstreamResponse;
use crate::transport::diagnostics::CodexUpstreamDiagnostics;
use crate::transport::diagnostics::CodexUpstreamSendPhase;
@@ -96,7 +97,7 @@ pub enum CodexWebSocketExchangeError {
}
/// WebSocket 上游错误帧载荷。
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Clone, PartialEq, Eq)]
pub struct CodexWebSocketUpstreamError {
/// 上游返回的 HTTP 风格状态码。
pub status_code: u16,
@@ -104,6 +105,8 @@ pub struct CodexWebSocketUpstreamError {
pub retry_after_seconds: Option<u64>,
/// 原始错误帧。
pub body: String,
/// opening 失败时可返回给当前客户端的原始 HTTP 响应。
pub client_response: Option<Box<CodexClientVisibleUpstreamResponse>>,
/// 上游透传的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
/// 上游诊断元数据。
@@ -112,13 +115,27 @@ pub struct CodexWebSocketUpstreamError {
pub send_phase: CodexUpstreamSendPhase,
}
impl fmt::Debug for CodexWebSocketUpstreamError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CodexWebSocketUpstreamError")
.field("status_code", &self.status_code)
.field("retry_after_seconds", &self.retry_after_seconds)
.field("body", &"<redacted>")
.field(
"client_response",
&self.client_response.as_ref().map(|_| "<present>"),
)
.field("set_cookie_headers", &self.set_cookie_headers.len())
.field("diagnostics", &self.diagnostics)
.field("send_phase", &self.send_phase)
.finish()
}
}
impl fmt::Display for CodexWebSocketUpstreamError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"websocket upstream error {}: {}",
self.status_code, self.body
)
write!(f, "websocket upstream returned status {}", self.status_code)
}
}
@@ -127,6 +144,7 @@ impl CodexWebSocketExchangeError {
status_code: u16,
retry_after_seconds: Option<u64>,
body: String,
client_response: Option<Box<CodexClientVisibleUpstreamResponse>>,
set_cookie_headers: Vec<String>,
diagnostics: CodexUpstreamDiagnostics,
send_phase: CodexUpstreamSendPhase,
@@ -135,6 +153,7 @@ impl CodexWebSocketExchangeError {
status_code,
retry_after_seconds,
body,
client_response,
set_cookie_headers,
diagnostics,
send_phase,
@@ -2,6 +2,7 @@
use std::time::Duration;
use bytes::Bytes;
use gateway_protocol::openai::events;
use tokio::time::timeout;
use tokio_tungstenite::{Connector, connect_async_tls_with_config};
@@ -18,8 +19,10 @@ use crate::{
responses::CodexResponsesRequest, websocket::websocket_response_create_payload_text,
},
transport::{
client::parse_retry_after, diagnostics::CodexUpstreamSendPhase,
endpoints::CODEX_RESPONSES_PATH, response_meta, tls,
client::{CodexClientVisibleUpstreamResponse, parse_retry_after},
diagnostics::CodexUpstreamSendPhase,
endpoints::CODEX_RESPONSES_PATH,
response_meta, tls,
},
};
@@ -171,17 +174,26 @@ fn websocket_config() -> WebSocketConfig {
extensions.permessage_deflate = Some(DeflateConfig::default());
let mut config = WebSocketConfig::default();
// 上游 Responses 事件属于 Codex 协议数据,不能沿用 tungstenite 的私有
// 64 MiB message / 16 MiB frame 默认限制提前中断本可继续的响应。
config.max_message_size = None;
config.max_frame_size = None;
config.extensions = extensions;
config
}
fn websocket_opening_error(response: &WsResponse<Option<Vec<u8>>>) -> CodexWebSocketExchangeError {
let status_code = response.status().as_u16();
let body = response
let raw_body = response
.body()
.as_ref()
.map(|body| String::from_utf8_lossy(body).into_owned())
.unwrap_or_default();
.map_or_else(Bytes::new, |body| Bytes::copy_from_slice(body));
let body = String::from_utf8_lossy(&raw_body).into_owned();
let content_type = response
.headers()
.get("content-type")
.map(|value| value.as_bytes().to_vec());
let client_headers = response_meta::client_headers(response.headers());
let retry_after_seconds = response
.headers()
.get("retry-after")
@@ -192,6 +204,16 @@ fn websocket_opening_error(response: &WsResponse<Option<Vec<u8>>>) -> CodexWebSo
status_code,
retry_after_seconds,
body,
reqwest::StatusCode::from_u16(status_code)
.ok()
.map(|status| {
Box::new(CodexClientVisibleUpstreamResponse::new(
status,
content_type,
client_headers,
raw_body,
))
}),
response_meta::set_cookie_headers(response.headers()),
response_meta::diagnostics(Some(status_code), response.headers()),
CodexUpstreamSendPhase::BeforePayload,
@@ -264,6 +264,7 @@ async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing()
let billing = admin
.calculated_billing(&ProviderBillingInput {
upstream_model_id: "gpt-4o".to_owned(),
service_tier: None,
input_tokens: Some(1_000_000),
output_tokens: Some(0),
cached_tokens: Some(0),
@@ -277,6 +278,26 @@ async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing()
.expect("known pricing");
assert_eq!(billing.total_amount.amount.as_str(), "2.5");
assert_eq!(billing.input_price_per_million.amount.as_str(), "2.5");
let fast_billing = admin
.calculated_billing(&ProviderBillingInput {
upstream_model_id: "gpt-4o".to_owned(),
service_tier: Some("priority".to_owned()),
input_tokens: Some(1_000_000),
output_tokens: Some(0),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
total: CurrencyCost {
currency: "USD".to_owned(),
amount: "5".parse().expect("fast amount"),
},
})
.expect("fast billing")
.expect("known fast pricing");
assert_eq!(fast_billing.service_tier.as_deref(), Some("priority"));
assert_eq!(fast_billing.multiplier_percent, 200);
assert_eq!(fast_billing.standard_amount.amount.as_str(), "2.5");
assert_eq!(fast_billing.total_amount.amount.as_str(), "5");
}
#[tokio::test]
@@ -281,6 +281,75 @@ async fn cross_plan_responses_api_conflict_still_fails_the_union() {
server.verify().await;
}
#[tokio::test]
async fn free_and_k12_catalog_entitlements_are_isolated_by_plan_scope() {
let store = Arc::new(MemoryAccountStore::default());
let first_free = seed_account_with_plan(&store, "acct_free_a", "free").await;
let second_free = seed_account_with_plan(&store, "acct_free_b", "free").await;
let k12 = seed_account_with_plan(&store, "acct_k12", "k12").await;
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/codex/models"))
.respond_with(SequencedCatalogResponder {
calls: Arc::new(AtomicUsize::new(0)),
bodies: [
br#"{"models":[{"slug":"gpt-5.4","display_name":"GPT-5.4","supported_in_api":true}]}"#,
br#"{"models":[{"slug":"gpt-5.4","display_name":"GPT-5.4","supported_in_api":true},{"slug":"gpt-5.6-sol","display_name":"GPT-5.6 Sol","supported_in_api":true}]}"#,
],
})
.expect(2)
.mount(&server)
.await;
let service = service_with_catalog_cache(&store, server.uri(), catalog_cache());
let snapshot = service.synchronize().await.expect("mixed-plan catalog");
assert_eq!(snapshot.models().len(), 2);
assert_eq!(
service
.cached_account_models(&first_free)
.expect("first free entitlement"),
Some(vec!["gpt-5.4".to_owned()])
);
assert_eq!(
service
.cached_account_models(&second_free)
.expect("second free entitlement"),
Some(vec!["gpt-5.4".to_owned()])
);
assert_eq!(
service
.observed_model_support(&first_free, "gpt-5.6-sol")
.expect("free support fact"),
Some(false)
);
assert_eq!(
service
.observed_model_support(&k12, "gpt-5.6-sol")
.expect("K12 support fact"),
Some(true)
);
let synchronized_generation = service.catalog_generation().get();
service.invalidate().expect("invalidate mixed-plan catalog");
assert!(
service
.cached()
.expect("cache read after invalidation")
.is_none()
);
assert_eq!(
service.catalog_generation().get(),
synchronized_generation + 1
);
assert_eq!(
service
.observed_model_support(&k12, "gpt-5.6-sol")
.expect("support fact after invalidation"),
None
);
server.verify().await;
}
#[tokio::test]
async fn response_etag_change_is_deduplicated_and_queued_once() {
let store = Arc::new(MemoryAccountStore::default());
@@ -374,7 +374,15 @@ async fn selector_should_reuse_the_account_bound_to_the_same_session() {
.await
.expect("select bound account");
assert_eq!(second.account_id(), &first_account);
assert_eq!(
(
second.account_id().as_str(),
second.affinity_hit(),
second.escape_reason(),
second.account_switch(),
),
(first_account.as_str(), true, None, false)
);
assert_eq!(
affinity
.load(&ProviderKind::new("openai").expect("provider"), &key)
@@ -421,7 +429,15 @@ async fn selector_should_replace_a_busy_affinity_binding_after_the_fallback_succ
.record_success(selected.account(), Some(&key))
.await;
assert_eq!(selected.account_id().as_str(), "acct_second");
assert_eq!(
(
selected.account_id().as_str(),
selected.affinity_hit(),
selected.escape_reason(),
selected.account_switch(),
),
("acct_second", false, Some("lease_saturated"), true)
);
assert_eq!(
affinity
.load(&provider, &key)
@@ -432,7 +448,7 @@ async fn selector_should_replace_a_busy_affinity_binding_after_the_fallback_succ
}
#[tokio::test]
async fn selector_should_escape_an_unhealthy_affinity_account() {
async fn selector_should_keep_a_schedulable_affinity_account_despite_soft_health_signals() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
@@ -473,9 +489,119 @@ async fn selector_should_escape_an_unhealthy_affinity_account() {
session_affinity_key: Some(&key),
})
.await
.expect("select healthy fallback");
.expect("select bound account");
assert_eq!(selected.account_id().as_str(), "acct_second");
assert_eq!(
(
selected.account_id().as_str(),
selected.affinity_hit(),
selected.escape_reason(),
selected.account_switch(),
),
("acct_first", true, None, false)
);
}
#[tokio::test]
async fn selector_should_escape_a_quota_exhausted_affinity_account() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
let first = store.account("acct_first").expect("first account");
store
.apply_state_change(AccountStateChange {
account_id: first.id().clone(),
expected_revision: first.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("mark first account exhausted");
let affinity = Arc::new(MemorySessionAffinity::default());
let provider = ProviderKind::new("openai").expect("provider");
let key = ProviderSessionAffinityKey::try_new("quota-session").expect("affinity key");
affinity
.bind(&provider, &key, first.id(), Duration::from_secs(60))
.await
.expect("seed affinity");
let selector =
selector_with_affinity(&store, Arc::new(TestLeaseCoordinator::default()), affinity);
let request_url =
Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL");
let request_attempt = attempt(BTreeSet::new());
let selected = selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &request_attempt,
session_affinity_key: Some(&key),
})
.await
.expect("select fallback account");
assert_eq!(
(
selected.account_id().as_str(),
selected.affinity_hit(),
selected.escape_reason(),
selected.account_switch(),
),
("acct_second", false, Some("quota_exhausted"), true)
);
}
#[tokio::test]
async fn selector_should_escape_an_affinity_account_during_cooldown() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
let first = store.account("acct_first").expect("first account");
store
.apply_state_change(AccountStateChange {
account_id: first.id().clone(),
expected_revision: first.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: Some(SystemTime::now() + Duration::from_secs(60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark first account cooling down");
let affinity = Arc::new(MemorySessionAffinity::default());
let provider = ProviderKind::new("openai").expect("provider");
let key = ProviderSessionAffinityKey::try_new("cooldown-session").expect("affinity key");
affinity
.bind(&provider, &key, first.id(), Duration::from_secs(60))
.await
.expect("seed affinity");
let selector =
selector_with_affinity(&store, Arc::new(TestLeaseCoordinator::default()), affinity);
let request_url =
Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL");
let request_attempt = attempt(BTreeSet::new());
let selected = selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &request_attempt,
session_affinity_key: Some(&key),
})
.await
.expect("select fallback account");
assert_eq!(
(
selected.account_id().as_str(),
selected.affinity_hit(),
selected.escape_reason(),
selected.account_switch(),
),
("acct_second", false, Some("cooldown"), true)
);
}
#[test]
@@ -361,6 +361,83 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
);
}
#[tokio::test]
async fn manual_quota_refresh_updates_account_state_from_provider_response() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_manual_quota_state";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header(
"authorization",
format!("Bearer token-{account_id}"),
))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": false,
"limit_reached": true,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let exhausted = service
.refresh_account(account.id())
.await
.expect("refresh exhausted quota");
assert!(exhausted.fact().exhausted());
assert_eq!(
store
.account(account_id)
.expect("exhausted account")
.availability(),
AccountAvailability::QuotaExhausted
);
server.reset().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header(
"authorization",
format!("Bearer token-{account_id}"),
))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 20, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let recovered = service
.refresh_account(account.id())
.await
.expect("refresh recovered quota");
assert!(!recovered.fact().exhausted());
assert_eq!(
store
.account(account_id)
.expect("recovered account")
.availability(),
AccountAvailability::Ready
);
}
#[tokio::test]
async fn passive_rate_limit_headers_replace_stale_secondary_window() {
let store = Arc::new(MemoryAccountStore::default());
@@ -3,6 +3,7 @@ use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use chrono::Utc;
use futures::StreamExt;
use gateway_core::engine::credential::{
@@ -11,11 +12,11 @@ use gateway_core::engine::credential::{
use gateway_core::engine::provider::{Provider as _, ProviderRequest};
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ContinuationAttempt, ModelRequestId,
RequestAttemptContext, UpstreamSendState,
ProviderAccountStateOwner, RequestAttemptContext, UpstreamSendState,
};
use gateway_core::error::ProviderErrorKind;
use gateway_core::event::GatewayEvent;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
@@ -45,6 +46,10 @@ use crate::support::{
const OFFICIAL_FIXTURE: &[u8] =
include_bytes!("../transport/fixtures/official_models_snapshot.json");
const CAPTURE_COMPLETED_SSE: &str = concat!(
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_scope_capture\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
);
fn wire_profile() -> CodexWireProfileState {
CodexWireProfileState::new(CodexWireProfile {
@@ -231,6 +236,78 @@ fn diagnostic_context(request_id: &str, account_id: &str) -> AttemptContext {
)
}
fn context_with_state_owner(request_id: &str, owner_account_id: &str) -> AttemptContext {
let owner = ProviderAccountStateOwner::new(
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new(owner_account_id).expect("owner account id"),
);
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).expect("request id"),
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
),
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, Some(owner)),
None,
CancellationToken::new(),
)
}
async fn capture_scoped_http_request(
request_id: &str,
selected_account_id: &str,
owner_account_id: &str,
body: Map<String, serde_json::Value>,
mut protocol_context: Map<String, serde_json::Value>,
) -> wiremock::Request {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, selected_account_id).await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(CAPTURE_COMPLETED_SSE),
)
.expect(1)
.mount(&server)
.await;
protocol_context.insert("use_websocket".to_owned(), json!(false));
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object("openai", body)
.expect("OpenAI payload")
.with_context(protocol_context),
));
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", operation),
context_with_state_owner(request_id, owner_account_id),
)
.await
.expect("prepare scoped provider stream");
while let Some(event) = stream.next().await {
event.expect("scoped provider response");
}
let mut requests = server
.received_requests()
.await
.expect("captured scoped request");
assert_eq!(requests.len(), 1);
requests.pop().expect("single scoped request")
}
fn captured_header_values(request: &wiremock::Request, name: &str) -> Vec<Vec<u8>> {
request
.headers
.get_all(name)
.iter()
.map(|value| value.as_bytes().to_vec())
.collect()
}
async fn paused_chunked_sse_server(
first_chunk: String,
second_chunk: String,
@@ -272,6 +349,28 @@ async fn paused_chunked_sse_server(
(base_url, release_sender, first_chunk_sent, server)
}
async fn truncated_chunked_sse_server() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind truncated SSE listener");
let base_url = format!(
"http://{}",
listener.local_addr().expect("listener address")
);
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept SSE request");
read_http_request(&mut stream).await;
stream
.write_all(
b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n20\r\npartial",
)
.await
.expect("write truncated chunked response");
stream.flush().await.expect("flush truncated response");
});
(base_url, server)
}
async fn read_http_request(stream: &mut TcpStream) {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
@@ -339,6 +438,34 @@ async fn cancelled_attempt_fails_before_account_selection_or_upstream_send() {
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
#[tokio::test]
async fn truncated_http_stream_allows_account_rotation_only_before_client_delivery() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_truncated_stream").await;
let (base_url, server) = truncated_chunked_sse_server().await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_truncated_stream", CancellationToken::new()),
)
.await
.expect("prepare HTTP stream");
let error = loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("truncated response must surface a transport error"),
}
};
server.await.expect("truncated SSE server");
assert_eq!(error.kind(), ProviderErrorKind::Transport);
assert_eq!(error.send_state(), UpstreamSendState::Sent);
assert!(error.allows_pre_delivery_retry());
assert!(!error.replay_is_safe());
}
#[tokio::test]
async fn openai_provider_keeps_a_compaction_trigger_as_a_regular_generate_request() {
let store = Arc::new(MemoryAccountStore::default());
@@ -438,6 +565,192 @@ async fn opaque_provider_options_do_not_change_openai_account_selection() {
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
#[tokio::test]
async fn same_account_scope_preserves_future_protocol_shapes() {
let request = capture_scoped_http_request(
"req_scope_same",
"acct_scope_same",
"acct_scope_same",
json!({
"model": "gpt-5.4",
"input": "hello",
"authorization": "client-extension-value",
"installation_id": "client-installation",
"turnState": {"future": true},
"turnMetadata": ["future", "shape"],
"x-codex-turn-state": 17,
"x-codex-turn-metadata": false,
"client_metadata": {
"token": "client-metadata-extension",
"x-codex-turn-state": {"future": "metadata"},
"x-codex-turn-metadata": [1, 2, 3],
"turnMetadata": {"opaque": true}
}
})
.as_object()
.expect("request object")
.clone(),
Map::new(),
)
.await;
let body: serde_json::Value =
serde_json::from_slice(&request.body).expect("captured JSON body");
assert_eq!(
body.get("authorization"),
Some(&json!("client-extension-value"))
);
assert_eq!(body.get("turnState"), Some(&json!({"future": true})));
assert_eq!(body.get("turnMetadata"), Some(&json!(["future", "shape"])));
assert_eq!(body.get("x-codex-turn-state"), Some(&json!(17)));
assert_eq!(body.get("x-codex-turn-metadata"), Some(&json!(false)));
assert_eq!(
body.pointer("/client_metadata/token"),
Some(&json!("client-metadata-extension"))
);
assert_eq!(
body.pointer("/client_metadata/x-codex-turn-state"),
Some(&json!({"future": "metadata"}))
);
assert_ne!(
body.get("installation_id"),
Some(&json!("client-installation"))
);
}
#[tokio::test]
async fn cross_account_scope_removes_only_account_bound_body_fields() {
let request = capture_scoped_http_request(
"req_scope_switch",
"acct_scope_new",
"acct_scope_old",
json!({
"model": "gpt-5.4",
"input": "hello",
"authorization": "client-extension-value",
"conversation": "upstream-account-handle",
"conversation_id": "client-correlation",
"installation_id": "client-installation",
"client_metadata": ["future", "shape"],
"future_field": {"keep": true}
})
.as_object()
.expect("request object")
.clone(),
Map::new(),
)
.await;
let body: serde_json::Value =
serde_json::from_slice(&request.body).expect("captured JSON body");
assert!(body.get("authorization").is_none());
assert!(body.get("conversation").is_none());
assert_eq!(
body.get("conversation_id"),
Some(&json!("client-correlation"))
);
assert_eq!(
body.get("client_metadata"),
Some(&json!(["future", "shape"]))
);
assert_eq!(body.get("future_field"), Some(&json!({"keep": true})));
assert_ne!(
body.get("installation_id"),
Some(&json!("client-installation"))
);
}
#[tokio::test]
async fn cross_account_scope_sanitizes_only_known_turn_metadata_fields() {
let request = capture_scoped_http_request(
"req_scope_metadata",
"acct_metadata_new",
"acct_metadata_old",
json!({
"model": "gpt-5.4",
"input": "hello",
"turnMetadata": r#"{"account_id":"old-account","future":{"keep":true}}"#,
"turn_metadata": "future-opaque-shape",
"x-codex-turn-metadata": r#"{"conversation":"old-conversation","safe":17}"#
})
.as_object()
.expect("request object")
.clone(),
Map::new(),
)
.await;
let body: serde_json::Value =
serde_json::from_slice(&request.body).expect("captured JSON body");
let turn_metadata = body
.get("turnMetadata")
.and_then(serde_json::Value::as_str)
.and_then(|value| serde_json::from_str::<serde_json::Value>(value).ok())
.expect("sanitized turnMetadata");
let codex_turn_metadata = body
.get("x-codex-turn-metadata")
.and_then(serde_json::Value::as_str)
.and_then(|value| serde_json::from_str::<serde_json::Value>(value).ok())
.expect("sanitized x-codex-turn-metadata");
assert_eq!(turn_metadata, json!({"future": {"keep": true}}));
assert_eq!(codex_turn_metadata, json!({"safe": 17}));
assert_eq!(
body.get("turn_metadata"),
Some(&json!("future-opaque-shape"))
);
}
#[tokio::test]
async fn account_change_drops_only_account_bound_opaque_headers() {
let protocol_context = Map::from_iter([(
"opaque_request_headers".to_owned(),
json!([
["x-codex-turn-state", STANDARD.encode(b"turn-first")],
["x-codex-turn-state", STANDARD.encode(b"turn-\x80")],
[
"x-codex-turn-metadata",
STANDARD.encode(br#"{"installation_id":"client-installation","safe":true}"#)
],
["x-openai-future", STANDARD.encode(b"keep-on-switch")]
]),
)]);
let body = json!({"model": "gpt-5.4", "input": "hello"})
.as_object()
.expect("request object")
.clone();
let same_account = capture_scoped_http_request(
"req_header_same",
"acct_header_same",
"acct_header_same",
body.clone(),
protocol_context.clone(),
)
.await;
let cross_account = capture_scoped_http_request(
"req_header_switch",
"acct_header_new",
"acct_header_old",
body,
protocol_context,
)
.await;
assert_eq!(
captured_header_values(&same_account, "x-codex-turn-state"),
vec![b"turn-first".to_vec(), b"turn-\x80".to_vec()]
);
assert_eq!(
captured_header_values(&same_account, "x-openai-future"),
vec![b"keep-on-switch".to_vec()]
);
assert!(captured_header_values(&cross_account, "x-codex-turn-state").is_empty());
assert!(captured_header_values(&cross_account, "x-codex-turn-metadata").is_empty());
assert_eq!(
captured_header_values(&cross_account, "x-openai-future"),
vec![b"keep-on-switch".to_vec()]
);
}
#[tokio::test]
async fn prompt_cache_key_should_become_an_opaque_session_affinity_lookup_key() {
let store = Arc::new(MemoryAccountStore::default());
@@ -484,7 +797,7 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_affinity\",\"model\":\"gpt-5.4\"}}\n\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_affinity\",\"model\":\"gpt-5.4\",\"service_tier\":\"default\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_affinity\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)),
@@ -498,6 +811,7 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("hello")),
("prompt_cache_key".to_owned(), json!("affinity-key")),
("service_tier".to_owned(), json!("priority")),
]),
)
.expect("OpenAI payload")
@@ -512,8 +826,15 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
.await
.expect("prepare provider stream");
let mut observed_service_tier = None;
while let Some(event) = stream.next().await {
let event = event.expect("provider event");
if let Some(service_tier) = event
.response_observation()
.and_then(|observation| observation.service_tier())
{
observed_service_tier = Some(service_tier.to_owned());
}
if event
.canonical_facts()
.iter()
@@ -525,6 +846,7 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
drop(stream);
assert_eq!(affinity.binding_count(), 1);
assert_eq!(observed_service_tier.as_deref(), Some("default"));
}
#[tokio::test]
@@ -776,7 +1098,7 @@ async fn bare_response_failed_should_remain_an_atomic_replay_safe_failure() {
}
#[tokio::test]
async fn continuation_prefetch_over_64_kib_without_a_boundary_should_fail_as_protocol_error() {
async fn continuation_prefetch_over_64_kib_should_commit_wire_without_protocol_failure() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_prefetch_limit").await;
let padding = "x".repeat(64 * 1024);
@@ -803,21 +1125,26 @@ async fn continuation_prefetch_over_64_kib_without_a_boundary_should_fail_as_pro
)
.await
.expect("prepare provider stream");
let mut visible_before_error = 0;
let failure = loop {
match stream.next().await {
Some(Ok(event)) => visible_before_error += usize::from(event.has_client_event()),
Some(Err(error)) => break error,
None => panic!("oversized prefetch must fail"),
let visible = loop {
let event = timeout(Duration::from_secs(1), stream.next())
.await
.expect("prefetch threshold must release buffered wire")
.expect("provider stream remains open")
.expect("threshold cannot create a protocol failure");
if event.has_client_event() {
break event;
}
};
assert_eq!(visible_before_error, 0);
assert_eq!(failure.kind(), ProviderErrorKind::Protocol);
assert!(!failure.has_atomic_client_events());
let _ = release.send(());
server.abort();
let _ = server.await;
assert_eq!(
visible.wire_event().and_then(|wire| wire.event_type()),
Some("response.created")
);
release.send(()).expect("finish upstream response");
while let Some(event) = stream.next().await {
event.expect("clean upstream EOF cannot become a protocol failure");
}
server.await.expect("chunked SSE server");
}
#[tokio::test]
@@ -943,6 +1270,59 @@ async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_st
assert_eq!(affinity.binding_count(), 0);
}
#[tokio::test]
async fn successful_response_preserves_quota_exhaustion_from_rate_limit_headers() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_success_exhausted";
create_account(&store, account_id).await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.insert_header("x-codex-active-limit", "codex")
.insert_header("x-codex-primary-used-percent", "100")
.insert_header("x-codex-primary-window-minutes", "300")
.insert_header("x-codex-primary-reset-at", "1900000000")
.insert_header("x-codex-limit-reached", "true")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_success_exhausted\",\"model\":\"gpt-5.4\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_success_exhausted\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)),
)
.mount(&server)
.await;
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", http_generate_operation()),
context("req_success_exhausted", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let mut completed = false;
while let Some(event) = stream.next().await {
let event = event.expect("successful upstream response");
completed |= event
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Completed(_)));
}
assert!(completed);
assert_eq!(
store
.account(account_id)
.expect("account after successful response")
.availability(),
AccountAvailability::QuotaExhausted
);
assert!(store.has_quota(account_id));
}
#[test]
fn request_observation_reads_openai_metadata_without_changing_the_operation() {
let store = Arc::new(MemoryAccountStore::default());
@@ -989,6 +1369,42 @@ fn request_observation_preserves_the_raw_reasoning_effort() {
);
}
#[test]
fn request_observation_ignores_future_session_state_without_rewriting_the_protocol_body() {
let store = Arc::new(MemoryAccountStore::default());
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-test")),
("previous_response_id".to_owned(), json!("resp_opaque")),
]),
)
.expect("OpenAI payload"),
)
.with_provider_session_state(
ProviderSessionState::new(
"openai",
Map::from_iter([("future_session_shape".to_owned(), json!([1, 2, 3]))]),
)
.expect("provider session state"),
);
let operation = Operation::Generate(generation);
let _observation = provider(&store).request_observation(&operation);
let Operation::Generate(generation) = &operation else {
panic!("operation should remain a generate request");
};
assert_eq!(
generation
.protocol_payload()
.body()
.get("previous_response_id"),
Some(&json!("resp_opaque"))
);
}
#[tokio::test]
async fn provider_compiles_catalog_presentation_for_codex_models() {
let store = Arc::new(MemoryAccountStore::default());
@@ -30,6 +30,27 @@ fn decoder_should_preserve_existing_metadata_fixture_as_openai_wire() {
));
}
#[test]
fn decoder_should_preserve_empty_opaque_response_id_in_canonical_facts() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"\",\"model\":\"gpt-test\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"\",\"model\":\"gpt-test\",\"status\":\"completed\",\"output\":[]}}\n\n",
);
let events = CodexCanonicalDecoder::new("fallback")
.push(body.as_bytes())
.expect("empty opaque response ID remains representable");
let facts = canonical_facts(&events);
assert!(matches!(
facts.as_slice(),
[GatewayEvent::Started(started), GatewayEvent::Completed(completed)]
if started.response_id().is_empty() && completed.response_id().is_empty()
));
}
#[test]
fn raw_sse_passthrough_should_keep_original_bytes_alongside_canonical_facts() {
let heartbeat = ": keep-alive\r\n\r\n";
@@ -129,16 +150,14 @@ fn raw_sse_passthrough_should_forward_unparseable_frames_without_failure() {
}
#[test]
fn raw_sse_passthrough_should_forward_frames_with_unsafe_event_metadata() {
// 上游事件名超出 wire 安全上限(256 字节)时剥离该元数据继续下发,
// 原始帧字节仍原样透传,而不是整帧静默丢弃。
fn raw_sse_passthrough_should_preserve_long_event_metadata() {
let event_name = "x".repeat(300);
let frame = format!("event: {event_name}\ndata: {{\"type\":\"noise\",\"marker\":1}}\n\n");
let events = CodexCanonicalDecoder::new("fallback")
.with_raw_sse_passthrough()
.push(frame.as_bytes())
.expect("unsafe event metadata must not abort transparent transport");
.expect("opaque event metadata must not abort transparent transport");
let wire = events[0].wire_event().expect("wire event");
assert_eq!(
@@ -146,22 +165,20 @@ fn raw_sse_passthrough_should_forward_frames_with_unsafe_event_metadata() {
Some(frame.as_bytes())
);
assert!(wire.has_json_data());
assert_eq!(wire.event_type(), None);
assert_eq!(wire.event_type(), Some(event_name.as_str()));
}
#[test]
fn decoder_should_strip_unsafe_event_metadata_instead_of_dropping_the_event() {
// 无 raw 帧的路径(WebSocket JSON 投影)同样不允许因元数据违规丢事件:
// 剥离事件名后 JSON 正文仍完整下发。
fn decoder_should_preserve_long_event_metadata_without_raw_frame() {
let event_name = "y".repeat(300);
let frame = format!("event: {event_name}\ndata: {{\"type\":\"noise\",\"marker\":1}}\n\n");
let events = CodexCanonicalDecoder::new("fallback")
.push(frame.as_bytes())
.expect("unsafe event metadata must not drop the event");
.expect("opaque event metadata must not drop the event");
let wire = events[0].wire_event().expect("wire event");
assert_eq!(wire.event_type(), None);
assert_eq!(wire.event_type(), Some(event_name.as_str()));
assert!(wire.has_json_data());
assert_eq!(wire.data()["marker"], serde_json::json!(1));
}
@@ -185,6 +202,27 @@ fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() {
)));
}
#[test]
fn decoder_should_keep_response_service_tier_when_terminal_event_omits_it() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"priority\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
);
let mut decoder = CodexCanonicalDecoder::new("fallback");
let events = decoder
.push(body.as_bytes())
.expect("canonical priority response");
assert_eq!(decoder.response_service_tier(), Some("priority"));
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
event,
GatewayEvent::CalculatedCost(cost)
if cost.total().amount().scaled() == 6_875_000
)));
}
#[test]
fn decoder_should_normalize_text_usage_and_completion() {
let body = concat!(
@@ -327,5 +327,6 @@ fn context() -> CodexRequestContext<'static> {
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
}
}
@@ -7,6 +7,7 @@ fn upstream_error_formatting_should_redact_body() {
let error = CodexClientError::Upstream {
status: reqwest::StatusCode::UNAUTHORIZED,
body: secret.to_owned(),
client_response: None,
retry_after_seconds: None,
diagnostics: Box::default(),
set_cookie_headers: vec!["session=secret".to_owned()],
@@ -40,6 +41,7 @@ fn request_context_debug_should_redact_all_identity_material() {
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
};
let debug = format!("{context:?}");
@@ -12,6 +12,7 @@ fn classify(status: u16, body: &str) -> CodexFailureCategory {
body,
None,
&CodexUpstreamDiagnostics::with_status(status.as_u16()),
None,
&[],
&[],
CodexUpstreamSendPhase::AfterPayload,
@@ -1,8 +1,106 @@
use base64::{Engine as _, engine::general_purpose::STANDARD};
use gateway_core::operation::{GenerateRequest, ProtocolPayload};
use provider_openai::encode_generate_request;
use provider_openai::transport::websocket::CodexWebSocketConnection;
use serde_json::Value;
use tokio_tungstenite::tungstenite::http::HeaderMap as WsHeaderMap;
use super::*;
fn request_with_opaque_headers(use_websocket: bool) -> CodexResponsesRequest {
let payload = ProtocolPayload::json_object(
"openai",
json!({"model": "gpt-test", "input": "hello"})
.as_object()
.expect("request object")
.clone(),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([
(
"opaque_request_headers".to_owned(),
json!([
["x-openai-future-mode", STANDARD.encode(b"future-ascii")],
["x-openai-future-mode", STANDARD.encode(b"\x80\xff")],
[
"accept",
STANDARD.encode(b"application/vnd.openai.future+json")
],
[
"content-type",
STANDARD.encode(b"application/vnd.openai.request+json")
],
["user-agent", STANDARD.encode(b"Codex future-client")],
["originator", STANDARD.encode(b"future-originator")],
["openai-beta", STANDARD.encode(b"future_responses=v2")],
["openai-beta", STANDARD.encode(b"future_tools=v3")],
[
"x-openai-internal-codex-residency",
STANDARD.encode(b"future-region")
],
["x-codex-turn-state", STANDARD.encode(b"turn-ascii")],
["x-codex-turn-state", STANDARD.encode(b"turn-\x80")],
["bad header name", STANDARD.encode(b"ignored")],
["x-invalid-base64", "%%%"],
["x-still-valid", STANDARD.encode(b"after-invalid")],
["authorization", STANDARD.encode(b"Bearer client-secret")],
["chatgpt-account-id", STANDARD.encode(b"client-account")],
[
"x-codex-installation-id",
STANDARD.encode(b"client-installation")
]
]),
),
("turn_state".to_owned(), json!("typed-turn-state")),
]));
let mut request = encode_generate_request(
&GenerateRequest::from_protocol_payload(payload),
"gpt-routed",
)
.expect("opaque request headers");
request.use_websocket = use_websocket;
request.force_http_sse = !use_websocket;
request
}
async fn read_http_request_head(stream: &mut TcpStream) -> Vec<u8> {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
loop {
let read = stream.read(&mut buffer).await.expect("read HTTP request");
if read == 0 {
break;
}
request.extend_from_slice(&buffer[..read]);
if let Some(end) = request.windows(4).position(|window| window == b"\r\n\r\n") {
request.truncate(end + 4);
break;
}
}
request
}
fn raw_header_values(request: &[u8], target: &str) -> Vec<Vec<u8>> {
request
.split(|byte| *byte == b'\n')
.skip(1)
.take_while(|line| *line != b"\r" && !line.is_empty())
.filter_map(|line| {
let line = line.strip_suffix(b"\r").unwrap_or(line);
let colon = line.iter().position(|byte| *byte == b':')?;
line[..colon]
.eq_ignore_ascii_case(target.as_bytes())
.then(|| {
line[colon + 1..]
.iter()
.copied()
.skip_while(|byte| matches!(byte, b' ' | b'\t'))
.collect()
})
})
.collect()
}
#[test]
fn websocket_connection_should_preserve_endpoint_and_header_order() {
let connection = CodexWebSocketConnection::new(
@@ -151,7 +249,7 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
request.memgen_request = Some("true".to_owned());
request.set_client_metadata(Some(json!({
"safe": "yes",
"x-openai-subagent": "review",
"x-openai-subagent": "future_codex_mode",
"ignored_non_string": 42
})));
let backend = CodexBackendClient::new(
@@ -184,6 +282,7 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
},
)
.await
@@ -208,6 +307,8 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
let headers = received_headers.lock().expect("headers lock");
for (name, expected) in [
("x-client-request-id", "req_ws_security"),
("content-type", "application/json"),
("accept", "text/event-stream"),
("x-codex-installation-id", "install-123"),
("x-openai-internal-codex-residency", "us"),
("x-codex-turn-state", "turn-state"),
@@ -217,7 +318,7 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
("version", "26.318.11754"),
("x-codex-window-id", "cw_derived"),
("x-codex-parent-thread-id", "parent-456"),
("x-openai-subagent", "review"),
("x-openai-subagent", "future_codex_mode"),
("x-openai-memgen-request", "true"),
("session-id", "cp_derived"),
] {
@@ -228,8 +329,6 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
);
}
for forbidden in [
"content-type",
"accept",
"session_id",
"thread-id",
"x-openai-internal-codex-responses-lite",
@@ -238,6 +337,186 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
}
}
#[tokio::test]
async fn backend_http_should_restore_opaque_multivalue_header_bytes_and_lease_identity() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind opaque HTTP server");
let address = listener.local_addr().expect("opaque HTTP server address");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept opaque HTTP client");
let request = read_http_request_head(&mut stream).await;
write_completed_sse_response(&mut stream).await;
request
});
let request = request_with_opaque_headers(false);
let client = CodexBackendClient::new(
reqwest::Client::builder()
.no_proxy()
.build()
.expect("HTTP client"),
format!("http://{address}"),
test_wire_profile(),
);
client
.create_response(
&request,
CodexRequestContext {
authorization: "Bearer lease-token",
account_id: Some("lease-account"),
installation_id: Some("lease-installation"),
turn_state: request.turn_state.as_deref(),
..request_context("req_opaque_http", Some("lease-account"))
},
)
.await
.expect("opaque HTTP response");
let raw = server.await.expect("opaque HTTP server task");
assert_eq!(
raw_header_values(&raw, "x-openai-future-mode"),
vec![b"future-ascii".to_vec(), b"\x80\xff".to_vec()]
);
assert_eq!(
raw_header_values(&raw, "openai-beta"),
vec![b"future_responses=v2".to_vec(), b"future_tools=v3".to_vec()]
);
assert_eq!(
raw_header_values(&raw, "x-codex-turn-state"),
vec![b"turn-ascii".to_vec(), b"turn-\x80".to_vec()]
);
for (name, value) in [
("accept", b"application/vnd.openai.future+json".as_slice()),
(
"content-type",
b"application/vnd.openai.request+json".as_slice(),
),
("user-agent", b"Codex future-client".as_slice()),
("originator", b"future-originator".as_slice()),
(
"x-openai-internal-codex-residency",
b"future-region".as_slice(),
),
("x-still-valid", b"after-invalid".as_slice()),
("authorization", b"Bearer lease-token".as_slice()),
("chatgpt-account-id", b"lease-account".as_slice()),
("x-codex-installation-id", b"lease-installation".as_slice()),
] {
assert_eq!(raw_header_values(&raw, name), vec![value.to_vec()]);
}
for secret in [
b"client-secret".as_slice(),
b"client-account",
b"client-installation",
] {
assert!(!raw.windows(secret.len()).any(|window| window == secret));
}
for omitted in ["bad header name", "x-invalid-base64"] {
assert!(raw_header_values(&raw, omitted).is_empty());
}
}
#[tokio::test]
async fn backend_websocket_should_drop_only_unrepresentable_opaque_header_values() {
let received = Arc::new(Mutex::new(WsHeaderMap::new()));
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind opaque WebSocket server");
let address = listener
.local_addr()
.expect("opaque WebSocket server address");
let received_for_server = Arc::clone(&received);
let server = tokio::spawn(async move {
let (stream, _) = listener
.accept()
.await
.expect("accept opaque WebSocket client");
let mut websocket = accept_codex_test_websocket_with(stream, move |request, response| {
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().expect("extension header"),
);
*received_for_server.lock().expect("opaque headers lock") = request.headers().clone();
})
.await;
let _ = websocket
.next()
.await
.expect("opaque response.create")
.expect("valid opaque response.create");
websocket
.send(Message::Text(
completed_websocket_response("resp_opaque_headers", 1, 1).into(),
))
.await
.expect("send opaque terminal event");
});
let request = request_with_opaque_headers(true);
let client = CodexBackendClient::new(
reqwest::Client::builder()
.no_proxy()
.build()
.expect("HTTP client"),
format!("http://{address}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))));
client
.create_response(
&request,
CodexRequestContext {
authorization: "Bearer lease-token",
account_id: Some("lease-account"),
installation_id: Some("lease-installation"),
turn_state: request.turn_state.as_deref(),
..request_context("req_opaque_ws", Some("lease-account"))
},
)
.await
.expect("opaque WebSocket response");
server.await.expect("opaque WebSocket server task");
let received = received.lock().expect("opaque headers lock");
let values = |name: &str| {
received
.get_all(name)
.iter()
.map(|value| value.as_bytes().to_vec())
.collect::<Vec<_>>()
};
assert_eq!(
values("x-openai-future-mode"),
vec![b"future-ascii".to_vec()]
);
assert_eq!(
values("openai-beta"),
vec![b"future_responses=v2".to_vec(), b"future_tools=v3".to_vec()]
);
assert_eq!(
values("accept"),
vec![b"application/vnd.openai.future+json".to_vec()]
);
assert_eq!(
values("content-type"),
vec![b"application/vnd.openai.request+json".to_vec()]
);
assert_eq!(
values("authorization"),
vec![b"Bearer lease-token".to_vec()]
);
assert_eq!(
values("chatgpt-account-id"),
vec![b"lease-account".to_vec()]
);
assert_eq!(
values("x-codex-installation-id"),
vec![b"lease-installation".to_vec()]
);
assert_eq!(values("x-codex-turn-state"), vec![b"turn-ascii".to_vec()]);
}
#[tokio::test]
async fn backend_http_should_send_codex_context_without_browser_headers() {
let listener = TcpListener::bind("127.0.0.1:0")
@@ -289,6 +568,7 @@ async fn backend_http_should_send_codex_context_without_browser_headers() {
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
},
)
.await
@@ -354,6 +634,84 @@ async fn backend_http_should_send_codex_context_without_browser_headers() {
}
}
#[tokio::test]
async fn backend_http_should_ignore_unrepresentable_protocol_headers_without_blocking_body() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind HTTP header server");
let address = listener.local_addr().expect("HTTP header server address");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept HTTP client");
let request = read_http_request(&mut stream).await;
write_completed_sse_response(&mut stream).await;
request
});
let mut request = codex_request("gpt-test", "", Vec::new());
request.force_http_sse = true;
request.turn_state = Some("opaque\nturn-state".to_owned());
request.turn_metadata = Some("opaque\nturn-metadata".to_owned());
request.responses_lite = Some("opaque\nresponses-lite".to_owned());
request.memgen_request = Some("opaque\nmemgen".to_owned());
request.set_client_metadata(Some(json!({
"x-openai-subagent": "opaque\nsubagent",
"future": "preserved"
})));
let original_metadata = request.client_metadata().cloned();
let client = CodexBackendClient::new(
reqwest::Client::builder()
.no_proxy()
.build()
.expect("HTTP client"),
format!("http://{address}"),
test_wire_profile(),
);
client
.create_response(
&request,
CodexRequestContext {
client_request_id: Some("opaque\nclient-request-id"),
session_id: Some("opaque\nsession"),
thread_id: Some("opaque\nthread"),
turn_id: Some("opaque\nturn"),
codex_window_id: Some("opaque\nwindow"),
parent_thread_id: Some("opaque\nparent"),
beta_features: Some("opaque\nbeta"),
include_timing_metrics: Some("opaque\ntiming"),
version: Some("opaque\nversion"),
turn_state: request.turn_state.as_deref(),
turn_metadata: request.turn_metadata.as_deref(),
..request_context("req_protocol_fallback", Some("chatgpt-account"))
},
)
.await
.expect("invalid optional projections must not block the request");
let raw_request = server.await.expect("HTTP header server task");
assert_eq!(
read_header_value(&raw_request, "x-client-request-id"),
Some("req_protocol_fallback")
);
for omitted in [
"session-id",
"thread-id",
"x-codex-turn-id",
"x-codex-window-id",
"x-codex-turn-state",
"x-codex-turn-metadata",
"x-codex-beta-features",
"x-responsesapi-include-timing-metrics",
"version",
"x-codex-parent-thread-id",
"x-openai-subagent",
"x-openai-internal-codex-responses-lite",
"x-openai-memgen-request",
] {
assert!(read_header_value(&raw_request, omitted).is_none());
}
assert_eq!(request.client_metadata(), original_metadata.as_ref());
}
#[tokio::test]
async fn websocket_should_keep_an_exact_chain_while_new_connections_adopt_the_latest_wire_profile()
{
@@ -104,12 +104,25 @@ fn custom_ca_should_report_environment_cache_key_consistently() {
}
#[tokio::test]
async fn codex_backend_client_should_bound_and_redact_oversized_error_body() {
async fn codex_backend_client_should_preserve_oversized_error_response() {
let server = wiremock::MockServer::start().await;
let large_error_body = "x".repeat(1024 * 1024 + 17);
wiremock::Mock::given(wiremock::matchers::method("POST"))
.and(wiremock::matchers::path("/codex/responses"))
.respond_with(wiremock::ResponseTemplate::new(500).set_body_string(large_error_body))
.respond_with(
wiremock::ResponseTemplate::new(500)
.append_header("x-future-error", "first")
.append_header("x-future-error", "second")
.insert_header("x-request-id", "req-upstream-error")
.insert_header("set-cookie", "account-secret=value")
.insert_header("authorization", "Bearer response-secret")
.insert_header("connection", "x-hop-secret")
.insert_header("x-hop-secret", "hop-secret")
.set_body_raw(
large_error_body.clone(),
"application/problem+json; charset=utf-8",
),
)
.mount(&server)
.await;
let client = CodexBackendClient::new(
@@ -140,15 +153,62 @@ async fn codex_backend_client_should_bound_and_redact_oversized_error_body() {
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
},
)
.await;
let Err(CodexClientError::Upstream { status, body, .. }) = result else {
let Err(CodexClientError::Upstream {
status,
body,
client_response: Some(client_response),
..
}) = result
else {
panic!("expected upstream error");
};
assert_eq!(status, reqwest::StatusCode::INTERNAL_SERVER_ERROR);
assert_eq!(body, "upstream error response exceeded the body limit");
assert_eq!(body, large_error_body);
assert_eq!(client_response.status(), 500);
assert_eq!(
client_response.content_type(),
Some(b"application/problem+json; charset=utf-8".as_slice())
);
let future_values = client_response
.client_headers()
.iter()
.filter(|(name, _)| name == "x-future-error")
.map(|(_, value)| value.as_ref())
.collect::<Vec<_>>();
assert_eq!(
future_values,
vec![b"first".as_slice(), b"second".as_slice()]
);
assert!(
client_response
.client_headers()
.iter()
.any(|(name, value)| {
name == "x-request-id" && value.as_ref() == b"req-upstream-error"
})
);
for excluded in [
"content-type",
"content-length",
"set-cookie",
"authorization",
"connection",
"x-hop-secret",
] {
assert!(
client_response
.client_headers()
.iter()
.all(|(name, _)| name != excluded),
"unexpected forwarded header {excluded}"
);
}
assert_eq!(client_response.body().as_ref(), large_error_body.as_bytes());
}
#[tokio::test]
@@ -192,6 +252,7 @@ async fn codex_backend_client_should_parse_retry_after_from_rate_limit_error_bod
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
},
)
.await;
@@ -247,7 +308,7 @@ async fn codex_backend_http_sse_should_fail_after_five_minutes_without_stream_da
}
#[tokio::test]
async fn codex_backend_client_should_capture_only_safe_response_metadata() {
async fn codex_backend_client_should_capture_forwardable_response_metadata() {
let server = wiremock::MockServer::start().await;
let body = concat!(
"event: response.completed\n",
@@ -263,7 +324,18 @@ async fn codex_backend_client_should_capture_only_safe_response_metadata() {
.insert_header("x-models-etag", "models-v2")
.insert_header("x-reasoning-included", "true")
.insert_header("openai-processing-ms", "42")
.append_header("x-future-multi", "first")
.append_header("x-future-multi", "second")
.insert_header(
"x-future-bytes",
reqwest::header::HeaderValue::from_bytes(b"\xffopaque")
.expect("non-UTF-8 header value"),
)
.insert_header("set-cookie", "secret=value")
.insert_header("authorization", "Bearer response-secret")
.insert_header("chatgpt-account-id", "account-secret")
.insert_header("connection", "x-hop-secret")
.insert_header("x-hop-secret", "hop-secret")
.insert_header("x-codex-primary-used-percent", "15")
.set_body_string(body),
)
@@ -297,6 +369,7 @@ async fn codex_backend_client_should_capture_only_safe_response_metadata() {
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
},
)
.await
@@ -321,13 +394,49 @@ async fn codex_backend_client_should_capture_only_safe_response_metadata() {
assert_eq!(
names,
vec![
"date",
"openai-model",
"openai-processing-ms",
"x-codex-primary-used-percent",
"x-future-bytes",
"x-future-multi",
"x-future-multi",
"x-models-etag",
"x-reasoning-included",
"x-request-id"
]
);
assert_eq!(
response
.response_metadata
.client_headers
.iter()
.filter(|(name, _)| name == "x-future-multi")
.map(|(_, value)| value.as_ref())
.collect::<Vec<_>>(),
vec![b"first".as_slice(), b"second".as_slice()]
);
assert!(
response
.response_metadata
.client_headers
.iter()
.any(|(name, value)| { name == "x-future-bytes" && value.as_ref() == b"\xffopaque" })
);
for blocked in [
"set-cookie",
"authorization",
"chatgpt-account-id",
"connection",
"x-hop-secret",
"content-type",
"content-length",
] {
assert!(
!names.contains(&blocked),
"leaked response header {blocked}"
);
}
assert_eq!(response.set_cookie_headers, vec!["secret=value"]);
assert_eq!(
response.rate_limit_headers,
@@ -145,6 +145,7 @@ fn request_context<'a>(
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
}
}
@@ -1,7 +1,7 @@
use gateway_protocol::openai::sse::parse_sse_events;
use provider_openai::transport::protocol::responses::{
CodexResponsesRequest, PreviousResponseScope, TransportRequirement, response_event_signals,
transport_requirement,
CodexResponsesRequest, PreviousResponseScope, ResponsesSseFailure, TransportRequirement,
response_event_signals, transport_requirement,
};
use serde_json::{Value, json};
@@ -171,6 +171,22 @@ fn semantic_output_should_ignore_an_empty_delta() {
assert!(!response_body_has_semantic_output(body));
}
#[test]
fn structured_failure_should_preserve_empty_opaque_fields() {
let failure = ResponsesSseFailure::from_event(
"response.failed",
&json!({
"response": {
"error": {"message": "", "code": "", "type": ""}
}
}),
);
assert_eq!(failure.message, "");
assert_eq!(failure.upstream_code.as_deref(), Some(""));
assert_eq!(failure.upstream_type.as_deref(), Some(""));
}
#[test]
fn semantic_output_should_ignore_the_done_control_frame() {
assert!(!response_body_has_semantic_output(b"data: [DONE]\n\n"));
@@ -172,6 +172,18 @@ fn websocket_metadata_turn_state_should_accept_case_insensitive_header() {
);
}
#[test]
fn websocket_metadata_turn_state_should_ignore_public_event_metadata() {
let event = json!({
"type": "future.business.event",
"metadata": {"turn_state": "business-value"},
"turn_state": "another-business-value"
});
assert_eq!(websocket_metadata_turn_state(&event), None);
assert!(websocket_event_to_sse_frame(&event.to_string()).is_some());
}
#[test]
fn websocket_completed_id_should_read_the_id_without_validating_the_response_shape() {
let valid = json!({
@@ -199,6 +211,16 @@ fn websocket_completed_id_should_read_the_id_without_validating_the_response_sha
);
}
#[test]
fn websocket_completed_id_should_preserve_an_empty_opaque_id() {
let event = json!({
"type": "response.completed",
"response": {"id": ""}
});
assert_eq!(websocket_response_completed_id(&event), Some(String::new()));
}
#[test]
fn websocket_typed_events_should_remain_transparent_without_schema_filtering() {
let malformed_delta = json!({
@@ -112,6 +112,7 @@ async fn fetch_should_use_wham_usage_headers_only() {
thread_id: Some("thread-1"),
client_request_id: Some("client-request-1"),
turn_id: Some("turn-1"),
account_selection: Default::default(),
})
.await
.expect("usage response");
@@ -415,5 +416,6 @@ fn context() -> CodexRequestContext<'static> {
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
}
}
@@ -286,6 +286,70 @@ async fn websocket_execute_response_create_request_should_collect_completed_sse(
assert_eq!(response.usage.expect("usage").input_tokens, 5);
}
#[tokio::test]
async fn websocket_exchange_should_accept_upstream_frame_above_removed_private_limit() {
const REMOVED_FRAME_LIMIT_BYTES: usize = 16 * 1024 * 1024;
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
// Do not negotiate compression: the wire frame itself must exceed the former limit.
let mut websocket =
accept_codex_test_websocket_with(stream, |_request, _response| {}).await;
let _message = websocket.next().await.unwrap().unwrap();
let oversized_event = json!({
"type": "response.output_text.delta",
"delta": "x".repeat(REMOVED_FRAME_LIMIT_BYTES + 1),
})
.to_string();
assert!(oversized_event.len() > REMOVED_FRAME_LIMIT_BYTES);
websocket
.send(Message::Text(oversized_event.into()))
.await
.unwrap();
websocket
.send(Message::Text(
json!({
"type": "response.completed",
"response": {
"id": "resp_after_oversized_frame",
"object": "response",
"output": [],
"usage": {
"input_tokens": 1,
"output_tokens": 1,
"total_tokens": 2
}
}
})
.to_string()
.into(),
))
.await
.unwrap();
});
let request = codex_request("gpt-5.5", "be brief", Vec::new());
let prepared = CodexWebSocketConnection::responses_create_request(
&format!("http://{addr}"),
"dGhlIHNhbXBsZSBub25jZQ==",
vec![(
"authorization".to_string(),
"Bearer access-token".to_string(),
)],
&request,
)
.expect("payload should serialize");
let response = execute_response_create_request(&prepared)
.await
.expect("upstream frame above the former private limit should remain forwardable");
server.await.unwrap();
assert!(response.body.len() > REMOVED_FRAME_LIMIT_BYTES);
assert!(response.body.contains("resp_after_oversized_frame"));
}
#[tokio::test]
async fn websocket_execute_response_create_request_should_return_business_coded_response_failed_as_terminal_sse_fact()
{
@@ -390,29 +454,29 @@ async fn websocket_execute_response_create_request_should_preserve_opening_error
let request = read_http_request(&mut stream).await;
assert!(request.starts_with("GET /codex/responses HTTP/1.1"));
assert!(request.contains("authorization: Bearer access-token"));
let body = r#"{"error":{"message":"rate limited"}}"#;
stream
.write_all(
format!(
"HTTP/1.1 429 Too Many Requests\r\nretry-after: 33\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
)
.as_bytes(),
)
.await
.unwrap();
let mut body = vec![b'x'; 3 * 1024 + 17];
body.extend_from_slice(b"\0\xffraw-opening-tail");
let head = format!(
"HTTP/1.1 429 Too Many Requests\r\nretry-after: 33\r\nx-request-id: req-ws-opening\r\nx-future-error: first\r\nx-future-error: second\r\nset-cookie: account-secret=value\r\ncontent-type: application/problem+json; profile=codex\r\ncontent-length: {}\r\nconnection: close, x-hop-secret\r\nx-hop-secret: hop-secret\r\n\r\n",
body.len()
);
let mut response = head.into_bytes();
response.extend_from_slice(&body);
stream.write_all(&response).await.unwrap();
body
});
let prepared = prepared_websocket_request(&format!("http://{addr}"));
let error = execute_response_create_request(&prepared)
.await
.expect_err("failed opening should surface upstream status");
server.await.unwrap();
let expected_body = server.await.unwrap();
let CodexClientError::Upstream {
status,
retry_after_seconds,
body,
client_response,
..
} = error
else {
@@ -420,7 +484,110 @@ async fn websocket_execute_response_create_request_should_preserve_opening_error
};
assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS);
assert_eq!(retry_after_seconds, Some(33));
assert!(body.contains("rate limited"));
assert_eq!(body.len(), String::from_utf8_lossy(&expected_body).len());
assert!(body.ends_with("\0\u{fffd}raw-opening-tail"));
let client_response = client_response.expect("raw opening response");
assert_eq!(client_response.status(), 429);
assert_eq!(
client_response.content_type(),
Some(b"application/problem+json; profile=codex".as_slice())
);
assert_eq!(client_response.body().len(), expected_body.len());
assert_eq!(client_response.body().as_ref(), expected_body.as_slice());
let future_values = client_response
.client_headers()
.iter()
.filter(|(name, _)| name == "x-future-error")
.map(|(_, value)| value.as_ref())
.collect::<Vec<_>>();
assert_eq!(
future_values,
vec![b"first".as_slice(), b"second".as_slice()]
);
assert!(
client_response
.client_headers()
.iter()
.any(|(name, value)| { name == "retry-after" && value.as_ref() == b"33" })
);
assert!(
client_response
.client_headers()
.iter()
.any(|(name, value)| { name == "x-request-id" && value.as_ref() == b"req-ws-opening" })
);
for excluded in [
"content-type",
"content-length",
"set-cookie",
"connection",
"x-hop-secret",
] {
assert!(
client_response
.client_headers()
.iter()
.all(|(name, _)| name != excluded),
"unexpected forwarded header {excluded}"
);
}
}
#[tokio::test]
async fn websocket_business_events_should_not_be_consumed_by_internal_observers() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
for event in [
json!({
"type": "future.business.rate_limits",
"rate_limits": {
"primary": {
"used_percent": 99,
"window_minutes": 5,
"reset_at": 1893456300
}
}
}),
json!({
"type": "future.business.metadata",
"metadata": {"turn_state": "business-value"}
}),
json!({
"type": "response.completed",
"response": {
"id": "resp_after_business_events",
"object": "response",
"output": [],
"usage": {
"input_tokens": 1,
"output_tokens": 1,
"total_tokens": 2
}
}
}),
] {
websocket
.send(Message::Text(event.to_string().into()))
.await
.unwrap();
}
websocket.close(None).await.unwrap();
});
let prepared = prepared_websocket_request(&format!("http://{addr}"));
let response = execute_response_create_request(&prepared)
.await
.expect("business events should remain on the wire");
server.await.unwrap();
assert!(response.body.contains("event: future.business.rate_limits"));
assert!(response.body.contains("event: future.business.metadata"));
assert_eq!(response.turn_state, None);
assert!(response.rate_limit_headers.is_empty());
}
#[tokio::test]
@@ -1710,6 +1877,7 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_
thread_id: None,
client_request_id: None,
turn_id: None,
account_selection: Default::default(),
},
)
.await
+1 -1
View File
@@ -21,7 +21,7 @@ getrandom = "0.4.3"
jsonwebtoken = { version = "=11.0.0", default-features = false, features = ["aws_lc_rs"] }
reqwest = { version = "=0.12.28", default-features = false, features = ["stream", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] }
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"] }
semver = "1.0.27"
sha2 = "0.11.0"
thiserror = "2.0.18"
@@ -182,7 +182,8 @@ impl GrokAccountSessionSelector {
};
let mut capacity_denied = false;
let mut retry_after = None;
while let Some(selected) = AccountSelector.select(&candidates, &context) {
while let Some(selection) = AccountSelector.select(&candidates, &context) {
let selected = selection.candidate();
let selected_id = selected.account.id().clone();
let selected_revision = selected.account.revision();
let allows_account_state_mutation = !diagnostic || selected.account.enabled();
+5 -12
View File
@@ -1224,7 +1224,7 @@ fn mark_transient_compaction_failure(error: ProviderError) -> ProviderError {
| ProviderErrorKind::Protocol
| ProviderErrorKind::Unavailable
) {
error.with_replay_safe().with_same_account_retry()
error.with_pre_delivery_retry()
} else {
error
}
@@ -1690,15 +1690,11 @@ fn map_request_error(error: GrokRequestEncodeError) -> ProviderError {
if kind != ProviderErrorKind::InvalidRequest {
return provider_error;
}
let detail = ClientVisibleUpstreamError::new(
provider_error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
error.to_string(),
Some("invalid_request_normalization".to_owned()),
Some("invalid_request_error".to_owned()),
);
match detail {
Ok(detail) => provider_error.with_client_visible_upstream_error(detail),
Err(_) => provider_error,
}
))
}
/// 将选择阶段失败映射为带结构化 code 与 retry_after 的 Provider 错误。
@@ -1738,14 +1734,11 @@ fn map_selection_error(error: GrokSessionSelectorError) -> ProviderError {
Some(retry_after) => error.with_retry_after(retry_after),
None => error,
};
match ClientVisibleUpstreamError::new(
error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
message,
Some(code.to_owned()),
Some("account_unavailable_error".to_owned()),
) {
Ok(detail) => error.with_client_visible_upstream_error(detail),
Err(_) => error,
}
))
}
fn cooling_down_message(retry_after: Option<Duration>) -> String {
@@ -8,7 +8,7 @@ use gateway_core::accounting::{
};
use gateway_core::engine::UpstreamSendState;
use gateway_core::error::{
ClientVisibleUpstreamError, ProviderError, ProviderErrorKind, SafeUpstreamValue,
ClientVisibleUpstreamError, OpaqueUpstreamValue, ProviderError, ProviderErrorKind,
};
use gateway_core::event::{
ContentItem, ContentKind, FinishReason, GatewayEvent, ProtocolWireEvent, ProviderEvent,
@@ -796,19 +796,17 @@ fn upstream_event_error(value: &Value) -> ProviderError {
};
let mut error = ProviderError::new(kind, UpstreamSendState::Sent)
.redact_sensitive_context("upstream event");
if let Some(code) = code.and_then(|code| SafeUpstreamValue::new(code.to_owned()).ok()) {
error = error.with_upstream_code(code);
if let Some(code) = code {
error = error.with_upstream_code(OpaqueUpstreamValue::new(code.to_owned()));
}
// 结构化 message/code/type 供原客户端展示与重试分类;message 先脱去
// 账号指纹(上游限流文案内嵌 team UUID),非结构化正文不透出。
if let Some(message) = upstream_error_field(value, "message")
&& let Ok(client_error) = ClientVisibleUpstreamError::new(
if let Some(message) = upstream_error_field(value, "message") {
error = error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
scrub_account_fingerprints(message),
code.map(str::to_owned),
upstream_error_field(value, "type").map(str::to_owned),
)
{
error = error.with_client_visible_upstream_error(client_error);
));
}
error
}
@@ -30,7 +30,7 @@ use crate::{
OAuthHttpResponse, OAuthHttpTransport, TransportFailure, TransportFailureKind, TransportFuture,
};
use gateway_core::engine::UpstreamSendState;
use gateway_core::error::{ClientVisibleUpstreamError, SafeUpstreamValue};
use gateway_core::error::{ClientVisibleUpstreamError, OpaqueUpstreamValue};
use gateway_core::event::UpstreamHttpVersion;
use super::scrub_account_fingerprints;
@@ -1017,17 +1017,15 @@ async fn classify_inference_status(response: Response) -> GrokInferenceTransport
} else {
metadata.code.as_deref().and_then(normalize_failure_code)
};
if let Some(message) = metadata.client_message.as_deref()
&& let Ok(detail) = ClientVisibleUpstreamError::new(
if let Some(message) = metadata.client_message.as_deref() {
error = error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
scrub_account_fingerprints(message),
upstream_code.clone(),
metadata.error_type.clone(),
)
{
error = error.with_client_visible_upstream_error(detail);
));
}
if let Some(code) = upstream_code.and_then(|code| SafeUpstreamValue::new(code).ok()) {
error = error.with_upstream_code(code);
if let Some(code) = upstream_code {
error = error.with_upstream_code(OpaqueUpstreamValue::new(code));
}
if credential_recovery_required {
error = error.with_credential_recovery();
@@ -1235,12 +1233,12 @@ fn upstream_http_version(version: reqwest::Version) -> UpstreamHttpVersion {
}
}
fn upstream_request_id(response: &Response) -> Option<SafeUpstreamValue> {
fn upstream_request_id(response: &Response) -> Option<OpaqueUpstreamValue> {
["x-request-id", "request-id", "cf-ray"]
.into_iter()
.find_map(|name| response.headers().get(name))
.and_then(|value| value.to_str().ok())
.and_then(|value| SafeUpstreamValue::new(value.to_owned()).ok())
.map(|value| OpaqueUpstreamValue::new(value.to_owned()))
}
async fn classify_model_catalog_status(response: Response) -> GrokModelCatalogTransportError {
@@ -5,7 +5,7 @@ use std::time::Duration;
use futures::Stream;
use gateway_core::engine::UpstreamSendState;
use gateway_core::error::{ClientVisibleUpstreamError, SafeUpstreamValue};
use gateway_core::error::{ClientVisibleUpstreamError, OpaqueUpstreamValue};
use gateway_core::event::UpstreamHttpVersion;
use url::Url;
use zeroize::Zeroizing;
@@ -196,7 +196,7 @@ pub struct GrokInferenceResponse {
body: GrokInferenceChunkStream,
http_version: UpstreamHttpVersion,
status_code: u16,
request_id: Option<SafeUpstreamValue>,
request_id: Option<OpaqueUpstreamValue>,
transport_metrics: GrokInferenceTransportMetrics,
}
@@ -207,7 +207,7 @@ impl GrokInferenceResponse {
body: GrokInferenceChunkStream,
http_version: UpstreamHttpVersion,
status_code: u16,
request_id: Option<SafeUpstreamValue>,
request_id: Option<OpaqueUpstreamValue>,
) -> Self {
Self {
body,
@@ -238,7 +238,7 @@ impl GrokInferenceResponse {
}
#[must_use]
pub const fn request_id(&self) -> Option<&SafeUpstreamValue> {
pub const fn request_id(&self) -> Option<&OpaqueUpstreamValue> {
self.request_id.as_ref()
}
@@ -310,8 +310,8 @@ pub struct GrokInferenceTransportError {
status: Option<u16>,
retry_after: Option<Duration>,
http_version: Option<UpstreamHttpVersion>,
request_id: Option<SafeUpstreamValue>,
upstream_code: Option<Box<SafeUpstreamValue>>,
request_id: Option<OpaqueUpstreamValue>,
upstream_code: Option<Box<OpaqueUpstreamValue>>,
client_visible_upstream_error: Option<Box<ClientVisibleUpstreamError>>,
transport_metrics: GrokInferenceTransportMetrics,
credential_recovery_required: bool,
@@ -361,7 +361,7 @@ impl GrokInferenceTransportError {
pub fn with_response_facts(
mut self,
http_version: UpstreamHttpVersion,
request_id: Option<SafeUpstreamValue>,
request_id: Option<OpaqueUpstreamValue>,
) -> Self {
self.http_version = Some(http_version);
self.request_id = request_id;
@@ -370,7 +370,7 @@ impl GrokInferenceTransportError {
/// 附着从错误 JSON 中提取并清洗后的稳定机器码。
#[must_use]
pub fn with_upstream_code(mut self, code: SafeUpstreamValue) -> Self {
pub fn with_upstream_code(mut self, code: OpaqueUpstreamValue) -> Self {
self.upstream_code = Some(Box::new(code));
self
}
@@ -434,12 +434,12 @@ impl GrokInferenceTransportError {
}
#[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 upstream_code(&self) -> Option<&SafeUpstreamValue> {
pub fn upstream_code(&self) -> Option<&OpaqueUpstreamValue> {
self.upstream_code.as_deref()
}
@@ -289,6 +289,7 @@ async fn xai_admin_provider_validates_known_billing_breakdown() {
let billing = admin
.calculated_billing(&ProviderBillingInput {
upstream_model_id: "grok-4.5".to_owned(),
service_tier: None,
input_tokens: Some(100),
output_tokens: Some(10),
cached_tokens: Some(25),
@@ -18,8 +18,8 @@ use gateway_core::engine::{
UpstreamSendState,
};
use gateway_core::error::{
ClientVisibleUpstreamError, ContinuationFailure, ProviderError, ProviderErrorKind,
SafeUpstreamValue,
ClientVisibleUpstreamError, ContinuationFailure, OpaqueUpstreamValue, ProviderError,
ProviderErrorKind,
};
use gateway_core::event::{GatewayEvent, UpstreamHttpVersion};
use gateway_core::operation::{
@@ -1099,12 +1099,13 @@ async fn compaction_stream_should_reject_degenerate_summary_before_commit() {
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::Protocol);
assert!(error.replay_is_safe());
assert!(error.retries_same_account());
assert!(!error.replay_is_safe());
assert!(error.allows_pre_delivery_retry());
assert!(!error.retries_same_account());
}
#[tokio::test]
async fn compaction_transport_failure_should_retry_only_the_same_account() {
async fn compaction_transport_failure_should_allow_pre_delivery_recovery() {
let transport = StubInferenceTransport::stream_error(GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::Transport,
UpstreamSendState::Sent,
@@ -1121,12 +1122,13 @@ async fn compaction_transport_failure_should_retry_only_the_same_account() {
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::Transport);
assert!(error.replay_is_safe());
assert!(error.retries_same_account());
assert!(!error.replay_is_safe());
assert!(error.allows_pre_delivery_retry());
assert!(!error.retries_same_account());
}
#[tokio::test]
async fn compaction_protocol_stream_failure_should_retry_only_the_same_account() {
async fn compaction_protocol_stream_failure_should_allow_pre_delivery_recovery() {
let transport =
StubInferenceTransport::sequence([InferenceMode::SuccessBody(malformed_compaction_sse())]);
let provider = provider(StubSelector::success(), transport).await;
@@ -1141,8 +1143,9 @@ async fn compaction_protocol_stream_failure_should_retry_only_the_same_account()
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::Protocol);
assert!(error.replay_is_safe());
assert!(error.retries_same_account());
assert!(!error.replay_is_safe());
assert!(error.allows_pre_delivery_retry());
assert!(!error.retries_same_account());
}
#[tokio::test]
@@ -1421,8 +1424,7 @@ async fn transport_error_should_preserve_client_visible_upstream_detail() {
"You have run out of credits",
Some("personal_team_blocked_spending_limit".to_owned()),
Some("insufficient_quota".to_owned()),
)
.expect("safe upstream detail");
);
let error = mapped_transport_error(
GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::QuotaExhausted,
@@ -1762,9 +1764,7 @@ async fn connection_state_inherits_session_and_recovers_reasoning_on_pinned_acco
UpstreamSendState::Sent,
)
.with_status(400)
.with_upstream_code(
SafeUpstreamValue::new("reasoning_decode_failed").expect("safe code"),
),
.with_upstream_code(OpaqueUpstreamValue::new("reasoning_decode_failed")),
),
InferenceMode::Success,
]);
@@ -1982,7 +1982,7 @@ async fn missing_native_response_should_be_replay_safe_for_the_same_account() {
UpstreamSendState::Sent,
)
.with_status(404)
.with_upstream_code(SafeUpstreamValue::new("not_found").expect("safe code")),
.with_upstream_code(OpaqueUpstreamValue::new("not_found")),
);
let provider = provider(StubSelector::success(), transport).await;
let state = ProviderSessionState::new(
+1
View File
@@ -1,3 +1,4 @@
0554fed953bd1440b50bc4033dc8bad96b51e132c57c2ccf1553728a991d5cb2 0001_initial.sql
0ae3945c02461347ab8dcfc64821027cd2366505778e0450d71bd82772c4e99d 0002_drop_opaque_response_id_index.sql
e608ed89991f9d6317491714447a00a61ede90caca4fb030eb30695654afb664 0003_store_opaque_response_ids_as_bytes.sql
df4e02cf2335733326c12429ced4cbd0b35b8a22b3fc24981c3c0549477d7c43 0004_add_model_request_service_tier.sql
@@ -0,0 +1,9 @@
alter table model_requests
add column service_tier text,
add constraint model_requests_service_tier_ck check (
service_tier is null
or (
octet_length(service_tier) between 1 and 64
and service_tier !~ '[[:cntrl:]]'
)
);