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

refactor: 拆分 gateway 单体为独立 crate，新增 systemd 部署方案

Browse Source 
将 gateway 内部的 model-fetch、provider-transport、scheduler-core、
usage-runtime、video-tasks-core 模块提取为独立 crate；重构 gateway
内部模块结构（state/router/cache/data/query 等）；移除大量遗留模块
文件；新增 systemd 二进制部署骨架及相关文档；更新前端 usage 相关
API 和组件。
This commit is contained in:
fawney19
2026-04-05 20:23:16 +08:00
parent cbc811f6ce
commit 763ff03a7b
777 changed files with 42659 additions and 21469 deletions
Download Patch File Download Diff File Expand all files Collapse all files

241
apps/aether-gateway/src/tests/usage/direct.rs Normal file
View File

@@ -0,0 +1,241 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, hash_api_key, json,
sample_local_openai_auth_snapshot, sample_local_openai_candidate_row,
sample_local_openai_endpoint, sample_local_openai_key, sample_local_openai_provider,
start_server, Arc, Body, GatewayDataState, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
InMemoryRequestCandidateRepository, InMemoryUsageReadRepository, Json, Request, Response,
Router, StatusCode, UsageReadRepository, UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY,
TRACE_ID_HEADER,
};
#[tokio::test]
async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let upstream = Router::new().route(
"/api/internal/gateway/report-sync",
any(|_request: Request| async move { Json(json!({"ok": true})) }),
);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(|_request: Request| async move {
Json(json!({
"request_id": "req-usage-sync-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-usage-sync-123",
"usage": {
"input_tokens": 3,
"output_tokens": 5,
"total_tokens": 8
}
}
},
"telemetry": {
"elapsed_ms": 45
}
}))
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-usage-sync")),
sample_local_openai_auth_snapshot("api-key-usage-sync-123", "user-usage-sync-123"),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state =
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-usage-sync",
)
.header(TRACE_ID_HEADER, "req-usage-sync-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-sync-123")
.await
.expect("usage lookup should succeed");
if stored.is_some() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be recorded");
assert_eq!(stored.status, "completed");
assert_eq!(stored.billing_status, "pending");
assert_eq!(stored.total_tokens, 8);
assert_eq!(stored.response_time_ms, Some(45));
assert_eq!(stored.user_id.as_deref(), Some("user-usage-sync-123"));
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
async fn gateway_records_usage_for_execution_runtime_stream_when_runtime_enabled() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let upstream = Router::new()
.route(
"/api/internal/gateway/plan-stream",
any(|_request: Request| async move {
Json(json!({
"action": "execution_runtime_stream",
"plan_kind": "openai_chat_stream",
"plan": {
"request_id": "req-usage-stream-123",
"provider_name": "openai",
"provider_id": "provider-usage-stream-123",
"endpoint_id": "endpoint-usage-stream-123",
"key_id": "key-usage-stream-123",
"method": "POST",
"url": "https://api.openai.example/v1/chat/completions",
"headers": {
"authorization": "Bearer upstream-key",
"content-type": "application/json"
},
"body": {
"json_body": {
"model": "gpt-5",
"messages": [],
"stream": true
}
},
"stream": true,
"client_api_format": "openai:chat",
"provider_api_format": "openai:chat",
"model_name": "gpt-5"
},
"report_kind": "openai_chat_stream_success",
"report_context": {
"user_id": "user-usage-stream-123",
"api_key_id": "api-key-usage-stream-123",
"provider_name": "openai",
"provider_id": "provider-usage-stream-123",
"endpoint_id": "endpoint-usage-stream-123",
"key_id": "key-usage-stream-123",
"client_api_format": "openai:chat",
"provider_api_format": "openai:chat",
"model": "gpt-5",
"mapped_model": "gpt-5"
}
}))
}),
)
.route(
"/api/internal/gateway/report-stream",
any(|_request: Request| async move { Json(json!({"ok": true})) }),
);
let execution_runtime = Router::new().route(
"/v1/execute/stream",
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":51,\"ttfb_ms\":19}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from(frames))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_usage_data_repository_for_tests(usage_repository.clone())
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let _ = response.text().await.expect("stream body should read");
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-stream-123")
.await
.expect("usage lookup should succeed");
if stored.is_some() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be recorded");
assert_eq!(stored.status, "completed");
assert_eq!(stored.billing_status, "pending");
assert_eq!(stored.total_tokens, 6);
assert_eq!(stored.first_byte_time_ms, Some(19));
assert_eq!(stored.is_stream, true);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}

