Isolate standalone key wallet handling

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

View File

@@ -129,6 +129,7 @@ pub(crate) fn build_execution_runtime_auth_context(
api_key_name: auth_context.api_key_name.clone(), api_key_name: auth_context.api_key_name.clone(),
balance_remaining: auth_context.balance_remaining, balance_remaining: auth_context.balance_remaining,
access_allowed: auth_context.access_allowed, access_allowed: auth_context.access_allowed,
api_key_is_standalone: auth_context.api_key_is_standalone,
} }
} }

View File

@@ -50,6 +50,10 @@ pub(crate) fn build_local_execution_report_context(
"api_key_id".to_string(), "api_key_id".to_string(),
Value::String(parts.auth_context.api_key_id.clone()), Value::String(parts.auth_context.api_key_id.clone()),
); );
object.insert(
"api_key_is_standalone".to_string(),
Value::Bool(parts.auth_context.api_key_is_standalone),
);
object.insert( object.insert(
"username".to_string(), "username".to_string(),
parts parts

View File

@@ -22,7 +22,10 @@ pub(crate) async fn resolve_wallet_auth_gate(
auth_snapshot.api_key_is_standalone, auth_snapshot.api_key_is_standalone,
) )
.await?; .await?;
let is_admin = auth_snapshot.user_role.eq_ignore_ascii_case("admin"); let is_admin = wallet_auth_allows_admin_bypass(
&auth_snapshot.user_role,
auth_snapshot.api_key_is_standalone,
);
Ok(Some(match wallet.as_ref() { Ok(Some(match wallet.as_ref() {
Some(wallet) => map_wallet_snapshot(wallet).access_decision(is_admin), Some(wallet) => map_wallet_snapshot(wallet).access_decision(is_admin),
@@ -60,12 +63,18 @@ fn map_wallet_snapshot(snapshot: &StoredWalletSnapshot) -> WalletSnapshot {
} }
} }
fn wallet_auth_allows_admin_bypass(user_role: &str, api_key_is_standalone: bool) -> bool {
user_role.eq_ignore_ascii_case("admin") && !api_key_is_standalone
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use aether_data::repository::wallet::StoredWalletSnapshot; use aether_data::repository::wallet::StoredWalletSnapshot;
use aether_wallet::{WalletAccessFailure, WalletLimitMode, WalletSnapshot, WalletStatus}; use aether_wallet::{WalletAccessFailure, WalletLimitMode, WalletSnapshot, WalletStatus};
use super::{local_rejection_from_wallet_access, map_wallet_snapshot}; use super::{
local_rejection_from_wallet_access, map_wallet_snapshot, wallet_auth_allows_admin_bypass,
};
use crate::control::GatewayLocalAuthRejection; use crate::control::GatewayLocalAuthRejection;
#[test] #[test]
@@ -114,4 +123,10 @@ mod tests {
assert!(decision.allowed); assert!(decision.allowed);
assert_eq!(decision.remaining, None); assert_eq!(decision.remaining, None);
} }
#[test]
fn standalone_key_never_uses_admin_wallet_bypass() {
assert!(wallet_auth_allows_admin_bypass("admin", false));
assert!(!wallet_auth_allows_admin_bypass("admin", true));
}
} }

View File

@@ -1,5 +1,9 @@
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
fn is_false(value: &bool) -> bool {
!*value
}
#[derive(Debug, Clone, Deserialize, Serialize)] #[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ExecutionRuntimeAuthContext { pub struct ExecutionRuntimeAuthContext {
pub user_id: String, pub user_id: String,
@@ -10,4 +14,6 @@ pub struct ExecutionRuntimeAuthContext {
pub api_key_name: Option<String>, pub api_key_name: Option<String>,
pub balance_remaining: Option<f64>, pub balance_remaining: Option<f64>,
pub access_allowed: bool, pub access_allowed: bool,
#[serde(default, skip_serializing_if = "is_false")]
pub api_key_is_standalone: bool,
} }

View File

@@ -268,6 +268,7 @@ mod tests {
api_key_name: None, api_key_name: None,
balance_remaining: Some(12.5), balance_remaining: Some(12.5),
access_allowed: true, access_allowed: true,
api_key_is_standalone: false,
}), }),
); );

