refactor: 优化调度候选排序与用量写入链路并改进 Fernet 缓存与前端批量列表

This commit is contained in:
fawney19
2026-04-21 16:19:07 +08:00
parent c5c56ff92f
commit 25a2b417be
77 changed files with 4801 additions and 2881 deletions
@@ -1,6 +1,6 @@
use std::collections::BTreeMap;
use aether_contracts::ExecutionResult;
use aether_contracts::ResponseBody;
use base64::Engine as _;
use crate::GatewayError;
@@ -8,14 +8,14 @@ use crate::GatewayError;
type DecodedBody = (Vec<u8>, Option<serde_json::Value>, Option<String>);
pub(super) fn decode_execution_result_body(
result: &ExecutionResult,
body: Option<ResponseBody>,
headers: &mut BTreeMap<String, String>,
) -> Result<DecodedBody, GatewayError> {
let Some(body) = result.body.as_ref() else {
let Some(body) = body else {
return Ok((Vec::new(), None, None));
};
if let Some(json_body) = body.json_body.clone() {
if let Some(json_body) = body.json_body {
headers
.entry("content-type".to_string())
.or_insert_with(|| "application/json".to_string());
@@ -25,7 +25,7 @@ pub(super) fn decode_execution_result_body(
return Ok((bytes, Some(json_body), None));
}
if let Some(body_bytes_b64) = body.body_bytes_b64.clone() {
if let Some(body_bytes_b64) = body.body_bytes_b64 {
let bytes = base64::engine::general_purpose::STANDARD
.decode(&body_bytes_b64)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
@@ -2,10 +2,13 @@ use std::collections::BTreeMap;
use aether_contracts::ExecutionPlan;
use axum::body::Body;
use axum::http::header::HeaderValue;
use axum::http::Response;
use serde_json::json;
use crate::api::response::build_client_response_from_parts;
use crate::api::response::{
build_client_response_from_parts, build_client_response_from_parts_with_mutator,
};
use crate::async_task::VideoTaskService;
use crate::control::GatewayControlDecision;
use crate::video_tasks::{
@@ -17,9 +20,15 @@ pub(crate) use crate::video_tasks::{
};
use crate::{usage::GatewaySyncReportRequest, GatewayError};
pub(crate) enum LocalVideoSyncSuccessBuild {
Handled(LocalVideoSyncSuccessOutcome),
NotHandled(GatewaySyncReportRequest),
}
pub(crate) struct LocalVideoSyncSuccessOutcome {
pub(crate) response: Response<Body>,
pub(crate) report_payload: GatewaySyncReportRequest,
pub(crate) original_report_context: Option<serde_json::Value>,
pub(crate) report_mode: VideoTaskSyncReportMode,
pub(crate) local_task_snapshot: Option<LocalVideoTaskSnapshot>,
}
@@ -29,8 +38,9 @@ fn cloned_report_context_object(
) -> serde_json::Map<String, serde_json::Value> {
payload
.report_context
.clone()
.and_then(|value| value.as_object().cloned())
.as_ref()
.and_then(serde_json::Value::as_object)
.cloned()
.unwrap_or_default()
}
@@ -56,54 +66,53 @@ fn build_local_video_success_response(
pub(crate) fn maybe_build_local_video_success_outcome(
trace_id: &str,
decision: &GatewayControlDecision,
payload: &GatewaySyncReportRequest,
mut payload: GatewaySyncReportRequest,
video_tasks: &VideoTaskService,
plan: &ExecutionPlan,
) -> Result<Option<LocalVideoSyncSuccessOutcome>, GatewayError> {
) -> Result<LocalVideoSyncSuccessBuild, GatewayError> {
if payload.status_code >= 400 {
return Ok(None);
return Ok(LocalVideoSyncSuccessBuild::NotHandled(payload));
}
let provider_body = match payload
.body_json
.as_ref()
.and_then(serde_json::Value::as_object)
{
Some(value) => value,
None => return Ok(None),
let mut report_context = cloned_report_context_object(&payload);
let prepared_plan = {
let provider_body = match payload
.body_json
.as_ref()
.and_then(serde_json::Value::as_object)
{
Some(value) => value,
None => return Ok(LocalVideoSyncSuccessBuild::NotHandled(payload)),
};
video_tasks.prepare_sync_success(
payload.report_kind.as_str(),
provider_body,
&report_context,
plan,
)
};
let mut report_context = cloned_report_context_object(payload);
let Some(plan) = video_tasks.prepare_sync_success(
payload.report_kind.as_str(),
provider_body,
&report_context,
plan,
) else {
return Ok(None);
let Some(plan) = prepared_plan else {
return Ok(LocalVideoSyncSuccessBuild::NotHandled(payload));
};
plan.apply_to_report_context(&mut report_context);
let client_body_json = plan.client_body_json();
let response = build_local_video_success_response(trace_id, decision, &client_body_json)?;
let report_payload = GatewaySyncReportRequest {
trace_id: payload.trace_id.clone(),
report_kind: plan.success_report_kind().to_string(),
report_context: Some(serde_json::Value::Object(report_context)),
status_code: payload.status_code,
headers: payload.headers.clone(),
body_json: payload.body_json.clone(),
client_body_json: Some(client_body_json),
body_base64: None,
telemetry: payload.telemetry.clone(),
};
let original_report_context = payload.report_context.take();
payload.report_kind = plan.success_report_kind().to_string();
payload.report_context = Some(serde_json::Value::Object(report_context));
payload.client_body_json = Some(client_body_json);
Ok(Some(LocalVideoSyncSuccessOutcome {
response,
report_payload,
report_mode: plan.report_mode(),
local_task_snapshot: matches!(plan.report_mode(), VideoTaskSyncReportMode::Background)
.then(|| plan.to_snapshot()),
}))
Ok(LocalVideoSyncSuccessBuild::Handled(
LocalVideoSyncSuccessOutcome {
response,
report_payload: payload,
original_report_context,
report_mode: plan.report_mode(),
local_task_snapshot: matches!(plan.report_mode(), VideoTaskSyncReportMode::Background)
.then(|| plan.to_snapshot()),
},
))
}
pub(crate) fn maybe_build_local_sync_finalize_response(
@@ -147,21 +156,114 @@ pub(crate) fn maybe_build_local_video_error_response(
return Ok(None);
}
let response_body = payload.body_json.clone().unwrap_or_else(|| json!({}));
let body_bytes = serde_json::to_vec(&response_body)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let empty_body = json!({});
let response_body = payload.body_json.as_ref().unwrap_or(&empty_body);
let body_bytes =
serde_json::to_vec(response_body).map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_len = body_bytes.len().to_string();
let mut response_headers = payload.headers.clone();
response_headers.remove("content-encoding");
response_headers.remove("content-length");
response_headers.insert("content-type".to_string(), "application/json".to_string());
response_headers.insert("content-length".to_string(), body_bytes.len().to_string());
Ok(Some(build_client_response_from_parts(
Ok(Some(build_client_response_from_parts_with_mutator(
payload.status_code,
&response_headers,
&payload.headers,
Body::from(body_bytes),
trace_id,
Some(decision),
|headers| {
headers.remove(http::header::CONTENT_ENCODING);
headers.remove(http::header::CONTENT_LENGTH);
headers.insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
headers.insert(
http::header::CONTENT_LENGTH,
HeaderValue::from_str(body_len.as_str())
.map_err(|err| GatewayError::Internal(err.to_string()))?,
);
Ok(())
},
)?))
}
#[cfg(test)]
mod tests {
use super::*;
use axum::body::to_bytes;
use serde_json::json;
#[tokio::test]
async fn local_video_error_response_rewrites_headers_without_mutating_payload() {
let decision = GatewayControlDecision::synthetic(
"/v1/videos",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("video".to_string()),
Some("openai:video".to_string()),
)
.with_execution_runtime_candidate(true);
let payload = GatewaySyncReportRequest {
trace_id: "trace-payload".to_string(),
report_kind: "openai_video_create_sync_finalize".to_string(),
report_context: Some(json!({
"request_id": "req_123",
})),
status_code: http::StatusCode::BAD_GATEWAY.as_u16(),
headers: BTreeMap::from([
("content-encoding".to_string(), "gzip".to_string()),
("content-length".to_string(), "999".to_string()),
("x-upstream-id".to_string(), "video-123".to_string()),
]),
body_json: Some(json!({
"error": {
"type": "video_backend_error",
"message": "backend failed",
}
})),
client_body_json: None,
body_base64: None,
telemetry: None,
};
let response =
maybe_build_local_video_error_response("trace-response", &decision, &payload)
.expect("video error response should build")
.expect("video error response should match local video error kinds");
assert_eq!(response.status(), http::StatusCode::BAD_GATEWAY);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("application/json")
);
assert_eq!(response.headers().get(http::header::CONTENT_ENCODING), None);
assert_eq!(
response
.headers()
.get("x-upstream-id")
.and_then(|value| value.to_str().ok()),
Some("video-123")
);
assert_eq!(
payload.headers.get("content-encoding").map(String::as_str),
Some("gzip")
);
assert_eq!(
payload.headers.get("content-length").map(String::as_str),
Some("999")
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("response body should parse"),
payload
.body_json
.clone()
.expect("payload body should exist")
);
}
}