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

Isolate standalone key wallet handling

Browse Source
This commit is contained in:
fawney19
2026-04-27 12:54:59 +08:00
parent 3b542434a2
commit 1d23bf4ecb
12 changed files with 240 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -115,6 +115,9 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
.map(|wallet| wallet.id.clone())
})
.or_else(|| {
if input.api_key_is_standalone {
return None;
}
input.user_id.as_deref().and_then(|user_id| {
wallets
.values()
@@ -215,6 +218,25 @@ mod tests {
.expect("wallet should build")
}
fn sample_user_wallet(wallet_id: &str, user_id: &str) -> StoredWalletSnapshot {
StoredWalletSnapshot::new(
wallet_id.to_string(),
Some(user_id.to_string()),
None,
10.0,
2.0,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
0.0,
0.0,
0.0,
0.0,
100,
)
.expect("wallet should build")
}
#[tokio::test]
async fn settles_usage_against_wallet_and_provider_quota() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
@@ -223,6 +245,7 @@ mod tests {
request_id: "req-1".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: "completed".to_string(),
billing_status: "pending".to_string(),
@@ -240,6 +263,60 @@ mod tests {
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5));
}
#[tokio::test]
async fn normal_key_settlement_falls_back_to_user_wallet() {
let repository =
InMemorySettlementRepository::seed(vec![sample_user_wallet("wallet-user-1", "user-1")]);
let settlement = repository
.settle_usage(UsageSettlementInput {
request_id: "req-user-wallet".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("normal-key-without-wallet".to_string()),
api_key_is_standalone: false,
provider_id: None,
status: "completed".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.wallet_id.as_deref(), Some("wallet-user-1"));
assert_eq!(settlement.wallet_balance_before, Some(12.0));
assert_eq!(settlement.wallet_balance_after, Some(9.0));
}
#[tokio::test]
async fn standalone_key_settlement_never_falls_back_to_owner_wallet() {
let repository = InMemorySettlementRepository::seed(vec![sample_user_wallet(
"wallet-admin-owner",
"admin-owner",
)]);
let settlement = repository
.settle_usage(UsageSettlementInput {
request_id: "req-standalone-no-key-wallet".to_string(),
user_id: Some("admin-owner".to_string()),
api_key_id: Some("standalone-key-without-wallet".to_string()),
api_key_is_standalone: true,
provider_id: None,
status: "completed".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.wallet_id, None);
assert_eq!(settlement.wallet_balance_before, None);
assert_eq!(settlement.wallet_balance_after, None);
}
#[tokio::test]
async fn returns_stored_snapshot_when_usage_is_already_finalized() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
@@ -248,6 +325,7 @@ mod tests {
request_id: "req-2".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: "completed".to_string(),
billing_status: "pending".to_string(),
@@ -264,6 +342,7 @@ mod tests {
request_id: "req-2".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: "completed".to_string(),
billing_status: "settled".to_string(),

59
crates/aether-data/src/repository/settlement/sql.rs
View File

@@ -247,11 +247,31 @@ impl SettlementWriteRepository for SqlxSettlementRepository {
};
if final_billing_status == "settled" {
let wallet_row = if let Some(api_key_id) = input
let api_key_id = input
.api_key_id
.as_deref()
.filter(|value| !value.is_empty())
{
.filter(|value| !value.is_empty());
let api_key_is_standalone = if input.api_key_is_standalone {
true
} else if let Some(api_key_id) = api_key_id {
sqlx::query_scalar::<_, bool>(
r#"
SELECT is_standalone
FROM api_keys
WHERE id = $1
LIMIT 1
"#,
)
.bind(api_key_id)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
.unwrap_or(false)
} else {
false
};
let wallet_row = if let Some(api_key_id) = api_key_id {
sqlx::query(
r#"
SELECT
@@ -275,11 +295,12 @@ LIMIT 1
let wallet_row = if wallet_row.is_some() {
wallet_row
} else if let Some(user_id) =
input.user_id.as_deref().filter(|value| !value.is_empty())
{
sqlx::query(
r#"
} else if !api_key_is_standalone {
if let Some(user_id) =
input.user_id.as_deref().filter(|value| !value.is_empty())
{
sqlx::query(
r#"
SELECT
id,
CAST(balance AS DOUBLE PRECISION) AS balance,
@@ -289,12 +310,15 @@ FROM wallets
WHERE user_id = $1
FOR UPDATE
LIMIT 1
"#,
)
.bind(user_id)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
"#,
)
.bind(user_id)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
} else {
None
}
} else {
None
};
@@ -417,4 +441,11 @@ mod tests {
let source = include_str!("sql.rs");
assert!(!source.contains("UPDATE \"usage\"\nSET\n wallet_id = $2"));
}
#[test]
fn settlement_sql_blocks_standalone_key_owner_wallet_fallback() {
let source = include_str!("sql.rs");
assert!(source.contains("SELECT is_standalone"));
assert!(source.contains("} else if !api_key_is_standalone {"));
}
}

1
crates/aether-data/src/repository/wallet/types.rs
View File

@@ -992,6 +992,7 @@ mod tests {
request_id: "".to_string(),
user_id: None,
api_key_id: None,
api_key_is_standalone: false,
provider_id: None,
status: "completed".to_string(),
billing_status: "pending".to_string(),