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

Implement independent provider pool scheduling runtime

Browse Source
This commit is contained in:
fawney19
2026-04-17 23:05:49 +08:00
parent dd4641d618
commit 3321bb3ccc
38 changed files with 3222 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
apps/aether-gateway/src/ai_pipeline/mod.rs
View File

@@ -33,10 +33,11 @@ pub(crate) use self::planner::{
build_openai_cli_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
build_standard_family_stream_plan_and_reports, build_standard_family_sync_plan_and_reports,
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
set_local_openai_chat_execution_exhausted_diagnostic, GatewayAuthApiKeySnapshot,
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
maybe_build_sync_plan_payload, set_local_openai_chat_execution_exhausted_diagnostic,
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
PlannerAppState,
};
pub(crate) use self::pure::*;
pub(crate) use crate::control::GatewayControlDecision;

5
apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs
View File

@@ -1040,6 +1040,7 @@ mod tests {
"openai:chat",
"gpt-4.1",
None,
None,
)
.await;
@@ -1115,6 +1116,7 @@ mod tests {
"openai:chat",
"gpt-4.1",
None,
None,
)
.await;
@@ -1186,6 +1188,7 @@ mod tests {
"openai:chat",
"gpt-4.1",
None,
None,
)
.await;
@@ -1244,6 +1247,7 @@ mod tests {
"openai:chat",
"gpt-4.1",
None,
None,
)
.await;
@@ -1320,6 +1324,7 @@ mod tests {
"openai:chat",
"gpt-4.1",
None,
None,
)
.await;

45
apps/aether-gateway/src/ai_pipeline/planner/candidate_eligibility.rs
View File

@@ -6,6 +6,7 @@ use std::collections::BTreeSet;
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
use super::candidate_affinity::rank_eligible_local_execution_candidates;
use super::pool_scheduler::apply_local_execution_pool_scheduler;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
@@ -28,6 +29,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates(
client_api_format: &str,
requested_model: &str,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
@@ -37,6 +39,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates(
candidates,
client_api_format,
required_capabilities,
sticky_session_token,
|candidate, transport| {
current_local_execution_candidate_skip_reason_with_transport(
candidate,
@@ -55,6 +58,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport
client_api_format: &str,
requested_model: Option<&str>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
@@ -64,6 +68,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport
candidates,
client_api_format,
required_capabilities,
sticky_session_token,
|candidate, transport| {
current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
@@ -80,6 +85,7 @@ async fn filter_and_rank_local_execution_candidates_with_gate<F>(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
runtime_skip_reason: F,
) -> (
Vec<EligibleLocalExecutionCandidate>,
@@ -149,10 +155,49 @@ where
required_capabilities,
)
.await;
let (ranked, pool_skipped) =
apply_local_execution_pool_scheduler(state, ranked, sticky_session_token).await;
skipped.extend(pool_skipped);
(ranked, skipped)
}
pub(crate) fn extract_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
fn non_empty_string(value: Option<&serde_json::Value>) -> Option<String> {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
let object = body_json.as_object()?;
non_empty_string(object.get("prompt_cache_key"))
.or_else(|| non_empty_string(object.get("conversation_id")))
.or_else(|| non_empty_string(object.get("conversationId")))
.or_else(|| non_empty_string(object.get("session_id")))
.or_else(|| non_empty_string(object.get("sessionId")))
.or_else(|| {
object
.get("metadata")
.and_then(serde_json::Value::as_object)
.and_then(|metadata| {
non_empty_string(metadata.get("session_id"))
.or_else(|| non_empty_string(metadata.get("conversation_id")))
})
})
.or_else(|| {
object
.get("conversationState")
.and_then(serde_json::Value::as_object)
.and_then(|state| {
non_empty_string(state.get("conversationId"))
.or_else(|| non_empty_string(state.get("sessionId")))
})
})
}
fn current_local_execution_candidate_common_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,

2
apps/aether-gateway/src/ai_pipeline/planner/mod.rs
View File

@@ -17,6 +17,7 @@ mod materialization_policy;
mod passthrough;
mod payload_metadata;
mod plan_builders;
mod pool_scheduler;
mod report_context;
mod route;
mod runtime_miss;
@@ -25,6 +26,7 @@ mod specialized;
mod standard;
mod state;
pub(crate) use self::candidate_eligibility::extract_pool_sticky_session_token;
pub(crate) use self::passthrough::{
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
};

14
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/build.rs
View File

@@ -55,9 +55,10 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
.await?;
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
state, trace_id, &input, body_json, spec,
)
.await?;
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,
@@ -121,9 +122,10 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
.await?;
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
state, trace_id, &input, body_json, spec,
)
.await?;
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,

6
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/candidates.rs
View File

@@ -1,7 +1,8 @@
use tracing::warn;
use crate::ai_pipeline::planner::candidate_eligibility::{
filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
persist_available_local_execution_candidates_with_context,
@@ -85,10 +86,12 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
state: &AppState,
trace_id: &str,
input: &LocalSameFormatProviderDecisionInput,
body_json: &serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Result<(Vec<LocalSameFormatProviderCandidateAttempt>, usize), GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
@@ -110,6 +113,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
spec_metadata.api_format,
&input.requested_model,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)
.await;
let skipped_candidates = preselection_skipped

14
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/plans.rs
View File

@@ -57,9 +57,10 @@ pub(crate) async fn build_local_sync_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
.await?;
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
state, trace_id, &input, body_json, spec,
)
.await?;
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,
@@ -141,9 +142,10 @@ pub(crate) async fn build_local_stream_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
.await?;
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
state, trace_id, &input, body_json, spec,
)
.await?;
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,

1449
apps/aether-gateway/src/ai_pipeline/planner/pool_scheduler.rs Normal file
View File

File diff suppressed because it is too large Load Diff

1
apps/aether-gateway/src/ai_pipeline/planner/specialized/files/support.rs
View File

@@ -94,6 +94,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
GEMINI_FILES_CLIENT_API_FORMAT,
None,
input.required_capabilities.as_ref(),
None,
)
.await;

2
apps/aether-gateway/src/ai_pipeline/planner/specialized/video.rs
View File

@@ -60,6 +60,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
state,
trace_id,
&input,
body_json,
spec_metadata.api_format,
spec_metadata.decision_kind,
)
@@ -102,6 +103,7 @@ async fn build_local_sync_plan_and_reports(
state,
trace_id,
&input,
body_json,
spec_metadata.api_format,
spec_metadata.decision_kind,
)

8
apps/aether-gateway/src/ai_pipeline/planner/specialized/video/support.rs
View File

