mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +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_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_family_stream_plan_and_reports, build_standard_family_sync_plan_and_reports,
|
||||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||||
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
|
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
|
||||||
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
|
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
|
||||||
set_local_openai_chat_execution_exhausted_diagnostic, GatewayAuthApiKeySnapshot,
|
maybe_build_sync_plan_payload, set_local_openai_chat_execution_exhausted_diagnostic,
|
||||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
|
||||||
|
PlannerAppState,
|
||||||
};
|
};
|
||||||
pub(crate) use self::pure::*;
|
pub(crate) use self::pure::*;
|
||||||
pub(crate) use crate::control::GatewayControlDecision;
|
pub(crate) use crate::control::GatewayControlDecision;
|
||||||
|
|||||||
@@ -1040,6 +1040,7 @@ mod tests {
|
|||||||
"openai:chat",
|
"openai:chat",
|
||||||
"gpt-4.1",
|
"gpt-4.1",
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -1115,6 +1116,7 @@ mod tests {
|
|||||||
"openai:chat",
|
"openai:chat",
|
||||||
"gpt-4.1",
|
"gpt-4.1",
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -1186,6 +1188,7 @@ mod tests {
|
|||||||
"openai:chat",
|
"openai:chat",
|
||||||
"gpt-4.1",
|
"gpt-4.1",
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -1244,6 +1247,7 @@ mod tests {
|
|||||||
"openai:chat",
|
"openai:chat",
|
||||||
"gpt-4.1",
|
"gpt-4.1",
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -1320,6 +1324,7 @@ mod tests {
|
|||||||
"openai:chat",
|
"openai:chat",
|
||||||
"gpt-4.1",
|
"gpt-4.1",
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ use std::collections::BTreeSet;
|
|||||||
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
|
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
|
||||||
|
|
||||||
use super::candidate_affinity::rank_eligible_local_execution_candidates;
|
use super::candidate_affinity::rank_eligible_local_execution_candidates;
|
||||||
|
use super::pool_scheduler::apply_local_execution_pool_scheduler;
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq)]
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
pub(crate) struct EligibleLocalExecutionCandidate {
|
pub(crate) struct EligibleLocalExecutionCandidate {
|
||||||
@@ -28,6 +29,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates(
|
|||||||
client_api_format: &str,
|
client_api_format: &str,
|
||||||
requested_model: &str,
|
requested_model: &str,
|
||||||
required_capabilities: Option<&serde_json::Value>,
|
required_capabilities: Option<&serde_json::Value>,
|
||||||
|
sticky_session_token: Option<&str>,
|
||||||
) -> (
|
) -> (
|
||||||
Vec<EligibleLocalExecutionCandidate>,
|
Vec<EligibleLocalExecutionCandidate>,
|
||||||
Vec<SkippedLocalExecutionCandidate>,
|
Vec<SkippedLocalExecutionCandidate>,
|
||||||
@@ -37,6 +39,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates(
|
|||||||
candidates,
|
candidates,
|
||||||
client_api_format,
|
client_api_format,
|
||||||
required_capabilities,
|
required_capabilities,
|
||||||
|
sticky_session_token,
|
||||||
|candidate, transport| {
|
|candidate, transport| {
|
||||||
current_local_execution_candidate_skip_reason_with_transport(
|
current_local_execution_candidate_skip_reason_with_transport(
|
||||||
candidate,
|
candidate,
|
||||||
@@ -55,6 +58,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport
|
|||||||
client_api_format: &str,
|
client_api_format: &str,
|
||||||
requested_model: Option<&str>,
|
requested_model: Option<&str>,
|
||||||
required_capabilities: Option<&serde_json::Value>,
|
required_capabilities: Option<&serde_json::Value>,
|
||||||
|
sticky_session_token: Option<&str>,
|
||||||
) -> (
|
) -> (
|
||||||
Vec<EligibleLocalExecutionCandidate>,
|
Vec<EligibleLocalExecutionCandidate>,
|
||||||
Vec<SkippedLocalExecutionCandidate>,
|
Vec<SkippedLocalExecutionCandidate>,
|
||||||
@@ -64,6 +68,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport
|
|||||||
candidates,
|
candidates,
|
||||||
client_api_format,
|
client_api_format,
|
||||||
required_capabilities,
|
required_capabilities,
|
||||||
|
sticky_session_token,
|
||||||
|candidate, transport| {
|
|candidate, transport| {
|
||||||
current_local_execution_candidate_common_skip_reason_with_transport(
|
current_local_execution_candidate_common_skip_reason_with_transport(
|
||||||
candidate,
|
candidate,
|
||||||
@@ -80,6 +85,7 @@ async fn filter_and_rank_local_execution_candidates_with_gate<F>(
|
|||||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||||
client_api_format: &str,
|
client_api_format: &str,
|
||||||
required_capabilities: Option<&serde_json::Value>,
|
required_capabilities: Option<&serde_json::Value>,
|
||||||
|
sticky_session_token: Option<&str>,
|
||||||
runtime_skip_reason: F,
|
runtime_skip_reason: F,
|
||||||
) -> (
|
) -> (
|
||||||
Vec<EligibleLocalExecutionCandidate>,
|
Vec<EligibleLocalExecutionCandidate>,
|
||||||
@@ -149,10 +155,49 @@ where
|
|||||||
required_capabilities,
|
required_capabilities,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
let (ranked, pool_skipped) =
|
||||||
|
apply_local_execution_pool_scheduler(state, ranked, sticky_session_token).await;
|
||||||
|
skipped.extend(pool_skipped);
|
||||||
|
|
||||||
(ranked, 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(
|
fn current_local_execution_candidate_common_skip_reason_with_transport(
|
||||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||||
transport: &GatewayProviderTransportSnapshot,
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ mod materialization_policy;
|
|||||||
mod passthrough;
|
mod passthrough;
|
||||||
mod payload_metadata;
|
mod payload_metadata;
|
||||||
mod plan_builders;
|
mod plan_builders;
|
||||||
|
mod pool_scheduler;
|
||||||
mod report_context;
|
mod report_context;
|
||||||
mod route;
|
mod route;
|
||||||
mod runtime_miss;
|
mod runtime_miss;
|
||||||
@@ -25,6 +26,7 @@ mod specialized;
|
|||||||
mod standard;
|
mod standard;
|
||||||
mod state;
|
mod state;
|
||||||
|
|
||||||
|
pub(crate) use self::candidate_eligibility::extract_pool_sticky_session_token;
|
||||||
pub(crate) use self::passthrough::{
|
pub(crate) use self::passthrough::{
|
||||||
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
|
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()),
|
Some(input.requested_model.as_str()),
|
||||||
"candidate_evaluation_incomplete",
|
"candidate_evaluation_incomplete",
|
||||||
);
|
);
|
||||||
let (attempts, candidate_count) =
|
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||||
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
|
state, trace_id, &input, body_json, spec,
|
||||||
|
)
|
||||||
.await?;
|
.await?;
|
||||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||||
state,
|
state,
|
||||||
@@ -121,8 +122,9 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
|
|||||||
Some(input.requested_model.as_str()),
|
Some(input.requested_model.as_str()),
|
||||||
"candidate_evaluation_incomplete",
|
"candidate_evaluation_incomplete",
|
||||||
);
|
);
|
||||||
let (attempts, candidate_count) =
|
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||||
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
|
state, trace_id, &input, body_json, spec,
|
||||||
|
)
|
||||||
.await?;
|
.await?;
|
||||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||||
state,
|
state,
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
use tracing::warn;
|
use tracing::warn;
|
||||||
|
|
||||||
use crate::ai_pipeline::planner::candidate_eligibility::{
|
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::{
|
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||||
persist_available_local_execution_candidates_with_context,
|
persist_available_local_execution_candidates_with_context,
|
||||||
@@ -85,10 +86,12 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
|||||||
state: &AppState,
|
state: &AppState,
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
input: &LocalSameFormatProviderDecisionInput,
|
input: &LocalSameFormatProviderDecisionInput,
|
||||||
|
body_json: &serde_json::Value,
|
||||||
spec: LocalSameFormatProviderSpec,
|
spec: LocalSameFormatProviderSpec,
|
||||||
) -> Result<(Vec<LocalSameFormatProviderCandidateAttempt>, usize), GatewayError> {
|
) -> Result<(Vec<LocalSameFormatProviderCandidateAttempt>, usize), GatewayError> {
|
||||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||||
let planner_state = PlannerAppState::new(state);
|
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(
|
let persistence_policy = build_local_candidate_persistence_policy(
|
||||||
&input.auth_context,
|
&input.auth_context,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
@@ -110,6 +113,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
|||||||
spec_metadata.api_format,
|
spec_metadata.api_format,
|
||||||
&input.requested_model,
|
&input.requested_model,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
|
sticky_session_token.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let skipped_candidates = preselection_skipped
|
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()),
|
Some(input.requested_model.as_str()),
|
||||||
"candidate_evaluation_incomplete",
|
"candidate_evaluation_incomplete",
|
||||||
);
|
);
|
||||||
let (attempts, candidate_count) =
|
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||||
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
|
state, trace_id, &input, body_json, spec,
|
||||||
|
)
|
||||||
.await?;
|
.await?;
|
||||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||||
state,
|
state,
|
||||||
@@ -141,8 +142,9 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
|||||||
Some(input.requested_model.as_str()),
|
Some(input.requested_model.as_str()),
|
||||||
"candidate_evaluation_incomplete",
|
"candidate_evaluation_incomplete",
|
||||||
);
|
);
|
||||||
let (attempts, candidate_count) =
|
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||||
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
|
state, trace_id, &input, body_json, spec,
|
||||||
|
)
|
||||||
.await?;
|
.await?;
|
||||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||||
state,
|
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,
|
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||||
None,
|
None,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
|
|||||||
@@ -60,6 +60,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
&input,
|
&input,
|
||||||
|
body_json,
|
||||||
spec_metadata.api_format,
|
spec_metadata.api_format,
|
||||||
spec_metadata.decision_kind,
|
spec_metadata.decision_kind,
|
||||||
)
|
)
|
||||||
@@ -102,6 +103,7 @@ async fn build_local_sync_plan_and_reports(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
&input,
|
&input,
|
||||||
|
body_json,
|
||||||
spec_metadata.api_format,
|
spec_metadata.api_format,
|
||||||
spec_metadata.decision_kind,
|
spec_metadata.decision_kind,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -4,7 +4,8 @@ use tracing::warn;
|
|||||||
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
|
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
|
||||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||||
use crate::ai_pipeline::planner::candidate_eligibility::{
|
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::{
|
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||||
mark_skipped_local_execution_candidate,
|
mark_skipped_local_execution_candidate,
|
||||||
@@ -96,6 +97,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
|
|||||||
state: &AppState,
|
state: &AppState,
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
input: &LocalVideoCreateDecisionInput,
|
input: &LocalVideoCreateDecisionInput,
|
||||||
|
body_json: &serde_json::Value,
|
||||||
api_format: &str,
|
api_format: &str,
|
||||||
decision_kind: &str,
|
decision_kind: &str,
|
||||||
) -> Option<Vec<LocalVideoCreateCandidateAttempt>> {
|
) -> Option<Vec<LocalVideoCreateCandidateAttempt>> {
|
||||||
@@ -128,6 +130,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
|
|||||||
planner_state,
|
planner_state,
|
||||||
trace_id,
|
trace_id,
|
||||||
input,
|
input,
|
||||||
|
body_json,
|
||||||
candidates,
|
candidates,
|
||||||
preselection_skipped
|
preselection_skipped
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -148,10 +151,12 @@ async fn materialize_local_video_create_candidate_attempts(
|
|||||||
state: PlannerAppState<'_>,
|
state: PlannerAppState<'_>,
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
input: &LocalVideoCreateDecisionInput,
|
input: &LocalVideoCreateDecisionInput,
|
||||||
|
body_json: &serde_json::Value,
|
||||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||||
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
||||||
api_format: &str,
|
api_format: &str,
|
||||||
) -> Vec<LocalVideoCreateCandidateAttempt> {
|
) -> Vec<LocalVideoCreateCandidateAttempt> {
|
||||||
|
let sticky_session_token = extract_pool_sticky_session_token(body_json);
|
||||||
let persistence_policy = build_local_candidate_persistence_policy(
|
let persistence_policy = build_local_candidate_persistence_policy(
|
||||||
&input.auth_context,
|
&input.auth_context,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
@@ -163,6 +168,7 @@ async fn materialize_local_video_create_candidate_attempts(
|
|||||||
api_format,
|
api_format,
|
||||||
&input.requested_model,
|
&input.requested_model,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
|
sticky_session_token.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let skipped_candidates = preselection_skipped
|
let skipped_candidates = preselection_skipped
|
||||||
|
|||||||
@@ -51,7 +51,8 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
|||||||
"candidate_evaluation_incomplete",
|
"candidate_evaluation_incomplete",
|
||||||
);
|
);
|
||||||
let (attempts, candidate_count) =
|
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);
|
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||||
|
|
||||||
for attempt in attempts {
|
for attempt in attempts {
|
||||||
@@ -99,7 +100,8 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
|
|||||||
"candidate_evaluation_incomplete",
|
"candidate_evaluation_incomplete",
|
||||||
);
|
);
|
||||||
let (attempts, candidate_count) =
|
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);
|
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||||
|
|
||||||
for attempt in attempts {
|
for attempt in attempts {
|
||||||
@@ -153,7 +155,8 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
|||||||
"candidate_evaluation_incomplete",
|
"candidate_evaluation_incomplete",
|
||||||
);
|
);
|
||||||
let (attempts, candidate_count) =
|
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);
|
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||||
if candidate_count == 0 {
|
if candidate_count == 0 {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
@@ -226,7 +229,8 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
|||||||
"candidate_evaluation_incomplete",
|
"candidate_evaluation_incomplete",
|
||||||
);
|
);
|
||||||
let (attempts, candidate_count) =
|
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);
|
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||||
if candidate_count == 0 {
|
if candidate_count == 0 {
|
||||||
return Ok(Vec::new());
|
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::conversion::{request_candidate_api_formats, request_conversion_kind};
|
||||||
use crate::ai_pipeline::planner::candidate_eligibility::{
|
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::{
|
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||||
persist_available_local_execution_candidates_with_context,
|
persist_available_local_execution_candidates_with_context,
|
||||||
@@ -87,10 +88,12 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
|||||||
state: &AppState,
|
state: &AppState,
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
input: &LocalStandardDecisionInput,
|
input: &LocalStandardDecisionInput,
|
||||||
|
body_json: &serde_json::Value,
|
||||||
spec: LocalStandardSpec,
|
spec: LocalStandardSpec,
|
||||||
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
|
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
|
||||||
let spec_metadata = local_standard_spec_metadata(spec);
|
let spec_metadata = local_standard_spec_metadata(spec);
|
||||||
let planner_state = PlannerAppState::new(state);
|
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(
|
let persistence_policy = build_local_candidate_persistence_policy(
|
||||||
&input.auth_context,
|
&input.auth_context,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
@@ -176,6 +179,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
|||||||
spec_metadata.api_format,
|
spec_metadata.api_format,
|
||||||
&input.requested_model,
|
&input.requested_model,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
|
sticky_session_token.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let skipped_candidates = preselection_skipped
|
let skipped_candidates = preselection_skipped
|
||||||
|
|||||||
@@ -2,7 +2,8 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
|||||||
|
|
||||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||||
use crate::ai_pipeline::planner::candidate_eligibility::{
|
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::{
|
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||||
mark_skipped_local_execution_candidate,
|
mark_skipped_local_execution_candidate,
|
||||||
@@ -55,10 +56,12 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
|||||||
state: &AppState,
|
state: &AppState,
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
input: &LocalOpenAiChatDecisionInput,
|
input: &LocalOpenAiChatDecisionInput,
|
||||||
|
body_json: &serde_json::Value,
|
||||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||||
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
||||||
) -> Vec<LocalOpenAiChatCandidateAttempt> {
|
) -> Vec<LocalOpenAiChatCandidateAttempt> {
|
||||||
let planner_state = PlannerAppState::new(state);
|
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 auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
|
||||||
let persistence_policy = build_local_candidate_persistence_policy(
|
let persistence_policy = build_local_candidate_persistence_policy(
|
||||||
auth_context,
|
auth_context,
|
||||||
@@ -71,6 +74,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
|||||||
"openai:chat",
|
"openai:chat",
|
||||||
&input.requested_model,
|
&input.requested_model,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
|
sticky_session_token.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let skipped_candidates = preselection_skipped
|
let skipped_candidates = preselection_skipped
|
||||||
|
|||||||
@@ -118,6 +118,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
&input,
|
&input,
|
||||||
|
body_json,
|
||||||
candidates,
|
candidates,
|
||||||
skipped_candidates,
|
skipped_candidates,
|
||||||
)
|
)
|
||||||
@@ -183,6 +184,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
&input,
|
&input,
|
||||||
|
body_json,
|
||||||
candidates,
|
candidates,
|
||||||
skipped_candidates,
|
skipped_candidates,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -80,6 +80,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
&input,
|
&input,
|
||||||
|
body_json,
|
||||||
candidates,
|
candidates,
|
||||||
skipped_candidates,
|
skipped_candidates,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -80,6 +80,7 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
&input,
|
&input,
|
||||||
|
body_json,
|
||||||
candidates,
|
candidates,
|
||||||
skipped_candidates,
|
skipped_candidates,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -6,7 +6,8 @@ use tracing::warn;
|
|||||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||||
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
|
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
|
||||||
use crate::ai_pipeline::planner::candidate_eligibility::{
|
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::{
|
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||||
mark_skipped_local_execution_candidate,
|
mark_skipped_local_execution_candidate,
|
||||||
@@ -82,11 +83,13 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
|
|||||||
state: &AppState,
|
state: &AppState,
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
input: &LocalOpenAiCliDecisionInput,
|
input: &LocalOpenAiCliDecisionInput,
|
||||||
|
body_json: &serde_json::Value,
|
||||||
spec: LocalOpenAiCliSpec,
|
spec: LocalOpenAiCliSpec,
|
||||||
) -> Result<Vec<LocalOpenAiCliCandidateAttempt>, GatewayError> {
|
) -> Result<Vec<LocalOpenAiCliCandidateAttempt>, GatewayError> {
|
||||||
let spec_metadata = local_openai_cli_spec_metadata(spec);
|
let spec_metadata = local_openai_cli_spec_metadata(spec);
|
||||||
let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
|
let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
|
||||||
let planner_state = PlannerAppState::new(state);
|
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 auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
|
||||||
let persistence_policy = build_local_candidate_persistence_policy(
|
let persistence_policy = build_local_candidate_persistence_policy(
|
||||||
auth_context,
|
auth_context,
|
||||||
@@ -173,6 +176,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
|
|||||||
spec_metadata.api_format,
|
spec_metadata.api_format,
|
||||||
&input.requested_model,
|
&input.requested_model,
|
||||||
input.required_capabilities.as_ref(),
|
input.required_capabilities.as_ref(),
|
||||||
|
sticky_session_token.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let skipped_candidates = preselection_skipped
|
let skipped_candidates = preselection_skipped
|
||||||
|
|||||||
@@ -66,7 +66,8 @@ pub(crate) async fn maybe_build_sync_local_openai_cli_decision_payload(
|
|||||||
};
|
};
|
||||||
|
|
||||||
let attempts =
|
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 {
|
for attempt in attempts {
|
||||||
if let Some(payload) = maybe_build_local_openai_cli_decision_payload_for_candidate(
|
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 =
|
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 {
|
for attempt in attempts {
|
||||||
if let Some(payload) = maybe_build_local_openai_cli_decision_payload_for_candidate(
|
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 =
|
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();
|
let mut plans = Vec::new();
|
||||||
for attempt in attempts {
|
for attempt in attempts {
|
||||||
@@ -78,7 +79,8 @@ pub(super) async fn build_local_stream_plan_and_reports(
|
|||||||
};
|
};
|
||||||
|
|
||||||
let attempts =
|
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();
|
let mut plans = Vec::new();
|
||||||
for attempt in attempts {
|
for attempt in attempts {
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ use serde_json::{Map, Value};
|
|||||||
mod constants;
|
mod constants;
|
||||||
mod fallback;
|
mod fallback;
|
||||||
pub(crate) mod ndjson;
|
pub(crate) mod ndjson;
|
||||||
|
mod pool_feedback;
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
pub(crate) mod remote_compat;
|
pub(crate) mod remote_compat;
|
||||||
mod server;
|
mod server;
|
||||||
@@ -28,6 +29,10 @@ pub(crate) use self::fallback::{
|
|||||||
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
|
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
|
||||||
LocalFailoverDecision,
|
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::{
|
pub use server::{
|
||||||
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
|
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
|
||||||
build_execution_runtime_router_with_request_gates, serve_execution_runtime_tcp,
|
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,
|
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
|
||||||
};
|
};
|
||||||
use crate::execution_runtime::{
|
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_core_stream_error_finalize_report_kind,
|
||||||
resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
|
resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
|
||||||
should_retry_next_local_candidate_stream, LocalFailoverDecision,
|
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 (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
|
||||||
let error_response_text =
|
let error_response_text =
|
||||||
local_failover_response_text(body_json.as_ref(), &error_body, None);
|
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(
|
let failover_decision = resolve_local_candidate_failover_decision_stream(
|
||||||
state,
|
state,
|
||||||
&plan,
|
&plan,
|
||||||
@@ -1492,6 +1502,13 @@ async fn execute_stream_from_frame_stream(
|
|||||||
}),
|
}),
|
||||||
telemetry: telemetry.clone(),
|
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(
|
record_stream_terminal_usage(
|
||||||
&state_for_report,
|
&state_for_report,
|
||||||
&plan_for_report,
|
&plan_for_report,
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ use crate::control::GatewayControlDecision;
|
|||||||
use crate::execution_runtime::submission::{
|
use crate::execution_runtime::submission::{
|
||||||
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
|
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::log_ids::short_request_id;
|
||||||
use crate::request_candidate_runtime::record_report_request_candidate_status;
|
use crate::request_candidate_runtime::record_report_request_candidate_status;
|
||||||
use crate::usage::submit_sync_report;
|
use crate::usage::submit_sync_report;
|
||||||
@@ -122,6 +123,22 @@ async fn record_stream_sync_failure(
|
|||||||
failure: &StreamFailureReport,
|
failure: &StreamFailureReport,
|
||||||
started_at_unix_ms: Option<u64>,
|
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 context_seed = build_terminal_usage_context_seed(plan, report_context);
|
||||||
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
||||||
state
|
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::submission::submit_local_core_error_or_sync_finalize;
|
||||||
use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
|
use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
|
||||||
use crate::execution_runtime::{
|
use crate::execution_runtime::{
|
||||||
local_failover_response_text, resolve_core_sync_error_finalize_report_kind,
|
local_failover_response_text, record_pool_error_feedback, record_sync_pool_success_feedback,
|
||||||
should_fallback_to_control_sync, should_finalize_sync_response,
|
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
|
||||||
should_retry_next_local_candidate_sync, should_stop_local_candidate_failover_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::log_ids::short_request_id;
|
||||||
use crate::request_candidate_runtime::{
|
use crate::request_candidate_runtime::{
|
||||||
@@ -255,6 +256,17 @@ pub(crate) async fn execute_execution_runtime_sync(
|
|||||||
local_failover_response_text.as_deref(),
|
local_failover_response_text.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.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(
|
if should_retry_next_local_candidate_sync(
|
||||||
state,
|
state,
|
||||||
&plan,
|
&plan,
|
||||||
@@ -393,6 +405,15 @@ pub(crate) async fn execute_execution_runtime_sync(
|
|||||||
body_base64: body_base64.clone(),
|
body_base64: body_base64.clone(),
|
||||||
telemetry: result.telemetry.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(finalize_report_kind) = finalize_report_kind {
|
||||||
if let Some(implicit_finalize) = implicit_finalize {
|
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};
|
use serde_json::{Map, Value};
|
||||||
|
|
||||||
fn json_u64(value: &Value) -> Option<u64> {
|
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()))
|
.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
|
if let Some(explicit) = raw_pool_advanced
|
||||||
.get("lru_enabled")
|
.get("lru_enabled")
|
||||||
.and_then(Value::as_bool)
|
.and_then(Value::as_bool)
|
||||||
@@ -15,56 +125,70 @@ fn admin_provider_pool_lru_enabled(raw_pool_advanced: &Map<String, Value>) -> bo
|
|||||||
return explicit;
|
return explicit;
|
||||||
}
|
}
|
||||||
|
|
||||||
let Some(presets) = raw_pool_advanced
|
scheduling_presets
|
||||||
.get("scheduling_presets")
|
.iter()
|
||||||
.and_then(Value::as_array)
|
.any(|item| item.enabled && item.preset.eq_ignore_ascii_case("lru"))
|
||||||
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)
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn admin_provider_pool_config(
|
pub(crate) fn admin_provider_pool_config(
|
||||||
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
|
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
|
||||||
) -> Option<AdminProviderPoolConfig> {
|
) -> Option<AdminProviderPoolConfig> {
|
||||||
let raw_pool_advanced = provider
|
admin_provider_pool_config_from_config_value(provider.config.as_ref())
|
||||||
.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(Value::as_object)
|
||||||
.and_then(|config| config.get("pool_advanced"))?;
|
.and_then(|config| config.get("pool_advanced"))?;
|
||||||
|
|
||||||
let Some(pool_advanced) = raw_pool_advanced.as_object() else {
|
let Some(pool_advanced) = raw_pool_advanced.as_object() else {
|
||||||
return Some(AdminProviderPoolConfig {
|
return Some(AdminProviderPoolConfig {
|
||||||
|
scheduling_presets: Vec::new(),
|
||||||
|
unschedulable_rules: Vec::new(),
|
||||||
lru_enabled: false,
|
lru_enabled: false,
|
||||||
skip_exhausted_accounts: false,
|
skip_exhausted_accounts: false,
|
||||||
|
sticky_session_ttl_seconds: 3600,
|
||||||
|
latency_window_seconds: 3600,
|
||||||
|
latency_sample_limit: 50,
|
||||||
cost_window_seconds: 18_000,
|
cost_window_seconds: 18_000,
|
||||||
cost_limit_per_key_tokens: None,
|
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 {
|
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
|
skip_exhausted_accounts: pool_advanced
|
||||||
.get("skip_exhausted_accounts")
|
.get("skip_exhausted_accounts")
|
||||||
.and_then(Value::as_bool)
|
.and_then(Value::as_bool)
|
||||||
.unwrap_or(false),
|
.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
|
cost_window_seconds: pool_advanced
|
||||||
.get("cost_window_seconds")
|
.get("cost_window_seconds")
|
||||||
.and_then(json_u64)
|
.and_then(json_u64)
|
||||||
@@ -73,12 +197,41 @@ pub(crate) fn admin_provider_pool_config(
|
|||||||
cost_limit_per_key_tokens: pool_advanced
|
cost_limit_per_key_tokens: pool_advanced
|
||||||
.get("cost_limit_per_key_tokens")
|
.get("cost_limit_per_key_tokens")
|
||||||
.and_then(json_u64),
|
.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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
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 aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
@@ -117,15 +270,91 @@ mod tests {
|
|||||||
"pool_advanced": {
|
"pool_advanced": {
|
||||||
"skip_exhausted_accounts": true,
|
"skip_exhausted_accounts": true,
|
||||||
"lru_enabled": true,
|
"lru_enabled": true,
|
||||||
|
"sticky_session_ttl_seconds": 600,
|
||||||
|
"latency_window_seconds": 900,
|
||||||
|
"latency_sample_limit": 75,
|
||||||
"cost_window_seconds": 7200,
|
"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");
|
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
|
||||||
|
|
||||||
assert!(config.skip_exhausted_accounts);
|
assert!(config.skip_exhausted_accounts);
|
||||||
assert!(config.lru_enabled);
|
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_window_seconds, 7200);
|
||||||
assert_eq!(config.cost_limit_per_key_tokens, Some(12_000));
|
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:*"))
|
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 {
|
pub(super) fn pool_lru_key(keyspace: &RedisKeyspace, provider_id: &str) -> String {
|
||||||
keyspace.key(&format!("ap:{provider_id}:lru"))
|
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}"))
|
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 {
|
pub(super) fn parse_pool_cost_member(member: &str) -> u64 {
|
||||||
member
|
member
|
||||||
.rsplit_once(':')
|
.rsplit_once(':')
|
||||||
@@ -31,6 +55,13 @@ pub(super) fn parse_pool_cost_member(member: &str) -> u64 {
|
|||||||
.unwrap_or(0)
|
.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(
|
pub(super) fn pool_cooldown_keys(
|
||||||
keyspace: &RedisKeyspace,
|
keyspace: &RedisKeyspace,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
@@ -52,3 +83,14 @@ pub(super) fn pool_cost_keys(
|
|||||||
.map(|key_id| pool_cost_key(keyspace, provider_id, key_id))
|
.map(|key_id| pool_cost_key(keyspace, provider_id, key_id))
|
||||||
.collect()
|
.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 mutations;
|
||||||
mod reads;
|
mod reads;
|
||||||
mod status;
|
mod status;
|
||||||
|
mod writes;
|
||||||
|
|
||||||
pub(crate) use self::mutations::{
|
pub(crate) use self::mutations::{
|
||||||
clear_admin_provider_pool_cooldown, reset_admin_provider_pool_cost,
|
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,
|
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::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::{
|
use super::keys::{
|
||||||
parse_pool_cost_member, pool_cooldown_index_key, pool_cooldown_keys, pool_cost_keys,
|
parse_pool_cost_member, parse_pool_latency_member, pool_cooldown_index_key, pool_cooldown_key,
|
||||||
pool_lru_key, pool_sticky_pattern,
|
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::{
|
use crate::handlers::admin::provider::shared::support::{
|
||||||
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, ADMIN_PROVIDER_POOL_SCAN_BATCH,
|
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, ADMIN_PROVIDER_POOL_SCAN_BATCH,
|
||||||
@@ -11,6 +12,13 @@ use std::collections::BTreeMap;
|
|||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
use tracing::warn;
|
use tracing::warn;
|
||||||
|
|
||||||
|
fn current_unix_secs() -> u64 {
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs()
|
||||||
|
}
|
||||||
|
|
||||||
async fn scan_redis_keys(
|
async fn scan_redis_keys(
|
||||||
connection: &mut redis::aio::MultiplexedConnection,
|
connection: &mut redis::aio::MultiplexedConnection,
|
||||||
pattern: &str,
|
pattern: &str,
|
||||||
@@ -77,7 +85,8 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
|
|||||||
runner: &RedisKvRunner,
|
runner: &RedisKvRunner,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
key_ids: &[String],
|
key_ids: &[String],
|
||||||
pool_config: AdminProviderPoolConfig,
|
pool_config: &AdminProviderPoolConfig,
|
||||||
|
sticky_session_token: Option<&str>,
|
||||||
) -> AdminProviderPoolRuntimeState {
|
) -> AdminProviderPoolRuntimeState {
|
||||||
let mut runtime = AdminProviderPoolRuntimeState::default();
|
let mut runtime = AdminProviderPoolRuntimeState::default();
|
||||||
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
|
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 keyspace = runner.keyspace().clone();
|
||||||
let cooldown_keys = pool_cooldown_keys(&keyspace, provider_id, key_ids);
|
let cooldown_keys = pool_cooldown_keys(&keyspace, provider_id, key_ids);
|
||||||
let cost_keys = pool_cost_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(
|
let sticky_keys = match scan_redis_keys(
|
||||||
&mut connection,
|
&mut connection,
|
||||||
@@ -174,11 +234,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if !cost_keys.is_empty() {
|
if !cost_keys.is_empty() {
|
||||||
let window_start = SystemTime::now()
|
let window_start = current_unix_secs().saturating_sub(pool_config.cost_window_seconds);
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap_or_default()
|
|
||||||
.as_secs()
|
|
||||||
.saturating_sub(pool_config.cost_window_seconds);
|
|
||||||
let mut cost_pipeline = redis::pipe();
|
let mut cost_pipeline = redis::pipe();
|
||||||
for cost_key in &cost_keys {
|
for cost_key in &cost_keys {
|
||||||
cost_pipeline
|
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");
|
let mut command = redis::cmd("ZMSCORE");
|
||||||
command.arg(pool_lru_key(&keyspace, provider_id));
|
command.arg(pool_lru_key(&keyspace, provider_id));
|
||||||
for key_id in key_ids {
|
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 key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
|
||||||
let runtime = match state.redis_kv_runner() {
|
let runtime = match state.redis_kv_runner() {
|
||||||
Some(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
|
.await
|
||||||
}
|
}
|
||||||
None => AdminProviderPoolRuntimeState::default(),
|
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_window_usage": runtime.cost_window_usage_by_key.get(&key.id).copied().unwrap_or(0),
|
||||||
"cost_limit": pool_config.cost_limit_per_key_tokens,
|
"cost_limit": pool_config.cost_limit_per_key_tokens,
|
||||||
"sticky_sessions": runtime.sticky_sessions_by_key.get(&key.id).copied().unwrap_or(0),
|
"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(),
|
"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 health_score = admin_pool_health_score(key);
|
||||||
let circuit_breaker_open = admin_pool_circuit_breaker_open(key);
|
let circuit_breaker_open = admin_pool_circuit_breaker_open(key);
|
||||||
let auth_semantics = provider_key_auth_semantics(key, provider_type);
|
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);
|
&& admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type);
|
||||||
let (scheduling_status, scheduling_reason, scheduling_label, scheduling_reasons) =
|
let (scheduling_status, scheduling_reason, scheduling_label, scheduling_reasons) =
|
||||||
admin_pool_scheduling_payload(
|
admin_pool_scheduling_payload(
|
||||||
@@ -988,7 +990,9 @@ pub(super) fn build_admin_pool_key_payload(
|
|||||||
);
|
);
|
||||||
payload.insert(
|
payload.insert(
|
||||||
"cost_limit".to_string(),
|
"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(
|
payload.insert(
|
||||||
"request_count".to_string(),
|
"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 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() => {
|
(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
|
.await
|
||||||
}
|
}
|
||||||
_ => AdminProviderPoolRuntimeState::default(),
|
_ => AdminProviderPoolRuntimeState::default(),
|
||||||
@@ -145,7 +151,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
|
|||||||
&provider.provider_type,
|
&provider.provider_type,
|
||||||
&key,
|
&key,
|
||||||
&runtime,
|
&runtime,
|
||||||
pool_config,
|
pool_config.clone(),
|
||||||
);
|
);
|
||||||
if let Some(summary) = usage_summary_by_key_id.get(&key.id) {
|
if let Some(summary) = usage_summary_by_key_id.get(&key.id) {
|
||||||
if let Some(object) = payload.as_object_mut() {
|
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 =
|
pub(crate) const ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL: &str =
|
||||||
"Admin provider OAuth data unavailable";
|
"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) struct AdminProviderPoolConfig {
|
||||||
|
pub(crate) scheduling_presets: Vec<AdminProviderPoolSchedulingPreset>,
|
||||||
|
pub(crate) unschedulable_rules: Vec<AdminProviderPoolUnschedulableRule>,
|
||||||
pub(crate) lru_enabled: bool,
|
pub(crate) lru_enabled: bool,
|
||||||
pub(crate) skip_exhausted_accounts: 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_window_seconds: u64,
|
||||||
pub(crate) cost_limit_per_key_tokens: Option<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)]
|
#[derive(Debug, Default)]
|
||||||
pub(crate) struct AdminProviderPoolRuntimeState {
|
pub(crate) struct AdminProviderPoolRuntimeState {
|
||||||
pub(crate) total_sticky_sessions: usize,
|
pub(crate) total_sticky_sessions: usize,
|
||||||
pub(crate) sticky_sessions_by_key: BTreeMap<String, 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_reason_by_key: BTreeMap<String, String>,
|
||||||
pub(crate) cooldown_ttl_by_key: BTreeMap<String, u64>,
|
pub(crate) cooldown_ttl_by_key: BTreeMap<String, u64>,
|
||||||
pub(crate) cost_window_usage_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>,
|
pub(crate) lru_score_by_key: BTreeMap<String, f64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ mod email_templates;
|
|||||||
mod external_models;
|
mod external_models;
|
||||||
mod normalize;
|
mod normalize;
|
||||||
mod payloads;
|
mod payloads;
|
||||||
|
pub(crate) mod provider_pool;
|
||||||
mod request_utils;
|
mod request_utils;
|
||||||
mod system_config_values;
|
mod system_config_values;
|
||||||
mod usage_stats;
|
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)
|
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(
|
async fn persist_local_oauth_refresh_entry(
|
||||||
&self,
|
&self,
|
||||||
transport: &provider_transport::GatewayProviderTransportSnapshot,
|
transport: &provider_transport::GatewayProviderTransportSnapshot,
|
||||||
|
|||||||
@@ -231,6 +231,10 @@ impl LocalOAuthRefreshCoordinator {
|
|||||||
self.cache.lock().await.insert(key_id.to_string(), entry);
|
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(
|
pub async fn resolve_with_result(
|
||||||
&self,
|
&self,
|
||||||
executor: &dyn LocalOAuthHttpExecutor,
|
executor: &dyn LocalOAuthHttpExecutor,
|
||||||
|
|||||||
Reference in New Issue
Block a user