feat(billing): 完成异步结算与请求执行闭环
统一任务成功、Attempt 与结算 Outbox 的事务边界，增加多实例安全结算、人工复核、请求幂等与执行租约。钱包决策使用九位精确金额，并通过审计保护约束保留流水事实；同时补充管理接口、指标与 PostgreSQL 集成测试。
This commit is contained in:
@@ -0,0 +1,637 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const (
|
||||
BillingSettlementBatchSize = 50
|
||||
BillingSettlementLockTimeout = 120 * time.Second
|
||||
BillingSettlementMaxAttempts = 20
|
||||
)
|
||||
|
||||
var ErrBillingSettlementNotRetryable = errors.New("billing settlement is not retryable")
|
||||
|
||||
type BillingSettlement struct {
|
||||
ID string `json:"id"`
|
||||
TaskID string `json:"taskId"`
|
||||
Action string `json:"action"`
|
||||
Amount float64 `json:"amount"`
|
||||
Currency string `json:"currency"`
|
||||
Status string `json:"status"`
|
||||
Attempts int `json:"attempts"`
|
||||
NextAttemptAt time.Time `json:"nextAttemptAt"`
|
||||
LockedBy string `json:"lockedBy,omitempty"`
|
||||
LockedAt string `json:"lockedAt,omitempty"`
|
||||
LastErrorCode string `json:"lastErrorCode,omitempty"`
|
||||
LastErrorMessage string `json:"lastErrorMessage,omitempty"`
|
||||
ManualReviewReason string `json:"manualReviewReason,omitempty"`
|
||||
PricingSnapshot map[string]any `json:"pricingSnapshot,omitempty"`
|
||||
Payload map[string]any `json:"payload,omitempty"`
|
||||
CompletedAt string `json:"completedAt,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
LockToken string `json:"-"`
|
||||
amountExact string
|
||||
}
|
||||
|
||||
type BillingSettlementListFilter struct {
|
||||
Status string
|
||||
Action string
|
||||
Page int
|
||||
PageSize int
|
||||
}
|
||||
|
||||
type BillingSettlementListResult struct {
|
||||
Items []BillingSettlement `json:"items"`
|
||||
Total int `json:"total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"pageSize"`
|
||||
}
|
||||
|
||||
type BillingMetricsSnapshot struct {
|
||||
SettlementBacklog int64
|
||||
SettlementDelaySecs float64
|
||||
ManualReview int64
|
||||
OrphanFrozen int64
|
||||
PricingUnavailable int64
|
||||
}
|
||||
|
||||
func (s *Store) BillingMetrics(ctx context.Context) (BillingMetricsSnapshot, error) {
|
||||
var snapshot BillingMetricsSnapshot
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
WITH open_outbox AS (
|
||||
SELECT created_at
|
||||
FROM settlement_outbox
|
||||
WHERE status IN ('pending', 'processing', 'retryable_failed')
|
||||
), active_reservations AS (
|
||||
SELECT reserve.account_id, COALESCE(SUM(reserve.amount), 0) AS amount
|
||||
FROM gateway_wallet_transactions reserve
|
||||
WHERE reserve.transaction_type = 'reserve'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions release
|
||||
WHERE release.account_id = reserve.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = reserve.idempotency_key || ':release'
|
||||
)
|
||||
GROUP BY reserve.account_id
|
||||
)
|
||||
SELECT
|
||||
(SELECT COUNT(*) FROM open_outbox),
|
||||
COALESCE((SELECT EXTRACT(EPOCH FROM now() - MIN(created_at)) FROM open_outbox), 0),
|
||||
(SELECT COUNT(*) FROM gateway_tasks WHERE billing_status = 'manual_review'),
|
||||
(SELECT COUNT(*) FROM gateway_wallet_accounts account
|
||||
LEFT JOIN active_reservations reservation ON reservation.account_id = account.id
|
||||
WHERE account.frozen_balance <> COALESCE(reservation.amount, 0)),
|
||||
(SELECT COUNT(*) FROM gateway_tasks WHERE error_code = 'pricing_unavailable')`).Scan(
|
||||
&snapshot.SettlementBacklog,
|
||||
&snapshot.SettlementDelaySecs,
|
||||
&snapshot.ManualReview,
|
||||
&snapshot.OrphanFrozen,
|
||||
&snapshot.PricingUnavailable,
|
||||
)
|
||||
return snapshot, err
|
||||
}
|
||||
|
||||
func (s *Store) ClaimBillingSettlements(ctx context.Context, workerID string, limit int, staleAfter time.Duration) ([]BillingSettlement, error) {
|
||||
if limit <= 0 || limit > BillingSettlementBatchSize {
|
||||
limit = BillingSettlementBatchSize
|
||||
}
|
||||
if staleAfter <= 0 {
|
||||
staleAfter = BillingSettlementLockTimeout
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
WITH picked AS (
|
||||
SELECT id
|
||||
FROM settlement_outbox
|
||||
WHERE (
|
||||
status IN ('pending', 'retryable_failed')
|
||||
AND next_attempt_at <= now()
|
||||
)
|
||||
OR (
|
||||
status = 'processing'
|
||||
AND locked_at < now() - ($3::int * interval '1 second')
|
||||
)
|
||||
ORDER BY next_attempt_at ASC, created_at ASC
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE settlement_outbox outbox
|
||||
SET status = 'processing',
|
||||
attempts = outbox.attempts + 1,
|
||||
locked_by = $1,
|
||||
lock_token = gen_random_uuid(),
|
||||
locked_at = now(),
|
||||
updated_at = now()
|
||||
FROM picked
|
||||
WHERE outbox.id = picked.id
|
||||
RETURNING outbox.id::text, outbox.task_id::text, outbox.action, outbox.amount::text,
|
||||
outbox.currency, outbox.status, outbox.attempts, outbox.next_attempt_at,
|
||||
COALESCE(outbox.locked_by, ''), COALESCE(outbox.lock_token::text, ''),
|
||||
COALESCE(outbox.locked_at::text, ''), COALESCE(outbox.last_error_code, ''),
|
||||
COALESCE(outbox.last_error_message, ''), COALESCE(outbox.manual_review_reason, ''),
|
||||
outbox.pricing_snapshot, outbox.payload, COALESCE(outbox.completed_at::text, ''),
|
||||
outbox.created_at, outbox.updated_at`,
|
||||
workerID, limit, int(staleAfter/time.Second))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]BillingSettlement, 0)
|
||||
for rows.Next() {
|
||||
item, err := scanBillingSettlement(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) ProcessBillingSettlement(ctx context.Context, settlement BillingSettlement) error {
|
||||
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
var action string
|
||||
var amount string
|
||||
var currency string
|
||||
var taskID string
|
||||
var taskStatus string
|
||||
var gatewayUserID string
|
||||
var gatewayTenantID string
|
||||
var userID string
|
||||
var tenantID string
|
||||
var tenantKey string
|
||||
err := tx.QueryRow(ctx, `
|
||||
SELECT outbox.action, outbox.amount::text, outbox.currency, outbox.task_id::text,
|
||||
task.status, COALESCE(task.gateway_user_id::text, ''),
|
||||
COALESCE(task.gateway_tenant_id::text, ''), task.user_id,
|
||||
COALESCE(task.tenant_id, ''), COALESCE(task.tenant_key, '')
|
||||
FROM settlement_outbox outbox
|
||||
JOIN gateway_tasks task ON task.id = outbox.task_id
|
||||
WHERE outbox.id = $1::uuid
|
||||
AND outbox.status = 'processing'
|
||||
AND outbox.lock_token = $2::uuid
|
||||
FOR UPDATE OF outbox, task`, settlement.ID, settlement.LockToken).Scan(
|
||||
&action, &amount, ¤cy, &taskID, &taskStatus, &gatewayUserID,
|
||||
&gatewayTenantID, &userID, &tenantID, &tenantKey,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET billing_status = 'processing', billing_updated_at = now(), updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID); err != nil {
|
||||
return err
|
||||
}
|
||||
switch action {
|
||||
case "settle":
|
||||
if taskStatus != "succeeded" {
|
||||
return fmt.Errorf("settle action requires succeeded task")
|
||||
}
|
||||
if err := settleBillingOutboxTx(ctx, tx, taskID, gatewayUserID, gatewayTenantID, userID, tenantID, tenantKey, currency, amount); err != nil {
|
||||
return err
|
||||
}
|
||||
return completeBillingOutboxTx(ctx, tx, settlement.ID, settlement.LockToken, taskID, "settled")
|
||||
case "release":
|
||||
if err := releaseBillingOutboxTx(ctx, tx, taskID, gatewayUserID, gatewayTenantID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
return completeBillingOutboxTx(ctx, tx, settlement.ID, settlement.LockToken, taskID, "released")
|
||||
default:
|
||||
return fmt.Errorf("unsupported billing settlement action %q", action)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func settleBillingOutboxTx(
|
||||
ctx context.Context,
|
||||
tx pgx.Tx,
|
||||
taskID string,
|
||||
gatewayUserID string,
|
||||
gatewayTenantID string,
|
||||
userID string,
|
||||
tenantID string,
|
||||
tenantKey string,
|
||||
currency string,
|
||||
amount string,
|
||||
) error {
|
||||
if gatewayUserID == "" {
|
||||
return fmt.Errorf("task %s has no gateway wallet user", taskID)
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_accounts (
|
||||
gateway_tenant_id, gateway_user_id, tenant_id, tenant_key, user_id, currency
|
||||
)
|
||||
VALUES (NULLIF($1, '')::uuid, $2::uuid, NULLIF($3, ''), NULLIF($4, ''), NULLIF($5, ''), $6)
|
||||
ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
|
||||
gatewayTenantID, gatewayUserID, tenantID, tenantKey, userID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := ensureWalletAccountAuditGuard(ctx, tx, gatewayUserID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
var accountID string
|
||||
var balanceBefore string
|
||||
var frozenBefore string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT id::text, balance::text, frozen_balance::text
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id = $1::uuid AND currency = $2
|
||||
FOR UPDATE`, gatewayUserID, currency).Scan(&accountID, &balanceBefore, &frozenBefore); err != nil {
|
||||
return err
|
||||
}
|
||||
var alreadySettled bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions
|
||||
WHERE account_id = $1::uuid AND idempotency_key = $2
|
||||
)`, accountID, billingIdempotencyKey(taskID)).Scan(&alreadySettled); err != nil {
|
||||
return err
|
||||
}
|
||||
if alreadySettled {
|
||||
return nil
|
||||
}
|
||||
|
||||
var reservationKey string
|
||||
var reservedAmount string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(reserve.idempotency_key, ''), reserve.amount::text
|
||||
FROM gateway_wallet_transactions reserve
|
||||
WHERE reserve.account_id = $1::uuid
|
||||
AND reserve.reference_type = 'gateway_task'
|
||||
AND reserve.reference_id = $2
|
||||
AND reserve.transaction_type = 'reserve'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions release
|
||||
WHERE release.account_id = reserve.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = reserve.idempotency_key || ':release'
|
||||
)
|
||||
ORDER BY reserve.created_at DESC
|
||||
LIMIT 1`, accountID, taskID).Scan(&reservationKey, &reservedAmount); errors.Is(err, pgx.ErrNoRows) {
|
||||
reservationKey = ""
|
||||
reservedAmount = "0"
|
||||
} else if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var balanceAfter string
|
||||
var frozenAfter string
|
||||
err := tx.QueryRow(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET balance = balance - $2::numeric,
|
||||
total_spent = total_spent + $2::numeric,
|
||||
frozen_balance = GREATEST(0, frozen_balance - $3::numeric),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND balance - frozen_balance + $3::numeric >= $2::numeric
|
||||
RETURNING balance::text, frozen_balance::text`,
|
||||
accountID, amount, reservedAmount).Scan(&balanceAfter, &frozenAfter)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return fmt.Errorf("%w: task %s settlement amount exceeds spendable balance", ErrInsufficientWalletBalance, taskID)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var reservedPositive bool
|
||||
if err := tx.QueryRow(ctx, `SELECT $1::numeric > 0`, reservedAmount).Scan(&reservedPositive); err != nil {
|
||||
return err
|
||||
}
|
||||
if reservedPositive {
|
||||
releaseMetadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "reason": "task_billing_settled",
|
||||
"reserved": reservedAmount, "frozenBefore": frozenBefore, "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::numeric, $5::numeric, $5::numeric, $6, 'gateway_task', $7, $8::jsonb
|
||||
)
|
||||
ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING`,
|
||||
accountID, gatewayTenantID, gatewayUserID, reservedAmount, balanceBefore,
|
||||
billingReservationReleaseIdempotencyKey(reservationKey), taskID, string(releaseMetadata)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "reservedAmount": reservedAmount,
|
||||
"frozenBefore": frozenBefore, "frozenAfter": frozenAfter,
|
||||
})
|
||||
_, 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', 'task_billing',
|
||||
$4::numeric, $5::numeric, $6::numeric, $7, 'gateway_task', $8, $9::jsonb
|
||||
)
|
||||
ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING`,
|
||||
accountID, gatewayTenantID, gatewayUserID, amount, balanceBefore, balanceAfter,
|
||||
billingIdempotencyKey(taskID), taskID, string(metadata))
|
||||
return err
|
||||
}
|
||||
|
||||
func releaseBillingOutboxTx(ctx context.Context, tx pgx.Tx, taskID string, gatewayUserID string, gatewayTenantID string, currency string) error {
|
||||
if gatewayUserID == "" {
|
||||
return nil
|
||||
}
|
||||
var accountID string
|
||||
var balance string
|
||||
var frozenBefore string
|
||||
err := tx.QueryRow(ctx, `
|
||||
SELECT id::text, balance::text, frozen_balance::text
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id = $1::uuid AND currency = $2
|
||||
FOR UPDATE`, gatewayUserID, currency).Scan(&accountID, &balance, &frozenBefore)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var reservationKey string
|
||||
var reservedAmount string
|
||||
err = tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(reserve.idempotency_key, ''), reserve.amount::text
|
||||
FROM gateway_wallet_transactions reserve
|
||||
WHERE reserve.account_id = $1::uuid
|
||||
AND reserve.reference_type = 'gateway_task'
|
||||
AND reserve.reference_id = $2
|
||||
AND reserve.transaction_type = 'reserve'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions release
|
||||
WHERE release.account_id = reserve.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = reserve.idempotency_key || ':release'
|
||||
)
|
||||
ORDER BY reserve.created_at DESC
|
||||
LIMIT 1`, accountID, taskID).Scan(&reservationKey, &reservedAmount)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var frozenAfter string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET frozen_balance = frozen_balance - $2::numeric,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND frozen_balance >= $2::numeric
|
||||
RETURNING frozen_balance::text`, accountID, reservedAmount).Scan(&frozenAfter); err != nil {
|
||||
return err
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "reason": "task_terminal_release",
|
||||
"reserved": reservedAmount, "frozenBefore": frozenBefore, "frozenAfter": frozenAfter,
|
||||
})
|
||||
_, 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::numeric, $5::numeric, $5::numeric, $6, 'gateway_task', $7, $8::jsonb
|
||||
)
|
||||
ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING`,
|
||||
accountID, gatewayTenantID, gatewayUserID, reservedAmount, balance,
|
||||
billingReservationReleaseIdempotencyKey(reservationKey), taskID, string(metadata))
|
||||
return err
|
||||
}
|
||||
|
||||
func completeBillingOutboxTx(ctx context.Context, tx pgx.Tx, settlementID string, lockToken string, taskID string, billingStatus string) error {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET billing_status = $2,
|
||||
reservation_amount = 0,
|
||||
billing_updated_at = now(),
|
||||
billing_settled_at = CASE WHEN $2 = 'settled' THEN now() ELSE billing_settled_at END,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, billingStatus); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
|
||||
SELECT task.id,
|
||||
COALESCE((SELECT MAX(event.seq) + 1 FROM gateway_task_events event WHERE event.task_id = task.id), 1),
|
||||
CASE WHEN $2 = 'settled' THEN 'task.billing.settled' ELSE 'task.billing.released' END,
|
||||
task.status,
|
||||
'billing',
|
||||
1,
|
||||
CASE WHEN $2 = 'settled' THEN 'task billing settled' ELSE 'task billing reservation released' END,
|
||||
jsonb_build_object('settlementId', $3::text, 'billingStatus', $2::text),
|
||||
task.run_mode = 'simulation'
|
||||
FROM gateway_tasks task
|
||||
WHERE task.id = $1::uuid`, taskID, billingStatus, settlementID); err != nil {
|
||||
return err
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE settlement_outbox
|
||||
SET status = 'completed',
|
||||
completed_at = now(),
|
||||
locked_by = NULL,
|
||||
lock_token = NULL,
|
||||
locked_at = NULL,
|
||||
last_error_code = NULL,
|
||||
last_error_message = NULL,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'processing'
|
||||
AND lock_token = $2::uuid`, settlementID, lockToken)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return pgx.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) MarkBillingSettlementFailed(ctx context.Context, settlement BillingSettlement, code string, message string, delay time.Duration) error {
|
||||
if delay < time.Second {
|
||||
delay = time.Second
|
||||
}
|
||||
if delay > 15*time.Minute {
|
||||
delay = 15 * time.Minute
|
||||
}
|
||||
manualReview := settlement.Attempts >= BillingSettlementMaxAttempts
|
||||
status := "retryable_failed"
|
||||
taskStatus := "retryable_failed"
|
||||
manualReason := ""
|
||||
if manualReview {
|
||||
status = "manual_review"
|
||||
taskStatus = "manual_review"
|
||||
manualReason = "maximum settlement attempts exceeded"
|
||||
}
|
||||
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE settlement_outbox
|
||||
SET status = $3,
|
||||
next_attempt_at = now() + ($4::int * interval '1 second'),
|
||||
locked_by = NULL,
|
||||
lock_token = NULL,
|
||||
locked_at = NULL,
|
||||
last_error_code = NULLIF($5, ''),
|
||||
last_error_message = NULLIF($6, ''),
|
||||
manual_review_reason = NULLIF($7, ''),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND lock_token = $2::uuid
|
||||
AND status = 'processing'`,
|
||||
settlement.ID, settlement.LockToken, status, int(delay/time.Second), code, message, manualReason)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return pgx.ErrNoRows
|
||||
}
|
||||
_, err = tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET billing_status = $2,
|
||||
billing_updated_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, settlement.TaskID, taskStatus)
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) ListBillingSettlements(ctx context.Context, filter BillingSettlementListFilter) (BillingSettlementListResult, error) {
|
||||
page := filter.Page
|
||||
if page <= 0 {
|
||||
page = 1
|
||||
}
|
||||
pageSize := filter.PageSize
|
||||
if pageSize <= 0 {
|
||||
pageSize = 50
|
||||
}
|
||||
if pageSize > 100 {
|
||||
pageSize = 100
|
||||
}
|
||||
status := strings.TrimSpace(filter.Status)
|
||||
action := strings.TrimSpace(filter.Action)
|
||||
var total int
|
||||
if err := s.pool.QueryRow(ctx, `
|
||||
SELECT count(*)
|
||||
FROM settlement_outbox
|
||||
WHERE (NULLIF($1, '') IS NULL OR status = $1)
|
||||
AND (NULLIF($2, '') IS NULL OR action = $2)`, status, action).Scan(&total); err != nil {
|
||||
return BillingSettlementListResult{}, err
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, task_id::text, action, amount::text, currency, status, attempts,
|
||||
next_attempt_at, COALESCE(locked_by, ''), COALESCE(lock_token::text, ''),
|
||||
COALESCE(locked_at::text, ''), COALESCE(last_error_code, ''),
|
||||
COALESCE(last_error_message, ''), COALESCE(manual_review_reason, ''),
|
||||
pricing_snapshot, payload, COALESCE(completed_at::text, ''), created_at, updated_at
|
||||
FROM settlement_outbox
|
||||
WHERE (NULLIF($1, '') IS NULL OR status = $1)
|
||||
AND (NULLIF($2, '') IS NULL OR action = $2)
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $3 OFFSET $4`, status, action, pageSize, (page-1)*pageSize)
|
||||
if err != nil {
|
||||
return BillingSettlementListResult{}, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]BillingSettlement, 0)
|
||||
for rows.Next() {
|
||||
item, err := scanBillingSettlement(rows)
|
||||
if err != nil {
|
||||
return BillingSettlementListResult{}, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return BillingSettlementListResult{}, err
|
||||
}
|
||||
return BillingSettlementListResult{Items: items, Total: total, Page: page, PageSize: pageSize}, nil
|
||||
}
|
||||
|
||||
func (s *Store) RetryBillingSettlementTx(ctx context.Context, tx Tx, id string, idempotencyKeyHash string) (BillingSettlement, bool, error) {
|
||||
current, err := scanBillingSettlement(tx.QueryRow(ctx, `
|
||||
SELECT id::text, task_id::text, action, amount::text, currency, status, attempts,
|
||||
next_attempt_at, COALESCE(locked_by, ''), COALESCE(lock_token::text, ''),
|
||||
COALESCE(locked_at::text, ''), COALESCE(last_error_code, ''),
|
||||
COALESCE(last_error_message, ''), COALESCE(manual_review_reason, ''),
|
||||
pricing_snapshot, payload, COALESCE(completed_at::text, ''), created_at, updated_at
|
||||
FROM settlement_outbox
|
||||
WHERE id = $1::uuid
|
||||
FOR UPDATE`, id))
|
||||
if err != nil {
|
||||
return BillingSettlement{}, false, err
|
||||
}
|
||||
var recordedHash string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(retry_idempotency_key_hash, '')
|
||||
FROM settlement_outbox
|
||||
WHERE id = $1::uuid`, id).Scan(&recordedHash); err != nil {
|
||||
return BillingSettlement{}, false, err
|
||||
}
|
||||
if recordedHash != "" && recordedHash == idempotencyKeyHash {
|
||||
return current, true, nil
|
||||
}
|
||||
if current.Status != "retryable_failed" && current.Status != "manual_review" {
|
||||
return BillingSettlement{}, false, ErrBillingSettlementNotRetryable
|
||||
}
|
||||
item, err := scanBillingSettlement(tx.QueryRow(ctx, `
|
||||
UPDATE settlement_outbox
|
||||
SET status = 'pending',
|
||||
next_attempt_at = now(),
|
||||
locked_by = NULL,
|
||||
lock_token = NULL,
|
||||
locked_at = NULL,
|
||||
manual_review_reason = NULL,
|
||||
retry_idempotency_key_hash = $2,
|
||||
retry_requested_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND status IN ('retryable_failed', 'manual_review')
|
||||
RETURNING id::text, task_id::text, action, amount::text, currency, status, attempts,
|
||||
next_attempt_at, COALESCE(locked_by, ''), COALESCE(lock_token::text, ''),
|
||||
COALESCE(locked_at::text, ''), COALESCE(last_error_code, ''),
|
||||
COALESCE(last_error_message, ''), COALESCE(manual_review_reason, ''),
|
||||
pricing_snapshot, payload, COALESCE(completed_at::text, ''), created_at, updated_at`, id, idempotencyKeyHash))
|
||||
if err != nil {
|
||||
return BillingSettlement{}, false, err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET billing_status = 'pending', billing_updated_at = now(), updated_at = now()
|
||||
WHERE id = $1::uuid`, item.TaskID); err != nil {
|
||||
return BillingSettlement{}, false, err
|
||||
}
|
||||
return item, false, nil
|
||||
}
|
||||
|
||||
func scanBillingSettlement(scanner taskScanner) (BillingSettlement, error) {
|
||||
var item BillingSettlement
|
||||
var pricingSnapshotBytes []byte
|
||||
var payloadBytes []byte
|
||||
if err := scanner.Scan(
|
||||
&item.ID, &item.TaskID, &item.Action, &item.amountExact, &item.Currency,
|
||||
&item.Status, &item.Attempts, &item.NextAttemptAt, &item.LockedBy, &item.LockToken,
|
||||
&item.LockedAt, &item.LastErrorCode, &item.LastErrorMessage, &item.ManualReviewReason,
|
||||
&pricingSnapshotBytes, &payloadBytes, &item.CompletedAt, &item.CreatedAt, &item.UpdatedAt,
|
||||
); err != nil {
|
||||
return BillingSettlement{}, err
|
||||
}
|
||||
item.Amount, _ = strconv.ParseFloat(item.amountExact, 64)
|
||||
item.PricingSnapshot = decodeObject(pricingSnapshotBytes)
|
||||
item.Payload = decodeObject(payloadBytes)
|
||||
return item, nil
|
||||
}
|
||||
@@ -0,0 +1,265 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
func TestTaskIdempotencyAndExecutionLease(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-v2-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
input := CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "simulation",
|
||||
Request: map[string]any{"model": "billing-v2-model", "prompt": "lease"},
|
||||
IdempotencyKeyHash: "key-hash-" + uuid.NewString(), IdempotencyRequestHash: "request-a",
|
||||
}
|
||||
created, err := db.CreateTaskIdempotent(ctx, input, user)
|
||||
if err != nil || created.Replayed {
|
||||
t.Fatalf("create task result=%+v err=%v", created, err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.Task.ID)
|
||||
})
|
||||
|
||||
replayed, err := db.CreateTaskIdempotent(ctx, input, user)
|
||||
if err != nil || !replayed.Replayed || replayed.Task.ID != created.Task.ID {
|
||||
t.Fatalf("replay result=%+v err=%v", replayed, err)
|
||||
}
|
||||
input.IdempotencyRequestHash = "request-b"
|
||||
if _, err := db.CreateTaskIdempotent(ctx, input, user); !errors.Is(err, ErrIdempotencyKeyReused) {
|
||||
t.Fatalf("different request error=%v", err)
|
||||
}
|
||||
|
||||
firstToken := uuid.NewString()
|
||||
claimed, err := db.ClaimTaskExecution(ctx, created.Task.ID, firstToken, 5*time.Minute)
|
||||
if err != nil || claimed.ExecutionToken != firstToken {
|
||||
t.Fatalf("first claim task=%+v err=%v", claimed, err)
|
||||
}
|
||||
if err := db.RenewTaskExecutionLease(ctx, created.Task.ID, firstToken, 5*time.Minute); err != nil {
|
||||
t.Fatalf("renew first lease: %v", err)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE gateway_tasks SET execution_lease_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, created.Task.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secondToken := uuid.NewString()
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.Task.ID, secondToken, 5*time.Minute); err != nil {
|
||||
t.Fatalf("take over expired lease: %v", err)
|
||||
}
|
||||
if _, err := db.FinishTaskFailure(ctx, FinishTaskFailureInput{TaskID: created.Task.ID, ExecutionToken: firstToken, Code: "old_worker", Message: "old"}); !errors.Is(err, ErrTaskExecutionLeaseLost) {
|
||||
t.Fatalf("old worker terminal error=%v", err)
|
||||
}
|
||||
finished, err := db.FinishTaskFailure(ctx, FinishTaskFailureInput{TaskID: created.Task.ID, ExecutionToken: secondToken, Code: "new_worker", Message: "new"})
|
||||
if err != nil || finished.ErrorCode != "new_worker" {
|
||||
t.Fatalf("new worker terminal task=%+v err=%v", finished, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingSettlementStaleTakeoverDebitsExactlyOnce(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-settle-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{GatewayUserID: gatewayUserID, Currency: "resource", Balance: 10, Reason: "billing v2 test"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{Kind: "images.generations", Model: "billing-v2-model", RunMode: "production", Request: map[string]any{"model": "billing-v2-model"}}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cleanupCtx := context.Background()
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
})
|
||||
token := uuid.NewString()
|
||||
claimed, err := db.ClaimTaskExecution(ctx, created.ID, token, 5*time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
amount := "1.000000001"
|
||||
reservations, err := db.ReserveTaskBilling(ctx, claimed, user, []any{map[string]any{"currency": "resource", "amount": amount}}, map[string]any{
|
||||
"pricingVersion": "effective-pricing-v2", "reservationAmount": amount, "currency": "resource", "requestFingerprint": "billing-v2-test",
|
||||
})
|
||||
if err != nil || len(reservations) != 1 {
|
||||
t.Fatalf("reserve=%+v err=%v", reservations, err)
|
||||
}
|
||||
if _, err := db.FinishTaskSuccess(ctx, FinishTaskSuccessInput{
|
||||
TaskID: created.ID, ExecutionToken: token, Result: map[string]any{"ok": true},
|
||||
FinalChargeAmountText: amount, BillingCurrency: "resource",
|
||||
PricingSnapshot: map[string]any{"pricingVersion": "effective-pricing-v2"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
firstClaims, err := db.ClaimBillingSettlements(ctx, "worker-one", BillingSettlementBatchSize, BillingSettlementLockTimeout)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
first := settlementForTask(t, firstClaims, created.ID)
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE settlement_outbox SET locked_at=now()-interval '3 minutes' WHERE id=$1::uuid`, first.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secondClaims, err := db.ClaimBillingSettlements(ctx, "worker-two", BillingSettlementBatchSize, BillingSettlementLockTimeout)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second := settlementForTask(t, secondClaims, created.ID)
|
||||
if err := db.ProcessBillingSettlement(ctx, first); !errors.Is(err, pgx.ErrNoRows) {
|
||||
t.Fatalf("stale settlement lock error=%v", err)
|
||||
}
|
||||
if err := db.ProcessBillingSettlement(ctx, second); err != nil {
|
||||
t.Fatalf("process takeover: %v", err)
|
||||
}
|
||||
|
||||
var walletExact bool
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT balance = 8.999999999::numeric
|
||||
AND frozen_balance = 0::numeric
|
||||
AND total_spent = 1.000000001::numeric
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id=$1::uuid AND currency='resource'`, gatewayUserID).Scan(&walletExact); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !walletExact {
|
||||
t.Fatal("wallet amounts did not preserve nine-decimal settlement")
|
||||
}
|
||||
var billingTransactions int
|
||||
if err := db.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_wallet_transactions WHERE reference_id=$1 AND transaction_type='task_billing'`, created.ID).Scan(&billingTransactions); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if billingTransactions != 1 {
|
||||
t.Fatalf("task billing transactions=%d", billingTransactions)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `
|
||||
DELETE FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id=$1::uuid AND currency='resource'`, gatewayUserID); err == nil {
|
||||
t.Fatal("wallet account with audit transactions must not be deletable")
|
||||
}
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT count(*) FROM gateway_wallet_transactions
|
||||
WHERE gateway_user_id=$1::uuid AND reference_id=$2`, gatewayUserID, created.ID).Scan(&billingTransactions); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if billingTransactions == 0 {
|
||||
t.Fatal("wallet audit transactions were lost after rejected account deletion")
|
||||
}
|
||||
settled, err := db.GetTask(ctx, created.ID)
|
||||
if err != nil || settled.BillingStatus != "settled" {
|
||||
t.Fatalf("settled task=%+v err=%v", settled, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseTaskBillingReservationsPreservesNineDecimalPlaces(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-release-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", Balance: 1, Reason: "billing v2 release test",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cleanupCtx := context.Background()
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
})
|
||||
amount := "0.000000001"
|
||||
reservations, err := db.ReserveTaskBilling(ctx, created, user, nil, map[string]any{
|
||||
"pricingVersion": "effective-pricing-v2", "reservationAmount": amount, "currency": "resource",
|
||||
})
|
||||
if err != nil || len(reservations) != 1 {
|
||||
t.Fatalf("reserve=%+v err=%v", reservations, err)
|
||||
}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", BalanceText: "0", Reason: "must not cross frozen balance",
|
||||
}); !errors.Is(err, ErrBalanceBelowFrozen) {
|
||||
t.Fatalf("balance below frozen error=%v", err)
|
||||
}
|
||||
if err := db.ReleaseTaskBillingReservations(ctx, reservations, "integration_test"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", BalanceText: "0.123456789", Reason: "exact adjustment",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.RechargeUserWalletBalance(ctx, WalletRechargeInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", AmountText: "0.000000001", Reason: "exact recharge",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var exact bool
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT account.balance = 0.123456790::numeric
|
||||
AND account.frozen_balance = 0::numeric
|
||||
AND task.reservation_amount = 0::numeric
|
||||
AND task.billing_status = 'not_started'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions transaction
|
||||
WHERE transaction.reference_id = task.id::text
|
||||
AND transaction.transaction_type = 'release'
|
||||
AND transaction.amount = 0.000000001::numeric
|
||||
)
|
||||
FROM gateway_wallet_accounts account
|
||||
JOIN gateway_tasks task ON task.gateway_user_id = account.gateway_user_id
|
||||
WHERE task.id=$1::uuid AND account.currency='resource'`, created.ID).Scan(&exact); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !exact {
|
||||
t.Fatal("nine-decimal reservation was not released exactly")
|
||||
}
|
||||
}
|
||||
|
||||
func billingV2IntegrationStore(t *testing.T) *Store {
|
||||
t.Helper()
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run billing v2 PostgreSQL integration tests")
|
||||
}
|
||||
db, err := Connect(context.Background(), databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var databaseName string
|
||||
if err := db.pool.QueryRow(context.Background(), `SELECT current_database()`).Scan(&databaseName); err != nil {
|
||||
db.Close()
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(databaseName), "test") {
|
||||
db.Close()
|
||||
t.Fatalf("refusing to use non-test database %q", databaseName)
|
||||
}
|
||||
t.Cleanup(db.Close)
|
||||
return db
|
||||
}
|
||||
|
||||
func settlementForTask(t *testing.T, items []BillingSettlement, taskID string) BillingSettlement {
|
||||
t.Helper()
|
||||
for _, item := range items {
|
||||
if item.TaskID == taskID {
|
||||
return item
|
||||
}
|
||||
}
|
||||
t.Fatalf("settlement for task %s not claimed", taskID)
|
||||
return BillingSettlement{}
|
||||
}
|
||||
+191
-112
@@ -47,16 +47,21 @@ func normalizeAPIKeyScopes(scopes []string) []string {
|
||||
}
|
||||
|
||||
var (
|
||||
ErrInvalidCredentials = errors.New("invalid account or password")
|
||||
ErrInvalidInvitation = errors.New("invalid or expired invitation code")
|
||||
ErrInvalidAPIKeyScopes = errors.New("api key scopes must not be empty")
|
||||
ErrAccessRuleResourceDenied = errors.New("access rule resource is not available")
|
||||
ErrInsufficientWalletBalance = errors.New("insufficient wallet balance")
|
||||
ErrLocalUserRequired = errors.New("local gateway user is required")
|
||||
ErrWalletBalanceUnchanged = errors.New("wallet balance unchanged")
|
||||
ErrProtectedDefault = errors.New("protected default resource cannot be deleted")
|
||||
ErrUserAlreadyExists = errors.New("user already exists")
|
||||
ErrWeakPassword = errors.New("password must be at least 8 characters")
|
||||
ErrInvalidCredentials = errors.New("invalid account or password")
|
||||
ErrInvalidInvitation = errors.New("invalid or expired invitation code")
|
||||
ErrInvalidAPIKeyScopes = errors.New("api key scopes must not be empty")
|
||||
ErrAccessRuleResourceDenied = errors.New("access rule resource is not available")
|
||||
ErrInsufficientWalletBalance = errors.New("insufficient wallet balance")
|
||||
ErrLocalUserRequired = errors.New("local gateway user is required")
|
||||
ErrWalletBalanceUnchanged = errors.New("wallet balance unchanged")
|
||||
ErrBalanceBelowFrozen = errors.New("wallet balance cannot be below frozen balance")
|
||||
ErrInvalidWalletAmount = errors.New("wallet amount must be a decimal with at most nine fractional digits")
|
||||
ErrIdempotencyKeyReused = errors.New("idempotency key was reused for a different request")
|
||||
ErrTaskExecutionLeaseUnavailable = errors.New("task execution lease is unavailable")
|
||||
ErrTaskExecutionLeaseLost = errors.New("task execution lease was lost")
|
||||
ErrProtectedDefault = errors.New("protected default resource cannot be deleted")
|
||||
ErrUserAlreadyExists = errors.New("user already exists")
|
||||
ErrWeakPassword = errors.New("password must be at least 8 characters")
|
||||
)
|
||||
|
||||
func Connect(ctx context.Context, databaseURL string) (*Store, error) {
|
||||
@@ -414,64 +419,81 @@ type RateLimitWindow struct {
|
||||
}
|
||||
|
||||
type CreateTaskInput struct {
|
||||
Kind string `json:"kind"`
|
||||
Model string `json:"model"`
|
||||
RunMode string `json:"runMode"`
|
||||
Async bool `json:"async"`
|
||||
Request map[string]any `json:"request"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
NewMessageCount int `json:"newMessageCount"`
|
||||
MessageRefs []TaskMessageRefInput `json:"messageRefs"`
|
||||
Kind string `json:"kind"`
|
||||
Model string `json:"model"`
|
||||
RunMode string `json:"runMode"`
|
||||
Async bool `json:"async"`
|
||||
Request map[string]any `json:"request"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
NewMessageCount int `json:"newMessageCount"`
|
||||
MessageRefs []TaskMessageRefInput `json:"messageRefs"`
|
||||
IdempotencyKeyHash string `json:"-"`
|
||||
IdempotencyRequestHash string `json:"-"`
|
||||
}
|
||||
|
||||
type CreateTaskResult struct {
|
||||
Task GatewayTask
|
||||
Replayed bool
|
||||
}
|
||||
|
||||
type GatewayTask struct {
|
||||
ID string `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
RunMode string `json:"runMode"`
|
||||
UserID string `json:"userId"`
|
||||
GatewayUserID string `json:"gatewayUserId,omitempty"`
|
||||
UserSource string `json:"userSource,omitempty"`
|
||||
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
TenantKey string `json:"tenantKey,omitempty"`
|
||||
APIKeyID string `json:"apiKeyId,omitempty"`
|
||||
APIKeyName string `json:"apiKeyName,omitempty"`
|
||||
APIKeyPrefix string `json:"apiKeyPrefix,omitempty"`
|
||||
UserGroupID string `json:"userGroupId,omitempty"`
|
||||
UserGroupKey string `json:"userGroupKey,omitempty"`
|
||||
Model string `json:"model"`
|
||||
ModelType string `json:"modelType,omitempty"`
|
||||
RequestedModel string `json:"requestedModel,omitempty"`
|
||||
ResolvedModel string `json:"resolvedModel,omitempty"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
NewMessageCount int `json:"newMessageCount,omitempty"`
|
||||
Request map[string]any `json:"request,omitempty"`
|
||||
AsyncMode bool `json:"asyncMode"`
|
||||
RiverJobID int64 `json:"riverJobId,omitempty"`
|
||||
Status string `json:"status"`
|
||||
Cancellable *bool `json:"cancellable,omitempty"`
|
||||
Submitted *bool `json:"submitted,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
AttemptCount int `json:"attemptCount"`
|
||||
RemoteTaskID string `json:"remoteTaskId,omitempty"`
|
||||
RemoteTaskPayload map[string]any `json:"remoteTaskPayload,omitempty"`
|
||||
Result map[string]any `json:"result,omitempty"`
|
||||
Billings []any `json:"billings,omitempty"`
|
||||
Usage map[string]any `json:"usage"`
|
||||
Metrics map[string]any `json:"metrics"`
|
||||
BillingSummary map[string]any `json:"billingSummary"`
|
||||
FinalChargeAmount float64 `json:"finalChargeAmount"`
|
||||
ResponseStartedAt string `json:"responseStartedAt,omitempty"`
|
||||
ResponseFinishedAt string `json:"responseFinishedAt,omitempty"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
Attempts []TaskAttempt `json:"attempts,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
ID string `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
RunMode string `json:"runMode"`
|
||||
UserID string `json:"userId"`
|
||||
GatewayUserID string `json:"gatewayUserId,omitempty"`
|
||||
UserSource string `json:"userSource,omitempty"`
|
||||
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
TenantKey string `json:"tenantKey,omitempty"`
|
||||
APIKeyID string `json:"apiKeyId,omitempty"`
|
||||
APIKeyName string `json:"apiKeyName,omitempty"`
|
||||
APIKeyPrefix string `json:"apiKeyPrefix,omitempty"`
|
||||
UserGroupID string `json:"userGroupId,omitempty"`
|
||||
UserGroupKey string `json:"userGroupKey,omitempty"`
|
||||
Model string `json:"model"`
|
||||
ModelType string `json:"modelType,omitempty"`
|
||||
RequestedModel string `json:"requestedModel,omitempty"`
|
||||
ResolvedModel string `json:"resolvedModel,omitempty"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
NewMessageCount int `json:"newMessageCount,omitempty"`
|
||||
Request map[string]any `json:"request,omitempty"`
|
||||
AsyncMode bool `json:"asyncMode"`
|
||||
RiverJobID int64 `json:"riverJobId,omitempty"`
|
||||
Status string `json:"status"`
|
||||
Cancellable *bool `json:"cancellable,omitempty"`
|
||||
Submitted *bool `json:"submitted,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
AttemptCount int `json:"attemptCount"`
|
||||
RemoteTaskID string `json:"remoteTaskId,omitempty"`
|
||||
RemoteTaskPayload map[string]any `json:"remoteTaskPayload,omitempty"`
|
||||
Result map[string]any `json:"result,omitempty"`
|
||||
Billings []any `json:"billings,omitempty"`
|
||||
Usage map[string]any `json:"usage"`
|
||||
Metrics map[string]any `json:"metrics"`
|
||||
BillingSummary map[string]any `json:"billingSummary"`
|
||||
FinalChargeAmount float64 `json:"finalChargeAmount"`
|
||||
BillingVersion string `json:"billingVersion"`
|
||||
BillingStatus string `json:"billingStatus"`
|
||||
BillingCurrency string `json:"billingCurrency"`
|
||||
PricingSnapshot map[string]any `json:"pricingSnapshot,omitempty"`
|
||||
RequestFingerprint string `json:"requestFingerprint,omitempty"`
|
||||
ReservationAmount float64 `json:"reservationAmount"`
|
||||
ExecutionToken string `json:"-"`
|
||||
ExecutionLeaseUntil string `json:"executionLeaseExpiresAt,omitempty"`
|
||||
BillingUpdatedAt string `json:"billingUpdatedAt,omitempty"`
|
||||
BillingSettledAt string `json:"billingSettledAt,omitempty"`
|
||||
ResponseStartedAt string `json:"responseStartedAt,omitempty"`
|
||||
ResponseFinishedAt string `json:"responseFinishedAt,omitempty"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
Attempts []TaskAttempt `json:"attempts,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
const gatewayTaskColumns = `
|
||||
@@ -485,7 +507,11 @@ request, COALESCE(async_mode, false), COALESCE(river_job_id, 0), status, COALESC
|
||||
COALESCE(remote_task_id, ''), COALESCE(remote_task_payload, '{}'::jsonb),
|
||||
COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
|
||||
COALESCE(usage, '{}'::jsonb), COALESCE(metrics, '{}'::jsonb), COALESCE(billing_summary, '{}'::jsonb),
|
||||
COALESCE(final_charge_amount, 0)::float8, COALESCE(response_started_at::text, ''),
|
||||
COALESCE(final_charge_amount, 0)::float8, billing_version, billing_status, billing_currency,
|
||||
COALESCE(pricing_snapshot, '{}'::jsonb), COALESCE(request_fingerprint, ''),
|
||||
COALESCE(reservation_amount, 0)::float8, COALESCE(execution_token::text, ''),
|
||||
COALESCE(execution_lease_expires_at::text, ''), COALESCE(billing_updated_at::text, ''),
|
||||
COALESCE(billing_settled_at::text, ''), COALESCE(response_started_at::text, ''),
|
||||
COALESCE(response_finished_at::text, ''), COALESCE(response_duration_ms, 0), COALESCE(error, ''),
|
||||
COALESCE(error_code, ''), COALESCE(error_message, ''),
|
||||
created_at, updated_at, COALESCE(finished_at::text, '')`
|
||||
@@ -505,35 +531,39 @@ type TaskEvent struct {
|
||||
}
|
||||
|
||||
type TaskAttempt struct {
|
||||
ID string `json:"id"`
|
||||
TaskID string `json:"taskId"`
|
||||
AttemptNo int `json:"attemptNo"`
|
||||
PlatformID string `json:"platformId,omitempty"`
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
PlatformModelID string `json:"platformModelId,omitempty"`
|
||||
ModelName string `json:"modelName,omitempty"`
|
||||
ProviderModelName string `json:"providerModelName,omitempty"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
ModelType string `json:"modelType,omitempty"`
|
||||
ClientID string `json:"clientId,omitempty"`
|
||||
QueueKey string `json:"queueKey"`
|
||||
Status string `json:"status"`
|
||||
Retryable bool `json:"retryable"`
|
||||
Simulated bool `json:"simulated"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
StatusCode int `json:"statusCode,omitempty"`
|
||||
Usage map[string]any `json:"usage,omitempty"`
|
||||
Metrics map[string]any `json:"metrics,omitempty"`
|
||||
RequestSnapshot map[string]any `json:"requestSnapshot,omitempty"`
|
||||
ResponseSnapshot map[string]any `json:"responseSnapshot,omitempty"`
|
||||
ResponseStartedAt string `json:"responseStartedAt,omitempty"`
|
||||
ResponseFinishedAt string `json:"responseFinishedAt,omitempty"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
StartedAt time.Time `json:"startedAt"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
ID string `json:"id"`
|
||||
TaskID string `json:"taskId"`
|
||||
AttemptNo int `json:"attemptNo"`
|
||||
PlatformID string `json:"platformId,omitempty"`
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
PlatformModelID string `json:"platformModelId,omitempty"`
|
||||
ModelName string `json:"modelName,omitempty"`
|
||||
ProviderModelName string `json:"providerModelName,omitempty"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
ModelType string `json:"modelType,omitempty"`
|
||||
ClientID string `json:"clientId,omitempty"`
|
||||
QueueKey string `json:"queueKey"`
|
||||
Status string `json:"status"`
|
||||
Retryable bool `json:"retryable"`
|
||||
Simulated bool `json:"simulated"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
StatusCode int `json:"statusCode,omitempty"`
|
||||
Usage map[string]any `json:"usage,omitempty"`
|
||||
Metrics map[string]any `json:"metrics,omitempty"`
|
||||
RequestSnapshot map[string]any `json:"requestSnapshot,omitempty"`
|
||||
ResponseSnapshot map[string]any `json:"responseSnapshot,omitempty"`
|
||||
ResponseStartedAt string `json:"responseStartedAt,omitempty"`
|
||||
ResponseFinishedAt string `json:"responseFinishedAt,omitempty"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
PricingSnapshot map[string]any `json:"pricingSnapshot,omitempty"`
|
||||
RequestFingerprint string `json:"requestFingerprint,omitempty"`
|
||||
UpstreamSubmissionStatus string `json:"upstreamSubmissionStatus"`
|
||||
UpstreamSubmissionUpdatedAt string `json:"upstreamSubmissionUpdatedAt,omitempty"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
StartedAt time.Time `json:"startedAt"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
}
|
||||
|
||||
type TaskParamPreprocessingLog struct {
|
||||
@@ -1716,6 +1746,9 @@ ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
|
||||
); err != nil {
|
||||
return GatewayUser{}, err
|
||||
}
|
||||
if err := ensureWalletAccountAuditGuard(ctx, tx, user.ID, "resource"); err != nil {
|
||||
return GatewayUser{}, err
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return GatewayUser{}, err
|
||||
}
|
||||
@@ -1826,6 +1859,11 @@ ORDER BY window_start DESC, scope_type ASC, scope_key ASC, metric ASC`)
|
||||
}
|
||||
|
||||
func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *auth.User) (GatewayTask, error) {
|
||||
result, err := s.CreateTaskIdempotent(ctx, input, user)
|
||||
return result.Task, err
|
||||
}
|
||||
|
||||
func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput, user *auth.User) (CreateTaskResult, error) {
|
||||
requestBody, _ := json.Marshal(input.Request)
|
||||
runMode := normalizeRunMode(input.RunMode, input.Request)
|
||||
status := "queued"
|
||||
@@ -1834,7 +1872,7 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
|
||||
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
@@ -1842,33 +1880,62 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
|
||||
INSERT INTO gateway_tasks (
|
||||
kind, run_mode, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
|
||||
api_key_id, api_key_name, api_key_prefix, user_group_id, user_group_key,
|
||||
model, requested_model, request, async_mode, status, result, billings, conversation_id, new_message_count, finished_at
|
||||
model, requested_model, request, async_mode, status, result, billings, conversation_id, new_message_count,
|
||||
idempotency_key_hash, idempotency_request_hash, finished_at
|
||||
)
|
||||
VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17, $18::jsonb, $19::jsonb, NULLIF($20, '')::uuid, $21, CASE WHEN $22 THEN now() ELSE NULL END)
|
||||
VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17, $18::jsonb, $19::jsonb, NULLIF($20, '')::uuid, $21, NULLIF($22, ''), NULLIF($23, ''), NULL)
|
||||
ON CONFLICT (user_id, idempotency_key_hash) WHERE idempotency_key_hash IS NOT NULL DO NOTHING
|
||||
RETURNING `+gatewayTaskColumns,
|
||||
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, input.ConversationID, input.NewMessageCount, false,
|
||||
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, input.ConversationID, input.NewMessageCount, strings.TrimSpace(input.IdempotencyKeyHash), strings.TrimSpace(input.IdempotencyRequestHash),
|
||||
))
|
||||
replayed := false
|
||||
if errors.Is(err, pgx.ErrNoRows) && strings.TrimSpace(input.IdempotencyKeyHash) != "" {
|
||||
var existingRequestHash string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(idempotency_request_hash, '')
|
||||
FROM gateway_tasks
|
||||
WHERE user_id = $1 AND idempotency_key_hash = $2
|
||||
FOR UPDATE`, user.ID, strings.TrimSpace(input.IdempotencyKeyHash)).Scan(&existingRequestHash); err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
if existingRequestHash != strings.TrimSpace(input.IdempotencyRequestHash) {
|
||||
return CreateTaskResult{}, ErrIdempotencyKeyReused
|
||||
}
|
||||
task, err = scanGatewayTask(tx.QueryRow(ctx, `
|
||||
SELECT `+gatewayTaskColumns+`
|
||||
FROM gateway_tasks
|
||||
WHERE user_id = $1 AND idempotency_key_hash = $2`, user.ID, strings.TrimSpace(input.IdempotencyKeyHash)))
|
||||
replayed = true
|
||||
}
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
if err := insertTaskMessageRefs(ctx, tx, task.ID, input.MessageRefs); err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
events := taskEventsForCreate(task.ID, runMode, status, nil)
|
||||
for _, event := range events {
|
||||
payload, _ := json.Marshal(event.Payload)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
if !replayed {
|
||||
if err := insertTaskMessageRefs(ctx, tx, task.ID, input.MessageRefs); err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
events := taskEventsForCreate(task.ID, runMode, status, nil)
|
||||
for _, event := range events {
|
||||
payload, _ := json.Marshal(event.Payload)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
|
||||
VALUES ($1::uuid, $2, $3::text, NULLIF($4::text, ''), NULLIF($5::text, ''), $6, NULLIF($7::text, ''), $8::jsonb, $9)`,
|
||||
task.ID, event.Seq, event.EventType, event.Status, event.Phase, event.Progress, event.Message, string(payload), event.Simulated,
|
||||
); err != nil {
|
||||
return GatewayTask{}, err
|
||||
task.ID, event.Seq, event.EventType, event.Status, event.Phase, event.Progress, event.Message, string(payload), event.Simulated,
|
||||
); err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return GatewayTask{}, err
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
return task, nil
|
||||
if replayed {
|
||||
task, err = s.GetTask(ctx, task.ID)
|
||||
if err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
}
|
||||
return CreateTaskResult{Task: task, Replayed: replayed}, nil
|
||||
}
|
||||
|
||||
func (s *Store) GetTask(ctx context.Context, taskID string) (GatewayTask, error) {
|
||||
@@ -1900,6 +1967,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
var usageBytes []byte
|
||||
var metricsBytes []byte
|
||||
var billingSummaryBytes []byte
|
||||
var pricingSnapshotBytes []byte
|
||||
var remoteTaskPayloadBytes []byte
|
||||
if err := scanner.Scan(
|
||||
&task.ID,
|
||||
@@ -1936,6 +2004,16 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
&metricsBytes,
|
||||
&billingSummaryBytes,
|
||||
&task.FinalChargeAmount,
|
||||
&task.BillingVersion,
|
||||
&task.BillingStatus,
|
||||
&task.BillingCurrency,
|
||||
&pricingSnapshotBytes,
|
||||
&task.RequestFingerprint,
|
||||
&task.ReservationAmount,
|
||||
&task.ExecutionToken,
|
||||
&task.ExecutionLeaseUntil,
|
||||
&task.BillingUpdatedAt,
|
||||
&task.BillingSettledAt,
|
||||
&task.ResponseStartedAt,
|
||||
&task.ResponseFinishedAt,
|
||||
&task.ResponseDurationMS,
|
||||
@@ -1955,6 +2033,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
task.Usage = decodeObject(usageBytes)
|
||||
task.Metrics = decodeObject(metricsBytes)
|
||||
task.BillingSummary = decodeObject(billingSummaryBytes)
|
||||
task.PricingSnapshot = decodeObject(pricingSnapshotBytes)
|
||||
return task, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -223,21 +223,24 @@ type RateLimitResult struct {
|
||||
}
|
||||
|
||||
type CreateTaskAttemptInput struct {
|
||||
TaskID string
|
||||
AttemptNo int
|
||||
PlatformID string
|
||||
PlatformModelID string
|
||||
ClientID string
|
||||
QueueKey string
|
||||
Status string
|
||||
Simulated bool
|
||||
RequestSnapshot map[string]any
|
||||
Metrics map[string]any
|
||||
TaskID string
|
||||
AttemptNo int
|
||||
PlatformID string
|
||||
PlatformModelID string
|
||||
ClientID string
|
||||
QueueKey string
|
||||
Status string
|
||||
Simulated bool
|
||||
RequestSnapshot map[string]any
|
||||
Metrics map[string]any
|
||||
PricingSnapshot map[string]any
|
||||
RequestFingerprint string
|
||||
}
|
||||
|
||||
type AsyncTaskQueueItem struct {
|
||||
ID string
|
||||
Priority int
|
||||
ID string
|
||||
Priority int
|
||||
NextRunAt time.Time
|
||||
}
|
||||
|
||||
type FinishTaskAttemptInput struct {
|
||||
@@ -256,15 +259,37 @@ type FinishTaskAttemptInput struct {
|
||||
}
|
||||
|
||||
type FinishTaskSuccessInput struct {
|
||||
TaskID string
|
||||
ExecutionToken string
|
||||
AttemptID string
|
||||
Result map[string]any
|
||||
Billings []any
|
||||
RequestID string
|
||||
ResolvedModel string
|
||||
Usage map[string]any
|
||||
Metrics map[string]any
|
||||
BillingSummary map[string]any
|
||||
FinalChargeAmount float64
|
||||
FinalChargeAmountText string
|
||||
BillingCurrency string
|
||||
PricingSnapshot map[string]any
|
||||
RequestFingerprint string
|
||||
ResponseStartedAt time.Time
|
||||
ResponseFinishedAt time.Time
|
||||
ResponseDurationMS int64
|
||||
}
|
||||
|
||||
type FinishTaskManualReviewInput struct {
|
||||
TaskID string
|
||||
ExecutionToken string
|
||||
AttemptID string
|
||||
TaskStatus string
|
||||
Code string
|
||||
Message string
|
||||
Result map[string]any
|
||||
Billings []any
|
||||
RequestID string
|
||||
ResolvedModel string
|
||||
Usage map[string]any
|
||||
Metrics map[string]any
|
||||
BillingSummary map[string]any
|
||||
FinalChargeAmount float64
|
||||
PricingSnapshot map[string]any
|
||||
RequestFingerprint string
|
||||
ResponseStartedAt time.Time
|
||||
ResponseFinishedAt time.Time
|
||||
ResponseDurationMS int64
|
||||
@@ -272,6 +297,7 @@ type FinishTaskSuccessInput struct {
|
||||
|
||||
type FinishTaskFailureInput struct {
|
||||
TaskID string
|
||||
ExecutionToken string
|
||||
Code string
|
||||
Message string
|
||||
Result map[string]any
|
||||
|
||||
@@ -3,6 +3,7 @@ package store
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -152,18 +153,71 @@ func nullableTaskListTime(value *time.Time) any {
|
||||
return *value
|
||||
}
|
||||
|
||||
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType string, normalizedRequest map[string]any) error {
|
||||
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
func (s *Store) ClaimTaskExecution(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) {
|
||||
if leaseTTL <= 0 {
|
||||
leaseTTL = 5 * time.Minute
|
||||
}
|
||||
task, err := scanGatewayTask(s.pool.QueryRow(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'running',
|
||||
model_type = NULLIF($2::text, ''),
|
||||
normalized_request = $3::jsonb,
|
||||
execution_token = $2::uuid,
|
||||
execution_lease_expires_at = now() + ($3::int * interval '1 second'),
|
||||
locked_at = now(),
|
||||
heartbeat_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, modelType, string(normalizedJSON))
|
||||
return err
|
||||
WHERE id = $1::uuid
|
||||
AND (
|
||||
(status = 'queued' AND next_run_at <= now())
|
||||
OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()))
|
||||
)
|
||||
RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return GatewayTask{}, ErrTaskExecutionLeaseUnavailable
|
||||
}
|
||||
return task, err
|
||||
}
|
||||
|
||||
func (s *Store) RenewTaskExecutionLease(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) error {
|
||||
if leaseTTL <= 0 {
|
||||
leaseTTL = 5 * time.Minute
|
||||
}
|
||||
tag, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET execution_lease_expires_at = now() + ($3::int * interval '1 second'),
|
||||
heartbeat_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'
|
||||
AND execution_token = $2::uuid`, taskID, executionToken, int(leaseTTL/time.Second))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, normalizedRequest map[string]any) error {
|
||||
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
|
||||
tag, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'running',
|
||||
model_type = NULLIF($3::text, ''),
|
||||
normalized_request = $4::jsonb,
|
||||
locked_at = now(),
|
||||
heartbeat_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'
|
||||
AND execution_token = $2::uuid`, taskID, executionToken, modelType, string(normalizedJSON))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) ClaimAsyncQueuedTask(ctx context.Context, workerID string) (GatewayTask, error) {
|
||||
@@ -189,7 +243,7 @@ WHERE t.id = picked.task_id
|
||||
RETURNING `+gatewayTaskColumns, workerID))
|
||||
}
|
||||
|
||||
func (s *Store) RequeueTask(ctx context.Context, taskID string, delay time.Duration, queueKey string) (GatewayTask, error) {
|
||||
func (s *Store) RequeueTask(ctx context.Context, taskID string, executionToken string, delay time.Duration, queueKey string) (GatewayTask, error) {
|
||||
if delay < time.Second {
|
||||
delay = time.Second
|
||||
}
|
||||
@@ -203,14 +257,18 @@ SET status = 'queued',
|
||||
locked_by = NULL,
|
||||
locked_at = NULL,
|
||||
heartbeat_at = NULL,
|
||||
next_run_at = $2::timestamptz,
|
||||
queue_key = COALESCE(NULLIF($3::text, ''), queue_key),
|
||||
execution_token = NULL,
|
||||
execution_lease_expires_at = NULL,
|
||||
next_run_at = $3::timestamptz,
|
||||
queue_key = COALESCE(NULLIF($4::text, ''), queue_key),
|
||||
error = NULL,
|
||||
error_code = NULL,
|
||||
error_message = NULL,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
RETURNING `+gatewayTaskColumns, taskID, nextRunAt, strings.TrimSpace(queueKey)))
|
||||
AND status = 'running'
|
||||
AND execution_token = $2::uuid
|
||||
RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt, strings.TrimSpace(queueKey)))
|
||||
}
|
||||
|
||||
func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error {
|
||||
@@ -225,25 +283,32 @@ WHERE id = $1::uuid`, taskID, riverJobID)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, attemptID string, remoteTaskID string, payload map[string]any) error {
|
||||
func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, executionToken string, attemptID string, remoteTaskID string, payload map[string]any) error {
|
||||
payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload))
|
||||
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET remote_task_id = NULLIF($2::text, ''),
|
||||
remote_task_payload = $3::jsonb,
|
||||
SET remote_task_id = NULLIF($3::text, ''),
|
||||
remote_task_payload = $4::jsonb,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'
|
||||
AND execution_token = $2::uuid`,
|
||||
taskID,
|
||||
executionToken,
|
||||
remoteTaskID,
|
||||
string(payloadJSON),
|
||||
); err != nil {
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
if strings.TrimSpace(attemptID) == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := tx.Exec(ctx, `
|
||||
_, err = tx.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET remote_task_id = NULLIF($2::text, ''),
|
||||
response_snapshot = COALESCE(response_snapshot, '{}'::jsonb) || jsonb_build_object('remote_task_payload', $3::jsonb)
|
||||
@@ -256,6 +321,20 @@ WHERE id = $1::uuid`,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) SetAttemptUpstreamSubmissionStatus(ctx context.Context, attemptID string, status string) error {
|
||||
switch status {
|
||||
case "not_submitted", "submitting", "response_received":
|
||||
default:
|
||||
return fmt.Errorf("invalid upstream submission status %q", status)
|
||||
}
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET upstream_submission_status = $2,
|
||||
upstream_submission_updated_at = now()
|
||||
WHERE id = $1::uuid`, attemptID, status)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message string) (GatewayTask, bool, error) {
|
||||
message = strings.TrimSpace(message)
|
||||
if message == "" {
|
||||
@@ -293,10 +372,13 @@ func (s *Store) ListRecoverableAsyncTasks(ctx context.Context, limit int) ([]Asy
|
||||
limit = 500
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, priority
|
||||
SELECT id::text, priority, next_run_at
|
||||
FROM gateway_tasks
|
||||
WHERE async_mode = true
|
||||
AND status IN ('queued', 'running')
|
||||
AND (
|
||||
status = 'queued'
|
||||
OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()))
|
||||
)
|
||||
ORDER BY priority ASC, created_at ASC
|
||||
LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
@@ -306,7 +388,7 @@ LIMIT $1`, limit)
|
||||
items := make([]AsyncTaskQueueItem, 0)
|
||||
for rows.Next() {
|
||||
var item AsyncTaskQueueItem
|
||||
if err := rows.Scan(&item.ID, &item.Priority); err != nil {
|
||||
if err := rows.Scan(&item.ID, &item.Priority, &item.NextRunAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
@@ -320,6 +402,7 @@ LIMIT $1`, limit)
|
||||
func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) {
|
||||
requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot))
|
||||
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
|
||||
pricingJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot))
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -330,11 +413,13 @@ func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptIn
|
||||
err = tx.QueryRow(ctx, `
|
||||
INSERT INTO gateway_task_attempts (
|
||||
task_id, attempt_no, platform_id, platform_model_id, client_id, queue_key,
|
||||
status, simulated, request_snapshot, metrics
|
||||
status, simulated, request_snapshot, metrics, pricing_snapshot, request_fingerprint,
|
||||
upstream_submission_status, upstream_submission_updated_at
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, $2::int, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, NULLIF($5::text, ''), $6,
|
||||
$7, $8, $9::jsonb, $10::jsonb
|
||||
$7, $8, $9::jsonb, $10::jsonb, $11::jsonb, NULLIF($12, ''),
|
||||
'not_submitted', now()
|
||||
)
|
||||
RETURNING id::text`,
|
||||
input.TaskID,
|
||||
@@ -347,6 +432,8 @@ RETURNING id::text`,
|
||||
input.Simulated,
|
||||
string(requestJSON),
|
||||
string(metricsJSON),
|
||||
string(pricingJSON),
|
||||
input.RequestFingerprint,
|
||||
).Scan(&attemptID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -489,6 +576,8 @@ SELECT a.id::text, a.task_id::text, a.attempt_no,
|
||||
a.request_snapshot, COALESCE(a.response_snapshot, '{}'::jsonb),
|
||||
COALESCE(a.response_started_at::text, ''), COALESCE(a.response_finished_at::text, ''),
|
||||
COALESCE(a.response_duration_ms, 0), COALESCE(a.error_code, ''), COALESCE(a.error_message, ''),
|
||||
COALESCE(a.pricing_snapshot, '{}'::jsonb), COALESCE(a.request_fingerprint, ''),
|
||||
a.upstream_submission_status, COALESCE(a.upstream_submission_updated_at::text, ''),
|
||||
a.started_at, COALESCE(a.finished_at::text, '')
|
||||
FROM gateway_task_attempts a
|
||||
LEFT JOIN integration_platforms p ON p.id = a.platform_id
|
||||
@@ -518,6 +607,7 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
|
||||
var metricsBytes []byte
|
||||
var requestBytes []byte
|
||||
var responseBytes []byte
|
||||
var pricingSnapshotBytes []byte
|
||||
if err := scanner.Scan(
|
||||
&item.ID,
|
||||
&item.TaskID,
|
||||
@@ -544,6 +634,10 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
|
||||
&item.ResponseDurationMS,
|
||||
&item.ErrorCode,
|
||||
&item.ErrorMessage,
|
||||
&pricingSnapshotBytes,
|
||||
&item.RequestFingerprint,
|
||||
&item.UpstreamSubmissionStatus,
|
||||
&item.UpstreamSubmissionUpdatedAt,
|
||||
&item.StartedAt,
|
||||
&item.FinishedAt,
|
||||
); err != nil {
|
||||
@@ -553,6 +647,7 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
|
||||
item.Metrics = decodeObject(metricsBytes)
|
||||
item.RequestSnapshot = decodeObject(requestBytes)
|
||||
item.ResponseSnapshot = decodeObject(responseBytes)
|
||||
item.PricingSnapshot = decodeObject(pricingSnapshotBytes)
|
||||
enrichTaskAttemptFromMetrics(&item)
|
||||
return item, nil
|
||||
}
|
||||
@@ -670,7 +765,41 @@ func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessIn
|
||||
usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage))
|
||||
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
|
||||
billingSummaryJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingSummary))
|
||||
if _, err := s.pool.Exec(ctx, `
|
||||
pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot))
|
||||
finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText)
|
||||
if finalChargeAmount == "" {
|
||||
finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64)
|
||||
}
|
||||
currency := normalizeWalletCurrency(input.BillingCurrency)
|
||||
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
if strings.TrimSpace(input.AttemptID) != "" {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET status = 'succeeded',
|
||||
retryable = false,
|
||||
request_id = NULLIF($2, ''),
|
||||
usage = $3::jsonb,
|
||||
metrics = $4::jsonb,
|
||||
response_snapshot = $5::jsonb,
|
||||
pricing_snapshot = $6::jsonb,
|
||||
request_fingerprint = NULLIF($7, ''),
|
||||
upstream_submission_status = 'response_received',
|
||||
upstream_submission_updated_at = now(),
|
||||
response_started_at = $8::timestamptz,
|
||||
response_finished_at = $9::timestamptz,
|
||||
response_duration_ms = $10,
|
||||
error_code = NULL,
|
||||
error_message = NULL,
|
||||
finished_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
input.AttemptID, input.RequestID, string(usageJSON), string(metricsJSON),
|
||||
string(resultJSON), string(pricingSnapshotJSON), input.RequestFingerprint,
|
||||
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS,
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'succeeded',
|
||||
result = $2::jsonb,
|
||||
@@ -680,32 +809,142 @@ SET status = 'succeeded',
|
||||
usage = $6::jsonb,
|
||||
metrics = $7::jsonb,
|
||||
billing_summary = $8::jsonb,
|
||||
final_charge_amount = $9,
|
||||
response_started_at = $10::timestamptz,
|
||||
response_finished_at = $11::timestamptz,
|
||||
response_duration_ms = $12,
|
||||
final_charge_amount = $9::numeric,
|
||||
billing_version = 'effective-pricing-v2',
|
||||
billing_status = CASE
|
||||
WHEN run_mode = 'production' AND gateway_user_id IS NOT NULL THEN 'pending'
|
||||
ELSE 'not_required'
|
||||
END,
|
||||
billing_currency = $10,
|
||||
pricing_snapshot = $11::jsonb,
|
||||
request_fingerprint = NULLIF($12, ''),
|
||||
billing_updated_at = now(),
|
||||
response_started_at = $13::timestamptz,
|
||||
response_finished_at = $14::timestamptz,
|
||||
response_duration_ms = $15,
|
||||
error = NULL,
|
||||
error_code = NULL,
|
||||
error_message = NULL,
|
||||
locked_by = NULL,
|
||||
locked_at = NULL,
|
||||
heartbeat_at = NULL,
|
||||
execution_token = NULL,
|
||||
execution_lease_expires_at = NULL,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
input.TaskID,
|
||||
string(resultJSON),
|
||||
string(billingsJSON),
|
||||
input.RequestID,
|
||||
input.ResolvedModel,
|
||||
string(usageJSON),
|
||||
string(metricsJSON),
|
||||
string(billingSummaryJSON),
|
||||
input.FinalChargeAmount,
|
||||
nullableTime(input.ResponseStartedAt),
|
||||
nullableTime(input.ResponseFinishedAt),
|
||||
input.ResponseDurationMS,
|
||||
); err != nil {
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'
|
||||
AND execution_token = $16::uuid`,
|
||||
input.TaskID,
|
||||
string(resultJSON),
|
||||
string(billingsJSON),
|
||||
input.RequestID,
|
||||
input.ResolvedModel,
|
||||
string(usageJSON),
|
||||
string(metricsJSON),
|
||||
string(billingSummaryJSON),
|
||||
finalChargeAmount,
|
||||
currency,
|
||||
string(pricingSnapshotJSON),
|
||||
input.RequestFingerprint,
|
||||
nullableTime(input.ResponseStartedAt),
|
||||
nullableTime(input.ResponseFinishedAt),
|
||||
input.ResponseDurationMS,
|
||||
input.ExecutionToken,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
payloadJSON, _ := json.Marshal(map[string]any{"taskId": input.TaskID, "pricingVersion": "effective-pricing-v2"})
|
||||
_, err = tx.Exec(ctx, `
|
||||
INSERT INTO settlement_outbox (
|
||||
task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at
|
||||
)
|
||||
SELECT id, 'task.billing.settle', 'settle', $2::numeric, $3, $4::jsonb, $5::jsonb, 'pending', now()
|
||||
FROM gateway_tasks
|
||||
WHERE id = $1::uuid
|
||||
AND run_mode = 'production'
|
||||
AND gateway_user_id IS NOT NULL
|
||||
ON CONFLICT (task_id, event_type) DO NOTHING`,
|
||||
input.TaskID, finalChargeAmount, currency, string(pricingSnapshotJSON), string(payloadJSON))
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
return s.GetTask(ctx, input.TaskID)
|
||||
}
|
||||
|
||||
func (s *Store) FinishTaskManualReview(ctx context.Context, input FinishTaskManualReviewInput) (GatewayTask, error) {
|
||||
status := strings.TrimSpace(input.TaskStatus)
|
||||
if status != "succeeded" && status != "failed" {
|
||||
return GatewayTask{}, fmt.Errorf("manual review task status must be succeeded or failed")
|
||||
}
|
||||
resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result))
|
||||
pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot))
|
||||
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
if strings.TrimSpace(input.AttemptID) != "" {
|
||||
attemptStatus := "failed"
|
||||
if status == "succeeded" {
|
||||
attemptStatus = "succeeded"
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET status = $2,
|
||||
request_id = NULLIF($3, ''),
|
||||
error_code = CASE WHEN $2 = 'failed' THEN NULLIF($4, '') ELSE NULL END,
|
||||
error_message = CASE WHEN $2 = 'failed' THEN NULLIF($5, '') ELSE NULL END,
|
||||
upstream_submission_status = CASE WHEN $2 = 'succeeded' THEN 'response_received' ELSE upstream_submission_status END,
|
||||
upstream_submission_updated_at = now(),
|
||||
response_started_at = $6::timestamptz,
|
||||
response_finished_at = $7::timestamptz,
|
||||
response_duration_ms = $8,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, input.Message,
|
||||
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = $2,
|
||||
result = $3::jsonb,
|
||||
request_id = NULLIF($4, ''),
|
||||
error = CASE WHEN $2 = 'failed' THEN NULLIF($6, '') ELSE NULL END,
|
||||
error_code = NULLIF($5, ''),
|
||||
error_message = NULLIF($6, ''),
|
||||
billing_status = 'manual_review',
|
||||
pricing_snapshot = CASE WHEN $7::jsonb = '{}'::jsonb THEN pricing_snapshot ELSE $7::jsonb END,
|
||||
request_fingerprint = COALESCE(NULLIF($8, ''), request_fingerprint),
|
||||
billing_updated_at = now(),
|
||||
response_started_at = $9::timestamptz,
|
||||
response_finished_at = $10::timestamptz,
|
||||
response_duration_ms = $11,
|
||||
locked_by = NULL,
|
||||
locked_at = NULL,
|
||||
heartbeat_at = NULL,
|
||||
execution_token = NULL,
|
||||
execution_lease_expires_at = NULL,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'
|
||||
AND execution_token = $12::uuid`, input.TaskID, status, string(resultJSON), input.RequestID, input.Code, input.Message,
|
||||
string(pricingSnapshotJSON), input.RequestFingerprint, nullableTime(input.ResponseStartedAt),
|
||||
nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ExecutionToken)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
return s.GetTask(ctx, input.TaskID)
|
||||
@@ -743,6 +982,9 @@ ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
|
||||
task.GatewayTenantID, task.GatewayUserID, task.TenantID, task.TenantKey, task.UserID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := ensureWalletAccountAuditGuard(ctx, tx, task.GatewayUserID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
var exists bool
|
||||
var accountID string
|
||||
var balanceBefore float64
|
||||
@@ -859,7 +1101,8 @@ func taskBillingString(value any) string {
|
||||
func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) {
|
||||
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
|
||||
resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result))
|
||||
if _, err := s.pool.Exec(ctx, `
|
||||
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'failed',
|
||||
error = NULLIF($2::text, ''),
|
||||
@@ -871,22 +1114,55 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
|
||||
response_finished_at = $7::timestamptz,
|
||||
response_duration_ms = $8,
|
||||
result = $9::jsonb,
|
||||
billing_status = CASE
|
||||
WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
|
||||
WHEN reservation_amount > 0 THEN 'pending'
|
||||
ELSE 'released'
|
||||
END,
|
||||
billing_updated_at = now(),
|
||||
locked_by = NULL,
|
||||
locked_at = NULL,
|
||||
heartbeat_at = NULL,
|
||||
execution_token = NULL,
|
||||
execution_lease_expires_at = NULL,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
input.TaskID,
|
||||
input.Message,
|
||||
input.Code,
|
||||
input.RequestID,
|
||||
string(metricsJSON),
|
||||
nullableTime(input.ResponseStartedAt),
|
||||
nullableTime(input.ResponseFinishedAt),
|
||||
input.ResponseDurationMS,
|
||||
string(resultJSON),
|
||||
); err != nil {
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'
|
||||
AND execution_token = $10::uuid`,
|
||||
input.TaskID,
|
||||
input.Message,
|
||||
input.Code,
|
||||
input.RequestID,
|
||||
string(metricsJSON),
|
||||
nullableTime(input.ResponseStartedAt),
|
||||
nullableTime(input.ResponseFinishedAt),
|
||||
input.ResponseDurationMS,
|
||||
string(resultJSON),
|
||||
input.ExecutionToken,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
payloadJSON, _ := json.Marshal(map[string]any{"taskId": input.TaskID, "reason": input.Code})
|
||||
_, err = tx.Exec(ctx, `
|
||||
INSERT INTO settlement_outbox (
|
||||
task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at
|
||||
)
|
||||
SELECT id, 'task.billing.release', 'release', reservation_amount, billing_currency,
|
||||
pricing_snapshot, $2::jsonb, 'pending', now()
|
||||
FROM gateway_tasks
|
||||
WHERE id = $1::uuid
|
||||
AND run_mode = 'production'
|
||||
AND gateway_user_id IS NOT NULL
|
||||
AND reservation_amount > 0
|
||||
ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, string(payloadJSON))
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
return s.GetTask(ctx, input.TaskID)
|
||||
|
||||
@@ -3,6 +3,7 @@ package store
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -82,6 +83,7 @@ type WalletBalanceAdjustmentInput struct {
|
||||
GatewayUserID string `json:"gatewayUserId"`
|
||||
Currency string `json:"currency"`
|
||||
Balance float64 `json:"balance"`
|
||||
BalanceText string `json:"-"`
|
||||
Reason string `json:"reason"`
|
||||
IdempotencyKey string `json:"idempotencyKey"`
|
||||
Metadata map[string]any `json:"metadata"`
|
||||
@@ -91,6 +93,7 @@ type WalletRechargeInput struct {
|
||||
GatewayUserID string `json:"gatewayUserId"`
|
||||
Currency string `json:"currency"`
|
||||
Amount float64 `json:"amount"`
|
||||
AmountText string `json:"-"`
|
||||
Reason string `json:"reason"`
|
||||
IdempotencyKey string `json:"idempotencyKey"`
|
||||
Metadata map[string]any `json:"metadata"`
|
||||
@@ -135,7 +138,7 @@ 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) {
|
||||
func (s *Store) ReserveTaskBilling(ctx context.Context, task GatewayTask, user *auth.User, billings []any, pricingSnapshots ...map[string]any) ([]WalletBillingReservation, error) {
|
||||
gatewayUserID := taskGatewayUserID(task, user)
|
||||
if gatewayUserID == "" {
|
||||
return nil, nil
|
||||
@@ -145,12 +148,21 @@ func (s *Store) ReserveTaskBilling(ctx context.Context, task GatewayTask, user *
|
||||
return nil, fmt.Errorf("task id is required for wallet reservation")
|
||||
}
|
||||
|
||||
pricingSnapshot := map[string]any{}
|
||||
if len(pricingSnapshots) > 0 {
|
||||
pricingSnapshot = emptyObjectIfNil(pricingSnapshots[0])
|
||||
}
|
||||
if exactAmount := walletString(pricingSnapshot["reservationAmount"]); exactAmount != "" {
|
||||
return s.reserveTaskBillingExact(ctx, task, gatewayUserID, exactAmount, pricingSnapshot)
|
||||
}
|
||||
amounts := walletBillingAmounts(billings)
|
||||
if len(amounts) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
reservations := make([]WalletBillingReservation, 0, len(amounts))
|
||||
pricingSnapshotJSON, _ := json.Marshal(pricingSnapshot)
|
||||
requestFingerprint := walletString(pricingSnapshot["requestFingerprint"])
|
||||
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
for currency, rawAmount := range amounts {
|
||||
amount := roundMoney(rawAmount)
|
||||
@@ -257,7 +269,30 @@ VALUES (
|
||||
}
|
||||
reservations = append(reservations, reservation)
|
||||
}
|
||||
return nil
|
||||
_, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks task
|
||||
SET reservation_amount = COALESCE((
|
||||
SELECT SUM(reserve.amount)
|
||||
FROM gateway_wallet_transactions reserve
|
||||
WHERE reserve.reference_type = 'gateway_task'
|
||||
AND reserve.reference_id = task.id::text
|
||||
AND reserve.transaction_type = 'reserve'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM gateway_wallet_transactions release
|
||||
WHERE release.account_id = reserve.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = reserve.idempotency_key || ':release'
|
||||
)
|
||||
), 0),
|
||||
billing_status = 'pending',
|
||||
billing_currency = 'resource',
|
||||
pricing_snapshot = $2::jsonb,
|
||||
request_fingerprint = NULLIF($3, ''),
|
||||
billing_updated_at = now(),
|
||||
updated_at = now()
|
||||
WHERE task.id = $1::uuid`, taskID, string(pricingSnapshotJSON), requestFingerprint)
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -265,6 +300,106 @@ VALUES (
|
||||
return reservations, err
|
||||
}
|
||||
|
||||
func (s *Store) reserveTaskBillingExact(ctx context.Context, task GatewayTask, gatewayUserID string, amount string, pricingSnapshot map[string]any) ([]WalletBillingReservation, error) {
|
||||
taskID := strings.TrimSpace(task.ID)
|
||||
currency := normalizeWalletCurrency(walletString(pricingSnapshot["currency"]))
|
||||
if currency != "resource" {
|
||||
return nil, fmt.Errorf("unsupported billing currency %q", currency)
|
||||
}
|
||||
pricingSnapshotJSON, _ := json.Marshal(pricingSnapshot)
|
||||
requestFingerprint := walletString(pricingSnapshot["requestFingerprint"])
|
||||
var reservations []WalletBillingReservation
|
||||
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
var positive bool
|
||||
if err := tx.QueryRow(ctx, `SELECT $1::numeric(38, 9) > 0`, amount).Scan(&positive); err != nil {
|
||||
return fmt.Errorf("invalid exact reservation amount: %w", err)
|
||||
}
|
||||
if !positive {
|
||||
return nil
|
||||
}
|
||||
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 := activeWalletReservationExact(ctx, tx, locked.ID, taskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if activeAmount != "" {
|
||||
presentationAmount, _ := strconv.ParseFloat(activeAmount, 64)
|
||||
reservations = []WalletBillingReservation{{
|
||||
TaskID: taskID, AccountID: locked.ID, GatewayUserID: gatewayUserID,
|
||||
GatewayTenantID: firstNonEmpty(locked.GatewayTenantID, task.GatewayTenantID),
|
||||
Currency: locked.Currency, Amount: presentationAmount, IdempotencyKey: activeKey,
|
||||
}}
|
||||
return nil
|
||||
}
|
||||
sequence, err := nextWalletReservationSequence(ctx, tx, locked.ID, taskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key := billingReservationIdempotencyKey(taskID, locked.Currency, sequence)
|
||||
var balanceBefore string
|
||||
var frozenBefore string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT balance::text, frozen_balance::text
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE id = $1::uuid`, locked.ID).Scan(&balanceBefore, &frozenBefore); err != nil {
|
||||
return err
|
||||
}
|
||||
var frozenAfter string
|
||||
err = tx.QueryRow(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET frozen_balance = frozen_balance + $2::numeric(38, 9), updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND balance - frozen_balance >= $2::numeric(38, 9)
|
||||
RETURNING frozen_balance::text`, locked.ID, amount).Scan(&frozenAfter)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return ErrInsufficientWalletBalance
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "reserved": amount, "balance": balanceBefore,
|
||||
"frozenBefore": frozenBefore, "frozenAfter": frozenAfter,
|
||||
"pricingVersion": walletString(pricingSnapshot["pricingVersion"]),
|
||||
})
|
||||
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::numeric(38, 9), $5::numeric, $5::numeric, $6, 'gateway_task', $7, $8::jsonb
|
||||
)`, locked.ID, firstNonEmpty(locked.GatewayTenantID, task.GatewayTenantID), gatewayUserID,
|
||||
amount, balanceBefore, key, taskID, string(metadata)); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET reservation_amount = $2::numeric(38, 9), billing_status = 'pending',
|
||||
billing_currency = $3, pricing_snapshot = $4::jsonb,
|
||||
request_fingerprint = NULLIF($5, ''), billing_updated_at = now(), updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, amount, currency, string(pricingSnapshotJSON), requestFingerprint); err != nil {
|
||||
return err
|
||||
}
|
||||
presentationAmount, _ := strconv.ParseFloat(amount, 64)
|
||||
reservations = []WalletBillingReservation{{
|
||||
TaskID: taskID, AccountID: locked.ID, GatewayUserID: gatewayUserID,
|
||||
GatewayTenantID: firstNonEmpty(locked.GatewayTenantID, task.GatewayTenantID),
|
||||
Currency: locked.Currency, Amount: presentationAmount, IdempotencyKey: key,
|
||||
}}
|
||||
return nil
|
||||
})
|
||||
return reservations, err
|
||||
}
|
||||
|
||||
func (s *Store) ReleaseTaskBillingReservations(ctx context.Context, reservations []WalletBillingReservation, reason string) error {
|
||||
if len(reservations) == 0 {
|
||||
return nil
|
||||
@@ -274,10 +409,14 @@ func (s *Store) ReleaseTaskBillingReservations(ctx context.Context, reservations
|
||||
reason = "task_not_settled"
|
||||
}
|
||||
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
taskIDs := map[string]struct{}{}
|
||||
for _, reservation := range reservations {
|
||||
if reservation.Amount <= 0 || strings.TrimSpace(reservation.AccountID) == "" {
|
||||
if strings.TrimSpace(reservation.AccountID) == "" {
|
||||
continue
|
||||
}
|
||||
if taskID := strings.TrimSpace(reservation.TaskID); taskID != "" {
|
||||
taskIDs[taskID] = struct{}{}
|
||||
}
|
||||
reserveKey := strings.TrimSpace(reservation.IdempotencyKey)
|
||||
if reserveKey == "" {
|
||||
reserveKey = billingReservationIdempotencyKey(reservation.TaskID, reservation.Currency, 1)
|
||||
@@ -303,39 +442,44 @@ SELECT EXISTS (
|
||||
if alreadyReleased {
|
||||
continue
|
||||
}
|
||||
var storedReservedAmount float64
|
||||
var storedReservedAmount string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE((
|
||||
SELECT amount::float8
|
||||
SELECT amount::text
|
||||
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 {
|
||||
), '0')`, reservation.AccountID, reserveKey).Scan(&storedReservedAmount); err != nil {
|
||||
return err
|
||||
}
|
||||
if storedReservedAmount <= 0 {
|
||||
var positive bool
|
||||
if err := tx.QueryRow(ctx, `SELECT $1::numeric > 0`, storedReservedAmount).Scan(&positive); err != nil {
|
||||
return err
|
||||
}
|
||||
if !positive {
|
||||
continue
|
||||
}
|
||||
|
||||
amount := roundMoney(storedReservedAmount)
|
||||
frozenAfter := roundMoney(locked.FrozenBalance - amount)
|
||||
if frozenAfter < 0 {
|
||||
frozenAfter = 0
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
var balanceBefore string
|
||||
var frozenBefore string
|
||||
var frozenAfter string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET frozen_balance = $2,
|
||||
SET frozen_balance = frozen_balance - $2::numeric,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, locked.ID, frozenAfter); err != nil {
|
||||
WHERE id = $1::uuid
|
||||
AND frozen_balance >= $2::numeric
|
||||
RETURNING balance::text, (frozen_balance + $2::numeric)::text, frozen_balance::text`,
|
||||
locked.ID, storedReservedAmount).Scan(&balanceBefore, &frozenBefore, &frozenAfter); err != nil {
|
||||
return err
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": reservation.TaskID,
|
||||
"reason": reason,
|
||||
"reserved": amount,
|
||||
"frozenBefore": roundMoney(locked.FrozenBalance),
|
||||
"reserved": storedReservedAmount,
|
||||
"frozenBefore": frozenBefore,
|
||||
"frozenAfter": frozenAfter,
|
||||
})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
@@ -345,15 +489,14 @@ INSERT INTO gateway_wallet_transactions (
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'credit', 'release',
|
||||
$4, $5, $6, $7, 'gateway_task', $8, $9::jsonb
|
||||
$4::numeric, $5::numeric, $5::numeric, $6, 'gateway_task', $7, $8::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),
|
||||
storedReservedAmount,
|
||||
balanceBefore,
|
||||
releaseKey,
|
||||
reservation.TaskID,
|
||||
string(metadata),
|
||||
@@ -361,6 +504,33 @@ ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO N
|
||||
return err
|
||||
}
|
||||
}
|
||||
for taskID := range taskIDs {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks task
|
||||
SET reservation_amount = COALESCE((
|
||||
SELECT SUM(reserve.amount)
|
||||
FROM gateway_wallet_transactions reserve
|
||||
WHERE reserve.reference_type = 'gateway_task'
|
||||
AND reserve.reference_id = task.id::text
|
||||
AND reserve.transaction_type = 'reserve'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions release
|
||||
WHERE release.account_id = reserve.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = reserve.idempotency_key || ':release'
|
||||
)
|
||||
), 0),
|
||||
billing_status = CASE
|
||||
WHEN status IN ('failed', 'cancelled') THEN 'released'
|
||||
WHEN status = 'succeeded' THEN billing_status
|
||||
ELSE 'not_started'
|
||||
END,
|
||||
billing_updated_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
@@ -588,8 +758,20 @@ func (s *Store) SetUserWalletBalanceTx(ctx context.Context, tx Tx, input WalletB
|
||||
if input.GatewayUserID == "" {
|
||||
return WalletAdjustmentResult{}, ErrLocalUserRequired
|
||||
}
|
||||
if input.Balance < 0 {
|
||||
return WalletAdjustmentResult{}, fmt.Errorf("wallet balance cannot be negative")
|
||||
targetInput := strings.TrimSpace(input.BalanceText)
|
||||
if targetInput == "" {
|
||||
targetInput = strconv.FormatFloat(input.Balance, 'f', 9, 64)
|
||||
}
|
||||
targetBalance, err := canonicalWalletAmount(ctx, tx, targetInput)
|
||||
if err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
var nonnegative bool
|
||||
if err := tx.QueryRow(ctx, `SELECT $1::numeric >= 0`, targetBalance).Scan(&nonnegative); err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
if !nonnegative {
|
||||
return WalletAdjustmentResult{}, ErrInvalidWalletAmount
|
||||
}
|
||||
account, err := s.ensureWalletAccount(ctx, tx, input.GatewayUserID, input.Currency)
|
||||
if err != nil {
|
||||
@@ -608,38 +790,68 @@ FOR UPDATE`, account.ID))
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
before := locked
|
||||
nextBalance := roundMoney(input.Balance)
|
||||
delta := roundMoney(nextBalance - locked.Balance)
|
||||
if delta == 0 {
|
||||
var balanceBefore string
|
||||
var frozenBalance string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT balance::text, frozen_balance::text
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE id = $1::uuid`, locked.ID).Scan(&balanceBefore, &frozenBalance); err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
var unchanged bool
|
||||
var belowFrozen bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT $1::numeric = $2::numeric, $1::numeric < $3::numeric`,
|
||||
targetBalance, balanceBefore, frozenBalance).Scan(&unchanged, &belowFrozen); err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
if unchanged {
|
||||
return WalletAdjustmentResult{}, ErrWalletBalanceUnchanged
|
||||
}
|
||||
if belowFrozen {
|
||||
return WalletAdjustmentResult{}, ErrBalanceBelowFrozen
|
||||
}
|
||||
direction := "credit"
|
||||
amount := delta
|
||||
if delta < 0 {
|
||||
if strings.HasPrefix(balanceBefore, "-") {
|
||||
return WalletAdjustmentResult{}, ErrInvalidWalletAmount
|
||||
}
|
||||
var targetGreater bool
|
||||
if err := tx.QueryRow(ctx, `SELECT $1::numeric > $2::numeric`, targetBalance, balanceBefore).Scan(&targetGreater); err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
if !targetGreater {
|
||||
direction = "debit"
|
||||
amount = -delta
|
||||
}
|
||||
var amount string
|
||||
if err := tx.QueryRow(ctx, `SELECT abs($1::numeric - $2::numeric)::text`, targetBalance, balanceBefore).Scan(&amount); err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
reason := strings.TrimSpace(input.Reason)
|
||||
if reason == "" {
|
||||
reason = "后台余额调整"
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET balance = $2,
|
||||
total_recharged = total_recharged + CASE WHEN $3 = 'credit' THEN $4 ELSE 0 END,
|
||||
SET balance = $2::numeric,
|
||||
total_recharged = total_recharged + CASE WHEN $3 = 'credit' THEN $4::numeric ELSE 0 END,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
WHERE id = $1::uuid
|
||||
AND $2::numeric >= frozen_balance`,
|
||||
locked.ID,
|
||||
nextBalance,
|
||||
targetBalance,
|
||||
direction,
|
||||
amount,
|
||||
); err != nil {
|
||||
)
|
||||
if err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return WalletAdjustmentResult{}, ErrBalanceBelowFrozen
|
||||
}
|
||||
metadata := mergeObjects(input.Metadata, map[string]any{
|
||||
"reason": reason,
|
||||
"previousBalance": roundMoney(before.Balance),
|
||||
"targetBalance": nextBalance,
|
||||
"previousBalance": balanceBefore,
|
||||
"targetBalance": targetBalance,
|
||||
})
|
||||
metadataJSON, _ := json.Marshal(emptyObjectIfNil(metadata))
|
||||
transaction, err := scanWalletTransaction(tx.QueryRow(ctx, `
|
||||
@@ -649,7 +861,7 @@ INSERT INTO gateway_wallet_transactions (
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, $4, 'admin_adjust',
|
||||
$5, $6, $7, NULLIF($8, ''), 'gateway_user', $9, $10::jsonb
|
||||
$5::numeric, $6::numeric, $7::numeric, NULLIF($8, ''), 'gateway_user', $9, $10::jsonb
|
||||
)
|
||||
RETURNING id::text, account_id::text, COALESCE(gateway_tenant_id::text, ''), COALESCE(gateway_user_id::text, ''),
|
||||
direction, transaction_type, amount::float8, balance_before::float8, balance_after::float8,
|
||||
@@ -660,8 +872,8 @@ RETURNING id::text, account_id::text, COALESCE(gateway_tenant_id::text, ''), COA
|
||||
locked.GatewayUserID,
|
||||
direction,
|
||||
amount,
|
||||
roundMoney(before.Balance),
|
||||
nextBalance,
|
||||
balanceBefore,
|
||||
targetBalance,
|
||||
strings.TrimSpace(input.IdempotencyKey),
|
||||
locked.GatewayUserID,
|
||||
string(metadataJSON),
|
||||
@@ -669,11 +881,10 @@ RETURNING id::text, account_id::text, COALESCE(gateway_tenant_id::text, ''), COA
|
||||
if err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
locked.Balance = nextBalance
|
||||
if direction == "credit" {
|
||||
locked.TotalRecharged = roundMoney(locked.TotalRecharged + amount)
|
||||
locked, err = s.ensureWalletAccount(ctx, tx, input.GatewayUserID, input.Currency)
|
||||
if err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
locked.UpdatedAt = time.Now()
|
||||
return WalletAdjustmentResult{Account: locked, Before: before, Transaction: transaction}, nil
|
||||
}
|
||||
|
||||
@@ -695,9 +906,20 @@ func (s *Store) RechargeUserWalletBalanceTx(ctx context.Context, tx Tx, input Wa
|
||||
if input.GatewayUserID == "" {
|
||||
return WalletAdjustmentResult{}, ErrLocalUserRequired
|
||||
}
|
||||
amount := roundMoney(input.Amount)
|
||||
if amount <= 0 {
|
||||
return WalletAdjustmentResult{}, fmt.Errorf("wallet recharge amount must be positive")
|
||||
amountInput := strings.TrimSpace(input.AmountText)
|
||||
if amountInput == "" {
|
||||
amountInput = strconv.FormatFloat(input.Amount, 'f', 9, 64)
|
||||
}
|
||||
amount, err := canonicalWalletAmount(ctx, tx, amountInput)
|
||||
if err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
var positive bool
|
||||
if err := tx.QueryRow(ctx, `SELECT $1::numeric > 0`, amount).Scan(&positive); err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
if !positive {
|
||||
return WalletAdjustmentResult{}, ErrInvalidWalletAmount
|
||||
}
|
||||
account, err := s.ensureWalletAccount(ctx, tx, input.GatewayUserID, input.Currency)
|
||||
if err != nil {
|
||||
@@ -715,26 +937,30 @@ FOR UPDATE`, account.ID))
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
before := locked
|
||||
nextBalance := roundMoney(locked.Balance + amount)
|
||||
var balanceBefore string
|
||||
if err := tx.QueryRow(ctx, `SELECT balance::text FROM gateway_wallet_accounts WHERE id = $1::uuid`, locked.ID).Scan(&balanceBefore); err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
reason := strings.TrimSpace(input.Reason)
|
||||
if reason == "" {
|
||||
reason = "后台余额充值"
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
var nextBalance string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET balance = $2,
|
||||
total_recharged = total_recharged + $3,
|
||||
SET balance = balance + $2::numeric,
|
||||
total_recharged = total_recharged + $2::numeric,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
WHERE id = $1::uuid
|
||||
RETURNING balance::text`,
|
||||
locked.ID,
|
||||
nextBalance,
|
||||
amount,
|
||||
); err != nil {
|
||||
).Scan(&nextBalance); err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
metadata := mergeObjects(input.Metadata, map[string]any{
|
||||
"reason": reason,
|
||||
"previousBalance": roundMoney(before.Balance),
|
||||
"previousBalance": balanceBefore,
|
||||
"rechargeAmount": amount,
|
||||
"targetBalance": nextBalance,
|
||||
})
|
||||
@@ -746,7 +972,7 @@ INSERT INTO gateway_wallet_transactions (
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'credit', 'recharge',
|
||||
$4, $5, $6, NULLIF($7, ''), 'gateway_user', $8, $9::jsonb
|
||||
$4::numeric, $5::numeric, $6::numeric, NULLIF($7, ''), 'gateway_user', $8, $9::jsonb
|
||||
)
|
||||
RETURNING id::text, account_id::text, COALESCE(gateway_tenant_id::text, ''), COALESCE(gateway_user_id::text, ''),
|
||||
direction, transaction_type, amount::float8, balance_before::float8, balance_after::float8,
|
||||
@@ -756,7 +982,7 @@ RETURNING id::text, account_id::text, COALESCE(gateway_tenant_id::text, ''), COA
|
||||
locked.GatewayTenantID,
|
||||
locked.GatewayUserID,
|
||||
amount,
|
||||
roundMoney(before.Balance),
|
||||
balanceBefore,
|
||||
nextBalance,
|
||||
strings.TrimSpace(input.IdempotencyKey),
|
||||
locked.GatewayUserID,
|
||||
@@ -765,12 +991,21 @@ RETURNING id::text, account_id::text, COALESCE(gateway_tenant_id::text, ''), COA
|
||||
if err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
locked.Balance = nextBalance
|
||||
locked.TotalRecharged = roundMoney(locked.TotalRecharged + amount)
|
||||
locked.UpdatedAt = time.Now()
|
||||
locked, err = s.ensureWalletAccount(ctx, tx, input.GatewayUserID, input.Currency)
|
||||
if err != nil {
|
||||
return WalletAdjustmentResult{}, err
|
||||
}
|
||||
return WalletAdjustmentResult{Account: locked, Before: before, Transaction: transaction}, nil
|
||||
}
|
||||
|
||||
func canonicalWalletAmount(ctx context.Context, q Tx, value string) (string, error) {
|
||||
var canonical string
|
||||
if err := q.QueryRow(ctx, `SELECT $1::numeric(38, 9)::text`, strings.TrimSpace(value)).Scan(&canonical); err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrInvalidWalletAmount, err)
|
||||
}
|
||||
return canonical, nil
|
||||
}
|
||||
|
||||
func (s *Store) ensureWalletAccount(ctx context.Context, q Tx, gatewayUserID string, currency string) (GatewayWalletAccount, error) {
|
||||
currency = normalizeWalletCurrency(currency)
|
||||
if _, err := q.Exec(ctx, `
|
||||
@@ -794,6 +1029,9 @@ WHERE gateway_wallet_accounts.gateway_tenant_id IS NULL
|
||||
OR COALESCE(gateway_wallet_accounts.user_id, '') = ''`, gatewayUserID, currency); err != nil {
|
||||
return GatewayWalletAccount{}, err
|
||||
}
|
||||
if err := ensureWalletAccountAuditGuard(ctx, q, gatewayUserID, currency); err != nil {
|
||||
return GatewayWalletAccount{}, err
|
||||
}
|
||||
account, err := scanWalletAccount(q.QueryRow(ctx, `
|
||||
SELECT id::text, COALESCE(gateway_tenant_id::text, ''), gateway_user_id::text,
|
||||
COALESCE(tenant_id, ''), COALESCE(tenant_key, ''), COALESCE(user_id, ''),
|
||||
@@ -811,6 +1049,17 @@ WHERE gateway_user_id = $1::uuid
|
||||
return account, nil
|
||||
}
|
||||
|
||||
func ensureWalletAccountAuditGuard(ctx context.Context, q Tx, gatewayUserID string, currency string) error {
|
||||
_, err := q.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_account_audit_guards (account_id)
|
||||
SELECT id
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id = $1::uuid
|
||||
AND currency = $2
|
||||
ON CONFLICT (account_id) DO NOTHING`, gatewayUserID, normalizeWalletCurrency(currency))
|
||||
return err
|
||||
}
|
||||
|
||||
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,
|
||||
@@ -851,6 +1100,31 @@ LIMIT 1`, accountID, taskID).Scan(&key, &amount)
|
||||
return key, roundMoney(amount), nil
|
||||
}
|
||||
|
||||
func activeWalletReservationExact(ctx context.Context, tx pgx.Tx, accountID string, taskID string) (string, string, error) {
|
||||
var key string
|
||||
var amount string
|
||||
err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(t.idempotency_key, ''), t.amount::text
|
||||
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 release
|
||||
WHERE release.account_id = t.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = t.idempotency_key || ':release'
|
||||
)
|
||||
ORDER BY t.created_at DESC
|
||||
LIMIT 1`, accountID, taskID).Scan(&key, &amount)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return "", "", nil
|
||||
}
|
||||
return key, amount, err
|
||||
}
|
||||
|
||||
func nextWalletReservationSequence(ctx context.Context, tx pgx.Tx, accountID string, taskID string) (int, error) {
|
||||
var count int
|
||||
if err := tx.QueryRow(ctx, `
|
||||
|
||||
@@ -132,6 +132,10 @@ RETURNING id::text`, "wallet-reservation-user-"+suffix, "wallet_reservation_"+su
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cleanupCtx := context.Background()
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id = $1::uuid`, userID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `
|
||||
DELETE FROM gateway_wallet_account_audit_guards
|
||||
WHERE account_id IN (SELECT id FROM gateway_wallet_accounts WHERE gateway_user_id = $1::uuid)`, userID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_users WHERE id = $1::uuid`, userID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tenants WHERE id = $1::uuid`, tenantID)
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user