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

fix(data): keep terminal usage state monotonic

Browse Source
This commit is contained in:
MMEXA
2026-07-12 05:45:37 +08:00
parent 72222d935c
commit 25c49dd804
3 changed files with 128 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/src/repository/usage/memory.rs
+68 -15
View File
@@ -2761,14 +2761,18 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
})
.unwrap_or_default();
if existing.as_ref().is_some_and(|existing| {
usage_status_is_finalized(existing.status.as_str())
&& usage_status_is_lifecycle(usage.status.as_str())
&& !usage_can_recover_terminal_failure(
existing.status.as_str(),
existing.billing_status.as_str(),
usage.status.as_str(),
usage.billing_status.as_str(),
)
let can_recover = usage_can_recover_terminal_failure(
existing.status.as_str(),
existing.billing_status.as_str(),
usage.status.as_str(),
usage.billing_status.as_str(),
);
let finalized_lifecycle_regression = usage_status_is_finalized(&existing.status)
&& usage_status_is_lifecycle(&usage.status);
let completed_terminal_failure_recovery = existing.billing_status == "void"
&& matches!(existing.status.as_str(), "failed" | "cancelled")
&& usage.status == "completed";
(finalized_lifecycle_regression || completed_terminal_failure_recovery) && !can_recover
}) {
return Ok(existing.expect("existing usage should be present").clone());
}
@@ -3589,7 +3593,7 @@ mod tests {
}
#[tokio::test]
async fn upsert_allows_streaming_recovery_after_void_failure() {
async fn upsert_allows_completed_recovery_after_void_failure() {
let repository = InMemoryUsageReadRepository::default();
repository
.upsert(UpsertUsageRecord {
@@ -3698,12 +3702,12 @@ mod tests {
output_price_per_1m: None,
total_cost_usd: None,
actual_total_cost_usd: None,
status_code: None,
status_code: Some(200),
error_message: None,
error_category: None,
response_time_ms: Some(45),
first_byte_time_ms: Some(12),
status: "streaming".to_string(),
status: "completed".to_string(),
billing_status: "pending".to_string(),
request_headers: None,
request_body: None,
@@ -3732,7 +3736,7 @@ mod tests {
request_metadata: Some(json!({
"trace_id": "trace-recovered"
})),
finalized_at_unix_secs: None,
finalized_at_unix_secs: Some(102),
created_at_unix_ms: Some(100),
updated_at_unix_secs: 102,
})
@@ -3744,11 +3748,11 @@ mod tests {
.await
.expect("usage lookup should succeed")
.expect("usage should exist");
assert_eq!(stored.status, "streaming");
assert_eq!(stored.status, "completed");
assert_eq!(stored.billing_status, "pending");
assert_eq!(stored.status_code, None);
assert_eq!(stored.status_code, Some(200));
assert_eq!(stored.error_message, None);
assert_eq!(stored.finalized_at_unix_secs, None);
assert_eq!(stored.finalized_at_unix_secs, Some(102));
assert_eq!(
stored.request_metadata,
Some(json!({ "trace_id": "trace-recovered" }))
@@ -3756,6 +3760,55 @@ mod tests {
assert_eq!(stored.total_tokens, 10);
}
#[tokio::test]
async fn upsert_rejects_non_authoritative_void_failure_recovery() {
let repository = InMemoryUsageReadRepository::default();
for (request_id, status, billing_status, status_code) in [
("req-late-active-1", "streaming", "pending", None),
(
"req-late-response-start-1",
"streaming",
"pending",
Some(200),
),
(
"req-settled-completion-1",
"completed",
"settled",
Some(200),
),
] {
repository
.upsert(UpsertUsageRecord {
status: "failed".to_string(),
billing_status: "void".to_string(),
status_code: Some(503),
finalized_at_unix_secs: Some(101),
updated_at_unix_secs: 101,
..sample_upsert_usage_record(request_id)
})
.await
.expect("failed usage should upsert");
let stored = repository
.upsert(UpsertUsageRecord {
status: status.to_string(),
billing_status: billing_status.to_string(),
status_code,
finalized_at_unix_secs: None,
updated_at_unix_secs: 102,
..sample_upsert_usage_record(request_id)
})
.await
.expect("non-authoritative recovery should be ignored");
assert_eq!(stored.status, "failed");
assert_eq!(stored.billing_status, "void");
assert_eq!(stored.status_code, Some(503));
assert_eq!(stored.finalized_at_unix_secs, Some(101));
}
}
#[tokio::test]
async fn stale_pending_update_does_not_reopen_void_failure() {
let repository = InMemoryUsageReadRepository::default();
crates/aether-data/src/repository/usage/mod.rs
+21 -18
View File
@@ -583,10 +583,7 @@ pub(crate) fn incoming_usage_can_recover_terminal_failure(
incoming_status: &str,
incoming_billing_status: &str,
) -> bool {
incoming_billing_status == "pending"
// Late pending placeholders are not authoritative enough to reopen a void terminal row;
// they can otherwise regress a real failure back to pending when background writes race.
&& matches!(incoming_status, "streaming" | "completed")
incoming_billing_status == "pending" && incoming_status == "completed"
}
pub(crate) fn usage_can_recover_terminal_failure(
@@ -827,52 +824,58 @@ mod tests {
}
#[test]
fn incoming_usage_recovery_requires_streaming_or_completed_state() {
fn incoming_usage_recovery_requires_completed_state() {
assert!(incoming_usage_can_recover_terminal_failure(
"completed",
"pending"
"pending",
));
assert!(incoming_usage_can_recover_terminal_failure(
assert!(!incoming_usage_can_recover_terminal_failure(
"streaming",
"pending"
"pending",
));
assert!(!incoming_usage_can_recover_terminal_failure(
"pending", "pending"
"pending", "pending",
));
assert!(!incoming_usage_can_recover_terminal_failure(
"failed", "void"
"failed", "void",
));
assert!(!incoming_usage_can_recover_terminal_failure(
"completed",
"settled"
"settled",
));
}
#[test]
fn usage_recovery_requires_void_failure_to_be_followed_by_streaming_or_completed_state() {
fn usage_recovery_requires_void_failure_and_completed_state() {
assert!(usage_can_recover_terminal_failure(
"failed",
"void",
"completed",
"pending"
"pending",
));
assert!(usage_can_recover_terminal_failure(
"cancelled",
"void",
"streaming",
"pending"
"completed",
"pending",
));
assert!(!usage_can_recover_terminal_failure(
"failed", "void", "pending", "pending"
"failed",
"void",
"streaming",
"pending",
));
assert!(!usage_can_recover_terminal_failure(
"failed", "void", "pending", "pending",
));
assert!(!usage_can_recover_terminal_failure(
"completed",
"pending",
"completed",
"pending"
"pending",
));
assert!(!usage_can_recover_terminal_failure(
"failed", "void", "failed", "void"
"failed", "void", "failed", "void",
));
}
crates/aether-data/src/repository/usage/sqlite.rs
+39
View File
@@ -4598,6 +4598,45 @@ mod tests {
assert_eq!(existing.updated_at_unix_secs, 1_000);
}
#[tokio::test]
async fn sqlite_usage_write_repository_does_not_reopen_void_failure_from_late_streaming() {
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);
for (request_id, status_code) in [
("request-late-active", None),
("request-late-response-start", Some(200)),
] {
let mut failed = sample_usage(request_id, "failed", "void", 1_000);
failed.status_code = Some(503);
repository
.upsert(failed)
.await
.expect("failed usage should upsert");
let mut late_streaming = sample_usage(request_id, "streaming", "pending", 1_001);
late_streaming.status_code = status_code;
late_streaming.finalized_at_unix_secs = None;
let current = repository
.upsert(late_streaming)
.await
.expect("late streaming usage should be ignored");
assert_eq!(current.status, "failed");
assert_eq!(current.billing_status, "void");
assert_eq!(current.status_code, Some(503));
assert_eq!(current.finalized_at_unix_secs, Some(1_000));
}
}
#[tokio::test]
async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late_streaming() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()