Genarrative/server-rs/crates/api-server/src/asset_billing.rs

use std::future::Future;

use axum::http::StatusCode;
use serde_json::json;
use spacetime_client::SpacetimeClientError;

use crate::{
    http_error::AppError,
    state::AppState,
    wallet_refund_outbox::{WalletRefundOutboxEnqueueOutcome, WalletRefundOutboxRecord},
};

pub(crate) const ASSET_OPERATION_POINTS_COST: u64 = 1;

/// 资产操作统一执行入口：业务层只声明操作类型与资源 ID，钱包扣退费由服务层收口。
pub(crate) async fn execute_billable_asset_operation<T, Fut>(
    state: &AppState,
    owner_user_id: &str,
    asset_kind: &str,
    asset_id: &str,
    operation: Fut,
) -> Result<T, AppError>
where
    Fut: Future<Output = Result<T, AppError>>,
{
    execute_billable_asset_operation_with_cost(
        state,
        owner_user_id,
        asset_kind,
        asset_id,
        ASSET_OPERATION_POINTS_COST,
        operation,
    )
    .await
}

/// 生图等特殊操作可声明独立泥点成本，避免修改全局资产操作默认价格。
pub(crate) async fn execute_billable_asset_operation_with_cost<T, Fut>(
    state: &AppState,
    owner_user_id: &str,
    asset_kind: &str,
    asset_id: &str,
    points_cost: u64,
    operation: Fut,
) -> Result<T, AppError>
where
    Fut: Future<Output = Result<T, AppError>>,
{
    let points_consumed =
        consume_asset_operation_points(state, owner_user_id, asset_kind, asset_id, points_cost)
            .await?;
    match operation.await {
        Ok(value) => Ok(value),
        Err(error) => {
            if points_consumed && should_refund_asset_operation_error(&error) {
                refund_asset_operation_points(
                    state,
                    owner_user_id,
                    asset_kind,
                    asset_id,
                    points_cost,
                )
                .await;
            }
            Err(error)
        }
    }
}

pub(crate) fn should_refund_asset_operation_error(error: &AppError) -> bool {
    let message = error.body_text();
    // 中文注释：worker lease guard 拒绝表示当前进程已失去队列写权限；
    // 这类 stale worker 失败不能补偿退款，否则可能冲掉后续合法 worker 的同一账本扣费。
    !(message.contains("external_generation_job")
        && (message.contains("lease")
            || message.contains("worker")
            || message.contains("job_kind")
            || message.contains("source_")
            || message.contains("owner_user_id")
            || message.contains("不存在")
            || message.contains("不是 running 状态")))
}

/// 资产操作统一预扣泥点；扣费流水 ID 由业务资源 ID 参与构造，保证重试幂等。
async fn consume_asset_operation_points(
    state: &AppState,
    owner_user_id: &str,
    asset_kind: &str,
    asset_id: &str,
    points_cost: u64,
) -> Result<bool, AppError> {
    if points_cost == 0 {
        return Ok(false);
    }

    let ledger_id = format!(
        "asset_operation_consume:{}:{}:{}",
        owner_user_id, asset_kind, asset_id
    );
    match state
        .spacetime_client()
        .consume_profile_wallet_points(
            owner_user_id.to_string(),
            points_cost,
            ledger_id,
            current_utc_micros(),
        )
        .await
    {
        Ok(_) => Ok(true),
        Err(error) => Err(map_asset_operation_wallet_error(error)),
    }
}

