i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: 修正用量失败状态判断与 Codex 会话头

Browse Source
This commit is contained in:
fawney19
2026-04-10 21:33:34 +08:00
parent 46f1507d44
commit 3f057628b7
9 changed files with 199 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
crates/aether-admin/src/observability/usage.rs
View File

@@ -181,6 +181,18 @@ pub fn admin_usage_matches_api_format(
.is_some_and(|value| value.eq_ignore_ascii_case(api_format))
}
pub fn admin_usage_is_failed(item: &StoredRequestUsageAudit) -> bool {
let status = item.status.trim();
if !status.is_empty() {
return status.eq_ignore_ascii_case("failed");
}
item.status_code.is_some_and(|value| value >= 400)
|| item
.error_message
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
}
pub fn admin_usage_matches_status(item: &StoredRequestUsageAudit, status: Option<&str>) -> bool {
let Some(status) = status.map(str::trim).filter(|value| !value.is_empty()) else {
return true;
@@ -192,11 +204,7 @@ pub fn admin_usage_matches_status(item: &StoredRequestUsageAudit, status: Option
item.status_code.is_some_and(|value| value >= 400) || item.error_message.is_some()
}
"pending" | "streaming" | "completed" | "cancelled" => item.status == status,
"failed" => {
item.status == "failed"
|| item.status_code.is_some_and(|value| value >= 400)
|| item.error_message.is_some()
}
"failed" => admin_usage_is_failed(item),
"active" => matches!(item.status.as_str(), "pending" | "streaming"),
_ => true,
}
@@ -227,6 +235,76 @@ pub fn admin_usage_provider_key_name(
.or_else(|| admin_usage_request_metadata_string(item, "key_name"))
}
#[cfg(test)]
mod tests {
use super::{admin_usage_is_failed, admin_usage_matches_status};
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
fn sample_usage(
status: &str,
status_code: Option<i32>,
error_message: Option<&str>,
) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
"OpenAI".to_string(),
"gpt-5".to_string(),
None,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
false,
false,
10,
20,
30,
0.0,
0.0,
status_code,
error_message.map(ToOwned::to_owned),
None,
Some(120),
None,
status.to_string(),
"settled".to_string(),
100,
101,
Some(102),
)
.expect("usage should build")
}
#[test]
fn explicit_completed_status_wins_over_legacy_failure_fields() {
let item = sample_usage("completed", Some(429), Some("rate limited on first attempt"));
assert!(!admin_usage_is_failed(&item));
assert!(!admin_usage_matches_status(&item, Some("failed")));
assert!(admin_usage_matches_status(&item, Some("completed")));
}
#[test]
fn legacy_failure_signals_still_work_when_status_is_missing() {
let item = StoredRequestUsageAudit {
status: String::new(),
..sample_usage("completed", Some(429), Some("rate limited"))
};
assert!(admin_usage_is_failed(&item));
assert!(admin_usage_matches_status(&item, Some("failed")));
}
}
pub fn admin_usage_record_json(
item: &StoredRequestUsageAudit,
users_by_id: &BTreeMap<String, StoredUserSummary>,

30
crates/aether-ai-pipeline/src/planner/standard/codex.rs
View File

@@ -181,41 +181,13 @@ pub fn apply_codex_openai_cli_special_headers(
return;
}
if let Some(account_id) = extract_codex_account_id(decrypted_auth_config_raw) {
provider_request_headers.insert("chatgpt-account-id".to_string(), account_id);
}
if !provider_request_headers
.get("x-client-request-id")
.map(|value| !value.trim().is_empty())
.unwrap_or(false)
{
if let Some(request_id) = request_id.map(str::trim).filter(|value| !value.is_empty()) {
provider_request_headers
.insert("x-client-request-id".to_string(), request_id.to_string());
}
}
let prompt_cache_key = provider_request_body
.get("prompt_cache_key")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let Some(short_id) = prompt_cache_key.and_then(build_short_codex_header_id) else {
return;
};
if !header_map_has_non_empty_value(original_headers, "session_id") {
provider_request_headers.insert("session_id".to_string(), short_id.clone());
}
if !provider_api_format
.trim()
.eq_ignore_ascii_case("openai:compact")
&& !header_map_has_non_empty_value(original_headers, "conversation_id")
{
let session_id = provider_request_headers.get("session_id").unwrap();
provider_request_headers.insert("conversation_id".to_string(), session_id.clone());
provider_request_headers.insert("session_id".to_string(), prompt_cache_key.unwrap().to_string());
}
}

28
crates/aether-data/src/repository/usage/sql.rs
View File

@@ -415,9 +415,18 @@ DO UPDATE SET
output_price_per_1m = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.output_price_per_1m, "usage".output_price_per_1m) ELSE "usage".output_price_per_1m END,
total_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.total_cost_usd, "usage".total_cost_usd) ELSE "usage".total_cost_usd END,
actual_total_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.actual_total_cost_usd, "usage".actual_total_cost_usd) ELSE "usage".actual_total_cost_usd END,
status_code = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.status_code, "usage".status_code) ELSE "usage".status_code END,
error_message = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.error_message, "usage".error_message) ELSE "usage".error_message END,
error_category = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.error_category, "usage".error_category) ELSE "usage".error_category END,
status_code = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') AND EXCLUDED.status_code IS NULL THEN NULL
ELSE COALESCE(EXCLUDED.status_code, "usage".status_code)
END ELSE "usage".status_code END,
error_message = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') THEN EXCLUDED.error_message
ELSE COALESCE(EXCLUDED.error_message, "usage".error_message)
END ELSE "usage".error_message END,
error_category = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') THEN EXCLUDED.error_category
ELSE COALESCE(EXCLUDED.error_category, "usage".error_category)
END ELSE "usage".error_category END,
response_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.response_time_ms, "usage".response_time_ms) ELSE "usage".response_time_ms END,
first_byte_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.first_byte_time_ms, "usage".first_byte_time_ms) ELSE "usage".first_byte_time_ms END,
status = CASE WHEN "usage".billing_status = 'pending' THEN EXCLUDED.status ELSE "usage".status END,
@@ -1074,4 +1083,17 @@ mod tests {
assert!(super::UPSERT_SQL.contains("WHEN $50 IS NULL THEN NULL"));
assert!(super::UPSERT_SQL.contains("TO_TIMESTAMP($51::double precision)"));
}
#[test]
fn usage_sql_clears_stale_failure_fields_for_non_failed_status_updates() {
assert!(super::UPSERT_SQL.contains(
"WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') AND EXCLUDED.status_code IS NULL THEN NULL"
));
assert!(super::UPSERT_SQL.contains(
"WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') THEN EXCLUDED.error_message"
));
assert!(super::UPSERT_SQL.contains(
"WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') THEN EXCLUDED.error_category"
));
}
}