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

Merge remote-tracking branch 'origin/main'

Browse Source
This commit is contained in:
fawney19
2026-05-22 14:16:07 +08:00
parent 8966fd6aac f6c3ebf7d3
commit 8ede01ad4e
23 changed files with 1550 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

173
apps/aether-gateway/src/bark_push.rs Normal file
View File

@@ -0,0 +1,173 @@
use crate::handlers::shared::{
decrypt_catalog_secret_with_fallbacks, system_config_bool, system_config_string,
};
use crate::{AppState, GatewayError};
use serde_json::{json, Value};
pub(crate) const BARK_PUSH_ENABLED_KEY: &str = "module.bark_push.enabled";
pub(crate) const BARK_PUSH_DEVICE_KEY_KEY: &str = "module.bark_push.device_key";
pub(crate) const BARK_PUSH_SERVER_URL_KEY: &str = "module.bark_push.server_url";
pub(crate) const BARK_PUSH_TEMPLATE_KEY: &str = "module.bark_push.template";
const DEFAULT_BARK_API_BASE: &str = "https://api.day.app";
#[derive(Debug, Clone)]
pub(crate) struct BarkPushConfig {
pub(crate) enabled: bool,
pub(crate) device_key: Option<String>,
pub(crate) server_url: String,
pub(crate) template: Option<String>,
}
pub(crate) async fn bark_push_module_enabled(state: &AppState) -> Result<bool, GatewayError> {
let value = state
.read_system_config_json_value(BARK_PUSH_ENABLED_KEY)
.await?;
Ok(system_config_bool(value.as_ref(), false))
}
pub(crate) async fn bark_push_configured(state: &AppState) -> Result<bool, GatewayError> {
let config = read_bark_push_config(state).await?;
Ok(config.device_key.is_some() && !config.server_url.trim().is_empty())
}
pub(crate) async fn read_bark_push_config(
state: &AppState,
) -> Result<BarkPushConfig, GatewayError> {
let enabled = bark_push_module_enabled(state).await?;
let device_key = state
.read_system_config_json_value(BARK_PUSH_DEVICE_KEY_KEY)
.await?
.and_then(|value| system_config_string(Some(&value)))
.map(|value| {
decrypt_catalog_secret_with_fallbacks(state.encryption_key(), &value).unwrap_or(value)
});
let server_url = state
.read_system_config_json_value(BARK_PUSH_SERVER_URL_KEY)
.await?
.and_then(|value| system_config_string(Some(&value)))
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| DEFAULT_BARK_API_BASE.to_string());
let template = state
.read_system_config_json_value(BARK_PUSH_TEMPLATE_KEY)
.await?
.and_then(|value| system_config_string(Some(&value)));
Ok(BarkPushConfig {
enabled,
device_key,
server_url,
template,
})
}
pub(crate) async fn send_bark_push(
state: &AppState,
config: &BarkPushConfig,
title: &str,
markdown_body: &str,
) -> Result<(), GatewayError> {
let Some(device_key) = config.device_key.as_deref() else {
return Err(GatewayError::Internal("未配置 Bark Device Key".to_string()));
};
let device_key = device_key.trim();
if device_key.is_empty() {
return Err(GatewayError::Internal(
"Bark Device Key 不能为空".to_string(),
));
}
let server_url = normalized_bark_server_url(&config.server_url)?;
let body = render_bark_body(config.template.as_deref(), title, markdown_body);
let response = state
.client
.post(format!("{server_url}/push"))
.json(&json!({
"device_key": device_key,
"title": title,
"body": body,
}))
.send()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if !status.is_success() {
return Err(GatewayError::Internal(format!(
"Bark 返回 HTTP {status}: {text}"
)));
}
if let Ok(payload) = serde_json::from_str::<Value>(&text) {
let code_is_ok = payload
.get("code")
.and_then(|value| {
value
.as_i64()
.map(|code| matches!(code, 0 | 200))
.or_else(|| {
value
.as_str()
.map(|code| matches!(code.trim(), "0" | "200"))
})
})
.unwrap_or(true);
if !code_is_ok {
return Err(GatewayError::Internal(format!("Bark 返回失败: {payload}")));
}
}
Ok(())
}
fn normalized_bark_server_url(server_url: &str) -> Result<String, GatewayError> {
let server_url = server_url.trim().trim_end_matches('/');
if server_url.is_empty() {
return Err(GatewayError::Internal(
"Bark 服务器地址不能为空".to_string(),
));
}
if !server_url.starts_with("https://") && !server_url.starts_with("http://") {
return Err(GatewayError::Internal(
"Bark 服务器地址必须以 http:// 或 https:// 开头".to_string(),
));
}
Ok(server_url.to_string())
}
fn render_bark_body(template: Option<&str>, title: &str, markdown_body: &str) -> String {
match template {
Some(template) if !template.trim().is_empty() => template
.replace("{title}", title)
.replace("{body}", markdown_body),
_ => markdown_body.to_string(),
}
}
#[cfg(test)]
mod tests {
use super::{normalized_bark_server_url, render_bark_body};
#[test]
fn bark_body_uses_template_when_provided() {
let rendered = render_bark_body(Some("{title}\n\n{body}"), "告警", "原始正文");
assert_eq!(rendered, "告警\n\n原始正文");
}
#[test]
fn bark_body_falls_back_to_markdown_body_for_empty_template() {
assert_eq!(render_bark_body(None, "告警", "原始正文"), "原始正文");
assert_eq!(
render_bark_body(Some(" "), "告警", "原始正文"),
"原始正文"
);
}
#[test]
fn bark_server_url_trims_trailing_slashes() {
assert_eq!(
normalized_bark_server_url(" https://api.day.app/ ").expect("url should parse"),
"https://api.day.app"
);
}
}

