feat(api): 添加多媒体内容支持并优化钱包计费系统

- 在 API 接口定义中为 video_url 和 audio_url 类型添加 mime_type 字段
- 实现 Google Gemini 客户端对视频和音频内容的支持,包括媒体类型检测和数据传输
- 添加 Gemini 客户端测试用例验证多媒体内容转换功能
- 重构 Playground 页面的媒体上传逻辑以支持 MIME 类型传递
- 实现钱包计费预留机制,确保任务执行前余额充足
- 添加钱包冻结余额管理,防止并发操作导致的超扣问题
- 实现计费预留释放逻辑,处理任务失败或取消情况下的资金返还
- 优化数据库事务处理,确保计费操作的原子性和一致性
- 添加数据库集成测试验证迁移脚本执行流程
- 统一 Google Gemini 相关模型提供商标识符映射
This commit is contained in:
2026-05-22 23:46:08 +08:00
parent af9b281d34
commit 8ad5b06c18
12 changed files with 1104 additions and 82 deletions
+346
View File
@@ -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 == "" {