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

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

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(

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,

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,

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()

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,