Compare commits

...
5 Commits
Author SHA1 Message Date
ZyyandCodex 0cb5a29699 chore: bump version to 3.2.12
Co-authored-by: Codex <noreply@openai.com>
2026-08-19 23:52:11 +08:00
ZyyandCodex 61788804b4 fix(openai): recover cross-account websocket continuations
Co-authored-by: Codex <noreply@openai.com>
2026-08-19 23:40:40 +08:00
ZyyandCodex 2751ea6aa4 fix(openai): preserve cross-account session recovery
Co-authored-by: Codex <noreply@openai.com>
2026-08-19 23:05:48 +08:00
ZyyandCodex 58fe66789c chore(deps): bump h2 to 0.4.16
Resolve RUSTSEC-2026-0258 reported by cargo audit.

Co-authored-by: Codex <noreply@openai.com>
2026-08-19 11:34:07 +08:00
ZyyandCodex 573b6b26db fix(openai): isolate quota buckets by limit id
Co-authored-by: Codex <noreply@openai.com>
2026-08-18 16:48:37 +08:00
34 changed files with 1780 additions and 359 deletions
+6 -6
View File
@@ -1246,7 +1246,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -1686,9 +1686,9 @@ dependencies = [
[[package]]
name = "h2"
version = "0.4.15"
version = "0.4.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155"
checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27"
dependencies = [
"atomic-waker",
"bytes",
@@ -2714,7 +2714,7 @@ dependencies = [
"once_cell",
"socket2",
"tracing",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -2984,7 +2984,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -3628,7 +3628,7 @@ dependencies = [
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -73,7 +73,6 @@ pub(super) fn response_header_is_forwardable(name: &str, connection_options: &[S
| "x-openai-organization"
| "x-openai-project"
| "x-codex-installation-id"
| "x-codex-turn-state"
| "x-codex-turn-metadata"
)
}
@@ -788,6 +788,10 @@ async fn streaming_response_should_append_ordinary_headers_and_skip_unrepresenta
let trace = Arc::new(Trace::default());
let headers = vec![
ProviderResponseHeader::new("x-models-etag", Bytes::from_static(b"models-v2")),
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-from-upstream"),
),
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")),
@@ -834,6 +838,13 @@ async fn streaming_response_should_append_ordinary_headers_and_skip_unrepresenta
.map(HeaderValue::as_bytes),
Some(b"\xffopaque".as_slice())
);
assert_eq!(
response
.headers()
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok()),
Some("turn-state-from-upstream")
);
assert!(response.headers().get("authorization").is_none());
assert!(response.headers().get("connection").is_none());
assert!(response.headers().get("x-hop-secret").is_none());
@@ -411,6 +411,10 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
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-codex-turn-state",
Bytes::from_static(b"turn-state-from-upstream"),
),
ProviderResponseHeader::new("x-future-bytes", Bytes::from_static(b"\xffopaque")),
ProviderResponseHeader::new("bad\0name", Bytes::from_static(b"unrepresentable")),
ProviderResponseHeader::new(
@@ -483,6 +487,10 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
// 官方 metadata 是 string map,同名多值只能保持既有的后值覆盖语义。
assert_eq!(headers.get("x-future-multi"), Some(&json!("second")));
assert_eq!(headers.get("x-request-id"), Some(&json!("req_upstream")));
assert_eq!(
headers.get("x-codex-turn-state"),
Some(&json!("turn-state-from-upstream"))
);
for omitted in [
"x-future-bytes",
"bad\0name",
@@ -1162,10 +1162,11 @@ where
true
}
/// 原生续写的原账号在本地调度阶段已不可用时,改用完整输入重放。
/// 原生续写的原账号在本地调度阶段已不可用时,交给 Provider 执行跨账号恢复。
///
/// 这不是强制 Smart:下一次 attempt 仍使用请求配置的调度策略。只有网关持有
/// 对应 Provider 的可携带会话状态时,才可安全移除 native handle 并换号。
/// 对应 Provider 的会话状态时才进入恢复;是否保留 native handle、执行 probe
/// 或使用完整 transcript,由 Provider 自己的协议边界决定。
fn prepare_unavailable_native_continuation_replay(&mut self, error: &ProviderError) -> bool {
if self.account_selection.required_account().is_some()
|| self.continuation_attempt != ContinuationAttempt::Native
@@ -362,9 +362,14 @@ pub enum ContinuationAttempt {
None,
/// 使用 Store 解析出的原生 handle 与账号绑定。
Native,
/// 在原账号上用完整连接内 transcript 重放。
/// 在原账号上执行 Provider 定义的恢复。
///
/// 只有持有可携带 transcript 的 Provider 才能清除 native handle 后重放;
/// 只持有 opaque state 的 Provider 必须保留 continuation 依赖或返回客户端错误。
ReplayOwner,
/// transcript 已可携带,允许选择其他账号。
/// 允许选择其他账号并执行 Provider 定义的恢复。
///
/// 该枚举本身不证明客户端 input 完整,也不授权 Provider 清除 previous-response。
ReplayAny,
}
@@ -493,6 +493,31 @@ fn bare_atomic_response_failed(response_id: &str) -> ProviderError {
.with_atomic_client_events(vec![failed])
}
fn atomic_history_unavailable(response_id: &str) -> ProviderError {
let failed = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": response_id,
"status": "failed",
"error": {
"code": "previous_response_not_found",
"message": "Previous response was not found. Retrying the full request."
}
}
}),
)
.expect("history unavailable wire"),
);
ProviderError::new(ProviderErrorKind::InvalidRequest, UpstreamSendState::Sent)
.with_status(400)
.with_continuation_failure(ContinuationFailure::HistoryUnavailable)
.with_atomic_client_events(vec![failed])
}
fn image_stream(image_output_tokens: Option<u64>) -> Vec<Result<GatewayEvent, ProviderError>> {
let mut usage = Usage::new();
usage.input_tokens = Some(12);
@@ -1481,7 +1506,7 @@ fn streaming_native_continuation_delivers_started_event_before_later_output() {
}
#[test]
fn native_continuation_replays_owner_before_safely_switching_account() {
fn native_continuation_uses_provider_defined_owner_then_cross_account_recovery() {
let Operation::Generate(generate) = generate_operation() else {
panic!("generate operation");
};
@@ -1536,7 +1561,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("portable replay succeeds");
block_on(session.collect_uncommitted()).expect("provider-defined recovery succeeds");
block_on(session.commit_downstream(Some(200))).expect("commit response");
let contexts = provider.contexts.lock().expect("contexts lock");
@@ -1561,7 +1586,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
}
#[test]
fn unavailable_native_continuation_replays_with_the_configured_account_policy() {
fn unavailable_native_continuation_uses_provider_recovery_with_the_configured_account_policy() {
let Operation::Generate(generate) = generate_operation() else {
panic!("generate operation");
};
@@ -1602,7 +1627,7 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("portable replay succeeds");
block_on(session.collect_uncommitted()).expect("provider-defined recovery succeeds");
block_on(session.commit_downstream(Some(200))).expect("commit response");
let contexts = provider.contexts.lock().expect("contexts lock");
@@ -1625,6 +1650,76 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn continuation_history_failure_without_replay_proof_is_delivered_without_hidden_retry() {
let Operation::Generate(generate) = generate_operation() else {
panic!("generate operation");
};
let operation = Operation::Generate(generate.with_provider_session_state(
ProviderSessionState::new("openai", Map::new()).expect("provider session state"),
));
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_one",
items: vec![Err(atomic_history_unavailable("response-history-missing"))],
},
Script::Stream {
account_id: "acct_two",
items: complete_stream(None),
},
]);
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
Some(ContinuationBinding::Pinned(continuation)),
CancellationToken::new(),
))
.expect("start execution");
let terminal = block_on(session.next_event())
.expect("history failure remains deliverable")
.expect("terminal history failure batch");
assert_eq!(
terminal.commit_requirement(),
CommitRequirement::CommitBeforeDelivery
);
let events = terminal.into_provider_events();
let failure = events
.iter()
.filter_map(ProviderEvent::wire_event)
.find(|wire| wire.event_type() == Some("response.failed"))
.expect("raw response.failed event");
assert_eq!(
failure
.data()
.pointer("/response/error/code")
.and_then(Value::as_str),
Some("previous_response_not_found")
);
block_on(session.commit_downstream(Some(200))).expect("commit terminal history failure");
assert!(matches!(
block_on(session.next_event()),
Err(EngineError::Provider(_))
));
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 1);
assert_eq!(
contexts[0].continuation_attempt(),
ContinuationAttempt::Native
);
}
#[test]
fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
let operation = operation(RetrySafety::Idempotent);
@@ -208,11 +208,24 @@ pub struct RateLimitWindow {
pub reset_at: Option<i64>,
}
/// 限流项键的来源。
///
/// HTTP headers 和事件里的 `metered_limit_name` 提供稳定 `limit_id`;
/// WebSocket `additional_rate_limits` 在部分协议版本里只以可读名称为 map key,
/// 此时必须由持久化层用既有名称元数据解析,不能把名称误当成稳定 ID。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RateLimitKeySource {
LimitId,
LimitName,
}
/// 单个计量项的限流信息。
#[derive(Debug, Clone, PartialEq)]
pub struct RateLimitDetails {
/// 稳定的计量项 ID。
pub limit_id: String,
/// [`limit_id`](Self::limit_id) 是稳定 ID,还是名称生成的临时解析键。
pub key_source: RateLimitKeySource,
/// 上游提供的可读名称。
pub limit_name: Option<String>,
/// 当前请求是否被允许。
@@ -341,20 +354,40 @@ pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
return None;
}
let limit_id = rate_limit_name(value)
.map(normalize_limit_id)
let explicit_limit_id = top_level_rate_limit_id(value).map(normalize_limit_id);
let explicit_limit_name = non_empty_string(value, "limit_name");
let default_limit_id = explicit_limit_id
.clone()
.unwrap_or_else(|| "codex".to_string());
let limit_name = value
.get("metered_limit_name")
.and_then(Value::as_str)
.and_then(|_| value.get("limit_name"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|name| !name.is_empty())
.map(ToString::to_string);
let details = value
.get("rate_limits")
.and_then(|details| parse_details_from_object(details, &limit_id, limit_name));
let default_details = value.get("rate_limits").and_then(|details| {
parse_details_from_object(
details,
&default_limit_id,
explicit_limit_name,
RateLimitKeySource::LimitId,
)
});
let mut limits = parse_additional_rate_limits(value);
let matching_additional_ids = default_details
.as_ref()
.map(|default| {
limits
.iter()
.filter(|(_, additional)| same_rate_limit_facts(default, additional))
.map(|(limit_id, _)| limit_id.clone())
.collect::<Vec<_>>()
})
.unwrap_or_default();
// 当前 WebSocket 协议会把活动具名桶同时镜像到顶层 `rate_limits`,但不附
// `metered_limit_name`。只要 additional 中存在相同事实,顶层就不是 core 桶。
let default_is_additional_alias =
explicit_limit_id.is_none() && !matching_additional_ids.is_empty();
if let Some(details) = default_details
&& !default_is_additional_alias
{
limits.insert(default_limit_id.clone(), details);
}
let credits = value.get("credits").and_then(parse_credits_from_object);
let plan_type = value
.get("plan_type")
@@ -363,17 +396,18 @@ pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
.filter(|plan| !plan.is_empty())
.map(ToString::to_string);
if details.is_none() && credits.is_none() && plan_type.is_none() {
if limits.is_empty() && credits.is_none() && plan_type.is_none() {
return None;
}
let mut limits = BTreeMap::new();
if let Some(details) = details {
limits.insert(limit_id.clone(), details);
}
let active_limit = explicit_limit_id.or_else(|| match matching_additional_ids.as_slice() {
[limit_id] => Some(limit_id.clone()),
[] if !limits.is_empty() => Some(default_limit_id),
_ => None,
});
Some(ParsedRateLimits {
limits,
active_limit: Some(limit_id),
active_limit,
credits,
plan_type,
promo_message: value
@@ -597,6 +631,7 @@ fn parse_details_from_lookup(
}
Some(RateLimitDetails {
limit_id: limit_id.to_string(),
key_source: RateLimitKeySource::LimitId,
limit_name,
allowed,
limit_reached,
@@ -687,6 +722,7 @@ fn parse_details_from_object(
value: &Value,
limit_id: &str,
limit_name: Option<String>,
key_source: RateLimitKeySource,
) -> Option<RateLimitDetails> {
let primary = value.get("primary").and_then(parse_window_from_object);
let secondary = value.get("secondary").and_then(parse_window_from_object);
@@ -697,6 +733,7 @@ fn parse_details_from_object(
}
Some(RateLimitDetails {
limit_id: limit_id.to_string(),
key_source,
limit_name,
allowed,
limit_reached,
@@ -752,11 +789,59 @@ fn value_as_positive_i64(value: &Value) -> Option<i64> {
.filter(|value| *value > 0)
}
fn rate_limit_name(value: &Value) -> Option<&str> {
fn rate_limit_id(value: &Value) -> Option<&str> {
value
.get("metered_limit_name")
.or_else(|| value.get("limit_name"))
.or_else(|| value.get("limit_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn top_level_rate_limit_id(value: &Value) -> Option<&str> {
rate_limit_id(value).or_else(|| {
value
.get("limit_name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
})
}
fn non_empty_string(value: &Value, field: &str) -> Option<String> {
value
.get(field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn parse_additional_rate_limits(value: &Value) -> BTreeMap<String, RateLimitDetails> {
value
.get("additional_rate_limits")
.and_then(Value::as_object)
.into_iter()
.flatten()
.filter_map(|(map_key, value)| {
let explicit_limit_id = rate_limit_id(value).map(normalize_limit_id);
let (limit_id, key_source) = explicit_limit_id.map_or_else(
|| (normalize_limit_id(map_key), RateLimitKeySource::LimitName),
|limit_id| (limit_id, RateLimitKeySource::LimitId),
);
let limit_name = non_empty_string(value, "limit_name")
.or_else(|| Some(map_key.trim().to_owned()).filter(|name| !name.is_empty()));
parse_details_from_object(value, &limit_id, limit_name, key_source)
.map(|details| (limit_id, details))
})
.collect()
}
fn same_rate_limit_facts(left: &RateLimitDetails, right: &RateLimitDetails) -> bool {
left.allowed == right.allowed
&& left.limit_reached == right.limit_reached
&& left.primary == right.primary
&& left.secondary == right.secondary
}
fn normalize_limit_id(value: impl AsRef<str>) -> String {
@@ -1,7 +1,7 @@
use gateway_protocol::openai::events::{
RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage, is_codex_quota_header_name,
is_rate_limit_header_name, parse_rate_limit_headers, parse_rate_limits_event,
rate_limits_to_header_pairs, retry_after_seconds_from_body,
RateLimitKeySource, RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage,
is_codex_quota_header_name, is_rate_limit_header_name, parse_rate_limit_headers,
parse_rate_limits_event, rate_limits_to_header_pairs, retry_after_seconds_from_body,
};
use serde_json::json;
@@ -483,6 +483,58 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
);
}
#[test]
fn parse_rate_limits_event_should_preserve_name_keyed_additional_limit() {
let event = json!({
"type": "codex.rate_limits",
"plan_type": "pro",
"rate_limits": {
"allowed": true,
"limit_reached": false,
"primary": {
"used_percent": 0,
"window_minutes": 10080,
"reset_after_seconds": 354916,
"reset_at": 1787396838
},
"secondary": null
},
"additional_rate_limits": {
"GPT-5.3-Codex-Spark": {
"allowed": true,
"limit_reached": false,
"primary": {
"used_percent": 0,
"window_minutes": 10080,
"reset_after_seconds": 354916,
"reset_at": 1787396838
},
"secondary": null
}
},
"credits": null
});
let parsed = parse_rate_limits_event(&event).expect("rate-limit event");
let details = parsed
.limits
.get("gpt_5.3_codex_spark")
.expect("name-keyed additional limit");
assert_eq!(parsed.active_limit.as_deref(), Some("gpt_5.3_codex_spark"));
assert_eq!(details.key_source, RateLimitKeySource::LimitName);
assert_eq!(details.limit_name.as_deref(), Some("GPT-5.3-Codex-Spark"));
assert_eq!(
details.primary,
Some(RateLimitWindow {
used_percent: 0.0,
window_minutes: Some(10080),
reset_at: Some(1787396838),
})
);
assert!(!parsed.limits.contains_key("codex"));
}
#[test]
fn parse_rate_limits_event_should_reject_unrelated_event_type() {
for event in [
@@ -0,0 +1,173 @@
//! OpenAI quota 文档的规范化 map 与上游 `/usage` 输入适配。
use std::collections::BTreeMap;
use gateway_protocol::openai::events::{RateLimitDetails, RateLimitKeySource};
use serde_json::{Map, Value};
pub(super) const DEFAULT_CODEX_LIMIT_ID: &str = "codex";
pub(super) const RATE_LIMITS_BY_LIMIT_ID: &str = "rate_limits_by_limit_id";
/// 官方新协议的多桶视图:每个 `limit_id` 对应一个独立快照。
#[derive(Debug, Default)]
pub(super) struct RateLimitSnapshotsByLimitId {
snapshots: BTreeMap<String, Map<String, Value>>,
}
impl RateLimitSnapshotsByLimitId {
/// 接收上游 `/usage` 形状或已经规范化的 map;输出始终只保留 map。
pub(super) fn take_from_document(quota: &mut Map<String, Value>) -> Self {
let mut result = Self::default();
if let Some(by_limit_id) = quota
.remove(RATE_LIMITS_BY_LIMIT_ID)
.and_then(|value| value.as_object().cloned())
{
for (raw_limit_id, value) in by_limit_id {
let Some(limit_id) = normalize_limit_id(&raw_limit_id) else {
continue;
};
let mut snapshot = snapshot_object(value);
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.clone()));
result.snapshots.insert(limit_id, snapshot);
}
}
if let Some(rate_limit) = quota.remove("rate_limit") {
result.insert_usage_snapshot(DEFAULT_CODEX_LIMIT_ID, None, rate_limit);
}
if let Some(rate_limit) = quota.remove("code_review_rate_limit") {
result.insert_usage_snapshot("code_review", Some("code_review"), rate_limit);
}
let additional = quota
.remove("additional_rate_limits")
.and_then(|value| value.as_array().cloned())
.unwrap_or_default();
for (index, value) in additional.into_iter().enumerate() {
let Value::Object(mut item) = value else {
result.insert_usage_snapshot(&format!("additional_{index}"), None, value);
continue;
};
let limit_id = item_limit_id(&item).unwrap_or_else(|| format!("additional_{index}"));
let limit_name = item
.get("limit_name")
.and_then(Value::as_str)
.map(str::to_owned);
let rate_limit = item.remove("rate_limit").unwrap_or(Value::Null);
result.insert_usage_snapshot(&limit_id, limit_name.as_deref(), rate_limit);
}
result
}
fn insert_usage_snapshot(
&mut self,
limit_id: &str,
limit_name: Option<&str>,
rate_limit: Value,
) {
let Some(limit_id) = normalize_limit_id(limit_id) else {
return;
};
let mut snapshot = snapshot_object(rate_limit);
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.clone()));
if let Some(limit_name) = limit_name {
snapshot.insert(
"limit_name".to_owned(),
Value::String(limit_name.to_owned()),
);
}
// map 协议中一个 ID 只能有一个快照;规范化 map 与顶层账号桶优先。
self.snapshots.entry(limit_id).or_insert(snapshot);
}
pub(super) fn upsert(
&mut self,
limit_id: &str,
details: &RateLimitDetails,
rate_limit: Map<String, Value>,
) {
let mut snapshot = self.snapshots.remove(limit_id).unwrap_or_default();
for field in [
"allowed",
"limit_reached",
"primary_window",
"secondary_window",
] {
snapshot.remove(field);
}
snapshot.extend(rate_limit);
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.to_owned()));
if let Some(name) = details.limit_name.as_ref() {
snapshot.insert("limit_name".to_owned(), Value::String(name.clone()));
}
self.snapshots.insert(limit_id.to_owned(), snapshot);
}
/// 把 wire 观察解析到稳定 `limit_id`。
///
/// WebSocket 的 `additional_rate_limits` 可能只返回可读名称;这时只接受
/// `/usage` 或 HTTP headers 已建立的唯一名称映射。无法唯一解析就不落库,
/// 避免凭名称猜 ID 或污染默认 `codex` 桶。
pub(super) fn resolve_limit_id(&self, details: &RateLimitDetails) -> Option<String> {
if details.key_source == RateLimitKeySource::LimitId {
return normalize_limit_id(&details.limit_id);
}
let limit_name = details.limit_name.as_deref()?.trim();
let mut matches = self
.snapshots
.iter()
.filter(|(_, snapshot)| {
snapshot
.get("limit_name")
.and_then(Value::as_str)
.is_some_and(|stored| stored.trim().eq_ignore_ascii_case(limit_name))
})
.map(|(limit_id, _)| limit_id.clone());
let resolved = matches.next()?;
matches.next().is_none().then_some(resolved)
}
pub(super) fn write_to_document(self, quota: &mut Map<String, Value>) {
let by_limit_id = self
.snapshots
.into_iter()
.map(|(limit_id, snapshot)| (limit_id, Value::Object(snapshot)))
.collect::<Map<_, _>>();
quota.insert(
RATE_LIMITS_BY_LIMIT_ID.to_owned(),
Value::Object(by_limit_id),
);
}
}
pub(super) fn canonicalize_rate_limit_document(
mut quota: Map<String, Value>,
) -> Map<String, Value> {
let snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota);
snapshots.write_to_document(&mut quota);
quota
}
fn snapshot_object(value: Value) -> Map<String, Value> {
value.as_object().cloned().unwrap_or_else(|| {
let mut snapshot = Map::new();
snapshot.insert("rate_limit".to_owned(), value);
snapshot
})
}
fn item_limit_id(item: &Map<String, Value>) -> Option<String> {
item.get("limit_id")
.or_else(|| item.get("metered_feature"))
.or_else(|| item.get("limit_name"))
.and_then(Value::as_str)
.and_then(normalize_limit_id)
}
fn normalize_limit_id(value: &str) -> Option<String> {
let value = value.trim();
(!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control))
.then(|| value.to_ascii_lowercase().replace('-', "_"))
}
@@ -1,8 +1,10 @@
//! Codex quota 服务编排:权威额度事实、展示投影、主动/被动同步与 429 冷却。
//!
//! - [`document`]:上游 `/usage` 输入到多桶 map 的单向规范化。
//! - [`snapshot`]:快照/窗口解析、聚合、O(1) 滚动与调度信号。
//! - [`evidence`]:额度接口错误到凭据/额度事实的分类。
mod document;
pub(crate) mod evidence;
pub(crate) mod snapshot;
@@ -38,6 +40,10 @@ use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext
use super::agent_identity::{CodexAgentIdentityTaskService, PreparedCodexRuntimeCredential};
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
use document::{
DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, RateLimitSnapshotsByLimitId,
canonicalize_rate_limit_document,
};
use evidence::{QuotaEndpointFailure, classify_quota_endpoint_failure};
use snapshot::{
parse_account_quota_snapshot, quota_projection_ttl, quota_snapshot_from_observation,
@@ -706,10 +712,25 @@ impl CodexCredentialQuotaService {
let Some(rate_limits) = parse_rate_limit_headers(headers) else {
return Ok(false);
};
let has_quota_facts = rate_limits
.limits
.values()
.any(|details| passive_rate_limit_snapshot(details).is_some());
self.synchronize_passive_rate_limits(account, std::slice::from_ref(&rate_limits))
.await
}
/// 把一次推理响应中采集的结构化限流观察合并后单次落库。
pub async fn synchronize_passive_rate_limits(
&self,
account: &ProviderAccount,
rate_limits: &[ParsedRateLimits],
) -> Result<bool, CodexCredentialQuotaError> {
if rate_limits.is_empty() {
return Ok(false);
}
let has_quota_facts = rate_limits.iter().any(|observation| {
observation
.limits
.values()
.any(|details| passive_rate_limit_snapshot(details).is_some())
});
let existing = self
.store
.get_quotas(std::slice::from_ref(account.id()))
@@ -734,14 +755,14 @@ impl CodexCredentialQuotaService {
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(merge_passive_quota(existing, &rate_limits)),
quota: OpaqueProviderData::new(merge_passive_quota(existing, rate_limits)),
observed_at: SystemTime::now(),
state,
})
.await?;
return Ok(outcome != QuotaWriteOutcome::Conflict);
}
let quota = merge_passive_quota(existing, &rate_limits);
let quota = merge_passive_quota(existing, rate_limits);
let observed_at = SystemTime::now();
let snapshot = parse_account_quota_snapshot(
account.id().clone(),
@@ -1177,18 +1198,48 @@ fn access_token_is_current(account: &ProviderAccount, now: SystemTime) -> bool {
fn merge_passive_quota(
mut quota: Map<String, Value>,
rate_limits: &ParsedRateLimits,
observations: &[ParsedRateLimits],
) -> Map<String, Value> {
let active_limit = rate_limits
.active_limit
.as_deref()
.or_else(|| rate_limits.limits.contains_key("codex").then_some("codex"));
if let Some(active_limit) = active_limit {
quota.insert(
"active_limit".to_owned(),
Value::String(active_limit.to_owned()),
);
let mut snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota);
for rate_limits in observations {
let default_is_named_alias = default_limit_is_named_alias(rate_limits);
let mut resolved_limit_ids = BTreeMap::new();
for (wire_limit_id, details) in &rate_limits.limits {
// HTTP 与 WebSocket 都可能把活动具名桶镜像为默认 `codex` 窗口。
// 同一 wire 观察里存在相同具名事实时丢弃镜像,不触碰 core 桶。
if wire_limit_id == DEFAULT_CODEX_LIMIT_ID && default_is_named_alias {
continue;
}
let Some(limit_id) = snapshots.resolve_limit_id(details) else {
continue;
};
let Some(rate_limit) = passive_rate_limit_snapshot(details) else {
continue;
};
snapshots.upsert(&limit_id, details, rate_limit);
resolved_limit_ids.insert(wire_limit_id.as_str(), limit_id);
}
let active_limit = rate_limits
.active_limit
.as_deref()
.and_then(|wire_limit_id| {
resolved_limit_ids.get(wire_limit_id).cloned().or_else(|| {
(!rate_limits.limits.contains_key(wire_limit_id))
.then(|| wire_limit_id.to_owned())
})
})
.or_else(|| resolved_limit_ids.get(DEFAULT_CODEX_LIMIT_ID).cloned());
if let Some(active_limit) = active_limit {
quota.insert("active_limit".to_owned(), Value::String(active_limit));
}
merge_passive_metadata(&mut quota, rate_limits);
}
snapshots.write_to_document(&mut quota);
quota
}
fn merge_passive_metadata(quota: &mut Map<String, Value>, rate_limits: &ParsedRateLimits) {
if let Some(plan_type) = rate_limits.plan_type.as_ref() {
quota.insert("plan_type".to_owned(), Value::String(plan_type.clone()));
}
@@ -1216,48 +1267,19 @@ fn merge_passive_quota(
}
quota.insert("credits".to_owned(), Value::Object(value));
}
}
let mut additional = quota
.remove("additional_rate_limits")
.and_then(|value| value.as_array().cloned())
.unwrap_or_default();
for (limit_id, details) in &rate_limits.limits {
let Some(rate_limit) = passive_rate_limit_snapshot(details) else {
continue;
};
if Some(limit_id.as_str()) == active_limit {
quota.insert("rate_limit".to_owned(), Value::Object(rate_limit));
continue;
}
let index = additional.iter().position(|item| {
item.get("metered_feature")
.or_else(|| item.get("limit_name"))
.and_then(Value::as_str)
.is_some_and(|existing| existing == limit_id)
});
let mut item = index
.and_then(|index| additional.get(index))
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
item.insert(
"metered_feature".to_owned(),
Value::String(limit_id.clone()),
);
if let Some(name) = details.limit_name.as_ref() {
item.insert("limit_name".to_owned(), Value::String(name.clone()));
}
item.insert("rate_limit".to_owned(), Value::Object(rate_limit));
match index {
Some(index) => additional[index] = Value::Object(item),
None => additional.push(Value::Object(item)),
}
}
quota.insert(
"additional_rate_limits".to_owned(),
Value::Array(additional),
);
quota
fn default_limit_is_named_alias(rate_limits: &ParsedRateLimits) -> bool {
let Some(default_primary) = rate_limits
.limits
.get(DEFAULT_CODEX_LIMIT_ID)
.and_then(|details| details.primary)
else {
return false;
};
rate_limits.limits.iter().any(|(limit_id, details)| {
limit_id != DEFAULT_CODEX_LIMIT_ID && details.primary == Some(default_primary)
})
}
/// 丢弃 core 限额中上游给出的无事实 `secondary_window` 占位。
@@ -1269,7 +1291,13 @@ fn merge_passive_quota(
/// 无事实值。带 reset、时长、正用量、触顶、未知或非法字段的次级窗口均完全按
/// 原有额度逻辑保留。
fn normalize_quota_window_placeholders(mut quota: Map<String, Value>) -> Map<String, Value> {
if let Some(rate_limit) = quota.get_mut("rate_limit").and_then(Value::as_object_mut) {
quota = canonicalize_rate_limit_document(quota);
if let Some(rate_limit) = quota
.get_mut(RATE_LIMITS_BY_LIMIT_ID)
.and_then(Value::as_object_mut)
.and_then(|limits| limits.get_mut(DEFAULT_CODEX_LIMIT_ID))
.and_then(Value::as_object_mut)
{
drop_secondary_window_placeholder(rate_limit);
}
quota
@@ -12,6 +12,9 @@ use gateway_core::engine::credential::{
};
use serde_json::{Map, Value};
use super::document::{
DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, canonicalize_rate_limit_document,
};
use super::{CodexCredentialQuotaError, QUOTA_SCHEDULING_TTL};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -173,9 +176,18 @@ impl CodexQuotaFact {
}
pub fn parse_codex_quota_usage(usage: &Value) -> Result<CodexQuotaFact, CodexCredentialQuotaError> {
let object = usage
.as_object()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
let object = canonicalize_rate_limit_document(
usage
.as_object()
.cloned()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?,
);
parse_codex_quota_object(&object)
}
fn parse_codex_quota_object(
object: &Map<String, Value>,
) -> Result<CodexQuotaFact, CodexCredentialQuotaError> {
let mut aggregate = QuotaAggregate::default();
for limit in canonical_rate_limits(object)? {
aggregate.observe_rate_limit(limit.rate_limit)?;
@@ -196,6 +208,15 @@ pub fn parse_codex_quota_usage(usage: &Value) -> Result<CodexQuotaFact, CodexCre
})
}
fn canonical_quota_object(usage: &Value) -> Result<Map<String, Value>, CodexCredentialQuotaError> {
Ok(canonicalize_rate_limit_document(
usage
.as_object()
.cloned()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?,
))
}
pub(crate) fn quota_snapshot_from_observation(
observation: &QuotaObservation,
) -> Option<CodexAccountQuotaSnapshot> {
@@ -277,13 +298,11 @@ pub(crate) fn parse_account_quota_snapshot(
observed_at: SystemTime,
usage: &Value,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
let fact = parse_codex_quota_usage(usage)?;
let object = usage
.as_object()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
let quota = authoritative_quota_state(object, observed_at)?;
let object = canonical_quota_object(usage)?;
let fact = parse_codex_quota_object(&object)?;
let quota = authoritative_quota_state(&object, observed_at)?;
let mut windows = Vec::new();
for limit in canonical_rate_limits(object)? {
for limit in canonical_rate_limits(&object)? {
parse_rate_limit_windows(
&limit.key,
&limit.source,
@@ -316,12 +335,16 @@ fn quota_source_order(source: &str) -> u8 {
}
}
/// 顶层 account-wide Codex 桶是唯一可改变账号额度访问结论的 quota 响应字段。
/// 默认 `codex` 桶是唯一可改变账号额度访问结论的 quota 响应字段。
fn authoritative_quota_state(
usage: &Map<String, Value>,
observed_at: SystemTime,
) -> Result<QuotaState, CodexCredentialQuotaError> {
let Some(rate_limit) = usage.get("rate_limit") else {
let Some(rate_limit) = usage
.get(RATE_LIMITS_BY_LIMIT_ID)
.and_then(Value::as_object)
.and_then(|limits| limits.get(DEFAULT_CODEX_LIMIT_ID))
else {
return Ok(QuotaState::observed_unknown(observed_at));
};
if rate_limit.is_null() {
@@ -398,64 +421,32 @@ fn canonical_rate_limits(
object: &Map<String, Value>,
) -> Result<Vec<CanonicalRateLimit<'_>>, CodexCredentialQuotaError> {
let mut limits = Vec::new();
let has_primary_codex_rate_limit = object
.get("rate_limit")
.is_some_and(|value| !value.is_null());
if let Some(rate_limit) = object.get("rate_limit") {
// 官方 Codex 将顶层 `rate_limit` 固定识别为 `codex`,而不把
// `active_limit` 投影成另一个展示桶。后者只是当前活动套餐的元数据。
limits.push(CanonicalRateLimit {
key: "core".to_owned(),
source: "codex".to_owned(),
account_wide: true,
limit_name: None,
rate_limit,
});
}
if let Some(rate_limit) = object
.get("code_review_rate_limit")
.filter(|value| !value.is_null())
{
limits.push(CanonicalRateLimit {
key: "code-review".to_owned(),
source: "code_review".to_owned(),
account_wide: false,
limit_name: Some("code_review".to_owned()),
rate_limit,
});
}
let Some(additional) = object
.get("additional_rate_limits")
.filter(|value| !value.is_null())
else {
let Some(by_limit_id) = object.get(RATE_LIMITS_BY_LIMIT_ID) else {
return Ok(limits);
};
for (index, value) in additional
.as_array()
for (limit_id, rate_limit) in by_limit_id
.as_object()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?
.iter()
.enumerate()
{
let item = value
let item = rate_limit
.as_object()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
let (source, limit_name) = canonical_rate_limit_source(item, index);
// 部分上游响应会在 `additional_rate_limits` 中重复默认的 `codex`
// 桶。它没有独立的 limit id,不能在账号面板再次展示。
if has_primary_codex_rate_limit && is_codex_limit_id(&source) {
continue;
}
if let Some(rate_limit) = item.get("rate_limit") {
// 非 Codex additional 桶保留为独立快照;额外用索引生成窗口键,
// 避免同名 additional 产生重复前端 key。
limits.push(CanonicalRateLimit {
key: format!("additional-{index}-{source}"),
account_wide: is_codex_limit_id(&source),
source,
limit_name,
rate_limit,
});
}
let source = item
.get("limit_id")
.and_then(Value::as_str)
.unwrap_or(limit_id)
.to_owned();
let limit_name = item
.get("limit_name")
.and_then(Value::as_str)
.map(str::to_owned);
limits.push(CanonicalRateLimit {
key: source.clone(),
account_wide: is_codex_limit_id(&source),
source,
limit_name,
rate_limit,
});
}
Ok(limits)
}
@@ -464,30 +455,6 @@ fn is_codex_limit_id(value: &str) -> bool {
value.trim().eq_ignore_ascii_case("codex")
}
fn canonical_rate_limit_source(
item: &Map<String, Value>,
index: usize,
) -> (String, Option<String>) {
let source = item
.get("metered_feature")
.or_else(|| item.get("limit_name"))
.and_then(Value::as_str)
.filter(|value| {
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
})
.map(str::to_owned)
.unwrap_or_else(|| format!("additional-{index}"));
let limit_name = item
.get("limit_name")
.or_else(|| item.get("metered_feature"))
.and_then(Value::as_str)
.filter(|value| {
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
})
.map(str::to_owned);
(source, limit_name)
}
fn parse_rate_limit_windows(
key_source: &str,
source: &str,
+198 -78
View File
@@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use async_trait::async_trait;
use bytes::Bytes;
use futures::{StreamExt, future::BoxFuture};
use gateway_core::engine::continuation::{ContinuationBinding, NativeContinuationScope};
use gateway_core::engine::credential::{AccountFeedbackStats, ProviderAccount};
@@ -33,6 +34,9 @@ use gateway_core::task::{
WorkerDefinitionError, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration,
WorkerRunnable, WorkerSchedule, WorkerTaskError,
};
use gateway_protocol::openai::events::{
ParsedRateLimits, parse_rate_limit_headers, rate_limits_to_header_pairs,
};
use reqwest::Client;
use secrecy::{ExposeSecret, SecretString};
use serde::{Deserialize, Serialize};
@@ -63,16 +67,16 @@ use crate::transport::protocol::responses::{
CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals,
};
use crate::transport::request::{
CodexRequestEncodeError, encode_generate_request, scope_request_to_account,
CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account,
};
use crate::transport::session::CodexSessionIdentity;
use crate::transport::usage::normalize_service_tier;
use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason};
use crate::transport::{
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError,
CodexRateLimitHeaderUpdates, CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics,
CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexRateLimitUpdates,
CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics, CodexUpstreamDiagnostics,
CodexWebSocketPool, endpoint_url,
};
const PROVIDER_NAME: &str = "openai";
@@ -247,7 +251,12 @@ impl Provider for CodexProvider {
if let Some(identity) = &self.session_identity {
identity.prepare_local_conversation(&mut upstream_request);
}
if upstream_request.turn_state.is_none() {
if previous_session.as_ref().is_some_and(|state| {
same_client_turn(
state.client_turn_id.as_deref(),
upstream_request.client_turn_id.as_deref(),
)
}) {
upstream_request.turn_state = previous_session
.as_ref()
.and_then(|state| state.turn_state.clone());
@@ -277,57 +286,48 @@ impl Provider for CodexProvider {
let output_started_at = Instant::now();
let provider_kind = ProviderKind::new(PROVIDER_NAME)
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?;
let cross_account = context
let state_owner_cross_account = context
.account_state_owner()
.is_some_and(|owner| !owner.matches(&provider_kind, lease.account_id()))
|| previous_session
.as_ref()
.is_some_and(|state| state.account_id != lease.account_id().as_str());
let replay_previous_response = matches!(
context.continuation_attempt(),
ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny
) || previous_session.as_ref().is_some_and(|state| {
state.continuation_scope == OpenAiContinuationScope::ReplayRequired
});
if replay_previous_response && previous_session.is_some() {
// Codex 客户端 input 是完整历史;恢复时只解除上游续链绑定,避免重复叠加代理副本。
upstream_request.set_previous_response_id(None);
upstream_request.previous_response_scope = None;
upstream_request.turn_state = None;
}
scope_request_to_account(
&mut upstream_request,
lease.installation_id(),
cross_account,
);
if context.continuation_attempt() == ContinuationAttempt::Native
&& !replay_previous_response
let has_explicit_state_owner =
context.account_state_owner().is_some() || previous_session.is_some();
let cross_account =
state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch());
if context.continuation_attempt() != ContinuationAttempt::None
&& let Some(continuation) = context.continuation()
{
match continuation {
ContinuationBinding::Pinned(continuation) => {
let previous_response_scope = match previous_session
.as_ref()
.map(|state| state.continuation_scope)
{
Some(OpenAiContinuationScope::Persisted) => {
PreviousResponseScope::Persisted
}
Some(OpenAiContinuationScope::ConnectionLocal) => {
PreviousResponseScope::ConnectionLocal
}
Some(OpenAiContinuationScope::ReplayRequired) => {
return Err(provider_error(
ProviderErrorKind::Protocol,
UpstreamSendState::NotSent,
));
}
None => match continuation.scope() {
NativeContinuationScope::Persisted => PreviousResponseScope::Persisted,
NativeContinuationScope::ConnectionLocal => {
let previous_response_scope = match context.continuation_attempt() {
ContinuationAttempt::Native => match previous_session
.as_ref()
.map(|state| state.continuation_scope)
{
Some(OpenAiContinuationScope::Persisted) => {
PreviousResponseScope::Persisted
}
Some(OpenAiContinuationScope::ConnectionLocal) => {
PreviousResponseScope::ConnectionLocal
}
Some(OpenAiContinuationScope::ReplayRequired) => {
PreviousResponseScope::ExternalUnknown
}
None => match continuation.scope() {
NativeContinuationScope::Persisted => {
PreviousResponseScope::Persisted
}
NativeContinuationScope::ConnectionLocal => {
PreviousResponseScope::ConnectionLocal
}
},
},
ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny => {
PreviousResponseScope::ExternalUnknown
}
ContinuationAttempt::None => PreviousResponseScope::ExternalUnknown,
};
upstream_request.set_previous_response_id(Some(
continuation.upstream_response_id().as_str().to_owned(),
@@ -342,6 +342,18 @@ impl Provider for CodexProvider {
}
}
}
let account_scope = if !cross_account {
RequestAccountScope::SameAccount
} else if upstream_request.previous_response_id().is_some() {
RequestAccountScope::CrossAccountContinuationProbe
} else {
RequestAccountScope::CrossAccountFullRequest
};
scope_request_to_account(
&mut upstream_request,
lease.installation_id(),
account_scope,
);
let metadata = ProviderCallMetadata::new(
provider_kind,
candidate.upstream_model().clone(),
@@ -359,6 +371,7 @@ impl Provider for CodexProvider {
account_id: lease.account_id().as_str().to_owned(),
conversation_id: upstream_request.local_conversation_id.clone(),
turn_state: upstream_request.turn_state.clone(),
client_turn_id: upstream_request.client_turn_id.clone(),
response_store,
continuation_scope: None,
});
@@ -411,6 +424,8 @@ struct OpenAiSessionState {
conversation_id: Option<String>,
#[serde(default)]
turn_state: Option<String>,
#[serde(default)]
client_turn_id: Option<String>,
continuation_scope: OpenAiContinuationScope,
}
@@ -426,10 +441,17 @@ struct OpenAiSessionCapture {
account_id: String,
conversation_id: Option<String>,
turn_state: Option<String>,
client_turn_id: Option<String>,
response_store: bool,
continuation_scope: Option<OpenAiContinuationScope>,
}
fn same_client_turn(previous: Option<&str>, current: Option<&str>) -> bool {
previous
.zip(current)
.is_some_and(|(previous, current)| !previous.is_empty() && previous == current)
}
fn decode_openai_session_state(request: &GenerateRequest) -> Option<OpenAiSessionState> {
request
.provider_session_state(PROVIDER_NAME)
@@ -471,6 +493,7 @@ fn attach_openai_session_update(
account_id: capture.account_id,
conversation_id: capture.conversation_id,
turn_state: capture.turn_state,
client_turn_id: capture.client_turn_id,
continuation_scope,
}) else {
return;
@@ -698,9 +721,6 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
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;
}
if let Some(etag) = response.response_metadata.models_etag.as_deref()
&& let Err(error) = catalog.observe_response_etag(etag)
{
@@ -728,7 +748,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let failure_diagnostics = response.diagnostics.clone();
let failure_set_cookie_headers = response.set_cookie_headers.clone();
let failure_rate_limit_headers = response.rate_limit_headers.clone();
let rate_limit_updates = response.rate_limit_header_updates;
let mut passive_quota_observation =
OpenAiPassiveQuotaObservation::new(response.rate_limit_headers);
let rate_limit_updates = response.rate_limit_updates;
let turn_state_updates = response.turn_state_update;
// OpenAI 线路为透明代理:HTTP SSE 与 WebSocket 两条上游均启用 raw 透传,
// 下游按字节转发上游原文,避免 serde 往返改写数值/精度(大整数→f64、logprobs 等)。
@@ -741,6 +763,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let mut pre_commit_events = PreCommitClientEvents::new();
loop {
let Some(stream_deadline) = remaining(context.deadline()) else {
if allows_account_state_mutation {
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
}
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
return;
};
@@ -773,15 +803,22 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(failure) => {
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
if !updates.is_empty() {
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
if observation_state.merge_rate_limit_headers(&updates)
passive_quota_observation.observe(&updates);
let update_headers = rate_limit_update_headers(&updates);
if observation_state.merge_rate_limit_headers(&update_headers)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
}
}
if allows_account_state_mutation {
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
}
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
@@ -793,18 +830,19 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let rate_limits_changed = if updates.is_empty() {
false
} else {
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
passive_quota_observation.observe(&updates);
observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates))
};
if let (Some(capture), Some(updates)) =
(session_capture.as_mut(), turn_state_updates.as_ref())
// take:仅在上游写入过新 turn_state 时更新 capture,避免逐 chunk 克隆。
let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref()
&& let Some(turn_state) = updates.lock().await.take()
{
capture.turn_state = Some(turn_state);
}
if let Some(capture) = session_capture.as_mut() {
capture.turn_state = Some(turn_state.clone());
}
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
} else {
false
};
let first_event_changed =
observation_state.observe_stream_chunk(&chunk, output_started_at);
let chunk_len = chunk.len();
@@ -843,6 +881,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.any(|event| matches!(event, GatewayEvent::Completed(_)));
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if allows_account_state_mutation && (completed || terminal_failure.is_some()) {
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
}
if let Some((failure, _)) = terminal_failure.as_ref() {
apply_failure(&failure_context, &active_account, failure)
.await;
@@ -861,6 +907,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if (rate_limits_changed
|| service_tier_changed
|| timing_changed
|| turn_state_changed
|| terminal_changed)
&& let Some(observation) = observation_state.observation()
{
@@ -920,21 +967,31 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let rate_limits_changed = if updates.is_empty() {
false
} else {
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
passive_quota_observation.observe(&updates);
observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates))
};
if allows_account_state_mutation {
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
}
if let Some((failure, _)) = terminal_failure.as_ref() {
apply_failure(&failure_context, &active_account, failure)
.await;
}
if let (Some(capture), Some(updates)) =
(session_capture.as_mut(), turn_state_updates.as_ref())
let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref()
&& let Some(turn_state) = updates.lock().await.take()
{
capture.turn_state = Some(turn_state);
}
if let Some(capture) = session_capture.as_mut() {
capture.turn_state = Some(turn_state.clone());
}
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
} else {
false
};
attach_openai_session_update(&mut events, &mut session_capture);
let completed = events
.iter()
@@ -951,7 +1008,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if (service_tier_changed || timing_changed || rate_limits_changed || terminal_changed)
if (service_tier_changed
|| timing_changed
|| rate_limits_changed
|| turn_state_changed
|| terminal_changed)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
@@ -1001,6 +1062,26 @@ struct OpenAiResponseObservationState {
terminal: Option<OpenAiResponseTerminal>,
}
struct OpenAiPassiveQuotaObservation {
rate_limits: Vec<ParsedRateLimits>,
}
impl OpenAiPassiveQuotaObservation {
fn new(headers: Vec<(String, String)>) -> Self {
Self {
rate_limits: parse_rate_limit_headers(&headers).into_iter().collect(),
}
}
fn observe(&mut self, updates: &[ParsedRateLimits]) {
self.rate_limits.extend_from_slice(updates);
}
fn rate_limits(&self) -> &[ParsedRateLimits] {
&self.rate_limits
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum OpenAiResponseTerminal {
Completed,
@@ -1120,6 +1201,26 @@ impl OpenAiResponseObservationState {
changed
}
fn merge_client_header(&mut self, name: &str, value: &str) -> bool {
let value = Bytes::copy_from_slice(value.as_bytes());
if let Some((_, existing)) = self
.response_metadata
.client_headers
.iter_mut()
.find(|(existing_name, _)| existing_name.eq_ignore_ascii_case(name))
{
if *existing == value {
return false;
}
*existing = value;
return true;
}
self.response_metadata
.client_headers
.push((name.to_owned(), value));
true
}
fn mark_completed(&mut self, incomplete: bool) -> bool {
let terminal = if incomplete {
OpenAiResponseTerminal::Incomplete
@@ -1432,15 +1533,20 @@ fn insert_optional_metadata_millis(
}
}
async fn take_rate_limit_updates(
updates: Option<&CodexRateLimitHeaderUpdates>,
) -> Vec<(String, String)> {
async fn take_rate_limit_updates(updates: Option<&CodexRateLimitUpdates>) -> Vec<ParsedRateLimits> {
let Some(updates) = updates else {
return Vec::new();
};
std::mem::take(&mut *updates.lock().await)
}
fn rate_limit_update_headers(updates: &[ParsedRateLimits]) -> Vec<(String, String)> {
updates
.iter()
.flat_map(rate_limits_to_header_pairs)
.collect()
}
fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool {
events
.iter()
@@ -1466,12 +1572,15 @@ fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool {
async fn synchronize_passive_quota(
quota: &CodexCredentialQuotaService,
account: &ProviderAccount,
headers: &[(String, String)],
rate_limits: &[ParsedRateLimits],
) {
if headers.is_empty() {
if rate_limits.is_empty() {
return;
}
if let Err(error) = quota.synchronize_passive_headers(account, headers).await {
if let Err(error) = quota
.synchronize_passive_rate_limits(account, rate_limits)
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
@@ -1480,6 +1589,17 @@ async fn synchronize_passive_quota(
}
}
async fn synchronize_passive_quota_headers(
quota: &CodexCredentialQuotaService,
account: &ProviderAccount,
headers: &[(String, String)],
) {
let Some(rate_limits) = parse_rate_limit_headers(headers) else {
return;
};
synchronize_passive_quota(quota, account, std::slice::from_ref(&rate_limits)).await;
}
const fn actual_transport_name(transport: CodexBackendTransport) -> &'static str {
match transport {
CodexBackendTransport::HttpSse => HTTP_SSE_TRANSPORT,
@@ -1873,7 +1993,7 @@ async fn apply_failure(
if !context.allows_account_state_mutation {
return;
}
synchronize_passive_quota(context.quota, account, &failure.rate_limit_headers).await;
synchronize_passive_quota_headers(context.quota, account, &failure.rate_limit_headers).await;
let needs_authoritative_quota_refresh = matches!(
failure.account_failure,
Some(CodexAccountFailure::QuotaExhausted | CodexAccountFailure::UsageLimitExhausted { .. })
@@ -2152,7 +2272,7 @@ fn map_upstream_failure(
error = error.with_status(status.as_u16());
}
if replay_boundary.permits_provider_proof()
&& (failure.replay_is_safe() || continuation_failure.is_some() || cyber_policy_failure)
&& (failure.replay_is_safe() || cyber_policy_failure)
{
error = error.with_replay_safe();
}
@@ -30,7 +30,7 @@ use super::response_meta::CodexResponseMetadata;
use super::tls::{CustomCaError, build_reqwest_client_with_custom_ca, custom_ca_env_cache_key};
use super::websocket::{
CodexWebSocketExchangeError, CodexWebSocketPool, CodexWebSocketPoolKey,
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketRequest, CodexWebSocketTurnStateUpdate,
CodexWebSocketRateLimitUpdates, CodexWebSocketRequest, CodexWebSocketTurnStateUpdate,
PreparedWebSocket, WebSocketOriginBreaker, WebSocketPoolDecision,
};
@@ -519,8 +519,8 @@ pub struct CodexTransportMetrics {
pub http_version: Option<String>,
}
/// 响应头之后在 live 流中采集的限流更新。
pub type CodexRateLimitHeaderUpdates = CodexWebSocketRateLimitHeaderUpdates;
/// 响应头之后在 live 流中采集的结构化限流更新。
pub type CodexRateLimitUpdates = CodexWebSocketRateLimitUpdates;
/// 响应头之后在 live 流中采集的 turn state 更新。
pub type CodexTurnStateUpdate = CodexWebSocketTurnStateUpdate;
@@ -537,8 +537,8 @@ pub struct CodexBackendStreamingResponse {
pub set_cookie_headers: Vec<String>,
/// 上游透传的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// live stream 期间捕获的限流头更新。
pub rate_limit_header_updates: Option<CodexRateLimitHeaderUpdates>,
/// live stream 期间捕获的结构化限流更新。
pub rate_limit_updates: Option<CodexRateLimitUpdates>,
/// live stream 期间捕获的 turn-state 更新。
pub turn_state_update: Option<CodexTurnStateUpdate>,
/// WebSocket 连接池决策。
@@ -147,14 +147,14 @@ impl CodexBackendClient {
});
}
let rate_limit_header_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
let rate_limit_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
Ok(CodexBackendStreamingResponse {
body: http_sse_stream(response, Arc::clone(&rate_limit_header_updates)),
body: http_sse_stream(response, Arc::clone(&rate_limit_updates)),
transport: CodexBackendTransport::HttpSse,
turn_state,
set_cookie_headers,
rate_limit_headers,
rate_limit_header_updates: Some(rate_limit_header_updates),
rate_limit_updates: Some(rate_limit_updates),
turn_state_update: None,
websocket_pool_decision: None,
diagnostics,
@@ -245,11 +245,12 @@ impl CodexBackendClient {
self.websocket_pool.as_deref().zip(pool_key)
};
let fast_path_budget = match requirement {
TransportRequirement::PersistedContinuation
| TransportRequirement::ExternalUnknown
| TransportRequirement::NewChain => Some(WEBSOCKET_FAST_PATH_BUDGET),
TransportRequirement::PersistedContinuation | TransportRequirement::NewChain => {
Some(WEBSOCKET_FAST_PATH_BUDGET)
}
TransportRequirement::ExplicitWebSocketWarmup
| TransportRequirement::ExactWebSocketContinuation => None,
| TransportRequirement::ExactWebSocketContinuation
| TransportRequirement::ExternalUnknown => None,
TransportRequirement::HttpRequired => None,
};
let prepare_started_at = Instant::now();
@@ -515,7 +516,7 @@ impl CodexBackendClient {
turn_state: exchange.turn_state,
set_cookie_headers: exchange.set_cookie_headers,
rate_limit_headers: exchange.rate_limit_headers,
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
rate_limit_updates: Some(exchange.rate_limit_updates),
turn_state_update: Some(exchange.turn_state_update),
websocket_pool_decision: exchange.pool_decision,
diagnostics: exchange.diagnostics,
@@ -878,7 +879,7 @@ fn websocket_connection_profile(headers: &HeaderMap) -> String {
fn http_sse_stream(
response: ReqwestResponse,
rate_limit_updates: CodexRateLimitHeaderUpdates,
rate_limit_updates: CodexRateLimitUpdates,
) -> CodexBackendSseStream {
let stream: CodexBackendSseStream =
Box::pin(response.bytes_stream().map_err(CodexClientError::Http));
@@ -901,7 +902,7 @@ fn http_sse_stream(
fn observe_http_sse_rate_limits(
stream: CodexBackendSseStream,
updates: CodexRateLimitHeaderUpdates,
updates: CodexRateLimitUpdates,
) -> CodexBackendSseStream {
Box::pin(futures::stream::unfold(
(stream, SseEventDecoder::default(), updates),
@@ -925,9 +926,9 @@ fn observe_http_sse_rate_limits(
async fn append_http_sse_rate_limit_updates(
frames: Vec<SseFrame>,
updates: &CodexRateLimitHeaderUpdates,
updates: &CodexRateLimitUpdates,
) {
let mut headers = Vec::new();
let mut observations = Vec::new();
for frame in frames {
for event in frame.events() {
if event
@@ -940,10 +941,10 @@ async fn append_http_sse_rate_limit_updates(
let Some(rate_limits) = events::parse_rate_limits_event_raw(&event.data) else {
continue;
};
headers.extend(events::rate_limits_to_header_pairs(&rate_limits));
observations.push(rate_limits);
}
}
if !headers.is_empty() {
updates.lock().await.extend(headers);
if !observations.is_empty() {
updates.lock().await.extend(observations);
}
}
@@ -27,8 +27,8 @@ pub use self::{
client::{
CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendSseStream,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexClientResult,
CodexRateLimitHeaderUpdates, CodexRequestContext, CodexTransportDecision,
CodexTransportMetrics, CodexTurnStateUpdate, build_reqwest_client,
CodexRateLimitUpdates, CodexRequestContext, CodexTransportDecision, CodexTransportMetrics,
CodexTurnStateUpdate, build_reqwest_client,
},
diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamSendPhase},
endpoints::{
@@ -5,6 +5,8 @@ use serde_json::{Map, Value};
use crate::transport::protocol::responses::{CodexResponsesRequest, transport_requirement};
const REDACTED_PAYLOAD_VALUE: &str = "<redacted>";
const PREVIOUS_RESPONSE_NOT_FOUND_CODE: &str = "previous_response_not_found";
const INVALID_PREVIOUS_RESPONSE_ID_MESSAGE: &str = "Invalid `previous_response_id`.";
/// WebSocket 握手审计快照。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
@@ -54,16 +56,49 @@ pub fn websocket_event_to_sse_frame(raw: &str) -> Option<String> {
websocket_event_frame(&value, raw)
}
/// 同上,复用已解析的事件 JSON;`data` 段逐字节内嵌 `raw` 原文。
/// 同上,复用已解析的事件 JSON;`data` 段通常逐字节内嵌 `raw` 原文。
///
/// 仅剥离传输层内部帧:`codex.rate_limits` 与承载 turn_state 的
/// `response.metadata` 不下发客户端;业务事件一律原样转发,不做 schema 校验。
/// `response.metadata` 不下发客户端。唯一的业务帧适配是为上游省略错误码的
/// `Invalid previous_response_id` 补齐官方客户端用于完整历史重放的稳定错误码;
/// 原始 status、type 与 message 均保持不变。
pub(crate) fn websocket_event_frame(value: &Value, raw: &str) -> Option<String> {
let event = websocket_event_type(value)?;
if is_internal_websocket_event(event) || event == "response.metadata" {
return None;
}
Some(encode_sse_event(event, raw))
let normalized = previous_response_not_found_event(value);
Some(encode_sse_event(
event,
normalized.as_deref().unwrap_or(raw),
))
}
fn previous_response_not_found_event(value: &Value) -> Option<String> {
if websocket_event_type(value) != Some("error")
|| value
.get("status")
.or_else(|| value.get("status_code"))
.and_then(Value::as_u64)
!= Some(400)
{
return None;
}
let error = value.get("error")?.as_object()?;
if error.get("code").is_some_and(|code| !code.is_null())
|| error.get("type").and_then(Value::as_str) != Some("invalid_request_error")
|| error.get("message").and_then(Value::as_str)
!= Some(INVALID_PREVIOUS_RESPONSE_ID_MESSAGE)
{
return None;
}
let mut normalized = value.clone();
normalized.get_mut("error")?.as_object_mut()?.insert(
"code".to_owned(),
Value::String(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()),
);
serde_json::to_string(&normalized).ok()
}
pub(crate) fn websocket_event_type(value: &Value) -> Option<&str> {
@@ -61,6 +61,23 @@ const INSTALLATION_ID_KEYS: &[&str] = &[
"x-codex-installation-id",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequestAccountScope {
SameAccount,
CrossAccountFullRequest,
CrossAccountContinuationProbe,
}
impl RequestAccountScope {
const fn is_cross_account(self) -> bool {
!matches!(self, Self::SameAccount)
}
const fn preserves_continuation(self) -> bool {
matches!(self, Self::CrossAccountContinuationProbe)
}
}
/// Provider 专属编码错误;不保存 prompt、schema 或 option 值。
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum CodexRequestEncodeError {
@@ -289,12 +306,18 @@ fn normalize_conversation_anchor_text(text: &str) -> String {
///
/// 真实 account ID 与 installation ID 由随后构造的 `CodexRequestContext` 注入
/// 请求头;正文只替换客户端原本声明过的 installation 字段,绝不接受客户端
/// 提供的 token、cookie 或账号身份。
/// 提供的 token、cookie 或账号身份。Continuation probe 保留 canonical
/// previous-response 和原始增量 input,不能把它们误当成完整请求清理。
pub(crate) fn scope_request_to_account(
request: &mut CodexResponsesRequest,
installation_id: &str,
cross_account: bool,
account_scope: RequestAccountScope,
) {
let cross_account = account_scope.is_cross_account();
let continuation = account_scope
.preserves_continuation()
.then(|| request.previous_response_id().map(ToOwned::to_owned))
.flatten();
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());
@@ -306,23 +329,38 @@ pub(crate) fn scope_request_to_account(
} else {
None
};
let turn_metadata = request
.turn_metadata
.as_deref()
.and_then(|metadata| scope_turn_metadata(metadata, installation_id, cross_account));
let client_metadata_turn_metadata = metadata_string(request, "x-codex-turn-metadata")
.and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account));
let turn_metadata = (!account_scope.preserves_continuation())
.then(|| {
request
.turn_metadata
.as_deref()
.and_then(|metadata| scope_turn_metadata(metadata, installation_id, cross_account))
})
.flatten();
let client_metadata_turn_metadata = (!account_scope.preserves_continuation())
.then(|| {
metadata_string(request, "x-codex-turn-metadata")
.and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account))
})
.flatten();
if cross_account {
request.passthrough_headers.remove("x-codex-turn-state");
request.passthrough_headers.remove("x-codex-turn-metadata");
sanitize_cross_account_input(request);
if !account_scope.preserves_continuation() {
sanitize_cross_account_input(request);
}
for key in CROSS_ACCOUNT_IDENTITY_KEYS
.iter()
.chain(ACCOUNT_BOUND_STATE_KEYS)
{
request.body_mut().remove(*key);
}
if account_scope.preserves_continuation() {
for key in TURN_METADATA_KEYS {
request.body_mut().remove(*key);
}
}
}
for key in INSTALLATION_ID_KEYS {
@@ -356,6 +394,11 @@ pub(crate) fn scope_request_to_account(
{
metadata.remove(*key);
}
if account_scope.preserves_continuation() {
for key in TURN_METADATA_KEYS {
metadata.remove(*key);
}
}
}
metadata.insert(
"x-codex-installation-id".to_owned(),
@@ -393,6 +436,9 @@ pub(crate) fn scope_request_to_account(
request.turn_state = turn_state;
request.turn_metadata = turn_metadata;
if let Some(continuation) = continuation {
request.set_previous_response_id(Some(continuation));
}
}
fn sanitize_cross_account_input(request: &mut CodexResponsesRequest) {
@@ -193,7 +193,6 @@ fn client_response_header_is_forwardable(name: &str, connection_options: &[Strin
| "x-openai-organization"
| "x-openai-project"
| "x-codex-installation-id"
| "x-codex-turn-state"
| "x-codex-turn-metadata"
)
}
@@ -8,6 +8,7 @@ use std::{pin::Pin, sync::Arc, time::Duration};
use bytes::Bytes;
use futures::Stream;
use gateway_protocol::openai::events::ParsedRateLimits;
use tokio::sync::Mutex;
use uuid::Uuid;
@@ -36,8 +37,8 @@ pub struct CodexWebSocketStreamingExchange {
pub set_cookie_headers: Vec<String>,
/// 上游握手响应里的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// 上游内部 rate-limit 事件里的动态更新。
pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
/// 上游内部 `codex.rate_limits` 事件里的结构化动态更新。
pub rate_limit_updates: CodexWebSocketRateLimitUpdates,
/// 上游内部 metadata 事件里的动态 turn state。
pub turn_state_update: CodexWebSocketTurnStateUpdate,
/// WebSocket 连接池决策。
@@ -53,8 +54,8 @@ pub struct CodexWebSocketStreamingExchange {
/// Responses WebSocket live SSE 字节流。
pub type CodexWebSocketSseStream =
Pin<Box<dyn Stream<Item = Result<Bytes, CodexWebSocketExchangeError>> + Send + 'static>>;
/// live 流中的限流头动态更新。
pub type CodexWebSocketRateLimitHeaderUpdates = Arc<Mutex<Vec<(String, String)>>>;
/// live 流中的结构化限流动态更新。
pub type CodexWebSocketRateLimitUpdates = Arc<Mutex<Vec<ParsedRateLimits>>>;
/// live 流中的 turn state 动态更新。
pub type CodexWebSocketTurnStateUpdate = Arc<Mutex<Option<String>>>;
@@ -20,7 +20,7 @@ pub(in crate::transport::websocket) enum WebSocketTerminalKind {
}
pub(super) enum ExchangeAction {
RateLimits(Vec<(String, String)>),
RateLimits(events::ParsedRateLimits),
TurnState(String),
Forward {
frame: String,
@@ -41,8 +41,8 @@ pub(super) fn reduce_websocket_event(
};
if let Some(parsed) = events::parse_rate_limits_event(&value) {
let headers = events::rate_limits_to_header_pairs(&parsed);
metadata.rate_limit_headers.extend(headers.iter().cloned());
return Ok(ExchangeAction::RateLimits(headers));
metadata.rate_limit_headers.extend(headers);
return Ok(ExchangeAction::RateLimits(parsed));
}
response_meta::merge_response_metadata(
@@ -16,9 +16,8 @@ use super::super::{
use super::io::{next_websocket_message, receive_idle_timeout, reused_stream_receive_error};
use super::reducer::{ExchangeAction, WebSocketTerminalKind, reduce_websocket_event};
use super::{
CodexWebSocketExchangeError, CodexWebSocketRateLimitHeaderUpdates,
CodexWebSocketStreamingExchange, CodexWebSocketTurnStateUpdate, WEBSOCKET_STREAM_BUFFER,
reusable_websocket_metadata,
CodexWebSocketExchangeError, CodexWebSocketRateLimitUpdates, CodexWebSocketStreamingExchange,
CodexWebSocketTurnStateUpdate, WEBSOCKET_STREAM_BUFFER, reusable_websocket_metadata,
};
pub(in crate::transport::websocket) struct WebSocketStreamPoolReturn {
@@ -63,8 +62,8 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
) -> CodexWebSocketStreamingExchange {
let websocket_connection_id = websocket.connection_id();
let response_metadata = metadata.clone();
let rate_limit_header_updates = Arc::new(Mutex::new(Vec::new()));
let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates);
let rate_limit_updates = Arc::new(Mutex::new(Vec::new()));
let rate_limit_updates_for_task = Arc::clone(&rate_limit_updates);
let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone()));
let turn_state_update_for_task = Arc::clone(&turn_state_update);
let (tx, rx) = mpsc::channel(WEBSOCKET_STREAM_BUFFER);
@@ -83,7 +82,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
reused_connection,
initial_event_timeout,
shutdown,
rate_limit_header_updates: rate_limit_header_updates_for_task,
rate_limit_updates: rate_limit_updates_for_task,
turn_state_update: turn_state_update_for_task,
tx,
})
@@ -105,7 +104,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
turn_state: response_metadata.turn_state,
set_cookie_headers: response_metadata.set_cookie_headers,
rate_limit_headers: response_metadata.rate_limit_headers,
rate_limit_header_updates,
rate_limit_updates,
turn_state_update,
pool_decision: None,
connection_local_continuation: false,
@@ -121,7 +120,7 @@ struct WebSocketStreamForwardState {
reused_connection: bool,
initial_event_timeout: Option<Duration>,
shutdown: CancellationToken,
rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
rate_limit_updates: CodexWebSocketRateLimitUpdates,
turn_state_update: CodexWebSocketTurnStateUpdate,
tx: mpsc::Sender<Result<Bytes, CodexWebSocketExchangeError>>,
}
@@ -134,7 +133,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
reused_connection,
initial_event_timeout,
shutdown,
rate_limit_header_updates,
rate_limit_updates,
turn_state_update,
tx,
} = state;
@@ -231,8 +230,8 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
};
let (frame, terminal) = match reduce_websocket_event(&raw, &mut metadata, &mut continuation)
{
Ok(ExchangeAction::RateLimits(headers)) => {
rate_limit_header_updates.lock().await.extend(headers);
Ok(ExchangeAction::RateLimits(rate_limits)) => {
rate_limit_updates.lock().await.push(rate_limits);
continue;
}
Ok(ExchangeAction::TurnState(turn_state)) => {
@@ -21,8 +21,8 @@ pub use self::{
},
error::{CodexWebSocketCloseError, CodexWebSocketExchangeError, CodexWebSocketUpstreamError},
exchange::{
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketSseStream,
CodexWebSocketStreamingExchange, CodexWebSocketTurnStateUpdate,
CodexWebSocketRateLimitUpdates, CodexWebSocketSseStream, CodexWebSocketStreamingExchange,
CodexWebSocketTurnStateUpdate,
},
handshake::responses_websocket_endpoint,
model::{
@@ -579,7 +579,7 @@ async fn openai_admin_provider_projects_official_codex_quota_and_independent_buc
let secondary = quota
.windows
.iter()
.find(|window| window.key == "core-secondary")
.find(|window| window.key == "codex-secondary")
.expect("core secondary quota");
assert_eq!(secondary.label, "次级额度");
assert_eq!(secondary.used_percent, Some(88.0));
@@ -11,6 +11,7 @@ use gateway_core::engine::credential::{
AccountStateChange, AccountStatus, CredentialState, ProviderAccount, ProviderAccountStore as _,
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaState,
};
use gateway_protocol::openai::events::parse_rate_limits_event;
use provider_openai::OFFICIAL_CODEX_BASE_URL;
use provider_openai::credential::{
CodexCredentialQuotaError, CodexCredentialQuotaService, ImportCodexOAuthCredential,
@@ -163,6 +164,224 @@ async fn successful_inference_headers_are_authoritative_even_at_full_display_usa
);
}
#[tokio::test]
async fn passive_quota_updates_keep_core_and_model_specific_buckets_independent() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_independent_passive_buckets";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("account");
let service = quota_service(&store);
let initial_headers = vec![
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-primary-used-percent".to_owned(), "81".to_owned()),
(
"x-codex-primary-window-minutes".to_owned(),
"10080".to_owned(),
),
(
"x-codex-bengalfox-primary-used-percent".to_owned(),
"0".to_owned(),
),
(
"x-codex-bengalfox-primary-window-minutes".to_owned(),
"10080".to_owned(),
),
(
"x-codex-bengalfox-limit-name".to_owned(),
"GPT-5.3-Codex-Spark".to_owned(),
),
];
service
.synchronize_passive_headers(&account, &initial_headers)
.await
.expect("initial passive quota");
let spark_headers = vec![
// 生产响应会把当前具名桶同时复制到旧版 `x-codex-*` 默认头;这个
// 兼容别名不能覆盖已有的账号 core 桶。
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-primary-used-percent".to_owned(), "1".to_owned()),
(
"x-codex-primary-window-minutes".to_owned(),
"10080".to_owned(),
),
(
"x-codex-primary-reset-at".to_owned(),
"1900000100".to_owned(),
),
(
"x-codex-bengalfox-primary-used-percent".to_owned(),
"1".to_owned(),
),
(
"x-codex-bengalfox-primary-window-minutes".to_owned(),
"10080".to_owned(),
),
(
"x-codex-bengalfox-primary-reset-at".to_owned(),
"1900000100".to_owned(),
),
(
"x-codex-bengalfox-limit-name".to_owned(),
"GPT-5.3-Codex-Spark".to_owned(),
),
];
service
.synchronize_passive_headers(&account, &spark_headers)
.await
.expect("Spark passive quota");
let snapshot = service
.read_account(account.id())
.await
.expect("read Spark quota")
.expect("Spark quota snapshot");
let core = snapshot
.windows()
.iter()
.find(|window| window.source() == "codex")
.expect("core quota window");
let spark = snapshot
.windows()
.iter()
.find(|window| window.source() == "codex_bengalfox")
.expect("Spark quota window");
assert_eq!(core.used_percent(), Some(81.0));
assert_eq!(spark.used_percent(), Some(1.0));
let core_headers = vec![
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-primary-used-percent".to_owned(), "82".to_owned()),
(
"x-codex-primary-window-minutes".to_owned(),
"10080".to_owned(),
),
];
service
.synchronize_passive_headers(&account, &core_headers)
.await
.expect("core passive quota");
let snapshot = service
.read_account(account.id())
.await
.expect("read core quota")
.expect("core quota snapshot");
let core = snapshot
.windows()
.iter()
.find(|window| window.source() == "codex")
.expect("core quota window");
let spark = snapshot
.windows()
.iter()
.find(|window| window.source() == "codex_bengalfox")
.expect("Spark quota window");
assert_eq!(core.used_percent(), Some(82.0));
assert_eq!(spark.used_percent(), Some(1.0));
let persisted = store.quota_json(account_id).expect("persisted quota");
let by_limit_id = persisted
.get("rate_limits_by_limit_id")
.and_then(serde_json::Value::as_object)
.expect("rate-limit map");
assert_eq!(
by_limit_id.keys().map(String::as_str).collect::<Vec<_>>(),
["codex", "codex_bengalfox"]
);
assert!(persisted.get("rate_limit").is_none());
assert!(persisted.get("additional_rate_limits").is_none());
}
#[tokio::test]
async fn websocket_additional_rate_limit_resolves_by_name_without_touching_core() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_websocket_additional_limit";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("account");
let service = quota_service(&store);
let initial_headers = vec![
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-primary-used-percent".to_owned(), "84".to_owned()),
(
"x-codex-primary-window-minutes".to_owned(),
"10080".to_owned(),
),
(
"x-codex-primary-reset-at".to_owned(),
"1900000000".to_owned(),
),
(
"x-codex-bengalfox-primary-used-percent".to_owned(),
"0".to_owned(),
),
(
"x-codex-bengalfox-primary-window-minutes".to_owned(),
"10080".to_owned(),
),
(
"x-codex-bengalfox-primary-reset-at".to_owned(),
"1900000100".to_owned(),
),
(
"x-codex-bengalfox-limit-name".to_owned(),
"GPT-5.3-Codex-Spark".to_owned(),
),
];
service
.synchronize_passive_headers(&account, &initial_headers)
.await
.expect("initial quota metadata");
let event = json!({
"type": "codex.rate_limits",
"plan_type": "pro",
"rate_limits": {
"allowed": true,
"limit_reached": false,
"primary": {
"used_percent": 2,
"window_minutes": 10080,
"reset_at": 1900000100
},
"secondary": null
},
"additional_rate_limits": {
"GPT-5.3-Codex-Spark": {
"allowed": true,
"limit_reached": false,
"primary": {
"used_percent": 2,
"window_minutes": 10080,
"reset_at": 1900000100
},
"secondary": null
}
},
"credits": null
});
let observation = parse_rate_limits_event(&event).expect("WebSocket rate-limit event");
service
.synchronize_passive_rate_limits(&account, std::slice::from_ref(&observation))
.await
.expect("WebSocket passive quota");
let persisted = store.quota_json(account_id).expect("persisted quota");
let by_limit_id = persisted["rate_limits_by_limit_id"]
.as_object()
.expect("rate-limit map");
assert_eq!(
by_limit_id["codex"]["primary_window"]["used_percent"],
json!(84.0)
);
assert_eq!(
by_limit_id["codex_bengalfox"]["primary_window"]["used_percent"],
json!(2.0)
);
assert_eq!(persisted["active_limit"], json!("codex_bengalfox"));
assert_eq!(by_limit_id.len(), 2);
}
#[tokio::test]
async fn quota_refresh_persists_explicit_provider_denial() {
let store = Arc::new(MemoryAccountStore::default());
@@ -209,6 +428,14 @@ async fn quota_refresh_persists_explicit_provider_denial() {
current.status_projection(SystemTime::now(), None).status,
AccountStatus::QuotaExhausted
);
let persisted = store.quota_json(account_id).expect("persisted quota");
assert!(
persisted
.get("rate_limits_by_limit_id")
.and_then(|limits| limits.get("codex"))
.is_some()
);
assert!(persisted.get("rate_limit").is_none());
}
#[tokio::test]
@@ -396,11 +396,11 @@ async fn persisted_quota_orders_core_window_before_additional_limit() {
.map(|window| window.source())
.collect::<Vec<_>>();
assert_eq!(sources, ["codex", "gpt-5.3-codex-spark"]);
assert_eq!(sources, ["codex", "gpt_5.3_codex_spark"]);
}
#[tokio::test]
async fn code_review_limit_projects_as_independent_window_with_limit_name() {
async fn code_review_limit_projects_as_one_snapshot_per_limit_id() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_code_review").await;
let account = store.account("acct_code_review").expect("account");
@@ -455,21 +455,16 @@ async fn code_review_limit_projects_as_independent_window_with_limit_name() {
.expect("read quota")
.expect("quota snapshot");
// spend_control 不再生成窗口(只作 exhaustion 信号);顶层和 additional
// 的 code_review 桶分别保留,避免一个快照覆盖另一个。
// spend_control 不生成窗口(只作 exhaustion 信号);官方 map 协议中同一个
// limit_id 只能保留一个快照,顶层 code_review 事实优先于重复 additional。
let review = snapshot
.windows()
.iter()
.filter(|window| window.source() == "code_review")
.collect::<Vec<_>>();
assert_eq!(review.len(), 2);
assert!(review.iter().any(|window| {
window.limit_name() == Some("code_review") && window.used_percent() == Some(80.0)
}));
assert!(review.iter().any(|window| {
window.limit_name() == Some("code review") && window.used_percent() == Some(55.0)
}));
assert_ne!(review[0].key(), review[1].key());
assert_eq!(snapshot.windows().len(), 3);
assert_eq!(review.len(), 1);
assert_eq!(review[0].limit_name(), Some("code_review"));
assert_eq!(review[0].used_percent(), Some(80.0));
assert_eq!(snapshot.windows().len(), 2);
assert_eq!(snapshot.quota().access(), QuotaAccessState::Unknown);
}
@@ -6,6 +6,9 @@ use std::time::{Duration, SystemTime};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use chrono::Utc;
use futures::{SinkExt, StreamExt};
use gateway_core::engine::continuation::{
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
};
use gateway_core::engine::credential::{
AccountFeedbackStats, CredentialState, OpaqueProviderData, ProviderAccountId,
ProviderAccountStore as _, QuotaAccessChange, QuotaAccessState, QuotaEvidence,
@@ -16,7 +19,7 @@ use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ContinuationAttempt, ModelRequestId,
ProviderAccountStateOwner, RequestAttemptContext, UpstreamSendState,
};
use gateway_core::error::ProviderErrorKind;
use gateway_core::error::{ContinuationFailure, ProviderErrorKind};
use gateway_core::event::GatewayEvent;
use gateway_core::operation::{
CapabilityRequirements, GenerateRequest, Operation, OperationKind, ProtocolPayload,
@@ -270,7 +273,8 @@ fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest
vec![ProviderModel::new(
provider,
upstream_model,
ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000)),
ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000))
.with_upstream_feature_validation(),
)],
Vec::new(),
)
@@ -291,6 +295,8 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
let provider = ProviderKind::new("openai").expect("provider");
let accounts = [
"acct_affinity",
"acct_affinity_switch_a",
"acct_affinity_switch_b",
"acct_atomic_failure",
"acct_bare_atomic_failure",
"acct_bounded_replay_grace",
@@ -322,6 +328,7 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
"acct_truncated_stream",
"acct_usage_limit_request_path",
"acct_websocket_close",
"acct_websocket_turn_state",
]
.into_iter()
.map(|id| {
@@ -392,6 +399,33 @@ fn context_with_state_owner(request_id: &str, owner_account_id: &str) -> Attempt
)
}
fn replay_any_context(request_id: &str, owner_account_id: &str) -> AttemptContext {
let owner_account = ProviderAccountId::new(owner_account_id).expect("owner account id");
let provider = ProviderKind::new("openai").expect("provider");
let owner = ProviderAccountStateOwner::new(provider.clone(), owner_account.clone());
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("client-previous-response"),
PreviousResponseId::new("upstream-previous-response"),
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
provider,
owner_account,
);
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).expect("request id"),
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
),
NonZeroU32::new(2).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, Some(owner))
.with_account_scope(contract_account_scope()),
Some(ContinuationBinding::Pinned(continuation)),
CancellationToken::new(),
)
.with_continuation_attempt(ContinuationAttempt::ReplayAny)
}
async fn capture_scoped_http_request(
request_id: &str,
selected_account_id: &str,
@@ -436,6 +470,76 @@ async fn capture_scoped_http_request(
requests.pop().expect("single scoped request")
}
async fn capture_turn_state_request(
request_id: &str,
previous_turn_id: Option<&str>,
current_turn_id: Option<&str>,
client_turn_state: Option<&str>,
) -> wiremock::Request {
let account_id = "acct_session_affinity";
let store = Arc::new(MemoryAccountStore::default());
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")
.set_body_string(CAPTURE_COMPLETED_SSE),
)
.expect(1)
.mount(&server)
.await;
let mut session_state = Map::from_iter([
("account_id".to_owned(), json!(account_id)),
("conversation_id".to_owned(), json!("conversation")),
("turn_state".to_owned(), json!("previous-turn-state")),
("continuation_scope".to_owned(), json!("persisted")),
]);
if let Some(previous_turn_id) = previous_turn_id {
session_state.insert("client_turn_id".to_owned(), json!(previous_turn_id));
}
let mut protocol_context = Map::from_iter([("use_websocket".to_owned(), json!(false))]);
if let Some(current_turn_id) = current_turn_id {
protocol_context.insert("turn_id".to_owned(), json!(current_turn_id));
}
if let Some(client_turn_state) = client_turn_state {
protocol_context.insert("turn_state".to_owned(), json!(client_turn_state));
}
let operation = Operation::Generate(
GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("current input")),
]),
)
.expect("OpenAI payload")
.with_context(protocol_context),
)
.with_provider_session_state(
ProviderSessionState::new("openai", session_state).expect("provider session state"),
),
);
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", operation),
context(request_id, CancellationToken::new()),
)
.await
.expect("prepare turn-state provider stream");
while let Some(event) = stream.next().await {
event.expect("turn-state provider response");
}
let mut requests = server
.received_requests()
.await
.expect("captured turn-state request");
assert_eq!(requests.len(), 1);
requests.pop().expect("single turn-state request")
}
fn captured_header_values(request: &wiremock::Request, name: &str) -> Vec<Vec<u8>> {
request
.headers
@@ -792,6 +896,94 @@ async fn websocket_close_details_are_only_exposed_through_the_client_error() {
assert_eq!(detail.error_type(), Some("websocket_close_error"));
}
#[tokio::test]
async fn websocket_turn_state_metadata_is_exposed_through_response_observation() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_websocket_turn_state").await;
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket listener");
let base_url = format!(
"http://{}",
listener.local_addr().expect("listener address")
);
let server = tokio::spawn(async move {
let (stream, _) = listener
.accept()
.await
.expect("accept WebSocket connection");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket
.next()
.await
.expect("WebSocket request")
.expect("valid WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "response.metadata",
"headers": {
"x-codex-turn-state": ["turn-state-from-websocket"]
}
})
.to_string()
.into(),
))
.await
.expect("send response metadata");
websocket
.send(Message::Text(
json!({
"type": "response.completed",
"response": {
"id": "resp_websocket_turn_state",
"model": "gpt-5.4",
"status": "completed",
"output": [],
"usage": {
"input_tokens": 1,
"output_tokens": 1,
"total_tokens": 2
}
}
})
.to_string()
.into(),
))
.await
.expect("send completed response");
websocket.close(None).await.expect("close WebSocket");
});
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", generate_operation()),
context("req_websocket_turn_state", CancellationToken::new()),
)
.await
.expect("prepare WebSocket provider stream");
let mut observed_turn_state = false;
while let Some(event) = stream.next().await {
let event = event.expect("provider event");
if let Some(observation) = event.response_observation() {
observed_turn_state |= observation.client_headers().iter().any(|header| {
header.name().eq_ignore_ascii_case("x-codex-turn-state")
&& header.value().as_ref() == b"turn-state-from-websocket"
});
}
if event
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Completed(_)))
{
break;
}
}
server.await.expect("WebSocket server");
assert!(observed_turn_state);
}
#[tokio::test]
async fn disabled_account_is_excluded_from_normal_scheduling() {
let store = Arc::new(MemoryAccountStore::default());
@@ -974,9 +1166,9 @@ async fn cross_account_scope_sanitizes_only_known_turn_metadata_fields() {
}
#[tokio::test]
async fn continuation_replay_should_use_the_client_full_input_without_legacy_transcript() {
async fn cross_account_continuation_probe_should_preserve_delta_input_and_previous_response() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_client_history").await;
create_account(&store, "acct_scope_new").await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
@@ -988,33 +1180,61 @@ async fn continuation_replay_should_use_the_client_full_input_without_legacy_tra
.expect(1)
.mount(&server)
.await;
let client_input = json!([
{"role": "user", "content": "complete client history"},
{"role": "assistant", "content": "previous answer"},
{"role": "user", "content": "current turn"}
let delta_input = json!([
{
"type": "message",
"id": "delta-message-id",
"encrypted_content": "opaque-delta-ciphertext",
"role": "user",
"content": "current turn only"
},
{
"type": "compaction",
"id": "delta-compaction-id",
"encrypted_content": "opaque-compaction-ciphertext"
}
]);
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), client_input.clone()),
("input".to_owned(), delta_input.clone()),
(
"previous_response_id".to_owned(),
json!("client-previous-response"),
),
("authorization".to_owned(), json!("client-extension-secret")),
(
"turnMetadata".to_owned(),
json!(r#"{"account_id":"acct_scope_old","safe":true}"#),
),
(
"client_metadata".to_owned(),
json!({
"x-codex-turn-state": "old-account-turn-state",
"x-codex-turn-metadata": "old-account-turn-metadata",
"future": "keep"
}),
),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
.with_context(Map::from_iter([
("use_websocket".to_owned(), json!(false)),
("turn_id".to_owned(), json!("turn-probe")),
("turn_state".to_owned(), json!("old-account-turn-state")),
])),
)
.with_provider_session_state(
ProviderSessionState::new(
"openai",
Map::from_iter([
("account_id".to_owned(), json!("acct_client_history")),
("account_id".to_owned(), json!("acct_scope_old")),
("conversation_id".to_owned(), json!("conversation")),
("continuation_scope".to_owned(), json!("replay_required")),
(
"transcript".to_owned(),
json!([{"client_input": {"role": "user", "content": "stale proxy copy"}}]),
),
("turn_state".to_owned(), json!("old-account-turn-state")),
("client_turn_id".to_owned(), json!("turn-probe")),
("continuation_scope".to_owned(), json!("persisted")),
]),
)
.expect("provider session state"),
@@ -1022,21 +1242,84 @@ async fn continuation_replay_should_use_the_client_full_input_without_legacy_tra
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", Operation::Generate(generation)),
context("req_client_history", CancellationToken::new())
.with_continuation_attempt(ContinuationAttempt::ReplayAny),
replay_any_context("req_cross_account_probe", "acct_scope_old"),
)
.await
.expect("prepare replay stream");
.expect("prepare continuation probe stream");
while let Some(event) = stream.next().await {
event.expect("replay response");
event.expect("continuation probe response");
}
let requests = server
.received_requests()
.await
.expect("captured replay request");
.expect("captured continuation probe request");
let body: Value = captured_request_body(&requests[0]);
assert_eq!(body.get("input"), Some(&client_input));
assert_eq!(body.get("input"), Some(&delta_input));
assert_eq!(
body.get("previous_response_id"),
Some(&json!("upstream-previous-response"))
);
assert!(body.get("authorization").is_none());
assert!(body.get("turnMetadata").is_none());
assert!(
body.pointer("/client_metadata/x-codex-turn-state")
.is_none()
);
assert!(
body.pointer("/client_metadata/x-codex-turn-metadata")
.is_none()
);
assert_eq!(
body.pointer("/client_metadata/future"),
Some(&json!("keep"))
);
assert!(captured_header_values(&requests[0], "x-codex-turn-state").is_empty());
}
#[tokio::test]
async fn matching_turn_id_should_restore_previous_turn_state() {
let request = capture_turn_state_request(
"req_restore_same_turn_state",
Some("turn-same"),
Some("turn-same"),
None,
)
.await;
assert_eq!(
captured_header_values(&request, "x-codex-turn-state"),
vec![b"previous-turn-state".to_vec()]
);
}
#[tokio::test]
async fn matching_turn_id_should_prefer_the_latest_provider_state_over_a_stale_client_echo() {
let request = capture_turn_state_request(
"req_replace_stale_client_turn_state",
Some("turn-same"),
Some("turn-same"),
Some("stale-client-turn-state"),
)
.await;
assert_eq!(
captured_header_values(&request, "x-codex-turn-state"),
vec![b"previous-turn-state".to_vec()]
);
}
#[tokio::test]
async fn new_or_unidentified_turn_should_not_restore_previous_turn_state() {
for (request_id, previous_turn_id, current_turn_id) in [
("req_new_turn_state", Some("turn-old"), Some("turn-new")),
("req_unidentified_turn_state", None, Some("turn-new")),
("req_missing_current_turn_state", Some("turn-old"), None),
] {
let request =
capture_turn_state_request(request_id, previous_turn_id, current_turn_id, None).await;
assert!(captured_header_values(&request, "x-codex-turn-state").is_empty());
}
}
#[tokio::test]
@@ -1262,6 +1545,134 @@ async fn explicit_session_id_should_override_turn_specific_prompt_cache_keys_for
assert_eq!(keys[0], keys[1]);
}
#[tokio::test]
async fn affinity_quota_switch_should_clear_old_turn_state_without_a_provider_state_owner() {
let store = Arc::new(MemoryAccountStore::default());
let first_account_id = "acct_affinity_switch_a";
let second_account_id = "acct_affinity_switch_b";
create_account(&store, first_account_id).await;
let affinity = Arc::new(MemorySessionAffinity::default());
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;
let (provider, quota) = provider_and_quota_with_affinity_and_base_url_and_leases(
&store,
Arc::clone(&affinity),
server.uri(),
Arc::new(TestLeaseCoordinator::default()),
);
let session_id = "stable-affinity-switch-session";
let first_payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("first request")),
("session_id".to_owned(), json!(session_id)),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))]));
let first = provider
.execute(
planned_request(
"openai",
Operation::Generate(GenerateRequest::from_protocol_payload(first_payload)),
),
context("req_affinity_switch_first", CancellationToken::new()),
)
.await
.expect("prepare first affinity request");
drop(first);
let affinity_keys = affinity.lookup_keys();
assert_eq!(affinity_keys.len(), 1);
affinity.seed_binding(
&ProviderKind::new("openai").expect("provider"),
&affinity_keys[0],
ProviderAccountId::new(first_account_id).expect("first account id"),
);
assert_eq!(affinity.binding_count(), 1);
create_account(&store, second_account_id).await;
let first_account = store.account(first_account_id).expect("first account");
let observed_at = SystemTime::now();
store
.compare_and_swap_quota(QuotaObservation {
account_id: first_account.id().clone(),
expected_revision: first_account.revision(),
quota: OpaqueProviderData::new(Map::new()),
observed_at,
state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None),
})
.await
.expect("seed exhausted affinity account");
let accounts = [
store
.account(first_account_id)
.expect("first account after quota"),
store.account(second_account_id).expect("second account"),
];
quota.prepare_scheduling(&accounts).await;
let second_payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("second request")),
("session_id".to_owned(), json!(session_id)),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([
("use_websocket".to_owned(), json!(false)),
("turn_id".to_owned(), json!("turn-affinity-switch")),
(
"turn_state".to_owned(),
json!("state-created-by-exhausted-account"),
),
]));
let mut second = provider
.execute(
planned_request(
"openai",
Operation::Generate(GenerateRequest::from_protocol_payload(second_payload)),
),
context("req_affinity_switch_second", CancellationToken::new()),
)
.await
.expect("prepare switched affinity request");
while let Some(event) = second.next().await {
event.expect("switched affinity response");
}
let requests = server
.received_requests()
.await
.expect("captured affinity requests");
assert_eq!(requests.len(), 1);
let second_request = &requests[0];
assert_eq!(
captured_request_body(second_request).get("input"),
Some(&json!("second request"))
);
assert!(captured_header_values(second_request, "x-codex-turn-state").is_empty());
let expected_account_id = format!("chatgpt-{second_account_id}");
assert_eq!(
second_request
.headers
.get("chatgpt-account-id")
.and_then(|value| value.to_str().ok()),
Some(expected_account_id.as_str())
);
}
#[tokio::test]
async fn thread_spawn_children_should_have_independent_session_affinity_keys() {
let store = Arc::new(MemoryAccountStore::default());
@@ -1934,6 +2345,77 @@ async fn bare_response_failed_should_remain_an_atomic_replay_safe_failure() {
let _ = server.await;
}
#[tokio::test]
async fn previous_response_not_found_should_remain_client_visible_without_internal_replay_proof() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_client_history").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(concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":400,\"response\":{\"id\":\"resp_history_missing\",\"status\":\"failed\",\"error\":{\"code\":\"previous_response_not_found\",\"message\":\"Previous response was not found. Retrying the full request.\"}}}\n\n"
)),
)
.expect(1)
.mount(&server)
.await;
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
(
"input".to_owned(),
json!([{"role": "user", "content": "delta"}]),
),
(
"previous_response_id".to_owned(),
json!("missing-previous-response"),
),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))]));
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request(
"openai",
Operation::Generate(GenerateRequest::from_protocol_payload(payload)),
),
context("req_history_missing", CancellationToken::new()),
)
.await
.expect("prepare missing-history stream");
let mut failure = loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("previous_response_not_found must produce a typed failure"),
}
};
assert_eq!(
failure.continuation_failure(),
Some(ContinuationFailure::HistoryUnavailable)
);
assert!(!failure.replay_is_safe());
assert_eq!(
failure
.upstream_code()
.map(gateway_core::error::OpaqueUpstreamValue::as_str),
Some("previous_response_not_found")
);
let events = failure.take_atomic_client_events();
assert_eq!(events.len(), 1);
assert_eq!(
events[0].wire_event().and_then(|wire| wire.event_type()),
Some("response.failed")
);
}
#[tokio::test]
async fn continuation_prefetch_over_64_kib_should_commit_wire_without_protocol_failure() {
let store = Arc::new(MemoryAccountStore::default());
@@ -622,6 +622,18 @@ impl MemorySessionAffinity {
pub(crate) fn binding_count(&self) -> usize {
self.bindings.lock().expect("session affinity lock").len()
}
pub(crate) fn seed_binding(
&self,
provider_kind: &ProviderKind,
key: &str,
account_id: ProviderAccountId,
) {
self.bindings.lock().expect("session affinity lock").insert(
(provider_kind.as_str().to_owned(), key.to_owned()),
account_id,
);
}
}
impl ProviderSessionAffinityPort for MemorySessionAffinity {
@@ -371,6 +371,7 @@ async fn codex_backend_client_should_capture_forwardable_response_metadata() {
.insert_header("openai-model", "gpt-5.5-2026-07-01")
.insert_header("x-models-etag", "models-v2")
.insert_header("x-reasoning-included", "true")
.insert_header("x-codex-turn-state", "turn-state-from-upstream")
.insert_header("openai-processing-ms", "42")
.append_header("x-future-multi", "first")
.append_header("x-future-multi", "second")
@@ -452,6 +453,7 @@ async fn codex_backend_client_should_capture_forwardable_response_metadata() {
"openai-processing-ms",
"retry-after",
"x-codex-safety-buffering-enabled",
"x-codex-turn-state",
"x-future-bytes",
"x-future-multi",
"x-future-multi",
@@ -478,6 +480,15 @@ async fn codex_backend_client_should_capture_forwardable_response_metadata() {
.iter()
.any(|(name, value)| { name == "x-future-bytes" && value.as_ref() == b"\xffopaque" })
);
assert!(
response
.response_metadata
.client_headers
.iter()
.any(|(name, value)| {
name == "x-codex-turn-state" && value.as_ref() == b"turn-state-from-upstream"
})
);
for blocked in [
"set-cookie",
"authorization",
@@ -16,6 +16,13 @@ fn new_chain_request(conversation_id: &str) -> CodexResponsesRequest {
request
}
fn external_continuation_request(conversation_id: &str) -> CodexResponsesRequest {
let mut request = new_chain_request(conversation_id);
request.set_previous_response_id(Some("resp_from_another_account".to_owned()));
request.previous_response_scope = Some(PreviousResponseScope::ExternalUnknown);
request
}
fn explicit_websocket_warmup_request(conversation_id: &str) -> CodexResponsesRequest {
let mut body = serde_json::Map::new();
body.insert("model".to_string(), json!("gpt-5.5"));
@@ -108,6 +115,43 @@ async fn cold_websocket_should_fall_back_without_recording_a_successful_connect(
pool.shutdown().await;
}
#[tokio::test(start_paused = true)]
async fn external_continuation_should_wait_for_a_cold_websocket() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (websocket_stream, _) = listener.accept().await.unwrap();
tokio::time::sleep(Duration::from_secs(1)).await;
let mut websocket = accept_codex_test_websocket(websocket_stream).await;
let _payload = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_cross_account", 2, 1).into(),
))
.await
.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::default());
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::clone(&pool));
let response = backend
.create_response(
&external_continuation_request("conversation-cross-account"),
request_context("req_cross_account", Some("chatgpt-account-b")),
)
.await
.expect("external continuation should complete over WebSocket");
server.await.unwrap();
assert_eq!(response.transport, CodexBackendTransport::WebSocket);
pool.shutdown().await;
}
#[tokio::test(start_paused = true)]
async fn timed_out_websocket_should_finish_in_background_and_serve_the_next_request() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -9,7 +9,9 @@ use std::{
use chrono::Utc;
use futures::{SinkExt, StreamExt};
use gateway_protocol::openai::events::{TokenUsage, extract_sse_usage};
use gateway_protocol::openai::events::{
TokenUsage, extract_sse_usage, rate_limits_to_header_pairs,
};
use provider_openai::transport::profile::{
CodexBundledReleaseProfile, CodexWireProfile, CodexWireProfileState,
};
@@ -233,7 +235,7 @@ async fn collect_backend_response(
mut turn_state,
set_cookie_headers,
mut rate_limit_headers,
rate_limit_header_updates,
rate_limit_updates,
turn_state_update,
websocket_pool_decision,
diagnostics: _,
@@ -251,8 +253,10 @@ async fn collect_backend_response(
});
body_bytes.extend_from_slice(&chunk);
}
if let Some(updates) = rate_limit_header_updates {
rate_limit_headers.extend(updates.lock().await.iter().cloned());
if let Some(updates) = rate_limit_updates {
for update in updates.lock().await.iter() {
rate_limit_headers.extend(rate_limits_to_header_pairs(update));
}
}
if let Some(update) = turn_state_update {
turn_state = update.lock().await.clone().or(turn_state);
@@ -159,6 +159,23 @@ fn websocket_event_to_sse_should_forward_public_events_and_strip_internal_events
assert!(websocket_event_to_sse_frame("not-json").is_none());
}
#[test]
fn websocket_event_to_sse_should_add_missing_previous_response_recovery_code() {
let event = json!({
"type": "error",
"status": 400,
"error": {
"type": "invalid_request_error",
"message": "Invalid `previous_response_id`."
}
})
.to_string();
let frame = websocket_event_to_sse_frame(&event).expect("public error event");
assert!(frame.contains(r#""code":"previous_response_not_found""#));
}
#[test]
fn websocket_metadata_turn_state_should_accept_case_insensitive_header() {
let event = json!({
+5 -1
View File
@@ -104,7 +104,11 @@ continuation 或 claim 表。
- `store=true`:使用 Provider 持久化的 native handle。
- `store=false`:opaque replay state 仅存在于活连接内,不落 PostgreSQL。
- OpenAI continuation 顺序为 native、replay owner、replay any。
- OpenAI continuation 顺序为 native、replay owner、replay any。`replay any` 是
Provider 定义的恢复阶段,不代表 Core 已持有完整 input:OpenAI 在跨账号时保留
canonical `previous_response_id` 与原始增量 input 做 continuation probe;上游返回
`previous_response_not_found` 时原样交给官方客户端触发 full retry。只有确实持有
可携带 transcript 的 Provider 才能清除 native handle 后在代理内部重放。
- xAI 使用客户端提交的完整历史作为已验证 continuation 路径。
- native continuation pin 同时绑定创建它的 Client Key、冻结账号范围、Provider 和账号;跨 Key 复用、
scope 外账号或计划中不存在该 Provider 都 fail closed。`ReplayOwner` 同样固定原 Provider 和账号,
+1 -1
View File
@@ -1 +1 @@
version: "3.2.11"
version: "3.2.12"