feat(proxy): record tunnel stability metrics

This commit is contained in:
fawney19
2026-05-08 22:03:17 +08:00
parent 9a84a6ff6c
commit a703acd1fe
81 changed files with 7229 additions and 1110 deletions

View File

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

View File

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

View File

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

View File

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

View File

@@ -7,6 +7,7 @@ use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
AiStreamAttempt, AiSyncAttempt,
};
@@ -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) => {

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -20,6 +20,8 @@ mod pending_cleanup;
mod pool_quota_probe;
#[path = "runtime/provider_checkin.rs"]
mod provider_checkin;
#[path = "runtime/proxy_node_metrics_cleanup.rs"]
mod proxy_node_metrics_cleanup;
#[path = "runtime/proxy_node_staleness.rs"]
mod proxy_node_staleness;
#[path = "runtime/proxy_upgrade_rollout.rs"]
@@ -59,6 +61,7 @@ pub(crate) use pool_quota_probe::{
PoolQuotaProbeWorkerConfig,
};
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
use proxy_node_metrics_cleanup::*;
use proxy_node_staleness::*;
use proxy_upgrade_rollout::*;
pub(crate) use proxy_upgrade_rollout::{
@@ -90,6 +93,8 @@ const AUDIT_LOG_CLEANUP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
const GEMINI_FILE_MAPPING_CLEANUP_INTERVAL: Duration = Duration::from_secs(60 * 60);
const PENDING_CLEANUP_INTERVAL: Duration = Duration::from_secs(5 * 60);
const PROXY_NODE_STALE_SWEEP_INTERVAL: Duration = Duration::from_secs(5);
const PROXY_NODE_METRICS_CLEANUP_HOUR: u32 = 2;
const PROXY_NODE_METRICS_CLEANUP_MINUTE: u32 = 10;
const PROXY_UPGRADE_ROLLOUT_INTERVAL: Duration = Duration::from_secs(15);
const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 15;
const PROXY_NODE_STALE_MISSED_HEARTBEATS: u64 = 3;

View File

@@ -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
}

View File

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

View File

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

View File

@@ -4,8 +4,10 @@ use std::sync::Mutex as StdMutex;
use std::time::Duration;
use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode, StoredProxyNodeEvent,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsStep, ProxyNodeTrafficMutation,
ProxyNodeTunnelStatusMutation, StoredProxyFleetMetricsBucket, StoredProxyNode,
StoredProxyNodeEvent, StoredProxyNodeMetricsBucket,
};
use aether_http::{build_http_client, HttpClientConfig};
use aether_runtime::{
@@ -46,6 +48,7 @@ use crate::maintenance::spawn_pending_cleanup_worker;
use crate::maintenance::spawn_pool_monitor_worker;
use crate::maintenance::spawn_pool_quota_probe_worker;
use crate::maintenance::spawn_provider_checkin_worker;
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
use crate::maintenance::spawn_proxy_upgrade_rollout_worker;
use crate::maintenance::spawn_request_candidate_cleanup_worker;
@@ -593,6 +596,44 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_node_events_filtered(
&self,
node_id: &str,
query: &ProxyNodeEventQuery,
) -> Result<Vec<StoredProxyNodeEvent>, GatewayError> {
self.data
.list_proxy_node_events_filtered(node_id, query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_node_metrics(
&self,
node_id: &str,
step: ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyNodeMetricsBucket>, GatewayError> {
self.data
.list_proxy_node_metrics(node_id, step, from_unix_secs, to_unix_secs, limit)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_fleet_metrics(
&self,
step: ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyFleetMetricsBucket>, GatewayError> {
self.data
.list_proxy_fleet_metrics(step, from_unix_secs, to_unix_secs, limit)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn register_proxy_node(
&self,
mutation: &aether_data::repository::proxy_nodes::ProxyNodeRegistrationMutation,
@@ -630,6 +671,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);
}

View File

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

View File

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

View File

@@ -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 keepalive0 关闭) |
| `--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]]` 条目;旧的顶层单服务器写法已不再支持。

View File

@@ -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,

View File

@@ -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")

View File

@@ -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() {

View File

@@ -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");

View File

@@ -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]

View File

@@ -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}")]

View File

@@ -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);

View File

@@ -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,

View File

@@ -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()),
})
}

View File

@@ -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,

View File

@@ -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,

View File

@@ -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
}

View File

@@ -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(());

View File

@@ -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,

View File

@@ -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,

View File

@@ -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")

File diff suppressed because it is too large Load Diff

View File

@@ -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;

View File

@@ -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(

View File

@@ -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);
};

View File

@@ -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)
);

View File

@@ -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);

View File

@@ -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);

View File

@@ -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);

View File

@@ -8,3 +8,4 @@
110_redeem_codes.sql
120_stats_rollups.sql
130_stats_cost_savings.sql
140_proxy_node_metrics.sql

View File

@@ -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`)
);

View File

@@ -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);

View File

@@ -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);

View File

@@ -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"]

View File

@@ -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)

View File

@@ -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()?)

View File

@@ -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)),
},
}
}

View File

@@ -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

View File

@@ -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,
]
);
}

View File

@@ -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| {

View File

@@ -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())?,
)
}

View File

@@ -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()?,

View File

@@ -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
)
"#,
)

View File

@@ -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),
},
],

View File

@@ -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,
};

View File

@@ -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;

View File

@@ -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)]

View File

@@ -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(
&registered.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(
&registered.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);
}
}

View File

@@ -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(&current), 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
);
}
}

View File

@@ -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)

View File

@@ -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()?;
}
}

View File

@@ -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()?

View File

@@ -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",

View File

@@ -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);
}

View File

@@ -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,
}
}
}

View File

@@ -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
}

View File

@@ -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')
})
})

View File

@@ -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')

View File

@@ -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) {