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

fix(usage): bill cancelled terminal usage and preserve terminal telemetry

Browse Source
This commit is contained in:
ZheFox
2026-05-16 23:43:40 +08:00
parent a2f91b4108
commit 2803e9317d
8 changed files with 205 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -1035,6 +1035,33 @@ fn should_refresh_stream_usage_telemetry(
|| (next_elapsed.is_some() && next_elapsed != previous_elapsed)
}
fn build_terminal_stream_telemetry(
stream_started_at: Instant,
telemetry: Option<&ExecutionTelemetry>,
usage_stream_telemetry: Option<&ExecutionTelemetry>,
upstream_bytes: u64,
) -> ExecutionTelemetry {
let current_elapsed_ms = stream_started_at
.elapsed()
.as_millis()
.min(u128::from(u64::MAX)) as u64;
let ttfb_ms = telemetry
.and_then(|telemetry| telemetry.ttfb_ms)
.or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.ttfb_ms));
let prior_elapsed_ms = telemetry
.and_then(|telemetry| telemetry.elapsed_ms)
.or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.elapsed_ms))
.unwrap_or(0);
let elapsed_ms = current_elapsed_ms
.max(prior_elapsed_ms)
.max(ttfb_ms.unwrap_or(0));
ExecutionTelemetry {
ttfb_ms,
elapsed_ms: Some(elapsed_ms),
upstream_bytes: Some(upstream_bytes),
}
}
fn should_skip_direct_finalize_prefetch(
direct_stream_finalize_kind: Option<&str>,
content_type: Option<&str>,
@@ -2738,6 +2765,12 @@ async fn execute_stream_from_frame_stream(
trace_id = %trace_id_owned,
"gateway skipped stream report because downstream disconnected before completion"
);
let terminal_telemetry = Some(build_terminal_stream_telemetry(
stream_started_at_for_report,
telemetry.as_ref(),
usage_stream_telemetry.as_ref(),
provider_stream_bytes.load(Ordering::Relaxed),
));
let usage_payload = build_stream_usage_payload(
trace_id_owned,
report_kind_owned.unwrap_or_default(),
@@ -2749,7 +2782,7 @@ async fn execute_stream_from_frame_stream(
&buffered_body,
client_body_truncated,
stream_terminal_summary,
telemetry,
terminal_telemetry,
);
record_stream_terminal_usage(
&state_for_report,
@@ -2781,6 +2814,12 @@ async fn execute_stream_from_frame_stream(
if let Some(failure) = terminal_failure {
record_manual_proxy_stream_error(&state_for_report, &plan_for_report).await;
let terminal_telemetry = Some(build_terminal_stream_telemetry(
stream_started_at_for_report,
telemetry.as_ref(),
usage_stream_telemetry.as_ref(),
provider_stream_bytes.load(Ordering::Relaxed),
));
submit_midstream_stream_failure(
&state_for_report,
&trace_id_owned,
@@ -2788,7 +2827,7 @@ async fn execute_stream_from_frame_stream(
direct_stream_finalize_kind_owned.as_deref(),
report_context_owned,
headers_for_report,
telemetry,
terminal_telemetry,
&provider_buffered_body,
candidate_started_unix_secs_for_report,
failure,
@@ -2798,6 +2837,12 @@ async fn execute_stream_from_frame_stream(
}
let should_submit_report = report_kind_owned.is_some();
let terminal_telemetry = Some(build_terminal_stream_telemetry(
stream_started_at_for_report,
telemetry.as_ref(),
usage_stream_telemetry.as_ref(),
provider_stream_bytes.load(Ordering::Relaxed),
));
let usage_payload = build_stream_usage_payload(
trace_id_owned.clone(),
report_kind_owned.unwrap_or_default(),
@@ -2809,7 +2854,7 @@ async fn execute_stream_from_frame_stream(
&buffered_body,
client_body_truncated,
stream_terminal_summary,
telemetry,
terminal_telemetry,
);
apply_local_execution_effect(
&state_for_report,
@@ -3312,6 +3357,7 @@ mod tests {
.await
.expect("first business chunk should arrive");
assert_eq!(first.as_ref(), b"data: {\"id\":\"first\"}\n\n");
tokio::time::sleep(Duration::from_millis(30)).await;
drop(body_stream);
tokio::time::timeout(Duration::from_secs(1), frame_stream_dropped.notified())
@@ -3339,6 +3385,36 @@ mod tests {
candidates[0].error_type.as_deref(),
Some("downstream_disconnect")
);
let stored_usage = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let usage = usage_repository
.find_by_request_id("req-client-drop-cancels-upstream")
.await
.expect("usage should read");
if usage
.as_ref()
.is_some_and(|usage| usage.status == "cancelled")
{
break usage.expect("cancelled usage should exist");
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("usage should be marked cancelled");
assert_eq!(stored_usage.billing_status, "pending");
assert_eq!(stored_usage.status_code, Some(499));
let first_byte_time_ms = stored_usage
.first_byte_time_ms
.expect("cancelled stream should retain first byte time");
let response_time_ms = stored_usage
.response_time_ms
.expect("cancelled stream should record terminal duration");
assert!(
response_time_ms > first_byte_time_ms,
"terminal duration should include time after the first byte"
);
}
#[tokio::test]

37
crates/aether-data/src/repository/settlement/memory.rs
View File

@@ -4,8 +4,8 @@ use std::sync::{Arc, RwLock};
use async_trait::async_trait;
use super::{
plan_finite_wallet_debit, SettlementWriteRepository, StoredUsageSettlement,
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
plan_finite_wallet_debit, settlement_billing_status_for_usage_status,
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
};
use crate::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
use crate::DataLayerError;
@@ -102,11 +102,8 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
})));
}
let mut final_billing_status = if input.status == "completed" {
"settled".to_string()
} else {
"void".to_string()
};
let mut final_billing_status =
settlement_billing_status_for_usage_status(&input.status).to_string();
let mut settlement = self.wallets.with_mut(|wallets| {
let wallet_id = input
.api_key_id
@@ -306,6 +303,32 @@ mod tests {
assert_eq!(settlement.wallet_balance_after, Some(9.0));
}
#[tokio::test]
async fn settles_cancelled_usage_against_wallet_and_provider_quota() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
let settlement = repository
.settle_usage(UsageSettlementInput {
request_id: "req-cancelled".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
api_key_is_standalone: false,
provider_id: Some("provider-1".to_string()),
status: "cancelled".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 1.5,
finalized_at_unix_secs: Some(200),
})
.await
.expect("settlement should succeed")
.expect("settlement should exist");
assert_eq!(settlement.billing_status, "settled");
assert_eq!(settlement.wallet_balance_before, Some(12.0));
assert_eq!(settlement.wallet_balance_after, Some(9.0));
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5));
}
#[tokio::test]
async fn standalone_key_settlement_never_falls_back_to_owner_wallet() {
let repository = InMemorySettlementRepository::seed(vec![sample_user_wallet(

25
crates/aether-data/src/repository/settlement/mod.rs
View File

@@ -36,6 +36,31 @@ fn plan_finite_wallet_debit(
}
}
fn settlement_billing_status_for_usage_status(status: &str) -> &'static str {
match status {
"completed" | "cancelled" => "settled",
_ => "void",
}
}
#[cfg(test)]
mod tests {
use super::settlement_billing_status_for_usage_status;
#[test]
fn cancelled_usage_status_is_billable() {
assert_eq!(
settlement_billing_status_for_usage_status("completed"),
"settled"
);
assert_eq!(
settlement_billing_status_for_usage_status("cancelled"),
"settled"
);
assert_eq!(settlement_billing_status_for_usage_status("failed"), "void");
}
}
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::settlement::{
SettlementRepository, SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,

12
crates/aether-data/src/repository/settlement/mysql.rs
View File

@@ -2,8 +2,9 @@ use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, Row};
use super::{
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
finite_wallet_available_usd, plan_finite_wallet_debit,
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
};
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
@@ -364,11 +365,8 @@ impl SettlementWriteRepository for MysqlSettlementRepository {
return Ok(Some(settlement));
}
let mut final_billing_status = if input.status == "completed" {
"settled".to_string()
} else {
"void".to_string()
};
let mut final_billing_status =
settlement_billing_status_for_usage_status(&input.status).to_string();
let mut settlement = StoredUsageSettlement {
request_id: input.request_id.clone(),
wallet_id: None,

12
crates/aether-data/src/repository/settlement/postgres.rs
View File

@@ -2,8 +2,9 @@ use async_trait::async_trait;
use sqlx::{PgPool, Row};
use super::{
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
finite_wallet_available_usd, plan_finite_wallet_debit,
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
};
use crate::driver::postgres::PostgresTransactionRunner;
use crate::error::SqlxResultExt;
@@ -408,11 +409,8 @@ impl SettlementWriteRepository for SqlxSettlementRepository {
return settlement_from_row(&usage_row).map(Some);
}
let mut final_billing_status = if input.status == "completed" {
"settled".to_string()
} else {
"void".to_string()
};
let mut final_billing_status =
settlement_billing_status_for_usage_status(&input.status).to_string();
let finalized_at =
i64::try_from(input.finalized_at_unix_secs.unwrap_or_else(|| {
std::time::SystemTime::now()

12
crates/aether-data/src/repository/settlement/sqlite.rs
View File

@@ -2,8 +2,9 @@ use async_trait::async_trait;
use sqlx::{sqlite::SqliteRow, Row};
use super::{
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
finite_wallet_available_usd, plan_finite_wallet_debit,
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
};
use crate::driver::sqlite::{sqlite_optional_real, sqlite_real, SqlitePool};
use crate::error::SqlResultExt;
@@ -377,11 +378,8 @@ impl SettlementWriteRepository for SqliteSettlementRepository {
return Ok(Some(settlement));
}
let mut final_billing_status = if input.status == "completed" {
"settled".to_string()
} else {
"void".to_string()
};
let mut final_billing_status =
settlement_billing_status_for_usage_status(&input.status).to_string();
let mut settlement = StoredUsageSettlement {
request_id: input.request_id.clone(),
wallet_id: None,

34
crates/aether-usage-runtime/src/record.rs
View File

@@ -136,7 +136,7 @@ fn lifecycle_status_and_billing(event_type: UsageEventType) -> (&'static str, &'
UsageEventType::Streaming => ("streaming", "pending"),
UsageEventType::Completed => ("completed", "pending"),
UsageEventType::Failed => ("failed", "void"),
UsageEventType::Cancelled => ("cancelled", "void"),
UsageEventType::Cancelled => ("cancelled", "pending"),
}
}
@@ -181,6 +181,38 @@ mod tests {
assert_eq!(record.finalized_at_unix_secs, Some(1_700_000_000));
}
#[test]
fn cancelled_terminal_record_stays_pending_for_settlement() {
let record = build_upsert_usage_record_from_event(&UsageEvent {
event_type: UsageEventType::Cancelled,
request_id: "req-cancelled".to_string(),
timestamp_ms: 1_700_000_000_000,
data: UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
input_tokens: Some(10),
output_tokens: Some(20),
total_tokens: Some(30),
total_cost_usd: Some(0.03),
actual_total_cost_usd: Some(0.02),
status_code: Some(499),
response_time_ms: Some(200),
first_byte_time_ms: Some(50),
..UsageEventData::default()
},
})
.expect("record should build");
assert_eq!(record.status, "cancelled");
assert_eq!(record.billing_status, "pending");
assert_eq!(record.total_tokens, Some(30));
assert_eq!(record.total_cost_usd, Some(0.03));
assert_eq!(record.actual_total_cost_usd, Some(0.02));
assert_eq!(record.status_code, Some(499));
assert_eq!(record.response_time_ms, Some(200));
assert_eq!(record.first_byte_time_ms, Some(50));
}
#[test]
fn sanitizes_request_metadata_before_building_upsert_record() {
let record = build_upsert_usage_record_from_event(&UsageEvent {

23
crates/aether-usage-runtime/src/settlement.rs
View File

@@ -164,6 +164,29 @@ mod tests {
assert!(!inputs[0].api_key_is_standalone);
}
#[tokio::test]
async fn settles_pending_cancelled_usage() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.status = "cancelled".to_string();
usage.status_code = Some(499);
settle_usage_if_needed(&writer, &usage)
.await
.expect("settlement should succeed");
let inputs = writer.inputs.lock().expect("settlement inputs lock");
assert_eq!(inputs.len(), 1);
assert_eq!(inputs[0].request_id, "req-1");
assert_eq!(inputs[0].status, "cancelled");
assert_eq!(inputs[0].billing_status, "pending");
assert_eq!(inputs[0].total_cost_usd, 1.25);
assert_eq!(inputs[0].actual_total_cost_usd, 0.75);
}
#[tokio::test]
async fn propagates_standalone_key_flag_from_usage_metadata() {
let writer = TestSettlementWriter {