mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge branch 'aether-rust-pioneer' of https://github.com/fawney19/Aether into codex/codex-image-progress-heartbeat
# Conflicts: # frontend/src/features/usage/components/__tests__/HorizontalRequestTimeline.spec.ts
This commit is contained in:
@@ -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)?
|
||||
|
||||
@@ -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()?;
|
||||
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()..]
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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::{
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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?;
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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<()>> {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -110,21 +110,24 @@ sudo aether-proxy uninstall
|
||||
| `--upstream-pool-idle-timeout-secs` | `AETHER_PROXY_UPSTREAM_POOL_IDLE_TIMEOUT_SECS` | `300` | 连接池空闲超时(秒) |
|
||||
| `--upstream-tcp-keepalive-secs` | `AETHER_PROXY_UPSTREAM_TCP_KEEPALIVE_SECS` | `60` | TCP keepalive(秒,0 关闭) |
|
||||
| `--upstream-tcp-nodelay` | `AETHER_PROXY_UPSTREAM_TCP_NODELAY` | `true` | 启用 TCP_NODELAY |
|
||||
| `--upstream-proxy-url` | `AETHER_PROXY_UPSTREAM_PROXY_URL` | 空 | 仅 provider 上游请求使用的出口代理,支持 `http://`、`socks5://`、`socks5h://` |
|
||||
| `--upstream-proxy-url` | `AETHER_PROXY_UPSTREAM_PROXY_URL` | 空 | 仅 provider 上游请求使用的出口代理 |
|
||||
| `--redirect-replay-budget-bytes` | `AETHER_PROXY_REDIRECT_REPLAY_BUDGET_BYTES` | `5M` | 307/308 请求体重放的预读预算,支持 `K/M/G`,`0` 表示禁用 body replay buffering |
|
||||
|
||||
`upstream_proxy_url` 只影响 `aether-proxy` 访问 OpenAI、Claude、Gemini 等 provider 的上游请求,不影响节点回连 Aether 服务器的 WebSocket tunnel。配合 WARP sidecar 时可填写:
|
||||
出口代理支持 `http://`、`socks5://`、`socks5h://`。配合 WARP sidecar 时可填写:
|
||||
|
||||
```toml
|
||||
upstream_proxy_url = "socks5h://microwarp:1080"
|
||||
```
|
||||
|
||||
如果需要让 Aether 管理 API 和 WebSocket tunnel 也走代理,使用 `aether_proxy_url`。
|
||||
|
||||
#### Aether API 客户端
|
||||
|
||||
| 参数 | 环境变量 | 默认值 | 说明 |
|
||||
|------|----------|--------|------|
|
||||
| `--aether-request-timeout-secs` | `AETHER_PROXY_AETHER_REQUEST_TIMEOUT_SECS` | `10` | 请求总超时(秒) |
|
||||
| `--aether-connect-timeout-secs` | `AETHER_PROXY_AETHER_CONNECT_TIMEOUT_SECS` | `10` | 建连超时(秒) |
|
||||
| `--aether-proxy-url` | `AETHER_PROXY_AETHER_PROXY_URL` | 空 | Aether 注册、心跳和 WebSocket tunnel 回连使用的出口代理(默认不走代理) |
|
||||
| `--aether-retry-max-attempts` | `AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS` | `3` | 最大重试次数 |
|
||||
|
||||
#### DNS 与安全
|
||||
@@ -154,6 +157,15 @@ upstream_proxy_url = "socks5h://microwarp:1080"
|
||||
- 以 `systemd` 或 `OpenRC` 安装时默认会额外打开文件日志到 `/var/log/aether-proxy`
|
||||
- OpenRC 安装时,`aether-proxy logs` 实际读取 `/var/log/aether-proxy/current.log` 和 `/var/log/aether-proxy/error.log`;这些文件通常需要用 `sudo aether-proxy logs` 查看
|
||||
|
||||
### 隧道健康上报(Heartbeat)
|
||||
|
||||
proxy 会在心跳 `proxy_metadata` 中主动上报隧道稳定性指标,便于后端直接入库/告警:
|
||||
|
||||
- `proxy_metadata.tunnel_metrics`:建连尝试/成功/失败、断开次数、累计在线时长、心跳 RTT、WebSocket 收发帧与字节等。
|
||||
- `proxy_metadata.recent_tunnel_errors`:最近隧道异常事件(时间戳、类别、错误摘要,环形缓冲)。
|
||||
|
||||
说明:仅主连接(`conn=0`)发送 heartbeat,避免多条 tunnel 重复上报同一份全局指标。
|
||||
|
||||
### 多服务器配置
|
||||
|
||||
在 `aether-proxy.toml` 中使用 `[[servers]]` 配置 Aether 服务器。即使只有一个服务器,也必须写成一个 `[[servers]]` 条目;旧的顶层单服务器写法已不再支持。
|
||||
|
||||
@@ -17,7 +17,7 @@ use crate::config::{Config, ServerEntry, TunnelPoolSizing};
|
||||
use crate::net;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::{self, DynamicConfig};
|
||||
use crate::state::{AppState, ProxyMetrics, ServerContext};
|
||||
use crate::state::{AppState, ProxyMetrics, ServerContext, TunnelMetrics};
|
||||
use crate::upstream_client;
|
||||
use crate::{hardware, target_filter, tunnel};
|
||||
|
||||
@@ -80,6 +80,14 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
|
||||
server_count = servers.len(),
|
||||
"aether-proxy starting (tunnel mode)"
|
||||
);
|
||||
if let Some(proxy_url) = config.effective_aether_proxy_url() {
|
||||
if let Ok(proxy) = crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url) {
|
||||
info!(
|
||||
aether_proxy_url = %proxy.redacted_url(),
|
||||
"Aether control and tunnel egress proxy configured"
|
||||
);
|
||||
}
|
||||
}
|
||||
if let Some(proxy_url) = config
|
||||
.upstream_proxy_url
|
||||
.as_deref()
|
||||
@@ -457,6 +465,7 @@ fn build_server_context(
|
||||
dynamic: Arc::new(ArcSwap::from_pointee(dynamic)),
|
||||
active_connections: Arc::new(AtomicU64::new(0)),
|
||||
metrics: Arc::new(ProxyMetrics::new()),
|
||||
tunnel_metrics: Arc::new(TunnelMetrics::new()),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -960,6 +969,7 @@ mod tests {
|
||||
aether_tcp_keepalive_secs: 60,
|
||||
aether_tcp_nodelay: true,
|
||||
aether_http2: true,
|
||||
aether_proxy_url: None,
|
||||
aether_retry_max_attempts: 1,
|
||||
aether_retry_base_delay_ms: 50,
|
||||
aether_retry_max_delay_ms: 100,
|
||||
|
||||
@@ -339,6 +339,11 @@ pub struct Config {
|
||||
#[arg(long, env = "AETHER_PROXY_AETHER_HTTP2", default_value_t = true)]
|
||||
pub aether_http2: bool,
|
||||
|
||||
/// Optional egress proxy used for Aether API registration and WebSocket tunnel reconnects.
|
||||
/// Supported schemes: http, socks5, socks5h.
|
||||
#[arg(long, env = "AETHER_PROXY_AETHER_PROXY_URL")]
|
||||
pub aether_proxy_url: Option<String>,
|
||||
|
||||
/// Aether API retry attempts (including initial)
|
||||
#[arg(
|
||||
long,
|
||||
@@ -680,12 +685,11 @@ impl Config {
|
||||
if self.upstream_connect_timeout_secs == 0 {
|
||||
anyhow::bail!("upstream_connect_timeout_secs must be > 0");
|
||||
}
|
||||
if let Some(proxy_url) = self
|
||||
.upstream_proxy_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
if let Some(proxy_url) = normalized_proxy_url(&self.aether_proxy_url) {
|
||||
crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url)
|
||||
.map_err(|err| anyhow::anyhow!("aether_proxy_url invalid: {err}"))?;
|
||||
}
|
||||
if let Some(proxy_url) = normalized_proxy_url(&self.upstream_proxy_url) {
|
||||
crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url)
|
||||
.map_err(|err| anyhow::anyhow!("upstream_proxy_url invalid: {err}"))?;
|
||||
}
|
||||
@@ -740,6 +744,10 @@ impl Config {
|
||||
Ok(Duration::from_millis(self.tunnel_stale_timeout_ms))
|
||||
}
|
||||
|
||||
pub fn effective_aether_proxy_url(&self) -> Option<&str> {
|
||||
normalized_proxy_url(&self.aether_proxy_url)
|
||||
}
|
||||
|
||||
pub fn resolve_tunnel_pool_sizing(
|
||||
&self,
|
||||
hw_info: &HardwareInfo,
|
||||
@@ -850,6 +858,8 @@ pub struct ConfigFile {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_http2: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_proxy_url: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_retry_max_attempts: Option<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_retry_base_delay_ms: Option<u64>,
|
||||
@@ -929,8 +939,7 @@ impl ConfigFile {
|
||||
/// Load from a TOML file.
|
||||
pub fn load(path: &Path) -> anyhow::Result<Self> {
|
||||
let content = std::fs::read_to_string(path)?;
|
||||
reject_removed_config_keys(&content)?;
|
||||
Ok(toml::from_str(&content)?)
|
||||
parse_config_file_content(&content)
|
||||
}
|
||||
|
||||
/// Save to a TOML file.
|
||||
@@ -1006,6 +1015,7 @@ impl ConfigFile {
|
||||
);
|
||||
set!("AETHER_PROXY_AETHER_TCP_NODELAY", self.aether_tcp_nodelay);
|
||||
set!("AETHER_PROXY_AETHER_HTTP2", self.aether_http2);
|
||||
set!("AETHER_PROXY_AETHER_PROXY_URL", self.aether_proxy_url);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS",
|
||||
self.aether_retry_max_attempts
|
||||
@@ -1124,6 +1134,59 @@ impl ConfigFile {
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_config_file_content(content: &str) -> anyhow::Result<ConfigFile> {
|
||||
reject_removed_config_keys(content)?;
|
||||
let mut value: toml::Value = toml::from_str(content)?;
|
||||
promote_server_scoped_upstream_proxy_url(&mut value)?;
|
||||
Ok(value.try_into()?)
|
||||
}
|
||||
|
||||
fn normalized_proxy_url(value: &Option<String>) -> Option<&str> {
|
||||
value
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn promote_server_scoped_upstream_proxy_url(value: &mut toml::Value) -> anyhow::Result<()> {
|
||||
const KEY: &str = "upstream_proxy_url";
|
||||
|
||||
let Some(root) = value.as_table_mut() else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let mut promoted = root.get(KEY).cloned();
|
||||
let Some(servers) = root.get_mut("servers").and_then(toml::Value::as_array_mut) else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
for (index, server) in servers.iter_mut().enumerate() {
|
||||
let Some(table) = server.as_table_mut() else {
|
||||
continue;
|
||||
};
|
||||
let Some(server_value) = table.remove(KEY) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match promoted.as_ref() {
|
||||
Some(existing) if existing != &server_value => {
|
||||
anyhow::bail!(
|
||||
"conflicting upstream_proxy_url values: top-level value and [[servers]] entry {} differ; configure it once at the top level",
|
||||
index + 1
|
||||
);
|
||||
}
|
||||
Some(_) => {}
|
||||
None => promoted = Some(server_value),
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(promoted) = promoted {
|
||||
root.insert(KEY.to_string(), promoted);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reject_removed_config_keys(content: &str) -> anyhow::Result<()> {
|
||||
let value: toml::Value = toml::from_str(content)?;
|
||||
let Some(table) = value.as_table() else {
|
||||
@@ -1233,6 +1296,92 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_file_deserializes_aether_proxy_url() {
|
||||
let cfg: ConfigFile = toml::from_str("aether_proxy_url = \"socks5h://127.0.0.1:1080\"")
|
||||
.expect("proxy URL toml");
|
||||
assert_eq!(
|
||||
cfg.aether_proxy_url.as_deref(),
|
||||
Some("socks5h://127.0.0.1:1080")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aether_proxy_url_requires_explicit_opt_in() {
|
||||
let default_direct = Config::parse_from([
|
||||
"aether-proxy",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"proxy-test",
|
||||
"--upstream-proxy-url",
|
||||
"socks5h://127.0.0.1:1080",
|
||||
]);
|
||||
assert_eq!(default_direct.effective_aether_proxy_url(), None);
|
||||
|
||||
let explicit = Config::parse_from([
|
||||
"aether-proxy",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"proxy-test",
|
||||
"--upstream-proxy-url",
|
||||
"socks5h://127.0.0.1:1080",
|
||||
"--aether-proxy-url",
|
||||
"http://127.0.0.1:8080",
|
||||
]);
|
||||
assert_eq!(
|
||||
explicit.effective_aether_proxy_url(),
|
||||
Some("http://127.0.0.1:8080")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_file_load_accepts_server_scoped_upstream_proxy_url() {
|
||||
let cfg = parse_config_file_content(
|
||||
r#"
|
||||
[[servers]]
|
||||
aether_url = "https://aether.example.com"
|
||||
upstream_proxy_url = "socks5://127.0.0.1:1080"
|
||||
management_token = "ae_test"
|
||||
node_name = "proxy-test"
|
||||
"#,
|
||||
)
|
||||
.expect("server-scoped proxy URL should be promoted");
|
||||
|
||||
assert_eq!(
|
||||
cfg.upstream_proxy_url.as_deref(),
|
||||
Some("socks5://127.0.0.1:1080")
|
||||
);
|
||||
assert_eq!(cfg.servers.len(), 1);
|
||||
assert_eq!(cfg.servers[0].aether_url, "https://aether.example.com");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_file_load_rejects_conflicting_server_scoped_upstream_proxy_url() {
|
||||
let error = parse_config_file_content(
|
||||
r#"
|
||||
upstream_proxy_url = "socks5://127.0.0.1:1080"
|
||||
|
||||
[[servers]]
|
||||
aether_url = "https://aether.example.com"
|
||||
upstream_proxy_url = "socks5://127.0.0.1:1081"
|
||||
management_token = "ae_test"
|
||||
node_name = "proxy-test"
|
||||
"#,
|
||||
)
|
||||
.expect_err("conflicting proxy URLs should be rejected");
|
||||
|
||||
assert!(
|
||||
error.to_string().contains("conflicting upstream_proxy_url"),
|
||||
"error should mention the conflicting key"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_file_rejects_removed_tunnel_seconds_keys() {
|
||||
let error = reject_removed_config_keys("tunnel_ping_interval_secs = 5")
|
||||
|
||||
@@ -1,4 +1,11 @@
|
||||
use std::io;
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::time::Duration;
|
||||
|
||||
use base64::Engine;
|
||||
use socket2::{SockRef, TcpKeepalive};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
use url::Url;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -114,6 +121,303 @@ impl UpstreamProxyConfig {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct ProxyConnectOptions {
|
||||
pub connect_timeout: Duration,
|
||||
pub tcp_nodelay: bool,
|
||||
pub tcp_keepalive: Option<Duration>,
|
||||
}
|
||||
|
||||
pub(crate) async fn connect_target_via_proxy(
|
||||
proxy: &UpstreamProxyConfig,
|
||||
target_host: &str,
|
||||
target_port: u16,
|
||||
options: ProxyConnectOptions,
|
||||
) -> io::Result<TcpStream> {
|
||||
let mut tcp = connect_proxy_tcp(
|
||||
proxy,
|
||||
options.connect_timeout,
|
||||
options.tcp_nodelay,
|
||||
options.tcp_keepalive,
|
||||
)
|
||||
.await?;
|
||||
|
||||
match proxy.scheme() {
|
||||
UpstreamProxyScheme::Http => {
|
||||
http_connect(&mut tcp, &target_authority(target_host, target_port), proxy).await?;
|
||||
}
|
||||
UpstreamProxyScheme::Socks5 | UpstreamProxyScheme::Socks5h => {
|
||||
socks5_connect(&mut tcp, proxy, target_host, target_port).await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(tcp)
|
||||
}
|
||||
|
||||
pub(crate) async fn connect_proxy_tcp(
|
||||
proxy: &UpstreamProxyConfig,
|
||||
connect_timeout: Duration,
|
||||
tcp_nodelay: bool,
|
||||
tcp_keepalive: Option<Duration>,
|
||||
) -> io::Result<TcpStream> {
|
||||
let resolved = tokio::time::timeout(
|
||||
connect_timeout,
|
||||
tokio::net::lookup_host((proxy.host(), proxy.port())),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "proxy DNS timeout"))?
|
||||
.map_err(|err| io::Error::other(format!("proxy DNS failed: {err}")))?;
|
||||
|
||||
let mut last_error = None;
|
||||
for addr in resolved {
|
||||
match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
|
||||
Ok(Ok(stream)) => {
|
||||
configure_tcp_stream(&stream, tcp_nodelay, tcp_keepalive)?;
|
||||
return Ok(stream);
|
||||
}
|
||||
Ok(Err(error)) => last_error = Some(error),
|
||||
Err(_) => {
|
||||
last_error = Some(io::Error::new(
|
||||
io::ErrorKind::TimedOut,
|
||||
format!("proxy connect timeout: {addr}"),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or_else(|| io::Error::other("proxy DNS returned no addresses")))
|
||||
}
|
||||
|
||||
fn configure_tcp_stream(
|
||||
stream: &TcpStream,
|
||||
tcp_nodelay: bool,
|
||||
tcp_keepalive: Option<Duration>,
|
||||
) -> io::Result<()> {
|
||||
stream.set_nodelay(tcp_nodelay)?;
|
||||
if let Some(keepalive) = tcp_keepalive {
|
||||
let keepalive = TcpKeepalive::new().with_time(keepalive);
|
||||
SockRef::from(stream).set_tcp_keepalive(&keepalive)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn http_connect(
|
||||
stream: &mut TcpStream,
|
||||
target_authority: &str,
|
||||
proxy: &UpstreamProxyConfig,
|
||||
) -> io::Result<()> {
|
||||
let mut request = format!(
|
||||
"CONNECT {target_authority} HTTP/1.1\r\nHost: {target_authority}\r\nProxy-Connection: Keep-Alive\r\n"
|
||||
);
|
||||
if let Some(auth) = proxy.basic_auth_header() {
|
||||
request.push_str("Proxy-Authorization: ");
|
||||
request.push_str(&auth);
|
||||
request.push_str("\r\n");
|
||||
}
|
||||
request.push_str("\r\n");
|
||||
stream.write_all(request.as_bytes()).await?;
|
||||
stream.flush().await?;
|
||||
|
||||
let mut response = Vec::with_capacity(1024);
|
||||
let mut chunk = [0u8; 1024];
|
||||
loop {
|
||||
if response.len() >= 16 * 1024 {
|
||||
return Err(io::Error::other("proxy CONNECT response too large"));
|
||||
}
|
||||
let n = stream.read(&mut chunk).await?;
|
||||
if n == 0 {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::UnexpectedEof,
|
||||
"proxy closed during CONNECT",
|
||||
));
|
||||
}
|
||||
response.extend_from_slice(&chunk[..n]);
|
||||
if response.windows(4).any(|window| window == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let status_line_end = response
|
||||
.windows(2)
|
||||
.position(|window| window == b"\r\n")
|
||||
.ok_or_else(|| io::Error::other("proxy CONNECT response missing status line"))?;
|
||||
let status_line = std::str::from_utf8(&response[..status_line_end])
|
||||
.map_err(|_| io::Error::other("proxy CONNECT status line is not UTF-8"))?;
|
||||
let status = status_line.split_whitespace().nth(1).unwrap_or_default();
|
||||
if status == "200" {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::other(format!(
|
||||
"proxy CONNECT failed: {status_line}"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn socks5_connect(
|
||||
stream: &mut TcpStream,
|
||||
proxy: &UpstreamProxyConfig,
|
||||
target_host: &str,
|
||||
target_port: u16,
|
||||
) -> io::Result<()> {
|
||||
let requires_auth = proxy.username().is_some();
|
||||
if requires_auth {
|
||||
stream.write_all(&[0x05, 0x02, 0x00, 0x02]).await?;
|
||||
} else {
|
||||
stream.write_all(&[0x05, 0x01, 0x00]).await?;
|
||||
}
|
||||
|
||||
let mut method_response = [0u8; 2];
|
||||
stream.read_exact(&mut method_response).await?;
|
||||
if method_response[0] != 0x05 {
|
||||
return Err(io::Error::other("invalid SOCKS5 method response"));
|
||||
}
|
||||
match method_response[1] {
|
||||
0x00 => {}
|
||||
0x02 => socks5_authenticate(stream, proxy).await?,
|
||||
0xff => return Err(io::Error::other("SOCKS5 proxy rejected all auth methods")),
|
||||
method => {
|
||||
return Err(io::Error::other(format!(
|
||||
"SOCKS5 proxy selected unsupported auth method 0x{method:02x}"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
let address = socks5_target_address(target_host, target_port, proxy.uses_remote_dns()).await?;
|
||||
stream.write_all(&address).await?;
|
||||
|
||||
let mut response = [0u8; 4];
|
||||
stream.read_exact(&mut response).await?;
|
||||
if response[0] != 0x05 {
|
||||
return Err(io::Error::other("invalid SOCKS5 connect response"));
|
||||
}
|
||||
if response[1] != 0x00 {
|
||||
return Err(io::Error::other(format!(
|
||||
"SOCKS5 connect failed: {}",
|
||||
socks5_reply_message(response[1])
|
||||
)));
|
||||
}
|
||||
|
||||
match response[3] {
|
||||
0x01 => {
|
||||
let mut ignored = [0u8; 4 + 2];
|
||||
stream.read_exact(&mut ignored).await?;
|
||||
}
|
||||
0x03 => {
|
||||
let mut len = [0u8; 1];
|
||||
stream.read_exact(&mut len).await?;
|
||||
let mut ignored = vec![0u8; len[0] as usize + 2];
|
||||
stream.read_exact(&mut ignored).await?;
|
||||
}
|
||||
0x04 => {
|
||||
let mut ignored = [0u8; 16 + 2];
|
||||
stream.read_exact(&mut ignored).await?;
|
||||
}
|
||||
atyp => {
|
||||
return Err(io::Error::other(format!(
|
||||
"SOCKS5 proxy returned unsupported address type 0x{atyp:02x}"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn socks5_authenticate(
|
||||
stream: &mut TcpStream,
|
||||
proxy: &UpstreamProxyConfig,
|
||||
) -> io::Result<()> {
|
||||
let username = proxy.username().unwrap_or_default().as_bytes();
|
||||
let password = proxy.password().unwrap_or_default().as_bytes();
|
||||
if username.len() > u8::MAX as usize || password.len() > u8::MAX as usize {
|
||||
return Err(io::Error::other(
|
||||
"SOCKS5 username/password must be at most 255 bytes",
|
||||
));
|
||||
}
|
||||
|
||||
let mut request = Vec::with_capacity(username.len() + password.len() + 3);
|
||||
request.push(0x01);
|
||||
request.push(username.len() as u8);
|
||||
request.extend_from_slice(username);
|
||||
request.push(password.len() as u8);
|
||||
request.extend_from_slice(password);
|
||||
stream.write_all(&request).await?;
|
||||
|
||||
let mut response = [0u8; 2];
|
||||
stream.read_exact(&mut response).await?;
|
||||
if response[0] != 0x01 || response[1] != 0x00 {
|
||||
return Err(io::Error::other("SOCKS5 username/password auth failed"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn socks5_target_address(
|
||||
target_host: &str,
|
||||
target_port: u16,
|
||||
remote_dns: bool,
|
||||
) -> io::Result<Vec<u8>> {
|
||||
let mut request = vec![0x05, 0x01, 0x00];
|
||||
if let Ok(ip) = target_host.parse::<IpAddr>() {
|
||||
push_socks5_ip_address(&mut request, ip);
|
||||
} else if remote_dns {
|
||||
let host = target_host.as_bytes();
|
||||
if host.len() > u8::MAX as usize {
|
||||
return Err(io::Error::other("SOCKS5 target hostname is too long"));
|
||||
}
|
||||
request.push(0x03);
|
||||
request.push(host.len() as u8);
|
||||
request.extend_from_slice(host);
|
||||
} else {
|
||||
let mut resolved = tokio::net::lookup_host((target_host, target_port))
|
||||
.await
|
||||
.map_err(|err| io::Error::other(format!("SOCKS5 target DNS failed: {err}")))?;
|
||||
let addr = resolved
|
||||
.next()
|
||||
.ok_or_else(|| io::Error::other("SOCKS5 target DNS returned no addresses"))?;
|
||||
push_socks5_socket_address(&mut request, addr);
|
||||
}
|
||||
request.extend_from_slice(&target_port.to_be_bytes());
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
fn push_socks5_socket_address(request: &mut Vec<u8>, addr: SocketAddr) {
|
||||
push_socks5_ip_address(request, addr.ip());
|
||||
}
|
||||
|
||||
fn push_socks5_ip_address(request: &mut Vec<u8>, ip: IpAddr) {
|
||||
match ip {
|
||||
IpAddr::V4(ip) => {
|
||||
request.push(0x01);
|
||||
request.extend_from_slice(&ip.octets());
|
||||
}
|
||||
IpAddr::V6(ip) => {
|
||||
request.push(0x04);
|
||||
request.extend_from_slice(&ip.octets());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn socks5_reply_message(reply: u8) -> &'static str {
|
||||
match reply {
|
||||
0x01 => "general failure",
|
||||
0x02 => "connection not allowed",
|
||||
0x03 => "network unreachable",
|
||||
0x04 => "host unreachable",
|
||||
0x05 => "connection refused",
|
||||
0x06 => "TTL expired",
|
||||
0x07 => "command not supported",
|
||||
0x08 => "address type not supported",
|
||||
_ => "unknown error",
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn target_authority(host: &str, port: u16) -> String {
|
||||
if host.contains(':') && !host.starts_with('[') {
|
||||
format!("[{host}]:{port}")
|
||||
} else {
|
||||
format!("{host}:{port}")
|
||||
}
|
||||
}
|
||||
|
||||
fn non_empty_url_part(value: &str) -> Option<String> {
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
|
||||
@@ -67,6 +67,7 @@ impl AetherClient {
|
||||
tcp_nodelay: config.aether_tcp_nodelay,
|
||||
http2_adaptive_window: config.aether_http2,
|
||||
user_agent: Some(format!("aether-proxy/{}", env!("CARGO_PKG_VERSION"))),
|
||||
proxy_url: config.effective_aether_proxy_url().map(str::to_string),
|
||||
..HttpClientConfig::default()
|
||||
})
|
||||
.expect("failed to create HTTP client");
|
||||
|
||||
@@ -93,15 +93,6 @@ impl ServerTab {
|
||||
required: true,
|
||||
help: "Node name for identification in Aether dashboard",
|
||||
},
|
||||
Field {
|
||||
label: "Upstream Proxy",
|
||||
key: "upstream_proxy_url",
|
||||
value: String::new(),
|
||||
kind: FieldKind::Text,
|
||||
required: false,
|
||||
help:
|
||||
"Optional provider egress proxy, e.g. http://127.0.0.1:8080 or socks5h://127.0.0.1:1080",
|
||||
},
|
||||
],
|
||||
}
|
||||
}
|
||||
@@ -115,12 +106,6 @@ impl ServerTab {
|
||||
}
|
||||
tab
|
||||
}
|
||||
|
||||
fn set_field_value(&mut self, key: &str, value: &str) {
|
||||
if let Some(field) = self.fields.iter_mut().find(|field| field.key == key) {
|
||||
field.value = value.to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -- App state ----------------------------------------------------------------
|
||||
@@ -153,6 +138,15 @@ impl App {
|
||||
server_tabs: vec![ServerTab::new()],
|
||||
active_tab: 0,
|
||||
global_fields: vec![
|
||||
Field {
|
||||
label: "Egress Proxy",
|
||||
key: "upstream_proxy_url",
|
||||
value: String::new(),
|
||||
kind: FieldKind::Text,
|
||||
required: false,
|
||||
help:
|
||||
"Optional egress proxy for Aether tunnel/API and provider requests, e.g. http://127.0.0.1:8080 or socks5h://127.0.0.1:1080",
|
||||
},
|
||||
Field {
|
||||
label: "Install Service",
|
||||
key: "install_service",
|
||||
@@ -282,6 +276,7 @@ impl App {
|
||||
"allow_private_targets" => cfg.allow_private_targets.map(|v| v.to_string()),
|
||||
"heartbeat_interval" => cfg.heartbeat_interval.map(|v| v.to_string()),
|
||||
"redirect_replay_budget_bytes" => cfg.redirect_replay_budget_bytes.clone(),
|
||||
"upstream_proxy_url" => cfg.upstream_proxy_url.clone(),
|
||||
_ => None,
|
||||
};
|
||||
if let Some(v) = val {
|
||||
@@ -296,11 +291,6 @@ impl App {
|
||||
} else {
|
||||
self.server_tabs = servers.iter().map(ServerTab::from_entry).collect();
|
||||
}
|
||||
if let Some(proxy_url) = cfg.upstream_proxy_url.as_deref() {
|
||||
for tab in &mut self.server_tabs {
|
||||
tab.set_field_value("upstream_proxy_url", proxy_url);
|
||||
}
|
||||
}
|
||||
self.active_tab = 0;
|
||||
self.selected = 0;
|
||||
self.scroll_offset = 0;
|
||||
@@ -322,12 +312,6 @@ impl App {
|
||||
.filter(|v| !v.is_empty())
|
||||
}
|
||||
|
||||
fn get_shared_server_field(&self, key: &str) -> Option<String> {
|
||||
self.server_tabs
|
||||
.iter()
|
||||
.find_map(|tab| Self::get_tab(tab, key))
|
||||
}
|
||||
|
||||
fn toggle_enabled(&self, key: &str) -> bool {
|
||||
self.get_global(key).as_deref() == Some("true")
|
||||
}
|
||||
@@ -373,7 +357,7 @@ impl App {
|
||||
}
|
||||
|
||||
fn parse_optional_upstream_proxy_url(&self) -> anyhow::Result<Option<String>> {
|
||||
let Some(raw) = self.get_shared_server_field("upstream_proxy_url") else {
|
||||
let Some(raw) = self.get_global("upstream_proxy_url") else {
|
||||
return Ok(None);
|
||||
};
|
||||
let trimmed = raw.trim();
|
||||
@@ -381,7 +365,7 @@ impl App {
|
||||
return Ok(None);
|
||||
}
|
||||
UpstreamProxyConfig::parse(trimmed)
|
||||
.map_err(|err| anyhow::anyhow!("upstream proxy URL invalid: {err}"))?;
|
||||
.map_err(|err| anyhow::anyhow!("egress proxy URL invalid: {err}"))?;
|
||||
Ok(Some(trimmed.to_string()))
|
||||
}
|
||||
|
||||
@@ -613,11 +597,7 @@ impl App {
|
||||
}
|
||||
// -- Add / remove server --
|
||||
KeyCode::Char('+') | KeyCode::Char('a') => {
|
||||
let upstream_proxy_url = self.get_shared_server_field("upstream_proxy_url");
|
||||
let mut tab = ServerTab::new();
|
||||
if let Some(proxy_url) = upstream_proxy_url.as_deref() {
|
||||
tab.set_field_value("upstream_proxy_url", proxy_url);
|
||||
}
|
||||
let tab = ServerTab::new();
|
||||
self.server_tabs.push(tab);
|
||||
self.active_tab = self.server_tabs.len() - 1;
|
||||
self.selected = 0;
|
||||
@@ -694,14 +674,7 @@ impl App {
|
||||
|
||||
fn commit_edit_buffer(&mut self) -> bool {
|
||||
if self.validate_edit() {
|
||||
let key = self.selected_field().key;
|
||||
if key == "upstream_proxy_url" {
|
||||
for tab in &mut self.server_tabs {
|
||||
tab.set_field_value(key, &self.edit_buffer);
|
||||
}
|
||||
} else {
|
||||
self.selected_field_mut().value = self.edit_buffer.clone();
|
||||
}
|
||||
self.selected_field_mut().value = self.edit_buffer.clone();
|
||||
self.modified = true;
|
||||
self.mode = Mode::Normal;
|
||||
true
|
||||
@@ -1124,23 +1097,20 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_app_places_upstream_proxy_under_node_name() {
|
||||
fn new_app_places_egress_proxy_in_global_fields() {
|
||||
let app = sample_app();
|
||||
let keys: Vec<&str> = app.server_tabs[0]
|
||||
let server_keys: Vec<&str> = app.server_tabs[0]
|
||||
.fields
|
||||
.iter()
|
||||
.map(|field| field.key)
|
||||
.collect();
|
||||
let global_keys: Vec<&str> = app.global_fields.iter().map(|field| field.key).collect();
|
||||
|
||||
assert_eq!(
|
||||
keys,
|
||||
vec![
|
||||
"aether_url",
|
||||
"management_token",
|
||||
"node_name",
|
||||
"upstream_proxy_url"
|
||||
]
|
||||
server_keys,
|
||||
vec!["aether_url", "management_token", "node_name"]
|
||||
);
|
||||
assert_eq!(global_keys.first().copied(), Some("upstream_proxy_url"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1172,7 +1142,7 @@ mod tests {
|
||||
set_global_field(&mut app, "allow_private_targets", "true");
|
||||
set_global_field(&mut app, "heartbeat_interval", "45");
|
||||
set_global_field(&mut app, "redirect_replay_budget_bytes", "6m");
|
||||
set_server_field(&mut app, "upstream_proxy_url", "socks5h://127.0.0.1:1080");
|
||||
set_global_field(&mut app, "upstream_proxy_url", "socks5h://127.0.0.1:1080");
|
||||
|
||||
let cfg = app.to_config().expect("config should serialize");
|
||||
assert_eq!(cfg.allow_private_targets, Some(true));
|
||||
@@ -1194,14 +1164,14 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_config_rejects_invalid_upstream_proxy_url() {
|
||||
fn to_config_rejects_invalid_egress_proxy_url() {
|
||||
let mut app = sample_app();
|
||||
set_server_field(&mut app, "upstream_proxy_url", "ftp://proxy.example");
|
||||
set_global_field(&mut app, "upstream_proxy_url", "ftp://proxy.example");
|
||||
|
||||
let error = app
|
||||
.to_config()
|
||||
.expect_err("invalid upstream proxy should be rejected");
|
||||
assert!(error.to_string().contains("upstream proxy URL"));
|
||||
.expect_err("invalid egress proxy should be rejected");
|
||||
assert!(error.to_string().contains("egress proxy URL"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
//! Shared application state passed to all subsystems.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::sync::{Arc, Mutex, RwLock};
|
||||
use std::time::Duration;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_runtime::{AdmissionPermit, ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot};
|
||||
use aether_runtime_state::{RuntimeSemaphore, RuntimeSemaphoreError, RuntimeSemaphoreSnapshot};
|
||||
@@ -50,6 +52,8 @@ pub struct ServerContext {
|
||||
pub active_connections: Arc<AtomicU64>,
|
||||
/// Per-server request/latency metrics.
|
||||
pub metrics: Arc<ProxyMetrics>,
|
||||
/// Per-server tunnel stability/traffic metrics.
|
||||
pub tunnel_metrics: Arc<TunnelMetrics>,
|
||||
}
|
||||
|
||||
/// Aggregate metrics for reporting to Aether.
|
||||
@@ -83,6 +87,221 @@ impl ProxyMetrics {
|
||||
}
|
||||
}
|
||||
|
||||
const RECENT_TUNNEL_ERROR_CAPACITY: usize = 64;
|
||||
const TUNNEL_ERROR_CATEGORY_MAX_CHARS: usize = 48;
|
||||
const TUNNEL_ERROR_MESSAGE_MAX_CHARS: usize = 320;
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct TunnelErrorEvent {
|
||||
pub timestamp_unix_secs: u64,
|
||||
pub category: String,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct TunnelMetricsSnapshot {
|
||||
pub connect_attempts: u64,
|
||||
pub connect_successes: u64,
|
||||
pub connect_errors: u64,
|
||||
pub disconnects: u64,
|
||||
pub last_connected_at_unix_secs: u64,
|
||||
pub last_disconnected_at_unix_secs: u64,
|
||||
pub last_connected_duration_ms: u64,
|
||||
pub connected_duration_total_ms: u64,
|
||||
pub heartbeat_sent: u64,
|
||||
pub heartbeat_ack: u64,
|
||||
pub heartbeat_rtt_last_ms: u64,
|
||||
pub heartbeat_rtt_total_ms: u64,
|
||||
pub ws_in_frames: u64,
|
||||
pub ws_in_bytes: u64,
|
||||
pub ws_out_frames: u64,
|
||||
pub ws_out_bytes: u64,
|
||||
pub error_events_total: u64,
|
||||
}
|
||||
|
||||
impl TunnelMetricsSnapshot {
|
||||
pub fn heartbeat_rtt_avg_ms(self) -> Option<f64> {
|
||||
if self.heartbeat_ack == 0 {
|
||||
None
|
||||
} else {
|
||||
Some(self.heartbeat_rtt_total_ms as f64 / self.heartbeat_ack as f64)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TunnelMetrics {
|
||||
connect_attempts: AtomicU64,
|
||||
connect_successes: AtomicU64,
|
||||
connect_errors: AtomicU64,
|
||||
disconnects: AtomicU64,
|
||||
last_connected_at_unix_secs: AtomicU64,
|
||||
last_disconnected_at_unix_secs: AtomicU64,
|
||||
last_connected_duration_ms: AtomicU64,
|
||||
connected_duration_total_ms: AtomicU64,
|
||||
heartbeat_sent: AtomicU64,
|
||||
heartbeat_ack: AtomicU64,
|
||||
heartbeat_rtt_last_ms: AtomicU64,
|
||||
heartbeat_rtt_total_ms: AtomicU64,
|
||||
ws_in_frames: AtomicU64,
|
||||
ws_in_bytes: AtomicU64,
|
||||
ws_out_frames: AtomicU64,
|
||||
ws_out_bytes: AtomicU64,
|
||||
error_events_total: AtomicU64,
|
||||
recent_errors: Mutex<VecDeque<TunnelErrorEvent>>,
|
||||
}
|
||||
|
||||
impl TunnelMetrics {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
connect_attempts: AtomicU64::new(0),
|
||||
connect_successes: AtomicU64::new(0),
|
||||
connect_errors: AtomicU64::new(0),
|
||||
disconnects: AtomicU64::new(0),
|
||||
last_connected_at_unix_secs: AtomicU64::new(0),
|
||||
last_disconnected_at_unix_secs: AtomicU64::new(0),
|
||||
last_connected_duration_ms: AtomicU64::new(0),
|
||||
connected_duration_total_ms: AtomicU64::new(0),
|
||||
heartbeat_sent: AtomicU64::new(0),
|
||||
heartbeat_ack: AtomicU64::new(0),
|
||||
heartbeat_rtt_last_ms: AtomicU64::new(0),
|
||||
heartbeat_rtt_total_ms: AtomicU64::new(0),
|
||||
ws_in_frames: AtomicU64::new(0),
|
||||
ws_in_bytes: AtomicU64::new(0),
|
||||
ws_out_frames: AtomicU64::new(0),
|
||||
ws_out_bytes: AtomicU64::new(0),
|
||||
error_events_total: AtomicU64::new(0),
|
||||
recent_errors: Mutex::new(VecDeque::with_capacity(RECENT_TUNNEL_ERROR_CAPACITY)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn record_connect_attempt(&self) {
|
||||
self.connect_attempts.fetch_add(1, Ordering::Release);
|
||||
}
|
||||
|
||||
pub fn record_connect_success(&self) {
|
||||
self.connect_successes.fetch_add(1, Ordering::Release);
|
||||
self.last_connected_at_unix_secs
|
||||
.store(now_unix_secs(), Ordering::Release);
|
||||
}
|
||||
|
||||
pub fn record_connect_error(&self) {
|
||||
self.connect_errors.fetch_add(1, Ordering::Release);
|
||||
}
|
||||
|
||||
pub fn record_disconnect(&self, connected_for: Duration) {
|
||||
let duration_ms = duration_to_millis_u64(connected_for);
|
||||
self.disconnects.fetch_add(1, Ordering::Release);
|
||||
self.last_disconnected_at_unix_secs
|
||||
.store(now_unix_secs(), Ordering::Release);
|
||||
self.last_connected_duration_ms
|
||||
.store(duration_ms, Ordering::Release);
|
||||
self.connected_duration_total_ms
|
||||
.fetch_add(duration_ms, Ordering::Release);
|
||||
}
|
||||
|
||||
pub fn record_heartbeat_sent(&self) {
|
||||
self.heartbeat_sent.fetch_add(1, Ordering::Release);
|
||||
}
|
||||
|
||||
pub fn record_heartbeat_ack(&self, rtt: Duration) {
|
||||
let rtt_ms = duration_to_millis_u64(rtt);
|
||||
self.heartbeat_ack.fetch_add(1, Ordering::Release);
|
||||
self.heartbeat_rtt_last_ms.store(rtt_ms, Ordering::Release);
|
||||
self.heartbeat_rtt_total_ms
|
||||
.fetch_add(rtt_ms, Ordering::Release);
|
||||
}
|
||||
|
||||
pub fn record_ws_incoming_frame(&self, payload_len: usize) {
|
||||
self.ws_in_frames.fetch_add(1, Ordering::Release);
|
||||
self.ws_in_bytes.fetch_add(
|
||||
u64::try_from(payload_len).unwrap_or(u64::MAX),
|
||||
Ordering::Release,
|
||||
);
|
||||
}
|
||||
|
||||
pub fn record_ws_outgoing_frame(&self, payload_len: usize) {
|
||||
self.ws_out_frames.fetch_add(1, Ordering::Release);
|
||||
self.ws_out_bytes.fetch_add(
|
||||
u64::try_from(payload_len).unwrap_or(u64::MAX),
|
||||
Ordering::Release,
|
||||
);
|
||||
}
|
||||
|
||||
pub fn record_error(&self, category: &str, message: &str) {
|
||||
self.error_events_total.fetch_add(1, Ordering::Release);
|
||||
|
||||
let event = TunnelErrorEvent {
|
||||
timestamp_unix_secs: now_unix_secs(),
|
||||
category: normalize_error_field(category, TUNNEL_ERROR_CATEGORY_MAX_CHARS, "unknown"),
|
||||
message: normalize_error_field(message, TUNNEL_ERROR_MESSAGE_MAX_CHARS, "n/a"),
|
||||
};
|
||||
|
||||
let mut recent_errors = match self.recent_errors.lock() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => poisoned.into_inner(),
|
||||
};
|
||||
if recent_errors.len() >= RECENT_TUNNEL_ERROR_CAPACITY {
|
||||
recent_errors.pop_front();
|
||||
}
|
||||
recent_errors.push_back(event);
|
||||
}
|
||||
|
||||
pub fn recent_errors(&self, limit: usize) -> Vec<TunnelErrorEvent> {
|
||||
if limit == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
let recent_errors = match self.recent_errors.lock() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => poisoned.into_inner(),
|
||||
};
|
||||
let start = recent_errors.len().saturating_sub(limit);
|
||||
recent_errors.iter().skip(start).cloned().collect()
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> TunnelMetricsSnapshot {
|
||||
TunnelMetricsSnapshot {
|
||||
connect_attempts: self.connect_attempts.load(Ordering::Acquire),
|
||||
connect_successes: self.connect_successes.load(Ordering::Acquire),
|
||||
connect_errors: self.connect_errors.load(Ordering::Acquire),
|
||||
disconnects: self.disconnects.load(Ordering::Acquire),
|
||||
last_connected_at_unix_secs: self.last_connected_at_unix_secs.load(Ordering::Acquire),
|
||||
last_disconnected_at_unix_secs: self
|
||||
.last_disconnected_at_unix_secs
|
||||
.load(Ordering::Acquire),
|
||||
last_connected_duration_ms: self.last_connected_duration_ms.load(Ordering::Acquire),
|
||||
connected_duration_total_ms: self.connected_duration_total_ms.load(Ordering::Acquire),
|
||||
heartbeat_sent: self.heartbeat_sent.load(Ordering::Acquire),
|
||||
heartbeat_ack: self.heartbeat_ack.load(Ordering::Acquire),
|
||||
heartbeat_rtt_last_ms: self.heartbeat_rtt_last_ms.load(Ordering::Acquire),
|
||||
heartbeat_rtt_total_ms: self.heartbeat_rtt_total_ms.load(Ordering::Acquire),
|
||||
ws_in_frames: self.ws_in_frames.load(Ordering::Acquire),
|
||||
ws_in_bytes: self.ws_in_bytes.load(Ordering::Acquire),
|
||||
ws_out_frames: self.ws_out_frames.load(Ordering::Acquire),
|
||||
ws_out_bytes: self.ws_out_bytes.load(Ordering::Acquire),
|
||||
error_events_total: self.error_events_total.load(Ordering::Acquire),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn duration_to_millis_u64(duration: Duration) -> u64 {
|
||||
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
|
||||
}
|
||||
|
||||
fn normalize_error_field(value: &str, max_chars: usize, fallback: &str) -> String {
|
||||
let normalized = value.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
if normalized.is_empty() {
|
||||
return fallback.to_string();
|
||||
}
|
||||
normalized.chars().take(max_chars).collect()
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum ProxyAdmissionError {
|
||||
#[error("proxy stream admission saturated at {limit} for gate {gate}")]
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! WebSocket tunnel client: connect, authenticate, and run the tunnel.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::watch;
|
||||
@@ -10,6 +10,7 @@ use tokio_tungstenite::tungstenite::http;
|
||||
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::egress_proxy::{connect_target_via_proxy, ProxyConnectOptions, UpstreamProxyConfig};
|
||||
use crate::state::{AppState, ServerContext};
|
||||
|
||||
use super::{dispatcher, heartbeat, writer};
|
||||
@@ -71,14 +72,7 @@ pub async fn connect_and_run(
|
||||
.config
|
||||
.tunnel_connect_timeout()
|
||||
.expect("validated config should resolve tunnel connect timeout");
|
||||
let tcp_stream = tokio::time::timeout(connect_timeout, TcpStream::connect((host, port)))
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow::anyhow!(
|
||||
"tunnel TCP connect timeout ({}ms)",
|
||||
connect_timeout.as_millis()
|
||||
)
|
||||
})??;
|
||||
let tcp_stream = connect_tunnel_tcp(state, host, port, connect_timeout).await?;
|
||||
|
||||
// Configure TCP parameters via socket2
|
||||
configure_tcp_socket(&tcp_stream, state);
|
||||
@@ -133,6 +127,8 @@ pub async fn connect_and_run(
|
||||
ping_interval_ms = ping_interval.as_millis(),
|
||||
"tunnel connected"
|
||||
);
|
||||
server.tunnel_metrics.record_connect_success();
|
||||
let connected_at = Instant::now();
|
||||
|
||||
// NOTE: reconnect_attempts reset is handled by the caller (mod.rs)
|
||||
// based on how long the connection stayed alive.
|
||||
@@ -141,7 +137,11 @@ pub async fn connect_and_run(
|
||||
let (ws_sink, ws_read) = futures_util::StreamExt::split(ws_stream);
|
||||
|
||||
// Spawn writer task (with WebSocket ping keepalive)
|
||||
let (frame_tx, mut writer_handle) = writer::spawn_writer(ws_sink, ping_interval);
|
||||
let (frame_tx, mut writer_handle) = writer::spawn_writer_with_metrics(
|
||||
ws_sink,
|
||||
ping_interval,
|
||||
Some(Arc::clone(&server.tunnel_metrics)),
|
||||
);
|
||||
let drain_signal = spawn_drain_signal(conn_idx, frame_tx.clone(), drain.clone());
|
||||
|
||||
// Spawn heartbeat task (only for primary connection to avoid
|
||||
@@ -174,8 +174,13 @@ pub async fn connect_and_run(
|
||||
drain.clone(),
|
||||
) => {
|
||||
match result {
|
||||
Ok(()) => TunnelOutcome::Disconnected,
|
||||
Err(e) => return Err(e),
|
||||
Ok(()) => Ok(TunnelOutcome::Disconnected),
|
||||
Err(e) => {
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("dispatcher_error", &e.to_string());
|
||||
Err(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
writer_result = &mut writer_handle => {
|
||||
@@ -184,16 +189,22 @@ pub async fn connect_and_run(
|
||||
Err(e) => {
|
||||
if e.is_panic() {
|
||||
tracing::error!(error = %e, "writer task panicked, triggering reconnect");
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("writer_task_panic", &e.to_string());
|
||||
} else {
|
||||
warn!(error = %e, "writer task cancelled, triggering reconnect");
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("writer_task_cancelled", &e.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
TunnelOutcome::Disconnected
|
||||
Ok(TunnelOutcome::Disconnected)
|
||||
}
|
||||
_ = shutdown.changed() => {
|
||||
debug!("shutdown during tunnel dispatch");
|
||||
TunnelOutcome::Shutdown
|
||||
Ok(TunnelOutcome::Shutdown)
|
||||
}
|
||||
};
|
||||
|
||||
@@ -213,8 +224,12 @@ pub async fn connect_and_run(
|
||||
let _ = tokio::time::timeout(Duration::from_secs(35), writer_handle).await;
|
||||
}
|
||||
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_disconnect(connected_at.elapsed());
|
||||
|
||||
debug!("tunnel disconnected");
|
||||
Ok(outcome)
|
||||
outcome
|
||||
}
|
||||
|
||||
fn spawn_drain_signal(
|
||||
@@ -246,6 +261,56 @@ fn spawn_drain_signal(
|
||||
})
|
||||
}
|
||||
|
||||
async fn connect_tunnel_tcp(
|
||||
state: &Arc<AppState>,
|
||||
host: &str,
|
||||
port: u16,
|
||||
connect_timeout: Duration,
|
||||
) -> Result<TcpStream, anyhow::Error> {
|
||||
if let Some(proxy_url) = state.config.effective_aether_proxy_url() {
|
||||
let proxy = UpstreamProxyConfig::parse(proxy_url)
|
||||
.map_err(|err| anyhow::anyhow!("aether proxy URL invalid: {err}"))?;
|
||||
debug!(
|
||||
proxy_url = %proxy.redacted_url(),
|
||||
host = %host,
|
||||
port = port,
|
||||
"connecting tunnel via Aether egress proxy"
|
||||
);
|
||||
return tokio::time::timeout(
|
||||
connect_timeout,
|
||||
connect_target_via_proxy(
|
||||
&proxy,
|
||||
host,
|
||||
port,
|
||||
ProxyConnectOptions {
|
||||
connect_timeout,
|
||||
tcp_nodelay: state.config.tunnel_tcp_nodelay,
|
||||
tcp_keepalive: (state.config.tunnel_tcp_keepalive_secs > 0)
|
||||
.then(|| Duration::from_secs(state.config.tunnel_tcp_keepalive_secs)),
|
||||
},
|
||||
),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow::anyhow!(
|
||||
"tunnel proxy TCP connect timeout ({}ms)",
|
||||
connect_timeout.as_millis()
|
||||
)
|
||||
})?
|
||||
.map_err(anyhow::Error::from);
|
||||
}
|
||||
|
||||
tokio::time::timeout(connect_timeout, TcpStream::connect((host, port)))
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow::anyhow!(
|
||||
"tunnel TCP connect timeout ({}ms)",
|
||||
connect_timeout.as_millis()
|
||||
)
|
||||
})?
|
||||
.map_err(anyhow::Error::from)
|
||||
}
|
||||
|
||||
/// Configure TCP keepalive and NODELAY on an established socket.
|
||||
fn configure_tcp_socket(stream: &TcpStream, state: &Arc<AppState>) {
|
||||
let sock_ref = socket2::SockRef::from(stream);
|
||||
|
||||
@@ -84,6 +84,10 @@ where
|
||||
stale_ms = stale_timeout.as_millis(),
|
||||
"tunnel connection stale, no data received"
|
||||
);
|
||||
server.tunnel_metrics.record_error(
|
||||
"stale_timeout",
|
||||
&format!("no tunnel frame received for {}ms", stale_timeout.as_millis()),
|
||||
);
|
||||
break None;
|
||||
}
|
||||
};
|
||||
@@ -92,6 +96,9 @@ where
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
error!(error = %e, "WebSocket read error");
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("ws_read_error", &e.to_string());
|
||||
break Some(e);
|
||||
}
|
||||
};
|
||||
@@ -100,7 +107,10 @@ where
|
||||
last_data_at = tokio::time::Instant::now();
|
||||
|
||||
let data = match msg {
|
||||
Message::Binary(data) => Bytes::from(data),
|
||||
Message::Binary(data) => {
|
||||
server.tunnel_metrics.record_ws_incoming_frame(data.len());
|
||||
Bytes::from(data)
|
||||
}
|
||||
Message::Ping(_) => continue,
|
||||
Message::Pong(_) => continue,
|
||||
Message::Close(_) => {
|
||||
@@ -114,6 +124,9 @@ where
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
warn!(error = %e, "failed to decode frame");
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("frame_decode_error", &e.to_string());
|
||||
continue;
|
||||
}
|
||||
};
|
||||
@@ -223,6 +236,10 @@ where
|
||||
if is_end || dispatch != StreamDispatchStatus::Delivered {
|
||||
streams.remove(&sid);
|
||||
if dispatch == StreamDispatchStatus::TimedOut {
|
||||
server.tunnel_metrics.record_error(
|
||||
"stream_dispatch_timeout",
|
||||
&format!("request body dispatch timed out for stream {}", sid),
|
||||
);
|
||||
try_send_stream_error(
|
||||
&frame_tx,
|
||||
sid,
|
||||
|
||||
@@ -8,6 +8,7 @@ use std::time::UNIX_EPOCH;
|
||||
|
||||
use bytes::Bytes;
|
||||
use tokio::sync::watch;
|
||||
use tokio::time::Instant;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::registration::client::RemoteConfig;
|
||||
@@ -60,6 +61,13 @@ struct HeartbeatSnapshot {
|
||||
stream_errors: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct PendingHeartbeat {
|
||||
heartbeat_id: u64,
|
||||
snapshot: HeartbeatSnapshot,
|
||||
sent_at: Option<Instant>,
|
||||
}
|
||||
|
||||
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
|
||||
pub fn spawn(
|
||||
state: Arc<AppState>,
|
||||
@@ -76,7 +84,7 @@ pub fn spawn(
|
||||
// At most one in-flight heartbeat snapshot is tracked at a time.
|
||||
// Snapshot is only cleared after receiving an ACK, which avoids losing
|
||||
// interval counters when ACK/frame delivery is temporarily unstable.
|
||||
let mut pending: Option<(u64, HeartbeatSnapshot)> = None;
|
||||
let mut pending: Option<PendingHeartbeat> = None;
|
||||
let mut next_heartbeat_id: u64 = 1;
|
||||
let heartbeat_session_id = format!(
|
||||
"{}-{}",
|
||||
@@ -93,8 +101,8 @@ pub fn spawn(
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(current_interval) => {
|
||||
let (heartbeat_id, snapshot) = if let Some((id, snap)) = pending {
|
||||
(id, snap)
|
||||
let pending_entry = if let Some(entry) = pending {
|
||||
entry
|
||||
} else {
|
||||
let snap = collect_snapshot(&server);
|
||||
let id = next_heartbeat_id;
|
||||
@@ -102,24 +110,34 @@ pub fn spawn(
|
||||
if next_heartbeat_id == 0 {
|
||||
next_heartbeat_id = 1;
|
||||
}
|
||||
pending = Some((id, snap));
|
||||
(id, snap)
|
||||
let entry = PendingHeartbeat {
|
||||
heartbeat_id: id,
|
||||
snapshot: snap,
|
||||
sent_at: None,
|
||||
};
|
||||
pending = Some(entry);
|
||||
entry
|
||||
};
|
||||
|
||||
let payload = build_heartbeat_payload(
|
||||
&state,
|
||||
&server,
|
||||
&heartbeat_session_id,
|
||||
heartbeat_id,
|
||||
snapshot
|
||||
pending_entry.heartbeat_id,
|
||||
pending_entry.snapshot
|
||||
).await;
|
||||
let frame = Frame::control(MsgType::HeartbeatData, payload);
|
||||
if frame_tx.send(frame).await.is_err() {
|
||||
if let Some((_, snap)) = pending.take() {
|
||||
restore_snapshot(&server, snap);
|
||||
if let Some(entry) = pending.take() {
|
||||
restore_snapshot(&server, entry.snapshot);
|
||||
}
|
||||
break; // Writer closed
|
||||
}
|
||||
server.tunnel_metrics.record_heartbeat_sent();
|
||||
if let Some(mut entry) = pending {
|
||||
entry.sent_at = Some(Instant::now());
|
||||
pending = Some(entry);
|
||||
}
|
||||
debug!("sent heartbeat data");
|
||||
|
||||
// Re-read interval from dynamic config (remote config may have
|
||||
@@ -142,8 +160,11 @@ pub fn spawn(
|
||||
heartbeat_id: ack_id,
|
||||
upgrade_to,
|
||||
} => {
|
||||
if let Some((pending_id, _)) = pending {
|
||||
if ack_id == pending_id {
|
||||
if let Some(entry) = pending {
|
||||
if ack_id == entry.heartbeat_id {
|
||||
if let Some(sent_at) = entry.sent_at {
|
||||
server.tunnel_metrics.record_heartbeat_ack(sent_at.elapsed());
|
||||
}
|
||||
pending = None;
|
||||
}
|
||||
}
|
||||
@@ -154,8 +175,8 @@ pub fn spawn(
|
||||
}
|
||||
_ = shutdown.changed() => {
|
||||
debug!("heartbeat task shutting down");
|
||||
if let Some((_, snap)) = pending.take() {
|
||||
restore_snapshot(&server, snap);
|
||||
if let Some(entry) = pending.take() {
|
||||
restore_snapshot(&server, entry.snapshot);
|
||||
}
|
||||
break;
|
||||
}
|
||||
@@ -217,6 +238,8 @@ async fn build_heartbeat_payload(
|
||||
snapshot: HeartbeatSnapshot,
|
||||
) -> Bytes {
|
||||
let node_id = server.node_id.read().unwrap().clone();
|
||||
let tunnel_snapshot = server.tunnel_metrics.snapshot();
|
||||
let recent_errors = server.tunnel_metrics.recent_errors(8);
|
||||
|
||||
let avg_latency_ms = if snapshot.requests > 0 {
|
||||
Some(snapshot.latency_ns as f64 / snapshot.requests as f64 / 1_000_000.0)
|
||||
@@ -268,6 +291,26 @@ async fn build_heartbeat_payload(
|
||||
"proxy_metadata": {
|
||||
"version": CURRENT_VERSION,
|
||||
"admission": admission,
|
||||
"tunnel_metrics": {
|
||||
"connect_attempts": tunnel_snapshot.connect_attempts,
|
||||
"connect_successes": tunnel_snapshot.connect_successes,
|
||||
"connect_errors": tunnel_snapshot.connect_errors,
|
||||
"disconnects": tunnel_snapshot.disconnects,
|
||||
"last_connected_at_unix_secs": tunnel_snapshot.last_connected_at_unix_secs,
|
||||
"last_disconnected_at_unix_secs": tunnel_snapshot.last_disconnected_at_unix_secs,
|
||||
"last_connected_duration_ms": tunnel_snapshot.last_connected_duration_ms,
|
||||
"connected_duration_total_ms": tunnel_snapshot.connected_duration_total_ms,
|
||||
"heartbeat_sent": tunnel_snapshot.heartbeat_sent,
|
||||
"heartbeat_ack": tunnel_snapshot.heartbeat_ack,
|
||||
"heartbeat_rtt_last_ms": tunnel_snapshot.heartbeat_rtt_last_ms,
|
||||
"heartbeat_rtt_avg_ms": tunnel_snapshot.heartbeat_rtt_avg_ms(),
|
||||
"ws_in_frames": tunnel_snapshot.ws_in_frames,
|
||||
"ws_in_bytes": tunnel_snapshot.ws_in_bytes,
|
||||
"ws_out_frames": tunnel_snapshot.ws_out_frames,
|
||||
"ws_out_bytes": tunnel_snapshot.ws_out_bytes,
|
||||
"error_events_total": tunnel_snapshot.error_events_total,
|
||||
},
|
||||
"recent_tunnel_errors": recent_errors,
|
||||
},
|
||||
});
|
||||
|
||||
@@ -277,6 +320,9 @@ async fn build_heartbeat_payload(
|
||||
fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
|
||||
if payload.is_empty() {
|
||||
warn!("received empty heartbeat ACK");
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("heartbeat_ack_empty", "received empty heartbeat ACK");
|
||||
return AckDecision::Ignore;
|
||||
}
|
||||
|
||||
@@ -303,6 +349,9 @@ fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(error = %e, "failed to parse heartbeat ACK");
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("heartbeat_ack_parse", &e.to_string());
|
||||
AckDecision::Ignore
|
||||
}
|
||||
}
|
||||
@@ -373,7 +422,7 @@ mod tests {
|
||||
use super::{handle_ack, AckDecision};
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::DynamicConfig;
|
||||
use crate::state::{ProxyMetrics, ServerContext};
|
||||
use crate::state::{ProxyMetrics, ServerContext, TunnelMetrics};
|
||||
|
||||
fn sample_server() -> Arc<ServerContext> {
|
||||
let config = Arc::new(crate::config::Config::parse_from([
|
||||
@@ -399,6 +448,7 @@ mod tests {
|
||||
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
|
||||
active_connections: Arc::new(AtomicU64::new(0)),
|
||||
metrics: Arc::new(ProxyMetrics::new()),
|
||||
tunnel_metrics: Arc::new(TunnelMetrics::new()),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -76,6 +76,7 @@ pub async fn run(
|
||||
info!(server = %server.server_label, conn = conn_idx, "tunnel drained, exiting slot");
|
||||
return;
|
||||
}
|
||||
server.tunnel_metrics.record_connect_attempt();
|
||||
let started_at = Instant::now();
|
||||
match client::connect_and_run(state, server, conn_idx, &mut shutdown, drain.clone()).await {
|
||||
Ok(client::TunnelOutcome::Shutdown) => {
|
||||
@@ -86,6 +87,10 @@ pub async fn run(
|
||||
debug!(server = %server.server_label, conn = conn_idx, "tunnel disconnected, reconnecting");
|
||||
}
|
||||
Err(e) => {
|
||||
server.tunnel_metrics.record_connect_error();
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("tunnel_connect_error", &e.to_string());
|
||||
error!(server = %server.server_label, conn = conn_idx, error = %e, "tunnel connection error, reconnecting");
|
||||
}
|
||||
}
|
||||
@@ -237,7 +242,7 @@ mod tests {
|
||||
use crate::config::Config;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::DynamicConfig;
|
||||
use crate::state::{AppState as ProxyAppState, ProxyMetrics, ServerContext};
|
||||
use crate::state::{AppState as ProxyAppState, ProxyMetrics, ServerContext, TunnelMetrics};
|
||||
use crate::target_filter::DnsCache;
|
||||
use crate::tunnel::protocol;
|
||||
use crate::upstream_client;
|
||||
@@ -478,6 +483,7 @@ mod tests {
|
||||
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
|
||||
active_connections: Arc::new(AtomicU64::new(0)),
|
||||
metrics: Arc::new(ProxyMetrics::new()),
|
||||
tunnel_metrics: Arc::new(TunnelMetrics::new()),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -498,6 +504,7 @@ mod tests {
|
||||
aether_tcp_keepalive_secs: 60,
|
||||
aether_tcp_nodelay: true,
|
||||
aether_http2: true,
|
||||
aether_proxy_url: None,
|
||||
aether_retry_max_attempts: 1,
|
||||
aether_retry_base_delay_ms: 50,
|
||||
aether_retry_max_delay_ms: 100,
|
||||
|
||||
@@ -1532,7 +1532,7 @@ mod tests {
|
||||
use crate::config::Config;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::DynamicConfig;
|
||||
use crate::state::ProxyMetrics;
|
||||
use crate::state::{ProxyMetrics, TunnelMetrics};
|
||||
use crate::target_filter::DnsCache;
|
||||
use crate::tunnel::client::build_tls_config;
|
||||
|
||||
@@ -2335,6 +2335,7 @@ mod tests {
|
||||
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
|
||||
active_connections: Arc::new(AtomicU64::new(0)),
|
||||
metrics: Arc::new(ProxyMetrics::new()),
|
||||
tunnel_metrics: Arc::new(TunnelMetrics::new()),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2355,6 +2356,7 @@ mod tests {
|
||||
aether_tcp_keepalive_secs: 60,
|
||||
aether_tcp_nodelay: true,
|
||||
aether_http2: true,
|
||||
aether_proxy_url: None,
|
||||
aether_retry_max_attempts: 3,
|
||||
aether_retry_base_delay_ms: 200,
|
||||
aether_retry_max_delay_ms: 2_000,
|
||||
|
||||
@@ -5,9 +5,10 @@
|
||||
//! periodic WebSocket Ping frames to keep the connection alive through
|
||||
//! intermediary proxies (Nginx, Cloudflare, etc.).
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_contracts::tunnel::MsgType;
|
||||
use aether_contracts::tunnel::{MsgType, HEADER_SIZE};
|
||||
#[cfg(test)]
|
||||
use aether_runtime::QueueSnapshot;
|
||||
use aether_runtime::{bounded_queue, BoundedQueueSender, QueueSendError};
|
||||
@@ -16,6 +17,8 @@ use tokio::task::JoinHandle;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tracing::{debug, error, trace};
|
||||
|
||||
use crate::state::TunnelMetrics;
|
||||
|
||||
use super::protocol::Frame;
|
||||
|
||||
const HIGH_PRIORITY_QUEUE_CAPACITY: usize = 64;
|
||||
@@ -77,7 +80,20 @@ impl FrameSender {
|
||||
///
|
||||
/// `ping_interval` controls WebSocket-level Ping frequency (typically 15s).
|
||||
/// This keeps the connection alive through intermediary proxies/load-balancers.
|
||||
pub fn spawn_writer<S>(mut sink: S, ping_interval: Duration) -> (FrameSender, JoinHandle<()>)
|
||||
#[cfg(test)]
|
||||
pub fn spawn_writer<S>(sink: S, ping_interval: Duration) -> (FrameSender, JoinHandle<()>)
|
||||
where
|
||||
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
||||
{
|
||||
spawn_writer_with_metrics(sink, ping_interval, None)
|
||||
}
|
||||
|
||||
/// Spawn the writer task with optional tunnel metrics instrumentation.
|
||||
pub fn spawn_writer_with_metrics<S>(
|
||||
mut sink: S,
|
||||
ping_interval: Duration,
|
||||
tunnel_metrics: Option<Arc<TunnelMetrics>>,
|
||||
) -> (FrameSender, JoinHandle<()>)
|
||||
where
|
||||
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
||||
{
|
||||
@@ -93,7 +109,7 @@ where
|
||||
|
||||
loop {
|
||||
if let Ok(frame) = high_rx.try_recv() {
|
||||
if !write_frame(&mut sink, frame).await {
|
||||
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
@@ -107,7 +123,7 @@ where
|
||||
frame = high_rx.recv(), if high_open => {
|
||||
match frame {
|
||||
Some(frame) => {
|
||||
if !write_frame(&mut sink, frame).await {
|
||||
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -117,6 +133,9 @@ where
|
||||
_ = ping_ticker.tick(), if high_open || normal_open => {
|
||||
if let Err(e) = sink.send(Message::Ping(vec![])).await {
|
||||
error!(error = %e, "failed to send WebSocket ping");
|
||||
if let Some(metrics) = tunnel_metrics.as_deref() {
|
||||
metrics.record_error("ws_ping_error", &e.to_string());
|
||||
}
|
||||
break;
|
||||
}
|
||||
trace!("sent WebSocket ping");
|
||||
@@ -124,7 +143,7 @@ where
|
||||
frame = normal_rx.recv(), if normal_open => {
|
||||
match frame {
|
||||
Some(frame) => {
|
||||
if !write_frame(&mut sink, frame).await {
|
||||
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -156,15 +175,22 @@ fn classify_frame_priority(frame: &Frame) -> FramePriority {
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_frame<S>(sink: &mut S, frame: Frame) -> bool
|
||||
async fn write_frame<S>(sink: &mut S, frame: Frame, tunnel_metrics: Option<&TunnelMetrics>) -> bool
|
||||
where
|
||||
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
||||
{
|
||||
let data = frame.encode();
|
||||
let wire_len = data.len().max(HEADER_SIZE);
|
||||
if let Err(e) = sink.send(Message::Binary(data.into())).await {
|
||||
error!(error = %e, "failed to write frame to WebSocket");
|
||||
if let Some(metrics) = tunnel_metrics {
|
||||
metrics.record_error("ws_write_error", &e.to_string());
|
||||
}
|
||||
return false;
|
||||
}
|
||||
if let Some(metrics) = tunnel_metrics {
|
||||
metrics.record_ws_outgoing_frame(wire_len);
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::collections::HashMap;
|
||||
use std::convert::Infallible;
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::net::IpAddr;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
@@ -28,14 +28,15 @@ use hyper_util::client::legacy::Client;
|
||||
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
|
||||
use rustls::pki_types::ServerName;
|
||||
use rustls::ClientConfig;
|
||||
use socket2::{SockRef, TcpKeepalive};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_rustls::TlsConnector;
|
||||
use tower_service::Service;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::egress_proxy::{UpstreamProxyConfig, UpstreamProxyScheme};
|
||||
use crate::egress_proxy::{
|
||||
connect_proxy_tcp, http_connect, socks5_connect, ProxyConnectOptions, UpstreamProxyConfig,
|
||||
UpstreamProxyScheme,
|
||||
};
|
||||
use crate::target_filter::{self, DnsCache};
|
||||
|
||||
type BoxError = Box<dyn std::error::Error + Send + Sync>;
|
||||
@@ -263,12 +264,6 @@ pub struct InstrumentedConnector {
|
||||
tcp_keepalive: Option<Duration>,
|
||||
}
|
||||
|
||||
struct ProxyConnectOptions {
|
||||
connect_timeout: Duration,
|
||||
tcp_nodelay: bool,
|
||||
tcp_keepalive: Option<Duration>,
|
||||
}
|
||||
|
||||
impl Service<Uri> for InstrumentedConnector {
|
||||
type Response = TimedConn;
|
||||
type Error = BoxError;
|
||||
@@ -402,262 +397,6 @@ async fn connect_via_proxy(
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect_proxy_tcp(
|
||||
proxy: &UpstreamProxyConfig,
|
||||
connect_timeout: Duration,
|
||||
tcp_nodelay: bool,
|
||||
tcp_keepalive: Option<Duration>,
|
||||
) -> io::Result<TcpStream> {
|
||||
let resolved = tokio::time::timeout(
|
||||
connect_timeout,
|
||||
tokio::net::lookup_host((proxy.host(), proxy.port())),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "proxy DNS timeout"))?
|
||||
.map_err(|err| io::Error::other(format!("proxy DNS failed: {err}")))?;
|
||||
|
||||
let mut last_error = None;
|
||||
for addr in resolved {
|
||||
match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
|
||||
Ok(Ok(stream)) => {
|
||||
configure_tcp_stream(&stream, tcp_nodelay, tcp_keepalive)?;
|
||||
return Ok(stream);
|
||||
}
|
||||
Ok(Err(error)) => last_error = Some(error),
|
||||
Err(_) => {
|
||||
last_error = Some(io::Error::new(
|
||||
io::ErrorKind::TimedOut,
|
||||
format!("proxy connect timeout: {addr}"),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or_else(|| io::Error::other("proxy DNS returned no addresses")))
|
||||
}
|
||||
|
||||
fn configure_tcp_stream(
|
||||
stream: &TcpStream,
|
||||
tcp_nodelay: bool,
|
||||
tcp_keepalive: Option<Duration>,
|
||||
) -> io::Result<()> {
|
||||
stream.set_nodelay(tcp_nodelay)?;
|
||||
if let Some(keepalive) = tcp_keepalive {
|
||||
let keepalive = TcpKeepalive::new().with_time(keepalive);
|
||||
SockRef::from(stream).set_tcp_keepalive(&keepalive)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn http_connect(
|
||||
stream: &mut TcpStream,
|
||||
target_authority: &str,
|
||||
proxy: &UpstreamProxyConfig,
|
||||
) -> io::Result<()> {
|
||||
let mut request = format!(
|
||||
"CONNECT {target_authority} HTTP/1.1\r\nHost: {target_authority}\r\nProxy-Connection: Keep-Alive\r\n"
|
||||
);
|
||||
if let Some(auth) = proxy.basic_auth_header() {
|
||||
request.push_str("Proxy-Authorization: ");
|
||||
request.push_str(&auth);
|
||||
request.push_str("\r\n");
|
||||
}
|
||||
request.push_str("\r\n");
|
||||
stream.write_all(request.as_bytes()).await?;
|
||||
stream.flush().await?;
|
||||
|
||||
let mut response = Vec::with_capacity(1024);
|
||||
let mut chunk = [0u8; 1024];
|
||||
loop {
|
||||
if response.len() >= 16 * 1024 {
|
||||
return Err(io::Error::other("proxy CONNECT response too large"));
|
||||
}
|
||||
let n = stream.read(&mut chunk).await?;
|
||||
if n == 0 {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::UnexpectedEof,
|
||||
"proxy closed during CONNECT",
|
||||
));
|
||||
}
|
||||
response.extend_from_slice(&chunk[..n]);
|
||||
if response.windows(4).any(|window| window == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let status_line_end = response
|
||||
.windows(2)
|
||||
.position(|window| window == b"\r\n")
|
||||
.ok_or_else(|| io::Error::other("proxy CONNECT response missing status line"))?;
|
||||
let status_line = std::str::from_utf8(&response[..status_line_end])
|
||||
.map_err(|_| io::Error::other("proxy CONNECT status line is not UTF-8"))?;
|
||||
let status = status_line.split_whitespace().nth(1).unwrap_or_default();
|
||||
if status == "200" {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::other(format!(
|
||||
"proxy CONNECT failed: {status_line}"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
async fn socks5_connect(
|
||||
stream: &mut TcpStream,
|
||||
proxy: &UpstreamProxyConfig,
|
||||
target_host: &str,
|
||||
target_port: u16,
|
||||
) -> io::Result<()> {
|
||||
let requires_auth = proxy.username().is_some();
|
||||
if requires_auth {
|
||||
stream.write_all(&[0x05, 0x02, 0x00, 0x02]).await?;
|
||||
} else {
|
||||
stream.write_all(&[0x05, 0x01, 0x00]).await?;
|
||||
}
|
||||
|
||||
let mut method_response = [0u8; 2];
|
||||
stream.read_exact(&mut method_response).await?;
|
||||
if method_response[0] != 0x05 {
|
||||
return Err(io::Error::other("invalid SOCKS5 method response"));
|
||||
}
|
||||
match method_response[1] {
|
||||
0x00 => {}
|
||||
0x02 => socks5_authenticate(stream, proxy).await?,
|
||||
0xff => return Err(io::Error::other("SOCKS5 proxy rejected all auth methods")),
|
||||
method => {
|
||||
return Err(io::Error::other(format!(
|
||||
"SOCKS5 proxy selected unsupported auth method 0x{method:02x}"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
let address = socks5_target_address(target_host, target_port, proxy.uses_remote_dns()).await?;
|
||||
stream.write_all(&address).await?;
|
||||
|
||||
let mut response = [0u8; 4];
|
||||
stream.read_exact(&mut response).await?;
|
||||
if response[0] != 0x05 {
|
||||
return Err(io::Error::other("invalid SOCKS5 connect response"));
|
||||
}
|
||||
if response[1] != 0x00 {
|
||||
return Err(io::Error::other(format!(
|
||||
"SOCKS5 connect failed: {}",
|
||||
socks5_reply_message(response[1])
|
||||
)));
|
||||
}
|
||||
|
||||
match response[3] {
|
||||
0x01 => {
|
||||
let mut ignored = [0u8; 4 + 2];
|
||||
stream.read_exact(&mut ignored).await?;
|
||||
}
|
||||
0x03 => {
|
||||
let mut len = [0u8; 1];
|
||||
stream.read_exact(&mut len).await?;
|
||||
let mut ignored = vec![0u8; len[0] as usize + 2];
|
||||
stream.read_exact(&mut ignored).await?;
|
||||
}
|
||||
0x04 => {
|
||||
let mut ignored = [0u8; 16 + 2];
|
||||
stream.read_exact(&mut ignored).await?;
|
||||
}
|
||||
atyp => {
|
||||
return Err(io::Error::other(format!(
|
||||
"SOCKS5 proxy returned unsupported address type 0x{atyp:02x}"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn socks5_authenticate(
|
||||
stream: &mut TcpStream,
|
||||
proxy: &UpstreamProxyConfig,
|
||||
) -> io::Result<()> {
|
||||
let username = proxy.username().unwrap_or_default().as_bytes();
|
||||
let password = proxy.password().unwrap_or_default().as_bytes();
|
||||
if username.len() > u8::MAX as usize || password.len() > u8::MAX as usize {
|
||||
return Err(io::Error::other(
|
||||
"SOCKS5 username/password must be at most 255 bytes",
|
||||
));
|
||||
}
|
||||
|
||||
let mut request = Vec::with_capacity(username.len() + password.len() + 3);
|
||||
request.push(0x01);
|
||||
request.push(username.len() as u8);
|
||||
request.extend_from_slice(username);
|
||||
request.push(password.len() as u8);
|
||||
request.extend_from_slice(password);
|
||||
stream.write_all(&request).await?;
|
||||
|
||||
let mut response = [0u8; 2];
|
||||
stream.read_exact(&mut response).await?;
|
||||
if response[0] != 0x01 || response[1] != 0x00 {
|
||||
return Err(io::Error::other("SOCKS5 username/password auth failed"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn socks5_target_address(
|
||||
target_host: &str,
|
||||
target_port: u16,
|
||||
remote_dns: bool,
|
||||
) -> io::Result<Vec<u8>> {
|
||||
let mut request = vec![0x05, 0x01, 0x00];
|
||||
if let Ok(ip) = target_host.parse::<IpAddr>() {
|
||||
push_socks5_ip_address(&mut request, ip);
|
||||
} else if remote_dns {
|
||||
let host = target_host.as_bytes();
|
||||
if host.len() > u8::MAX as usize {
|
||||
return Err(io::Error::other("SOCKS5 target hostname is too long"));
|
||||
}
|
||||
request.push(0x03);
|
||||
request.push(host.len() as u8);
|
||||
request.extend_from_slice(host);
|
||||
} else {
|
||||
let mut resolved = tokio::net::lookup_host((target_host, target_port))
|
||||
.await
|
||||
.map_err(|err| io::Error::other(format!("SOCKS5 target DNS failed: {err}")))?;
|
||||
let addr = resolved
|
||||
.next()
|
||||
.ok_or_else(|| io::Error::other("SOCKS5 target DNS returned no addresses"))?;
|
||||
push_socks5_socket_address(&mut request, addr);
|
||||
}
|
||||
request.extend_from_slice(&target_port.to_be_bytes());
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
fn push_socks5_socket_address(request: &mut Vec<u8>, addr: SocketAddr) {
|
||||
push_socks5_ip_address(request, addr.ip());
|
||||
}
|
||||
|
||||
fn push_socks5_ip_address(request: &mut Vec<u8>, ip: IpAddr) {
|
||||
match ip {
|
||||
IpAddr::V4(ip) => {
|
||||
request.push(0x01);
|
||||
request.extend_from_slice(&ip.octets());
|
||||
}
|
||||
IpAddr::V6(ip) => {
|
||||
request.push(0x04);
|
||||
request.extend_from_slice(&ip.octets());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn socks5_reply_message(reply: u8) -> &'static str {
|
||||
match reply {
|
||||
0x01 => "general failure",
|
||||
0x02 => "connection not allowed",
|
||||
0x03 => "network unreachable",
|
||||
0x04 => "host unreachable",
|
||||
0x05 => "connection refused",
|
||||
0x06 => "TTL expired",
|
||||
0x07 => "command not supported",
|
||||
0x08 => "address type not supported",
|
||||
_ => "unknown error",
|
||||
}
|
||||
}
|
||||
|
||||
fn uri_host(uri: &Uri) -> Result<String, io::Error> {
|
||||
uri.host()
|
||||
.map(|host| {
|
||||
@@ -934,8 +673,11 @@ mod tests {
|
||||
use clap::Parser;
|
||||
use http_body_util::BodyExt;
|
||||
use hyper::Response;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
use crate::egress_proxy::socks5_target_address;
|
||||
|
||||
#[test]
|
||||
fn fresh_connection_uses_connector_breakdown() {
|
||||
let mut response = Response::new(());
|
||||
|
||||
Reference in New Issue
Block a user