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

fix: carry session affinity through Gemini files

Browse Source
This commit is contained in:
RWDai
2026-05-05 12:25:07 +08:00
parent 628329e493
commit 657f9f0be1
6 changed files with 131 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
View File

@@ -28,6 +28,7 @@ pub(crate) struct LocalAuthenticatedDecisionInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
}
struct GatewayAuthenticatedDecisionInputPort<'a> {
@@ -108,6 +109,7 @@ pub(crate) fn build_local_authenticated_decision_input(
auth_context: resolved_input.auth_context,
auth_snapshot: resolved_input.auth_snapshot,
required_capabilities: resolved_input.required_capabilities,
client_session_affinity: None,
}
}

36
apps/aether-gateway/src/ai_serving/planner/specialized/files.rs
View File

@@ -103,7 +103,14 @@ pub(crate) async fn build_local_gemini_files_sync_attempt_source_for_kind<'a>(
return Ok(None);
};
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
let Some(input) = resolve_local_gemini_files_decision_input(
state,
parts,
Some(body_json),
trace_id,
decision,
)
.await
else {
return Ok(None);
};
@@ -140,7 +147,8 @@ pub(crate) async fn build_local_gemini_files_stream_attempt_source_for_kind<'a>(
return Ok(None);
};
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
let Some(input) =
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
else {
return Ok(None);
};
@@ -298,7 +306,14 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
return Ok(None);
};
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
let Some(input) = resolve_local_gemini_files_decision_input(
state,
parts,
Some(body_json),
trace_id,
decision,
)
.await
else {
return Ok(None);
};
@@ -338,7 +353,8 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
return Ok(None);
};
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
let Some(input) =
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
else {
return Ok(None);
};
@@ -379,7 +395,14 @@ async fn build_local_sync_plan_and_reports(
spec: LocalGeminiFilesSpec,
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
let Some(input) = resolve_local_gemini_files_decision_input(
state,
parts,
Some(body_json),
trace_id,
decision,
)
.await
else {
return Ok(Vec::new());
};
@@ -430,7 +453,8 @@ async fn build_local_stream_plan_and_reports(
spec: LocalGeminiFilesSpec,
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
let Some(input) =
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
else {
return Ok(Vec::new());
};

13
apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs
View File

@@ -23,6 +23,7 @@ use crate::ai_serving::{
resolve_local_decision_execution_runtime_auth_context, CandidateFailureDiagnostic,
ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayError};
@@ -36,6 +37,8 @@ pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
pub(super) async fn resolve_local_gemini_files_decision_input(
state: &AppState,
parts: &http::request::Parts,
body_json: Option<&serde_json::Value>,
trace_id: &str,
decision: &GatewayControlDecision,
) -> Option<LocalGeminiFilesDecisionInput> {
@@ -64,7 +67,9 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
}
};
Some(build_local_authenticated_decision_input(resolved_input))
let mut input = build_local_authenticated_decision_input(resolved_input);
input.client_session_affinity = client_session_affinity_from_parts(parts, body_json);
Some(input)
}
pub(super) async fn materialize_local_gemini_files_candidate_attempts(
@@ -84,6 +89,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
GEMINI_FILES_REQUIRED_CAPABILITY,
false,
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await?;
@@ -93,7 +99,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
GEMINI_FILES_CLIENT_API_FORMAT,
None,
Some(&input.auth_snapshot),
None,
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
None,
None,
@@ -155,6 +161,7 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
GEMINI_FILES_REQUIRED_CAPABILITY,
false,
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await?;
@@ -164,7 +171,7 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
GEMINI_FILES_CLIENT_API_FORMAT,
None,
Some(&input.auth_snapshot),
None,
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
None,
None,

2
apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs
View File

@@ -138,6 +138,7 @@ impl<'a> PlannerAppState<'a> {
required_capability: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
@@ -153,6 +154,7 @@ impl<'a> PlannerAppState<'a> {
required_capability,
require_streaming,
auth_snapshot,
client_session_affinity,
attempt_now_unix_secs,
)
.await?;

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

@@ -158,6 +158,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
required_capability: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
Ok(
@@ -168,6 +169,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
required_capability,
require_streaming,
auth_snapshot,
client_session_affinity,
now_unix_secs,
)
.await?
@@ -182,6 +184,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
required_capability: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
let normalized_api_format = normalize_api_format(candidate_api_format);
@@ -224,7 +227,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
require_streaming,
required_capabilities.as_ref(),
auth_snapshot,
None,
client_session_affinity,
now_unix_secs,
false,
)

83
apps/aether-gateway/src/scheduler/candidate/tests/required_capability.rs
View File

@@ -7,10 +7,14 @@ use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_scheduler_core::ClientSessionAffinity;
use crate::cache::SchedulerAffinityTarget;
use crate::data::GatewayDataState;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::AppState;
use super::super::affinity::build_scheduler_affinity_cache_key;
use super::super::{
list_selectable_candidates_for_required_capability_without_requested_model,
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal,
@@ -59,6 +63,7 @@ async fn compatible_required_capability_prefers_matching_keys_without_hard_filte
"cache_1h",
false,
None,
None,
100,
)
.await
@@ -113,6 +118,7 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
"gemini_files",
false,
None,
None,
100,
)
.await
@@ -123,6 +129,82 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
assert_eq!(selection[0].key_id, "key-b");
}
#[tokio::test]
async fn required_capability_without_model_uses_session_scoped_affinity() {
let mut fallback = sample_row();
fallback.provider_id = "provider-a".to_string();
fallback.provider_name = "provider-a".to_string();
fallback.endpoint_id = "endpoint-a".to_string();
fallback.endpoint_api_format = "gemini:generate_content".to_string();
fallback.endpoint_api_family = Some("gemini".to_string());
fallback.key_api_formats = Some(vec!["gemini:generate_content".to_string()]);
fallback.key_id = "key-a".to_string();
fallback.key_name = "alpha".to_string();
fallback.global_model_name = "gemini-2.5-pro".to_string();
fallback.key_capabilities = Some(serde_json::json!({"gemini_files": true}));
fallback.key_global_priority_by_format =
Some(serde_json::json!({"gemini:generate_content": 0}));
let mut session_target = fallback.clone();
session_target.provider_id = "provider-b".to_string();
session_target.provider_name = "provider-b".to_string();
session_target.endpoint_id = "endpoint-b".to_string();
session_target.key_id = "key-b".to_string();
session_target.key_name = "beta".to_string();
session_target.key_global_priority_by_format =
Some(serde_json::json!({"gemini:generate_content": 10}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
fallback,
session_target,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
);
let auth_snapshot = sample_auth_snapshot("api-key-1");
let client_session_affinity = ClientSessionAffinity::new(
Some("generic".to_string()),
Some("session=conversation-1".to_string()),
);
let cache_key = build_scheduler_affinity_cache_key(
Some(&auth_snapshot),
"gemini:generate_content",
"gemini-2.5-pro",
Some(&client_session_affinity),
)
.expect("session affinity cache key should build");
state.scheduler_affinity_cache.insert(
cache_key,
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),
key_id: "key-b".to_string(),
},
SCHEDULER_AFFINITY_TTL,
16,
);
let selection = list_selectable_candidates_for_required_capability_without_requested_model(
state.data.as_ref(),
&state,
"gemini:generate_content",
"gemini_files",
false,
Some(&auth_snapshot),
Some(&client_session_affinity),
100,
)
.await
.expect("selection should succeed");
assert_eq!(selection.len(), 2);
assert_eq!(selection[0].provider_id, "provider-b");
assert_eq!(selection[0].key_id, "key-b");
}
#[tokio::test]
async fn required_capability_reports_auth_limit_signal_when_every_model_is_blocked_by_api_key_concurrency(
) {
@@ -189,6 +271,7 @@ async fn required_capability_reports_auth_limit_signal_when_every_model_is_block
"cache_1h",
false,
Some(&auth_snapshot),
None,
100,
)
.await