mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat(proxy): record tunnel stability metrics
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,
|
||||
};
|
||||
@@ -284,7 +285,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!(
|
||||
@@ -320,7 +327,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!(
|
||||
@@ -508,7 +521,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) => {
|
||||
@@ -578,7 +597,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,21 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn cleanup_proxy_node_metrics(
|
||||
&self,
|
||||
retain_1m_from_unix_secs: u64,
|
||||
retain_1h_from_unix_secs: u64,
|
||||
) -> 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)
|
||||
.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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
use aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary;
|
||||
use aether_data_contracts::DataLayerError;
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
use super::now_unix_secs;
|
||||
|
||||
const PROXY_NODE_METRICS_1M_RETENTION_SECS: u64 = 30 * 24 * 60 * 60;
|
||||
const PROXY_NODE_METRICS_1H_RETENTION_SECS: u64 = 180 * 24 * 60 * 60;
|
||||
|
||||
pub(super) async fn cleanup_proxy_node_metrics_once(
|
||||
data: &GatewayDataState,
|
||||
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
|
||||
let now = now_unix_secs();
|
||||
data.cleanup_proxy_node_metrics(
|
||||
now.saturating_sub(PROXY_NODE_METRICS_1M_RETENTION_SECS),
|
||||
now.saturating_sub(PROXY_NODE_METRICS_1H_RETENTION_SECS),
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -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?;
|
||||
|
||||
@@ -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,18 @@ 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,
|
||||
) -> Result<aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary, GatewayError>
|
||||
{
|
||||
self.data
|
||||
.cleanup_proxy_node_metrics(retain_1m_from_unix_secs, retain_1h_from_unix_secs)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn apply_proxy_node_heartbeat(
|
||||
&self,
|
||||
mutation: &ProxyNodeHeartbeatMutation,
|
||||
@@ -979,6 +1032,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(());
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
use aether_data::repository::{
|
||||
auth_modules::{StoredLdapModuleConfig, StoredOAuthProviderModuleConfig},
|
||||
proxy_nodes::{StoredProxyNode, StoredProxyNodeEvent},
|
||||
proxy_nodes::{
|
||||
ProxyNodeMetricsStep, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
StoredProxyNodeMetricsBucket,
|
||||
},
|
||||
system::StoredSystemConfigEntry,
|
||||
wallet::StoredWalletSnapshot,
|
||||
};
|
||||
@@ -1832,11 +1835,13 @@ pub fn build_admin_proxy_node_event_payload(event: &StoredProxyNodeEvent) -> ser
|
||||
"id": event.id,
|
||||
"event_type": event.event_type,
|
||||
"detail": event.detail,
|
||||
"event_metadata": event.event_metadata,
|
||||
"created_at": event.created_at_unix_ms.and_then(unix_secs_to_rfc3339),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn admin_proxy_node_event_node_id_from_path(request_path: &str) -> Option<&str> {
|
||||
let request_path = request_path.trim_end_matches('/');
|
||||
let node_id = request_path.strip_prefix("/api/admin/proxy-nodes/")?;
|
||||
let node_id = node_id.strip_suffix("/events")?;
|
||||
if node_id.is_empty() || node_id.contains('/') {
|
||||
@@ -1846,6 +1851,219 @@ pub fn admin_proxy_node_event_node_id_from_path(request_path: &str) -> Option<&s
|
||||
}
|
||||
}
|
||||
|
||||
pub fn admin_proxy_node_metrics_node_id_from_path(request_path: &str) -> Option<&str> {
|
||||
let request_path = request_path.trim_end_matches('/');
|
||||
let node_id = request_path.strip_prefix("/api/admin/proxy-nodes/")?;
|
||||
let node_id = node_id.strip_suffix("/metrics")?;
|
||||
if node_id.is_empty() || node_id.contains('/') {
|
||||
None
|
||||
} else {
|
||||
Some(node_id)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_admin_proxy_node_metrics_payload_response(
|
||||
step: ProxyNodeMetricsStep,
|
||||
from_unix_secs: u64,
|
||||
to_unix_secs: u64,
|
||||
items: Vec<StoredProxyNodeMetricsBucket>,
|
||||
) -> Response<Body> {
|
||||
let summary = summarize_proxy_node_metric_buckets(items.iter().map(|item| {
|
||||
(
|
||||
item.samples,
|
||||
item.uptime_samples,
|
||||
item.active_connections_sum,
|
||||
item.active_connections_max,
|
||||
item.heartbeat_rtt_ms_sum,
|
||||
item.heartbeat_rtt_ms_max,
|
||||
item.connect_errors_delta,
|
||||
item.disconnects_delta,
|
||||
item.error_events_delta,
|
||||
item.ws_in_bytes_delta,
|
||||
item.ws_out_bytes_delta,
|
||||
item.ws_in_frames_delta,
|
||||
item.ws_out_frames_delta,
|
||||
)
|
||||
}));
|
||||
let items = items
|
||||
.into_iter()
|
||||
.map(build_admin_proxy_node_metrics_bucket_payload)
|
||||
.collect::<Vec<_>>();
|
||||
Json(json!({
|
||||
"step": step.as_api_value(),
|
||||
"from": from_unix_secs,
|
||||
"to": to_unix_secs,
|
||||
"items": items,
|
||||
"summary": summary,
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
pub fn build_admin_proxy_fleet_metrics_payload_response(
|
||||
step: ProxyNodeMetricsStep,
|
||||
from_unix_secs: u64,
|
||||
to_unix_secs: u64,
|
||||
items: Vec<StoredProxyFleetMetricsBucket>,
|
||||
) -> Response<Body> {
|
||||
let summary = summarize_proxy_node_metric_buckets(items.iter().map(|item| {
|
||||
(
|
||||
item.samples,
|
||||
item.uptime_samples,
|
||||
item.active_connections_sum,
|
||||
item.active_connections_max,
|
||||
item.heartbeat_rtt_ms_sum,
|
||||
item.heartbeat_rtt_ms_max,
|
||||
item.connect_errors_delta,
|
||||
item.disconnects_delta,
|
||||
item.error_events_delta,
|
||||
item.ws_in_bytes_delta,
|
||||
item.ws_out_bytes_delta,
|
||||
item.ws_in_frames_delta,
|
||||
item.ws_out_frames_delta,
|
||||
)
|
||||
}));
|
||||
let items = items
|
||||
.into_iter()
|
||||
.map(build_admin_proxy_fleet_metrics_bucket_payload)
|
||||
.collect::<Vec<_>>();
|
||||
Json(json!({
|
||||
"step": step.as_api_value(),
|
||||
"from": from_unix_secs,
|
||||
"to": to_unix_secs,
|
||||
"items": items,
|
||||
"summary": summary,
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn build_admin_proxy_node_metrics_bucket_payload(
|
||||
item: StoredProxyNodeMetricsBucket,
|
||||
) -> serde_json::Value {
|
||||
json!({
|
||||
"node_id": item.node_id,
|
||||
"bucket_start_unix_secs": item.bucket_start_unix_secs,
|
||||
"bucket_start": unix_secs_to_rfc3339(item.bucket_start_unix_secs),
|
||||
"samples": item.samples,
|
||||
"uptime_samples": item.uptime_samples,
|
||||
"uptime_ratio": ratio(item.uptime_samples, item.samples),
|
||||
"active_connections_sum": item.active_connections_sum,
|
||||
"active_connections_max": item.active_connections_max,
|
||||
"active_connections_avg": ratio(item.active_connections_sum, item.samples),
|
||||
"heartbeat_rtt_ms_sum": item.heartbeat_rtt_ms_sum,
|
||||
"heartbeat_rtt_ms_max": item.heartbeat_rtt_ms_max,
|
||||
"heartbeat_rtt_ms_avg": ratio(item.heartbeat_rtt_ms_sum, item.samples),
|
||||
"connect_errors_delta": item.connect_errors_delta,
|
||||
"disconnects_delta": item.disconnects_delta,
|
||||
"error_events_delta": item.error_events_delta,
|
||||
"ws_in_bytes_delta": item.ws_in_bytes_delta,
|
||||
"ws_out_bytes_delta": item.ws_out_bytes_delta,
|
||||
"ws_in_frames_delta": item.ws_in_frames_delta,
|
||||
"ws_out_frames_delta": item.ws_out_frames_delta,
|
||||
})
|
||||
}
|
||||
|
||||
fn build_admin_proxy_fleet_metrics_bucket_payload(
|
||||
item: StoredProxyFleetMetricsBucket,
|
||||
) -> serde_json::Value {
|
||||
json!({
|
||||
"bucket_start_unix_secs": item.bucket_start_unix_secs,
|
||||
"bucket_start": unix_secs_to_rfc3339(item.bucket_start_unix_secs),
|
||||
"samples": item.samples,
|
||||
"uptime_samples": item.uptime_samples,
|
||||
"uptime_ratio": ratio(item.uptime_samples, item.samples),
|
||||
"active_connections_sum": item.active_connections_sum,
|
||||
"active_connections_max": item.active_connections_max,
|
||||
"active_connections_avg": ratio(item.active_connections_sum, item.samples),
|
||||
"heartbeat_rtt_ms_sum": item.heartbeat_rtt_ms_sum,
|
||||
"heartbeat_rtt_ms_max": item.heartbeat_rtt_ms_max,
|
||||
"heartbeat_rtt_ms_avg": ratio(item.heartbeat_rtt_ms_sum, item.samples),
|
||||
"connect_errors_delta": item.connect_errors_delta,
|
||||
"disconnects_delta": item.disconnects_delta,
|
||||
"error_events_delta": item.error_events_delta,
|
||||
"ws_in_bytes_delta": item.ws_in_bytes_delta,
|
||||
"ws_out_bytes_delta": item.ws_out_bytes_delta,
|
||||
"ws_in_frames_delta": item.ws_in_frames_delta,
|
||||
"ws_out_frames_delta": item.ws_out_frames_delta,
|
||||
})
|
||||
}
|
||||
|
||||
fn summarize_proxy_node_metric_buckets<I>(items: I) -> serde_json::Value
|
||||
where
|
||||
I: IntoIterator<
|
||||
Item = (
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
),
|
||||
>,
|
||||
{
|
||||
let mut samples = 0;
|
||||
let mut uptime_samples = 0;
|
||||
let mut active_connections_sum = 0;
|
||||
let mut active_connections_max = 0;
|
||||
let mut heartbeat_rtt_ms_sum = 0;
|
||||
let mut heartbeat_rtt_ms_max = 0;
|
||||
let mut connect_errors_delta = 0;
|
||||
let mut disconnects_delta = 0;
|
||||
let mut error_events_delta = 0;
|
||||
let mut ws_in_bytes_delta = 0;
|
||||
let mut ws_out_bytes_delta = 0;
|
||||
let mut ws_in_frames_delta = 0;
|
||||
let mut ws_out_frames_delta = 0;
|
||||
|
||||
for item in items {
|
||||
samples += item.0;
|
||||
uptime_samples += item.1;
|
||||
active_connections_sum += item.2;
|
||||
active_connections_max = active_connections_max.max(item.3);
|
||||
heartbeat_rtt_ms_sum += item.4;
|
||||
heartbeat_rtt_ms_max = heartbeat_rtt_ms_max.max(item.5);
|
||||
connect_errors_delta += item.6;
|
||||
disconnects_delta += item.7;
|
||||
error_events_delta += item.8;
|
||||
ws_in_bytes_delta += item.9;
|
||||
ws_out_bytes_delta += item.10;
|
||||
ws_in_frames_delta += item.11;
|
||||
ws_out_frames_delta += item.12;
|
||||
}
|
||||
|
||||
json!({
|
||||
"samples": samples,
|
||||
"uptime_samples": uptime_samples,
|
||||
"uptime_ratio": ratio(uptime_samples, samples),
|
||||
"active_connections_sum": active_connections_sum,
|
||||
"active_connections_max": active_connections_max,
|
||||
"active_connections_avg": ratio(active_connections_sum, samples),
|
||||
"heartbeat_rtt_ms_sum": heartbeat_rtt_ms_sum,
|
||||
"heartbeat_rtt_ms_max": heartbeat_rtt_ms_max,
|
||||
"heartbeat_rtt_ms_avg": ratio(heartbeat_rtt_ms_sum, samples),
|
||||
"connect_errors_delta": connect_errors_delta,
|
||||
"disconnects_delta": disconnects_delta,
|
||||
"error_events_delta": error_events_delta,
|
||||
"ws_in_bytes_delta": ws_in_bytes_delta,
|
||||
"ws_out_bytes_delta": ws_out_bytes_delta,
|
||||
"ws_in_frames_delta": ws_in_frames_delta,
|
||||
"ws_out_frames_delta": ws_out_frames_delta,
|
||||
})
|
||||
}
|
||||
|
||||
fn ratio(numerator: i64, denominator: i64) -> Option<f64> {
|
||||
if denominator <= 0 {
|
||||
return None;
|
||||
}
|
||||
Some(numerator as f64 / denominator as f64)
|
||||
}
|
||||
|
||||
pub fn build_admin_proxy_nodes_list_payload_response(
|
||||
items: Vec<serde_json::Value>,
|
||||
total: usize,
|
||||
|
||||
@@ -58,6 +58,17 @@ pub use crate::formats::openai::shared::{
|
||||
pub use crate::formats::shared::error_body::{
|
||||
build_core_error_body_for_client_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind,
|
||||
};
|
||||
pub use crate::formats::shared::image_bridge::{
|
||||
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_openai_image_provider_body_from_response_stream_sync_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, gemini_request_is_image_generation,
|
||||
resolve_requested_gemini_image_model_for_request, GeminiImageRequestForOpenAi,
|
||||
OpenAiImageRequestForGemini,
|
||||
};
|
||||
pub use crate::formats::shared::model_directives::{
|
||||
apply_model_directive_mapping_patch, apply_model_directive_overrides_from_model,
|
||||
apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort,
|
||||
|
||||
@@ -368,6 +368,7 @@ pub fn maybe_build_openai_image_sync_finalize_product(
|
||||
report_kind: &str,
|
||||
status_code: u16,
|
||||
report_context: Option<&Value>,
|
||||
body_json: Option<&Value>,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<OpenAiImageSyncFinalizeProduct>, AiSurfaceFinalizeError> {
|
||||
if report_kind != OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND || status_code >= 400 {
|
||||
@@ -384,6 +385,45 @@ pub fn maybe_build_openai_image_sync_finalize_product(
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let provider_api_format = report_context
|
||||
.get("provider_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if provider_api_format == "gemini:generate_content" {
|
||||
let Some(provider_body_json) = body_json else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(client_body_json) =
|
||||
crate::formats::shared::image_bridge::build_openai_image_response_from_gemini_response(
|
||||
provider_body_json,
|
||||
Some(report_context),
|
||||
)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
return Ok(Some(OpenAiImageSyncFinalizeProduct {
|
||||
client_body_json,
|
||||
provider_body_json: provider_body_json.clone(),
|
||||
}));
|
||||
}
|
||||
if provider_api_format != "openai:image" {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(provider_body_json) = body_json {
|
||||
if provider_body_json.get("output").is_some() && provider_body_json.get("data").is_none() {
|
||||
let Some(client_body_json) = crate::formats::shared::image_bridge::build_openai_image_response_from_response_stream_sync_body(
|
||||
provider_body_json,
|
||||
Some(report_context),
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
return Ok(Some(OpenAiImageSyncFinalizeProduct {
|
||||
client_body_json,
|
||||
provider_body_json: provider_body_json.clone(),
|
||||
}));
|
||||
}
|
||||
}
|
||||
let Some(body_base64) = body_base64 else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -689,6 +729,7 @@ mod tests {
|
||||
"openai_image_sync_finalize",
|
||||
200,
|
||||
Some(&report_context),
|
||||
None,
|
||||
Some(&body_base64),
|
||||
)
|
||||
.expect("finalize should succeed")
|
||||
|
||||
1043
crates/aether-ai-formats/src/formats/shared/image_bridge.rs
Normal file
1043
crates/aether-ai-formats/src/formats/shared/image_bridge.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -2,6 +2,7 @@ use std::fmt;
|
||||
|
||||
pub mod error_body;
|
||||
pub mod family;
|
||||
pub mod image_bridge;
|
||||
pub mod model_directives;
|
||||
pub mod passthrough;
|
||||
pub mod request;
|
||||
|
||||
@@ -595,6 +595,17 @@ pub fn maybe_build_standard_cross_format_sync_product(
|
||||
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
|
||||
let client_api_format = client_api_format.trim().to_ascii_lowercase();
|
||||
|
||||
if provider_api_format == "openai:image" && client_api_format == "gemini:generate_content" {
|
||||
let client_body_json = crate::formats::shared::image_bridge::build_gemini_image_response_from_openai_responses_image_response(
|
||||
&provider_body_json,
|
||||
Some(report_context),
|
||||
)?;
|
||||
return Some(StandardCrossFormatSyncProduct {
|
||||
client_body_json,
|
||||
provider_body_json,
|
||||
});
|
||||
}
|
||||
|
||||
let client_body_json = if is_standard_chat_finalize_kind(report_kind) {
|
||||
sync_chat_response_conversion_kind(&provider_api_format, &client_api_format)?;
|
||||
convert_standard_chat_response(
|
||||
|
||||
@@ -27,7 +27,12 @@ pub fn maybe_bridge_standard_sync_json_to_stream(
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
|
||||
let provider_api_format = normalize_api_format(provider_api_format);
|
||||
let client_api_format = normalize_api_format(client_api_format);
|
||||
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
|
||||
if client_api_format == "openai:image"
|
||||
&& matches!(
|
||||
provider_api_format.as_str(),
|
||||
"openai:image" | "gemini:generate_content"
|
||||
)
|
||||
{
|
||||
return maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context);
|
||||
}
|
||||
if !is_standard_api_format(provider_api_format.as_str())
|
||||
@@ -67,6 +72,36 @@ fn maybe_bridge_openai_image_sync_json_to_stream(
|
||||
provider_body_json: &Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
|
||||
let provider_api_format = report_context
|
||||
.and_then(|value| value.get("provider_api_format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or("openai:image");
|
||||
let owned_response;
|
||||
let provider_body_json = if provider_api_format == "gemini:generate_content" {
|
||||
let Some(converted) =
|
||||
crate::formats::shared::image_bridge::build_openai_image_response_from_gemini_response(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
owned_response = converted;
|
||||
&owned_response
|
||||
} else if provider_body_json.get("output").is_some() && provider_body_json.get("data").is_none()
|
||||
{
|
||||
let Some(converted) = crate::formats::shared::image_bridge::build_openai_image_response_from_response_stream_sync_body(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
owned_response = converted;
|
||||
&owned_response
|
||||
} else {
|
||||
provider_body_json
|
||||
};
|
||||
let Some(response) = provider_body_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
ALTER TABLE proxy_node_events
|
||||
ADD COLUMN event_metadata TEXT NULL AFTER detail;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS proxy_node_metrics_1m (
|
||||
node_id VARCHAR(64) NOT NULL,
|
||||
bucket_start_unix_secs BIGINT NOT NULL,
|
||||
samples BIGINT NOT NULL DEFAULT 0,
|
||||
uptime_samples BIGINT NOT NULL DEFAULT 0,
|
||||
active_connections_sum BIGINT NOT NULL DEFAULT 0,
|
||||
active_connections_max BIGINT NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_sum BIGINT NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_max BIGINT NOT NULL DEFAULT 0,
|
||||
connect_errors_delta BIGINT NOT NULL DEFAULT 0,
|
||||
disconnects_delta BIGINT NOT NULL DEFAULT 0,
|
||||
error_events_delta BIGINT NOT NULL DEFAULT 0,
|
||||
ws_in_bytes_delta BIGINT NOT NULL DEFAULT 0,
|
||||
ws_out_bytes_delta BIGINT NOT NULL DEFAULT 0,
|
||||
ws_in_frames_delta BIGINT NOT NULL DEFAULT 0,
|
||||
ws_out_frames_delta BIGINT NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs),
|
||||
INDEX idx_proxy_node_metrics_1m_bucket_start (bucket_start_unix_secs)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS proxy_node_metrics_1h (
|
||||
node_id VARCHAR(64) NOT NULL,
|
||||
bucket_start_unix_secs BIGINT NOT NULL,
|
||||
samples BIGINT NOT NULL DEFAULT 0,
|
||||
uptime_samples BIGINT NOT NULL DEFAULT 0,
|
||||
active_connections_sum BIGINT NOT NULL DEFAULT 0,
|
||||
active_connections_max BIGINT NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_sum BIGINT NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_max BIGINT NOT NULL DEFAULT 0,
|
||||
connect_errors_delta BIGINT NOT NULL DEFAULT 0,
|
||||
disconnects_delta BIGINT NOT NULL DEFAULT 0,
|
||||
error_events_delta BIGINT NOT NULL DEFAULT 0,
|
||||
ws_in_bytes_delta BIGINT NOT NULL DEFAULT 0,
|
||||
ws_out_bytes_delta BIGINT NOT NULL DEFAULT 0,
|
||||
ws_in_frames_delta BIGINT NOT NULL DEFAULT 0,
|
||||
ws_out_frames_delta BIGINT NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs),
|
||||
INDEX idx_proxy_node_metrics_1h_bucket_start (bucket_start_unix_secs)
|
||||
);
|
||||
@@ -0,0 +1,50 @@
|
||||
ALTER TABLE public.proxy_node_events
|
||||
ADD COLUMN IF NOT EXISTS event_metadata json;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.proxy_node_metrics_1m (
|
||||
node_id character varying(36) NOT NULL,
|
||||
bucket_start_unix_secs bigint NOT NULL,
|
||||
samples bigint DEFAULT 0 NOT NULL,
|
||||
uptime_samples bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_sum bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_max bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_sum bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_max bigint DEFAULT 0 NOT NULL,
|
||||
connect_errors_delta bigint DEFAULT 0 NOT NULL,
|
||||
disconnects_delta bigint DEFAULT 0 NOT NULL,
|
||||
error_events_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs),
|
||||
CONSTRAINT proxy_node_metrics_1m_node_id_fkey
|
||||
FOREIGN KEY (node_id) REFERENCES public.proxy_nodes(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.proxy_node_metrics_1h (
|
||||
node_id character varying(36) NOT NULL,
|
||||
bucket_start_unix_secs bigint NOT NULL,
|
||||
samples bigint DEFAULT 0 NOT NULL,
|
||||
uptime_samples bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_sum bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_max bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_sum bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_max bigint DEFAULT 0 NOT NULL,
|
||||
connect_errors_delta bigint DEFAULT 0 NOT NULL,
|
||||
disconnects_delta bigint DEFAULT 0 NOT NULL,
|
||||
error_events_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs),
|
||||
CONSTRAINT proxy_node_metrics_1h_node_id_fkey
|
||||
FOREIGN KEY (node_id) REFERENCES public.proxy_nodes(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1m_bucket_start
|
||||
ON public.proxy_node_metrics_1m (bucket_start_unix_secs);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1h_bucket_start
|
||||
ON public.proxy_node_metrics_1h (bucket_start_unix_secs);
|
||||
@@ -0,0 +1,46 @@
|
||||
ALTER TABLE proxy_node_events
|
||||
ADD COLUMN event_metadata TEXT;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS proxy_node_metrics_1m (
|
||||
node_id TEXT NOT NULL,
|
||||
bucket_start_unix_secs INTEGER NOT NULL,
|
||||
samples INTEGER NOT NULL DEFAULT 0,
|
||||
uptime_samples INTEGER NOT NULL DEFAULT 0,
|
||||
active_connections_sum INTEGER NOT NULL DEFAULT 0,
|
||||
active_connections_max INTEGER NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_sum INTEGER NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_max INTEGER NOT NULL DEFAULT 0,
|
||||
connect_errors_delta INTEGER NOT NULL DEFAULT 0,
|
||||
disconnects_delta INTEGER NOT NULL DEFAULT 0,
|
||||
error_events_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_in_bytes_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_out_bytes_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_in_frames_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_out_frames_delta INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS proxy_node_metrics_1h (
|
||||
node_id TEXT NOT NULL,
|
||||
bucket_start_unix_secs INTEGER NOT NULL,
|
||||
samples INTEGER NOT NULL DEFAULT 0,
|
||||
uptime_samples INTEGER NOT NULL DEFAULT 0,
|
||||
active_connections_sum INTEGER NOT NULL DEFAULT 0,
|
||||
active_connections_max INTEGER NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_sum INTEGER NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_max INTEGER NOT NULL DEFAULT 0,
|
||||
connect_errors_delta INTEGER NOT NULL DEFAULT 0,
|
||||
disconnects_delta INTEGER NOT NULL DEFAULT 0,
|
||||
error_events_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_in_bytes_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_out_bytes_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_in_frames_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_out_frames_delta INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1m_bucket_start
|
||||
ON proxy_node_metrics_1m (bucket_start_unix_secs);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1h_bucket_start
|
||||
ON proxy_node_metrics_1h (bucket_start_unix_secs);
|
||||
@@ -0,0 +1,50 @@
|
||||
ALTER TABLE public.proxy_node_events
|
||||
ADD COLUMN IF NOT EXISTS event_metadata json;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.proxy_node_metrics_1m (
|
||||
node_id character varying(36) NOT NULL,
|
||||
bucket_start_unix_secs bigint NOT NULL,
|
||||
samples bigint DEFAULT 0 NOT NULL,
|
||||
uptime_samples bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_sum bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_max bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_sum bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_max bigint DEFAULT 0 NOT NULL,
|
||||
connect_errors_delta bigint DEFAULT 0 NOT NULL,
|
||||
disconnects_delta bigint DEFAULT 0 NOT NULL,
|
||||
error_events_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs),
|
||||
CONSTRAINT proxy_node_metrics_1m_node_id_fkey
|
||||
FOREIGN KEY (node_id) REFERENCES public.proxy_nodes(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.proxy_node_metrics_1h (
|
||||
node_id character varying(36) NOT NULL,
|
||||
bucket_start_unix_secs bigint NOT NULL,
|
||||
samples bigint DEFAULT 0 NOT NULL,
|
||||
uptime_samples bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_sum bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_max bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_sum bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_max bigint DEFAULT 0 NOT NULL,
|
||||
connect_errors_delta bigint DEFAULT 0 NOT NULL,
|
||||
disconnects_delta bigint DEFAULT 0 NOT NULL,
|
||||
error_events_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs),
|
||||
CONSTRAINT proxy_node_metrics_1h_node_id_fkey
|
||||
FOREIGN KEY (node_id) REFERENCES public.proxy_nodes(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1m_bucket_start
|
||||
ON public.proxy_node_metrics_1m (bucket_start_unix_secs);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1h_bucket_start
|
||||
ON public.proxy_node_metrics_1h (bucket_start_unix_secs);
|
||||
@@ -8,3 +8,4 @@
|
||||
110_redeem_codes.sql
|
||||
120_stats_rollups.sql
|
||||
130_stats_cost_savings.sql
|
||||
140_proxy_node_metrics.sql
|
||||
|
||||
@@ -39,7 +39,48 @@ CREATE TABLE IF NOT EXISTS proxy_node_events (
|
||||
`node_id` VARCHAR(64) NOT NULL,
|
||||
`event_type` VARCHAR(64) NOT NULL,
|
||||
`detail` VARCHAR(500),
|
||||
`event_metadata` JSON,
|
||||
`created_at` BIGINT NOT NULL,
|
||||
PRIMARY KEY (`id`)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS proxy_node_metrics_1m (
|
||||
`node_id` VARCHAR(64) NOT NULL,
|
||||
`bucket_start_unix_secs` BIGINT NOT NULL,
|
||||
`samples` BIGINT NOT NULL DEFAULT 0,
|
||||
`uptime_samples` BIGINT NOT NULL DEFAULT 0,
|
||||
`active_connections_sum` BIGINT NOT NULL DEFAULT 0,
|
||||
`active_connections_max` BIGINT NOT NULL DEFAULT 0,
|
||||
`heartbeat_rtt_ms_sum` BIGINT NOT NULL DEFAULT 0,
|
||||
`heartbeat_rtt_ms_max` BIGINT NOT NULL DEFAULT 0,
|
||||
`connect_errors_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`disconnects_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`error_events_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`ws_in_bytes_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`ws_out_bytes_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`ws_in_frames_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`ws_out_frames_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (`node_id`, `bucket_start_unix_secs`),
|
||||
KEY idx_proxy_node_metrics_1m_bucket_start (`bucket_start_unix_secs`)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS proxy_node_metrics_1h (
|
||||
`node_id` VARCHAR(64) NOT NULL,
|
||||
`bucket_start_unix_secs` BIGINT NOT NULL,
|
||||
`samples` BIGINT NOT NULL DEFAULT 0,
|
||||
`uptime_samples` BIGINT NOT NULL DEFAULT 0,
|
||||
`active_connections_sum` BIGINT NOT NULL DEFAULT 0,
|
||||
`active_connections_max` BIGINT NOT NULL DEFAULT 0,
|
||||
`heartbeat_rtt_ms_sum` BIGINT NOT NULL DEFAULT 0,
|
||||
`heartbeat_rtt_ms_max` BIGINT NOT NULL DEFAULT 0,
|
||||
`connect_errors_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`disconnects_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`error_events_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`ws_in_bytes_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`ws_out_bytes_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`ws_in_frames_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
`ws_out_frames_delta` BIGINT NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (`node_id`, `bucket_start_unix_secs`),
|
||||
KEY idx_proxy_node_metrics_1h_bucket_start (`bucket_start_unix_secs`)
|
||||
);
|
||||
|
||||
|
||||
@@ -40,8 +40,51 @@ CREATE TABLE IF NOT EXISTS public.proxy_node_events (
|
||||
node_id character varying(64) NOT NULL,
|
||||
event_type character varying(64) NOT NULL,
|
||||
detail character varying(500),
|
||||
event_metadata jsonb,
|
||||
created_at bigint NOT NULL
|
||||
);
|
||||
|
||||
ALTER TABLE ONLY public.proxy_node_events ADD CONSTRAINT proxy_node_events_pkey PRIMARY KEY (id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.proxy_node_metrics_1m (
|
||||
node_id character varying(64) NOT NULL,
|
||||
bucket_start_unix_secs bigint NOT NULL,
|
||||
samples bigint DEFAULT 0 NOT NULL,
|
||||
uptime_samples bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_sum bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_max bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_sum bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_max bigint DEFAULT 0 NOT NULL,
|
||||
connect_errors_delta bigint DEFAULT 0 NOT NULL,
|
||||
disconnects_delta bigint DEFAULT 0 NOT NULL,
|
||||
error_events_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_frames_delta bigint DEFAULT 0 NOT NULL
|
||||
);
|
||||
|
||||
ALTER TABLE ONLY public.proxy_node_metrics_1m ADD CONSTRAINT proxy_node_metrics_1m_pkey PRIMARY KEY (node_id, bucket_start_unix_secs);
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1m_bucket_start ON public.proxy_node_metrics_1m USING btree (bucket_start_unix_secs);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.proxy_node_metrics_1h (
|
||||
node_id character varying(64) NOT NULL,
|
||||
bucket_start_unix_secs bigint NOT NULL,
|
||||
samples bigint DEFAULT 0 NOT NULL,
|
||||
uptime_samples bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_sum bigint DEFAULT 0 NOT NULL,
|
||||
active_connections_max bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_sum bigint DEFAULT 0 NOT NULL,
|
||||
heartbeat_rtt_ms_max bigint DEFAULT 0 NOT NULL,
|
||||
connect_errors_delta bigint DEFAULT 0 NOT NULL,
|
||||
disconnects_delta bigint DEFAULT 0 NOT NULL,
|
||||
error_events_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_bytes_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_in_frames_delta bigint DEFAULT 0 NOT NULL,
|
||||
ws_out_frames_delta bigint DEFAULT 0 NOT NULL
|
||||
);
|
||||
|
||||
ALTER TABLE ONLY public.proxy_node_metrics_1h ADD CONSTRAINT proxy_node_metrics_1h_pkey PRIMARY KEY (node_id, bucket_start_unix_secs);
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1h_bucket_start ON public.proxy_node_metrics_1h USING btree (bucket_start_unix_secs);
|
||||
|
||||
|
||||
@@ -38,6 +38,47 @@ CREATE TABLE IF NOT EXISTS proxy_node_events (
|
||||
node_id TEXT NOT NULL,
|
||||
event_type TEXT NOT NULL,
|
||||
detail TEXT,
|
||||
event_metadata TEXT,
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS proxy_node_metrics_1m (
|
||||
node_id TEXT NOT NULL,
|
||||
bucket_start_unix_secs INTEGER NOT NULL,
|
||||
samples INTEGER NOT NULL DEFAULT 0,
|
||||
uptime_samples INTEGER NOT NULL DEFAULT 0,
|
||||
active_connections_sum INTEGER NOT NULL DEFAULT 0,
|
||||
active_connections_max INTEGER NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_sum INTEGER NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_max INTEGER NOT NULL DEFAULT 0,
|
||||
connect_errors_delta INTEGER NOT NULL DEFAULT 0,
|
||||
disconnects_delta INTEGER NOT NULL DEFAULT 0,
|
||||
error_events_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_in_bytes_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_out_bytes_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_in_frames_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_out_frames_delta INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1m_bucket_start ON proxy_node_metrics_1m (bucket_start_unix_secs);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS proxy_node_metrics_1h (
|
||||
node_id TEXT NOT NULL,
|
||||
bucket_start_unix_secs INTEGER NOT NULL,
|
||||
samples INTEGER NOT NULL DEFAULT 0,
|
||||
uptime_samples INTEGER NOT NULL DEFAULT 0,
|
||||
active_connections_sum INTEGER NOT NULL DEFAULT 0,
|
||||
active_connections_max INTEGER NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_sum INTEGER NOT NULL DEFAULT 0,
|
||||
heartbeat_rtt_ms_max INTEGER NOT NULL DEFAULT 0,
|
||||
connect_errors_delta INTEGER NOT NULL DEFAULT 0,
|
||||
disconnects_delta INTEGER NOT NULL DEFAULT 0,
|
||||
error_events_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_in_bytes_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_out_bytes_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_in_frames_delta INTEGER NOT NULL DEFAULT 0,
|
||||
ws_out_frames_delta INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (node_id, bucket_start_unix_secs)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1h_bucket_start ON proxy_node_metrics_1h (bucket_start_unix_secs);
|
||||
|
||||
|
||||
@@ -177,6 +177,177 @@ type = "text"
|
||||
length = 500
|
||||
nullable = true
|
||||
|
||||
[[table.proxy_node_events.columns]]
|
||||
name = "event_metadata"
|
||||
type = "json"
|
||||
nullable = true
|
||||
|
||||
[[table.proxy_node_events.columns]]
|
||||
name = "created_at"
|
||||
type = "unix_seconds"
|
||||
|
||||
[table.proxy_node_metrics_1m]
|
||||
domain = "proxy_nodes"
|
||||
order = 30
|
||||
primary_key = ["node_id", "bucket_start_unix_secs"]
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "node_id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "bucket_start_unix_secs"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "samples"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "uptime_samples"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "active_connections_sum"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "active_connections_max"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "heartbeat_rtt_ms_sum"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "heartbeat_rtt_ms_max"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "connect_errors_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "disconnects_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "error_events_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "ws_in_bytes_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "ws_out_bytes_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "ws_in_frames_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.columns]]
|
||||
name = "ws_out_frames_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1m.indexes]]
|
||||
name = "idx_proxy_node_metrics_1m_bucket_start"
|
||||
columns = ["bucket_start_unix_secs"]
|
||||
|
||||
[table.proxy_node_metrics_1h]
|
||||
domain = "proxy_nodes"
|
||||
order = 40
|
||||
primary_key = ["node_id", "bucket_start_unix_secs"]
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "node_id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "bucket_start_unix_secs"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "samples"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "uptime_samples"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "active_connections_sum"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "active_connections_max"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "heartbeat_rtt_ms_sum"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "heartbeat_rtt_ms_max"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "connect_errors_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "disconnects_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "error_events_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "ws_in_bytes_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "ws_out_bytes_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "ws_in_frames_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.columns]]
|
||||
name = "ws_out_frames_delta"
|
||||
type = "int64"
|
||||
default = 0
|
||||
|
||||
[[table.proxy_node_metrics_1h.indexes]]
|
||||
name = "idx_proxy_node_metrics_1h_bucket_start"
|
||||
columns = ["bucket_start_unix_secs"]
|
||||
|
||||
@@ -3,7 +3,7 @@ use sqlx::Row;
|
||||
|
||||
use crate::backend::stats_common::{stats_id, unix_ms, unix_secs, utc_from_unix_secs};
|
||||
use crate::backend::SqliteBackend;
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::{
|
||||
DataLayerError, StatsDailyAggregationInput, StatsDailyAggregationSummary,
|
||||
@@ -124,9 +124,9 @@ SELECT
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(cache_creation_input_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_input_tokens), 0) AS cache_read_tokens,
|
||||
COALESCE(SUM(total_cost_usd), 0.0) AS total_cost,
|
||||
COALESCE(SUM(actual_total_cost_usd), 0.0) AS actual_total_cost,
|
||||
COALESCE(AVG(response_time_ms), 0.0) AS avg_response_time_ms
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL) AS total_cost,
|
||||
CAST(COALESCE(SUM(actual_total_cost_usd), 0) AS REAL) AS actual_total_cost,
|
||||
CAST(COALESCE(AVG(response_time_ms), 0) AS REAL) AS avg_response_time_ms
|
||||
FROM "usage"
|
||||
WHERE created_at_unix_ms >= ?
|
||||
AND created_at_unix_ms < ?
|
||||
@@ -188,12 +188,9 @@ ON CONFLICT (hour_utc) DO UPDATE SET
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("total_cost").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("actual_total_cost").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("avg_response_time_ms")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(sqlite_real(&row, "total_cost")?)
|
||||
.bind(sqlite_real(&row, "actual_total_cost")?)
|
||||
.bind(sqlite_real(&row, "avg_response_time_ms")?)
|
||||
.bind(aggregated_at_unix_secs)
|
||||
.bind(aggregated_at_unix_secs)
|
||||
.bind(aggregated_at_unix_secs)
|
||||
@@ -277,12 +274,9 @@ ON CONFLICT ("date") DO UPDATE SET
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("total_cost").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("actual_total_cost").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("avg_response_time_ms")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(sqlite_real(&row, "total_cost")?)
|
||||
.bind(sqlite_real(&row, "actual_total_cost")?)
|
||||
.bind(sqlite_real(&row, "avg_response_time_ms")?)
|
||||
.bind(unique_models)
|
||||
.bind(unique_providers)
|
||||
.bind(aggregated_at_unix_secs)
|
||||
|
||||
@@ -2,6 +2,7 @@ use sha2::{Digest, Sha256};
|
||||
use sqlx::Row;
|
||||
|
||||
use crate::backend::{MysqlBackend, PostgresBackend, SqliteBackend};
|
||||
use crate::driver::sqlite::sqlite_real;
|
||||
use crate::error::{SqlResultExt, SqlxResultExt};
|
||||
use crate::{DataLayerError, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult};
|
||||
|
||||
@@ -119,7 +120,7 @@ const SQLITE_SELECT_WALLET_DAILY_USAGE_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
usage_settlement_snapshots.wallet_id AS wallet_id,
|
||||
COUNT(*) AS total_requests,
|
||||
COALESCE(SUM("usage".total_cost_usd), 0) AS total_cost_usd,
|
||||
CAST(COALESCE(SUM("usage".total_cost_usd), 0) AS REAL) AS total_cost_usd,
|
||||
COALESCE(SUM("usage".input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM("usage".output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM("usage".cache_creation_input_tokens), 0) AS cache_creation_tokens,
|
||||
@@ -400,7 +401,7 @@ INSERT INTO wallet_daily_usage_ledgers (
|
||||
.bind(&wallet_id)
|
||||
.bind(&input.billing_date)
|
||||
.bind(&input.billing_timezone)
|
||||
.bind(row.try_get::<f64, _>("total_cost_usd").map_sql_err()?)
|
||||
.bind(sqlite_real(&row, "total_cost_usd")?)
|
||||
.bind(row.try_get::<i64, _>("total_requests").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("input_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens").map_sql_err()?)
|
||||
|
||||
@@ -1,3 +1,29 @@
|
||||
mod pool;
|
||||
|
||||
pub use pool::{SqlitePool, SqlitePoolConfig, SqlitePoolFactory};
|
||||
|
||||
use crate::DataLayerError;
|
||||
use sqlx::{sqlite::SqliteRow, Row};
|
||||
|
||||
pub(crate) fn sqlite_real(row: &SqliteRow, field: &str) -> Result<f64, DataLayerError> {
|
||||
match row.try_get::<f64, _>(field) {
|
||||
Ok(value) => Ok(value),
|
||||
Err(real_err) => match row.try_get::<i64, _>(field) {
|
||||
Ok(value) => Ok(value as f64),
|
||||
Err(_) => Err(DataLayerError::sql(real_err)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn sqlite_optional_real(
|
||||
row: &SqliteRow,
|
||||
field: &str,
|
||||
) -> Result<Option<f64>, DataLayerError> {
|
||||
match row.try_get::<Option<f64>, _>(field) {
|
||||
Ok(value) => Ok(value),
|
||||
Err(real_err) => match row.try_get::<Option<i64>, _>(field) {
|
||||
Ok(value) => Ok(value.map(|value| value as f64)),
|
||||
Err(_) => Err(DataLayerError::sql(real_err)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ use tracing::info;
|
||||
// Generated by build.rs from schema/bootstrap/postgres.
|
||||
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
|
||||
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
|
||||
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260507120000;
|
||||
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260508000000;
|
||||
|
||||
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
||||
SELECT COUNT(*)::BIGINT
|
||||
|
||||
@@ -293,6 +293,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
|
||||
20260505130000,
|
||||
20260507000000,
|
||||
20260507120000,
|
||||
20260508000000,
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -510,8 +511,14 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
||||
.map(|migration| migration.version)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(mysql_versions, vec![20260403000000, 20260507120000]);
|
||||
assert_eq!(sqlite_versions, vec![20260403000000, 20260507120000]);
|
||||
assert_eq!(
|
||||
mysql_versions,
|
||||
vec![20260403000000, 20260507120000, 20260508000000]
|
||||
);
|
||||
assert_eq!(
|
||||
sqlite_versions,
|
||||
vec![20260403000000, 20260507120000, 20260508000000]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1014,6 +1021,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
||||
20260505130000,
|
||||
20260507000000,
|
||||
20260507120000,
|
||||
20260508000000,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ use super::types::{
|
||||
StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -57,7 +57,7 @@ SELECT
|
||||
api_keys.auto_delete_on_expiry,
|
||||
api_keys.total_requests,
|
||||
COALESCE(api_keys.total_tokens, 0) AS total_tokens,
|
||||
COALESCE(api_keys.total_cost_usd, 0) AS total_cost_usd,
|
||||
CAST(COALESCE(api_keys.total_cost_usd, 0) AS REAL) AS total_cost_usd,
|
||||
api_keys.last_used_at AS last_used_at_unix_secs,
|
||||
api_keys.created_at AS created_at_unix_secs,
|
||||
api_keys.updated_at AS updated_at_unix_secs,
|
||||
@@ -904,7 +904,7 @@ fn map_auth_api_key_export_row(
|
||||
row.try_get("auto_delete_on_expiry").map_sql_err()?,
|
||||
row.try_get("total_requests").map_sql_err()?,
|
||||
row.try_get("total_tokens").map_sql_err()?,
|
||||
row.try_get("total_cost_usd").map_sql_err()?,
|
||||
sqlite_real(row, "total_cost_usd")?,
|
||||
row.try_get("is_standalone").map_sql_err()?,
|
||||
)
|
||||
.and_then(|record| {
|
||||
|
||||
@@ -6,7 +6,7 @@ use super::{
|
||||
AdminBillingPresetApplyResult, AdminBillingRuleRecord, AdminBillingRuleWriteInput,
|
||||
BillingReadRepository, StoredBillingModelContext,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_optional_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -20,12 +20,12 @@ SELECT
|
||||
gm.id AS global_model_id,
|
||||
gm.name AS global_model_name,
|
||||
gm.config AS global_model_config,
|
||||
gm.default_price_per_request AS default_price_per_request,
|
||||
CAST(gm.default_price_per_request AS REAL) AS default_price_per_request,
|
||||
gm.default_tiered_pricing AS default_tiered_pricing,
|
||||
m.id AS model_id,
|
||||
m.provider_model_name AS model_provider_model_name,
|
||||
m.config AS model_config,
|
||||
m.price_per_request AS model_price_per_request,
|
||||
CAST(m.price_per_request AS REAL) AS model_price_per_request,
|
||||
m.tiered_pricing AS model_tiered_pricing,
|
||||
m.provider_model_mappings AS provider_model_mappings,
|
||||
m.is_available AS model_is_available,
|
||||
@@ -641,19 +641,13 @@ fn match_rank(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let has_model_price = row
|
||||
.try_get::<Option<f64>, _>("model_price_per_request")
|
||||
.map_sql_err()?
|
||||
.is_some()
|
||||
let has_model_price = sqlite_optional_real(row, "model_price_per_request")?.is_some()
|
||||
|| row
|
||||
.try_get::<Option<String>, _>("model_tiered_pricing")
|
||||
.ok()
|
||||
.flatten()
|
||||
.is_some();
|
||||
let has_default_price = row
|
||||
.try_get::<Option<f64>, _>("default_price_per_request")
|
||||
.map_sql_err()?
|
||||
.is_some()
|
||||
let has_default_price = sqlite_optional_real(row, "default_price_per_request")?.is_some()
|
||||
|| row
|
||||
.try_get::<Option<String>, _>("default_tiered_pricing")
|
||||
.ok()
|
||||
@@ -717,12 +711,12 @@ fn map_row(row: &SqliteRow) -> Result<StoredBillingModelContext, DataLayerError>
|
||||
row.try_get("global_model_id").map_sql_err()?,
|
||||
row.try_get("global_model_name").map_sql_err()?,
|
||||
parse_json(row.try_get("global_model_config").ok().flatten())?,
|
||||
row.try_get("default_price_per_request").map_sql_err()?,
|
||||
sqlite_optional_real(row, "default_price_per_request")?,
|
||||
parse_json(row.try_get("default_tiered_pricing").ok().flatten())?,
|
||||
row.try_get("model_id").map_sql_err()?,
|
||||
row.try_get("model_provider_model_name").map_sql_err()?,
|
||||
parse_json(row.try_get("model_config").ok().flatten())?,
|
||||
row.try_get("model_price_per_request").map_sql_err()?,
|
||||
sqlite_optional_real(row, "model_price_per_request")?,
|
||||
parse_json(row.try_get("model_tiered_pricing").ok().flatten())?,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ use super::{
|
||||
StoredPublicCatalogModel, StoredPublicGlobalModel, StoredPublicGlobalModelPage,
|
||||
UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_optional_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -50,7 +50,7 @@ SELECT
|
||||
name,
|
||||
display_name,
|
||||
is_active,
|
||||
default_price_per_request,
|
||||
CAST(default_price_per_request AS REAL) AS default_price_per_request,
|
||||
default_tiered_pricing,
|
||||
supported_capabilities,
|
||||
config,
|
||||
@@ -74,7 +74,7 @@ SELECT
|
||||
name,
|
||||
COALESCE(NULLIF(display_name, ''), name) AS display_name,
|
||||
is_active,
|
||||
default_price_per_request,
|
||||
CAST(default_price_per_request AS REAL) AS default_price_per_request,
|
||||
default_tiered_pricing,
|
||||
supported_capabilities,
|
||||
config,
|
||||
@@ -101,7 +101,7 @@ SELECT
|
||||
m.global_model_id,
|
||||
m.provider_model_name,
|
||||
m.provider_model_mappings,
|
||||
m.price_per_request,
|
||||
CAST(m.price_per_request AS REAL) AS price_per_request,
|
||||
m.tiered_pricing,
|
||||
m.supports_vision,
|
||||
m.supports_function_calling,
|
||||
@@ -115,7 +115,7 @@ SELECT
|
||||
m.updated_at AS updated_at_unix_secs,
|
||||
gm.name AS global_model_name,
|
||||
gm.display_name AS global_model_display_name,
|
||||
gm.default_price_per_request AS global_model_default_price_per_request,
|
||||
CAST(gm.default_price_per_request AS REAL) AS global_model_default_price_per_request,
|
||||
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
|
||||
gm.supported_capabilities AS global_model_supported_capabilities,
|
||||
gm.config AS global_model_config
|
||||
@@ -711,7 +711,7 @@ fn map_public_global_model_row(row: &SqliteRow) -> Result<StoredPublicGlobalMode
|
||||
row.try_get("name").map_sql_err()?,
|
||||
row.try_get("display_name").map_sql_err()?,
|
||||
row.try_get("is_active").map_sql_err()?,
|
||||
row.try_get("default_price_per_request").map_sql_err()?,
|
||||
sqlite_optional_real(row, "default_price_per_request")?,
|
||||
optional_json_from_string(
|
||||
row.try_get("default_tiered_pricing").map_sql_err()?,
|
||||
"global_models.default_tiered_pricing",
|
||||
@@ -731,7 +731,7 @@ fn map_admin_global_model_row(row: &SqliteRow) -> Result<StoredAdminGlobalModel,
|
||||
row.try_get("name").map_sql_err()?,
|
||||
row.try_get("display_name").map_sql_err()?,
|
||||
row.try_get("is_active").map_sql_err()?,
|
||||
row.try_get("default_price_per_request").map_sql_err()?,
|
||||
sqlite_optional_real(row, "default_price_per_request")?,
|
||||
optional_json_from_string(
|
||||
row.try_get("default_tiered_pricing").map_sql_err()?,
|
||||
"global_models.default_tiered_pricing",
|
||||
@@ -767,7 +767,7 @@ fn map_admin_provider_model_row(
|
||||
row.try_get("provider_model_mappings").map_sql_err()?,
|
||||
"models.provider_model_mappings",
|
||||
)?,
|
||||
row.try_get("price_per_request").map_sql_err()?,
|
||||
sqlite_optional_real(row, "price_per_request")?,
|
||||
optional_json_from_string(
|
||||
row.try_get("tiered_pricing").map_sql_err()?,
|
||||
"models.tiered_pricing",
|
||||
@@ -790,8 +790,7 @@ fn map_admin_provider_model_row(
|
||||
)?,
|
||||
row.try_get("global_model_name").map_sql_err()?,
|
||||
row.try_get("global_model_display_name").map_sql_err()?,
|
||||
row.try_get("global_model_default_price_per_request")
|
||||
.map_sql_err()?,
|
||||
sqlite_optional_real(row, "global_model_default_price_per_request")?,
|
||||
optional_json_from_string(
|
||||
row.try_get("global_model_default_tiered_pricing")
|
||||
.map_sql_err()?,
|
||||
|
||||
@@ -7,7 +7,7 @@ use super::{
|
||||
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_optional_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -34,7 +34,9 @@ impl SqliteProviderCatalogReadRepository {
|
||||
r#"
|
||||
SELECT
|
||||
id, name, description, website, provider_type, billing_type,
|
||||
monthly_quota_usd, monthly_used_usd, quota_reset_day,
|
||||
CAST(monthly_quota_usd AS REAL) AS monthly_quota_usd,
|
||||
CAST(monthly_used_usd AS REAL) AS monthly_used_usd,
|
||||
quota_reset_day,
|
||||
quota_last_reset_at AS quota_last_reset_at_unix_secs,
|
||||
quota_expires_at AS quota_expires_at_unix_secs,
|
||||
provider_priority, is_active, keep_priority_on_conversion,
|
||||
@@ -1258,8 +1260,8 @@ fn map_provider_row(row: &SqliteRow) -> Result<StoredProviderCatalogProvider, Da
|
||||
.with_description(row.try_get("description").map_sql_err()?)
|
||||
.with_billing_fields(
|
||||
row.try_get("billing_type").map_sql_err()?,
|
||||
row.try_get("monthly_quota_usd").map_sql_err()?,
|
||||
row.try_get("monthly_used_usd").map_sql_err()?,
|
||||
sqlite_optional_real(row, "monthly_quota_usd")?,
|
||||
sqlite_optional_real(row, "monthly_used_usd")?,
|
||||
optional_u64(
|
||||
row.try_get("quota_reset_day").map_sql_err()?,
|
||||
"providers.quota_reset_day",
|
||||
@@ -1316,11 +1318,7 @@ fn map_endpoint_row(row: &SqliteRow) -> Result<StoredProviderCatalogEndpoint, Da
|
||||
"provider_endpoints.updated_at",
|
||||
)?,
|
||||
)
|
||||
.with_health_score(
|
||||
row.try_get::<Option<f64>, _>("health_score")
|
||||
.map_sql_err()?
|
||||
.unwrap_or(1.0),
|
||||
)
|
||||
.with_health_score(sqlite_optional_real(row, "health_score")?.unwrap_or(1.0))
|
||||
.with_transport_fields(
|
||||
row.try_get("base_url").map_sql_err()?,
|
||||
optional_json_from_string(
|
||||
@@ -1349,10 +1347,7 @@ fn map_endpoint_row(row: &SqliteRow) -> Result<StoredProviderCatalogEndpoint, Da
|
||||
}
|
||||
|
||||
fn map_key_row(row: &SqliteRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
|
||||
let total_cost_usd = row
|
||||
.try_get::<Option<f64>, _>("total_cost_usd")
|
||||
.map_sql_err()?
|
||||
.unwrap_or(0.0);
|
||||
let total_cost_usd = sqlite_optional_real(row, "total_cost_usd")?.unwrap_or(0.0);
|
||||
if !total_cost_usd.is_finite() {
|
||||
return Err(DataLayerError::UnexpectedValue(
|
||||
"invalid provider_api_keys.total_cost_usd".to_string(),
|
||||
@@ -1571,6 +1566,8 @@ mod tests {
|
||||
.expect("providers should list");
|
||||
assert_eq!(providers.len(), 1);
|
||||
assert_eq!(providers[0].provider_priority, 10);
|
||||
assert_eq!(providers[0].monthly_quota_usd, Some(0.0));
|
||||
assert_eq!(providers[0].monthly_used_usd, Some(0.0));
|
||||
|
||||
let endpoints = repository
|
||||
.list_endpoints_by_provider_ids(&["provider-1".to_string()])
|
||||
@@ -1827,11 +1824,12 @@ mod tests {
|
||||
r#"
|
||||
INSERT INTO providers (
|
||||
id, name, description, website, provider_type, provider_priority,
|
||||
monthly_quota_usd, monthly_used_usd,
|
||||
is_active, keep_priority_on_conversion, enable_format_conversion,
|
||||
config, created_at, updated_at
|
||||
) VALUES (
|
||||
'provider-1', 'Provider One', 'test provider', 'https://example.com',
|
||||
'custom', 10, 1, 1, 1, '{"region":"us"}', 1, 2
|
||||
'custom', 10, 0, 0, 1, 1, 1, '{"region":"us"}', 1, 2
|
||||
)
|
||||
"#,
|
||||
)
|
||||
|
||||
@@ -7,10 +7,14 @@ use serde_json::{json, Map, Value};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::types::{
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery,
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
|
||||
ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep, ProxyNodeReadRepository,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket,
|
||||
StoredProxyNode, StoredProxyNodeEvent, StoredProxyNodeMetricsBucket, TunnelMetricsSample,
|
||||
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -18,6 +22,8 @@ use crate::DataLayerError;
|
||||
pub struct InMemoryProxyNodeRepository {
|
||||
nodes: RwLock<BTreeMap<String, StoredProxyNode>>,
|
||||
events: RwLock<Vec<StoredProxyNodeEvent>>,
|
||||
metrics_1m: RwLock<BTreeMap<(String, u64), StoredProxyNodeMetricsBucket>>,
|
||||
metrics_1h: RwLock<BTreeMap<(String, u64), StoredProxyNodeMetricsBucket>>,
|
||||
}
|
||||
|
||||
impl InMemoryProxyNodeRepository {
|
||||
@@ -33,6 +39,8 @@ impl InMemoryProxyNodeRepository {
|
||||
.collect(),
|
||||
),
|
||||
events: RwLock::new(Vec::new()),
|
||||
metrics_1m: RwLock::new(BTreeMap::new()),
|
||||
metrics_1h: RwLock::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -49,6 +57,8 @@ impl InMemoryProxyNodeRepository {
|
||||
.collect(),
|
||||
),
|
||||
events: RwLock::new(events.into_iter().collect()),
|
||||
metrics_1m: RwLock::new(BTreeMap::new()),
|
||||
metrics_1h: RwLock::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,6 +73,50 @@ impl InMemoryProxyNodeRepository {
|
||||
events.iter().map(|event| event.id).max().unwrap_or(0) + 1
|
||||
}
|
||||
|
||||
fn upsert_metrics_bucket(
|
||||
metrics: &mut BTreeMap<(String, u64), StoredProxyNodeMetricsBucket>,
|
||||
node_id: &str,
|
||||
bucket_start_unix_secs: u64,
|
||||
sample: &TunnelMetricsSample,
|
||||
) {
|
||||
let key = (node_id.to_string(), bucket_start_unix_secs);
|
||||
let bucket = metrics
|
||||
.entry(key)
|
||||
.or_insert_with(|| StoredProxyNodeMetricsBucket {
|
||||
node_id: node_id.to_string(),
|
||||
bucket_start_unix_secs,
|
||||
samples: 0,
|
||||
uptime_samples: 0,
|
||||
active_connections_sum: 0,
|
||||
active_connections_max: 0,
|
||||
heartbeat_rtt_ms_sum: 0,
|
||||
heartbeat_rtt_ms_max: 0,
|
||||
connect_errors_delta: 0,
|
||||
disconnects_delta: 0,
|
||||
error_events_delta: 0,
|
||||
ws_in_bytes_delta: 0,
|
||||
ws_out_bytes_delta: 0,
|
||||
ws_in_frames_delta: 0,
|
||||
ws_out_frames_delta: 0,
|
||||
});
|
||||
|
||||
bucket.samples += sample.samples;
|
||||
bucket.uptime_samples += sample.uptime_samples;
|
||||
bucket.active_connections_sum += sample.active_connections_sum;
|
||||
bucket.active_connections_max = bucket
|
||||
.active_connections_max
|
||||
.max(sample.active_connections_max);
|
||||
bucket.heartbeat_rtt_ms_sum += sample.heartbeat_rtt_ms_sum;
|
||||
bucket.heartbeat_rtt_ms_max = bucket.heartbeat_rtt_ms_max.max(sample.heartbeat_rtt_ms_max);
|
||||
bucket.connect_errors_delta += sample.connect_errors_delta;
|
||||
bucket.disconnects_delta += sample.disconnects_delta;
|
||||
bucket.error_events_delta += sample.error_events_delta;
|
||||
bucket.ws_in_bytes_delta += sample.ws_in_bytes_delta;
|
||||
bucket.ws_out_bytes_delta += sample.ws_out_bytes_delta;
|
||||
bucket.ws_in_frames_delta += sample.ws_in_frames_delta;
|
||||
bucket.ws_out_frames_delta += sample.ws_out_frames_delta;
|
||||
}
|
||||
|
||||
fn normalize_remote_config(
|
||||
mutation: &ProxyNodeRemoteConfigMutation,
|
||||
existing: Option<&Value>,
|
||||
@@ -154,6 +208,130 @@ impl ProxyNodeReadRepository for InMemoryProxyNodeRepository {
|
||||
items.truncate(limit);
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn list_proxy_node_events_filtered(
|
||||
&self,
|
||||
node_id: &str,
|
||||
query: &ProxyNodeEventQuery,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
||||
let events = self.events.read().expect("proxy node repository lock");
|
||||
let mut items = events
|
||||
.iter()
|
||||
.filter(|event| event.node_id == node_id)
|
||||
.filter(|event| {
|
||||
query
|
||||
.from_unix_secs
|
||||
.map(|from| event.created_at_unix_ms.unwrap_or(0) >= from)
|
||||
.unwrap_or(true)
|
||||
})
|
||||
.filter(|event| {
|
||||
query
|
||||
.to_unix_secs
|
||||
.map(|to| event.created_at_unix_ms.unwrap_or(u64::MAX) <= to)
|
||||
.unwrap_or(true)
|
||||
})
|
||||
.filter(|event| {
|
||||
query
|
||||
.event_type
|
||||
.as_deref()
|
||||
.map(|event_type| event.event_type.eq_ignore_ascii_case(event_type))
|
||||
.unwrap_or(true)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_ms
|
||||
.unwrap_or(0)
|
||||
.cmp(&left.created_at_unix_ms.unwrap_or(0))
|
||||
.then(right.id.cmp(&left.id))
|
||||
});
|
||||
items.truncate(query.limit);
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
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>, DataLayerError> {
|
||||
let metrics = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => self.metrics_1m.read(),
|
||||
ProxyNodeMetricsStep::OneHour => self.metrics_1h.read(),
|
||||
}
|
||||
.expect("proxy node repository lock");
|
||||
let mut items = metrics
|
||||
.values()
|
||||
.filter(|bucket| bucket.node_id == node_id)
|
||||
.filter(|bucket| bucket.bucket_start_unix_secs >= from_unix_secs)
|
||||
.filter(|bucket| bucket.bucket_start_unix_secs <= to_unix_secs)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
items.sort_by_key(|bucket| bucket.bucket_start_unix_secs);
|
||||
items.truncate(limit);
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn list_proxy_fleet_metrics(
|
||||
&self,
|
||||
step: ProxyNodeMetricsStep,
|
||||
from_unix_secs: u64,
|
||||
to_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyFleetMetricsBucket>, DataLayerError> {
|
||||
let metrics = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => self.metrics_1m.read(),
|
||||
ProxyNodeMetricsStep::OneHour => self.metrics_1h.read(),
|
||||
}
|
||||
.expect("proxy node repository lock");
|
||||
let mut grouped = BTreeMap::<u64, StoredProxyFleetMetricsBucket>::new();
|
||||
for bucket in metrics.values() {
|
||||
if bucket.bucket_start_unix_secs < from_unix_secs
|
||||
|| bucket.bucket_start_unix_secs > to_unix_secs
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let item = grouped
|
||||
.entry(bucket.bucket_start_unix_secs)
|
||||
.or_insert_with(|| StoredProxyFleetMetricsBucket {
|
||||
bucket_start_unix_secs: bucket.bucket_start_unix_secs,
|
||||
samples: 0,
|
||||
uptime_samples: 0,
|
||||
active_connections_sum: 0,
|
||||
active_connections_max: 0,
|
||||
heartbeat_rtt_ms_sum: 0,
|
||||
heartbeat_rtt_ms_max: 0,
|
||||
connect_errors_delta: 0,
|
||||
disconnects_delta: 0,
|
||||
error_events_delta: 0,
|
||||
ws_in_bytes_delta: 0,
|
||||
ws_out_bytes_delta: 0,
|
||||
ws_in_frames_delta: 0,
|
||||
ws_out_frames_delta: 0,
|
||||
});
|
||||
item.samples += bucket.samples;
|
||||
item.uptime_samples += bucket.uptime_samples;
|
||||
item.active_connections_sum += bucket.active_connections_sum;
|
||||
item.active_connections_max = item
|
||||
.active_connections_max
|
||||
.max(bucket.active_connections_max);
|
||||
item.heartbeat_rtt_ms_sum += bucket.heartbeat_rtt_ms_sum;
|
||||
item.heartbeat_rtt_ms_max = item.heartbeat_rtt_ms_max.max(bucket.heartbeat_rtt_ms_max);
|
||||
item.connect_errors_delta += bucket.connect_errors_delta;
|
||||
item.disconnects_delta += bucket.disconnects_delta;
|
||||
item.error_events_delta += bucket.error_events_delta;
|
||||
item.ws_in_bytes_delta += bucket.ws_in_bytes_delta;
|
||||
item.ws_out_bytes_delta += bucket.ws_out_bytes_delta;
|
||||
item.ws_in_frames_delta += bucket.ws_in_frames_delta;
|
||||
item.ws_out_frames_delta += bucket.ws_out_frames_delta;
|
||||
}
|
||||
let mut items = grouped.into_values().collect::<Vec<_>>();
|
||||
items.truncate(limit);
|
||||
Ok(items)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -376,65 +554,114 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
|
||||
&self,
|
||||
mutation: &ProxyNodeHeartbeatMutation,
|
||||
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
||||
let mut nodes = self.nodes.write().expect("proxy node repository lock");
|
||||
let Some(node) = nodes.get_mut(&mutation.node_id) else {
|
||||
return Ok(None);
|
||||
let (node, sample, now_unix_secs) = {
|
||||
let mut nodes = self.nodes.write().expect("proxy node repository lock");
|
||||
let Some(node) = nodes.get_mut(&mutation.node_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !node.tunnel_mode {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let previous_proxy_metadata = node.proxy_metadata.clone();
|
||||
let now_unix_secs = Self::now_unix_secs().unwrap_or(0);
|
||||
let now = Some(now_unix_secs);
|
||||
node.last_heartbeat_at_unix_secs = now;
|
||||
if node.status != "online" || !node.tunnel_connected {
|
||||
node.status = "online".to_string();
|
||||
node.tunnel_connected = true;
|
||||
node.tunnel_connected_at_unix_secs = now;
|
||||
node.updated_at_unix_secs = now;
|
||||
}
|
||||
|
||||
if let Some(value) = mutation.heartbeat_interval {
|
||||
node.heartbeat_interval = value;
|
||||
}
|
||||
if let Some(value) = mutation.active_connections {
|
||||
node.active_connections = value;
|
||||
}
|
||||
if let Some(value) = mutation.avg_latency_ms {
|
||||
node.avg_latency_ms = Some(value);
|
||||
}
|
||||
let normalized_proxy_metadata = normalize_proxy_metadata(
|
||||
mutation.proxy_metadata.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
);
|
||||
if let Some(value) = normalized_proxy_metadata {
|
||||
node.proxy_metadata = Some(value);
|
||||
}
|
||||
if let Some(value) = mutation.total_requests_delta.filter(|value| *value > 0) {
|
||||
node.total_requests += value;
|
||||
}
|
||||
if let Some(value) = mutation.failed_requests_delta.filter(|value| *value > 0) {
|
||||
node.failed_requests += value;
|
||||
}
|
||||
if let Some(value) = mutation.dns_failures_delta.filter(|value| *value > 0) {
|
||||
node.dns_failures += value;
|
||||
}
|
||||
if let Some(value) = mutation.stream_errors_delta.filter(|value| *value > 0) {
|
||||
node.stream_errors += value;
|
||||
}
|
||||
let reconciled_remote_config = reconcile_remote_config_after_heartbeat(
|
||||
node.remote_config.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
);
|
||||
if reconciled_remote_config != node.remote_config {
|
||||
node.remote_config = reconciled_remote_config;
|
||||
node.config_version = node.config_version.saturating_add(1);
|
||||
node.updated_at_unix_secs = now;
|
||||
}
|
||||
|
||||
let sample = build_tunnel_metrics_sample(
|
||||
previous_proxy_metadata.as_ref(),
|
||||
node.proxy_metadata.as_ref(),
|
||||
node.active_connections,
|
||||
node.tunnel_connected,
|
||||
);
|
||||
(node.clone(), sample, now_unix_secs)
|
||||
};
|
||||
if !node.tunnel_mode {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
|
||||
.to_string(),
|
||||
));
|
||||
|
||||
if let Some(sample) = sample.as_ref() {
|
||||
Self::upsert_metrics_bucket(
|
||||
&mut self.metrics_1m.write().expect("proxy node repository lock"),
|
||||
&node.id,
|
||||
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneMinute),
|
||||
sample,
|
||||
);
|
||||
Self::upsert_metrics_bucket(
|
||||
&mut self.metrics_1h.write().expect("proxy node repository lock"),
|
||||
&node.id,
|
||||
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneHour),
|
||||
sample,
|
||||
);
|
||||
|
||||
let mut events = self.events.write().expect("proxy node repository lock");
|
||||
for error in &sample.recent_error_events {
|
||||
let event_id = Self::next_event_id(&events);
|
||||
events.push(StoredProxyNodeEvent {
|
||||
id: event_id,
|
||||
node_id: node.id.clone(),
|
||||
event_type: PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR.to_string(),
|
||||
detail: Some(build_tunnel_error_event_detail(error)),
|
||||
event_metadata: Some(json!({
|
||||
"source": "heartbeat",
|
||||
"category": error.category,
|
||||
"message": error.message,
|
||||
"timestamp_unix_secs": error.timestamp_unix_secs,
|
||||
})),
|
||||
created_at_unix_ms: Some(if error.timestamp_unix_secs == 0 {
|
||||
now_unix_secs
|
||||
} else {
|
||||
error.timestamp_unix_secs
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let now = Self::now_unix_secs();
|
||||
node.last_heartbeat_at_unix_secs = now;
|
||||
if node.status != "online" || !node.tunnel_connected {
|
||||
node.status = "online".to_string();
|
||||
node.tunnel_connected = true;
|
||||
node.tunnel_connected_at_unix_secs = now;
|
||||
node.updated_at_unix_secs = now;
|
||||
}
|
||||
|
||||
if let Some(value) = mutation.heartbeat_interval {
|
||||
node.heartbeat_interval = value;
|
||||
}
|
||||
if let Some(value) = mutation.active_connections {
|
||||
node.active_connections = value;
|
||||
}
|
||||
if let Some(value) = mutation.avg_latency_ms {
|
||||
node.avg_latency_ms = Some(value);
|
||||
}
|
||||
let normalized_proxy_metadata = normalize_proxy_metadata(
|
||||
mutation.proxy_metadata.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
);
|
||||
if let Some(value) = normalized_proxy_metadata {
|
||||
node.proxy_metadata = Some(value);
|
||||
}
|
||||
if let Some(value) = mutation.total_requests_delta.filter(|value| *value > 0) {
|
||||
node.total_requests += value;
|
||||
}
|
||||
if let Some(value) = mutation.failed_requests_delta.filter(|value| *value > 0) {
|
||||
node.failed_requests += value;
|
||||
}
|
||||
if let Some(value) = mutation.dns_failures_delta.filter(|value| *value > 0) {
|
||||
node.dns_failures += value;
|
||||
}
|
||||
if let Some(value) = mutation.stream_errors_delta.filter(|value| *value > 0) {
|
||||
node.stream_errors += value;
|
||||
}
|
||||
let reconciled_remote_config = reconcile_remote_config_after_heartbeat(
|
||||
node.remote_config.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
);
|
||||
if reconciled_remote_config != node.remote_config {
|
||||
node.remote_config = reconciled_remote_config;
|
||||
node.config_version = node.config_version.saturating_add(1);
|
||||
node.updated_at_unix_secs = now;
|
||||
}
|
||||
|
||||
Ok(Some(node.clone()))
|
||||
Ok(Some(node))
|
||||
}
|
||||
|
||||
async fn record_traffic(
|
||||
@@ -490,6 +717,7 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
|
||||
node_id: mutation.node_id.clone(),
|
||||
event_type: event_type.to_string(),
|
||||
detail: Some(format!("[stale_ignored] {event_detail}")),
|
||||
event_metadata: None,
|
||||
created_at_unix_ms: Self::now_unix_secs(),
|
||||
});
|
||||
return Ok(Some(node.clone()));
|
||||
@@ -513,6 +741,7 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
|
||||
node_id: mutation.node_id.clone(),
|
||||
event_type: event_type.to_string(),
|
||||
detail: Some(event_detail),
|
||||
event_metadata: None,
|
||||
created_at_unix_ms: Some(event_time),
|
||||
});
|
||||
Ok(Some(node.clone()))
|
||||
@@ -546,6 +775,14 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
|
||||
.write()
|
||||
.expect("proxy node repository lock")
|
||||
.retain(|event| event.node_id != node_id);
|
||||
self.metrics_1m
|
||||
.write()
|
||||
.expect("proxy node repository lock")
|
||||
.retain(|(metric_node_id, _), _| metric_node_id != node_id);
|
||||
self.metrics_1h
|
||||
.write()
|
||||
.expect("proxy node repository lock")
|
||||
.retain(|(metric_node_id, _), _| metric_node_id != node_id);
|
||||
}
|
||||
Ok(removed)
|
||||
}
|
||||
@@ -598,6 +835,28 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup_proxy_node_metrics(
|
||||
&self,
|
||||
retain_1m_from_unix_secs: u64,
|
||||
retain_1h_from_unix_secs: u64,
|
||||
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
|
||||
let mut metrics_1m = self.metrics_1m.write().expect("proxy node repository lock");
|
||||
let before_1m = metrics_1m.len();
|
||||
metrics_1m.retain(|(_, bucket_start), _| *bucket_start >= retain_1m_from_unix_secs);
|
||||
let deleted_1m_rows = before_1m.saturating_sub(metrics_1m.len());
|
||||
drop(metrics_1m);
|
||||
|
||||
let mut metrics_1h = self.metrics_1h.write().expect("proxy node repository lock");
|
||||
let before_1h = metrics_1h.len();
|
||||
metrics_1h.retain(|(_, bucket_start), _| *bucket_start >= retain_1h_from_unix_secs);
|
||||
let deleted_1h_rows = before_1h.saturating_sub(metrics_1h.len());
|
||||
|
||||
Ok(ProxyNodeMetricsCleanupSummary {
|
||||
deleted_1m_rows,
|
||||
deleted_1h_rows,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -750,6 +1009,7 @@ mod tests {
|
||||
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 {
|
||||
@@ -757,6 +1017,7 @@ mod tests {
|
||||
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),
|
||||
},
|
||||
],
|
||||
|
||||
@@ -9,11 +9,14 @@ pub use mysql::MysqlProxyNodeReadRepository;
|
||||
pub use postgres::SqlxProxyNodeRepository;
|
||||
pub use sqlite::SqliteProxyNodeReadRepository;
|
||||
pub use types::{
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
normalize_proxy_node_scheduling_state, proxy_node_accepts_new_tunnels, proxy_reported_version,
|
||||
reconcile_remote_config_after_heartbeat, remote_config_scheduling_state,
|
||||
remote_config_upgrade_target, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
|
||||
ProxyNodeManualUpdateMutation, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
|
||||
remote_config_upgrade_target, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
|
||||
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
|
||||
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
StoredProxyNodeMetricsBucket, TunnelErrorEventRecord, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
PROXY_NODE_SCHEDULING_STATE_CORDONED, PROXY_NODE_SCHEDULING_STATE_DRAINING,
|
||||
};
|
||||
|
||||
@@ -2,10 +2,14 @@ use async_trait::async_trait;
|
||||
use sqlx::{mysql::MySqlRow, Row};
|
||||
|
||||
use super::types::{
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery,
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
|
||||
ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep, ProxyNodeReadRepository,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket,
|
||||
StoredProxyNode, StoredProxyNodeEvent, StoredProxyNodeMetricsBucket,
|
||||
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
use crate::error::SqlResultExt;
|
||||
@@ -148,17 +152,22 @@ ON DUPLICATE KEY UPDATE
|
||||
node_id: &str,
|
||||
event_type: &str,
|
||||
detail: Option<&str>,
|
||||
event_metadata: Option<&serde_json::Value>,
|
||||
created_at_unix_secs: Option<u64>,
|
||||
) -> Result<(), DataLayerError> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, event_metadata, created_at)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(node_id)
|
||||
.bind(event_type)
|
||||
.bind(detail)
|
||||
.bind(optional_json_to_string(
|
||||
&event_metadata.cloned(),
|
||||
"proxy_node_events.event_metadata",
|
||||
)?)
|
||||
.bind(created_at_unix_secs.unwrap_or_else(current_unix_secs) as i64)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
@@ -166,6 +175,70 @@ VALUES (?, ?, ?, ?)
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn upsert_metrics_bucket(
|
||||
&self,
|
||||
table: &str,
|
||||
node_id: &str,
|
||||
bucket_start: u64,
|
||||
sample: &super::types::TunnelMetricsSample,
|
||||
) -> Result<(), DataLayerError> {
|
||||
sqlx::query(&format!(
|
||||
r#"
|
||||
INSERT INTO {table} (
|
||||
node_id,
|
||||
bucket_start_unix_secs,
|
||||
samples,
|
||||
uptime_samples,
|
||||
active_connections_sum,
|
||||
active_connections_max,
|
||||
heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max,
|
||||
connect_errors_delta,
|
||||
disconnects_delta,
|
||||
error_events_delta,
|
||||
ws_in_bytes_delta,
|
||||
ws_out_bytes_delta,
|
||||
ws_in_frames_delta,
|
||||
ws_out_frames_delta
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON DUPLICATE KEY UPDATE
|
||||
samples = samples + VALUES(samples),
|
||||
uptime_samples = uptime_samples + VALUES(uptime_samples),
|
||||
active_connections_sum = active_connections_sum + VALUES(active_connections_sum),
|
||||
active_connections_max = GREATEST(active_connections_max, VALUES(active_connections_max)),
|
||||
heartbeat_rtt_ms_sum = heartbeat_rtt_ms_sum + VALUES(heartbeat_rtt_ms_sum),
|
||||
heartbeat_rtt_ms_max = GREATEST(heartbeat_rtt_ms_max, VALUES(heartbeat_rtt_ms_max)),
|
||||
connect_errors_delta = connect_errors_delta + VALUES(connect_errors_delta),
|
||||
disconnects_delta = disconnects_delta + VALUES(disconnects_delta),
|
||||
error_events_delta = error_events_delta + VALUES(error_events_delta),
|
||||
ws_in_bytes_delta = ws_in_bytes_delta + VALUES(ws_in_bytes_delta),
|
||||
ws_out_bytes_delta = ws_out_bytes_delta + VALUES(ws_out_bytes_delta),
|
||||
ws_in_frames_delta = ws_in_frames_delta + VALUES(ws_in_frames_delta),
|
||||
ws_out_frames_delta = ws_out_frames_delta + VALUES(ws_out_frames_delta)
|
||||
"#
|
||||
))
|
||||
.bind(node_id)
|
||||
.bind(i64::try_from(bucket_start).unwrap_or(i64::MAX))
|
||||
.bind(sample.samples)
|
||||
.bind(sample.uptime_samples)
|
||||
.bind(sample.active_connections_sum)
|
||||
.bind(sample.active_connections_max)
|
||||
.bind(sample.heartbeat_rtt_ms_sum)
|
||||
.bind(sample.heartbeat_rtt_ms_max)
|
||||
.bind(sample.connect_errors_delta)
|
||||
.bind(sample.disconnects_delta)
|
||||
.bind(sample.error_events_delta)
|
||||
.bind(sample.ws_in_bytes_delta)
|
||||
.bind(sample.ws_out_bytes_delta)
|
||||
.bind(sample.ws_in_frames_delta)
|
||||
.bind(sample.ws_out_frames_delta)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn normalize_remote_config(
|
||||
mutation: &ProxyNodeRemoteConfigMutation,
|
||||
existing: Option<&serde_json::Value>,
|
||||
@@ -298,6 +371,7 @@ SELECT
|
||||
node_id,
|
||||
event_type,
|
||||
detail,
|
||||
event_metadata,
|
||||
created_at AS created_at_unix_ms
|
||||
FROM proxy_node_events
|
||||
WHERE node_id = ?
|
||||
@@ -312,6 +386,152 @@ LIMIT ?
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_proxy_node_event_row).collect()
|
||||
}
|
||||
|
||||
async fn list_proxy_node_events_filtered(
|
||||
&self,
|
||||
node_id: &str,
|
||||
query: &ProxyNodeEventQuery,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
||||
let rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
node_id,
|
||||
event_type,
|
||||
detail,
|
||||
event_metadata,
|
||||
created_at AS created_at_unix_ms
|
||||
FROM proxy_node_events
|
||||
WHERE node_id = ?
|
||||
AND (? IS NULL OR created_at >= ?)
|
||||
AND (? IS NULL OR created_at <= ?)
|
||||
AND (? IS NULL OR LOWER(event_type) = LOWER(?))
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT ?
|
||||
"#,
|
||||
)
|
||||
.bind(node_id)
|
||||
.bind(
|
||||
query
|
||||
.from_unix_secs
|
||||
.map(|v| i64::try_from(v).unwrap_or(i64::MAX)),
|
||||
)
|
||||
.bind(
|
||||
query
|
||||
.from_unix_secs
|
||||
.map(|v| i64::try_from(v).unwrap_or(i64::MAX)),
|
||||
)
|
||||
.bind(
|
||||
query
|
||||
.to_unix_secs
|
||||
.map(|v| i64::try_from(v).unwrap_or(i64::MAX)),
|
||||
)
|
||||
.bind(
|
||||
query
|
||||
.to_unix_secs
|
||||
.map(|v| i64::try_from(v).unwrap_or(i64::MAX)),
|
||||
)
|
||||
.bind(query.event_type.as_deref())
|
||||
.bind(query.event_type.as_deref())
|
||||
.bind(i64::try_from(query.limit).unwrap_or(i64::MAX))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_proxy_node_event_row).collect()
|
||||
}
|
||||
|
||||
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>, DataLayerError> {
|
||||
let table = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => "proxy_node_metrics_1m",
|
||||
ProxyNodeMetricsStep::OneHour => "proxy_node_metrics_1h",
|
||||
};
|
||||
let rows = sqlx::query(&format!(
|
||||
r#"
|
||||
SELECT
|
||||
node_id,
|
||||
bucket_start_unix_secs,
|
||||
samples,
|
||||
uptime_samples,
|
||||
active_connections_sum,
|
||||
active_connections_max,
|
||||
heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max,
|
||||
connect_errors_delta,
|
||||
disconnects_delta,
|
||||
error_events_delta,
|
||||
ws_in_bytes_delta,
|
||||
ws_out_bytes_delta,
|
||||
ws_in_frames_delta,
|
||||
ws_out_frames_delta
|
||||
FROM {table}
|
||||
WHERE node_id = ?
|
||||
AND bucket_start_unix_secs >= ?
|
||||
AND bucket_start_unix_secs <= ?
|
||||
ORDER BY bucket_start_unix_secs ASC
|
||||
LIMIT ?
|
||||
"#
|
||||
))
|
||||
.bind(node_id)
|
||||
.bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_proxy_node_metric_row).collect()
|
||||
}
|
||||
|
||||
async fn list_proxy_fleet_metrics(
|
||||
&self,
|
||||
step: ProxyNodeMetricsStep,
|
||||
from_unix_secs: u64,
|
||||
to_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyFleetMetricsBucket>, DataLayerError> {
|
||||
let table = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => "proxy_node_metrics_1m",
|
||||
ProxyNodeMetricsStep::OneHour => "proxy_node_metrics_1h",
|
||||
};
|
||||
let rows = sqlx::query(&format!(
|
||||
r#"
|
||||
SELECT
|
||||
bucket_start_unix_secs,
|
||||
SUM(samples) AS samples,
|
||||
SUM(uptime_samples) AS uptime_samples,
|
||||
SUM(active_connections_sum) AS active_connections_sum,
|
||||
MAX(active_connections_max) AS active_connections_max,
|
||||
SUM(heartbeat_rtt_ms_sum) AS heartbeat_rtt_ms_sum,
|
||||
MAX(heartbeat_rtt_ms_max) AS heartbeat_rtt_ms_max,
|
||||
SUM(connect_errors_delta) AS connect_errors_delta,
|
||||
SUM(disconnects_delta) AS disconnects_delta,
|
||||
SUM(error_events_delta) AS error_events_delta,
|
||||
SUM(ws_in_bytes_delta) AS ws_in_bytes_delta,
|
||||
SUM(ws_out_bytes_delta) AS ws_out_bytes_delta,
|
||||
SUM(ws_in_frames_delta) AS ws_in_frames_delta,
|
||||
SUM(ws_out_frames_delta) AS ws_out_frames_delta
|
||||
FROM {table}
|
||||
WHERE bucket_start_unix_secs >= ?
|
||||
AND bucket_start_unix_secs <= ?
|
||||
GROUP BY bucket_start_unix_secs
|
||||
ORDER BY bucket_start_unix_secs ASC
|
||||
LIMIT ?
|
||||
"#
|
||||
))
|
||||
.bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_proxy_fleet_metric_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -540,7 +760,9 @@ WHERE is_manual = 0
|
||||
));
|
||||
}
|
||||
|
||||
let now = Some(current_unix_secs());
|
||||
let previous_proxy_metadata = node.proxy_metadata.clone();
|
||||
let now_unix_secs = current_unix_secs();
|
||||
let now = Some(now_unix_secs);
|
||||
node.last_heartbeat_at_unix_secs = now;
|
||||
if node.status != "online" || !node.tunnel_connected {
|
||||
node.status = "online".to_string();
|
||||
@@ -584,7 +806,52 @@ WHERE is_manual = 0
|
||||
node.config_version = node.config_version.saturating_add(1);
|
||||
node.updated_at_unix_secs = now;
|
||||
}
|
||||
let tunnel_metrics_sample = build_tunnel_metrics_sample(
|
||||
previous_proxy_metadata.as_ref(),
|
||||
node.proxy_metadata.as_ref(),
|
||||
node.active_connections,
|
||||
node.tunnel_connected,
|
||||
);
|
||||
self.upsert_node(&node).await?;
|
||||
|
||||
if let Some(sample) = tunnel_metrics_sample.as_ref() {
|
||||
self.upsert_metrics_bucket(
|
||||
"proxy_node_metrics_1m",
|
||||
&node.id,
|
||||
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneMinute),
|
||||
sample,
|
||||
)
|
||||
.await?;
|
||||
self.upsert_metrics_bucket(
|
||||
"proxy_node_metrics_1h",
|
||||
&node.id,
|
||||
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneHour),
|
||||
sample,
|
||||
)
|
||||
.await?;
|
||||
|
||||
for error in &sample.recent_error_events {
|
||||
let detail = build_tunnel_error_event_detail(error);
|
||||
let event_metadata = serde_json::json!({
|
||||
"source": "heartbeat",
|
||||
"category": error.category,
|
||||
"message": error.message,
|
||||
"timestamp_unix_secs": error.timestamp_unix_secs,
|
||||
});
|
||||
self.insert_event(
|
||||
&node.id,
|
||||
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
Some(detail.as_str()),
|
||||
Some(&event_metadata),
|
||||
Some(if error.timestamp_unix_secs == 0 {
|
||||
now_unix_secs
|
||||
} else {
|
||||
error.timestamp_unix_secs
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
Ok(Some(node))
|
||||
}
|
||||
|
||||
@@ -638,6 +905,7 @@ WHERE is_manual = 0
|
||||
&mutation.node_id,
|
||||
event_type,
|
||||
Some(&format!("[stale_ignored] {event_detail}")),
|
||||
None,
|
||||
Some(current_unix_secs()),
|
||||
)
|
||||
.await?;
|
||||
@@ -660,6 +928,7 @@ WHERE is_manual = 0
|
||||
&mutation.node_id,
|
||||
event_type,
|
||||
Some(&event_detail),
|
||||
None,
|
||||
Some(event_time),
|
||||
)
|
||||
.await?;
|
||||
@@ -691,6 +960,16 @@ WHERE is_manual = 0
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1m WHERE node_id = ?")
|
||||
.bind(node_id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1h WHERE node_id = ?")
|
||||
.bind(node_id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM proxy_nodes WHERE id = ?")
|
||||
.bind(node_id)
|
||||
.execute(&self.pool)
|
||||
@@ -747,6 +1026,33 @@ WHERE is_manual = 0
|
||||
node.updated_at_unix_secs = Some(current_unix_secs());
|
||||
self.upsert_node(&node).await
|
||||
}
|
||||
|
||||
async fn cleanup_proxy_node_metrics(
|
||||
&self,
|
||||
retain_1m_from_unix_secs: u64,
|
||||
retain_1h_from_unix_secs: u64,
|
||||
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
|
||||
let deleted_1m =
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1m WHERE bucket_start_unix_secs < ?")
|
||||
.bind(i64::try_from(retain_1m_from_unix_secs).unwrap_or(i64::MAX))
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected() as usize;
|
||||
|
||||
let deleted_1h =
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1h WHERE bucket_start_unix_secs < ?")
|
||||
.bind(i64::try_from(retain_1h_from_unix_secs).unwrap_or(i64::MAX))
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected() as usize;
|
||||
|
||||
Ok(ProxyNodeMetricsCleanupSummary {
|
||||
deleted_1m_rows: deleted_1m,
|
||||
deleted_1h_rows: deleted_1h,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_unix_secs(value: Option<i64>) -> Option<u64> {
|
||||
@@ -861,10 +1167,63 @@ fn map_proxy_node_event_row(row: &MySqlRow) -> Result<StoredProxyNodeEvent, Data
|
||||
node_id: row.try_get("node_id").map_sql_err()?,
|
||||
event_type: row.try_get("event_type").map_sql_err()?,
|
||||
detail: row.try_get("detail").map_sql_err()?,
|
||||
event_metadata: optional_json_from_string(
|
||||
row.try_get("event_metadata").map_sql_err()?,
|
||||
"proxy_node_events.event_metadata",
|
||||
)?,
|
||||
created_at_unix_ms: optional_unix_secs(row.try_get("created_at_unix_ms").map_sql_err()?),
|
||||
})
|
||||
}
|
||||
|
||||
fn map_proxy_node_metric_row(
|
||||
row: &MySqlRow,
|
||||
) -> Result<StoredProxyNodeMetricsBucket, DataLayerError> {
|
||||
Ok(StoredProxyNodeMetricsBucket {
|
||||
node_id: row.try_get("node_id").map_sql_err()?,
|
||||
bucket_start_unix_secs: optional_unix_secs(
|
||||
row.try_get("bucket_start_unix_secs").map_sql_err()?,
|
||||
)
|
||||
.unwrap_or_default(),
|
||||
samples: row.try_get("samples").map_sql_err()?,
|
||||
uptime_samples: row.try_get("uptime_samples").map_sql_err()?,
|
||||
active_connections_sum: row.try_get("active_connections_sum").map_sql_err()?,
|
||||
active_connections_max: row.try_get("active_connections_max").map_sql_err()?,
|
||||
heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_sql_err()?,
|
||||
heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_sql_err()?,
|
||||
connect_errors_delta: row.try_get("connect_errors_delta").map_sql_err()?,
|
||||
disconnects_delta: row.try_get("disconnects_delta").map_sql_err()?,
|
||||
error_events_delta: row.try_get("error_events_delta").map_sql_err()?,
|
||||
ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_sql_err()?,
|
||||
ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_sql_err()?,
|
||||
ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_sql_err()?,
|
||||
ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_proxy_fleet_metric_row(
|
||||
row: &MySqlRow,
|
||||
) -> Result<StoredProxyFleetMetricsBucket, DataLayerError> {
|
||||
Ok(StoredProxyFleetMetricsBucket {
|
||||
bucket_start_unix_secs: optional_unix_secs(
|
||||
row.try_get("bucket_start_unix_secs").map_sql_err()?,
|
||||
)
|
||||
.unwrap_or_default(),
|
||||
samples: row.try_get("samples").map_sql_err()?,
|
||||
uptime_samples: row.try_get("uptime_samples").map_sql_err()?,
|
||||
active_connections_sum: row.try_get("active_connections_sum").map_sql_err()?,
|
||||
active_connections_max: row.try_get("active_connections_max").map_sql_err()?,
|
||||
heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_sql_err()?,
|
||||
heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_sql_err()?,
|
||||
connect_errors_delta: row.try_get("connect_errors_delta").map_sql_err()?,
|
||||
disconnects_delta: row.try_get("disconnects_delta").map_sql_err()?,
|
||||
error_events_delta: row.try_get("error_events_delta").map_sql_err()?,
|
||||
ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_sql_err()?,
|
||||
ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_sql_err()?,
|
||||
ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_sql_err()?,
|
||||
ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::MysqlProxyNodeReadRepository;
|
||||
|
||||
@@ -4,10 +4,14 @@ use sha2::{Digest, Sha256};
|
||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery,
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
|
||||
ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep, ProxyNodeReadRepository,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket,
|
||||
StoredProxyNode, StoredProxyNodeEvent, StoredProxyNodeMetricsBucket, TunnelMetricsSample,
|
||||
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use crate::{
|
||||
error::{postgres_error, SqlxResultExt},
|
||||
@@ -91,6 +95,7 @@ SELECT
|
||||
node_id,
|
||||
CAST(event_type AS TEXT) AS event_type,
|
||||
detail,
|
||||
event_metadata,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms
|
||||
FROM proxy_node_events
|
||||
WHERE node_id = $1
|
||||
@@ -98,6 +103,23 @@ ORDER BY created_at DESC, id DESC
|
||||
LIMIT $2
|
||||
"#;
|
||||
|
||||
const LIST_PROXY_NODE_EVENTS_FILTERED_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
node_id,
|
||||
CAST(event_type AS TEXT) AS event_type,
|
||||
detail,
|
||||
event_metadata,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms
|
||||
FROM proxy_node_events
|
||||
WHERE node_id = $1
|
||||
AND ($2::double precision IS NULL OR created_at >= TO_TIMESTAMP($2::double precision))
|
||||
AND ($3::double precision IS NULL OR created_at <= TO_TIMESTAMP($3::double precision))
|
||||
AND ($4::text IS NULL OR LOWER(CAST(event_type AS TEXT)) = LOWER($4::text))
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT $5
|
||||
"#;
|
||||
|
||||
const APPLY_HEARTBEAT_SQL: &str = r#"
|
||||
UPDATE proxy_nodes
|
||||
SET
|
||||
@@ -381,6 +403,188 @@ WHERE id = $4
|
||||
AND is_manual = TRUE
|
||||
"#;
|
||||
|
||||
const INSERT_PROXY_NODE_EVENT_SQL: &str = r#"
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, event_metadata, created_at)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4::json,
|
||||
CASE
|
||||
WHEN $5::double precision IS NULL THEN NOW()
|
||||
ELSE TO_TIMESTAMP($5::double precision)
|
||||
END
|
||||
)
|
||||
"#;
|
||||
|
||||
const UPSERT_PROXY_NODE_METRICS_1M_SQL: &str = r#"
|
||||
INSERT INTO proxy_node_metrics_1m (
|
||||
node_id,
|
||||
bucket_start_unix_secs,
|
||||
samples,
|
||||
uptime_samples,
|
||||
active_connections_sum,
|
||||
active_connections_max,
|
||||
heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max,
|
||||
connect_errors_delta,
|
||||
disconnects_delta,
|
||||
error_events_delta,
|
||||
ws_in_bytes_delta,
|
||||
ws_out_bytes_delta,
|
||||
ws_in_frames_delta,
|
||||
ws_out_frames_delta
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
|
||||
ON CONFLICT (node_id, bucket_start_unix_secs) DO UPDATE SET
|
||||
samples = proxy_node_metrics_1m.samples + EXCLUDED.samples,
|
||||
uptime_samples = proxy_node_metrics_1m.uptime_samples + EXCLUDED.uptime_samples,
|
||||
active_connections_sum = proxy_node_metrics_1m.active_connections_sum + EXCLUDED.active_connections_sum,
|
||||
active_connections_max = GREATEST(proxy_node_metrics_1m.active_connections_max, EXCLUDED.active_connections_max),
|
||||
heartbeat_rtt_ms_sum = proxy_node_metrics_1m.heartbeat_rtt_ms_sum + EXCLUDED.heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max = GREATEST(proxy_node_metrics_1m.heartbeat_rtt_ms_max, EXCLUDED.heartbeat_rtt_ms_max),
|
||||
connect_errors_delta = proxy_node_metrics_1m.connect_errors_delta + EXCLUDED.connect_errors_delta,
|
||||
disconnects_delta = proxy_node_metrics_1m.disconnects_delta + EXCLUDED.disconnects_delta,
|
||||
error_events_delta = proxy_node_metrics_1m.error_events_delta + EXCLUDED.error_events_delta,
|
||||
ws_in_bytes_delta = proxy_node_metrics_1m.ws_in_bytes_delta + EXCLUDED.ws_in_bytes_delta,
|
||||
ws_out_bytes_delta = proxy_node_metrics_1m.ws_out_bytes_delta + EXCLUDED.ws_out_bytes_delta,
|
||||
ws_in_frames_delta = proxy_node_metrics_1m.ws_in_frames_delta + EXCLUDED.ws_in_frames_delta,
|
||||
ws_out_frames_delta = proxy_node_metrics_1m.ws_out_frames_delta + EXCLUDED.ws_out_frames_delta
|
||||
"#;
|
||||
|
||||
const UPSERT_PROXY_NODE_METRICS_1H_SQL: &str = r#"
|
||||
INSERT INTO proxy_node_metrics_1h (
|
||||
node_id,
|
||||
bucket_start_unix_secs,
|
||||
samples,
|
||||
uptime_samples,
|
||||
active_connections_sum,
|
||||
active_connections_max,
|
||||
heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max,
|
||||
connect_errors_delta,
|
||||
disconnects_delta,
|
||||
error_events_delta,
|
||||
ws_in_bytes_delta,
|
||||
ws_out_bytes_delta,
|
||||
ws_in_frames_delta,
|
||||
ws_out_frames_delta
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
|
||||
ON CONFLICT (node_id, bucket_start_unix_secs) DO UPDATE SET
|
||||
samples = proxy_node_metrics_1h.samples + EXCLUDED.samples,
|
||||
uptime_samples = proxy_node_metrics_1h.uptime_samples + EXCLUDED.uptime_samples,
|
||||
active_connections_sum = proxy_node_metrics_1h.active_connections_sum + EXCLUDED.active_connections_sum,
|
||||
active_connections_max = GREATEST(proxy_node_metrics_1h.active_connections_max, EXCLUDED.active_connections_max),
|
||||
heartbeat_rtt_ms_sum = proxy_node_metrics_1h.heartbeat_rtt_ms_sum + EXCLUDED.heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max = GREATEST(proxy_node_metrics_1h.heartbeat_rtt_ms_max, EXCLUDED.heartbeat_rtt_ms_max),
|
||||
connect_errors_delta = proxy_node_metrics_1h.connect_errors_delta + EXCLUDED.connect_errors_delta,
|
||||
disconnects_delta = proxy_node_metrics_1h.disconnects_delta + EXCLUDED.disconnects_delta,
|
||||
error_events_delta = proxy_node_metrics_1h.error_events_delta + EXCLUDED.error_events_delta,
|
||||
ws_in_bytes_delta = proxy_node_metrics_1h.ws_in_bytes_delta + EXCLUDED.ws_in_bytes_delta,
|
||||
ws_out_bytes_delta = proxy_node_metrics_1h.ws_out_bytes_delta + EXCLUDED.ws_out_bytes_delta,
|
||||
ws_in_frames_delta = proxy_node_metrics_1h.ws_in_frames_delta + EXCLUDED.ws_in_frames_delta,
|
||||
ws_out_frames_delta = proxy_node_metrics_1h.ws_out_frames_delta + EXCLUDED.ws_out_frames_delta
|
||||
"#;
|
||||
|
||||
const LIST_PROXY_NODE_METRICS_1M_SQL: &str = r#"
|
||||
SELECT
|
||||
node_id,
|
||||
bucket_start_unix_secs,
|
||||
samples,
|
||||
uptime_samples,
|
||||
active_connections_sum,
|
||||
active_connections_max,
|
||||
heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max,
|
||||
connect_errors_delta,
|
||||
disconnects_delta,
|
||||
error_events_delta,
|
||||
ws_in_bytes_delta,
|
||||
ws_out_bytes_delta,
|
||||
ws_in_frames_delta,
|
||||
ws_out_frames_delta
|
||||
FROM proxy_node_metrics_1m
|
||||
WHERE node_id = $1
|
||||
AND bucket_start_unix_secs >= $2
|
||||
AND bucket_start_unix_secs <= $3
|
||||
ORDER BY bucket_start_unix_secs ASC
|
||||
LIMIT $4
|
||||
"#;
|
||||
|
||||
const LIST_PROXY_NODE_METRICS_1H_SQL: &str = r#"
|
||||
SELECT
|
||||
node_id,
|
||||
bucket_start_unix_secs,
|
||||
samples,
|
||||
uptime_samples,
|
||||
active_connections_sum,
|
||||
active_connections_max,
|
||||
heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max,
|
||||
connect_errors_delta,
|
||||
disconnects_delta,
|
||||
error_events_delta,
|
||||
ws_in_bytes_delta,
|
||||
ws_out_bytes_delta,
|
||||
ws_in_frames_delta,
|
||||
ws_out_frames_delta
|
||||
FROM proxy_node_metrics_1h
|
||||
WHERE node_id = $1
|
||||
AND bucket_start_unix_secs >= $2
|
||||
AND bucket_start_unix_secs <= $3
|
||||
ORDER BY bucket_start_unix_secs ASC
|
||||
LIMIT $4
|
||||
"#;
|
||||
|
||||
const LIST_PROXY_FLEET_METRICS_1M_SQL: &str = r#"
|
||||
SELECT
|
||||
bucket_start_unix_secs,
|
||||
SUM(samples) AS samples,
|
||||
SUM(uptime_samples) AS uptime_samples,
|
||||
SUM(active_connections_sum) AS active_connections_sum,
|
||||
MAX(active_connections_max) AS active_connections_max,
|
||||
SUM(heartbeat_rtt_ms_sum) AS heartbeat_rtt_ms_sum,
|
||||
MAX(heartbeat_rtt_ms_max) AS heartbeat_rtt_ms_max,
|
||||
SUM(connect_errors_delta) AS connect_errors_delta,
|
||||
SUM(disconnects_delta) AS disconnects_delta,
|
||||
SUM(error_events_delta) AS error_events_delta,
|
||||
SUM(ws_in_bytes_delta) AS ws_in_bytes_delta,
|
||||
SUM(ws_out_bytes_delta) AS ws_out_bytes_delta,
|
||||
SUM(ws_in_frames_delta) AS ws_in_frames_delta,
|
||||
SUM(ws_out_frames_delta) AS ws_out_frames_delta
|
||||
FROM proxy_node_metrics_1m
|
||||
WHERE bucket_start_unix_secs >= $1
|
||||
AND bucket_start_unix_secs <= $2
|
||||
GROUP BY bucket_start_unix_secs
|
||||
ORDER BY bucket_start_unix_secs ASC
|
||||
LIMIT $3
|
||||
"#;
|
||||
|
||||
const LIST_PROXY_FLEET_METRICS_1H_SQL: &str = r#"
|
||||
SELECT
|
||||
bucket_start_unix_secs,
|
||||
SUM(samples) AS samples,
|
||||
SUM(uptime_samples) AS uptime_samples,
|
||||
SUM(active_connections_sum) AS active_connections_sum,
|
||||
MAX(active_connections_max) AS active_connections_max,
|
||||
SUM(heartbeat_rtt_ms_sum) AS heartbeat_rtt_ms_sum,
|
||||
MAX(heartbeat_rtt_ms_max) AS heartbeat_rtt_ms_max,
|
||||
SUM(connect_errors_delta) AS connect_errors_delta,
|
||||
SUM(disconnects_delta) AS disconnects_delta,
|
||||
SUM(error_events_delta) AS error_events_delta,
|
||||
SUM(ws_in_bytes_delta) AS ws_in_bytes_delta,
|
||||
SUM(ws_out_bytes_delta) AS ws_out_bytes_delta,
|
||||
SUM(ws_in_frames_delta) AS ws_in_frames_delta,
|
||||
SUM(ws_out_frames_delta) AS ws_out_frames_delta
|
||||
FROM proxy_node_metrics_1h
|
||||
WHERE bucket_start_unix_secs >= $1
|
||||
AND bucket_start_unix_secs <= $2
|
||||
GROUP BY bucket_start_unix_secs
|
||||
ORDER BY bucket_start_unix_secs ASC
|
||||
LIMIT $3
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxProxyNodeRepository {
|
||||
pool: PgPool,
|
||||
@@ -446,12 +650,111 @@ impl SqlxProxyNodeRepository {
|
||||
node_id: row.try_get("node_id").map_postgres_err()?,
|
||||
event_type: row.try_get("event_type").map_postgres_err()?,
|
||||
detail: row.try_get("detail").map_postgres_err()?,
|
||||
event_metadata: row.try_get("event_metadata").map_postgres_err()?,
|
||||
created_at_unix_ms: Self::optional_unix_secs(
|
||||
row.try_get("created_at_unix_ms").map_postgres_err()?,
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
fn row_to_node_metric(row: &PgRow) -> Result<StoredProxyNodeMetricsBucket, DataLayerError> {
|
||||
Ok(StoredProxyNodeMetricsBucket {
|
||||
node_id: row.try_get("node_id").map_postgres_err()?,
|
||||
bucket_start_unix_secs: Self::optional_unix_secs(
|
||||
row.try_get("bucket_start_unix_secs").map_postgres_err()?,
|
||||
)
|
||||
.unwrap_or_default(),
|
||||
samples: row.try_get("samples").map_postgres_err()?,
|
||||
uptime_samples: row.try_get("uptime_samples").map_postgres_err()?,
|
||||
active_connections_sum: row.try_get("active_connections_sum").map_postgres_err()?,
|
||||
active_connections_max: row.try_get("active_connections_max").map_postgres_err()?,
|
||||
heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_postgres_err()?,
|
||||
heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_postgres_err()?,
|
||||
connect_errors_delta: row.try_get("connect_errors_delta").map_postgres_err()?,
|
||||
disconnects_delta: row.try_get("disconnects_delta").map_postgres_err()?,
|
||||
error_events_delta: row.try_get("error_events_delta").map_postgres_err()?,
|
||||
ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_postgres_err()?,
|
||||
ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_postgres_err()?,
|
||||
ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_postgres_err()?,
|
||||
ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_postgres_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn row_to_fleet_metric(row: &PgRow) -> Result<StoredProxyFleetMetricsBucket, DataLayerError> {
|
||||
Ok(StoredProxyFleetMetricsBucket {
|
||||
bucket_start_unix_secs: Self::optional_unix_secs(
|
||||
row.try_get("bucket_start_unix_secs").map_postgres_err()?,
|
||||
)
|
||||
.unwrap_or_default(),
|
||||
samples: row.try_get("samples").map_postgres_err()?,
|
||||
uptime_samples: row.try_get("uptime_samples").map_postgres_err()?,
|
||||
active_connections_sum: row.try_get("active_connections_sum").map_postgres_err()?,
|
||||
active_connections_max: row.try_get("active_connections_max").map_postgres_err()?,
|
||||
heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_postgres_err()?,
|
||||
heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_postgres_err()?,
|
||||
connect_errors_delta: row.try_get("connect_errors_delta").map_postgres_err()?,
|
||||
disconnects_delta: row.try_get("disconnects_delta").map_postgres_err()?,
|
||||
error_events_delta: row.try_get("error_events_delta").map_postgres_err()?,
|
||||
ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_postgres_err()?,
|
||||
ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_postgres_err()?,
|
||||
ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_postgres_err()?,
|
||||
ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_postgres_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
async fn insert_event(
|
||||
&self,
|
||||
node_id: &str,
|
||||
event_type: &str,
|
||||
detail: Option<&str>,
|
||||
event_metadata: Option<&serde_json::Value>,
|
||||
created_at_unix_secs: Option<u64>,
|
||||
) -> Result<(), DataLayerError> {
|
||||
sqlx::query(INSERT_PROXY_NODE_EVENT_SQL)
|
||||
.bind(node_id)
|
||||
.bind(event_type)
|
||||
.bind(detail)
|
||||
.bind(event_metadata)
|
||||
.bind(created_at_unix_secs.map(|value| value as f64))
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn upsert_metrics_bucket(
|
||||
&self,
|
||||
step: ProxyNodeMetricsStep,
|
||||
node_id: &str,
|
||||
bucket_start: u64,
|
||||
sample: &TunnelMetricsSample,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let sql = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => UPSERT_PROXY_NODE_METRICS_1M_SQL,
|
||||
ProxyNodeMetricsStep::OneHour => UPSERT_PROXY_NODE_METRICS_1H_SQL,
|
||||
};
|
||||
sqlx::query(sql)
|
||||
.bind(node_id)
|
||||
.bind(i64::try_from(bucket_start).unwrap_or(i64::MAX))
|
||||
.bind(sample.samples)
|
||||
.bind(sample.uptime_samples)
|
||||
.bind(sample.active_connections_sum)
|
||||
.bind(sample.active_connections_max)
|
||||
.bind(sample.heartbeat_rtt_ms_sum)
|
||||
.bind(sample.heartbeat_rtt_ms_max)
|
||||
.bind(sample.connect_errors_delta)
|
||||
.bind(sample.disconnects_delta)
|
||||
.bind(sample.error_events_delta)
|
||||
.bind(sample.ws_in_bytes_delta)
|
||||
.bind(sample.ws_out_bytes_delta)
|
||||
.bind(sample.ws_in_frames_delta)
|
||||
.bind(sample.ws_out_frames_delta)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn registration_lock_key(ip: &str, port: i32) -> i64 {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(ip.as_bytes());
|
||||
@@ -602,6 +905,73 @@ impl ProxyNodeReadRepository for SqlxProxyNodeRepository {
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn list_proxy_node_events_filtered(
|
||||
&self,
|
||||
node_id: &str,
|
||||
query: &ProxyNodeEventQuery,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
||||
let mut rows = sqlx::query(LIST_PROXY_NODE_EVENTS_FILTERED_SQL)
|
||||
.bind(node_id)
|
||||
.bind(query.from_unix_secs.map(|value| value as f64))
|
||||
.bind(query.to_unix_secs.map(|value| value as f64))
|
||||
.bind(query.event_type.as_deref())
|
||||
.bind(i64::try_from(query.limit).unwrap_or(i64::MAX))
|
||||
.fetch(&self.pool);
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
items.push(Self::row_to_event(&row)?);
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
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>, DataLayerError> {
|
||||
let sql = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => LIST_PROXY_NODE_METRICS_1M_SQL,
|
||||
ProxyNodeMetricsStep::OneHour => LIST_PROXY_NODE_METRICS_1H_SQL,
|
||||
};
|
||||
let mut rows = sqlx::query(sql)
|
||||
.bind(node_id)
|
||||
.bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.fetch(&self.pool);
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
items.push(Self::row_to_node_metric(&row)?);
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn list_proxy_fleet_metrics(
|
||||
&self,
|
||||
step: ProxyNodeMetricsStep,
|
||||
from_unix_secs: u64,
|
||||
to_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyFleetMetricsBucket>, DataLayerError> {
|
||||
let sql = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => LIST_PROXY_FLEET_METRICS_1M_SQL,
|
||||
ProxyNodeMetricsStep::OneHour => LIST_PROXY_FLEET_METRICS_1H_SQL,
|
||||
};
|
||||
let mut rows = sqlx::query(sql)
|
||||
.bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.fetch(&self.pool);
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
items.push(Self::row_to_fleet_metric(&row)?);
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -822,6 +1192,54 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
|
||||
let Some(updated) = updated else {
|
||||
return Ok(None);
|
||||
};
|
||||
let now_unix_secs = updated
|
||||
.last_heartbeat_at_unix_secs
|
||||
.unwrap_or_else(|| chrono::Utc::now().timestamp().max(0) as u64);
|
||||
let tunnel_metrics_sample = build_tunnel_metrics_sample(
|
||||
existing.proxy_metadata.as_ref(),
|
||||
updated.proxy_metadata.as_ref(),
|
||||
updated.active_connections,
|
||||
updated.tunnel_connected,
|
||||
);
|
||||
|
||||
if let Some(sample) = tunnel_metrics_sample.as_ref() {
|
||||
self.upsert_metrics_bucket(
|
||||
ProxyNodeMetricsStep::OneMinute,
|
||||
&updated.id,
|
||||
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneMinute),
|
||||
sample,
|
||||
)
|
||||
.await?;
|
||||
self.upsert_metrics_bucket(
|
||||
ProxyNodeMetricsStep::OneHour,
|
||||
&updated.id,
|
||||
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneHour),
|
||||
sample,
|
||||
)
|
||||
.await?;
|
||||
|
||||
for error in &sample.recent_error_events {
|
||||
let detail = build_tunnel_error_event_detail(error);
|
||||
let event_metadata = serde_json::json!({
|
||||
"source": "heartbeat",
|
||||
"category": error.category,
|
||||
"message": error.message,
|
||||
"timestamp_unix_secs": error.timestamp_unix_secs,
|
||||
});
|
||||
self.insert_event(
|
||||
&updated.id,
|
||||
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
Some(detail.as_str()),
|
||||
Some(&event_metadata),
|
||||
Some(if error.timestamp_unix_secs == 0 {
|
||||
now_unix_secs
|
||||
} else {
|
||||
error.timestamp_unix_secs
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
if reconcile_remote_config_after_heartbeat(
|
||||
updated.remote_config.as_ref(),
|
||||
@@ -890,23 +1308,15 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
|
||||
.zip(observed_at_unix_secs)
|
||||
.is_some_and(|(last_transition, observed_at)| observed_at < last_transition)
|
||||
{
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, created_at)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
NOW()
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(event_type)
|
||||
.bind(format!("[stale_ignored] {event_detail}"))
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
sqlx::query(INSERT_PROXY_NODE_EVENT_SQL)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(event_type)
|
||||
.bind(format!("[stale_ignored] {event_detail}"))
|
||||
.bind(None::<serde_json::Value>)
|
||||
.bind(None::<f64>)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
tx.commit().await.map_err(postgres_error)?;
|
||||
return self.find_proxy_node(&mutation.node_id).await;
|
||||
}
|
||||
@@ -942,27 +1352,15 @@ WHERE id = $1
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, created_at)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
CASE
|
||||
WHEN $4::double precision IS NULL THEN NOW()
|
||||
ELSE TO_TIMESTAMP($4::double precision)
|
||||
END
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(event_type)
|
||||
.bind(event_detail)
|
||||
.bind(observed_at_unix_secs.map(|value| value as f64))
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
sqlx::query(INSERT_PROXY_NODE_EVENT_SQL)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(event_type)
|
||||
.bind(event_detail)
|
||||
.bind(None::<serde_json::Value>)
|
||||
.bind(observed_at_unix_secs.map(|value| value as f64))
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
tx.commit().await.map_err(postgres_error)?;
|
||||
self.find_proxy_node(&mutation.node_id).await
|
||||
@@ -1045,6 +1443,33 @@ VALUES (
|
||||
.map_postgres_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup_proxy_node_metrics(
|
||||
&self,
|
||||
retain_1m_from_unix_secs: u64,
|
||||
retain_1h_from_unix_secs: u64,
|
||||
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
|
||||
let deleted_1m =
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1m WHERE bucket_start_unix_secs < $1")
|
||||
.bind(i64::try_from(retain_1m_from_unix_secs).unwrap_or(i64::MAX))
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected() as usize;
|
||||
|
||||
let deleted_1h =
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1h WHERE bucket_start_unix_secs < $1")
|
||||
.bind(i64::try_from(retain_1h_from_unix_secs).unwrap_or(i64::MAX))
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected() as usize;
|
||||
|
||||
Ok(ProxyNodeMetricsCleanupSummary {
|
||||
deleted_1m_rows: deleted_1m,
|
||||
deleted_1h_rows: deleted_1h,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -2,10 +2,14 @@ use async_trait::async_trait;
|
||||
use sqlx::{sqlite::SqliteRow, Row};
|
||||
|
||||
use super::types::{
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery,
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
|
||||
ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep, ProxyNodeReadRepository,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket,
|
||||
StoredProxyNode, StoredProxyNodeEvent, StoredProxyNodeMetricsBucket,
|
||||
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::error::SqlResultExt;
|
||||
@@ -148,17 +152,22 @@ ON CONFLICT(id) DO UPDATE SET
|
||||
node_id: &str,
|
||||
event_type: &str,
|
||||
detail: Option<&str>,
|
||||
event_metadata: Option<&serde_json::Value>,
|
||||
created_at_unix_secs: Option<u64>,
|
||||
) -> Result<(), DataLayerError> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, event_metadata, created_at)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(node_id)
|
||||
.bind(event_type)
|
||||
.bind(detail)
|
||||
.bind(optional_json_to_string(
|
||||
&event_metadata.cloned(),
|
||||
"proxy_node_events.event_metadata",
|
||||
)?)
|
||||
.bind(created_at_unix_secs.unwrap_or_else(current_unix_secs) as i64)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
@@ -166,6 +175,70 @@ VALUES (?, ?, ?, ?)
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn upsert_metrics_bucket(
|
||||
&self,
|
||||
table: &str,
|
||||
node_id: &str,
|
||||
bucket_start: u64,
|
||||
sample: &super::types::TunnelMetricsSample,
|
||||
) -> Result<(), DataLayerError> {
|
||||
sqlx::query(&format!(
|
||||
r#"
|
||||
INSERT INTO {table} (
|
||||
node_id,
|
||||
bucket_start_unix_secs,
|
||||
samples,
|
||||
uptime_samples,
|
||||
active_connections_sum,
|
||||
active_connections_max,
|
||||
heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max,
|
||||
connect_errors_delta,
|
||||
disconnects_delta,
|
||||
error_events_delta,
|
||||
ws_in_bytes_delta,
|
||||
ws_out_bytes_delta,
|
||||
ws_in_frames_delta,
|
||||
ws_out_frames_delta
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(node_id, bucket_start_unix_secs) DO UPDATE SET
|
||||
samples = {table}.samples + excluded.samples,
|
||||
uptime_samples = {table}.uptime_samples + excluded.uptime_samples,
|
||||
active_connections_sum = {table}.active_connections_sum + excluded.active_connections_sum,
|
||||
active_connections_max = MAX({table}.active_connections_max, excluded.active_connections_max),
|
||||
heartbeat_rtt_ms_sum = {table}.heartbeat_rtt_ms_sum + excluded.heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max = MAX({table}.heartbeat_rtt_ms_max, excluded.heartbeat_rtt_ms_max),
|
||||
connect_errors_delta = {table}.connect_errors_delta + excluded.connect_errors_delta,
|
||||
disconnects_delta = {table}.disconnects_delta + excluded.disconnects_delta,
|
||||
error_events_delta = {table}.error_events_delta + excluded.error_events_delta,
|
||||
ws_in_bytes_delta = {table}.ws_in_bytes_delta + excluded.ws_in_bytes_delta,
|
||||
ws_out_bytes_delta = {table}.ws_out_bytes_delta + excluded.ws_out_bytes_delta,
|
||||
ws_in_frames_delta = {table}.ws_in_frames_delta + excluded.ws_in_frames_delta,
|
||||
ws_out_frames_delta = {table}.ws_out_frames_delta + excluded.ws_out_frames_delta
|
||||
"#
|
||||
))
|
||||
.bind(node_id)
|
||||
.bind(i64::try_from(bucket_start).unwrap_or(i64::MAX))
|
||||
.bind(sample.samples)
|
||||
.bind(sample.uptime_samples)
|
||||
.bind(sample.active_connections_sum)
|
||||
.bind(sample.active_connections_max)
|
||||
.bind(sample.heartbeat_rtt_ms_sum)
|
||||
.bind(sample.heartbeat_rtt_ms_max)
|
||||
.bind(sample.connect_errors_delta)
|
||||
.bind(sample.disconnects_delta)
|
||||
.bind(sample.error_events_delta)
|
||||
.bind(sample.ws_in_bytes_delta)
|
||||
.bind(sample.ws_out_bytes_delta)
|
||||
.bind(sample.ws_in_frames_delta)
|
||||
.bind(sample.ws_out_frames_delta)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn normalize_remote_config(
|
||||
mutation: &ProxyNodeRemoteConfigMutation,
|
||||
existing: Option<&serde_json::Value>,
|
||||
@@ -298,6 +371,7 @@ SELECT
|
||||
node_id,
|
||||
event_type,
|
||||
detail,
|
||||
event_metadata,
|
||||
created_at AS created_at_unix_ms
|
||||
FROM proxy_node_events
|
||||
WHERE node_id = ?
|
||||
@@ -312,6 +386,152 @@ LIMIT ?
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_proxy_node_event_row).collect()
|
||||
}
|
||||
|
||||
async fn list_proxy_node_events_filtered(
|
||||
&self,
|
||||
node_id: &str,
|
||||
query: &ProxyNodeEventQuery,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
||||
let rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
node_id,
|
||||
event_type,
|
||||
detail,
|
||||
event_metadata,
|
||||
created_at AS created_at_unix_ms
|
||||
FROM proxy_node_events
|
||||
WHERE node_id = ?
|
||||
AND (? IS NULL OR created_at >= ?)
|
||||
AND (? IS NULL OR created_at <= ?)
|
||||
AND (? IS NULL OR LOWER(event_type) = LOWER(?))
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT ?
|
||||
"#,
|
||||
)
|
||||
.bind(node_id)
|
||||
.bind(
|
||||
query
|
||||
.from_unix_secs
|
||||
.map(|v| i64::try_from(v).unwrap_or(i64::MAX)),
|
||||
)
|
||||
.bind(
|
||||
query
|
||||
.from_unix_secs
|
||||
.map(|v| i64::try_from(v).unwrap_or(i64::MAX)),
|
||||
)
|
||||
.bind(
|
||||
query
|
||||
.to_unix_secs
|
||||
.map(|v| i64::try_from(v).unwrap_or(i64::MAX)),
|
||||
)
|
||||
.bind(
|
||||
query
|
||||
.to_unix_secs
|
||||
.map(|v| i64::try_from(v).unwrap_or(i64::MAX)),
|
||||
)
|
||||
.bind(query.event_type.as_deref())
|
||||
.bind(query.event_type.as_deref())
|
||||
.bind(i64::try_from(query.limit).unwrap_or(i64::MAX))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_proxy_node_event_row).collect()
|
||||
}
|
||||
|
||||
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>, DataLayerError> {
|
||||
let table = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => "proxy_node_metrics_1m",
|
||||
ProxyNodeMetricsStep::OneHour => "proxy_node_metrics_1h",
|
||||
};
|
||||
let rows = sqlx::query(&format!(
|
||||
r#"
|
||||
SELECT
|
||||
node_id,
|
||||
bucket_start_unix_secs,
|
||||
samples,
|
||||
uptime_samples,
|
||||
active_connections_sum,
|
||||
active_connections_max,
|
||||
heartbeat_rtt_ms_sum,
|
||||
heartbeat_rtt_ms_max,
|
||||
connect_errors_delta,
|
||||
disconnects_delta,
|
||||
error_events_delta,
|
||||
ws_in_bytes_delta,
|
||||
ws_out_bytes_delta,
|
||||
ws_in_frames_delta,
|
||||
ws_out_frames_delta
|
||||
FROM {table}
|
||||
WHERE node_id = ?
|
||||
AND bucket_start_unix_secs >= ?
|
||||
AND bucket_start_unix_secs <= ?
|
||||
ORDER BY bucket_start_unix_secs ASC
|
||||
LIMIT ?
|
||||
"#
|
||||
))
|
||||
.bind(node_id)
|
||||
.bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_proxy_node_metric_row).collect()
|
||||
}
|
||||
|
||||
async fn list_proxy_fleet_metrics(
|
||||
&self,
|
||||
step: ProxyNodeMetricsStep,
|
||||
from_unix_secs: u64,
|
||||
to_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyFleetMetricsBucket>, DataLayerError> {
|
||||
let table = match step {
|
||||
ProxyNodeMetricsStep::OneMinute => "proxy_node_metrics_1m",
|
||||
ProxyNodeMetricsStep::OneHour => "proxy_node_metrics_1h",
|
||||
};
|
||||
let rows = sqlx::query(&format!(
|
||||
r#"
|
||||
SELECT
|
||||
bucket_start_unix_secs,
|
||||
SUM(samples) AS samples,
|
||||
SUM(uptime_samples) AS uptime_samples,
|
||||
SUM(active_connections_sum) AS active_connections_sum,
|
||||
MAX(active_connections_max) AS active_connections_max,
|
||||
SUM(heartbeat_rtt_ms_sum) AS heartbeat_rtt_ms_sum,
|
||||
MAX(heartbeat_rtt_ms_max) AS heartbeat_rtt_ms_max,
|
||||
SUM(connect_errors_delta) AS connect_errors_delta,
|
||||
SUM(disconnects_delta) AS disconnects_delta,
|
||||
SUM(error_events_delta) AS error_events_delta,
|
||||
SUM(ws_in_bytes_delta) AS ws_in_bytes_delta,
|
||||
SUM(ws_out_bytes_delta) AS ws_out_bytes_delta,
|
||||
SUM(ws_in_frames_delta) AS ws_in_frames_delta,
|
||||
SUM(ws_out_frames_delta) AS ws_out_frames_delta
|
||||
FROM {table}
|
||||
WHERE bucket_start_unix_secs >= ?
|
||||
AND bucket_start_unix_secs <= ?
|
||||
GROUP BY bucket_start_unix_secs
|
||||
ORDER BY bucket_start_unix_secs ASC
|
||||
LIMIT ?
|
||||
"#
|
||||
))
|
||||
.bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_proxy_fleet_metric_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -540,7 +760,9 @@ WHERE is_manual = 0
|
||||
));
|
||||
}
|
||||
|
||||
let now = Some(current_unix_secs());
|
||||
let previous_proxy_metadata = node.proxy_metadata.clone();
|
||||
let now_unix_secs = current_unix_secs();
|
||||
let now = Some(now_unix_secs);
|
||||
node.last_heartbeat_at_unix_secs = now;
|
||||
if node.status != "online" || !node.tunnel_connected {
|
||||
node.status = "online".to_string();
|
||||
@@ -584,7 +806,55 @@ WHERE is_manual = 0
|
||||
node.config_version = node.config_version.saturating_add(1);
|
||||
node.updated_at_unix_secs = now;
|
||||
}
|
||||
|
||||
let tunnel_metrics_sample = build_tunnel_metrics_sample(
|
||||
previous_proxy_metadata.as_ref(),
|
||||
node.proxy_metadata.as_ref(),
|
||||
node.active_connections,
|
||||
node.tunnel_connected,
|
||||
);
|
||||
|
||||
self.upsert_node(&node).await?;
|
||||
|
||||
if let Some(sample) = tunnel_metrics_sample.as_ref() {
|
||||
self.upsert_metrics_bucket(
|
||||
"proxy_node_metrics_1m",
|
||||
&node.id,
|
||||
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneMinute),
|
||||
sample,
|
||||
)
|
||||
.await?;
|
||||
self.upsert_metrics_bucket(
|
||||
"proxy_node_metrics_1h",
|
||||
&node.id,
|
||||
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneHour),
|
||||
sample,
|
||||
)
|
||||
.await?;
|
||||
|
||||
for error in &sample.recent_error_events {
|
||||
let detail = build_tunnel_error_event_detail(error);
|
||||
let event_metadata = serde_json::json!({
|
||||
"source": "heartbeat",
|
||||
"category": error.category,
|
||||
"message": error.message,
|
||||
"timestamp_unix_secs": error.timestamp_unix_secs,
|
||||
});
|
||||
self.insert_event(
|
||||
&node.id,
|
||||
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
Some(detail.as_str()),
|
||||
Some(&event_metadata),
|
||||
Some(if error.timestamp_unix_secs == 0 {
|
||||
now_unix_secs
|
||||
} else {
|
||||
error.timestamp_unix_secs
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Some(node))
|
||||
}
|
||||
|
||||
@@ -638,6 +908,7 @@ WHERE is_manual = 0
|
||||
&mutation.node_id,
|
||||
event_type,
|
||||
Some(&format!("[stale_ignored] {event_detail}")),
|
||||
None,
|
||||
Some(current_unix_secs()),
|
||||
)
|
||||
.await?;
|
||||
@@ -660,6 +931,7 @@ WHERE is_manual = 0
|
||||
&mutation.node_id,
|
||||
event_type,
|
||||
Some(&event_detail),
|
||||
None,
|
||||
Some(event_time),
|
||||
)
|
||||
.await?;
|
||||
@@ -691,6 +963,16 @@ WHERE is_manual = 0
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1m WHERE node_id = ?")
|
||||
.bind(node_id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1h WHERE node_id = ?")
|
||||
.bind(node_id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM proxy_nodes WHERE id = ?")
|
||||
.bind(node_id)
|
||||
.execute(&self.pool)
|
||||
@@ -747,6 +1029,33 @@ WHERE is_manual = 0
|
||||
node.updated_at_unix_secs = Some(current_unix_secs());
|
||||
self.upsert_node(&node).await
|
||||
}
|
||||
|
||||
async fn cleanup_proxy_node_metrics(
|
||||
&self,
|
||||
retain_1m_from_unix_secs: u64,
|
||||
retain_1h_from_unix_secs: u64,
|
||||
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
|
||||
let deleted_1m =
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1m WHERE bucket_start_unix_secs < ?")
|
||||
.bind(i64::try_from(retain_1m_from_unix_secs).unwrap_or(i64::MAX))
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected() as usize;
|
||||
|
||||
let deleted_1h =
|
||||
sqlx::query("DELETE FROM proxy_node_metrics_1h WHERE bucket_start_unix_secs < ?")
|
||||
.bind(i64::try_from(retain_1h_from_unix_secs).unwrap_or(i64::MAX))
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected() as usize;
|
||||
|
||||
Ok(ProxyNodeMetricsCleanupSummary {
|
||||
deleted_1m_rows: deleted_1m,
|
||||
deleted_1h_rows: deleted_1h,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_unix_secs(value: Option<i64>) -> Option<u64> {
|
||||
@@ -861,18 +1170,73 @@ fn map_proxy_node_event_row(row: &SqliteRow) -> Result<StoredProxyNodeEvent, Dat
|
||||
node_id: row.try_get("node_id").map_sql_err()?,
|
||||
event_type: row.try_get("event_type").map_sql_err()?,
|
||||
detail: row.try_get("detail").map_sql_err()?,
|
||||
event_metadata: optional_json_from_string(
|
||||
row.try_get("event_metadata").map_sql_err()?,
|
||||
"proxy_node_events.event_metadata",
|
||||
)?,
|
||||
created_at_unix_ms: optional_unix_secs(row.try_get("created_at_unix_ms").map_sql_err()?),
|
||||
})
|
||||
}
|
||||
|
||||
fn map_proxy_node_metric_row(
|
||||
row: &SqliteRow,
|
||||
) -> Result<StoredProxyNodeMetricsBucket, DataLayerError> {
|
||||
Ok(StoredProxyNodeMetricsBucket {
|
||||
node_id: row.try_get("node_id").map_sql_err()?,
|
||||
bucket_start_unix_secs: optional_unix_secs(
|
||||
row.try_get("bucket_start_unix_secs").map_sql_err()?,
|
||||
)
|
||||
.unwrap_or_default(),
|
||||
samples: row.try_get("samples").map_sql_err()?,
|
||||
uptime_samples: row.try_get("uptime_samples").map_sql_err()?,
|
||||
active_connections_sum: row.try_get("active_connections_sum").map_sql_err()?,
|
||||
active_connections_max: row.try_get("active_connections_max").map_sql_err()?,
|
||||
heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_sql_err()?,
|
||||
heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_sql_err()?,
|
||||
connect_errors_delta: row.try_get("connect_errors_delta").map_sql_err()?,
|
||||
disconnects_delta: row.try_get("disconnects_delta").map_sql_err()?,
|
||||
error_events_delta: row.try_get("error_events_delta").map_sql_err()?,
|
||||
ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_sql_err()?,
|
||||
ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_sql_err()?,
|
||||
ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_sql_err()?,
|
||||
ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_proxy_fleet_metric_row(
|
||||
row: &SqliteRow,
|
||||
) -> Result<StoredProxyFleetMetricsBucket, DataLayerError> {
|
||||
Ok(StoredProxyFleetMetricsBucket {
|
||||
bucket_start_unix_secs: optional_unix_secs(
|
||||
row.try_get("bucket_start_unix_secs").map_sql_err()?,
|
||||
)
|
||||
.unwrap_or_default(),
|
||||
samples: row.try_get("samples").map_sql_err()?,
|
||||
uptime_samples: row.try_get("uptime_samples").map_sql_err()?,
|
||||
active_connections_sum: row.try_get("active_connections_sum").map_sql_err()?,
|
||||
active_connections_max: row.try_get("active_connections_max").map_sql_err()?,
|
||||
heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_sql_err()?,
|
||||
heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_sql_err()?,
|
||||
connect_errors_delta: row.try_get("connect_errors_delta").map_sql_err()?,
|
||||
disconnects_delta: row.try_get("disconnects_delta").map_sql_err()?,
|
||||
error_events_delta: row.try_get("error_events_delta").map_sql_err()?,
|
||||
ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_sql_err()?,
|
||||
ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_sql_err()?,
|
||||
ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_sql_err()?,
|
||||
ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqliteProxyNodeReadRepository;
|
||||
use crate::lifecycle::migrate::run_sqlite_migrations;
|
||||
use crate::repository::proxy_nodes::{
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
|
||||
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
|
||||
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
|
||||
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
|
||||
ProxyNodeManualUpdateMutation, ProxyNodeMetricsStep, ProxyNodeReadRepository,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
|
||||
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
@@ -1142,4 +1506,131 @@ VALUES ('node-1', 'registered', 'ok', 3)
|
||||
.expect("manual node should delete")
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_repository_aggregates_proxy_node_metrics_and_filters_events() {
|
||||
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.expect("sqlite pool should connect");
|
||||
run_sqlite_migrations(&pool)
|
||||
.await
|
||||
.expect("sqlite migrations should run");
|
||||
|
||||
let repository = SqliteProxyNodeReadRepository::new(pool);
|
||||
let registered = repository
|
||||
.register_node(&ProxyNodeRegistrationMutation {
|
||||
name: "tunnel-1".to_string(),
|
||||
ip: "10.0.0.1".to_string(),
|
||||
port: 7000,
|
||||
region: None,
|
||||
heartbeat_interval: 30,
|
||||
active_connections: Some(0),
|
||||
total_requests: Some(0),
|
||||
avg_latency_ms: None,
|
||||
hardware_info: None,
|
||||
estimated_max_concurrency: None,
|
||||
proxy_metadata: None,
|
||||
proxy_version: Some("1.0.0".to_string()),
|
||||
registered_by: None,
|
||||
tunnel_mode: true,
|
||||
})
|
||||
.await
|
||||
.expect("node should register");
|
||||
let now = super::current_unix_secs();
|
||||
repository
|
||||
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
||||
node_id: registered.id.clone(),
|
||||
heartbeat_interval: Some(30),
|
||||
active_connections: Some(5),
|
||||
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": 4,
|
||||
"disconnects": 1,
|
||||
"error_events_total": 1,
|
||||
"ws_in_bytes": 100,
|
||||
"ws_out_bytes": 200,
|
||||
"ws_in_frames": 3,
|
||||
"ws_out_frames": 6,
|
||||
"heartbeat_rtt_last_ms": 33
|
||||
},
|
||||
"recent_tunnel_errors": [{
|
||||
"timestamp_unix_secs": now,
|
||||
"category": "tcp_connect_timeout",
|
||||
"message": "timeout"
|
||||
}]
|
||||
})),
|
||||
proxy_version: Some("1.0.0".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect("heartbeat should apply")
|
||||
.expect("node should exist");
|
||||
|
||||
let metrics = repository
|
||||
.list_proxy_node_metrics(
|
||||
®istered.id,
|
||||
ProxyNodeMetricsStep::OneMinute,
|
||||
now.saturating_sub(120),
|
||||
now.saturating_add(120),
|
||||
10,
|
||||
)
|
||||
.await
|
||||
.expect("metrics should list");
|
||||
assert_eq!(metrics.len(), 1);
|
||||
assert_eq!(metrics[0].samples, 1);
|
||||
assert_eq!(metrics[0].uptime_samples, 1);
|
||||
assert_eq!(metrics[0].active_connections_max, 5);
|
||||
assert_eq!(metrics[0].heartbeat_rtt_ms_sum, 33);
|
||||
assert_eq!(metrics[0].connect_errors_delta, 4);
|
||||
assert_eq!(metrics[0].ws_out_frames_delta, 6);
|
||||
|
||||
let fleet = repository
|
||||
.list_proxy_fleet_metrics(
|
||||
ProxyNodeMetricsStep::OneMinute,
|
||||
now.saturating_sub(120),
|
||||
now.saturating_add(120),
|
||||
10,
|
||||
)
|
||||
.await
|
||||
.expect("fleet metrics should list");
|
||||
assert_eq!(fleet.len(), 1);
|
||||
assert_eq!(fleet[0].samples, 1);
|
||||
assert_eq!(fleet[0].error_events_delta, 1);
|
||||
|
||||
let events = repository
|
||||
.list_proxy_node_events_filtered(
|
||||
®istered.id,
|
||||
&ProxyNodeEventQuery {
|
||||
limit: 10,
|
||||
from_unix_secs: Some(now.saturating_sub(120)),
|
||||
to_unix_secs: Some(now.saturating_add(120)),
|
||||
event_type: Some(PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR.to_string()),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("events should list");
|
||||
assert_eq!(events.len(), 1);
|
||||
assert_eq!(events[0].event_type, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR);
|
||||
assert_eq!(
|
||||
events[0]
|
||||
.event_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("category"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("tcp_connect_timeout")
|
||||
);
|
||||
|
||||
let cleanup = repository
|
||||
.cleanup_proxy_node_metrics(now.saturating_add(1), now.saturating_add(1))
|
||||
.await
|
||||
.expect("cleanup should run");
|
||||
assert_eq!(cleanup.deleted_1m_rows, 1);
|
||||
assert_eq!(cleanup.deleted_1h_rows, 1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProxyNode {
|
||||
@@ -239,9 +240,186 @@ pub struct StoredProxyNodeEvent {
|
||||
pub node_id: String,
|
||||
pub event_type: String,
|
||||
pub detail: Option<String>,
|
||||
pub event_metadata: Option<serde_json::Value>,
|
||||
pub created_at_unix_ms: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProxyNodeEventQuery {
|
||||
pub limit: usize,
|
||||
pub from_unix_secs: Option<u64>,
|
||||
pub to_unix_secs: Option<u64>,
|
||||
pub event_type: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum ProxyNodeMetricsStep {
|
||||
OneMinute,
|
||||
OneHour,
|
||||
}
|
||||
|
||||
impl ProxyNodeMetricsStep {
|
||||
pub fn bucket_size_secs(self) -> u64 {
|
||||
match self {
|
||||
Self::OneMinute => 60,
|
||||
Self::OneHour => 3_600,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_api_value(self) -> &'static str {
|
||||
match self {
|
||||
Self::OneMinute => "1m",
|
||||
Self::OneHour => "1h",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProxyNodeMetricsBucket {
|
||||
pub node_id: String,
|
||||
pub bucket_start_unix_secs: u64,
|
||||
pub samples: i64,
|
||||
pub uptime_samples: i64,
|
||||
pub active_connections_sum: i64,
|
||||
pub active_connections_max: i64,
|
||||
pub heartbeat_rtt_ms_sum: i64,
|
||||
pub heartbeat_rtt_ms_max: i64,
|
||||
pub connect_errors_delta: i64,
|
||||
pub disconnects_delta: i64,
|
||||
pub error_events_delta: i64,
|
||||
pub ws_in_bytes_delta: i64,
|
||||
pub ws_out_bytes_delta: i64,
|
||||
pub ws_in_frames_delta: i64,
|
||||
pub ws_out_frames_delta: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProxyFleetMetricsBucket {
|
||||
pub bucket_start_unix_secs: u64,
|
||||
pub samples: i64,
|
||||
pub uptime_samples: i64,
|
||||
pub active_connections_sum: i64,
|
||||
pub active_connections_max: i64,
|
||||
pub heartbeat_rtt_ms_sum: i64,
|
||||
pub heartbeat_rtt_ms_max: i64,
|
||||
pub connect_errors_delta: i64,
|
||||
pub disconnects_delta: i64,
|
||||
pub error_events_delta: i64,
|
||||
pub ws_in_bytes_delta: i64,
|
||||
pub ws_out_bytes_delta: i64,
|
||||
pub ws_in_frames_delta: i64,
|
||||
pub ws_out_frames_delta: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProxyNodeMetricsCleanupSummary {
|
||||
pub deleted_1m_rows: usize,
|
||||
pub deleted_1h_rows: usize,
|
||||
}
|
||||
|
||||
pub const PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR: &str = "tunnel_err";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct TunnelErrorEventRecord {
|
||||
pub timestamp_unix_secs: u64,
|
||||
pub category: String,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub struct TunnelMetricsCounters {
|
||||
pub connect_errors: u64,
|
||||
pub disconnects: u64,
|
||||
pub error_events_total: u64,
|
||||
pub ws_in_bytes: u64,
|
||||
pub ws_out_bytes: u64,
|
||||
pub ws_in_frames: u64,
|
||||
pub ws_out_frames: u64,
|
||||
pub heartbeat_rtt_last_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct TunnelMetricsSample {
|
||||
pub samples: i64,
|
||||
pub uptime_samples: i64,
|
||||
pub active_connections_sum: i64,
|
||||
pub active_connections_max: i64,
|
||||
pub heartbeat_rtt_ms_sum: i64,
|
||||
pub heartbeat_rtt_ms_max: i64,
|
||||
pub connect_errors_delta: i64,
|
||||
pub disconnects_delta: i64,
|
||||
pub error_events_delta: i64,
|
||||
pub ws_in_bytes_delta: i64,
|
||||
pub ws_out_bytes_delta: i64,
|
||||
pub ws_in_frames_delta: i64,
|
||||
pub ws_out_frames_delta: i64,
|
||||
pub recent_error_events: Vec<TunnelErrorEventRecord>,
|
||||
}
|
||||
|
||||
pub fn bucket_start_unix_secs(timestamp_unix_secs: u64, step: ProxyNodeMetricsStep) -> u64 {
|
||||
let size = step.bucket_size_secs();
|
||||
timestamp_unix_secs / size * size
|
||||
}
|
||||
|
||||
pub fn build_tunnel_metrics_sample(
|
||||
previous_proxy_metadata: Option<&Value>,
|
||||
current_proxy_metadata: Option<&Value>,
|
||||
active_connections: i32,
|
||||
tunnel_connected: bool,
|
||||
) -> Option<TunnelMetricsSample> {
|
||||
let current = extract_tunnel_metrics_counters(current_proxy_metadata)?;
|
||||
let previous = extract_tunnel_metrics_counters(previous_proxy_metadata);
|
||||
let current_recent_errors = extract_recent_tunnel_errors(current_proxy_metadata);
|
||||
|
||||
let connect_errors_delta =
|
||||
counter_delta_u64(previous.map(|v| v.connect_errors), current.connect_errors);
|
||||
let disconnects_delta = counter_delta_u64(previous.map(|v| v.disconnects), current.disconnects);
|
||||
let error_events_delta = counter_delta_u64(
|
||||
previous.map(|v| v.error_events_total),
|
||||
current.error_events_total,
|
||||
);
|
||||
let ws_in_bytes_delta = counter_delta_u64(previous.map(|v| v.ws_in_bytes), current.ws_in_bytes);
|
||||
let ws_out_bytes_delta =
|
||||
counter_delta_u64(previous.map(|v| v.ws_out_bytes), current.ws_out_bytes);
|
||||
let ws_in_frames_delta =
|
||||
counter_delta_u64(previous.map(|v| v.ws_in_frames), current.ws_in_frames);
|
||||
let ws_out_frames_delta =
|
||||
counter_delta_u64(previous.map(|v| v.ws_out_frames), current.ws_out_frames);
|
||||
|
||||
let take_recent = usize::try_from(error_events_delta).unwrap_or(usize::MAX);
|
||||
let recent_error_events = if take_recent == 0 {
|
||||
Vec::new()
|
||||
} else {
|
||||
let capture = take_recent.min(current_recent_errors.len());
|
||||
let from = current_recent_errors.len().saturating_sub(capture);
|
||||
current_recent_errors[from..].to_vec()
|
||||
};
|
||||
|
||||
let active_connections = i64::from(active_connections.max(0));
|
||||
let heartbeat_rtt_last_ms = i64::try_from(current.heartbeat_rtt_last_ms).unwrap_or(i64::MAX);
|
||||
|
||||
Some(TunnelMetricsSample {
|
||||
samples: 1,
|
||||
uptime_samples: if tunnel_connected { 1 } else { 0 },
|
||||
active_connections_sum: active_connections,
|
||||
active_connections_max: active_connections,
|
||||
heartbeat_rtt_ms_sum: heartbeat_rtt_last_ms,
|
||||
heartbeat_rtt_ms_max: heartbeat_rtt_last_ms,
|
||||
connect_errors_delta: i64::try_from(connect_errors_delta).unwrap_or(i64::MAX),
|
||||
disconnects_delta: i64::try_from(disconnects_delta).unwrap_or(i64::MAX),
|
||||
error_events_delta: i64::try_from(error_events_delta).unwrap_or(i64::MAX),
|
||||
ws_in_bytes_delta: i64::try_from(ws_in_bytes_delta).unwrap_or(i64::MAX),
|
||||
ws_out_bytes_delta: i64::try_from(ws_out_bytes_delta).unwrap_or(i64::MAX),
|
||||
ws_in_frames_delta: i64::try_from(ws_in_frames_delta).unwrap_or(i64::MAX),
|
||||
ws_out_frames_delta: i64::try_from(ws_out_frames_delta).unwrap_or(i64::MAX),
|
||||
recent_error_events,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_tunnel_error_event_detail(event: &TunnelErrorEventRecord) -> String {
|
||||
format!("[{}] {}", event.category, event.message)
|
||||
}
|
||||
|
||||
pub fn normalize_proxy_metadata(
|
||||
proxy_metadata: Option<&serde_json::Value>,
|
||||
proxy_version: Option<&str>,
|
||||
@@ -276,6 +454,71 @@ pub fn normalize_proxy_metadata(
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_tunnel_metrics_counters(
|
||||
proxy_metadata: Option<&Value>,
|
||||
) -> Option<TunnelMetricsCounters> {
|
||||
let tunnel_metrics = proxy_metadata
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| metadata.get("tunnel_metrics"))
|
||||
.and_then(Value::as_object)?;
|
||||
|
||||
Some(TunnelMetricsCounters {
|
||||
connect_errors: json_u64(tunnel_metrics.get("connect_errors")).unwrap_or(0),
|
||||
disconnects: json_u64(tunnel_metrics.get("disconnects")).unwrap_or(0),
|
||||
error_events_total: json_u64(tunnel_metrics.get("error_events_total")).unwrap_or(0),
|
||||
ws_in_bytes: json_u64(tunnel_metrics.get("ws_in_bytes")).unwrap_or(0),
|
||||
ws_out_bytes: json_u64(tunnel_metrics.get("ws_out_bytes")).unwrap_or(0),
|
||||
ws_in_frames: json_u64(tunnel_metrics.get("ws_in_frames")).unwrap_or(0),
|
||||
ws_out_frames: json_u64(tunnel_metrics.get("ws_out_frames")).unwrap_or(0),
|
||||
heartbeat_rtt_last_ms: json_u64(tunnel_metrics.get("heartbeat_rtt_last_ms")).unwrap_or(0),
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_recent_tunnel_errors(proxy_metadata: Option<&Value>) -> Vec<TunnelErrorEventRecord> {
|
||||
proxy_metadata
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| metadata.get("recent_tunnel_errors"))
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(|item| {
|
||||
let item = item.as_object()?;
|
||||
Some(TunnelErrorEventRecord {
|
||||
timestamp_unix_secs: json_u64(item.get("timestamp_unix_secs"))
|
||||
.unwrap_or_default(),
|
||||
category: item
|
||||
.get("category")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
message: item
|
||||
.get("message")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("n/a")
|
||||
.to_string(),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn json_u64(value: Option<&Value>) -> Option<u64> {
|
||||
value.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_i64().and_then(|n| (n >= 0).then_some(n as u64)))
|
||||
})
|
||||
}
|
||||
|
||||
fn counter_delta_u64(previous: Option<u64>, current: u64) -> u64 {
|
||||
match previous {
|
||||
Some(previous) if current >= previous => current - previous,
|
||||
Some(_) | None => current,
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_proxy_version_label(value: &str) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.is_empty() {
|
||||
@@ -375,6 +618,42 @@ pub trait ProxyNodeReadRepository: Send + Sync {
|
||||
node_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, crate::DataLayerError>;
|
||||
|
||||
async fn list_proxy_node_events_filtered(
|
||||
&self,
|
||||
node_id: &str,
|
||||
query: &ProxyNodeEventQuery,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, crate::DataLayerError> {
|
||||
let mut items = self.list_proxy_node_events(node_id, query.limit).await?;
|
||||
if let Some(from_unix_secs) = query.from_unix_secs {
|
||||
items.retain(|item| item.created_at_unix_ms.unwrap_or(0) >= from_unix_secs);
|
||||
}
|
||||
if let Some(to_unix_secs) = query.to_unix_secs {
|
||||
items.retain(|item| item.created_at_unix_ms.unwrap_or(u64::MAX) <= to_unix_secs);
|
||||
}
|
||||
if let Some(event_type) = query.event_type.as_deref() {
|
||||
items.retain(|item| item.event_type.eq_ignore_ascii_case(event_type));
|
||||
}
|
||||
items.truncate(query.limit);
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
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>, crate::DataLayerError>;
|
||||
|
||||
async fn list_proxy_fleet_metrics(
|
||||
&self,
|
||||
step: ProxyNodeMetricsStep,
|
||||
from_unix_secs: u64,
|
||||
to_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyFleetMetricsBucket>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -433,6 +712,12 @@ pub trait ProxyNodeWriteRepository: Send + Sync {
|
||||
failed_delta: i64,
|
||||
latency_ms: Option<i64>,
|
||||
) -> Result<(), crate::DataLayerError>;
|
||||
|
||||
async fn cleanup_proxy_node_metrics(
|
||||
&self,
|
||||
retain_1m_from_unix_secs: u64,
|
||||
retain_1h_from_unix_secs: u64,
|
||||
) -> Result<ProxyNodeMetricsCleanupSummary, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -440,9 +725,10 @@ mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
normalize_proxy_node_scheduling_state, proxy_node_accepts_new_tunnels,
|
||||
proxy_reported_version, reconcile_remote_config_after_heartbeat,
|
||||
remote_config_scheduling_state, remote_config_upgrade_target, StoredProxyNode,
|
||||
bucket_start_unix_secs, build_tunnel_metrics_sample, normalize_proxy_node_scheduling_state,
|
||||
proxy_node_accepts_new_tunnels, proxy_reported_version,
|
||||
reconcile_remote_config_after_heartbeat, remote_config_scheduling_state,
|
||||
remote_config_upgrade_target, ProxyNodeMetricsStep, StoredProxyNode,
|
||||
};
|
||||
|
||||
#[test]
|
||||
@@ -527,4 +813,63 @@ mod tests {
|
||||
|
||||
assert!(!proxy_node_accepts_new_tunnels(&node));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_tunnel_metrics_sample_with_reset_safe_counter_deltas() {
|
||||
let previous = json!({
|
||||
"tunnel_metrics": {
|
||||
"connect_errors": 10,
|
||||
"disconnects": 5,
|
||||
"error_events_total": 7,
|
||||
"ws_in_bytes": 1_000,
|
||||
"ws_out_bytes": 2_000,
|
||||
"ws_in_frames": 10,
|
||||
"ws_out_frames": 20,
|
||||
"heartbeat_rtt_last_ms": 30
|
||||
}
|
||||
});
|
||||
let current = json!({
|
||||
"tunnel_metrics": {
|
||||
"connect_errors": 12,
|
||||
"disconnects": 2,
|
||||
"error_events_total": 9,
|
||||
"ws_in_bytes": 1_500,
|
||||
"ws_out_bytes": 100,
|
||||
"ws_in_frames": 11,
|
||||
"ws_out_frames": 3,
|
||||
"heartbeat_rtt_last_ms": 44
|
||||
},
|
||||
"recent_tunnel_errors": [
|
||||
{"timestamp_unix_secs": 100, "category": "older", "message": "old"},
|
||||
{"timestamp_unix_secs": 101, "category": "newer", "message": "new"}
|
||||
]
|
||||
});
|
||||
|
||||
let sample = build_tunnel_metrics_sample(Some(&previous), Some(¤t), 4, true)
|
||||
.expect("sample should build");
|
||||
assert_eq!(sample.samples, 1);
|
||||
assert_eq!(sample.uptime_samples, 1);
|
||||
assert_eq!(sample.active_connections_sum, 4);
|
||||
assert_eq!(sample.heartbeat_rtt_ms_sum, 44);
|
||||
assert_eq!(sample.connect_errors_delta, 2);
|
||||
assert_eq!(sample.disconnects_delta, 2);
|
||||
assert_eq!(sample.error_events_delta, 2);
|
||||
assert_eq!(sample.ws_in_bytes_delta, 500);
|
||||
assert_eq!(sample.ws_out_bytes_delta, 100);
|
||||
assert_eq!(sample.ws_out_frames_delta, 3);
|
||||
assert_eq!(sample.recent_error_events.len(), 2);
|
||||
assert_eq!(sample.recent_error_events[0].category, "older");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_timestamps_to_metric_buckets() {
|
||||
assert_eq!(
|
||||
bucket_start_unix_secs(1_710_000_119, ProxyNodeMetricsStep::OneMinute),
|
||||
1_710_000_060
|
||||
);
|
||||
assert_eq!(
|
||||
bucket_start_unix_secs(1_710_003_999, ProxyNodeMetricsStep::OneHour),
|
||||
1_710_003_600
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,7 +4,7 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
|
||||
use super::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_optional_real, sqlite_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -12,8 +12,8 @@ const QUOTA_COLUMNS: &str = r#"
|
||||
SELECT
|
||||
id AS provider_id,
|
||||
billing_type,
|
||||
monthly_quota_usd,
|
||||
COALESCE(monthly_used_usd, 0) AS monthly_used_usd,
|
||||
CAST(monthly_quota_usd AS REAL) AS monthly_quota_usd,
|
||||
CAST(COALESCE(monthly_used_usd, 0) AS REAL) AS monthly_used_usd,
|
||||
quota_reset_day,
|
||||
quota_last_reset_at AS quota_last_reset_at_unix_secs,
|
||||
quota_expires_at AS quota_expires_at_unix_secs,
|
||||
@@ -77,7 +77,7 @@ impl ProviderQuotaWriteRepository for SqliteProviderQuotaRepository {
|
||||
let rows_affected = sqlx::query(
|
||||
r#"
|
||||
UPDATE providers
|
||||
SET monthly_used_usd = 0,
|
||||
SET monthly_used_usd = 0.0,
|
||||
quota_last_reset_at = ?,
|
||||
updated_at = ?
|
||||
WHERE billing_type = 'monthly_quota'
|
||||
@@ -103,8 +103,8 @@ fn map_row(row: &SqliteRow) -> Result<StoredProviderQuotaSnapshot, DataLayerErro
|
||||
StoredProviderQuotaSnapshot::new(
|
||||
row.try_get("provider_id").map_sql_err()?,
|
||||
row.try_get("billing_type").map_sql_err()?,
|
||||
row.try_get("monthly_quota_usd").map_sql_err()?,
|
||||
row.try_get("monthly_used_usd").map_sql_err()?,
|
||||
sqlite_optional_real(row, "monthly_quota_usd")?,
|
||||
sqlite_real(row, "monthly_used_usd")?,
|
||||
row.try_get("quota_reset_day").map_sql_err()?,
|
||||
row.try_get("quota_last_reset_at_unix_secs").map_sql_err()?,
|
||||
row.try_get("quota_expires_at_unix_secs").map_sql_err()?,
|
||||
@@ -138,6 +138,13 @@ mod tests {
|
||||
.expect("quota should exist");
|
||||
assert_eq!(quota.monthly_used_usd, 5.0);
|
||||
|
||||
let quota = repository
|
||||
.find_by_provider_id("provider-null-used")
|
||||
.await
|
||||
.expect("quota with null usage should load")
|
||||
.expect("quota with null usage should exist");
|
||||
assert_eq!(quota.monthly_used_usd, 0.0);
|
||||
|
||||
let quotas = repository
|
||||
.find_by_provider_ids(&["provider-2".to_string(), "provider-1".to_string()])
|
||||
.await
|
||||
@@ -173,7 +180,8 @@ INSERT INTO providers (
|
||||
)
|
||||
VALUES
|
||||
('provider-1', 'Provider One', 'openai', 'monthly_quota', 20.0, 5.0, 7, 1000, 1, 1, 1),
|
||||
('provider-2', 'Provider Two', 'openai', 'payg', NULL, 1.5, NULL, NULL, 1, 1, 1)
|
||||
('provider-2', 'Provider Two', 'openai', 'payg', NULL, 1.5, NULL, NULL, 1, 1, 1),
|
||||
('provider-null-used', 'Provider Null Used', 'openai', 'payg', NULL, NULL, NULL, NULL, 1, 1, 1)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
|
||||
@@ -2,7 +2,7 @@ use async_trait::async_trait;
|
||||
use sqlx::{sqlite::SqliteRow, Row};
|
||||
|
||||
use super::{SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_optional_real, sqlite_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -35,7 +35,7 @@ SELECT
|
||||
usage_settlement_snapshots.wallet_gift_balance_after,
|
||||
usage_record.wallet_gift_balance_after
|
||||
) AS wallet_gift_balance_after,
|
||||
usage_settlement_snapshots.provider_monthly_used_usd AS provider_monthly_used_usd,
|
||||
CAST(usage_settlement_snapshots.provider_monthly_used_usd AS REAL) AS provider_monthly_used_usd,
|
||||
usage_record.provider_id,
|
||||
COALESCE(usage_settlement_snapshots.finalized_at, usage_record.finalized_at) AS finalized_at_unix_secs
|
||||
FROM "usage" AS usage_record
|
||||
@@ -120,17 +120,16 @@ fn settlement_from_row(row: &SqliteRow) -> Result<StoredUsageSettlement, DataLay
|
||||
request_id: row.try_get("request_id").map_sql_err()?,
|
||||
wallet_id: row.try_get("wallet_id").map_sql_err()?,
|
||||
billing_status: row.try_get("billing_status").map_sql_err()?,
|
||||
wallet_balance_before: row.try_get("wallet_balance_before").map_sql_err()?,
|
||||
wallet_balance_after: row.try_get("wallet_balance_after").map_sql_err()?,
|
||||
wallet_recharge_balance_before: row
|
||||
.try_get("wallet_recharge_balance_before")
|
||||
.map_sql_err()?,
|
||||
wallet_recharge_balance_after: row
|
||||
.try_get("wallet_recharge_balance_after")
|
||||
.map_sql_err()?,
|
||||
wallet_gift_balance_before: row.try_get("wallet_gift_balance_before").map_sql_err()?,
|
||||
wallet_gift_balance_after: row.try_get("wallet_gift_balance_after").map_sql_err()?,
|
||||
provider_monthly_used_usd: row.try_get("provider_monthly_used_usd").map_sql_err()?,
|
||||
wallet_balance_before: sqlite_optional_real(row, "wallet_balance_before")?,
|
||||
wallet_balance_after: sqlite_optional_real(row, "wallet_balance_after")?,
|
||||
wallet_recharge_balance_before: sqlite_optional_real(
|
||||
row,
|
||||
"wallet_recharge_balance_before",
|
||||
)?,
|
||||
wallet_recharge_balance_after: sqlite_optional_real(row, "wallet_recharge_balance_after")?,
|
||||
wallet_gift_balance_before: sqlite_optional_real(row, "wallet_gift_balance_before")?,
|
||||
wallet_gift_balance_after: sqlite_optional_real(row, "wallet_gift_balance_after")?,
|
||||
provider_monthly_used_usd: sqlite_optional_real(row, "provider_monthly_used_usd")?,
|
||||
finalized_at_unix_secs: row
|
||||
.try_get::<Option<i64>, _>("finalized_at_unix_secs")
|
||||
.map_sql_err()?
|
||||
@@ -268,8 +267,8 @@ LIMIT 1
|
||||
|
||||
if let Some(wallet_row) = wallet_row {
|
||||
let wallet_id: String = wallet_row.try_get("id").map_sql_err()?;
|
||||
let before_recharge: f64 = wallet_row.try_get("balance").map_sql_err()?;
|
||||
let before_gift: f64 = wallet_row.try_get("gift_balance").map_sql_err()?;
|
||||
let before_recharge = sqlite_real(&wallet_row, "balance")?;
|
||||
let before_gift = sqlite_real(&wallet_row, "gift_balance")?;
|
||||
let limit_mode: String = wallet_row.try_get("limit_mode").map_sql_err()?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let mut after_recharge = before_recharge;
|
||||
@@ -318,7 +317,7 @@ WHERE id = ?
|
||||
r#"
|
||||
UPDATE providers
|
||||
SET
|
||||
monthly_used_usd = COALESCE(monthly_used_usd, 0) + ?,
|
||||
monthly_used_usd = CAST(COALESCE(monthly_used_usd, 0) AS REAL) + ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
"#,
|
||||
@@ -330,14 +329,15 @@ WHERE id = ?
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
settlement.provider_monthly_used_usd = sqlx::query_scalar::<_, Option<f64>>(
|
||||
"SELECT monthly_used_usd FROM providers WHERE id = ? LIMIT 1",
|
||||
settlement.provider_monthly_used_usd = sqlx::query(
|
||||
"SELECT CAST(monthly_used_usd AS REAL) AS monthly_used_usd FROM providers WHERE id = ? LIMIT 1",
|
||||
)
|
||||
.bind(provider_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.flatten();
|
||||
.map(|row| sqlite_real(&row, "monthly_used_usd"))
|
||||
.transpose()?;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ use super::{
|
||||
InMemoryUsageReadRepository, PendingUsageCleanupSummary, StoredRequestUsageAudit,
|
||||
UpsertUsageRecord, UsageWriteRepository,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_optional_real, sqlite_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -42,11 +42,11 @@ SELECT
|
||||
cache_creation_ephemeral_5m_input_tokens,
|
||||
cache_creation_ephemeral_1h_input_tokens,
|
||||
cache_read_input_tokens,
|
||||
cache_creation_cost_usd,
|
||||
cache_read_cost_usd,
|
||||
output_price_per_1m,
|
||||
total_cost_usd,
|
||||
actual_total_cost_usd,
|
||||
CAST(cache_creation_cost_usd AS REAL) AS cache_creation_cost_usd,
|
||||
CAST(cache_read_cost_usd AS REAL) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS REAL) AS output_price_per_1m,
|
||||
CAST(total_cost_usd AS REAL) AS total_cost_usd,
|
||||
CAST(actual_total_cost_usd AS REAL) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
@@ -291,7 +291,7 @@ impl UsageWriteRepository for SqliteUsageWriteRepository {
|
||||
UPDATE api_keys
|
||||
SET total_requests = 0,
|
||||
total_tokens = 0,
|
||||
total_cost_usd = 0,
|
||||
total_cost_usd = 0.0,
|
||||
last_used_at = NULL
|
||||
"#,
|
||||
)
|
||||
@@ -305,7 +305,7 @@ SELECT
|
||||
api_key_id,
|
||||
COUNT(*) AS total_requests,
|
||||
COALESCE(SUM(total_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(total_cost_usd), 0) AS total_cost_usd,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL) AS total_cost_usd,
|
||||
MAX(updated_at_unix_secs) AS last_used_at
|
||||
FROM "usage"
|
||||
WHERE api_key_id IS NOT NULL AND api_key_id <> ''
|
||||
@@ -329,7 +329,7 @@ WHERE id = ?
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("total_requests").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("total_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("total_cost_usd").map_sql_err()?)
|
||||
.bind(sqlite_real(row, "total_cost_usd")?)
|
||||
.bind(
|
||||
row.try_get::<Option<i64>, _>("last_used_at")
|
||||
.map_sql_err()?,
|
||||
@@ -351,7 +351,7 @@ SET request_count = 0,
|
||||
success_count = 0,
|
||||
error_count = 0,
|
||||
total_tokens = 0,
|
||||
total_cost_usd = 0,
|
||||
total_cost_usd = 0.0,
|
||||
total_response_time_ms = 0,
|
||||
last_used_at = NULL
|
||||
"#,
|
||||
@@ -368,7 +368,7 @@ SELECT
|
||||
status_code,
|
||||
error_message,
|
||||
total_tokens,
|
||||
total_cost_usd,
|
||||
CAST(total_cost_usd AS REAL) AS total_cost_usd,
|
||||
response_time_ms,
|
||||
updated_at_unix_secs
|
||||
FROM "usage"
|
||||
@@ -396,7 +396,7 @@ WHERE provider_api_key_id IS NOT NULL AND provider_api_key_id <> ''
|
||||
entry.error_count += 1;
|
||||
}
|
||||
entry.total_tokens += row.try_get::<i64, _>("total_tokens").map_sql_err()?;
|
||||
entry.total_cost_usd += row.try_get::<f64, _>("total_cost_usd").map_sql_err()?;
|
||||
entry.total_cost_usd += sqlite_real(&row, "total_cost_usd")?;
|
||||
entry.total_response_time_ms += row
|
||||
.try_get::<Option<i64>, _>("response_time_ms")
|
||||
.map_sql_err()?
|
||||
@@ -528,8 +528,8 @@ SET status = 'failed',
|
||||
error_message = ?,
|
||||
billing_status = 'void',
|
||||
finalized_at = ?,
|
||||
total_cost_usd = 0,
|
||||
actual_total_cost_usd = 0
|
||||
total_cost_usd = 0.0,
|
||||
actual_total_cost_usd = 0.0
|
||||
WHERE request_id = ?
|
||||
"#,
|
||||
)
|
||||
@@ -818,8 +818,8 @@ fn map_usage_row(row: &SqliteRow) -> Result<StoredRequestUsageAudit, DataLayerEr
|
||||
row_i32(row, "input_tokens")?,
|
||||
row_i32(row, "output_tokens")?,
|
||||
row_i32(row, "total_tokens")?,
|
||||
row.try_get("total_cost_usd").map_sql_err()?,
|
||||
row.try_get("actual_total_cost_usd").map_sql_err()?,
|
||||
sqlite_real(row, "total_cost_usd")?,
|
||||
sqlite_real(row, "actual_total_cost_usd")?,
|
||||
row_optional_i32(row, "status_code")?,
|
||||
row.try_get("error_message").map_sql_err()?,
|
||||
row.try_get("error_category").map_sql_err()?,
|
||||
@@ -837,9 +837,10 @@ fn map_usage_row(row: &SqliteRow) -> Result<StoredRequestUsageAudit, DataLayerEr
|
||||
audit.cache_creation_ephemeral_1h_input_tokens =
|
||||
row_u64(row, "cache_creation_ephemeral_1h_input_tokens")?;
|
||||
audit.cache_read_input_tokens = row_u64(row, "cache_read_input_tokens")?;
|
||||
audit.cache_creation_cost_usd = row.try_get("cache_creation_cost_usd").map_sql_err()?;
|
||||
audit.cache_read_cost_usd = row.try_get("cache_read_cost_usd").map_sql_err()?;
|
||||
audit.output_price_per_1m = row.try_get("output_price_per_1m").map_sql_err()?;
|
||||
audit.cache_creation_cost_usd =
|
||||
sqlite_optional_real(row, "cache_creation_cost_usd")?.unwrap_or(0.0);
|
||||
audit.cache_read_cost_usd = sqlite_optional_real(row, "cache_read_cost_usd")?.unwrap_or(0.0);
|
||||
audit.output_price_per_1m = sqlite_optional_real(row, "output_price_per_1m")?;
|
||||
audit.request_metadata = row
|
||||
.try_get::<Option<String>, _>("request_metadata")
|
||||
.map_sql_err()?
|
||||
|
||||
@@ -27,7 +27,7 @@ use super::{
|
||||
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, WalletLookupKey, WalletMutationOutcome,
|
||||
WalletReadRepository, WalletWriteRepository,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::driver::sqlite::{sqlite_optional_real, sqlite_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -715,7 +715,7 @@ LIMIT 1
|
||||
tx.commit().await.map_sql_err()?;
|
||||
return Ok(CreateWalletRefundRequestOutcome::WalletMissing);
|
||||
};
|
||||
let wallet_recharge_balance: f64 = get(&wallet_row, "balance")?;
|
||||
let wallet_recharge_balance = sqlite_real(&wallet_row, "balance")?;
|
||||
let wallet_reserved_amount: f64 = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT COALESCE(SUM(amount_usd), 0.0)
|
||||
@@ -779,7 +779,7 @@ WHERE payment_order_id = ?
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let refundable_amount: f64 = get(&order_row, "refundable_amount_usd")?;
|
||||
let refundable_amount = sqlite_real(&order_row, "refundable_amount_usd")?;
|
||||
if input.amount_usd > (refundable_amount - order_reserved_amount) {
|
||||
tx.commit().await.map_sql_err()?;
|
||||
return Ok(
|
||||
@@ -946,7 +946,7 @@ VALUES (?, NULL, ?, ?, ?, ?, ?, ?, 'received', ?, NULL, ?, NULL)
|
||||
let order_no: String = get(&order_row, "order_no")?;
|
||||
let order_wallet_id: String = get(&order_row, "wallet_id")?;
|
||||
let order_payment_method: String = get(&order_row, "payment_method")?;
|
||||
let order_amount_usd: f64 = get(&order_row, "amount_usd")?;
|
||||
let order_amount_usd = sqlite_real(&order_row, "amount_usd")?;
|
||||
let order_status: String = get(&order_row, "status")?;
|
||||
let expires_at_unix_secs: Option<i64> = get(&order_row, "expires_at_unix_secs")?;
|
||||
|
||||
@@ -1070,8 +1070,8 @@ LIMIT 1
|
||||
});
|
||||
}
|
||||
|
||||
let before_recharge: f64 = get(&wallet_row, "balance")?;
|
||||
let before_gift: f64 = get(&wallet_row, "gift_balance")?;
|
||||
let before_recharge = sqlite_real(&wallet_row, "balance")?;
|
||||
let before_gift = sqlite_real(&wallet_row, "gift_balance")?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let after_recharge = before_recharge + order_amount_usd;
|
||||
let after_total = after_recharge + before_gift;
|
||||
@@ -1174,8 +1174,8 @@ WHERE id = ?
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let before_recharge: f64 = get(&row, "balance")?;
|
||||
let before_gift: f64 = get(&row, "gift_balance")?;
|
||||
let before_recharge = sqlite_real(&row, "balance")?;
|
||||
let before_gift = sqlite_real(&row, "gift_balance")?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let mut after_recharge = before_recharge;
|
||||
let mut after_gift = before_gift;
|
||||
@@ -1278,8 +1278,8 @@ VALUES (?, ?, 'adjust', 'adjust_admin', ?, ?, ?, ?, ?, ?, ?, 'admin_action', ?,
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let before_recharge: f64 = get(&wallet_row, "balance")?;
|
||||
let before_gift: f64 = get(&wallet_row, "gift_balance")?;
|
||||
let before_recharge = sqlite_real(&wallet_row, "balance")?;
|
||||
let before_gift = sqlite_real(&wallet_row, "gift_balance")?;
|
||||
let user_id: Option<String> = get(&wallet_row, "user_id")?;
|
||||
let order_id = uuid::Uuid::new_v4().to_string();
|
||||
let gateway_response = json_string(
|
||||
@@ -1416,8 +1416,8 @@ VALUES (?, ?, 'recharge', ?, ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, ?, ?, ?)
|
||||
"wallet not found".to_string(),
|
||||
));
|
||||
};
|
||||
let before_recharge: f64 = get(&wallet_row, "balance")?;
|
||||
let before_gift: f64 = get(&wallet_row, "gift_balance")?;
|
||||
let before_recharge = sqlite_real(&wallet_row, "balance")?;
|
||||
let before_gift = sqlite_real(&wallet_row, "gift_balance")?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let amount_usd = refund.amount_usd;
|
||||
let after_recharge = before_recharge - amount_usd;
|
||||
@@ -1437,7 +1437,7 @@ VALUES (?, ?, 'recharge', ?, ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, ?, ?, ?)
|
||||
"payment order not found".to_string(),
|
||||
));
|
||||
};
|
||||
let refundable_amount: f64 = get(&order_row, "refundable_amount_usd")?;
|
||||
let refundable_amount = sqlite_real(&order_row, "refundable_amount_usd")?;
|
||||
if amount_usd > refundable_amount {
|
||||
tx.commit().await.map_sql_err()?;
|
||||
return Ok(WalletMutationOutcome::Invalid(
|
||||
@@ -1665,8 +1665,8 @@ WHERE id = ? AND wallet_id = ?
|
||||
));
|
||||
};
|
||||
let amount_usd = refund.amount_usd;
|
||||
let before_recharge: f64 = get(&wallet_row, "balance")?;
|
||||
let before_gift: f64 = get(&wallet_row, "gift_balance")?;
|
||||
let before_recharge = sqlite_real(&wallet_row, "balance")?;
|
||||
let before_gift = sqlite_real(&wallet_row, "gift_balance")?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let after_recharge = before_recharge + amount_usd;
|
||||
|
||||
@@ -1919,8 +1919,8 @@ WHERE id = ? AND wallet_id = ?
|
||||
));
|
||||
}
|
||||
|
||||
let before_recharge: f64 = get(&wallet_row, "balance")?;
|
||||
let before_gift: f64 = get(&wallet_row, "gift_balance")?;
|
||||
let before_recharge = sqlite_real(&wallet_row, "balance")?;
|
||||
let before_gift = sqlite_real(&wallet_row, "gift_balance")?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let after_recharge = before_recharge + order.amount_usd;
|
||||
sqlx::query(
|
||||
@@ -2358,7 +2358,7 @@ LIMIT 1
|
||||
let batch_id: String = get(&code_row, "batch_id")?;
|
||||
let batch_name: String = get(&code_row, "batch_name")?;
|
||||
let balance_bucket: String = get(&code_row, "balance_bucket")?;
|
||||
let amount_usd: f64 = get(&code_row, "amount_usd")?;
|
||||
let amount_usd = sqlite_real(&code_row, "amount_usd")?;
|
||||
let credits_recharge_balance = redeem_code_credits_recharge_balance(&balance_bucket);
|
||||
|
||||
let wallet_row = sqlite_wallet_by_user_id(&mut tx, &input.user_id).await?;
|
||||
@@ -2374,7 +2374,10 @@ LIMIT 1
|
||||
};
|
||||
|
||||
let (before_recharge, before_gift) = if let Some(row) = wallet_row.as_ref() {
|
||||
(get(row, "balance")?, get(row, "gift_balance")?)
|
||||
(
|
||||
sqlite_real(row, "balance")?,
|
||||
sqlite_real(row, "gift_balance")?,
|
||||
)
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
@@ -2383,7 +2386,7 @@ INSERT INTO wallets (
|
||||
total_recharged, total_consumed, total_refunded, total_adjusted,
|
||||
created_at, updated_at
|
||||
)
|
||||
VALUES (?, ?, 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, ?, ?)
|
||||
VALUES (?, ?, 0.0, 0.0, 'finite', 'USD', 'active', 0.0, 0.0, 0.0, 0.0, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&wallet_id)
|
||||
@@ -2557,15 +2560,15 @@ fn map_wallet_row(row: &SqliteRow) -> Result<StoredWalletSnapshot, DataLayerErro
|
||||
get(row, "id")?,
|
||||
get(row, "user_id")?,
|
||||
get(row, "api_key_id")?,
|
||||
get(row, "balance")?,
|
||||
get(row, "gift_balance")?,
|
||||
sqlite_real(row, "balance")?,
|
||||
sqlite_real(row, "gift_balance")?,
|
||||
get(row, "limit_mode")?,
|
||||
get(row, "currency")?,
|
||||
get(row, "status")?,
|
||||
get(row, "total_recharged")?,
|
||||
get(row, "total_consumed")?,
|
||||
get(row, "total_refunded")?,
|
||||
get(row, "total_adjusted")?,
|
||||
sqlite_real(row, "total_recharged")?,
|
||||
sqlite_real(row, "total_consumed")?,
|
||||
sqlite_real(row, "total_refunded")?,
|
||||
sqlite_real(row, "total_adjusted")?,
|
||||
get(row, "updated_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
@@ -3159,12 +3162,12 @@ fn map_payment_order_row(row: &SqliteRow) -> Result<StoredAdminPaymentOrder, Dat
|
||||
order_no: get(row, "order_no")?,
|
||||
wallet_id: get(row, "wallet_id")?,
|
||||
user_id: get(row, "user_id")?,
|
||||
amount_usd: get(row, "amount_usd")?,
|
||||
pay_amount: get(row, "pay_amount")?,
|
||||
amount_usd: sqlite_real(row, "amount_usd")?,
|
||||
pay_amount: sqlite_optional_real(row, "pay_amount")?,
|
||||
pay_currency: get(row, "pay_currency")?,
|
||||
exchange_rate: get(row, "exchange_rate")?,
|
||||
refunded_amount_usd: get(row, "refunded_amount_usd")?,
|
||||
refundable_amount_usd: get(row, "refundable_amount_usd")?,
|
||||
exchange_rate: sqlite_optional_real(row, "exchange_rate")?,
|
||||
refunded_amount_usd: sqlite_real(row, "refunded_amount_usd")?,
|
||||
refundable_amount_usd: sqlite_real(row, "refundable_amount_usd")?,
|
||||
payment_method: get(row, "payment_method")?,
|
||||
gateway_order_id: get(row, "gateway_order_id")?,
|
||||
gateway_response: optional_json(
|
||||
@@ -3223,13 +3226,13 @@ fn map_wallet_transaction_row(
|
||||
wallet_id: get(row, "wallet_id")?,
|
||||
category: get(row, "category")?,
|
||||
reason_code: get(row, "reason_code")?,
|
||||
amount: get(row, "amount")?,
|
||||
balance_before: get(row, "balance_before")?,
|
||||
balance_after: get(row, "balance_after")?,
|
||||
recharge_balance_before: get(row, "recharge_balance_before")?,
|
||||
recharge_balance_after: get(row, "recharge_balance_after")?,
|
||||
gift_balance_before: get(row, "gift_balance_before")?,
|
||||
gift_balance_after: get(row, "gift_balance_after")?,
|
||||
amount: sqlite_real(row, "amount")?,
|
||||
balance_before: sqlite_real(row, "balance_before")?,
|
||||
balance_after: sqlite_real(row, "balance_after")?,
|
||||
recharge_balance_before: sqlite_real(row, "recharge_balance_before")?,
|
||||
recharge_balance_after: sqlite_real(row, "recharge_balance_after")?,
|
||||
gift_balance_before: sqlite_real(row, "gift_balance_before")?,
|
||||
gift_balance_after: sqlite_real(row, "gift_balance_after")?,
|
||||
link_type: get(row, "link_type")?,
|
||||
link_id: get(row, "link_id")?,
|
||||
operator_id: get(row, "operator_id")?,
|
||||
@@ -3253,7 +3256,7 @@ fn map_refund_row(row: &SqliteRow) -> Result<StoredAdminWalletRefund, DataLayerE
|
||||
source_type: get(row, "source_type")?,
|
||||
source_id: get(row, "source_id")?,
|
||||
refund_mode: get(row, "refund_mode")?,
|
||||
amount_usd: get(row, "amount_usd")?,
|
||||
amount_usd: sqlite_real(row, "amount_usd")?,
|
||||
status: get(row, "status")?,
|
||||
reason: get(row, "reason")?,
|
||||
failure_reason: get(row, "failure_reason")?,
|
||||
@@ -3287,7 +3290,7 @@ fn map_redeem_batch_row(row: &SqliteRow) -> Result<StoredAdminRedeemCodeBatch, D
|
||||
Ok(StoredAdminRedeemCodeBatch {
|
||||
id: get(row, "id")?,
|
||||
name: get(row, "name")?,
|
||||
amount_usd: get(row, "amount_usd")?,
|
||||
amount_usd: sqlite_real(row, "amount_usd")?,
|
||||
currency: get(row, "currency")?,
|
||||
balance_bucket: get(row, "balance_bucket")?,
|
||||
total_count: nonnegative_u64(get(row, "total_count")?, "redeem_code_batches.total_count")?,
|
||||
@@ -3352,7 +3355,7 @@ fn map_daily_usage_row(row: &SqliteRow) -> Result<StoredWalletDailyUsageLedger,
|
||||
id: get(row, "id")?,
|
||||
billing_date: get(row, "billing_date")?,
|
||||
billing_timezone: get(row, "billing_timezone")?,
|
||||
total_cost_usd: get(row, "total_cost_usd")?,
|
||||
total_cost_usd: sqlite_real(row, "total_cost_usd")?,
|
||||
total_requests: nonnegative_u64(
|
||||
get(row, "total_requests")?,
|
||||
"wallet_daily_usage_ledgers.total_requests",
|
||||
|
||||
@@ -33,7 +33,6 @@ pub fn apply_http_client_config(
|
||||
if let Some(user_agent) = &config.user_agent {
|
||||
builder = builder.user_agent(user_agent.clone());
|
||||
}
|
||||
|
||||
builder
|
||||
}
|
||||
|
||||
@@ -46,6 +45,14 @@ pub fn build_http_client_with_headers(
|
||||
default_headers: HeaderMap,
|
||||
) -> Result<reqwest::Client, reqwest::Error> {
|
||||
let mut builder = apply_http_client_config(reqwest::Client::builder(), config);
|
||||
if let Some(proxy_url) = config
|
||||
.proxy_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
builder = builder.proxy(reqwest::Proxy::all(proxy_url)?);
|
||||
}
|
||||
if !default_headers.is_empty() {
|
||||
builder = builder.default_headers(default_headers);
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ pub struct HttpClientConfig {
|
||||
pub http2_adaptive_window: bool,
|
||||
pub use_rustls_tls: bool,
|
||||
pub user_agent: Option<String>,
|
||||
pub proxy_url: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for HttpClientConfig {
|
||||
@@ -23,6 +24,7 @@ impl Default for HttpClientConfig {
|
||||
http2_adaptive_window: false,
|
||||
use_rustls_tls: true,
|
||||
user_agent: None,
|
||||
proxy_url: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -215,31 +215,12 @@
|
||||
</span>
|
||||
</div>
|
||||
<div
|
||||
v-if="currentAttemptSchedulerInfo"
|
||||
v-if="currentAttemptRequestPathDisplay"
|
||||
class="info-item"
|
||||
>
|
||||
<span class="info-label">调度顺位</span>
|
||||
<span class="info-value info-value-stacked">
|
||||
<code class="format-code">
|
||||
全局 {{ currentAttemptSchedulerInfo.globalPriorityLabel }}
|
||||
/ Provider {{ currentAttemptSchedulerInfo.providerPriorityLabel }}
|
||||
/ Key {{ currentAttemptSchedulerInfo.keyPriorityLabel }}
|
||||
</code>
|
||||
<span class="text-xs text-muted-foreground">
|
||||
{{ currentAttemptSchedulerInfo.hint }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div
|
||||
v-if="currentAttemptRankingInfo"
|
||||
class="info-item"
|
||||
>
|
||||
<span class="info-label">排序原因</span>
|
||||
<span class="info-value info-value-stacked">
|
||||
<code class="format-code">{{ currentAttemptRankingInfo.summary }}</code>
|
||||
<span class="text-xs text-muted-foreground">
|
||||
{{ currentAttemptRankingInfo.hint }}
|
||||
</span>
|
||||
<span class="info-label">请求路径</span>
|
||||
<span class="info-value">
|
||||
<code class="format-code request-path-code">{{ currentAttemptRequestPathDisplay }}</code>
|
||||
</span>
|
||||
</div>
|
||||
<div
|
||||
@@ -269,19 +250,7 @@
|
||||
<span class="info-value info-value-stacked">
|
||||
<code class="format-code">{{ currentAttemptKeyFormatsDisplay }}</code>
|
||||
<span class="text-xs text-muted-foreground">
|
||||
同一 Key 的不同 endpoint 会分别参与候选与转换判定
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div
|
||||
v-if="currentAttemptConversionInfo"
|
||||
class="info-item"
|
||||
>
|
||||
<span class="info-label">转换策略</span>
|
||||
<span class="info-value info-value-stacked">
|
||||
<code class="format-code">{{ currentAttemptConversionInfo.summary }}</code>
|
||||
<span class="text-xs text-muted-foreground">
|
||||
{{ currentAttemptConversionInfo.hint }}
|
||||
Key 声明的可用 endpoint 格式
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
@@ -1181,59 +1150,11 @@ const extractStringList = (value: unknown): string[] => {
|
||||
return []
|
||||
}
|
||||
|
||||
const normalizePriorityNumber = (value: unknown): number | null => {
|
||||
if (typeof value === 'number' && Number.isFinite(value)) {
|
||||
return Math.trunc(value)
|
||||
}
|
||||
if (typeof value === 'string' && value.trim()) {
|
||||
const parsed = Number(value)
|
||||
if (Number.isFinite(parsed)) {
|
||||
return Math.trunc(parsed)
|
||||
}
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
const resolveClientApiFormat = (attempt: CandidateRecord): string => {
|
||||
const extra = (
|
||||
attempt.extra_data && typeof attempt.extra_data === 'object' && !Array.isArray(attempt.extra_data)
|
||||
? attempt.extra_data
|
||||
: {}
|
||||
) as Record<string, unknown>
|
||||
const fromExtra = typeof extra.client_api_format === 'string' ? extra.client_api_format.trim() : ''
|
||||
if (fromExtra) return fromExtra
|
||||
if (typeof props.requestApiFormat === 'string' && props.requestApiFormat.trim()) {
|
||||
return props.requestApiFormat.trim()
|
||||
}
|
||||
return ''
|
||||
}
|
||||
|
||||
const resolveProviderApiFormat = (attempt: CandidateRecord): string => {
|
||||
const extra = (
|
||||
attempt.extra_data && typeof attempt.extra_data === 'object' && !Array.isArray(attempt.extra_data)
|
||||
? attempt.extra_data
|
||||
: {}
|
||||
) as Record<string, unknown>
|
||||
const fromExtra = typeof extra.provider_api_format === 'string' ? extra.provider_api_format.trim() : ''
|
||||
if (fromExtra) return fromExtra
|
||||
if (typeof attempt.endpoint_name === 'string' && attempt.endpoint_name.trim()) {
|
||||
return attempt.endpoint_name.trim()
|
||||
}
|
||||
return ''
|
||||
}
|
||||
|
||||
const resolveTransportDiagnostics = (attempt: CandidateRecord): Record<string, unknown> | null => {
|
||||
const extra = extractObject(attempt.extra_data)
|
||||
return extractObject(extra?.transport_diagnostics)
|
||||
}
|
||||
|
||||
const resolveEndpointFormatAcceptanceConfig = (attempt: CandidateRecord): Record<string, unknown> | null => {
|
||||
const fromAttempt = extractObject(attempt.endpoint_format_acceptance_config)
|
||||
if (fromAttempt) return fromAttempt
|
||||
const transport = resolveTransportDiagnostics(attempt)
|
||||
return extractObject(transport?.endpoint_format_acceptance_config)
|
||||
}
|
||||
|
||||
const currentAttemptFormatDisplay = computed(() => {
|
||||
const attempt = currentAttempt.value
|
||||
if (!attempt) return ''
|
||||
@@ -1261,120 +1182,52 @@ const currentAttemptFormatDisplay = computed(() => {
|
||||
return providerText || requestText
|
||||
})
|
||||
|
||||
const currentAttemptSchedulerInfo = computed<{
|
||||
globalPriorityLabel: string
|
||||
providerPriorityLabel: string
|
||||
keyPriorityLabel: string
|
||||
hint: string
|
||||
} | null>(() => {
|
||||
const attempt = currentAttempt.value
|
||||
if (!attempt) return null
|
||||
const normalizeQueryString = (value: string): string => {
|
||||
const trimmed = value.trim()
|
||||
if (!trimmed) return ''
|
||||
return trimmed.startsWith('?') ? trimmed.slice(1) : trimmed
|
||||
}
|
||||
|
||||
const clientApiFormat = resolveClientApiFormat(attempt)
|
||||
const providerApiFormat = resolveProviderApiFormat(attempt)
|
||||
const providerPriority = normalizePriorityNumber(attempt.provider_priority)
|
||||
const keyInternalPriority = normalizePriorityNumber(attempt.key_internal_priority)
|
||||
const resolveRequestPathFromObject = (value: unknown): string => {
|
||||
const object = extractObject(value)
|
||||
if (!object) return ''
|
||||
|
||||
let globalPriority: number | null = null
|
||||
const globalPriorityMap = attempt.key_global_priority_by_format
|
||||
if (globalPriorityMap && typeof globalPriorityMap === 'object' && !Array.isArray(globalPriorityMap) && clientApiFormat) {
|
||||
const match = Object.entries(globalPriorityMap).find(([format]) => (
|
||||
normalizeFormatSignature(format) === normalizeFormatSignature(clientApiFormat)
|
||||
))
|
||||
globalPriority = match ? normalizePriorityNumber(match[1]) : null
|
||||
}
|
||||
|
||||
const isCrossFormat = Boolean(
|
||||
clientApiFormat &&
|
||||
providerApiFormat &&
|
||||
normalizeFormatSignature(clientApiFormat) !== normalizeFormatSignature(providerApiFormat),
|
||||
const pathWithQuery = (
|
||||
readStringField(object, 'request_path_and_query')
|
||||
|| readStringField(object, 'public_path_and_query')
|
||||
|| readStringField(object, 'path_and_query')
|
||||
|| readStringField(object, 'request_uri')
|
||||
|| readStringField(object, 'public_uri')
|
||||
)
|
||||
const keepPriorityOnConversion = attempt.provider_keep_priority_on_conversion === true
|
||||
if (pathWithQuery) return pathWithQuery
|
||||
|
||||
let hint = '顺位展示 Provider / Key 优先级;不同 endpoint 独立参与候选'
|
||||
if (globalPriority !== null) {
|
||||
hint = `当前格式 ${formatApiFormat(clientApiFormat)} 先看全局 Key 优先级;不同 endpoint 单独判定`
|
||||
}
|
||||
if (isCrossFormat) {
|
||||
hint = keepPriorityOnConversion
|
||||
? '跨格式候选已开启保持优先级;同一 Key 的不同 endpoint 独立参与'
|
||||
: '跨格式候选默认排在同格式候选之后;同一 Key 的不同 endpoint 独立参与'
|
||||
}
|
||||
const path = (
|
||||
readStringField(object, 'request_path')
|
||||
|| readStringField(object, 'public_path')
|
||||
|| readStringField(object, 'path')
|
||||
)
|
||||
if (!path) return ''
|
||||
|
||||
return {
|
||||
globalPriorityLabel: globalPriority !== null ? String(globalPriority) : '-',
|
||||
providerPriorityLabel: providerPriority !== null ? String(providerPriority) : '-',
|
||||
keyPriorityLabel: keyInternalPriority !== null ? String(keyInternalPriority) : '-',
|
||||
hint,
|
||||
}
|
||||
})
|
||||
|
||||
const normalizeMetadataText = (value: unknown): string => {
|
||||
return typeof value === 'string' ? value.trim() : ''
|
||||
const query = normalizeQueryString(
|
||||
readStringField(object, 'request_query_string')
|
||||
|| readStringField(object, 'public_query_string')
|
||||
|| readStringField(object, 'query_string')
|
||||
|| readStringField(object, 'query')
|
||||
|| '',
|
||||
)
|
||||
if (!query || path.includes('?')) return path
|
||||
return `${path}?${query}`
|
||||
}
|
||||
|
||||
const formatRankingModeLabel = (value: string): string => {
|
||||
const normalized = value.replace(/([a-z])([A-Z])/g, '$1_$2').toLowerCase()
|
||||
const labels: Record<string, string> = {
|
||||
fixed_order: '固定顺序',
|
||||
cache_affinity: '亲和性优先',
|
||||
load_balance: '负载均衡',
|
||||
}
|
||||
return labels[normalized] || value
|
||||
}
|
||||
|
||||
const formatPriorityModeLabel = (value: string): string => {
|
||||
const normalized = value.replace(/([a-z])([A-Z])/g, '$1_$2').toLowerCase()
|
||||
const labels: Record<string, string> = {
|
||||
provider: 'Provider 优先级',
|
||||
global_key: '全局 Key 优先级',
|
||||
}
|
||||
return labels[normalized] || value
|
||||
}
|
||||
|
||||
const formatRankingReasonLabel = (value: string): string => {
|
||||
const labels: Record<string, string> = {
|
||||
cached_affinity: '缓存亲和性命中',
|
||||
local_tunnel: '本地隧道优先',
|
||||
cross_format: '跨格式降级',
|
||||
}
|
||||
return labels[value] || value
|
||||
}
|
||||
|
||||
const currentAttemptRankingInfo = computed<{
|
||||
summary: string
|
||||
hint: string
|
||||
} | null>(() => {
|
||||
const currentAttemptRequestPathDisplay = computed(() => {
|
||||
const attempt = currentAttempt.value
|
||||
if (!attempt) return null
|
||||
const ranking = extractObject(attempt.ranking)
|
||||
const extra = extractObject(attempt.extra_data)
|
||||
if (!ranking && !extra) return null
|
||||
const fromAttempt = resolveRequestPathFromObject(attempt?.extra_data)
|
||||
if (fromAttempt) return fromAttempt
|
||||
|
||||
const rankingMode = normalizeMetadataText(ranking?.mode ?? extra?.ranking_mode)
|
||||
const priorityMode = normalizeMetadataText(ranking?.priority_mode ?? extra?.priority_mode)
|
||||
const promotedBy = normalizeMetadataText(ranking?.promoted_by ?? extra?.promoted_by)
|
||||
const demotedBy = normalizeMetadataText(ranking?.demoted_by ?? extra?.demoted_by)
|
||||
const rankingIndex = normalizePriorityNumber(ranking?.index ?? extra?.ranking_index)
|
||||
const prioritySlot = normalizePriorityNumber(ranking?.priority_slot ?? extra?.priority_slot)
|
||||
if (!rankingMode && !priorityMode && !promotedBy && !demotedBy && rankingIndex === null && prioritySlot === null) {
|
||||
return null
|
||||
}
|
||||
const fromRequestMetadata = resolveRequestPathFromObject(props.requestMetadata)
|
||||
if (fromRequestMetadata) return fromRequestMetadata
|
||||
|
||||
const summaryParts: string[] = []
|
||||
if (rankingMode) summaryParts.push(formatRankingModeLabel(rankingMode))
|
||||
if (promotedBy) summaryParts.push(formatRankingReasonLabel(promotedBy))
|
||||
if (demotedBy) summaryParts.push(formatRankingReasonLabel(demotedBy))
|
||||
|
||||
const hintParts: string[] = []
|
||||
if (rankingIndex !== null) hintParts.push(`排序 #${rankingIndex + 1}`)
|
||||
if (priorityMode) hintParts.push(formatPriorityModeLabel(priorityMode))
|
||||
if (prioritySlot !== null) hintParts.push(`槽位 ${prioritySlot}`)
|
||||
|
||||
return {
|
||||
summary: summaryParts.length > 0 ? summaryParts.join(' / ') : '排序元数据',
|
||||
hint: hintParts.length > 0 ? hintParts.join(' · ') : '候选排序由 scheduler ranking engine 生成',
|
||||
}
|
||||
return ''
|
||||
})
|
||||
|
||||
const currentAttemptKeyFormatsDisplay = computed(() => {
|
||||
@@ -1388,66 +1241,6 @@ const currentAttemptKeyFormatsDisplay = computed(() => {
|
||||
.map(format => formatApiFormat(format))
|
||||
.join(' / ')
|
||||
})
|
||||
|
||||
const currentAttemptConversionInfo = computed<{
|
||||
summary: string
|
||||
hint: string
|
||||
} | null>(() => {
|
||||
const attempt = currentAttempt.value
|
||||
if (!attempt) return null
|
||||
|
||||
const clientApiFormat = resolveClientApiFormat(attempt)
|
||||
const providerApiFormat = resolveProviderApiFormat(attempt)
|
||||
if (!clientApiFormat || !providerApiFormat) return null
|
||||
|
||||
const isCrossFormat = normalizeFormatSignature(clientApiFormat) !== normalizeFormatSignature(providerApiFormat)
|
||||
if (!isCrossFormat) {
|
||||
return {
|
||||
summary: '同格式直连',
|
||||
hint: '当前候选直接命中 endpoint 原生格式',
|
||||
}
|
||||
}
|
||||
|
||||
const transportDiagnostics = resolveTransportDiagnostics(attempt)
|
||||
const providerEnabled = attempt.provider_enable_format_conversion === true
|
||||
|| transportDiagnostics?.provider_enable_format_conversion === true
|
||||
const endpointConfig = resolveEndpointFormatAcceptanceConfig(attempt)
|
||||
const endpointRuleEnabled = endpointConfig
|
||||
? endpointConfig.enabled !== false
|
||||
: false
|
||||
const acceptFormats = extractStringList(endpointConfig?.accept_formats)
|
||||
const rejectFormats = extractStringList(endpointConfig?.reject_formats)
|
||||
const normalizedClientFormat = normalizeFormatSignature(clientApiFormat)
|
||||
const endpointAcceptsClient = acceptFormats.some(
|
||||
format => normalizeFormatSignature(format) === normalizedClientFormat,
|
||||
)
|
||||
const endpointRejectsClient = rejectFormats.some(
|
||||
format => normalizeFormatSignature(format) === normalizedClientFormat,
|
||||
)
|
||||
|
||||
let summary = '未开启格式转换'
|
||||
if (providerEnabled && endpointRuleEnabled) {
|
||||
summary = 'Provider 总开关 + Endpoint 规则'
|
||||
} else if (providerEnabled) {
|
||||
summary = 'Provider 总格式转换'
|
||||
} else if (endpointRuleEnabled) {
|
||||
summary = 'Endpoint 独立格式转换'
|
||||
}
|
||||
|
||||
let hint = '同一 Key 的不同 endpoint 会分别判定格式转换'
|
||||
if (endpointRejectsClient) {
|
||||
hint = `当前 endpoint 明确拒绝 ${formatApiFormat(clientApiFormat)}`
|
||||
} else if (endpointAcceptsClient) {
|
||||
hint = `当前 endpoint 明确接受 ${formatApiFormat(clientApiFormat)}`
|
||||
} else if (providerEnabled) {
|
||||
hint = '当前跨格式由 Provider 总开关放行'
|
||||
} else if (attempt.skip_reason === 'format_conversion_disabled') {
|
||||
hint = 'Provider 总开关关闭,且当前 endpoint 未单独放行'
|
||||
}
|
||||
|
||||
return { summary, hint }
|
||||
})
|
||||
|
||||
const currentAttemptSkipReasonDisplay = computed(() => {
|
||||
const attempt = currentAttempt.value
|
||||
if (!attempt?.skip_reason) return ''
|
||||
@@ -1870,8 +1663,16 @@ const getStatusColorClass = (status: string) => {
|
||||
// 展示状态:进行中态优先(包括 started 但未 finished 的中间态),再按 HTTP 状态码兜底
|
||||
function getDisplayStatus(attempt: CandidateRecord | null | undefined): string {
|
||||
if (!attempt) return 'available'
|
||||
const code = attempt.status_code
|
||||
const isTerminalSuccessCode = typeof code === 'number' && code >= 200 && code < 300
|
||||
|
||||
if (attempt.status === 'success') {
|
||||
if (typeof code === 'number' && !isTerminalSuccessCode) {
|
||||
return 'failed'
|
||||
}
|
||||
return 'success'
|
||||
}
|
||||
if (
|
||||
attempt.status === 'success' ||
|
||||
attempt.status === 'failed' ||
|
||||
attempt.status === 'cancelled' ||
|
||||
attempt.status === 'skipped' ||
|
||||
@@ -1890,10 +1691,9 @@ function getDisplayStatus(attempt: CandidateRecord | null | undefined): string {
|
||||
if (isExplicitPending || isImplicitPending) {
|
||||
return 'pending'
|
||||
}
|
||||
const code = attempt.status_code
|
||||
if (typeof code === 'number') {
|
||||
if (code >= 200 && code < 300) return 'success'
|
||||
if (code >= 400) return 'failed'
|
||||
if (isTerminalSuccessCode) return 'success'
|
||||
if (code >= 300) return 'failed'
|
||||
}
|
||||
return attempt.status
|
||||
}
|
||||
|
||||
@@ -99,12 +99,16 @@ function buildTrace(candidates: CandidateRecord[]): RequestTrace {
|
||||
}
|
||||
}
|
||||
|
||||
function mountTimeline(traceData: RequestTrace) {
|
||||
function mountTimeline(
|
||||
traceData: RequestTrace,
|
||||
extraProps: Record<string, unknown> = {},
|
||||
) {
|
||||
const root = document.createElement('div')
|
||||
document.body.appendChild(root)
|
||||
const app = createApp(HorizontalRequestTimeline, {
|
||||
requestId: traceData.request_id,
|
||||
traceData,
|
||||
...extraProps,
|
||||
})
|
||||
app.mount(root)
|
||||
mountedApps.push({ app, root })
|
||||
@@ -269,4 +273,52 @@ describe('HorizontalRequestTimeline', () => {
|
||||
expect(nodeDots[0].classList.contains('status-success')).toBe(false)
|
||||
expect(nodeDots[1].classList.contains('status-success')).toBe(true)
|
||||
})
|
||||
|
||||
it('treats 3xx terminal responses as failed for node display', async () => {
|
||||
const trace = buildTrace([
|
||||
buildCandidate({
|
||||
id: 'cand-redirect',
|
||||
provider_id: 'provider-redirect',
|
||||
provider_name: 'Provider Redirect',
|
||||
key_id: 'key-redirect',
|
||||
key_name: 'Redirect Key',
|
||||
candidate_index: 0,
|
||||
status: 'success',
|
||||
status_code: 302,
|
||||
}),
|
||||
])
|
||||
|
||||
const root = mountTimeline(trace)
|
||||
await nextTick()
|
||||
|
||||
const nodeDot = root.querySelector<HTMLElement>('.node-dot')
|
||||
expect(nodeDot?.classList.contains('status-failed')).toBe(true)
|
||||
expect(nodeDot?.classList.contains('status-success')).toBe(false)
|
||||
})
|
||||
|
||||
it('shows request path from request metadata', async () => {
|
||||
const trace = buildTrace([
|
||||
buildCandidate({
|
||||
id: 'cand-path',
|
||||
provider_id: 'provider-path',
|
||||
provider_name: 'Provider Path',
|
||||
key_id: 'key-path',
|
||||
key_name: 'Path Key',
|
||||
candidate_index: 0,
|
||||
status: 'failed',
|
||||
}),
|
||||
])
|
||||
|
||||
const root = mountTimeline(trace, {
|
||||
requestMetadata: {
|
||||
request_path: '/v1beta/models/gemini-2.5-pro:generateContent',
|
||||
request_query_string: 'alt=sse',
|
||||
},
|
||||
})
|
||||
await nextTick()
|
||||
|
||||
expect(root.textContent).toContain('请求路径')
|
||||
const requestPathCode = root.querySelector<HTMLElement>('.request-path-code')
|
||||
expect(requestPathCode?.textContent).toContain('/v1beta/models/gemini-2.5-pro:generateContent?alt=sse')
|
||||
})
|
||||
})
|
||||
|
||||
@@ -106,6 +106,14 @@ describe('usage status helpers', () => {
|
||||
})).toBe('failed')
|
||||
})
|
||||
|
||||
it('downgrades terminal success to failed when status code is 3xx', () => {
|
||||
expect(resolveTimelineFinalStatus({
|
||||
traceFinalStatus: 'success',
|
||||
requestStatus: 'completed',
|
||||
statusCode: 302,
|
||||
})).toBe('failed')
|
||||
})
|
||||
|
||||
it('falls back to request lifecycle status when status code and trace are missing', () => {
|
||||
expect(resolveTimelineFinalStatus({
|
||||
requestStatus: 'failed',
|
||||
@@ -189,6 +197,9 @@ describe('usage status helpers', () => {
|
||||
expect(resolveTimelineFinalStatus({
|
||||
statusCode: 200,
|
||||
})).toBe('success')
|
||||
expect(resolveTimelineFinalStatus({
|
||||
statusCode: 302,
|
||||
})).toBe('failed')
|
||||
expect(resolveTimelineFinalStatus({
|
||||
statusCode: 503,
|
||||
})).toBe('failed')
|
||||
|
||||
@@ -157,7 +157,7 @@ function hasTerminalSuccessStatusCode(
|
||||
): boolean {
|
||||
return typeof record.status_code === 'number' &&
|
||||
record.status_code >= 200 &&
|
||||
record.status_code < 400
|
||||
record.status_code < 300
|
||||
}
|
||||
|
||||
export function isUsageRecordFailed(
|
||||
@@ -272,18 +272,28 @@ export function resolveTimelineFinalStatus(params: {
|
||||
requestStatus?: RequestStatusLike
|
||||
statusCode?: number
|
||||
}): TimelineFinalStatus {
|
||||
const hasTerminalSuccessStatusCode = typeof params.statusCode === 'number'
|
||||
? params.statusCode >= 200 && params.statusCode < 300
|
||||
: undefined
|
||||
|
||||
const requestStatus = mapRequestStatusToTimelineStatus(params.requestStatus)
|
||||
if (requestStatus === 'success' || requestStatus === 'failed' || requestStatus === 'cancelled') {
|
||||
if (requestStatus === 'success' && hasTerminalSuccessStatusCode === false) {
|
||||
return 'failed'
|
||||
}
|
||||
return requestStatus
|
||||
}
|
||||
|
||||
const traceStatus = normalizeTimelineFinalStatus(params.traceFinalStatus)
|
||||
if (traceStatus === 'success' || traceStatus === 'failed' || traceStatus === 'cancelled') {
|
||||
if (traceStatus === 'success' && hasTerminalSuccessStatusCode === false) {
|
||||
return 'failed'
|
||||
}
|
||||
return traceStatus
|
||||
}
|
||||
|
||||
if (typeof params.statusCode === 'number') {
|
||||
return params.statusCode >= 200 && params.statusCode < 400 ? 'success' : 'failed'
|
||||
if (hasTerminalSuccessStatusCode !== undefined) {
|
||||
return hasTerminalSuccessStatusCode ? 'success' : 'failed'
|
||||
}
|
||||
|
||||
if (params.hasPendingCandidates) {
|
||||
|
||||
Reference in New Issue
Block a user