mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 10:57:03 +08:00
refactor: 优化调度候选排序与用量写入链路并改进 Fernet 缓存与前端批量列表
This commit is contained in:
@@ -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")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user