Merge remote-tracking branch 'origin/pr-484'

This commit is contained in:
fawney19
2026-05-18 18:01:57 +08:00
17 changed files with 925 additions and 160 deletions

View File

@@ -11,6 +11,9 @@ use crate::tests::{
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
EXECUTION_PATH_HEADER,
};
use aether_data::repository::billing::InMemoryBillingReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
use base64::Engine as _;
#[tokio::test]
@@ -1154,6 +1157,118 @@ async fn gateway_handles_internal_gateway_decision_sync_locally_with_supplied_au
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_internal_decision_sync_revalidates_supplied_auth_context_wallet() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("proxied"))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
None,
unrestricted_models_snapshot("api-key-empty-wallet", "user-empty-wallet"),
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_models_candidate_row("provider-1", "openai", "openai:chat", "gpt-5", 10),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-1", "openai", 10)],
vec![sample_endpoint(
"endpoint-provider-1",
"provider-1",
"openai:chat",
"https://api.openai.example",
)],
vec![sample_key(
"key-provider-1",
"provider-1",
"openai:chat",
"sk-upstream-openai",
)],
));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(Vec::new()));
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
StoredWalletSnapshot::new(
"wallet-empty".to_string(),
Some("user-empty-wallet".to_string()),
None,
0.0,
0.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_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_usage_billing_and_wallet_for_tests(
auth_repository,
candidate_repository,
provider_catalog_repository,
request_candidate_repository,
usage_repository,
billing_repository,
wallet_repository,
DEVELOPMENT_ENCRYPTION_KEY,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/gateway/decision-sync"))
.json(&json!({
"trace_id": "trace-internal-decision-sync-empty-wallet",
"method": "POST",
"path": "/v1/chat/completions",
"headers": {
"content-type": "application/json",
},
"body_json": {
"model": "gpt-5",
"messages": [],
},
"auth_context": {
"user_id": "user-empty-wallet",
"api_key_id": "api-key-empty-wallet",
"access_allowed": true,
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["action"], "fallback_plan");
assert_eq!(payload["auth_context"], serde_json::Value::Null);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_internal_gateway_decision_sync_fallback_with_resolved_auth_context() {
let upstream_hits = Arc::new(Mutex::new(0usize));

View File

@@ -4133,13 +4133,7 @@ async fn gateway_handles_wallet_balance_locally_without_proxying_upstream() {
#[tokio::test]
async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
let auth_now = Utc::now();
let usage_now = chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
auth_now
.date_naive()
.and_hms_opt(12, 0, 0)
.expect("midday should be valid"),
chrono::Utc,
);
let usage_now = auth_now - chrono::Duration::minutes(30);
let user = sample_auth_user(auth_now);
let access_token = build_test_auth_token(
"access",
@@ -4211,6 +4205,72 @@ async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_wallet_flow_today_entry_uses_live_settled_usage() {
let auth_now = Utc::now();
let user = sample_auth_user(auth_now);
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(user.id)),
("role".to_string(), json!(user.role)),
(
"created_at".to_string(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
("session_id".to_string(), json!("session-wallet-flow-today")),
]),
auth_now + chrono::Duration::hours(1),
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_user_usage_audit(
"usage-wallet-flow-today",
"req-wallet-flow-today",
"user-auth-1",
"gpt-4.1",
"OpenAI",
"completed",
auth_now - chrono::Duration::minutes(5),
),
]));
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_usage_state(
user,
sample_auth_wallet("user-auth-1", auth_now),
[sample_auth_session(
"user-auth-1",
"session-wallet-flow-today",
"device-wallet-flow-today",
"refresh-token-placeholder",
auth_now,
)],
usage_repository,
)
.await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/wallet/flow?limit=20&offset=0"))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-wallet-flow-today")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["today_entry"]["is_today"], true);
assert_eq!(payload["today_entry"]["timezone"], "Asia/Shanghai");
assert_eq!(payload["today_entry"]["total_requests"], 1);
assert_eq!(payload["today_entry"]["input_tokens"], 120);
assert_eq!(payload["today_entry"]["cache_read_tokens"], 15);
assert_eq!(payload["today_entry"]["total_cost"], 1.25);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_service_unavailable_for_wallet_today_cost_without_usage_reader() {
let now = Utc::now();