355
apps/aether-gateway/src/tests/usage/local.rs Normal file
View File

@@ -0,0 +1,355 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, hash_api_key, json,
sample_local_openai_auth_snapshot, sample_local_openai_candidate_row,
sample_local_openai_endpoint, sample_local_openai_key, sample_local_openai_provider,
start_server, Arc, Body, GatewayDataState, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
InMemoryRequestCandidateRepository, InMemoryUsageReadRepository, Json, Mutex, Request,
RequestCandidateReadRepository, RequestCandidateStatus, Response, Router, StatusCode,
UsageReadRepository, UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
};
#[tokio::test]
async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled(
) {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let report_hits = Arc::new(Mutex::new(0usize));
let report_hits_clone = Arc::clone(&report_hits);
let decision_hits = Arc::new(Mutex::new(0usize));
let decision_hits_clone = Arc::clone(&decision_hits);
let plan_hits = Arc::new(Mutex::new(0usize));
let plan_hits_clone = Arc::clone(&plan_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/decision-sync",
any(move |_request: Request| {
let decision_hits_inner = Arc::clone(&decision_hits_clone);
async move {
*decision_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"action": "proxy_public"}))
}
}),
)
.route(
"/api/internal/gateway/plan-sync",
any(move |_request: Request| {
let plan_hits_inner = Arc::clone(&plan_hits_clone);
async move {
*plan_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"action": "proxy_public"}))
}
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |_request: Request| {
let report_hits_inner = Arc::clone(&report_hits_clone);
async move {
*report_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"ok": true}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(|_request: Request| async move {
Json(json!({
"request_id": "trace-openai-chat-local-report-sync-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-local-report-sync-123",
"object": "chat.completion",
"model": "gpt-5-upstream",
"choices": [],
"usage": {
"prompt_tokens": 2,
"completion_tokens": 3,
"total_tokens": 5
}
}
},
"telemetry": {
"elapsed_ms": 25
}
}))
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-local-report-sync")),
sample_local_openai_auth_snapshot(
"api-key-openai-usage-local-1",
"user-openai-usage-local-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state =
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-local-report-sync",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-local-report-sync-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let body_json: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(body_json["model"], "gpt-5-upstream");
let mut stored_usage = None;
for _ in 0..50 {
stored_usage = usage_repository
.find_by_request_id("trace-openai-chat-local-report-sync-123")
.await
.expect("usage lookup should succeed");
if stored_usage.is_some() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored_usage = stored_usage.expect("usage should be recorded");
assert_eq!(stored_usage.status, "completed");
assert_eq!(stored_usage.total_tokens, 5);
assert_eq!(stored_usage.response_time_ms, Some(25));
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-local-report-sync-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(*report_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_when_usage_runtime_enabled(
) {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let report_hits = Arc::new(Mutex::new(0usize));
let report_hits_clone = Arc::clone(&report_hits);
let decision_hits = Arc::new(Mutex::new(0usize));
let decision_hits_clone = Arc::clone(&decision_hits);
let plan_hits = Arc::new(Mutex::new(0usize));
let plan_hits_clone = Arc::clone(&plan_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/decision-stream",
any(move |_request: Request| {
let decision_hits_inner = Arc::clone(&decision_hits_clone);
async move {
*decision_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"action": "proxy_public"}))
}
}),
)
.route(
"/api/internal/gateway/plan-stream",
any(move |_request: Request| {
let plan_hits_inner = Arc::clone(&plan_hits_clone);
async move {
*plan_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"action": "proxy_public"}))
}
}),
)
.route(
"/api/internal/gateway/report-stream",
any(move |_request: Request| {
let report_hits_inner = Arc::clone(&report_hits_clone);
async move {
*report_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"ok": true}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let execution_runtime = Router::new().route(
"/v1/execute/stream",
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-local-report-stream-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from(frames))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-local-report-stream")),
sample_local_openai_auth_snapshot(
"api-key-openai-usage-local-1",
"user-openai-usage-local-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state =
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-local-report-stream",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-local-report-stream-123")
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let body_text = response.text().await.expect("stream body should read");
assert_eq!(
body_text,
"data: {\"id\":\"chatcmpl-local-report-stream-123\",\"usage\":{\"input_tokens\":2,\"output_tokens\":4,\"total_tokens\":6}}\n\ndata: [DONE]\n\n"
);
let mut stored_usage = None;
for _ in 0..50 {
stored_usage = usage_repository
.find_by_request_id("trace-openai-chat-local-report-stream-123")
.await
.expect("usage lookup should succeed");
if stored_usage.is_some() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored_usage = stored_usage.expect("usage should be recorded");
assert_eq!(stored_usage.status, "completed");
assert_eq!(stored_usage.total_tokens, 6);
assert_eq!(stored_usage.first_byte_time_ms, Some(11));
assert!(stored_usage.is_stream);
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-local-report-stream-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(*report_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}

168
apps/aether-gateway/src/tests/usage/wallet.rs Normal file
View File

@@ -0,0 +1,168 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, hash_api_key, json,
sample_local_openai_auth_snapshot, sample_local_openai_candidate_row,
sample_local_openai_endpoint, sample_local_openai_key, sample_local_openai_provider,
start_server, Arc, GatewayDataState, InMemoryAuthApiKeySnapshotRepository,
InMemoryBillingReadRepository, InMemoryMinimalCandidateSelectionReadRepository,
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository,
InMemoryUsageReadRepository, InMemoryWalletRepository, Json, Request, Router, StatusCode,
StoredBillingModelContext, StoredWalletSnapshot, UsageReadRepository, UsageRuntimeConfig,
WalletLookupKey, WalletReadRepository, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
};
use aether_data::repository::settlement::InMemorySettlementRepository;
#[tokio::test]
async fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(vec![
StoredBillingModelContext::new(
"provider-openai-usage-local-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-openai-usage-local-1".to_string()),
Some(json!({"openai:chat": 1.0})),
Some(60),
"global-model-openai-usage-local-1".to_string(),
"gpt-5".to_string(),
None,
Some(0.02),
Some(
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]}),
),
Some("model-openai-usage-local-1".to_string()),
Some("gpt-5".to_string()),
None,
None,
None,
)
.expect("billing context should build"),
]));
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
StoredWalletSnapshot::new(
"wallet-usage-sync-123".to_string(),
Some("user-usage-sync-123".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"),
]));
let upstream = Router::new().route(
"/api/internal/gateway/report-sync",
any(|_request: Request| async move { Json(json!({"ok": true})) }),
);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(|_request: Request| async move {
Json(json!({
"request_id": "req-usage-wallet-sync-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-usage-wallet-sync-123",
"usage": {
"input_tokens": 1000,
"output_tokens": 500,
"total_tokens": 1500
}
}
},
"telemetry": {
"elapsed_ms": 45
}
}))
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-wallet-sync")),
sample_local_openai_auth_snapshot("api-key-usage-sync-123", "user-usage-sync-123"),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let data_state = GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_usage_billing_and_wallet_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
usage_repository.clone(),
billing_repository,
wallet_repository.clone(),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_settlement_writer_for_tests(Arc::new(
InMemorySettlementRepository::from_wallet_repository(wallet_repository.clone()),
));
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(data_state)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-wallet-sync",
)
.header(TRACE_ID_HEADER, "req-usage-wallet-sync-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-wallet-sync-123")
.await
.expect("usage lookup should succeed");
if stored.is_some() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be recorded");
assert_eq!(stored.status, "completed");
assert_eq!(stored.total_tokens, 1500);
let wallet = wallet_repository
.find(WalletLookupKey::UserId("user-usage-sync-123"))
.await
.expect("wallet lookup should succeed")
.expect("wallet should exist");
assert!(wallet.balance < 10.0 || wallet.gift_balance < 2.0);
assert!(wallet.total_consumed > 0.0);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}