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

Merge branch 'aether-rust-pioneer' of https://github.com/fawney19/Aether into codex/codex-image-progress-heartbeat

Browse Source 
# Conflicts:
#	frontend/src/features/usage/components/__tests__/HorizontalRequestTimeline.spec.ts
This commit is contained in:
Entropy.Xu
2026-05-09 01:26:44 +08:00
parent 3b4f27f767 4cf0de681a
commit eda3738d59
90 changed files with 8542 additions and 1252 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
apps/aether-gateway/src/ai_serving/finalize/internal/sync_finalize.rs
View File

@@ -92,6 +92,7 @@ fn maybe_build_local_openai_image_sync_finalize_response(
payload.report_kind.as_str(),
payload.status_code,
payload.report_context.as_ref(),
payload.body_json.as_ref(),
payload.body_base64.as_deref(),
)
.map_err(GatewayError::from)?

12
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
View File

@@ -75,6 +75,18 @@ enum LocalExecutionCandidateAttemptSourceItem<'a> {
}
impl<'a> LocalExecutionCandidateAttemptSource<'a> {
pub(crate) fn from_static_attempts_for_image_bridge(
attempts: Vec<LocalExecutionCandidateAttempt>,
) -> Self {
let mut items = VecDeque::new();
if !attempts.is_empty() {
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static {
attempts: VecDeque::from(attempts),
});
}
Self { items }
}
pub(crate) async fn next_attempt(&mut self) -> Option<LocalExecutionCandidateAttempt> {
loop {
let front = self.items.front_mut()?;

47
apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs
View File

@@ -398,15 +398,11 @@ pub(crate) fn candidate_auth_channel_skip_reason(
) -> Option<&'static str> {
let request_auth_channel = normalize_request_auth_channel(request_auth_channel?)?;
let upstream_auth_channel = resolve_transport_request_auth_channel(transport)?;
if request_auth_channel == upstream_auth_channel
|| provider_runtime_policy(&transport.provider.provider_type)
.allow_auth_channel_mismatch_by_default
|| allow_auth_channel_mismatch_for_format(transport)
{
None
} else {
Some("auth_channel_mismatch")
if request_auth_channel == upstream_auth_channel {
return None;
}
auth_channel_mismatch_is_explicitly_disabled_for_format(transport)
.then_some("auth_channel_mismatch")
}
fn normalize_request_auth_channel(value: &str) -> Option<&'static str> {
@@ -452,19 +448,22 @@ fn resolve_transport_auth_type_for_endpoint_format(
.unwrap_or(default_auth_type)
}
fn allow_auth_channel_mismatch_for_format(transport: &GatewayProviderTransportSnapshot) -> bool {
fn auth_channel_mismatch_is_explicitly_disabled_for_format(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
let api_format = crate::ai_serving::normalize_api_format_alias(&transport.endpoint.api_format);
transport
let Some(items) = transport
.key
.allow_auth_channel_mismatch_formats
.as_ref()
.and_then(serde_json::Value::as_array)
.is_some_and(|items| {
items
.iter()
.filter_map(serde_json::Value::as_str)
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
})
else {
return false;
};
!items
.iter()
.filter_map(serde_json::Value::as_str)
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
}
pub(crate) async fn read_candidate_transport_snapshot(
@@ -585,11 +584,11 @@ mod tests {
}
#[test]
fn auth_channel_gate_skips_mismatched_raw_secret_auth() {
fn auth_channel_gate_allows_mismatched_raw_secret_auth_by_default() {
let transport = sample_transport("bearer");
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
None
);
}
@@ -603,6 +602,16 @@ mod tests {
);
}
#[test]
fn auth_channel_gate_blocks_explicitly_disabled_mismatch_format() {
let mut transport = sample_transport("bearer");
transport.key.allow_auth_channel_mismatch_formats = Some(json!(["openai:responses"]));
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
);
}
#[test]
fn auth_channel_gate_treats_cli_oauth_provider_as_bearer_like() {
let mut transport = sample_transport("oauth");
@@ -613,7 +622,7 @@ mod tests {
);
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
None
);
}

34
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
View File

@@ -159,6 +159,40 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
.into_iter()
.map(str::to_string)
.collect::<Vec<_>>();
preselect_local_execution_candidates_for_api_formats_with_serving(
state,
client_api_format,
requested_model,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
use_api_format_alias_match,
key_mode,
candidate_api_formats,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_serving(
state: PlannerAppState<'_>,
client_api_format: &str,
requested_model: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: &GatewayAuthApiKeySnapshot,
client_session_affinity: Option<&ClientSessionAffinity>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
candidate_api_formats: Vec<String>,
) -> Result<
AiCandidatePreselectionOutcome<
SchedulerMinimalCandidateSelectionCandidate,
SkippedLocalExecutionCandidate,
>,
GatewayError,
> {
let mut model_directive_enabled_api_formats = BTreeSet::new();
for api_format in &candidate_api_formats {
if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(

33
apps/aether-gateway/src/ai_serving/planner/specialized/image.rs
View File

@@ -7,6 +7,7 @@ use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
AiStreamAttempt, AiSyncAttempt,
};
@@ -314,7 +315,13 @@ impl LocalOpenAiImageSyncAttemptSource<'_> {
return Ok(None);
};
match build_passthrough_sync_plan_from_decision(self.parts, payload) {
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_sync_plan_from_decision(self.parts, self.body_json, payload)
} else {
build_passthrough_sync_plan_from_decision(self.parts, payload)
};
match built {
Ok(value) => Ok(value),
Err(err) => {
warn!(
@@ -350,7 +357,13 @@ impl LocalOpenAiImageStreamAttemptSource<'_> {
return Ok(None);
};
match build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false) {
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_stream_plan_from_decision(self.parts, self.body_json, payload)
} else {
build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false)
};
match built {
Ok(value) => Ok(value),
Err(err) => {
warn!(
@@ -538,7 +551,13 @@ async fn build_local_sync_plan_and_reports(
continue;
};
match build_passthrough_sync_plan_from_decision(parts, payload) {
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_sync_plan_from_decision(parts, body_json, payload)
} else {
build_passthrough_sync_plan_from_decision(parts, payload)
};
match built {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
@@ -608,7 +627,13 @@ async fn build_local_stream_plan_and_reports(
continue;
};
match build_standard_stream_plan_from_decision(parts, body_json, payload, false) {
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_stream_plan_from_decision(parts, body_json, payload)
} else {
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
};
match built {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {

88
apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs
View File

@@ -10,7 +10,7 @@ use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
};
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
use crate::{AiExecutionDecision, AppState};
use crate::{append_execution_contract_fields_to_value, AiExecutionDecision, AppState};
use super::request::resolve_local_openai_image_candidate_payload_parts;
use super::support::{LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput};
@@ -47,8 +47,12 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
} = attempt;
let candidate = eligible.candidate;
let transport = resolved.transport;
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
let provider_api_format = resolved.provider_api_format.clone();
let needs_conversion = provider_api_format != spec_metadata.api_format;
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
provider_api_format.as_str(),
);
let proxy = planner_state
.app()
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
@@ -84,41 +88,47 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
.get("stream")
.and_then(serde_json::Value::as_bool)
.unwrap_or(spec_metadata.require_streaming);
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
attempt_identity,
model: &resolved.requested_model,
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),
body_rules: transport.endpoint.body_rules.as_ref(),
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,
needs_conversion: false,
extra_fields,
});
let report_context = append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
attempt_identity,
model: &resolved.requested_model,
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),
body_rules: transport.endpoint.body_rules.as_ref(),
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,
needs_conversion,
extra_fields,
}),
execution_strategy,
conversion_mode,
spec_metadata.api_format,
provider_api_format.as_str(),
);
Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
@@ -137,7 +147,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(resolved.auth_header),
auth_value: Some(resolved.auth_value),
provider_api_format: spec_metadata.api_format.to_string(),
provider_api_format,
client_api_format: spec_metadata.api_format.to_string(),
model_name: resolved.requested_model,
mapped_model: resolved.mapped_model,

215
apps/aether-gateway/src/ai_serving/planner/specialized/image/request.rs
View File

