Merge remote-tracking branch 'origin/aether-rust-pioneer' into codex/async-cleanup-records

# Conflicts:
#	apps/aether-gateway/src/maintenance/mod.rs
#	apps/aether-gateway/src/maintenance/runtime/runners.rs
This commit is contained in:
fawney19
2026-05-10 00:28:39 +08:00
315 changed files with 19750 additions and 3183 deletions

View File

@@ -23,6 +23,7 @@ aether-provider-transport.workspace = true
aether-scheduler-core.workspace = true
aether-runtime.workspace = true
aether-runtime-state.workspace = true
aether-task-runtime.workspace = true
aether-usage-runtime.workspace = true
aether-video-tasks-core.workspace = true
aether-wallet.workspace = true

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

@@ -1,8 +1,11 @@
use axum::body::Bytes;
use crate::ai_serving::{
endpoint_config_forces_upstream_stream_policy as endpoint_config_forces_upstream_stream_policy_impl,
enforce_request_body_stream_field as enforce_request_body_stream_field_impl,
force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl,
is_json_request, parse_direct_request_body as parse_direct_request_body_impl,
resolve_upstream_is_stream_from_endpoint_config as resolve_upstream_is_stream_from_endpoint_config_impl,
};
pub(crate) use crate::ai_serving::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
@@ -37,6 +40,52 @@ pub(crate) fn force_upstream_streaming_for_provider(
force_upstream_streaming_for_provider_impl(provider_type, provider_api_format)
}
pub(crate) fn resolve_upstream_is_stream_for_provider(
endpoint_config: Option<&serde_json::Value>,
provider_type: &str,
provider_api_format: &str,
client_is_stream: bool,
hard_requires_streaming: bool,
) -> bool {
let hard_requires_streaming = hard_requires_streaming
|| force_upstream_streaming_for_provider(provider_type, provider_api_format);
resolve_upstream_is_stream_from_endpoint_config_impl(
endpoint_config,
client_is_stream,
hard_requires_streaming,
)
}
pub(crate) fn endpoint_config_forces_body_stream_field(
endpoint_config: Option<&serde_json::Value>,
) -> bool {
endpoint_config_forces_upstream_stream_policy_impl(endpoint_config)
}
pub(crate) fn request_requires_body_stream_field(
body_json: &serde_json::Value,
force_body_stream_field: bool,
) -> bool {
force_body_stream_field
|| body_json
.as_object()
.is_some_and(|object| object.contains_key("stream"))
}
pub(crate) fn enforce_provider_body_stream_policy(
provider_request_body: &mut serde_json::Value,
provider_api_format: &str,
upstream_is_stream: bool,
require_body_stream_field: bool,
) {
enforce_request_body_stream_field_impl(
provider_request_body,
provider_api_format,
upstream_is_stream,
require_body_stream_field,
);
}
pub(crate) fn extract_standard_requested_model(body_json: &serde_json::Value) -> Option<String> {
aether_ai_serving::extract_ai_standard_requested_model(body_json)
}
@@ -56,8 +105,10 @@ pub(crate) fn extract_requested_model_from_request(
#[cfg(test)]
mod tests {
use super::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
extract_requested_model_from_request, extract_standard_requested_model,
force_upstream_streaming_for_provider, RequestedModelFamily,
force_upstream_streaming_for_provider, resolve_upstream_is_stream_for_provider,
RequestedModelFamily,
};
use axum::http::Request;
use serde_json::json;
@@ -90,6 +141,67 @@ mod tests {
));
}
#[test]
fn resolves_endpoint_upstream_stream_policy_with_provider_hard_constraints() {
assert!(resolve_upstream_is_stream_for_provider(
Some(&json!({"upstream_stream_policy": "force_stream"})),
"openai",
"openai:chat",
false,
false,
));
assert!(!resolve_upstream_is_stream_for_provider(
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
"openai",
"openai:chat",
true,
false,
));
assert!(resolve_upstream_is_stream_for_provider(
Some(&json!({"upstream_stream_policy": "auto"})),
"openai",
"openai:chat",
true,
false,
));
assert!(resolve_upstream_is_stream_for_provider(
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
"codex",
"openai:responses",
true,
false,
));
}
#[test]
fn enforces_provider_body_stream_policy_for_body_and_streamless_formats() {
let mut openai_chat = json!({"stream": true});
enforce_provider_body_stream_policy(&mut openai_chat, "openai:chat", false, false);
assert_eq!(openai_chat.get("stream"), Some(&json!(false)));
let mut ordinary_sync = json!({"messages": []});
enforce_provider_body_stream_policy(&mut ordinary_sync, "openai:chat", false, false);
assert!(ordinary_sync.get("stream").is_none());
let mut compact = json!({"stream": true});
enforce_provider_body_stream_policy(&mut compact, "openai:responses:compact", true, true);
assert!(compact.get("stream").is_none());
}
#[test]
fn detects_endpoint_configs_that_force_body_stream_field() {
assert!(endpoint_config_forces_body_stream_field(Some(
&json!({"upstream_stream_policy": "force_stream"})
)));
assert!(endpoint_config_forces_body_stream_field(Some(
&json!({"upstream_stream_policy": "force_non_stream"})
)));
assert!(!endpoint_config_forces_body_stream_field(Some(
&json!({"upstream_stream_policy": "auto"})
)));
assert!(!endpoint_config_forces_body_stream_field(None));
}
#[test]
fn extracts_standard_requested_model_from_request_body() {
let requested_model =

View File

@@ -112,6 +112,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -3,6 +3,9 @@ use std::sync::Arc;
use serde_json::Value;
use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field,
};
use crate::ai_serving::transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
@@ -50,6 +53,7 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac
};
let behavior = policy::classify_same_format_provider_request_behavior(
transport,
provider_api_format,
crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata {
api_format: provider_api_format,
require_streaming: false,
@@ -131,6 +135,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
prepared.transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
prepared.upstream_is_stream,
prepared.force_body_stream_field,
prepared.kiro_auth.as_ref(),
prepared.is_claude_code,
enable_model_directives,
@@ -170,6 +175,18 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
&mut base_provider_request_body,
&mapping,
);
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
// Kiro behavior classification already hard-requires upstream streaming,
// and the Kiro envelope does not use a top-level body stream field.
if prepared.kiro_auth.is_none() {
enforce_provider_body_stream_policy(
&mut base_provider_request_body,
prepared.provider_api_format.as_str(),
prepared.upstream_is_stream,
request_requires_body_stream_field(body_json, prepared.force_body_stream_field),
);
}
}
let antigravity_auth = if prepared.is_antigravity {

View File

@@ -13,12 +13,14 @@ use super::super::LocalSameFormatProviderFamily;
pub(super) fn classify_same_format_provider_request_behavior(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
spec_metadata: LocalExecutionSurfaceSpecMetadata,
) -> SameFormatProviderRequestBehavior {
classify_same_format_provider_request_behavior_impl(
transport,
SameFormatProviderRequestBehaviorParams {
require_streaming: spec_metadata.require_streaming,
provider_api_format,
report_kind: spec_metadata
.report_kind
.expect("same-format provider specs should declare report kind"),

View File

@@ -35,6 +35,7 @@ pub(super) struct PreparedSameFormatProviderCandidate {
pub(super) mapped_model: String,
pub(super) report_kind: &'static str,
pub(super) upstream_is_stream: bool,
pub(super) force_body_stream_field: bool,
}
pub(super) async fn prepare_local_same_format_provider_candidate(
@@ -51,7 +52,11 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
let candidate = &eligible.candidate;
let transport = Arc::clone(&eligible.transport);
let provider_api_format = eligible.provider_api_format.as_str();
let behavior = classify_same_format_provider_request_behavior(&transport, spec_metadata);
let behavior = classify_same_format_provider_request_behavior(
&transport,
provider_api_format,
spec_metadata,
);
if !same_format_provider_transport_supported(
&behavior,
@@ -174,5 +179,6 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
mapped_model,
report_kind: behavior.report_kind,
upstream_is_stream: behavior.upstream_is_stream,
force_body_stream_field: behavior.force_body_stream_field,
})
}

View File

@@ -14,6 +14,7 @@ pub(crate) fn build_same_format_provider_request_body(
body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
upstream_is_stream: bool,
force_body_stream_field: bool,
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
is_claude_code: bool,
enable_model_directives: bool,
@@ -28,6 +29,7 @@ pub(crate) fn build_same_format_provider_request_body(
body_rules,
request_headers,
upstream_is_stream,
force_body_stream_field,
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
is_claude_code,
enable_model_directives,

View File

@@ -9,7 +9,11 @@ use aether_ai_serving::{
use aether_scheduler_core::{ClientSessionAffinity, SchedulerRankingOutcome};
use serde_json::{Map, Value};
use crate::ai_serving::{request_origin_from_headers, ExecutionRuntimeAuthContext, RequestOrigin};
use crate::ai_serving::{
request_origin_from_headers, request_path_implies_stream_request, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string, ExecutionRuntimeAuthContext,
RequestOrigin,
};
use crate::client_session_affinity::{
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
};
@@ -40,6 +44,8 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) provider_request_method: Option<Value>,
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
pub(crate) original_headers: &'a http::HeaderMap,
pub(crate) request_path: Option<&'a str>,
pub(crate) request_query_string: Option<&'a str>,
pub(crate) request_origin: Option<RequestOrigin>,
pub(crate) original_request_body_json: Option<&'a Value>,
pub(crate) original_request_body_base64: Option<&'a str>,
@@ -80,6 +86,15 @@ pub(crate) fn build_local_execution_report_context(
{
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
}
insert_request_path_fields(
&mut extra_fields,
parts.request_path,
parts.request_query_string,
);
let client_requested_stream = parts.client_requested_stream
|| parts
.request_path
.is_some_and(request_path_implies_stream_request);
build_ai_execution_report_context(AiExecutionReportContextParts {
auth_context: parts.auth_context,
@@ -113,7 +128,7 @@ pub(crate) fn build_local_execution_report_context(
client_ip,
user_agent,
},
client_requested_stream: parts.client_requested_stream,
client_requested_stream,
upstream_is_stream: parts.upstream_is_stream,
has_envelope: parts.has_envelope,
needs_conversion: parts.needs_conversion,
@@ -121,6 +136,30 @@ pub(crate) fn build_local_execution_report_context(
})
}
fn insert_request_path_fields(
extra_fields: &mut Map<String, Value>,
request_path: Option<&str>,
request_query_string: Option<&str>,
) {
let Some(path) = request_path.and_then(sanitize_request_path) else {
return;
};
let query = request_query_string.and_then(sanitize_request_query_string);
let path_and_query = sanitize_request_path_and_query(path.as_str(), query.as_deref())
.unwrap_or_else(|| path.clone());
extra_fields
.entry("request_path".to_string())
.or_insert_with(|| Value::String(path.clone()));
if let Some(query) = query.clone() {
extra_fields
.entry("request_query_string".to_string())
.or_insert_with(|| Value::String(query.to_string()));
}
extra_fields
.entry("request_path_and_query".to_string())
.or_insert_with(|| Value::String(path_and_query));
}
pub(crate) fn provider_stream_event_api_format_for_provider_type(
provider_type: &str,
) -> Option<&'static str> {
@@ -226,6 +265,8 @@ mod tests {
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: Some("/v1/chat/completions"),
request_query_string: Some("debug=true&limit=10"),
request_origin: Some(RequestOrigin {
client_ip: Some("203.0.113.8".to_string()),
user_agent: Some("Claude-Code/1.0".to_string()),
@@ -255,6 +296,77 @@ mod tests {
"session_key": "account=account-1;session=session-1"
})
);
assert_eq!(report_context["request_path"], "/v1/chat/completions");
assert_eq!(report_context["request_query_string"], "limit=10");
assert_eq!(
report_context["request_path_and_query"],
"/v1/chat/completions?limit=10"
);
}
#[test]
fn local_execution_report_context_treats_stream_generate_content_path_as_client_stream() {
let auth_context = ExecutionRuntimeAuthContext {
user_id: "user-1".to_string(),
api_key_id: "api-key-1".to_string(),
username: None,
api_key_name: None,
balance_remaining: None,
access_allowed: true,
api_key_is_standalone: false,
};
let original_headers = http::HeaderMap::new();
let provider_request_headers = BTreeMap::new();
let report_context =
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &auth_context,
request_id: "trace-1",
candidate_id: "candidate-1",
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
model: "gemini-3.1-flash-image-preview",
provider_name: "Gemini",
provider_id: "provider-1",
endpoint_id: "endpoint-1",
key_id: "key-1",
key_name: None,
model_id: None,
global_model_id: None,
global_model_name: None,
provider_api_format: "gemini:generate_content",
client_api_format: "gemini:generate_content",
mapped_model: None,
candidate_group_id: None,
ranking: None,
upstream_url: None,
header_rules: None,
body_rules: None,
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: Some(
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent",
),
request_query_string: Some("key=secret&alt=sse"),
request_origin: None,
original_request_body_json: Some(&json!({
"contents": [{"role": "user", "parts": [{"text": "hi"}]}]
})),
original_request_body_base64: None,
client_session_affinity: None,
client_requested_stream: false,
upstream_is_stream: true,
has_envelope: false,
needs_conversion: false,
extra_fields: Map::new(),
});
assert_eq!(report_context["client_requested_stream"], true);
assert_eq!(report_context["request_query_string"], "alt=sse");
assert_eq!(
report_context["request_path_and_query"],
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent?alt=sse"
);
}
#[test]
@@ -300,6 +412,8 @@ mod tests {
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: None,
request_query_string: None,
request_origin: None,
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,

View File

@@ -94,6 +94,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),

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,49 @@ 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_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
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 +149,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

@@ -75,6 +75,8 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

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

@@ -120,6 +120,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -7,7 +7,10 @@ use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
OauthPreparationContext,
};
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider;
use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
};
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, request_body_build_failure_extra_data,
@@ -18,10 +21,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 +65,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)
@@ -163,11 +178,15 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
}
};
let upstream_is_stream = spec_metadata.require_streaming
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
provider_api_format,
);
let upstream_is_stream = resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
provider_api_format,
spec_metadata.require_streaming,
is_kiro_claude_cli,
);
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let enable_model_directives =
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state,
@@ -213,6 +232,12 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
return None;
}
};
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
if let Some(mapping) =
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
state,
@@ -225,6 +250,14 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
&mut provider_request_body,
&mapping,
);
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
}
if let Some(kiro_auth) = kiro_auth.as_ref() {
@@ -331,6 +364,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

@@ -116,9 +116,9 @@ pub(crate) fn build_gemini_stream_plan_from_decision(
content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_api_format.as_str(),
client_api_format: core.client_api_format.as_str(),
upstream_is_stream: true,
upstream_is_stream: payload.upstream_is_stream,
build_from_request_when_empty: false,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
});
let content_type = payload
.content_type

View File

@@ -15,3 +15,6 @@ pub(crate) use self::responses::{
build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body,
build_local_openai_responses_upstream_url,
};
pub(super) use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field,
};

View File

@@ -9,10 +9,13 @@ use crate::ai_serving::{
GatewayProviderTransportSnapshot,
};
use super::{enforce_provider_body_stream_policy, request_requires_body_stream_field};
pub(crate) fn build_local_openai_chat_request_body(
body_json: &Value,
mapped_model: &str,
upstream_is_stream: bool,
force_body_stream_field: bool,
body_rules: Option<&Value>,
request_headers: &http::HeaderMap,
enable_model_directives: bool,
@@ -23,12 +26,20 @@ pub(crate) fn build_local_openai_chat_request_body(
upstream_is_stream,
enable_model_directives,
)?;
apply_standard_provider_request_body_rules_with_request_headers(
provider_request_body,
body_rules,
body_json,
request_headers,
)
let mut provider_request_body =
apply_standard_provider_request_body_rules_with_request_headers(
provider_request_body,
body_rules,
body_json,
request_headers,
)?;
enforce_provider_body_stream_policy(
&mut provider_request_body,
"openai:chat",
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
Some(provider_request_body)
}
pub(crate) fn build_local_openai_chat_upstream_url(
@@ -44,6 +55,7 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
provider_type: &str,
provider_api_format: &str,
upstream_is_stream: bool,
force_body_stream_field: bool,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap,
@@ -74,6 +86,12 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
&mut provider_request_body,
provider_api_format,
);
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
Some(provider_request_body)
}

View File

@@ -9,10 +9,13 @@ use crate::ai_serving::{
GatewayProviderTransportSnapshot,
};
use super::{enforce_provider_body_stream_policy, request_requires_body_stream_field};
pub(crate) fn build_local_openai_responses_request_body(
body_json: &Value,
mapped_model: &str,
require_streaming: bool,
force_body_stream_field: bool,
provider_type: &str,
provider_api_format: &str,
body_rules: Option<&Value>,
@@ -44,6 +47,12 @@ pub(crate) fn build_local_openai_responses_request_body(
&mut provider_request_body,
provider_api_format,
);
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
require_streaming,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
Some(provider_request_body)
}
@@ -53,6 +62,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
client_api_format: &str,
provider_api_format: &str,
upstream_is_stream: bool,
force_body_stream_field: bool,
provider_type: &str,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
@@ -85,6 +95,12 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
&mut provider_request_body,
provider_api_format,
);
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
Some(provider_request_body)
}

View File

@@ -87,6 +87,7 @@ fn builds_openai_chat_cross_format_request_body_from_openai_responses_source() {
"openai:responses",
"openai:chat",
false,
false,
"openai",
None,
None,
@@ -121,6 +122,7 @@ fn local_openai_responses_wrapper_preserves_body_order_after_edits() {
&body_json,
"gpt-5.4",
true,
false,
"codex",
"openai:responses",
None,
@@ -160,6 +162,7 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests(
&body_json,
"gpt-5.4",
false,
false,
"openai",
"openai:responses:compact",
None,
@@ -170,6 +173,7 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests(
.expect("local openai compact body should build");
assert!(provider_request_body.get("store").is_none());
assert!(provider_request_body.get("stream").is_none());
}
#[test]
@@ -187,6 +191,7 @@ fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() {
&body_json,
"gpt-5.4",
false,
false,
"openai",
"openai:responses",
Some(&body_rules),
@@ -237,6 +242,7 @@ fn strips_metadata_for_codex_openai_responses_requests() {
"claude:messages",
"openai:responses",
true,
false,
"codex",
None,
None,
@@ -271,6 +277,7 @@ fn applies_codex_defaults_unless_body_rules_handle_the_field() {
"claude:messages",
"openai:responses",
true,
false,
"codex",
Some(&body_rules),
None,
@@ -300,6 +307,7 @@ fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() {
"claude:messages",
"openai:responses",
true,
false,
"codex",
None,
Some("key-123"),
@@ -330,6 +338,7 @@ fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
"codex",
"openai:responses",
false,
false,
None,
Some("key-123"),
&http::HeaderMap::new(),

View File

@@ -1,4 +1,5 @@
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, insert_provider_stream_event_api_format,
LocalExecutionReportContextParts,
@@ -29,6 +30,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
report_kind: &str,
upstream_is_stream: bool,
) -> Option<AiExecutionDecision> {
let decision_is_stream = decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND;
let attempt_identity = attempt.attempt_identity();
let LocalOpenAiChatCandidateAttempt {
eligible,
@@ -107,6 +109,8 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,
@@ -147,7 +151,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream: upstream_is_stream,
decision_is_stream,
decision_kind: decision_kind.to_string(),
execution_strategy,
conversion_mode,

View File

@@ -8,7 +8,10 @@ use crate::ai_serving::planner::candidate_preparation::{
OauthPreparationContext,
};
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, OPENAI_CHAT_STREAM_PLAN_KIND,
};
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
@@ -72,6 +75,8 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let candidate = &eligible.candidate;
let provider_api_format = eligible.provider_api_format.as_str();
let transport = &eligible.transport;
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let enable_model_directives =
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state,
@@ -128,6 +133,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
body_json,
&prepared_candidate.mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.endpoint.body_rules.as_ref(),
&parts.headers,
enable_model_directives,
@@ -214,6 +220,13 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", "openai:chat");
let resolved_report_kind =
if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND || !upstream_is_stream {
report_kind.to_string()
} else {
"openai_chat_sync_finalize".to_string()
};
return Some(LocalOpenAiChatCandidatePayloadParts {
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
@@ -224,7 +237,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
upstream_url,
execution_strategy,
conversion_mode,
report_kind: report_kind.to_string(),
report_kind: resolved_report_kind,
envelope_name: None,
transport: Arc::clone(transport),
});
@@ -349,6 +362,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
transport.provider.provider_type.as_str(),
provider_api_format.as_str(),
upstream_is_stream,
force_body_stream_field,
if is_kiro_claude_cli {
None
} else {
@@ -387,6 +401,14 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
&mut provider_request_body,
&mapping,
);
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format.as_str(),
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
}
if let Some(kiro_auth) = kiro_auth.as_ref() {

View File

@@ -9,6 +9,9 @@ mod stream;
#[path = "plans/sync.rs"]
mod sync;
use crate::ai_serving::planner::common::resolve_upstream_is_stream_for_provider;
use crate::ai_serving::GatewayProviderTransportSnapshot;
pub(super) use self::candidates::list_local_openai_chat_candidates;
pub(super) use self::diagnostic::set_local_openai_chat_miss_diagnostic;
pub(super) use self::resolve::resolve_local_openai_chat_decision_input;
@@ -18,3 +21,147 @@ pub(super) use self::stream::{
pub(super) use self::sync::{
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
};
fn openai_chat_upstream_is_stream_for_candidate(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
client_is_stream: bool,
) -> bool {
let hard_requires_streaming =
crate::ai_serving::transport::kiro::is_kiro_claude_messages_transport(
transport,
provider_api_format,
);
resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
provider_api_format,
client_is_stream,
hard_requires_streaming,
)
}
#[cfg(test)]
mod tests {
use super::openai_chat_upstream_is_stream_for_candidate;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use serde_json::{json, Value};
fn sample_transport(
provider_type: &str,
api_format: &str,
endpoint_config: Option<Value>,
) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider".to_string(),
provider_type: provider_type.to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: api_format.to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://api.example.test".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: endpoint_config,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn openai_chat_policy_resolver_supports_force_stream_force_non_stream_and_auto() {
let force_stream = sample_transport(
"openai",
"openai:chat",
Some(json!({"upstream_stream_policy": "force_stream"})),
);
assert!(openai_chat_upstream_is_stream_for_candidate(
&force_stream,
"openai:chat",
false,
));
let force_non_stream = sample_transport(
"openai",
"openai:chat",
Some(json!({"upstream_stream_policy": "force_non_stream"})),
);
assert!(!openai_chat_upstream_is_stream_for_candidate(
&force_non_stream,
"openai:chat",
true,
));
let auto = sample_transport(
"openai",
"openai:chat",
Some(json!({"upstream_stream_policy": "auto"})),
);
assert!(openai_chat_upstream_is_stream_for_candidate(
&auto,
"openai:chat",
true,
));
assert!(!openai_chat_upstream_is_stream_for_candidate(
&auto,
"openai:chat",
false,
));
}
#[test]
fn openai_chat_policy_resolver_preserves_provider_hard_streaming() {
let codex = sample_transport(
"codex",
"openai:responses",
Some(json!({"upstream_stream_policy": "force_non_stream"})),
);
assert!(openai_chat_upstream_is_stream_for_candidate(
&codex,
"openai:responses",
false,
));
}
}

View File

@@ -13,6 +13,7 @@ use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
};
use super::openai_chat_upstream_is_stream_for_candidate;
use super::resolve::resolve_local_openai_chat_decision_input;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
@@ -120,6 +121,11 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
&self,
attempt: LocalOpenAiChatCandidateAttempt,
) -> Result<Option<AiStreamAttempt>, GatewayError> {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
true,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
self.state,
self.parts,
@@ -129,7 +135,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
true,
upstream_is_stream,
)
.await
else {
@@ -222,6 +228,11 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
let mut plans = Vec::new();
for attempt in attempts {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
true,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
@@ -231,7 +242,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
true,
upstream_is_stream,
)
.await
else {

View File

@@ -13,11 +13,10 @@ use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
};
use super::openai_chat_upstream_is_stream_for_candidate;
use super::resolve::resolve_local_openai_chat_decision_input;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::common::{
force_upstream_streaming_for_provider, OPENAI_CHAT_SYNC_PLAN_KIND,
};
use crate::ai_serving::planner::common::OPENAI_CHAT_SYNC_PLAN_KIND;
use crate::ai_serving::planner::plan_builders::{
build_openai_chat_sync_plan_from_decision, AiSyncAttempt,
};
@@ -32,13 +31,6 @@ pub(crate) struct LocalOpenAiChatSyncAttemptSource<'a> {
candidates: LocalOpenAiChatCandidateAttemptSource<'a>,
}
fn openai_chat_sync_upstream_is_stream_for_candidate(
provider_type: &str,
provider_api_format: &str,
) -> bool {
force_upstream_streaming_for_provider(provider_type, provider_api_format)
}
pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
@@ -129,9 +121,10 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
&self,
attempt: LocalOpenAiChatCandidateAttempt,
) -> Result<Option<AiSyncAttempt>, GatewayError> {
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate(
attempt.eligible.transport.provider.provider_type.as_str(),
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
false,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
self.state,
@@ -235,9 +228,10 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
let mut plans = Vec::new();
for attempt in attempts {
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate(
attempt.eligible.transport.provider.provider_type.as_str(),
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
false,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
@@ -272,28 +266,3 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
Ok(plans)
}
#[cfg(test)]
mod tests {
use super::openai_chat_sync_upstream_is_stream_for_candidate;
#[test]
fn openai_chat_sync_forces_streaming_for_codex_openai_responses_candidates() {
assert!(openai_chat_sync_upstream_is_stream_for_candidate(
"codex",
"openai:responses"
));
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
"codex",
"openai:responses:compact"
));
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
"openai",
"openai:responses"
));
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
"codex",
"openai:chat"
));
}
}