110
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -1464,6 +1464,8 @@ fn build_sse_body_stream(
#[derive(Default)]
struct SseControlBlockFilter {
buffered: Vec<u8>,
emitted_len: usize,
passthrough_current_block: bool,
}
impl SseControlBlockFilter {
@@ -1475,17 +1477,39 @@ impl SseControlBlockFilter {
self.buffered.extend_from_slice(chunk);
let mut output = Vec::new();
while let Some((block_end, separator_len)) = find_sse_block_boundary(&self.buffered) {
let block = self
.buffered
.drain(..block_end + separator_len)
.collect::<Vec<_>>();
if sse_block_has_data_line(&block) {
output.extend(block);
let block_len = block_end + separator_len;
let block = self.buffered.drain(..block_len).collect::<Vec<_>>();
if self.passthrough_current_block {
let emitted_len = self.emitted_len.min(block.len());
output.extend_from_slice(&block[emitted_len..]);
} else if sse_block_has_data_line(&block) {
output.extend_from_slice(&block);
}
self.emitted_len = 0;
self.passthrough_current_block = false;
}
if self.passthrough_current_block {
if self.buffered.len() > self.emitted_len {
output.extend_from_slice(&self.buffered[self.emitted_len..]);
self.emitted_len = self.buffered.len();
}
} else if sse_buffer_has_data_line(&self.buffered) {
self.passthrough_current_block = true;
output.extend_from_slice(&self.buffered);
self.emitted_len = self.buffered.len();
}
if self.buffered.len() > SSE_CONTROL_FILTER_MAX_BUFFER_BYTES {
output.extend(std::mem::take(&mut self.buffered));
let buffered = std::mem::take(&mut self.buffered);
if self.passthrough_current_block {
let emitted_len = self.emitted_len.min(buffered.len());
output.extend_from_slice(&buffered[emitted_len..]);
} else {
output.extend(buffered);
}
self.emitted_len = 0;
self.passthrough_current_block = false;
}
output
@@ -1497,7 +1521,13 @@ impl SseControlBlockFilter {
}
let block = std::mem::take(&mut self.buffered);
if sse_block_has_data_line(&block) {
let emitted_len = self.emitted_len.min(block.len());
let passthrough_current_block = self.passthrough_current_block;
self.emitted_len = 0;
self.passthrough_current_block = false;
if passthrough_current_block {
block[emitted_len..].to_vec()
} else if sse_block_has_data_line(&block) {
block
} else {
Vec::new()
@@ -1549,6 +1579,15 @@ fn sse_block_has_data_line(block: &[u8]) -> bool {
.any(|line| line.trim_start().starts_with("data:"))
}
fn sse_buffer_has_data_line(buffer: &[u8]) -> bool {
let Ok(text) = std::str::from_utf8(buffer) else {
return true;
};
text.lines()
.any(|line| line.trim_start().starts_with("data:"))
}
fn stream_chunk_contains_sse_done(chunk: &[u8]) -> bool {
std::str::from_utf8(chunk).ok().is_some_and(|text| {
text.lines().any(|line| {
@@ -4691,6 +4730,61 @@ mod tests {
);
}
#[tokio::test]
async fn sse_body_stream_forwards_data_line_before_block_boundary() {
let (tx, rx) = mpsc::channel::<Result<Bytes, std::io::Error>>(4);
let mut body_stream = Box::pin(build_sse_body_stream(
Vec::new(),
rx,
true,
Duration::from_secs(60),
));
let keepalive = tokio::time::timeout(Duration::from_millis(50), body_stream.next())
.await
.expect("initial keepalive should be immediate")
.expect("stream should yield initial keepalive")
.expect("initial keepalive should be ok");
assert_eq!(keepalive.as_ref(), b": aether-keepalive\n\n");
tx.send(Ok(Bytes::from_static(
b"event: response.output_text.delta\n",
)))
.await
.expect("event line should send");
assert!(
tokio::time::timeout(Duration::from_millis(20), body_stream.next())
.await
.is_err(),
"event-only partial block should remain buffered"
);
tx.send(Ok(Bytes::from_static(
b"data: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n",
)))
.await
.expect("data line should send");
let data_chunk = tokio::time::timeout(Duration::from_millis(50), body_stream.next())
.await
.expect("data-bearing block should stream before terminator")
.expect("stream should yield data-bearing block")
.expect("data-bearing block should be ok");
assert_eq!(
data_chunk.as_ref(),
b"event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n"
);
tx.send(Ok(Bytes::from_static(b"\n")))
.await
.expect("terminator should send");
let terminator = tokio::time::timeout(Duration::from_millis(50), body_stream.next())
.await
.expect("terminator should stream")
.expect("stream should yield terminator")
.expect("terminator should be ok");
assert_eq!(terminator.as_ref(), b"\n");
}
#[tokio::test]
async fn sse_body_stream_uses_local_keepalive_when_prefetched_blocks_are_control_only() {
let (_tx, rx) = mpsc::channel::<Result<Bytes, std::io::Error>>(1);

17
apps/aether-gateway/src/handlers/admin/system/shared/modules.rs
View File

@@ -1,3 +1,4 @@
use crate::bark_push::bark_push_configured;
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::shared::{module_available_from_env, system_config_bool};
use crate::important_notification::{
@@ -96,6 +97,18 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
admin_menu_group: Some("system"),
admin_menu_order: 59,
},
AdminModuleDefinition {
name: "bark_push",
display_name: "Bark 推送",
description: "第三方推送服务，配置 Bark Device Key 并测试 iOS 推送",
category: "integration",
env_key: "BARK_PUSH_AVAILABLE",
default_available: true,
admin_route: Some("/admin/modules/bark"),
admin_menu_icon: Some("Send"),
admin_menu_group: Some("system"),
admin_menu_order: 59,
},
AdminModuleDefinition {
name: "model_directives",
display_name: "模型后缀参数",
@@ -169,6 +182,7 @@ pub(crate) struct AdminModuleRuntimeState {
gemini_files_has_capable_key: bool,
important_notification_configured: bool,
server_chan_push_configured: bool,
bark_push_configured: bool,
}
pub(crate) fn admin_module_by_name(name: &str) -> Option<&'static AdminModuleDefinition> {
@@ -257,6 +271,7 @@ pub(crate) async fn build_admin_module_runtime_state(
let notification_configured = important_notification_configured(state.app()).await?;
let server_chan_configured = server_chan_push_configured(state.app()).await?;
let bark_configured = bark_push_configured(state.app()).await?;
Ok(AdminModuleRuntimeState {
oauth_providers,
@@ -264,6 +279,7 @@ pub(crate) async fn build_admin_module_runtime_state(
gemini_files_has_capable_key,
important_notification_configured: notification_configured,
server_chan_push_configured: server_chan_configured,
bark_push_configured: bark_configured,
})
}
@@ -278,6 +294,7 @@ pub(crate) fn build_admin_module_validation_result(
runtime.gemini_files_has_capable_key,
runtime.important_notification_configured,
runtime.server_chan_push_configured,
runtime.bark_push_configured,
)
}

80
apps/aether-gateway/src/important_notification.rs
View File

@@ -1,4 +1,5 @@
use crate::admin_api::AdminAppState;
use crate::bark_push::{read_bark_push_config, send_bark_push, BarkPushConfig};
use crate::email_delivery::{
read_smtp_delivery_config, send_smtp_email, ComposedEmail, SmtpDeliveryConfig,
};
@@ -35,6 +36,7 @@ pub(crate) enum ImportantNotificationChannelFilter {
All,
Email,
ServerChan,
Bark,
}
#[derive(Debug, Clone)]
@@ -45,6 +47,7 @@ struct ImportantNotificationConfig {
default_channel: ImportantNotificationChannelFilter,
items: Vec<ImportantNotificationItemConfig>,
server_chan: ServerChanPushConfig,
bark: BarkPushConfig,
}
#[derive(Debug, Clone)]
@@ -63,6 +66,7 @@ struct ImportantNotificationItemConfig {
struct NotificationChannelReadiness {
email: bool,
server_chan: bool,
bark: bool,
}
#[derive(Debug, Clone, Serialize)]
@@ -226,6 +230,7 @@ async fn read_notification_channel_readiness(
Ok(NotificationChannelReadiness {
email: config.email_enabled && !config.email_recipients.is_empty() && smtp_config.is_some(),
server_chan: config.server_chan.enabled && config.server_chan.send_key.is_some(),
bark: config.bark.enabled && config.bark.device_key.is_some(),
})
}
@@ -234,9 +239,12 @@ fn channel_filter_has_ready_channel(
readiness: NotificationChannelReadiness,
) -> bool {
match filter {
ImportantNotificationChannelFilter::All => readiness.email || readiness.server_chan,
ImportantNotificationChannelFilter::All => {
readiness.email || readiness.server_chan || readiness.bark
}
ImportantNotificationChannelFilter::Email => readiness.email,
ImportantNotificationChannelFilter::ServerChan => readiness.server_chan,
ImportantNotificationChannelFilter::Bark => readiness.bark,
}
}
@@ -292,6 +300,20 @@ async fn dispatch_important_notification(
.await;
}
if matches!(
channel_filter,
ImportantNotificationChannelFilter::All | ImportantNotificationChannelFilter::Bark
) {
maybe_send_bark_notification(
state,
&config,
&notification,
bypass_enable_checks,
&mut reports,
)
.await;
}
if reports.is_empty() {
reports.push(ImportantNotificationChannelReport {
channel: "none",
@@ -385,6 +407,7 @@ async fn read_important_notification_config(
.unwrap_or(ImportantNotificationChannelFilter::All),
items: parse_notification_items(items.as_ref()),
server_chan: read_server_chan_push_config(state).await?,
bark: read_bark_push_config(state).await?,
})
}
@@ -395,6 +418,7 @@ fn parse_channel_filter(raw: &str) -> Option<ImportantNotificationChannelFilter>
"server_chan" | "serverchan" | "serve_chan" => {
Some(ImportantNotificationChannelFilter::ServerChan)
}
"bark" => Some(ImportantNotificationChannelFilter::Bark),
"global" | "" => None,
_ => None,
}
@@ -680,6 +704,48 @@ async fn maybe_send_server_chan_notification(
}
}
async fn maybe_send_bark_notification(
state: &AppState,
config: &ImportantNotificationConfig,
notification: &ImportantNotification,
bypass_channel_toggle: bool,
reports: &mut Vec<ImportantNotificationChannelReport>,
) {
if !bypass_channel_toggle && !config.bark.enabled {
return;
}
if config.bark.device_key.is_none() {
reports.push(ImportantNotificationChannelReport {
channel: "bark",
success: false,
message: "未配置 Bark Device Key".to_string(),
});
return;
};
match send_bark_push(
state,
&config.bark,
&notification.title,
&notification.markdown_body,
)
.await
{
Ok(()) => reports.push(ImportantNotificationChannelReport {
channel: "bark",
success: true,
message: "Bark 通知已发送".to_string(),
}),
Err(err) => {
warn!(error = ?err, "failed to send bark important notification");
reports.push(ImportantNotificationChannelReport {
channel: "bark",
success: false,
message: format!("Bark 通知发送失败: {err:?}"),
});
}
}
}
fn single_report(
channel: &'static str,
success: bool,
@@ -743,8 +809,8 @@ fn escape_html(value: &str) -> String {
#[cfg(test)]
mod tests {
use super::{
apply_notification_item_template, parse_notification_items, parse_recipient_list,
ImportantNotification, ImportantNotificationChannelFilter,
apply_notification_item_template, parse_channel_filter, parse_notification_items,
parse_recipient_list, ImportantNotification, ImportantNotificationChannelFilter,
};
use serde_json::json;
@@ -785,6 +851,14 @@ mod tests {
assert!(items[0].user_email_enabled);
}
#[test]
fn parse_channel_filter_accepts_bark() {
assert_eq!(
parse_channel_filter("bark"),
Some(ImportantNotificationChannelFilter::Bark)
);
}
#[test]
fn item_template_renders_fallback_and_variables() {
let items = parse_notification_items(Some(&json!([

1
apps/aether-gateway/src/lib.rs
View File

@@ -30,6 +30,7 @@ mod api;
mod async_task;
mod audit;
mod auth;
mod bark_push;
mod cache;
mod client_session_affinity;
mod clock;

290
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -243,6 +243,35 @@ fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffi
})
}
async fn local_execution_plan_uses_pool(state: &AppState, plan: &ExecutionPlan) -> bool {
let Ok(Some(transport)) = state
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
.await
else {
return false;
};
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref()).is_some()
}
async fn local_scheduler_affinity_matches_failed_target(
state: &AppState,
plan: &ExecutionPlan,
cached_target: &SchedulerAffinityTarget,
failed_target: &SchedulerAffinityTarget,
) -> bool {
if cached_target == failed_target {
return true;
}
if cached_target.provider_id != failed_target.provider_id
|| cached_target.endpoint_id != failed_target.endpoint_id
{
return false;
}
local_execution_plan_uses_pool(state, plan).await
}
async fn scheduler_cache_affinity_enabled(state: &AppState) -> bool {
match read_scheduler_ordering_config(state).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
@@ -360,7 +389,24 @@ async fn record_attempt_failure_effect(
}
if let Some(cache_key) = local_scheduler_affinity_cache_key(context.report_context) {
let _ = state.remove_scheduler_affinity_cache_entry(&cache_key);
let Some(failed_target) = local_scheduler_affinity_target(context.plan) else {
return;
};
let Some(cached_target) =
state.read_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
else {
return;
};
if local_scheduler_affinity_matches_failed_target(
state,
context.plan,
&cached_target,
&failed_target,
)
.await
{
let _ = state.remove_scheduler_affinity_cache_entry(&cache_key);
}
}
}
@@ -448,7 +494,10 @@ async fn record_adaptive_rate_limit_effect(
updated_key.status_snapshot = Some(projection.status_snapshot);
updated_key.updated_at_unix_secs = Some(observed_at_unix_secs);
if let Err(err) = state.update_provider_catalog_key(&updated_key).await {
if let Err(err) = state
.update_provider_catalog_key_runtime_state(&updated_key)
.await
{
warn!(
"gateway orchestration effects: failed to persist adaptive rate-limit projection for provider {} endpoint {} key {}: {:?}",
context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
@@ -496,7 +545,10 @@ async fn record_adaptive_success_effect(
updated_key.status_snapshot = Some(projection.status_snapshot);
updated_key.updated_at_unix_secs = Some(observed_at_unix_secs);
if let Err(err) = state.update_provider_catalog_key(&updated_key).await {
if let Err(err) = state
.update_provider_catalog_key_runtime_state(&updated_key)
.await
{
warn!(
"gateway orchestration effects: failed to persist adaptive success projection for provider {} endpoint {} key {}: {:?}",
context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
@@ -1204,6 +1256,20 @@ mod tests {
.expect("provider should build")
}
fn sample_pool_health_provider() -> StoredProviderCatalogProvider {
sample_health_provider().with_transport_fields(
true,
false,
false,
None,
None,
None,
None,
None,
Some(json!({"pool_advanced": {}})),
)
}
fn sample_health_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"ep-1".to_string(),
@@ -1266,6 +1332,20 @@ mod tests {
)
}
fn pool_health_state() -> AppState {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_pool_health_provider()],
vec![sample_health_endpoint()],
vec![sample_health_key()],
));
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
}
fn health_state_with_key(key: StoredProviderCatalogKey) -> AppState {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_health_provider()],
@@ -1429,6 +1509,137 @@ mod tests {
.is_some());
}
#[tokio::test]
async fn attempt_failure_keeps_scheduler_affinity_for_non_affinity_candidate() {
let state = AppState::new().expect("gateway state should build");
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let affinity_target = SchedulerAffinityTarget {
provider_id: "prov-2".to_string(),
endpoint_id: "ep-2".to_string(),
key_id: "key-2".to_string(),
};
state.remember_scheduler_affinity_target(
&cache_key,
affinity_target.clone(),
SCHEDULER_AFFINITY_TTL,
16,
);
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 524,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(affinity_target)
);
}
#[tokio::test]
async fn attempt_failure_keeps_scheduler_affinity_for_non_pool_sibling_key() {
let state = health_state();
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let affinity_target = SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-2".to_string(),
};
state.remember_scheduler_affinity_target(
&cache_key,
affinity_target.clone(),
SCHEDULER_AFFINITY_TTL,
16,
);
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 524,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(affinity_target)
);
}
#[tokio::test]
async fn attempt_failure_invalidates_scheduler_affinity_for_same_pool_candidate() {
let state = pool_health_state();
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-2".to_string(),
},
SCHEDULER_AFFINITY_TTL,
16,
);
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 524,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
None
);
}
#[tokio::test]
async fn attempt_failure_keeps_scheduler_affinity_for_non_failure_status() {
let state = AppState::new().expect("gateway state should build");
@@ -1546,6 +1757,39 @@ mod tests {
);
}
#[tokio::test]
async fn health_success_keeps_scheduler_affinity_after_health_state_update() {
let state = health_state();
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-1".to_string(),
})
);
}
#[tokio::test]
async fn load_balance_success_does_not_remember_scheduler_affinity_cache() {
let state = AppState::new()
@@ -1985,6 +2229,21 @@ mod tests {
async fn adaptive_rate_limit_effect_updates_adaptive_key_observation() {
let state = adaptive_state();
let plan = sample_plan();
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let target = SchedulerAffinityTarget {
provider_id: plan.provider_id.clone(),
endpoint_id: plan.endpoint_id.clone(),
key_id: plan.key_id.clone(),
};
state.remember_scheduler_affinity_target(
&cache_key,
target.clone(),
SCHEDULER_AFFINITY_TTL,
16,
);
let initial_epoch = state.scheduler_affinity_epoch();
apply_local_execution_effect(
&state,
@@ -2048,6 +2307,11 @@ mod tests {
.and_then(|value| value.get("enforcement_active")),
Some(&json!(false))
);
assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(target)
);
}
#[tokio::test]
@@ -2178,6 +2442,21 @@ mod tests {
.expect("request candidate should build")],
);
let plan = sample_plan();
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let target = SchedulerAffinityTarget {
provider_id: plan.provider_id.clone(),
endpoint_id: plan.endpoint_id.clone(),
key_id: plan.key_id.clone(),
};
state.remember_scheduler_affinity_target(
&cache_key,
target.clone(),
SCHEDULER_AFFINITY_TTL,
16,
);
let initial_epoch = state.scheduler_affinity_epoch();
apply_local_execution_effect(
&state,
@@ -2209,5 +2488,10 @@ mod tests {
.and_then(Value::as_str),
Some("high_utilization")
);
assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(target)
);
}
}

4
apps/aether-gateway/src/orchestration/report_effects.rs
View File

@@ -424,7 +424,7 @@ async fn sync_grok_quota_from_report_context(
updated_key.updated_at_unix_secs = Some(now_unix_secs);
Ok(state
.update_provider_catalog_key(&updated_key)
.update_provider_catalog_key_runtime_state(&updated_key)
.await?
.is_some())
}
@@ -721,7 +721,7 @@ async fn sync_codex_quota_from_response_headers(
updated_key.updated_at_unix_secs = Some(now_unix_secs);
let updated = state
.update_provider_catalog_key(&updated_key)
.update_provider_catalog_key_runtime_state(&updated_key)
.await?
.is_some();
if updated {

60
apps/aether-gateway/src/state/catalog.rs
View File

@@ -615,6 +615,21 @@ impl AppState {
Ok(updated)
}
pub(crate) async fn update_provider_catalog_key_runtime_state(
&self,
key: &provider_catalog::StoredProviderCatalogKey,
) -> Result<Option<provider_catalog::StoredProviderCatalogKey>, GatewayError> {
let updated = self
.data
.update_provider_catalog_key(key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated.is_some() {
self.invalidate_provider_health_routing_caches();
}
Ok(updated)
}
pub(crate) async fn update_provider_catalog_key_upstream_metadata(
&self,
key_id: &str,
@@ -806,7 +821,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated {
self.invalidate_provider_routing_caches();
self.invalidate_provider_health_routing_caches();
}
Ok(updated)
}
@@ -930,4 +945,47 @@ mod tests {
.expect("provider transport should exist after update");
assert!(snapshot.provider.keep_priority_on_conversion);
}
#[tokio::test]
async fn provider_catalog_health_update_keeps_scheduler_affinity_cache() {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_encryption_key_for_tests("test-encryption-key"),
);
let cache_key = "scheduler_affinity:api-key-1:openai:chat:gpt-5";
let ttl = Duration::from_secs(300);
let target = SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
};
state.remember_scheduler_affinity_target(cache_key, target.clone(), ttl, 128);
let initial_epoch = state.scheduler_affinity_epoch();
let health_by_format = serde_json::json!({
"openai:chat": {
"last_success_at_unix_secs": 1,
"consecutive_failures": 0
}
});
let updated = state
.update_provider_catalog_key_health_state("key-1", true, Some(&health_by_format), None)
.await
.expect("key health update should succeed");
assert!(updated);
assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
assert_eq!(
state.read_scheduler_affinity_target(cache_key, ttl),
Some(target)
);
}
}

5
apps/aether-gateway/src/state/core.rs
View File

@@ -572,6 +572,11 @@ impl AppState {
self.invalidate_scheduler_affinity_cache();
}
pub(crate) fn invalidate_provider_health_routing_caches(&self) {
self.data.clear_minimal_candidate_selection_cache();
self.clear_provider_transport_snapshot_cache();
}
pub(crate) fn invalidate_auth_context_cache(&self) {
self.auth_context_cache.clear();
}

2
apps/aether-gateway/src/tests/control/admin/health_access.rs
View File

@@ -816,6 +816,8 @@ async fn gateway_handles_admin_modules_status_locally_with_trusted_admin_princip
payload["server_chan_push"]["admin_route"],
"/admin/modules/server-chan"
);
assert_eq!(payload["bark_push"]["display_name"], "Bark 推送");
assert_eq!(payload["bark_push"]["admin_route"], "/admin/modules/bark");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();