feat(api): 添加多媒体内容支持并优化钱包计费系统
- 在 API 接口定义中为 video_url 和 audio_url 类型添加 mime_type 字段 - 实现 Google Gemini 客户端对视频和音频内容的支持,包括媒体类型检测和数据传输 - 添加 Gemini 客户端测试用例验证多媒体内容转换功能 - 重构 Playground 页面的媒体上传逻辑以支持 MIME 类型传递 - 实现钱包计费预留机制,确保任务执行前余额充足 - 添加钱包冻结余额管理,防止并发操作导致的超扣问题 - 实现计费预留释放逻辑,处理任务失败或取消情况下的资金返还 - 优化数据库事务处理,确保计费操作的原子性和一致性 - 添加数据库集成测试验证迁移脚本执行流程 - 统一 Google Gemini 相关模型提供商标识符映射
This commit is contained in:
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -92,6 +93,16 @@ type WalletAdjustmentResult struct {
|
||||
Transaction GatewayWalletTransaction `json:"transaction"`
|
||||
}
|
||||
|
||||
type WalletBillingReservation struct {
|
||||
TaskID string `json:"taskId"`
|
||||
AccountID string `json:"accountId"`
|
||||
GatewayUserID string `json:"gatewayUserId"`
|
||||
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
|
||||
Currency string `json:"currency"`
|
||||
Amount float64 `json:"amount"`
|
||||
IdempotencyKey string `json:"idempotencyKey"`
|
||||
}
|
||||
|
||||
func (s *Store) WalletAvailability(ctx context.Context, user *auth.User, currency string, requiredAmount float64) (WalletAvailability, error) {
|
||||
gatewayUserID := localGatewayUserID(user)
|
||||
if gatewayUserID == "" {
|
||||
@@ -115,6 +126,223 @@ func (s *Store) WalletAvailability(ctx context.Context, user *auth.User, currenc
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *Store) ReserveTaskBilling(ctx context.Context, task GatewayTask, user *auth.User, billings []any) ([]WalletBillingReservation, error) {
|
||||
gatewayUserID := taskGatewayUserID(task, user)
|
||||
if gatewayUserID == "" {
|
||||
return nil, nil
|
||||
}
|
||||
taskID := strings.TrimSpace(task.ID)
|
||||
if taskID == "" {
|
||||
return nil, fmt.Errorf("task id is required for wallet reservation")
|
||||
}
|
||||
|
||||
amounts := walletBillingAmounts(billings)
|
||||
if len(amounts) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
reservations := make([]WalletBillingReservation, 0, len(amounts))
|
||||
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
for currency, rawAmount := range amounts {
|
||||
amount := roundMoney(rawAmount)
|
||||
if amount <= 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
account, err := s.ensureWalletAccount(ctx, tx, gatewayUserID, currency)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
locked, err := lockWalletAccount(ctx, tx, account.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
activeKey, activeAmount, err := activeWalletReservation(ctx, tx, locked.ID, taskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if activeAmount > 0 {
|
||||
reservation := WalletBillingReservation{
|
||||
TaskID: taskID,
|
||||
AccountID: locked.ID,
|
||||
GatewayUserID: gatewayUserID,
|
||||
GatewayTenantID: firstNonEmpty(locked.GatewayTenantID, task.GatewayTenantID),
|
||||
Currency: locked.Currency,
|
||||
Amount: activeAmount,
|
||||
IdempotencyKey: activeKey,
|
||||
}
|
||||
reservations = append(reservations, reservation)
|
||||
continue
|
||||
}
|
||||
|
||||
sequence, err := nextWalletReservationSequence(ctx, tx, locked.ID, taskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key := billingReservationIdempotencyKey(taskID, locked.Currency, sequence)
|
||||
reservation := WalletBillingReservation{
|
||||
TaskID: taskID,
|
||||
AccountID: locked.ID,
|
||||
GatewayUserID: gatewayUserID,
|
||||
GatewayTenantID: firstNonEmpty(locked.GatewayTenantID, task.GatewayTenantID),
|
||||
Currency: locked.Currency,
|
||||
Amount: amount,
|
||||
IdempotencyKey: key,
|
||||
}
|
||||
available := roundMoney(locked.Balance - locked.FrozenBalance)
|
||||
if available+0.000001 < amount {
|
||||
return fmt.Errorf("%w: required %.6f %s, available %.6f", ErrInsufficientWalletBalance, amount, locked.Currency, available)
|
||||
}
|
||||
|
||||
frozenAfter := roundMoney(locked.FrozenBalance + amount)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET frozen_balance = $2,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, locked.ID, frozenAfter); err != nil {
|
||||
return err
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID,
|
||||
"kind": task.Kind,
|
||||
"model": task.Model,
|
||||
"reserved": amount,
|
||||
"balance": roundMoney(locked.Balance),
|
||||
"frozenBefore": roundMoney(locked.FrozenBalance),
|
||||
"frozenAfter": frozenAfter,
|
||||
})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_transactions (
|
||||
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
|
||||
amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'debit', 'reserve',
|
||||
$4, $5, $6, $7, 'gateway_task', $8, $9::jsonb
|
||||
)`,
|
||||
locked.ID,
|
||||
firstNonEmpty(locked.GatewayTenantID, task.GatewayTenantID),
|
||||
gatewayUserID,
|
||||
amount,
|
||||
roundMoney(locked.Balance),
|
||||
roundMoney(locked.Balance),
|
||||
key,
|
||||
taskID,
|
||||
string(metadata),
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
reservations = append(reservations, reservation)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return reservations, err
|
||||
}
|
||||
|
||||
func (s *Store) ReleaseTaskBillingReservations(ctx context.Context, reservations []WalletBillingReservation, reason string) error {
|
||||
if len(reservations) == 0 {
|
||||
return nil
|
||||
}
|
||||
reason = strings.TrimSpace(reason)
|
||||
if reason == "" {
|
||||
reason = "task_not_settled"
|
||||
}
|
||||
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
for _, reservation := range reservations {
|
||||
if reservation.Amount <= 0 || strings.TrimSpace(reservation.AccountID) == "" {
|
||||
continue
|
||||
}
|
||||
reserveKey := strings.TrimSpace(reservation.IdempotencyKey)
|
||||
if reserveKey == "" {
|
||||
reserveKey = billingReservationIdempotencyKey(reservation.TaskID, reservation.Currency, 1)
|
||||
}
|
||||
releaseKey := billingReservationReleaseIdempotencyKey(reserveKey)
|
||||
locked, err := lockWalletAccount(ctx, tx, reservation.AccountID)
|
||||
if err != nil {
|
||||
if err == pgx.ErrNoRows {
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
var alreadyReleased bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT EXISTS (
|
||||
SELECT 1
|
||||
FROM gateway_wallet_transactions
|
||||
WHERE account_id = $1::uuid
|
||||
AND idempotency_key = $2
|
||||
)`, reservation.AccountID, releaseKey).Scan(&alreadyReleased); err != nil {
|
||||
return err
|
||||
}
|
||||
if alreadyReleased {
|
||||
continue
|
||||
}
|
||||
var storedReservedAmount float64
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE((
|
||||
SELECT amount::float8
|
||||
FROM gateway_wallet_transactions
|
||||
WHERE account_id = $1::uuid
|
||||
AND idempotency_key = $2
|
||||
AND transaction_type = 'reserve'
|
||||
LIMIT 1
|
||||
), 0)::float8`, reservation.AccountID, reserveKey).Scan(&storedReservedAmount); err != nil {
|
||||
return err
|
||||
}
|
||||
if storedReservedAmount <= 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
amount := roundMoney(storedReservedAmount)
|
||||
frozenAfter := roundMoney(locked.FrozenBalance - amount)
|
||||
if frozenAfter < 0 {
|
||||
frozenAfter = 0
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET frozen_balance = $2,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, locked.ID, frozenAfter); err != nil {
|
||||
return err
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": reservation.TaskID,
|
||||
"reason": reason,
|
||||
"reserved": amount,
|
||||
"frozenBefore": roundMoney(locked.FrozenBalance),
|
||||
"frozenAfter": frozenAfter,
|
||||
})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_transactions (
|
||||
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
|
||||
amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'credit', 'release',
|
||||
$4, $5, $6, $7, 'gateway_task', $8, $9::jsonb
|
||||
)
|
||||
ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING`,
|
||||
locked.ID,
|
||||
locked.GatewayTenantID,
|
||||
locked.GatewayUserID,
|
||||
amount,
|
||||
roundMoney(locked.Balance),
|
||||
roundMoney(locked.Balance),
|
||||
releaseKey,
|
||||
reservation.TaskID,
|
||||
string(metadata),
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) GetWalletSummary(ctx context.Context, user *auth.User, currency string) (WalletSummary, error) {
|
||||
gatewayUserID := localGatewayUserID(user)
|
||||
if gatewayUserID == "" {
|
||||
@@ -465,6 +693,124 @@ WHERE gateway_user_id = $1::uuid
|
||||
return account, nil
|
||||
}
|
||||
|
||||
func lockWalletAccount(ctx context.Context, tx pgx.Tx, accountID string) (GatewayWalletAccount, error) {
|
||||
return scanWalletAccount(tx.QueryRow(ctx, `
|
||||
SELECT id::text, COALESCE(gateway_tenant_id::text, ''), gateway_user_id::text,
|
||||
COALESCE(tenant_id, ''), COALESCE(tenant_key, ''), COALESCE(user_id, ''),
|
||||
currency, balance::float8, frozen_balance::float8, total_recharged::float8,
|
||||
total_spent::float8, status, metadata, created_at, updated_at
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE id = $1::uuid
|
||||
FOR UPDATE`, accountID))
|
||||
}
|
||||
|
||||
func activeWalletReservation(ctx context.Context, tx pgx.Tx, accountID string, taskID string) (string, float64, error) {
|
||||
var key string
|
||||
var amount float64
|
||||
err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(t.idempotency_key, ''), t.amount::float8
|
||||
FROM gateway_wallet_transactions t
|
||||
WHERE t.account_id = $1::uuid
|
||||
AND t.reference_type = 'gateway_task'
|
||||
AND t.reference_id = $2
|
||||
AND t.transaction_type = 'reserve'
|
||||
AND COALESCE(t.idempotency_key, '') <> ''
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM gateway_wallet_transactions r
|
||||
WHERE r.account_id = t.account_id
|
||||
AND r.transaction_type = 'release'
|
||||
AND r.idempotency_key = t.idempotency_key || ':release'
|
||||
)
|
||||
ORDER BY t.created_at DESC
|
||||
LIMIT 1`, accountID, taskID).Scan(&key, &amount)
|
||||
if err == pgx.ErrNoRows {
|
||||
return "", 0, nil
|
||||
}
|
||||
if err != nil {
|
||||
return "", 0, err
|
||||
}
|
||||
return key, roundMoney(amount), nil
|
||||
}
|
||||
|
||||
func nextWalletReservationSequence(ctx context.Context, tx pgx.Tx, accountID string, taskID string) (int, error) {
|
||||
var count int
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COUNT(*)::int
|
||||
FROM gateway_wallet_transactions
|
||||
WHERE account_id = $1::uuid
|
||||
AND reference_type = 'gateway_task'
|
||||
AND reference_id = $2
|
||||
AND transaction_type = 'reserve'`, accountID, taskID).Scan(&count); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return count + 1, nil
|
||||
}
|
||||
|
||||
func walletBillingAmounts(billings []any) map[string]float64 {
|
||||
amounts := map[string]float64{}
|
||||
for _, raw := range billings {
|
||||
line, _ := raw.(map[string]any)
|
||||
if line == nil {
|
||||
continue
|
||||
}
|
||||
amount := roundMoney(walletFloat(line["amount"]))
|
||||
if amount <= 0 {
|
||||
continue
|
||||
}
|
||||
currency := normalizeWalletCurrency(walletString(line["currency"]))
|
||||
amounts[currency] = roundMoney(amounts[currency] + amount)
|
||||
}
|
||||
return amounts
|
||||
}
|
||||
|
||||
func taskGatewayUserID(task GatewayTask, user *auth.User) string {
|
||||
return firstNonEmpty(strings.TrimSpace(task.GatewayUserID), localGatewayUserID(user))
|
||||
}
|
||||
|
||||
func billingReservationIdempotencyKey(taskID string, currency string, sequence int) string {
|
||||
if sequence <= 0 {
|
||||
sequence = 1
|
||||
}
|
||||
return "task:" + strings.TrimSpace(taskID) + ":wallet-reservation:" + normalizeWalletCurrency(currency) + ":" + strconv.Itoa(sequence)
|
||||
}
|
||||
|
||||
func billingReservationReleaseIdempotencyKey(reservationKey string) string {
|
||||
return strings.TrimSpace(reservationKey) + ":release"
|
||||
}
|
||||
|
||||
func walletString(value any) string {
|
||||
if text, ok := value.(string); ok {
|
||||
return strings.TrimSpace(text)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func walletFloat(value any) float64 {
|
||||
switch typed := value.(type) {
|
||||
case float64:
|
||||
return typed
|
||||
case float32:
|
||||
return float64(typed)
|
||||
case int:
|
||||
return float64(typed)
|
||||
case int64:
|
||||
return float64(typed)
|
||||
case json.Number:
|
||||
next, _ := typed.Float64()
|
||||
return next
|
||||
case string:
|
||||
next := strings.TrimSpace(typed)
|
||||
if next == "" {
|
||||
return 0
|
||||
}
|
||||
parsed, _ := strconv.ParseFloat(next, 64)
|
||||
return parsed
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeWalletCurrency(currency string) string {
|
||||
currency = strings.TrimSpace(currency)
|
||||
if currency == "" {
|
||||
|
||||
Reference in New Issue
Block a user