diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index 64f40dd..da9685a 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -730,6 +730,12 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, var firstPreprocessing parameterPreprocessingLog var walletReservations []store.WalletBillingReservation walletReservationFinalized := false + retainWalletReservationForQueue := func() { + // A queued task still owns its billing hold. Keeping the same active + // reservation lets the next Worker claim reuse its idempotency key + // instead of emitting another reserve/release pair. + walletReservationFinalized = true + } defer func() { if !walletReservationFinalized && len(walletReservations) > 0 { _ = s.store.ReleaseTaskBillingReservations(context.WithoutCancel(ctx), walletReservations, "task_not_settled") @@ -844,6 +850,7 @@ candidatesLoop: if queueErr != nil { return Result{}, queueErr } + retainWalletReservationForQueue() return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: delay} } return Result{}, admissionErr @@ -1098,6 +1105,7 @@ candidatesLoop: if queueErr != nil { return Result{}, queueErr } + retainWalletReservationForQueue() return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: delay} } attemptNo = s.recordFailedAttempt(ctx, failedAttemptRecord{ @@ -1222,6 +1230,7 @@ candidatesLoop: if queueErr != nil { return Result{}, queueErr } + retainWalletReservationForQueue() return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: 0} } if task.AsyncMode && errors.Is(lastErr, store.ErrRateLimited) && store.RateLimitRetryable(lastErr) { @@ -1232,6 +1241,7 @@ candidatesLoop: if queueErr != nil { return Result{}, queueErr } + retainWalletReservationForQueue() return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: delay} } extraMetrics := []map[string]any{} diff --git a/apps/api/internal/store/task_history_test.go b/apps/api/internal/store/task_history_test.go index 6f2cc1a..2fe9c2a 100644 --- a/apps/api/internal/store/task_history_test.go +++ b/apps/api/internal/store/task_history_test.go @@ -36,6 +36,21 @@ func TestMinimalTaskResultDropsProviderRawCopy(t *testing.T) { } } +func TestMinimalTaskAttemptMetricsKeepsBoundedRoutingSnapshot(t *testing.T) { + metrics := minimalTaskAttemptMetrics(map[string]any{ + "loadMetrics": map[string]any{"concurrent": map[string]any{"softRatio": 0.5}}, + "fullReasons": []any{"waiting"}, + "selectionReason": "full_avoided", + "requestBody": map[string]any{"secret": "must not persist"}, + }) + if metrics["loadMetrics"] == nil || metrics["fullReasons"] == nil || metrics["selectionReason"] != "full_avoided" { + t.Fatalf("routing snapshot was dropped: %+v", metrics) + } + if metrics["requestBody"] != nil { + t.Fatalf("unbounded attempt payload was retained: %+v", metrics) + } +} + func TestTaskEventsAreMinimalAndConsecutiveDuplicatesAreSkipped(t *testing.T) { db := billingV2IntegrationStore(t) ctx := context.Background() diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index 5064b6d..18c79f7 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -901,6 +901,12 @@ SET status = 'cancelled', error = NULL, error_code = 'task_cancelled', error_message = NULLIF($2::text, ''), + billing_status = CASE + WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required' + WHEN reservation_amount > 0 THEN 'pending' + ELSE 'released' + END, + billing_updated_at = now(), remote_task_payload = '{}'::jsonb, locked_by = NULL, locked_at = NULL, @@ -930,7 +936,20 @@ WHERE task_id = $1::uuid if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil { return err } - return nil + payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": "queued_cancelled"}) + _, 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 IN ('production', 'acceptance', 'acceptance_canary') + AND gateway_user_id IS NOT NULL + AND reservation_amount > 0 +ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)) + return err }) if err != nil { return GatewayTask{}, false, err @@ -1568,6 +1587,7 @@ func minimalTaskAttemptUsage(usage map[string]any) map[string]any { func minimalTaskAttemptMetrics(metrics map[string]any) map[string]any { return whitelistedTaskAttemptMap(metrics, []string{ "platformPriority", "currentPriority", "loadRatio", "loadAvoided", + "loadMetrics", "fullReasons", "selectionReason", "cacheAffinityKey", "cacheAdjustedPriority", "cacheAffinitySamples", "cacheAffinityScore", "cacheAffinityConfidence", "cacheAffinityHitRatio", "cacheAffinityEMAHitRatio", "cacheAffinityLastHitRatio", diff --git a/apps/api/internal/store/wallet_reservation_test.go b/apps/api/internal/store/wallet_reservation_test.go index 5aadd46..7aa42c1 100644 --- a/apps/api/internal/store/wallet_reservation_test.go +++ b/apps/api/internal/store/wallet_reservation_test.go @@ -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)