mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-09 12:40:20 +08:00
983 lines
34 KiB
Rust
983 lines
34 KiB
Rust
use super::{SqliteUsageReadRepository, SqliteUsageWriteRepository};
|
|
use crate::run_migrations;
|
|
use aether_data_contracts::repository::usage::{
|
|
UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
|
|
UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageProviderPerformanceQuery,
|
|
UsageReadRepository, UsageTimeSeriesGranularity, UsageWriteRepository,
|
|
};
|
|
|
|
#[test]
|
|
fn sqlite_usage_upsert_guards_candidate_identity_metadata_and_routing_from_late_lifecycle() {
|
|
for field in [
|
|
"provider_name",
|
|
"model",
|
|
"target_model",
|
|
"provider_id",
|
|
"provider_endpoint_id",
|
|
"provider_api_key_id",
|
|
"request_type",
|
|
"api_format",
|
|
"api_family",
|
|
"endpoint_kind",
|
|
"endpoint_api_format",
|
|
"provider_api_family",
|
|
"provider_endpoint_kind",
|
|
"has_format_conversion",
|
|
"is_stream",
|
|
"upstream_is_stream",
|
|
"request_metadata",
|
|
"candidate_id",
|
|
"candidate_index",
|
|
"key_name",
|
|
"planner_kind",
|
|
"route_family",
|
|
"route_kind",
|
|
"execution_path",
|
|
"local_execution_runtime_miss_reason",
|
|
] {
|
|
let assignment = format!("{field} = CASE WHEN (");
|
|
assert!(
|
|
super::UPSERT_USAGE_SQL.contains(&assignment),
|
|
"missing lifecycle guard for {field}"
|
|
);
|
|
assert!(
|
|
super::UPSERT_USAGE_SQL.contains(&format!("THEN \"usage\".{field}")),
|
|
"late lifecycle must preserve {field}"
|
|
);
|
|
}
|
|
assert!(super::UPSERT_USAGE_SQL
|
|
.contains("OR (\"usage\".status = 'streaming' AND excluded.status = 'pending')"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_provider_performance_can_skip_timeline() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
SqliteUsageWriteRepository::new(pool.clone())
|
|
.upsert(sample_usage(
|
|
"provider-performance",
|
|
"completed",
|
|
"pending",
|
|
1_000,
|
|
))
|
|
.await
|
|
.expect("usage should upsert");
|
|
let reader = SqliteUsageReadRepository::new(pool);
|
|
let mut query = UsageProviderPerformanceQuery {
|
|
created_from_unix_secs: 0,
|
|
created_until_unix_secs: 2_000,
|
|
granularity: UsageTimeSeriesGranularity::Hour,
|
|
tz_offset_minutes: 0,
|
|
limit: 1,
|
|
provider_id: None,
|
|
model: None,
|
|
api_format: None,
|
|
endpoint_kind: None,
|
|
is_stream: None,
|
|
has_format_conversion: None,
|
|
slow_threshold_ms: 10_000,
|
|
include_timeline: true,
|
|
};
|
|
|
|
let with_timeline = reader
|
|
.summarize_usage_provider_performance(&query)
|
|
.await
|
|
.expect("provider performance should load");
|
|
assert_eq!(with_timeline.summary.request_count, 1);
|
|
assert_eq!(with_timeline.providers.len(), 1);
|
|
assert_eq!(with_timeline.timeline.len(), 1);
|
|
|
|
query.include_timeline = false;
|
|
let without_timeline = reader
|
|
.summarize_usage_provider_performance(&query)
|
|
.await
|
|
.expect("provider performance without timeline should load");
|
|
assert_eq!(without_timeline.summary, with_timeline.summary);
|
|
assert_eq!(without_timeline.providers, with_timeline.providers);
|
|
assert!(without_timeline.timeline.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_upserts_and_rebuilds_stats() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool.clone());
|
|
let record = repository
|
|
.upsert(sample_usage("request-1", "completed", "pending", 1_000))
|
|
.await
|
|
.expect("usage should upsert");
|
|
|
|
assert_eq!(record.request_id, "request-1");
|
|
assert_eq!(record.api_key_id.as_deref(), Some("api-key-1"));
|
|
assert_eq!(record.total_tokens, 7);
|
|
assert_eq!(record.cache_read_input_tokens, 2);
|
|
assert_eq!(
|
|
record.request_metadata.as_ref().unwrap()["trace_id"],
|
|
"trace-1"
|
|
);
|
|
assert_eq!(
|
|
record.request_metadata.as_ref().unwrap()["upstream_is_stream"],
|
|
true
|
|
);
|
|
let upstream_is_stream: Option<i64> =
|
|
sqlx::query_scalar("SELECT upstream_is_stream FROM \"usage\" WHERE request_id = ?")
|
|
.bind("request-1")
|
|
.fetch_one(&pool)
|
|
.await
|
|
.expect("usage stream mode should load");
|
|
assert_eq!(upstream_is_stream, Some(1));
|
|
|
|
let loaded = repository
|
|
.find_by_request_id("request-1")
|
|
.await
|
|
.expect("usage should load")
|
|
.expect("usage should exist");
|
|
assert_eq!(
|
|
loaded.provider_api_key_id.as_deref(),
|
|
Some("provider-key-1")
|
|
);
|
|
|
|
let stats = sqlx::query_as::<_, (i64, i64, f64, Option<i64>)>(
|
|
"SELECT total_requests, total_tokens, total_cost_usd, last_used_at FROM api_keys WHERE id = 'api-key-1'",
|
|
)
|
|
.fetch_one(&pool)
|
|
.await
|
|
.expect("api key stats should load");
|
|
assert_eq!(stats, (1, 7, 0.5, Some(1_000)));
|
|
|
|
let provider_stats = sqlx::query_as::<_, (i64, i64, i64, i64, f64, i64, Option<i64>)>(
|
|
"SELECT request_count, success_count, error_count, total_tokens, total_cost_usd, total_response_time_ms, last_used_at FROM provider_api_keys WHERE id = 'provider-key-1'",
|
|
)
|
|
.fetch_one(&pool)
|
|
.await
|
|
.expect("provider key stats should load");
|
|
assert_eq!(provider_stats, (1, 1, 0, 7, 0.5, 42, Some(1_000)));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_does_not_regress_void_usage() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool.clone());
|
|
repository
|
|
.upsert(sample_usage("request-1", "failed", "void", 1_000))
|
|
.await
|
|
.expect("void usage should upsert");
|
|
let existing = repository
|
|
.upsert(sample_usage("request-1", "pending", "pending", 1_001))
|
|
.await
|
|
.expect("stale usage should be ignored");
|
|
|
|
assert_eq!(existing.status, "failed");
|
|
assert_eq!(existing.billing_status, "void");
|
|
assert_eq!(existing.updated_at_unix_secs, 1_000);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_does_not_reopen_void_failure_from_late_streaming() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool);
|
|
for (request_id, status_code) in [
|
|
("request-late-active", None),
|
|
("request-late-response-start", Some(200)),
|
|
] {
|
|
let mut failed = sample_usage(request_id, "failed", "void", 1_000);
|
|
failed.status_code = Some(503);
|
|
repository
|
|
.upsert(failed)
|
|
.await
|
|
.expect("failed usage should upsert");
|
|
|
|
let mut late_streaming = sample_usage(request_id, "streaming", "pending", 1_001);
|
|
late_streaming.status_code = status_code;
|
|
late_streaming.finalized_at_unix_secs = None;
|
|
let current = repository
|
|
.upsert(late_streaming)
|
|
.await
|
|
.expect("late streaming usage should be ignored");
|
|
|
|
assert_eq!(current.status, "failed");
|
|
assert_eq!(current.billing_status, "void");
|
|
assert_eq!(current.status_code, Some(503));
|
|
assert_eq!(current.finalized_at_unix_secs, Some(1_000));
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late_streaming() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool.clone());
|
|
let mut terminal = sample_usage("request-1", "completed", "pending", 1_000);
|
|
terminal.request_metadata = Some(serde_json::json!({"trace_id": "terminal-trace"}));
|
|
repository
|
|
.upsert(terminal)
|
|
.await
|
|
.expect("terminal usage should upsert");
|
|
|
|
let mut late_streaming = sample_usage("request-1", "streaming", "pending", 1_001);
|
|
late_streaming.input_tokens = Some(0);
|
|
late_streaming.output_tokens = Some(0);
|
|
late_streaming.total_tokens = Some(0);
|
|
late_streaming.cache_read_input_tokens = Some(0);
|
|
late_streaming.cache_read_cost_usd = Some(0.0);
|
|
late_streaming.total_cost_usd = Some(0.0);
|
|
late_streaming.actual_total_cost_usd = Some(0.0);
|
|
late_streaming.response_time_ms = Some(9_999);
|
|
late_streaming.first_byte_time_ms = Some(9_999);
|
|
late_streaming.finalized_at_unix_secs = None;
|
|
late_streaming.provider_name = "Late Provider".to_string();
|
|
late_streaming.model = "late-model".to_string();
|
|
late_streaming.target_model = Some("late-target".to_string());
|
|
late_streaming.request_type = Some("late-request".to_string());
|
|
late_streaming.api_format = Some("late:api".to_string());
|
|
late_streaming.api_family = Some("late-family".to_string());
|
|
late_streaming.endpoint_kind = Some("late-endpoint".to_string());
|
|
late_streaming.endpoint_api_format = Some("late:endpoint".to_string());
|
|
late_streaming.provider_api_family = Some("late-provider-family".to_string());
|
|
late_streaming.provider_endpoint_kind = Some("late-provider-endpoint".to_string());
|
|
late_streaming.has_format_conversion = Some(false);
|
|
late_streaming.is_stream = Some(true);
|
|
late_streaming.candidate_id = Some("late-candidate".to_string());
|
|
late_streaming.candidate_index = Some(99);
|
|
late_streaming.key_name = Some("late-key".to_string());
|
|
late_streaming.planner_kind = Some("late-planner".to_string());
|
|
late_streaming.route_family = Some("late-route-family".to_string());
|
|
late_streaming.route_kind = Some("late-route-kind".to_string());
|
|
late_streaming.execution_path = Some("late-path".to_string());
|
|
late_streaming.local_execution_runtime_miss_reason = Some("late-miss".to_string());
|
|
late_streaming.request_metadata = Some(serde_json::json!({
|
|
"provider_service_tier": "priority",
|
|
"upstream_is_stream": true
|
|
}));
|
|
|
|
let current = repository
|
|
.upsert(late_streaming)
|
|
.await
|
|
.expect("late streaming usage should not regress terminal usage");
|
|
|
|
assert_eq!(current.status, "completed");
|
|
assert_eq!(current.billing_status, "pending");
|
|
assert_eq!(current.total_tokens, 7);
|
|
assert_eq!(current.cache_read_input_tokens, 2);
|
|
assert_eq!(current.total_cost_usd, 0.5);
|
|
assert_eq!(current.actual_total_cost_usd, 0.4);
|
|
assert_eq!(current.response_time_ms, Some(42));
|
|
assert_eq!(current.first_byte_time_ms, Some(12));
|
|
assert_eq!(current.finalized_at_unix_secs, Some(1_000));
|
|
assert_eq!(current.updated_at_unix_secs, 1_000);
|
|
assert_eq!(current.provider_name, "Provider One");
|
|
assert_eq!(current.model, "model-1");
|
|
assert_eq!(current.target_model.as_deref(), Some("target-model"));
|
|
assert_eq!(current.request_type.as_deref(), Some("chat"));
|
|
assert_eq!(current.api_format.as_deref(), Some("openai"));
|
|
assert!(current.has_format_conversion);
|
|
assert!(!current.is_stream);
|
|
assert_eq!(current.candidate_id.as_deref(), Some("candidate-1"));
|
|
assert_eq!(current.candidate_index, Some(1));
|
|
assert_eq!(current.key_name.as_deref(), Some("key-one"));
|
|
assert_eq!(current.planner_kind.as_deref(), Some("default"));
|
|
assert_eq!(current.route_family.as_deref(), Some("chat"));
|
|
assert_eq!(current.route_kind.as_deref(), Some("completion"));
|
|
assert_eq!(current.execution_path.as_deref(), Some("remote"));
|
|
assert_eq!(current.provider_service_tier(), None);
|
|
assert_eq!(
|
|
current
|
|
.request_metadata
|
|
.as_ref()
|
|
.and_then(|value| value.get("trace_id"))
|
|
.and_then(serde_json::Value::as_str),
|
|
Some("terminal-trace")
|
|
);
|
|
|
|
let listed = SqliteUsageReadRepository::new(pool)
|
|
.list_usage_audits(&UsageAuditListQuery {
|
|
limit: Some(10),
|
|
newest_first: true,
|
|
..UsageAuditListQuery::default()
|
|
})
|
|
.await
|
|
.expect("usage list should load")
|
|
.into_iter()
|
|
.find(|item| item.request_id == "request-1")
|
|
.expect("terminal usage should be listed");
|
|
assert_eq!(listed.provider_name, "Provider One");
|
|
assert_eq!(listed.model, "model-1");
|
|
assert_eq!(listed.candidate_id.as_deref(), Some("candidate-1"));
|
|
assert_eq!(listed.provider_service_tier(), None);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_allows_authoritative_completed_recovery() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool);
|
|
repository
|
|
.upsert(sample_usage("request-recovery", "failed", "void", 1_000))
|
|
.await
|
|
.expect("void failure should upsert");
|
|
|
|
let mut recovery = sample_usage("request-recovery", "completed", "pending", 1_001);
|
|
recovery.provider_name = "Recovered Provider".to_string();
|
|
recovery.model = "recovered-model".to_string();
|
|
recovery.target_model = Some("recovered-target".to_string());
|
|
recovery.api_format = Some("recovered:api".to_string());
|
|
recovery.candidate_id = Some("recovered-candidate".to_string());
|
|
recovery.request_metadata = Some(serde_json::json!({"provider_service_tier": "priority"}));
|
|
let recovered = repository
|
|
.upsert(recovery)
|
|
.await
|
|
.expect("completed recovery should upsert");
|
|
|
|
assert_eq!(recovered.status, "completed");
|
|
assert_eq!(recovered.billing_status, "pending");
|
|
assert_eq!(recovered.provider_name, "Recovered Provider");
|
|
assert_eq!(recovered.model, "recovered-model");
|
|
assert_eq!(recovered.target_model.as_deref(), Some("recovered-target"));
|
|
assert_eq!(recovered.api_format.as_deref(), Some("recovered:api"));
|
|
assert_eq!(
|
|
recovered.candidate_id.as_deref(),
|
|
Some("recovered-candidate")
|
|
);
|
|
assert_eq!(
|
|
recovered.provider_service_tier().as_deref(),
|
|
Some("priority")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_preserves_streaming_response_start_from_late_active() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool);
|
|
repository
|
|
.upsert(sample_usage(
|
|
"request-late-active",
|
|
"streaming",
|
|
"pending",
|
|
1_000,
|
|
))
|
|
.await
|
|
.expect("response-start usage should upsert");
|
|
|
|
let mut late_active = sample_usage("request-late-active", "streaming", "pending", 1_001);
|
|
late_active.status_code = None;
|
|
late_active.response_time_ms = None;
|
|
late_active.first_byte_time_ms = None;
|
|
|
|
let current = repository
|
|
.upsert(late_active)
|
|
.await
|
|
.expect("late active usage should not clear response-start fields");
|
|
|
|
assert_eq!(current.status, "streaming");
|
|
assert_eq!(current.status_code, Some(200));
|
|
assert_eq!(current.response_time_ms, Some(42));
|
|
assert_eq!(current.first_byte_time_ms, Some(12));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_keeps_streaming_capture_from_late_pending() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool);
|
|
let mut streaming = sample_usage("request-streaming-capture", "streaming", "pending", 1_000);
|
|
streaming.request_metadata = Some(serde_json::json!({"trace_id": "streaming-final"}));
|
|
repository
|
|
.upsert(streaming)
|
|
.await
|
|
.expect("streaming usage should upsert");
|
|
|
|
let mut late_pending = sample_usage("request-streaming-capture", "pending", "pending", 1_001);
|
|
late_pending.provider_name = "Late Provider".to_string();
|
|
late_pending.model = "late-model".to_string();
|
|
late_pending.candidate_id = Some("late-candidate".to_string());
|
|
late_pending.request_metadata = Some(serde_json::json!({"provider_service_tier": "priority"}));
|
|
let current = repository
|
|
.upsert(late_pending)
|
|
.await
|
|
.expect("late pending usage should not regress streaming capture");
|
|
|
|
assert_eq!(current.status, "streaming");
|
|
assert_eq!(current.provider_name, "Provider One");
|
|
assert_eq!(current.model, "model-1");
|
|
assert_eq!(current.candidate_id.as_deref(), Some("candidate-1"));
|
|
assert_eq!(current.provider_service_tier(), None);
|
|
assert_eq!(
|
|
current
|
|
.request_metadata
|
|
.as_ref()
|
|
.and_then(|value| value.get("trace_id"))
|
|
.and_then(serde_json::Value::as_str),
|
|
Some("streaming-final")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_cleans_stale_pending_requests() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool.clone());
|
|
repository
|
|
.upsert(sample_usage("request-recovered", "streaming", "pending", 1))
|
|
.await
|
|
.expect("streaming usage should upsert");
|
|
repository
|
|
.upsert(sample_usage("request-failed", "pending", "pending", 1))
|
|
.await
|
|
.expect("pending usage should upsert");
|
|
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO request_candidates (
|
|
id,
|
|
request_id,
|
|
candidate_index,
|
|
retry_index,
|
|
status,
|
|
is_cached,
|
|
created_at
|
|
) VALUES
|
|
('candidate-recovered', 'request-recovered', 0, 0, 'streaming', 0, 1),
|
|
('candidate-failed', 'request-failed', 0, 0, 'pending', 0, 1)
|
|
"#,
|
|
)
|
|
.execute(&pool)
|
|
.await
|
|
.expect("request candidates should seed");
|
|
|
|
let summary = repository
|
|
.cleanup_stale_pending_requests(2, 10, 5, 1)
|
|
.await
|
|
.expect("cleanup should run");
|
|
assert_eq!(summary.recovered, 1);
|
|
assert_eq!(summary.failed, 1);
|
|
|
|
let recovered = repository
|
|
.find_by_request_id("request-recovered")
|
|
.await
|
|
.expect("recovered usage should load")
|
|
.expect("recovered usage should exist");
|
|
assert_eq!(recovered.status, "completed");
|
|
assert_eq!(recovered.status_code, Some(200));
|
|
|
|
let failed = repository
|
|
.find_by_request_id("request-failed")
|
|
.await
|
|
.expect("failed usage should load")
|
|
.expect("failed usage should exist");
|
|
assert_eq!(failed.status, "failed");
|
|
assert_eq!(failed.status_code, Some(504));
|
|
assert_eq!(failed.billing_status, "void");
|
|
assert_eq!(failed.total_cost_usd, 0.0);
|
|
assert_eq!(failed.finalized_at_unix_secs, Some(10));
|
|
|
|
let candidate_statuses = sqlx::query_as::<_, (String, String, Option<i64>)>(
|
|
r#"
|
|
SELECT request_id, status, finished_at
|
|
FROM request_candidates
|
|
ORDER BY request_id
|
|
"#,
|
|
)
|
|
.fetch_all(&pool)
|
|
.await
|
|
.expect("candidate statuses should load");
|
|
assert_eq!(
|
|
candidate_statuses,
|
|
vec![
|
|
(
|
|
"request-failed".to_string(),
|
|
"failed".to_string(),
|
|
Some(10_000)
|
|
),
|
|
(
|
|
"request-recovered".to_string(),
|
|
"success".to_string(),
|
|
Some(10_000)
|
|
),
|
|
]
|
|
);
|
|
|
|
let snapshot = sqlx::query_as::<_, (String, Option<i64>)>(
|
|
"SELECT billing_status, finalized_at FROM usage_settlement_snapshots WHERE request_id = 'request-failed'",
|
|
)
|
|
.fetch_one(&pool)
|
|
.await
|
|
.expect("void settlement snapshot should load");
|
|
assert_eq!(snapshot, ("void".to_string(), Some(10)));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_write_repository_cleanup_uses_failed_candidate_status_when_present() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let repository = SqliteUsageWriteRepository::new(pool.clone());
|
|
repository
|
|
.upsert(sample_usage(
|
|
"request-upstream-reset",
|
|
"pending",
|
|
"pending",
|
|
1,
|
|
))
|
|
.await
|
|
.expect("pending usage should upsert");
|
|
repository
|
|
.upsert(sample_usage("request-stuck", "pending", "pending", 1))
|
|
.await
|
|
.expect("pending usage should upsert");
|
|
|
|
// request-upstream-reset has a failed candidate carrying a concrete 502 status
|
|
// and a connection-reset message — cleanup should use them instead of 504.
|
|
// request-stuck has only a still-pending candidate, so cleanup should fall back to 504.
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO request_candidates (
|
|
id,
|
|
request_id,
|
|
candidate_index,
|
|
retry_index,
|
|
status,
|
|
status_code,
|
|
error_message,
|
|
is_cached,
|
|
created_at,
|
|
started_at,
|
|
finished_at
|
|
) VALUES
|
|
('candidate-reset', 'request-upstream-reset', 0, 0, 'failed', 502, 'upstream connection reset by peer', 0, 1, 2, 3),
|
|
('candidate-stuck', 'request-stuck', 0, 0, 'pending', NULL, NULL, 0, 1, NULL, NULL)
|
|
"#,
|
|
)
|
|
.execute(&pool)
|
|
.await
|
|
.expect("request candidates should seed");
|
|
|
|
let summary = repository
|
|
.cleanup_stale_pending_requests(2, 10, 5, 5)
|
|
.await
|
|
.expect("cleanup should run");
|
|
assert_eq!(summary.recovered, 0);
|
|
assert_eq!(summary.failed, 2);
|
|
|
|
let reset = repository
|
|
.find_by_request_id("request-upstream-reset")
|
|
.await
|
|
.expect("upstream-reset usage should load")
|
|
.expect("upstream-reset usage should exist");
|
|
assert_eq!(reset.status, "failed");
|
|
assert_eq!(reset.status_code, Some(502));
|
|
assert_eq!(
|
|
reset.error_message.as_deref(),
|
|
Some("upstream connection reset by peer")
|
|
);
|
|
|
|
let stuck = repository
|
|
.find_by_request_id("request-stuck")
|
|
.await
|
|
.expect("stuck usage should load")
|
|
.expect("stuck usage should exist");
|
|
assert_eq!(stuck.status, "failed");
|
|
assert_eq!(stuck.status_code, Some(504));
|
|
assert!(stuck
|
|
.error_message
|
|
.as_deref()
|
|
.is_some_and(|message| message.contains("超过 5 分钟未完成")));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_read_repository_reads_usage_contract_views() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
seed_stats_targets(&pool).await;
|
|
|
|
let writer = SqliteUsageWriteRepository::new(pool.clone());
|
|
writer
|
|
.upsert(sample_usage("request-1", "completed", "settled", 1_000))
|
|
.await
|
|
.expect("usage should upsert");
|
|
writer
|
|
.upsert(sample_usage("request-2", "failed", "void", 1_010))
|
|
.await
|
|
.expect("usage should upsert");
|
|
|
|
let reader = SqliteUsageReadRepository::new(pool);
|
|
let loaded = reader
|
|
.find_by_request_id("request-1")
|
|
.await
|
|
.expect("usage should load")
|
|
.expect("usage should exist");
|
|
assert_eq!(loaded.total_tokens, 7);
|
|
assert_eq!(loaded.billing_status, "settled");
|
|
|
|
let listed = reader
|
|
.list_usage_audits(&UsageAuditListQuery {
|
|
provider_name: Some("Provider One".to_string()),
|
|
newest_first: true,
|
|
..UsageAuditListQuery::default()
|
|
})
|
|
.await
|
|
.expect("usage list should load");
|
|
assert_eq!(listed.len(), 2);
|
|
assert_eq!(listed[0].request_id, "request-2");
|
|
|
|
let summary = reader
|
|
.summarize_dashboard_usage(&UsageDashboardSummaryQuery {
|
|
created_from_unix_secs: 999,
|
|
created_until_unix_secs: 1_020,
|
|
user_id: Some("user-1".to_string()),
|
|
})
|
|
.await
|
|
.expect("dashboard summary should load");
|
|
assert_eq!(summary.total_requests, 2);
|
|
assert_eq!(summary.error_requests, 1);
|
|
assert_eq!(summary.total_tokens, 10);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_daily_heatmap_reads_imported_daily_aggregates() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO stats_daily (
|
|
id, "date", total_requests, success_requests, error_requests,
|
|
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
|
|
total_cost, actual_total_cost, is_complete, created_at, updated_at
|
|
) VALUES (
|
|
'daily-1', 86400, 9, 8, 1, 10, 20, 3, 4, 1.25, 1.0, 1, 1, 1
|
|
);
|
|
INSERT INTO stats_user_daily (
|
|
id, user_id, username, "date", total_requests, success_requests, error_requests,
|
|
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
|
|
total_cost, created_at, updated_at
|
|
) VALUES (
|
|
'user-daily-1', 'user-1', 'user one', 86400, 5, 5, 0, 7, 8, 2, 1, 0.75, 1, 1
|
|
);
|
|
"#,
|
|
)
|
|
.execute(&pool)
|
|
.await
|
|
.expect("daily aggregates should seed");
|
|
|
|
let reader = SqliteUsageReadRepository::new(pool);
|
|
let admin = reader
|
|
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
|
|
created_from_unix_secs: 0,
|
|
user_id: None,
|
|
admin_mode: true,
|
|
})
|
|
.await
|
|
.expect("admin heatmap should load");
|
|
assert_eq!(admin.len(), 1);
|
|
assert_eq!(admin[0].date, "1970-01-02");
|
|
assert_eq!(admin[0].requests, 9);
|
|
assert_eq!(admin[0].total_tokens, 37);
|
|
assert_eq!(admin[0].actual_total_cost_usd, 1.0);
|
|
|
|
let user = reader
|
|
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
|
|
created_from_unix_secs: 0,
|
|
user_id: Some("user-1".to_string()),
|
|
admin_mode: false,
|
|
})
|
|
.await
|
|
.expect("user heatmap should load");
|
|
assert_eq!(user.len(), 1);
|
|
assert_eq!(user[0].date, "1970-01-02");
|
|
assert_eq!(user[0].requests, 5);
|
|
assert_eq!(user[0].total_tokens, 18);
|
|
assert_eq!(user[0].actual_total_cost_usd, 0.75);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_usage_totals_by_user_ids_reads_imported_user_daily_aggregates() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO stats_user_daily (
|
|
id, user_id, username, "date", total_requests, success_requests, error_requests,
|
|
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
|
|
total_cost, created_at, updated_at
|
|
) VALUES (
|
|
'user-daily-1', 'user-1', 'user one', 86400, 5, 5, 0, 7, 8, 2, 1, 0.75, 1, 1
|
|
);
|
|
INSERT INTO "usage" (
|
|
request_id, id, user_id, api_key_id, provider_name, model, total_tokens,
|
|
status, billing_status, created_at_unix_ms, updated_at_unix_secs
|
|
) VALUES
|
|
('raw-before-cutoff', 'usage-1', 'user-1', 'api-key-1', 'Provider One', 'model-1', 99,
|
|
'completed', 'settled', 90000, 90000),
|
|
('raw-after-cutoff', 'usage-2', 'user-1', 'api-key-1', 'Provider One', 'model-1', 7,
|
|
'completed', 'settled', 172800, 172800);
|
|
"#,
|
|
)
|
|
.execute(&pool)
|
|
.await
|
|
.expect("usage totals fixtures should seed");
|
|
|
|
let reader = SqliteUsageReadRepository::new(pool);
|
|
let totals = reader
|
|
.summarize_usage_totals_by_user_ids(&["user-1".to_string()])
|
|
.await
|
|
.expect("user totals should load");
|
|
|
|
assert_eq!(totals.len(), 1);
|
|
assert_eq!(totals[0].user_id, "user-1");
|
|
assert_eq!(totals[0].request_count, 6);
|
|
assert_eq!(totals[0].total_tokens, 25);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sqlite_dashboard_daily_stats_reads_imported_daily_aggregates() {
|
|
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.expect("sqlite pool should connect");
|
|
run_migrations(&pool)
|
|
.await
|
|
.expect("sqlite migrations should run");
|
|
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO stats_daily (
|
|
id, "date", total_requests, success_requests, error_requests,
|
|
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
|
|
total_cost, actual_total_cost, is_complete, created_at, updated_at
|
|
) VALUES (
|
|
'daily-1', 86400, 9, 8, 1, 10, 20, 3, 4, 1.25, 1.0, 1, 1, 1
|
|
);
|
|
"#,
|
|
)
|
|
.execute(&pool)
|
|
.await
|
|
.expect("daily aggregates should seed");
|
|
|
|
let reader = SqliteUsageReadRepository::new(pool);
|
|
let summary = reader
|
|
.summarize_dashboard_usage(&UsageDashboardSummaryQuery {
|
|
created_from_unix_secs: 0,
|
|
created_until_unix_secs: 172800,
|
|
user_id: None,
|
|
})
|
|
.await
|
|
.expect("dashboard summary should load");
|
|
assert_eq!(summary.total_requests, 9);
|
|
assert_eq!(summary.total_tokens, 37);
|
|
|
|
let rows = reader
|
|
.list_dashboard_daily_breakdown(&UsageDashboardDailyBreakdownQuery {
|
|
created_from_unix_secs: 0,
|
|
created_until_unix_secs: 172800,
|
|
tz_offset_minutes: 480,
|
|
user_id: None,
|
|
})
|
|
.await
|
|
.expect("dashboard daily breakdown should load");
|
|
assert_eq!(rows.len(), 1);
|
|
assert_eq!(rows[0].date, "1970-01-02");
|
|
assert_eq!(rows[0].model, "aggregate");
|
|
assert_eq!(rows[0].requests, 9);
|
|
assert_eq!(rows[0].total_tokens, 37);
|
|
}
|
|
|
|
async fn seed_stats_targets(pool: &sqlx::SqlitePool) {
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO users (id, auth_source, created_at, updated_at)
|
|
VALUES ('user-1', 'local', 1, 1);
|
|
INSERT INTO api_keys (id, user_id, key_hash, created_at, updated_at)
|
|
VALUES ('api-key-1', 'user-1', 'hash-1', 1, 1);
|
|
INSERT INTO providers (id, name, provider_type, created_at, updated_at)
|
|
VALUES ('provider-1', 'Provider One', 'openai', 1, 1);
|
|
INSERT INTO provider_api_keys (id, provider_id, name, created_at, updated_at)
|
|
VALUES ('provider-key-1', 'provider-1', 'Provider Key One', 1, 1);
|
|
"#,
|
|
)
|
|
.execute(pool)
|
|
.await
|
|
.expect("stats targets should seed");
|
|
}
|
|
|
|
fn sample_usage(
|
|
request_id: &str,
|
|
status: &str,
|
|
billing_status: &str,
|
|
updated_at: u64,
|
|
) -> UpsertUsageRecord {
|
|
UpsertUsageRecord {
|
|
request_id: request_id.to_string(),
|
|
user_id: Some("user-1".to_string()),
|
|
api_key_id: Some("api-key-1".to_string()),
|
|
username: Some("legacy-user".to_string()),
|
|
api_key_name: Some("legacy-key".to_string()),
|
|
provider_name: "Provider One".to_string(),
|
|
model: "model-1".to_string(),
|
|
target_model: Some("target-model".to_string()),
|
|
provider_id: Some("provider-1".to_string()),
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
|
request_type: Some("chat".to_string()),
|
|
api_format: Some("openai".to_string()),
|
|
api_family: Some("chat".to_string()),
|
|
endpoint_kind: Some("chat".to_string()),
|
|
endpoint_api_format: Some("openai".to_string()),
|
|
provider_api_family: Some("chat".to_string()),
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
|
has_format_conversion: Some(true),
|
|
is_stream: Some(false),
|
|
input_tokens: Some(2),
|
|
output_tokens: Some(3),
|
|
total_tokens: None,
|
|
cache_creation_input_tokens: None,
|
|
cache_creation_ephemeral_5m_input_tokens: Some(0),
|
|
cache_creation_ephemeral_1h_input_tokens: Some(0),
|
|
cache_read_input_tokens: Some(2),
|
|
cache_creation_cost_usd: Some(0.0),
|
|
cache_read_cost_usd: Some(0.1),
|
|
output_price_per_1m: Some(2.0),
|
|
total_cost_usd: Some(0.5),
|
|
actual_total_cost_usd: Some(0.4),
|
|
status_code: Some(200),
|
|
error_message: None,
|
|
error_category: None,
|
|
response_time_ms: Some(42),
|
|
first_byte_time_ms: Some(12),
|
|
status: status.to_string(),
|
|
billing_status: billing_status.to_string(),
|
|
request_headers: None,
|
|
request_body: None,
|
|
request_body_ref: None,
|
|
request_body_state: None,
|
|
provider_request_headers: None,
|
|
provider_request_body: None,
|
|
provider_request_body_ref: None,
|
|
provider_request_body_state: None,
|
|
response_headers: None,
|
|
response_body: None,
|
|
response_body_ref: None,
|
|
response_body_state: None,
|
|
client_response_headers: None,
|
|
client_response_body: None,
|
|
client_response_body_ref: None,
|
|
client_response_body_state: None,
|
|
candidate_id: Some("candidate-1".to_string()),
|
|
candidate_index: Some(1),
|
|
key_name: Some("key-one".to_string()),
|
|
planner_kind: Some("default".to_string()),
|
|
route_family: Some("chat".to_string()),
|
|
route_kind: Some("completion".to_string()),
|
|
execution_path: Some("remote".to_string()),
|
|
local_execution_runtime_miss_reason: None,
|
|
request_metadata: Some(serde_json::json!({
|
|
"trace_id": "trace-1",
|
|
"upstream_is_stream": true,
|
|
})),
|
|
finalized_at_unix_secs: Some(updated_at),
|
|
created_at_unix_ms: Some(updated_at),
|
|
updated_at_unix_secs: updated_at,
|
|
}
|
|
}
|