@@ -4,7 +4,8 @@ use tracing::warn;
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::{
filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
@@ -96,6 +97,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
state: &AppState,
trace_id: &str,
input: &LocalVideoCreateDecisionInput,
body_json: &serde_json::Value,
api_format: &str,
decision_kind: &str,
) -> Option<Vec<LocalVideoCreateCandidateAttempt>> {
@@ -128,6 +130,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
planner_state,
trace_id,
input,
body_json,
candidates,
preselection_skipped
.into_iter()
@@ -148,10 +151,12 @@ async fn materialize_local_video_create_candidate_attempts(
state: PlannerAppState<'_>,
trace_id: &str,
input: &LocalVideoCreateDecisionInput,
body_json: &serde_json::Value,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
api_format: &str,
) -> Vec<LocalVideoCreateCandidateAttempt> {
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
@@ -163,6 +168,7 @@ async fn materialize_local_video_create_candidate_attempts(
api_format,
&input.requested_model,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)
.await;
let skipped_candidates = preselection_skipped

12
apps/aether-gateway/src/ai_pipeline/planner/standard/family/build.rs
View File

@@ -51,7 +51,8 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
for attempt in attempts {
@@ -99,7 +100,8 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
for attempt in attempts {
@@ -153,7 +155,8 @@ pub(crate) async fn build_local_sync_plan_and_reports(
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
if candidate_count == 0 {
return Ok(Vec::new());
@@ -226,7 +229,8 @@ pub(crate) async fn build_local_stream_plan_and_reports(
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
if candidate_count == 0 {
return Ok(Vec::new());

6
apps/aether-gateway/src/ai_pipeline/planner/standard/family/candidates.rs
View File

@@ -4,7 +4,8 @@ use tracing::warn;
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::{
filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
persist_available_local_execution_candidates_with_context,
@@ -87,10 +88,12 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
state: &AppState,
trace_id: &str,
input: &LocalStandardDecisionInput,
body_json: &serde_json::Value,
spec: LocalStandardSpec,
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
@@ -176,6 +179,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
spec_metadata.api_format,
&input.requested_model,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)
.await;
let skipped_candidates = preselection_skipped

6
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/support.rs
View File

@@ -2,7 +2,8 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::{
filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
@@ -55,10 +56,12 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
state: &AppState,
trace_id: &str,
input: &LocalOpenAiChatDecisionInput,
body_json: &serde_json::Value,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
) -> Vec<LocalOpenAiChatCandidateAttempt> {
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
@@ -71,6 +74,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
"openai:chat",
&input.requested_model,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)
.await;
let skipped_candidates = preselection_skipped

2
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/mod.rs
View File

@@ -118,6 +118,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
@@ -183,6 +184,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)

1
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/plans/stream.rs
View File

@@ -80,6 +80,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)

1
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/plans/sync.rs
View File

@@ -80,6 +80,7 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)

6
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/decision/support.rs
View File

@@ -6,7 +6,8 @@ use tracing::warn;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::{
filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
@@ -82,11 +83,13 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
state: &AppState,
trace_id: &str,
input: &LocalOpenAiCliDecisionInput,
body_json: &serde_json::Value,
spec: LocalOpenAiCliSpec,
) -> Result<Vec<LocalOpenAiCliCandidateAttempt>, GatewayError> {
let spec_metadata = local_openai_cli_spec_metadata(spec);
let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
@@ -173,6 +176,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
spec_metadata.api_format,
&input.requested_model,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)
.await;
let skipped_candidates = preselection_skipped

6
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/mod.rs
View File

@@ -66,7 +66,8 @@ pub(crate) async fn maybe_build_sync_local_openai_cli_decision_payload(
};
let attempts =
materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, spec).await?;
materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, body_json, spec)
.await?;
for attempt in attempts {
if let Some(payload) = maybe_build_local_openai_cli_decision_payload_for_candidate(
@@ -100,7 +101,8 @@ pub(crate) async fn maybe_build_stream_local_openai_cli_decision_payload(
};
let attempts =
materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, spec).await?;
materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, body_json, spec)
.await?;
for attempt in attempts {
if let Some(payload) = maybe_build_local_openai_cli_decision_payload_for_candidate(

6
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/plans.rs
View File

@@ -33,7 +33,8 @@ pub(super) async fn build_local_sync_plan_and_reports(
};
let attempts =
materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, spec).await?;
materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, body_json, spec)
.await?;
let mut plans = Vec::new();
for attempt in attempts {
@@ -78,7 +79,8 @@ pub(super) async fn build_local_stream_plan_and_reports(
};
let attempts =
materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, spec).await?;
materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, body_json, spec)
.await?;
let mut plans = Vec::new();
for attempt in attempts {

5
apps/aether-gateway/src/execution_runtime/mod.rs
View File

@@ -6,6 +6,7 @@ use serde_json::{Map, Value};
mod constants;
mod fallback;
pub(crate) mod ndjson;
mod pool_feedback;
#[cfg(test)]
pub(crate) mod remote_compat;
mod server;
@@ -28,6 +29,10 @@ pub(crate) use self::fallback::{
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
LocalFailoverDecision,
};
pub(crate) use pool_feedback::{
record_pool_error_feedback, record_pool_stream_timeout_feedback,
record_stream_pool_success_feedback, record_sync_pool_success_feedback,
};
pub use server::{
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
build_execution_runtime_router_with_request_gates, serve_execution_runtime_tcp,

188
apps/aether-gateway/src/execution_runtime/pool_feedback.rs Normal file
View File

@@ -0,0 +1,188 @@
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_usage_runtime::{
build_stream_terminal_usage_outcome, build_sync_terminal_usage_outcome, TerminalUsageOutcome,
};
use serde_json::Value;
use std::collections::BTreeMap;
use tracing::warn;
use crate::ai_pipeline::extract_pool_sticky_session_token;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::handlers::shared::provider_pool::{
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success, AdminProviderPoolConfig,
};
use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest};
use crate::AppState;
struct PoolFeedbackContext {
runner: aether_data::redis::RedisKvRunner,
pool_config: AdminProviderPoolConfig,
sticky_session_token: Option<String>,
}
fn pool_feedback_request_body<'a>(
plan: &'a ExecutionPlan,
report_context: Option<&'a Value>,
) -> Option<&'a Value> {
report_context
.and_then(Value::as_object)
.and_then(|object| object.get("original_request_body"))
.filter(|value| !value.is_null())
.or(plan.body.json_body.as_ref())
}
async fn resolve_pool_feedback_context(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
) -> Option<PoolFeedbackContext> {
let Some(runner) = state.redis_kv_runner() else {
return None;
};
let transport = match state
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
.await
{
Ok(Some(transport)) => transport,
Ok(None) => return None,
Err(err) => {
warn!(
"gateway execution runtime pool feedback: failed to read transport snapshot for provider {} endpoint {} key {}: {:?}",
plan.provider_id, plan.endpoint_id, plan.key_id, err
);
return None;
}
};
let Some(pool_config) =
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref())
else {
return None;
};
let sticky_session_token = pool_feedback_request_body(plan, report_context)
.and_then(extract_pool_sticky_session_token);
Some(PoolFeedbackContext {
runner,
pool_config,
sticky_session_token,
})
}
fn total_tokens_used(outcome: &TerminalUsageOutcome) -> u64 {
outcome
.standardized_usage
.as_ref()
.map(|usage| {
usage
.input_tokens
.saturating_add(usage.output_tokens)
.max(0) as u64
})
.unwrap_or(0)
}
fn resolve_ttfb_ms(telemetry: Option<&ExecutionTelemetry>) -> Option<u64> {
telemetry.and_then(|telemetry| telemetry.ttfb_ms.or(telemetry.elapsed_ms))
}
pub(crate) async fn record_sync_pool_success_feedback(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
payload: &GatewaySyncReportRequest,
) {
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
return;
};
let usage_outcome = build_sync_terminal_usage_outcome(plan, report_context, payload);
record_admin_provider_pool_success(
&context.runner,
&plan.provider_id,
&plan.key_id,
&context.pool_config,
context.sticky_session_token.as_deref(),
total_tokens_used(&usage_outcome),
resolve_ttfb_ms(payload.telemetry.as_ref()),
)
.await;
}
pub(crate) async fn record_stream_pool_success_feedback(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
payload: &GatewayStreamReportRequest,
) {
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
return;
};
let usage_outcome = build_stream_terminal_usage_outcome(plan, report_context, payload);
record_admin_provider_pool_success(
&context.runner,
&plan.provider_id,
&plan.key_id,
&context.pool_config,
context.sticky_session_token.as_deref(),
total_tokens_used(&usage_outcome),
resolve_ttfb_ms(payload.telemetry.as_ref()),
)
.await;
}
pub(crate) async fn record_pool_error_feedback(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
status_code: u16,
headers: &BTreeMap<String, String>,
error_body: Option<&str>,
) {
if status_code < 400 {
return;
}
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
return;
};
if status_code == 401 {
let _ = state
.invalidate_local_oauth_refresh_entry(&plan.key_id)
.await;
}
record_admin_provider_pool_error(
&context.runner,
&plan.provider_id,
&plan.key_id,
&context.pool_config,
status_code,
error_body,
Some(headers),
)
.await;
}
pub(crate) async fn record_pool_stream_timeout_feedback(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
) {
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
return;
};
record_admin_provider_pool_stream_timeout(
&context.runner,
&plan.provider_id,
&plan.key_id,
&context.pool_config,
)
.await;
}

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

