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

Preserve usage data in system imports

Browse Source
This commit is contained in:
fawney19
2026-05-24 21:40:48 +08:00
parent 18d9004f22
commit e2d5fc9dfb
21 changed files with 2683 additions and 309 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/control/admin/system.rs
+15
View File
@@ -1043,6 +1043,7 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["group_names"],
json!(["Restricted GPT"])
);
assert_eq!(payload["users"][0]["id"], json!("user-1"));
assert_eq!(payload["users"][0]["wallet"]["balance"], json!(12.5));
assert_eq!(
payload["users"][0]["wallet"]["recharge_balance"],
@@ -1063,6 +1064,10 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["api_keys"][0]["is_standalone"],
json!(false)
);
assert_eq!(
payload["users"][0]["api_keys"][0]["api_key_id"],
json!("key-user-1")
);
assert_eq!(
payload["users"][0]["api_keys"][0]["total_tokens"],
json!(420)
@@ -1071,12 +1076,22 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["standalone_keys"][0]["key"],
json!("ak-standalone-live-1")
);
assert_eq!(
payload["standalone_keys"][0]["api_key_id"],
json!("key-standalone-1")
);
assert_eq!(payload["standalone_keys"][0]["total_tokens"], json!(84));
assert_eq!(
payload["standalone_keys"][0]["wallet"]["unlimited"],
json!(true)
);
assert_eq!(payload["standalone_keys"][0].get("is_standalone"), None,);
assert_eq!(payload["usage_aggregates"]["stats_daily"], json!([]));
assert_eq!(payload["usage_aggregates"]["stats_user_daily"], json!([]));
assert_eq!(
payload["usage_aggregates"]["stats_daily_api_key"],
json!([])
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
apps/aether-gateway/src/tests/control/admin/system_import.rs
+222 -42
View File
@@ -4,7 +4,9 @@ use aether_contracts::ExecutionPlan;
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
use aether_data::repository::auth_modules::{
AuthModuleReadRepository, InMemoryAuthModuleReadRepository, StoredOAuthProviderModuleConfig,
};
@@ -27,7 +29,7 @@ use axum::{extract::Request, Json, Router};
use http::StatusCode;
use serde_json::{json, Value};
use super::super::helpers::{sample_endpoint, sample_key, sample_provider};
use super::super::helpers::{hash_api_key, sample_endpoint, sample_key, sample_provider};
use super::super::{
build_router_with_state, build_state_with_execution_runtime_override, start_server, AppState,
};
@@ -1052,6 +1054,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
.expect("user api keys should load");
assert_eq!(user_api_keys.len(), 1);
assert_eq!(user_api_keys[0].name.as_deref(), Some("Alice CLI"));
assert_eq!(user_api_keys[0].total_requests, 12);
assert_eq!(user_api_keys[0].total_tokens, 3456);
assert_eq!(user_api_keys[0].total_cost_usd, 1.25);
assert_eq!(
@@ -1079,6 +1082,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
standalone_keys[0].name.as_deref(),
Some("Imported Standalone")
);
assert_eq!(standalone_keys[0].total_requests, 3);
assert_eq!(standalone_keys[0].total_tokens, 789);
assert_eq!(standalone_keys[0].total_cost_usd, 0.75);
assert_eq!(
@@ -1119,6 +1123,211 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
let _ = upstream_url;
}
#[tokio::test]
async fn gateway_overwrites_existing_admin_system_user_key_usage_totals() {
let user_key_hash = hash_api_key("sk-existing-user-key");
let standalone_key_hash = hash_api_key("sk-existing-standalone-key");
let existing_user = StoredUserAuthRecord::new(
"user-existing".to_string(),
Some("[email protected]".to_string()),
true,
"existing".to_string(),
Some("existing-hash".to_string()),
"user".to_string(),
"local".to_string(),
None,
None,
None,
true,
false,
Some(chrono::Utc::now()),
Some(chrono::Utc::now()),
)
.expect("existing user should build");
let user_key_snapshot = StoredAuthApiKeySnapshot::new(
"user-existing".to_string(),
"existing".to_string(),
Some("[email protected]".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-user-existing".to_string(),
Some("Existing User Key".to_string()),
true,
false,
false,
Some(10),
None,
None,
None,
None,
None,
)
.expect("user key snapshot should build");
let standalone_key_snapshot = StoredAuthApiKeySnapshot::new(
"admin-user-123".to_string(),
"admin".to_string(),
Some("[email protected]".to_string()),
"admin".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-standalone-existing".to_string(),
Some("Existing Standalone Key".to_string()),
true,
false,
true,
Some(20),
None,
None,
None,
None,
None,
)
.expect("standalone key snapshot should build");
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![
(Some(user_key_hash.clone()), user_key_snapshot),
(Some(standalone_key_hash.clone()), standalone_key_snapshot),
])
.with_export_records(vec![
StoredAuthApiKeyExportRecord::new(
"user-existing".to_string(),
"key-user-existing".to_string(),
user_key_hash.clone(),
None,
Some("Existing User Key".to_string()),
None,
None,
None,
Some(10),
None,
None,
true,
None,
false,
1,
2,
0.03,
false,
)
.expect("existing user key export should build"),
StoredAuthApiKeyExportRecord::new(
"admin-user-123".to_string(),
"key-standalone-existing".to_string(),
standalone_key_hash.clone(),
None,
Some("Existing Standalone Key".to_string()),
None,
None,
None,
Some(20),
None,
None,
true,
None,
false,
4,
5,
0.06,
true,
)
.expect("existing standalone key export should build"),
]),
);
let user_repository =
Arc::new(aether_data::repository::users::InMemoryUserReadRepository::default());
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_auth_api_key_repository_for_tests(Arc::clone(&auth_repository))
.with_user_reader(user_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_auth_users_for_tests([sample_import_admin_user("admin-user-123"), existing_user])
.with_auth_wallets_for_tests(Vec::<StoredWalletSnapshot>::new());
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/users/import"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"version": "1.4",
"merge_mode": "overwrite",
"users": [{
"id": "source-user-existing",
"email": "[email protected]",
"username": "existing",
"password_hash": "existing-hash",
"role": "user",
"is_active": true,
"api_keys": [{
"api_key_id": "source-user-key",
"key_hash": user_key_hash,
"name": "Imported User Key",
"is_active": true,
"total_requests": 222,
"total_tokens": 3333,
"total_cost_usd": 4.56
}]
}],
"standalone_keys": [{
"api_key_id": "source-standalone-key",
"key_hash": standalone_key_hash,
"name": "Imported Standalone Key",
"is_active": true,
"total_requests": 444,
"total_tokens": 5555,
"total_cost_usd": 6.78
}]
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["stats"]["users"]["updated"], json!(1));
assert_eq!(payload["stats"]["api_keys"]["updated"], json!(1));
assert_eq!(payload["stats"]["standalone_keys"]["updated"], json!(1));
let updated_records = state
.list_auth_api_key_export_records_by_ids(&[
"key-user-existing".to_string(),
"key-standalone-existing".to_string(),
])
.await
.expect("api key export records should load");
let user_key = updated_records
.iter()
.find(|record| record.api_key_id == "key-user-existing")
.expect("updated user key should exist");
assert_eq!(user_key.total_requests, 222);
assert_eq!(user_key.total_tokens, 3333);
assert_eq!(user_key.total_cost_usd, 4.56);
let standalone_key = updated_records
.iter()
.find(|record| record.api_key_id == "key-standalone-existing")
.expect("updated standalone key should exist");
assert_eq!(standalone_key.total_requests, 444);
assert_eq!(standalone_key.total_tokens, 5555);
assert_eq!(standalone_key.total_cost_usd, 6.78);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_imports_admin_system_config_fixture_v22() {
let gateway = build_router_with_state(
@@ -1671,33 +1880,20 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
}
#[tokio::test]
async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_import_and_forces_refresh(
async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_import_without_refresh(
) {
#[derive(Debug, Clone)]
struct SeenRefreshRequest {
content_type: String,
body: String,
}
let seen_refresh = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
let seen_refresh = Arc::new(Mutex::new(false));
let seen_refresh_clone = Arc::clone(&seen_refresh);
let refresh_hits = Arc::new(Mutex::new(0usize));
let refresh_hits_clone = Arc::clone(&refresh_hits);
let refresh_server = Router::new().route(
"/oauth/token",
post(move |headers: HeaderMap, body: Bytes| {
post(move |_headers: HeaderMap, _body: Bytes| {
let seen_refresh_inner = Arc::clone(&seen_refresh_clone);
let refresh_hits_inner = Arc::clone(&refresh_hits_clone);
async move {
*refresh_hits_inner.lock().expect("mutex should lock") += 1;
*seen_refresh_inner.lock().expect("mutex should lock") = Some(SeenRefreshRequest {
content_type: headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(body.to_vec()).unwrap_or_default(),
});
*seen_refresh_inner.lock().expect("mutex should lock") = true;
axum::Json(json!({
"access_token": "oauth-access-token-refreshed",
"refresh_token": "oauth-refresh-token-refreshed",
@@ -1792,21 +1988,8 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 1);
let seen_refresh = seen_refresh
.lock()
.expect("mutex should lock")
.clone()
.expect("refresh request should be captured");
assert_eq!(
seen_refresh.content_type,
"application/x-www-form-urlencoded"
);
assert!(seen_refresh.body.contains("grant_type=refresh_token"));
assert!(seen_refresh
.body
.contains("refresh_token=oauth-refresh-token-new"));
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 0);
assert!(!*seen_refresh.lock().expect("mutex should lock"));
let providers = provider_catalog_repository
.list_providers(false)
@@ -1823,7 +2006,7 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
assert_eq!(key.name, "oauth-primary");
assert_eq!(key.oauth_invalid_at_unix_secs, None);
assert_eq!(key.oauth_invalid_reason, None);
assert!(key.expires_at_unix_secs.is_some());
assert_eq!(key.expires_at_unix_secs, None);
assert_eq!(
decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
@@ -1832,7 +2015,7 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
.expect("api key should be present"),
)
.expect("oauth access token should decrypt"),
"oauth-access-token-refreshed"
"oauth-access-token-new"
);
let auth_config = decrypt_python_fernet_ciphertext(
@@ -1845,15 +2028,12 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
let auth_config: Value =
serde_json::from_str(&auth_config).expect("oauth auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(
auth_config["refresh_token"],
"oauth-refresh-token-refreshed"
);
assert_eq!(auth_config["refresh_token"], "oauth-refresh-token-new");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-codex-123");
assert_eq!(auth_config["plan_type"], "plus");
assert_eq!(auth_config["token_type"], "Bearer");
assert_eq!(auth_config["expires_at"].as_u64(), key.expires_at_unix_secs);
assert!(auth_config.get("token_type").is_none());
assert!(auth_config.get("expires_at").is_none());
gateway_handle.abort();
refresh_handle.abort();