mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 10:27:46 +08:00
修复 ChatGPT Web 生图调用起始预扣额度
This commit is contained in:
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user