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

fix: scope scheduler effects by session

Browse Source
This commit is contained in:
RWDai
2026-05-05 12:25:07 +08:00
parent 82e7a0080e
commit 628329e493
1 changed files with 150 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,8 +3,8 @@ use std::collections::BTreeMap;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, count_recent_rpm_requests_for_provider_key,
SchedulerAffinityTarget,
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
count_recent_rpm_requests_for_provider_key, ClientSessionAffinity, SchedulerAffinityTarget,
};
use aether_usage_runtime::{
build_stream_terminal_usage_outcome, build_sync_terminal_usage_outcome,
@@ -152,13 +152,47 @@ fn report_context_string_field<'a>(
}
fn local_scheduler_affinity_cache_key(report_context: Option<&Value>) -> Option<String> {
build_scheduler_affinity_cache_key_for_api_key_id(
let client_session_affinity = local_client_session_affinity(report_context);
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
report_context_string_field(report_context, "api_key_id")?,
report_context_string_field(report_context, "client_api_format")?,
report_context_string_field(report_context, "model")?,
client_session_affinity.as_ref(),
)
}
fn local_client_session_affinity(report_context: Option<&Value>) -> Option<ClientSessionAffinity> {
let report_context = report_context?;
let headers = header_map_from_report_context(report_context.get("original_headers"));
let body_json = report_context
.get("original_request_body")
.filter(|value| !value.is_null());
crate::client_session_affinity::client_session_affinity_from_request(&headers, body_json)
}
fn header_map_from_report_context(headers: Option<&Value>) -> http::HeaderMap {
let mut header_map = http::HeaderMap::new();
let Some(headers) = headers.and_then(Value::as_object) else {
return header_map;
};
for (name, value) in headers {
let Some(value) = value.as_str() else {
continue;
};
let Ok(name) = http::header::HeaderName::from_bytes(name.as_bytes()) else {
continue;
};
let Ok(value) = http::HeaderValue::from_str(value) else {
continue;
};
header_map.insert(name, value);
}
header_map
}
fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffinityTarget> {
let provider_id = plan.provider_id.trim();
let endpoint_id = plan.endpoint_id.trim();
@@ -737,7 +771,9 @@ mod tests {
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::AppState;
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
build_scheduler_affinity_cache_key_for_api_key_id,
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
ClientSessionAffinity, SchedulerAffinityTarget,
};
async fn start_managed_redis_or_skip() -> Option<ManagedRedisServer> {
@@ -775,6 +811,38 @@ mod tests {
}
}
fn session_affinity() -> ClientSessionAffinity {
ClientSessionAffinity::new(
Some("generic".to_string()),
Some("session=session-1;agent=coder".to_string()),
)
}
fn session_report_context() -> Value {
json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
"original_headers": {
"x-aether-session-id": "session-1",
"x-aether-agent-id": "coder"
},
"original_request_body": {
"model": "gpt-5"
}
})
}
fn session_scheduler_affinity_cache_key() -> String {
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
"api-key-1",
"openai:chat",
"gpt-5",
Some(&session_affinity()),
)
.expect("session scheduler affinity cache key should build")
}
fn sample_codex_plan() -> ExecutionPlan {
ExecutionPlan {
request_id: "req-codex-1".to_string(),
@@ -1096,6 +1164,50 @@ mod tests {
.is_none());
}
#[tokio::test]
async fn attempt_failure_invalidates_session_scoped_scheduler_affinity_cache() {
let state = AppState::new().expect("gateway state should build");
let plan = sample_plan();
let report_context = session_report_context();
let session_cache_key = session_scheduler_affinity_cache_key();
let legacy_cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("legacy scheduler affinity cache key should build");
for cache_key in [&session_cache_key, &legacy_cache_key] {
state.scheduler_affinity_cache.insert(
cache_key.to_string(),
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-1".to_string(),
},
SCHEDULER_AFFINITY_TTL,
16,
);
}
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 429,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert!(state
.read_scheduler_affinity_target(session_cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
.is_none());
assert!(state
.read_scheduler_affinity_target(legacy_cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
.is_some());
}
#[tokio::test]
async fn attempt_failure_keeps_scheduler_affinity_for_non_failure_status() {
let state = AppState::new().expect("gateway state should build");
@@ -1213,6 +1325,40 @@ mod tests {
);
}
#[tokio::test]
async fn success_remembers_session_scoped_scheduler_affinity_cache() {
let state = AppState::new().expect("gateway state should build");
let plan = sample_plan();
let report_context = session_report_context();
let session_cache_key = session_scheduler_affinity_cache_key();
let legacy_cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("legacy 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(session_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(),
})
);
assert!(state
.read_scheduler_affinity_target(legacy_cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
.is_none());
}
#[tokio::test]
async fn fallback_success_rewarms_scheduler_affinity_after_failed_candidate_invalidates() {
let state = AppState::new().expect("gateway state should build");