修复 ChatGPT Web 生图调用起始预扣额度

This commit is contained in:
Codex
2026-05-24 18:50:14 +08:00
parent 9dba6ec1d9
commit 14b182d09b
6 changed files with 218 additions and 8 deletions
@@ -130,6 +130,17 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_stream(
}))
}
pub(crate) async fn maybe_apply_chatgpt_web_image_quota_request_delta_at_candidate_start(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
) {
if !is_chatgpt_web_image_plan(plan, report_context) {
return;
}
apply_chatgpt_web_image_quota_request_delta_at_start(state, plan).await;
}
fn is_chatgpt_web_image_plan(plan: &ExecutionPlan, report_context: Option<&Value>) -> bool {
if !plan
.provider_api_format
@@ -1062,10 +1073,12 @@ async fn apply_chatgpt_web_image_quota_request_delta(
.unwrap_or_default();
let now_unix_secs = current_unix_secs();
let request_dedup_key = chatgpt_web_image_quota_request_delta_dedup_key(plan);
if !apply_chatgpt_web_image_quota_request_delta_to_metadata(
&mut metadata,
latest_key.status_snapshot.as_ref(),
now_unix_secs,
request_dedup_key.as_deref(),
) {
return Ok(false);
}
@@ -1097,7 +1110,21 @@ fn apply_chatgpt_web_image_quota_request_delta_to_metadata(
metadata: &mut Map<String, Value>,
status_snapshot: Option<&Value>,
now_unix_secs: u64,
request_dedup_key: Option<&str>,
) -> bool {
let request_dedup_key = request_dedup_key
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(request_dedup_key) = request_dedup_key {
if metadata
.get("image_quota_last_local_request_key")
.and_then(Value::as_str)
.is_some_and(|value| value == request_dedup_key)
{
return false;
}
}
let snapshot_window = chatgpt_web_image_quota_snapshot_window(status_snapshot);
let metadata_limit =
chatgpt_web_image_quota_f64(metadata.get("image_quota_total")).filter(|value| *value > 0.0);
@@ -1163,6 +1190,12 @@ fn apply_chatgpt_web_image_quota_request_delta_to_metadata(
"image_quota_last_local_request_at".to_string(),
json!(now_unix_secs),
);
if let Some(request_dedup_key) = request_dedup_key {
metadata.insert(
"image_quota_last_local_request_key".to_string(),
json!(request_dedup_key),
);
}
let local_request_count =
chatgpt_web_image_quota_u64(metadata.get("image_quota_local_request_count")).unwrap_or(0);
metadata.insert(
@@ -1172,6 +1205,22 @@ fn apply_chatgpt_web_image_quota_request_delta_to_metadata(
true
}
fn chatgpt_web_image_quota_request_delta_dedup_key(plan: &ExecutionPlan) -> Option<String> {
let request_id = plan.request_id.trim();
if request_id.is_empty() {
return None;
}
let candidate_id = plan
.candidate_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
Some(match candidate_id {
Some(candidate_id) => format!("{request_id}:{candidate_id}"),
None => request_id.to_string(),
})
}
#[derive(Debug, Clone)]
struct ChatGptWebImageQuotaRequestLimit {
value: f64,
@@ -2754,6 +2803,7 @@ mod tests {
&mut metadata,
None,
1_000,
None,
));
assert_eq!(metadata["image_quota_remaining"], json!(24.0));
@@ -2783,6 +2833,7 @@ mod tests {
&mut metadata,
Some(&status_snapshot),
1_000,
None,
));
assert_eq!(metadata["image_quota_remaining"], json!(18.0));
@@ -2804,6 +2855,7 @@ mod tests {
&mut metadata,
None,
1_000,
None,
));
assert_eq!(metadata["image_quota_remaining"], json!(18.0));
@@ -2815,6 +2867,44 @@ mod tests {
);
}
#[test]
fn chatgpt_web_image_quota_request_delta_dedupes_same_candidate_start() {
let mut metadata = Map::from_iter([
("plan_type".to_string(), json!("free")),
("image_quota_remaining".to_string(), json!(25.0)),
("image_quota_total".to_string(), json!(25.0)),
("image_quota_used".to_string(), json!(0.0)),
]);
assert!(apply_chatgpt_web_image_quota_request_delta_to_metadata(
&mut metadata,
None,
1_000,
Some("request-1:candidate-1"),
));
assert!(!apply_chatgpt_web_image_quota_request_delta_to_metadata(
&mut metadata,
None,
1_001,
Some("request-1:candidate-1"),
));
assert!(apply_chatgpt_web_image_quota_request_delta_to_metadata(
&mut metadata,
None,
1_002,
Some("request-1:candidate-2"),
));
assert_eq!(metadata["image_quota_remaining"], json!(23.0));
assert_eq!(metadata["image_quota_total"], json!(25.0));
assert_eq!(metadata["image_quota_used"], json!(2.0));
assert_eq!(metadata["image_quota_local_request_count"], json!(2u64));
assert_eq!(
metadata["image_quota_last_local_request_key"],
json!("request-1:candidate-2")
);
}
async fn start_mock_chatgpt_web() -> (String, tokio::task::JoinHandle<()>) {
let app = Router::new().fallback(any(|request: Request| async move {
let path = request.uri().path().to_string();
@@ -61,7 +61,10 @@ use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
use crate::control::GatewayControlDecision;
use crate::execution_runtime::build_direct_execution_frame_stream;
use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_stream;
use crate::execution_runtime::chatgpt_web_image::{
maybe_apply_chatgpt_web_image_quota_request_delta_at_candidate_start,
maybe_execute_chatgpt_web_image_stream,
};
use crate::execution_runtime::grok::maybe_execute_grok_stream;
use crate::execution_runtime::kiro_cache::{
billed_input_tokens as kiro_billed_input_tokens, build_kiro_prompt_cache_profile,
@@ -858,6 +861,12 @@ pub(crate) async fn execute_execution_runtime_stream(
.await;
});
}
maybe_apply_chatgpt_web_image_quota_request_delta_at_candidate_start(
state,
&plan,
report_context.as_ref(),
)
.await;
let plan_request_id_for_log = short_request_id(plan.request_id.as_str());
let provider_name = plan
.provider_name
@@ -39,7 +39,10 @@ use crate::api::response::{
};
use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync;
use crate::execution_runtime::chatgpt_web_image::{
maybe_apply_chatgpt_web_image_quota_request_delta_at_candidate_start,
maybe_execute_chatgpt_web_image_sync,
};
use crate::execution_runtime::grok::maybe_execute_grok_sync;
use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
#[cfg(test)]
@@ -1335,6 +1338,12 @@ async fn execute_execution_runtime_sync_impl(
},
)
.await;
maybe_apply_chatgpt_web_image_quota_request_delta_at_candidate_start(
state,
&plan,
report_context.as_ref(),
)
.await;
let mut terminal_guard = SyncAttemptTerminalGuard::new(
state,
&plan,
@@ -442,13 +442,49 @@ fn chatgpt_web_image_quota_limit(
.get("image_quota_total")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.filter(|value| *value > 0.0);
let plan_type = metadata
.get("plan_type")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase());
let limit_source = metadata
.get("image_quota_limit_source")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(limit) = explicit_limit {
return Some(limit);
if !chatgpt_web_image_quota_limit_is_legacy_free_default(
limit,
limit_source,
plan_type.as_deref(),
remaining,
) {
return Some(limit);
}
}
remaining.filter(|value| *value > 0.0)
}
fn chatgpt_web_image_quota_limit_is_legacy_free_default(
limit: f64,
limit_source: Option<&str>,
plan_type: Option<&str>,
remaining: Option<f64>,
) -> bool {
let plan_type_is_free = plan_type
.map(str::trim)
.is_some_and(|value| value.eq_ignore_ascii_case("free"));
if !plan_type_is_free || limit_source.is_some() {
return false;
}
if (limit - 25.0).abs() > f64::EPSILON {
return false;
}
remaining.is_some_and(|value| value < limit)
}
fn model_quota_window_snapshot(
model_name: &str,
item: &Map<String, Value>,
@@ -2533,6 +2569,33 @@ mod tests {
assert_eq!(window.get("remaining_ratio"), Some(&json!(1.0)));
}
#[test]
fn provider_key_status_snapshot_payload_ignores_chatgpt_web_legacy_free_25_limit() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"chatgpt_web": {
"updated_at": 1_778_067_246u64,
"plan_type": "free",
"image_quota_remaining": 19.0,
"image_quota_total": 25.0
}
}));
let payload = provider_key_status_snapshot_payload(&key, "chatgpt_web");
let window = payload
.get("quota")
.and_then(Value::as_object)
.and_then(|quota| quota.get("windows"))
.and_then(Value::as_array)
.and_then(|windows| windows.first())
.and_then(Value::as_object)
.expect("image quota window should exist");
assert_eq!(window.get("remaining_value"), Some(&json!(19.0)));
assert_eq!(window.get("limit_value"), Some(&json!(19.0)));
assert_eq!(window.get("used_value"), Some(&json!(0.0)));
}
#[test]
fn provider_key_status_snapshot_payload_backfills_grok_model_quota() {
let mut key = sample_catalog_key();