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

Merge remote-tracking branch 'origin/pr/538'

Browse Source
This commit is contained in:
fawney19
2026-05-21 23:12:42 +08:00
parent d6eb41aa78 cadc45c5b8
commit 0e6d7539ad
16 changed files with 422 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
crates/aether-data/src/repository/usage/memory.rs
View File

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::RwLock;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_data_contracts::repository::usage::{
parse_usage_body_ref, usage_body_ref, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
@@ -952,7 +953,7 @@ fn usage_output_tps_uses_generation_time(item: &StoredRequestUsageAudit) -> bool
item.request_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.unwrap_or(item.is_stream)
}

81
crates/aether-data/src/repository/usage/mysql.rs
View File

@@ -1,5 +1,6 @@
use std::collections::{BTreeMap, HashSet};
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, Row};
@@ -518,13 +519,20 @@ WHERE request_id = ?
continue;
}
let error_message = stale_pending_error_message(&row.status, timeout_minutes);
let candidate_info =
latest_failed_candidate_mysql(&mut tx, &row.request_id).await?;
let (status_code, error_message) = resolve_stale_pending_failure(
candidate_info.as_ref(),
&row.status,
timeout_minutes,
);
let status_code_i64 = i64::from(status_code);
if row.billing_status == "pending" {
sqlx::query(
r#"
UPDATE `usage`
SET status = 'failed',
status_code = 504,
status_code = ?,
error_message = ?,
billing_status = 'void',
finalized_at = ?,
@@ -533,6 +541,7 @@ SET status = 'failed',
WHERE request_id = ?
"#,
)
.bind(status_code_i64)
.bind(&error_message)
.bind(to_i64(now_unix_secs, "usage finalized_at")?)
.bind(&row.request_id)
@@ -550,11 +559,12 @@ WHERE request_id = ?
r#"
UPDATE `usage`
SET status = 'failed',
status_code = 504,
status_code = ?,
error_message = ?
WHERE request_id = ?
"#,
)
.bind(status_code_i64)
.bind(&error_message)
.bind(&row.request_id)
.execute(&mut *tx)
@@ -687,6 +697,67 @@ fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String {
format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成")
}
struct FailedCandidateCleanupInfo {
status_code: Option<u16>,
error_message: Option<String>,
}
fn resolve_stale_pending_failure(
candidate: Option<&FailedCandidateCleanupInfo>,
status: &str,
timeout_minutes: u64,
) -> (u16, String) {
match candidate {
Some(info) => (
info.status_code.unwrap_or(502),
info.error_message
.clone()
.unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)),
),
None => (504, stale_pending_error_message(status, timeout_minutes)),
}
}
async fn latest_failed_candidate_mysql(
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
request_id: &str,
) -> Result<Option<FailedCandidateCleanupInfo>, DataLayerError> {
let row = sqlx::query(
r#"
SELECT status_code, error_message
FROM request_candidates
WHERE request_id = ?
AND status IN ('failed', 'cancelled')
ORDER BY
COALESCE(finished_at, started_at, created_at) DESC,
retry_index DESC,
candidate_index DESC
LIMIT 1
"#,
)
.bind(request_id)
.fetch_optional(&mut **tx)
.await
.map_sql_err()?;
let Some(row) = row else {
return Ok(None);
};
let status_code = row
.try_get::<Option<i64>, _>("status_code")
.map_sql_err()?
.and_then(|value| u16::try_from(value).ok());
let error_message = row
.try_get::<Option<String>, _>("error_message")
.map_sql_err()?
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
Ok(Some(FailedCandidateCleanupInfo {
status_code,
error_message,
}))
}
fn bind_upsert<'q>(
mut query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>,
usage: &'q UpsertUsageRecord,
@@ -881,7 +952,7 @@ fn usage_upstream_is_stream(usage: &UpsertUsageRecord) -> bool {
.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or_else(|| usage.is_stream.unwrap_or(false))
}
@@ -895,7 +966,7 @@ fn merge_usage_stream_metadata(metadata: &mut Option<serde_json::Value>, upstrea
return;
};
object
.entry("upstream_is_stream")
.entry(UPSTREAM_IS_STREAM_KEY)
.or_insert(serde_json::Value::Bool(upstream));
}

