starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-13 08:06:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): preserve Codex quota after image generation

Browse Source
This commit is contained in:
Zyy
2026-11-08 18:04:45 +08:00
parent ce620f1b53
commit e770c94a9e
4 changed files with 151 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-protocol/src/openai/events.rs
+11 -1
View File
@@ -328,7 +328,7 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option<ParsedRa
limit_ids.insert(active_limit.clone());
}
limit_ids.insert("codex".to_string());
let limits = limit_ids
let mut limits = limit_ids
.into_iter()
.filter_map(|limit_id| {
let prefix = format!("x-{}", limit_id.replace('_', "-"));
@@ -336,6 +336,16 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option<ParsedRa
.map(|details| (limit_id, details))
})
.collect::<BTreeMap<_, _>>();
// 生图等响应只通过 active-limit 标明通用 x-codex 窗口的归属
// 有独立具名头时保留各桶原始事实，避免把真正的 Codex 窗口一起迁移
if let Some(active_limit) = active_limit.as_ref()
&& active_limit != "codex"
&& !limits.contains_key(active_limit)
&& let Some(mut details) = limits.remove("codex")
{
details.limit_id.clone_from(active_limit);
limits.insert(active_limit.clone(), details);
}
let credits = parse_credits_from_lookup(&normalized);
let plan_type = lookup_non_empty(&normalized, "x-codex-plan-type");
let promo_message = lookup_non_empty(&normalized, "x-codex-promo-message");
backend/crates/gateway-protocol/tests/openai/events.rs
+40
View File
@@ -540,6 +540,7 @@ fn parse_rate_limit_headers_should_extract_primary_and_secondary_windows() {
#[test]
fn parse_rate_limit_headers_should_preserve_dynamic_limit_and_account_metadata() {
let headers = vec![
("x-codex-primary-used-percent".to_owned(), "95".to_owned()),
(
"x-codex-other-primary-used-percent".to_owned(),
"63.5".to_owned(),
@@ -556,6 +557,13 @@ fn parse_rate_limit_headers_should_preserve_dynamic_limit_and_account_metadata()
];
let parsed = parse_rate_limit_headers(&headers).expect("rate-limit headers");
assert_eq!(
parsed.limits["codex"]
.primary
.expect("core window")
.used_percent,
95.0
);
let details = parsed.limits.get("codex_other").expect("dynamic limit");
let credits = parsed.credits.expect("credits");
@@ -579,6 +587,38 @@ fn parse_rate_limit_headers_should_preserve_dynamic_limit_and_account_metadata()
);
}
#[test]
fn parse_rate_limit_headers_should_assign_generic_windows_to_the_active_limit() {
let headers = [
("x-codex-active-limit", "imagegen_premium"),
("x-codex-primary-used-percent", "0"),
("x-codex-primary-window-minutes", "1440"),
("x-codex-primary-reset-at", "1900000000"),
("x-codex-secondary-used-percent", "0"),
("x-codex-secondary-window-minutes", "0"),
("x-codex-secondary-reset-at", ""),
]
.map(|(name, value)| (name.to_owned(), value.to_owned()));
let parsed = parse_rate_limit_headers(&headers).expect("image rate-limit headers");
let image = parsed.limits.get("imagegen_premium").expect("image limit");
assert_eq!(image.limit_id, "imagegen_premium");
assert_eq!(image.key_source, RateLimitKeySource::LimitId);
assert_eq!(
image.primary,
Some(RateLimitWindow {
used_percent: 0.0,
window_minutes: Some(1440),
reset_at: Some(1_900_000_000),
})
);
assert!(!parsed.limits.contains_key("codex"));
assert_eq!(
parse_rate_limit_headers(&rate_limits_to_header_pairs(&parsed)),
Some(parsed)
);
}
#[test]
fn parse_rate_limit_headers_should_preserve_a_code_review_window() {
let headers = vec![
backend/crates/providers/openai/tests/provider/contract/mod.rs
+96
View File
@@ -2888,6 +2888,102 @@ async fn image_endpoints_bypass_only_the_text_catalog_and_preserve_the_current_c
server.verify().await;
}
#[tokio::test]
async fn image_quota_headers_keep_the_codex_weekly_window_and_access_state() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_provider_contract";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("account");
let server = MockServer::start().await;
let (provider, quota) = provider_and_quota_with_affinity_and_base_url_and_leases(
&store,
Arc::new(MemorySessionAffinity::default()),
server.uri(),
Arc::new(TestLeaseCoordinator::default()),
0,
);
let core_headers = [
("x-codex-active-limit", "codex"),
("x-codex-primary-used-percent", "95"),
("x-codex-primary-window-minutes", "10080"),
("x-codex-primary-reset-at", "1900604800"),
("x-codex-allowed", "true"),
]
.map(|(name, value)| (name.to_owned(), value.to_owned()));
quota
.synchronize_passive_headers(&account, &core_headers)
.await
.expect("initial weekly quota");
let initial = quota.read_account(account.id()).await.unwrap().unwrap();
let response = br#"{"data":[{"b64_json":"AAEC"}]}"#;
// 生图响应仅用 active-limit 标明归属，窗口仍使用通用 x-codex 头
Mock::given(method("POST"))
.and(path("/codex/images/generations"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("x-codex-active-limit", "imagegen_premium")
.insert_header("x-codex-primary-used-percent", "0")
.insert_header("x-codex-primary-window-minutes", "1440")
.insert_header("x-codex-primary-reset-at", "1900000000")
.insert_header("x-codex-secondary-used-percent", "0")
.insert_header("x-codex-secondary-window-minutes", "0")
.insert_header("x-codex-secondary-reset-at", "")
.set_body_raw(response.as_slice(), "application/json"),
)
.expect(1)
.mount(&server)
.await;
let payload = RawJsonPayload::new(
"openai",
Bytes::from_static(br#"{"model":"gpt-image-2.6-flare","prompt":"a tree"}"#),
)
.unwrap();
let operation = Operation::GenerateImage(ImageRequest::from_raw_json(
ImageRequestKind::Generation,
payload,
));
let mut stream = provider
.execute(
planned_provider_endpoint_request("openai", operation),
context("req_image_quota", CancellationToken::new()),
)
.await
.expect("image stream");
let mut delivered = None;
while let Some(event) = stream.next().await {
let event = event.expect("image event");
if let Some(body) = event.wire_event().and_then(|wire| wire.raw_json_body()) {
delivered = Some(body.clone());
}
}
assert_eq!(delivered.as_deref(), Some(response.as_slice()));
let snapshot = quota.read_account(account.id()).await.unwrap().unwrap();
let core = snapshot
.windows()
.iter()
.find(|window| window.is_account_wide())
.unwrap();
let image = snapshot
.windows()
.iter()
.find(|window| window.source() == "imagegen_premium")
.unwrap();
assert_eq!(core, &initial.windows()[0]);
assert_eq!(image.window_seconds(), Some(86_400));
assert_eq!(image.used_percent(), Some(0.0));
assert!(!image.is_account_wide());
assert_eq!(store.account(account_id).unwrap().quota(), initial.quota());
quota
.synchronize_passive_headers(&account, &core_headers)
.await
.expect("next Codex quota observation");
let refreshed = quota.read_account(account.id()).await.unwrap().unwrap();
assert_eq!(refreshed.windows(), snapshot.windows());
server.verify().await;
}
#[tokio::test]
async fn image_endpoints_should_report_usage_before_delivering_the_unchanged_body() {
let response = br#"{ "created":1778832973,"data":[{"b64_json":"AAEC"}],"quality":"medium","size":"1024x1536","usage":{"input_tokens":1474,"input_tokens_details":{"image_tokens":1457,"text_tokens":17,"cached_tokens":100},"output_tokens":1372,"output_tokens_details":{"image_tokens":1372,"text_tokens":0},"total_tokens":2846},"future":9007199254740993 }"#;
docs/architecture.md
+4
View File
@@ -1029,6 +1029,10 @@ credential 与 quota 是两组独立事实：credential refresh 不等于 quota
OpenAI Provider 将其中明确的套餐变更与额度原子提交，共用凭据版本和观察时间保护；
空值及 `unknown` 不覆盖套餐，同族泛化值保留具体子类型。Token 刷新保留提交时的账号资料
OpenAI 响应头的额度窗口由 `gateway-protocol` 按上游计量项归属解析；仅通过 `x-codex-active-limit`
标明具名额度时，通用 `x-codex-*` 窗口归入该计量项。Provider 按计量项独立保存，生图等具名额度
不覆盖默认 Codex 窗口，也不改变账号级套餐访问结论
OpenAI Provider 统一解释点数可用性，主动查询、耗尽恢复与失败回读复用同一规则。
Store 在账号事务内保存点数放行依据与快照；关闭自动使用时撤销该依据，迟到观测不能绕过开关。
放行条件与默认值见[账号额度 API](api.md#查询与账号字段)。