mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Implement independent provider pool scheduling runtime
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -55,8 +55,9 @@ 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)
|
||||
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,
|
||||
@@ -121,8 +122,9 @@ 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)
|
||||
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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -57,8 +57,9 @@ 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)
|
||||
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,
|
||||
@@ -141,8 +142,9 @@ 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)
|
||||
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,
|
||||
|
||||
1449
apps/aether-gateway/src/ai_pipeline/planner/pool_scheduler.rs
Normal file
1449
apps/aether-gateway/src/ai_pipeline/planner/pool_scheduler.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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;
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
188
apps/aether-gateway/src/execution_runtime/pool_feedback.rs
Normal 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;
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -36,7 +36,13 @@ 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)
|
||||
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(),
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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(),
|
||||
|
||||
@@ -125,9 +125,15 @@ 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)
|
||||
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() {
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
|
||||
@@ -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
9
apps/aether-gateway/src/handlers/shared/provider_pool.rs
Normal 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,
|
||||
};
|
||||
@@ -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,
|
||||
|
||||
@@ -231,6 +231,10 @@ impl LocalOAuthRefreshCoordinator {
|
||||
self.cache.lock().await.insert(key_id.to_string(), entry);
|
||||
}
|
||||
|
||||
pub async fn invalidate_cached_entry(&self, key_id: &str) -> bool {
|
||||
self.cache.lock().await.remove(key_id).is_some()
|
||||
}
|
||||
|
||||
pub async fn resolve_with_result(
|
||||
&self,
|
||||
executor: &dyn LocalOAuthHttpExecutor,
|
||||
|
||||
Reference in New Issue
Block a user