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

refactor: isolate dispatch scheduling core

Browse Source
This commit is contained in:
fawney19
2026-05-12 13:15:41 +08:00
parent 81ee27cdea
commit fa73655134
50 changed files with 4467 additions and 3202 deletions
Download Patch File Download Diff File Expand all files Collapse all files

99
apps/aether-gateway/src/state/runtime/api_key_exports.rs
View File

@@ -182,10 +182,15 @@ impl AppState {
record: aether_data::repository::auth::CreateUserApiKeyRecord,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.create_user_api_key(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn create_standalone_api_key(
@@ -193,10 +198,15 @@ impl AppState {
record: aether_data::repository::auth::CreateStandaloneApiKeyRecord,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.create_standalone_api_key(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn update_user_api_key_basic(
@@ -204,10 +214,15 @@ impl AppState {
record: aether_data::repository::auth::UpdateUserApiKeyBasicRecord,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.update_user_api_key_basic(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn update_standalone_api_key_basic(
@@ -215,10 +230,15 @@ impl AppState {
record: aether_data::repository::auth::UpdateStandaloneApiKeyBasicRecord,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.update_standalone_api_key_basic(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_user_api_key_active(
@@ -228,10 +248,15 @@ impl AppState {
is_active: bool,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.set_user_api_key_active(user_id, api_key_id, is_active)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_standalone_api_key_active(
@@ -240,10 +265,15 @@ impl AppState {
is_active: bool,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.set_standalone_api_key_active(api_key_id, is_active)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_user_api_key_locked(
@@ -252,10 +282,15 @@ impl AppState {
api_key_id: &str,
is_locked: bool,
) -> Result<bool, GatewayError> {
self.data
let updated = self
.data
.set_user_api_key_locked(user_id, api_key_id, is_locked)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated {
self.invalidate_auth_context_cache();
}
Ok(updated)
}
pub(crate) async fn set_user_api_key_allowed_providers(
@@ -265,10 +300,15 @@ impl AppState {
allowed_providers: Option<Vec<String>>,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.set_user_api_key_allowed_providers(user_id, api_key_id, allowed_providers)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_user_api_key_force_capabilities(
@@ -278,10 +318,15 @@ impl AppState {
force_capabilities: Option<serde_json::Value>,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.set_user_api_key_force_capabilities(user_id, api_key_id, force_capabilities)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn delete_user_api_key(
@@ -289,19 +334,29 @@ impl AppState {
user_id: &str,
api_key_id: &str,
) -> Result<bool, GatewayError> {
self.data
let deleted = self
.data
.delete_user_api_key(user_id, api_key_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
self.invalidate_auth_context_cache();
}
Ok(deleted)
}
pub(crate) async fn delete_standalone_api_key(
&self,
api_key_id: &str,
) -> Result<bool, GatewayError> {
self.data
let deleted = self
.data
.delete_standalone_api_key(api_key_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
self.invalidate_auth_context_cache();
}
Ok(deleted)
}
}

175
apps/aether-gateway/src/state/runtime/auth/user_lifecycle.rs
View File

@@ -248,10 +248,15 @@ impl AppState {
&self,
record: aether_data::repository::users::UpsertUserGroupRecord,
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
self.data
let group = self
.data
.create_user_group(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if group.is_some() {
self.invalidate_auth_context_cache();
}
Ok(group)
}
pub(crate) async fn update_user_group(
@@ -259,17 +264,27 @@ impl AppState {
group_id: &str,
record: aether_data::repository::users::UpsertUserGroupRecord,
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
self.data
let group = self
.data
.update_user_group(group_id, record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if group.is_some() {
self.invalidate_auth_context_cache();
}
Ok(group)
}
pub(crate) async fn delete_user_group(&self, group_id: &str) -> Result<bool, GatewayError> {
self.data
let deleted = self
.data
.delete_user_group(group_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
self.invalidate_auth_context_cache();
}
Ok(deleted)
}
pub(crate) async fn list_user_group_members(
@@ -287,10 +302,13 @@ impl AppState {
group_id: &str,
user_ids: &[String],
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMember>, GatewayError> {
self.data
let members = self
.data
.replace_user_group_members(group_id, user_ids)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.invalidate_auth_context_cache();
Ok(members)
}
pub(crate) async fn list_user_groups_for_user(
@@ -318,10 +336,13 @@ impl AppState {
user_id: &str,
group_ids: &[String],
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
self.data
let groups = self
.data
.replace_user_groups_for_user(user_id, group_ids)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.invalidate_auth_context_cache();
Ok(groups)
}
pub(crate) async fn add_user_to_group(
@@ -329,10 +350,15 @@ impl AppState {
group_id: &str,
user_id: &str,
) -> Result<bool, GatewayError> {
self.data
let added = self
.data
.add_user_to_group(group_id, user_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if added {
self.invalidate_auth_context_cache();
}
Ok(added)
}
pub(crate) async fn is_other_user_auth_email_taken(
@@ -417,13 +443,19 @@ impl AppState {
.lock()
.expect("auth user store should lock")
.insert(user.id.clone(), user.clone());
self.invalidate_auth_context_cache();
return Ok(Some(user));
}
self.data
let user = self
.data
.update_local_auth_user_profile(user_id, email, username)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if user.is_some() {
self.invalidate_auth_context_cache();
}
Ok(user)
}
pub(crate) async fn update_local_auth_user_password_hash(
@@ -606,10 +638,14 @@ impl AppState {
user.is_active = is_active;
}
let _ = (rate_limit_present, rate_limit);
return Ok(Some(user.clone()));
let user = user.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(user));
}
self.data
let user = self
.data
.update_local_auth_user_admin_fields(
user_id,
role,
@@ -624,7 +660,11 @@ impl AppState {
is_active,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if user.is_some() {
self.invalidate_auth_context_cache();
}
Ok(user)
}
pub(crate) async fn update_local_auth_user_policy_modes(
@@ -651,10 +691,14 @@ impl AppState {
user.allowed_models_mode = mode;
}
let _ = rate_limit_mode;
return Ok(Some(user.clone()));
let user = user.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(user));
}
self.data
let user = self
.data
.update_local_auth_user_policy_modes(
user_id,
allowed_providers_mode,
@@ -663,7 +707,11 @@ impl AppState {
rate_limit_mode,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if user.is_some() {
self.invalidate_auth_context_cache();
}
Ok(user)
}
pub(crate) async fn touch_auth_user_last_login(
@@ -821,3 +869,90 @@ fn normalized_user_group_ids(group_ids: &[String]) -> BTreeSet<String> {
.map(ToOwned::to_owned)
.collect()
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use aether_data::repository::users::{InMemoryUserReadRepository, UpsertUserGroupRecord};
use crate::control::GatewayControlAuthContext;
use crate::data::GatewayDataState;
use crate::AppState;
fn user_group_record(
allowed_models: Option<Vec<&str>>,
allowed_models_mode: &str,
) -> UpsertUserGroupRecord {
UpsertUserGroupRecord {
name: "Team".to_string(),
description: None,
priority: 0,
allowed_providers: None,
allowed_providers_mode: "unrestricted".to_string(),
allowed_api_formats: None,
allowed_api_formats_mode: "unrestricted".to_string(),
allowed_models: allowed_models.map(|values| {
values
.into_iter()
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
}),
allowed_models_mode: allowed_models_mode.to_string(),
rate_limit: None,
rate_limit_mode: "inherit".to_string(),
}
}
fn cached_auth_context() -> GatewayControlAuthContext {
GatewayControlAuthContext {
user_id: "user-1".to_string(),
api_key_id: "key-1".to_string(),
username: Some("alice".to_string()),
api_key_name: Some("default".to_string()),
balance_remaining: None,
access_allowed: true,
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: Some(vec!["gpt-4.1".to_string()]),
}
}
#[tokio::test]
async fn updating_user_group_invalidates_cached_auth_context() {
let repository = Arc::new(InMemoryUserReadRepository::default());
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(GatewayDataState::with_user_reader_for_tests(repository));
let group = state
.create_user_group(user_group_record(Some(vec!["gpt-4.1"]), "specific"))
.await
.expect("group should create")
.expect("group should exist");
let cache_key = "auth-context-cache-key".to_string();
let ttl = Duration::from_secs(60);
state
.auth_context_cache
.insert(cache_key.clone(), cached_auth_context(), ttl, 10);
assert!(state
.auth_context_cache
.get_fresh(&cache_key, ttl)
.is_some());
state
.update_user_group(&group.id, user_group_record(None, "unrestricted"))
.await
.expect("group should update")
.expect("group should exist after update");
assert!(state
.auth_context_cache
.get_fresh(&cache_key, ttl)
.is_none());
}
}

76
apps/aether-gateway/src/state/runtime/auth/user_provisioning.rs
View File

@@ -234,13 +234,19 @@ impl AppState {
.lock()
.expect("auth wallet store should lock")
.insert(wallet.id.clone(), wallet.clone());
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.initialize_auth_user_wallet(user_id, initial_gift_usd, unlimited)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
pub(crate) async fn initialize_auth_api_key_wallet(
@@ -284,13 +290,19 @@ impl AppState {
.lock()
.expect("auth wallet store should lock")
.insert(wallet.id.clone(), wallet.clone());
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.initialize_auth_api_key_wallet(api_key_id, initial_gift_usd, unlimited)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
pub(crate) async fn update_auth_user_wallet_limit_mode(
@@ -310,13 +322,21 @@ impl AppState {
let _ = wallet_id;
wallet.limit_mode = limit_mode.to_string();
wallet.updated_at_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
return Ok(Some(wallet.clone()));
let wallet = wallet.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.update_auth_user_wallet_limit_mode(user_id, limit_mode)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
pub(crate) async fn update_auth_api_key_wallet_limit_mode(
@@ -336,13 +356,21 @@ impl AppState {
let _ = wallet_id;
wallet.limit_mode = limit_mode.to_string();
wallet.updated_at_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
return Ok(Some(wallet.clone()));
let wallet = wallet.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.update_auth_api_key_wallet_limit_mode(api_key_id, limit_mode)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
#[allow(clippy::too_many_arguments)]
@@ -381,10 +409,14 @@ impl AppState {
if let Some(updated_at_unix_secs) = updated_at_unix_secs {
wallet.updated_at_unix_secs = updated_at_unix_secs;
}
return Ok(Some(wallet.clone()));
let wallet = wallet.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.update_auth_user_wallet_snapshot(
user_id,
balance,
@@ -399,7 +431,11 @@ impl AppState {
updated_at_unix_secs,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
#[allow(clippy::too_many_arguments)]
@@ -438,10 +474,14 @@ impl AppState {
if let Some(updated_at_unix_secs) = updated_at_unix_secs {
wallet.updated_at_unix_secs = updated_at_unix_secs;
}
return Ok(Some(wallet.clone()));
let wallet = wallet.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.update_auth_api_key_wallet_snapshot(
api_key_id,
balance,
@@ -456,6 +496,10 @@ impl AppState {
updated_at_unix_secs,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
}