View File

@@ -5,6 +5,8 @@ pub struct UsageSettlementInput {
pub request_id: String, pub request_id: String,
pub user_id: Option<String>, pub user_id: Option<String>,
pub api_key_id: Option<String>, pub api_key_id: Option<String>,
#[serde(default)]
pub api_key_is_standalone: bool,
pub provider_id: Option<String>, pub provider_id: Option<String>,
pub status: String, pub status: String,
pub billing_status: String, pub billing_status: String,
@@ -71,6 +73,7 @@ mod tests {
request_id: "".to_string(), request_id: "".to_string(),
user_id: None, user_id: None,
api_key_id: None, api_key_id: None,
api_key_is_standalone: false,
provider_id: None, provider_id: None,
status: "completed".to_string(), status: "completed".to_string(),
billing_status: "pending".to_string(), billing_status: "pending".to_string(),

View File

@@ -115,6 +115,9 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
.map(|wallet| wallet.id.clone()) .map(|wallet| wallet.id.clone())
}) })
.or_else(|| { .or_else(|| {
if input.api_key_is_standalone {
return None;
}
input.user_id.as_deref().and_then(|user_id| { input.user_id.as_deref().and_then(|user_id| {
wallets wallets
.values() .values()
@@ -215,6 +218,25 @@ mod tests {
.expect("wallet should build") .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] #[tokio::test]
async fn settles_usage_against_wallet_and_provider_quota() { async fn settles_usage_against_wallet_and_provider_quota() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]); let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
@@ -223,6 +245,7 @@ mod tests {
request_id: "req-1".to_string(), request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()), user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()), api_key_id: Some("key-1".to_string()),
api_key_is_standalone: false,
provider_id: Some("provider-1".to_string()), provider_id: Some("provider-1".to_string()),
status: "completed".to_string(), status: "completed".to_string(),
billing_status: "pending".to_string(), billing_status: "pending".to_string(),
@@ -240,6 +263,60 @@ mod tests {
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5)); 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] #[tokio::test]
async fn returns_stored_snapshot_when_usage_is_already_finalized() { async fn returns_stored_snapshot_when_usage_is_already_finalized() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]); let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
@@ -248,6 +325,7 @@ mod tests {
request_id: "req-2".to_string(), request_id: "req-2".to_string(),
user_id: Some("user-1".to_string()), user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()), api_key_id: Some("key-1".to_string()),
api_key_is_standalone: false,
provider_id: Some("provider-1".to_string()), provider_id: Some("provider-1".to_string()),
status: "completed".to_string(), status: "completed".to_string(),
billing_status: "pending".to_string(), billing_status: "pending".to_string(),
@@ -264,6 +342,7 @@ mod tests {
request_id: "req-2".to_string(), request_id: "req-2".to_string(),
user_id: Some("user-1".to_string()), user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()), api_key_id: Some("key-1".to_string()),
api_key_is_standalone: false,
provider_id: Some("provider-1".to_string()), provider_id: Some("provider-1".to_string()),
status: "completed".to_string(), status: "completed".to_string(),
billing_status: "settled".to_string(), billing_status: "settled".to_string(),

View File

@@ -247,11 +247,31 @@ impl SettlementWriteRepository for SqlxSettlementRepository {
}; };
if final_billing_status == "settled" { if final_billing_status == "settled" {
let wallet_row = if let Some(api_key_id) = input let api_key_id = input
.api_key_id .api_key_id
.as_deref() .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( sqlx::query(
r#" r#"
SELECT SELECT
@@ -275,7 +295,8 @@ LIMIT 1
let wallet_row = if wallet_row.is_some() { let wallet_row = if wallet_row.is_some() {
wallet_row wallet_row
} else if let Some(user_id) = } else if !api_key_is_standalone {
if let Some(user_id) =
input.user_id.as_deref().filter(|value| !value.is_empty()) input.user_id.as_deref().filter(|value| !value.is_empty())
{ {
sqlx::query( sqlx::query(
@@ -297,6 +318,9 @@ LIMIT 1
.map_postgres_err()? .map_postgres_err()?
} else { } else {
None None
}
} else {
None
}; };
if let Some(wallet_row) = wallet_row { if let Some(wallet_row) = wallet_row {
@@ -417,4 +441,11 @@ mod tests {
let source = include_str!("sql.rs"); let source = include_str!("sql.rs");
assert!(!source.contains("UPDATE \"usage\"\nSET\n wallet_id = $2")); 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 {"));
}
} }

View File

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

View File

@@ -69,6 +69,7 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_non_empty_string(source, target, "trace_id"); copy_non_empty_string(source, target, "trace_id");
copy_bool(source, target, "client_requested_stream"); copy_bool(source, target, "client_requested_stream");
copy_bool(source, target, "upstream_is_stream"); copy_bool(source, target, "upstream_is_stream");
copy_bool(source, target, "api_key_is_standalone");
copy_number(source, target, "provider_request_body_base64_bytes"); copy_number(source, target, "provider_request_body_base64_bytes");
copy_number(source, target, "provider_response_body_base64_bytes"); copy_number(source, target, "provider_response_body_base64_bytes");
copy_number(source, target, "client_response_body_base64_bytes"); copy_number(source, target, "client_response_body_base64_bytes");
@@ -97,6 +98,7 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
remove_non_empty_string(&mut source, target, "trace_id"); remove_non_empty_string(&mut source, target, "trace_id");
remove_bool(&mut source, target, "client_requested_stream"); remove_bool(&mut source, target, "client_requested_stream");
remove_bool(&mut source, target, "upstream_is_stream"); remove_bool(&mut source, target, "upstream_is_stream");
remove_bool(&mut source, target, "api_key_is_standalone");
remove_number(&mut source, target, "provider_request_body_base64_bytes"); remove_number(&mut source, target, "provider_request_body_base64_bytes");
remove_number(&mut source, target, "provider_response_body_base64_bytes"); remove_number(&mut source, target, "provider_response_body_base64_bytes");
remove_number(&mut source, target, "client_response_body_base64_bytes"); remove_number(&mut source, target, "client_response_body_base64_bytes");
@@ -384,6 +386,7 @@ mod tests {
"trace_id": "trace-1", "trace_id": "trace-1",
"client_requested_stream": false, "client_requested_stream": false,
"upstream_is_stream": true, "upstream_is_stream": true,
"api_key_is_standalone": true,
"provider_request_body_base64_bytes": 512, "provider_request_body_base64_bytes": 512,
"provider_response_body_base64_bytes": 1024, "provider_response_body_base64_bytes": 1024,
"client_response_body_base64_bytes": 2048, "client_response_body_base64_bytes": 2048,
@@ -414,6 +417,7 @@ mod tests {
"trace_id": "trace-1", "trace_id": "trace-1",
"client_requested_stream": false, "client_requested_stream": false,
"upstream_is_stream": true, "upstream_is_stream": true,
"api_key_is_standalone": true,
"provider_request_body_base64_bytes": 512, "provider_request_body_base64_bytes": 512,
"provider_response_body_base64_bytes": 1024, "provider_response_body_base64_bytes": 1024,
"client_response_body_base64_bytes": 2048, "client_response_body_base64_bytes": 2048,
@@ -472,6 +476,7 @@ mod tests {
"candidate_index": 0, "candidate_index": 0,
"client_requested_stream": false, "client_requested_stream": false,
"upstream_is_stream": true, "upstream_is_stream": true,
"api_key_is_standalone": true,
"provider_id": "provider-1", "provider_id": "provider-1",
"model_id": "model-1", "model_id": "model-1",
"global_model_id": "global-model-1", "global_model_id": "global-model-1",
@@ -489,6 +494,7 @@ mod tests {
json!({ json!({
"client_requested_stream": false, "client_requested_stream": false,
"upstream_is_stream": true, "upstream_is_stream": true,
"api_key_is_standalone": true,
"model_id": "model-1", "model_id": "model-1",
"global_model_id": "global-model-1", "global_model_id": "global-model-1",
"global_model_name": "gpt-5", "global_model_name": "gpt-5",

View File

@@ -31,6 +31,7 @@ pub async fn settle_usage_if_needed(
request_id: usage.request_id.clone(), request_id: usage.request_id.clone(),
user_id: usage.user_id.clone(), user_id: usage.user_id.clone(),
api_key_id: usage.api_key_id.clone(), api_key_id: usage.api_key_id.clone(),
api_key_is_standalone: usage_api_key_is_standalone(usage),
provider_id: usage.provider_id.clone(), provider_id: usage.provider_id.clone(),
status: usage.status.clone(), status: usage.status.clone(),
billing_status: usage.billing_status.clone(), billing_status: usage.billing_status.clone(),
@@ -42,6 +43,15 @@ pub async fn settle_usage_if_needed(
Ok(()) Ok(())
} }
fn usage_api_key_is_standalone(usage: &StoredRequestUsageAudit) -> bool {
usage
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("api_key_is_standalone"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
fn finite_cost(value: f64) -> Result<f64, DataLayerError> { fn finite_cost(value: f64) -> Result<f64, DataLayerError> {
if value.is_finite() { if value.is_finite() {
Ok(value) Ok(value)
@@ -60,6 +70,7 @@ mod tests {
use aether_data_contracts::repository::settlement::UsageSettlementInput; use aether_data_contracts::repository::settlement::UsageSettlementInput;
use aether_data_contracts::repository::usage::StoredRequestUsageAudit; use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::json;
#[derive(Default)] #[derive(Default)]
struct TestSettlementWriter { struct TestSettlementWriter {
@@ -150,6 +161,25 @@ mod tests {
assert_eq!(inputs[0].finalized_at_unix_secs, Some(200)); assert_eq!(inputs[0].finalized_at_unix_secs, Some(200));
assert_eq!(inputs[0].total_cost_usd, 1.25); assert_eq!(inputs[0].total_cost_usd, 1.25);
assert_eq!(inputs[0].actual_total_cost_usd, 0.75); assert_eq!(inputs[0].actual_total_cost_usd, 0.75);
assert!(!inputs[0].api_key_is_standalone);
}
#[tokio::test]
async fn propagates_standalone_key_flag_from_usage_metadata() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.request_metadata = Some(json!({ "api_key_is_standalone": true }));
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!(inputs[0].api_key_is_standalone);
} }
#[tokio::test] #[tokio::test]

View File

@@ -1543,6 +1543,12 @@ fn build_runtime_request_metadata_seed_from_parts(
Value::Bool(upstream_is_stream), Value::Bool(upstream_is_stream),
); );
} }
if let Some(api_key_is_standalone) = context_bool(context, "api_key_is_standalone") {
metadata.insert(
"api_key_is_standalone".to_string(),
Value::Bool(api_key_is_standalone),
);
}
let provider_source_bytes = provider_request_body_base64.and_then(decoded_base64_len_hint); let provider_source_bytes = provider_request_body_base64.and_then(decoded_base64_len_hint);
append_runtime_body_capture_metadata( append_runtime_body_capture_metadata(
&mut metadata, &mut metadata,
@@ -2498,6 +2504,47 @@ mod tests {
assert_eq!(record.request_metadata, None); assert_eq!(record.request_metadata, None);
} }
#[test]
fn pending_usage_record_preserves_standalone_key_metadata() {
let plan = ExecutionPlan {
request_id: "req-pending-standalone-1".to_string(),
candidate_id: None,
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/responses".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({"model": "gpt-5.4"})),
stream: false,
client_api_format: "openai:responses".to_string(),
provider_api_format: "openai:responses".to_string(),
model_name: Some("gpt-5.4".to_string()),
proxy: None,
tls_profile: None,
timeouts: None,
};
let record = build_pending_usage_record(
&plan,
Some(&json!({
"api_key_is_standalone": true
})),
1_700_000_000,
)
.expect("pending usage should build");
assert_eq!(
record.request_metadata,
Some(json!({
"api_key_is_standalone": true
}))
);
}
#[test] #[test]
fn streaming_usage_records_stay_lightweight_by_default() { fn streaming_usage_records_stay_lightweight_by_default() {
let plan = ExecutionPlan { let plan = ExecutionPlan {