@@ -53,7 +53,8 @@ use crate::execution_runtime::transport::{
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
};
use crate::execution_runtime::{
local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
local_failover_response_text, record_pool_error_feedback, record_stream_pool_success_feedback,
resolve_core_stream_direct_finalize_report_kind,
resolve_core_stream_error_finalize_report_kind,
resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
should_retry_next_local_candidate_stream, LocalFailoverDecision,
@@ -565,6 +566,15 @@ async fn execute_stream_from_frame_stream(
let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
let error_response_text =
local_failover_response_text(body_json.as_ref(), &error_body, None);
record_pool_error_feedback(
state,
&plan,
report_context.as_ref(),
status_code,
&headers,
error_response_text.as_deref(),
)
.await;
let failover_decision = resolve_local_candidate_failover_decision_stream(
state,
&plan,
@@ -1492,6 +1502,13 @@ async fn execute_stream_from_frame_stream(
}),
telemetry: telemetry.clone(),
};
record_stream_pool_success_feedback(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
&usage_payload,
)
.await;
record_stream_terminal_usage(
&state_for_report,
&plan_for_report,

17
apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs
View File

@@ -16,6 +16,7 @@ use crate::control::GatewayControlDecision;
use crate::execution_runtime::submission::{
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
};
use crate::execution_runtime::{record_pool_error_feedback, record_pool_stream_timeout_feedback};
use crate::log_ids::short_request_id;
use crate::request_candidate_runtime::record_report_request_candidate_status;
use crate::usage::submit_sync_report;
@@ -122,6 +123,22 @@ async fn record_stream_sync_failure(
failure: &StreamFailureReport,
started_at_unix_ms: Option<u64>,
) {
if matches!(
failure.error_type.as_str(),
"first_byte_timeout" | "read_timeout"
) {
record_pool_stream_timeout_feedback(state, plan, report_context).await;
}
let error_body = serde_json::to_string(&failure.body_json).ok();
record_pool_error_feedback(
state,
plan,
report_context,
failure.status_code,
&payload.headers,
error_body.as_deref(),
)
.await;
let context_seed = build_terminal_usage_context_seed(plan, report_context);
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
state

27
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -30,9 +30,10 @@ use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
use crate::execution_runtime::{
local_failover_response_text, resolve_core_sync_error_finalize_report_kind,
should_fallback_to_control_sync, should_finalize_sync_response,
should_retry_next_local_candidate_sync, should_stop_local_candidate_failover_sync,
local_failover_response_text, record_pool_error_feedback, record_sync_pool_success_feedback,
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
should_finalize_sync_response, should_retry_next_local_candidate_sync,
should_stop_local_candidate_failover_sync,
};
use crate::log_ids::short_request_id;
use crate::request_candidate_runtime::{
@@ -255,6 +256,17 @@ pub(crate) async fn execute_execution_runtime_sync(
local_failover_response_text.as_deref(),
)
.await;
if result.status_code >= 400 {
record_pool_error_feedback(
state,
&plan,
report_context.as_ref(),
result.status_code,
&headers,
local_failover_response_text.as_deref(),
)
.await;
}
if should_retry_next_local_candidate_sync(
state,
&plan,
@@ -393,6 +405,15 @@ pub(crate) async fn execute_execution_runtime_sync(
body_base64: body_base64.clone(),
telemetry: result.telemetry.clone(),
};
if result.status_code < 400 {
record_sync_pool_success_feedback(
state,
&plan,
report_context.as_ref(),
&base_usage_payload,
)
.await;
}
if let Some(finalize_report_kind) = finalize_report_kind {
if let Some(implicit_finalize) = implicit_finalize {

293
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -1,4 +1,6 @@
use crate::handlers::admin::provider::shared::support::AdminProviderPoolConfig;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
};
use serde_json::{Map, Value};
fn json_u64(value: &Value) -> Option<u64> {
@@ -7,7 +9,115 @@ fn json_u64(value: &Value) -> Option<u64> {
.or_else(|| value.as_i64().and_then(|raw| u64::try_from(raw).ok()))
}
fn admin_provider_pool_lru_enabled(raw_pool_advanced: &Map<String, Value>) -> bool {
fn parse_pool_scheduling_presets(
raw_pool_advanced: &Map<String, Value>,
) -> Vec<AdminProviderPoolSchedulingPreset> {
let Some(presets) = raw_pool_advanced
.get("scheduling_presets")
.and_then(Value::as_array)
else {
return raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.filter(|enabled| *enabled)
.map(|_| {
vec![AdminProviderPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: true,
mode: None,
}]
})
.unwrap_or_default();
};
let mut normalized = Vec::new();
for item in presets {
if let Some(preset) = item
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
{
normalized.push(AdminProviderPoolSchedulingPreset {
preset: preset.to_ascii_lowercase(),
enabled: true,
mode: None,
});
continue;
}
let Some(object) = item.as_object() else {
continue;
};
let Some(preset) = object
.get("preset")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
normalized.push(AdminProviderPoolSchedulingPreset {
preset: preset.to_ascii_lowercase(),
enabled: object
.get("enabled")
.and_then(Value::as_bool)
.unwrap_or(true),
mode: object
.get("mode")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
});
}
if normalized.is_empty()
&& raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.unwrap_or(false)
{
normalized.push(AdminProviderPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: true,
mode: None,
});
}
normalized
}
fn parse_pool_unschedulable_rules(
raw_pool_advanced: &Map<String, Value>,
) -> Vec<AdminProviderPoolUnschedulableRule> {
raw_pool_advanced
.get("unschedulable_rules")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|item| {
let object = item.as_object()?;
let keyword = object
.get("keyword")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
Some(AdminProviderPoolUnschedulableRule {
keyword: keyword.to_string(),
duration_minutes: object
.get("duration_minutes")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(5),
})
})
.collect()
}
fn admin_provider_pool_lru_enabled(
raw_pool_advanced: &Map<String, Value>,
scheduling_presets: &[AdminProviderPoolSchedulingPreset],
) -> bool {
if let Some(explicit) = raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
@@ -15,56 +125,70 @@ fn admin_provider_pool_lru_enabled(raw_pool_advanced: &Map<String, Value>) -> bo
return explicit;
}
let Some(presets) = raw_pool_advanced
.get("scheduling_presets")
.and_then(Value::as_array)
else {
return false;
};
let Some(first) = presets.first() else {
return false;
};
if first.is_string() {
return raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.unwrap_or(true);
}
presets.iter().filter_map(Value::as_object).any(|item| {
item.get("preset")
.and_then(Value::as_str)
.is_some_and(|preset| preset.eq_ignore_ascii_case("lru"))
&& item.get("enabled").and_then(Value::as_bool).unwrap_or(true)
})
scheduling_presets
.iter()
.any(|item| item.enabled && item.preset.eq_ignore_ascii_case("lru"))
}
pub(crate) fn admin_provider_pool_config(
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
) -> Option<AdminProviderPoolConfig> {
let raw_pool_advanced = provider
.config
.as_ref()
admin_provider_pool_config_from_config_value(provider.config.as_ref())
}
pub(crate) fn admin_provider_pool_config_from_config_value(
config: Option<&serde_json::Value>,
) -> Option<AdminProviderPoolConfig> {
let raw_pool_advanced = config
.and_then(Value::as_object)
.and_then(|config| config.get("pool_advanced"))?;
let Some(pool_advanced) = raw_pool_advanced.as_object() else {
return Some(AdminProviderPoolConfig {
scheduling_presets: Vec::new(),
unschedulable_rules: Vec::new(),
lru_enabled: false,
skip_exhausted_accounts: false,
sticky_session_ttl_seconds: 3600,
latency_window_seconds: 3600,
latency_sample_limit: 50,
cost_window_seconds: 18_000,
cost_limit_per_key_tokens: None,
rate_limit_cooldown_seconds: 300,
overload_cooldown_seconds: 30,
health_policy_enabled: true,
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,
});
};
let scheduling_presets = parse_pool_scheduling_presets(pool_advanced);
let unschedulable_rules = parse_pool_unschedulable_rules(pool_advanced);
Some(AdminProviderPoolConfig {
lru_enabled: admin_provider_pool_lru_enabled(pool_advanced),
lru_enabled: admin_provider_pool_lru_enabled(pool_advanced, &scheduling_presets),
scheduling_presets,
unschedulable_rules,
skip_exhausted_accounts: pool_advanced
.get("skip_exhausted_accounts")
.and_then(Value::as_bool)
.unwrap_or(false),
sticky_session_ttl_seconds: pool_advanced
.get("sticky_session_ttl_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(3600),
latency_window_seconds: pool_advanced
.get("latency_window_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(3600),
latency_sample_limit: pool_advanced
.get("latency_sample_limit")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(50),
cost_window_seconds: pool_advanced
.get("cost_window_seconds")
.and_then(json_u64)
@@ -73,12 +197,41 @@ pub(crate) fn admin_provider_pool_config(
cost_limit_per_key_tokens: pool_advanced
.get("cost_limit_per_key_tokens")
.and_then(json_u64),
rate_limit_cooldown_seconds: pool_advanced
.get("rate_limit_cooldown_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(300),
overload_cooldown_seconds: pool_advanced
.get("overload_cooldown_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(30),
health_policy_enabled: pool_advanced
.get("health_policy_enabled")
.and_then(Value::as_bool)
.unwrap_or(true),
stream_timeout_threshold: pool_advanced
.get("stream_timeout_threshold")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(3),
stream_timeout_window_seconds: pool_advanced
.get("stream_timeout_window_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(1800),
stream_timeout_cooldown_seconds: pool_advanced
.get("stream_timeout_cooldown_seconds")
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(300),
})
}
#[cfg(test)]
mod tests {
use super::admin_provider_pool_config;
use super::{admin_provider_pool_config, admin_provider_pool_config_from_config_value};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use serde_json::json;
@@ -117,15 +270,91 @@ mod tests {
"pool_advanced": {
"skip_exhausted_accounts": true,
"lru_enabled": true,
"sticky_session_ttl_seconds": 600,
"latency_window_seconds": 900,
"latency_sample_limit": 75,
"cost_window_seconds": 7200,
"cost_limit_per_key_tokens": 12000
"cost_limit_per_key_tokens": 12000,
"rate_limit_cooldown_seconds": 420,
"overload_cooldown_seconds": 45,
"health_policy_enabled": false,
"stream_timeout_threshold": 4,
"stream_timeout_window_seconds": 900,
"stream_timeout_cooldown_seconds": 180
}
}));
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
assert!(config.skip_exhausted_accounts);
assert!(config.lru_enabled);
assert_eq!(config.sticky_session_ttl_seconds, 600);
assert_eq!(config.latency_window_seconds, 900);
assert_eq!(config.latency_sample_limit, 75);
assert_eq!(config.cost_window_seconds, 7200);
assert_eq!(config.cost_limit_per_key_tokens, Some(12_000));
assert_eq!(config.rate_limit_cooldown_seconds, 420);
assert_eq!(config.overload_cooldown_seconds, 45);
assert!(!config.health_policy_enabled);
assert_eq!(config.stream_timeout_threshold, 4);
assert_eq!(config.stream_timeout_window_seconds, 900);
assert_eq!(config.stream_timeout_cooldown_seconds, 180);
}
#[test]
fn parses_pool_config_from_generic_config_value() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [{"preset": "lru", "enabled": true}],
"cost_limit_per_key_tokens": 4096
}
})))
.expect("pool config should parse");
assert!(config.lru_enabled);
assert_eq!(config.scheduling_presets.len(), 1);
assert_eq!(config.scheduling_presets[0].preset, "lru");
assert_eq!(config.cost_limit_per_key_tokens, Some(4096));
}
#[test]
fn parses_object_style_scheduling_presets_with_modes() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "cache_affinity", "enabled": false},
{"preset": "plus_first", "enabled": true, "mode": "plus_only"}
]
}
})))
.expect("pool config should parse");
assert!(!config.lru_enabled);
assert_eq!(config.scheduling_presets.len(), 2);
assert_eq!(config.scheduling_presets[0].preset, "cache_affinity");
assert!(!config.scheduling_presets[0].enabled);
assert_eq!(config.scheduling_presets[1].preset, "plus_first");
assert_eq!(
config.scheduling_presets[1].mode.as_deref(),
Some("plus_only")
);
}
#[test]
fn parses_unschedulable_rules_from_pool_advanced() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"unschedulable_rules": [
{"keyword": "suspended", "duration_minutes": 15},
{"keyword": "review_required"}
]
}
})))
.expect("pool config should parse");
assert_eq!(config.unschedulable_rules.len(), 2);
assert_eq!(config.unschedulable_rules[0].keyword, "suspended");
assert_eq!(config.unschedulable_rules[0].duration_minutes, 15);
assert_eq!(config.unschedulable_rules[1].keyword, "review_required");
assert_eq!(config.unschedulable_rules[1].duration_minutes, 5);
}
}