@@ -9,14 +9,17 @@ use crate::ai_serving::planner::candidate_preparation::{
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::transport::{
build_openai_image_headers, build_openai_image_upstream_url,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput,
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
resolve_openai_image_auth, ProviderOpenAiImageHeadersInput,
StandardProviderRequestHeadersInput,
};
use crate::ai_serving::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
build_chatgpt_web_image_request_body, build_openai_image_provider_request_body,
default_model_for_openai_image_operation, normalize_openai_image_request,
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, PlannerAppState,
build_chatgpt_web_image_request_body,
build_gemini_image_request_body_from_openai_image_request,
build_openai_image_provider_request_body, default_model_for_openai_image_operation,
normalize_openai_image_request, request_conversion_direct_auth, CandidateFailureDiagnostic,
GatewayProviderTransportSnapshot, PlannerAppState, RequestConversionKind,
};
use crate::AppState;
@@ -35,6 +38,7 @@ pub(super) struct LocalOpenAiImageCandidatePayloadParts {
pub(super) auth_value: String,
pub(super) requested_model: String,
pub(super) mapped_model: String,
pub(super) provider_api_format: String,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Value,
pub(super) upstream_url: String,
@@ -54,6 +58,21 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
let spec_metadata = local_openai_image_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = attempt.eligible.provider_api_format.as_str();
if provider_api_format == "gemini:generate_content" {
return resolve_local_openai_image_to_gemini_candidate_payload_parts(
state,
parts,
body_json,
body_base64,
trace_id,
input,
attempt,
spec,
)
.await;
}
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, spec_metadata.api_format)
@@ -70,7 +89,6 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
transport,
@@ -216,6 +234,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
auth_value,
requested_model,
mapped_model,
provider_api_format: spec_metadata.api_format.to_string(),
provider_request_headers,
provider_request_body,
upstream_url,
@@ -223,6 +242,190 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
})
}
async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
body_json: &Value,
body_base64: Option<&str>,
trace_id: &str,
input: &LocalOpenAiImageDecisionInput,
attempt: &LocalOpenAiImageCandidateAttempt,
spec: LocalOpenAiImageSpec,
) -> Option<LocalOpenAiImageCandidatePayloadParts> {
let spec_metadata = local_openai_image_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = "gemini:generate_content";
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
transport,
candidate,
request_conversion_direct_auth(transport, RequestConversionKind::ToGeminiStandard),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_image_to_gemini_candidate_request",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_image_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(normalized_request) = normalize_openai_image_request(parts, body_json, body_base64)
else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_request_normalize",
),
)
.await;
return None;
};
let Some(mut converted) = build_gemini_image_request_body_from_openai_image_request(
&normalized_request,
&prepared_candidate.mapped_model,
) else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_request_body",
),
)
.await;
return None;
};
converted.body_json =
match crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers(
converted.body_json,
transport.endpoint.body_rules.as_ref(),
body_json,
&parts.headers,
) {
Some(body) => body,
None => {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_body_rules",
),
)
.await;
return None;
}
};
let upstream_is_stream = spec_metadata.require_streaming;
let Some(upstream_url) = crate::ai_serving::planner::standard::build_standard_upstream_url(
parts,
transport,
&converted.mapped_model,
provider_api_format,
upstream_is_stream,
) else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_url",
),
)
.await;
return None;
};
let Some(resolved_headers) =
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
transport,
provider_api_format,
same_format: false,
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
extra_headers: &BTreeMap::new(),
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &converted.body_json,
original_request_body: body_json,
upstream_is_stream,
})
else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_headers",
),
)
.await;
return None;
};
Some(LocalOpenAiImageCandidatePayloadParts {
transport: Arc::clone(transport),
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
requested_model: converted.requested_model,
mapped_model: converted.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_headers: resolved_headers.headers,
provider_request_body: converted.body_json,
upstream_url,
input_summary: converted.summary_json,
})
}
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
let base_url = if base_url.is_empty() {

225
apps/aether-gateway/src/ai_serving/planner/specialized/image/support.rs
View File

@@ -11,6 +11,7 @@ use crate::ai_serving::planner::candidate_metadata::{
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_serving::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
@@ -19,9 +20,9 @@ use crate::ai_serving::planner::materialization_policy::{
};
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::{
extract_pool_sticky_session_token, resolve_local_decision_execution_runtime_auth_context,
CandidateFailureDiagnostic, ExecutionRuntimeAuthContext, GatewayControlDecision,
PlannerAppState,
extract_pool_sticky_session_token, request_candidate_api_formats,
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;
@@ -88,52 +89,69 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
api_format: &str,
decision_kind: &str,
) -> Option<Vec<LocalOpenAiImageCandidateAttempt>> {
let planner_state = PlannerAppState::new(state);
let (candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
return None;
let candidate_api_formats = image_candidate_api_formats(api_format);
let mut attempts = Vec::new();
for candidate_api_format in candidate_api_formats {
let matches_client_format = candidate_api_format == api_format;
let planner_state = PlannerAppState::new(state);
let (mut candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
candidate_api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
matches_client_format.then_some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
api_format = candidate_api_format,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
continue;
}
};
if !matches_client_format {
candidates.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
candidate,
false,
)
});
}
};
attempts.extend(
materialize_local_openai_image_candidate_attempts(
planner_state,
trace_id,
input,
body_json,
candidates,
preselection_skipped
.into_iter()
.map(|item| SkippedLocalExecutionCandidate {
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),
api_format,
)
.await,
);
}
Some(
materialize_local_openai_image_candidate_attempts(
planner_state,
trace_id,
input,
body_json,
candidates,
preselection_skipped
.into_iter()
.map(|item| SkippedLocalExecutionCandidate {
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),
api_format,
)
.await,
)
Some(attempts)
}
pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
@@ -145,29 +163,56 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
decision_kind: &str,
) -> Result<Option<(LocalOpenAiImageCandidateAttemptSource<'a>, usize)>, GatewayError> {
let planner_state = PlannerAppState::new(state);
let (candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
return Ok(None);
let mut candidates = Vec::new();
let mut preselection_skipped = Vec::new();
for candidate_api_format in image_candidate_api_formats(api_format) {
let matches_client_format = candidate_api_format == api_format;
match planner_state
.list_selectable_candidates_with_skip_reasons(
candidate_api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
matches_client_format.then_some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok((mut format_candidates, mut format_skipped)) => {
if !matches_client_format {
format_candidates.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
candidate,
false,
)
});
format_skipped.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
&candidate.candidate,
false,
)
});
}
candidates.append(&mut format_candidates);
preselection_skipped.append(&mut format_skipped);
}
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
api_format = candidate_api_format,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
continue;
}
}
};
}
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
@@ -198,23 +243,34 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
extra_data: None,
})
.collect(),
LocalCandidateResolutionMode::Standard,
LocalCandidateResolutionMode::WithoutTransportPairGate,
|eligible| {
Some(build_local_execution_candidate_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: api_format,
provider_api_format: eligible.provider_api_format.as_str(),
client_api_format: api_format,
extra_fields: serde_json::Map::new(),
},
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
skipped_candidate.extra_data =
Some(build_local_execution_candidate_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
api_format,
provider_api_format.as_str(),
api_format,
serde_json::Map::new(),
));
@@ -254,23 +310,34 @@ async fn materialize_local_openai_image_candidate_attempts(
persistence_policy,
candidates,
preselection_skipped,
LocalCandidateResolutionMode::Standard,
LocalCandidateResolutionMode::WithoutTransportPairGate,
|eligible| {
Some(build_local_execution_candidate_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: api_format,
provider_api_format: eligible.provider_api_format.as_str(),
client_api_format: api_format,
extra_fields: serde_json::Map::new(),
},
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
skipped_candidate.extra_data =
Some(build_local_execution_candidate_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
api_format,
provider_api_format.as_str(),
api_format,
serde_json::Map::new(),
));
@@ -282,6 +349,14 @@ async fn materialize_local_openai_image_candidate_attempts(
outcome.attempts
}
fn image_candidate_api_formats(api_format: &str) -> Vec<&'static str> {
if api_format.trim().eq_ignore_ascii_case("openai:image") {
vec!["openai:image", "gemini:generate_content"]
} else {
request_candidate_api_formats(api_format, false)
}
}
pub(super) async fn mark_skipped_local_openai_image_candidate(
state: &AppState,
input: &LocalOpenAiImageDecisionInput,

89
apps/aether-gateway/src/ai_serving/planner/standard/family/candidates.rs
View File

@@ -11,6 +11,7 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_source::{
preselect_local_execution_candidates_for_api_formats_with_serving,
preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
};
use crate::ai_serving::planner::common::extract_requested_model_from_request;
@@ -23,7 +24,8 @@ use crate::ai_serving::planner::materialization_policy::{
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
gemini_request_is_image_generation, resolve_local_decision_execution_runtime_auth_context,
GatewayControlDecision, PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::{AppState, GatewayError};
@@ -86,6 +88,8 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
spec: LocalStandardSpec,
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
let is_gemini_image_bridge = spec_metadata.api_format == "gemini:generate_content"
&& gemini_request_is_image_generation(body_json);
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(
@@ -105,6 +109,17 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
)
.await?;
let (candidates, skipped_candidates) = maybe_append_gemini_image_openai_image_preselection(
state,
planner_state,
trace_id,
input,
body_json,
spec,
preselection.candidates,
preselection.skipped_candidates,
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
planner_state,
trace_id,
@@ -116,9 +131,13 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::Standard,
candidates,
skipped_candidates,
if is_gemini_image_bridge {
LocalCandidateResolutionMode::WithoutTransportPairGate
} else {
LocalCandidateResolutionMode::Standard
},
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
@@ -188,7 +207,18 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::StandardDecision,
);
Ok(
if spec_metadata.api_format == "gemini:generate_content"
&& gemini_request_is_image_generation(body_json)
{
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, input, body_json, spec)
.await?;
let source =
LocalExecutionCandidateAttemptSource::from_static_attempts_for_image_bridge(attempts);
return Ok((source, candidate_count));
}
let (source, candidate_count) =
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
@@ -253,6 +283,51 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
skipped_candidate
},
)
.await,
)
.await;
Ok((source, candidate_count))
}
#[allow(clippy::too_many_arguments)]
async fn maybe_append_gemini_image_openai_image_preselection(
state: &AppState,
planner_state: PlannerAppState<'_>,
trace_id: &str,
input: &LocalStandardDecisionInput,
body_json: &serde_json::Value,
spec: LocalStandardSpec,
mut candidates: Vec<aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
mut skipped_candidates: Vec<
crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate,
>,
) -> Result<
(
Vec<aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
Vec<crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate>,
),
GatewayError,
> {
let spec_metadata = local_standard_spec_metadata(spec);
if spec_metadata.api_format != "gemini:generate_content"
|| !gemini_request_is_image_generation(body_json)
{
return Ok((candidates, skipped_candidates));
}
let image_preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.client_session_affinity.as_ref(),
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
vec!["openai:image".to_string()],
)
.await?;
candidates.extend(image_preselection.candidates);
skipped_candidates.extend(image_preselection.skipped_candidates);
let _ = (state, trace_id);
Ok((candidates, skipped_candidates))
}

151
apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs
View File

