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

feat(routing): make client disconnect behavior strategy-scoped

Browse Source
This commit is contained in:
elky
2026-09-08 23:11:37 +08:00
parent 99f6499b2b
commit e58570d79d
22 changed files with 1156 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-billing/src/event_enrichment.rs
+124 -62
View File
@@ -1,8 +1,8 @@
use aether_data_contracts::repository::billing::StoredBillingModelContext;
use aether_data_contracts::repository::usage::{
extract_provider_cache_ttl_minutes_from_metadata, resolve_provider_cache_ttl_minutes,
resolve_provider_service_tier_from_request_capture, USAGE_AVAILABLE_METADATA_KEY,
USAGE_PRICING_AVAILABLE_METADATA_KEY,
resolve_provider_service_tier_from_request_capture, CANCELLED_REQUEST_FEE_METADATA_KEY,
USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY,
};
use aether_data_contracts::DataLayerError;
use aether_usage_runtime::{UsageEvent, UsageEventType};
@@ -40,6 +40,18 @@ pub async fn enrich_usage_event_with_billing(
data: &dyn BillingModelContextLookup,
event: &mut UsageEvent,
) -> Result<(), DataLayerError> {
if matches!(event.event_type, UsageEventType::Cancelled) {
event.data.total_cost_usd = Some(0.0);
event.data.actual_total_cost_usd = Some(0.0);
if let Some(metadata) = event
.data
.request_metadata
.as_mut()
.and_then(Value::as_object_mut)
{
metadata.remove(CANCELLED_REQUEST_FEE_METADATA_KEY);
}
}
// Session transports such as Codex Live expose lifecycle telemetry but no
// authoritative token/cost object. Do not run request-based pricing with
// zero default tokens: that would turn "unknown" into a fabricated charge.
@@ -65,7 +77,10 @@ pub async fn enrich_usage_event_with_billing(
clear_usage_costs(event);
return Ok(());
}
if !matches!(event.event_type, UsageEventType::Completed) {
if !matches!(
event.event_type,
UsageEventType::Completed | UsageEventType::Cancelled
) {
event.data.total_cost_usd = Some(0.0);
event.data.actual_total_cost_usd = Some(0.0);
return Ok(());
@@ -189,7 +204,10 @@ fn calculate_billing_computation(
} else {
usage_event_image_count(&event.data).unwrap_or(0)
};
let request_count = if failed {
let cancelled = matches!(event.event_type, UsageEventType::Cancelled);
let request_count = if cancelled {
1
} else if failed {
0
} else if is_image_usage && image_count > 0 {
image_count
@@ -197,7 +215,7 @@ fn calculate_billing_computation(
1
};
let processing_tiers = usage_event_processing_tiers(&event.data);
let input = BillingUsageInput {
let mut input = BillingUsageInput {
task_type: if is_image_usage {
"image".to_string()
} else {
@@ -237,6 +255,16 @@ fn calculate_billing_computation(
.or(pricing.provider_api_key_cache_ttl_minutes),
};
if cancelled {
input.input_tokens = 0;
input.output_tokens = 0;
input.cache_creation_tokens = 0;
input.cache_creation_ephemeral_5m_tokens = 0;
input.cache_creation_ephemeral_1h_tokens = 0;
input.cache_read_tokens = 0;
input.image_count = 0;
}
BillingService::new()
.calculate(pricing, &input)
.map_err(|err| {
@@ -356,9 +384,32 @@ fn apply_billing_computation(
pricing: &BillingModelPricingSnapshot,
computation: BillingComputation,
) -> Result<(), DataLayerError> {
let cancelled = matches!(event.event_type, UsageEventType::Cancelled);
if cancelled
&& !computation
.pricing_resolution
.price_per_request
.is_some_and(|price| price > 0.0)
{
return Ok(());
}
event.data.total_cost_usd = Some(computation.cost_result.cost);
event.data.actual_total_cost_usd = Some(computation.actual_total_cost);
merge_billing_snapshot_metadata(&mut event.data.request_metadata, pricing, &computation)
merge_billing_snapshot_metadata(&mut event.data.request_metadata, pricing, &computation)?;
if cancelled {
if let Some(metadata) = event
.data
.request_metadata
.as_mut()
.and_then(Value::as_object_mut)
{
metadata.insert(
CANCELLED_REQUEST_FEE_METADATA_KEY.to_string(),
Value::Bool(true),
);
}
}
Ok(())
}
fn map_pricing_context(context: StoredBillingModelContext) -> BillingModelPricingSnapshot {
@@ -1272,8 +1323,11 @@ mod tests {
}
#[tokio::test]
async fn cancelled_usage_event_remains_unbilled() {
let lookup = TestLookup {
async fn cancelled_usage_bills_only_configured_request_fee() {
for (request_type, request_price) in
[("chat", None), ("chat", Some(0.02)), ("image", Some(0.02))]
{
let lookup = TestLookup {
name_context: Some(
StoredBillingModelContext::new(
"provider-1".to_string(),
@@ -1284,7 +1338,7 @@ mod tests {
"global-model-1".to_string(),
"gpt-5".to_string(),
None,
Some(0.02),
request_price,
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0,"cache_creation_price_per_1m":3.75,"cache_read_price_per_1m":0.30}]})),
Some("model-1".to_string()),
Some("gpt-5-upstream".to_string()),
@@ -1296,61 +1350,69 @@ mod tests {
),
model_id_context: None,
};
let mut event = UsageEvent::new(
UsageEventType::Cancelled,
"req-billing-cancelled-1",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
provider_id: Some("provider-1".to_string()),
provider_api_key_id: Some("key-1".to_string()),
request_type: Some("chat".to_string()),
api_format: Some("openai:responses".to_string()),
endpoint_api_format: Some("openai:responses".to_string()),
input_tokens: Some(1_000),
output_tokens: Some(500),
cache_read_input_tokens: Some(100),
status_code: Some(499),
..UsageEventData::default()
},
);
let mut event = UsageEvent::new(
UsageEventType::Cancelled,
"req-billing-cancelled-1",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
provider_id: Some("provider-1".to_string()),
provider_api_key_id: Some("key-1".to_string()),
request_type: Some(request_type.to_string()),
api_format: Some("openai:responses".to_string()),
endpoint_api_format: Some("openai:responses".to_string()),
input_tokens: Some(1_000),
output_tokens: Some(500),
cache_read_input_tokens: Some(100),
status_code: Some(499),
request_metadata: Some(
json!({"cancelled_request_fee": true, "image_count": 3}),
),
..UsageEventData::default()
},
);
enrich_usage_event_with_billing(&lookup, &mut event)
.await
.expect("billing should succeed");
enrich_usage_event_with_billing(&lookup, &mut event)
.await
.expect("billing should succeed");
assert_eq!(event.data.total_cost_usd, Some(0.0));
assert_eq!(event.data.actual_total_cost_usd, Some(0.0));
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_snapshot"))
.and_then(|value| value.get("status"))
.and_then(Value::as_str),
None
);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_dimensions"))
.and_then(|value| value.get("input_tokens"))
.and_then(Value::as_i64),
None
);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_dimensions"))
.and_then(|value| value.get("cache_read_tokens"))
.and_then(Value::as_i64),
None
);
let expected_cost = request_price.unwrap_or(0.0);
assert_eq!(event.data.total_cost_usd, Some(expected_cost));
assert_eq!(event.data.actual_total_cost_usd, Some(expected_cost * 0.5));
assert_eq!(event.data.input_tokens, Some(1_000));
assert_eq!(event.data.output_tokens, Some(500));
let metadata = event.data.request_metadata.as_ref().unwrap();
assert_eq!(
aether_data_contracts::repository::usage::cancelled_request_fee_is_billable(Some(
metadata
)),
request_price.is_some()
);
if request_price.is_some() {
assert_eq!(
metadata.pointer("/billing_snapshot/cost_breakdown/request_cost"),
Some(&json!(expected_cost))
);
assert_eq!(
metadata.pointer("/billing_dimensions/input_tokens"),
Some(&json!(0))
);
assert_eq!(
metadata.pointer("/billing_dimensions/output_tokens"),
Some(&json!(0))
);
assert_eq!(
metadata.pointer("/billing_dimensions/cache_read_tokens"),
Some(&json!(0))
);
assert_eq!(
metadata.pointer("/billing_dimensions/request_count"),
Some(&json!(1))
);
} else {
assert!(metadata.get("billing_snapshot").is_none());
}
}
}
#[tokio::test]
crates/aether-data/contracts/src/repository/usage/metadata_policy.rs
+9
View File
@@ -20,6 +20,14 @@ use super::{
const UPSTREAM_IS_STREAM_KEY: &str = "upstream_is_stream";
const PLAN_USAGE_RESERVATION_TOKEN_KEY: &str = "plan_usage_reservation_token";
const BODY_SIZE_BASIS: &str = "serialized gateway request bodies after normalization";
pub const CANCELLED_REQUEST_FEE_METADATA_KEY: &str = "cancelled_request_fee";
pub fn cancelled_request_fee_is_billable(metadata: Option<&Value>) -> bool {
metadata
.and_then(|metadata| metadata.get(CANCELLED_REQUEST_FEE_METADATA_KEY))
.and_then(Value::as_bool)
.unwrap_or(false)
}
/// Projects request metadata onto the persistence contract. Unknown fields and malformed values
/// are discarded instead of being recursively copied into an audit row.
@@ -48,6 +56,7 @@ pub fn sanitize_usage_request_metadata_object(source: &Map<String, Value>) -> Op
PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY,
"transport_error",
"is_free_tier",
CANCELLED_REQUEST_FEE_METADATA_KEY,
USAGE_AVAILABLE_METADATA_KEY,
USAGE_PRICING_AVAILABLE_METADATA_KEY,
] {
crates/aether-routing-core/src/model.rs
+35
View File
@@ -39,6 +39,8 @@ pub struct RoutingExecutionPolicy {
pub enable_cf_heartbeat: bool,
#[serde(default, skip_serializing_if = "is_false")]
pub cyber_continue_failover: bool,
#[serde(default, skip_serializing_if = "is_false")]
pub cancel_on_client_disconnect: bool,
}
impl<'de> Deserialize<'de> for RoutingExecutionPolicy {
@@ -56,6 +58,8 @@ impl<'de> Deserialize<'de> for RoutingExecutionPolicy {
enable_standard_text_sync_heartbeat: bool,
#[serde(default)]
cyber_continue_failover: bool,
#[serde(default)]
cancel_on_client_disconnect: bool,
}
let value = LegacyCompatibleExecutionPolicy::deserialize(deserializer)?;
@@ -64,6 +68,7 @@ impl<'de> Deserialize<'de> for RoutingExecutionPolicy {
|| value.enable_openai_image_sync_heartbeat
|| value.enable_standard_text_sync_heartbeat,
cyber_continue_failover: value.cyber_continue_failover,
cancel_on_client_disconnect: value.cancel_on_client_disconnect,
})
}
}
@@ -107,6 +112,36 @@ fn is_false(value: &bool) -> bool {
!*value
}
#[cfg(test)]
mod execution_policy_tests {
use super::*;
#[test]
fn cancellation_defaults_off_and_round_trips_with_legacy_heartbeat() {
let default: RoutingDefaultPolicy = serde_json::from_str("{}").unwrap();
assert!(!default.execution_policy.cancel_on_client_disconnect);
let policy: RoutingDefaultPolicy = serde_json::from_value(serde_json::json!({
"cancel_on_client_disconnect": true,
"enable_standard_text_sync_heartbeat": true
}))
.unwrap();
assert!(policy.execution_policy.cancel_on_client_disconnect);
assert!(policy.execution_policy.enable_cf_heartbeat);
let encoded = serde_json::to_value(&policy).unwrap();
assert_eq!(encoded["cancel_on_client_disconnect"], true);
assert_eq!(
serde_json::from_value::<RoutingDefaultPolicy>(encoded).unwrap(),
policy
);
assert!(
serde_json::from_value::<RoutingDefaultPolicy>(serde_json::json!({
"cancel_on_client_disconnect": "true"
}))
.is_err()
);
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingModelPolicy {
pub model: String,
crates/aether-testing/integration/tests/responses_websocket_e2e.rs
+137 -3
View File
@@ -25,6 +25,9 @@ use aether_data_contracts::repository::global_models::{
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::routing_profiles::{
RoutingGroupLookupKey, UpdateRoutingGroupRecord,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageAuditListQuery};
use aether_gateway::{build_router_with_state, AppState, GatewayDataConfig, UsageRuntimeConfig};
use aether_testkit::{ManagedPostgresServer, SpawnedServer};
@@ -491,8 +494,7 @@ async fn disabling_the_downstream_key_is_enforced_on_the_next_turn_of_the_same_s
Ok(())
}
/// A client that walks away before the provider produced anything must settle
/// as a void row: nothing was produced, so nothing is billed.
/// Immediate cancellation voids token billing before the provider completes.
///
/// This is the path with no protocol event to announce it: the relay loop owns
/// the turn, and losing the client is an exit the upstream never reports.
@@ -506,7 +508,14 @@ async fn disabling_the_downstream_key_is_enforced_on_the_next_turn_of_the_same_s
/// `a_closed_client_socket_before_any_terminal_still_voids_the_bill`.
#[tokio::test]
async fn client_disconnect_before_any_provider_output_settles_a_void_row() -> Result<(), BoxError> {
let harness = Harness::start(UpstreamBehavior::StallAfterCreated).await?;
let harness = Harness::start_configured(
UpstreamBehavior::StallAfterCreated,
ProviderFixture::SingleOpenAiKey,
PiiRedaction::Disabled,
true,
None,
)
.await?;
let mut client = harness.connect().await?;
client
@@ -542,6 +551,73 @@ async fn client_disconnect_before_any_provider_output_settles_a_void_row() -> Re
Ok(())
}
#[tokio::test]
async fn client_disconnect_defaults_to_completing_and_billing_the_turn() -> Result<(), BoxError> {
let harness = Harness::start(UpstreamBehavior::CompleteAfterRelease).await?;
let mut client = harness.connect().await?;
client
.send(response_create(json!({"input": "finish without client"})))
.await?;
receive_event(&mut client, "response.created").await?;
drop(client);
tokio::time::sleep(Duration::from_millis(50)).await;
harness.upstream.release_completion.notify_one();
let audits = harness
.usage_audits_where(1, "completed disconnected turn", |audit| {
audit.status == "completed" && audit.billing_status == "settled"
})
.await?;
assert_eq!(audits.len(), 1);
assert_eq!(audits[0].input_tokens, INPUT_TOKENS);
assert_eq!(audits[0].output_tokens, OUTPUT_TOKENS);
assert_eq!(audits[0].status_code, Some(200));
Ok(())
}
#[tokio::test]
async fn client_disconnect_still_settles_the_per_request_fee_when_aborted() -> Result<(), BoxError>
{
let harness = Harness::start_configured(
UpstreamBehavior::StallAfterCreated,
ProviderFixture::SingleOpenAiKey,
PiiRedaction::Disabled,
true,
Some(0.02),
)
.await?;
let mut client = harness.connect().await?;
client
.send(response_create(json!({"input": "cancel with request fee"})))
.await?;
receive_event(&mut client, "response.created").await?;
drop(client);
let audits = harness
.usage_audits_where(1, "cancelled request fee settlement", |audit| {
audit.status == "cancelled" && audit.billing_status == "settled"
})
.await?;
assert_eq!(audits.len(), 1);
assert_eq!(audits[0].status_code, Some(499));
assert_eq!(audits[0].total_tokens, 0);
assert_eq!(audits[0].total_cost_usd, 0.02);
assert_eq!(audits[0].actual_total_cost_usd, 0.02);
let backends = DataBackends::from_config(DataLayerConfig::from_database(
harness.database.config.clone(),
))?;
let detail = backends
.read()
.usage()
.ok_or("usage reader unavailable")?
.find_by_request_id(&audits[0].request_id)
.await?
.ok_or("usage detail unavailable")?;
assert_eq!(
detail.request_metadata.as_ref().unwrap()["cancelled_request_fee"],
true
);
Ok(())
}
/// An upstream that dies mid-turn must surface an error and still settle.
#[tokio::test]
async fn upstream_drop_mid_turn_reports_an_error_and_settles_the_usage_row() -> Result<(), BoxError>
@@ -850,6 +926,16 @@ impl Harness {
behavior: UpstreamBehavior,
fixture: ProviderFixture,
redaction: PiiRedaction,
) -> Result<Self, BoxError> {
Self::start_configured(behavior, fixture, redaction, false, None).await
}
async fn start_configured(
behavior: UpstreamBehavior,
fixture: ProviderFixture,
redaction: PiiRedaction,
cancel_on_client_disconnect: bool,
request_price: Option<f64>,
) -> Result<Self, BoxError> {
let upstream = Arc::new(MockUpstreamState::new(behavior));
let upstream_server =
@@ -864,6 +950,16 @@ impl Harness {
)
.await?;
if let Some(price) = request_price {
let pool = sqlx::PgPool::connect(&database.config.url).await?;
sqlx::query("UPDATE models SET price_per_request = $1 WHERE id = $2")
.bind(price)
.bind(PROVIDER_MODEL_ID)
.execute(&pool)
.await?;
pool.close().await;
}
let data_config = GatewayDataConfig::from_database_config(database.config.clone())
.with_encryption_key(DEVELOPMENT_ENCRYPTION_KEY);
let state = AppState::new()?
@@ -877,6 +973,32 @@ impl Harness {
..UsageRuntimeConfig::default()
})?;
state.ensure_system_default_routing_group().await?;
if cancel_on_client_disconnect {
let backends =
DataBackends::from_config(DataLayerConfig::from_database(database.config.clone()))?;
let mut group = backends
.read()
.routing_groups()
.ok_or("routing reader unavailable")?
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
.await?
.ok_or("default routing group unavailable")?;
group.config_json["default_policy"]["cancel_on_client_disconnect"] = json!(true);
backends
.write()
.routing_groups()
.ok_or("routing writer unavailable")?
.update_routing_group(
&group.id,
UpdateRoutingGroupRecord {
config_json: Some(group.config_json),
version: Some(group.version + 1),
updated_at: group.updated_at + 1,
..Default::default()
},
)
.await?;
}
let gateway_server = SpawnedServer::start(build_router_with_state(state)).await?;
let websocket_url = format!(
"{}/v1/responses",
@@ -1093,6 +1215,7 @@ where
enum UpstreamBehavior {
/// Announce, stream one delta, and complete — the ordinary turn.
CompleteEveryTurn,
CompleteAfterRelease,
/// Announce the response and then go quiet, leaving the turn in flight.
StallAfterCreated,
/// Announce the response and then hang up mid-turn.
@@ -1118,6 +1241,7 @@ struct MockUpstreamState {
events: Mutex<Vec<Value>>,
authorization_headers: Mutex<Vec<Option<String>>>,
handshakes: Mutex<Vec<ObservedUpstreamHandshake>>,
release_completion: tokio::sync::Notify,
}
#[derive(Debug, Clone)]
@@ -1134,6 +1258,7 @@ impl MockUpstreamState {
events: Mutex::new(Vec::new()),
authorization_headers: Mutex::new(Vec::new()),
handshakes: Mutex::new(Vec::new()),
release_completion: tokio::sync::Notify::new(),
}
}
@@ -1215,6 +1340,15 @@ async fn run_mock_upstream(
break;
}
}
UpstreamBehavior::CompleteAfterRelease => {
if send_mock_created(&mut socket, &response_id).await.is_err() {
break;
}
state.release_completion.notified().await;
if send_mock_turn(&mut socket, &response_id).await.is_err() {
break;
}
}
UpstreamBehavior::StallAfterCreated => {
if send_mock_created(&mut socket, &response_id).await.is_err() {
break;
crates/aether-usage/runtime/src/record.rs
+32
View File
@@ -221,6 +221,13 @@ fn lifecycle_status_and_billing(
}
UsageEventType::Completed => ("completed", "pending"),
UsageEventType::Failed => ("failed", "void"),
UsageEventType::Cancelled
if aether_data_contracts::repository::usage::cancelled_request_fee_is_billable(
request_metadata,
) =>
{
("cancelled", "pending")
}
UsageEventType::Cancelled => ("cancelled", "void"),
}
}
@@ -491,6 +498,31 @@ mod tests {
assert_eq!(record.first_byte_time_ms, Some(50));
}
#[test]
fn cancelled_request_fee_keeps_cancelled_status_and_pending_billing() {
let event = UsageEvent::new(
UsageEventType::Cancelled,
"req-cancelled-fee",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
total_cost_usd: Some(0.02),
actual_total_cost_usd: Some(0.01),
request_metadata: Some(serde_json::json!({"cancelled_request_fee": true})),
..Default::default()
},
);
let record = build_upsert_usage_record_from_event(&event).unwrap();
assert_eq!(record.status, "cancelled");
assert_eq!(record.billing_status, "pending");
assert_eq!(record.total_cost_usd, Some(0.02));
assert_eq!(record.actual_total_cost_usd, Some(0.01));
assert_eq!(
record.request_metadata.unwrap()["cancelled_request_fee"],
true
);
}
#[test]
fn completed_unmetered_session_audit_is_void_without_fabricated_usage() {
let record = build_upsert_usage_record_from_event(&UsageEvent {
crates/aether-usage/runtime/src/settlement.rs
+71 -3
View File
@@ -5,8 +5,10 @@ use aether_data_contracts::repository::settlement::{
ReconcileUsagePolicyCostInput, StoredUsagePolicyCostReservation, StoredUsageSettlement,
UsagePolicyCostReservationState, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use aether_data_contracts::repository::usage::PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY;
use aether_data_contracts::repository::usage::{
cancelled_request_fee_is_billable, StoredRequestUsageAudit,
};
use aether_data_contracts::{DataLayerError, DataLayerError::InvalidInput};
use async_trait::async_trait;
@@ -39,6 +41,11 @@ pub async fn reconcile_usage_policy_cost_for_event(
}
let terminal_state = match event.event_type {
UsageEventType::Completed => UsagePolicyCostReservationState::Finalized,
UsageEventType::Cancelled
if cancelled_request_fee_is_billable(event.data.request_metadata.as_ref()) =>
{
UsagePolicyCostReservationState::Finalized
}
UsageEventType::Failed | UsageEventType::Cancelled => {
UsagePolicyCostReservationState::Released
}
@@ -107,7 +114,10 @@ pub async fn settle_usage_if_needed(
usage.user_id.as_deref().and_then(non_empty_trimmed),
usage_policy_reservation_token(usage),
) {
let (terminal_state, actual_cost_units) = if usage.status == "completed" {
let (terminal_state, actual_cost_units) = if usage.status == "completed"
|| (usage.status == "cancelled"
&& cancelled_request_fee_is_billable(usage.request_metadata.as_ref()))
{
(
UsagePolicyCostReservationState::Finalized,
nonnegative_usd_to_usage_policy_cost_units(
@@ -136,7 +146,10 @@ pub async fn settle_usage_if_needed(
}
}
if usage.status == "cancelled" || usage.billing_status != "pending" {
if usage.billing_status != "pending"
|| (usage.status == "cancelled"
&& !cancelled_request_fee_is_billable(usage.request_metadata.as_ref()))
{
return Ok(());
}
let input = UsageSettlementInput {
@@ -429,6 +442,61 @@ mod tests {
);
}
#[tokio::test]
async fn cancelled_request_fee_settles_wallet_and_finalizes_cost_reservation() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.status = "cancelled".to_string();
usage.status_code = Some(499);
usage.request_metadata.as_mut().unwrap()["cancelled_request_fee"] = json!(true);
settle_usage_if_needed(&writer, &usage).await.unwrap();
let inputs = writer.inputs.lock().unwrap();
assert_eq!(inputs.len(), 1);
assert_eq!(inputs[0].status, "cancelled");
assert_eq!(inputs[0].actual_total_cost_usd, usage.actual_total_cost_usd);
let reconciliations = writer.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 1);
assert_eq!(
reconciliations[0].terminal_state,
UsagePolicyCostReservationState::Finalized
);
assert_eq!(reconciliations[0].actual_cost_units, 75_000_000);
}
#[tokio::test]
async fn cancelled_request_fee_event_finalizes_cost_reservation() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let event = UsageEvent::new(
UsageEventType::Cancelled,
"req-cancelled-fee",
UsageEventData {
user_id: Some("user-1".to_string()),
actual_total_cost_usd: Some(0.01),
request_metadata: Some(json!({
"cancelled_request_fee": true,
"plan_usage_reservation_token": "server-token"
})),
..Default::default()
},
);
reconcile_usage_policy_cost_for_event(&writer, &event)
.await
.unwrap();
let reconciliations = writer.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 1);
assert_eq!(
reconciliations[0].terminal_state,
UsagePolicyCostReservationState::Finalized
);
assert_eq!(reconciliations[0].actual_cost_units, 1_000_000);
}
#[tokio::test]
async fn releases_failed_usage_before_void_settlement() {
let writer = TestSettlementWriter {