mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 05:00:19 +08:00
elky
7323d41fbe
Replace the provider/endpoint max_retries fields as the source of same-key retries with a routing policy setting, sticky_key_attempts (default 2). Only the first-ranked candidate is retried on the same key; every failover candidate gets a single attempt so failover keeps advancing instead of retrying each fallback key. Materialize exactly one attempt per candidate and derive same-key retries in the attempt loop after a candidate-scoped failure, so the retry budget no longer inflates up-front materialization and needs no upper bound. The budget travels in the report context; retries reuse the plan with a fresh candidate id and incremented retry index. Pool groups only retry their first key within the retry-index stride. Expose the setting in the routing profile editor and the set_scheduling rule action, and drop the max_retries input from the provider form.
513 lines
20 KiB
Rust
513 lines
20 KiB
Rust
use std::collections::BTreeMap;
|
|
|
|
use aether_ai_serving::{
|
|
build_ai_execution_report_context,
|
|
insert_provider_stream_event_api_format as insert_ai_provider_stream_event_api_format,
|
|
provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type,
|
|
AiExecutionReportContextParts, AiRequestOrigin, STICKY_KEY_ATTEMPTS_REPORT_FIELD,
|
|
};
|
|
use aether_routing_core::ResolvedRoutingPolicy;
|
|
use aether_runtime_state::RuntimeLockLease;
|
|
use aether_scheduler_core::{ClientSessionAffinity, SchedulerRankingOutcome};
|
|
use serde_json::{Map, Value};
|
|
|
|
use crate::ai_serving::{
|
|
request_origin_from_headers, request_path_implies_stream_request, sanitize_request_path,
|
|
sanitize_request_path_and_query, sanitize_request_query_string, ExecutionRuntimeAuthContext,
|
|
RequestOrigin,
|
|
};
|
|
use crate::client_session_affinity::{
|
|
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
|
|
};
|
|
use crate::orchestration::{
|
|
insert_pool_key_lease_report_context_fields, ExecutionAttemptIdentity,
|
|
ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD, SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
|
|
};
|
|
use crate::scheduler::affinity::insert_scheduler_affinity_policy_report_context_field;
|
|
|
|
pub(crate) struct LocalExecutionReportContextParts<'a> {
|
|
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
|
|
pub(crate) request_id: &'a str,
|
|
pub(crate) candidate_id: &'a str,
|
|
pub(crate) attempt_identity: ExecutionAttemptIdentity,
|
|
pub(crate) model: &'a str,
|
|
pub(crate) provider_name: &'a str,
|
|
pub(crate) provider_id: &'a str,
|
|
pub(crate) endpoint_id: &'a str,
|
|
pub(crate) key_id: &'a str,
|
|
pub(crate) key_name: Option<&'a str>,
|
|
pub(crate) model_id: Option<&'a str>,
|
|
pub(crate) global_model_id: Option<&'a str>,
|
|
pub(crate) global_model_name: Option<&'a str>,
|
|
pub(crate) provider_api_format: &'a str,
|
|
pub(crate) client_api_format: &'a str,
|
|
pub(crate) mapped_model: Option<&'a str>,
|
|
pub(crate) candidate_group_id: Option<&'a str>,
|
|
pub(crate) pool_key_lease: Option<&'a RuntimeLockLease>,
|
|
pub(crate) ranking: Option<&'a SchedulerRankingOutcome>,
|
|
pub(crate) upstream_url: Option<&'a str>,
|
|
pub(crate) header_rules: Option<&'a Value>,
|
|
pub(crate) body_rules: Option<&'a Value>,
|
|
pub(crate) provider_request_method: Option<Value>,
|
|
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
|
|
pub(crate) original_headers: &'a http::HeaderMap,
|
|
pub(crate) request_path: Option<&'a str>,
|
|
pub(crate) request_query_string: Option<&'a str>,
|
|
pub(crate) request_origin: Option<RequestOrigin>,
|
|
pub(crate) original_request_body_json: Option<&'a Value>,
|
|
pub(crate) original_request_body_base64: Option<&'a str>,
|
|
pub(crate) client_session_affinity: Option<&'a ClientSessionAffinity>,
|
|
pub(crate) routing_policy: Option<&'a ResolvedRoutingPolicy>,
|
|
pub(crate) scheduler_affinity_epoch: Option<u64>,
|
|
/// Routing policy sticky-key attempt budget; read back by the attempt
|
|
/// loop to derive same-key retries lazily.
|
|
pub(crate) sticky_key_attempts: Option<u32>,
|
|
pub(crate) client_requested_stream: bool,
|
|
pub(crate) upstream_is_stream: bool,
|
|
pub(crate) has_envelope: bool,
|
|
pub(crate) needs_conversion: bool,
|
|
pub(crate) extra_fields: Map<String, Value>,
|
|
}
|
|
|
|
pub(crate) fn build_local_execution_report_context(
|
|
parts: LocalExecutionReportContextParts<'_>,
|
|
) -> Value {
|
|
let RequestOrigin {
|
|
client_ip,
|
|
user_agent,
|
|
} = parts
|
|
.request_origin
|
|
.unwrap_or_else(|| request_origin_from_headers(parts.original_headers));
|
|
let original_headers = crate::ai_serving::collect_control_headers(parts.original_headers);
|
|
let original_request_body = crate::ai_serving::build_report_context_original_request_echo(
|
|
parts.original_request_body_json,
|
|
parts.original_request_body_base64,
|
|
);
|
|
let mut extra_fields = parts.extra_fields;
|
|
if let Some(value) = parts
|
|
.client_session_affinity
|
|
.and_then(client_session_affinity_report_context_value)
|
|
{
|
|
if let Some(client_family) = value
|
|
.as_object()
|
|
.and_then(|object| object.get("client_family"))
|
|
.and_then(Value::as_str)
|
|
.map(str::trim)
|
|
.filter(|client_family| !client_family.is_empty())
|
|
{
|
|
extra_fields.insert(
|
|
"client_family".to_string(),
|
|
Value::String(client_family.to_ascii_lowercase()),
|
|
);
|
|
}
|
|
extra_fields.insert(
|
|
CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD.to_string(),
|
|
value,
|
|
);
|
|
}
|
|
if let Some(incoming_tls) =
|
|
crate::ai_serving::tls_fingerprint_from_headers(parts.original_headers)
|
|
{
|
|
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
|
|
}
|
|
insert_pool_key_lease_report_context_fields(&mut extra_fields, parts.pool_key_lease);
|
|
insert_scheduler_affinity_policy_report_context_field(&mut extra_fields, parts.routing_policy);
|
|
if let Some(override_policy) = parts
|
|
.routing_policy
|
|
.and_then(|policy| policy.pool_policy_overrides.get(parts.provider_id))
|
|
.filter(|override_policy| !override_policy.scheduling_presets.is_empty())
|
|
{
|
|
if let Ok(value) = serde_json::to_value(override_policy) {
|
|
extra_fields.insert(ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD.to_string(), value);
|
|
}
|
|
}
|
|
if let Some(epoch) = parts.scheduler_affinity_epoch {
|
|
extra_fields.insert(
|
|
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD.to_string(),
|
|
Value::Number(epoch.into()),
|
|
);
|
|
}
|
|
if let Some(sticky_key_attempts) = parts.sticky_key_attempts {
|
|
extra_fields.insert(
|
|
STICKY_KEY_ATTEMPTS_REPORT_FIELD.to_string(),
|
|
Value::Number(sticky_key_attempts.into()),
|
|
);
|
|
}
|
|
insert_request_path_fields(
|
|
&mut extra_fields,
|
|
parts.request_path,
|
|
parts.request_query_string,
|
|
);
|
|
let client_requested_stream = parts.client_requested_stream
|
|
|| parts
|
|
.request_path
|
|
.is_some_and(request_path_implies_stream_request);
|
|
|
|
build_ai_execution_report_context(AiExecutionReportContextParts {
|
|
auth_context: parts.auth_context,
|
|
request_id: parts.request_id,
|
|
candidate_id: parts.candidate_id,
|
|
candidate_index: parts.attempt_identity.candidate_index,
|
|
retry_index: parts.attempt_identity.retry_index,
|
|
pool_key_index: parts.attempt_identity.pool_key_index,
|
|
model: parts.model,
|
|
provider_name: parts.provider_name,
|
|
provider_id: parts.provider_id,
|
|
endpoint_id: parts.endpoint_id,
|
|
key_id: parts.key_id,
|
|
key_name: parts.key_name,
|
|
model_id: parts.model_id,
|
|
global_model_id: parts.global_model_id,
|
|
global_model_name: parts.global_model_name,
|
|
provider_api_format: parts.provider_api_format,
|
|
client_api_format: parts.client_api_format,
|
|
mapped_model: parts.mapped_model,
|
|
candidate_group_id: parts.candidate_group_id,
|
|
ranking: parts.ranking,
|
|
upstream_url: parts.upstream_url,
|
|
header_rules: parts.header_rules,
|
|
body_rules: parts.body_rules,
|
|
provider_request_method: parts.provider_request_method,
|
|
provider_request_headers: parts.provider_request_headers,
|
|
original_headers: &original_headers,
|
|
original_request_body,
|
|
request_origin: AiRequestOrigin {
|
|
client_ip,
|
|
user_agent,
|
|
},
|
|
client_requested_stream,
|
|
upstream_is_stream: parts.upstream_is_stream,
|
|
has_envelope: parts.has_envelope,
|
|
needs_conversion: parts.needs_conversion,
|
|
extra_fields,
|
|
})
|
|
}
|
|
|
|
fn insert_request_path_fields(
|
|
extra_fields: &mut Map<String, Value>,
|
|
request_path: Option<&str>,
|
|
request_query_string: Option<&str>,
|
|
) {
|
|
let Some(path) = request_path.and_then(sanitize_request_path) else {
|
|
return;
|
|
};
|
|
let query = request_query_string.and_then(sanitize_request_query_string);
|
|
let path_and_query = sanitize_request_path_and_query(path.as_str(), query.as_deref())
|
|
.unwrap_or_else(|| path.clone());
|
|
extra_fields
|
|
.entry("request_path".to_string())
|
|
.or_insert_with(|| Value::String(path.clone()));
|
|
if let Some(query) = query.clone() {
|
|
extra_fields
|
|
.entry("request_query_string".to_string())
|
|
.or_insert_with(|| Value::String(query.to_string()));
|
|
}
|
|
extra_fields
|
|
.entry("request_path_and_query".to_string())
|
|
.or_insert_with(|| Value::String(path_and_query));
|
|
}
|
|
|
|
pub(crate) fn provider_stream_event_api_format_for_provider_type(
|
|
provider_type: &str,
|
|
) -> Option<&'static str> {
|
|
ai_provider_stream_event_api_format_for_provider_type(provider_type)
|
|
}
|
|
|
|
pub(crate) fn insert_provider_stream_event_api_format(
|
|
extra_fields: &mut Map<String, Value>,
|
|
provider_type: &str,
|
|
) {
|
|
insert_ai_provider_stream_event_api_format(extra_fields, provider_type);
|
|
}
|
|
|
|
pub(crate) fn insert_native_client_envelope_name(
|
|
extra_fields: &mut Map<String, Value>,
|
|
envelope_name: &str,
|
|
request_path: &str,
|
|
) {
|
|
if envelope_name.eq_ignore_ascii_case("antigravity:v1internal")
|
|
&& request_path == "/v1internal:streamGenerateContent"
|
|
{
|
|
extra_fields.insert(
|
|
"client_envelope_name".to_string(),
|
|
Value::String(envelope_name.to_string()),
|
|
);
|
|
}
|
|
}
|
|
|
|
fn merge_incoming_tls_fingerprint(extra_fields: &mut Map<String, Value>, incoming_tls: Value) {
|
|
let entry = extra_fields
|
|
.entry("tls_fingerprint".to_string())
|
|
.or_insert_with(|| Value::Object(Map::new()));
|
|
if let Value::Object(object) = entry {
|
|
object.insert("incoming".to_string(), incoming_tls);
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::collections::BTreeMap;
|
|
|
|
use aether_scheduler_core::ClientSessionAffinity;
|
|
use serde_json::{json, Map, Value};
|
|
|
|
use super::{
|
|
build_local_execution_report_context, provider_stream_event_api_format_for_provider_type,
|
|
LocalExecutionReportContextParts,
|
|
};
|
|
use crate::ai_serving::ExecutionRuntimeAuthContext;
|
|
use crate::ai_serving::RequestOrigin;
|
|
use crate::orchestration::ExecutionAttemptIdentity;
|
|
|
|
#[test]
|
|
fn codex_provider_uses_openai_responses_stream_event_format() {
|
|
assert_eq!(
|
|
provider_stream_event_api_format_for_provider_type("codex"),
|
|
Some("openai:responses")
|
|
);
|
|
assert_eq!(
|
|
provider_stream_event_api_format_for_provider_type("CODEX"),
|
|
Some("openai:responses")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn ordinary_providers_do_not_override_stream_event_format() {
|
|
assert_eq!(
|
|
provider_stream_event_api_format_for_provider_type("openai"),
|
|
None
|
|
);
|
|
assert_eq!(
|
|
provider_stream_event_api_format_for_provider_type("anthropic"),
|
|
None
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn local_execution_report_context_records_request_origin_and_session_affinity() {
|
|
let auth_context = ExecutionRuntimeAuthContext {
|
|
user_id: "user-1".to_string(),
|
|
api_key_id: "api-key-1".to_string(),
|
|
username: None,
|
|
api_key_name: None,
|
|
balance_remaining: None,
|
|
access_allowed: true,
|
|
api_key_is_standalone: false,
|
|
};
|
|
let original_headers = http::HeaderMap::new();
|
|
let provider_request_headers = BTreeMap::new();
|
|
let client_session_affinity = ClientSessionAffinity::new(
|
|
Some("codex".to_string()),
|
|
Some("account=account-1;session=session-1".to_string()),
|
|
);
|
|
|
|
let report_context =
|
|
build_local_execution_report_context(LocalExecutionReportContextParts {
|
|
auth_context: &auth_context,
|
|
request_id: "trace-1",
|
|
candidate_id: "candidate-1",
|
|
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
|
|
model: "gpt-5",
|
|
provider_name: "OpenAI",
|
|
provider_id: "provider-1",
|
|
endpoint_id: "endpoint-1",
|
|
key_id: "key-1",
|
|
key_name: None,
|
|
model_id: None,
|
|
global_model_id: None,
|
|
global_model_name: None,
|
|
provider_api_format: "openai:chat",
|
|
client_api_format: "openai:chat",
|
|
mapped_model: None,
|
|
candidate_group_id: None,
|
|
pool_key_lease: None,
|
|
ranking: None,
|
|
upstream_url: None,
|
|
header_rules: None,
|
|
body_rules: None,
|
|
provider_request_method: None,
|
|
provider_request_headers: Some(&provider_request_headers),
|
|
original_headers: &original_headers,
|
|
request_path: Some("/v1/chat/completions"),
|
|
request_query_string: Some("debug=true&limit=10"),
|
|
request_origin: Some(RequestOrigin {
|
|
client_ip: Some("203.0.113.8".to_string()),
|
|
user_agent: Some("Claude-Code/1.0".to_string()),
|
|
}),
|
|
original_request_body_json: Some(&json!({"model": "gpt-5"})),
|
|
original_request_body_base64: None,
|
|
client_session_affinity: Some(&client_session_affinity),
|
|
routing_policy: None,
|
|
scheduler_affinity_epoch: None,
|
|
sticky_key_attempts: None,
|
|
client_requested_stream: false,
|
|
upstream_is_stream: false,
|
|
has_envelope: false,
|
|
needs_conversion: false,
|
|
extra_fields: Map::new(),
|
|
});
|
|
|
|
assert_eq!(
|
|
report_context["client_ip"],
|
|
Value::String("203.0.113.8".to_string())
|
|
);
|
|
assert_eq!(
|
|
report_context["user_agent"],
|
|
Value::String("Claude-Code/1.0".to_string())
|
|
);
|
|
assert_eq!(
|
|
report_context["client_session_affinity"],
|
|
json!({
|
|
"client_family": "codex",
|
|
"session_key": "account=account-1;session=session-1"
|
|
})
|
|
);
|
|
assert_eq!(report_context["request_path"], "/v1/chat/completions");
|
|
assert_eq!(report_context["request_query_string"], "limit=10");
|
|
assert_eq!(
|
|
report_context["request_path_and_query"],
|
|
"/v1/chat/completions?limit=10"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn local_execution_report_context_treats_stream_generate_content_path_as_client_stream() {
|
|
let auth_context = ExecutionRuntimeAuthContext {
|
|
user_id: "user-1".to_string(),
|
|
api_key_id: "api-key-1".to_string(),
|
|
username: None,
|
|
api_key_name: None,
|
|
balance_remaining: None,
|
|
access_allowed: true,
|
|
api_key_is_standalone: false,
|
|
};
|
|
let original_headers = http::HeaderMap::new();
|
|
let provider_request_headers = BTreeMap::new();
|
|
|
|
let report_context =
|
|
build_local_execution_report_context(LocalExecutionReportContextParts {
|
|
auth_context: &auth_context,
|
|
request_id: "trace-1",
|
|
candidate_id: "candidate-1",
|
|
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
|
|
model: "gemini-3.1-flash-image-preview",
|
|
provider_name: "Gemini",
|
|
provider_id: "provider-1",
|
|
endpoint_id: "endpoint-1",
|
|
key_id: "key-1",
|
|
key_name: None,
|
|
model_id: None,
|
|
global_model_id: None,
|
|
global_model_name: None,
|
|
provider_api_format: "gemini:generate_content",
|
|
client_api_format: "gemini:generate_content",
|
|
mapped_model: None,
|
|
candidate_group_id: None,
|
|
pool_key_lease: None,
|
|
ranking: None,
|
|
upstream_url: None,
|
|
header_rules: None,
|
|
body_rules: None,
|
|
provider_request_method: None,
|
|
provider_request_headers: Some(&provider_request_headers),
|
|
original_headers: &original_headers,
|
|
request_path: Some(
|
|
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent",
|
|
),
|
|
request_query_string: Some("key=secret&alt=sse"),
|
|
request_origin: None,
|
|
original_request_body_json: Some(&json!({
|
|
"contents": [{"role": "user", "parts": [{"text": "hi"}]}]
|
|
})),
|
|
original_request_body_base64: None,
|
|
client_session_affinity: None,
|
|
routing_policy: None,
|
|
scheduler_affinity_epoch: None,
|
|
sticky_key_attempts: None,
|
|
client_requested_stream: false,
|
|
upstream_is_stream: true,
|
|
has_envelope: false,
|
|
needs_conversion: false,
|
|
extra_fields: Map::new(),
|
|
});
|
|
|
|
assert_eq!(report_context["client_requested_stream"], true);
|
|
assert_eq!(report_context["request_query_string"], "alt=sse");
|
|
assert_eq!(
|
|
report_context["request_path_and_query"],
|
|
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent?alt=sse"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn local_execution_report_context_records_forwarded_tls_fingerprint() {
|
|
let auth_context = ExecutionRuntimeAuthContext {
|
|
user_id: "user-1".to_string(),
|
|
api_key_id: "api-key-1".to_string(),
|
|
username: None,
|
|
api_key_name: None,
|
|
balance_remaining: None,
|
|
access_allowed: true,
|
|
api_key_is_standalone: false,
|
|
};
|
|
let mut original_headers = http::HeaderMap::new();
|
|
original_headers.insert("x-aether-tls-ja3", "ja3-value".parse().unwrap());
|
|
original_headers.insert("x-aether-tls-ja4", "ja4-value".parse().unwrap());
|
|
original_headers.insert("x-aether-tls-protocol", "TLSv1.3".parse().unwrap());
|
|
let provider_request_headers = BTreeMap::new();
|
|
|
|
let report_context =
|
|
build_local_execution_report_context(LocalExecutionReportContextParts {
|
|
auth_context: &auth_context,
|
|
request_id: "trace-1",
|
|
candidate_id: "candidate-1",
|
|
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
|
|
model: "gpt-5",
|
|
provider_name: "OpenAI",
|
|
provider_id: "provider-1",
|
|
endpoint_id: "endpoint-1",
|
|
key_id: "key-1",
|
|
key_name: None,
|
|
model_id: None,
|
|
global_model_id: None,
|
|
global_model_name: None,
|
|
provider_api_format: "openai:chat",
|
|
client_api_format: "openai:chat",
|
|
mapped_model: None,
|
|
candidate_group_id: None,
|
|
pool_key_lease: None,
|
|
ranking: None,
|
|
upstream_url: None,
|
|
header_rules: None,
|
|
body_rules: None,
|
|
provider_request_method: None,
|
|
provider_request_headers: Some(&provider_request_headers),
|
|
original_headers: &original_headers,
|
|
request_path: None,
|
|
request_query_string: None,
|
|
request_origin: None,
|
|
original_request_body_json: Some(&json!({"model": "gpt-5"})),
|
|
original_request_body_base64: None,
|
|
client_session_affinity: None,
|
|
routing_policy: None,
|
|
scheduler_affinity_epoch: None,
|
|
sticky_key_attempts: None,
|
|
client_requested_stream: false,
|
|
upstream_is_stream: false,
|
|
has_envelope: false,
|
|
needs_conversion: false,
|
|
extra_fields: Map::new(),
|
|
});
|
|
|
|
assert_eq!(
|
|
report_context["tls_fingerprint"]["incoming"],
|
|
json!({
|
|
"source": "forwarded_header",
|
|
"ja3": "ja3-value",
|
|
"ja4": "ja4-value",
|
|
"protocol": "TLSv1.3"
|
|
})
|
|
);
|
|
}
|
|
}
|