@@ -18,10 +18,13 @@ use crate::ai_serving::transport::kiro::{
KIRO_ENVELOPE_NAME,
};
use crate::ai_serving::transport::{
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
StandardProviderRequestHeadersInput,
build_kiro_cross_format_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_standard_provider_request_headers,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput,
};
use crate::ai_serving::{
build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation,
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use crate::AppState;
@@ -59,6 +62,15 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = attempt.eligible.provider_api_format.as_str();
if spec_metadata.api_format == "gemini:generate_content"
&& provider_api_format == "openai:image"
&& gemini_request_is_image_generation(body_json)
{
return resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state, parts, trace_id, body_json, input, attempt,
)
.await;
}
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
let Some(conversion_kind) =
crate::ai_serving::request_conversion_kind(spec_metadata.api_format, provider_api_format)
@@ -331,6 +343,141 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
})
}
async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalStandardDecisionInput,
attempt: &LocalStandardCandidateAttempt,
) -> Option<LocalStandardCandidatePayloadParts> {
let client_api_format = "gemini:generate_content";
let provider_api_format = "openai:image";
let planner_state = crate::ai_serving::PlannerAppState::new(state);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, provider_api_format)
{
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
resolve_openai_image_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "gemini_image_to_openai_image_candidate_request",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(converted) = build_openai_image_request_body_from_gemini_image_request(
body_json,
parts.uri.path(),
&prepared_candidate.mapped_model,
) else {
mark_skipped_local_standard_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
body_json,
client_api_format,
provider_api_format,
),
)
.await;
return None;
};
let upstream_is_stream = true;
let upstream_url = build_openai_image_upstream_url(transport, None);
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &converted.body_json,
original_request_body: body_json,
})
else {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
client_api_format,
provider_api_format,
"gemini_image_to_openai_image_headers",
),
)
.await;
return None;
};
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&converted.body_json,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
Some(LocalStandardCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: converted.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body: converted.body_json,
provider_request_headers,
upstream_url,
upstream_is_stream,
envelope_name: None,
transport: Arc::clone(transport),
})
}
#[allow(clippy::too_many_arguments)]
async fn build_kiro_cross_format_payload_parts(
state: &AppState,

42
apps/aether-gateway/src/ai_serving/pure/mod.rs
View File

@@ -10,14 +10,20 @@ pub(crate) use aether_ai_formats::api::{
build_cross_format_openai_chat_request_body_with_model_directives,
build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_request_body_with_model_directives,
build_generated_tool_call_id, build_kiro_final_message_sse_events,
build_kiro_initial_sse_events, build_kiro_stream_error_sse_events,
build_local_openai_chat_request_body,
build_gemini_image_request_body_from_openai_image_request,
build_gemini_image_response_from_openai_image_response,
build_gemini_image_response_from_openai_responses_image_response, build_generated_tool_call_id,
build_kiro_final_message_sse_events, build_kiro_initial_sse_events,
build_kiro_stream_error_sse_events, build_local_openai_chat_request_body,
build_local_openai_chat_request_body_with_model_directives,
build_local_openai_responses_request_body,
build_local_openai_responses_request_body_with_model_directives,
build_local_success_background_report, build_local_success_conversion_background_report,
build_openai_image_provider_request_body, build_openai_responses_response,
build_openai_image_provider_body_from_response_stream_sync_body,
build_openai_image_provider_request_body,
build_openai_image_request_body_from_gemini_image_request,
build_openai_image_response_from_gemini_response,
build_openai_image_response_from_response_stream_sync_body, build_openai_responses_response,
build_standard_request_body, build_standard_request_body_from_canonical,
build_standard_request_body_from_canonical_with_model_directives,
build_standard_request_body_with_model_directives,
@@ -37,10 +43,10 @@ pub(crate) use aether_ai_formats::api::{
encode_kiro_sse_events, estimate_kiro_tokens, extract_openai_text_content,
find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end,
find_kiro_real_thinking_start_tag, force_upstream_streaming_for_provider,
implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
is_matching_stream_http_request, is_matching_stream_request, is_openai_image_stream_request,
is_openai_responses_family_format, is_openai_responses_format, kiro_crc32,
map_claude_stop_reason, map_openai_reasoning_effort_to_claude_output,
gemini_request_is_image_generation, implicit_sync_finalize_report_kind,
is_core_error_finalize_kind, is_matching_stream_http_request, is_matching_stream_request,
is_openai_image_stream_request, is_openai_responses_family_format, is_openai_responses_format,
kiro_crc32, map_claude_stop_reason, map_openai_reasoning_effort_to_claude_output,
map_openai_reasoning_effort_to_gemini_budget, maybe_bridge_standard_sync_json_to_stream,
maybe_build_ai_surface_stream_rewriter,
maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload,
@@ -72,6 +78,7 @@ pub(crate) use aether_ai_formats::api::{
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens,
resolve_openai_responses_stream_spec, resolve_openai_responses_sync_spec,
resolve_requested_gemini_image_model_for_request,
resolve_requested_openai_image_model_for_request, stream_body_contains_error_event,
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
@@ -79,15 +86,16 @@ pub(crate) use aether_ai_formats::api::{
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
ChatGptWebImageRequestError, ClaudeClientEmitter, ClaudeProviderState,
ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode, FormatContext, GeminiClientEmitter,
GeminiProviderState, KiroToClaudeCliStreamState, LocalCoreSyncErrorKind, LocalGeminiFilesSpec,
LocalOpenAiImageSpec, LocalOpenAiResponsesSpec, LocalSameFormatProviderFamily,
LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode,
LocalStandardSpec, LocalSyncReportParts, LocalVideoCreateFamily, LocalVideoCreateSpec,
NormalizedOpenAiImageRequest, OpenAIChatClientEmitter, OpenAIChatProviderState,
OpenAIResponsesClientEmitter, OpenAIResponsesProviderState, OpenAiImageOperation,
OpenAiImageResponseFormat, OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct,
ProviderAdaptationDescriptor, ProviderAdaptationSurface, ProviderPrivateStreamNormalizer,
RequestConversionKind, StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
GeminiImageRequestForOpenAi, GeminiProviderState, KiroToClaudeCliStreamState,
LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiImageSpec, LocalOpenAiResponsesSpec,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
LocalStandardSourceMode, LocalStandardSpec, LocalSyncReportParts, LocalVideoCreateFamily,
LocalVideoCreateSpec, NormalizedOpenAiImageRequest, OpenAIChatClientEmitter,
OpenAIChatProviderState, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
OpenAiImageOperation, OpenAiImageRequestForGemini, OpenAiImageResponseFormat,
OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct, ProviderAdaptationDescriptor,
ProviderAdaptationSurface, ProviderPrivateStreamNormalizer, RequestConversionKind,
StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
StreamingStandardFormatMatrix, SyncChatResponseConversionKind, SyncCliResponseConversionKind,
SyncToStreamBridgeOutcome, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, CLAUDE_CHAT_STREAM_PLAN_KIND,
CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND,

28
apps/aether-gateway/src/control/route/admin/operations_families.rs
View File

@@ -113,6 +113,34 @@ pub(super) fn classify_admin_operations_family_route(
"admin:proxy_nodes",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
"/api/admin/proxy-nodes/metrics/fleet" | "/api/admin/proxy-nodes/metrics/fleet/"
)
{
Some(classified(
"admin_proxy",
"proxy_nodes_manage",
"list_fleet_metrics",
"admin:proxy_nodes",
false,
))
} else if method == http::Method::GET
&& normalized_path_no_trailing.starts_with("/api/admin/proxy-nodes/")
&& normalized_path_no_trailing.ends_with("/metrics")
&& normalized_path_no_trailing["/api/admin/proxy-nodes/".len()..]
.split('/')
.count()
== 2
{
Some(classified(
"admin_proxy",
"proxy_nodes_manage",
"list_node_metrics",
"admin:proxy_nodes",
false,
))
} else if method == http::Method::GET
&& normalized_path_no_trailing.starts_with("/api/admin/proxy-nodes/")
&& normalized_path_no_trailing["/api/admin/proxy-nodes/".len()..]

18
apps/aether-gateway/src/control/tests/admin_proxy_nodes.rs
View File

@@ -100,6 +100,24 @@ fn classifies_admin_proxy_nodes_detail_as_admin_proxy_route() {
);
}
#[test]
fn classifies_admin_proxy_node_metrics_as_admin_proxy_route() {
assert_proxy_nodes_admin_route(
http::Method::GET,
"/api/admin/proxy-nodes/node-1/metrics?from=1700000000&to=1700003600&step=1m",
"list_node_metrics",
);
}
#[test]
fn classifies_admin_proxy_fleet_metrics_as_admin_proxy_route() {
assert_proxy_nodes_admin_route(
http::Method::GET,
"/api/admin/proxy-nodes/metrics/fleet?from=1700000000&to=1700003600&step=1m",
"list_fleet_metrics",
);
}
#[test]
fn classifies_admin_proxy_nodes_delete_as_admin_proxy_route() {
assert_proxy_nodes_admin_route(

73
apps/aether-gateway/src/data/state/auth.rs
View File

@@ -6,7 +6,8 @@ use super::{
RegenerateManagementTokenSecret, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
StoredLdapModuleConfig, StoredManagementToken, StoredManagementTokenListPage,
StoredManagementTokenWithUser, StoredOAuthProviderConfig, StoredOAuthProviderModuleConfig,
StoredProxyNode, StoredProxyNodeEvent, StoredUserAuthRecord, StoredUserOAuthLinkSummary,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket, StoredUserAuthRecord, StoredUserOAuthLinkSummary,
StoredUserPreferenceRecord, StoredUserSessionRecord, StoredWalletSnapshot,
UpdateManagementTokenRecord, UpsertOAuthProviderConfigRecord,
};
@@ -979,6 +980,56 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_proxy_node_events_filtered(
&self,
node_id: &str,
query: &super::ProxyNodeEventQuery,
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
match &self.proxy_node_reader {
Some(repository) => {
repository
.list_proxy_node_events_filtered(node_id, query)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_proxy_node_metrics(
&self,
node_id: &str,
step: super::ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyNodeMetricsBucket>, DataLayerError> {
match &self.proxy_node_reader {
Some(repository) => {
repository
.list_proxy_node_metrics(node_id, step, from_unix_secs, to_unix_secs, limit)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_proxy_fleet_metrics(
&self,
step: super::ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyFleetMetricsBucket>, DataLayerError> {
match &self.proxy_node_reader {
Some(repository) => {
repository
.list_proxy_fleet_metrics(step, from_unix_secs, to_unix_secs, limit)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn register_proxy_node(
&self,
mutation: &ProxyNodeRegistrationMutation,
@@ -1018,6 +1069,26 @@ impl GatewayDataState {
}
}
pub(crate) async fn cleanup_proxy_node_metrics(
&self,
retain_1m_from_unix_secs: u64,
retain_1h_from_unix_secs: u64,
delete_limit: usize,
) -> Result<super::ProxyNodeMetricsCleanupSummary, DataLayerError> {
match &self.proxy_node_writer {
Some(repository) => {
repository
.cleanup_proxy_node_metrics(
retain_1m_from_unix_secs,
retain_1h_from_unix_secs,
delete_limit,
)
.await
}
None => Ok(super::ProxyNodeMetricsCleanupSummary::default()),
}
}
pub(crate) async fn apply_proxy_node_heartbeat(
&self,
mutation: &ProxyNodeHeartbeatMutation,

6
apps/aether-gateway/src/data/state/mod.rs
View File

@@ -42,10 +42,12 @@ use aether_data::repository::oauth_providers::{
UpsertOAuthProviderConfigRecord,
};
use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
StoredProxyNode, StoredProxyNodeEvent,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket,
};
pub(crate) use aether_data::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
use aether_data::repository::users::{

42
apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs
View File

@@ -104,11 +104,7 @@ pub(crate) fn normalize_allow_auth_channel_mismatch_formats(
normalized.push(serde_json::Value::String(canonical));
}
}
if normalized.is_empty() {
Ok(None)
} else {
Ok(Some(serde_json::Value::Array(normalized)))
}
Ok(Some(serde_json::Value::Array(normalized)))
}
pub(crate) fn normalize_auth_type(value: Option<&str>) -> Result<String, String> {
@@ -179,9 +175,9 @@ fn normalize_json_like_object(
#[cfg(test)]
mod tests {
use super::{
normalize_api_format_json_object_keys, normalize_api_format_list, normalize_auth_type,
normalize_auth_type_by_format, normalize_pool_advanced_config,
normalize_provider_type_input, validate_vertex_api_formats,
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
normalize_pool_advanced_config, normalize_provider_type_input, validate_vertex_api_formats,
};
use serde_json::json;
@@ -280,6 +276,36 @@ mod tests {
);
}
#[test]
fn normalize_allow_auth_channel_mismatch_formats_preserves_explicit_empty_array() {
assert_eq!(
normalize_allow_auth_channel_mismatch_formats(
Some(Vec::new()),
"allow_auth_channel_mismatch_formats",
&["claude:messages".to_string()],
)
.expect("empty array should normalize"),
Some(json!([]))
);
}
#[test]
fn normalize_allow_auth_channel_mismatch_formats_normalizes_and_dedupes_values() {
assert_eq!(
normalize_allow_auth_channel_mismatch_formats(
Some(vec![
"claude:messages".to_string(),
"CLAUDE:MESSAGES".to_string(),
" claude:messages ".to_string(),
]),
"allow_auth_channel_mismatch_formats",
&["claude:messages".to_string()],
)
.expect("format list should normalize"),
Some(json!(["claude:messages"]))
);
}
#[test]
fn validate_vertex_api_formats_uses_canonical_message_formats() {
assert!(validate_vertex_api_formats(

67
apps/aether-gateway/src/handlers/admin/request/system/proxy_nodes.rs
View File

@@ -3,8 +3,10 @@ use crate::handlers::shared::unix_secs_to_rfc3339;
use crate::maintenance::{inspect_proxy_upgrade_rollout, ProxyUpgradeRolloutStatus};
use crate::GatewayError;
use aether_admin::system::{
build_admin_proxy_node_event_payload, build_admin_proxy_node_events_payload_response,
build_admin_proxy_node_payload, build_admin_proxy_nodes_data_unavailable_response,
build_admin_proxy_fleet_metrics_payload_response, build_admin_proxy_node_event_payload,
build_admin_proxy_node_events_payload_response,
build_admin_proxy_node_metrics_payload_response, build_admin_proxy_node_payload,
build_admin_proxy_nodes_data_unavailable_response,
build_admin_proxy_nodes_invalid_status_response, build_admin_proxy_nodes_list_payload_response,
build_admin_proxy_nodes_not_found_response,
};
@@ -84,6 +86,30 @@ impl<'a> AdminAppState<'a> {
pub(crate) async fn build_admin_proxy_node_events_response(
&self,
node_id: &str,
query: &aether_data::repository::proxy_nodes::ProxyNodeEventQuery,
) -> Result<Response<Body>, GatewayError> {
if !self.has_proxy_node_reader() {
return Ok(build_admin_proxy_nodes_data_unavailable_response());
}
if self.find_proxy_node(node_id).await?.is_none() {
return Ok(build_admin_proxy_nodes_not_found_response());
}
let items = self
.app
.list_proxy_node_events_filtered(node_id, query)
.await?
.into_iter()
.map(|event| build_admin_proxy_node_event_payload(&event))
.collect::<Vec<_>>();
Ok(build_admin_proxy_node_events_payload_response(items))
}
pub(crate) async fn build_admin_proxy_node_metrics_response(
&self,
node_id: &str,
step: aether_data::repository::proxy_nodes::ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Response<Body>, GatewayError> {
if !self.has_proxy_node_reader() {
@@ -93,12 +119,37 @@ impl<'a> AdminAppState<'a> {
return Ok(build_admin_proxy_nodes_not_found_response());
}
let items = self
.list_proxy_node_events(node_id, limit)
.await?
.into_iter()
.map(|event| build_admin_proxy_node_event_payload(&event))
.collect::<Vec<_>>();
Ok(build_admin_proxy_node_events_payload_response(items))
.app
.list_proxy_node_metrics(node_id, step, from_unix_secs, to_unix_secs, limit)
.await?;
Ok(build_admin_proxy_node_metrics_payload_response(
step,
from_unix_secs,
to_unix_secs,
items,
))
}
pub(crate) async fn build_admin_proxy_fleet_metrics_response(
&self,
step: aether_data::repository::proxy_nodes::ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Response<Body>, GatewayError> {
if !self.has_proxy_node_reader() {
return Ok(build_admin_proxy_nodes_data_unavailable_response());
}
let items = self
.app
.list_proxy_fleet_metrics(step, from_unix_secs, to_unix_secs, limit)
.await?;
Ok(build_admin_proxy_fleet_metrics_payload_response(
step,
from_unix_secs,
to_unix_secs,
items,
))
}
pub(crate) async fn unregister_proxy_node(

4
apps/aether-gateway/src/handlers/admin/system/core/system_routes.rs
View File

@@ -533,6 +533,8 @@ async fn build_admin_system_cleanup_payload(
let cleaned = json!({
"audit_logs": summary.audit_logs_deleted,
"request_candidates": summary.request_candidates_deleted,
"proxy_node_metrics_1m": summary.proxy_node_metrics.deleted_1m_rows,
"proxy_node_metrics_1h": summary.proxy_node_metrics.deleted_1h_rows,
"pending_failed": summary.pending_failed,
"pending_recovered": summary.pending_recovered,
"usage_body_externalized": summary.usage.body_externalized,
@@ -545,6 +547,8 @@ async fn build_admin_system_cleanup_payload(
let total = summary
.audit_logs_deleted
.saturating_add(summary.request_candidates_deleted)
.saturating_add(summary.proxy_node_metrics.deleted_1m_rows)
.saturating_add(summary.proxy_node_metrics.deleted_1h_rows)
.saturating_add(summary.pending_failed)
.saturating_add(summary.pending_recovered)
.saturating_add(summary.usage.body_externalized)

139
apps/aether-gateway/src/handlers/admin/system/proxy_nodes.rs
View File

@@ -10,12 +10,14 @@ use crate::maintenance::{
};
use crate::GatewayError;
use aether_admin::system::{
admin_proxy_node_event_node_id_from_path, build_admin_proxy_node_payload,
build_admin_proxy_nodes_data_unavailable_response, build_admin_proxy_nodes_not_found_response,
admin_proxy_node_event_node_id_from_path, admin_proxy_node_metrics_node_id_from_path,
build_admin_proxy_node_payload, build_admin_proxy_nodes_data_unavailable_response,
build_admin_proxy_nodes_not_found_response,
};
use aether_contracts::tunnel::{
TUNNEL_RELAY_FORWARDED_BY_HEADER, TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
};
use aether_data::repository::proxy_nodes::{ProxyNodeEventQuery, ProxyNodeMetricsStep};
use axum::{
body::{Body, Bytes},
http,
@@ -148,6 +150,9 @@ const DEFAULT_PROXY_CONNECTIVITY_PROBE_URL: &str = "https://www.cloudflare.com/c
const PROXY_CONNECTIVITY_TIMEOUT_SECS: u64 = 10;
const TUNNEL_RELAY_ENVELOPE_CONTENT_TYPE: &str = "application/vnd.aether.tunnel-envelope";
const MAX_PROXY_CONNECTIVITY_RESPONSE_BYTES: usize = 64 * 1024;
const PROXY_NODE_METRICS_MAX_POINTS: usize = 50_000;
const PROXY_NODE_METRICS_1M_MAX_WINDOW_SECS: u64 = 30 * 24 * 60 * 60;
const PROXY_NODE_METRICS_1H_MAX_WINDOW_SECS: u64 = 365 * 24 * 60 * 60;
#[cfg(test)]
fn manual_proxy_connectivity_probe_url_override() -> &'static std::sync::RwLock<Option<String>> {
@@ -244,13 +249,53 @@ pub(crate) async fn maybe_build_local_admin_proxy_nodes_response(
return Ok(Some(build_admin_proxy_nodes_not_found_response()));
};
let limit = query_param_value(request_context.query_string(), "limit")
.and_then(|value| value.parse::<usize>().ok())
.filter(|value| *value > 0 && *value <= 200)
.unwrap_or(50);
let query = match parse_proxy_node_event_query(request_context.query_string()) {
Ok(query) => query,
Err(response) => return Ok(Some(response)),
};
return Ok(Some(
state
.build_admin_proxy_node_events_response(node_id, limit)
.build_admin_proxy_node_events_response(node_id, &query)
.await?,
));
}
if decision.route_kind.as_deref() == Some("list_node_metrics")
&& request_context.method() == http::Method::GET
{
let Some(node_id) = admin_proxy_node_metrics_node_id_from_path(request_context.path())
else {
return Ok(Some(build_admin_proxy_nodes_not_found_response()));
};
let (step, from_unix_secs, to_unix_secs, limit) =
match parse_proxy_node_metrics_query(request_context.query_string()) {
Ok(query) => query,
Err(response) => return Ok(Some(response)),
};
return Ok(Some(
state
.build_admin_proxy_node_metrics_response(
node_id,
step,
from_unix_secs,
to_unix_secs,
limit,
)
.await?,
));
}
if decision.route_kind.as_deref() == Some("list_fleet_metrics")
&& request_context.method() == http::Method::GET
{
let (step, from_unix_secs, to_unix_secs, limit) =
match parse_proxy_node_metrics_query(request_context.query_string()) {
Ok(query) => query,
Err(response) => return Ok(Some(response)),
};
return Ok(Some(
state
.build_admin_proxy_fleet_metrics_response(step, from_unix_secs, to_unix_secs, limit)
.await?,
));
}
@@ -1976,6 +2021,86 @@ fn validate_optional_object(value: Option<&Value>, field: &str) -> Result<(), Re
Ok(())
}
fn parse_proxy_node_event_query(
query: Option<&str>,
) -> Result<ProxyNodeEventQuery, Response<Body>> {
let limit = query_param_value(query, "limit")
.map(|value| parse_query_u64("limit", &value))
.transpose()?
.and_then(|value| usize::try_from(value).ok())
.filter(|value| *value > 0 && *value <= 200)
.unwrap_or(50);
let from_unix_secs = query_param_value(query, "from")
.map(|value| parse_query_u64("from", &value))
.transpose()?;
let to_unix_secs = query_param_value(query, "to")
.map(|value| parse_query_u64("to", &value))
.transpose()?;
if from_unix_secs
.zip(to_unix_secs)
.is_some_and(|(from, to)| from > to)
{
return Err(bad_request_response("from 不能大于 to"));
}
let event_type = query_param_value(query, "event_type")
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty());
Ok(ProxyNodeEventQuery {
limit,
from_unix_secs,
to_unix_secs,
event_type,
})
}
fn parse_proxy_node_metrics_query(
query: Option<&str>,
) -> Result<(ProxyNodeMetricsStep, u64, u64, usize), Response<Body>> {
let step = match query_param_value(query, "step")
.unwrap_or_else(|| "1m".to_string())
.trim()
{
"1m" => ProxyNodeMetricsStep::OneMinute,
"1h" => ProxyNodeMetricsStep::OneHour,
_ => return Err(bad_request_response("step 仅支持 1m 或 1h")),
};
let from_unix_secs = query_param_value(query, "from")
.ok_or_else(|| bad_request_response("from 为必填 Unix 秒时间戳"))?;
let from_unix_secs = parse_query_u64("from", &from_unix_secs)?;
let to_unix_secs = query_param_value(query, "to")
.ok_or_else(|| bad_request_response("to 为必填 Unix 秒时间戳"))?;
let to_unix_secs = parse_query_u64("to", &to_unix_secs)?;
if from_unix_secs > to_unix_secs {
return Err(bad_request_response("from 不能大于 to"));
}
let window_secs = to_unix_secs.saturating_sub(from_unix_secs);
let max_window_secs = match step {
ProxyNodeMetricsStep::OneMinute => PROXY_NODE_METRICS_1M_MAX_WINDOW_SECS,
ProxyNodeMetricsStep::OneHour => PROXY_NODE_METRICS_1H_MAX_WINDOW_SECS,
};
if window_secs > max_window_secs {
return Err(bad_request_response(match step {
ProxyNodeMetricsStep::OneMinute => "1m 最大查询窗口为 30 天",
ProxyNodeMetricsStep::OneHour => "1h 最大查询窗口为 365 天",
}));
}
let points = window_secs / step.bucket_size_secs() + 1;
let limit = usize::try_from(points)
.ok()
.filter(|value| *value > 0 && *value <= PROXY_NODE_METRICS_MAX_POINTS)
.ok_or_else(|| bad_request_response("查询点数过多"))?;
Ok((step, from_unix_secs, to_unix_secs, limit))
}
fn parse_query_u64(field: &str, value: &str) -> Result<u64, Response<Body>> {
value
.parse::<u64>()
.map_err(|_| bad_request_response(format!("{field} 必须是非负 Unix 秒时间戳")))
}
fn bad_request_response(detail: impl Into<String>) -> Response<Body> {
(
http::StatusCode::BAD_REQUEST,

19
apps/aether-gateway/src/maintenance/mod.rs
View File

@@ -11,13 +11,14 @@ pub(crate) use runtime::{
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
spawn_request_candidate_cleanup_worker, spawn_stats_aggregation_worker,
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, OAuthTokenRefreshRunSummary,
PoolQuotaProbeRunSummary, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_proxy_upgrade_rollout, AdminStatsRebuildSummary, AdminSystemCleanupSummary,
OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary, ProviderCheckinRunSummary,
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus,
ProxyUpgradeRolloutTrackedNodeState,
};

9
apps/aether-gateway/src/maintenance/runtime.rs
View File

@@ -20,6 +20,8 @@ mod pending_cleanup;
mod pool_quota_probe;
#[path = "runtime/provider_checkin.rs"]
mod provider_checkin;
#[path = "runtime/proxy_node_metrics_cleanup.rs"]
mod proxy_node_metrics_cleanup;
#[path = "runtime/proxy_node_staleness.rs"]
mod proxy_node_staleness;
#[path = "runtime/proxy_upgrade_rollout.rs"]
@@ -59,6 +61,7 @@ pub(crate) use pool_quota_probe::{
PoolQuotaProbeWorkerConfig,
};
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
use proxy_node_metrics_cleanup::*;
use proxy_node_staleness::*;
use proxy_upgrade_rollout::*;
pub(crate) use proxy_upgrade_rollout::{
@@ -90,6 +93,8 @@ const AUDIT_LOG_CLEANUP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
const GEMINI_FILE_MAPPING_CLEANUP_INTERVAL: Duration = Duration::from_secs(60 * 60);
const PENDING_CLEANUP_INTERVAL: Duration = Duration::from_secs(5 * 60);
const PROXY_NODE_STALE_SWEEP_INTERVAL: Duration = Duration::from_secs(5);
const PROXY_NODE_METRICS_CLEANUP_HOUR: u32 = 2;
const PROXY_NODE_METRICS_CLEANUP_MINUTE: u32 = 10;
const PROXY_UPGRADE_ROLLOUT_INTERVAL: Duration = Duration::from_secs(15);
const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 15;
const PROXY_NODE_STALE_MISSED_HEARTBEATS: u64 = 3;
@@ -127,6 +132,8 @@ struct UsageCleanupSettings {
pub(crate) struct AdminSystemCleanupSummary {
pub(crate) audit_logs_deleted: usize,
pub(crate) request_candidates_deleted: usize,
pub(crate) proxy_node_metrics:
aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary,
pub(crate) pending_failed: usize,
pub(crate) pending_recovered: usize,
pub(crate) usage: UsageCleanupSummary,
@@ -144,12 +151,14 @@ pub(crate) async fn run_admin_system_cleanup_once(
) -> Result<AdminSystemCleanupSummary, aether_data::DataLayerError> {
let audit_logs_deleted = cleanup_audit_logs_once(data).await?;
let request_candidates_deleted = cleanup_request_candidates_once(data).await?;
let proxy_node_metrics = cleanup_proxy_node_metrics_once(data).await?;
let pending = cleanup_stale_pending_requests_once(data).await?;
let usage = perform_usage_cleanup_once(data).await?;
Ok(AdminSystemCleanupSummary {
audit_logs_deleted,
request_candidates_deleted,
proxy_node_metrics,
pending_failed: pending.failed,
pending_recovered: pending.recovered,
usage,

109
apps/aether-gateway/src/maintenance/runtime/proxy_node_metrics_cleanup.rs Normal file
View File

@@ -0,0 +1,109 @@
use aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary;
use aether_data_contracts::DataLayerError;
use crate::data::GatewayDataState;
use super::{now_unix_secs, system_config_bool, system_config_u64, system_config_usize};
const SECS_PER_DAY: u64 = 24 * 60 * 60;
const PROXY_NODE_METRICS_1M_RETENTION_DAYS_DEFAULT: u64 = 30;
const PROXY_NODE_METRICS_1H_RETENTION_DAYS_DEFAULT: u64 = 180;
const PROXY_NODE_METRICS_RETENTION_DAYS_MIN: u64 = 1;
const PROXY_NODE_METRICS_1M_RETENTION_DAYS_MAX: u64 = 365;
const PROXY_NODE_METRICS_1H_RETENTION_DAYS_MAX: u64 = 1_095;
const PROXY_NODE_METRICS_CLEANUP_BATCH_SIZE_DEFAULT: usize = 5_000;
const PROXY_NODE_METRICS_CLEANUP_BATCH_SIZE_MAX: usize = 50_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct ProxyNodeMetricsCleanupSettings {
pub retain_1m_days: u64,
pub retain_1h_days: u64,
pub batch_size: usize,
}
pub(super) async fn proxy_node_metrics_cleanup_settings(
data: &GatewayDataState,
) -> Result<ProxyNodeMetricsCleanupSettings, DataLayerError> {
let retain_1m_days = system_config_u64(
data,
"proxy_node_metrics_1m_retention_days",
PROXY_NODE_METRICS_1M_RETENTION_DAYS_DEFAULT,
)
.await?
.clamp(
PROXY_NODE_METRICS_RETENTION_DAYS_MIN,
PROXY_NODE_METRICS_1M_RETENTION_DAYS_MAX,
);
let retain_1h_days = system_config_u64(
data,
"proxy_node_metrics_1h_retention_days",
PROXY_NODE_METRICS_1H_RETENTION_DAYS_DEFAULT,
)
.await?
.clamp(retain_1m_days, PROXY_NODE_METRICS_1H_RETENTION_DAYS_MAX);
let cleanup_batch_size =
system_config_usize(data, "proxy_node_metrics_cleanup_batch_size", 0).await?;
let fallback_batch_size = system_config_usize(
data,
"cleanup_batch_size",
PROXY_NODE_METRICS_CLEANUP_BATCH_SIZE_DEFAULT,
)
.await?;
let batch_size = (if cleanup_batch_size > 0 {
cleanup_batch_size
} else {
fallback_batch_size
})
.clamp(1, PROXY_NODE_METRICS_CLEANUP_BATCH_SIZE_MAX);
Ok(ProxyNodeMetricsCleanupSettings {
retain_1m_days,
retain_1h_days,
batch_size,
})
}
pub(super) async fn cleanup_proxy_node_metrics_once(
data: &GatewayDataState,
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
cleanup_proxy_node_metrics_at(data, now_unix_secs()).await
}
pub(super) async fn cleanup_proxy_node_metrics_at(
data: &GatewayDataState,
now_unix_secs: u64,
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
if !system_config_bool(data, "enable_auto_cleanup", true).await? {
return Ok(ProxyNodeMetricsCleanupSummary::default());
}
let settings = proxy_node_metrics_cleanup_settings(data).await?;
let retain_1m_from_unix_secs =
now_unix_secs.saturating_sub(settings.retain_1m_days.saturating_mul(SECS_PER_DAY));
let retain_1h_from_unix_secs =
now_unix_secs.saturating_sub(settings.retain_1h_days.saturating_mul(SECS_PER_DAY));
let mut summary = ProxyNodeMetricsCleanupSummary::default();
loop {
let deleted = data
.cleanup_proxy_node_metrics(
retain_1m_from_unix_secs,
retain_1h_from_unix_secs,
settings.batch_size,
)
.await?;
summary.deleted_1m_rows = summary
.deleted_1m_rows
.saturating_add(deleted.deleted_1m_rows);
summary.deleted_1h_rows = summary
.deleted_1h_rows
.saturating_add(deleted.deleted_1h_rows);
if deleted.deleted_1m_rows < settings.batch_size
&& deleted.deleted_1h_rows < settings.batch_size
{
break;
}
}
Ok(summary)
}

25
apps/aether-gateway/src/maintenance/runtime/runners.rs
View File

@@ -6,10 +6,10 @@ use crate::{AppState, GatewayError};
use super::{
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_once,
cleanup_expired_gemini_file_mappings_once, cleanup_request_candidates_once,
cleanup_stale_pending_requests_once, cleanup_stale_proxy_nodes_once,
collect_proxy_upgrade_rollout_probes, perform_db_maintenance_once,
perform_provider_checkin_once, perform_stats_aggregation_once,
cleanup_expired_gemini_file_mappings_once, cleanup_proxy_node_metrics_once,
cleanup_request_candidates_once, cleanup_stale_pending_requests_once,
cleanup_stale_proxy_nodes_once, collect_proxy_upgrade_rollout_probes,
perform_db_maintenance_once, perform_provider_checkin_once, perform_stats_aggregation_once,
perform_stats_hourly_aggregation_once, perform_usage_cleanup_once,
perform_wallet_daily_usage_aggregation_once, record_proxy_upgrade_traffic_success,
summarize_database_pool,
@@ -61,6 +61,23 @@ pub(super) async fn run_proxy_node_stale_cleanup_once(
Ok(())
}
pub(super) async fn run_proxy_node_metrics_cleanup_once(
data: &GatewayDataState,
) -> Result<(), DataLayerError> {
let summary = cleanup_proxy_node_metrics_once(data).await?;
if summary.deleted_1m_rows > 0 || summary.deleted_1h_rows > 0 {
info!(
event_name = "proxy_node_metrics_cleanup_completed",
log_type = "ops",
worker = "proxy_node_metrics_cleanup",
deleted_1m_rows = summary.deleted_1m_rows,
deleted_1h_rows = summary.deleted_1h_rows,
"gateway deleted expired proxy node metrics buckets"
);
}
Ok(())
}
pub(super) async fn run_proxy_upgrade_rollout_once(state: &AppState) -> Result<(), DataLayerError> {
let mut summary = advance_proxy_upgrade_rollout_once(&state.data).await?;
let probes = collect_proxy_upgrade_rollout_probes(&state.data).await?;

198
apps/aether-gateway/src/maintenance/runtime/tests.rs
View File

@@ -3,8 +3,8 @@ use std::sync::{Arc, Mutex};
use std::time::Duration;
use aether_data::repository::proxy_nodes::{
InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation, ProxyNodeReadRepository,
ProxyNodeWriteRepository, StoredProxyNode,
bucket_start_unix_secs, InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation,
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeWriteRepository, StoredProxyNode,
};
use aether_runtime::bounded_queue;
use axum::extract::ws::Message;
@@ -14,23 +14,25 @@ use serde_json::json;
use tokio::sync::watch;
use super::{
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_with, cleanup_stale_proxy_nodes_once,
inspect_proxy_upgrade_rollout, next_daily_run_after, next_db_maintenance_run_after,
next_stats_aggregation_run_after, next_stats_hourly_aggregation_run_after,
pending_cleanup_batch_size, pending_cleanup_timeout_minutes, plan_pending_cleanup_batch,
provider_checkin_schedule, record_proxy_upgrade_traffic_success, run_db_maintenance_with,
run_proxy_upgrade_rollout_once, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_with, cleanup_proxy_node_metrics_at,
cleanup_proxy_node_metrics_once, cleanup_stale_proxy_nodes_once, inspect_proxy_upgrade_rollout,
next_daily_run_after, next_db_maintenance_run_after, next_stats_aggregation_run_after,
next_stats_hourly_aggregation_run_after, pending_cleanup_batch_size,
pending_cleanup_timeout_minutes, plan_pending_cleanup_batch, provider_checkin_schedule,
proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success,
run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_audit_cleanup_worker,
spawn_db_maintenance_worker, spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker,
spawn_pool_monitor_worker, spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_proxy_upgrade_rollout, stats_aggregation_target_day,
stats_hourly_aggregation_target_hour, summarize_database_pool, usage_cleanup_settings,
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
UsageCleanupSettings, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
FailedPendingUsageRow, GatewayDataState, ProxyNodeMetricsCleanupSettings,
ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
#[tokio::test]
@@ -734,6 +736,176 @@ async fn usage_cleanup_settings_resolve_batch_and_delete_toggle() {
);
}
#[tokio::test]
async fn proxy_node_metrics_cleanup_settings_use_dedicated_retention_and_batch_limits() {
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
("cleanup_batch_size".to_string(), json!(250)),
("proxy_node_metrics_1m_retention_days".to_string(), json!(0)),
("proxy_node_metrics_1h_retention_days".to_string(), json!(7)),
(
"proxy_node_metrics_cleanup_batch_size".to_string(),
json!(100_000),
),
]);
let settings = proxy_node_metrics_cleanup_settings(&data)
.await
.expect("proxy metrics cleanup settings should resolve");
assert_eq!(
settings,
ProxyNodeMetricsCleanupSettings {
retain_1m_days: 1,
retain_1h_days: 7,
batch_size: 50_000,
}
);
}
#[tokio::test]
async fn proxy_node_metrics_cleanup_settings_fallback_to_global_batch_size() {
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
("cleanup_batch_size".to_string(), json!(250)),
(
"proxy_node_metrics_cleanup_batch_size".to_string(),
json!(0),
),
]);
let settings = proxy_node_metrics_cleanup_settings(&data)
.await
.expect("proxy metrics cleanup settings should resolve");
assert_eq!(
settings,
ProxyNodeMetricsCleanupSettings {
retain_1m_days: 30,
retain_1h_days: 180,
batch_size: 250,
}
);
}
#[tokio::test]
async fn proxy_node_metrics_cleanup_deletes_expired_buckets_in_batches() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
sample_connected_proxy_node("node-metrics-1", 30, 1),
sample_connected_proxy_node("node-metrics-2", 30, 1),
sample_connected_proxy_node("node-metrics-3", 30, 1),
]));
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
.with_system_config_values_for_tests([
("enable_auto_cleanup".to_string(), json!(true)),
("proxy_node_metrics_1m_retention_days".to_string(), json!(1)),
("proxy_node_metrics_1h_retention_days".to_string(), json!(1)),
(
"proxy_node_metrics_cleanup_batch_size".to_string(),
json!(1),
),
]);
for (idx, node_id) in ["node-metrics-1", "node-metrics-2", "node-metrics-3"]
.into_iter()
.enumerate()
{
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: node_id.to_string(),
heartbeat_interval: Some(30),
active_connections: Some(i32::try_from(idx + 1).unwrap()),
total_requests_delta: None,
avg_latency_ms: None,
failed_requests_delta: None,
dns_failures_delta: None,
stream_errors_delta: None,
proxy_metadata: Some(json!({
"tunnel_metrics": {
"connect_errors": idx + 1,
"disconnects": 0,
"error_events_total": 0,
"ws_in_bytes": idx + 1,
"ws_out_bytes": idx + 1,
"ws_in_frames": idx + 1,
"ws_out_frames": idx + 1,
"heartbeat_rtt_last_ms": 10
}
})),
proxy_version: Some("1.0.0".to_string()),
})
.await
.expect("heartbeat should write metrics");
}
let now = chrono::Utc::now().timestamp().max(0) as u64;
let old_bucket = bucket_start_unix_secs(now, ProxyNodeMetricsStep::OneMinute);
let cleanup = data
.cleanup_proxy_node_metrics(
old_bucket.saturating_add(60),
old_bucket.saturating_add(3_600),
1,
)
.await
.expect("direct cleanup should delete a limited batch");
assert_eq!(cleanup.deleted_1m_rows, 1);
assert_eq!(cleanup.deleted_1h_rows, 1);
let cleanup = cleanup_proxy_node_metrics_at(&data, now.saturating_add(2 * 86_400))
.await
.expect("runtime cleanup should loop over batches");
assert_eq!(cleanup.deleted_1m_rows, 2);
assert_eq!(cleanup.deleted_1h_rows, 2);
}
#[tokio::test]
async fn proxy_node_metrics_cleanup_respects_auto_cleanup_toggle() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
sample_connected_proxy_node("node-metrics-disabled", 30, 1),
]));
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
.with_system_config_values_for_tests([
("enable_auto_cleanup".to_string(), json!(false)),
("proxy_node_metrics_1m_retention_days".to_string(), json!(1)),
("proxy_node_metrics_1h_retention_days".to_string(), json!(1)),
(
"proxy_node_metrics_cleanup_batch_size".to_string(),
json!(1),
),
]);
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: "node-metrics-disabled".to_string(),
heartbeat_interval: Some(30),
active_connections: Some(1),
total_requests_delta: None,
avg_latency_ms: None,
failed_requests_delta: None,
dns_failures_delta: None,
stream_errors_delta: None,
proxy_metadata: Some(json!({
"tunnel_metrics": {
"connect_errors": 1,
"disconnects": 0,
"error_events_total": 0,
"ws_in_bytes": 1,
"ws_out_bytes": 1,
"ws_in_frames": 1,
"ws_out_frames": 1,
"heartbeat_rtt_last_ms": 10
}
})),
proxy_version: Some("1.0.0".to_string()),
})
.await
.expect("heartbeat should write metrics");
let cleanup = cleanup_proxy_node_metrics_once(&data)
.await
.expect("runtime cleanup should short-circuit");
assert_eq!(cleanup.deleted_1m_rows, 0);
assert_eq!(cleanup.deleted_1h_rows, 0);
}
#[test]
fn usage_cleanup_window_uses_non_overlapping_ranges() {
let now_utc = "2026-03-18T03:00:00Z"

30
apps/aether-gateway/src/maintenance/runtime/workers.rs
View File

@@ -12,12 +12,14 @@ use super::{
maintenance_timezone, parse_hhmm_time, perform_oauth_token_refresh_once,
provider_checkin_schedule, run_audit_cleanup_once, run_db_maintenance_once,
run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once, run_pool_monitor_once,
run_provider_checkin_once, run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
run_provider_checkin_once, run_proxy_node_metrics_cleanup_once,
run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
run_request_candidate_cleanup_once, run_stats_aggregation_once,
run_stats_hourly_aggregation_once, run_usage_cleanup_once,
run_wallet_daily_usage_aggregation_once, AUDIT_LOG_CLEANUP_INTERVAL,
GEMINI_FILE_MAPPING_CLEANUP_INTERVAL, OAUTH_TOKEN_REFRESH_INTERVAL, PENDING_CLEANUP_INTERVAL,
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_STALE_SWEEP_INTERVAL,
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_METRICS_CLEANUP_HOUR,
PROXY_NODE_METRICS_CLEANUP_MINUTE, PROXY_NODE_STALE_SWEEP_INTERVAL,
PROXY_UPGRADE_ROLLOUT_INTERVAL, REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
@@ -292,6 +294,30 @@ pub(crate) fn spawn_proxy_node_stale_cleanup_worker(
}))
}
pub(crate) fn spawn_proxy_node_metrics_cleanup_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if !data.has_proxy_node_writer() {
return None;
}
let timezone = maintenance_timezone();
Some(tokio::spawn(async move {
loop {
tokio::time::sleep(duration_until_next_daily_run(
Utc::now(),
timezone,
PROXY_NODE_METRICS_CLEANUP_HOUR,
PROXY_NODE_METRICS_CLEANUP_MINUTE,
))
.await;
if let Err(err) = run_proxy_node_metrics_cleanup_once(&data).await {
log_maintenance_worker_failure("proxy_node_metrics_cleanup", "tick", &err);
}
}
}))
}
pub(crate) fn spawn_proxy_upgrade_rollout_worker(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {

65
apps/aether-gateway/src/state/core.rs
View File

@@ -4,8 +4,10 @@ use std::sync::Mutex as StdMutex;
use std::time::Duration;
use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode, StoredProxyNodeEvent,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsStep, ProxyNodeTrafficMutation,
ProxyNodeTunnelStatusMutation, StoredProxyFleetMetricsBucket, StoredProxyNode,
StoredProxyNodeEvent, StoredProxyNodeMetricsBucket,
};
use aether_http::{build_http_client, HttpClientConfig};
use aether_runtime::{
@@ -46,6 +48,7 @@ use crate::maintenance::spawn_pending_cleanup_worker;
use crate::maintenance::spawn_pool_monitor_worker;
use crate::maintenance::spawn_pool_quota_probe_worker;
use crate::maintenance::spawn_provider_checkin_worker;
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
use crate::maintenance::spawn_proxy_upgrade_rollout_worker;
use crate::maintenance::spawn_request_candidate_cleanup_worker;
@@ -593,6 +596,44 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_node_events_filtered(
&self,
node_id: &str,
query: &ProxyNodeEventQuery,
) -> Result<Vec<StoredProxyNodeEvent>, GatewayError> {
self.data
.list_proxy_node_events_filtered(node_id, query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_node_metrics(
&self,
node_id: &str,
step: ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyNodeMetricsBucket>, GatewayError> {
self.data
.list_proxy_node_metrics(node_id, step, from_unix_secs, to_unix_secs, limit)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_fleet_metrics(
&self,
step: ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyFleetMetricsBucket>, GatewayError> {
self.data
.list_proxy_fleet_metrics(step, from_unix_secs, to_unix_secs, limit)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn register_proxy_node(
&self,
mutation: &aether_data::repository::proxy_nodes::ProxyNodeRegistrationMutation,
@@ -630,6 +671,23 @@ impl AppState {
.map_err(|err| std::io::Error::other(err.to_string()))
}
pub(crate) async fn cleanup_proxy_node_metrics(
&self,
retain_1m_from_unix_secs: u64,
retain_1h_from_unix_secs: u64,
delete_limit: usize,
) -> Result<aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary, GatewayError>
{
self.data
.cleanup_proxy_node_metrics(
retain_1m_from_unix_secs,
retain_1h_from_unix_secs,
delete_limit,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn apply_proxy_node_heartbeat(
&self,
mutation: &ProxyNodeHeartbeatMutation,
@@ -979,6 +1037,9 @@ impl AppState {
if let Some(handle) = spawn_proxy_node_stale_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_node_metrics_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_upgrade_rollout_worker(self.clone()) {
tasks.push(handle);
}

715
apps/aether-gateway/src/tests/ai_execute/sync/image.rs
View File

@@ -19,6 +19,721 @@ use aether_data_contracts::repository::provider_catalog::{
use base64::Engine as _;
use sha2::{Digest, Sha256};
#[tokio::test]
async fn gateway_converts_openai_image_sync_to_gemini_image_provider() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
url: String,
auth_header_value: String,
has_model_field: bool,
prompt: String,
response_modalities: Vec<String>,
image_size: String,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai", "google"])),
Some(serde_json::json!(["openai:image"])),
Some(serde_json::json!(["gpt-image-2"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
Some(serde_json::json!(["openai", "google"])),
Some(serde_json::json!(["openai:image"])),
Some(serde_json::json!(["gpt-image-2"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-gemini-image-bridge-1".to_string(),
provider_name: "google".to_string(),
provider_type: "google".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-gemini-image-bridge-1".to_string(),
endpoint_api_format: "gemini:generate_content".to_string(),
endpoint_api_family: Some("gemini".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-gemini-image-bridge-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["gemini:generate_content".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({
"gemini:generate_content": 1
})),
model_id: "model-gemini-image-bridge-1".to_string(),
global_model_id: "global-model-gemini-image-bridge-1".to_string(),
global_model_name: "gpt-image-2".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gemini-2.5-flash-image-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gemini-2.5-flash-image-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
endpoint_ids: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-gemini-image-bridge-1".to_string(),
"google".to_string(),
Some("https://generativelanguage.googleapis.com".to_string()),
"google".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-gemini-image-bridge-1".to_string(),
"provider-gemini-image-bridge-1".to_string(),
"gemini:generate_content".to_string(),
Some("gemini".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
None,
Some(serde_json::json!([
{"action":"drop","path":"model"}
])),
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-gemini-image-bridge-1".to_string(),
"provider-gemini-image-bridge-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["gemini:generate_content"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-gemini-image")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"gemini:generate_content": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
.expect("execution runtime payload should parse");
let body_json = payload
.get("body")
.and_then(|value| value.get("json_body"))
.cloned()
.unwrap_or_else(|| json!({}));
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
auth_header_value: payload
.get("headers")
.and_then(|value| value.get("x-goog-api-key"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
has_model_field: body_json.get("model").is_some(),
prompt: body_json
.get("contents")
.and_then(|value| value.get(0))
.and_then(|value| value.get("parts"))
.and_then(|value| value.get(0))
.and_then(|value| value.get("text"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
response_modalities: body_json
.get("generationConfig")
.and_then(|value| value.get("responseModalities"))
.and_then(|value| value.as_array())
.map(|items| {
items
.iter()
.filter_map(|item| item.as_str().map(ToOwned::to_owned))
.collect()
})
.unwrap_or_default(),
image_size: body_json
.get("generationConfig")
.and_then(|value| value.get("imageSize"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"request_id": "trace-openai-image-to-gemini-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"modelVersion": "gemini-2.5-flash-image-upstream",
"usageMetadata": {
"promptTokenCount": 11,
"candidatesTokenCount": 22,
"totalTokenCount": 33
},
"candidates": [{
"content": {
"role": "model",
"parts": [
{"text": "revised kite prompt"},
{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}
]
},
"finishReason": "STOP"
}]
}
},
"telemetry": {
"elapsed_ms": 37
}
}))
}
}),
);
let client_api_key = "sk-client-openai-image-to-gemini";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(client_api_key)),
sample_auth_snapshot(
"key-openai-image-client-bridge-1",
"user-openai-image-bridge-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::new(InMemoryRequestCandidateRepository::default()),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/images/generations"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
format!("Bearer {client_api_key}"),
)
.header(TRACE_ID_HEADER, "trace-openai-image-to-gemini-123")
.body(
"{\"model\":\"gpt-image-2\",\"prompt\":\"Draw a red kite\",\"size\":\"1024x1024\",\"response_format\":\"b64_json\"}",
)
.send()
.await
.expect("request should succeed");
let response_status = response.status();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK, "{response_body}");
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(response_json["data"][0]["b64_json"], "aGVsbG8=");
assert_eq!(
response_json["data"][0]["revised_prompt"],
"revised kite prompt"
);
assert_eq!(response_json["model"], "gemini-2.5-flash-image-upstream");
assert_eq!(response_json["usage"]["input_tokens"], 11);
assert_eq!(response_json["usage"]["output_tokens"], 22);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime sync should be captured");
assert_eq!(
seen_execution_runtime_request.trace_id,
"trace-openai-image-to-gemini-123"
);
assert_eq!(
seen_execution_runtime_request.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-upstream:generateContent"
);
assert_eq!(
seen_execution_runtime_request.auth_header_value,
"sk-upstream-gemini-image"
);
assert!(!seen_execution_runtime_request.has_model_field);
assert_eq!(seen_execution_runtime_request.prompt, "Draw a red kite");
assert_eq!(
seen_execution_runtime_request.response_modalities,
vec!["TEXT".to_string(), "IMAGE".to_string()]
);
assert_eq!(seen_execution_runtime_request.image_size, "1024x1024");
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_converts_gemini_image_sync_to_openai_image_provider() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
url: String,
authorization: String,
model: String,
action: String,
prompt: String,
image_url: String,
request_stream: bool,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["gemini", "openai"])),
Some(serde_json::json!(["gemini:generate_content"])),
Some(serde_json::json!(["gemini-image"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
Some(serde_json::json!(["gemini", "openai"])),
Some(serde_json::json!(["gemini:generate_content"])),
Some(serde_json::json!(["gemini-image"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-image-bridge-1".to_string(),
provider_name: "openai".to_string(),
provider_type: "openai".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-image-bridge-1".to_string(),
endpoint_api_format: "openai:image".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("image".to_string()),
endpoint_is_active: true,
key_id: "key-openai-image-bridge-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:image".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})),
model_id: "model-openai-image-bridge-1".to_string(),
global_model_id: "global-model-openai-image-bridge-1".to_string(),
global_model_name: "gemini-image".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-image-2-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-image-2-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:image".to_string()]),
endpoint_ids: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-image-bridge-1".to_string(),
"openai".to_string(),
Some("https://api.openai.com".to_string()),
"openai".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-image-bridge-1".to_string(),
"provider-openai-image-bridge-1".to_string(),
"openai:image".to_string(),
Some("openai".to_string()),
Some("image".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com".to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-image-bridge-1".to_string(),
"provider-openai-image-bridge-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:image"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-openai-image")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:image": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
.expect("execution runtime payload should parse");
let body_json = payload
.get("body")
.and_then(|value| value.get("json_body"))
.cloned()
.unwrap_or_else(|| json!({}));
let content = body_json
.get("input")
.and_then(|value| value.get(0))
.and_then(|value| value.get("content"))
.cloned()
.unwrap_or_else(|| json!([]));
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
authorization: payload
.get("headers")
.and_then(|value| value.get("authorization"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
model: body_json
.get("model")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
action: body_json
.get("tools")
.and_then(|value| value.get(0))
.and_then(|value| value.get("action"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
prompt: content
.as_array()
.into_iter()
.flatten()
.find(|item| {
item.get("type").and_then(|value| value.as_str()) == Some("input_text")
})
.and_then(|item| item.get("text"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
image_url: content
.as_array()
.into_iter()
.flatten()
.find(|item| {
item.get("type").and_then(|value| value.as_str()) == Some("input_image")
})
.and_then(|item| item.get("image_url"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
request_stream: body_json
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(true),
});
Json(json!({
"request_id": "trace-gemini-image-to-openai-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "resp_img_bridge_123",
"object": "response",
"model": "gpt-image-2-upstream",
"status": "completed",
"usage": {
"input_tokens": 3,
"output_tokens": 4,
"total_tokens": 7
},
"output": [{
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"revised_prompt": "converted gemini prompt",
"result": "aGVsbG8="
}]
}
},
"telemetry": {
"elapsed_ms": 43
}
}))
}
}),
);
let client_api_key = "client-gemini-image-to-openai";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(client_api_key)),
sample_auth_snapshot(
"key-gemini-image-client-bridge-1",
"user-gemini-image-bridge-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::new(InMemoryRequestCandidateRepository::default()),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1beta/models/gemini-image:generateContent?key={client_api_key}"
))
.header(http::header::CONTENT_TYPE, "application/json")
.header(TRACE_ID_HEADER, "trace-gemini-image-to-openai-123")
.body(
"{\"generationConfig\":{\"responseModalities\":[\"TEXT\",\"IMAGE\"]},\"contents\":[{\"role\":\"user\",\"parts\":[{\"text\":\"Change the background\"},{\"inlineData\":{\"mimeType\":\"image/png\",\"data\":\"aGVsbG8=\"}}]}]}",
)
.send()
.await
.expect("request should succeed");
let response_status = response.status();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK, "{response_body}");
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(response_json["modelVersion"], "gpt-image-2-upstream");
assert_eq!(
response_json["candidates"][0]["content"]["parts"][0]["text"],
"converted gemini prompt"
);
assert_eq!(
response_json["candidates"][0]["content"]["parts"][1]["inlineData"]["mimeType"],
"image/png"
);
assert_eq!(
response_json["candidates"][0]["content"]["parts"][1]["inlineData"]["data"],
"aGVsbG8="
);
assert_eq!(response_json["usageMetadata"]["promptTokenCount"], 3);
assert_eq!(response_json["usageMetadata"]["candidatesTokenCount"], 4);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime sync should be captured");
assert_eq!(
seen_execution_runtime_request.trace_id,
"trace-gemini-image-to-openai-123"
);
assert_eq!(
seen_execution_runtime_request.url,
"https://api.openai.com/v1/responses"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer sk-upstream-openai-image"
);
assert_eq!(seen_execution_runtime_request.model, "gpt-image-2-upstream");
assert_eq!(seen_execution_runtime_request.action, "edit");
assert_eq!(
seen_execution_runtime_request.prompt,
"Change the background"
);
assert_eq!(
seen_execution_runtime_request.image_url,
"data:image/png;base64,aGVsbG8="
);
assert!(!seen_execution_runtime_request.request_stream);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_refresh() {
#[derive(Debug, Clone)]

147
apps/aether-gateway/src/tests/control/admin/proxy_nodes.rs
View File

@@ -1,4 +1,5 @@
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data::repository::management_tokens::InMemoryManagementTokenRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
@@ -1442,6 +1443,7 @@ async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_prin
node_id: "node-1".to_string(),
event_type: "connected".to_string(),
detail: Some("older".to_string()),
event_metadata: None,
created_at_unix_ms: Some(1_710_000_000),
},
StoredProxyNodeEvent {
@@ -1449,6 +1451,7 @@ async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_prin
node_id: "node-1".to_string(),
event_type: "disconnected".to_string(),
detail: Some("newer".to_string()),
event_metadata: None,
created_at_unix_ms: Some(1_710_000_100),
},
],
@@ -1490,6 +1493,150 @@ async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_prin
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_reports_proxy_node_metrics_and_filters_events_locally() {
let proxy_node_repository =
Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
"node-1",
)]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
proxy_node_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system time should be after epoch")
.as_secs();
let heartbeat_response = client
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
.json(&json!({
"node_id": "node-1",
"heartbeat_id": 91,
"heartbeat_interval": 30,
"active_connections": 7,
"proxy_metadata": {
"tunnel_metrics": {
"connect_errors": 3,
"disconnects": 1,
"error_events_total": 1,
"ws_in_bytes": 1000,
"ws_out_bytes": 2000,
"ws_in_frames": 10,
"ws_out_frames": 20,
"heartbeat_rtt_last_ms": 42
},
"recent_tunnel_errors": [{
"timestamp_unix_secs": now_unix_secs,
"category": "tcp_connect_timeout",
"message": "tunnel TCP connect timeout"
}]
},
"proxy_version": "2.0.0"
}))
.send()
.await
.expect("heartbeat request should succeed");
assert_eq!(heartbeat_response.status(), StatusCode::OK);
let from = now_unix_secs.saturating_sub(120);
let to = now_unix_secs.saturating_add(120);
let metrics_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/node-1/metrics?from={from}&to={to}&step=1m"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("metrics request should succeed");
assert_eq!(metrics_response.status(), StatusCode::OK);
let metrics_payload: serde_json::Value = metrics_response
.json()
.await
.expect("metrics json should parse");
assert_eq!(metrics_payload["step"], "1m");
assert_eq!(metrics_payload["summary"]["samples"], 1);
assert_eq!(metrics_payload["summary"]["uptime_samples"], 1);
assert_eq!(metrics_payload["summary"]["active_connections_max"], 7);
assert_eq!(metrics_payload["summary"]["heartbeat_rtt_ms_avg"], 42.0);
assert_eq!(metrics_payload["summary"]["connect_errors_delta"], 3);
assert_eq!(metrics_payload["summary"]["ws_out_frames_delta"], 20);
let metric_items = metrics_payload["items"]
.as_array()
.expect("metrics items should be array");
assert_eq!(metric_items.len(), 1);
assert_eq!(metric_items[0]["node_id"], "node-1");
assert!(metric_items[0]["bucket_start"].is_string());
let fleet_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/metrics/fleet?from={from}&to={to}&step=1m"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("fleet metrics request should succeed");
assert_eq!(fleet_response.status(), StatusCode::OK);
let fleet_payload: serde_json::Value = fleet_response
.json()
.await
.expect("fleet json should parse");
assert_eq!(fleet_payload["summary"]["samples"], 1);
assert_eq!(fleet_payload["summary"]["error_events_delta"], 1);
let events_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/node-1/events?from={from}&to={to}&event_type=tunnel_err"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("events request should succeed");
assert_eq!(events_response.status(), StatusCode::OK);
let events_payload: serde_json::Value = events_response
.json()
.await
.expect("events json should parse");
let event_items = events_payload["items"]
.as_array()
.expect("event items should be array");
assert_eq!(event_items.len(), 1);
assert_eq!(event_items[0]["event_type"], "tunnel_err");
assert_eq!(
event_items[0]["event_metadata"]["category"],
"tcp_connect_timeout"
);
let invalid_metrics_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/node-1/metrics?from={from}&to={to}&step=5m"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("invalid metrics request should succeed");
assert_eq!(invalid_metrics_response.status(), StatusCode::BAD_REQUEST);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_updates_proxy_node_config_and_dispatches_upgrade_targets_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));