mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +08:00
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:
@@ -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
|
||||
|
||||
@@ -92,6 +92,7 @@ fn maybe_build_local_openai_image_sync_finalize_response(
|
||||
payload.report_kind.as_str(),
|
||||
payload.status_code,
|
||||
payload.report_context.as_ref(),
|
||||
payload.body_json.as_ref(),
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.map_err(GatewayError::from)?
|
||||
|
||||
@@ -75,6 +75,18 @@ enum LocalExecutionCandidateAttemptSourceItem<'a> {
|
||||
}
|
||||
|
||||
impl<'a> LocalExecutionCandidateAttemptSource<'a> {
|
||||
pub(crate) fn from_static_attempts_for_image_bridge(
|
||||
attempts: Vec<LocalExecutionCandidateAttempt>,
|
||||
) -> Self {
|
||||
let mut items = VecDeque::new();
|
||||
if !attempts.is_empty() {
|
||||
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static {
|
||||
attempts: VecDeque::from(attempts),
|
||||
});
|
||||
}
|
||||
Self { items }
|
||||
}
|
||||
|
||||
pub(crate) async fn next_attempt(&mut self) -> Option<LocalExecutionCandidateAttempt> {
|
||||
loop {
|
||||
let front = self.items.front_mut()?;
|
||||
|
||||
@@ -398,15 +398,11 @@ pub(crate) fn candidate_auth_channel_skip_reason(
|
||||
) -> Option<&'static str> {
|
||||
let request_auth_channel = normalize_request_auth_channel(request_auth_channel?)?;
|
||||
let upstream_auth_channel = resolve_transport_request_auth_channel(transport)?;
|
||||
if request_auth_channel == upstream_auth_channel
|
||||
|| provider_runtime_policy(&transport.provider.provider_type)
|
||||
.allow_auth_channel_mismatch_by_default
|
||||
|| allow_auth_channel_mismatch_for_format(transport)
|
||||
{
|
||||
None
|
||||
} else {
|
||||
Some("auth_channel_mismatch")
|
||||
if request_auth_channel == upstream_auth_channel {
|
||||
return None;
|
||||
}
|
||||
auth_channel_mismatch_is_explicitly_disabled_for_format(transport)
|
||||
.then_some("auth_channel_mismatch")
|
||||
}
|
||||
|
||||
fn normalize_request_auth_channel(value: &str) -> Option<&'static str> {
|
||||
@@ -452,19 +448,22 @@ fn resolve_transport_auth_type_for_endpoint_format(
|
||||
.unwrap_or(default_auth_type)
|
||||
}
|
||||
|
||||
fn allow_auth_channel_mismatch_for_format(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||
fn auth_channel_mismatch_is_explicitly_disabled_for_format(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
let api_format = crate::ai_serving::normalize_api_format_alias(&transport.endpoint.api_format);
|
||||
transport
|
||||
let Some(items) = transport
|
||||
.key
|
||||
.allow_auth_channel_mismatch_formats
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.is_some_and(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(serde_json::Value::as_str)
|
||||
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
|
||||
})
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
!items
|
||||
.iter()
|
||||
.filter_map(serde_json::Value::as_str)
|
||||
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
|
||||
}
|
||||
|
||||
pub(crate) async fn read_candidate_transport_snapshot(
|
||||
@@ -585,11 +584,11 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_channel_gate_skips_mismatched_raw_secret_auth() {
|
||||
fn auth_channel_gate_allows_mismatched_raw_secret_auth_by_default() {
|
||||
let transport = sample_transport("bearer");
|
||||
assert_eq!(
|
||||
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
|
||||
Some("auth_channel_mismatch")
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
@@ -603,6 +602,16 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_channel_gate_blocks_explicitly_disabled_mismatch_format() {
|
||||
let mut transport = sample_transport("bearer");
|
||||
transport.key.allow_auth_channel_mismatch_formats = Some(json!(["openai:responses"]));
|
||||
assert_eq!(
|
||||
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
|
||||
Some("auth_channel_mismatch")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_channel_gate_treats_cli_oauth_provider_as_bearer_like() {
|
||||
let mut transport = sample_transport("oauth");
|
||||
@@ -613,7 +622,7 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
|
||||
Some("auth_channel_mismatch")
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -159,6 +159,40 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
|
||||
.into_iter()
|
||||
.map(str::to_string)
|
||||
.collect::<Vec<_>>();
|
||||
preselect_local_execution_candidates_for_api_formats_with_serving(
|
||||
state,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
use_api_format_alias_match,
|
||||
key_mode,
|
||||
candidate_api_formats,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_serving(
|
||||
state: PlannerAppState<'_>,
|
||||
client_api_format: &str,
|
||||
requested_model: &str,
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
candidate_api_formats: Vec<String>,
|
||||
) -> Result<
|
||||
AiCandidatePreselectionOutcome<
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
SkippedLocalExecutionCandidate,
|
||||
>,
|
||||
GatewayError,
|
||||
> {
|
||||
let mut model_directive_enabled_api_formats = BTreeSet::new();
|
||||
for api_format in &candidate_api_formats {
|
||||
if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -7,6 +7,7 @@ use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
@@ -284,7 +285,13 @@ impl LocalOpenAiImageSyncAttemptSource<'_> {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match build_passthrough_sync_plan_from_decision(self.parts, payload) {
|
||||
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
|
||||
let built = if provider_api_format == "gemini:generate_content" {
|
||||
build_gemini_sync_plan_from_decision(self.parts, self.body_json, payload)
|
||||
} else {
|
||||
build_passthrough_sync_plan_from_decision(self.parts, payload)
|
||||
};
|
||||
match built {
|
||||
Ok(value) => Ok(value),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -320,7 +327,13 @@ impl LocalOpenAiImageStreamAttemptSource<'_> {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false) {
|
||||
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
|
||||
let built = if provider_api_format == "gemini:generate_content" {
|
||||
build_gemini_stream_plan_from_decision(self.parts, self.body_json, payload)
|
||||
} else {
|
||||
build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false)
|
||||
};
|
||||
match built {
|
||||
Ok(value) => Ok(value),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -508,7 +521,13 @@ async fn build_local_sync_plan_and_reports(
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_passthrough_sync_plan_from_decision(parts, payload) {
|
||||
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
|
||||
let built = if provider_api_format == "gemini:generate_content" {
|
||||
build_gemini_sync_plan_from_decision(parts, body_json, payload)
|
||||
} else {
|
||||
build_passthrough_sync_plan_from_decision(parts, payload)
|
||||
};
|
||||
match built {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
@@ -578,7 +597,13 @@ async fn build_local_stream_plan_and_reports(
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_standard_stream_plan_from_decision(parts, body_json, payload, false) {
|
||||
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
|
||||
let built = if provider_api_format == "gemini:generate_content" {
|
||||
build_gemini_stream_plan_from_decision(parts, body_json, payload)
|
||||
} else {
|
||||
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
|
||||
};
|
||||
match built {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
|
||||
@@ -10,7 +10,7 @@ use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_profile,
|
||||
};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{AiExecutionDecision, AppState};
|
||||
use crate::{append_execution_contract_fields_to_value, AiExecutionDecision, AppState};
|
||||
|
||||
use super::request::resolve_local_openai_image_candidate_payload_parts;
|
||||
use super::support::{LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput};
|
||||
@@ -47,8 +47,12 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
} = attempt;
|
||||
let candidate = eligible.candidate;
|
||||
let transport = resolved.transport;
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
|
||||
let provider_api_format = resolved.provider_api_format.clone();
|
||||
let needs_conversion = provider_api_format != spec_metadata.api_format;
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format.as_str(),
|
||||
);
|
||||
let proxy = planner_state
|
||||
.app()
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
||||
@@ -84,41 +88,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,
|
||||
|
||||
@@ -9,14 +9,17 @@ use crate::ai_serving::planner::candidate_preparation::{
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
|
||||
use crate::ai_serving::transport::{
|
||||
build_openai_image_headers, build_openai_image_upstream_url,
|
||||
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
|
||||
ProviderOpenAiImageHeadersInput,
|
||||
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
|
||||
resolve_openai_image_auth, ProviderOpenAiImageHeadersInput,
|
||||
StandardProviderRequestHeadersInput,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
|
||||
build_chatgpt_web_image_request_body, build_openai_image_provider_request_body,
|
||||
default_model_for_openai_image_operation, normalize_openai_image_request,
|
||||
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, PlannerAppState,
|
||||
build_chatgpt_web_image_request_body,
|
||||
build_gemini_image_request_body_from_openai_image_request,
|
||||
build_openai_image_provider_request_body, default_model_for_openai_image_operation,
|
||||
normalize_openai_image_request, request_conversion_direct_auth, CandidateFailureDiagnostic,
|
||||
GatewayProviderTransportSnapshot, PlannerAppState, RequestConversionKind,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
@@ -35,6 +38,7 @@ pub(super) struct LocalOpenAiImageCandidatePayloadParts {
|
||||
pub(super) auth_value: String,
|
||||
pub(super) requested_model: String,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) provider_api_format: String,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) provider_request_body: Value,
|
||||
pub(super) upstream_url: String,
|
||||
@@ -54,6 +58,21 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let provider_api_format = attempt.eligible.provider_api_format.as_str();
|
||||
|
||||
if provider_api_format == "gemini:generate_content" {
|
||||
return resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
if let Some(skip_reason) =
|
||||
openai_image_transport_unsupported_reason(transport, spec_metadata.api_format)
|
||||
@@ -70,7 +89,6 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let prepared_candidate = match prepare_header_authenticated_candidate(
|
||||
PlannerAppState::new(state),
|
||||
transport,
|
||||
@@ -216,6 +234,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
auth_value,
|
||||
requested_model,
|
||||
mapped_model,
|
||||
provider_api_format: spec_metadata.api_format.to_string(),
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
upstream_url,
|
||||
@@ -223,6 +242,190 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
})
|
||||
}
|
||||
|
||||
async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
attempt: &LocalOpenAiImageCandidateAttempt,
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Option<LocalOpenAiImageCandidatePayloadParts> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let provider_api_format = "gemini:generate_content";
|
||||
|
||||
let prepared_candidate = match prepare_header_authenticated_candidate(
|
||||
PlannerAppState::new(state),
|
||||
transport,
|
||||
candidate,
|
||||
request_conversion_direct_auth(transport, RequestConversionKind::ToGeminiStandard),
|
||||
OauthPreparationContext {
|
||||
trace_id,
|
||||
api_format: provider_api_format,
|
||||
operation: "openai_image_to_gemini_candidate_request",
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prepared) => prepared,
|
||||
Err(skip_reason) => {
|
||||
mark_skipped_local_openai_image_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let Some(normalized_request) = normalize_openai_image_request(parts, body_json, body_base64)
|
||||
else {
|
||||
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_missing",
|
||||
CandidateFailureDiagnostic::provider_request_body_missing(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"openai_image_request_normalize",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(mut converted) = build_gemini_image_request_body_from_openai_image_request(
|
||||
&normalized_request,
|
||||
&prepared_candidate.mapped_model,
|
||||
) else {
|
||||
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_missing",
|
||||
CandidateFailureDiagnostic::provider_request_body_missing(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"openai_image_to_gemini_request_body",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
converted.body_json =
|
||||
match crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers(
|
||||
converted.body_json,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
body_json,
|
||||
&parts.headers,
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_missing",
|
||||
CandidateFailureDiagnostic::provider_request_body_missing(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"openai_image_to_gemini_body_rules",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let upstream_is_stream = spec_metadata.require_streaming;
|
||||
let Some(upstream_url) = crate::ai_serving::planner::standard::build_standard_upstream_url(
|
||||
parts,
|
||||
transport,
|
||||
&converted.mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
) else {
|
||||
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"openai_image_to_gemini_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let Some(resolved_headers) =
|
||||
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format: false,
|
||||
headers: &parts.headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
provider_request_body: &converted.body_json,
|
||||
original_request_body: body_json,
|
||||
upstream_is_stream,
|
||||
})
|
||||
else {
|
||||
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::header_rules_apply_failed(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"openai_image_to_gemini_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
Some(LocalOpenAiImageCandidatePayloadParts {
|
||||
transport: Arc::clone(transport),
|
||||
auth_header: resolved_headers.auth_header,
|
||||
auth_value: resolved_headers.auth_value,
|
||||
requested_model: converted.requested_model,
|
||||
mapped_model: converted.mapped_model,
|
||||
provider_api_format: provider_api_format.to_string(),
|
||||
provider_request_headers: resolved_headers.headers,
|
||||
provider_request_body: converted.body_json,
|
||||
upstream_url,
|
||||
input_summary: converted.summary_json,
|
||||
})
|
||||
}
|
||||
|
||||
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
|
||||
let base_url = base_url.trim().trim_end_matches('/');
|
||||
let base_url = if base_url.is_empty() {
|
||||
|
||||
@@ -11,6 +11,7 @@ use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
@@ -19,9 +20,9 @@ use crate::ai_serving::planner::materialization_policy::{
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
|
||||
use crate::ai_serving::{
|
||||
extract_pool_sticky_session_token, resolve_local_decision_execution_runtime_auth_context,
|
||||
CandidateFailureDiagnostic, ExecutionRuntimeAuthContext, GatewayControlDecision,
|
||||
PlannerAppState,
|
||||
extract_pool_sticky_session_token, request_candidate_api_formats,
|
||||
resolve_local_decision_execution_runtime_auth_context, CandidateFailureDiagnostic,
|
||||
ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
|
||||
};
|
||||
use crate::client_session_affinity::client_session_affinity_from_parts;
|
||||
use crate::clock::current_unix_secs;
|
||||
@@ -88,52 +89,69 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
|
||||
api_format: &str,
|
||||
decision_kind: &str,
|
||||
) -> Option<Vec<LocalOpenAiImageCandidateAttempt>> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let (candidates, preselection_skipped) = match planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind,
|
||||
error = ?err,
|
||||
"gateway local openai image decision scheduler selection failed"
|
||||
);
|
||||
return None;
|
||||
let candidate_api_formats = image_candidate_api_formats(api_format);
|
||||
let mut attempts = Vec::new();
|
||||
for candidate_api_format in candidate_api_formats {
|
||||
let matches_client_format = candidate_api_format == api_format;
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let (mut candidates, preselection_skipped) = match planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
candidate_api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
input.required_capabilities.as_ref(),
|
||||
matches_client_format.then_some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind,
|
||||
api_format = candidate_api_format,
|
||||
error = ?err,
|
||||
"gateway local openai image decision scheduler selection failed"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if !matches_client_format {
|
||||
candidates.retain(|candidate| {
|
||||
auth_snapshot_allows_cross_format_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
candidate,
|
||||
false,
|
||||
)
|
||||
});
|
||||
}
|
||||
};
|
||||
attempts.extend(
|
||||
materialize_local_openai_image_candidate_attempts(
|
||||
planner_state,
|
||||
trace_id,
|
||||
input,
|
||||
body_json,
|
||||
candidates,
|
||||
preselection_skipped
|
||||
.into_iter()
|
||||
.map(|item| SkippedLocalExecutionCandidate {
|
||||
candidate: item.candidate,
|
||||
skip_reason: item.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
})
|
||||
.collect(),
|
||||
api_format,
|
||||
)
|
||||
.await,
|
||||
);
|
||||
}
|
||||
|
||||
Some(
|
||||
materialize_local_openai_image_candidate_attempts(
|
||||
planner_state,
|
||||
trace_id,
|
||||
input,
|
||||
body_json,
|
||||
candidates,
|
||||
preselection_skipped
|
||||
.into_iter()
|
||||
.map(|item| SkippedLocalExecutionCandidate {
|
||||
candidate: item.candidate,
|
||||
skip_reason: item.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
})
|
||||
.collect(),
|
||||
api_format,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
Some(attempts)
|
||||
}
|
||||
|
||||
pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
|
||||
@@ -145,29 +163,56 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
|
||||
decision_kind: &str,
|
||||
) -> Result<Option<(LocalOpenAiImageCandidateAttemptSource<'a>, usize)>, GatewayError> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let (candidates, preselection_skipped) = match planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind,
|
||||
error = ?err,
|
||||
"gateway local openai image decision scheduler selection failed"
|
||||
);
|
||||
return Ok(None);
|
||||
let mut candidates = Vec::new();
|
||||
let mut preselection_skipped = Vec::new();
|
||||
for candidate_api_format in image_candidate_api_formats(api_format) {
|
||||
let matches_client_format = candidate_api_format == api_format;
|
||||
match planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
candidate_api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
input.required_capabilities.as_ref(),
|
||||
matches_client_format.then_some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok((mut format_candidates, mut format_skipped)) => {
|
||||
if !matches_client_format {
|
||||
format_candidates.retain(|candidate| {
|
||||
auth_snapshot_allows_cross_format_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
candidate,
|
||||
false,
|
||||
)
|
||||
});
|
||||
format_skipped.retain(|candidate| {
|
||||
auth_snapshot_allows_cross_format_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
&candidate.candidate,
|
||||
false,
|
||||
)
|
||||
});
|
||||
}
|
||||
candidates.append(&mut format_candidates);
|
||||
preselection_skipped.append(&mut format_skipped);
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind,
|
||||
api_format = candidate_api_format,
|
||||
error = ?err,
|
||||
"gateway local openai image decision scheduler selection failed"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
let sticky_session_token = extract_pool_sticky_session_token(body_json);
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
@@ -198,23 +243,34 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
|
||||
extra_data: None,
|
||||
})
|
||||
.collect(),
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
LocalCandidateResolutionMode::WithoutTransportPairGate,
|
||||
|eligible| {
|
||||
Some(build_local_execution_candidate_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: api_format,
|
||||
provider_api_format: eligible.provider_api_format.as_str(),
|
||||
client_api_format: api_format,
|
||||
extra_fields: serde_json::Map::new(),
|
||||
},
|
||||
))
|
||||
},
|
||||
|mut skipped_candidate| {
|
||||
let provider_api_format = skipped_candidate
|
||||
.transport
|
||||
.as_ref()
|
||||
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
|
||||
.unwrap_or_else(|| {
|
||||
skipped_candidate
|
||||
.candidate
|
||||
.endpoint_api_format
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
});
|
||||
skipped_candidate.extra_data =
|
||||
Some(build_local_execution_candidate_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
api_format,
|
||||
provider_api_format.as_str(),
|
||||
api_format,
|
||||
serde_json::Map::new(),
|
||||
));
|
||||
@@ -254,23 +310,34 @@ async fn materialize_local_openai_image_candidate_attempts(
|
||||
persistence_policy,
|
||||
candidates,
|
||||
preselection_skipped,
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
LocalCandidateResolutionMode::WithoutTransportPairGate,
|
||||
|eligible| {
|
||||
Some(build_local_execution_candidate_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: api_format,
|
||||
provider_api_format: eligible.provider_api_format.as_str(),
|
||||
client_api_format: api_format,
|
||||
extra_fields: serde_json::Map::new(),
|
||||
},
|
||||
))
|
||||
},
|
||||
|mut skipped_candidate| {
|
||||
let provider_api_format = skipped_candidate
|
||||
.transport
|
||||
.as_ref()
|
||||
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
|
||||
.unwrap_or_else(|| {
|
||||
skipped_candidate
|
||||
.candidate
|
||||
.endpoint_api_format
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
});
|
||||
skipped_candidate.extra_data =
|
||||
Some(build_local_execution_candidate_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
api_format,
|
||||
provider_api_format.as_str(),
|
||||
api_format,
|
||||
serde_json::Map::new(),
|
||||
));
|
||||
@@ -282,6 +349,14 @@ async fn materialize_local_openai_image_candidate_attempts(
|
||||
outcome.attempts
|
||||
}
|
||||
|
||||
fn image_candidate_api_formats(api_format: &str) -> Vec<&'static str> {
|
||||
if api_format.trim().eq_ignore_ascii_case("openai:image") {
|
||||
vec!["openai:image", "gemini:generate_content"]
|
||||
} else {
|
||||
request_candidate_api_formats(api_format, false)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_openai_image_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()..]
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(&[]);
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
) {
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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}")))
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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,
|
||||
)));
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
561
apps/aether-gateway/src/handlers/public/support/install.rs
Normal file
561
apps/aether-gateway/src/handlers/public/support/install.rs
Normal 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)
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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) =>
|
||||
{
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -58,6 +58,7 @@ mod router;
|
||||
mod scheduler;
|
||||
mod state;
|
||||
mod system_features;
|
||||
mod task_runtime;
|
||||
mod tunnel;
|
||||
mod usage;
|
||||
mod video_tasks;
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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?;
|
||||
|
||||
@@ -3,8 +3,8 @@ use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_data::repository::proxy_nodes::{
|
||||
InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation, ProxyNodeReadRepository,
|
||||
ProxyNodeWriteRepository, StoredProxyNode,
|
||||
bucket_start_unix_secs, InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeWriteRepository, StoredProxyNode,
|
||||
};
|
||||
use aether_runtime::bounded_queue;
|
||||
use axum::extract::ws::Message;
|
||||
@@ -14,23 +14,25 @@ use serde_json::json;
|
||||
use tokio::sync::watch;
|
||||
|
||||
use super::{
|
||||
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_with, cleanup_stale_proxy_nodes_once,
|
||||
inspect_proxy_upgrade_rollout, next_daily_run_after, next_db_maintenance_run_after,
|
||||
next_stats_aggregation_run_after, next_stats_hourly_aggregation_run_after,
|
||||
pending_cleanup_batch_size, pending_cleanup_timeout_minutes, plan_pending_cleanup_batch,
|
||||
provider_checkin_schedule, record_proxy_upgrade_traffic_success, run_db_maintenance_with,
|
||||
run_proxy_upgrade_rollout_once, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
|
||||
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
|
||||
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
|
||||
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_with, cleanup_proxy_node_metrics_at,
|
||||
cleanup_proxy_node_metrics_once, cleanup_stale_proxy_nodes_once, inspect_proxy_upgrade_rollout,
|
||||
next_daily_run_after, next_db_maintenance_run_after, next_stats_aggregation_run_after,
|
||||
next_stats_hourly_aggregation_run_after, pending_cleanup_batch_size,
|
||||
pending_cleanup_timeout_minutes, plan_pending_cleanup_batch, provider_checkin_schedule,
|
||||
proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success,
|
||||
run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_audit_cleanup_worker,
|
||||
spawn_db_maintenance_worker, spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker,
|
||||
spawn_pool_monitor_worker, spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
|
||||
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
|
||||
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
|
||||
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
|
||||
start_proxy_upgrade_rollout, stats_aggregation_target_day,
|
||||
stats_hourly_aggregation_target_hour, summarize_database_pool, usage_cleanup_settings,
|
||||
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
|
||||
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
|
||||
UsageCleanupSettings, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
FailedPendingUsageRow, GatewayDataState, ProxyNodeMetricsCleanupSettings,
|
||||
ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings, USAGE_CLEANUP_HOUR,
|
||||
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -734,6 +736,176 @@ async fn usage_cleanup_settings_resolve_batch_and_delete_toggle() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_node_metrics_cleanup_settings_use_dedicated_retention_and_batch_limits() {
|
||||
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
|
||||
("cleanup_batch_size".to_string(), json!(250)),
|
||||
("proxy_node_metrics_1m_retention_days".to_string(), json!(0)),
|
||||
("proxy_node_metrics_1h_retention_days".to_string(), json!(7)),
|
||||
(
|
||||
"proxy_node_metrics_cleanup_batch_size".to_string(),
|
||||
json!(100_000),
|
||||
),
|
||||
]);
|
||||
|
||||
let settings = proxy_node_metrics_cleanup_settings(&data)
|
||||
.await
|
||||
.expect("proxy metrics cleanup settings should resolve");
|
||||
|
||||
assert_eq!(
|
||||
settings,
|
||||
ProxyNodeMetricsCleanupSettings {
|
||||
retain_1m_days: 1,
|
||||
retain_1h_days: 7,
|
||||
batch_size: 50_000,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_node_metrics_cleanup_settings_fallback_to_global_batch_size() {
|
||||
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
|
||||
("cleanup_batch_size".to_string(), json!(250)),
|
||||
(
|
||||
"proxy_node_metrics_cleanup_batch_size".to_string(),
|
||||
json!(0),
|
||||
),
|
||||
]);
|
||||
|
||||
let settings = proxy_node_metrics_cleanup_settings(&data)
|
||||
.await
|
||||
.expect("proxy metrics cleanup settings should resolve");
|
||||
|
||||
assert_eq!(
|
||||
settings,
|
||||
ProxyNodeMetricsCleanupSettings {
|
||||
retain_1m_days: 30,
|
||||
retain_1h_days: 180,
|
||||
batch_size: 250,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_node_metrics_cleanup_deletes_expired_buckets_in_batches() {
|
||||
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
|
||||
sample_connected_proxy_node("node-metrics-1", 30, 1),
|
||||
sample_connected_proxy_node("node-metrics-2", 30, 1),
|
||||
sample_connected_proxy_node("node-metrics-3", 30, 1),
|
||||
]));
|
||||
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
||||
.with_system_config_values_for_tests([
|
||||
("enable_auto_cleanup".to_string(), json!(true)),
|
||||
("proxy_node_metrics_1m_retention_days".to_string(), json!(1)),
|
||||
("proxy_node_metrics_1h_retention_days".to_string(), json!(1)),
|
||||
(
|
||||
"proxy_node_metrics_cleanup_batch_size".to_string(),
|
||||
json!(1),
|
||||
),
|
||||
]);
|
||||
|
||||
for (idx, node_id) in ["node-metrics-1", "node-metrics-2", "node-metrics-3"]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
{
|
||||
repository
|
||||
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
||||
node_id: node_id.to_string(),
|
||||
heartbeat_interval: Some(30),
|
||||
active_connections: Some(i32::try_from(idx + 1).unwrap()),
|
||||
total_requests_delta: None,
|
||||
avg_latency_ms: None,
|
||||
failed_requests_delta: None,
|
||||
dns_failures_delta: None,
|
||||
stream_errors_delta: None,
|
||||
proxy_metadata: Some(json!({
|
||||
"tunnel_metrics": {
|
||||
"connect_errors": idx + 1,
|
||||
"disconnects": 0,
|
||||
"error_events_total": 0,
|
||||
"ws_in_bytes": idx + 1,
|
||||
"ws_out_bytes": idx + 1,
|
||||
"ws_in_frames": idx + 1,
|
||||
"ws_out_frames": idx + 1,
|
||||
"heartbeat_rtt_last_ms": 10
|
||||
}
|
||||
})),
|
||||
proxy_version: Some("1.0.0".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect("heartbeat should write metrics");
|
||||
}
|
||||
|
||||
let now = chrono::Utc::now().timestamp().max(0) as u64;
|
||||
let old_bucket = bucket_start_unix_secs(now, ProxyNodeMetricsStep::OneMinute);
|
||||
let cleanup = data
|
||||
.cleanup_proxy_node_metrics(
|
||||
old_bucket.saturating_add(60),
|
||||
old_bucket.saturating_add(3_600),
|
||||
1,
|
||||
)
|
||||
.await
|
||||
.expect("direct cleanup should delete a limited batch");
|
||||
assert_eq!(cleanup.deleted_1m_rows, 1);
|
||||
assert_eq!(cleanup.deleted_1h_rows, 1);
|
||||
|
||||
let cleanup = cleanup_proxy_node_metrics_at(&data, now.saturating_add(2 * 86_400))
|
||||
.await
|
||||
.expect("runtime cleanup should loop over batches");
|
||||
assert_eq!(cleanup.deleted_1m_rows, 2);
|
||||
assert_eq!(cleanup.deleted_1h_rows, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_node_metrics_cleanup_respects_auto_cleanup_toggle() {
|
||||
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
|
||||
sample_connected_proxy_node("node-metrics-disabled", 30, 1),
|
||||
]));
|
||||
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
||||
.with_system_config_values_for_tests([
|
||||
("enable_auto_cleanup".to_string(), json!(false)),
|
||||
("proxy_node_metrics_1m_retention_days".to_string(), json!(1)),
|
||||
("proxy_node_metrics_1h_retention_days".to_string(), json!(1)),
|
||||
(
|
||||
"proxy_node_metrics_cleanup_batch_size".to_string(),
|
||||
json!(1),
|
||||
),
|
||||
]);
|
||||
|
||||
repository
|
||||
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
||||
node_id: "node-metrics-disabled".to_string(),
|
||||
heartbeat_interval: Some(30),
|
||||
active_connections: Some(1),
|
||||
total_requests_delta: None,
|
||||
avg_latency_ms: None,
|
||||
failed_requests_delta: None,
|
||||
dns_failures_delta: None,
|
||||
stream_errors_delta: None,
|
||||
proxy_metadata: Some(json!({
|
||||
"tunnel_metrics": {
|
||||
"connect_errors": 1,
|
||||
"disconnects": 0,
|
||||
"error_events_total": 0,
|
||||
"ws_in_bytes": 1,
|
||||
"ws_out_bytes": 1,
|
||||
"ws_in_frames": 1,
|
||||
"ws_out_frames": 1,
|
||||
"heartbeat_rtt_last_ms": 10
|
||||
}
|
||||
})),
|
||||
proxy_version: Some("1.0.0".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect("heartbeat should write metrics");
|
||||
|
||||
let cleanup = cleanup_proxy_node_metrics_once(&data)
|
||||
.await
|
||||
.expect("runtime cleanup should short-circuit");
|
||||
assert_eq!(cleanup.deleted_1m_rows, 0);
|
||||
assert_eq!(cleanup.deleted_1h_rows, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_cleanup_window_uses_non_overlapping_ranges() {
|
||||
let now_utc = "2026-03-18T03:00:00Z"
|
||||
|
||||
@@ -12,12 +12,14 @@ use super::{
|
||||
maintenance_timezone, parse_hhmm_time, perform_oauth_token_refresh_once,
|
||||
provider_checkin_schedule, run_audit_cleanup_once, run_db_maintenance_once,
|
||||
run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once, run_pool_monitor_once,
|
||||
run_provider_checkin_once, run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
|
||||
run_provider_checkin_once, run_proxy_node_metrics_cleanup_once,
|
||||
run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
|
||||
run_request_candidate_cleanup_once, run_stats_aggregation_once,
|
||||
run_stats_hourly_aggregation_once, run_usage_cleanup_once,
|
||||
run_wallet_daily_usage_aggregation_once, AUDIT_LOG_CLEANUP_INTERVAL,
|
||||
GEMINI_FILE_MAPPING_CLEANUP_INTERVAL, OAUTH_TOKEN_REFRESH_INTERVAL, PENDING_CLEANUP_INTERVAL,
|
||||
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_STALE_SWEEP_INTERVAL,
|
||||
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_METRICS_CLEANUP_HOUR,
|
||||
PROXY_NODE_METRICS_CLEANUP_MINUTE, PROXY_NODE_STALE_SWEEP_INTERVAL,
|
||||
PROXY_UPGRADE_ROLLOUT_INTERVAL, REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR,
|
||||
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
@@ -292,6 +294,30 @@ pub(crate) fn spawn_proxy_node_stale_cleanup_worker(
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_proxy_node_metrics_cleanup_worker(
|
||||
data: Arc<GatewayDataState>,
|
||||
) -> Option<tokio::task::JoinHandle<()>> {
|
||||
if !data.has_proxy_node_writer() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let timezone = maintenance_timezone();
|
||||
Some(tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(duration_until_next_daily_run(
|
||||
Utc::now(),
|
||||
timezone,
|
||||
PROXY_NODE_METRICS_CLEANUP_HOUR,
|
||||
PROXY_NODE_METRICS_CLEANUP_MINUTE,
|
||||
))
|
||||
.await;
|
||||
if let Err(err) = run_proxy_node_metrics_cleanup_once(&data).await {
|
||||
log_maintenance_worker_failure("proxy_node_metrics_cleanup", "tick", &err);
|
||||
}
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_proxy_upgrade_rollout_worker(
|
||||
state: AppState,
|
||||
) -> Option<tokio::task::JoinHandle<()>> {
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
80
apps/aether-gateway/src/state/runtime/background_tasks.rs
Normal file
80
apps/aether-gateway/src/state/runtime/background_tasks.rs
Normal 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()))
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
647
apps/aether-gateway/src/task_runtime/mod.rs
Normal file
647
apps/aether-gateway/src/task_runtime/mod.rs
Normal 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))
|
||||
}
|
||||
@@ -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(
|
||||
) {
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -110,21 +110,24 @@ sudo aether-proxy uninstall
|
||||
| `--upstream-pool-idle-timeout-secs` | `AETHER_PROXY_UPSTREAM_POOL_IDLE_TIMEOUT_SECS` | `300` | 连接池空闲超时(秒) |
|
||||
| `--upstream-tcp-keepalive-secs` | `AETHER_PROXY_UPSTREAM_TCP_KEEPALIVE_SECS` | `60` | TCP keepalive(秒,0 关闭) |
|
||||
| `--upstream-tcp-nodelay` | `AETHER_PROXY_UPSTREAM_TCP_NODELAY` | `true` | 启用 TCP_NODELAY |
|
||||
| `--upstream-proxy-url` | `AETHER_PROXY_UPSTREAM_PROXY_URL` | 空 | 仅 provider 上游请求使用的出口代理,支持 `http://`、`socks5://`、`socks5h://` |
|
||||
| `--upstream-proxy-url` | `AETHER_PROXY_UPSTREAM_PROXY_URL` | 空 | 仅 provider 上游请求使用的出口代理 |
|
||||
| `--redirect-replay-budget-bytes` | `AETHER_PROXY_REDIRECT_REPLAY_BUDGET_BYTES` | `5M` | 307/308 请求体重放的预读预算,支持 `K/M/G`,`0` 表示禁用 body replay buffering |
|
||||
|
||||
`upstream_proxy_url` 只影响 `aether-proxy` 访问 OpenAI、Claude、Gemini 等 provider 的上游请求,不影响节点回连 Aether 服务器的 WebSocket tunnel。配合 WARP sidecar 时可填写:
|
||||
出口代理支持 `http://`、`socks5://`、`socks5h://`。配合 WARP sidecar 时可填写:
|
||||
|
||||
```toml
|
||||
upstream_proxy_url = "socks5h://microwarp:1080"
|
||||
```
|
||||
|
||||
如果需要让 Aether 管理 API 和 WebSocket tunnel 也走代理,使用 `aether_proxy_url`。
|
||||
|
||||
#### Aether API 客户端
|
||||
|
||||
| 参数 | 环境变量 | 默认值 | 说明 |
|
||||
|------|----------|--------|------|
|
||||
| `--aether-request-timeout-secs` | `AETHER_PROXY_AETHER_REQUEST_TIMEOUT_SECS` | `10` | 请求总超时(秒) |
|
||||
| `--aether-connect-timeout-secs` | `AETHER_PROXY_AETHER_CONNECT_TIMEOUT_SECS` | `10` | 建连超时(秒) |
|
||||
| `--aether-proxy-url` | `AETHER_PROXY_AETHER_PROXY_URL` | 空 | Aether 注册、心跳和 WebSocket tunnel 回连使用的出口代理(默认不走代理) |
|
||||
| `--aether-retry-max-attempts` | `AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS` | `3` | 最大重试次数 |
|
||||
|
||||
#### DNS 与安全
|
||||
@@ -154,6 +157,15 @@ upstream_proxy_url = "socks5h://microwarp:1080"
|
||||
- 以 `systemd` 或 `OpenRC` 安装时默认会额外打开文件日志到 `/var/log/aether-proxy`
|
||||
- OpenRC 安装时,`aether-proxy logs` 实际读取 `/var/log/aether-proxy/current.log` 和 `/var/log/aether-proxy/error.log`;这些文件通常需要用 `sudo aether-proxy logs` 查看
|
||||
|
||||
### 隧道健康上报(Heartbeat)
|
||||
|
||||
proxy 会在心跳 `proxy_metadata` 中主动上报隧道稳定性指标,便于后端直接入库/告警:
|
||||
|
||||
- `proxy_metadata.tunnel_metrics`:建连尝试/成功/失败、断开次数、累计在线时长、心跳 RTT、WebSocket 收发帧与字节等。
|
||||
- `proxy_metadata.recent_tunnel_errors`:最近隧道异常事件(时间戳、类别、错误摘要,环形缓冲)。
|
||||
|
||||
说明:仅主连接(`conn=0`)发送 heartbeat,避免多条 tunnel 重复上报同一份全局指标。
|
||||
|
||||
### 多服务器配置
|
||||
|
||||
在 `aether-proxy.toml` 中使用 `[[servers]]` 配置 Aether 服务器。即使只有一个服务器,也必须写成一个 `[[servers]]` 条目;旧的顶层单服务器写法已不再支持。
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user