42
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/keys.rs
View File

@@ -4,6 +4,14 @@ pub(super) fn pool_sticky_pattern(keyspace: &RedisKeyspace, provider_id: &str) -
keyspace.key(&format!("ap:{provider_id}:sticky:*"))
}
pub(super) fn pool_sticky_key(
keyspace: &RedisKeyspace,
provider_id: &str,
session_token: &str,
) -> String {
keyspace.key(&format!("ap:{provider_id}:sticky:{session_token}"))
}
pub(super) fn pool_lru_key(keyspace: &RedisKeyspace, provider_id: &str) -> String {
keyspace.key(&format!("ap:{provider_id}:lru"))
}
@@ -24,6 +32,22 @@ pub(super) fn pool_cost_key(keyspace: &RedisKeyspace, provider_id: &str, key_id:
keyspace.key(&format!("ap:{provider_id}:cost:{key_id}"))
}
pub(super) fn pool_latency_key(
keyspace: &RedisKeyspace,
provider_id: &str,
key_id: &str,
) -> String {
keyspace.key(&format!("ap:{provider_id}:latency:{key_id}"))
}
pub(super) fn pool_stream_timeout_key(
keyspace: &RedisKeyspace,
provider_id: &str,
key_id: &str,
) -> String {
keyspace.key(&format!("ap:{provider_id}:stream_timeout:{key_id}"))
}
pub(super) fn parse_pool_cost_member(member: &str) -> u64 {
member
.rsplit_once(':')
@@ -31,6 +55,13 @@ pub(super) fn parse_pool_cost_member(member: &str) -> u64 {
.unwrap_or(0)
}
pub(super) fn parse_pool_latency_member(member: &str) -> u64 {
member
.rsplit_once(':')
.and_then(|(_, suffix)| suffix.parse::<u64>().ok())
.unwrap_or(0)
}
pub(super) fn pool_cooldown_keys(
keyspace: &RedisKeyspace,
provider_id: &str,
@@ -52,3 +83,14 @@ pub(super) fn pool_cost_keys(
.map(|key_id| pool_cost_key(keyspace, provider_id, key_id))
.collect()
}
pub(super) fn pool_latency_keys(
keyspace: &RedisKeyspace,
provider_id: &str,
key_ids: &[String],
) -> Vec<String> {
key_ids
.iter()
.map(|key_id| pool_latency_key(keyspace, provider_id, key_id))
.collect()
}

5
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/mod.rs
View File

@@ -2,6 +2,7 @@ mod keys;
mod mutations;
mod reads;
mod status;
mod writes;
pub(crate) use self::mutations::{
clear_admin_provider_pool_cooldown, reset_admin_provider_pool_cost,
@@ -11,3 +12,7 @@ pub(crate) use self::reads::{
read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_runtime_state,
};
pub(crate) use self::status::build_admin_provider_pool_status_payload;
pub(crate) use self::writes::{
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success,
};

119
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -1,6 +1,7 @@
use super::keys::{
parse_pool_cost_member, pool_cooldown_index_key, pool_cooldown_keys, pool_cost_keys,
pool_lru_key, pool_sticky_pattern,
parse_pool_cost_member, parse_pool_latency_member, pool_cooldown_index_key, pool_cooldown_key,
pool_cooldown_keys, pool_cost_keys, pool_latency_keys, pool_lru_key, pool_sticky_key,
pool_sticky_pattern,
};
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, ADMIN_PROVIDER_POOL_SCAN_BATCH,
@@ -11,6 +12,13 @@ use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
async fn scan_redis_keys(
connection: &mut redis::aio::MultiplexedConnection,
pattern: &str,
@@ -77,7 +85,8 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
runner: &RedisKvRunner,
provider_id: &str,
key_ids: &[String],
pool_config: AdminProviderPoolConfig,
pool_config: &AdminProviderPoolConfig,
sticky_session_token: Option<&str>,
) -> AdminProviderPoolRuntimeState {
let mut runtime = AdminProviderPoolRuntimeState::default();
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
@@ -87,6 +96,57 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
let keyspace = runner.keyspace().clone();
let cooldown_keys = pool_cooldown_keys(&keyspace, provider_id, key_ids);
let cost_keys = pool_cost_keys(&keyspace, provider_id, key_ids);
let latency_keys = pool_latency_keys(&keyspace, provider_id, key_ids);
if let Some(sticky_session_token) = sticky_session_token
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|_| pool_config.sticky_session_ttl_seconds > 0)
{
let sticky_key = pool_sticky_key(&keyspace, provider_id, sticky_session_token);
let sticky_bound_key_id = redis::cmd("GET")
.arg(&sticky_key)
.query_async::<Option<String>>(&mut connection)
.await
.unwrap_or_else(|err| {
warn!(
"gateway admin provider pool: failed to read sticky binding for provider {provider_id}: {:?}",
err
);
None
});
if let Some(bound_key_id) = sticky_bound_key_id {
let cooldown_key = pool_cooldown_key(&keyspace, provider_id, &bound_key_id);
runtime.sticky_bound_key_id = match redis::cmd("EXISTS")
.arg(&cooldown_key)
.query_async::<u64>(&mut connection)
.await
{
Ok(0) => {
let _: Result<bool, _> = redis::cmd("EXPIRE")
.arg(&sticky_key)
.arg(pool_config.sticky_session_ttl_seconds)
.query_async(&mut connection)
.await;
Some(bound_key_id)
}
Ok(_) => {
let _: Result<i64, _> = redis::cmd("DEL")
.arg(&sticky_key)
.query_async(&mut connection)
.await;
None
}
Err(err) => {
warn!(
"gateway admin provider pool: failed to validate sticky cooldown for provider {provider_id}: {:?}",
err
);
Some(bound_key_id)
}
};
}
}
let sticky_keys = match scan_redis_keys(
&mut connection,
@@ -174,11 +234,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
if !cost_keys.is_empty() {
let window_start = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
.saturating_sub(pool_config.cost_window_seconds);
let window_start = current_unix_secs().saturating_sub(pool_config.cost_window_seconds);
let mut cost_pipeline = redis::pipe();
for cost_key in &cost_keys {
cost_pipeline
@@ -208,7 +264,52 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
}
if pool_config.lru_enabled && !key_ids.is_empty() {
if !latency_keys.is_empty() {
let window_start = current_unix_secs().saturating_sub(pool_config.latency_window_seconds);
let mut latency_pipeline = redis::pipe();
for latency_key in &latency_keys {
latency_pipeline
.cmd("ZRANGEBYSCORE")
.arg(latency_key)
.arg(window_start)
.arg("+inf");
}
let members_by_key = latency_pipeline
.query_async::<Vec<Vec<String>>>(&mut connection)
.await
.unwrap_or_else(|err| {
warn!(
"gateway admin provider pool: failed to batch read latency windows for provider {provider_id}: {:?}",
err
);
vec![Vec::new(); latency_keys.len()]
});
for (key_id, members) in key_ids.iter().zip(members_by_key) {
let samples = members
.iter()
.map(|member| parse_pool_latency_member(member))
.filter(|value| *value > 0)
.collect::<Vec<_>>();
if samples.is_empty() {
continue;
}
let total = samples.iter().sum::<u64>() as f64;
let average = total / samples.len() as f64;
if average.is_finite() && average >= 0.0 {
runtime
.latency_avg_ms_by_key
.insert(key_id.clone(), average);
}
}
}
if (pool_config.lru_enabled
|| pool_config
.scheduling_presets
.iter()
.any(|item| item.enabled))
&& !key_ids.is_empty()
{
let mut command = redis::cmd("ZMSCORE");
command.arg(pool_lru_key(&keyspace, provider_id));
for key_id in key_ids {

11
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/status.rs
View File

@@ -36,8 +36,14 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let runtime = match state.redis_kv_runner() {
Some(runner) => {
read_admin_provider_pool_runtime_state(&runner, &provider.id, &key_ids, pool_config)
.await
read_admin_provider_pool_runtime_state(
&runner,
&provider.id,
&key_ids,
&pool_config,
None,
)
.await
}
None => AdminProviderPoolRuntimeState::default(),
};
@@ -56,6 +62,7 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
"cost_window_usage": runtime.cost_window_usage_by_key.get(&key.id).copied().unwrap_or(0),
"cost_limit": pool_config.cost_limit_per_key_tokens,
"sticky_sessions": runtime.sticky_sessions_by_key.get(&key.id).copied().unwrap_or(0),
"latency_avg_ms": runtime.latency_avg_ms_by_key.get(&key.id).copied(),
"lru_score": runtime.lru_score_by_key.get(&key.id).copied(),
})
})

910
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs Normal file
View File

@@ -0,0 +1,910 @@
use super::keys::{
pool_cooldown_index_key, pool_cooldown_key, pool_cost_key, pool_latency_key, pool_lru_key,
pool_sticky_key, pool_stream_timeout_key,
};
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolUnschedulableRule,
};
use aether_data::redis::RedisKvRunner;
use regex::Regex;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
use uuid::Uuid;
const ACCOUNT_DISABLE_PATTERNS: &[&str] = &[
"organization has been disabled",
"organization_disabled",
"account has been disabled",
"account_disabled",
"account has been deactivated",
"account_deactivated",
"account deactivated",
];
const FORBIDDEN_ACCOUNT_PATTERNS: &[&str] = &[
"account suspended",
"account banned",
"account deactivated",
"subscription inactive",
"suspended",
"banned",
"deactivated",
];
fn current_unix_secs_f64() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs_f64()
}
fn enabled_pool_presets(pool_config: &AdminProviderPoolConfig) -> impl Iterator<Item = &str> {
pool_config
.scheduling_presets
.iter()
.filter(|item| item.enabled)
.map(|item| item.preset.as_str())
}
fn should_touch_lru(pool_config: &AdminProviderPoolConfig) -> bool {
pool_config.lru_enabled || enabled_pool_presets(pool_config).next().is_some()
}
fn should_record_latency(pool_config: &AdminProviderPoolConfig) -> bool {
enabled_pool_presets(pool_config).any(|preset| preset.eq_ignore_ascii_case("latency_first"))
}
fn oauth_cache_key(key_id: &str) -> String {
format!("provider_oauth_token_cache:{key_id}")
}
fn parse_retry_after_seconds(headers: Option<&BTreeMap<String, String>>) -> Option<u64> {
let raw = headers.and_then(|headers| {
headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case("retry-after"))
.map(|(_, value)| value.trim())
.filter(|value| !value.is_empty())
})?;
let seconds = raw.parse::<u64>().ok()?;
Some(seconds.clamp(1, 3600))
}
fn parse_google_quota_duration_seconds(raw: &serde_json::Value) -> Option<u64> {
match raw {
serde_json::Value::Number(number) => {
let seconds = number.as_f64()?;
Some((seconds.max(1.0).floor() as u64).max(1))
}
serde_json::Value::String(text) => {
let trimmed = text.trim().to_ascii_lowercase();
if trimmed.is_empty() {
return None;
}
let regex = Regex::new(r"(\d+(?:\.\d+)?)([dhms])").ok()?;
let mut total_seconds = 0.0;
let mut matched = false;
for capture in regex.captures_iter(&trimmed) {
let amount = capture.get(1)?.as_str().parse::<f64>().ok()?;
let unit = capture.get(2)?.as_str();
matched = true;
total_seconds += match unit {
"d" => amount * 86_400.0,
"h" => amount * 3_600.0,
"m" => amount * 60.0,
"s" => amount,
_ => 0.0,
};
}
matched.then(|| (total_seconds.max(1.0).floor() as u64).max(1))
}
_ => None,
}
}
fn parse_google_quota_cooldown_seconds_at(
error_body: Option<&str>,
now_unix_secs: u64,
) -> Option<u64> {
let error_body = error_body
.map(str::trim)
.filter(|value| !value.is_empty())?;
let payload = serde_json::from_str::<serde_json::Value>(error_body).ok()?;
let error = payload.get("error")?.as_object()?;
if let Some(details) = error.get("details").and_then(serde_json::Value::as_array) {
for detail in details {
let Some(metadata) = detail
.get("metadata")
.and_then(serde_json::Value::as_object)
else {
continue;
};
if let Some(reset_at_text) = metadata
.get("quotaResetTimeStamp")
.or_else(|| metadata.get("quotaResetTimestamp"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
if let Ok(parsed) = chrono::DateTime::parse_from_rfc3339(reset_at_text) {
let reset_at = parsed.timestamp().max(0) as u64;
return Some(reset_at.saturating_sub(now_unix_secs).max(1));
}
}
if let Some(delay) = metadata
.get("quotaResetDelay")
.and_then(parse_google_quota_duration_seconds)
{
return Some(delay.max(1));
}
}
}
let message = error
.get("message")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let regex = Regex::new(r"(?i)reset after\s+([^.,;]+)").ok()?;
let capture = regex.captures(message)?;
parse_google_quota_duration_seconds(&serde_json::Value::String(
capture.get(1)?.as_str().to_string(),
))
}
fn parse_google_quota_cooldown_seconds(error_body: Option<&str>) -> Option<u64> {
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
parse_google_quota_cooldown_seconds_at(error_body, now_unix_secs)
}
fn extract_error_message(error_body: Option<&str>) -> String {
let Some(error_body) = error_body.map(str::trim).filter(|value| !value.is_empty()) else {
return String::new();
};
serde_json::from_str::<serde_json::Value>(error_body)
.ok()
.and_then(|value| {
value
.as_object()
.and_then(|object| object.get("error").or_else(|| object.get("message")))
.and_then(|error| match error {
serde_json::Value::Object(object) => object
.get("message")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned),
serde_json::Value::String(text) => Some(text.clone()),
_ => None,
})
})
.unwrap_or_else(|| error_body.chars().take(500).collect())
}
fn resolve_transient_cooldown_ttl(
status_code: u16,
retry_after_seconds: Option<u64>,
pool_config: &AdminProviderPoolConfig,
) -> u64 {
if matches!(status_code, 429 | 503) {
if let Some(retry_after_seconds) = retry_after_seconds {
return retry_after_seconds;
}
}
if status_code == 429 {
return pool_config.rate_limit_cooldown_seconds;
}
pool_config.overload_cooldown_seconds
}
async fn set_pool_cooldown(
runner: &RedisKvRunner,
provider_id: &str,
key_id: &str,
reason: &str,
ttl_seconds: u64,
) {
if ttl_seconds == 0 {
return;
}
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
warn!(
"gateway admin provider pool: failed to connect redis to set cooldown for key {key_id}"
);
return;
};
let keyspace = runner.keyspace().clone();
let result: Result<(), _> = redis::pipe()
.cmd("SETEX")
.arg(pool_cooldown_key(&keyspace, provider_id, key_id))
.arg(ttl_seconds)
.arg(reason)
.ignore()
.cmd("SADD")
.arg(pool_cooldown_index_key(&keyspace, provider_id))
.arg(key_id)
.ignore()
.cmd("EXPIRE")
.arg(pool_cooldown_index_key(&keyspace, provider_id))
.arg(ttl_seconds.saturating_add(60))
.ignore()
.query_async(&mut connection)
.await;
if let Err(err) = result {
warn!(
"gateway admin provider pool: failed to set cooldown for provider {provider_id} key {key_id}: {:?}",
err
);
}
}
async fn invalidate_pool_oauth_cache(runner: &RedisKvRunner, key_id: &str) {
if let Err(err) = runner.del(&oauth_cache_key(key_id)).await {
warn!(
"gateway admin provider pool: failed to invalidate oauth cache for key {key_id}: {:?}",
err
);
}
}
fn matching_unschedulable_rule<'a>(
rules: &'a [AdminProviderPoolUnschedulableRule],
error_message: &str,
) -> Option<&'a AdminProviderPoolUnschedulableRule> {
rules.iter().find(|rule| {
let keyword = rule.keyword.trim().to_ascii_lowercase();
!keyword.is_empty() && error_message.contains(keyword.as_str())
})
}
pub(crate) async fn record_admin_provider_pool_success(
runner: &RedisKvRunner,
provider_id: &str,
key_id: &str,
pool_config: &AdminProviderPoolConfig,
sticky_session_token: Option<&str>,
tokens_used: u64,
ttfb_ms: Option<u64>,
) {
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
warn!("gateway admin provider pool: failed to connect redis to record success for key {key_id}");
return;
};
let keyspace = runner.keyspace().clone();
let now = current_unix_secs_f64();
let mut pipeline = redis::pipe();
let mut has_commands = false;
if let Some(sticky_session_token) = sticky_session_token
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|_| pool_config.sticky_session_ttl_seconds > 0)
{
pipeline
.cmd("SETEX")
.arg(pool_sticky_key(
&keyspace,
provider_id,
sticky_session_token,
))
.arg(pool_config.sticky_session_ttl_seconds)
.arg(key_id)
.ignore();
has_commands = true;
}
if should_touch_lru(pool_config) {
pipeline
.cmd("ZADD")
.arg(pool_lru_key(&keyspace, provider_id))
.arg(now)
.arg(key_id)
.ignore();
has_commands = true;
}
if tokens_used > 0 && pool_config.cost_limit_per_key_tokens.is_some() {
let cost_key = pool_cost_key(&keyspace, provider_id, key_id);
let window_seconds = pool_config.cost_window_seconds.max(1);
let member = format!("{}:{tokens_used}", Uuid::new_v4().simple());
pipeline
.cmd("ZADD")
.arg(&cost_key)
.arg(now)
.arg(member)
.ignore()
.cmd("ZREMRANGEBYSCORE")
.arg(&cost_key)
.arg("-inf")
.arg(now - window_seconds as f64)
.ignore()
.cmd("EXPIRE")
.arg(&cost_key)
.arg(window_seconds.saturating_add(600))
.ignore();
has_commands = true;
}
if let Some(ttfb_ms) = ttfb_ms
.filter(|value| should_record_latency(pool_config))
.filter(|_| pool_config.latency_window_seconds > 0)
{
let latency_key = pool_latency_key(&keyspace, provider_id, key_id);
let window_seconds = pool_config.latency_window_seconds.max(1);
let sample_limit = pool_config.latency_sample_limit.max(1);
let member = format!("{}:{ttfb_ms}", Uuid::new_v4().simple());
pipeline
.cmd("ZADD")
.arg(&latency_key)
.arg(now)
.arg(member)
.ignore()
.cmd("ZREMRANGEBYSCORE")
.arg(&latency_key)
.arg("-inf")
.arg(now - window_seconds as f64)
.ignore()
.cmd("ZREMRANGEBYRANK")
.arg(&latency_key)
.arg(0)
.arg(-((sample_limit as i64) + 1))
.ignore()
.cmd("EXPIRE")
.arg(&latency_key)
.arg(window_seconds.saturating_add(600))
.ignore();
has_commands = true;
}
if !has_commands {
return;
}
let result: Result<(), _> = pipeline.query_async(&mut connection).await;
if let Err(err) = result {
warn!(
"gateway admin provider pool: failed to record success feedback for provider {provider_id} key {key_id}: {:?}",
err
);
}
}
pub(crate) async fn record_admin_provider_pool_error(
runner: &RedisKvRunner,
provider_id: &str,
key_id: &str,
pool_config: &AdminProviderPoolConfig,
status_code: u16,
error_body: Option<&str>,
response_headers: Option<&BTreeMap<String, String>>,
) {
if !pool_config.health_policy_enabled {
return;
}
let error_message = extract_error_message(error_body).to_ascii_lowercase();
if status_code == 401 {
invalidate_pool_oauth_cache(runner, key_id).await;
if ACCOUNT_DISABLE_PATTERNS
.iter()
.any(|pattern| error_message.contains(pattern))
{
set_pool_cooldown(runner, provider_id, key_id, "account_deactivated_401", 3600).await;
}
return;
}
if status_code == 402 {
set_pool_cooldown(runner, provider_id, key_id, "payment_required_402", 3600).await;
return;
}
if status_code == 403 {
let severe = FORBIDDEN_ACCOUNT_PATTERNS
.iter()
.any(|pattern| error_message.contains(pattern));
let ttl_seconds = if severe {
3600
} else {
pool_config.rate_limit_cooldown_seconds.max(300)
};
set_pool_cooldown(runner, provider_id, key_id, "forbidden_403", ttl_seconds).await;
return;
}
if status_code == 400 {
if let Some(pattern) = ACCOUNT_DISABLE_PATTERNS
.iter()
.find(|pattern| error_message.contains(**pattern))
{
set_pool_cooldown(
runner,
provider_id,
key_id,
&format!("account_disabled_400:{pattern}"),
3600,
)
.await;
return;
}
}
if let Some(rule) =
matching_unschedulable_rule(&pool_config.unschedulable_rules, &error_message)
{
let ttl_seconds = (rule.duration_minutes.max(1)).saturating_mul(60).max(60);
set_pool_cooldown(
runner,
provider_id,
key_id,
&format!("rule:{}", rule.keyword),
ttl_seconds,
)
.await;
return;
}
if status_code == 429 {
let ttl_seconds = resolve_transient_cooldown_ttl(
status_code,
parse_retry_after_seconds(response_headers)
.or_else(|| parse_google_quota_cooldown_seconds(error_body)),
pool_config,
);
set_pool_cooldown(runner, provider_id, key_id, "rate_limited_429", ttl_seconds).await;
return;
}
if status_code == 529 {
set_pool_cooldown(
runner,
provider_id,
key_id,
"overloaded_529",
pool_config.overload_cooldown_seconds,
)
.await;
return;
}
let transient_reason = match status_code {
408 => Some("request_timeout_408".to_string()),
409 => Some("conflict_409".to_string()),
423 => Some("locked_423".to_string()),
425 => Some("too_early_425".to_string()),
500 => Some("server_error_500".to_string()),
502 => Some("bad_gateway_502".to_string()),
503 => Some("service_unavailable_503".to_string()),
504 => Some("gateway_timeout_504".to_string()),
501 | 505..=599 => Some(format!("server_error_{status_code}")),
_ => None,
};
if let Some(reason) = transient_reason {
let ttl_seconds = resolve_transient_cooldown_ttl(
status_code,
parse_retry_after_seconds(response_headers),
pool_config,
);
set_pool_cooldown(runner, provider_id, key_id, &reason, ttl_seconds).await;
}
}
pub(crate) async fn record_admin_provider_pool_stream_timeout(
runner: &RedisKvRunner,
provider_id: &str,
key_id: &str,
pool_config: &AdminProviderPoolConfig,
) {
if !pool_config.health_policy_enabled || pool_config.stream_timeout_threshold == 0 {
return;
}
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
warn!("gateway admin provider pool: failed to connect redis to record stream timeout for key {key_id}");
return;
};
let keyspace = runner.keyspace().clone();
let timeout_key = pool_stream_timeout_key(&keyspace, provider_id, key_id);
let now = current_unix_secs_f64();
let window_seconds = pool_config.stream_timeout_window_seconds.max(1);
let member = Uuid::new_v4().simple().to_string();
let results = redis::pipe()
.cmd("ZREMRANGEBYSCORE")
.arg(&timeout_key)
.arg("-inf")
.arg(now - window_seconds as f64)
.cmd("ZADD")
.arg(&timeout_key)
.arg(now)
.arg(member)
.cmd("ZCARD")
.arg(&timeout_key)
.cmd("EXPIRE")
.arg(&timeout_key)
.arg(window_seconds.saturating_add(60))
.query_async::<Vec<redis::Value>>(&mut connection)
.await;
let count = match results
.ok()
.and_then(|values| values.get(2).cloned())
.and_then(|value| redis::from_redis_value::<u64>(&value).ok())
{
Some(count) => count,
None => {
warn!(
"gateway admin provider pool: failed to compute stream timeout count for provider {provider_id} key {key_id}"
);
return;
}
};
if count >= pool_config.stream_timeout_threshold {
set_pool_cooldown(
runner,
provider_id,
key_id,
&format!("stream_timeout_x{count}"),
pool_config.stream_timeout_cooldown_seconds.max(1),
)
.await;
}
}
#[cfg(test)]
mod tests {
use super::{
parse_google_quota_cooldown_seconds_at, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
};
use crate::data::{GatewayDataConfig, GatewayDataState};
use crate::handlers::admin::provider::pool::runtime::reads::read_admin_provider_pool_runtime_state;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset,
AdminProviderPoolUnschedulableRule,
};
use crate::AppState;
use aether_testkit::ManagedRedisServer;
use std::collections::BTreeMap;
async fn start_managed_redis_or_skip() -> Option<ManagedRedisServer> {
match ManagedRedisServer::start().await {
Ok(server) => Some(server),
Err(err) if err.to_string().contains("No such file or directory") => {
eprintln!("skipping redis-backed pool runtime test: {err}");
None
}
Err(err) => panic!("redis server should start: {err}"),
}
}
fn sample_pool_config() -> AdminProviderPoolConfig {
AdminProviderPoolConfig {
scheduling_presets: vec![AdminProviderPoolSchedulingPreset {
preset: "latency_first".to_string(),
enabled: true,
mode: None,
}],
unschedulable_rules: Vec::new(),
lru_enabled: true,
skip_exhausted_accounts: false,
sticky_session_ttl_seconds: 120,
latency_window_seconds: 600,
latency_sample_limit: 10,
cost_window_seconds: 600,
cost_limit_per_key_tokens: Some(10_000),
rate_limit_cooldown_seconds: 300,
overload_cooldown_seconds: 30,
health_policy_enabled: true,
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,
}
}
fn build_runner_app(redis_url: &str, key_prefix: &str) -> AppState {
let data_state = GatewayDataState::from_config(
GatewayDataConfig::disabled().with_redis_url(redis_url, Some(key_prefix)),
)
.expect("data state should build");
AppState::new()
.expect("app state should build")
.with_data_state_for_tests(data_state)
}
#[test]
fn parses_google_quota_cooldown_from_reset_timestamp() {
let now_unix_secs = chrono::DateTime::parse_from_rfc3339("2026-04-17T10:00:00Z")
.expect("timestamp should parse")
.timestamp()
.max(0) as u64;
let cooldown = parse_google_quota_cooldown_seconds_at(
Some(
r#"{
"error": {
"message": "Quota exhausted.",
"details": [{
"metadata": {
"quotaResetTimeStamp": "2026-04-17T10:01:30Z"
}
}]
}
}"#,
),
now_unix_secs,
);
assert_eq!(cooldown, Some(90));
}
#[test]
fn parses_google_quota_cooldown_from_reset_delay_and_message() {
let now_unix_secs = chrono::DateTime::parse_from_rfc3339("2026-04-17T10:00:00Z")
.expect("timestamp should parse")
.timestamp()
.max(0) as u64;
let delay_cooldown = parse_google_quota_cooldown_seconds_at(
Some(
r#"{
"error": {
"message": "Quota exhausted.",
"details": [{
"metadata": {
"quotaResetDelay": "1h30m15s"
}
}]
}
}"#,
),
now_unix_secs,
);
let message_cooldown = parse_google_quota_cooldown_seconds_at(
Some(
r#"{
"error": {
"message": "Too many requests, reset after 45m."
}
}"#,
),
now_unix_secs,
);
assert_eq!(delay_cooldown, Some(5_415));
assert_eq!(message_cooldown, Some(2_700));
}
#[tokio::test]
async fn success_feedback_writes_sticky_lru_cost_and_latency() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_success_feedback");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let pool_config = sample_pool_config();
let key_ids = vec!["key-1".to_string()];
record_admin_provider_pool_success(
&runner,
"provider-1",
"key-1",
&pool_config,
Some("session-1"),
120,
Some(80),
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
Some("session-1"),
)
.await;
assert_eq!(runtime.total_sticky_sessions, 1);
assert_eq!(runtime.sticky_bound_key_id.as_deref(), Some("key-1"));
assert_eq!(runtime.sticky_sessions_by_key.get("key-1"), Some(&1));
assert_eq!(runtime.cost_window_usage_by_key.get("key-1"), Some(&120));
assert_eq!(runtime.latency_avg_ms_by_key.get("key-1"), Some(&80.0));
assert!(runtime.lru_score_by_key.contains_key("key-1"));
}
#[tokio::test]
async fn error_feedback_respects_retry_after_for_rate_limits() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_error_feedback");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let pool_config = sample_pool_config();
let key_ids = vec!["key-2".to_string()];
record_admin_provider_pool_error(
&runner,
"provider-1",
"key-2",
&pool_config,
429,
Some(r#"{"error":{"message":"rate limited"}}"#),
Some(&BTreeMap::from([(
"Retry-After".to_string(),
"120".to_string(),
)])),
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-2")
.map(String::as_str),
Some("rate_limited_429")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-2")
.is_some_and(|ttl| *ttl <= 120 && *ttl >= 100));
}
#[tokio::test]
async fn error_feedback_uses_google_quota_cooldown_when_retry_after_missing() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_google_quota_cooldown");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let pool_config = sample_pool_config();
let key_ids = vec!["key-google-429".to_string()];
record_admin_provider_pool_error(
&runner,
"provider-1",
"key-google-429",
&pool_config,
429,
Some(
r#"{
"error": {
"message": "Quota exhausted. reset after 45s.",
"status": "RESOURCE_EXHAUSTED",
"details": [{
"metadata": {
"quotaResetDelay": "45s"
}
}]
}
}"#,
),
None,
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-google-429")
.map(String::as_str),
Some("rate_limited_429")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-google-429")
.is_some_and(|ttl| *ttl <= 45 && *ttl >= 30));
}
#[tokio::test]
async fn error_feedback_applies_unschedulable_rule_cooldown() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_unschedulable_rule");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let mut pool_config = sample_pool_config();
pool_config.unschedulable_rules = vec![AdminProviderPoolUnschedulableRule {
keyword: "review required".to_string(),
duration_minutes: 7,
}];
let key_ids = vec!["key-3".to_string()];
record_admin_provider_pool_error(
&runner,
"provider-1",
"key-3",
&pool_config,
418,
Some(r#"{"error":{"message":"manual review required before reuse"}}"#),
None,
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-3")
.map(String::as_str),
Some("rule:review required")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-3")
.is_some_and(|ttl| *ttl <= 420 && *ttl >= 380));
}
#[tokio::test]
async fn stream_timeout_policy_cools_down_after_threshold() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_stream_timeout");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let mut pool_config = sample_pool_config();
pool_config.stream_timeout_threshold = 2;
pool_config.stream_timeout_window_seconds = 300;
pool_config.stream_timeout_cooldown_seconds = 90;
let key_ids = vec!["key-4".to_string()];
record_admin_provider_pool_stream_timeout(&runner, "provider-1", "key-4", &pool_config)
.await;
record_admin_provider_pool_stream_timeout(&runner, "provider-1", "key-4", &pool_config)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-4")
.map(String::as_str),
Some("stream_timeout_x2")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-4")
.is_some_and(|ttl| *ttl <= 90 && *ttl >= 70));
}
}

