Fix Codex image progress heartbeat merge regressions

This commit is contained in:
fawney19
2026-05-10 02:10:23 +08:00
156 changed files with 11741 additions and 943 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -114,6 +114,8 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
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,

View File

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

View File

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

View File

@@ -7,7 +7,10 @@ use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
OauthPreparationContext,
};
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider;
use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
};
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, request_body_build_failure_extra_data,
@@ -175,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,
@@ -225,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,
@@ -237,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() {

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -40,7 +40,8 @@ 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,
gemini_request_is_image_generation, implicit_sync_finalize_report_kind,
@@ -70,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,
@@ -79,11 +81,13 @@ pub(crate) use aether_ai_formats::api::{
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens,
resolve_openai_responses_stream_spec, resolve_openai_responses_sync_spec,
resolve_requested_gemini_image_model_for_request,
resolve_requested_openai_image_model_for_request, stream_body_contains_error_event,
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
transform_provider_private_stream_line, value_as_u64, AiControlPlanRequest,
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
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,
GeminiImageRequestForOpenAi, GeminiProviderState, KiroToClaudeCliStreamState,

View File

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

View File

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

View File

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

View File

@@ -95,6 +95,14 @@ pub(super) fn classify_admin_system_family_route(
"admin:system",
false,
))
} else if method == http::Method::GET && normalized_path == "/api/admin/system/cleanup/runs" {
Some(classified(
"admin_proxy",
"system_manage",
"cleanup_runs",
"admin:system",
false,
))
} else if method == http::Method::POST && normalized_path == "/api/admin/system/purge/config" {
Some(classified(
"admin_proxy",
@@ -139,6 +147,16 @@ pub(super) fn classify_admin_system_family_route(
"admin:system",
false,
))
} else if method == http::Method::POST
&& normalized_path == "/api/admin/system/purge/request-bodies/task"
{
Some(classified(
"admin_proxy",
"system_manage",
"purge_request_bodies_task",
"admin:system",
false,
))
} else if method == http::Method::POST && normalized_path == "/api/admin/system/purge/stats" {
Some(classified(
"admin_proxy",

View File

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

View File

@@ -172,6 +172,10 @@ fn classifies_admin_system_maintenance_write_routes_as_admin_proxy_route() {
"/api/admin/system/purge/request-bodies",
"purge_request_bodies",
),
(
"/api/admin/system/purge/request-bodies/task",
"purge_request_bodies_task",
),
("/api/admin/system/purge/stats", "purge_stats"),
];
@@ -246,6 +250,25 @@ fn classifies_admin_system_stats_as_admin_proxy_route() {
assert!(!decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_admin_system_cleanup_runs_as_admin_proxy_route() {
let headers = headers(&[]);
let uri: Uri = "/api/admin/system/cleanup/runs"
.parse()
.expect("uri should parse");
let decision =
classify_control_route(&http::Method::GET, &uri, &headers).expect("route should classify");
assert_eq!(decision.route_class.as_deref(), Some("admin_proxy"));
assert_eq!(decision.route_family.as_deref(), Some("system_manage"));
assert_eq!(decision.route_kind.as_deref(), Some("cleanup_runs"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("admin:system")
);
assert!(!decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_admin_system_settings_set_as_admin_proxy_route() {
let headers = headers(&[]);

View File

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

View File

@@ -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()
@@ -447,4 +461,14 @@ impl GatewayDataState {
None => Ok(aether_data::repository::system::AdminSystemPurgeSummary::default()),
}
}
pub(crate) async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
) -> Result<aether_data::repository::system::AdminSystemPurgeSummary, DataLayerError> {
match self.backends.as_ref() {
Some(backends) => backends.purge_admin_request_bodies_batch(batch_size).await,
None => Ok(aether_data::repository::system::AdminSystemPurgeSummary::default()),
}
}
}

View File

@@ -81,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,
@@ -140,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>>,
@@ -206,6 +213,14 @@ impl fmt::Debug for GatewayDataState {
.field("has_proxy_node_reader", &self.proxy_node_reader.is_some())
.field("has_proxy_node_writer", &self.proxy_node_writer.is_some())
.field("has_billing_reader", &self.billing_reader.is_some())
.field(
"has_background_task_reader",
&self.background_task_reader.is_some(),
)
.field(
"has_background_task_writer",
&self.background_task_writer.is_some(),
)
.field(
"has_gemini_file_mapping_reader",
&self.gemini_file_mapping_reader.is_some(),

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -27,7 +27,7 @@ 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::time::MissedTickBehavior;
use tokio_util::codec::{FramedRead, LinesCodec};
@@ -35,8 +35,11 @@ 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;
@@ -81,9 +84,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,
@@ -91,7 +95,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,
};
const OPENAI_IMAGE_STREAM_PLAN_KIND: &str = "openai_image_stream";
const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
@@ -137,6 +143,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,
@@ -170,6 +222,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,
@@ -1069,11 +1170,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,
@@ -1164,18 +1283,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 {
@@ -1215,18 +1333,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 {
@@ -1245,46 +1362,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),
@@ -1303,13 +1434,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,
@@ -2746,7 +2876,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;
@@ -3063,9 +3195,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()
@@ -3267,6 +3405,213 @@ mod tests {
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_rewrites_redirect_to_structured_failure() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/v1/execute/stream",
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":302,\"headers\":{\"location\":\"/\",\"content-type\":\"text/html\",\"content-length\":\"0\"}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = axum::http::Response::new(Body::from(frames));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}),
);
axum::serve(listener, app)
.await
.expect("server should start");
});
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
})
.with_execution_runtime_override_base_url(format!("http://{addr}"));
let plan = ExecutionPlan {
request_id: "req-remote-runtime-stream-redirect".into(),
candidate_id: Some("cand-remote-runtime-stream-redirect".into()),
provider_name: Some("ChatGPTWeb".into()),
provider_id: "prov-redirect".into(),
endpoint_id: "ep-redirect".into(),
key_id: "key-redirect".into(),
method: "POST".into(),
url: "https://chatgpt.com/backend-api/codex/responses".into(),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
("accept".into(), "text/event-stream".into()),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5.4",
"input": "hello",
"stream": true
})),
stream: true,
client_api_format: "gemini:generate_content".into(),
provider_api_format: "openai:responses".into(),
model_name: Some("gemini-3.1-flash-image-preview".into()),
proxy: None,
transport_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let decision = GatewayControlDecision::synthetic(
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent",
Some("ai_public".to_string()),
Some("gemini".to_string()),
Some("generate_content".to_string()),
Some("gemini:generate_content".to_string()),
)
.with_execution_runtime_candidate(true);
let response = execute_execution_runtime_stream(
&state,
plan,
"trace-remote-runtime-stream-redirect",
&decision,
"gemini_chat_stream",
None,
Some(json!({
"request_id": "req-remote-runtime-stream-redirect",
"candidate_id": "cand-remote-runtime-stream-redirect",
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "openai:responses",
"client_api_format": "gemini:generate_content",
"needs_conversion": true
})),
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
assert_eq!(response.status(), axum::http::StatusCode::BAD_GATEWAY);
assert_eq!(
response
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("application/json")
);
assert_eq!(
response
.headers()
.get("x-aether-upstream-status")
.and_then(|value| value.to_str().ok()),
Some("302")
);
assert!(
response.headers().get(header::LOCATION).is_none(),
"redirect location should not be forwarded to AI clients"
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let body_json: Value =
serde_json::from_slice(&body).expect("response body should decode as json");
assert_eq!(
body_json["error"]["type"],
json!("execution_runtime_non_success_status")
);
assert_eq!(body_json["error"]["upstream_status"], json!(302));
assert_eq!(body_json["error"]["location"], json!("/"));
assert!(body_json["error"]["message"]
.as_str()
.is_some_and(|value| value.contains("non-success status 302")));
let usage = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if let Some(usage) = usage_repository
.find_by_request_id("req-remote-runtime-stream-redirect")
.await
.expect("usage should read")
.filter(|usage| usage.status == "failed")
{
break usage;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("usage should be written");
assert_eq!(usage.status_code, Some(302));
assert_eq!(usage.error_category.as_deref(), Some("redirect"));
assert!(usage
.error_message
.as_deref()
.is_some_and(|value| value.contains("non-success status 302")));
assert_eq!(
usage
.client_response_headers
.as_ref()
.and_then(|headers| headers.get("x-aether-upstream-status")),
Some(&json!("302"))
);
assert_eq!(
usage
.response_headers
.as_ref()
.and_then(|headers| headers.get("location")),
Some(&json!("/"))
);
assert!(
usage.response_body.is_none(),
"upstream redirect did not include a body"
);
assert_eq!(
usage
.client_response_body
.as_ref()
.and_then(|body| body.pointer("/error/upstream_status")),
Some(&json!(302))
);
let candidates = request_candidate_repository
.list_by_request_id("req-remote-runtime-stream-redirect")
.await
.expect("candidate trace should read");
let candidate_extra = candidates
.first()
.and_then(|candidate| candidate.extra_data.as_ref())
.expect("failed candidate extra_data should exist");
assert_eq!(
candidate_extra["upstream_response"]["status_code"],
json!(302)
);
assert_eq!(
candidate_extra["upstream_response"]["headers"]["location"],
json!("/")
);
assert!(candidate_extra["upstream_response"].get("body").is_none());
assert!(candidate_extra.get("client_response").is_none());
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_bridges_openai_image_sync_json_from_remote_runtime_to_image_sse(
) {

View File

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

View File

@@ -54,9 +54,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_extra_data,
@@ -134,6 +135,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,
@@ -383,11 +405,11 @@ impl<'a> OpenAiImageSyncProgressRecorder<'a> {
.await;
}
async fn fail(&mut self, elapsed_ms: u64) {
async fn fail(&mut self, status_code: Option<u16>, elapsed_ms: u64) {
self.snapshot.lock().await.phase = "failed";
self.persist(
RequestCandidateStatus::Failed,
Some(StatusCode::GATEWAY_TIMEOUT.as_u16()),
status_code,
Some(elapsed_ms),
true,
)
@@ -574,10 +596,12 @@ async fn execute_direct_sync_runtime_candidate(
Ok(result) => return Ok(result),
Err(err) => {
let elapsed_ms = started_at.elapsed().as_millis() as u64;
let status_code = err.status_code;
record_openai_image_sync_failed_progress(
state,
plan,
report_context,
status_code,
elapsed_ms,
progress_snapshot.clone(),
)
@@ -609,6 +633,7 @@ async fn execute_direct_sync_runtime_candidate(
state,
plan,
report_context,
Some(StatusCode::GATEWAY_TIMEOUT.as_u16()),
elapsed_ms,
progress_snapshot.clone(),
)
@@ -719,12 +744,13 @@ async fn record_openai_image_sync_failed_progress(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
status_code: Option<u16>,
elapsed_ms: u64,
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
) {
let mut progress =
OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot);
progress.fail(elapsed_ms).await;
progress.fail(status_code, elapsed_ms).await;
}
fn resolve_openai_image_sync_total_timeout_ms(plan: &ExecutionPlan) -> u64 {
@@ -752,11 +778,11 @@ fn should_track_openai_image_sync_upstream_sse(
}
fn should_enable_openai_image_sync_json_heartbeat(
plan_kind: &str,
plan: &ExecutionPlan,
report_context: Option<&Value>,
_plan_kind: &str,
_plan: &ExecutionPlan,
_report_context: Option<&Value>,
) -> bool {
should_track_openai_image_sync_upstream_sse(plan_kind, plan, report_context)
false
}
#[allow(clippy::too_many_arguments)]
@@ -1445,18 +1471,17 @@ async fn execute_execution_runtime_sync_impl(
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 {
@@ -1522,18 +1547,17 @@ async fn execute_execution_runtime_sync_impl(
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 {
@@ -1553,13 +1577,12 @@ async fn execute_execution_runtime_sync_impl(
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();
@@ -2103,7 +2126,7 @@ mod tests {
}
#[test]
fn openai_image_sync_progress_uses_report_context_upstream_stream_flag() {
fn openai_image_sync_progress_tracks_upstream_stream_without_json_heartbeat_wrapper() {
let plan = test_openai_image_plan(false);
let report_context = json!({"upstream_is_stream": true});
@@ -2112,7 +2135,7 @@ mod tests {
&plan,
Some(&report_context),
));
assert!(should_enable_openai_image_sync_json_heartbeat(
assert!(!should_enable_openai_image_sync_json_heartbeat(
OPENAI_IMAGE_SYNC_PLAN_KIND,
&plan,
Some(&report_context),

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -18,7 +18,6 @@ use crate::handlers::admin::system::shared::settings::{
};
use crate::handlers::admin::system::shared::smtp::build_admin_smtp_test_payload;
use crate::GatewayError;
use aether_data::repository::system::AdminSystemPurgeTarget;
use axum::{
body::{Body, Bytes},
http,
@@ -26,6 +25,7 @@ use axum::{
Json,
};
use serde_json::json;
use std::time::Instant;
pub(super) async fn maybe_build_local_admin_core_system_response(
state: &AdminAppState<'_>,
@@ -194,15 +194,31 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
)));
}
if let Some((target, action, object_type, object_id)) =
admin_system_purge_target_for_route_kind(decision.route_kind.as_deref())
if decision.route_kind.as_deref() == Some("cleanup_runs")
&& request_method == http::Method::GET
&& request_path == "/api/admin/system/cleanup/runs"
{
let records = crate::maintenance::list_admin_cleanup_run_records(&state.app().data)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
return Ok(Some(Json(json!({ "items": records })).into_response()));
}
if let Some((task_kind, action, object_type, object_id)) =
admin_system_purge_task_for_route_kind(decision.route_kind.as_deref())
{
if request_method != http::Method::POST {
return Ok(None);
}
let task = crate::maintenance::start_admin_system_purge_task(state.cloned_app(), task_kind)
.await?;
return Ok(Some(attach_admin_audit_response(
Json(build_admin_system_purge_payload(state, target).await?).into_response(),
"admin_system_data_purged",
Json(json!({
"message": task.message.clone(),
"task": task,
}))
.into_response(),
"admin_system_purge_task_started",
action,
object_type,
object_id,
@@ -435,100 +451,60 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
Ok(None)
}
fn admin_system_purge_target_for_route_kind(
fn admin_system_purge_task_for_route_kind(
route_kind: Option<&str>,
) -> Option<(
AdminSystemPurgeTarget,
crate::maintenance::AdminCleanupTaskKind,
&'static str,
&'static str,
&'static str,
)> {
match route_kind {
Some("purge_config") => Some((
AdminSystemPurgeTarget::Config,
"purge_system_config",
crate::maintenance::AdminCleanupTaskKind::Config,
"purge_system_config_async",
"system_config",
"global",
)),
Some("purge_users") => Some((
AdminSystemPurgeTarget::Users,
"purge_non_admin_users",
crate::maintenance::AdminCleanupTaskKind::Users,
"purge_non_admin_users_async",
"users",
"non_admin",
)),
Some("purge_usage") => Some((
AdminSystemPurgeTarget::Usage,
"purge_usage_records",
crate::maintenance::AdminCleanupTaskKind::Usage,
"purge_usage_records_async",
"usage",
"all",
)),
Some("purge_audit_logs") => Some((
AdminSystemPurgeTarget::AuditLogs,
"purge_audit_logs",
crate::maintenance::AdminCleanupTaskKind::AuditLogs,
"purge_audit_logs_async",
"audit_logs",
"all",
)),
Some("purge_request_bodies") => Some((
AdminSystemPurgeTarget::RequestBodies,
"purge_request_bodies",
Some("purge_request_bodies") | Some("purge_request_bodies_task") => Some((
crate::maintenance::AdminCleanupTaskKind::RequestBodies,
"purge_request_bodies_async",
"request_bodies",
"all",
)),
Some("purge_stats") => Some((AdminSystemPurgeTarget::Stats, "purge_stats", "stats", "all")),
Some("purge_stats") => Some((
crate::maintenance::AdminCleanupTaskKind::Stats,
"purge_stats_async",
"stats",
"all",
)),
_ => None,
}
}
async fn build_admin_system_purge_payload(
state: &AdminAppState<'_>,
target: AdminSystemPurgeTarget,
) -> Result<serde_json::Value, GatewayError> {
let summary = state.purge_admin_system_data(target).await?;
let total = summary.total();
let affected = summary.affected.clone();
if target == AdminSystemPurgeTarget::Stats {
let rebuild = state.rebuild_admin_stats_once().await?;
let message = if rebuild.capped {
format!(
"统计聚合已清空,已重建 {} 个小时桶和 {} 个日桶,仍有历史统计待后台任务继续重建",
rebuild.hourly_buckets, rebuild.daily_buckets
)
} else {
format!(
"统计聚合已清空并重建,删除 {} 行,重建 {} 个小时桶和 {} 个日桶",
total, rebuild.hourly_buckets, rebuild.daily_buckets
)
};
return Ok(json!({
"message": message,
"deleted": affected,
"rebuilt": {
"hourly_buckets": rebuild.hourly_buckets,
"daily_buckets": rebuild.daily_buckets,
"capped": rebuild.capped,
},
}));
}
let (message, count_key) = match target {
AdminSystemPurgeTarget::Config => ("系统配置已清空", "deleted"),
AdminSystemPurgeTarget::Users => ("非管理员用户已清空", "deleted"),
AdminSystemPurgeTarget::Usage => ("使用记录已清空", "deleted"),
AdminSystemPurgeTarget::AuditLogs => ("审计日志已清空", "deleted"),
AdminSystemPurgeTarget::RequestBodies => ("请求/响应体已清空", "cleaned"),
AdminSystemPurgeTarget::Stats => unreachable!("stats handled above"),
};
Ok(json!({
"message": format!("{message},影响 {} 行", total),
count_key: affected,
}))
}
async fn build_admin_system_cleanup_payload(
state: &AdminAppState<'_>,
) -> Result<serde_json::Value, GatewayError> {
let started_at_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
let started_at = Instant::now();
let summary = state.run_admin_system_cleanup_once().await?;
let cleaned = json!({
"audit_logs": summary.audit_logs_deleted,
@@ -558,6 +534,17 @@ async fn build_admin_system_cleanup_payload(
.saturating_add(summary.usage.keys_cleaned)
.saturating_add(summary.usage.records_deleted);
crate::maintenance::record_completed_cleanup_run(
&state.app().data,
"system_cleanup",
"manual",
started_at_unix_secs,
started_at,
cleaned.clone(),
format!("系统清理已执行,影响 {total}"),
)
.await;
Ok(json!({
"message": format!("系统清理已执行,影响 {} 项", total),
"cleaned": cleaned,

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -4,8 +4,9 @@ mod tests;
pub(crate) use runtime::{
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
inspect_proxy_upgrade_rollout, perform_oauth_token_refresh_once, perform_pool_quota_probe_once,
perform_provider_checkin_once, rebuild_admin_stats_once, record_proxy_upgrade_traffic_success,
inspect_proxy_upgrade_rollout, list_admin_cleanup_run_records,
perform_oauth_token_refresh_once, perform_pool_quota_probe_once, perform_provider_checkin_once,
rebuild_admin_stats_once, record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
run_admin_system_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,
@@ -15,10 +16,11 @@ pub(crate) use runtime::{
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_proxy_upgrade_rollout, AdminStatsRebuildSummary, AdminSystemCleanupSummary,
OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary, ProviderCheckinRunSummary,
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus,
ProxyUpgradeRolloutTrackedNodeState,
start_admin_request_body_cleanup_task, start_admin_system_purge_task,
start_proxy_upgrade_rollout, AdminCleanupRunRecord, AdminCleanupTaskKind,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, OAuthTokenRefreshRunSummary,
PoolQuotaProbeRunSummary, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
};

View File

@@ -8,6 +8,8 @@ use crate::{AppState, GatewayError};
#[path = "runtime/audit_cleanup.rs"]
mod audit_cleanup;
#[path = "runtime/cleanup_runs.rs"]
mod cleanup_runs;
#[path = "runtime/config.rs"]
mod config;
#[path = "runtime/db_maintenance.rs"]
@@ -49,6 +51,11 @@ pub(crate) use aether_data_contracts::repository::usage::{
UsageCleanupSummary, UsageCleanupWindow,
};
use audit_cleanup::*;
pub(crate) use cleanup_runs::{
list_admin_cleanup_run_records, record_completed_cleanup_run, record_failed_cleanup_run,
start_admin_request_body_cleanup_task, start_admin_system_purge_task, AdminCleanupRunRecord,
AdminCleanupTaskKind,
};
use config::*;
use db_maintenance::*;
pub(crate) use oauth_token_refresh::{

View File

@@ -0,0 +1,518 @@
use std::time::Instant;
use aether_data::repository::system::{AdminSystemPurgeSummary, AdminSystemPurgeTarget};
use aether_data_contracts::DataLayerError;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tracing::{info, warn};
use crate::data::GatewayDataState;
use crate::{AppState, GatewayError};
use super::{now_unix_secs, system_config_usize};
const CLEANUP_RUN_HISTORY_KEY: &str = "admin_cleanup_run_history";
const CLEANUP_RUN_HISTORY_LIMIT: usize = 50;
const REQUEST_BODY_PROGRESS_UPDATE_BATCHES: usize = 10;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct AdminCleanupRunRecord {
pub(crate) id: String,
pub(crate) kind: String,
pub(crate) trigger: String,
pub(crate) status: String,
pub(crate) message: String,
pub(crate) started_at_unix_secs: u64,
pub(crate) completed_at_unix_secs: Option<u64>,
pub(crate) duration_ms: Option<u64>,
pub(crate) summary: Value,
pub(crate) error: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum AdminCleanupTaskKind {
Config,
Users,
RequestBodies,
Usage,
AuditLogs,
Stats,
}
impl AdminCleanupTaskKind {
fn record_kind(self) -> &'static str {
match self {
Self::Config => "config_purge",
Self::Users => "users_purge",
Self::RequestBodies => "request_bodies",
Self::Usage => "usage_purge",
Self::AuditLogs => "audit_logs_purge",
Self::Stats => "stats_purge",
}
}
fn target(self) -> Option<AdminSystemPurgeTarget> {
match self {
Self::RequestBodies => None,
Self::Config => Some(AdminSystemPurgeTarget::Config),
Self::Users => Some(AdminSystemPurgeTarget::Users),
Self::Usage => Some(AdminSystemPurgeTarget::Usage),
Self::AuditLogs => Some(AdminSystemPurgeTarget::AuditLogs),
Self::Stats => Some(AdminSystemPurgeTarget::Stats),
}
}
fn start_message(self, batch_size: Option<usize>) -> String {
match self {
Self::Config => "系统配置后台清空已开始".to_string(),
Self::Users => "非管理员用户后台清空已开始".to_string(),
Self::RequestBodies => format!(
"请求/响应体后台清理已开始,每批 {}",
batch_size.unwrap_or(1)
),
Self::Usage => "使用记录后台清空已开始".to_string(),
Self::AuditLogs => "审计日志后台清空已开始".to_string(),
Self::Stats => "统计聚合后台清空和重建已开始".to_string(),
}
}
fn failure_message(self) -> &'static str {
match self {
Self::Config => "系统配置后台清空失败",
Self::Users => "非管理员用户后台清空失败",
Self::RequestBodies => "请求/响应体后台清理失败",
Self::Usage => "使用记录后台清空失败",
Self::AuditLogs => "审计日志后台清空失败",
Self::Stats => "统计聚合后台清空或重建失败",
}
}
}
pub(crate) async fn list_admin_cleanup_run_records(
data: &GatewayDataState,
) -> Result<Vec<AdminCleanupRunRecord>, DataLayerError> {
let Some(value) = data
.find_system_config_value(CLEANUP_RUN_HISTORY_KEY)
.await?
else {
return Ok(Vec::new());
};
Ok(parse_cleanup_run_records(value))
}
pub(crate) async fn start_admin_request_body_cleanup_task(
app: AppState,
) -> Result<AdminCleanupRunRecord, GatewayError> {
start_admin_system_purge_task(app, AdminCleanupTaskKind::RequestBodies).await
}
pub(crate) async fn start_admin_system_purge_task(
app: AppState,
kind: AdminCleanupTaskKind,
) -> Result<AdminCleanupRunRecord, GatewayError> {
let data = app.data.clone();
let records = list_admin_cleanup_run_records(&data)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if let Some(existing) = records
.iter()
.find(|record| record.kind == kind.record_kind() && record.status == "processing")
.cloned()
{
return Ok(existing);
}
let batch_size = if kind == AdminCleanupTaskKind::RequestBodies {
Some(
system_config_usize(&data, "cleanup_batch_size", 1_000)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
.max(1),
)
} else {
None
};
let started_at = now_unix_secs();
let record = AdminCleanupRunRecord {
id: uuid::Uuid::new_v4().to_string(),
kind: kind.record_kind().to_string(),
trigger: "manual".to_string(),
status: "processing".to_string(),
message: kind.start_message(batch_size),
started_at_unix_secs: started_at,
completed_at_unix_secs: None,
duration_ms: None,
summary: initial_task_summary(kind, batch_size),
error: None,
};
record_cleanup_run(&data, record.clone())
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
match kind {
AdminCleanupTaskKind::RequestBodies => {
tokio::spawn(run_request_body_cleanup_task(
data,
record.clone(),
batch_size.unwrap_or(1),
));
}
AdminCleanupTaskKind::Usage
| AdminCleanupTaskKind::AuditLogs
| AdminCleanupTaskKind::Config
| AdminCleanupTaskKind::Users
| AdminCleanupTaskKind::Stats => {
tokio::spawn(run_admin_system_purge_task(app, record.clone(), kind));
}
}
Ok(record)
}
pub(crate) async fn record_completed_cleanup_run(
data: &GatewayDataState,
kind: &str,
trigger: &str,
started_at_unix_secs: u64,
started_at: Instant,
summary: Value,
message: impl Into<String>,
) {
let record = AdminCleanupRunRecord {
id: uuid::Uuid::new_v4().to_string(),
kind: kind.to_string(),
trigger: trigger.to_string(),
status: "completed".to_string(),
message: message.into(),
started_at_unix_secs,
completed_at_unix_secs: Some(now_unix_secs()),
duration_ms: Some(
started_at
.elapsed()
.as_millis()
.try_into()
.unwrap_or(u64::MAX),
),
summary,
error: None,
};
if let Err(err) = record_cleanup_run(data, record).await {
warn!(error = %err, kind, "failed to record cleanup run");
}
}
pub(crate) async fn record_failed_cleanup_run(
data: &GatewayDataState,
kind: &str,
trigger: &str,
started_at_unix_secs: u64,
started_at: Instant,
error: &DataLayerError,
) {
let record = AdminCleanupRunRecord {
id: uuid::Uuid::new_v4().to_string(),
kind: kind.to_string(),
trigger: trigger.to_string(),
status: "failed".to_string(),
message: "清理执行失败".to_string(),
started_at_unix_secs,
completed_at_unix_secs: Some(now_unix_secs()),
duration_ms: Some(
started_at
.elapsed()
.as_millis()
.try_into()
.unwrap_or(u64::MAX),
),
summary: json!({}),
error: Some(error.to_string()),
};
if let Err(err) = record_cleanup_run(data, record).await {
warn!(error = %err, kind, "failed to record failed cleanup run");
}
}
async fn run_request_body_cleanup_task(
data: std::sync::Arc<GatewayDataState>,
initial_record: AdminCleanupRunRecord,
batch_size: usize,
) {
let started_at = Instant::now();
let mut total = AdminSystemPurgeSummary::default();
let mut batches = 0usize;
loop {
match data.purge_admin_request_bodies_batch(batch_size).await {
Ok(batch) if batch.total() == 0 => break,
Ok(batch) => {
batches = batches.saturating_add(1);
total.merge(&batch);
if batches.is_multiple_of(REQUEST_BODY_PROGRESS_UPDATE_BATCHES) {
let progress = request_body_cleanup_record(
&initial_record,
"processing",
format!(
"请求/响应体后台清理中,已处理 {} 批,影响 {}",
batches,
total.total()
),
batches,
batch_size,
&total,
Some(started_at),
None,
);
if let Err(err) = record_cleanup_run(&data, progress).await {
warn!(error = %err, "failed to update request body cleanup progress");
}
}
tokio::task::yield_now().await;
}
Err(err) => {
let failed = request_body_cleanup_record(
&initial_record,
"failed",
"请求/响应体后台清理失败".to_string(),
batches,
batch_size,
&total,
Some(started_at),
Some(err.to_string()),
);
if let Err(record_err) = record_cleanup_run(&data, failed).await {
warn!(error = %record_err, "failed to record request body cleanup failure");
}
warn!(error = %err, "request body cleanup task failed");
return;
}
}
}
let completed = request_body_cleanup_record(
&initial_record,
"completed",
format!(
"请求/响应体后台清理完成,影响 {} 行,共 {}",
total.total(),
batches
),
batches,
batch_size,
&total,
Some(started_at),
None,
);
if let Err(err) = record_cleanup_run(&data, completed).await {
warn!(error = %err, "failed to record request body cleanup completion");
}
info!(
event_name = "request_body_cleanup_task_completed",
log_type = "ops",
worker = "request_body_cleanup_task",
batches,
affected = total.total(),
"gateway finished request body cleanup task"
);
}
async fn run_admin_system_purge_task(
app: AppState,
initial_record: AdminCleanupRunRecord,
kind: AdminCleanupTaskKind,
) {
let started_at = Instant::now();
let data = app.data.clone();
match run_admin_system_purge_task_once(&app, kind).await {
Ok((summary, message)) => {
let completed = cleanup_task_record(
&initial_record,
"completed",
message,
summary,
Some(started_at),
None,
);
if let Err(err) = record_cleanup_run(&data, completed).await {
warn!(error = %err, "failed to record admin system purge task completion");
}
info!(
event_name = "admin_system_purge_task_completed",
log_type = "ops",
worker = "admin_system_purge_task",
kind = ?kind,
"gateway finished admin system purge task"
);
}
Err(err) => {
let failed = cleanup_task_record(
&initial_record,
"failed",
kind.failure_message().to_string(),
json!({}),
Some(started_at),
Some(format!("{err:?}")),
);
if let Err(record_err) = record_cleanup_run(&data, failed).await {
warn!(error = %record_err, "failed to record admin system purge task failure");
}
warn!(error = ?err, kind = ?kind, "admin system purge task failed");
}
}
}
async fn run_admin_system_purge_task_once(
app: &AppState,
kind: AdminCleanupTaskKind,
) -> Result<(Value, String), GatewayError> {
let target = kind
.target()
.expect("non-request-body cleanup task should map to purge target");
let purge = app.purge_admin_system_data(target).await?;
let deleted_total = purge.total();
let affected = purge.affected.clone();
if kind == AdminCleanupTaskKind::Stats {
let rebuild = app.rebuild_admin_stats_once().await?;
let message = if rebuild.capped {
format!(
"统计聚合后台清空完成,已重建 {} 个小时桶和 {} 个日桶,仍有历史统计待后台任务继续重建",
rebuild.hourly_buckets, rebuild.daily_buckets
)
} else {
format!(
"统计聚合后台清空并重建完成,删除 {} 行,重建 {} 个小时桶和 {} 个日桶",
deleted_total, rebuild.hourly_buckets, rebuild.daily_buckets
)
};
return Ok((
json!({
"deleted": affected,
"rebuilt": {
"hourly_buckets": rebuild.hourly_buckets,
"daily_buckets": rebuild.daily_buckets,
"capped": rebuild.capped,
},
"total": deleted_total,
}),
message,
));
}
let message = match kind {
AdminCleanupTaskKind::Config => format!("系统配置后台清空完成,影响 {deleted_total}"),
AdminCleanupTaskKind::Users => format!("非管理员用户后台清空完成,影响 {deleted_total}"),
AdminCleanupTaskKind::Usage => format!("使用记录后台清空完成,影响 {deleted_total}"),
AdminCleanupTaskKind::AuditLogs => format!("审计日志后台清空完成,影响 {deleted_total}"),
AdminCleanupTaskKind::RequestBodies | AdminCleanupTaskKind::Stats => unreachable!(),
};
Ok((
json!({
"deleted": affected,
"total": deleted_total,
}),
message,
))
}
fn initial_task_summary(kind: AdminCleanupTaskKind, batch_size: Option<usize>) -> Value {
match kind {
AdminCleanupTaskKind::RequestBodies => json!({
"batch_size": batch_size.unwrap_or(1),
"batches": 0,
"cleaned": {},
}),
AdminCleanupTaskKind::Stats => json!({
"deleted": {},
"rebuilt": {},
}),
AdminCleanupTaskKind::Config
| AdminCleanupTaskKind::Users
| AdminCleanupTaskKind::Usage
| AdminCleanupTaskKind::AuditLogs => json!({
"deleted": {},
}),
}
}
fn cleanup_task_record(
initial: &AdminCleanupRunRecord,
status: &str,
message: String,
summary: Value,
started_at: Option<Instant>,
error: Option<String>,
) -> AdminCleanupRunRecord {
let completed = matches!(status, "completed" | "failed");
AdminCleanupRunRecord {
id: initial.id.clone(),
kind: initial.kind.clone(),
trigger: initial.trigger.clone(),
status: status.to_string(),
message,
started_at_unix_secs: initial.started_at_unix_secs,
completed_at_unix_secs: completed.then(now_unix_secs),
duration_ms: started_at
.map(|value| value.elapsed().as_millis().try_into().unwrap_or(u64::MAX)),
summary,
error,
}
}
fn request_body_cleanup_record(
initial: &AdminCleanupRunRecord,
status: &str,
message: String,
batches: usize,
batch_size: usize,
total: &AdminSystemPurgeSummary,
started_at: Option<Instant>,
error: Option<String>,
) -> AdminCleanupRunRecord {
let completed = matches!(status, "completed" | "failed");
AdminCleanupRunRecord {
id: initial.id.clone(),
kind: initial.kind.clone(),
trigger: initial.trigger.clone(),
status: status.to_string(),
message,
started_at_unix_secs: initial.started_at_unix_secs,
completed_at_unix_secs: completed.then(now_unix_secs),
duration_ms: started_at
.map(|value| value.elapsed().as_millis().try_into().unwrap_or(u64::MAX)),
summary: json!({
"batch_size": batch_size,
"batches": batches,
"cleaned": total.affected,
"total": total.total(),
}),
error,
}
}
async fn record_cleanup_run(
data: &GatewayDataState,
record: AdminCleanupRunRecord,
) -> Result<(), DataLayerError> {
let mut records = list_admin_cleanup_run_records(data).await?;
records.retain(|existing| existing.id != record.id);
records.insert(0, record);
records.truncate(CLEANUP_RUN_HISTORY_LIMIT);
let value = serde_json::to_value(records).map_err(|err| {
DataLayerError::UnexpectedValue(format!("invalid cleanup run history: {err}"))
})?;
data.upsert_system_config_entry(
CLEANUP_RUN_HISTORY_KEY,
&value,
Some("最近的系统清理执行记录"),
)
.await?;
Ok(())
}
fn parse_cleanup_run_records(value: Value) -> Vec<AdminCleanupRunRecord> {
value
.as_array()
.into_iter()
.flat_map(|items| items.iter())
.filter_map(|item| serde_json::from_value::<AdminCleanupRunRecord>(item.clone()).ok())
.collect()
}

View File

@@ -1,9 +1,14 @@
use aether_data_contracts::DataLayerError;
use serde_json::json;
use std::time::Instant;
use tracing::info;
use crate::data::GatewayDataState;
use super::{now_unix_secs, system_config_bool, system_config_u64, system_config_usize};
use super::{
now_unix_secs, record_completed_cleanup_run, record_failed_cleanup_run, system_config_bool,
system_config_u64, system_config_usize,
};
pub(crate) async fn cleanup_request_candidates_once(
data: &GatewayDataState,
@@ -47,7 +52,33 @@ pub(crate) async fn cleanup_request_candidates_once(
pub(super) async fn run_request_candidate_cleanup_once(
data: &GatewayDataState,
) -> Result<(), DataLayerError> {
let deleted = cleanup_request_candidates_once(data).await?;
let started_at_unix_secs = now_unix_secs();
let started_at = Instant::now();
let deleted = match cleanup_request_candidates_once(data).await {
Ok(deleted) => deleted,
Err(err) => {
record_failed_cleanup_run(
data,
"request_candidate_cleanup",
"auto",
started_at_unix_secs,
started_at,
&err,
)
.await;
return Err(err);
}
};
record_completed_cleanup_run(
data,
"request_candidate_cleanup",
"auto",
started_at_unix_secs,
started_at,
json!({ "request_candidates_deleted": deleted }),
format!("候选记录自动清理完成,删除 {deleted}"),
)
.await;
if deleted > 0 {
info!(deleted, "gateway deleted expired request candidates");
}

View File

@@ -1,4 +1,6 @@
use aether_data_contracts::DataLayerError;
use serde_json::json;
use std::time::Instant;
use tracing::{info, warn};
use crate::data::GatewayDataState;
@@ -8,15 +10,41 @@ use super::{
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_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,
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,
perform_stats_hourly_aggregation_once, perform_usage_cleanup_once,
perform_wallet_daily_usage_aggregation_once, record_proxy_upgrade_traffic_success,
summarize_database_pool,
perform_wallet_daily_usage_aggregation_once, record_completed_cleanup_run,
record_failed_cleanup_run, record_proxy_upgrade_traffic_success, summarize_database_pool,
};
pub(super) async fn run_audit_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
let deleted = cleanup_audit_logs_once(data).await?;
let started_at_unix_secs = now_unix_secs();
let started_at = Instant::now();
let deleted = match cleanup_audit_logs_once(data).await {
Ok(deleted) => deleted,
Err(err) => {
record_failed_cleanup_run(
data,
"audit_cleanup",
"auto",
started_at_unix_secs,
started_at,
&err,
)
.await;
return Err(err);
}
};
record_completed_cleanup_run(
data,
"audit_cleanup",
"auto",
started_at_unix_secs,
started_at,
json!({ "audit_logs_deleted": deleted }),
format!("审计日志自动清理完成,删除 {deleted}"),
)
.await;
if deleted > 0 {
info!(
event_name = "audit_cleanup_completed",
@@ -204,7 +232,49 @@ pub(super) async fn run_stats_aggregation_once(
}
pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
let summary = perform_usage_cleanup_once(data).await?;
let started_at_unix_secs = now_unix_secs();
let started_at = Instant::now();
let summary = match perform_usage_cleanup_once(data).await {
Ok(summary) => summary,
Err(err) => {
record_failed_cleanup_run(
data,
"usage_cleanup",
"auto",
started_at_unix_secs,
started_at,
&err,
)
.await;
return Err(err);
}
};
record_completed_cleanup_run(
data,
"usage_cleanup",
"auto",
started_at_unix_secs,
started_at,
json!({
"body_externalized": summary.body_externalized,
"legacy_body_refs_migrated": summary.legacy_body_refs_migrated,
"body_cleaned": summary.body_cleaned,
"header_cleaned": summary.header_cleaned,
"keys_cleaned": summary.keys_cleaned,
"records_deleted": summary.records_deleted,
}),
format!(
"请求记录自动清理完成,影响 {}",
summary
.body_externalized
.saturating_add(summary.legacy_body_refs_migrated)
.saturating_add(summary.body_cleaned)
.saturating_add(summary.header_cleaned)
.saturating_add(summary.keys_cleaned)
.saturating_add(summary.records_deleted)
),
)
.await;
if summary.body_externalized > 0
|| summary.legacy_body_refs_migrated > 0
|| summary.body_cleaned > 0

View File

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

View File

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

View File

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

View File

@@ -19,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};
@@ -997,71 +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_node_metrics_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_upgrade_rollout_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_provider_checkin_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_oauth_token_refresh_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_request_candidate_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_gemini_file_mapping_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_model_fetch_worker(self.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_video_task_poller(self.clone()) {
tasks.push(handle);
}
tasks
pub fn spawn_background_tasks(&self) -> crate::task_runtime::TaskSupervisor {
let mut supervisor = crate::task_runtime::TaskSupervisor::new();
let record_boot = |task_key: &'static str| {
if !self.has_background_task_data_writer() {
return;
}
let Some(definition) = crate::task_runtime::task_definition(task_key) else {
return;
};
std::mem::drop(crate::task_runtime::spawn_record_worker_boot(
self.clone(),
task_key,
crate::task_runtime::background_task_kind(definition.kind),
definition.trigger,
));
};
let mut supervise_worker =
|task_key: &'static str, handle: Option<tokio::task::JoinHandle<()>>| {
if let Some(handle) = handle {
supervisor.supervise_handle(task_key, handle);
record_boot(task_key);
}
};
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER,
self.usage_runtime.spawn_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROVIDER_QUOTA_RESET,
crate::wallet_runtime::spawn_provider_quota_reset_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_AUDIT_CLEANUP,
spawn_audit_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_DB_MAINTENANCE,
spawn_db_maintenance_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_WALLET_DAILY_USAGE_AGG,
spawn_wallet_daily_usage_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_STATS_DAILY_AGG,
spawn_stats_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_CLEANUP,
spawn_usage_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_MONITOR,
spawn_pool_monitor_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
spawn_pool_quota_probe_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_STATS_HOURLY_AGG,
spawn_stats_hourly_aggregation_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PENDING_CLEANUP,
spawn_pending_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_NODE_STALE_CLEANUP,
spawn_proxy_node_stale_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_NODE_METRICS_CLEANUP,
spawn_proxy_node_metrics_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROXY_UPGRADE_ROLLOUT,
spawn_proxy_upgrade_rollout_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROVIDER_CHECKIN,
spawn_provider_checkin_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_OAUTH_TOKEN_REFRESH,
spawn_oauth_token_refresh_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_REQUEST_CANDIDATE_CLEANUP,
spawn_request_candidate_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_GEMINI_FILES_CLEANUP,
spawn_gemini_file_mapping_cleanup_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_MODEL_FETCH_WORKER,
spawn_model_fetch_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_VIDEO_TASK_POLLER,
spawn_video_task_poller(self.clone()),
);
supervisor
}
}

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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