View File

@@ -9,6 +9,7 @@ use super::super::{
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
AiStreamAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::provider_adaptation_requires_eventstream_accept;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
@@ -17,6 +18,18 @@ use crate::ai_serving::transport::{
};
use crate::{AiExecutionDecision, GatewayError};
fn effective_stream_accept_mode(
payload_upstream_is_stream: bool,
provider_request_body: &serde_json::Value,
) -> bool {
payload_upstream_is_stream
|| provider_request_body
.as_object()
.and_then(|body| body.get("stream"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
pub(crate) fn build_openai_chat_stream_plan_from_decision(
parts: &http::request::Parts,
body_json: &serde_json::Value,
@@ -53,22 +66,34 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
provider_request_body
.insert("model".to_string(), serde_json::Value::String(mapped_model));
}
provider_request_body.insert("stream".to_string(), serde_json::Value::Bool(true));
let require_body_stream_field = provider_request_body.contains_key("stream");
let mut provider_request_body = serde_json::Value::Object(provider_request_body);
enforce_provider_body_stream_policy(
&mut provider_request_body,
core.provider_api_format.as_str(),
payload.upstream_is_stream,
require_body_stream_field,
);
let Some(provider_request_object) = provider_request_body.as_object_mut() else {
return Ok(None);
};
if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) {
let existing = provider_request_body
let existing = provider_request_object
.get("prompt_cache_key")
.and_then(|value| value.as_str())
.map(str::trim)
.unwrap_or_default();
if existing.is_empty() {
provider_request_body.insert(
provider_request_object.insert(
"prompt_cache_key".to_string(),
serde_json::Value::String(prompt_cache_key),
);
}
}
serde_json::Value::Object(provider_request_body)
provider_request_body
};
let effective_upstream_is_stream =
effective_stream_accept_mode(payload.upstream_is_stream, &provider_request_body_value);
let extra_headers = std::mem::take(&mut payload.extra_headers);
let mut provider_request_headers =
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
@@ -82,9 +107,9 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_api_format.as_str(),
client_api_format: core.client_api_format.as_str(),
upstream_is_stream: true,
upstream_is_stream: effective_upstream_is_stream,
build_from_request_when_empty: true,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreamingOrWildcard,
});
let content_type = payload
.content_type
@@ -157,13 +182,18 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
.as_ref()
.and_then(|context| context.get("envelope_name"))
.and_then(serde_json::Value::as_str);
let accept_policy = if provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
let effective_upstream_is_stream =
effective_stream_accept_mode(payload.upstream_is_stream, &provider_request_body_value);
let accept_policy = if effective_upstream_is_stream
&& provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
} else if envelope_name.is_some() {
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
} else {
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreamingOrWildcard
};
let mut provider_request_headers =
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
@@ -177,7 +207,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_api_format.as_str(),
client_api_format: core.client_api_format.as_str(),
upstream_is_stream: true,
upstream_is_stream: effective_upstream_is_stream,
build_from_request_when_empty: false,
accept_policy,
});
@@ -219,7 +249,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
plan_url = %plan.url,
client_api_format = %plan.client_api_format,
provider_api_format = %plan.provider_api_format,
upstream_is_stream = payload.upstream_is_stream,
upstream_is_stream = effective_upstream_is_stream,
compact,
"gateway built local openai responses stream execution plan"
);
@@ -427,6 +457,104 @@ mod tests {
);
}
#[test]
fn build_openai_chat_stream_plan_keeps_downstream_stream_for_force_non_stream_upstream() {
fn force_non_stream_payload(provider_request_body: Option<Value>) -> AiExecutionDecision {
AiExecutionDecision {
action: "stream".to_string(),
decision_kind: Some("openai_chat_stream".to_string()),
execution_strategy: None,
conversion_mode: None,
request_id: Some("req_force_non_stream".to_string()),
candidate_id: Some("cand_force_non_stream".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: Some("prov_force_non_stream".to_string()),
endpoint_id: Some("ep_force_non_stream".to_string()),
key_id: Some("key_force_non_stream".to_string()),
upstream_base_url: Some("https://example.com".to_string()),
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
provider_request_method: None,
auth_header: Some("authorization".to_string()),
auth_value: Some("Bearer upstream-token".to_string()),
provider_api_format: Some("openai:chat".to_string()),
client_api_format: Some("openai:chat".to_string()),
provider_contract: Some("openai:chat".to_string()),
client_contract: Some("openai:chat".to_string()),
model_name: Some("gpt-5.4".to_string()),
mapped_model: Some("gpt-5.4".to_string()),
prompt_cache_key: None,
extra_headers: BTreeMap::new(),
provider_request_headers: BTreeMap::new(),
provider_request_body,
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
upstream_is_stream: false,
report_kind: Some("openai_chat_stream_success".to_string()),
report_context: Some(json!({})),
auth_context: None,
}
}
let parts = http::Request::builder()
.uri("http://localhost/v1/chat/completions")
.body(())
.expect("request should build")
.into_parts()
.0;
let built = build_openai_chat_stream_plan_from_decision(
&parts,
&json!({}),
force_non_stream_payload(Some(json!({
"model": "gpt-5.4",
"messages": [],
"stream": false
}))),
)
.expect("plan build should succeed")
.expect("plan should be produced");
assert!(built.plan.stream);
assert_eq!(
built
.plan
.body
.json_body
.as_ref()
.and_then(|body| body.get("stream"))
.and_then(Value::as_bool),
Some(false)
);
let fallback_body = json!({
"model": "client-model",
"messages": [],
"stream": true
});
let built = build_openai_chat_stream_plan_from_decision(
&parts,
&fallback_body,
force_non_stream_payload(None),
)
.expect("fallback plan build should succeed")
.expect("fallback plan should be produced");
assert!(built.plan.stream);
assert_eq!(
built
.plan
.body
.json_body
.as_ref()
.and_then(|body| body.get("stream"))
.and_then(Value::as_bool),
Some(false)
);
}
#[test]
fn build_openai_chat_stream_plan_fallback_restores_claude_headers_for_cross_format() {
let parts = http::Request::builder()

View File

@@ -9,6 +9,7 @@ use super::super::{
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
AiSyncAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
@@ -51,23 +52,31 @@ pub(crate) fn build_openai_chat_sync_plan_from_decision(
provider_request_body
.insert("model".to_string(), serde_json::Value::String(mapped_model));
}
if payload.upstream_is_stream {
provider_request_body.insert("stream".to_string(), serde_json::Value::Bool(true));
}
let require_body_stream_field = provider_request_body.contains_key("stream");
let mut provider_request_body = serde_json::Value::Object(provider_request_body);
enforce_provider_body_stream_policy(
&mut provider_request_body,
core.provider_api_format.as_str(),
payload.upstream_is_stream,
require_body_stream_field,
);
let Some(provider_request_object) = provider_request_body.as_object_mut() else {
return Ok(None);
};
if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) {
let existing = provider_request_body
let existing = provider_request_object
.get("prompt_cache_key")
.and_then(|value| value.as_str())
.map(str::trim)
.unwrap_or_default();
if existing.is_empty() {
provider_request_body.insert(
provider_request_object.insert(
"prompt_cache_key".to_string(),
serde_json::Value::String(prompt_cache_key),
);
}
}
serde_json::Value::Object(provider_request_body)
provider_request_body
};
let extra_headers = std::mem::take(&mut payload.extra_headers);
let mut provider_request_headers =

View File

@@ -105,6 +105,8 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -9,7 +9,10 @@ use crate::ai_serving::planner::candidate_preparation::{
OauthPreparationContext,
};
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider;
use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_responses_request_body,
@@ -224,12 +227,15 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
.await;
let needs_bidirectional_conversion = !same_format && conversion_kind.is_some();
let upstream_is_stream = spec_metadata.require_streaming
|| is_antigravity
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
provider_api_format,
);
let upstream_is_stream = resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
provider_api_format,
spec_metadata.require_streaming,
is_antigravity || is_kiro_claude_cli,
);
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let Some(mut base_provider_request_body) = (if needs_bidirectional_conversion {
build_cross_format_openai_responses_request_body(
body_json,
@@ -237,6 +243,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
spec_metadata.api_format,
provider_api_format,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
if is_kiro_claude_cli {
None
@@ -252,6 +259,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
body_json,
&mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
provider_api_format,
if is_kiro_claude_cli {
@@ -293,6 +301,14 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
&mut base_provider_request_body,
&mapping,
);
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
enforce_provider_body_stream_policy(
&mut base_provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
}
let antigravity_auth = if is_antigravity {
match classify_local_antigravity_request_support(

View File

@@ -112,13 +112,14 @@ pub(crate) fn build_standard_stream_plan_from_decision(
.as_ref()
.and_then(|context| context.get("envelope_name"))
.and_then(serde_json::Value::as_str);
let accept_policy = if provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
let accept_policy = if payload.upstream_is_stream
&& provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
} else {
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
};
let mut provider_request_headers =
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
@@ -132,7 +133,7 @@ pub(crate) fn build_standard_stream_plan_from_decision(
content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_api_format.as_str(),
client_api_format: core.client_api_format.as_str(),
upstream_is_stream: true,
upstream_is_stream: payload.upstream_is_stream,
build_from_request_when_empty: false,
accept_policy,
});

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,
@@ -34,13 +40,14 @@ pub(crate) use aether_ai_formats::api::{
copy_request_number_field_as, core_error_background_report_kind,
core_error_default_client_api_format, core_success_background_report_kind,
default_model_for_openai_image_operation, encode_done_sse, encode_json_sse,
encode_kiro_sse_events, estimate_kiro_tokens, extract_openai_text_content,
encode_kiro_sse_events, endpoint_config_forces_upstream_stream_policy,
enforce_request_body_stream_field, 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,
@@ -64,7 +71,8 @@ pub(crate) use aether_ai_formats::api::{
provider_adaptation_should_unwrap_stream_envelope,
provider_private_response_allows_sync_finalize, request_candidate_api_format_preference,
request_candidate_api_formats, request_conversion_kind,
request_conversion_requires_enable_flag, resolve_claude_stream_spec, resolve_claude_sync_spec,
request_conversion_requires_enable_flag, request_path_implies_stream_request,
resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec,
resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
@@ -72,22 +80,26 @@ 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_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,
transform_provider_private_stream_line, value_as_u64, AiControlPlanRequest,
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
resolve_requested_gemini_image_model_for_request,
resolve_requested_openai_image_model_for_request,
resolve_upstream_is_stream_from_endpoint_config, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string,
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, transform_provider_private_stream_line, value_as_u64,
AiControlPlanRequest, 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

@@ -24,5 +24,7 @@ pub(crate) fn mount_public_support_routes(router: Router<AppState>) -> Router<Ap
.route("/api/capabilities", get(proxy_request))
.route("/api/capabilities/user-configurable", get(proxy_request))
.route("/api/capabilities/model/{*model_path}", get(proxy_request))
.route("/install/{*install_path}", get(proxy_request))
.route("/i/{*install_path}", get(proxy_request))
.route("/", get(proxy_request))
}

View File

@@ -152,6 +152,11 @@ const PERMISSION_GROUPS: &[PermissionGroup] = &[
label: "系统",
assignable: true,
},
PermissionGroup {
scope: "tasks",
label: "后台任务",
assignable: true,
},
PermissionGroup {
scope: "usage",
label: "用量",
@@ -434,6 +439,9 @@ fn permission_key(scope: &str, access: &str) -> &'static str {
("system", "read") => "admin:system:read",
("system", "write") => "admin:system:write",
("system", "admin") => "admin:system:admin",
("tasks", "read") => "admin:tasks:read",
("tasks", "write") => "admin:tasks:write",
("tasks", "admin") => "admin:tasks:admin",
("usage", "read") => "admin:usage:read",
("usage", "write") => "admin:usage:write",
("usage", "admin") => "admin:usage:admin",
@@ -492,6 +500,7 @@ mod tests {
"security",
"stats",
"system",
"tasks",
"usage",
"users",
"video_tasks",

View File

@@ -46,6 +46,83 @@ pub(super) fn classify_admin_operations_family_route(
"admin:video_tasks",
false,
))
} else if method == http::Method::GET
&& matches!(normalized_path, "/api/admin/tasks" | "/api/admin/tasks/")
{
Some(classified(
"admin_proxy",
"tasks_manage",
"list_tasks",
"admin:tasks",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
"/api/admin/tasks/stats" | "/api/admin/tasks/stats/"
)
{
Some(classified(
"admin_proxy",
"tasks_manage",
"stats",
"admin:tasks",
false,
))
} else if method == http::Method::GET
&& normalized_path.starts_with("/api/admin/tasks/")
&& normalized_path.ends_with("/events")
&& normalized_path.matches('/').count() == 5
{
Some(classified(
"admin_proxy",
"tasks_manage",
"events",
"admin:tasks",
false,
))
} else if method == http::Method::POST
&& normalized_path.starts_with("/api/admin/tasks/")
&& normalized_path.ends_with("/cancel")
&& normalized_path.matches('/').count() == 5
{
Some(classified(
"admin_proxy",
"tasks_manage",
"cancel",
"admin:tasks",
false,
))
} else if method == http::Method::POST
&& normalized_path.starts_with("/api/admin/tasks/")
&& normalized_path.ends_with("/trigger")
&& normalized_path.matches('/').count() == 5
{
Some(classified(
"admin_proxy",
"tasks_manage",
"trigger",
"admin:tasks",
false,
))
} else if method == http::Method::GET
&& normalized_path.starts_with("/api/admin/tasks/")
&& normalized_path["/api/admin/tasks/".len()..]
.split('/')
.count()
== 1
&& !matches!(
normalized_path,
"/api/admin/tasks/stats" | "/api/admin/tasks/stats/"
)
{
Some(classified(
"admin_proxy",
"tasks_manage",
"detail",
"admin:tasks",
false,
))
} else if method == http::Method::GET
&& normalized_path.starts_with("/api/admin/video-tasks/")
&& normalized_path.ends_with("/video")
@@ -113,6 +190,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

@@ -468,6 +468,20 @@ pub(super) fn classify_public_support_route(
"user:self",
false,
))
} else if method == http::Method::POST
&& has_single_nested_suffix_after_prefix(
normalized_path,
"/api/users/me/api-keys/",
"install-sessions",
)
{
Some(classified(
"public_support",
"users_me",
"api_key_install_session_create",
"user:self",
false,
))
} else if method == http::Method::POST
&& normalized_path.starts_with("/api/me/management-tokens/")
&& normalized_path.ends_with("/regenerate")
@@ -659,6 +673,17 @@ pub(super) fn classify_public_support_route(
"public:system_catalog",
false,
))
} else if method == http::Method::GET
&& (has_single_segment_after_prefix(normalized_path, "/install/")
|| has_single_segment_after_prefix(normalized_path, "/i/"))
{
Some(classified(
"public_support",
"install",
"script",
"public:install",
false,
))
} else if method == http::Method::GET && normalized_path == "/test-connection" {
Some(classified(
"public_support",

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

@@ -1,6 +1,7 @@
use http::Uri;
use super::{classify_control_route, headers};
use super::{classify_control_route, headers, GatewayPublicRequestContext};
use crate::handlers::shared::local_proxy_route_requires_buffered_body;
#[test]
fn classifies_models_list_as_public_support_route() {
@@ -366,6 +367,11 @@ fn classifies_users_me_routes_as_public_support_route() {
"/api/users/me/api-keys",
"api_keys_create",
),
(
http::Method::POST,
"/api/users/me/api-keys/key-1/install-sessions",
"api_key_install_session_create",
),
(
http::Method::PUT,
"/api/users/me/api-keys/key-1",
@@ -452,6 +458,25 @@ fn classifies_users_me_routes_as_public_support_route() {
}
}
#[test]
fn user_api_key_install_session_create_buffers_request_body() {
let headers = headers(&[]);
let uri: Uri = "/api/users/me/api-keys/key-1/install-sessions"
.parse()
.expect("uri should parse");
let decision =
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
let context = GatewayPublicRequestContext::from_request_parts(
"trace-install-session",
&http::Method::POST,
&uri,
&headers,
Some(decision),
);
assert!(local_proxy_route_requires_buffered_body(&context));
}
#[test]
fn classifies_payment_callback_as_public_support_route() {
let headers = headers(&[]);

View File

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

View File

@@ -34,6 +34,8 @@ impl GatewayDataState {
proxy_node_reader: None,
proxy_node_writer: None,
billing_reader: None,
background_task_reader: None,
background_task_writer: None,
gemini_file_mapping_reader: None,
gemini_file_mapping_writer: None,
global_model_reader: None,
@@ -73,6 +75,8 @@ impl GatewayDataState {
let proxy_node_reader = backends.read().proxy_nodes();
let proxy_node_writer = backends.write().proxy_nodes();
let billing_reader = backends.read().billing();
let background_task_reader = backends.read().background_tasks();
let background_task_writer = backends.write().background_tasks();
let gemini_file_mapping_reader = backends.read().gemini_file_mappings();
let global_model_reader = backends.read().global_models();
let global_model_writer = backends.write().global_models();
@@ -110,6 +114,8 @@ impl GatewayDataState {
proxy_node_reader,
proxy_node_writer,
billing_reader,
background_task_reader,
background_task_writer,
gemini_file_mapping_reader,
gemini_file_mapping_writer,
global_model_reader,
@@ -197,6 +203,14 @@ impl GatewayDataState {
self.announcement_writer.is_some()
}
pub(crate) fn has_background_task_reader(&self) -> bool {
self.background_task_reader.is_some()
}
pub(crate) fn has_background_task_writer(&self) -> bool {
self.background_task_writer.is_some()
}
pub(crate) fn has_audit_log_reader(&self) -> bool {
self.backends
.as_ref()

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::{
@@ -79,6 +81,11 @@ use aether_data::{
DataBackends, DataLayerError, DatabaseMaintenanceSummary, WalletDailyUsageAggregationInput,
WalletDailyUsageAggregationResult,
};
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskListQuery, BackgroundTaskReadRepository, BackgroundTaskSummary,
BackgroundTaskWriteRepository, StoredBackgroundTaskEvent, StoredBackgroundTaskRun,
StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
};
use aether_data_contracts::repository::billing::{
AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingMutationOutcome,
AdminBillingPresetApplyResult, AdminBillingRuleRecord, AdminBillingRuleWriteInput,
@@ -138,6 +145,8 @@ pub(crate) struct GatewayDataState {
proxy_node_reader: Option<Arc<dyn ProxyNodeReadRepository>>,
proxy_node_writer: Option<Arc<dyn ProxyNodeWriteRepository>>,
billing_reader: Option<Arc<dyn BillingReadRepository>>,
background_task_reader: Option<Arc<dyn BackgroundTaskReadRepository>>,
background_task_writer: Option<Arc<dyn BackgroundTaskWriteRepository>>,
gemini_file_mapping_reader: Option<Arc<dyn GeminiFileMappingReadRepository>>,
gemini_file_mapping_writer: Option<Arc<dyn GeminiFileMappingWriteRepository>>,
global_model_reader: Option<Arc<dyn GlobalModelReadRepository>>,
@@ -204,6 +213,14 @@ impl fmt::Debug for GatewayDataState {
.field("has_proxy_node_reader", &self.proxy_node_reader.is_some())
.field("has_proxy_node_writer", &self.proxy_node_writer.is_some())
.field("has_billing_reader", &self.billing_reader.is_some())
.field(
"has_background_task_reader",
&self.background_task_reader.is_some(),
)
.field(
"has_background_task_writer",
&self.background_task_writer.is_some(),
)
.field(
"has_gemini_file_mapping_reader",
&self.gemini_file_mapping_reader.is_some(),

View File

@@ -5,11 +5,12 @@ use super::{
AdminBillingRuleWriteInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery,
AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery,
AdminWalletRefundRequestListQuery, AnnouncementListQuery, AuditLogListQuery,
CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord, CreateManualWalletRechargeInput,
CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome,
CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput,
DataLayerError, DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
BackgroundTaskListQuery, BackgroundTaskSummary, CompleteAdminWalletRefundInput,
CreateAdminRedeemCodeBatchInput, CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord,
CreateManualWalletRechargeInput, CreateWalletRechargeOrderInput,
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, DataLayerError,
DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput,
GatewayDataState, GatewayProviderTransportSnapshot, LocalVideoTaskReadResponse,
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
@@ -19,15 +20,16 @@ use super::{
StoredAdminRedeemCodePage, StoredAdminWalletLedgerPage, StoredAdminWalletListPage,
StoredAdminWalletRefund, StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage,
StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredAnnouncement,
StoredAnnouncementPage, StoredBillingModelContext, StoredProviderQuotaSnapshot,
StoredAnnouncementPage, StoredBackgroundTaskEvent, StoredBackgroundTaskRun,
StoredBackgroundTaskRunPage, StoredBillingModelContext, StoredProviderQuotaSnapshot,
StoredProviderUsageSummary, StoredRequestUsageAudit, StoredSuspiciousActivity,
StoredUsageSettlement, StoredUserAuditLogPage, StoredUserAuthRecord, StoredUserExportRow,
StoredUserSummary, StoredVideoTask, StoredWalletDailyUsageLedger,
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAnnouncementRecord,
UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput, VideoTaskLookupKey,
VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount,
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult, WalletLookupKey,
WalletMutationOutcome,
UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun, UpsertUsageRecord, UpsertVideoTask,
UsageSettlementInput, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter,
VideoTaskStatusCount, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
WalletLookupKey, WalletMutationOutcome,
};
use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
@@ -1716,6 +1718,82 @@ impl GatewayDataState {
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
read_data_backed_video_task_response(self, route_family, request_path).await
}
pub(crate) async fn find_background_task_run(
&self,
run_id: &str,
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
match &self.background_task_reader {
Some(repository) => repository.find_run(run_id).await,
None => Ok(None),
}
}
pub(crate) async fn list_background_task_runs(
&self,
query: &BackgroundTaskListQuery,
) -> Result<StoredBackgroundTaskRunPage, DataLayerError> {
match &self.background_task_reader {
Some(repository) => repository.list_runs(query).await,
None => Ok(StoredBackgroundTaskRunPage::default()),
}
}
pub(crate) async fn list_background_task_events(
&self,
run_id: &str,
offset: usize,
limit: usize,
) -> Result<Vec<StoredBackgroundTaskEvent>, DataLayerError> {
match &self.background_task_reader {
Some(repository) => repository.list_events(run_id, offset, limit).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn summarize_background_task_runs(
&self,
) -> Result<BackgroundTaskSummary, DataLayerError> {
match &self.background_task_reader {
Some(repository) => repository.summarize_runs().await,
None => Ok(BackgroundTaskSummary::default()),
}
}
pub(crate) async fn upsert_background_task_run(
&self,
run: UpsertBackgroundTaskRun,
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
match &self.background_task_writer {
Some(repository) => repository.upsert_run(run).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn request_cancel_background_task_run(
&self,
run_id: &str,
updated_at_unix_secs: u64,
) -> Result<bool, DataLayerError> {
match &self.background_task_writer {
Some(repository) => {
repository
.request_cancel(run_id, updated_at_unix_secs)
.await
}
None => Ok(false),
}
}
pub(crate) async fn upsert_background_task_event(
&self,
event: UpsertBackgroundTaskEvent,
) -> Result<Option<StoredBackgroundTaskEvent>, DataLayerError> {
match &self.background_task_writer {
Some(repository) => repository.upsert_event(event).await.map(Some),
None => Ok(None),
}
}
}
#[cfg(test)]

View File

@@ -43,6 +43,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -93,6 +95,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,

View File

@@ -76,6 +76,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -125,6 +127,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -170,6 +174,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -303,6 +309,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -364,6 +372,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -418,6 +428,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -481,6 +493,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -526,6 +540,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -572,6 +588,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -629,6 +647,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -688,6 +708,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -731,6 +753,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -789,6 +813,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -840,6 +866,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -896,6 +924,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -953,6 +983,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -1009,6 +1041,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1054,6 +1088,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1099,6 +1135,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1156,6 +1194,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1218,6 +1258,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1263,6 +1305,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1313,6 +1357,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1380,6 +1426,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1442,6 +1490,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1488,6 +1538,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1534,6 +1586,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1582,6 +1636,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1628,6 +1684,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1674,6 +1732,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1728,6 +1788,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1783,6 +1845,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1841,6 +1905,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1905,6 +1971,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -1970,6 +2038,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2039,6 +2109,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2115,6 +2187,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -2173,6 +2247,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2222,6 +2298,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2267,6 +2345,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -2318,6 +2398,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -2373,6 +2455,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -2429,6 +2513,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: Some(wallet_reader),
wallet_writer: Some(wallet_writer),
settlement_writer: None,
@@ -2478,6 +2564,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,

View File

@@ -47,6 +47,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(repository),
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -99,6 +101,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -148,6 +152,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -201,6 +207,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -258,6 +266,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
@@ -324,6 +334,8 @@ impl GatewayDataState {
usage_worker_queue: None,
video_task_reader: Some(video_task_reader),
video_task_writer: Some(video_task_writer),
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,

View File

@@ -1,23 +1,17 @@
use std::collections::BTreeMap;
use aether_contracts::{StreamFrame, StreamFramePayload};
use axum::body::Body;
use axum::http::Response;
use axum::http::StatusCode;
use base64::Engine as _;
use futures_util::StreamExt;
use serde_json::json;
use serde_json::{json, Map, Value};
use tokio_util::codec::{FramedRead, LinesCodec};
use tracing::warn;
use crate::api::response::build_client_response_from_parts;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::ndjson::decode_stream_frame_ndjson;
use crate::execution_runtime::submission::{has_nested_error, strip_utf8_bom_and_ws};
use crate::GatewayError;
use crate::{
GEMINI_FILES_DOWNLOAD_PLAN_KIND, MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_FRAMES,
OPENAI_VIDEO_CONTENT_PLAN_KIND,
};
use crate::{MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_FRAMES};
#[derive(Debug)]
pub(super) enum StreamPrefetchInspection {
@@ -51,6 +45,88 @@ pub(super) fn decode_stream_error_body(
)
}
fn header_value_case_insensitive<'a>(
headers: &'a BTreeMap<String, String>,
name: &str,
) -> Option<&'a str> {
headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case(name))
.map(|(_, value)| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn remove_header_case_insensitive(headers: &mut BTreeMap<String, String>, name: &str) {
let keys = headers
.keys()
.filter(|key| key.eq_ignore_ascii_case(name))
.cloned()
.collect::<Vec<_>>();
for key in keys {
headers.remove(&key);
}
}
pub(super) fn should_synthesize_non_success_stream_error_body(
status_code: u16,
error_body: &[u8],
) -> bool {
!(200..300).contains(&status_code)
&& ((300..400).contains(&status_code) || error_body.is_empty())
}
pub(super) fn build_synthetic_non_success_stream_error_body(
status_code: u16,
headers: &BTreeMap<String, String>,
) -> Value {
let mut error = Map::from_iter([
(
"type".to_string(),
Value::String("execution_runtime_non_success_status".to_string()),
),
(
"message".to_string(),
Value::String(format!(
"execution runtime stream returned non-success status {status_code}"
)),
),
("code".to_string(), Value::from(status_code)),
("upstream_status".to_string(), Value::from(status_code)),
]);
if let Some(location) = header_value_case_insensitive(headers, "location") {
error.insert("location".to_string(), Value::String(location.to_string()));
}
Value::Object(Map::from_iter([(
"error".to_string(),
Value::Object(error),
)]))
}
pub(super) fn synthetic_error_response_headers(
mut headers: BTreeMap<String, String>,
) -> BTreeMap<String, String> {
remove_header_case_insensitive(&mut headers, "content-encoding");
remove_header_case_insensitive(&mut headers, "content-length");
remove_header_case_insensitive(&mut headers, "content-type");
remove_header_case_insensitive(&mut headers, "location");
headers.insert("content-type".to_string(), "application/json".to_string());
headers
}
fn client_error_status_code_for_upstream_status(status_code: u16) -> u16 {
if (300..400).contains(&status_code) || status_code < 200 {
StatusCode::BAD_GATEWAY.as_u16()
} else {
status_code
}
}
pub(super) fn stream_client_error_status_code_for_upstream_status(status_code: u16) -> u16 {
client_error_status_code_for_upstream_status(status_code)
}
pub(super) fn inspect_prefetched_stream_body(
headers: &BTreeMap<String, String>,
body: &[u8],
@@ -166,62 +242,3 @@ where
}
Ok(None)
}
pub(super) fn build_execution_runtime_error_response(
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
status_code: u16,
headers: BTreeMap<String, String>,
error_body: Vec<u8>,
) -> Result<Response<Body>, GatewayError> {
let content_type = headers
.get("content-type")
.map(|value| value.to_ascii_lowercase())
.unwrap_or_default();
if plan_kind == GEMINI_FILES_DOWNLOAD_PLAN_KIND && !content_type.starts_with("application/json")
{
let wrapped = serde_json::to_vec(&json!({
"error": String::from_utf8_lossy(&error_body).to_string(),
}))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let wrapped_headers =
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
return build_client_response_from_parts(
status_code,
&wrapped_headers,
Body::from(wrapped),
trace_id,
Some(decision),
);
}
if plan_kind == OPENAI_VIDEO_CONTENT_PLAN_KIND && !content_type.starts_with("application/json")
{
let wrapped = serde_json::to_vec(&json!({
"error": {
"type": "upstream_error",
"message": "Video not available",
}
}))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let wrapped_headers =
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
return build_client_response_from_parts(
status_code,
&wrapped_headers,
Body::from(wrapped),
trace_id,
Some(decision),
);
}
build_client_response_from_parts(
status_code,
&headers,
Body::from(error_body),
trace_id,
Some(decision),
)
}

View File

@@ -23,15 +23,18 @@ use axum::http::Response;
use base64::Engine as _;
use futures_util::stream::BoxStream;
use futures_util::{StreamExt, TryStreamExt};
use serde_json::Value;
use serde_json::{json, Value};
use tokio::sync::mpsc;
use tokio_util::codec::{FramedRead, LinesCodec};
use tokio_util::io::StreamReader;
use tracing::{debug, info, warn};
use super::error::{
build_execution_runtime_error_response, collect_error_body, decode_stream_error_body,
inspect_prefetched_stream_body, read_next_frame, StreamPrefetchInspection,
build_synthetic_non_success_stream_error_body, collect_error_body, decode_stream_error_body,
inspect_prefetched_stream_body, read_next_frame,
should_synthesize_non_success_stream_error_body,
stream_client_error_status_code_for_upstream_status, synthetic_error_response_headers,
StreamPrefetchInspection,
};
#[path = "execution_failures.rs"]
mod execution_failures;
@@ -76,9 +79,10 @@ use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FR
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, build_local_error_flow_metadata, with_error_flow_report_context,
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
with_upstream_response_report_context, LocalAdaptiveRateLimitEffect,
LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect,
LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
@@ -86,7 +90,9 @@ use crate::request_candidate_runtime::{
};
use crate::usage::submit_stream_report;
use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest};
use crate::{AppState, GatewayError};
use crate::{
AppState, GatewayError, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
};
fn record_sync_terminal_usage(
state: &AppState,
@@ -124,6 +130,52 @@ fn build_stream_sync_payload(
}
}
#[allow(clippy::too_many_arguments)]
fn build_stream_error_sync_payload(
trace_id: &str,
report_kind: String,
report_context: Option<Value>,
upstream_status_code: u16,
provider_headers: BTreeMap<String, String>,
provider_body_json: Option<Value>,
provider_body_base64: Option<String>,
client_headers: BTreeMap<String, String>,
client_body_json: Option<Value>,
telemetry: Option<ExecutionTelemetry>,
) -> GatewaySyncReportRequest {
let client_status_code =
stream_client_error_status_code_for_upstream_status(upstream_status_code);
let mut report_context = report_context;
if client_status_code != upstream_status_code || client_headers != provider_headers {
let mut object = match report_context {
Some(Value::Object(object)) => object,
Some(other) => serde_json::Map::from_iter([("seed".to_string(), other)]),
None => serde_json::Map::new(),
};
object.insert(
"client_response_status_code".to_string(),
Value::from(client_status_code),
);
object.insert(
"client_response_headers".to_string(),
serde_json::to_value(client_headers).unwrap_or(Value::Null),
);
report_context = Some(Value::Object(object));
}
GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind,
report_context,
status_code: upstream_status_code,
headers: provider_headers,
body_json: provider_body_json,
client_body_json,
body_base64: provider_body_base64,
telemetry,
}
}
fn record_stream_terminal_usage(
state: &AppState,
plan: &ExecutionPlan,
@@ -157,6 +209,55 @@ fn build_stream_body_capture(
(body_base64, body_state)
}
fn wrap_non_json_binary_stream_error_for_client(
plan_kind: &str,
headers: &BTreeMap<String, String>,
error_body: &[u8],
) -> Result<Option<Value>, GatewayError> {
let content_type = headers
.get("content-type")
.map(|value| value.to_ascii_lowercase())
.unwrap_or_default();
if content_type.starts_with("application/json") {
return Ok(None);
}
let body = match plan_kind {
GEMINI_FILES_DOWNLOAD_PLAN_KIND => json!({
"error": String::from_utf8_lossy(error_body).to_string(),
}),
OPENAI_VIDEO_CONTENT_PLAN_KIND => json!({
"error": {
"type": "upstream_error",
"message": "Video not available",
}
}),
_ => return Ok(None),
};
Ok(Some(body))
}
fn with_stream_error_trace_context(
report_context: Option<&Value>,
status_code: u16,
headers: &BTreeMap<String, String>,
response_text: Option<&str>,
local_failover_analysis: crate::orchestration::LocalFailoverAnalysis,
) -> Option<Value> {
let upstream_context = with_upstream_response_report_context(
report_context,
status_code,
Some(headers),
None,
None,
None,
);
with_error_flow_report_context(
upstream_context.as_ref().or(report_context),
build_local_error_flow_metadata(status_code, response_text, local_failover_analysis),
)
}
#[allow(clippy::too_many_arguments)] // stream report payload assembly mirrors runtime state
fn build_stream_usage_payload(
trace_id: String,
@@ -948,11 +1049,29 @@ async fn execute_stream_from_frame_stream(
let stream_error_finalize_kind =
resolve_core_stream_error_finalize_report_kind(plan_kind, status_code);
if status_code >= 400 {
let error_body = collect_error_body(&mut lines).await?;
let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
if !(200..300).contains(&status_code) {
let provider_error_body = collect_error_body(&mut lines).await?;
let synthetic_body_json =
should_synthesize_non_success_stream_error_body(status_code, &provider_error_body)
.then(|| build_synthetic_non_success_stream_error_body(status_code, &headers));
let (provider_body_json, provider_body_base64) =
decode_stream_error_body(&headers, &provider_error_body);
let client_status_code = stream_client_error_status_code_for_upstream_status(status_code);
let wrapped_binary_body_json = wrap_non_json_binary_stream_error_for_client(
plan_kind,
&headers,
&provider_error_body,
)?;
let (client_body_json, client_error_body) =
if let Some(body_json) = synthetic_body_json.or(wrapped_binary_body_json) {
let body_bytes = serde_json::to_vec(&body_json)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
(Some(body_json), body_bytes)
} else {
(provider_body_json.clone(), provider_error_body.clone())
};
let error_response_text =
local_failover_response_text(body_json.as_ref(), &error_body, None);
local_failover_response_text(client_body_json.as_ref(), &client_error_body, None);
let failover_analysis = resolve_local_candidate_failover_analysis_stream(
state,
&plan,
@@ -1043,18 +1162,17 @@ async fn execute_stream_from_frame_stream(
);
if matches!(failover_decision, LocalFailoverDecision::RetryNextCandidate) {
let terminal_unix_secs = current_request_candidate_unix_ms();
let error_flow_report_context = with_error_flow_report_context(
let error_trace_report_context = with_stream_error_trace_context(
report_context.as_ref(),
build_local_error_flow_metadata(
status_code,
error_response_text.as_deref(),
failover_analysis,
),
status_code,
&headers,
error_response_text.as_deref(),
failover_analysis,
);
record_local_request_candidate_status(
state,
&plan,
error_flow_report_context
error_trace_report_context
.as_ref()
.or(report_context.as_ref()),
SchedulerRequestCandidateStatusUpdate {
@@ -1094,18 +1212,17 @@ async fn execute_stream_from_frame_stream(
)
{
let terminal_unix_secs = current_request_candidate_unix_ms();
let error_flow_report_context = with_error_flow_report_context(
let error_trace_report_context = with_stream_error_trace_context(
report_context.as_ref(),
build_local_error_flow_metadata(
status_code,
error_response_text.as_deref(),
failover_analysis,
),
status_code,
&headers,
error_response_text.as_deref(),
failover_analysis,
);
record_local_request_candidate_status(
state,
&plan,
error_flow_report_context
error_trace_report_context
.as_ref()
.or(report_context.as_ref()),
SchedulerRequestCandidateStatusUpdate {
@@ -1124,46 +1241,60 @@ async fn execute_stream_from_frame_stream(
return Ok(None);
}
let mut client_headers = headers.clone();
apply_endpoint_response_header_rules(state, &plan, &mut client_headers, body_json.as_ref())
.await?;
let mut client_headers = if (300..400).contains(&status_code) {
let mut headers = synthetic_error_response_headers(headers.clone());
headers.insert(
"x-aether-upstream-status".to_string(),
status_code.to_string(),
);
headers
} else {
headers.clone()
};
apply_endpoint_response_header_rules(
state,
&plan,
&mut client_headers,
client_body_json.as_ref(),
)
.await?;
let payload = build_stream_sync_payload(
let client_response_headers = client_headers.clone();
let error_trace_report_context = with_stream_error_trace_context(
report_context.as_ref(),
status_code,
&headers,
error_response_text.as_deref(),
failover_analysis,
);
let payload = build_stream_error_sync_payload(
trace_id,
stream_error_finalize_kind
.as_deref()
.or(report_kind.as_deref())
.unwrap_or_default()
.to_string(),
report_context,
error_trace_report_context.or(report_context),
status_code,
headers.clone(),
provider_body_json,
provider_body_base64,
client_headers,
body_json,
body_base64,
client_body_json,
None,
);
record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload);
let terminal_unix_secs = current_request_candidate_unix_ms();
let error_flow_report_context = with_error_flow_report_context(
payload.report_context.as_ref(),
build_local_error_flow_metadata(
status_code,
error_response_text.as_deref(),
failover_analysis,
),
);
record_local_request_candidate_status(
state,
&plan,
error_flow_report_context
.as_ref()
.or(payload.report_context.as_ref()),
payload.report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("execution_runtime_stream_error".to_string()),
error_type: Some("execution_runtime_stream_non_success_status".to_string()),
error_message: Some(format!(
"execution runtime stream returned error status {status_code}"
"execution runtime stream returned non-success status {status_code}"
)),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
@@ -1182,13 +1313,12 @@ async fn execute_stream_from_frame_stream(
)?));
}
return Ok(Some(attach_control_metadata_headers(
build_execution_runtime_error_response(
build_client_response_from_parts(
client_status_code,
&client_response_headers,
Body::from(client_error_body),
trace_id,
decision,
plan_kind,
status_code,
payload.headers,
error_body,
Some(decision),
)?,
Some(request_id),
candidate_id,
@@ -2343,7 +2473,9 @@ mod tests {
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTimeouts, RequestBody,
StandardizedUsage,
};
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
use aether_data_contracts::repository::usage::UsageReadRepository;
use aether_usage_runtime::UsageRuntimeConfig;
use async_stream::stream;
@@ -2532,9 +2664,15 @@ mod tests {
});
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_usage_data_repository_for_tests(Arc::clone(&usage_repository))
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
@@ -2736,6 +2874,213 @@ mod tests {
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_rewrites_redirect_to_structured_failure() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/v1/execute/stream",
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":302,\"headers\":{\"location\":\"/\",\"content-type\":\"text/html\",\"content-length\":\"0\"}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = axum::http::Response::new(Body::from(frames));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}),
);
axum::serve(listener, app)
.await
.expect("server should start");
});
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
})
.with_execution_runtime_override_base_url(format!("http://{addr}"));
let plan = ExecutionPlan {
request_id: "req-remote-runtime-stream-redirect".into(),
candidate_id: Some("cand-remote-runtime-stream-redirect".into()),
provider_name: Some("ChatGPTWeb".into()),
provider_id: "prov-redirect".into(),
endpoint_id: "ep-redirect".into(),
key_id: "key-redirect".into(),
method: "POST".into(),
url: "https://chatgpt.com/backend-api/codex/responses".into(),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
("accept".into(), "text/event-stream".into()),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5.4",
"input": "hello",
"stream": true
})),
stream: true,
client_api_format: "gemini:generate_content".into(),
provider_api_format: "openai:responses".into(),
model_name: Some("gemini-3.1-flash-image-preview".into()),
proxy: None,
transport_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let decision = GatewayControlDecision::synthetic(
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent",
Some("ai_public".to_string()),
Some("gemini".to_string()),
Some("generate_content".to_string()),
Some("gemini:generate_content".to_string()),
)
.with_execution_runtime_candidate(true);
let response = execute_execution_runtime_stream(
&state,
plan,
"trace-remote-runtime-stream-redirect",
&decision,
"gemini_chat_stream",
None,
Some(json!({
"request_id": "req-remote-runtime-stream-redirect",
"candidate_id": "cand-remote-runtime-stream-redirect",
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "openai:responses",
"client_api_format": "gemini:generate_content",
"needs_conversion": true
})),
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
assert_eq!(response.status(), axum::http::StatusCode::BAD_GATEWAY);
assert_eq!(
response
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("application/json")
);
assert_eq!(
response
.headers()
.get("x-aether-upstream-status")
.and_then(|value| value.to_str().ok()),
Some("302")
);
assert!(
response.headers().get(header::LOCATION).is_none(),
"redirect location should not be forwarded to AI clients"
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let body_json: Value =
serde_json::from_slice(&body).expect("response body should decode as json");
assert_eq!(
body_json["error"]["type"],
json!("execution_runtime_non_success_status")
);
assert_eq!(body_json["error"]["upstream_status"], json!(302));
assert_eq!(body_json["error"]["location"], json!("/"));
assert!(body_json["error"]["message"]
.as_str()
.is_some_and(|value| value.contains("non-success status 302")));
let usage = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if let Some(usage) = usage_repository
.find_by_request_id("req-remote-runtime-stream-redirect")
.await
.expect("usage should read")
.filter(|usage| usage.status == "failed")
{
break usage;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("usage should be written");
assert_eq!(usage.status_code, Some(302));
assert_eq!(usage.error_category.as_deref(), Some("redirect"));
assert!(usage
.error_message
.as_deref()
.is_some_and(|value| value.contains("non-success status 302")));
assert_eq!(
usage
.client_response_headers
.as_ref()
.and_then(|headers| headers.get("x-aether-upstream-status")),
Some(&json!("302"))
);
assert_eq!(
usage
.response_headers
.as_ref()
.and_then(|headers| headers.get("location")),
Some(&json!("/"))
);
assert!(
usage.response_body.is_none(),
"upstream redirect did not include a body"
);
assert_eq!(
usage
.client_response_body
.as_ref()
.and_then(|body| body.pointer("/error/upstream_status")),
Some(&json!(302))
);
let candidates = request_candidate_repository
.list_by_request_id("req-remote-runtime-stream-redirect")
.await
.expect("candidate trace should read");
let candidate_extra = candidates
.first()
.and_then(|candidate| candidate.extra_data.as_ref())
.expect("failed candidate extra_data should exist");
assert_eq!(
candidate_extra["upstream_response"]["status_code"],
json!(302)
);
assert_eq!(
candidate_extra["upstream_response"]["headers"]["location"],
json!("/")
);
assert!(candidate_extra["upstream_response"].get("body").is_none());
assert!(candidate_extra.get("client_response").is_none());
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_bridges_openai_image_sync_json_from_remote_runtime_to_image_sse(
) {

View File

@@ -20,9 +20,9 @@ use crate::execution_runtime::submission::{
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, resolve_local_failover_analysis_for_attempt,
LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalOAuthInvalidationEffect,
LocalPoolErrorEffect,
with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::request_candidate_runtime::record_report_request_candidate_status;
use crate::usage::submit_sync_report;
@@ -135,6 +135,15 @@ fn build_stream_failure_sync_payload(
failure: StreamFailureReport,
) -> GatewaySyncReportRequest {
let status_code = failure.status_code;
let report_context = with_upstream_response_report_context(
report_context.as_ref(),
status_code,
Some(&headers),
None,
None,
None,
)
.or(report_context);
headers.remove("content-encoding");
headers.remove("content-length");
headers.insert("content-type".to_string(), "application/json".to_string());

View File

@@ -39,9 +39,10 @@ use crate::execution_runtime::{
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, build_local_error_flow_metadata, with_error_flow_report_context,
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
with_upstream_response_report_context, LocalAdaptiveRateLimitEffect,
LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect,
LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
@@ -81,6 +82,27 @@ fn record_sync_terminal_usage(
.record_sync_terminal(state.data.as_ref(), context_seed, payload_seed);
}
fn with_sync_error_trace_context(
report_context: Option<&serde_json::Value>,
status_code: u16,
headers: &BTreeMap<String, String>,
response_text: Option<&str>,
local_failover_analysis: crate::orchestration::LocalFailoverAnalysis,
) -> Option<serde_json::Value> {
let upstream_context = with_upstream_response_report_context(
report_context,
status_code,
Some(headers),
None,
None,
None,
);
with_error_flow_report_context(
upstream_context.as_ref().or(report_context),
build_local_error_flow_metadata(status_code, response_text, local_failover_analysis),
)
}
fn build_sync_report_payload(
trace_id: &str,
report_kind: String,
@@ -572,18 +594,17 @@ pub(crate) async fn execute_execution_runtime_sync(
LocalFailoverDecision::RetryNextCandidate
) {
let terminal_unix_secs = current_request_candidate_unix_ms();
let error_flow_report_context = with_error_flow_report_context(
let error_trace_report_context = with_sync_error_trace_context(
report_context.as_ref(),
build_local_error_flow_metadata(
result.status_code,
local_failover_response_text.as_deref(),
local_failover_analysis,
),
result.status_code,
&headers,
local_failover_response_text.as_deref(),
local_failover_analysis,
);
record_local_request_candidate_status(
state,
&plan,
error_flow_report_context
error_trace_report_context
.as_ref()
.or(report_context.as_ref()),
SchedulerRequestCandidateStatusUpdate {
@@ -649,18 +670,17 @@ pub(crate) async fn execute_execution_runtime_sync(
mapped_error_finalize_kind.is_some(),
) {
let terminal_unix_secs = current_request_candidate_unix_ms();
let error_flow_report_context = with_error_flow_report_context(
let error_trace_report_context = with_sync_error_trace_context(
report_context.as_ref(),
build_local_error_flow_metadata(
result.status_code,
local_failover_response_text.as_deref(),
local_failover_analysis,
),
result.status_code,
&headers,
local_failover_response_text.as_deref(),
local_failover_analysis,
);
record_local_request_candidate_status(
state,
&plan,
error_flow_report_context
error_trace_report_context
.as_ref()
.or(report_context.as_ref()),
SchedulerRequestCandidateStatusUpdate {
@@ -680,13 +700,12 @@ pub(crate) async fn execute_execution_runtime_sync(
let terminal_unix_secs = current_request_candidate_unix_ms();
let error_flow_report_context = (result.status_code >= 400)
.then(|| {
with_error_flow_report_context(
with_sync_error_trace_context(
report_context.as_ref(),
build_local_error_flow_metadata(
result.status_code,
local_failover_response_text.as_deref(),
local_failover_analysis,
),
result.status_code,
&headers,
local_failover_response_text.as_deref(),
local_failover_analysis,
)
})
.flatten();

View File

@@ -0,0 +1,15 @@
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use axum::body::{Body, Bytes};
use axum::response::Response;
mod routes;
pub(crate) async fn maybe_build_local_admin_background_tasks_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
request_body: Option<&Bytes>,
) -> Result<Option<Response<Body>>, GatewayError> {
routes::maybe_build_local_admin_background_tasks_response(state, request_context, request_body)
.await
}

View File

@@ -0,0 +1,373 @@
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{
attach_admin_audit_response, query_param_value, unix_secs_to_rfc3339,
};
use crate::task_runtime::{
self, set_cancel_signal, TASK_KEY_PROVIDER_DELETE, TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
};
use crate::GatewayError;
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskKind, BackgroundTaskListQuery, BackgroundTaskStatus,
};
use axum::{
body::{Body, Bytes},
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
const DEFAULT_PAGE_SIZE: usize = 20;
const MAX_PAGE_SIZE: usize = 100;
const DEFAULT_EVENTS_PAGE_SIZE: usize = 50;
pub(super) async fn maybe_build_local_admin_background_tasks_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
request_body: Option<&Bytes>,
) -> Result<Option<Response<Body>>, GatewayError> {
if request_context.route_family() != Some("tasks_manage") {
return Ok(None);
}
match request_context.route_kind() {
Some("list_tasks") if request_context.method() == http::Method::GET => {
let query = request_context.query_string();
let page = query_param_value(query, "page")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(1)
.max(1);
let page_size = query_param_value(query, "page_size")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(DEFAULT_PAGE_SIZE)
.clamp(1, MAX_PAGE_SIZE);
let kind = query_param_value(query, "kind")
.map(|value| BackgroundTaskKind::from_database(value.as_str()))
.transpose()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let status = query_param_value(query, "status")
.map(|value| BackgroundTaskStatus::from_database(value.as_str()))
.transpose()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let trigger = query_param_value(query, "trigger");
let task_key_substring = query_param_value(query, "task_key");
let offset = (page - 1).saturating_mul(page_size);
let response = state
.list_background_task_runs(&BackgroundTaskListQuery {
task_key_substring,
kind,
status,
trigger,
offset,
limit: page_size,
})
.await?;
let pages = if response.total == 0 {
0
} else {
(response.total + page_size - 1) / page_size
};
let items = response
.items
.iter()
.map(|run| {
json!({
"id": run.id,
"task_key": run.task_key,
"kind": run.kind.as_database(),
"trigger": run.trigger,
"status": run.status.as_database(),
"attempt": run.attempt,
"max_attempts": run.max_attempts,
"owner_instance": run.owner_instance,
"progress_percent": run.progress_percent,
"progress_message": run.progress_message,
"payload": run.payload_json,
"result": run.result_json,
"error_message": run.error_message,
"cancel_requested": run.cancel_requested,
"created_by": run.created_by,
"created_at": unix_secs_to_rfc3339(run.created_at_unix_secs),
"started_at": run.started_at_unix_secs.and_then(unix_secs_to_rfc3339),
"finished_at": run.finished_at_unix_secs.and_then(unix_secs_to_rfc3339),
"updated_at": unix_secs_to_rfc3339(run.updated_at_unix_secs),
})
})
.collect::<Vec<_>>();
let definitions = task_runtime::task_definitions()
.iter()
.map(|definition| {
json!({
"task_key": definition.key,
"kind": definition.kind.as_str(),
"trigger": definition.trigger,
"max_attempts": definition.retry_policy.max_attempts,
"singleton": definition.singleton,
"persist_history": definition.persist_history,
})
})
.collect::<Vec<_>>();
return Ok(Some(
Json(json!({
"items": items,
"total": response.total,
"page": page,
"page_size": page_size,
"pages": pages,
"definitions": definitions,
}))
.into_response(),
));
}
Some("stats") if request_context.method() == http::Method::GET => {
let stats = state.summarize_background_task_runs().await?;
return Ok(Some(
Json(json!({
"total": stats.total,
"running_count": stats.running_count,
"by_status": stats.by_status,
"by_kind": stats.by_kind,
"registered_tasks": task_runtime::task_definitions().len(),
}))
.into_response(),
));
}
Some("detail") if request_context.method() == http::Method::GET => {
let Some(run_id) = task_id_from_path(request_context.path()) else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
let Some(run) = state.find_background_task_run(run_id).await? else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
return Ok(Some(attach_admin_audit_response(
Json(json!({
"id": run.id,
"task_key": run.task_key,
"kind": run.kind.as_database(),
"trigger": run.trigger,
"status": run.status.as_database(),
"attempt": run.attempt,
"max_attempts": run.max_attempts,
"owner_instance": run.owner_instance,
"progress_percent": run.progress_percent,
"progress_message": run.progress_message,
"payload": run.payload_json,
"result": run.result_json,
"error_message": run.error_message,
"cancel_requested": run.cancel_requested,
"created_by": run.created_by,
"created_at": unix_secs_to_rfc3339(run.created_at_unix_secs),
"started_at": run.started_at_unix_secs.and_then(unix_secs_to_rfc3339),
"finished_at": run.finished_at_unix_secs.and_then(unix_secs_to_rfc3339),
"updated_at": unix_secs_to_rfc3339(run.updated_at_unix_secs),
}))
.into_response(),
"admin_task_detail_viewed",
"view_task_detail",
"background_task",
run_id,
)));
}
Some("events") if request_context.method() == http::Method::GET => {
let Some(run_id) = nested_task_id_from_path(request_context.path(), "/events") else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
let query = request_context.query_string();
let page = query_param_value(query, "page")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(1)
.max(1);
let page_size = query_param_value(query, "page_size")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(DEFAULT_EVENTS_PAGE_SIZE)
.clamp(1, MAX_PAGE_SIZE);
let offset = (page - 1).saturating_mul(page_size);
let events = state
.list_background_task_events(run_id, offset, page_size)
.await?;
return Ok(Some(
Json(json!({
"items": events.into_iter().map(|event| {
json!({
"id": event.id,
"run_id": event.run_id,
"event_type": event.event_type,
"message": event.message,
"payload": event.payload_json,
"created_at": unix_secs_to_rfc3339(event.created_at_unix_secs),
})
}).collect::<Vec<_>>(),
"page": page,
"page_size": page_size,
}))
.into_response(),
));
}
Some("cancel") if request_context.method() == http::Method::POST => {
let Some(run_id) = nested_task_id_from_path(request_context.path(), "/cancel") else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
let now = task_runtime::now_unix_secs();
let cancelled = state
.request_cancel_background_task_run(run_id, now)
.await?;
if !cancelled {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({ "detail": "Task not found" })),
)
.into_response(),
));
}
let _ = set_cancel_signal(state.app(), run_id).await;
task_runtime::append_event_with_logging(
state.app(),
run_id,
"cancel_requested",
"cancel requested by admin",
None,
)
.await;
return Ok(Some(attach_admin_audit_response(
Json(json!({
"id": run_id,
"status": "cancel_requested",
"message": "Task cancellation requested",
}))
.into_response(),
"admin_task_cancel_requested",
"cancel_task",
"background_task",
run_id,
)));
}
Some("trigger") if request_context.method() == http::Method::POST => {
let Some(task_key) = nested_task_id_from_path(request_context.path(), "/trigger")
else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Task not found"})),
)
.into_response(),
));
};
let payload = parse_json_payload(request_body)?;
if task_key == TASK_KEY_PROVIDER_DELETE {
let provider_id = payload
.get("provider_id")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| {
GatewayError::Internal(
"admin task trigger provider delete requires provider_id".to_string(),
)
})?;
let Some(run_id) =
task_runtime::submit_provider_delete_task(state, provider_id, Some("admin"))
.await?
else {
return Ok(Some(
(
http::StatusCode::NOT_FOUND,
Json(json!({"detail":"Provider 不存在"})),
)
.into_response(),
));
};
return Ok(Some(attach_admin_audit_response(
Json(json!({
"task_key": task_key,
"run_id": run_id,
"status": "queued",
}))
.into_response(),
"admin_task_triggered",
"trigger_task",
"background_task",
task_key,
)));
}
if task_key == TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT {
return Ok(Some(
(
http::StatusCode::BAD_REQUEST,
Json(json!({
"detail": "请使用 provider oauth batch import 专用接口触发该任务",
})),
)
.into_response(),
));
}
return Ok(Some(
(
http::StatusCode::BAD_REQUEST,
Json(json!({
"detail": format!("Unsupported task_key: {task_key}"),
})),
)
.into_response(),
));
}
_ => {}
}
Ok(None)
}
fn task_id_from_path(request_path: &str) -> Option<&str> {
let task_id = request_path.strip_prefix("/api/admin/tasks/")?;
if task_id.is_empty() || task_id.contains('/') || task_id == "stats" {
return None;
}
Some(task_id)
}
fn nested_task_id_from_path<'a>(request_path: &'a str, suffix: &str) -> Option<&'a str> {
let task_id = request_path
.strip_prefix("/api/admin/tasks/")?
.strip_suffix(suffix)?;
if task_id.is_empty() || task_id.contains('/') {
return None;
}
Some(task_id)
}
fn parse_json_payload(request_body: Option<&Bytes>) -> Result<serde_json::Value, GatewayError> {
let Some(body) = request_body else {
return Ok(json!({}));
};
if body.is_empty() {
return Ok(json!({}));
}
serde_json::from_slice::<serde_json::Value>(body)
.map_err(|err| GatewayError::Internal(format!("invalid json body: {err}")))
}

View File

@@ -1,7 +1,9 @@
mod background_tasks;
mod gemini_files;
mod routes;
mod video_tasks;
pub(super) use self::background_tasks::maybe_build_local_admin_background_tasks_response;
pub(super) use self::gemini_files::maybe_build_local_admin_gemini_files_response;
pub(super) use self::routes::maybe_build_local_admin_features_response;
pub(crate) use self::video_tasks::maybe_build_local_admin_video_tasks_response;

View File

@@ -1,9 +1,19 @@
use super::{gemini_files, video_tasks};
use super::{background_tasks, gemini_files, video_tasks};
use crate::handlers::admin::request::{AdminRouteRequest, AdminRouteResult};
pub(crate) async fn maybe_build_local_admin_features_response(
request: AdminRouteRequest<'_>,
) -> AdminRouteResult {
if let Some(response) = background_tasks::maybe_build_local_admin_background_tasks_response(
&request.state(),
&request.request_context(),
request.request_body(),
)
.await?
{
return Ok(Some(response));
}
if let Some(response) = video_tasks::maybe_build_local_admin_video_tasks_response(
&request.state(),
&request.request_context(),

View File

@@ -356,6 +356,164 @@ async fn admin_monitoring_trace_request_enriches_proxy_timing_from_usage_audit()
payload["candidates"][0]["extra_data"]["proxy"]["timing"]["response_wait_ms"],
json!(475)
);
assert!(payload["candidates"][0]["extra_data"]
.get("upstream_response")
.is_none());
}
#[tokio::test]
async fn admin_monitoring_trace_request_exposes_request_path_from_usage_audit() {
let mut candidate = sample_candidate(
"cand-used",
"request-1",
0,
RequestCandidateStatus::Failed,
Some(101),
Some(33),
Some(403),
);
candidate.extra_data = Some(json!({
"client_api_format": "gemini:generate_content"
}));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![candidate]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let mut usage = sample_usage(
"request-1",
"provider-1",
"OpenAI",
40,
0.02,
"failed",
Some(403),
100,
);
usage.candidate_id = Some("cand-used".to_string());
usage.request_metadata = Some(json!({
"request_path": "/v1beta/models/gemini-2.5-pro:generateContent",
"request_query_string": "alt=sse"
}));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
let data_state =
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
request_candidates,
usage_repository,
)
.with_provider_catalog_reader(provider_catalog);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let context = request_context(http::Method::GET, "/api/admin/monitoring/trace/request-1");
let response = local_monitoring_response(&state, &context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(response.status(), http::StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(
payload["request_path"],
json!("/v1beta/models/gemini-2.5-pro:generateContent")
);
assert_eq!(payload["request_query_string"], json!("alt=sse"));
assert_eq!(
payload["request_path_and_query"],
json!("/v1beta/models/gemini-2.5-pro:generateContent?alt=sse")
);
assert_eq!(
payload["candidates"][0]["extra_data"]["request_path_and_query"],
json!("/v1beta/models/gemini-2.5-pro:generateContent?alt=sse")
);
}
#[tokio::test]
async fn admin_monitoring_trace_request_exposes_failed_candidate_upstream_response_boundary() {
let candidate = sample_candidate(
"cand-used",
"request-1",
0,
RequestCandidateStatus::Failed,
Some(101),
Some(33),
Some(302),
);
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![candidate]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let mut usage = sample_usage(
"request-1",
"provider-1",
"OpenAI",
40,
0.02,
"failed",
Some(302),
100,
);
usage.candidate_id = Some("cand-used".to_string());
usage.response_headers = Some(json!({
"location": "/",
"content-type": "text/html"
}));
usage.client_response_headers = Some(json!({
"content-type": "application/json",
"x-aether-upstream-status": "302"
}));
usage.client_response_body = Some(json!({
"error": {
"type": "execution_runtime_non_success_status",
"message": "execution runtime stream returned non-success status 302",
"upstream_status": 302,
"location": "/"
}
}));
usage.request_metadata = Some(json!({
"client_response_status_code": 502
}));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
let data_state =
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
request_candidates,
usage_repository,
)
.with_provider_catalog_reader(provider_catalog);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let context = request_context(http::Method::GET, "/api/admin/monitoring/trace/request-1");
let response = local_monitoring_response(&state, &context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(response.status(), http::StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
let extra = &payload["candidates"][0]["extra_data"];
assert_eq!(extra["upstream_response"]["status_code"], json!(302));
assert_eq!(
extra["upstream_response"]["headers"]["location"],
json!("/")
);
assert!(extra["upstream_response"]["body"].is_null());
assert!(extra.get("client_response").is_none());
assert!(extra.get("provider_response").is_none());
}
#[tokio::test]

View File

@@ -1,11 +1,10 @@
use crate::handlers::admin::provider::delete_task::run_admin_provider_delete_task;
use crate::handlers::admin::provider::shared::paths::{
admin_provider_delete_task_parts, admin_provider_id_for_manage_path,
};
use crate::handlers::admin::provider::shared::support::build_admin_provider_delete_task_payload;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::attach_admin_audit_response;
use crate::{GatewayError, LocalProviderDeleteTaskState};
use crate::GatewayError;
use axum::{
body::Body,
http,
@@ -13,8 +12,6 @@ use axum::{
Json,
};
use serde_json::json;
use tracing::warn;
use uuid::Uuid;
fn build_admin_provider_not_found_response(detail: impl Into<String>) -> Response<Body> {
(
@@ -60,52 +57,18 @@ pub(crate) async fn maybe_build_local_admin_provider_delete_task_response(
"Provider 不存在",
)));
};
let Some(provider) = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
.await?
.into_iter()
.next()
let Some(task_id) =
crate::task_runtime::submit_provider_delete_task(state, &provider_id, Some("admin"))
.await?
else {
return Ok(Some(build_admin_provider_not_found_response(
"提供商不存在",
)));
};
let task_id = Uuid::new_v4().simple().to_string()[..16].to_string();
let pending_task = LocalProviderDeleteTaskState {
task_id: task_id.clone(),
provider_id: provider.id.clone(),
status: "pending".to_string(),
stage: "queued".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "delete task submitted".to_string(),
};
state.put_provider_delete_task(pending_task.clone());
if let Err(err) = state
.run_admin_provider_delete_task(&provider.id, &task_id)
.await
{
warn!(
"gateway admin provider delete task failed for provider {}: {:?}",
provider.id, err
);
state.put_provider_delete_task(LocalProviderDeleteTaskState {
task_id: task_id.clone(),
provider_id: provider.id.clone(),
status: "failed".to_string(),
stage: "failed".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: format!("provider delete failed: {err:?}"),
});
}
return Ok(Some(attach_admin_audit_response(
Json(json!({
"task_id": task_id,
"run_id": task_id,
"status": "pending",
"message": "删除任务已提交,提供商已进入后台删除队列",
}))
@@ -113,7 +76,7 @@ pub(crate) async fn maybe_build_local_admin_provider_delete_task_response(
"admin_provider_delete_queued",
"delete_provider",
"provider",
&provider.id,
&provider_id,
)));
}

View File

@@ -15,7 +15,14 @@ use crate::handlers::admin::provider::oauth::state::{
};
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_task_provider_id;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::task_runtime::{
append_event_with_logging, now_unix_secs, task_definition, update_run_status,
upsert_run_with_logging, TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
};
use crate::GatewayError;
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskKind, BackgroundTaskStatus, UpsertBackgroundTaskRun,
};
use axum::{
body::Bytes,
http,
@@ -133,11 +140,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
}
let task_id = Uuid::new_v4().to_string();
let created_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
.unwrap_or(0);
let created_at = now_unix_secs();
let submitted_state = build_admin_provider_oauth_batch_task_state(
&task_id,
&provider_id,
@@ -167,6 +170,50 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
));
}
if state.has_background_task_data_writer() {
let max_attempts = task_definition(TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT)
.map(|item| item.retry_policy.max_attempts)
.unwrap_or(1);
let run = UpsertBackgroundTaskRun {
id: task_id.clone(),
task_key: TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT.to_string(),
kind: BackgroundTaskKind::OnDemand,
trigger: "manual".to_string(),
status: BackgroundTaskStatus::Queued,
attempt: 1,
max_attempts,
owner_instance: Some(state.app().tunnel.local_instance_id().to_string()),
progress_percent: 0,
progress_message: Some("provider oauth batch import queued".to_string()),
payload_json: Some(json!({
"provider_id": provider_id.clone(),
"provider_type": provider_type.clone(),
"total": total,
})),
result_json: None,
error_message: None,
cancel_requested: false,
created_by: Some("admin".to_string()),
created_at_unix_secs: created_at,
started_at_unix_secs: None,
finished_at_unix_secs: None,
updated_at_unix_secs: created_at,
};
let _ = upsert_run_with_logging(state.app(), run).await;
append_event_with_logging(
state.app(),
&task_id,
"queued",
"provider oauth batch import queued",
Some(json!({
"provider_id": provider_id.clone(),
"provider_type": provider_type.clone(),
"total": total,
})),
)
.await;
}
let task_state = state.cloned_app();
let task_id_for_worker = task_id.clone();
let provider_id_for_worker = provider_id.clone();
@@ -201,6 +248,27 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &processing_state)
.await;
let _ = update_run_status(
&task_state,
&task_id_for_worker,
BackgroundTaskStatus::Running,
Some(1),
Some("provider oauth batch import started".to_string()),
None,
None,
Some(started_at),
None,
)
.await;
append_event_with_logging(
&task_state,
&task_id_for_worker,
"running",
"provider oauth batch import started",
None,
)
.await;
let mut progress_reporter = BatchTaskProgressReporter {
app: task_state.clone(),
task_id: task_id_for_worker.clone(),
@@ -270,6 +338,34 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
let _ = AdminAppState::new(&task_state)
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &completed_state)
.await;
let _ = update_run_status(
&task_state,
&task_id_for_worker,
BackgroundTaskStatus::Succeeded,
Some(100),
Some(message),
Some(json!({
"provider_id": provider_id_for_worker,
"provider_type": provider_type_for_worker,
"total": outcome.total,
"success": outcome.success,
"failed": outcome.failed,
"created_count": created_count,
"replaced_count": replaced_count,
})),
None,
None,
Some(finished_at),
)
.await;
append_event_with_logging(
&task_state,
&task_id_for_worker,
"succeeded",
"provider oauth batch import completed",
None,
)
.await;
}
Err(err) => {
let finished_at = SystemTime::now()
@@ -299,6 +395,26 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
let _ = AdminAppState::new(&task_state)
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &failed_state)
.await;
let _ = update_run_status(
&task_state,
&task_id_for_worker,
BackgroundTaskStatus::Failed,
Some(100),
Some("provider oauth batch import failed".to_string()),
None,
Some(error_message.clone()),
None,
Some(finished_at),
)
.await;
append_event_with_logging(
&task_state,
&task_id_for_worker,
"failed",
"provider oauth batch import failed",
Some(json!({ "error": error_message.clone() })),
)
.await;
tracing::warn!(
task_id = %task_id_for_worker,
provider_id = %provider_id_for_worker,

View File

@@ -4,6 +4,7 @@ use super::quota::codex::refresh_codex_provider_quota_locally;
use super::quota::kiro::refresh_kiro_provider_quota_locally;
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH};
use crate::{AppState, GatewayError};
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
@@ -170,7 +171,7 @@ pub(crate) fn spawn_provider_oauth_account_state_refresh_after_update(
key_id: String,
proxy_override: Option<ProxySnapshot>,
) {
tokio::spawn(async move {
spawn_fire_and_forget(TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH, async move {
let _ = refresh_provider_oauth_account_state_after_update(
&AdminAppState::new(&app),
&provider,

View File

@@ -1,5 +1,6 @@
use super::actions::admin_provider_ops_local_action_response;
use crate::handlers::admin::request::AdminAppState;
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_PROVIDER_BALANCE_REFRESH};
use serde_json::{json, Value};
use std::collections::HashSet;
use std::time::Duration;
@@ -130,7 +131,7 @@ pub(super) async fn spawn_admin_provider_ops_balance_refresh(
let app = state.cloned_app();
let provider_id = provider_id.to_string();
tokio::spawn(async move {
spawn_fire_and_forget(TASK_KEY_PROVIDER_BALANCE_REFRESH, async move {
let permit = match tokio::time::timeout(
Duration::from_secs(5),
ADMIN_PROVIDER_OPS_BALANCE_REFRESH_SEMAPHORE.acquire(),

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

@@ -76,6 +76,14 @@ impl<'a> AdminAppState<'a> {
self.app.has_gemini_file_mapping_data_writer()
}
pub(crate) fn has_background_task_data_reader(&self) -> bool {
self.app.has_background_task_data_reader()
}
pub(crate) fn has_background_task_data_writer(&self) -> bool {
self.app.has_background_task_data_writer()
}
pub(crate) fn has_auth_api_key_data_reader(&self) -> bool {
self.app.has_auth_api_key_data_reader()
}

View File

@@ -2,6 +2,59 @@ use super::{AdminAppState, AdminCancelVideoTaskError};
use crate::GatewayError;
impl<'a> AdminAppState<'a> {
pub(crate) async fn list_background_task_runs(
&self,
query: &aether_data_contracts::repository::background_tasks::BackgroundTaskListQuery,
) -> Result<
aether_data_contracts::repository::background_tasks::StoredBackgroundTaskRunPage,
GatewayError,
> {
self.app.list_background_task_runs(query).await
}
pub(crate) async fn find_background_task_run(
&self,
run_id: &str,
) -> Result<
Option<aether_data_contracts::repository::background_tasks::StoredBackgroundTaskRun>,
GatewayError,
> {
self.app.find_background_task_run(run_id).await
}
pub(crate) async fn list_background_task_events(
&self,
run_id: &str,
offset: usize,
limit: usize,
) -> Result<
Vec<aether_data_contracts::repository::background_tasks::StoredBackgroundTaskEvent>,
GatewayError,
> {
self.app
.list_background_task_events(run_id, offset, limit)
.await
}
pub(crate) async fn summarize_background_task_runs(
&self,
) -> Result<
aether_data_contracts::repository::background_tasks::BackgroundTaskSummary,
GatewayError,
> {
self.app.summarize_background_task_runs().await
}
pub(crate) async fn request_cancel_background_task_run(
&self,
run_id: &str,
updated_at_unix_secs: u64,
) -> Result<bool, GatewayError> {
self.app
.request_cancel_background_task_run(run_id, updated_at_unix_secs)
.await
}
pub(crate) async fn list_gemini_file_mappings(
&self,
query: &aether_data::repository::gemini_file_mappings::GeminiFileMappingListQuery,

View File

@@ -516,6 +516,13 @@ fn imported_optional_f64(value: Option<&Value>, field_name: &str) -> Result<Opti
.filter(|value| value.is_finite())
.ok_or_else(|| format!("{field_name} 必须是有限数值"))
.map(Some),
Some(Value::String(value)) => value
.trim()
.parse::<f64>()
.ok()
.filter(|value| value.is_finite())
.ok_or_else(|| format!("{field_name} 必须是有限数值"))
.map(Some),
_ => Err(format!("{field_name} 必须是有限数值")),
}
}
@@ -542,9 +549,14 @@ fn imported_rfc3339_to_unix_secs(
let Some(value) = imported_optional_string(value)? else {
return Ok(None);
};
let parsed = chrono::DateTime::parse_from_rfc3339(&value)
let parsed_timestamp = chrono::DateTime::parse_from_rfc3339(&value)
.map(|parsed| parsed.timestamp())
.or_else(|_| {
chrono::NaiveDateTime::parse_from_str(&value, "%Y-%m-%dT%H:%M:%S%.f")
.map(|parsed| parsed.and_utc().timestamp())
})
.map_err(|_| format!("{field_name} 必须是 RFC3339 时间"))?;
Ok(Some(parsed.timestamp().max(0) as u64))
Ok(Some(parsed_timestamp.max(0) as u64))
}
fn imported_string_list_from_value(
@@ -609,12 +621,11 @@ fn normalize_imported_wallet_target(
} else {
let total_balance =
imported_optional_f64(map.get("balance"), "wallet.balance")?.unwrap_or(0.0);
(total_balance - gift_balance).max(0.0)
total_balance - gift_balance
}
} else {
0.0
}
.max(0.0);
};
let limit_mode = if let Some(map) = wallet {
if let Some(mode) = imported_optional_string(map.get("limit_mode"))? {
match mode.to_ascii_lowercase().as_str() {
@@ -2436,8 +2447,8 @@ mod tests {
use super::{
imported_optional_bool, imported_optional_f64, imported_optional_i32,
imported_optional_u64, imported_string_list_from_value,
validate_imported_system_users_export_version,
imported_optional_u64, imported_rfc3339_to_unix_secs, imported_string_list_from_value,
normalize_imported_wallet_target, validate_imported_system_users_export_version,
};
#[test]
@@ -2454,7 +2465,7 @@ mod tests {
}
#[test]
fn import_rejects_legacy_string_scalars() {
fn import_handles_legacy_string_scalars() {
assert_eq!(
imported_optional_bool(Some(&json!("true"))).unwrap_err(),
"字段必须是布尔值"
@@ -2468,11 +2479,63 @@ mod tests {
"total_requests 必须是非负整数"
);
assert_eq!(
imported_optional_f64(Some(&json!("1.25")), "total_cost_usd").unwrap_err(),
imported_optional_f64(Some(&json!("1.25000000")), "total_cost_usd").unwrap(),
Some(1.25)
);
assert_eq!(
imported_optional_f64(Some(&json!("not-a-number")), "total_cost_usd").unwrap_err(),
"total_cost_usd 必须是有限数值"
);
}
#[test]
fn import_handles_python_isoformat_timestamps() {
assert_eq!(
imported_rfc3339_to_unix_secs(Some(&json!("2099-01-01T00:00:00+00:00")), "expires_at")
.unwrap(),
Some(4_070_908_800)
);
assert_eq!(
imported_rfc3339_to_unix_secs(Some(&json!("2099-01-01T00:00:00")), "expires_at")
.unwrap(),
Some(4_070_908_800)
);
assert_eq!(
imported_rfc3339_to_unix_secs(Some(&json!("invalid")), "expires_at").unwrap_err(),
"expires_at 必须是 RFC3339 时间"
);
}
#[test]
fn import_preserves_python_wallet_negative_recharge_balance() {
let wallet = json!({
"balance": -4.5,
"recharge_balance": -5.25,
"gift_balance": 0.75,
"limit_mode": "finite"
});
let wallet = wallet.as_object().expect("wallet should be object");
let target = normalize_imported_wallet_target(Some(wallet), false).unwrap();
assert_eq!(target.recharge_balance, -5.25);
assert_eq!(target.gift_balance, 0.75);
assert_eq!(target.total_recharged, -5.25);
}
#[test]
fn import_preserves_python_wallet_negative_balance_fallback() {
let wallet = json!({
"balance": -4.5,
"gift_balance": 0.75,
"limit_mode": "finite"
});
let wallet = wallet.as_object().expect("wallet should be object");
let target = normalize_imported_wallet_target(Some(wallet), false).unwrap();
assert_eq!(target.recharge_balance, -5.25);
assert_eq!(target.gift_balance, 0.75);
}
#[test]
fn import_rejects_legacy_string_lists() {
assert_eq!(

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

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

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

@@ -28,6 +28,8 @@ mod support_announcements;
mod support_auth;
#[path = "support/dashboard.rs"]
mod support_dashboard;
#[path = "support/install.rs"]
mod support_install;
#[path = "support/models.rs"]
mod support_models;
#[path = "support/monitoring.rs"]
@@ -61,6 +63,10 @@ use self::support_auth::{
build_auth_settings_payload, extract_client_device_id, maybe_build_local_auth_response,
};
use self::support_dashboard::maybe_build_local_dashboard_response;
use self::support_install::{
handle_users_me_api_key_install_session_create, maybe_build_local_install_response,
users_me_api_key_install_sessions_path_matches,
};
use self::support_models::{
build_models_auth_error_response, maybe_build_local_models_response, models_api_format,
};
@@ -146,6 +152,10 @@ pub(crate) async fn maybe_build_local_public_support_response(
.await;
}
if decision.route_family.as_deref() == Some("install") {
return maybe_build_local_install_response(state, request_context).await;
}
if decision.route_family.as_deref() == Some("payment_callback") {
return maybe_build_local_payment_callback_response(
state,

View File

@@ -0,0 +1,561 @@
use axum::{
body::Body,
http,
response::{IntoResponse, Response},
Json,
};
use serde::{Deserialize, Serialize};
use serde_json::json;
use super::{
build_auth_error_response, decrypt_catalog_secret_with_fallbacks,
resolve_authenticated_local_user, AppState, GatewayPublicRequestContext,
};
const INSTALL_SESSION_TTL_SECS: u64 = 15 * 60;
const INSTALL_SESSION_KEY_PREFIX: &str = "install:session:";
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
enum InstallTargetCli {
ClaudeCode,
CodexCli,
GeminiCli,
}
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
enum InstallTargetSystem {
Macos,
Linux,
Windows,
Auto,
}
#[derive(Debug, Deserialize)]
struct UsersMeCreateInstallSessionRequest {
target_cli: InstallTargetCli,
target_system: InstallTargetSystem,
}
#[derive(Debug, Serialize, Deserialize)]
struct StoredInstallSession {
api_key_id: String,
api_key_name: String,
api_key: String,
base_url: String,
target_cli: InstallTargetCli,
target_system: InstallTargetSystem,
expires_at_unix_secs: u64,
}
pub(super) fn users_me_api_key_install_sessions_path_matches(request_path: &str) -> bool {
users_me_api_key_install_session_id_from_path(request_path).is_some()
}
fn users_me_api_key_install_session_id_from_path(request_path: &str) -> Option<String> {
let raw = request_path
.strip_prefix("/api/users/me/api-keys/")?
.trim()
.trim_matches('/');
let mut segments = raw.split('/').map(str::trim);
let api_key_id = segments.next()?.to_string();
let suffix = segments.next()?;
(suffix == "install-sessions" && segments.next().is_none()).then_some(api_key_id)
}
fn install_code_from_path(request_path: &str) -> Option<(String, bool)> {
let raw = request_path
.strip_prefix("/install/")
.or_else(|| request_path.strip_prefix("/i/"))?
.trim()
.trim_matches('/');
if raw.is_empty() || raw.contains('/') {
return None;
}
let is_powershell = raw.ends_with(".ps1");
let code = raw.strip_suffix(".ps1").unwrap_or(raw).trim();
(!code.is_empty()).then(|| (code.to_string(), is_powershell))
}
fn install_session_runtime_key(code: &str) -> String {
format!("{INSTALL_SESSION_KEY_PREFIX}{code}")
}
fn generate_install_code() -> String {
uuid::Uuid::new_v4()
.simple()
.to_string()
.chars()
.take(24)
.collect()
}
fn unix_secs_now() -> u64 {
chrono::Utc::now().timestamp().max(0) as u64
}
fn base_url_from_request(
headers: &http::HeaderMap,
request_context: &GatewayPublicRequestContext,
) -> String {
if let Some(value) = std::env::var("AETHER_PUBLIC_BASE_URL")
.ok()
.or_else(|| std::env::var("PUBLIC_BASE_URL").ok())
.map(|value| value.trim().trim_end_matches('/').to_string())
.filter(|value| value.starts_with("https://") || value.starts_with("http://"))
{
return value;
}
let host = crate::headers::header_value_str(headers, "x-forwarded-host")
.or_else(|| request_context.host_header.clone())
.map(|value| value.trim().trim_end_matches('/').to_string())
.filter(|value| {
!value.is_empty()
&& !value.contains('/')
&& !value.contains('\\')
&& !value.contains('@')
&& !value.contains(char::is_whitespace)
})
.unwrap_or_else(|| "localhost".to_string());
let proto = crate::headers::header_value_str(headers, "x-forwarded-proto")
.map(|value| value.trim().trim_end_matches(':').to_ascii_lowercase())
.filter(|value| value == "http" || value == "https")
.unwrap_or_else(|| "http".to_string());
format!("{proto}://{host}")
}
fn shell_single_quote(value: &str) -> String {
format!("'{}'", value.replace('\'', "'\\''"))
}
fn powershell_single_quote(value: &str) -> String {
format!("'{}'", value.replace('\'', "''"))
}
fn cli_label(target_cli: InstallTargetCli) -> &'static str {
match target_cli {
InstallTargetCli::ClaudeCode => "Claude Code",
InstallTargetCli::CodexCli => "Codex CLI",
InstallTargetCli::GeminiCli => "Gemini CLI",
}
}
fn system_label(target_system: InstallTargetSystem) -> &'static str {
match target_system {
InstallTargetSystem::Macos => "macOS",
InstallTargetSystem::Linux => "Linux",
InstallTargetSystem::Windows => "Windows",
InstallTargetSystem::Auto => "Auto",
}
}
fn npm_package(target_cli: InstallTargetCli) -> &'static str {
match target_cli {
InstallTargetCli::ClaudeCode => "@anthropic-ai/claude-code",
InstallTargetCli::CodexCli => "@openai/codex",
InstallTargetCli::GeminiCli => "@google/gemini-cli",
}
}
fn cli_binary(target_cli: InstallTargetCli) -> &'static str {
match target_cli {
InstallTargetCli::ClaudeCode => "claude",
InstallTargetCli::CodexCli => "codex",
InstallTargetCli::GeminiCli => "gemini",
}
}
fn build_unix_script(session: &StoredInstallSession) -> String {
let target_cli = match session.target_cli {
InstallTargetCli::ClaudeCode => "claude_code",
InstallTargetCli::CodexCli => "codex_cli",
InstallTargetCli::GeminiCli => "gemini_cli",
};
let target_system = match session.target_system {
InstallTargetSystem::Macos => "macos",
InstallTargetSystem::Linux => "linux",
InstallTargetSystem::Windows => "windows",
InstallTargetSystem::Auto => "auto",
};
format!(
r###"#!/bin/sh
set -eu
TARGET_CLI={target_cli}
TARGET_SYSTEM={target_system}
AETHER_BASE_URL={base_url}
AETHER_API_KEY={api_key}
CLI_LABEL={label}
CLI_BIN={binary}
NPM_PACKAGE={npm_package}
say() {{ printf '%s\n' "[Aether] $1"; }}
fail() {{ printf '%s\n' "[Aether] $1" >&2; exit 1; }}
os="$(uname -s 2>/dev/null || printf unknown)"
case "$os" in
Darwin) actual_system=macos ;;
Linux) actual_system=linux ;;
MINGW*|MSYS*|CYGWIN*) fail "检测到 Windows shell请在 PowerShell 中使用 Windows 命令irm <url>.ps1 | iex" ;;
*) fail "不支持的系统:$os" ;;
esac
if [ "$TARGET_SYSTEM" = "windows" ]; then
fail "该 install code 绑定 Windows请复制 PowerShell 命令执行。"
fi
if [ "$TARGET_SYSTEM" != "auto" ] && [ "$TARGET_SYSTEM" != "$actual_system" ]; then
fail "所选系统 $TARGET_SYSTEM 与当前系统 $actual_system 不一致,请回到 Aether 重新选择目标系统。"
fi
say "准备安装/复用 $CLI_LABEL"
if ! command -v "$CLI_BIN" >/dev/null 2>&1; then
command -v npm >/dev/null 2>&1 || fail "未找到 $CLI_BIN也未找到 npm。请先安装 Node.js/npm 后重试。"
say "未找到 $CLI_BIN正在通过 npm 安装 $NPM_PACKAGE"
npm install -g "$NPM_PACKAGE"
fi
umask 077
mkdir -p "$HOME/.aether"
cat > "$HOME/.aether/client.env" <<EOF
AETHER_BASE_URL=$AETHER_BASE_URL
AETHER_API_KEY=$AETHER_API_KEY
EOF
chmod 600 "$HOME/.aether/client.env" 2>/dev/null || true
case "$TARGET_CLI" in
claude_code)
mkdir -p "$HOME/.claude"
python3 - "$HOME/.claude/settings.json" "$AETHER_BASE_URL" "$AETHER_API_KEY" <<'PY'
import json, pathlib, sys
path = pathlib.Path(sys.argv[1])
data = json.loads(path.read_text() or '{{}}') if path.exists() else {{}}
env = data.setdefault('env', {{}})
env['ANTHROPIC_BASE_URL'] = sys.argv[2]
env['ANTHROPIC_AUTH_TOKEN'] = sys.argv[3]
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + '\n')
PY
chmod 600 "$HOME/.claude/settings.json" 2>/dev/null || true
;;
codex_cli)
mkdir -p "$HOME/.codex"
cat > "$HOME/.codex/auth.json" <<EOF
{{"OPENAI_API_KEY":"$AETHER_API_KEY"}}
EOF
cat > "$HOME/.codex/config.toml" <<EOF
# Managed by Aether
model_provider = "aether"
[model_providers.aether]
name = "Aether"
base_url = "$AETHER_BASE_URL/v1"
env_key = "OPENAI_API_KEY"
wire_api = "chat"
EOF
chmod 600 "$HOME/.codex/auth.json" "$HOME/.codex/config.toml" 2>/dev/null || true
;;
gemini_cli)
mkdir -p "$HOME/.gemini"
cat > "$HOME/.gemini/.env" <<EOF
GEMINI_API_KEY=$AETHER_API_KEY
GOOGLE_API_KEY=$AETHER_API_KEY
GOOGLE_GEMINI_BASE_URL=$AETHER_BASE_URL
AETHER_BASE_URL=$AETHER_BASE_URL
EOF
python3 - "$HOME/.gemini/settings.json" "$AETHER_BASE_URL" <<'PY'
import json, pathlib, sys
path = pathlib.Path(sys.argv[1])
data = json.loads(path.read_text() or '{{}}') if path.exists() else {{}}
data.setdefault('aether', {{}})['baseUrl'] = sys.argv[2]
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + '\n')
PY
chmod 600 "$HOME/.gemini/.env" "$HOME/.gemini/settings.json" 2>/dev/null || true
;;
esac
say "$CLI_LABEL 已配置到 Aether。执行 $CLI_BIN --version 验证安装。"
"###,
target_cli = target_cli,
target_system = target_system,
base_url = shell_single_quote(&session.base_url),
api_key = shell_single_quote(&session.api_key),
label = shell_single_quote(cli_label(session.target_cli)),
binary = shell_single_quote(cli_binary(session.target_cli)),
npm_package = shell_single_quote(npm_package(session.target_cli)),
)
}
fn build_powershell_script(session: &StoredInstallSession) -> String {
let target_cli = match session.target_cli {
InstallTargetCli::ClaudeCode => "claude_code",
InstallTargetCli::CodexCli => "codex_cli",
InstallTargetCli::GeminiCli => "gemini_cli",
};
format!(
r###"$ErrorActionPreference = 'Stop'
$TargetCli = {target_cli}
$TargetSystem = {target_system}
$AetherBaseUrl = {base_url}
$AetherApiKey = {api_key}
$CliLabel = {label}
$CliBin = {binary}
$NpmPackage = {npm_package}
function Say($Message) {{ Write-Host "[Aether] $Message" }}
function Fail($Message) {{ Write-Error "[Aether] $Message"; exit 1 }}
if ($TargetSystem -ne 'auto' -and $TargetSystem -ne 'windows') {{ Fail "该 install code 绑定 $TargetSystem请复制 macOS/Linux 命令执行。" }}
Say "准备安装/复用 $CliLabel"
if (-not (Get-Command $CliBin -ErrorAction SilentlyContinue)) {{
if (-not (Get-Command npm -ErrorAction SilentlyContinue)) {{ Fail "未找到 $CliBin也未找到 npm。请先安装 Node.js/npm 后重试。" }}
Say "未找到 $CliBin正在通过 npm 安装 $NpmPackage"
npm install -g $NpmPackage
}}
$HomeDir = [Environment]::GetFolderPath('UserProfile')
$AetherDir = Join-Path $HomeDir '.aether'
New-Item -ItemType Directory -Force -Path $AetherDir | Out-Null
Set-Content -Path (Join-Path $AetherDir 'client.env') -Value "AETHER_BASE_URL=$AetherBaseUrl`nAETHER_API_KEY=$AetherApiKey`n" -Encoding UTF8
if ($TargetCli -eq 'claude_code') {{
$Dir = Join-Path $HomeDir '.claude'; New-Item -ItemType Directory -Force -Path $Dir | Out-Null
$Path = Join-Path $Dir 'settings.json'
$Data = if (Test-Path $Path) {{ Get-Content $Path -Raw | ConvertFrom-Json -AsHashtable }} else {{ @{{}} }}
if (-not $Data.ContainsKey('env')) {{ $Data.env = @{{}} }}
$Data.env.ANTHROPIC_BASE_URL = $AetherBaseUrl
$Data.env.ANTHROPIC_AUTH_TOKEN = $AetherApiKey
$Data | ConvertTo-Json -Depth 8 | Set-Content $Path -Encoding UTF8
}} elseif ($TargetCli -eq 'codex_cli') {{
$Dir = Join-Path $HomeDir '.codex'; New-Item -ItemType Directory -Force -Path $Dir | Out-Null
Set-Content (Join-Path $Dir 'auth.json') -Value (@{{ OPENAI_API_KEY = $AetherApiKey }} | ConvertTo-Json) -Encoding UTF8
Set-Content (Join-Path $Dir 'config.toml') -Value "# Managed by Aether`nmodel_provider = \"aether\"`n`n[model_providers.aether]`nname = \"Aether\"`nbase_url = \"$AetherBaseUrl/v1\"`nenv_key = \"OPENAI_API_KEY\"`nwire_api = \"chat\"`n" -Encoding UTF8
}} elseif ($TargetCli -eq 'gemini_cli') {{
$Dir = Join-Path $HomeDir '.gemini'; New-Item -ItemType Directory -Force -Path $Dir | Out-Null
Set-Content (Join-Path $Dir '.env') -Value "GEMINI_API_KEY=$AetherApiKey`nGOOGLE_API_KEY=$AetherApiKey`nGOOGLE_GEMINI_BASE_URL=$AetherBaseUrl`nAETHER_BASE_URL=$AetherBaseUrl`n" -Encoding UTF8
$Path = Join-Path $Dir 'settings.json'
$Data = if (Test-Path $Path) {{ Get-Content $Path -Raw | ConvertFrom-Json -AsHashtable }} else {{ @{{}} }}
$Data.aether = @{{ baseUrl = $AetherBaseUrl }}
$Data | ConvertTo-Json -Depth 8 | Set-Content $Path -Encoding UTF8
}}
Say "$CliLabel 已配置到 Aether。执行 $CliBin --version 验证安装。"
"###,
target_cli = powershell_single_quote(target_cli),
target_system = powershell_single_quote(match session.target_system {
InstallTargetSystem::Macos => "macos",
InstallTargetSystem::Linux => "linux",
InstallTargetSystem::Windows => "windows",
InstallTargetSystem::Auto => "auto",
}),
base_url = powershell_single_quote(&session.base_url),
api_key = powershell_single_quote(&session.api_key),
label = powershell_single_quote(cli_label(session.target_cli)),
binary = powershell_single_quote(cli_binary(session.target_cli)),
npm_package = powershell_single_quote(npm_package(session.target_cli)),
)
}
pub(super) async fn handle_users_me_api_key_install_session_create(
state: &AppState,
request_context: &GatewayPublicRequestContext,
headers: &http::HeaderMap,
request_body: Option<&axum::body::Bytes>,
) -> Response<Body> {
let auth = match resolve_authenticated_local_user(state, request_context, headers).await {
Ok(value) => value,
Err(response) => return response,
};
let Some(api_key_id) =
users_me_api_key_install_session_id_from_path(&request_context.request_path)
else {
return build_auth_error_response(http::StatusCode::NOT_FOUND, "API密钥不存在", false);
};
let Some(request_body) = request_body else {
return build_auth_error_response(http::StatusCode::BAD_REQUEST, "请求数据验证失败", false);
};
let payload = match serde_json::from_slice::<UsersMeCreateInstallSessionRequest>(request_body) {
Ok(value) => value,
Err(_) => {
return build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"请求数据验证失败",
false,
)
}
};
let records = match state
.list_auth_api_key_export_records_by_user_ids(std::slice::from_ref(&auth.user.id))
.await
{
Ok(value) => value,
Err(err) => {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user api key lookup failed: {err:?}"),
false,
)
}
};
let Some(record) = records
.into_iter()
.find(|record| !record.is_standalone && record.api_key_id == api_key_id)
else {
return build_auth_error_response(http::StatusCode::NOT_FOUND, "API密钥不存在", false);
};
let Some(ciphertext) = record
.key_encrypted
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"该密钥没有存储完整密钥信息",
false,
);
};
let Some(api_key) = decrypt_catalog_secret_with_fallbacks(state.encryption_key(), ciphertext)
else {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
"解密密钥失败",
false,
);
};
let code = generate_install_code();
let expires_at_unix_secs = unix_secs_now().saturating_add(INSTALL_SESSION_TTL_SECS);
let session = StoredInstallSession {
api_key_id: record.api_key_id.clone(),
api_key_name: record.name.unwrap_or_else(|| "API Key".to_string()),
api_key,
base_url: base_url_from_request(headers, request_context),
target_cli: payload.target_cli,
target_system: payload.target_system,
expires_at_unix_secs,
};
let serialized = match serde_json::to_string(&session) {
Ok(value) => value,
Err(err) => {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("install session serialize failed: {err:?}"),
false,
)
}
};
if let Err(err) = state
.runtime_kv_setex(
&install_session_runtime_key(&code),
&serialized,
INSTALL_SESSION_TTL_SECS,
)
.await
{
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("install session create failed: {err:?}"),
false,
);
}
let base_url = session.base_url.trim_end_matches('/');
Json(json!({
"install_code": code,
"expires_at_unix_secs": expires_at_unix_secs,
"expires_in_seconds": INSTALL_SESSION_TTL_SECS,
"target_cli": session.target_cli,
"target_cli_label": cli_label(session.target_cli),
"target_system": session.target_system,
"target_system_label": system_label(session.target_system),
"unix_command": format!("curl -fsSL {base_url}/install/{code} | sh"),
"powershell_command": format!("irm {base_url}/install/{code}.ps1 | iex"),
}))
.into_response()
}
pub(super) async fn maybe_build_local_install_response(
state: &AppState,
request_context: &GatewayPublicRequestContext,
) -> Option<Response<Body>> {
let decision = request_context.control_decision.as_ref()?;
if decision.route_family.as_deref() != Some("install") {
return None;
}
let Some((code, wants_powershell)) = install_code_from_path(&request_context.request_path)
else {
return Some(build_auth_error_response(
http::StatusCode::NOT_FOUND,
"install code 不存在或已失效",
false,
));
};
let raw = match state
.runtime_kv_getdel(&install_session_runtime_key(&code))
.await
{
Ok(Some(value)) => value,
Ok(None) => {
return Some(build_auth_error_response(
http::StatusCode::NOT_FOUND,
"install code 不存在、已过期或已使用",
false,
))
}
Err(err) => {
return Some(build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("install session lookup failed: {err:?}"),
false,
))
}
};
let session = match serde_json::from_str::<StoredInstallSession>(&raw) {
Ok(value) => value,
Err(_) => {
return Some(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"install code 数据无效",
false,
))
}
};
if session.expires_at_unix_secs <= unix_secs_now() {
return Some(build_auth_error_response(
http::StatusCode::NOT_FOUND,
"install code 已过期",
false,
));
}
let body = if wants_powershell {
build_powershell_script(&session)
} else {
build_unix_script(&session)
};
let content_type = if wants_powershell {
"text/plain; charset=utf-8"
} else {
"text/x-shellscript; charset=utf-8"
};
let mut response = Response::new(Body::from(body));
response.headers_mut().insert(
http::header::CONTENT_TYPE,
http::HeaderValue::from_static(content_type),
);
response.headers_mut().insert(
http::header::CACHE_CONTROL,
http::HeaderValue::from_static("no-store"),
);
response.headers_mut().insert(
http::header::PRAGMA,
http::HeaderValue::from_static("no-cache"),
);
response.headers_mut().insert(
http::header::HeaderName::from_static("x-content-type-options"),
http::HeaderValue::from_static("nosniff"),
);
Some(response)
}

View File

@@ -1,9 +1,10 @@
use super::{
auth_password_policy_level, build_auth_error_response, build_auth_wallet_summary_payload,
decrypt_catalog_secret_with_fallbacks, encrypt_catalog_secret_with_fallbacks, handle_auth_me,
query_param_optional_bool, query_param_value, resolve_authenticated_local_user,
unix_secs_to_rfc3339, validate_auth_register_password, AppState, AuthenticatedLocalUserContext,
GatewayPublicRequestContext, PUBLIC_CAPABILITY_DEFINITIONS,
handle_users_me_api_key_install_session_create, query_param_optional_bool, query_param_value,
resolve_authenticated_local_user, unix_secs_to_rfc3339,
users_me_api_key_install_sessions_path_matches, validate_auth_register_password, AppState,
AuthenticatedLocalUserContext, GatewayPublicRequestContext, PUBLIC_CAPABILITY_DEFINITIONS,
};
use crate::admin_api::build_admin_endpoint_health_status_payload;
use crate::handlers::internal::build_management_token_payload;

View File

@@ -4,8 +4,9 @@ use axum::{body::Body, http, response::Response};
use super::{
handle_auth_me, handle_users_me_api_key_capabilities_put, handle_users_me_api_key_create,
handle_users_me_api_key_delete, handle_users_me_api_key_detail_get,
handle_users_me_api_key_patch, handle_users_me_api_key_providers_put,
handle_users_me_api_key_update, handle_users_me_api_keys_get, handle_users_me_available_models,
handle_users_me_api_key_install_session_create, handle_users_me_api_key_patch,
handle_users_me_api_key_providers_put, handle_users_me_api_key_update,
handle_users_me_api_keys_get, handle_users_me_available_models,
handle_users_me_delete_other_sessions, handle_users_me_delete_session,
handle_users_me_detail_put, handle_users_me_endpoint_status_get,
handle_users_me_management_token_create, handle_users_me_management_token_delete,
@@ -17,8 +18,8 @@ use super::{
handle_users_me_providers_get, handle_users_me_sessions_get, handle_users_me_update_session,
handle_users_me_usage_active_get, handle_users_me_usage_get, handle_users_me_usage_heatmap_get,
handle_users_me_usage_interval_timeline_get, users_me_api_key_capabilities_path_matches,
users_me_api_key_detail_path_matches, users_me_api_key_providers_path_matches,
users_me_management_token_detail_path_matches,
users_me_api_key_detail_path_matches, users_me_api_key_install_sessions_path_matches,
users_me_api_key_providers_path_matches, users_me_management_token_detail_path_matches,
users_me_management_token_regenerate_path_matches,
users_me_management_token_toggle_path_matches, users_me_management_tokens_root,
users_me_session_detail_path_matches, AppState, GatewayPublicRequestContext,
@@ -167,6 +168,19 @@ pub(crate) async fn maybe_build_local_users_me_response(
.await,
)
}
Some("api_key_install_session_create")
if users_me_api_key_install_sessions_path_matches(&request_context.request_path) =>
{
Some(
handle_users_me_api_key_install_session_create(
state,
request_context,
headers,
request_body,
)
.await,
)
}
Some("management_token_regenerate")
if users_me_management_token_regenerate_path_matches(&request_context.request_path) =>
{

View File

@@ -443,7 +443,11 @@ pub(crate) fn public_support_local_requires_buffered_body(
) | (
Some("users_me"),
http::Method::POST,
Some("api_keys_create" | "management_tokens_create"),
Some(
"api_keys_create"
| "api_key_install_session_create"
| "management_tokens_create",
),
) | (
Some("wallet"),
http::Method::POST,

View File

@@ -58,6 +58,7 @@ mod router;
mod scheduler;
mod state;
mod system_features;
mod task_runtime;
mod tunnel;
mod usage;
mod video_tasks;

View File

@@ -1202,13 +1202,13 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
state.bootstrap_admin_from_env().await?;
let background_tasks = if args.node_role.spawns_background_tasks() {
state.spawn_background_tasks()
Some(state.spawn_background_tasks())
} else {
info!(
node_role = args.node_role.as_str(),
"background workers disabled for this node role"
);
Vec::new()
None
};
let listener = tokio::net::TcpListener::bind(bind_addr).await?;
let public_base_url = resolve_local_http_base_url(app_port)?;
@@ -1241,8 +1241,8 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
router.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.await?;
for handle in background_tasks {
handle.abort();
if let Some(background_tasks) = background_tasks {
background_tasks.shutdown().await;
}
Ok(())
}

View File

@@ -12,13 +12,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_admin_request_body_cleanup_task,
start_proxy_upgrade_rollout, AdminCleanupRunRecord, AdminStatsRebuildSummary,
AdminSystemCleanupSummary, OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary,
ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
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_admin_request_body_cleanup_task, start_admin_system_purge_task,
start_proxy_upgrade_rollout, AdminCleanupRunRecord, AdminCleanupTaskKind,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, OAuthTokenRefreshRunSummary,
PoolQuotaProbeRunSummary, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,

View File

@@ -22,6 +22,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"]
@@ -65,6 +67,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::{
@@ -96,6 +99,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;
@@ -133,6 +138,8 @@ struct UsageCleanupSettings {
pub(crate) struct AdminSystemCleanupSummary {
pub(crate) audit_logs_deleted: usize,
pub(crate) request_candidates_deleted: usize,
pub(crate) proxy_node_metrics:
aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary,
pub(crate) pending_failed: usize,
pub(crate) pending_recovered: usize,
pub(crate) usage: UsageCleanupSummary,
@@ -150,12 +157,14 @@ pub(crate) async fn run_admin_system_cleanup_once(
) -> Result<AdminSystemCleanupSummary, aether_data::DataLayerError> {
let audit_logs_deleted = cleanup_audit_logs_once(data).await?;
let request_candidates_deleted = cleanup_request_candidates_once(data).await?;
let proxy_node_metrics = cleanup_proxy_node_metrics_once(data).await?;
let pending = cleanup_stale_pending_requests_once(data).await?;
let usage = perform_usage_cleanup_once(data).await?;
Ok(AdminSystemCleanupSummary {
audit_logs_deleted,
request_candidates_deleted,
proxy_node_metrics,
pending_failed: pending.failed,
pending_recovered: pending.recovered,
usage,

View File

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

View File

@@ -8,10 +8,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, now_unix_secs, 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_completed_cleanup_run,
record_failed_cleanup_run, record_proxy_upgrade_traffic_success, summarize_database_pool,
@@ -89,6 +89,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

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

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

@@ -119,9 +119,7 @@ async fn gateway_background_request_candidate_cleanup_deletes_expired_entries_in
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
}
#[tokio::test]

View File

@@ -539,8 +539,6 @@ async fn gateway_background_model_fetch_updates_key_and_syncs_provider_model_whi
.expect("execution runtime plan should be captured");
assert_eq!(seen_plan.url, "https://api.openai.example/v1/models");
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
execution_runtime_handle.abort();
}

View File

@@ -1,5 +1,5 @@
use aether_contracts::ExecutionPlan;
use serde_json::{json, Value};
use serde_json::{json, Map, Value};
use crate::AppState;
@@ -120,3 +120,66 @@ pub(crate) fn with_error_flow_report_context(
object.insert("error_flow".to_string(), error_flow);
Some(Value::Object(object))
}
pub(crate) fn with_upstream_response_report_context(
report_context: Option<&Value>,
status_code: u16,
headers: Option<&std::collections::BTreeMap<String, String>>,
body: Option<&Value>,
body_ref: Option<&str>,
body_state: Option<&str>,
) -> Option<Value> {
let mut object = report_context?.as_object()?.clone();
let mut upstream_response = serde_json::Map::new();
upstream_response.insert("status_code".to_string(), json!(status_code));
if let Some(headers) = headers {
upstream_response.insert("headers".to_string(), trace_headers_to_json(headers));
}
if let Some(body) = body {
upstream_response.insert("body".to_string(), body.clone());
}
if let Some(body_ref) = body_ref {
upstream_response.insert("body_ref".to_string(), json!(body_ref));
}
if let Some(body_state) = body_state {
upstream_response.insert("body_state".to_string(), json!(body_state));
}
object.insert(
"upstream_response".to_string(),
Value::Object(upstream_response),
);
Some(Value::Object(object))
}
fn trace_headers_to_json(headers: &std::collections::BTreeMap<String, String>) -> Value {
Value::Object(Map::from_iter(headers.iter().map(|(key, value)| {
(
key.clone(),
Value::String(mask_trace_header_value(key, value)),
)
})))
}
fn mask_trace_header_value(name: &str, value: &str) -> String {
if !trace_header_is_sensitive(name) {
return value.to_string();
}
if value.len() <= 8 {
return "****".to_string();
}
format!("{}****{}", &value[..4], &value[value.len() - 4..])
}
fn trace_header_is_sensitive(name: &str) -> bool {
[
"authorization",
"x-api-key",
"api-key",
"x-goog-api-key",
"cookie",
"set-cookie",
"proxy-authorization",
]
.iter()
.any(|candidate| name.trim().eq_ignore_ascii_case(candidate))
}

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::{
@@ -17,7 +19,6 @@ use aether_runtime_state::{
RuntimeSemaphoreSnapshot, RuntimeState,
};
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
use tokio::task::JoinHandle;
use super::{AppState, FrontdoorCorsConfig, LocalExecutionRuntimeMissDiagnostic};
@@ -46,6 +47,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 +595,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 +670,23 @@ impl AppState {
.map_err(|err| std::io::Error::other(err.to_string()))
}
pub(crate) async fn cleanup_proxy_node_metrics(
&self,
retain_1m_from_unix_secs: u64,
retain_1h_from_unix_secs: u64,
delete_limit: usize,
) -> Result<aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary, GatewayError>
{
self.data
.cleanup_proxy_node_metrics(
retain_1m_from_unix_secs,
retain_1h_from_unix_secs,
delete_limit,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn apply_proxy_node_heartbeat(
&self,
mutation: &ProxyNodeHeartbeatMutation,
@@ -939,68 +996,112 @@ impl AppState {
Ok(self)
}
pub fn spawn_background_tasks(&self) -> Vec<JoinHandle<()>> {
let mut tasks = Vec::new();
if let Some(handle) = self.usage_runtime.spawn_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) =
crate::wallet_runtime::spawn_provider_quota_reset_worker(self.data.clone())
{
tasks.push(handle);
}
if let Some(handle) = spawn_audit_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_db_maintenance_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_wallet_daily_usage_aggregation_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_stats_aggregation_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_usage_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_pool_monitor_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_pool_quota_probe_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_stats_hourly_aggregation_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_pending_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_node_stale_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_upgrade_rollout_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_provider_checkin_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_oauth_token_refresh_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_request_candidate_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_gemini_file_mapping_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_model_fetch_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_video_task_poller(self.clone()) {
tasks.push(handle);
}
tasks
pub fn spawn_background_tasks(&self) -> crate::task_runtime::TaskSupervisor {
let mut supervisor = crate::task_runtime::TaskSupervisor::new();
let record_boot = |task_key: &'static str| {
if !self.has_background_task_data_writer() {
return;
}
let Some(definition) = crate::task_runtime::task_definition(task_key) else {
return;
};
std::mem::drop(crate::task_runtime::spawn_record_worker_boot(
self.clone(),
task_key,
crate::task_runtime::background_task_kind(definition.kind),
definition.trigger,
));
};
let mut supervise_worker =
|task_key: &'static str, handle: Option<tokio::task::JoinHandle<()>>| {
if let Some(handle) = handle {
supervisor.supervise_handle(task_key, handle);
record_boot(task_key);
}
};
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER,
self.usage_runtime.spawn_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROVIDER_QUOTA_RESET,
crate::wallet_runtime::spawn_provider_quota_reset_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_AUDIT_CLEANUP,
spawn_audit_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_DB_MAINTENANCE,
spawn_db_maintenance_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_WALLET_DAILY_USAGE_AGG,
spawn_wallet_daily_usage_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_STATS_DAILY_AGG,
spawn_stats_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_CLEANUP,
spawn_usage_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_MONITOR,
spawn_pool_monitor_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
spawn_pool_quota_probe_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_STATS_HOURLY_AGG,
spawn_stats_hourly_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PENDING_CLEANUP,
spawn_pending_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_NODE_STALE_CLEANUP,
spawn_proxy_node_stale_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_NODE_METRICS_CLEANUP,
spawn_proxy_node_metrics_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_UPGRADE_ROLLOUT,
spawn_proxy_upgrade_rollout_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROVIDER_CHECKIN,
spawn_provider_checkin_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_OAUTH_TOKEN_REFRESH,
spawn_oauth_token_refresh_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_REQUEST_CANDIDATE_CLEANUP,
spawn_request_candidate_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_GEMINI_FILES_CLEANUP,
spawn_gemini_file_mapping_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_MODEL_FETCH_WORKER,
spawn_model_fetch_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_VIDEO_TASK_POLLER,
spawn_video_task_poller(self.clone()),
);
supervisor
}
}

View File

@@ -0,0 +1,80 @@
use crate::{AppState, GatewayError};
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskListQuery, BackgroundTaskSummary, StoredBackgroundTaskEvent,
StoredBackgroundTaskRun, StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent,
UpsertBackgroundTaskRun,
};
impl AppState {
pub(crate) async fn find_background_task_run(
&self,
run_id: &str,
) -> Result<Option<StoredBackgroundTaskRun>, GatewayError> {
self.data
.find_background_task_run(run_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_background_task_runs(
&self,
query: &BackgroundTaskListQuery,
) -> Result<StoredBackgroundTaskRunPage, GatewayError> {
self.data
.list_background_task_runs(query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_background_task_events(
&self,
run_id: &str,
offset: usize,
limit: usize,
) -> Result<Vec<StoredBackgroundTaskEvent>, GatewayError> {
self.data
.list_background_task_events(run_id, offset, limit)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn summarize_background_task_runs(
&self,
) -> Result<BackgroundTaskSummary, GatewayError> {
self.data
.summarize_background_task_runs()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn upsert_background_task_run(
&self,
run: UpsertBackgroundTaskRun,
) -> Result<Option<StoredBackgroundTaskRun>, GatewayError> {
self.data
.upsert_background_task_run(run)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn request_cancel_background_task_run(
&self,
run_id: &str,
updated_at_unix_secs: u64,
) -> Result<bool, GatewayError> {
self.data
.request_cancel_background_task_run(run_id, updated_at_unix_secs)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn upsert_background_task_event(
&self,
event: UpsertBackgroundTaskEvent,
) -> Result<Option<StoredBackgroundTaskEvent>, GatewayError> {
self.data
.upsert_background_task_event(event)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
}

View File

@@ -10,6 +10,7 @@ use super::{
mod announcements;
mod api_key_exports;
mod auth;
mod background_tasks;
mod billing;
mod candidate_queries;
mod gemini_files;
@@ -29,6 +30,14 @@ impl AppState {
self.data.has_announcement_writer()
}
pub fn has_background_task_data_reader(&self) -> bool {
self.data.has_background_task_reader()
}
pub fn has_background_task_data_writer(&self) -> bool {
self.data.has_background_task_writer()
}
pub fn has_video_task_data_reader(&self) -> bool {
self.data.has_video_task_reader()
}

View File

@@ -0,0 +1,647 @@
use std::future::Future;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::background_tasks::{
BackgroundTaskKind, BackgroundTaskStatus, StoredBackgroundTaskRun, UpsertBackgroundTaskEvent,
UpsertBackgroundTaskRun,
};
use aether_runtime::task::spawn_named;
pub(crate) use aether_task_runtime::TaskSupervisor;
use aether_task_runtime::{RetryPolicy, TaskDefinition, TaskKind};
use serde_json::Value;
use tokio::task::JoinHandle;
use tracing::warn;
use uuid::Uuid;
use crate::{AppState, GatewayError};
pub(crate) const TASK_KEY_PROVIDER_DELETE: &str = "admin.provider.delete";
pub(crate) const TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT: &str = "admin.provider.oauth.batch_import";
pub(crate) const TASK_KEY_USAGE_QUEUE_WORKER: &str = "usage.queue.worker";
pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
pub(crate) const TASK_KEY_PENDING_CLEANUP: &str = "maintenance.pending.cleanup";
pub(crate) const TASK_KEY_REQUEST_CANDIDATE_CLEANUP: &str = "maintenance.request.candidate.cleanup";
pub(crate) const TASK_KEY_GEMINI_FILES_CLEANUP: &str = "maintenance.gemini.files.cleanup";
pub(crate) const TASK_KEY_OAUTH_TOKEN_REFRESH: &str = "maintenance.oauth.token.refresh";
pub(crate) const TASK_KEY_PROXY_NODE_STALE_CLEANUP: &str = "maintenance.proxy.node.stale.cleanup";
pub(crate) const TASK_KEY_PROXY_NODE_METRICS_CLEANUP: &str =
"maintenance.proxy.node.metrics.cleanup";
pub(crate) const TASK_KEY_PROXY_UPGRADE_ROLLOUT: &str = "maintenance.proxy.upgrade.rollout";
pub(crate) const TASK_KEY_PROVIDER_CHECKIN: &str = "maintenance.provider.checkin";
pub(crate) const TASK_KEY_USAGE_CLEANUP: &str = "maintenance.usage.cleanup";
pub(crate) const TASK_KEY_WALLET_DAILY_USAGE_AGG: &str = "maintenance.wallet.daily.usage.agg";
pub(crate) const TASK_KEY_STATS_DAILY_AGG: &str = "maintenance.stats.daily.agg";
pub(crate) const TASK_KEY_STATS_HOURLY_AGG: &str = "maintenance.stats.hourly.agg";
pub(crate) const TASK_KEY_USAGE_SYNC_REPORT: &str = "usage.sync.report";
pub(crate) const TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH: &str = "provider.oauth.account.refresh";
pub(crate) const TASK_KEY_PROVIDER_BALANCE_REFRESH: &str = "provider.ops.balance.refresh";
const RETRY_ONCE: RetryPolicy = RetryPolicy { max_attempts: 1 };
const TASK_DEFINITIONS: &[TaskDefinition] = &[
TaskDefinition::new(
TASK_KEY_PROVIDER_DELETE,
TaskKind::OnDemand,
"manual",
false,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
TaskKind::OnDemand,
"manual",
false,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_USAGE_QUEUE_WORKER,
TaskKind::Daemon,
"daemon",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_VIDEO_TASK_POLLER,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_MODEL_FETCH_WORKER,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_QUOTA_RESET,
TaskKind::Scheduled,
"daily",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_QUOTA_PROBE,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_MONITOR,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_AUDIT_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_DB_MAINTENANCE,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PENDING_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_REQUEST_CANDIDATE_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_GEMINI_FILES_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_OAUTH_TOKEN_REFRESH,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROXY_NODE_STALE_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROXY_NODE_METRICS_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROXY_UPGRADE_ROLLOUT,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_CHECKIN,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_USAGE_CLEANUP,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_WALLET_DAILY_USAGE_AGG,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_STATS_DAILY_AGG,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_STATS_HOURLY_AGG,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_USAGE_SYNC_REPORT,
TaskKind::FireAndForget,
"internal",
false,
false,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH,
TaskKind::FireAndForget,
"internal",
false,
false,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_BALANCE_REFRESH,
TaskKind::FireAndForget,
"internal",
false,
false,
RETRY_ONCE,
),
];
pub(crate) fn task_definitions() -> &'static [TaskDefinition] {
TASK_DEFINITIONS
}
pub(crate) fn task_definition(task_key: &str) -> Option<TaskDefinition> {
task_definitions()
.iter()
.copied()
.find(|definition| definition.key == task_key)
}
pub(crate) const fn background_task_kind(kind: TaskKind) -> BackgroundTaskKind {
match kind {
TaskKind::Scheduled => BackgroundTaskKind::Scheduled,
TaskKind::Daemon => BackgroundTaskKind::Daemon,
TaskKind::OnDemand => BackgroundTaskKind::OnDemand,
TaskKind::FireAndForget => BackgroundTaskKind::FireAndForget,
}
}
pub(crate) fn task_cancel_kv_key(run_id: &str) -> String {
format!("task_runtime:run:{run_id}:cancel")
}
pub(crate) fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default()
}
pub(crate) fn build_task_run_id() -> String {
Uuid::new_v4().to_string()
}
pub(crate) fn spawn_fire_and_forget<F>(task_name: &'static str, future: F) -> JoinHandle<()>
where
F: Future<Output = ()> + Send + 'static,
{
spawn_named(task_name, future)
}
fn stored_run_to_upsert(run: StoredBackgroundTaskRun) -> UpsertBackgroundTaskRun {
UpsertBackgroundTaskRun {
id: run.id,
task_key: run.task_key,
kind: run.kind,
trigger: run.trigger,
status: run.status,
attempt: run.attempt,
max_attempts: run.max_attempts,
owner_instance: run.owner_instance,
progress_percent: run.progress_percent,
progress_message: run.progress_message,
payload_json: run.payload_json,
result_json: run.result_json,
error_message: run.error_message,
cancel_requested: run.cancel_requested,
created_by: run.created_by,
created_at_unix_secs: run.created_at_unix_secs,
started_at_unix_secs: run.started_at_unix_secs,
finished_at_unix_secs: run.finished_at_unix_secs,
updated_at_unix_secs: run.updated_at_unix_secs,
}
}
pub(crate) async fn upsert_run_with_logging(
app: &AppState,
run: UpsertBackgroundTaskRun,
) -> Option<StoredBackgroundTaskRun> {
match app.upsert_background_task_run(run).await {
Ok(result) => result,
Err(error) => {
warn!(error = ?error, "failed to upsert background task run");
None
}
}
}
pub(crate) async fn update_run_status(
app: &AppState,
run_id: &str,
status: BackgroundTaskStatus,
progress_percent: Option<u16>,
progress_message: Option<String>,
result_json: Option<Value>,
error_message: Option<String>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
) -> Option<StoredBackgroundTaskRun> {
let Some(mut existing) = app.find_background_task_run(run_id).await.ok().flatten() else {
return None;
};
existing.status = status;
if let Some(progress_percent) = progress_percent {
existing.progress_percent = progress_percent.min(100);
}
if let Some(progress_message) = progress_message {
existing.progress_message = Some(progress_message);
}
if result_json.is_some() {
existing.result_json = result_json;
}
if error_message.is_some() {
existing.error_message = error_message;
}
if started_at_unix_secs.is_some() {
existing.started_at_unix_secs = started_at_unix_secs;
}
if finished_at_unix_secs.is_some() {
existing.finished_at_unix_secs = finished_at_unix_secs;
}
existing.updated_at_unix_secs = now_unix_secs();
upsert_run_with_logging(app, stored_run_to_upsert(existing)).await
}
pub(crate) async fn append_event_with_logging(
app: &AppState,
run_id: &str,
event_type: &str,
message: &str,
payload_json: Option<Value>,
) {
let event = UpsertBackgroundTaskEvent {
id: Uuid::new_v4().to_string(),
run_id: run_id.to_string(),
event_type: event_type.to_string(),
message: message.to_string(),
payload_json,
created_at_unix_secs: now_unix_secs(),
};
if let Err(error) = app.upsert_background_task_event(event).await {
warn!(error = ?error, run_id = %run_id, "failed to upsert background task event");
}
}
pub(crate) fn spawn_record_worker_boot(
app: AppState,
task_key: &'static str,
kind: BackgroundTaskKind,
trigger: &'static str,
) -> JoinHandle<()> {
spawn_named("task-runtime-record-worker-boot", async move {
let now = now_unix_secs();
let run_id = format!("boot:{}:{}", task_key, app.tunnel.local_instance_id());
let run = UpsertBackgroundTaskRun {
id: run_id.clone(),
task_key: task_key.to_string(),
kind,
trigger: trigger.to_string(),
status: BackgroundTaskStatus::Running,
attempt: 1,
max_attempts: 1,
owner_instance: Some(app.tunnel.local_instance_id().to_string()),
progress_percent: 0,
progress_message: Some("worker booted".to_string()),
payload_json: None,
result_json: None,
error_message: None,
cancel_requested: false,
created_by: Some("system".to_string()),
created_at_unix_secs: now,
started_at_unix_secs: Some(now),
finished_at_unix_secs: None,
updated_at_unix_secs: now,
};
let _ = upsert_run_with_logging(&app, run).await;
append_event_with_logging(
&app,
&run_id,
"worker_boot",
"background worker started",
None,
)
.await;
})
}
pub(crate) async fn set_cancel_signal(app: &AppState, run_id: &str) -> Result<(), GatewayError> {
app.runtime_kv_setex(&task_cancel_kv_key(run_id), "1", 60 * 60)
.await
}
pub(crate) async fn is_cancel_requested(app: &AppState, run_id: &str) -> bool {
if let Ok(Some(run)) = app.find_background_task_run(run_id).await {
if run.cancel_requested {
return true;
}
}
app.runtime_kv_exists(&task_cancel_kv_key(run_id))
.await
.unwrap_or(false)
}
pub(crate) async fn submit_provider_delete_task(
state: &crate::admin_api::AdminAppState<'_>,
provider_id: &str,
created_by: Option<&str>,
) -> Result<Option<String>, GatewayError> {
let Some(provider) = state
.read_provider_catalog_providers_by_ids(&[provider_id.to_string()])
.await?
.into_iter()
.next()
else {
return Ok(None);
};
let task_id = Uuid::new_v4().simple().to_string()[..16].to_string();
state.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
task_id: task_id.clone(),
provider_id: provider.id.clone(),
status: "pending".to_string(),
stage: "queued".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "delete task submitted".to_string(),
});
let app = state.cloned_app();
let provider_id = provider.id.clone();
let run_id = task_id.clone();
let created_at = now_unix_secs();
let max_attempts = task_definition(TASK_KEY_PROVIDER_DELETE)
.map(|item| item.retry_policy.max_attempts)
.unwrap_or(1);
if state.has_background_task_data_writer() {
let run = UpsertBackgroundTaskRun {
id: run_id.clone(),
task_key: TASK_KEY_PROVIDER_DELETE.to_string(),
kind: BackgroundTaskKind::OnDemand,
trigger: "manual".to_string(),
status: BackgroundTaskStatus::Queued,
attempt: 1,
max_attempts,
owner_instance: Some(app.tunnel.local_instance_id().to_string()),
progress_percent: 0,
progress_message: Some("delete task queued".to_string()),
payload_json: Some(serde_json::json!({ "provider_id": provider_id.clone() })),
result_json: None,
error_message: None,
cancel_requested: false,
created_by: Some(created_by.unwrap_or("admin").to_string()),
created_at_unix_secs: created_at,
started_at_unix_secs: None,
finished_at_unix_secs: None,
updated_at_unix_secs: created_at,
};
let _ = upsert_run_with_logging(&app, run).await;
append_event_with_logging(
&app,
&run_id,
"queued",
"provider delete task queued",
Some(serde_json::json!({ "provider_id": provider_id.clone() })),
)
.await;
}
spawn_named("task-runtime-provider-delete", async move {
let lock_key = format!("task_runtime:lock:{TASK_KEY_PROVIDER_DELETE}:{provider_id}");
let lock_ttl = std::time::Duration::from_secs(60 * 15);
let lock = app
.runtime_state
.lock_try_acquire(&lock_key, app.tunnel.local_instance_id(), lock_ttl)
.await
.ok()
.flatten();
if lock.is_none() {
let _ = update_run_status(
&app,
&run_id,
BackgroundTaskStatus::Skipped,
Some(0),
Some("provider delete skipped: another node is running this task".to_string()),
None,
None,
None,
Some(now_unix_secs()),
)
.await;
append_event_with_logging(
&app,
&run_id,
"skipped",
"provider delete skipped by singleton lock",
None,
)
.await;
return;
}
let started_at = now_unix_secs();
let _ = update_run_status(
&app,
&run_id,
BackgroundTaskStatus::Running,
Some(5),
Some("provider delete task started".to_string()),
None,
None,
Some(started_at),
None,
)
.await;
append_event_with_logging(
&app,
&run_id,
"running",
"provider delete task started",
None,
)
.await;
let admin_state = crate::admin_api::AdminAppState::new(&app);
let result = admin_state
.run_admin_provider_delete_task(&provider_id, &run_id)
.await;
match result {
Ok(task_state) => {
let _ = update_run_status(
&app,
&run_id,
BackgroundTaskStatus::Succeeded,
Some(100),
Some(task_state.message.clone()),
Some(serde_json::json!({
"provider_id": task_state.provider_id,
"status": task_state.status,
"stage": task_state.stage,
"deleted_keys": task_state.deleted_keys,
"total_keys": task_state.total_keys,
"deleted_endpoints": task_state.deleted_endpoints,
"total_endpoints": task_state.total_endpoints,
"message": task_state.message,
})),
None,
None,
Some(now_unix_secs()),
)
.await;
append_event_with_logging(
&app,
&run_id,
"succeeded",
"provider delete task completed",
None,
)
.await;
}
Err(err) => {
warn!(
"gateway admin provider delete task failed for provider {}: {:?}",
provider_id, err
);
app.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
task_id: run_id.clone(),
provider_id: provider_id.clone(),
status: "failed".to_string(),
stage: "failed".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: format!("provider delete failed: {err:?}"),
});
let _ = update_run_status(
&app,
&run_id,
BackgroundTaskStatus::Failed,
Some(100),
Some("provider delete task failed".to_string()),
None,
Some(format!("{err:?}")),
None,
Some(now_unix_secs()),
)
.await;
append_event_with_logging(
&app,
&run_id,
"failed",
"provider delete task failed",
Some(serde_json::json!({ "error": format!("{err:?}") })),
)
.await;
}
}
if let Some(lock) = lock {
let _ = app.runtime_state.lock_release(&lock).await;
}
});
Ok(Some(task_id))
}

View File

@@ -928,6 +928,428 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_openai_chat_stream_via_local_cross_format_gemini_candidate_with_stream_path_rewrite(
) {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeStreamRequest {
trace_id: String,
url: String,
provider_model: String,
auth_header_value: String,
accept: String,
endpoint_tag: 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", "gemini"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(serde_json::json!(["openai", "gemini"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-chat-gemini-stream-local-1".to_string(),
provider_name: "gemini".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-chat-gemini-stream-local-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-openai-chat-gemini-stream-local-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-openai-chat-gemini-stream-local-1".to_string(),
global_model_id: "global-model-openai-chat-gemini-stream-local-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gemini-2.5-pro-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gemini-2.5-pro-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-openai-chat-gemini-stream-local-1".to_string(),
"gemini".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
true,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-gemini-stream-local-1".to_string(),
"provider-openai-chat-gemini-stream-local-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(),
Some(serde_json::json!([
{"action":"set","key":"x-endpoint-tag","value":"openai-chat-gemini-cross-format-stream"}
])),
None,
Some(2),
Some("/custom/v1beta/models/gemini-2.5-pro-upstream:generateContent".to_string()),
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-chat-gemini-stream-local-1".to_string(),
"provider-openai-chat-gemini-stream-local-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-openai-chat-gemini-stream",
)
.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::<SeenExecutionRuntimeStreamRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let seen_report = Arc::new(Mutex::new(false));
let seen_report_clone = Arc::clone(&seen_report);
let decision_hits = Arc::new(Mutex::new(0usize));
let decision_hits_clone = Arc::clone(&decision_hits);
let plan_hits = Arc::new(Mutex::new(0usize));
let plan_hits_clone = Arc::clone(&plan_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "openai",
"route_kind": "chat",
"auth_endpoint_signature": "openai:chat",
"execution_runtime_candidate": true,
"auth_context": {
"user_id": "user-openai-chat-gemini-stream-local-1",
"api_key_id": "api-key-openai-chat-gemini-stream-local-1",
"access_allowed": true
},
"public_path": "/v1/chat/completions"
}))
}),
)
.route(
"/api/internal/gateway/decision-stream",
any(move |_request: Request| {
let decision_hits_inner = Arc::clone(&decision_hits_clone);
async move {
*decision_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"action": "proxy_public"}))
}
}),
)
.route(
"/api/internal/gateway/plan-stream",
any(move |_request: Request| {
let plan_hits_inner = Arc::clone(&plan_hits_clone);
async move {
*plan_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"action": "proxy_public"}))
}
}),
)
.route(
"/api/internal/gateway/report-stream",
any(move |request: Request| {
let seen_report_inner = Arc::clone(&seen_report_clone);
async move {
let (_parts, body) = request.into_parts();
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
*seen_report_inner.lock().expect("mutex should lock") = true;
Json(json!({"ok": true}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let execution_runtime = Router::new().route(
"/v1/execute/stream",
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");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest {
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(),
provider_model: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("model"))
.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(),
accept: payload
.get("headers")
.and_then(|value| value.get("accept"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
endpoint_tag: payload
.get("headers")
.and_then(|value| value.get("x-endpoint-tag"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"responseId\\\":\\\"resp_openai_chat_gemini_stream_123\\\",\\\"candidates\\\":[{\\\"content\\\":{\\\"parts\\\":[{\\\"text\\\":\\\"Hello \\\"}],\\\"role\\\":\\\"model\\\"},\\\"index\\\":0}],\\\"modelVersion\\\":\\\"gemini-2.5-pro-upstream\\\"}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"responseId\\\":\\\"resp_openai_chat_gemini_stream_123\\\",\\\"candidates\\\":[{\\\"content\\\":{\\\"parts\\\":[{\\\"text\\\":\\\"Hello Gemini stream\\\"}],\\\"role\\\":\\\"model\\\"},\\\"finishReason\\\":\\\"STOP\\\",\\\"index\\\":0}],\\\"modelVersion\\\":\\\"gemini-2.5-pro-upstream\\\",\\\"usageMetadata\\\":{\\\"promptTokenCount\\\":1,\\\"candidatesTokenCount\\\":2,\\\"totalTokenCount\\\":3}}\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from(frames))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-chat-gemini-stream")),
sample_auth_snapshot(
"api-key-openai-chat-gemini-stream-local-1",
"user-openai-chat-gemini-stream-local-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 (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state =
build_state_with_execution_runtime_override(execution_runtime_url.clone())
.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::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![(
"provider_priority_mode".to_string(),
json!("global_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/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-chat-gemini-stream",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-gemini-stream-123")
.body(
"{\"model\":\"gpt-5\",\"messages\":[{\"role\":\"system\",\"content\":\"You are terse.\"},{\"role\":\"user\",\"content\":\"Say hello\"}],\"stream\":true}",
)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
let response_text = response.text().await.expect("body should read");
assert!(response_text.contains("\"object\":\"chat.completion.chunk\""));
assert!(response_text.contains("data: [DONE]"));
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime stream should be captured");
assert_eq!(
seen_execution_runtime_request.trace_id,
"trace-openai-chat-gemini-stream-123"
);
assert_eq!(
seen_execution_runtime_request.url,
"https://generativelanguage.googleapis.com/custom/v1beta/models/gemini-2.5-pro-upstream:streamGenerateContent?alt=sse"
);
assert_eq!(
seen_execution_runtime_request.provider_model,
"gemini-2.5-pro-upstream"
);
assert_eq!(
seen_execution_runtime_request.auth_header_value,
"sk-upstream-openai-chat-gemini-stream"
);
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
assert_eq!(
seen_execution_runtime_request.endpoint_tag,
"openai-chat-gemini-cross-format-stream"
);
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-gemini-stream-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
let extra_data = stored_candidates[0]
.extra_data
.as_ref()
.expect("request candidate extra_data should exist");
assert_eq!(extra_data["client_api_format"], "openai:chat");
assert_eq!(extra_data["provider_api_format"], "gemini:generate_content");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert!(
!*seen_report.lock().expect("mutex should lock"),
"report-stream should stay local when request candidate persistence is available"
);
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_openai_chat_stream_with_custom_path_via_local_decision_gate_with_local_stream_decision(
) {

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

@@ -701,13 +701,13 @@ fn gateway_ai_serving_api_module_delegates_pure_ownership_to_format_crate() {
"pub use crate::provider_compat::kiro_stream::{",
"pub use crate::provider_compat::private_envelope::{",
"pub use crate::provider_compat::surfaces::{",
"pub use crate::request::common::{",
"pub use crate::request::route::{",
"pub use crate::response::common::{",
"pub use crate::response::error_body::{",
"pub use crate::formats::shared::request::{",
"pub use crate::formats::shared::routing::{",
"pub use crate::formats::shared::response::{",
"pub use crate::formats::shared::error_body::{",
"pub use aether_ai_formats::{",
"pub use aether_ai_formats::protocol::conversion::request::{",
"pub use aether_ai_formats::protocol::conversion::response::{",
"pub use aether_ai_formats::formats::conversion::request::{",
"pub use aether_ai_formats::formats::conversion::response::{",
] {
assert!(
format_crate_api.contains(pattern),

View File

@@ -2153,7 +2153,7 @@ fn ai_serving_image_routes_split_surface_normalization_and_transport_policy() {
}
let surface_image =
read_workspace_file("crates/aether-ai-formats/src/request/specialized/image.rs");
read_workspace_file("crates/aether-ai-formats/src/formats/openai/image/request.rs");
for pattern in [
"pub enum OpenAiImageOperation",
"pub fn is_openai_image_stream_request(",
@@ -3220,7 +3220,7 @@ fn ai_serving_leaf_planner_owners_route_contract_specs_through_gateway_seams() {
);
assert!(
!specialized_files_decision
.contains("aether_ai_formats::request::specialized::files::LocalGeminiFilesSpec"),
.contains("aether_ai_formats::formats::gemini::files::spec::LocalGeminiFilesSpec"),
"planner/specialized/files/decision.rs should consume LocalGeminiFilesSpec through the local specialized seam"
);
assert!(
@@ -3243,7 +3243,7 @@ fn ai_serving_m5_moves_contracts_and_route_logic_into_format_crate() {
"crates/aether-ai-formats/src/contracts/actions.rs",
"crates/aether-ai-formats/src/contracts/plan_kinds.rs",
"crates/aether-ai-formats/src/contracts/report_kinds.rs",
"crates/aether-ai-formats/src/request/route.rs",
"crates/aether-ai-formats/src/formats/shared/routing.rs",
] {
assert!(
workspace_file_exists(path),
@@ -3304,7 +3304,8 @@ fn ai_serving_m5_moves_contracts_and_route_logic_into_format_crate() {
);
}
let surface_route = read_workspace_file("crates/aether-ai-formats/src/request/route.rs");
let surface_route =
read_workspace_file("crates/aether-ai-formats/src/formats/shared/routing.rs");
for pattern in [
"pub fn is_matching_stream_http_request(",
"is_openai_image_stream_request(parts, body_json, body_base64)",
@@ -3553,11 +3554,11 @@ fn ai_serving_error_body_is_owned_by_format_finalize_module() {
"ai_serving/conversion/error.rs should stay removed"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/response/error_body.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/error_body.rs"),
"format error response-body helpers should live under finalize/error_body.rs"
);
assert!(
!workspace_file_exists("crates/aether-ai-formats/src/protocol/conversion/error.rs"),
!workspace_file_exists("crates/aether-ai-formats/src/formats/conversion/error.rs"),
"aether-ai-formats should not keep error response-body helpers under conversion"
);
@@ -3582,7 +3583,7 @@ fn ai_serving_error_body_is_owned_by_format_finalize_module() {
#[test]
fn ai_serving_conversion_request_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/protocol/conversion/request/mod.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/conversion/request.rs"),
"request conversion should live in aether-ai-formats"
);
assert!(
@@ -3609,7 +3610,7 @@ fn ai_serving_conversion_request_is_owned_by_format_crate() {
let surface_api = read_workspace_file("crates/aether-ai-formats/src/api.rs");
assert!(
surface_api.contains("pub use aether_ai_formats::protocol::conversion::request::{"),
surface_api.contains("pub use aether_ai_formats::formats::conversion::request::{"),
"format API facade should re-export request conversion directly from aether-ai-formats"
);
}
@@ -3617,7 +3618,7 @@ fn ai_serving_conversion_request_is_owned_by_format_crate() {
#[test]
fn ai_serving_conversion_response_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/protocol/conversion/response/mod.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/conversion/response.rs"),
"response conversion should live in aether-ai-formats"
);
assert!(
@@ -3644,7 +3645,7 @@ fn ai_serving_conversion_response_is_owned_by_format_crate() {
let surface_api = read_workspace_file("crates/aether-ai-formats/src/api.rs");
assert!(
surface_api.contains("pub use aether_ai_formats::protocol::conversion::response::{"),
surface_api.contains("pub use aether_ai_formats::formats::conversion::response::{"),
"format API facade should re-export response conversion directly from aether-ai-formats"
);
}
@@ -3652,7 +3653,7 @@ fn ai_serving_conversion_response_is_owned_by_format_crate() {
#[test]
fn ai_format_crate_owns_conversion_and_surface_facade() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/protocol/conversion"),
workspace_file_exists("crates/aether-ai-formats/src/formats/conversion"),
"aether-ai-formats should own the conversion directory"
);
@@ -3665,9 +3666,9 @@ fn ai_format_crate_owns_conversion_and_surface_facade() {
let surface_api = read_workspace_file("crates/aether-ai-formats/src/api.rs");
for pattern in [
"pub use aether_ai_formats::{",
"pub use aether_ai_formats::protocol::conversion::request::{",
"pub use aether_ai_formats::protocol::conversion::response::{",
"pub use crate::response::error_body::{",
"pub use aether_ai_formats::formats::conversion::request::{",
"pub use aether_ai_formats::formats::conversion::response::{",
"pub use crate::formats::shared::error_body::{",
] {
assert!(
surface_api.contains(pattern),
@@ -3746,12 +3747,12 @@ fn ai_serving_finalize_standard_sync_response_converters_are_owned_by_format_cra
#[test]
fn ai_serving_finalize_stream_engine_is_owned_by_format_crate() {
for path in [
"crates/aether-ai-formats/src/response/sse.rs",
"crates/aether-ai-formats/src/response/standard/stream_core/common.rs",
"crates/aether-ai-formats/src/response/standard/stream_core/format_matrix.rs",
"crates/aether-ai-formats/src/response/standard/openai/stream.rs",
"crates/aether-ai-formats/src/response/standard/claude/stream.rs",
"crates/aether-ai-formats/src/response/standard/gemini/stream.rs",
"crates/aether-ai-formats/src/formats/shared/sse.rs",
"crates/aether-ai-formats/src/formats/shared/stream_core/common.rs",
"crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs",
"crates/aether-ai-formats/src/formats/openai/chat/stream.rs",
"crates/aether-ai-formats/src/formats/claude/messages/stream.rs",
"crates/aether-ai-formats/src/formats/gemini/generate_content/stream.rs",
] {
assert!(
workspace_file_exists(path),
@@ -3787,7 +3788,7 @@ fn ai_serving_finalize_stream_engine_is_owned_by_format_crate() {
}
let surface_format_matrix = read_workspace_file(
"crates/aether-ai-formats/src/response/standard/stream_core/format_matrix.rs",
"crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs",
);
for pattern in [
"pub struct StreamingStandardFormatMatrix",
@@ -3811,16 +3812,16 @@ fn ai_serving_finalize_stream_engine_is_owned_by_format_crate() {
#[test]
fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/response/sync_products.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/sync_products.rs"),
"finalize sync_products should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/response/sync_to_stream.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/sync_to_stream.rs"),
"finalize sync-to-stream bridge should live in aether-ai-formats"
);
let surface_sync_products =
read_workspace_file("crates/aether-ai-formats/src/response/sync_products.rs");
read_workspace_file("crates/aether-ai-formats/src/formats/shared/sync_products.rs");
for expected in [
"pub fn maybe_build_standard_cross_format_sync_product_from_normalized_payload(",
"pub fn maybe_build_standard_same_format_sync_body_from_normalized_payload(",
@@ -3947,7 +3948,7 @@ fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
}
let surface_openai_image_stream =
read_workspace_file("crates/aether-ai-formats/src/response/openai_image_stream.rs");
read_workspace_file("crates/aether-ai-formats/src/formats/openai/image/stream.rs");
for expected in [
"pub fn maybe_build_openai_image_sync_finalize_product(",
"pub struct OpenAiImageSyncFinalizeProduct",
@@ -3962,7 +3963,7 @@ fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
}
let surface_sync_to_stream =
read_workspace_file("crates/aether-ai-formats/src/response/sync_to_stream.rs");
read_workspace_file("crates/aether-ai-formats/src/formats/shared/sync_to_stream.rs");
for expected in [
"pub fn maybe_bridge_standard_sync_json_to_stream(",
"pub struct SyncToStreamBridgeOutcome",
@@ -3999,11 +4000,11 @@ fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
#[test]
fn ai_serving_finalize_stream_rewrite_matrix_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/response/stream_rewrite.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/stream_rewrite.rs"),
"finalize stream rewrite matrix should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/response/openai_image_stream.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/openai/image/stream.rs"),
"OpenAI image stream rewrite state should live in aether-ai-formats"
);
@@ -4052,7 +4053,7 @@ fn ai_serving_finalize_stream_rewrite_matrix_is_owned_by_format_crate() {
#[test]
fn ai_serving_planner_common_parser_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/common.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/request.rs"),
"planner/common pure parser should exist in aether-ai-formats"
);
@@ -4195,7 +4196,7 @@ fn ai_serving_root_owns_shared_gemini_request_path_parser() {
#[test]
fn ai_serving_planner_standard_normalize_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/standard/normalize.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/standard_normalize.rs"),
"planner/standard/normalize should live in aether-ai-formats"
);
@@ -4258,7 +4259,7 @@ fn ai_serving_planner_standard_normalize_is_owned_by_format_crate() {
#[test]
fn ai_serving_openai_helpers_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/openai.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/openai/shared.rs"),
"planner/openai helper owner should exist in aether-ai-formats"
);
@@ -4291,17 +4292,17 @@ fn ai_serving_openai_helpers_are_owned_by_format_crate() {
#[test]
fn ai_serving_standard_matrix_delegates_format_conversion_to_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/matrix.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/request_matrix.rs"),
"planner/matrix facade should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/standard/matrix.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/standard_matrix.rs"),
"format standard request-body planner should live in aether-ai-formats"
);
for path in [
"crates/aether-ai-formats/src/protocol/canonical.rs",
"crates/aether-ai-formats/src/protocol/matrix.rs",
"crates/aether-ai-formats/src/protocol/registry.rs",
"crates/aether-ai-formats/src/formats/matrix.rs",
"crates/aether-ai-formats/src/formats/registry.rs",
] {
assert!(
workspace_file_exists(path),
@@ -4309,12 +4310,12 @@ fn ai_serving_standard_matrix_delegates_format_conversion_to_format_crate() {
);
}
let surface_matrix =
read_workspace_file("crates/aether-ai-formats/src/request/standard/matrix.rs");
read_workspace_file("crates/aether-ai-formats/src/formats/shared/standard_matrix.rs");
assert!(
surface_matrix.contains("use aether_ai_formats::protocol::registry::{")
surface_matrix.contains("use aether_ai_formats::formats::registry::{")
&& surface_matrix.contains("convert_request")
&& surface_matrix.contains("FormatContext")
&& surface_matrix.contains("aether_ai_formats::protocol::conversion::request::{"),
&& surface_matrix.contains("aether_ai_formats::formats::conversion::request::{"),
"format standard matrix should delegate format conversion to aether-ai-formats"
);
for forbidden in [
@@ -4361,23 +4362,27 @@ fn ai_serving_standard_matrix_delegates_format_conversion_to_format_crate() {
#[test]
fn ai_serving_standard_family_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/standard/family.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/family.rs"),
"planner/standard/family pure spec owner should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/standard/claude/chat.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/claude/messages/chat_spec.rs"),
"planner/standard/claude/chat pure spec resolver should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/standard/claude/cli.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/claude/messages/cli_spec.rs"),
"planner/standard/claude/cli pure spec resolver should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/standard/gemini/chat.rs"),
workspace_file_exists(
"crates/aether-ai-formats/src/formats/gemini/generate_content/chat_spec.rs"
),
"planner/standard/gemini/chat pure spec resolver should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/standard/gemini/cli.rs"),
workspace_file_exists(
"crates/aether-ai-formats/src/formats/gemini/generate_content/cli_spec.rs"
),
"planner/standard/gemini/cli pure spec resolver should live in aether-ai-formats"
);
@@ -4444,7 +4449,7 @@ fn ai_serving_standard_family_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_same_format_provider_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/passthrough/provider.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/passthrough.rs"),
"planner/passthrough/provider pure spec owner should live in aether-ai-formats"
);
@@ -4495,7 +4500,7 @@ fn ai_serving_same_format_provider_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_passthrough_provider_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/passthrough/provider.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/passthrough.rs"),
"planner/passthrough/provider pure spec owner should live in aether-ai-formats"
);
@@ -4540,7 +4545,7 @@ fn ai_serving_passthrough_provider_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_specialized_files_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/specialized/files.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/gemini/files/spec.rs"),
"planner/specialized/files pure spec owner should live in aether-ai-formats"
);
@@ -4570,9 +4575,18 @@ fn ai_serving_specialized_files_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_specialized_video_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/specialized/video.rs"),
"planner/specialized/video pure spec owner should live in aether-ai-formats"
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/video.rs"),
"planner/specialized/video shared spec seam should live in aether-ai-formats"
);
for path in [
"crates/aether-ai-formats/src/formats/openai/video/spec.rs",
"crates/aether-ai-formats/src/formats/gemini/video/spec.rs",
] {
assert!(
workspace_file_exists(path),
"{path} should own provider-specific video create spec resolution"
);
}
let video =
read_workspace_file("apps/aether-gateway/src/ai_serving/planner/specialized/video.rs");
@@ -4598,7 +4612,7 @@ fn ai_serving_specialized_video_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_openai_responses_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/request/standard/openai_responses.rs"),
workspace_file_exists("crates/aether-ai-formats/src/formats/openai/responses/spec.rs"),
"planner/standard/openai_responses pure spec owner should live in aether-ai-formats"
);
@@ -4639,8 +4653,8 @@ fn ai_serving_openai_responses_specs_are_owned_by_format_crate() {
fn ai_serving_legacy_api_format_names_stay_out_of_primary_paths() {
for path in [
"crates/aether-ai-formats/src/contracts/plan_kinds.rs",
"crates/aether-ai-formats/src/request/route.rs",
"crates/aether-ai-formats/src/request/standard/openai_responses.rs",
"crates/aether-ai-formats/src/formats/shared/routing.rs",
"crates/aether-ai-formats/src/formats/openai/responses/spec.rs",
"apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs",
"apps/aether-gateway/src/execution_runtime/fallback.rs",
] {
@@ -4664,7 +4678,7 @@ fn ai_serving_legacy_api_format_names_stay_out_of_primary_paths() {
}
}
let registry = read_workspace_file("crates/aether-ai-formats/src/protocol/registry.rs");
let registry = read_workspace_file("crates/aether-ai-formats/src/formats/registry.rs");
let implementation = registry
.split("#[cfg(test)]")
.next()
@@ -4697,9 +4711,9 @@ fn retired_api_format_occurrences_are_whitelisted() {
let allowed_paths = [
"apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs",
"crates/aether-ai-formats/src/protocol/formats.rs",
"crates/aether-ai-formats/src/protocol/matrix.rs",
"crates/aether-ai-formats/src/protocol/registry.rs",
"crates/aether-ai-formats/src/formats/id.rs",
"crates/aether-ai-formats/src/formats/matrix.rs",
"crates/aether-ai-formats/src/formats/registry.rs",
"crates/aether-data/src/migrate.rs",
"crates/aether-data/src/lifecycle/migrate/tests.rs",
"crates/aether-usage-runtime/src/report.rs",

View File

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

View File

@@ -614,7 +614,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
"auto_delete_on_expiry": false,
"total_requests": 12,
"total_tokens": 3456,
"total_cost_usd": 1.25
"total_cost_usd": "1.25000000"
}]
}],
"standalone_keys": [{
@@ -630,7 +630,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
"auto_delete_on_expiry": false,
"total_requests": 3,
"total_tokens": 789,
"total_cost_usd": 0.75,
"total_cost_usd": "0.75000000",
"wallet": {
"balance": 30.0,
"recharge_balance": 20.0,
@@ -713,6 +713,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
assert_eq!(user_api_keys.len(), 1);
assert_eq!(user_api_keys[0].name.as_deref(), Some("Alice CLI"));
assert_eq!(user_api_keys[0].total_tokens, 3456);
assert_eq!(user_api_keys[0].total_cost_usd, 1.25);
assert_eq!(
user_api_keys[0].allowed_api_formats,
Some(vec!["openai:chat".to_string()])
@@ -739,6 +740,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
Some("Imported Standalone")
);
assert_eq!(standalone_keys[0].total_tokens, 789);
assert_eq!(standalone_keys[0].total_cost_usd, 0.75);
assert_eq!(
decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,

View File

@@ -72,7 +72,5 @@ async fn gateway_background_gemini_file_mapping_cleanup_deletes_expired_entries(
"active mapping should remain"
);
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
}

View File

@@ -231,9 +231,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
}]
);
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
execution_runtime_handle.abort();
}
@@ -321,8 +319,6 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
}]
);
for handle in background_tasks {
handle.abort();
}
background_tasks.shutdown().await;
upstream_handle.abort();
}

View File

@@ -9,6 +9,7 @@ use crate::clock::current_unix_ms;
use crate::log_ids::short_request_id;
use crate::orchestration::{apply_local_report_effect, LocalReportEffect};
use crate::request_candidate_runtime::record_report_request_candidate_status;
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_USAGE_SYNC_REPORT};
use crate::{AppState, GatewayError};
mod context;
@@ -142,7 +143,7 @@ pub(crate) async fn submit_sync_report(
pub(crate) fn spawn_sync_report(state: AppState, payload: GatewaySyncReportRequest) {
let report_request_id_for_log =
short_request_id(report_request_id(payload.report_context.as_ref()));
tokio::spawn(async move {
spawn_fire_and_forget(TASK_KEY_USAGE_SYNC_REPORT, async move {
let trace_id = payload.trace_id.clone();
if let Err(err) = submit_sync_report(&state, payload).await {
warn!(

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

Some files were not shown because too many files have changed in this diff Show More