mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: 移除独立 hub/proxy/executor/gateway crate,统一为 gateway tunnel 架构
- 删除 aether-hub、aether-proxy 独立项目及其 Dockerfile/配置 - 删除 crates/aether-executor 和 crates/aether-gateway 全部模块 - 新增 apps/ 目录作为应用入口 - 将 hub 概念重构为 gateway tunnel transport - 将 executor 重构为 execution runtime - 新增 tunnel.rs 合约定义和 testkit tunnel/execution_runtime 模块 - 更新 Python 服务层和测试适配新架构命名
This commit is contained in:
281
apps/aether-gateway/src/async_task/http.rs
Normal file
281
apps/aether-gateway/src/async_task/http.rs
Normal file
@@ -0,0 +1,281 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_contracts::ExecutionResult;
|
||||
use aether_data::repository::video_tasks::{
|
||||
StoredVideoTask, VideoTaskQueryFilter, VideoTaskStatus,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::extract::{Path, Query, State};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::response::Redirect;
|
||||
use axum::Json;
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
mod cancel;
|
||||
|
||||
use super::query::VideoTaskVideoSource;
|
||||
use super::{
|
||||
read_video_task_detail, read_video_task_page, read_video_task_stats,
|
||||
read_video_task_video_source,
|
||||
};
|
||||
use crate::gateway::{AppState, GatewayError};
|
||||
|
||||
pub(crate) use self::cancel::{cancel_video_task_record, CancelVideoTaskError};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ListVideoTasksQuery {
|
||||
pub(crate) status: Option<String>,
|
||||
pub(crate) user_id: Option<String>,
|
||||
pub(crate) model: Option<String>,
|
||||
pub(crate) client_api_format: Option<String>,
|
||||
pub(crate) page: Option<usize>,
|
||||
pub(crate) page_size: Option<usize>,
|
||||
}
|
||||
|
||||
pub(crate) async fn list_video_tasks(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ListVideoTasksQuery>,
|
||||
) -> Result<Json<super::query::VideoTaskPageResponse>, axum::response::Response> {
|
||||
let filter = parse_filter(&query)?;
|
||||
let response = read_video_task_page(
|
||||
&state,
|
||||
&filter,
|
||||
query.page.unwrap_or(1),
|
||||
query.page_size.unwrap_or(20),
|
||||
)
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?;
|
||||
Ok(Json(response))
|
||||
}
|
||||
|
||||
pub(crate) async fn get_video_task_stats(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ListVideoTasksQuery>,
|
||||
) -> Result<Json<super::query::VideoTaskStatsResponse>, axum::response::Response> {
|
||||
let filter = parse_filter(&query)?;
|
||||
let response = read_video_task_stats(&state, &filter, current_unix_secs())
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?;
|
||||
Ok(Json(response))
|
||||
}
|
||||
|
||||
pub(crate) async fn get_video_task_detail(
|
||||
State(state): State<AppState>,
|
||||
Path(task_id): Path<String>,
|
||||
) -> Result<Json<aether_data::repository::video_tasks::StoredVideoTask>, axum::response::Response> {
|
||||
let task = read_video_task_detail(&state, &task_id)
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?;
|
||||
|
||||
match task {
|
||||
Some(task) => Ok(Json(task)),
|
||||
None => Err((
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Video task not found",
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn cancel_video_task(
|
||||
State(state): State<AppState>,
|
||||
Path(task_id): Path<String>,
|
||||
) -> Result<Json<Value>, axum::response::Response> {
|
||||
let stored = cancel_video_task_record(&state, &task_id)
|
||||
.await
|
||||
.map_err(|err| match err {
|
||||
CancelVideoTaskError::NotFound => (
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Video task not found",
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response(),
|
||||
CancelVideoTaskError::InvalidStatus(status) => (
|
||||
axum::http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": format!(
|
||||
"Cannot cancel task with status: {}",
|
||||
video_task_status_name(status),
|
||||
),
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response(),
|
||||
CancelVideoTaskError::Response(response) => response,
|
||||
CancelVideoTaskError::Gateway(err) => err.into_response(),
|
||||
})?;
|
||||
|
||||
Ok(Json(json!({
|
||||
"id": stored.id,
|
||||
"status": "cancelled",
|
||||
"message": "Task cancelled successfully",
|
||||
})))
|
||||
}
|
||||
|
||||
pub(crate) async fn get_video_task_video(
|
||||
State(state): State<AppState>,
|
||||
Path(task_id): Path<String>,
|
||||
) -> Result<axum::response::Response, axum::response::Response> {
|
||||
let Some(source) = read_video_task_video_source(&state, &task_id)
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?
|
||||
else {
|
||||
return Err((
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Video task or video not found",
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response());
|
||||
};
|
||||
|
||||
build_video_task_video_response(&state, &task_id, source)
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_video_task_video_response(
|
||||
state: &AppState,
|
||||
task_id: &str,
|
||||
source: VideoTaskVideoSource,
|
||||
) -> Result<axum::response::Response, GatewayError> {
|
||||
match source {
|
||||
VideoTaskVideoSource::Redirect { url } => Ok(Redirect::temporary(&url).into_response()),
|
||||
VideoTaskVideoSource::Proxy {
|
||||
url,
|
||||
header_name,
|
||||
header_value,
|
||||
filename,
|
||||
} => proxy_video_stream(state, task_id, &url, &header_name, &header_value, &filename).await,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_filter(
|
||||
query: &ListVideoTasksQuery,
|
||||
) -> Result<VideoTaskQueryFilter, axum::response::Response> {
|
||||
let status = match query.status.as_deref() {
|
||||
Some(value) => Some(VideoTaskStatus::from_database(value).map_err(|err| {
|
||||
(
|
||||
axum::http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": err.to_string(),
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
})?),
|
||||
None => None,
|
||||
};
|
||||
|
||||
Ok(VideoTaskQueryFilter {
|
||||
user_id: query.user_id.clone(),
|
||||
status,
|
||||
model_substring: query.model.clone(),
|
||||
client_api_format: query.client_api_format.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
fn video_task_status_name(status: VideoTaskStatus) -> &'static str {
|
||||
match status {
|
||||
VideoTaskStatus::Pending => "pending",
|
||||
VideoTaskStatus::Submitted => "submitted",
|
||||
VideoTaskStatus::Queued => "queued",
|
||||
VideoTaskStatus::Processing => "processing",
|
||||
VideoTaskStatus::Completed => "completed",
|
||||
VideoTaskStatus::Failed => "failed",
|
||||
VideoTaskStatus::Cancelled => "cancelled",
|
||||
VideoTaskStatus::Expired => "expired",
|
||||
VideoTaskStatus::Deleted => "deleted",
|
||||
}
|
||||
}
|
||||
|
||||
async fn proxy_video_stream(
|
||||
state: &AppState,
|
||||
task_id: &str,
|
||||
url: &str,
|
||||
header_name: &str,
|
||||
header_value: &str,
|
||||
filename: &str,
|
||||
) -> Result<axum::response::Response, GatewayError> {
|
||||
let response = state
|
||||
.client
|
||||
.get(url)
|
||||
.header(header_name, header_value)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| GatewayError::UpstreamUnavailable {
|
||||
trace_id: task_id.to_string(),
|
||||
message: err.to_string(),
|
||||
})?;
|
||||
|
||||
if response.status().is_client_error() || response.status().is_server_error() {
|
||||
return Err(GatewayError::UpstreamUnavailable {
|
||||
trace_id: task_id.to_string(),
|
||||
message: format!("video upstream returned HTTP {}", response.status()),
|
||||
});
|
||||
}
|
||||
|
||||
let status = response.status();
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(axum::http::header::CONTENT_TYPE)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| axum::http::HeaderValue::from_static("video/mp4"));
|
||||
let content_length = response
|
||||
.headers()
|
||||
.get(axum::http::header::CONTENT_LENGTH)
|
||||
.cloned();
|
||||
let cache_control = response
|
||||
.headers()
|
||||
.get(axum::http::header::CACHE_CONTROL)
|
||||
.cloned();
|
||||
let body = Body::from_stream(response.bytes_stream());
|
||||
|
||||
let mut outbound = axum::http::Response::builder()
|
||||
.status(status)
|
||||
.body(body)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
outbound
|
||||
.headers_mut()
|
||||
.insert(axum::http::header::CONTENT_TYPE, content_type);
|
||||
outbound.headers_mut().insert(
|
||||
axum::http::header::CONTENT_DISPOSITION,
|
||||
axum::http::HeaderValue::from_str(&format!("inline; filename=\"{filename}\""))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
if let Some(content_length) = content_length {
|
||||
outbound
|
||||
.headers_mut()
|
||||
.insert(axum::http::header::CONTENT_LENGTH, content_length);
|
||||
}
|
||||
if let Some(cache_control) = cache_control {
|
||||
outbound
|
||||
.headers_mut()
|
||||
.insert(axum::http::header::CACHE_CONTROL, cache_control);
|
||||
} else {
|
||||
outbound.headers_mut().insert(
|
||||
axum::http::header::CACHE_CONTROL,
|
||||
axum::http::HeaderValue::from_static("private, max-age=3600"),
|
||||
);
|
||||
}
|
||||
Ok(outbound)
|
||||
}
|
||||
264
apps/aether-gateway/src/async_task/http/cancel.rs
Normal file
264
apps/aether-gateway/src/async_task/http/cancel.rs
Normal file
@@ -0,0 +1,264 @@
|
||||
use aether_data::repository::video_tasks::{StoredVideoTask, UpsertVideoTask, VideoTaskStatus};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::Json;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::gateway::{AppState, GatewayError};
|
||||
|
||||
use super::super::finalize_video_task_if_terminal;
|
||||
use super::super::read_video_task_detail;
|
||||
use super::current_unix_secs;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum CancelVideoTaskError {
|
||||
NotFound,
|
||||
InvalidStatus(VideoTaskStatus),
|
||||
Response(axum::response::Response),
|
||||
Gateway(GatewayError),
|
||||
}
|
||||
|
||||
impl From<GatewayError> for CancelVideoTaskError {
|
||||
fn from(value: GatewayError) -> Self {
|
||||
Self::Gateway(value)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn cancel_video_task_record(
|
||||
state: &AppState,
|
||||
task_id: &str,
|
||||
) -> Result<StoredVideoTask, CancelVideoTaskError> {
|
||||
let Some(task) = read_video_task_detail(state, task_id).await? else {
|
||||
return Err(CancelVideoTaskError::NotFound);
|
||||
};
|
||||
|
||||
if matches!(
|
||||
task.status,
|
||||
VideoTaskStatus::Completed
|
||||
| VideoTaskStatus::Failed
|
||||
| VideoTaskStatus::Cancelled
|
||||
| VideoTaskStatus::Expired
|
||||
| VideoTaskStatus::Deleted
|
||||
) {
|
||||
return Err(CancelVideoTaskError::InvalidStatus(task.status));
|
||||
}
|
||||
|
||||
let trace_id = format!("async-task-admin-cancel-{task_id}");
|
||||
if let Some(cancel_plan) = build_video_task_cancel_plan(&task) {
|
||||
state
|
||||
.hydrate_video_task_for_route(Some(cancel_plan.route_family), &cancel_plan.request_path)
|
||||
.await?;
|
||||
|
||||
let body_json = json!({});
|
||||
let follow_up = state.video_tasks.prepare_follow_up_sync_plan(
|
||||
cancel_plan.plan_kind,
|
||||
&cancel_plan.request_path,
|
||||
Some(&body_json),
|
||||
None,
|
||||
&trace_id,
|
||||
);
|
||||
|
||||
if let Some(follow_up) = follow_up {
|
||||
execute_video_task_cancel_plan(state, &trace_id, follow_up.plan)
|
||||
.await
|
||||
.map_err(CancelVideoTaskError::Response)?;
|
||||
}
|
||||
|
||||
state
|
||||
.video_tasks
|
||||
.apply_finalize_mutation(&cancel_plan.request_path, cancel_plan.report_kind);
|
||||
}
|
||||
|
||||
let request_metadata = build_cancelled_request_metadata(state, &task).await?;
|
||||
let stored = persist_cancelled_video_task(state, &task, request_metadata)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
CancelVideoTaskError::Gateway(GatewayError::Internal(
|
||||
"video task repository is unavailable".to_string(),
|
||||
))
|
||||
})?;
|
||||
finalize_video_task_if_terminal(state, &stored).await;
|
||||
Ok(stored)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct VideoTaskCancelPlan<'a> {
|
||||
route_family: &'a str,
|
||||
plan_kind: &'a str,
|
||||
report_kind: &'a str,
|
||||
request_path: String,
|
||||
}
|
||||
|
||||
fn build_video_task_cancel_plan(task: &StoredVideoTask) -> Option<VideoTaskCancelPlan<'_>> {
|
||||
let provider_api_format = task
|
||||
.provider_api_format
|
||||
.as_deref()
|
||||
.or(task.client_api_format.as_deref())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
|
||||
match provider_api_format {
|
||||
"openai:video" => Some(VideoTaskCancelPlan {
|
||||
route_family: "openai",
|
||||
plan_kind: "openai_video_cancel_sync",
|
||||
report_kind: "openai_video_cancel_sync_finalize",
|
||||
request_path: format!("/v1/videos/{}/cancel", task.id),
|
||||
}),
|
||||
"gemini:video" => {
|
||||
let short_id = task.short_id.as_deref().unwrap_or(task.id.as_str()).trim();
|
||||
let model = task
|
||||
.model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
Some(VideoTaskCancelPlan {
|
||||
route_family: "gemini",
|
||||
plan_kind: "gemini_video_cancel_sync",
|
||||
report_kind: "gemini_video_cancel_sync_finalize",
|
||||
request_path: format!("/v1beta/models/{model}/operations/{short_id}:cancel"),
|
||||
})
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_video_task_cancel_plan(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
plan: aether_contracts::ExecutionPlan,
|
||||
) -> Result<(), axum::response::Response> {
|
||||
let result = crate::gateway::execute_execution_runtime_sync_plan(state, Some(trace_id), &plan)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
GatewayError::UpstreamUnavailable {
|
||||
trace_id: trace_id.to_string(),
|
||||
message: format!("{err:?}"),
|
||||
}
|
||||
.into_response()
|
||||
})?;
|
||||
|
||||
if result.status_code >= 400 {
|
||||
let status = axum::http::StatusCode::from_u16(result.status_code)
|
||||
.unwrap_or(axum::http::StatusCode::BAD_GATEWAY);
|
||||
let body_json = result
|
||||
.body
|
||||
.and_then(|body| body.json_body)
|
||||
.unwrap_or_else(|| {
|
||||
json!({
|
||||
"error": {
|
||||
"message": result
|
||||
.error
|
||||
.as_ref()
|
||||
.map(|error| error.message.clone())
|
||||
.unwrap_or_else(|| {
|
||||
format!("execution runtime returned {}", result.status_code)
|
||||
}),
|
||||
}
|
||||
})
|
||||
});
|
||||
return Err((status, Json(body_json)).into_response());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn build_cancelled_request_metadata(
|
||||
state: &AppState,
|
||||
task: &StoredVideoTask,
|
||||
) -> Result<Option<Value>, GatewayError> {
|
||||
let mut metadata = match task.request_metadata.clone() {
|
||||
Some(Value::Object(object)) => object,
|
||||
_ => Map::new(),
|
||||
};
|
||||
let mut snapshot_value = metadata.get("rust_local_snapshot").cloned();
|
||||
if snapshot_value.is_none() {
|
||||
snapshot_value = state
|
||||
.reconstruct_video_task_snapshot(task)
|
||||
.await?
|
||||
.map(|snapshot| {
|
||||
serde_json::to_value(snapshot)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
})
|
||||
.transpose()?;
|
||||
}
|
||||
if let Some(snapshot_value_ref) = snapshot_value.as_mut() {
|
||||
mark_snapshot_value_cancelled(snapshot_value_ref);
|
||||
metadata.insert(
|
||||
"rust_owner".to_string(),
|
||||
Value::String("async_task".to_string()),
|
||||
);
|
||||
metadata.insert(
|
||||
"rust_local_snapshot".to_string(),
|
||||
snapshot_value_ref.clone(),
|
||||
);
|
||||
return Ok(Some(Value::Object(metadata)));
|
||||
}
|
||||
|
||||
Ok(task.request_metadata.clone())
|
||||
}
|
||||
|
||||
fn mark_snapshot_value_cancelled(snapshot_value: &mut Value) {
|
||||
if let Some(object) = snapshot_value
|
||||
.get_mut("OpenAi")
|
||||
.and_then(Value::as_object_mut)
|
||||
{
|
||||
object.insert("status".to_string(), Value::String("Cancelled".to_string()));
|
||||
return;
|
||||
}
|
||||
if let Some(object) = snapshot_value
|
||||
.get_mut("Gemini")
|
||||
.and_then(Value::as_object_mut)
|
||||
{
|
||||
object.insert("status".to_string(), Value::String("Cancelled".to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_cancelled_video_task(
|
||||
state: &AppState,
|
||||
task: &StoredVideoTask,
|
||||
request_metadata: Option<Value>,
|
||||
) -> Result<Option<StoredVideoTask>, GatewayError> {
|
||||
let now_unix_secs = current_unix_secs();
|
||||
state
|
||||
.data
|
||||
.upsert_video_task(UpsertVideoTask {
|
||||
id: task.id.clone(),
|
||||
short_id: task.short_id.clone(),
|
||||
request_id: task.request_id.clone(),
|
||||
user_id: task.user_id.clone(),
|
||||
api_key_id: task.api_key_id.clone(),
|
||||
username: task.username.clone(),
|
||||
api_key_name: task.api_key_name.clone(),
|
||||
external_task_id: task.external_task_id.clone(),
|
||||
provider_id: task.provider_id.clone(),
|
||||
endpoint_id: task.endpoint_id.clone(),
|
||||
key_id: task.key_id.clone(),
|
||||
client_api_format: task.client_api_format.clone(),
|
||||
provider_api_format: task.provider_api_format.clone(),
|
||||
format_converted: task.format_converted,
|
||||
model: task.model.clone(),
|
||||
prompt: task.prompt.clone(),
|
||||
original_request_body: task.original_request_body.clone(),
|
||||
duration_seconds: task.duration_seconds,
|
||||
resolution: task.resolution.clone(),
|
||||
aspect_ratio: task.aspect_ratio.clone(),
|
||||
size: task.size.clone(),
|
||||
status: VideoTaskStatus::Cancelled,
|
||||
progress_percent: task.progress_percent,
|
||||
progress_message: task.progress_message.clone(),
|
||||
retry_count: task.retry_count,
|
||||
poll_interval_seconds: task.poll_interval_seconds,
|
||||
next_poll_at_unix_secs: None,
|
||||
poll_count: task.poll_count,
|
||||
max_poll_count: task.max_poll_count,
|
||||
created_at_unix_secs: task.created_at_unix_secs,
|
||||
submitted_at_unix_secs: task.submitted_at_unix_secs,
|
||||
completed_at_unix_secs: Some(now_unix_secs),
|
||||
updated_at_unix_secs: now_unix_secs,
|
||||
error_code: task.error_code.clone(),
|
||||
error_message: task.error_message.clone(),
|
||||
video_url: task.video_url.clone(),
|
||||
request_metadata,
|
||||
})
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
22
apps/aether-gateway/src/async_task/mod.rs
Normal file
22
apps/aether-gateway/src/async_task/mod.rs
Normal file
@@ -0,0 +1,22 @@
|
||||
#[path = "../video_tasks/mod.rs"]
|
||||
pub(crate) mod video;
|
||||
|
||||
mod http;
|
||||
mod query;
|
||||
mod runtime;
|
||||
|
||||
pub(crate) use http::{
|
||||
build_video_task_video_response, cancel_video_task, cancel_video_task_record,
|
||||
get_video_task_detail, get_video_task_stats, get_video_task_video, list_video_tasks,
|
||||
CancelVideoTaskError,
|
||||
};
|
||||
pub(crate) use query::{
|
||||
read_video_task_detail, read_video_task_page, read_video_task_stats,
|
||||
read_video_task_video_source,
|
||||
};
|
||||
pub(crate) use runtime::{
|
||||
execute_video_task_refresh_plan, finalize_video_task_if_terminal, spawn_video_task_poller,
|
||||
VideoTaskPollerConfig,
|
||||
};
|
||||
pub(crate) use video::VideoTaskService;
|
||||
pub use video::VideoTaskTruthSourceMode;
|
||||
201
apps/aether-gateway/src/async_task/query.rs
Normal file
201
apps/aether-gateway/src/async_task/query.rs
Normal file
@@ -0,0 +1,201 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_data::repository::video_tasks::{
|
||||
StoredVideoTask, VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount,
|
||||
};
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::gateway::{AppState, GatewayError};
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(crate) struct VideoTaskPageResponse {
|
||||
pub(crate) items: Vec<StoredVideoTask>,
|
||||
pub(crate) total: u64,
|
||||
pub(crate) page: usize,
|
||||
pub(crate) page_size: usize,
|
||||
pub(crate) pages: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(crate) struct VideoTaskStatsResponse {
|
||||
pub(crate) total: u64,
|
||||
pub(crate) by_status: BTreeMap<String, u64>,
|
||||
pub(crate) by_model: BTreeMap<String, u64>,
|
||||
pub(crate) today_count: u64,
|
||||
pub(crate) processing_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) enum VideoTaskVideoSource {
|
||||
Redirect {
|
||||
url: String,
|
||||
},
|
||||
Proxy {
|
||||
url: String,
|
||||
header_name: String,
|
||||
header_value: String,
|
||||
filename: String,
|
||||
},
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_page(
|
||||
state: &AppState,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
page: usize,
|
||||
page_size: usize,
|
||||
) -> Result<VideoTaskPageResponse, GatewayError> {
|
||||
let page = page.max(1);
|
||||
let page_size = page_size.clamp(1, 100);
|
||||
let total = state.count_video_tasks(filter).await?;
|
||||
let offset = page_size.saturating_mul(page.saturating_sub(1));
|
||||
let items = state
|
||||
.list_video_task_page(filter, offset, page_size)
|
||||
.await?;
|
||||
let pages = if total == 0 {
|
||||
0
|
||||
} else {
|
||||
((total as usize) + page_size - 1) / page_size
|
||||
};
|
||||
|
||||
Ok(VideoTaskPageResponse {
|
||||
items,
|
||||
total,
|
||||
page,
|
||||
page_size,
|
||||
pages,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_detail(
|
||||
state: &AppState,
|
||||
task_id: &str,
|
||||
) -> Result<Option<StoredVideoTask>, GatewayError> {
|
||||
state.find_video_task_by_id(task_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_video_source(
|
||||
state: &AppState,
|
||||
task_id: &str,
|
||||
) -> Result<Option<VideoTaskVideoSource>, GatewayError> {
|
||||
let Some(task) = read_video_task_detail(state, task_id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(video_url) = task
|
||||
.video_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if !video_url.contains("generativelanguage.googleapis.com") {
|
||||
return Ok(Some(VideoTaskVideoSource::Redirect { url: video_url }));
|
||||
}
|
||||
|
||||
let Some(provider_id) = task.provider_id.as_deref() else {
|
||||
return Err(GatewayError::Internal(
|
||||
"video task is missing provider_id for proxied video".to_string(),
|
||||
));
|
||||
};
|
||||
let Some(endpoint_id) = task.endpoint_id.as_deref() else {
|
||||
return Err(GatewayError::Internal(
|
||||
"video task is missing endpoint_id for proxied video".to_string(),
|
||||
));
|
||||
};
|
||||
let Some(key_id) = task.key_id.as_deref() else {
|
||||
return Err(GatewayError::Internal(
|
||||
"video task is missing key_id for proxied video".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
let Some(transport) = state
|
||||
.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
|
||||
.await?
|
||||
else {
|
||||
return Err(GatewayError::Internal(
|
||||
"provider transport snapshot is unavailable for proxied video".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
let api_key = transport.key.decrypted_api_key.trim();
|
||||
if api_key.is_empty() {
|
||||
return Err(GatewayError::Internal(
|
||||
"provider transport key is unavailable for proxied video".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Some(VideoTaskVideoSource::Proxy {
|
||||
url: video_url,
|
||||
header_name: "x-goog-api-key".to_string(),
|
||||
header_value: api_key.to_string(),
|
||||
filename: format!("video_{task_id}.mp4"),
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_stats(
|
||||
state: &AppState,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<VideoTaskStatsResponse, GatewayError> {
|
||||
let total = state.count_video_tasks(filter).await?;
|
||||
let by_status = state.count_video_tasks_by_status(filter).await?;
|
||||
let by_model = state.top_video_task_models(filter, 10).await?;
|
||||
let today_count = state
|
||||
.count_video_tasks_created_since(filter, start_of_utc_day(now_unix_secs))
|
||||
.await?;
|
||||
let processing_count = by_status
|
||||
.iter()
|
||||
.filter(|entry| {
|
||||
matches!(
|
||||
entry.status,
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Submitted
|
||||
| aether_data::repository::video_tasks::VideoTaskStatus::Queued
|
||||
| aether_data::repository::video_tasks::VideoTaskStatus::Processing
|
||||
)
|
||||
})
|
||||
.map(|entry| entry.count)
|
||||
.sum();
|
||||
|
||||
Ok(VideoTaskStatsResponse {
|
||||
total,
|
||||
by_status: map_status_counts(by_status),
|
||||
by_model: map_model_counts(by_model),
|
||||
today_count,
|
||||
processing_count,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_status_counts(counts: Vec<VideoTaskStatusCount>) -> BTreeMap<String, u64> {
|
||||
counts
|
||||
.into_iter()
|
||||
.map(|entry| (status_key(entry.status), entry.count))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn map_model_counts(counts: Vec<VideoTaskModelCount>) -> BTreeMap<String, u64> {
|
||||
counts
|
||||
.into_iter()
|
||||
.map(|entry| (entry.model, entry.count))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn status_key(status: aether_data::repository::video_tasks::VideoTaskStatus) -> String {
|
||||
match status {
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Pending => "pending",
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Submitted => "submitted",
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Queued => "queued",
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Processing => "processing",
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Completed => "completed",
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Failed => "failed",
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Cancelled => "cancelled",
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Expired => "expired",
|
||||
aether_data::repository::video_tasks::VideoTaskStatus::Deleted => "deleted",
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn start_of_utc_day(now_unix_secs: u64) -> u64 {
|
||||
now_unix_secs - (now_unix_secs % 86_400)
|
||||
}
|
||||
497
apps/aether-gateway/src/async_task/runtime.rs
Normal file
497
apps/aether-gateway/src/async_task/runtime.rs
Normal file
@@ -0,0 +1,497 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_contracts::{ExecutionErrorKind, ExecutionResult};
|
||||
use aether_data::repository::video_tasks::{StoredVideoTask, UpsertVideoTask, VideoTaskStatus};
|
||||
use serde_json::{Map, Value};
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::warn;
|
||||
|
||||
use super::video::{LocalVideoTaskReadRefreshPlan, LocalVideoTaskSnapshot};
|
||||
use crate::gateway::billing_runtime::enrich_usage_event_with_billing;
|
||||
use crate::gateway::usage::{
|
||||
build_upsert_usage_record_from_event, UsageEvent, UsageEventData, UsageEventType,
|
||||
};
|
||||
use crate::gateway::wallet_runtime::settle_usage_if_needed;
|
||||
use crate::gateway::{AppState, GatewayError};
|
||||
|
||||
const MAX_VIDEO_TASK_POLL_BACKOFF_SECONDS: u64 = 300;
|
||||
const VIDEO_TASK_POLL_CLAIM_SECONDS: u64 = 30;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct VideoTaskRefreshError {
|
||||
message: String,
|
||||
permanent: bool,
|
||||
}
|
||||
|
||||
enum VideoTaskRefreshAttempt {
|
||||
Success { provider_body: Map<String, Value> },
|
||||
Error(VideoTaskRefreshError),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct VideoTaskPollerConfig {
|
||||
pub(crate) interval: Duration,
|
||||
pub(crate) batch_size: usize,
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_video_task_refresh_plan(
|
||||
state: &AppState,
|
||||
refresh_plan: &LocalVideoTaskReadRefreshPlan,
|
||||
) -> Result<bool, GatewayError> {
|
||||
match fetch_video_task_refresh_attempt(state, refresh_plan).await? {
|
||||
VideoTaskRefreshAttempt::Success { provider_body } => {
|
||||
let projected = state
|
||||
.video_tasks
|
||||
.apply_read_refresh_projection(refresh_plan, &provider_body);
|
||||
if projected {
|
||||
if let Some(snapshot) = state.video_tasks.snapshot_for_refresh_plan(refresh_plan) {
|
||||
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
|
||||
}
|
||||
}
|
||||
Ok(projected)
|
||||
}
|
||||
VideoTaskRefreshAttempt::Error(err) => {
|
||||
warn!(error = %err.message, permanent = err.permanent, "gateway video task refresh failed");
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn poll_video_tasks_once(state: &AppState, batch_size: usize) -> Result<usize, GatewayError> {
|
||||
if !state.video_tasks.is_rust_authoritative() {
|
||||
return Ok(0);
|
||||
}
|
||||
let now_unix_secs = now_unix_secs();
|
||||
let tasks = state
|
||||
.claim_due_video_tasks(
|
||||
now_unix_secs,
|
||||
now_unix_secs.saturating_add(VIDEO_TASK_POLL_CLAIM_SECONDS),
|
||||
batch_size,
|
||||
)
|
||||
.await?;
|
||||
let mut refreshed = 0usize;
|
||||
for (index, task) in tasks.into_iter().enumerate() {
|
||||
let trace_id = format!("video-task-poller-{index}");
|
||||
let Some(refresh_plan) = state
|
||||
.video_tasks
|
||||
.prepare_poll_refresh_plan_for_stored_task(&task, &trace_id)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match fetch_video_task_refresh_attempt(state, &refresh_plan).await? {
|
||||
VideoTaskRefreshAttempt::Success { provider_body } => {
|
||||
let Some(updated) =
|
||||
build_successful_poll_update(&task, &provider_body, now_unix_secs)?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
match state.update_active_video_task(updated).await? {
|
||||
Some(stored) => {
|
||||
if let Some(snapshot) = LocalVideoTaskSnapshot::from_stored_task(&stored) {
|
||||
state.video_tasks.record_snapshot(snapshot);
|
||||
}
|
||||
finalize_video_task_if_terminal(state, &stored).await;
|
||||
refreshed += 1;
|
||||
}
|
||||
None => continue,
|
||||
}
|
||||
}
|
||||
VideoTaskRefreshAttempt::Error(err) => {
|
||||
let updated = build_failed_poll_update(&task, &err, now_unix_secs);
|
||||
match state.update_active_video_task(updated).await? {
|
||||
Some(stored) => {
|
||||
if let Some(snapshot) = LocalVideoTaskSnapshot::from_stored_task(&stored) {
|
||||
state.video_tasks.record_snapshot(snapshot);
|
||||
}
|
||||
finalize_video_task_if_terminal(state, &stored).await;
|
||||
refreshed += 1;
|
||||
}
|
||||
None => continue,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(refreshed)
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_video_task_poller(state: AppState) -> Option<JoinHandle<()>> {
|
||||
let config = state.video_task_poller?;
|
||||
if !state.video_tasks.is_rust_authoritative() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(config.interval);
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
interval.tick().await;
|
||||
loop {
|
||||
interval.tick().await;
|
||||
if let Err(err) = poll_video_tasks_once(&state, config.batch_size).await {
|
||||
warn!(error = ?err, "gateway video task poller tick failed");
|
||||
}
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
async fn fetch_video_task_refresh_attempt(
|
||||
state: &AppState,
|
||||
refresh_plan: &LocalVideoTaskReadRefreshPlan,
|
||||
) -> Result<VideoTaskRefreshAttempt, GatewayError> {
|
||||
let result =
|
||||
match crate::gateway::execute_execution_runtime_sync_plan(state, None, &refresh_plan.plan)
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
return Ok(VideoTaskRefreshAttempt::Error(VideoTaskRefreshError {
|
||||
message: format!("{err:?}"),
|
||||
permanent: false,
|
||||
}));
|
||||
}
|
||||
};
|
||||
if result.status_code >= 400 {
|
||||
return Ok(VideoTaskRefreshAttempt::Error(
|
||||
classify_refresh_result_error(&result),
|
||||
));
|
||||
}
|
||||
|
||||
let Some(provider_body) = result
|
||||
.body
|
||||
.and_then(|body| body.json_body)
|
||||
.and_then(|body| body.as_object().cloned())
|
||||
else {
|
||||
return Ok(VideoTaskRefreshAttempt::Error(VideoTaskRefreshError {
|
||||
message: "video task refresh missing json provider body".to_string(),
|
||||
permanent: false,
|
||||
}));
|
||||
};
|
||||
|
||||
Ok(VideoTaskRefreshAttempt::Success { provider_body })
|
||||
}
|
||||
|
||||
fn classify_refresh_result_error(result: &ExecutionResult) -> VideoTaskRefreshError {
|
||||
let status_code = result
|
||||
.error
|
||||
.as_ref()
|
||||
.and_then(|error| error.upstream_status)
|
||||
.unwrap_or(result.status_code);
|
||||
let message = result
|
||||
.error
|
||||
.as_ref()
|
||||
.map(|error| error.message.clone())
|
||||
.or_else(|| {
|
||||
result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
.and_then(|value| value.get("error"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
})
|
||||
.unwrap_or_else(|| format!("upstream returned {status_code}"));
|
||||
let permanent = result.error.as_ref().map_or(
|
||||
matches!(status_code, 400 | 401 | 403 | 404 | 422),
|
||||
|error| match error.kind {
|
||||
ExecutionErrorKind::Upstream4xx => !matches!(status_code, 408 | 409 | 429),
|
||||
ExecutionErrorKind::Upstream5xx
|
||||
| ExecutionErrorKind::ConnectTimeout
|
||||
| ExecutionErrorKind::FirstByteTimeout
|
||||
| ExecutionErrorKind::ReadTimeout
|
||||
| ExecutionErrorKind::TlsError
|
||||
| ExecutionErrorKind::ProxyError
|
||||
| ExecutionErrorKind::ProtocolError
|
||||
| ExecutionErrorKind::Internal => false,
|
||||
ExecutionErrorKind::Cancelled => true,
|
||||
},
|
||||
);
|
||||
|
||||
VideoTaskRefreshError { message, permanent }
|
||||
}
|
||||
|
||||
fn build_successful_poll_update(
|
||||
task: &StoredVideoTask,
|
||||
provider_body: &Map<String, Value>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<UpsertVideoTask>, GatewayError> {
|
||||
let Some(mut snapshot) = LocalVideoTaskSnapshot::from_stored_task(task) else {
|
||||
return Ok(None);
|
||||
};
|
||||
snapshot.apply_provider_body(provider_body);
|
||||
|
||||
let mut record = snapshot.to_upsert_record();
|
||||
record.id = task.id.clone();
|
||||
record.short_id = task.short_id.clone().or(record.short_id);
|
||||
record.request_id = task.request_id.clone();
|
||||
record.user_id = task.user_id.clone();
|
||||
record.api_key_id = task.api_key_id.clone();
|
||||
record.username = task.username.clone();
|
||||
record.api_key_name = task.api_key_name.clone();
|
||||
record.external_task_id = task.external_task_id.clone().or(record.external_task_id);
|
||||
record.provider_id = task.provider_id.clone();
|
||||
record.endpoint_id = task.endpoint_id.clone();
|
||||
record.key_id = task.key_id.clone();
|
||||
record.client_api_format = task.client_api_format.clone();
|
||||
record.provider_api_format = task.provider_api_format.clone();
|
||||
record.format_converted = task.format_converted;
|
||||
record.model = task.model.clone().or(record.model);
|
||||
record.prompt = task.prompt.clone().or(record.prompt);
|
||||
record.original_request_body = task.original_request_body.clone();
|
||||
record.duration_seconds = task.duration_seconds;
|
||||
record.resolution = task.resolution.clone();
|
||||
record.aspect_ratio = task.aspect_ratio.clone();
|
||||
record.size = task.size.clone();
|
||||
record.created_at_unix_secs = task.created_at_unix_secs;
|
||||
record.submitted_at_unix_secs = task.submitted_at_unix_secs;
|
||||
record.updated_at_unix_secs = now_unix_secs;
|
||||
record.retry_count = task.retry_count;
|
||||
record.poll_interval_seconds = task.poll_interval_seconds.max(1);
|
||||
record.poll_count = task.poll_count.saturating_add(1);
|
||||
record.max_poll_count = task.max_poll_count.max(1);
|
||||
record.next_poll_at_unix_secs = if record.status.is_active() {
|
||||
Some(now_unix_secs.saturating_add(u64::from(record.poll_interval_seconds)))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if !record.status.is_active() && record.completed_at_unix_secs.is_none() {
|
||||
record.completed_at_unix_secs = Some(now_unix_secs);
|
||||
}
|
||||
if record.status.is_active() && record.poll_count >= record.max_poll_count {
|
||||
record.status = VideoTaskStatus::Failed;
|
||||
record.error_code = Some("poll_timeout".to_string());
|
||||
record.error_message = Some(format!("Task timed out after {} polls", record.poll_count));
|
||||
record.completed_at_unix_secs = Some(now_unix_secs);
|
||||
record.next_poll_at_unix_secs = None;
|
||||
}
|
||||
record.request_metadata = merge_video_task_request_metadata(
|
||||
task.request_metadata.clone(),
|
||||
&snapshot,
|
||||
Some(provider_body),
|
||||
None,
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
|
||||
Ok(Some(record))
|
||||
}
|
||||
|
||||
fn build_failed_poll_update(
|
||||
task: &StoredVideoTask,
|
||||
err: &VideoTaskRefreshError,
|
||||
now_unix_secs: u64,
|
||||
) -> UpsertVideoTask {
|
||||
let mut record = stored_task_to_upsert(task);
|
||||
record.updated_at_unix_secs = now_unix_secs;
|
||||
record.poll_count = task.poll_count.saturating_add(1);
|
||||
record.progress_message = Some(format!("Poll error: {}", err.message));
|
||||
if err.permanent {
|
||||
record.status = VideoTaskStatus::Failed;
|
||||
record.error_code = Some("poll_permanent_error".to_string());
|
||||
record.error_message = Some(err.message.clone());
|
||||
record.completed_at_unix_secs = Some(now_unix_secs);
|
||||
record.next_poll_at_unix_secs = None;
|
||||
} else {
|
||||
let backoff =
|
||||
compute_poll_backoff_seconds(task.poll_interval_seconds.max(1), task.retry_count);
|
||||
record.retry_count = task.retry_count.saturating_add(1);
|
||||
record.next_poll_at_unix_secs = Some(now_unix_secs.saturating_add(backoff));
|
||||
}
|
||||
if record.status.is_active() && record.poll_count >= record.max_poll_count {
|
||||
record.status = VideoTaskStatus::Failed;
|
||||
record.error_code = Some("poll_timeout".to_string());
|
||||
record.error_message = Some(format!("Task timed out after {} polls", record.poll_count));
|
||||
record.completed_at_unix_secs = Some(now_unix_secs);
|
||||
record.next_poll_at_unix_secs = None;
|
||||
}
|
||||
record.request_metadata = LocalVideoTaskSnapshot::from_stored_task(task)
|
||||
.and_then(|snapshot| {
|
||||
merge_video_task_request_metadata(
|
||||
task.request_metadata.clone(),
|
||||
&snapshot,
|
||||
None,
|
||||
Some(err),
|
||||
)
|
||||
.ok()
|
||||
.flatten()
|
||||
})
|
||||
.or(task.request_metadata.clone());
|
||||
record
|
||||
}
|
||||
|
||||
fn stored_task_to_upsert(task: &StoredVideoTask) -> UpsertVideoTask {
|
||||
UpsertVideoTask {
|
||||
id: task.id.clone(),
|
||||
short_id: task.short_id.clone(),
|
||||
request_id: task.request_id.clone(),
|
||||
user_id: task.user_id.clone(),
|
||||
api_key_id: task.api_key_id.clone(),
|
||||
username: task.username.clone(),
|
||||
api_key_name: task.api_key_name.clone(),
|
||||
external_task_id: task.external_task_id.clone(),
|
||||
provider_id: task.provider_id.clone(),
|
||||
endpoint_id: task.endpoint_id.clone(),
|
||||
key_id: task.key_id.clone(),
|
||||
client_api_format: task.client_api_format.clone(),
|
||||
provider_api_format: task.provider_api_format.clone(),
|
||||
format_converted: task.format_converted,
|
||||
model: task.model.clone(),
|
||||
prompt: task.prompt.clone(),
|
||||
original_request_body: task.original_request_body.clone(),
|
||||
duration_seconds: task.duration_seconds,
|
||||
resolution: task.resolution.clone(),
|
||||
aspect_ratio: task.aspect_ratio.clone(),
|
||||
size: task.size.clone(),
|
||||
status: task.status,
|
||||
progress_percent: task.progress_percent,
|
||||
progress_message: task.progress_message.clone(),
|
||||
retry_count: task.retry_count,
|
||||
poll_interval_seconds: task.poll_interval_seconds.max(1),
|
||||
next_poll_at_unix_secs: task.next_poll_at_unix_secs,
|
||||
poll_count: task.poll_count,
|
||||
max_poll_count: task.max_poll_count.max(1),
|
||||
created_at_unix_secs: task.created_at_unix_secs,
|
||||
submitted_at_unix_secs: task.submitted_at_unix_secs,
|
||||
completed_at_unix_secs: task.completed_at_unix_secs,
|
||||
updated_at_unix_secs: task.updated_at_unix_secs,
|
||||
error_code: task.error_code.clone(),
|
||||
error_message: task.error_message.clone(),
|
||||
video_url: task.video_url.clone(),
|
||||
request_metadata: task.request_metadata.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_poll_backoff_seconds(poll_interval_seconds: u32, retry_count: u32) -> u64 {
|
||||
let exponent = retry_count.min(5);
|
||||
let multiplier = 1u64 << exponent;
|
||||
u64::from(poll_interval_seconds)
|
||||
.saturating_mul(multiplier)
|
||||
.min(MAX_VIDEO_TASK_POLL_BACKOFF_SECONDS)
|
||||
}
|
||||
|
||||
fn merge_video_task_request_metadata(
|
||||
existing: Option<Value>,
|
||||
snapshot: &LocalVideoTaskSnapshot,
|
||||
provider_body: Option<&Map<String, Value>>,
|
||||
poll_error: Option<&VideoTaskRefreshError>,
|
||||
) -> Result<Option<Value>, serde_json::Error> {
|
||||
let mut metadata = match existing {
|
||||
Some(Value::Object(object)) => object,
|
||||
_ => Map::new(),
|
||||
};
|
||||
metadata.insert(
|
||||
"rust_owner".to_string(),
|
||||
Value::String("async_task".to_string()),
|
||||
);
|
||||
metadata.insert(
|
||||
"rust_local_snapshot".to_string(),
|
||||
serde_json::to_value(snapshot)?,
|
||||
);
|
||||
if let Some(provider_body) = provider_body {
|
||||
metadata.insert(
|
||||
"poll_raw_response".to_string(),
|
||||
Value::Object(provider_body.clone()),
|
||||
);
|
||||
metadata.remove("poll_error");
|
||||
}
|
||||
if let Some(poll_error) = poll_error {
|
||||
metadata.insert(
|
||||
"poll_error".to_string(),
|
||||
serde_json::json!({
|
||||
"message": poll_error.message,
|
||||
"permanent": poll_error.permanent,
|
||||
"observed_at_unix_secs": now_unix_secs(),
|
||||
}),
|
||||
);
|
||||
}
|
||||
Ok(Some(Value::Object(metadata)))
|
||||
}
|
||||
|
||||
pub(crate) async fn finalize_video_task_if_terminal(state: &AppState, task: &StoredVideoTask) {
|
||||
let Some(event) = build_video_task_terminal_usage_event(task) else {
|
||||
return;
|
||||
};
|
||||
let mut event = event;
|
||||
if let Err(err) = enrich_usage_event_with_billing(&state.data, &mut event).await {
|
||||
warn!(error = %err, request_id = %task.request_id, "gateway video task finalize failed to enrich billing");
|
||||
}
|
||||
match build_upsert_usage_record_from_event(&event) {
|
||||
Ok(record) => match state.data.upsert_usage(record).await {
|
||||
Ok(Some(stored)) => {
|
||||
if let Err(err) = settle_usage_if_needed(&state.data, &stored).await {
|
||||
warn!(error = %err, request_id = %task.request_id, "gateway video task finalize failed to settle usage");
|
||||
}
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(error = %err, request_id = %task.request_id, "gateway video task finalize failed to upsert usage");
|
||||
}
|
||||
},
|
||||
Err(err) => {
|
||||
warn!(error = %err, request_id = %task.request_id, "gateway video task finalize failed to build usage record");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_video_task_terminal_usage_event(task: &StoredVideoTask) -> Option<UsageEvent> {
|
||||
let event_type = match task.status {
|
||||
VideoTaskStatus::Completed => UsageEventType::Completed,
|
||||
VideoTaskStatus::Failed | VideoTaskStatus::Expired => UsageEventType::Failed,
|
||||
VideoTaskStatus::Cancelled | VideoTaskStatus::Deleted => UsageEventType::Cancelled,
|
||||
VideoTaskStatus::Pending
|
||||
| VideoTaskStatus::Submitted
|
||||
| VideoTaskStatus::Queued
|
||||
| VideoTaskStatus::Processing => {
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let provider_name = LocalVideoTaskSnapshot::from_stored_task(task)
|
||||
.and_then(|snapshot| snapshot.provider_name().map(str::to_string))
|
||||
.or_else(|| task.provider_id.clone())
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let response_time_ms = task
|
||||
.submitted_at_unix_secs
|
||||
.zip(
|
||||
task.completed_at_unix_secs
|
||||
.or(Some(task.updated_at_unix_secs)),
|
||||
)
|
||||
.map(|(submitted, completed)| completed.saturating_sub(submitted).saturating_mul(1_000));
|
||||
let status_code = match event_type {
|
||||
UsageEventType::Completed => Some(200),
|
||||
UsageEventType::Cancelled => Some(499),
|
||||
UsageEventType::Failed => Some(500),
|
||||
UsageEventType::Pending | UsageEventType::Streaming => None,
|
||||
};
|
||||
|
||||
Some(UsageEvent::new(
|
||||
event_type,
|
||||
task.request_id.clone(),
|
||||
UsageEventData {
|
||||
user_id: task.user_id.clone(),
|
||||
api_key_id: task.api_key_id.clone(),
|
||||
username: task.username.clone(),
|
||||
api_key_name: task.api_key_name.clone(),
|
||||
provider_name,
|
||||
model: task.model.clone().unwrap_or_else(|| "unknown".to_string()),
|
||||
provider_id: task.provider_id.clone(),
|
||||
provider_endpoint_id: task.endpoint_id.clone(),
|
||||
provider_api_key_id: task.key_id.clone(),
|
||||
request_type: Some("video".to_string()),
|
||||
api_format: task.client_api_format.clone(),
|
||||
endpoint_api_format: task.provider_api_format.clone(),
|
||||
has_format_conversion: Some(task.format_converted),
|
||||
is_stream: Some(false),
|
||||
status_code,
|
||||
error_message: task.error_message.clone().or(task.error_code.clone()),
|
||||
response_time_ms,
|
||||
request_body: task.original_request_body.clone(),
|
||||
request_metadata: task.request_metadata.clone(),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
fn now_unix_secs() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
Reference in New Issue
Block a user