/// 外部生成或发布 mutation 失败后补偿退款；立即退款失败会进入 outbox，避免覆盖原始业务错误。
async fn refund_asset_operation_points(
    state: &AppState,
    owner_user_id: &str,
    asset_kind: &str,
    asset_id: &str,
    points_cost: u64,
) {
    let ledger_id = format!(
        "asset_operation_refund:{}:{}:{}",
        owner_user_id, asset_kind, asset_id
    );
    let created_at_micros = current_utc_micros();
    if let Err(error) = state
        .spacetime_client()
        .refund_profile_wallet_points(
            owner_user_id.to_string(),
            points_cost,
            ledger_id.clone(),
            created_at_micros,
        )
        .await
    {
        let refund_error = error.to_string();
        if let Some(outbox) = state.wallet_refund_outbox() {
            match outbox
                .enqueue(WalletRefundOutboxRecord {
                    owner_user_id: owner_user_id.to_string(),
                    amount: points_cost,
                    ledger_id: ledger_id.clone(),
                    created_at_micros,
                    asset_kind: asset_kind.to_string(),
                    asset_id: asset_id.to_string(),
                })
                .await
            {
                Ok(WalletRefundOutboxEnqueueOutcome::Enqueued) => {
                    tracing::warn!(
                        owner_user_id,
                        asset_kind,
                        asset_id,
                        ledger_id,
                        error = %refund_error,
                        "资产操作失败后的泥点退款立即执行失败，已写入 wallet refund outbox"
                    );
                    return;
                }
                Ok(WalletRefundOutboxEnqueueOutcome::Dropped { reason }) => {
                    tracing::error!(
                        owner_user_id,
                        asset_kind,
                        asset_id,
                        ledger_id,
                        reason,
                        error = %refund_error,
                        "资产操作失败后的泥点退款立即执行失败，且 wallet refund outbox 因容量限制丢弃"
                    );
                    return;
                }
                Err(outbox_error) => {
                    tracing::error!(
                        owner_user_id,
                        asset_kind,
                        asset_id,
                        ledger_id,
                        refund_error = %refund_error,
                        outbox_error = %outbox_error,
                        "资产操作失败后的泥点退款立即执行失败，且写入 wallet refund outbox 失败"
                    );
                    return;
                }
            }
        }
        tracing::error!(
            owner_user_id,
            asset_kind,
            asset_id,
            ledger_id,
            error = %refund_error,
            "资产操作失败后的泥点退款失败，且 wallet refund outbox 未启用"
        );
    }
}

pub(crate) fn map_asset_operation_wallet_error(error: SpacetimeClientError) -> AppError {
    let message = error.to_string();
    tracing::warn!(
        provider = "profile-wallet",
        error = %message,
        "资产操作泥点预扣失败"
    );
    let status = match &error {
        SpacetimeClientError::Procedure(message) if message.contains("泥点余额不足") => {
            StatusCode::CONFLICT
        }
        _ => StatusCode::BAD_GATEWAY,
    };

    AppError::from_status(status).with_details(json!({
        "provider": "profile-wallet",
        "message": message,
    }))
}

pub(crate) fn should_skip_asset_operation_billing_for_connectivity(
    error: &SpacetimeClientError,
) -> bool {
    match error {
        SpacetimeClientError::ConnectDropped | SpacetimeClientError::Timeout => true,
        SpacetimeClientError::Build(message)
        | SpacetimeClientError::Procedure(message)
        | SpacetimeClientError::Runtime(message) => {
            message.contains("503")
                || message.contains("Service Unavailable")
                || message.contains("Failed to connect")
                || message.contains("WebSocket")
                || message.contains("No such procedure")
                || message.contains("连接已断开")
                || message.contains("连接在返回结果前已断开")
        }
    }
}

fn current_utc_micros() -> i64 {
    time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64 / 1_000
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn asset_operation_connectivity_errors_are_classified_for_non_billing_fallbacks() {
        assert_eq!(ASSET_OPERATION_POINTS_COST, 1);
        assert!(should_skip_asset_operation_billing_for_connectivity(
            &SpacetimeClientError::ConnectDropped
        ));
        assert!(should_skip_asset_operation_billing_for_connectivity(
            &SpacetimeClientError::Runtime(
                "Failed to connect: HTTP error: 503 Service Unavailable".to_string(),
            ),
        ));
        assert!(should_skip_asset_operation_billing_for_connectivity(
            &SpacetimeClientError::Procedure(
                "No such procedure: consume_profile_wallet_points_and_return".to_string(),
            ),
        ));
        assert!(!should_skip_asset_operation_billing_for_connectivity(
            &SpacetimeClientError::Procedure("泥点余额不足".to_string()),
        ));
    }

    #[test]
    fn asset_operation_billing_does_not_refund_stale_worker_lease_errors() {
        let stale_error = AppError::from_status(StatusCode::BAD_GATEWAY).with_details(json!({
            "provider": "spacetimedb",
            "message": "external_generation_job lease 已过期",
        }));
        let completed_job_error =
            AppError::from_status(StatusCode::BAD_GATEWAY).with_details(json!({
                "provider": "spacetimedb",
                "message": "external_generation_job 当前不是 running 状态",
            }));
        let missing_job_error =
            AppError::from_status(StatusCode::BAD_GATEWAY).with_details(json!({
                "provider": "spacetimedb",
                "message": "external_generation_job 不存在",
            }));
        let ordinary_error = AppError::from_status(StatusCode::BAD_GATEWAY).with_details(json!({
            "provider": "vector-engine",
            "message": "图片生成失败",
        }));

        assert!(!should_refund_asset_operation_error(&stale_error));
        assert!(!should_refund_asset_operation_error(&completed_job_error));
        assert!(!should_refund_asset_operation_error(&missing_job_error));
        assert!(should_refund_asset_operation_error(&ordinary_error));
    }
}