mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
Fix Codex image progress heartbeat merge regressions
This commit is contained in:
@@ -23,6 +23,7 @@ aether-provider-transport.workspace = true
|
||||
aether-scheduler-core.workspace = true
|
||||
aether-runtime.workspace = true
|
||||
aether-runtime-state.workspace = true
|
||||
aether-task-runtime.workspace = true
|
||||
aether-usage-runtime.workspace = true
|
||||
aether-video-tasks-core.workspace = true
|
||||
aether-wallet.workspace = true
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
use axum::body::Bytes;
|
||||
|
||||
use crate::ai_serving::{
|
||||
endpoint_config_forces_upstream_stream_policy as endpoint_config_forces_upstream_stream_policy_impl,
|
||||
enforce_request_body_stream_field as enforce_request_body_stream_field_impl,
|
||||
force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl,
|
||||
is_json_request, parse_direct_request_body as parse_direct_request_body_impl,
|
||||
resolve_upstream_is_stream_from_endpoint_config as resolve_upstream_is_stream_from_endpoint_config_impl,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
|
||||
@@ -37,6 +40,52 @@ pub(crate) fn force_upstream_streaming_for_provider(
|
||||
force_upstream_streaming_for_provider_impl(provider_type, provider_api_format)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_upstream_is_stream_for_provider(
|
||||
endpoint_config: Option<&serde_json::Value>,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
client_is_stream: bool,
|
||||
hard_requires_streaming: bool,
|
||||
) -> bool {
|
||||
let hard_requires_streaming = hard_requires_streaming
|
||||
|| force_upstream_streaming_for_provider(provider_type, provider_api_format);
|
||||
resolve_upstream_is_stream_from_endpoint_config_impl(
|
||||
endpoint_config,
|
||||
client_is_stream,
|
||||
hard_requires_streaming,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn endpoint_config_forces_body_stream_field(
|
||||
endpoint_config: Option<&serde_json::Value>,
|
||||
) -> bool {
|
||||
endpoint_config_forces_upstream_stream_policy_impl(endpoint_config)
|
||||
}
|
||||
|
||||
pub(crate) fn request_requires_body_stream_field(
|
||||
body_json: &serde_json::Value,
|
||||
force_body_stream_field: bool,
|
||||
) -> bool {
|
||||
force_body_stream_field
|
||||
|| body_json
|
||||
.as_object()
|
||||
.is_some_and(|object| object.contains_key("stream"))
|
||||
}
|
||||
|
||||
pub(crate) fn enforce_provider_body_stream_policy(
|
||||
provider_request_body: &mut serde_json::Value,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
require_body_stream_field: bool,
|
||||
) {
|
||||
enforce_request_body_stream_field_impl(
|
||||
provider_request_body,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
require_body_stream_field,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn extract_standard_requested_model(body_json: &serde_json::Value) -> Option<String> {
|
||||
aether_ai_serving::extract_ai_standard_requested_model(body_json)
|
||||
}
|
||||
@@ -56,8 +105,10 @@ pub(crate) fn extract_requested_model_from_request(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
|
||||
extract_requested_model_from_request, extract_standard_requested_model,
|
||||
force_upstream_streaming_for_provider, RequestedModelFamily,
|
||||
force_upstream_streaming_for_provider, resolve_upstream_is_stream_for_provider,
|
||||
RequestedModelFamily,
|
||||
};
|
||||
use axum::http::Request;
|
||||
use serde_json::json;
|
||||
@@ -90,6 +141,67 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_endpoint_upstream_stream_policy_with_provider_hard_constraints() {
|
||||
assert!(resolve_upstream_is_stream_for_provider(
|
||||
Some(&json!({"upstream_stream_policy": "force_stream"})),
|
||||
"openai",
|
||||
"openai:chat",
|
||||
false,
|
||||
false,
|
||||
));
|
||||
assert!(!resolve_upstream_is_stream_for_provider(
|
||||
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
|
||||
"openai",
|
||||
"openai:chat",
|
||||
true,
|
||||
false,
|
||||
));
|
||||
assert!(resolve_upstream_is_stream_for_provider(
|
||||
Some(&json!({"upstream_stream_policy": "auto"})),
|
||||
"openai",
|
||||
"openai:chat",
|
||||
true,
|
||||
false,
|
||||
));
|
||||
assert!(resolve_upstream_is_stream_for_provider(
|
||||
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
|
||||
"codex",
|
||||
"openai:responses",
|
||||
true,
|
||||
false,
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enforces_provider_body_stream_policy_for_body_and_streamless_formats() {
|
||||
let mut openai_chat = json!({"stream": true});
|
||||
enforce_provider_body_stream_policy(&mut openai_chat, "openai:chat", false, false);
|
||||
assert_eq!(openai_chat.get("stream"), Some(&json!(false)));
|
||||
|
||||
let mut ordinary_sync = json!({"messages": []});
|
||||
enforce_provider_body_stream_policy(&mut ordinary_sync, "openai:chat", false, false);
|
||||
assert!(ordinary_sync.get("stream").is_none());
|
||||
|
||||
let mut compact = json!({"stream": true});
|
||||
enforce_provider_body_stream_policy(&mut compact, "openai:responses:compact", true, true);
|
||||
assert!(compact.get("stream").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_endpoint_configs_that_force_body_stream_field() {
|
||||
assert!(endpoint_config_forces_body_stream_field(Some(
|
||||
&json!({"upstream_stream_policy": "force_stream"})
|
||||
)));
|
||||
assert!(endpoint_config_forces_body_stream_field(Some(
|
||||
&json!({"upstream_stream_policy": "force_non_stream"})
|
||||
)));
|
||||
assert!(!endpoint_config_forces_body_stream_field(Some(
|
||||
&json!({"upstream_stream_policy": "auto"})
|
||||
)));
|
||||
assert!(!endpoint_config_forces_body_stream_field(None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_standard_requested_model_from_request_body() {
|
||||
let requested_model =
|
||||
|
||||
@@ -112,6 +112,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
|
||||
@@ -3,6 +3,9 @@ use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
enforce_provider_body_stream_policy, request_requires_body_stream_field,
|
||||
};
|
||||
use crate::ai_serving::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
|
||||
@@ -50,6 +53,7 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac
|
||||
};
|
||||
let behavior = policy::classify_same_format_provider_request_behavior(
|
||||
transport,
|
||||
provider_api_format,
|
||||
crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata {
|
||||
api_format: provider_api_format,
|
||||
require_streaming: false,
|
||||
@@ -131,6 +135,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
prepared.transport.endpoint.body_rules.as_ref(),
|
||||
Some(&parts.headers),
|
||||
prepared.upstream_is_stream,
|
||||
prepared.force_body_stream_field,
|
||||
prepared.kiro_auth.as_ref(),
|
||||
prepared.is_claude_code,
|
||||
enable_model_directives,
|
||||
@@ -170,6 +175,18 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
&mut base_provider_request_body,
|
||||
&mapping,
|
||||
);
|
||||
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
|
||||
// re-enforce stream-field policy afterward.
|
||||
// Kiro behavior classification already hard-requires upstream streaming,
|
||||
// and the Kiro envelope does not use a top-level body stream field.
|
||||
if prepared.kiro_auth.is_none() {
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut base_provider_request_body,
|
||||
prepared.provider_api_format.as_str(),
|
||||
prepared.upstream_is_stream,
|
||||
request_requires_body_stream_field(body_json, prepared.force_body_stream_field),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let antigravity_auth = if prepared.is_antigravity {
|
||||
|
||||
@@ -13,12 +13,14 @@ use super::super::LocalSameFormatProviderFamily;
|
||||
|
||||
pub(super) fn classify_same_format_provider_request_behavior(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
spec_metadata: LocalExecutionSurfaceSpecMetadata,
|
||||
) -> SameFormatProviderRequestBehavior {
|
||||
classify_same_format_provider_request_behavior_impl(
|
||||
transport,
|
||||
SameFormatProviderRequestBehaviorParams {
|
||||
require_streaming: spec_metadata.require_streaming,
|
||||
provider_api_format,
|
||||
report_kind: spec_metadata
|
||||
.report_kind
|
||||
.expect("same-format provider specs should declare report kind"),
|
||||
|
||||
@@ -35,6 +35,7 @@ pub(super) struct PreparedSameFormatProviderCandidate {
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) report_kind: &'static str,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
pub(super) force_body_stream_field: bool,
|
||||
}
|
||||
|
||||
pub(super) async fn prepare_local_same_format_provider_candidate(
|
||||
@@ -51,7 +52,11 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
|
||||
let candidate = &eligible.candidate;
|
||||
let transport = Arc::clone(&eligible.transport);
|
||||
let provider_api_format = eligible.provider_api_format.as_str();
|
||||
let behavior = classify_same_format_provider_request_behavior(&transport, spec_metadata);
|
||||
let behavior = classify_same_format_provider_request_behavior(
|
||||
&transport,
|
||||
provider_api_format,
|
||||
spec_metadata,
|
||||
);
|
||||
|
||||
if !same_format_provider_transport_supported(
|
||||
&behavior,
|
||||
@@ -174,5 +179,6 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
|
||||
mapped_model,
|
||||
report_kind: behavior.report_kind,
|
||||
upstream_is_stream: behavior.upstream_is_stream,
|
||||
force_body_stream_field: behavior.force_body_stream_field,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ pub(crate) fn build_same_format_provider_request_body(
|
||||
body_rules: Option<&Value>,
|
||||
request_headers: Option<&http::HeaderMap>,
|
||||
upstream_is_stream: bool,
|
||||
force_body_stream_field: bool,
|
||||
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
|
||||
is_claude_code: bool,
|
||||
enable_model_directives: bool,
|
||||
@@ -28,6 +29,7 @@ pub(crate) fn build_same_format_provider_request_body(
|
||||
body_rules,
|
||||
request_headers,
|
||||
upstream_is_stream,
|
||||
force_body_stream_field,
|
||||
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
|
||||
is_claude_code,
|
||||
enable_model_directives,
|
||||
|
||||
@@ -9,7 +9,11 @@ use aether_ai_serving::{
|
||||
use aether_scheduler_core::{ClientSessionAffinity, SchedulerRankingOutcome};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::ai_serving::{request_origin_from_headers, ExecutionRuntimeAuthContext, RequestOrigin};
|
||||
use crate::ai_serving::{
|
||||
request_origin_from_headers, request_path_implies_stream_request, sanitize_request_path,
|
||||
sanitize_request_path_and_query, sanitize_request_query_string, ExecutionRuntimeAuthContext,
|
||||
RequestOrigin,
|
||||
};
|
||||
use crate::client_session_affinity::{
|
||||
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
|
||||
};
|
||||
@@ -40,6 +44,8 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||
pub(crate) provider_request_method: Option<Value>,
|
||||
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
|
||||
pub(crate) original_headers: &'a http::HeaderMap,
|
||||
pub(crate) request_path: Option<&'a str>,
|
||||
pub(crate) request_query_string: Option<&'a str>,
|
||||
pub(crate) request_origin: Option<RequestOrigin>,
|
||||
pub(crate) original_request_body_json: Option<&'a Value>,
|
||||
pub(crate) original_request_body_base64: Option<&'a str>,
|
||||
@@ -80,6 +86,15 @@ pub(crate) fn build_local_execution_report_context(
|
||||
{
|
||||
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
|
||||
}
|
||||
insert_request_path_fields(
|
||||
&mut extra_fields,
|
||||
parts.request_path,
|
||||
parts.request_query_string,
|
||||
);
|
||||
let client_requested_stream = parts.client_requested_stream
|
||||
|| parts
|
||||
.request_path
|
||||
.is_some_and(request_path_implies_stream_request);
|
||||
|
||||
build_ai_execution_report_context(AiExecutionReportContextParts {
|
||||
auth_context: parts.auth_context,
|
||||
@@ -113,7 +128,7 @@ pub(crate) fn build_local_execution_report_context(
|
||||
client_ip,
|
||||
user_agent,
|
||||
},
|
||||
client_requested_stream: parts.client_requested_stream,
|
||||
client_requested_stream,
|
||||
upstream_is_stream: parts.upstream_is_stream,
|
||||
has_envelope: parts.has_envelope,
|
||||
needs_conversion: parts.needs_conversion,
|
||||
@@ -121,6 +136,30 @@ pub(crate) fn build_local_execution_report_context(
|
||||
})
|
||||
}
|
||||
|
||||
fn insert_request_path_fields(
|
||||
extra_fields: &mut Map<String, Value>,
|
||||
request_path: Option<&str>,
|
||||
request_query_string: Option<&str>,
|
||||
) {
|
||||
let Some(path) = request_path.and_then(sanitize_request_path) else {
|
||||
return;
|
||||
};
|
||||
let query = request_query_string.and_then(sanitize_request_query_string);
|
||||
let path_and_query = sanitize_request_path_and_query(path.as_str(), query.as_deref())
|
||||
.unwrap_or_else(|| path.clone());
|
||||
extra_fields
|
||||
.entry("request_path".to_string())
|
||||
.or_insert_with(|| Value::String(path.clone()));
|
||||
if let Some(query) = query.clone() {
|
||||
extra_fields
|
||||
.entry("request_query_string".to_string())
|
||||
.or_insert_with(|| Value::String(query.to_string()));
|
||||
}
|
||||
extra_fields
|
||||
.entry("request_path_and_query".to_string())
|
||||
.or_insert_with(|| Value::String(path_and_query));
|
||||
}
|
||||
|
||||
pub(crate) fn provider_stream_event_api_format_for_provider_type(
|
||||
provider_type: &str,
|
||||
) -> Option<&'static str> {
|
||||
@@ -226,6 +265,8 @@ mod tests {
|
||||
provider_request_method: None,
|
||||
provider_request_headers: Some(&provider_request_headers),
|
||||
original_headers: &original_headers,
|
||||
request_path: Some("/v1/chat/completions"),
|
||||
request_query_string: Some("debug=true&limit=10"),
|
||||
request_origin: Some(RequestOrigin {
|
||||
client_ip: Some("203.0.113.8".to_string()),
|
||||
user_agent: Some("Claude-Code/1.0".to_string()),
|
||||
@@ -255,6 +296,77 @@ mod tests {
|
||||
"session_key": "account=account-1;session=session-1"
|
||||
})
|
||||
);
|
||||
assert_eq!(report_context["request_path"], "/v1/chat/completions");
|
||||
assert_eq!(report_context["request_query_string"], "limit=10");
|
||||
assert_eq!(
|
||||
report_context["request_path_and_query"],
|
||||
"/v1/chat/completions?limit=10"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_execution_report_context_treats_stream_generate_content_path_as_client_stream() {
|
||||
let auth_context = ExecutionRuntimeAuthContext {
|
||||
user_id: "user-1".to_string(),
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
balance_remaining: None,
|
||||
access_allowed: true,
|
||||
api_key_is_standalone: false,
|
||||
};
|
||||
let original_headers = http::HeaderMap::new();
|
||||
let provider_request_headers = BTreeMap::new();
|
||||
|
||||
let report_context =
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &auth_context,
|
||||
request_id: "trace-1",
|
||||
candidate_id: "candidate-1",
|
||||
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
|
||||
model: "gemini-3.1-flash-image-preview",
|
||||
provider_name: "Gemini",
|
||||
provider_id: "provider-1",
|
||||
endpoint_id: "endpoint-1",
|
||||
key_id: "key-1",
|
||||
key_name: None,
|
||||
model_id: None,
|
||||
global_model_id: None,
|
||||
global_model_name: None,
|
||||
provider_api_format: "gemini:generate_content",
|
||||
client_api_format: "gemini:generate_content",
|
||||
mapped_model: None,
|
||||
candidate_group_id: None,
|
||||
ranking: None,
|
||||
upstream_url: None,
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
provider_request_method: None,
|
||||
provider_request_headers: Some(&provider_request_headers),
|
||||
original_headers: &original_headers,
|
||||
request_path: Some(
|
||||
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent",
|
||||
),
|
||||
request_query_string: Some("key=secret&alt=sse"),
|
||||
request_origin: None,
|
||||
original_request_body_json: Some(&json!({
|
||||
"contents": [{"role": "user", "parts": [{"text": "hi"}]}]
|
||||
})),
|
||||
original_request_body_base64: None,
|
||||
client_session_affinity: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: true,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields: Map::new(),
|
||||
});
|
||||
|
||||
assert_eq!(report_context["client_requested_stream"], true);
|
||||
assert_eq!(report_context["request_query_string"], "alt=sse");
|
||||
assert_eq!(
|
||||
report_context["request_path_and_query"],
|
||||
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent?alt=sse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -300,6 +412,8 @@ mod tests {
|
||||
provider_request_method: None,
|
||||
provider_request_headers: Some(&provider_request_headers),
|
||||
original_headers: &original_headers,
|
||||
request_path: None,
|
||||
request_query_string: None,
|
||||
request_origin: None,
|
||||
original_request_body_json: Some(&json!({"model": "gpt-5"})),
|
||||
original_request_body_base64: None,
|
||||
|
||||
@@ -94,6 +94,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
provider_request_method: None,
|
||||
provider_request_headers: None,
|
||||
original_headers: &parts.headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -120,6 +120,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
|
||||
@@ -7,7 +7,10 @@ use crate::ai_serving::planner::candidate_preparation::{
|
||||
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
|
||||
OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider;
|
||||
use crate::ai_serving::planner::common::{
|
||||
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
|
||||
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
|
||||
use crate::ai_serving::planner::standard::{
|
||||
apply_codex_openai_responses_special_headers, request_body_build_failure_extra_data,
|
||||
@@ -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() {
|
||||
|
||||
@@ -116,9 +116,9 @@ pub(crate) fn build_gemini_stream_plan_from_decision(
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: true,
|
||||
upstream_is_stream: payload.upstream_is_stream,
|
||||
build_from_request_when_empty: false,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
|
||||
});
|
||||
let content_type = payload
|
||||
.content_type
|
||||
|
||||
@@ -15,3 +15,6 @@ pub(crate) use self::responses::{
|
||||
build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body,
|
||||
build_local_openai_responses_upstream_url,
|
||||
};
|
||||
pub(super) use crate::ai_serving::planner::common::{
|
||||
enforce_provider_body_stream_policy, request_requires_body_stream_field,
|
||||
};
|
||||
|
||||
@@ -9,10 +9,13 @@ use crate::ai_serving::{
|
||||
GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
use super::{enforce_provider_body_stream_policy, request_requires_body_stream_field};
|
||||
|
||||
pub(crate) fn build_local_openai_chat_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
upstream_is_stream: bool,
|
||||
force_body_stream_field: bool,
|
||||
body_rules: Option<&Value>,
|
||||
request_headers: &http::HeaderMap,
|
||||
enable_model_directives: bool,
|
||||
@@ -23,12 +26,20 @@ pub(crate) fn build_local_openai_chat_request_body(
|
||||
upstream_is_stream,
|
||||
enable_model_directives,
|
||||
)?;
|
||||
apply_standard_provider_request_body_rules_with_request_headers(
|
||||
provider_request_body,
|
||||
body_rules,
|
||||
body_json,
|
||||
request_headers,
|
||||
)
|
||||
let mut provider_request_body =
|
||||
apply_standard_provider_request_body_rules_with_request_headers(
|
||||
provider_request_body,
|
||||
body_rules,
|
||||
body_json,
|
||||
request_headers,
|
||||
)?;
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut provider_request_body,
|
||||
"openai:chat",
|
||||
upstream_is_stream,
|
||||
request_requires_body_stream_field(body_json, force_body_stream_field),
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_openai_chat_upstream_url(
|
||||
@@ -44,6 +55,7 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
force_body_stream_field: bool,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
request_headers: &http::HeaderMap,
|
||||
@@ -74,6 +86,12 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
request_requires_body_stream_field(body_json, force_body_stream_field),
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
|
||||
@@ -9,10 +9,13 @@ use crate::ai_serving::{
|
||||
GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
use super::{enforce_provider_body_stream_policy, request_requires_body_stream_field};
|
||||
|
||||
pub(crate) fn build_local_openai_responses_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
require_streaming: bool,
|
||||
force_body_stream_field: bool,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
body_rules: Option<&Value>,
|
||||
@@ -44,6 +47,12 @@ pub(crate) fn build_local_openai_responses_request_body(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
require_streaming,
|
||||
request_requires_body_stream_field(body_json, force_body_stream_field),
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
@@ -53,6 +62,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
force_body_stream_field: bool,
|
||||
provider_type: &str,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
@@ -85,6 +95,12 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
request_requires_body_stream_field(body_json, force_body_stream_field),
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
|
||||
@@ -87,6 +87,7 @@ fn builds_openai_chat_cross_format_request_body_from_openai_responses_source() {
|
||||
"openai:responses",
|
||||
"openai:chat",
|
||||
false,
|
||||
false,
|
||||
"openai",
|
||||
None,
|
||||
None,
|
||||
@@ -121,6 +122,7 @@ fn local_openai_responses_wrapper_preserves_body_order_after_edits() {
|
||||
&body_json,
|
||||
"gpt-5.4",
|
||||
true,
|
||||
false,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
None,
|
||||
@@ -160,6 +162,7 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests(
|
||||
&body_json,
|
||||
"gpt-5.4",
|
||||
false,
|
||||
false,
|
||||
"openai",
|
||||
"openai:responses:compact",
|
||||
None,
|
||||
@@ -170,6 +173,7 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests(
|
||||
.expect("local openai compact body should build");
|
||||
|
||||
assert!(provider_request_body.get("store").is_none());
|
||||
assert!(provider_request_body.get("stream").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -187,6 +191,7 @@ fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() {
|
||||
&body_json,
|
||||
"gpt-5.4",
|
||||
false,
|
||||
false,
|
||||
"openai",
|
||||
"openai:responses",
|
||||
Some(&body_rules),
|
||||
@@ -237,6 +242,7 @@ fn strips_metadata_for_codex_openai_responses_requests() {
|
||||
"claude:messages",
|
||||
"openai:responses",
|
||||
true,
|
||||
false,
|
||||
"codex",
|
||||
None,
|
||||
None,
|
||||
@@ -271,6 +277,7 @@ fn applies_codex_defaults_unless_body_rules_handle_the_field() {
|
||||
"claude:messages",
|
||||
"openai:responses",
|
||||
true,
|
||||
false,
|
||||
"codex",
|
||||
Some(&body_rules),
|
||||
None,
|
||||
@@ -300,6 +307,7 @@ fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() {
|
||||
"claude:messages",
|
||||
"openai:responses",
|
||||
true,
|
||||
false,
|
||||
"codex",
|
||||
None,
|
||||
Some("key-123"),
|
||||
@@ -330,6 +338,7 @@ fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
|
||||
"codex",
|
||||
"openai:responses",
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some("key-123"),
|
||||
&http::HeaderMap::new(),
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, insert_provider_stream_event_api_format,
|
||||
LocalExecutionReportContextParts,
|
||||
@@ -29,6 +30,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
report_kind: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let decision_is_stream = decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let LocalOpenAiChatCandidateAttempt {
|
||||
eligible,
|
||||
@@ -107,6 +109,8 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
@@ -147,7 +151,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: upstream_is_stream,
|
||||
decision_is_stream,
|
||||
decision_kind: decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
|
||||
@@ -8,7 +8,10 @@ use crate::ai_serving::planner::candidate_preparation::{
|
||||
OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
use crate::ai_serving::planner::common::{
|
||||
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
|
||||
request_requires_body_stream_field, OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::standard::{
|
||||
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,
|
||||
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
|
||||
@@ -72,6 +75,8 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
let candidate = &eligible.candidate;
|
||||
let provider_api_format = eligible.provider_api_format.as_str();
|
||||
let transport = &eligible.transport;
|
||||
let force_body_stream_field =
|
||||
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
|
||||
let enable_model_directives =
|
||||
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
|
||||
state,
|
||||
@@ -128,6 +133,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
body_json,
|
||||
&prepared_candidate.mapped_model,
|
||||
upstream_is_stream,
|
||||
force_body_stream_field,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
&parts.headers,
|
||||
enable_model_directives,
|
||||
@@ -214,6 +220,13 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats("openai:chat", "openai:chat");
|
||||
let resolved_report_kind =
|
||||
if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND || !upstream_is_stream {
|
||||
report_kind.to_string()
|
||||
} else {
|
||||
"openai_chat_sync_finalize".to_string()
|
||||
};
|
||||
|
||||
return Some(LocalOpenAiChatCandidatePayloadParts {
|
||||
auth_header: resolved_headers.auth_header,
|
||||
auth_value: resolved_headers.auth_value,
|
||||
@@ -224,7 +237,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
upstream_url,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
report_kind: report_kind.to_string(),
|
||||
report_kind: resolved_report_kind,
|
||||
envelope_name: None,
|
||||
transport: Arc::clone(transport),
|
||||
});
|
||||
@@ -349,6 +362,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
upstream_is_stream,
|
||||
force_body_stream_field,
|
||||
if is_kiro_claude_cli {
|
||||
None
|
||||
} else {
|
||||
@@ -387,6 +401,14 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
&mut provider_request_body,
|
||||
&mapping,
|
||||
);
|
||||
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
|
||||
// re-enforce stream-field policy afterward.
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut provider_request_body,
|
||||
provider_api_format.as_str(),
|
||||
upstream_is_stream,
|
||||
request_requires_body_stream_field(body_json, force_body_stream_field),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
|
||||
@@ -9,6 +9,9 @@ mod stream;
|
||||
#[path = "plans/sync.rs"]
|
||||
mod sync;
|
||||
|
||||
use crate::ai_serving::planner::common::resolve_upstream_is_stream_for_provider;
|
||||
use crate::ai_serving::GatewayProviderTransportSnapshot;
|
||||
|
||||
pub(super) use self::candidates::list_local_openai_chat_candidates;
|
||||
pub(super) use self::diagnostic::set_local_openai_chat_miss_diagnostic;
|
||||
pub(super) use self::resolve::resolve_local_openai_chat_decision_input;
|
||||
@@ -18,3 +21,147 @@ pub(super) use self::stream::{
|
||||
pub(super) use self::sync::{
|
||||
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
|
||||
};
|
||||
|
||||
fn openai_chat_upstream_is_stream_for_candidate(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
client_is_stream: bool,
|
||||
) -> bool {
|
||||
let hard_requires_streaming =
|
||||
crate::ai_serving::transport::kiro::is_kiro_claude_messages_transport(
|
||||
transport,
|
||||
provider_api_format,
|
||||
);
|
||||
resolve_upstream_is_stream_for_provider(
|
||||
transport.endpoint.config.as_ref(),
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
client_is_stream,
|
||||
hard_requires_streaming,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::openai_chat_upstream_is_stream_for_candidate;
|
||||
use aether_provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn sample_transport(
|
||||
provider_type: &str,
|
||||
api_format: &str,
|
||||
endpoint_config: Option<Value>,
|
||||
) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "Provider".to_string(),
|
||||
provider_type: provider_type.to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: api_format.to_string(),
|
||||
api_family: None,
|
||||
endpoint_kind: None,
|
||||
is_active: true,
|
||||
base_url: "https://api.example.test".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: endpoint_config,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "key".to_string(),
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
allow_auth_channel_mismatch_formats: None,
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "secret".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_policy_resolver_supports_force_stream_force_non_stream_and_auto() {
|
||||
let force_stream = sample_transport(
|
||||
"openai",
|
||||
"openai:chat",
|
||||
Some(json!({"upstream_stream_policy": "force_stream"})),
|
||||
);
|
||||
assert!(openai_chat_upstream_is_stream_for_candidate(
|
||||
&force_stream,
|
||||
"openai:chat",
|
||||
false,
|
||||
));
|
||||
|
||||
let force_non_stream = sample_transport(
|
||||
"openai",
|
||||
"openai:chat",
|
||||
Some(json!({"upstream_stream_policy": "force_non_stream"})),
|
||||
);
|
||||
assert!(!openai_chat_upstream_is_stream_for_candidate(
|
||||
&force_non_stream,
|
||||
"openai:chat",
|
||||
true,
|
||||
));
|
||||
|
||||
let auto = sample_transport(
|
||||
"openai",
|
||||
"openai:chat",
|
||||
Some(json!({"upstream_stream_policy": "auto"})),
|
||||
);
|
||||
assert!(openai_chat_upstream_is_stream_for_candidate(
|
||||
&auto,
|
||||
"openai:chat",
|
||||
true,
|
||||
));
|
||||
assert!(!openai_chat_upstream_is_stream_for_candidate(
|
||||
&auto,
|
||||
"openai:chat",
|
||||
false,
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_policy_resolver_preserves_provider_hard_streaming() {
|
||||
let codex = sample_transport(
|
||||
"codex",
|
||||
"openai:responses",
|
||||
Some(json!({"upstream_stream_policy": "force_non_stream"})),
|
||||
);
|
||||
|
||||
assert!(openai_chat_upstream_is_stream_for_candidate(
|
||||
&codex,
|
||||
"openai:responses",
|
||||
false,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ use super::candidates::list_local_openai_chat_candidates;
|
||||
use super::diagnostic::{
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
use super::openai_chat_upstream_is_stream_for_candidate;
|
||||
use super::resolve::resolve_local_openai_chat_decision_input;
|
||||
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
|
||||
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
@@ -120,6 +121,11 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
|
||||
&self,
|
||||
attempt: LocalOpenAiChatCandidateAttempt,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
|
||||
&attempt.eligible.transport,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
true,
|
||||
);
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
self.state,
|
||||
self.parts,
|
||||
@@ -129,7 +135,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
"openai_chat_stream_success",
|
||||
true,
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
@@ -222,6 +228,11 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
|
||||
&attempt.eligible.transport,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
true,
|
||||
);
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
@@ -231,7 +242,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
"openai_chat_stream_success",
|
||||
true,
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
|
||||
@@ -13,11 +13,10 @@ use super::candidates::list_local_openai_chat_candidates;
|
||||
use super::diagnostic::{
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
use super::openai_chat_upstream_is_stream_for_candidate;
|
||||
use super::resolve::resolve_local_openai_chat_decision_input;
|
||||
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
|
||||
use crate::ai_serving::planner::common::{
|
||||
force_upstream_streaming_for_provider, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::common::OPENAI_CHAT_SYNC_PLAN_KIND;
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_openai_chat_sync_plan_from_decision, AiSyncAttempt,
|
||||
};
|
||||
@@ -32,13 +31,6 @@ pub(crate) struct LocalOpenAiChatSyncAttemptSource<'a> {
|
||||
candidates: LocalOpenAiChatCandidateAttemptSource<'a>,
|
||||
}
|
||||
|
||||
fn openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
) -> bool {
|
||||
force_upstream_streaming_for_provider(provider_type, provider_api_format)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
@@ -129,9 +121,10 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
|
||||
&self,
|
||||
attempt: LocalOpenAiChatCandidateAttempt,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
attempt.eligible.transport.provider.provider_type.as_str(),
|
||||
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
|
||||
&attempt.eligible.transport,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
false,
|
||||
);
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
self.state,
|
||||
@@ -235,9 +228,10 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
attempt.eligible.transport.provider.provider_type.as_str(),
|
||||
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
|
||||
&attempt.eligible.transport,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
false,
|
||||
);
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
@@ -272,28 +266,3 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::openai_chat_sync_upstream_is_stream_for_candidate;
|
||||
|
||||
#[test]
|
||||
fn openai_chat_sync_forces_streaming_for_codex_openai_responses_candidates() {
|
||||
assert!(openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
"codex",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
"openai",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
"codex",
|
||||
"openai:chat"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ use super::super::{
|
||||
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
|
||||
AiStreamAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
|
||||
use crate::ai_serving::provider_adaptation_requires_eventstream_accept;
|
||||
use crate::ai_serving::transport::{
|
||||
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
|
||||
@@ -17,6 +18,18 @@ use crate::ai_serving::transport::{
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
fn effective_stream_accept_mode(
|
||||
payload_upstream_is_stream: bool,
|
||||
provider_request_body: &serde_json::Value,
|
||||
) -> bool {
|
||||
payload_upstream_is_stream
|
||||
|| provider_request_body
|
||||
.as_object()
|
||||
.and_then(|body| body.get("stream"))
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub(crate) fn build_openai_chat_stream_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
@@ -53,22 +66,34 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
|
||||
provider_request_body
|
||||
.insert("model".to_string(), serde_json::Value::String(mapped_model));
|
||||
}
|
||||
provider_request_body.insert("stream".to_string(), serde_json::Value::Bool(true));
|
||||
let require_body_stream_field = provider_request_body.contains_key("stream");
|
||||
let mut provider_request_body = serde_json::Value::Object(provider_request_body);
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut provider_request_body,
|
||||
core.provider_api_format.as_str(),
|
||||
payload.upstream_is_stream,
|
||||
require_body_stream_field,
|
||||
);
|
||||
let Some(provider_request_object) = provider_request_body.as_object_mut() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) {
|
||||
let existing = provider_request_body
|
||||
let existing = provider_request_object
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if existing.is_empty() {
|
||||
provider_request_body.insert(
|
||||
provider_request_object.insert(
|
||||
"prompt_cache_key".to_string(),
|
||||
serde_json::Value::String(prompt_cache_key),
|
||||
);
|
||||
}
|
||||
}
|
||||
serde_json::Value::Object(provider_request_body)
|
||||
provider_request_body
|
||||
};
|
||||
let effective_upstream_is_stream =
|
||||
effective_stream_accept_mode(payload.upstream_is_stream, &provider_request_body_value);
|
||||
let extra_headers = std::mem::take(&mut payload.extra_headers);
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
@@ -82,9 +107,9 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: true,
|
||||
upstream_is_stream: effective_upstream_is_stream,
|
||||
build_from_request_when_empty: true,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreamingOrWildcard,
|
||||
});
|
||||
let content_type = payload
|
||||
.content_type
|
||||
@@ -157,13 +182,18 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("envelope_name"))
|
||||
.and_then(serde_json::Value::as_str);
|
||||
let accept_policy = if provider_adaptation_requires_eventstream_accept(
|
||||
envelope_name,
|
||||
core.provider_api_format.as_str(),
|
||||
) {
|
||||
let effective_upstream_is_stream =
|
||||
effective_stream_accept_mode(payload.upstream_is_stream, &provider_request_body_value);
|
||||
let accept_policy = if effective_upstream_is_stream
|
||||
&& provider_adaptation_requires_eventstream_accept(
|
||||
envelope_name,
|
||||
core.provider_api_format.as_str(),
|
||||
) {
|
||||
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
|
||||
} else if envelope_name.is_some() {
|
||||
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
|
||||
} else {
|
||||
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired
|
||||
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreamingOrWildcard
|
||||
};
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
@@ -177,7 +207,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: true,
|
||||
upstream_is_stream: effective_upstream_is_stream,
|
||||
build_from_request_when_empty: false,
|
||||
accept_policy,
|
||||
});
|
||||
@@ -219,7 +249,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
|
||||
plan_url = %plan.url,
|
||||
client_api_format = %plan.client_api_format,
|
||||
provider_api_format = %plan.provider_api_format,
|
||||
upstream_is_stream = payload.upstream_is_stream,
|
||||
upstream_is_stream = effective_upstream_is_stream,
|
||||
compact,
|
||||
"gateway built local openai responses stream execution plan"
|
||||
);
|
||||
@@ -427,6 +457,104 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_stream_plan_keeps_downstream_stream_for_force_non_stream_upstream() {
|
||||
fn force_non_stream_payload(provider_request_body: Option<Value>) -> AiExecutionDecision {
|
||||
AiExecutionDecision {
|
||||
action: "stream".to_string(),
|
||||
decision_kind: Some("openai_chat_stream".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_force_non_stream".to_string()),
|
||||
candidate_id: Some("cand_force_non_stream".to_string()),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_id: Some("prov_force_non_stream".to_string()),
|
||||
endpoint_id: Some("ep_force_non_stream".to_string()),
|
||||
key_id: Some("key_force_non_stream".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer upstream-token".to_string()),
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("openai:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body,
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some("openai_chat_stream_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
|
||||
let built = build_openai_chat_stream_plan_from_decision(
|
||||
&parts,
|
||||
&json!({}),
|
||||
force_non_stream_payload(Some(json!({
|
||||
"model": "gpt-5.4",
|
||||
"messages": [],
|
||||
"stream": false
|
||||
}))),
|
||||
)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert!(built.plan.stream);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("stream"))
|
||||
.and_then(Value::as_bool),
|
||||
Some(false)
|
||||
);
|
||||
|
||||
let fallback_body = json!({
|
||||
"model": "client-model",
|
||||
"messages": [],
|
||||
"stream": true
|
||||
});
|
||||
let built = build_openai_chat_stream_plan_from_decision(
|
||||
&parts,
|
||||
&fallback_body,
|
||||
force_non_stream_payload(None),
|
||||
)
|
||||
.expect("fallback plan build should succeed")
|
||||
.expect("fallback plan should be produced");
|
||||
|
||||
assert!(built.plan.stream);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("stream"))
|
||||
.and_then(Value::as_bool),
|
||||
Some(false)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_stream_plan_fallback_restores_claude_headers_for_cross_format() {
|
||||
let parts = http::Request::builder()
|
||||
|
||||
@@ -9,6 +9,7 @@ use super::super::{
|
||||
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
|
||||
AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
|
||||
use crate::ai_serving::transport::{
|
||||
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
|
||||
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
|
||||
@@ -51,23 +52,31 @@ pub(crate) fn build_openai_chat_sync_plan_from_decision(
|
||||
provider_request_body
|
||||
.insert("model".to_string(), serde_json::Value::String(mapped_model));
|
||||
}
|
||||
if payload.upstream_is_stream {
|
||||
provider_request_body.insert("stream".to_string(), serde_json::Value::Bool(true));
|
||||
}
|
||||
let require_body_stream_field = provider_request_body.contains_key("stream");
|
||||
let mut provider_request_body = serde_json::Value::Object(provider_request_body);
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut provider_request_body,
|
||||
core.provider_api_format.as_str(),
|
||||
payload.upstream_is_stream,
|
||||
require_body_stream_field,
|
||||
);
|
||||
let Some(provider_request_object) = provider_request_body.as_object_mut() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) {
|
||||
let existing = provider_request_body
|
||||
let existing = provider_request_object
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if existing.is_empty() {
|
||||
provider_request_body.insert(
|
||||
provider_request_object.insert(
|
||||
"prompt_cache_key".to_string(),
|
||||
serde_json::Value::String(prompt_cache_key),
|
||||
);
|
||||
}
|
||||
}
|
||||
serde_json::Value::Object(provider_request_body)
|
||||
provider_request_body
|
||||
};
|
||||
let extra_headers = std::mem::take(&mut payload.extra_headers);
|
||||
let mut provider_request_headers =
|
||||
|
||||
@@ -105,6 +105,8 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
|
||||
@@ -9,7 +9,10 @@ use crate::ai_serving::planner::candidate_preparation::{
|
||||
OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider;
|
||||
use crate::ai_serving::planner::common::{
|
||||
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
|
||||
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
|
||||
use crate::ai_serving::planner::standard::{
|
||||
apply_codex_openai_responses_special_headers, build_cross_format_openai_responses_request_body,
|
||||
@@ -224,12 +227,15 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
.await;
|
||||
|
||||
let needs_bidirectional_conversion = !same_format && conversion_kind.is_some();
|
||||
let upstream_is_stream = spec_metadata.require_streaming
|
||||
|| is_antigravity
|
||||
|| force_upstream_streaming_for_provider(
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
);
|
||||
let upstream_is_stream = resolve_upstream_is_stream_for_provider(
|
||||
transport.endpoint.config.as_ref(),
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
spec_metadata.require_streaming,
|
||||
is_antigravity || is_kiro_claude_cli,
|
||||
);
|
||||
let force_body_stream_field =
|
||||
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
|
||||
let Some(mut base_provider_request_body) = (if needs_bidirectional_conversion {
|
||||
build_cross_format_openai_responses_request_body(
|
||||
body_json,
|
||||
@@ -237,6 +243,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
force_body_stream_field,
|
||||
transport.provider.provider_type.as_str(),
|
||||
if is_kiro_claude_cli {
|
||||
None
|
||||
@@ -252,6 +259,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
body_json,
|
||||
&mapped_model,
|
||||
upstream_is_stream,
|
||||
force_body_stream_field,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
if is_kiro_claude_cli {
|
||||
@@ -293,6 +301,14 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
&mut base_provider_request_body,
|
||||
&mapping,
|
||||
);
|
||||
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
|
||||
// re-enforce stream-field policy afterward.
|
||||
enforce_provider_body_stream_policy(
|
||||
&mut base_provider_request_body,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
request_requires_body_stream_field(body_json, force_body_stream_field),
|
||||
);
|
||||
}
|
||||
let antigravity_auth = if is_antigravity {
|
||||
match classify_local_antigravity_request_support(
|
||||
|
||||
@@ -112,13 +112,14 @@ pub(crate) fn build_standard_stream_plan_from_decision(
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("envelope_name"))
|
||||
.and_then(serde_json::Value::as_str);
|
||||
let accept_policy = if provider_adaptation_requires_eventstream_accept(
|
||||
envelope_name,
|
||||
core.provider_api_format.as_str(),
|
||||
) {
|
||||
let accept_policy = if payload.upstream_is_stream
|
||||
&& provider_adaptation_requires_eventstream_accept(
|
||||
envelope_name,
|
||||
core.provider_api_format.as_str(),
|
||||
) {
|
||||
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
|
||||
} else {
|
||||
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired
|
||||
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
|
||||
};
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
@@ -132,7 +133,7 @@ pub(crate) fn build_standard_stream_plan_from_decision(
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: true,
|
||||
upstream_is_stream: payload.upstream_is_stream,
|
||||
build_from_request_when_empty: false,
|
||||
accept_policy,
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -24,5 +24,7 @@ pub(crate) fn mount_public_support_routes(router: Router<AppState>) -> Router<Ap
|
||||
.route("/api/capabilities", get(proxy_request))
|
||||
.route("/api/capabilities/user-configurable", get(proxy_request))
|
||||
.route("/api/capabilities/model/{*model_path}", get(proxy_request))
|
||||
.route("/install/{*install_path}", get(proxy_request))
|
||||
.route("/i/{*install_path}", get(proxy_request))
|
||||
.route("/", get(proxy_request))
|
||||
}
|
||||
|
||||
@@ -152,6 +152,11 @@ const PERMISSION_GROUPS: &[PermissionGroup] = &[
|
||||
label: "系统",
|
||||
assignable: true,
|
||||
},
|
||||
PermissionGroup {
|
||||
scope: "tasks",
|
||||
label: "后台任务",
|
||||
assignable: true,
|
||||
},
|
||||
PermissionGroup {
|
||||
scope: "usage",
|
||||
label: "用量",
|
||||
@@ -434,6 +439,9 @@ fn permission_key(scope: &str, access: &str) -> &'static str {
|
||||
("system", "read") => "admin:system:read",
|
||||
("system", "write") => "admin:system:write",
|
||||
("system", "admin") => "admin:system:admin",
|
||||
("tasks", "read") => "admin:tasks:read",
|
||||
("tasks", "write") => "admin:tasks:write",
|
||||
("tasks", "admin") => "admin:tasks:admin",
|
||||
("usage", "read") => "admin:usage:read",
|
||||
("usage", "write") => "admin:usage:write",
|
||||
("usage", "admin") => "admin:usage:admin",
|
||||
@@ -492,6 +500,7 @@ mod tests {
|
||||
"security",
|
||||
"stats",
|
||||
"system",
|
||||
"tasks",
|
||||
"usage",
|
||||
"users",
|
||||
"video_tasks",
|
||||
|
||||
@@ -46,6 +46,83 @@ pub(super) fn classify_admin_operations_family_route(
|
||||
"admin:video_tasks",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& matches!(normalized_path, "/api/admin/tasks" | "/api/admin/tasks/")
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"tasks_manage",
|
||||
"list_tasks",
|
||||
"admin:tasks",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& matches!(
|
||||
normalized_path,
|
||||
"/api/admin/tasks/stats" | "/api/admin/tasks/stats/"
|
||||
)
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"tasks_manage",
|
||||
"stats",
|
||||
"admin:tasks",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& normalized_path.starts_with("/api/admin/tasks/")
|
||||
&& normalized_path.ends_with("/events")
|
||||
&& normalized_path.matches('/').count() == 5
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"tasks_manage",
|
||||
"events",
|
||||
"admin:tasks",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& normalized_path.starts_with("/api/admin/tasks/")
|
||||
&& normalized_path.ends_with("/cancel")
|
||||
&& normalized_path.matches('/').count() == 5
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"tasks_manage",
|
||||
"cancel",
|
||||
"admin:tasks",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& normalized_path.starts_with("/api/admin/tasks/")
|
||||
&& normalized_path.ends_with("/trigger")
|
||||
&& normalized_path.matches('/').count() == 5
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"tasks_manage",
|
||||
"trigger",
|
||||
"admin:tasks",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& normalized_path.starts_with("/api/admin/tasks/")
|
||||
&& normalized_path["/api/admin/tasks/".len()..]
|
||||
.split('/')
|
||||
.count()
|
||||
== 1
|
||||
&& !matches!(
|
||||
normalized_path,
|
||||
"/api/admin/tasks/stats" | "/api/admin/tasks/stats/"
|
||||
)
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"tasks_manage",
|
||||
"detail",
|
||||
"admin:tasks",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& normalized_path.starts_with("/api/admin/video-tasks/")
|
||||
&& normalized_path.ends_with("/video")
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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(&[]);
|
||||
|
||||
@@ -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(&[]);
|
||||
|
||||
@@ -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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -5,11 +5,12 @@ use super::{
|
||||
AdminBillingRuleWriteInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery,
|
||||
AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery,
|
||||
AdminWalletRefundRequestListQuery, AnnouncementListQuery, AuditLogListQuery,
|
||||
CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
|
||||
CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord, CreateManualWalletRechargeInput,
|
||||
CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome,
|
||||
CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput,
|
||||
DataLayerError, DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
|
||||
BackgroundTaskListQuery, BackgroundTaskSummary, CompleteAdminWalletRefundInput,
|
||||
CreateAdminRedeemCodeBatchInput, CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord,
|
||||
CreateManualWalletRechargeInput, CreateWalletRechargeOrderInput,
|
||||
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
|
||||
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, DataLayerError,
|
||||
DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
|
||||
DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput,
|
||||
GatewayDataState, GatewayProviderTransportSnapshot, LocalVideoTaskReadResponse,
|
||||
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
|
||||
@@ -19,15 +20,16 @@ use super::{
|
||||
StoredAdminRedeemCodePage, StoredAdminWalletLedgerPage, StoredAdminWalletListPage,
|
||||
StoredAdminWalletRefund, StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage,
|
||||
StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredAnnouncement,
|
||||
StoredAnnouncementPage, StoredBillingModelContext, StoredProviderQuotaSnapshot,
|
||||
StoredAnnouncementPage, StoredBackgroundTaskEvent, StoredBackgroundTaskRun,
|
||||
StoredBackgroundTaskRunPage, StoredBillingModelContext, StoredProviderQuotaSnapshot,
|
||||
StoredProviderUsageSummary, StoredRequestUsageAudit, StoredSuspiciousActivity,
|
||||
StoredUsageSettlement, StoredUserAuditLogPage, StoredUserAuthRecord, StoredUserExportRow,
|
||||
StoredUserSummary, StoredVideoTask, StoredWalletDailyUsageLedger,
|
||||
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAnnouncementRecord,
|
||||
UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput, VideoTaskLookupKey,
|
||||
VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount,
|
||||
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult, WalletLookupKey,
|
||||
WalletMutationOutcome,
|
||||
UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun, UpsertUsageRecord, UpsertVideoTask,
|
||||
UsageSettlementInput, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter,
|
||||
VideoTaskStatusCount, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
|
||||
WalletLookupKey, WalletMutationOutcome,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::{
|
||||
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
|
||||
@@ -1716,6 +1718,82 @@ impl GatewayDataState {
|
||||
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
|
||||
read_data_backed_video_task_response(self, route_family, request_path).await
|
||||
}
|
||||
|
||||
pub(crate) async fn find_background_task_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
|
||||
match &self.background_task_reader {
|
||||
Some(repository) => repository.find_run(run_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_background_task_runs(
|
||||
&self,
|
||||
query: &BackgroundTaskListQuery,
|
||||
) -> Result<StoredBackgroundTaskRunPage, DataLayerError> {
|
||||
match &self.background_task_reader {
|
||||
Some(repository) => repository.list_runs(query).await,
|
||||
None => Ok(StoredBackgroundTaskRunPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_background_task_events(
|
||||
&self,
|
||||
run_id: &str,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredBackgroundTaskEvent>, DataLayerError> {
|
||||
match &self.background_task_reader {
|
||||
Some(repository) => repository.list_events(run_id, offset, limit).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_background_task_runs(
|
||||
&self,
|
||||
) -> Result<BackgroundTaskSummary, DataLayerError> {
|
||||
match &self.background_task_reader {
|
||||
Some(repository) => repository.summarize_runs().await,
|
||||
None => Ok(BackgroundTaskSummary::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_background_task_run(
|
||||
&self,
|
||||
run: UpsertBackgroundTaskRun,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
|
||||
match &self.background_task_writer {
|
||||
Some(repository) => repository.upsert_run(run).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn request_cancel_background_task_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.background_task_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.request_cancel(run_id, updated_at_unix_secs)
|
||||
.await
|
||||
}
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_background_task_event(
|
||||
&self,
|
||||
event: UpsertBackgroundTaskEvent,
|
||||
) -> Result<Option<StoredBackgroundTaskEvent>, DataLayerError> {
|
||||
match &self.background_task_writer {
|
||||
Some(repository) => repository.upsert_event(event).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -43,6 +43,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -93,6 +95,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
|
||||
@@ -76,6 +76,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -125,6 +127,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -170,6 +174,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -303,6 +309,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -364,6 +372,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -418,6 +428,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -481,6 +493,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -526,6 +540,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -572,6 +588,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -629,6 +647,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -688,6 +708,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -731,6 +753,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -789,6 +813,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -840,6 +866,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: Some(wallet_reader),
|
||||
wallet_writer: Some(wallet_writer),
|
||||
settlement_writer: None,
|
||||
@@ -896,6 +924,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: Some(wallet_reader),
|
||||
wallet_writer: Some(wallet_writer),
|
||||
settlement_writer: None,
|
||||
@@ -953,6 +983,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: Some(wallet_reader),
|
||||
wallet_writer: Some(wallet_writer),
|
||||
settlement_writer: None,
|
||||
@@ -1009,6 +1041,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1054,6 +1088,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1099,6 +1135,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1156,6 +1194,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1218,6 +1258,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1263,6 +1305,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1313,6 +1357,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1380,6 +1426,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1442,6 +1490,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1488,6 +1538,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1534,6 +1586,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1582,6 +1636,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1628,6 +1684,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1674,6 +1732,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1728,6 +1788,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1783,6 +1845,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1841,6 +1905,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1905,6 +1971,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -1970,6 +2038,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -2039,6 +2109,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -2115,6 +2187,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: Some(wallet_reader),
|
||||
wallet_writer: Some(wallet_writer),
|
||||
settlement_writer: None,
|
||||
@@ -2173,6 +2247,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -2222,6 +2298,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -2267,6 +2345,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -2318,6 +2398,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: Some(wallet_reader),
|
||||
wallet_writer: Some(wallet_writer),
|
||||
settlement_writer: None,
|
||||
@@ -2373,6 +2455,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: Some(wallet_reader),
|
||||
wallet_writer: Some(wallet_writer),
|
||||
settlement_writer: None,
|
||||
@@ -2429,6 +2513,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: Some(wallet_reader),
|
||||
wallet_writer: Some(wallet_writer),
|
||||
settlement_writer: None,
|
||||
@@ -2478,6 +2564,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
|
||||
@@ -47,6 +47,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: Some(repository),
|
||||
video_task_writer: None,
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -99,6 +101,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -148,6 +152,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -201,6 +207,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -258,6 +266,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
@@ -324,6 +334,8 @@ impl GatewayDataState {
|
||||
usage_worker_queue: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
background_task_reader: None,
|
||||
background_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
|
||||
@@ -1,23 +1,17 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::{StreamFrame, StreamFramePayload};
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use axum::http::StatusCode;
|
||||
use base64::Engine as _;
|
||||
use futures_util::StreamExt;
|
||||
use serde_json::json;
|
||||
use serde_json::{json, Map, Value};
|
||||
use tokio_util::codec::{FramedRead, LinesCodec};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::api::response::build_client_response_from_parts;
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::ndjson::decode_stream_frame_ndjson;
|
||||
use crate::execution_runtime::submission::{has_nested_error, strip_utf8_bom_and_ws};
|
||||
use crate::GatewayError;
|
||||
use crate::{
|
||||
GEMINI_FILES_DOWNLOAD_PLAN_KIND, MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_FRAMES,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
};
|
||||
use crate::{MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_FRAMES};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(super) enum StreamPrefetchInspection {
|
||||
@@ -51,6 +45,88 @@ pub(super) fn decode_stream_error_body(
|
||||
)
|
||||
}
|
||||
|
||||
fn header_value_case_insensitive<'a>(
|
||||
headers: &'a BTreeMap<String, String>,
|
||||
name: &str,
|
||||
) -> Option<&'a str> {
|
||||
headers
|
||||
.iter()
|
||||
.find(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||
.map(|(_, value)| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn remove_header_case_insensitive(headers: &mut BTreeMap<String, String>, name: &str) {
|
||||
let keys = headers
|
||||
.keys()
|
||||
.filter(|key| key.eq_ignore_ascii_case(name))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
for key in keys {
|
||||
headers.remove(&key);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn should_synthesize_non_success_stream_error_body(
|
||||
status_code: u16,
|
||||
error_body: &[u8],
|
||||
) -> bool {
|
||||
!(200..300).contains(&status_code)
|
||||
&& ((300..400).contains(&status_code) || error_body.is_empty())
|
||||
}
|
||||
|
||||
pub(super) fn build_synthetic_non_success_stream_error_body(
|
||||
status_code: u16,
|
||||
headers: &BTreeMap<String, String>,
|
||||
) -> Value {
|
||||
let mut error = Map::from_iter([
|
||||
(
|
||||
"type".to_string(),
|
||||
Value::String("execution_runtime_non_success_status".to_string()),
|
||||
),
|
||||
(
|
||||
"message".to_string(),
|
||||
Value::String(format!(
|
||||
"execution runtime stream returned non-success status {status_code}"
|
||||
)),
|
||||
),
|
||||
("code".to_string(), Value::from(status_code)),
|
||||
("upstream_status".to_string(), Value::from(status_code)),
|
||||
]);
|
||||
if let Some(location) = header_value_case_insensitive(headers, "location") {
|
||||
error.insert("location".to_string(), Value::String(location.to_string()));
|
||||
}
|
||||
|
||||
Value::Object(Map::from_iter([(
|
||||
"error".to_string(),
|
||||
Value::Object(error),
|
||||
)]))
|
||||
}
|
||||
|
||||
pub(super) fn synthetic_error_response_headers(
|
||||
mut headers: BTreeMap<String, String>,
|
||||
) -> BTreeMap<String, String> {
|
||||
remove_header_case_insensitive(&mut headers, "content-encoding");
|
||||
remove_header_case_insensitive(&mut headers, "content-length");
|
||||
remove_header_case_insensitive(&mut headers, "content-type");
|
||||
remove_header_case_insensitive(&mut headers, "location");
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
headers
|
||||
}
|
||||
|
||||
fn client_error_status_code_for_upstream_status(status_code: u16) -> u16 {
|
||||
if (300..400).contains(&status_code) || status_code < 200 {
|
||||
StatusCode::BAD_GATEWAY.as_u16()
|
||||
} else {
|
||||
status_code
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn stream_client_error_status_code_for_upstream_status(status_code: u16) -> u16 {
|
||||
client_error_status_code_for_upstream_status(status_code)
|
||||
}
|
||||
|
||||
pub(super) fn inspect_prefetched_stream_body(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body: &[u8],
|
||||
@@ -166,62 +242,3 @@ where
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(super) fn build_execution_runtime_error_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
status_code: u16,
|
||||
headers: BTreeMap<String, String>,
|
||||
error_body: Vec<u8>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let content_type = headers
|
||||
.get("content-type")
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.unwrap_or_default();
|
||||
|
||||
if plan_kind == GEMINI_FILES_DOWNLOAD_PLAN_KIND && !content_type.starts_with("application/json")
|
||||
{
|
||||
let wrapped = serde_json::to_vec(&json!({
|
||||
"error": String::from_utf8_lossy(&error_body).to_string(),
|
||||
}))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let wrapped_headers =
|
||||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||||
return build_client_response_from_parts(
|
||||
status_code,
|
||||
&wrapped_headers,
|
||||
Body::from(wrapped),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
);
|
||||
}
|
||||
|
||||
if plan_kind == OPENAI_VIDEO_CONTENT_PLAN_KIND && !content_type.starts_with("application/json")
|
||||
{
|
||||
let wrapped = serde_json::to_vec(&json!({
|
||||
"error": {
|
||||
"type": "upstream_error",
|
||||
"message": "Video not available",
|
||||
}
|
||||
}))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let wrapped_headers =
|
||||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||||
return build_client_response_from_parts(
|
||||
status_code,
|
||||
&wrapped_headers,
|
||||
Body::from(wrapped),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
);
|
||||
}
|
||||
|
||||
build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
Body::from(error_body),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
) {
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::GatewayError;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::response::Response;
|
||||
|
||||
mod routes;
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_background_tasks_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
routes::maybe_build_local_admin_background_tasks_response(state, request_context, request_body)
|
||||
.await
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::{
|
||||
attach_admin_audit_response, query_param_value, unix_secs_to_rfc3339,
|
||||
};
|
||||
use crate::task_runtime::{
|
||||
self, set_cancel_signal, TASK_KEY_PROVIDER_DELETE, TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use aether_data_contracts::repository::background_tasks::{
|
||||
BackgroundTaskKind, BackgroundTaskListQuery, BackgroundTaskStatus,
|
||||
};
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
http,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
const DEFAULT_PAGE_SIZE: usize = 20;
|
||||
const MAX_PAGE_SIZE: usize = 100;
|
||||
const DEFAULT_EVENTS_PAGE_SIZE: usize = 50;
|
||||
|
||||
pub(super) async fn maybe_build_local_admin_background_tasks_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if request_context.route_family() != Some("tasks_manage") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
match request_context.route_kind() {
|
||||
Some("list_tasks") if request_context.method() == http::Method::GET => {
|
||||
let query = request_context.query_string();
|
||||
let page = query_param_value(query, "page")
|
||||
.and_then(|value| value.parse::<usize>().ok())
|
||||
.unwrap_or(1)
|
||||
.max(1);
|
||||
let page_size = query_param_value(query, "page_size")
|
||||
.and_then(|value| value.parse::<usize>().ok())
|
||||
.unwrap_or(DEFAULT_PAGE_SIZE)
|
||||
.clamp(1, MAX_PAGE_SIZE);
|
||||
let kind = query_param_value(query, "kind")
|
||||
.map(|value| BackgroundTaskKind::from_database(value.as_str()))
|
||||
.transpose()
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let status = query_param_value(query, "status")
|
||||
.map(|value| BackgroundTaskStatus::from_database(value.as_str()))
|
||||
.transpose()
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let trigger = query_param_value(query, "trigger");
|
||||
let task_key_substring = query_param_value(query, "task_key");
|
||||
let offset = (page - 1).saturating_mul(page_size);
|
||||
let response = state
|
||||
.list_background_task_runs(&BackgroundTaskListQuery {
|
||||
task_key_substring,
|
||||
kind,
|
||||
status,
|
||||
trigger,
|
||||
offset,
|
||||
limit: page_size,
|
||||
})
|
||||
.await?;
|
||||
let pages = if response.total == 0 {
|
||||
0
|
||||
} else {
|
||||
(response.total + page_size - 1) / page_size
|
||||
};
|
||||
|
||||
let items = response
|
||||
.items
|
||||
.iter()
|
||||
.map(|run| {
|
||||
json!({
|
||||
"id": run.id,
|
||||
"task_key": run.task_key,
|
||||
"kind": run.kind.as_database(),
|
||||
"trigger": run.trigger,
|
||||
"status": run.status.as_database(),
|
||||
"attempt": run.attempt,
|
||||
"max_attempts": run.max_attempts,
|
||||
"owner_instance": run.owner_instance,
|
||||
"progress_percent": run.progress_percent,
|
||||
"progress_message": run.progress_message,
|
||||
"payload": run.payload_json,
|
||||
"result": run.result_json,
|
||||
"error_message": run.error_message,
|
||||
"cancel_requested": run.cancel_requested,
|
||||
"created_by": run.created_by,
|
||||
"created_at": unix_secs_to_rfc3339(run.created_at_unix_secs),
|
||||
"started_at": run.started_at_unix_secs.and_then(unix_secs_to_rfc3339),
|
||||
"finished_at": run.finished_at_unix_secs.and_then(unix_secs_to_rfc3339),
|
||||
"updated_at": unix_secs_to_rfc3339(run.updated_at_unix_secs),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let definitions = task_runtime::task_definitions()
|
||||
.iter()
|
||||
.map(|definition| {
|
||||
json!({
|
||||
"task_key": definition.key,
|
||||
"kind": definition.kind.as_str(),
|
||||
"trigger": definition.trigger,
|
||||
"max_attempts": definition.retry_policy.max_attempts,
|
||||
"singleton": definition.singleton,
|
||||
"persist_history": definition.persist_history,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
return Ok(Some(
|
||||
Json(json!({
|
||||
"items": items,
|
||||
"total": response.total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"pages": pages,
|
||||
"definitions": definitions,
|
||||
}))
|
||||
.into_response(),
|
||||
));
|
||||
}
|
||||
Some("stats") if request_context.method() == http::Method::GET => {
|
||||
let stats = state.summarize_background_task_runs().await?;
|
||||
return Ok(Some(
|
||||
Json(json!({
|
||||
"total": stats.total,
|
||||
"running_count": stats.running_count,
|
||||
"by_status": stats.by_status,
|
||||
"by_kind": stats.by_kind,
|
||||
"registered_tasks": task_runtime::task_definitions().len(),
|
||||
}))
|
||||
.into_response(),
|
||||
));
|
||||
}
|
||||
Some("detail") if request_context.method() == http::Method::GET => {
|
||||
let Some(run_id) = task_id_from_path(request_context.path()) else {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({"detail":"Task not found"})),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
};
|
||||
let Some(run) = state.find_background_task_run(run_id).await? else {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({"detail":"Task not found"})),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
};
|
||||
return Ok(Some(attach_admin_audit_response(
|
||||
Json(json!({
|
||||
"id": run.id,
|
||||
"task_key": run.task_key,
|
||||
"kind": run.kind.as_database(),
|
||||
"trigger": run.trigger,
|
||||
"status": run.status.as_database(),
|
||||
"attempt": run.attempt,
|
||||
"max_attempts": run.max_attempts,
|
||||
"owner_instance": run.owner_instance,
|
||||
"progress_percent": run.progress_percent,
|
||||
"progress_message": run.progress_message,
|
||||
"payload": run.payload_json,
|
||||
"result": run.result_json,
|
||||
"error_message": run.error_message,
|
||||
"cancel_requested": run.cancel_requested,
|
||||
"created_by": run.created_by,
|
||||
"created_at": unix_secs_to_rfc3339(run.created_at_unix_secs),
|
||||
"started_at": run.started_at_unix_secs.and_then(unix_secs_to_rfc3339),
|
||||
"finished_at": run.finished_at_unix_secs.and_then(unix_secs_to_rfc3339),
|
||||
"updated_at": unix_secs_to_rfc3339(run.updated_at_unix_secs),
|
||||
}))
|
||||
.into_response(),
|
||||
"admin_task_detail_viewed",
|
||||
"view_task_detail",
|
||||
"background_task",
|
||||
run_id,
|
||||
)));
|
||||
}
|
||||
Some("events") if request_context.method() == http::Method::GET => {
|
||||
let Some(run_id) = nested_task_id_from_path(request_context.path(), "/events") else {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({"detail":"Task not found"})),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
};
|
||||
let query = request_context.query_string();
|
||||
let page = query_param_value(query, "page")
|
||||
.and_then(|value| value.parse::<usize>().ok())
|
||||
.unwrap_or(1)
|
||||
.max(1);
|
||||
let page_size = query_param_value(query, "page_size")
|
||||
.and_then(|value| value.parse::<usize>().ok())
|
||||
.unwrap_or(DEFAULT_EVENTS_PAGE_SIZE)
|
||||
.clamp(1, MAX_PAGE_SIZE);
|
||||
let offset = (page - 1).saturating_mul(page_size);
|
||||
let events = state
|
||||
.list_background_task_events(run_id, offset, page_size)
|
||||
.await?;
|
||||
return Ok(Some(
|
||||
Json(json!({
|
||||
"items": events.into_iter().map(|event| {
|
||||
json!({
|
||||
"id": event.id,
|
||||
"run_id": event.run_id,
|
||||
"event_type": event.event_type,
|
||||
"message": event.message,
|
||||
"payload": event.payload_json,
|
||||
"created_at": unix_secs_to_rfc3339(event.created_at_unix_secs),
|
||||
})
|
||||
}).collect::<Vec<_>>(),
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
}))
|
||||
.into_response(),
|
||||
));
|
||||
}
|
||||
Some("cancel") if request_context.method() == http::Method::POST => {
|
||||
let Some(run_id) = nested_task_id_from_path(request_context.path(), "/cancel") else {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({"detail":"Task not found"})),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
};
|
||||
let now = task_runtime::now_unix_secs();
|
||||
let cancelled = state
|
||||
.request_cancel_background_task_run(run_id, now)
|
||||
.await?;
|
||||
if !cancelled {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({ "detail": "Task not found" })),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
}
|
||||
let _ = set_cancel_signal(state.app(), run_id).await;
|
||||
task_runtime::append_event_with_logging(
|
||||
state.app(),
|
||||
run_id,
|
||||
"cancel_requested",
|
||||
"cancel requested by admin",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
return Ok(Some(attach_admin_audit_response(
|
||||
Json(json!({
|
||||
"id": run_id,
|
||||
"status": "cancel_requested",
|
||||
"message": "Task cancellation requested",
|
||||
}))
|
||||
.into_response(),
|
||||
"admin_task_cancel_requested",
|
||||
"cancel_task",
|
||||
"background_task",
|
||||
run_id,
|
||||
)));
|
||||
}
|
||||
Some("trigger") if request_context.method() == http::Method::POST => {
|
||||
let Some(task_key) = nested_task_id_from_path(request_context.path(), "/trigger")
|
||||
else {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({"detail":"Task not found"})),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
};
|
||||
let payload = parse_json_payload(request_body)?;
|
||||
if task_key == TASK_KEY_PROVIDER_DELETE {
|
||||
let provider_id = payload
|
||||
.get("provider_id")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| {
|
||||
GatewayError::Internal(
|
||||
"admin task trigger provider delete requires provider_id".to_string(),
|
||||
)
|
||||
})?;
|
||||
let Some(run_id) =
|
||||
task_runtime::submit_provider_delete_task(state, provider_id, Some("admin"))
|
||||
.await?
|
||||
else {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({"detail":"Provider 不存在"})),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
};
|
||||
return Ok(Some(attach_admin_audit_response(
|
||||
Json(json!({
|
||||
"task_key": task_key,
|
||||
"run_id": run_id,
|
||||
"status": "queued",
|
||||
}))
|
||||
.into_response(),
|
||||
"admin_task_triggered",
|
||||
"trigger_task",
|
||||
"background_task",
|
||||
task_key,
|
||||
)));
|
||||
}
|
||||
if task_key == TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({
|
||||
"detail": "请使用 provider oauth batch import 专用接口触发该任务",
|
||||
})),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
}
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({
|
||||
"detail": format!("Unsupported task_key: {task_key}"),
|
||||
})),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn task_id_from_path(request_path: &str) -> Option<&str> {
|
||||
let task_id = request_path.strip_prefix("/api/admin/tasks/")?;
|
||||
if task_id.is_empty() || task_id.contains('/') || task_id == "stats" {
|
||||
return None;
|
||||
}
|
||||
Some(task_id)
|
||||
}
|
||||
|
||||
fn nested_task_id_from_path<'a>(request_path: &'a str, suffix: &str) -> Option<&'a str> {
|
||||
let task_id = request_path
|
||||
.strip_prefix("/api/admin/tasks/")?
|
||||
.strip_suffix(suffix)?;
|
||||
if task_id.is_empty() || task_id.contains('/') {
|
||||
return None;
|
||||
}
|
||||
Some(task_id)
|
||||
}
|
||||
|
||||
fn parse_json_payload(request_body: Option<&Bytes>) -> Result<serde_json::Value, GatewayError> {
|
||||
let Some(body) = request_body else {
|
||||
return Ok(json!({}));
|
||||
};
|
||||
if body.is_empty() {
|
||||
return Ok(json!({}));
|
||||
}
|
||||
serde_json::from_slice::<serde_json::Value>(body)
|
||||
.map_err(|err| GatewayError::Internal(format!("invalid json body: {err}")))
|
||||
}
|
||||
@@ -1,7 +1,9 @@
|
||||
mod background_tasks;
|
||||
mod gemini_files;
|
||||
mod routes;
|
||||
mod video_tasks;
|
||||
|
||||
pub(super) use self::background_tasks::maybe_build_local_admin_background_tasks_response;
|
||||
pub(super) use self::gemini_files::maybe_build_local_admin_gemini_files_response;
|
||||
pub(super) use self::routes::maybe_build_local_admin_features_response;
|
||||
pub(crate) use self::video_tasks::maybe_build_local_admin_video_tasks_response;
|
||||
|
||||
@@ -1,9 +1,19 @@
|
||||
use super::{gemini_files, video_tasks};
|
||||
use super::{background_tasks, gemini_files, video_tasks};
|
||||
use crate::handlers::admin::request::{AdminRouteRequest, AdminRouteResult};
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_features_response(
|
||||
request: AdminRouteRequest<'_>,
|
||||
) -> AdminRouteResult {
|
||||
if let Some(response) = background_tasks::maybe_build_local_admin_background_tasks_response(
|
||||
&request.state(),
|
||||
&request.request_context(),
|
||||
request.request_body(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
if let Some(response) = video_tasks::maybe_build_local_admin_video_tasks_response(
|
||||
&request.state(),
|
||||
&request.request_context(),
|
||||
|
||||
@@ -356,6 +356,164 @@ async fn admin_monitoring_trace_request_enriches_proxy_timing_from_usage_audit()
|
||||
payload["candidates"][0]["extra_data"]["proxy"]["timing"]["response_wait_ms"],
|
||||
json!(475)
|
||||
);
|
||||
assert!(payload["candidates"][0]["extra_data"]
|
||||
.get("upstream_response")
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_monitoring_trace_request_exposes_request_path_from_usage_audit() {
|
||||
let mut candidate = sample_candidate(
|
||||
"cand-used",
|
||||
"request-1",
|
||||
0,
|
||||
RequestCandidateStatus::Failed,
|
||||
Some(101),
|
||||
Some(33),
|
||||
Some(403),
|
||||
);
|
||||
candidate.extra_data = Some(json!({
|
||||
"client_api_format": "gemini:generate_content"
|
||||
}));
|
||||
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![candidate]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider()],
|
||||
vec![sample_endpoint()],
|
||||
vec![sample_key()],
|
||||
));
|
||||
let mut usage = sample_usage(
|
||||
"request-1",
|
||||
"provider-1",
|
||||
"OpenAI",
|
||||
40,
|
||||
0.02,
|
||||
"failed",
|
||||
Some(403),
|
||||
100,
|
||||
);
|
||||
usage.candidate_id = Some("cand-used".to_string());
|
||||
usage.request_metadata = Some(json!({
|
||||
"request_path": "/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
"request_query_string": "alt=sse"
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
|
||||
let data_state =
|
||||
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
||||
request_candidates,
|
||||
usage_repository,
|
||||
)
|
||||
.with_provider_catalog_reader(provider_catalog);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
let context = request_context(http::Method::GET, "/api/admin/monitoring/trace/request-1");
|
||||
|
||||
let response = local_monitoring_response(&state, &context)
|
||||
.await
|
||||
.expect("handler should not error")
|
||||
.expect("route should be handled locally");
|
||||
|
||||
assert_eq!(response.status(), http::StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
|
||||
|
||||
assert_eq!(
|
||||
payload["request_path"],
|
||||
json!("/v1beta/models/gemini-2.5-pro:generateContent")
|
||||
);
|
||||
assert_eq!(payload["request_query_string"], json!("alt=sse"));
|
||||
assert_eq!(
|
||||
payload["request_path_and_query"],
|
||||
json!("/v1beta/models/gemini-2.5-pro:generateContent?alt=sse")
|
||||
);
|
||||
assert_eq!(
|
||||
payload["candidates"][0]["extra_data"]["request_path_and_query"],
|
||||
json!("/v1beta/models/gemini-2.5-pro:generateContent?alt=sse")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_monitoring_trace_request_exposes_failed_candidate_upstream_response_boundary() {
|
||||
let candidate = sample_candidate(
|
||||
"cand-used",
|
||||
"request-1",
|
||||
0,
|
||||
RequestCandidateStatus::Failed,
|
||||
Some(101),
|
||||
Some(33),
|
||||
Some(302),
|
||||
);
|
||||
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![candidate]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider()],
|
||||
vec![sample_endpoint()],
|
||||
vec![sample_key()],
|
||||
));
|
||||
let mut usage = sample_usage(
|
||||
"request-1",
|
||||
"provider-1",
|
||||
"OpenAI",
|
||||
40,
|
||||
0.02,
|
||||
"failed",
|
||||
Some(302),
|
||||
100,
|
||||
);
|
||||
usage.candidate_id = Some("cand-used".to_string());
|
||||
usage.response_headers = Some(json!({
|
||||
"location": "/",
|
||||
"content-type": "text/html"
|
||||
}));
|
||||
usage.client_response_headers = Some(json!({
|
||||
"content-type": "application/json",
|
||||
"x-aether-upstream-status": "302"
|
||||
}));
|
||||
usage.client_response_body = Some(json!({
|
||||
"error": {
|
||||
"type": "execution_runtime_non_success_status",
|
||||
"message": "execution runtime stream returned non-success status 302",
|
||||
"upstream_status": 302,
|
||||
"location": "/"
|
||||
}
|
||||
}));
|
||||
usage.request_metadata = Some(json!({
|
||||
"client_response_status_code": 502
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
|
||||
let data_state =
|
||||
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
||||
request_candidates,
|
||||
usage_repository,
|
||||
)
|
||||
.with_provider_catalog_reader(provider_catalog);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
let context = request_context(http::Method::GET, "/api/admin/monitoring/trace/request-1");
|
||||
|
||||
let response = local_monitoring_response(&state, &context)
|
||||
.await
|
||||
.expect("handler should not error")
|
||||
.expect("route should be handled locally");
|
||||
|
||||
assert_eq!(response.status(), http::StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
|
||||
let extra = &payload["candidates"][0]["extra_data"];
|
||||
assert_eq!(extra["upstream_response"]["status_code"], json!(302));
|
||||
assert_eq!(
|
||||
extra["upstream_response"]["headers"]["location"],
|
||||
json!("/")
|
||||
);
|
||||
assert!(extra["upstream_response"]["body"].is_null());
|
||||
assert!(extra.get("client_response").is_none());
|
||||
assert!(extra.get("provider_response").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
use crate::handlers::admin::provider::delete_task::run_admin_provider_delete_task;
|
||||
use crate::handlers::admin::provider::shared::paths::{
|
||||
admin_provider_delete_task_parts, admin_provider_id_for_manage_path,
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::support::build_admin_provider_delete_task_payload;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::attach_admin_audit_response;
|
||||
use crate::{GatewayError, LocalProviderDeleteTaskState};
|
||||
use crate::GatewayError;
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
@@ -13,8 +12,6 @@ use axum::{
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
fn build_admin_provider_not_found_response(detail: impl Into<String>) -> Response<Body> {
|
||||
(
|
||||
@@ -60,52 +57,18 @@ pub(crate) async fn maybe_build_local_admin_provider_delete_task_response(
|
||||
"Provider 不存在",
|
||||
)));
|
||||
};
|
||||
let Some(provider) = state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
let Some(task_id) =
|
||||
crate::task_runtime::submit_provider_delete_task(state, &provider_id, Some("admin"))
|
||||
.await?
|
||||
else {
|
||||
return Ok(Some(build_admin_provider_not_found_response(
|
||||
"提供商不存在",
|
||||
)));
|
||||
};
|
||||
let task_id = Uuid::new_v4().simple().to_string()[..16].to_string();
|
||||
let pending_task = LocalProviderDeleteTaskState {
|
||||
task_id: task_id.clone(),
|
||||
provider_id: provider.id.clone(),
|
||||
status: "pending".to_string(),
|
||||
stage: "queued".to_string(),
|
||||
total_keys: 0,
|
||||
deleted_keys: 0,
|
||||
total_endpoints: 0,
|
||||
deleted_endpoints: 0,
|
||||
message: "delete task submitted".to_string(),
|
||||
};
|
||||
state.put_provider_delete_task(pending_task.clone());
|
||||
if let Err(err) = state
|
||||
.run_admin_provider_delete_task(&provider.id, &task_id)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
"gateway admin provider delete task failed for provider {}: {:?}",
|
||||
provider.id, err
|
||||
);
|
||||
state.put_provider_delete_task(LocalProviderDeleteTaskState {
|
||||
task_id: task_id.clone(),
|
||||
provider_id: provider.id.clone(),
|
||||
status: "failed".to_string(),
|
||||
stage: "failed".to_string(),
|
||||
total_keys: 0,
|
||||
deleted_keys: 0,
|
||||
total_endpoints: 0,
|
||||
deleted_endpoints: 0,
|
||||
message: format!("provider delete failed: {err:?}"),
|
||||
});
|
||||
}
|
||||
return Ok(Some(attach_admin_audit_response(
|
||||
Json(json!({
|
||||
"task_id": task_id,
|
||||
"run_id": task_id,
|
||||
"status": "pending",
|
||||
"message": "删除任务已提交,提供商已进入后台删除队列",
|
||||
}))
|
||||
@@ -113,7 +76,7 @@ pub(crate) async fn maybe_build_local_admin_provider_delete_task_response(
|
||||
"admin_provider_delete_queued",
|
||||
"delete_provider",
|
||||
"provider",
|
||||
&provider.id,
|
||||
&provider_id,
|
||||
)));
|
||||
}
|
||||
|
||||
|
||||
@@ -15,7 +15,14 @@ use crate::handlers::admin::provider::oauth::state::{
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_task_provider_id;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::task_runtime::{
|
||||
append_event_with_logging, now_unix_secs, task_definition, update_run_status,
|
||||
upsert_run_with_logging, TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use aether_data_contracts::repository::background_tasks::{
|
||||
BackgroundTaskKind, BackgroundTaskStatus, UpsertBackgroundTaskRun,
|
||||
};
|
||||
use axum::{
|
||||
body::Bytes,
|
||||
http,
|
||||
@@ -133,11 +140,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
}
|
||||
|
||||
let task_id = Uuid::new_v4().to_string();
|
||||
let created_at = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let created_at = now_unix_secs();
|
||||
let submitted_state = build_admin_provider_oauth_batch_task_state(
|
||||
&task_id,
|
||||
&provider_id,
|
||||
@@ -167,6 +170,50 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
));
|
||||
}
|
||||
|
||||
if state.has_background_task_data_writer() {
|
||||
let max_attempts = task_definition(TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT)
|
||||
.map(|item| item.retry_policy.max_attempts)
|
||||
.unwrap_or(1);
|
||||
let run = UpsertBackgroundTaskRun {
|
||||
id: task_id.clone(),
|
||||
task_key: TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT.to_string(),
|
||||
kind: BackgroundTaskKind::OnDemand,
|
||||
trigger: "manual".to_string(),
|
||||
status: BackgroundTaskStatus::Queued,
|
||||
attempt: 1,
|
||||
max_attempts,
|
||||
owner_instance: Some(state.app().tunnel.local_instance_id().to_string()),
|
||||
progress_percent: 0,
|
||||
progress_message: Some("provider oauth batch import queued".to_string()),
|
||||
payload_json: Some(json!({
|
||||
"provider_id": provider_id.clone(),
|
||||
"provider_type": provider_type.clone(),
|
||||
"total": total,
|
||||
})),
|
||||
result_json: None,
|
||||
error_message: None,
|
||||
cancel_requested: false,
|
||||
created_by: Some("admin".to_string()),
|
||||
created_at_unix_secs: created_at,
|
||||
started_at_unix_secs: None,
|
||||
finished_at_unix_secs: None,
|
||||
updated_at_unix_secs: created_at,
|
||||
};
|
||||
let _ = upsert_run_with_logging(state.app(), run).await;
|
||||
append_event_with_logging(
|
||||
state.app(),
|
||||
&task_id,
|
||||
"queued",
|
||||
"provider oauth batch import queued",
|
||||
Some(json!({
|
||||
"provider_id": provider_id.clone(),
|
||||
"provider_type": provider_type.clone(),
|
||||
"total": total,
|
||||
})),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let task_state = state.cloned_app();
|
||||
let task_id_for_worker = task_id.clone();
|
||||
let provider_id_for_worker = provider_id.clone();
|
||||
@@ -201,6 +248,27 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &processing_state)
|
||||
.await;
|
||||
|
||||
let _ = update_run_status(
|
||||
&task_state,
|
||||
&task_id_for_worker,
|
||||
BackgroundTaskStatus::Running,
|
||||
Some(1),
|
||||
Some("provider oauth batch import started".to_string()),
|
||||
None,
|
||||
None,
|
||||
Some(started_at),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
append_event_with_logging(
|
||||
&task_state,
|
||||
&task_id_for_worker,
|
||||
"running",
|
||||
"provider oauth batch import started",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut progress_reporter = BatchTaskProgressReporter {
|
||||
app: task_state.clone(),
|
||||
task_id: task_id_for_worker.clone(),
|
||||
@@ -270,6 +338,34 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
let _ = AdminAppState::new(&task_state)
|
||||
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &completed_state)
|
||||
.await;
|
||||
let _ = update_run_status(
|
||||
&task_state,
|
||||
&task_id_for_worker,
|
||||
BackgroundTaskStatus::Succeeded,
|
||||
Some(100),
|
||||
Some(message),
|
||||
Some(json!({
|
||||
"provider_id": provider_id_for_worker,
|
||||
"provider_type": provider_type_for_worker,
|
||||
"total": outcome.total,
|
||||
"success": outcome.success,
|
||||
"failed": outcome.failed,
|
||||
"created_count": created_count,
|
||||
"replaced_count": replaced_count,
|
||||
})),
|
||||
None,
|
||||
None,
|
||||
Some(finished_at),
|
||||
)
|
||||
.await;
|
||||
append_event_with_logging(
|
||||
&task_state,
|
||||
&task_id_for_worker,
|
||||
"succeeded",
|
||||
"provider oauth batch import completed",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(err) => {
|
||||
let finished_at = SystemTime::now()
|
||||
@@ -299,6 +395,26 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
let _ = AdminAppState::new(&task_state)
|
||||
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &failed_state)
|
||||
.await;
|
||||
let _ = update_run_status(
|
||||
&task_state,
|
||||
&task_id_for_worker,
|
||||
BackgroundTaskStatus::Failed,
|
||||
Some(100),
|
||||
Some("provider oauth batch import failed".to_string()),
|
||||
None,
|
||||
Some(error_message.clone()),
|
||||
None,
|
||||
Some(finished_at),
|
||||
)
|
||||
.await;
|
||||
append_event_with_logging(
|
||||
&task_state,
|
||||
&task_id_for_worker,
|
||||
"failed",
|
||||
"provider oauth batch import failed",
|
||||
Some(json!({ "error": error_message.clone() })),
|
||||
)
|
||||
.await;
|
||||
tracing::warn!(
|
||||
task_id = %task_id_for_worker,
|
||||
provider_id = %provider_id_for_worker,
|
||||
|
||||
@@ -4,6 +4,7 @@ use super::quota::codex::refresh_codex_provider_quota_locally;
|
||||
use super::quota::kiro::refresh_kiro_provider_quota_locally;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::provider_key_auth::provider_key_is_oauth_managed;
|
||||
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH};
|
||||
use crate::{AppState, GatewayError};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
@@ -170,7 +171,7 @@ pub(crate) fn spawn_provider_oauth_account_state_refresh_after_update(
|
||||
key_id: String,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
spawn_fire_and_forget(TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH, async move {
|
||||
let _ = refresh_provider_oauth_account_state_after_update(
|
||||
&AdminAppState::new(&app),
|
||||
&provider,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use super::actions::admin_provider_ops_local_action_response;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_PROVIDER_BALANCE_REFRESH};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashSet;
|
||||
use std::time::Duration;
|
||||
@@ -130,7 +131,7 @@ pub(super) async fn spawn_admin_provider_ops_balance_refresh(
|
||||
|
||||
let app = state.cloned_app();
|
||||
let provider_id = provider_id.to_string();
|
||||
tokio::spawn(async move {
|
||||
spawn_fire_and_forget(TASK_KEY_PROVIDER_BALANCE_REFRESH, async move {
|
||||
let permit = match tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
ADMIN_PROVIDER_OPS_BALANCE_REFRESH_SEMAPHORE.acquire(),
|
||||
|
||||
@@ -76,6 +76,14 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.has_gemini_file_mapping_data_writer()
|
||||
}
|
||||
|
||||
pub(crate) fn has_background_task_data_reader(&self) -> bool {
|
||||
self.app.has_background_task_data_reader()
|
||||
}
|
||||
|
||||
pub(crate) fn has_background_task_data_writer(&self) -> bool {
|
||||
self.app.has_background_task_data_writer()
|
||||
}
|
||||
|
||||
pub(crate) fn has_auth_api_key_data_reader(&self) -> bool {
|
||||
self.app.has_auth_api_key_data_reader()
|
||||
}
|
||||
|
||||
@@ -2,6 +2,59 @@ use super::{AdminAppState, AdminCancelVideoTaskError};
|
||||
use crate::GatewayError;
|
||||
|
||||
impl<'a> AdminAppState<'a> {
|
||||
pub(crate) async fn list_background_task_runs(
|
||||
&self,
|
||||
query: &aether_data_contracts::repository::background_tasks::BackgroundTaskListQuery,
|
||||
) -> Result<
|
||||
aether_data_contracts::repository::background_tasks::StoredBackgroundTaskRunPage,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app.list_background_task_runs(query).await
|
||||
}
|
||||
|
||||
pub(crate) async fn find_background_task_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::background_tasks::StoredBackgroundTaskRun>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app.find_background_task_run(run_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_background_task_events(
|
||||
&self,
|
||||
run_id: &str,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<
|
||||
Vec<aether_data_contracts::repository::background_tasks::StoredBackgroundTaskEvent>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app
|
||||
.list_background_task_events(run_id, offset, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_background_task_runs(
|
||||
&self,
|
||||
) -> Result<
|
||||
aether_data_contracts::repository::background_tasks::BackgroundTaskSummary,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app.summarize_background_task_runs().await
|
||||
}
|
||||
|
||||
pub(crate) async fn request_cancel_background_task_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.app
|
||||
.request_cancel_background_task_run(run_id, updated_at_unix_secs)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_gemini_file_mappings(
|
||||
&self,
|
||||
query: &aether_data::repository::gemini_file_mappings::GeminiFileMappingListQuery,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -28,6 +28,8 @@ mod support_announcements;
|
||||
mod support_auth;
|
||||
#[path = "support/dashboard.rs"]
|
||||
mod support_dashboard;
|
||||
#[path = "support/install.rs"]
|
||||
mod support_install;
|
||||
#[path = "support/models.rs"]
|
||||
mod support_models;
|
||||
#[path = "support/monitoring.rs"]
|
||||
@@ -61,6 +63,10 @@ use self::support_auth::{
|
||||
build_auth_settings_payload, extract_client_device_id, maybe_build_local_auth_response,
|
||||
};
|
||||
use self::support_dashboard::maybe_build_local_dashboard_response;
|
||||
use self::support_install::{
|
||||
handle_users_me_api_key_install_session_create, maybe_build_local_install_response,
|
||||
users_me_api_key_install_sessions_path_matches,
|
||||
};
|
||||
use self::support_models::{
|
||||
build_models_auth_error_response, maybe_build_local_models_response, models_api_format,
|
||||
};
|
||||
@@ -146,6 +152,10 @@ pub(crate) async fn maybe_build_local_public_support_response(
|
||||
.await;
|
||||
}
|
||||
|
||||
if decision.route_family.as_deref() == Some("install") {
|
||||
return maybe_build_local_install_response(state, request_context).await;
|
||||
}
|
||||
|
||||
if decision.route_family.as_deref() == Some("payment_callback") {
|
||||
return maybe_build_local_payment_callback_response(
|
||||
state,
|
||||
|
||||
561
apps/aether-gateway/src/handlers/public/support/install.rs
Normal file
561
apps/aether-gateway/src/handlers/public/support/install.rs
Normal file
@@ -0,0 +1,561 @@
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
build_auth_error_response, decrypt_catalog_secret_with_fallbacks,
|
||||
resolve_authenticated_local_user, AppState, GatewayPublicRequestContext,
|
||||
};
|
||||
|
||||
const INSTALL_SESSION_TTL_SECS: u64 = 15 * 60;
|
||||
const INSTALL_SESSION_KEY_PREFIX: &str = "install:session:";
|
||||
|
||||
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum InstallTargetCli {
|
||||
ClaudeCode,
|
||||
CodexCli,
|
||||
GeminiCli,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum InstallTargetSystem {
|
||||
Macos,
|
||||
Linux,
|
||||
Windows,
|
||||
Auto,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct UsersMeCreateInstallSessionRequest {
|
||||
target_cli: InstallTargetCli,
|
||||
target_system: InstallTargetSystem,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct StoredInstallSession {
|
||||
api_key_id: String,
|
||||
api_key_name: String,
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
target_cli: InstallTargetCli,
|
||||
target_system: InstallTargetSystem,
|
||||
expires_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
pub(super) fn users_me_api_key_install_sessions_path_matches(request_path: &str) -> bool {
|
||||
users_me_api_key_install_session_id_from_path(request_path).is_some()
|
||||
}
|
||||
|
||||
fn users_me_api_key_install_session_id_from_path(request_path: &str) -> Option<String> {
|
||||
let raw = request_path
|
||||
.strip_prefix("/api/users/me/api-keys/")?
|
||||
.trim()
|
||||
.trim_matches('/');
|
||||
let mut segments = raw.split('/').map(str::trim);
|
||||
let api_key_id = segments.next()?.to_string();
|
||||
let suffix = segments.next()?;
|
||||
(suffix == "install-sessions" && segments.next().is_none()).then_some(api_key_id)
|
||||
}
|
||||
|
||||
fn install_code_from_path(request_path: &str) -> Option<(String, bool)> {
|
||||
let raw = request_path
|
||||
.strip_prefix("/install/")
|
||||
.or_else(|| request_path.strip_prefix("/i/"))?
|
||||
.trim()
|
||||
.trim_matches('/');
|
||||
if raw.is_empty() || raw.contains('/') {
|
||||
return None;
|
||||
}
|
||||
let is_powershell = raw.ends_with(".ps1");
|
||||
let code = raw.strip_suffix(".ps1").unwrap_or(raw).trim();
|
||||
(!code.is_empty()).then(|| (code.to_string(), is_powershell))
|
||||
}
|
||||
|
||||
fn install_session_runtime_key(code: &str) -> String {
|
||||
format!("{INSTALL_SESSION_KEY_PREFIX}{code}")
|
||||
}
|
||||
|
||||
fn generate_install_code() -> String {
|
||||
uuid::Uuid::new_v4()
|
||||
.simple()
|
||||
.to_string()
|
||||
.chars()
|
||||
.take(24)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn unix_secs_now() -> u64 {
|
||||
chrono::Utc::now().timestamp().max(0) as u64
|
||||
}
|
||||
|
||||
fn base_url_from_request(
|
||||
headers: &http::HeaderMap,
|
||||
request_context: &GatewayPublicRequestContext,
|
||||
) -> String {
|
||||
if let Some(value) = std::env::var("AETHER_PUBLIC_BASE_URL")
|
||||
.ok()
|
||||
.or_else(|| std::env::var("PUBLIC_BASE_URL").ok())
|
||||
.map(|value| value.trim().trim_end_matches('/').to_string())
|
||||
.filter(|value| value.starts_with("https://") || value.starts_with("http://"))
|
||||
{
|
||||
return value;
|
||||
}
|
||||
|
||||
let host = crate::headers::header_value_str(headers, "x-forwarded-host")
|
||||
.or_else(|| request_context.host_header.clone())
|
||||
.map(|value| value.trim().trim_end_matches('/').to_string())
|
||||
.filter(|value| {
|
||||
!value.is_empty()
|
||||
&& !value.contains('/')
|
||||
&& !value.contains('\\')
|
||||
&& !value.contains('@')
|
||||
&& !value.contains(char::is_whitespace)
|
||||
})
|
||||
.unwrap_or_else(|| "localhost".to_string());
|
||||
let proto = crate::headers::header_value_str(headers, "x-forwarded-proto")
|
||||
.map(|value| value.trim().trim_end_matches(':').to_ascii_lowercase())
|
||||
.filter(|value| value == "http" || value == "https")
|
||||
.unwrap_or_else(|| "http".to_string());
|
||||
format!("{proto}://{host}")
|
||||
}
|
||||
|
||||
fn shell_single_quote(value: &str) -> String {
|
||||
format!("'{}'", value.replace('\'', "'\\''"))
|
||||
}
|
||||
|
||||
fn powershell_single_quote(value: &str) -> String {
|
||||
format!("'{}'", value.replace('\'', "''"))
|
||||
}
|
||||
|
||||
fn cli_label(target_cli: InstallTargetCli) -> &'static str {
|
||||
match target_cli {
|
||||
InstallTargetCli::ClaudeCode => "Claude Code",
|
||||
InstallTargetCli::CodexCli => "Codex CLI",
|
||||
InstallTargetCli::GeminiCli => "Gemini CLI",
|
||||
}
|
||||
}
|
||||
|
||||
fn system_label(target_system: InstallTargetSystem) -> &'static str {
|
||||
match target_system {
|
||||
InstallTargetSystem::Macos => "macOS",
|
||||
InstallTargetSystem::Linux => "Linux",
|
||||
InstallTargetSystem::Windows => "Windows",
|
||||
InstallTargetSystem::Auto => "Auto",
|
||||
}
|
||||
}
|
||||
|
||||
fn npm_package(target_cli: InstallTargetCli) -> &'static str {
|
||||
match target_cli {
|
||||
InstallTargetCli::ClaudeCode => "@anthropic-ai/claude-code",
|
||||
InstallTargetCli::CodexCli => "@openai/codex",
|
||||
InstallTargetCli::GeminiCli => "@google/gemini-cli",
|
||||
}
|
||||
}
|
||||
|
||||
fn cli_binary(target_cli: InstallTargetCli) -> &'static str {
|
||||
match target_cli {
|
||||
InstallTargetCli::ClaudeCode => "claude",
|
||||
InstallTargetCli::CodexCli => "codex",
|
||||
InstallTargetCli::GeminiCli => "gemini",
|
||||
}
|
||||
}
|
||||
|
||||
fn build_unix_script(session: &StoredInstallSession) -> String {
|
||||
let target_cli = match session.target_cli {
|
||||
InstallTargetCli::ClaudeCode => "claude_code",
|
||||
InstallTargetCli::CodexCli => "codex_cli",
|
||||
InstallTargetCli::GeminiCli => "gemini_cli",
|
||||
};
|
||||
let target_system = match session.target_system {
|
||||
InstallTargetSystem::Macos => "macos",
|
||||
InstallTargetSystem::Linux => "linux",
|
||||
InstallTargetSystem::Windows => "windows",
|
||||
InstallTargetSystem::Auto => "auto",
|
||||
};
|
||||
|
||||
format!(
|
||||
r###"#!/bin/sh
|
||||
set -eu
|
||||
TARGET_CLI={target_cli}
|
||||
TARGET_SYSTEM={target_system}
|
||||
AETHER_BASE_URL={base_url}
|
||||
AETHER_API_KEY={api_key}
|
||||
CLI_LABEL={label}
|
||||
CLI_BIN={binary}
|
||||
NPM_PACKAGE={npm_package}
|
||||
|
||||
say() {{ printf '%s\n' "[Aether] $1"; }}
|
||||
fail() {{ printf '%s\n' "[Aether] $1" >&2; exit 1; }}
|
||||
|
||||
os="$(uname -s 2>/dev/null || printf unknown)"
|
||||
case "$os" in
|
||||
Darwin) actual_system=macos ;;
|
||||
Linux) actual_system=linux ;;
|
||||
MINGW*|MSYS*|CYGWIN*) fail "检测到 Windows shell,请在 PowerShell 中使用 Windows 命令:irm <url>.ps1 | iex" ;;
|
||||
*) fail "不支持的系统:$os" ;;
|
||||
esac
|
||||
|
||||
if [ "$TARGET_SYSTEM" = "windows" ]; then
|
||||
fail "该 install code 绑定 Windows,请复制 PowerShell 命令执行。"
|
||||
fi
|
||||
if [ "$TARGET_SYSTEM" != "auto" ] && [ "$TARGET_SYSTEM" != "$actual_system" ]; then
|
||||
fail "所选系统 $TARGET_SYSTEM 与当前系统 $actual_system 不一致,请回到 Aether 重新选择目标系统。"
|
||||
fi
|
||||
|
||||
say "准备安装/复用 $CLI_LABEL"
|
||||
if ! command -v "$CLI_BIN" >/dev/null 2>&1; then
|
||||
command -v npm >/dev/null 2>&1 || fail "未找到 $CLI_BIN,也未找到 npm。请先安装 Node.js/npm 后重试。"
|
||||
say "未找到 $CLI_BIN,正在通过 npm 安装 $NPM_PACKAGE"
|
||||
npm install -g "$NPM_PACKAGE"
|
||||
fi
|
||||
|
||||
umask 077
|
||||
mkdir -p "$HOME/.aether"
|
||||
cat > "$HOME/.aether/client.env" <<EOF
|
||||
AETHER_BASE_URL=$AETHER_BASE_URL
|
||||
AETHER_API_KEY=$AETHER_API_KEY
|
||||
EOF
|
||||
chmod 600 "$HOME/.aether/client.env" 2>/dev/null || true
|
||||
|
||||
case "$TARGET_CLI" in
|
||||
claude_code)
|
||||
mkdir -p "$HOME/.claude"
|
||||
python3 - "$HOME/.claude/settings.json" "$AETHER_BASE_URL" "$AETHER_API_KEY" <<'PY'
|
||||
import json, pathlib, sys
|
||||
path = pathlib.Path(sys.argv[1])
|
||||
data = json.loads(path.read_text() or '{{}}') if path.exists() else {{}}
|
||||
env = data.setdefault('env', {{}})
|
||||
env['ANTHROPIC_BASE_URL'] = sys.argv[2]
|
||||
env['ANTHROPIC_AUTH_TOKEN'] = sys.argv[3]
|
||||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + '\n')
|
||||
PY
|
||||
chmod 600 "$HOME/.claude/settings.json" 2>/dev/null || true
|
||||
;;
|
||||
codex_cli)
|
||||
mkdir -p "$HOME/.codex"
|
||||
cat > "$HOME/.codex/auth.json" <<EOF
|
||||
{{"OPENAI_API_KEY":"$AETHER_API_KEY"}}
|
||||
EOF
|
||||
cat > "$HOME/.codex/config.toml" <<EOF
|
||||
# Managed by Aether
|
||||
model_provider = "aether"
|
||||
|
||||
[model_providers.aether]
|
||||
name = "Aether"
|
||||
base_url = "$AETHER_BASE_URL/v1"
|
||||
env_key = "OPENAI_API_KEY"
|
||||
wire_api = "chat"
|
||||
EOF
|
||||
chmod 600 "$HOME/.codex/auth.json" "$HOME/.codex/config.toml" 2>/dev/null || true
|
||||
;;
|
||||
gemini_cli)
|
||||
mkdir -p "$HOME/.gemini"
|
||||
cat > "$HOME/.gemini/.env" <<EOF
|
||||
GEMINI_API_KEY=$AETHER_API_KEY
|
||||
GOOGLE_API_KEY=$AETHER_API_KEY
|
||||
GOOGLE_GEMINI_BASE_URL=$AETHER_BASE_URL
|
||||
AETHER_BASE_URL=$AETHER_BASE_URL
|
||||
EOF
|
||||
python3 - "$HOME/.gemini/settings.json" "$AETHER_BASE_URL" <<'PY'
|
||||
import json, pathlib, sys
|
||||
path = pathlib.Path(sys.argv[1])
|
||||
data = json.loads(path.read_text() or '{{}}') if path.exists() else {{}}
|
||||
data.setdefault('aether', {{}})['baseUrl'] = sys.argv[2]
|
||||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + '\n')
|
||||
PY
|
||||
chmod 600 "$HOME/.gemini/.env" "$HOME/.gemini/settings.json" 2>/dev/null || true
|
||||
;;
|
||||
esac
|
||||
|
||||
say "$CLI_LABEL 已配置到 Aether。执行 $CLI_BIN --version 验证安装。"
|
||||
"###,
|
||||
target_cli = target_cli,
|
||||
target_system = target_system,
|
||||
base_url = shell_single_quote(&session.base_url),
|
||||
api_key = shell_single_quote(&session.api_key),
|
||||
label = shell_single_quote(cli_label(session.target_cli)),
|
||||
binary = shell_single_quote(cli_binary(session.target_cli)),
|
||||
npm_package = shell_single_quote(npm_package(session.target_cli)),
|
||||
)
|
||||
}
|
||||
|
||||
fn build_powershell_script(session: &StoredInstallSession) -> String {
|
||||
let target_cli = match session.target_cli {
|
||||
InstallTargetCli::ClaudeCode => "claude_code",
|
||||
InstallTargetCli::CodexCli => "codex_cli",
|
||||
InstallTargetCli::GeminiCli => "gemini_cli",
|
||||
};
|
||||
format!(
|
||||
r###"$ErrorActionPreference = 'Stop'
|
||||
$TargetCli = {target_cli}
|
||||
$TargetSystem = {target_system}
|
||||
$AetherBaseUrl = {base_url}
|
||||
$AetherApiKey = {api_key}
|
||||
$CliLabel = {label}
|
||||
$CliBin = {binary}
|
||||
$NpmPackage = {npm_package}
|
||||
|
||||
function Say($Message) {{ Write-Host "[Aether] $Message" }}
|
||||
function Fail($Message) {{ Write-Error "[Aether] $Message"; exit 1 }}
|
||||
|
||||
if ($TargetSystem -ne 'auto' -and $TargetSystem -ne 'windows') {{ Fail "该 install code 绑定 $TargetSystem,请复制 macOS/Linux 命令执行。" }}
|
||||
|
||||
Say "准备安装/复用 $CliLabel"
|
||||
if (-not (Get-Command $CliBin -ErrorAction SilentlyContinue)) {{
|
||||
if (-not (Get-Command npm -ErrorAction SilentlyContinue)) {{ Fail "未找到 $CliBin,也未找到 npm。请先安装 Node.js/npm 后重试。" }}
|
||||
Say "未找到 $CliBin,正在通过 npm 安装 $NpmPackage"
|
||||
npm install -g $NpmPackage
|
||||
}}
|
||||
|
||||
$HomeDir = [Environment]::GetFolderPath('UserProfile')
|
||||
$AetherDir = Join-Path $HomeDir '.aether'
|
||||
New-Item -ItemType Directory -Force -Path $AetherDir | Out-Null
|
||||
Set-Content -Path (Join-Path $AetherDir 'client.env') -Value "AETHER_BASE_URL=$AetherBaseUrl`nAETHER_API_KEY=$AetherApiKey`n" -Encoding UTF8
|
||||
|
||||
if ($TargetCli -eq 'claude_code') {{
|
||||
$Dir = Join-Path $HomeDir '.claude'; New-Item -ItemType Directory -Force -Path $Dir | Out-Null
|
||||
$Path = Join-Path $Dir 'settings.json'
|
||||
$Data = if (Test-Path $Path) {{ Get-Content $Path -Raw | ConvertFrom-Json -AsHashtable }} else {{ @{{}} }}
|
||||
if (-not $Data.ContainsKey('env')) {{ $Data.env = @{{}} }}
|
||||
$Data.env.ANTHROPIC_BASE_URL = $AetherBaseUrl
|
||||
$Data.env.ANTHROPIC_AUTH_TOKEN = $AetherApiKey
|
||||
$Data | ConvertTo-Json -Depth 8 | Set-Content $Path -Encoding UTF8
|
||||
}} elseif ($TargetCli -eq 'codex_cli') {{
|
||||
$Dir = Join-Path $HomeDir '.codex'; New-Item -ItemType Directory -Force -Path $Dir | Out-Null
|
||||
Set-Content (Join-Path $Dir 'auth.json') -Value (@{{ OPENAI_API_KEY = $AetherApiKey }} | ConvertTo-Json) -Encoding UTF8
|
||||
Set-Content (Join-Path $Dir 'config.toml') -Value "# Managed by Aether`nmodel_provider = \"aether\"`n`n[model_providers.aether]`nname = \"Aether\"`nbase_url = \"$AetherBaseUrl/v1\"`nenv_key = \"OPENAI_API_KEY\"`nwire_api = \"chat\"`n" -Encoding UTF8
|
||||
}} elseif ($TargetCli -eq 'gemini_cli') {{
|
||||
$Dir = Join-Path $HomeDir '.gemini'; New-Item -ItemType Directory -Force -Path $Dir | Out-Null
|
||||
Set-Content (Join-Path $Dir '.env') -Value "GEMINI_API_KEY=$AetherApiKey`nGOOGLE_API_KEY=$AetherApiKey`nGOOGLE_GEMINI_BASE_URL=$AetherBaseUrl`nAETHER_BASE_URL=$AetherBaseUrl`n" -Encoding UTF8
|
||||
$Path = Join-Path $Dir 'settings.json'
|
||||
$Data = if (Test-Path $Path) {{ Get-Content $Path -Raw | ConvertFrom-Json -AsHashtable }} else {{ @{{}} }}
|
||||
$Data.aether = @{{ baseUrl = $AetherBaseUrl }}
|
||||
$Data | ConvertTo-Json -Depth 8 | Set-Content $Path -Encoding UTF8
|
||||
}}
|
||||
|
||||
Say "$CliLabel 已配置到 Aether。执行 $CliBin --version 验证安装。"
|
||||
"###,
|
||||
target_cli = powershell_single_quote(target_cli),
|
||||
target_system = powershell_single_quote(match session.target_system {
|
||||
InstallTargetSystem::Macos => "macos",
|
||||
InstallTargetSystem::Linux => "linux",
|
||||
InstallTargetSystem::Windows => "windows",
|
||||
InstallTargetSystem::Auto => "auto",
|
||||
}),
|
||||
base_url = powershell_single_quote(&session.base_url),
|
||||
api_key = powershell_single_quote(&session.api_key),
|
||||
label = powershell_single_quote(cli_label(session.target_cli)),
|
||||
binary = powershell_single_quote(cli_binary(session.target_cli)),
|
||||
npm_package = powershell_single_quote(npm_package(session.target_cli)),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) async fn handle_users_me_api_key_install_session_create(
|
||||
state: &AppState,
|
||||
request_context: &GatewayPublicRequestContext,
|
||||
headers: &http::HeaderMap,
|
||||
request_body: Option<&axum::body::Bytes>,
|
||||
) -> Response<Body> {
|
||||
let auth = match resolve_authenticated_local_user(state, request_context, headers).await {
|
||||
Ok(value) => value,
|
||||
Err(response) => return response,
|
||||
};
|
||||
let Some(api_key_id) =
|
||||
users_me_api_key_install_session_id_from_path(&request_context.request_path)
|
||||
else {
|
||||
return build_auth_error_response(http::StatusCode::NOT_FOUND, "API密钥不存在", false);
|
||||
};
|
||||
let Some(request_body) = request_body else {
|
||||
return build_auth_error_response(http::StatusCode::BAD_REQUEST, "请求数据验证失败", false);
|
||||
};
|
||||
let payload = match serde_json::from_slice::<UsersMeCreateInstallSessionRequest>(request_body) {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"请求数据验证失败",
|
||||
false,
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
let records = match state
|
||||
.list_auth_api_key_export_records_by_user_ids(std::slice::from_ref(&auth.user.id))
|
||||
.await
|
||||
{
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("user api key lookup failed: {err:?}"),
|
||||
false,
|
||||
)
|
||||
}
|
||||
};
|
||||
let Some(record) = records
|
||||
.into_iter()
|
||||
.find(|record| !record.is_standalone && record.api_key_id == api_key_id)
|
||||
else {
|
||||
return build_auth_error_response(http::StatusCode::NOT_FOUND, "API密钥不存在", false);
|
||||
};
|
||||
let Some(ciphertext) = record
|
||||
.key_encrypted
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"该密钥没有存储完整密钥信息",
|
||||
false,
|
||||
);
|
||||
};
|
||||
let Some(api_key) = decrypt_catalog_secret_with_fallbacks(state.encryption_key(), ciphertext)
|
||||
else {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"解密密钥失败",
|
||||
false,
|
||||
);
|
||||
};
|
||||
|
||||
let code = generate_install_code();
|
||||
let expires_at_unix_secs = unix_secs_now().saturating_add(INSTALL_SESSION_TTL_SECS);
|
||||
let session = StoredInstallSession {
|
||||
api_key_id: record.api_key_id.clone(),
|
||||
api_key_name: record.name.unwrap_or_else(|| "API Key".to_string()),
|
||||
api_key,
|
||||
base_url: base_url_from_request(headers, request_context),
|
||||
target_cli: payload.target_cli,
|
||||
target_system: payload.target_system,
|
||||
expires_at_unix_secs,
|
||||
};
|
||||
let serialized = match serde_json::to_string(&session) {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("install session serialize failed: {err:?}"),
|
||||
false,
|
||||
)
|
||||
}
|
||||
};
|
||||
if let Err(err) = state
|
||||
.runtime_kv_setex(
|
||||
&install_session_runtime_key(&code),
|
||||
&serialized,
|
||||
INSTALL_SESSION_TTL_SECS,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("install session create failed: {err:?}"),
|
||||
false,
|
||||
);
|
||||
}
|
||||
|
||||
let base_url = session.base_url.trim_end_matches('/');
|
||||
Json(json!({
|
||||
"install_code": code,
|
||||
"expires_at_unix_secs": expires_at_unix_secs,
|
||||
"expires_in_seconds": INSTALL_SESSION_TTL_SECS,
|
||||
"target_cli": session.target_cli,
|
||||
"target_cli_label": cli_label(session.target_cli),
|
||||
"target_system": session.target_system,
|
||||
"target_system_label": system_label(session.target_system),
|
||||
"unix_command": format!("curl -fsSL {base_url}/install/{code} | sh"),
|
||||
"powershell_command": format!("irm {base_url}/install/{code}.ps1 | iex"),
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
pub(super) async fn maybe_build_local_install_response(
|
||||
state: &AppState,
|
||||
request_context: &GatewayPublicRequestContext,
|
||||
) -> Option<Response<Body>> {
|
||||
let decision = request_context.control_decision.as_ref()?;
|
||||
if decision.route_family.as_deref() != Some("install") {
|
||||
return None;
|
||||
}
|
||||
let Some((code, wants_powershell)) = install_code_from_path(&request_context.request_path)
|
||||
else {
|
||||
return Some(build_auth_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"install code 不存在或已失效",
|
||||
false,
|
||||
));
|
||||
};
|
||||
let raw = match state
|
||||
.runtime_kv_getdel(&install_session_runtime_key(&code))
|
||||
.await
|
||||
{
|
||||
Ok(Some(value)) => value,
|
||||
Ok(None) => {
|
||||
return Some(build_auth_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"install code 不存在、已过期或已使用",
|
||||
false,
|
||||
))
|
||||
}
|
||||
Err(err) => {
|
||||
return Some(build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("install session lookup failed: {err:?}"),
|
||||
false,
|
||||
))
|
||||
}
|
||||
};
|
||||
let session = match serde_json::from_str::<StoredInstallSession>(&raw) {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
return Some(build_auth_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"install code 数据无效",
|
||||
false,
|
||||
))
|
||||
}
|
||||
};
|
||||
if session.expires_at_unix_secs <= unix_secs_now() {
|
||||
return Some(build_auth_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"install code 已过期",
|
||||
false,
|
||||
));
|
||||
}
|
||||
let body = if wants_powershell {
|
||||
build_powershell_script(&session)
|
||||
} else {
|
||||
build_unix_script(&session)
|
||||
};
|
||||
let content_type = if wants_powershell {
|
||||
"text/plain; charset=utf-8"
|
||||
} else {
|
||||
"text/x-shellscript; charset=utf-8"
|
||||
};
|
||||
let mut response = Response::new(Body::from(body));
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
http::HeaderValue::from_static(content_type),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
http::header::CACHE_CONTROL,
|
||||
http::HeaderValue::from_static("no-store"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
http::header::PRAGMA,
|
||||
http::HeaderValue::from_static("no-cache"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
http::header::HeaderName::from_static("x-content-type-options"),
|
||||
http::HeaderValue::from_static("nosniff"),
|
||||
);
|
||||
Some(response)
|
||||
}
|
||||
@@ -1,9 +1,10 @@
|
||||
use super::{
|
||||
auth_password_policy_level, build_auth_error_response, build_auth_wallet_summary_payload,
|
||||
decrypt_catalog_secret_with_fallbacks, encrypt_catalog_secret_with_fallbacks, handle_auth_me,
|
||||
query_param_optional_bool, query_param_value, resolve_authenticated_local_user,
|
||||
unix_secs_to_rfc3339, validate_auth_register_password, AppState, AuthenticatedLocalUserContext,
|
||||
GatewayPublicRequestContext, PUBLIC_CAPABILITY_DEFINITIONS,
|
||||
handle_users_me_api_key_install_session_create, query_param_optional_bool, query_param_value,
|
||||
resolve_authenticated_local_user, unix_secs_to_rfc3339,
|
||||
users_me_api_key_install_sessions_path_matches, validate_auth_register_password, AppState,
|
||||
AuthenticatedLocalUserContext, GatewayPublicRequestContext, PUBLIC_CAPABILITY_DEFINITIONS,
|
||||
};
|
||||
use crate::admin_api::build_admin_endpoint_health_status_payload;
|
||||
use crate::handlers::internal::build_management_token_payload;
|
||||
|
||||
@@ -4,8 +4,9 @@ use axum::{body::Body, http, response::Response};
|
||||
use super::{
|
||||
handle_auth_me, handle_users_me_api_key_capabilities_put, handle_users_me_api_key_create,
|
||||
handle_users_me_api_key_delete, handle_users_me_api_key_detail_get,
|
||||
handle_users_me_api_key_patch, handle_users_me_api_key_providers_put,
|
||||
handle_users_me_api_key_update, handle_users_me_api_keys_get, handle_users_me_available_models,
|
||||
handle_users_me_api_key_install_session_create, handle_users_me_api_key_patch,
|
||||
handle_users_me_api_key_providers_put, handle_users_me_api_key_update,
|
||||
handle_users_me_api_keys_get, handle_users_me_available_models,
|
||||
handle_users_me_delete_other_sessions, handle_users_me_delete_session,
|
||||
handle_users_me_detail_put, handle_users_me_endpoint_status_get,
|
||||
handle_users_me_management_token_create, handle_users_me_management_token_delete,
|
||||
@@ -17,8 +18,8 @@ use super::{
|
||||
handle_users_me_providers_get, handle_users_me_sessions_get, handle_users_me_update_session,
|
||||
handle_users_me_usage_active_get, handle_users_me_usage_get, handle_users_me_usage_heatmap_get,
|
||||
handle_users_me_usage_interval_timeline_get, users_me_api_key_capabilities_path_matches,
|
||||
users_me_api_key_detail_path_matches, users_me_api_key_providers_path_matches,
|
||||
users_me_management_token_detail_path_matches,
|
||||
users_me_api_key_detail_path_matches, users_me_api_key_install_sessions_path_matches,
|
||||
users_me_api_key_providers_path_matches, users_me_management_token_detail_path_matches,
|
||||
users_me_management_token_regenerate_path_matches,
|
||||
users_me_management_token_toggle_path_matches, users_me_management_tokens_root,
|
||||
users_me_session_detail_path_matches, AppState, GatewayPublicRequestContext,
|
||||
@@ -167,6 +168,19 @@ pub(crate) async fn maybe_build_local_users_me_response(
|
||||
.await,
|
||||
)
|
||||
}
|
||||
Some("api_key_install_session_create")
|
||||
if users_me_api_key_install_sessions_path_matches(&request_context.request_path) =>
|
||||
{
|
||||
Some(
|
||||
handle_users_me_api_key_install_session_create(
|
||||
state,
|
||||
request_context,
|
||||
headers,
|
||||
request_body,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
Some("management_token_regenerate")
|
||||
if users_me_management_token_regenerate_path_matches(&request_context.request_path) =>
|
||||
{
|
||||
|
||||
@@ -443,7 +443,11 @@ pub(crate) fn public_support_local_requires_buffered_body(
|
||||
) | (
|
||||
Some("users_me"),
|
||||
http::Method::POST,
|
||||
Some("api_keys_create" | "management_tokens_create"),
|
||||
Some(
|
||||
"api_keys_create"
|
||||
| "api_key_install_session_create"
|
||||
| "management_tokens_create",
|
||||
),
|
||||
) | (
|
||||
Some("wallet"),
|
||||
http::Method::POST,
|
||||
|
||||
@@ -58,6 +58,7 @@ mod router;
|
||||
mod scheduler;
|
||||
mod state;
|
||||
mod system_features;
|
||||
mod task_runtime;
|
||||
mod tunnel;
|
||||
mod usage;
|
||||
mod video_tasks;
|
||||
|
||||
@@ -1202,13 +1202,13 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
state.bootstrap_admin_from_env().await?;
|
||||
|
||||
let background_tasks = if args.node_role.spawns_background_tasks() {
|
||||
state.spawn_background_tasks()
|
||||
Some(state.spawn_background_tasks())
|
||||
} else {
|
||||
info!(
|
||||
node_role = args.node_role.as_str(),
|
||||
"background workers disabled for this node role"
|
||||
);
|
||||
Vec::new()
|
||||
None
|
||||
};
|
||||
let listener = tokio::net::TcpListener::bind(bind_addr).await?;
|
||||
let public_base_url = resolve_local_http_base_url(app_port)?;
|
||||
@@ -1241,8 +1241,8 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
router.into_make_service_with_connect_info::<std::net::SocketAddr>(),
|
||||
)
|
||||
.await?;
|
||||
for handle in background_tasks {
|
||||
handle.abort();
|
||||
if let Some(background_tasks) = background_tasks {
|
||||
background_tasks.shutdown().await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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::{
|
||||
|
||||
518
apps/aether-gateway/src/maintenance/runtime/cleanup_runs.rs
Normal file
518
apps/aether-gateway/src/maintenance/runtime/cleanup_runs.rs
Normal 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()
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -119,9 +119,7 @@ async fn gateway_background_request_candidate_cleanup_deletes_expired_entries_in
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
for handle in background_tasks {
|
||||
handle.abort();
|
||||
}
|
||||
background_tasks.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -539,8 +539,6 @@ async fn gateway_background_model_fetch_updates_key_and_syncs_provider_model_whi
|
||||
.expect("execution runtime plan should be captured");
|
||||
assert_eq!(seen_plan.url, "https://api.openai.example/v1/models");
|
||||
|
||||
for handle in background_tasks {
|
||||
handle.abort();
|
||||
}
|
||||
background_tasks.shutdown().await;
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use serde_json::{json, Value};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::AppState;
|
||||
|
||||
@@ -120,3 +120,66 @@ pub(crate) fn with_error_flow_report_context(
|
||||
object.insert("error_flow".to_string(), error_flow);
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
|
||||
pub(crate) fn with_upstream_response_report_context(
|
||||
report_context: Option<&Value>,
|
||||
status_code: u16,
|
||||
headers: Option<&std::collections::BTreeMap<String, String>>,
|
||||
body: Option<&Value>,
|
||||
body_ref: Option<&str>,
|
||||
body_state: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let mut object = report_context?.as_object()?.clone();
|
||||
let mut upstream_response = serde_json::Map::new();
|
||||
upstream_response.insert("status_code".to_string(), json!(status_code));
|
||||
if let Some(headers) = headers {
|
||||
upstream_response.insert("headers".to_string(), trace_headers_to_json(headers));
|
||||
}
|
||||
if let Some(body) = body {
|
||||
upstream_response.insert("body".to_string(), body.clone());
|
||||
}
|
||||
if let Some(body_ref) = body_ref {
|
||||
upstream_response.insert("body_ref".to_string(), json!(body_ref));
|
||||
}
|
||||
if let Some(body_state) = body_state {
|
||||
upstream_response.insert("body_state".to_string(), json!(body_state));
|
||||
}
|
||||
object.insert(
|
||||
"upstream_response".to_string(),
|
||||
Value::Object(upstream_response),
|
||||
);
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
|
||||
fn trace_headers_to_json(headers: &std::collections::BTreeMap<String, String>) -> Value {
|
||||
Value::Object(Map::from_iter(headers.iter().map(|(key, value)| {
|
||||
(
|
||||
key.clone(),
|
||||
Value::String(mask_trace_header_value(key, value)),
|
||||
)
|
||||
})))
|
||||
}
|
||||
|
||||
fn mask_trace_header_value(name: &str, value: &str) -> String {
|
||||
if !trace_header_is_sensitive(name) {
|
||||
return value.to_string();
|
||||
}
|
||||
if value.len() <= 8 {
|
||||
return "****".to_string();
|
||||
}
|
||||
format!("{}****{}", &value[..4], &value[value.len() - 4..])
|
||||
}
|
||||
|
||||
fn trace_header_is_sensitive(name: &str) -> bool {
|
||||
[
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"api-key",
|
||||
"x-goog-api-key",
|
||||
"cookie",
|
||||
"set-cookie",
|
||||
"proxy-authorization",
|
||||
]
|
||||
.iter()
|
||||
.any(|candidate| name.trim().eq_ignore_ascii_case(candidate))
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
80
apps/aether-gateway/src/state/runtime/background_tasks.rs
Normal file
80
apps/aether-gateway/src/state/runtime/background_tasks.rs
Normal file
@@ -0,0 +1,80 @@
|
||||
use crate::{AppState, GatewayError};
|
||||
use aether_data_contracts::repository::background_tasks::{
|
||||
BackgroundTaskListQuery, BackgroundTaskSummary, StoredBackgroundTaskEvent,
|
||||
StoredBackgroundTaskRun, StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent,
|
||||
UpsertBackgroundTaskRun,
|
||||
};
|
||||
|
||||
impl AppState {
|
||||
pub(crate) async fn find_background_task_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, GatewayError> {
|
||||
self.data
|
||||
.find_background_task_run(run_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_background_task_runs(
|
||||
&self,
|
||||
query: &BackgroundTaskListQuery,
|
||||
) -> Result<StoredBackgroundTaskRunPage, GatewayError> {
|
||||
self.data
|
||||
.list_background_task_runs(query)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_background_task_events(
|
||||
&self,
|
||||
run_id: &str,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredBackgroundTaskEvent>, GatewayError> {
|
||||
self.data
|
||||
.list_background_task_events(run_id, offset, limit)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_background_task_runs(
|
||||
&self,
|
||||
) -> Result<BackgroundTaskSummary, GatewayError> {
|
||||
self.data
|
||||
.summarize_background_task_runs()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_background_task_run(
|
||||
&self,
|
||||
run: UpsertBackgroundTaskRun,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, GatewayError> {
|
||||
self.data
|
||||
.upsert_background_task_run(run)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn request_cancel_background_task_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.data
|
||||
.request_cancel_background_task_run(run_id, updated_at_unix_secs)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_background_task_event(
|
||||
&self,
|
||||
event: UpsertBackgroundTaskEvent,
|
||||
) -> Result<Option<StoredBackgroundTaskEvent>, GatewayError> {
|
||||
self.data
|
||||
.upsert_background_task_event(event)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
}
|
||||
@@ -10,6 +10,7 @@ use super::{
|
||||
mod announcements;
|
||||
mod api_key_exports;
|
||||
mod auth;
|
||||
mod background_tasks;
|
||||
mod billing;
|
||||
mod candidate_queries;
|
||||
mod gemini_files;
|
||||
@@ -29,6 +30,14 @@ impl AppState {
|
||||
self.data.has_announcement_writer()
|
||||
}
|
||||
|
||||
pub fn has_background_task_data_reader(&self) -> bool {
|
||||
self.data.has_background_task_reader()
|
||||
}
|
||||
|
||||
pub fn has_background_task_data_writer(&self) -> bool {
|
||||
self.data.has_background_task_writer()
|
||||
}
|
||||
|
||||
pub fn has_video_task_data_reader(&self) -> bool {
|
||||
self.data.has_video_task_reader()
|
||||
}
|
||||
|
||||
647
apps/aether-gateway/src/task_runtime/mod.rs
Normal file
647
apps/aether-gateway/src/task_runtime/mod.rs
Normal file
@@ -0,0 +1,647 @@
|
||||
use std::future::Future;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_data_contracts::repository::background_tasks::{
|
||||
BackgroundTaskKind, BackgroundTaskStatus, StoredBackgroundTaskRun, UpsertBackgroundTaskEvent,
|
||||
UpsertBackgroundTaskRun,
|
||||
};
|
||||
use aether_runtime::task::spawn_named;
|
||||
pub(crate) use aether_task_runtime::TaskSupervisor;
|
||||
use aether_task_runtime::{RetryPolicy, TaskDefinition, TaskKind};
|
||||
use serde_json::Value;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(crate) const TASK_KEY_PROVIDER_DELETE: &str = "admin.provider.delete";
|
||||
pub(crate) const TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT: &str = "admin.provider.oauth.batch_import";
|
||||
pub(crate) const TASK_KEY_USAGE_QUEUE_WORKER: &str = "usage.queue.worker";
|
||||
pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
|
||||
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
|
||||
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
|
||||
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
|
||||
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
|
||||
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
|
||||
pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
|
||||
pub(crate) const TASK_KEY_PENDING_CLEANUP: &str = "maintenance.pending.cleanup";
|
||||
pub(crate) const TASK_KEY_REQUEST_CANDIDATE_CLEANUP: &str = "maintenance.request.candidate.cleanup";
|
||||
pub(crate) const TASK_KEY_GEMINI_FILES_CLEANUP: &str = "maintenance.gemini.files.cleanup";
|
||||
pub(crate) const TASK_KEY_OAUTH_TOKEN_REFRESH: &str = "maintenance.oauth.token.refresh";
|
||||
pub(crate) const TASK_KEY_PROXY_NODE_STALE_CLEANUP: &str = "maintenance.proxy.node.stale.cleanup";
|
||||
pub(crate) const TASK_KEY_PROXY_NODE_METRICS_CLEANUP: &str =
|
||||
"maintenance.proxy.node.metrics.cleanup";
|
||||
pub(crate) const TASK_KEY_PROXY_UPGRADE_ROLLOUT: &str = "maintenance.proxy.upgrade.rollout";
|
||||
pub(crate) const TASK_KEY_PROVIDER_CHECKIN: &str = "maintenance.provider.checkin";
|
||||
pub(crate) const TASK_KEY_USAGE_CLEANUP: &str = "maintenance.usage.cleanup";
|
||||
pub(crate) const TASK_KEY_WALLET_DAILY_USAGE_AGG: &str = "maintenance.wallet.daily.usage.agg";
|
||||
pub(crate) const TASK_KEY_STATS_DAILY_AGG: &str = "maintenance.stats.daily.agg";
|
||||
pub(crate) const TASK_KEY_STATS_HOURLY_AGG: &str = "maintenance.stats.hourly.agg";
|
||||
pub(crate) const TASK_KEY_USAGE_SYNC_REPORT: &str = "usage.sync.report";
|
||||
pub(crate) const TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH: &str = "provider.oauth.account.refresh";
|
||||
pub(crate) const TASK_KEY_PROVIDER_BALANCE_REFRESH: &str = "provider.ops.balance.refresh";
|
||||
|
||||
const RETRY_ONCE: RetryPolicy = RetryPolicy { max_attempts: 1 };
|
||||
|
||||
const TASK_DEFINITIONS: &[TaskDefinition] = &[
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROVIDER_DELETE,
|
||||
TaskKind::OnDemand,
|
||||
"manual",
|
||||
false,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT,
|
||||
TaskKind::OnDemand,
|
||||
"manual",
|
||||
false,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_USAGE_QUEUE_WORKER,
|
||||
TaskKind::Daemon,
|
||||
"daemon",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_VIDEO_TASK_POLLER,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_MODEL_FETCH_WORKER,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROVIDER_QUOTA_RESET,
|
||||
TaskKind::Scheduled,
|
||||
"daily",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_POOL_QUOTA_PROBE,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_POOL_MONITOR,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_AUDIT_CLEANUP,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_DB_MAINTENANCE,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PENDING_CLEANUP,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_REQUEST_CANDIDATE_CLEANUP,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_GEMINI_FILES_CLEANUP,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_OAUTH_TOKEN_REFRESH,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROXY_NODE_STALE_CLEANUP,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROXY_NODE_METRICS_CLEANUP,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROXY_UPGRADE_ROLLOUT,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROVIDER_CHECKIN,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_USAGE_CLEANUP,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_WALLET_DAILY_USAGE_AGG,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_STATS_DAILY_AGG,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_STATS_HOURLY_AGG,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_USAGE_SYNC_REPORT,
|
||||
TaskKind::FireAndForget,
|
||||
"internal",
|
||||
false,
|
||||
false,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH,
|
||||
TaskKind::FireAndForget,
|
||||
"internal",
|
||||
false,
|
||||
false,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROVIDER_BALANCE_REFRESH,
|
||||
TaskKind::FireAndForget,
|
||||
"internal",
|
||||
false,
|
||||
false,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
];
|
||||
|
||||
pub(crate) fn task_definitions() -> &'static [TaskDefinition] {
|
||||
TASK_DEFINITIONS
|
||||
}
|
||||
|
||||
pub(crate) fn task_definition(task_key: &str) -> Option<TaskDefinition> {
|
||||
task_definitions()
|
||||
.iter()
|
||||
.copied()
|
||||
.find(|definition| definition.key == task_key)
|
||||
}
|
||||
|
||||
pub(crate) const fn background_task_kind(kind: TaskKind) -> BackgroundTaskKind {
|
||||
match kind {
|
||||
TaskKind::Scheduled => BackgroundTaskKind::Scheduled,
|
||||
TaskKind::Daemon => BackgroundTaskKind::Daemon,
|
||||
TaskKind::OnDemand => BackgroundTaskKind::OnDemand,
|
||||
TaskKind::FireAndForget => BackgroundTaskKind::FireAndForget,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn task_cancel_kv_key(run_id: &str) -> String {
|
||||
format!("task_runtime:run:{run_id}:cancel")
|
||||
}
|
||||
|
||||
pub(crate) fn now_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub(crate) fn build_task_run_id() -> String {
|
||||
Uuid::new_v4().to_string()
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_fire_and_forget<F>(task_name: &'static str, future: F) -> JoinHandle<()>
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
spawn_named(task_name, future)
|
||||
}
|
||||
|
||||
fn stored_run_to_upsert(run: StoredBackgroundTaskRun) -> UpsertBackgroundTaskRun {
|
||||
UpsertBackgroundTaskRun {
|
||||
id: run.id,
|
||||
task_key: run.task_key,
|
||||
kind: run.kind,
|
||||
trigger: run.trigger,
|
||||
status: run.status,
|
||||
attempt: run.attempt,
|
||||
max_attempts: run.max_attempts,
|
||||
owner_instance: run.owner_instance,
|
||||
progress_percent: run.progress_percent,
|
||||
progress_message: run.progress_message,
|
||||
payload_json: run.payload_json,
|
||||
result_json: run.result_json,
|
||||
error_message: run.error_message,
|
||||
cancel_requested: run.cancel_requested,
|
||||
created_by: run.created_by,
|
||||
created_at_unix_secs: run.created_at_unix_secs,
|
||||
started_at_unix_secs: run.started_at_unix_secs,
|
||||
finished_at_unix_secs: run.finished_at_unix_secs,
|
||||
updated_at_unix_secs: run.updated_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_run_with_logging(
|
||||
app: &AppState,
|
||||
run: UpsertBackgroundTaskRun,
|
||||
) -> Option<StoredBackgroundTaskRun> {
|
||||
match app.upsert_background_task_run(run).await {
|
||||
Ok(result) => result,
|
||||
Err(error) => {
|
||||
warn!(error = ?error, "failed to upsert background task run");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_run_status(
|
||||
app: &AppState,
|
||||
run_id: &str,
|
||||
status: BackgroundTaskStatus,
|
||||
progress_percent: Option<u16>,
|
||||
progress_message: Option<String>,
|
||||
result_json: Option<Value>,
|
||||
error_message: Option<String>,
|
||||
started_at_unix_secs: Option<u64>,
|
||||
finished_at_unix_secs: Option<u64>,
|
||||
) -> Option<StoredBackgroundTaskRun> {
|
||||
let Some(mut existing) = app.find_background_task_run(run_id).await.ok().flatten() else {
|
||||
return None;
|
||||
};
|
||||
existing.status = status;
|
||||
if let Some(progress_percent) = progress_percent {
|
||||
existing.progress_percent = progress_percent.min(100);
|
||||
}
|
||||
if let Some(progress_message) = progress_message {
|
||||
existing.progress_message = Some(progress_message);
|
||||
}
|
||||
if result_json.is_some() {
|
||||
existing.result_json = result_json;
|
||||
}
|
||||
if error_message.is_some() {
|
||||
existing.error_message = error_message;
|
||||
}
|
||||
if started_at_unix_secs.is_some() {
|
||||
existing.started_at_unix_secs = started_at_unix_secs;
|
||||
}
|
||||
if finished_at_unix_secs.is_some() {
|
||||
existing.finished_at_unix_secs = finished_at_unix_secs;
|
||||
}
|
||||
existing.updated_at_unix_secs = now_unix_secs();
|
||||
upsert_run_with_logging(app, stored_run_to_upsert(existing)).await
|
||||
}
|
||||
|
||||
pub(crate) async fn append_event_with_logging(
|
||||
app: &AppState,
|
||||
run_id: &str,
|
||||
event_type: &str,
|
||||
message: &str,
|
||||
payload_json: Option<Value>,
|
||||
) {
|
||||
let event = UpsertBackgroundTaskEvent {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
run_id: run_id.to_string(),
|
||||
event_type: event_type.to_string(),
|
||||
message: message.to_string(),
|
||||
payload_json,
|
||||
created_at_unix_secs: now_unix_secs(),
|
||||
};
|
||||
if let Err(error) = app.upsert_background_task_event(event).await {
|
||||
warn!(error = ?error, run_id = %run_id, "failed to upsert background task event");
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_record_worker_boot(
|
||||
app: AppState,
|
||||
task_key: &'static str,
|
||||
kind: BackgroundTaskKind,
|
||||
trigger: &'static str,
|
||||
) -> JoinHandle<()> {
|
||||
spawn_named("task-runtime-record-worker-boot", async move {
|
||||
let now = now_unix_secs();
|
||||
let run_id = format!("boot:{}:{}", task_key, app.tunnel.local_instance_id());
|
||||
let run = UpsertBackgroundTaskRun {
|
||||
id: run_id.clone(),
|
||||
task_key: task_key.to_string(),
|
||||
kind,
|
||||
trigger: trigger.to_string(),
|
||||
status: BackgroundTaskStatus::Running,
|
||||
attempt: 1,
|
||||
max_attempts: 1,
|
||||
owner_instance: Some(app.tunnel.local_instance_id().to_string()),
|
||||
progress_percent: 0,
|
||||
progress_message: Some("worker booted".to_string()),
|
||||
payload_json: None,
|
||||
result_json: None,
|
||||
error_message: None,
|
||||
cancel_requested: false,
|
||||
created_by: Some("system".to_string()),
|
||||
created_at_unix_secs: now,
|
||||
started_at_unix_secs: Some(now),
|
||||
finished_at_unix_secs: None,
|
||||
updated_at_unix_secs: now,
|
||||
};
|
||||
let _ = upsert_run_with_logging(&app, run).await;
|
||||
append_event_with_logging(
|
||||
&app,
|
||||
&run_id,
|
||||
"worker_boot",
|
||||
"background worker started",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn set_cancel_signal(app: &AppState, run_id: &str) -> Result<(), GatewayError> {
|
||||
app.runtime_kv_setex(&task_cancel_kv_key(run_id), "1", 60 * 60)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn is_cancel_requested(app: &AppState, run_id: &str) -> bool {
|
||||
if let Ok(Some(run)) = app.find_background_task_run(run_id).await {
|
||||
if run.cancel_requested {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
app.runtime_kv_exists(&task_cancel_kv_key(run_id))
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub(crate) async fn submit_provider_delete_task(
|
||||
state: &crate::admin_api::AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
created_by: Option<&str>,
|
||||
) -> Result<Option<String>, GatewayError> {
|
||||
let Some(provider) = state
|
||||
.read_provider_catalog_providers_by_ids(&[provider_id.to_string()])
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let task_id = Uuid::new_v4().simple().to_string()[..16].to_string();
|
||||
state.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
|
||||
task_id: task_id.clone(),
|
||||
provider_id: provider.id.clone(),
|
||||
status: "pending".to_string(),
|
||||
stage: "queued".to_string(),
|
||||
total_keys: 0,
|
||||
deleted_keys: 0,
|
||||
total_endpoints: 0,
|
||||
deleted_endpoints: 0,
|
||||
message: "delete task submitted".to_string(),
|
||||
});
|
||||
|
||||
let app = state.cloned_app();
|
||||
let provider_id = provider.id.clone();
|
||||
let run_id = task_id.clone();
|
||||
let created_at = now_unix_secs();
|
||||
let max_attempts = task_definition(TASK_KEY_PROVIDER_DELETE)
|
||||
.map(|item| item.retry_policy.max_attempts)
|
||||
.unwrap_or(1);
|
||||
if state.has_background_task_data_writer() {
|
||||
let run = UpsertBackgroundTaskRun {
|
||||
id: run_id.clone(),
|
||||
task_key: TASK_KEY_PROVIDER_DELETE.to_string(),
|
||||
kind: BackgroundTaskKind::OnDemand,
|
||||
trigger: "manual".to_string(),
|
||||
status: BackgroundTaskStatus::Queued,
|
||||
attempt: 1,
|
||||
max_attempts,
|
||||
owner_instance: Some(app.tunnel.local_instance_id().to_string()),
|
||||
progress_percent: 0,
|
||||
progress_message: Some("delete task queued".to_string()),
|
||||
payload_json: Some(serde_json::json!({ "provider_id": provider_id.clone() })),
|
||||
result_json: None,
|
||||
error_message: None,
|
||||
cancel_requested: false,
|
||||
created_by: Some(created_by.unwrap_or("admin").to_string()),
|
||||
created_at_unix_secs: created_at,
|
||||
started_at_unix_secs: None,
|
||||
finished_at_unix_secs: None,
|
||||
updated_at_unix_secs: created_at,
|
||||
};
|
||||
let _ = upsert_run_with_logging(&app, run).await;
|
||||
append_event_with_logging(
|
||||
&app,
|
||||
&run_id,
|
||||
"queued",
|
||||
"provider delete task queued",
|
||||
Some(serde_json::json!({ "provider_id": provider_id.clone() })),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
spawn_named("task-runtime-provider-delete", async move {
|
||||
let lock_key = format!("task_runtime:lock:{TASK_KEY_PROVIDER_DELETE}:{provider_id}");
|
||||
let lock_ttl = std::time::Duration::from_secs(60 * 15);
|
||||
let lock = app
|
||||
.runtime_state
|
||||
.lock_try_acquire(&lock_key, app.tunnel.local_instance_id(), lock_ttl)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
if lock.is_none() {
|
||||
let _ = update_run_status(
|
||||
&app,
|
||||
&run_id,
|
||||
BackgroundTaskStatus::Skipped,
|
||||
Some(0),
|
||||
Some("provider delete skipped: another node is running this task".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(now_unix_secs()),
|
||||
)
|
||||
.await;
|
||||
append_event_with_logging(
|
||||
&app,
|
||||
&run_id,
|
||||
"skipped",
|
||||
"provider delete skipped by singleton lock",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
let started_at = now_unix_secs();
|
||||
let _ = update_run_status(
|
||||
&app,
|
||||
&run_id,
|
||||
BackgroundTaskStatus::Running,
|
||||
Some(5),
|
||||
Some("provider delete task started".to_string()),
|
||||
None,
|
||||
None,
|
||||
Some(started_at),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
append_event_with_logging(
|
||||
&app,
|
||||
&run_id,
|
||||
"running",
|
||||
"provider delete task started",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let admin_state = crate::admin_api::AdminAppState::new(&app);
|
||||
let result = admin_state
|
||||
.run_admin_provider_delete_task(&provider_id, &run_id)
|
||||
.await;
|
||||
match result {
|
||||
Ok(task_state) => {
|
||||
let _ = update_run_status(
|
||||
&app,
|
||||
&run_id,
|
||||
BackgroundTaskStatus::Succeeded,
|
||||
Some(100),
|
||||
Some(task_state.message.clone()),
|
||||
Some(serde_json::json!({
|
||||
"provider_id": task_state.provider_id,
|
||||
"status": task_state.status,
|
||||
"stage": task_state.stage,
|
||||
"deleted_keys": task_state.deleted_keys,
|
||||
"total_keys": task_state.total_keys,
|
||||
"deleted_endpoints": task_state.deleted_endpoints,
|
||||
"total_endpoints": task_state.total_endpoints,
|
||||
"message": task_state.message,
|
||||
})),
|
||||
None,
|
||||
None,
|
||||
Some(now_unix_secs()),
|
||||
)
|
||||
.await;
|
||||
append_event_with_logging(
|
||||
&app,
|
||||
&run_id,
|
||||
"succeeded",
|
||||
"provider delete task completed",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"gateway admin provider delete task failed for provider {}: {:?}",
|
||||
provider_id, err
|
||||
);
|
||||
app.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
|
||||
task_id: run_id.clone(),
|
||||
provider_id: provider_id.clone(),
|
||||
status: "failed".to_string(),
|
||||
stage: "failed".to_string(),
|
||||
total_keys: 0,
|
||||
deleted_keys: 0,
|
||||
total_endpoints: 0,
|
||||
deleted_endpoints: 0,
|
||||
message: format!("provider delete failed: {err:?}"),
|
||||
});
|
||||
let _ = update_run_status(
|
||||
&app,
|
||||
&run_id,
|
||||
BackgroundTaskStatus::Failed,
|
||||
Some(100),
|
||||
Some("provider delete task failed".to_string()),
|
||||
None,
|
||||
Some(format!("{err:?}")),
|
||||
None,
|
||||
Some(now_unix_secs()),
|
||||
)
|
||||
.await;
|
||||
append_event_with_logging(
|
||||
&app,
|
||||
&run_id,
|
||||
"failed",
|
||||
"provider delete task failed",
|
||||
Some(serde_json::json!({ "error": format!("{err:?}") })),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(lock) = lock {
|
||||
let _ = app.runtime_state.lock_release(&lock).await;
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Some(task_id))
|
||||
}
|
||||
@@ -928,6 +928,428 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_chat_stream_via_local_cross_format_gemini_candidate_with_stream_path_rewrite(
|
||||
) {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
provider_model: String,
|
||||
auth_header_value: String,
|
||||
accept: String,
|
||||
endpoint_tag: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai", "gemini"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["openai", "gemini"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-openai-chat-gemini-stream-local-1".to_string(),
|
||||
provider_name: "gemini".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-openai-chat-gemini-stream-local-1".to_string(),
|
||||
endpoint_api_format: "gemini:generate_content".to_string(),
|
||||
endpoint_api_family: Some("gemini".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-openai-chat-gemini-stream-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["gemini:generate_content".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"gemini:generate_content": 1})),
|
||||
model_id: "model-openai-chat-gemini-stream-local-1".to_string(),
|
||||
global_model_id: "global-model-openai-chat-gemini-stream-local-1".to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gemini-2.5-pro-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gemini-2.5-pro-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["gemini:generate_content".to_string()]),
|
||||
endpoint_ids: None,
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-openai-chat-gemini-stream-local-1".to_string(),
|
||||
"gemini".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-openai-chat-gemini-stream-local-1".to_string(),
|
||||
"provider-openai-chat-gemini-stream-local-1".to_string(),
|
||||
"gemini:generate_content".to_string(),
|
||||
Some("gemini".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://generativelanguage.googleapis.com".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"openai-chat-gemini-cross-format-stream"}
|
||||
])),
|
||||
None,
|
||||
Some(2),
|
||||
Some("/custom/v1beta/models/gemini-2.5-pro-upstream:generateContent".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-openai-chat-gemini-stream-local-1".to_string(),
|
||||
"provider-openai-chat-gemini-stream-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["gemini:generate_content"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-openai-chat-gemini-stream",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"gemini:generate_content": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeStreamRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-openai-chat-gemini-stream-local-1",
|
||||
"api_key_id": "api-key-openai-chat-gemini-stream-local-1",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/decision-stream",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-stream",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-stream",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
provider_model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
auth_header_value: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-goog-api-key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
accept: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("accept"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
let frames = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"responseId\\\":\\\"resp_openai_chat_gemini_stream_123\\\",\\\"candidates\\\":[{\\\"content\\\":{\\\"parts\\\":[{\\\"text\\\":\\\"Hello \\\"}],\\\"role\\\":\\\"model\\\"},\\\"index\\\":0}],\\\"modelVersion\\\":\\\"gemini-2.5-pro-upstream\\\"}\\n\\n\"}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"responseId\\\":\\\"resp_openai_chat_gemini_stream_123\\\",\\\"candidates\\\":[{\\\"content\\\":{\\\"parts\\\":[{\\\"text\\\":\\\"Hello Gemini stream\\\"}],\\\"role\\\":\\\"model\\\"},\\\"finishReason\\\":\\\"STOP\\\",\\\"index\\\":0}],\\\"modelVersion\\\":\\\"gemini-2.5-pro-upstream\\\",\\\"usageMetadata\\\":{\\\"promptTokenCount\\\":1,\\\"candidatesTokenCount\\\":2,\\\"totalTokenCount\\\":3}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(frames))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-chat-gemini-stream")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-openai-chat-gemini-stream-local-1",
|
||||
"user-openai-chat-gemini-stream-local-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"provider_priority_mode".to_string(),
|
||||
json!("global_key"),
|
||||
)]),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-chat-gemini-stream",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-openai-chat-gemini-stream-123")
|
||||
.body(
|
||||
"{\"model\":\"gpt-5\",\"messages\":[{\"role\":\"system\",\"content\":\"You are terse.\"},{\"role\":\"user\",\"content\":\"Say hello\"}],\"stream\":true}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
|
||||
);
|
||||
let response_text = response.text().await.expect("body should read");
|
||||
assert!(response_text.contains("\"object\":\"chat.completion.chunk\""));
|
||||
assert!(response_text.contains("data: [DONE]"));
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime stream should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-openai-chat-gemini-stream-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://generativelanguage.googleapis.com/custom/v1beta/models/gemini-2.5-pro-upstream:streamGenerateContent?alt=sse"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.provider_model,
|
||||
"gemini-2.5-pro-upstream"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.auth_header_value,
|
||||
"sk-upstream-openai-chat-gemini-stream"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"openai-chat-gemini-cross-format-stream"
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-chat-gemini-stream-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
let extra_data = stored_candidates[0]
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.expect("request candidate extra_data should exist");
|
||||
assert_eq!(extra_data["client_api_format"], "openai:chat");
|
||||
assert_eq!(extra_data["provider_api_format"], "gemini:generate_content");
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-stream should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_chat_stream_with_custom_path_via_local_decision_gate_with_local_stream_decision(
|
||||
) {
|
||||
|
||||
@@ -72,7 +72,5 @@ async fn gateway_background_gemini_file_mapping_cleanup_deletes_expired_entries(
|
||||
"active mapping should remain"
|
||||
);
|
||||
|
||||
for handle in background_tasks {
|
||||
handle.abort();
|
||||
}
|
||||
background_tasks.shutdown().await;
|
||||
}
|
||||
|
||||
@@ -231,9 +231,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
|
||||
}]
|
||||
);
|
||||
|
||||
for handle in background_tasks {
|
||||
handle.abort();
|
||||
}
|
||||
background_tasks.shutdown().await;
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
@@ -321,8 +319,6 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
|
||||
}]
|
||||
);
|
||||
|
||||
for handle in background_tasks {
|
||||
handle.abort();
|
||||
}
|
||||
background_tasks.shutdown().await;
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ use crate::clock::current_unix_ms;
|
||||
use crate::log_ids::short_request_id;
|
||||
use crate::orchestration::{apply_local_report_effect, LocalReportEffect};
|
||||
use crate::request_candidate_runtime::record_report_request_candidate_status;
|
||||
use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_USAGE_SYNC_REPORT};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
mod context;
|
||||
@@ -142,7 +143,7 @@ pub(crate) async fn submit_sync_report(
|
||||
pub(crate) fn spawn_sync_report(state: AppState, payload: GatewaySyncReportRequest) {
|
||||
let report_request_id_for_log =
|
||||
short_request_id(report_request_id(payload.report_context.as_ref()));
|
||||
tokio::spawn(async move {
|
||||
spawn_fire_and_forget(TASK_KEY_USAGE_SYNC_REPORT, async move {
|
||||
let trace_id = payload.trace_id.clone();
|
||||
if let Err(err) = submit_sync_report(&state, payload).await {
|
||||
warn!(
|
||||
|
||||
Reference in New Issue
Block a user