fix(billing): 重排任务复用钱包预留

平台额度跨窗口重排时保留任务已有的钱包预留，后续 Worker 领取复用同一幂等键，避免重复 reserve/release 事务。\n\n补齐排队任务取消后的原子结算释放，并允许任务 attempt 持久化软负载、满载原因和选择原因等有界路由快照。\n\n验证：go test ./...；go vet ./...；钱包预留复用与排队取消 PostgreSQL 集成测试。
2026-08-03 18:09:47 +08:00
parent 13faf1d072
commit c8d04ca731
4 changed files with 147 additions and 1 deletions
@@ -84,6 +84,28 @@ func TestReserveTaskBillingSerializesConcurrentWalletReservations(t *testing.T)
 	if len(successReservations) != 1 || !walletFloatNear(successReservations[0].Amount, 10) {
 		t.Fatalf("unexpected successful reservations: %+v", successReservations)
 	}
+	reusedReservations, err := db.ReserveTaskBilling(ctx, GatewayTask{
+		ID:              successReservations[0].TaskID,
+		GatewayUserID:   userID,
+		GatewayTenantID: tenantID,
+		Kind:            "images.generations",
+		Model:           "mock-image",
+	}, user, billings)
+	if err != nil || len(reusedReservations) != 1 || reusedReservations[0].IdempotencyKey != successReservations[0].IdempotencyKey {
+		t.Fatalf("active reservation was not reused: reservations=%+v err=%v", reusedReservations, err)
+	}
+	var reserveTransactions int
+	if err := db.pool.QueryRow(ctx, `
+SELECT count(*)
+FROM gateway_wallet_transactions
+WHERE reference_type='gateway_task'
+  AND reference_id=$1
+  AND transaction_type='reserve'`, successReservations[0].TaskID).Scan(&reserveTransactions); err != nil {
+		t.Fatalf("count reused wallet reservations: %v", err)
+	}
+	if reserveTransactions != 1 {
+		t.Fatalf("active reservation reuse created %d reserve transactions, want 1", reserveTransactions)
+	}

 	balance, frozen, spent := readWalletReservationAccount(t, ctx, db, userID)
 	if !walletFloatNear(balance, 10) || !walletFloatNear(frozen, 10) || !walletFloatNear(spent, 0) {
@@ -138,6 +160,85 @@ func TestWalletReservationLockSetSerializesSameWallet(t *testing.T) {
 	<-released
 }

+func TestCancelQueuedTaskReleasesRetainedBillingReservation(t *testing.T) {
+	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
+	if databaseURL == "" {
+		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the queued cancellation billing integration test")
+	}
+
+	ctx := context.Background()
+	db, err := Connect(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect store: %v", err)
+	}
+	defer db.Close()
+
+	tenantID, userID := seedWalletReservationUser(t, ctx, db)
+	if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
+		GatewayUserID: userID,
+		Currency:      "resource",
+		Balance:       10,
+		Reason:        "seed queued cancellation reservation test",
+	}); err != nil {
+		t.Fatalf("seed wallet balance: %v", err)
+	}
+	user := &auth.User{GatewayUserID: userID, GatewayTenantID: tenantID}
+	task, err := db.CreateTask(ctx, CreateTaskInput{
+		Kind: "videos.generations", Model: "queued-cancellation", RunMode: "production", Async: true,
+		Request: map[string]any{"model": "queued-cancellation"},
+	}, user)
+	if err != nil {
+		t.Fatalf("create queued task: %v", err)
+	}
+	t.Cleanup(func() {
+		cleanupCtx := context.Background()
+		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM settlement_outbox WHERE task_id=$1::uuid`, task.ID)
+		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
+		_, _ = 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_wallet_accounts WHERE gateway_user_id=$1::uuid`, userID)
+	})
+
+	reservations, err := db.ReserveTaskBilling(ctx, task, user, []any{map[string]any{"currency": "resource", "amount": float64(1)}})
+	if err != nil || len(reservations) != 1 {
+		t.Fatalf("reserve queued task billing: reservations=%+v err=%v", reservations, err)
+	}
+	if _, changed, err := db.CancelQueuedTask(ctx, task.ID, "cancel retained reservation"); err != nil || !changed {
+		t.Fatalf("cancel queued task: changed=%v err=%v", changed, err)
+	}
+
+	claims, err := db.ClaimBillingSettlements(ctx, "queued-cancellation-test", BillingSettlementBatchSize, BillingSettlementLockTimeout)
+	if err != nil {
+		t.Fatalf("claim cancellation release: %v", err)
+	}
+	settlement := settlementForTask(t, claims, task.ID)
+	if settlement.Action != "release" {
+		t.Fatalf("queued cancellation settlement=%+v, want release", settlement)
+	}
+	if err := db.ProcessBillingSettlement(ctx, settlement); err != nil {
+		t.Fatalf("process cancellation release: %v", err)
+	}
+
+	_, frozen, spent := readWalletReservationAccount(t, ctx, db, userID)
+	if !walletFloatNear(frozen, 0) || !walletFloatNear(spent, 0) {
+		t.Fatalf("queued cancellation did not release hold: frozen=%f spent=%f", frozen, spent)
+	}
+	var reserveCount int
+	var releaseCount int
+	if err := db.pool.QueryRow(ctx, `
+SELECT count(*) FILTER (WHERE transaction_type='reserve'),
+       count(*) FILTER (WHERE transaction_type='release')
+FROM gateway_wallet_transactions
+WHERE reference_type='gateway_task' AND reference_id=$1`, task.ID).Scan(&reserveCount, &releaseCount); err != nil {
+		t.Fatalf("count queued cancellation transactions: %v", err)
+	}
+	if reserveCount != 1 || releaseCount != 1 {
+		t.Fatalf("queued cancellation transactions reserve/release=%d/%d, want 1/1", reserveCount, releaseCount)
+	}
+}
+
 func seedWalletReservationUser(t *testing.T, ctx context.Context, db *Store) (string, string) {
 	t.Helper()
 	suffix := strconv.FormatInt(time.Now().UnixNano(), 10)