8
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
View File

@@ -772,7 +772,9 @@ pub(super) fn build_admin_pool_key_payload(
let health_score = admin_pool_health_score(key);
let circuit_breaker_open = admin_pool_circuit_breaker_open(key);
let auth_semantics = provider_key_auth_semantics(key, provider_type);
let account_quota_exhausted = pool_config.is_some_and(|config| config.skip_exhausted_accounts)
let account_quota_exhausted = pool_config
.as_ref()
.is_some_and(|config| config.skip_exhausted_accounts)
&& admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type);
let (scheduling_status, scheduling_reason, scheduling_label, scheduling_reasons) =
admin_pool_scheduling_payload(
@@ -988,7 +990,9 @@ pub(super) fn build_admin_pool_key_payload(
);
payload.insert(
"cost_limit".to_string(),
json!(pool_config.map(|config| config.cost_limit_per_key_tokens)),
json!(pool_config
.as_ref()
.map(|config| config.cost_limit_per_key_tokens)),
);
payload.insert(
"request_count".to_string(),

14
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/keys.rs
View File

@@ -125,10 +125,16 @@ pub(super) async fn build_admin_pool_list_keys_response(
};
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let runtime = match (state.redis_kv_runner(), pool_config) {
let runtime = match (state.redis_kv_runner(), pool_config.as_ref()) {
(Some(runner), Some(pool_config)) if !key_ids.is_empty() => {
read_admin_provider_pool_runtime_state(&runner, &provider.id, &key_ids, pool_config)
.await
read_admin_provider_pool_runtime_state(
&runner,
&provider.id,
&key_ids,
pool_config,
None,
)
.await
}
_ => AdminProviderPoolRuntimeState::default(),
};
@@ -145,7 +151,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
&provider.provider_type,
&key,
&runtime,
pool_config,
pool_config.clone(),
);
if let Some(summary) = usage_summary_by_key_id.get(&key.id) {
if let Some(object) = payload.as_object_mut() {

28
apps/aether-gateway/src/handlers/admin/provider/shared/support.rs
View File

@@ -10,21 +10,47 @@ pub(crate) const ADMIN_PROVIDER_POOL_SCAN_BATCH: u64 = 200;
pub(crate) const ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL: &str =
"Admin provider OAuth data unavailable";
#[derive(Debug, Clone, Copy)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct AdminProviderPoolSchedulingPreset {
pub(crate) preset: String,
pub(crate) enabled: bool,
pub(crate) mode: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct AdminProviderPoolUnschedulableRule {
pub(crate) keyword: String,
pub(crate) duration_minutes: u64,
}
#[derive(Debug, Clone)]
pub(crate) struct AdminProviderPoolConfig {
pub(crate) scheduling_presets: Vec<AdminProviderPoolSchedulingPreset>,
pub(crate) unschedulable_rules: Vec<AdminProviderPoolUnschedulableRule>,
pub(crate) lru_enabled: bool,
pub(crate) skip_exhausted_accounts: bool,
pub(crate) sticky_session_ttl_seconds: u64,
pub(crate) latency_window_seconds: u64,
pub(crate) latency_sample_limit: u64,
pub(crate) cost_window_seconds: u64,
pub(crate) cost_limit_per_key_tokens: Option<u64>,
pub(crate) rate_limit_cooldown_seconds: u64,
pub(crate) overload_cooldown_seconds: u64,
pub(crate) health_policy_enabled: bool,
pub(crate) stream_timeout_threshold: u64,
pub(crate) stream_timeout_window_seconds: u64,
pub(crate) stream_timeout_cooldown_seconds: u64,
}
#[derive(Debug, Default)]
pub(crate) struct AdminProviderPoolRuntimeState {
pub(crate) total_sticky_sessions: usize,
pub(crate) sticky_sessions_by_key: BTreeMap<String, usize>,
pub(crate) sticky_bound_key_id: Option<String>,
pub(crate) cooldown_reason_by_key: BTreeMap<String, String>,
pub(crate) cooldown_ttl_by_key: BTreeMap<String, u64>,
pub(crate) cost_window_usage_by_key: BTreeMap<String, u64>,
pub(crate) latency_avg_ms_by_key: BTreeMap<String, f64>,
pub(crate) lru_score_by_key: BTreeMap<String, f64>,
}

1
apps/aether-gateway/src/handlers/shared/mod.rs
View File

@@ -5,6 +5,7 @@ mod email_templates;
mod external_models;
mod normalize;
mod payloads;
pub(crate) mod provider_pool;
mod request_utils;
mod system_config_values;
mod usage_stats;

9
apps/aether-gateway/src/handlers/shared/provider_pool.rs Normal file
View File

@@ -0,0 +1,9 @@
pub(crate) use super::super::admin::provider::pool::config::admin_provider_pool_config_from_config_value;
pub(crate) use super::super::admin::provider::pool::runtime::{
read_admin_provider_pool_runtime_state, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
};
pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
AdminProviderPoolUnschedulableRule, ADMIN_PROVIDER_POOL_SCAN_BATCH,
};

4
apps/aether-gateway/src/state/oauth.rs
View File

@@ -672,6 +672,10 @@ impl AppState {
Ok(None)
}
pub(crate) async fn invalidate_local_oauth_refresh_entry(&self, key_id: &str) -> bool {
self.oauth_refresh.invalidate_cached_entry(key_id).await
}
async fn persist_local_oauth_refresh_entry(
&self,
transport: &provider_transport::GatewayProviderTransportSnapshot,