84
crates/aether-data/src/repository/usage/postgres/mod.rs
View File

@@ -1690,10 +1690,25 @@ SET status = 'completed',
WHERE request_id = $1
"#;
const SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL: &str = r#"
SELECT DISTINCT ON (request_id)
request_id,
status_code,
error_message
FROM request_candidates
WHERE request_id = ANY($1)
AND status IN ('failed', 'cancelled')
ORDER BY request_id,
COALESCE(finished_at, started_at, created_at) DESC,
retry_index DESC,
candidate_index DESC,
created_at DESC
"#;
const UPDATE_FAILED_STALE_USAGE_SQL: &str = r#"
UPDATE usage
SET status = 'failed',
status_code = 504,
status_code = $3,
error_message = $2
WHERE request_id = $1
"#;
@@ -1702,7 +1717,7 @@ const UPDATE_FAILED_VOID_STALE_USAGE_SQL: &str = r#"
WITH updated_usage AS (
UPDATE usage
SET status = 'failed',
status_code = 504,
status_code = $4,
error_message = $2,
billing_status = 'void',
finalized_at = $3,
@@ -8292,17 +8307,44 @@ ORDER BY "usage".user_id ASC
.iter()
.map(|row| row.request_id.clone())
.collect::<Vec<_>>();
let completed_request_ids = if request_ids.is_empty() {
Vec::new()
let (completed_request_ids, failed_candidate_info) = if request_ids.is_empty() {
(Vec::new(), std::collections::HashMap::new())
} else {
sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
.bind(request_ids)
let completed = sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
.bind(&request_ids)
.fetch_all(&mut *tx)
.await
.map_postgres_err()?
.iter()
.map(|row| row.try_get("request_id").map_postgres_err())
.collect::<Result<Vec<String>, DataLayerError>>()?
.collect::<Result<Vec<String>, DataLayerError>>()?;
let failed_rows =
sqlx::query(SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL)
.bind(&request_ids)
.fetch_all(&mut *tx)
.await
.map_postgres_err()?;
let mut failed_map = std::collections::HashMap::new();
for row in failed_rows {
let request_id: String = row.try_get("request_id").map_postgres_err()?;
let status_code = row
.try_get::<Option<i32>, _>("status_code")
.map_postgres_err()?
.and_then(|value| u16::try_from(value).ok());
let error_message = row
.try_get::<Option<String>, _>("error_message")
.map_postgres_err()?
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
failed_map.insert(
request_id,
FailedCandidateCleanupInfo {
status_code,
error_message,
},
);
}
(completed, failed_map)
};
for row in stale_rows {
@@ -8322,12 +8364,16 @@ ORDER BY "usage".user_id ASC
continue;
}
let error_message = stale_pending_error_message(&row.status, timeout_minutes);
let candidate_info = failed_candidate_info.get(&row.request_id);
let (status_code, error_message) =
resolve_stale_pending_failure(candidate_info, &row.status, timeout_minutes);
let status_code_i32 = i32::from(status_code);
if row.billing_status == "pending" {
sqlx::query(UPDATE_FAILED_VOID_STALE_USAGE_SQL)
.bind(&row.request_id)
.bind(&error_message)
.bind(now)
.bind(status_code_i32)
.execute(&mut *tx)
.await
.map_postgres_err()?;
@@ -8335,6 +8381,7 @@ ORDER BY "usage".user_id ASC
sqlx::query(UPDATE_FAILED_STALE_USAGE_SQL)
.bind(&row.request_id)
.bind(&error_message)
.bind(status_code_i32)
.execute(&mut *tx)
.await
.map_postgres_err()?;
@@ -8785,10 +8832,31 @@ struct StalePendingUsageRow {
billing_status: String,
}
struct FailedCandidateCleanupInfo {
status_code: Option<u16>,
error_message: Option<String>,
}
fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String {
format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成")
}
fn resolve_stale_pending_failure(
candidate: Option<&FailedCandidateCleanupInfo>,
status: &str,
timeout_minutes: u64,
) -> (u16, String) {
match candidate {
Some(info) => (
info.status_code.unwrap_or(502),
info.error_message
.clone()
.unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)),
),
None => (504, stale_pending_error_message(status, timeout_minutes)),
}
}
async fn find_usage_by_request_id_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
request_id: &str,

8
crates/aether-data/src/repository/usage/postgres/tests.rs
View File

@@ -825,6 +825,14 @@ fn usage_sql_clears_stale_failure_fields_for_non_failed_status_updates() {
));
}
#[test]
fn stale_cleanup_failed_candidate_sql_orders_by_effective_timestamp() {
let sql = super::SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL;
assert!(sql.contains("COALESCE(finished_at, started_at, created_at) DESC"));
assert!(!sql.contains("finished_at DESC NULLS LAST"));
assert!(!sql.contains("started_at DESC NULLS LAST"));
}
#[test]
fn usage_sql_does_not_allow_streaming_to_regress_back_to_pending() {
assert!(super::UPSERT_SQL.contains(

166
crates/aether-data/src/repository/usage/sqlite.rs
View File

@@ -1,6 +1,7 @@
use std::collections::{BTreeMap, HashSet};
use std::io::Read;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_data_contracts::repository::usage::{parse_usage_body_ref, UsageBodyField};
use async_trait::async_trait;
use flate2::read::GzDecoder;
@@ -3548,13 +3549,20 @@ WHERE request_id = ?
continue;
}
let error_message = stale_pending_error_message(&row.status, timeout_minutes);
let candidate_info =
latest_failed_candidate_sqlite(&mut tx, &row.request_id).await?;
let (status_code, error_message) = resolve_stale_pending_failure(
candidate_info.as_ref(),
&row.status,
timeout_minutes,
);
let status_code_i64 = i64::from(status_code);
if row.billing_status == "pending" {
sqlx::query(
r#"
UPDATE "usage"
SET status = 'failed',
status_code = 504,
status_code = ?,
error_message = ?,
billing_status = 'void',
finalized_at = ?,
@@ -3563,6 +3571,7 @@ SET status = 'failed',
WHERE request_id = ?
"#,
)
.bind(status_code_i64)
.bind(&error_message)
.bind(to_i64(now_unix_secs, "usage finalized_at")?)
.bind(&row.request_id)
@@ -3580,11 +3589,12 @@ WHERE request_id = ?
r#"
UPDATE "usage"
SET status = 'failed',
status_code = 504,
status_code = ?,
error_message = ?
WHERE request_id = ?
"#,
)
.bind(status_code_i64)
.bind(&error_message)
.bind(&row.request_id)
.execute(&mut *tx)
@@ -3707,6 +3717,67 @@ fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String {
format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成")
}
struct FailedCandidateCleanupInfo {
status_code: Option<u16>,
error_message: Option<String>,
}
fn resolve_stale_pending_failure(
candidate: Option<&FailedCandidateCleanupInfo>,
status: &str,
timeout_minutes: u64,
) -> (u16, String) {
match candidate {
Some(info) => (
info.status_code.unwrap_or(502),
info.error_message
.clone()
.unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)),
),
None => (504, stale_pending_error_message(status, timeout_minutes)),
}
}
async fn latest_failed_candidate_sqlite(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
request_id: &str,
) -> Result<Option<FailedCandidateCleanupInfo>, DataLayerError> {
let row = sqlx::query(
r#"
SELECT status_code, error_message
FROM request_candidates
WHERE request_id = ?
AND status IN ('failed', 'cancelled')
ORDER BY
COALESCE(finished_at, started_at, created_at) DESC,
retry_index DESC,
candidate_index DESC
LIMIT 1
"#,
)
.bind(request_id)
.fetch_optional(&mut **tx)
.await
.map_sql_err()?;
let Some(row) = row else {
return Ok(None);
};
let status_code = row
.try_get::<Option<i64>, _>("status_code")
.map_sql_err()?
.and_then(|value| u16::try_from(value).ok());
let error_message = row
.try_get::<Option<String>, _>("error_message")
.map_sql_err()?
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
Ok(Some(FailedCandidateCleanupInfo {
status_code,
error_message,
}))
}
fn bind_upsert<'q>(
mut query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>,
usage: &'q UpsertUsageRecord,
@@ -3898,7 +3969,7 @@ fn usage_upstream_is_stream(usage: &UpsertUsageRecord) -> bool {
.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or_else(|| usage.is_stream.unwrap_or(false))
}
@@ -3912,7 +3983,7 @@ fn merge_usage_stream_metadata(metadata: &mut Option<serde_json::Value>, upstrea
return;
};
object
.entry("upstream_is_stream")
.entry(UPSTREAM_IS_STREAM_KEY)
.or_insert(serde_json::Value::Bool(upstream));
}
@@ -4139,6 +4210,91 @@ ORDER BY request_id
assert_eq!(snapshot, ("void".to_string(), Some(10)));
}
#[tokio::test]
async fn sqlite_usage_write_repository_cleanup_uses_failed_candidate_status_when_present() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool should connect");
run_sqlite_migrations(&pool)
.await
.expect("sqlite migrations should run");
seed_stats_targets(&pool).await;
let repository = SqliteUsageWriteRepository::new(pool.clone());
repository
.upsert(sample_usage(
"request-upstream-reset",
"pending",
"pending",
1,
))
.await
.expect("pending usage should upsert");
repository
.upsert(sample_usage("request-stuck", "pending", "pending", 1))
.await
.expect("pending usage should upsert");
// request-upstream-reset has a failed candidate carrying a concrete 502 status
// and a connection-reset message — cleanup should use them instead of 504.
// request-stuck has only a still-pending candidate, so cleanup should fall back to 504.
sqlx::query(
r#"
INSERT INTO request_candidates (
id,
request_id,
candidate_index,
retry_index,
status,
status_code,
error_message,
is_cached,
created_at,
started_at,
finished_at
) VALUES
('candidate-reset', 'request-upstream-reset', 0, 0, 'failed', 502, 'upstream connection reset by peer', 0, 1, 2, 3),
('candidate-stuck', 'request-stuck', 0, 0, 'pending', NULL, NULL, 0, 1, NULL, NULL)
"#,
)
.execute(&pool)
.await
.expect("request candidates should seed");
let summary = repository
.cleanup_stale_pending_requests(2, 10, 5, 5)
.await
.expect("cleanup should run");
assert_eq!(summary.recovered, 0);
assert_eq!(summary.failed, 2);
let reset = repository
.find_by_request_id("request-upstream-reset")
.await
.expect("upstream-reset usage should load")
.expect("upstream-reset usage should exist");
assert_eq!(reset.status, "failed");
assert_eq!(reset.status_code, Some(502));
assert_eq!(
reset.error_message.as_deref(),
Some("upstream connection reset by peer")
);
let stuck = repository
.find_by_request_id("request-stuck")
.await
.expect("stuck usage should load")
.expect("stuck usage should exist");
assert_eq!(stuck.status, "failed");
assert_eq!(stuck.status_code, Some(504));
assert!(stuck
.error_message
.as_deref()
.is_some_and(|message| message.contains("超过 5 分钟未完成")));
}
#[tokio::test]
async fn sqlite_usage_read_repository_reads_usage_contract_views() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()