easyai-ai-gateway/apps/api/internal/store/rate_limits.go

package store

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"time"

	"github.com/jackc/pgx/v5"
)

type RuntimeRecoveryResult struct {
	ReleasedConcurrencyLeases int64 `json:"releasedConcurrencyLeases"`
	ReleasedRateReservations  int64 `json:"releasedRateReservations"`
	FailedAttempts            int64 `json:"failedAttempts"`
	FailedTasks               int64 `json:"failedTasks"`
	RequeuedAsyncTasks        int64 `json:"requeuedAsyncTasks"`
	CleanedTaskAdmissions     int64 `json:"cleanedTaskAdmissions"`
}

const (
	runtimeRecoveryBatchSize  = 50
	runtimeWorkerStaleAfter   = 60 * time.Second
	runtimeCounterRepairBatch = 100
)

var ErrConcurrencyLeaseLost = errors.New("concurrency lease lost")

// CheckRateLimits performs a non-consuming admission preflight for fixed-window
// limits. The same limits are checked and reserved again when execution starts.
func (s *Store) CheckRateLimits(ctx context.Context, reservations []RateLimitReservation) error {
	for _, reservation := range reservations {
		if reservation.Limit <= 0 || reservation.Amount <= 0 || reservation.Metric == "concurrent" || reservation.Metric == "queue_size" {
			continue
		}
		if reservation.Metric == "" || reservation.Amount > reservation.Limit {
			return &RateLimitExceededError{
				ScopeType:     reservation.ScopeType,
				ScopeKey:      reservation.ScopeKey,
				ScopeName:     reservation.ScopeName,
				ScopeMetadata: reservation.ScopeMetadata,
				Metric:        reservation.Metric,
				Limit:         reservation.Limit,
				Amount:        reservation.Amount,
				Projected:     reservation.Amount,
				WindowSeconds: reservation.WindowSeconds,
				Policy:        reservation.Policy,
				Message:       fmt.Sprintf("rate limit exceeded: %s request amount %.0f is greater than limit %.0f", reservation.Metric, reservation.Amount, reservation.Limit),
				Retryable:     false,
			}
		}
		windowSeconds := reservation.WindowSeconds
		if windowSeconds <= 0 {
			windowSeconds = 60
		}
		var used float64
		var reserved float64
		var resetAt time.Time
		err := s.pool.QueryRow(ctx, `
WITH bounds AS (
  SELECT to_timestamp(floor(extract(epoch FROM now()) / $4::int) * $4::int) AS window_start,
         to_timestamp(floor(extract(epoch FROM now()) / $4::int) * $4::int) + ($4::int * interval '1 second') AS reset_at
)
SELECT COALESCE(counters.used_value, 0)::float8,
       COALESCE(counters.reserved_value, 0)::float8,
       bounds.reset_at
FROM bounds
LEFT JOIN gateway_rate_limit_counters counters
  ON counters.scope_type = $1
 AND counters.scope_key = $2
 AND counters.metric = $3
 AND counters.window_start = bounds.window_start`,
			reservation.ScopeType,
			reservation.ScopeKey,
			reservation.Metric,
			windowSeconds,
		).Scan(&used, &reserved, &resetAt)
		if err != nil {
			return err
		}
		current := used + reserved
		if current+reservation.Amount > reservation.Limit {
			return &RateLimitExceededError{
				ScopeType:     reservation.ScopeType,
				ScopeKey:      reservation.ScopeKey,
				ScopeName:     reservation.ScopeName,
				ScopeMetadata: reservation.ScopeMetadata,
				Metric:        reservation.Metric,
				Limit:         reservation.Limit,
				Amount:        reservation.Amount,
				Current:       current,
				Used:          used,
				Reserved:      reserved,
				Projected:     current + reservation.Amount,
				WindowSeconds: windowSeconds,
				ResetAt:       resetAt,
				Policy:        reservation.Policy,
				Message:       fmt.Sprintf("rate limit exceeded: %s window has no remaining capacity", reservation.Metric),
				RetryAfter:    retryAfterUntil(resetAt),
				Retryable:     true,
			}
		}
	}
	return nil
}

func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID string, reservations []RateLimitReservation) (RateLimitResult, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return RateLimitResult{}, err
	}
	defer rollbackTransaction(tx)

	lockKeys := make([]string, 0)
	lockKeySet := make(map[string]struct{})
	for _, reservation := range reservations {
		if reservation.Metric == "queue_size" {
			continue
		}
		if reservation.Metric != "concurrent" || reservation.Limit <= 0 || reservation.Amount <= 0 {
			continue
		}
		key := fmt.Sprintf("%d:%s%d:%s", len(reservation.ScopeType), reservation.ScopeType, len(reservation.ScopeKey), reservation.ScopeKey)
		if _, exists := lockKeySet[key]; exists {
			continue
		}
		lockKeySet[key] = struct{}{}
		lockKeys = append(lockKeys, key)
	}
	sort.Strings(lockKeys)
	for _, key := range lockKeys {
		if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, key); err != nil {
			return RateLimitResult{}, err
		}
	}

	result := RateLimitResult{}
	for _, reservation := range reservations {
		if reservation.Limit <= 0 || reservation.Amount <= 0 {
			continue
		}
		if reservation.Metric == "" || reservation.Amount > reservation.Limit {
			return RateLimitResult{}, &RateLimitExceededError{
				ScopeType:     reservation.ScopeType,
				ScopeKey:      reservation.ScopeKey,
				ScopeName:     reservation.ScopeName,
				ScopeMetadata: reservation.ScopeMetadata,
				Metric:        reservation.Metric,
				Limit:         reservation.Limit,
				Amount:        reservation.Amount,
				Projected:     reservation.Amount,
				WindowSeconds: reservation.WindowSeconds,
				Policy:        reservation.Policy,
				Message:       fmt.Sprintf("rate limit exceeded: %s request amount %.0f is greater than limit %.0f", reservation.Metric, reservation.Amount, reservation.Limit),
				Retryable:     false,
			}
		}
		if reservation.WindowSeconds <= 0 {
			reservation.WindowSeconds = 60
		}
		if reservation.Metric == "concurrent" {
			lease, err := reserveConcurrencyLease(ctx, tx, taskID, attemptID, reservation)
			if err != nil {
				return RateLimitResult{}, err
			}
			result.Leases = append(result.Leases, lease)
			continue
		}
		normalized, err := reserveCounterWindow(ctx, tx, taskID, attemptID, reservation)
		if err != nil {
			return RateLimitResult{}, err
		}
		result.Reservations = append(result.Reservations, normalized)
	}
	return result, tx.Commit(ctx)
}

func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (ConcurrencyLease, error) {
	if reservation.LeaseTTLSeconds <= 0 {
		reservation.LeaseTTLSeconds = 120
	}
	var active float64
	var nextAvailableAt time.Time
	if err := tx.QueryRow(ctx, `
SELECT COALESCE(SUM(lease_value), 0)::float8,
       COALESCE(MIN(expires_at), now() + ($3::int * interval '1 second'))
FROM gateway_concurrency_leases
WHERE scope_type = $1
  AND scope_key = $2
  AND released_at IS NULL
  AND expires_at > now()`,
		reservation.ScopeType,
		reservation.ScopeKey,
		reservation.LeaseTTLSeconds,
	).Scan(&active, &nextAvailableAt); err != nil {
		return ConcurrencyLease{}, err
	}
	if active+reservation.Amount > reservation.Limit {
		return ConcurrencyLease{}, &RateLimitExceededError{
			ScopeType:     reservation.ScopeType,
			ScopeKey:      reservation.ScopeKey,
			ScopeName:     reservation.ScopeName,
			ScopeMetadata: reservation.ScopeMetadata,
			Metric:        reservation.Metric,
			Limit:         reservation.Limit,
			Amount:        reservation.Amount,
			Current:       active,
			Used:          active,
			Projected:     active + reservation.Amount,
			WindowSeconds: reservation.WindowSeconds,
			ResetAt:       nextAvailableAt,
			Policy:        reservation.Policy,
			Message:       fmt.Sprintf("rate limit exceeded: concurrent active %.0f plus request %.0f is greater than limit %.0f", active, reservation.Amount, reservation.Limit),
			RetryAfter:    concurrencyRetryAfter(nextAvailableAt),
			Retryable:     true,
		}
	}
	var leaseID string
	if err := tx.QueryRow(ctx, `
INSERT INTO gateway_concurrency_leases (task_id, attempt_id, scope_type, scope_key, lease_value, expires_at)
VALUES ($1::uuid, NULLIF($2, '')::uuid, $3, $4, $5, now() + ($6::int * interval '1 second'))
RETURNING id::text`,
		taskID,
		attemptID,
		reservation.ScopeType,
		reservation.ScopeKey,
		reservation.Amount,
		reservation.LeaseTTLSeconds,
	).Scan(&leaseID); err != nil {
		return ConcurrencyLease{}, err
	}
	return ConcurrencyLease{ID: leaseID, TTL: time.Duration(reservation.LeaseTTLSeconds) * time.Second}, nil
}

func reserveCounterWindow(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (RateLimitReservation, error) {
	usedAmount := 0.0
	reservedAmount := reservation.Amount
	var windowStart time.Time
	err := tx.QueryRow(ctx, `
WITH bounds AS (
  SELECT
    to_timestamp(floor(extract(epoch FROM now()) / $7::int) * $7::int) AS window_start,
    to_timestamp(floor(extract(epoch FROM now()) / $7::int) * $7::int) + ($7::int * interval '1 second') AS reset_at
)
INSERT INTO gateway_rate_limit_counters (
  scope_type, scope_key, metric, window_start, limit_value, used_value, reserved_value, reset_at
)
SELECT $1, $2, $3, bounds.window_start, $4, $5, $6, bounds.reset_at
FROM bounds
ON CONFLICT (scope_type, scope_key, metric, window_start) DO UPDATE
SET limit_value = EXCLUDED.limit_value,
    used_value = gateway_rate_limit_counters.used_value + EXCLUDED.used_value,
    reserved_value = gateway_rate_limit_counters.reserved_value + EXCLUDED.reserved_value,
    reset_at = EXCLUDED.reset_at,
    updated_at = now()
WHERE gateway_rate_limit_counters.used_value + gateway_rate_limit_counters.reserved_value + EXCLUDED.used_value + EXCLUDED.reserved_value <= EXCLUDED.limit_value
RETURNING window_start`,
		reservation.ScopeType,
		reservation.ScopeKey,
		reservation.Metric,
		reservation.Limit,
		usedAmount,
		reservedAmount,
		reservation.WindowSeconds,
	).Scan(&windowStart)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			resetAt := time.Now().Add(time.Duration(reservation.WindowSeconds) * time.Second)
			currentUsed := 0.0
			currentReserved := 0.0
			_ = tx.QueryRow(ctx, `
WITH bounds AS (
  SELECT to_timestamp(floor(extract(epoch FROM now()) / $4::int) * $4::int) AS window_start
)
SELECT counters.used_value::float8, counters.reserved_value::float8, counters.reset_at
FROM gateway_rate_limit_counters counters
JOIN bounds ON counters.window_start = bounds.window_start
WHERE scope_type = $1
  AND scope_key = $2
  AND metric = $3`,
				reservation.ScopeType,
				reservation.ScopeKey,
				reservation.Metric,
				reservation.WindowSeconds,
			).Scan(&currentUsed, &currentReserved, &resetAt)
			current := currentUsed + currentReserved
			return RateLimitReservation{}, &RateLimitExceededError{
				ScopeType:     reservation.ScopeType,
				ScopeKey:      reservation.ScopeKey,
				ScopeName:     reservation.ScopeName,
				ScopeMetadata: reservation.ScopeMetadata,
				Metric:        reservation.Metric,
				Limit:         reservation.Limit,
				Amount:        reservation.Amount,
				Current:       current,
				Used:          currentUsed,
				Reserved:      currentReserved,
				Projected:     current + reservation.Amount,
				WindowSeconds: reservation.WindowSeconds,
				ResetAt:       resetAt,
				Policy:        reservation.Policy,
				Message:       fmt.Sprintf("rate limit exceeded: %s window has no remaining capacity", reservation.Metric),
				RetryAfter:    retryAfterUntil(resetAt),
				Retryable:     true,
			}
		}
		return RateLimitReservation{}, err
	}
	reservation.WindowStart = windowStart
	if err := tx.QueryRow(ctx, `
INSERT INTO gateway_rate_limit_reservations (
  task_id, attempt_id, scope_type, scope_key, metric, window_start, limit_value, reserved_amount, status
)
VALUES (
  $1::uuid, NULLIF($2, '')::uuid, $3, $4, $5, $6, $7, $8, 'reserved'
)
RETURNING id::text`,
		taskID,
		attemptID,
		reservation.ScopeType,
		reservation.ScopeKey,
		reservation.Metric,
		windowStart,
		reservation.Limit,
		reservedAmount,
	).Scan(&reservation.ReservationID); err != nil {
		return RateLimitReservation{}, err
	}
	return reservation, nil
}

func retryAfterUntil(when time.Time) time.Duration {
	if when.IsZero() {
		return 0
	}
	duration := time.Until(when)
	if duration < time.Second {
		return time.Second
	}
	return duration
}

func concurrencyRetryAfter(leaseExpiresAt time.Time) time.Duration {
	if leaseExpiresAt.IsZero() {
		return 5 * time.Second
	}
	duration := time.Until(leaseExpiresAt)
	if duration <= time.Second {
		return time.Second
	}
	if duration > 5*time.Second {
		return 5 * time.Second
	}
	return duration
}

func (s *Store) CommitRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64) error {
	return s.finishRateLimitReservations(ctx, reservations, actualByMetric, "committed", "success")
}

func (s *Store) ReleaseRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, reason string) error {
	return s.finishRateLimitReservations(ctx, reservations, nil, "released", reason)
}

func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error {
	leaseIDs := concurrencyLeaseIDs(leases)
	if len(leaseIDs) == 0 {
		return nil
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollbackTransaction(tx)
	// A lost release after database failover is conservative: the lease
	// remains unavailable until its TTL or task-admission cleanup. It cannot
	// create duplicate execution, so avoid a dedicated synchronous wait.
	if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE id = ANY($1::uuid[]) AND released_at IS NULL`, leaseIDs); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (s *Store) RenewConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error {
	leaseIDs := make([]string, 0, len(leases))
	ttlSeconds := make([]int32, 0, len(leases))
	for _, lease := range leases {
		if lease.ID == "" {
			continue
		}
		ttl := lease.TTL
		if ttl <= 0 {
			ttl = 120 * time.Second
		}
		seconds := int32(ttl / time.Second)
		if seconds < 1 {
			seconds = 1
		}
		leaseIDs = append(leaseIDs, lease.ID)
		ttlSeconds = append(ttlSeconds, seconds)
	}
	if len(leaseIDs) == 0 {
		return nil
	}
	tag, err := s.pool.Exec(ctx, `
UPDATE gateway_concurrency_leases lease
SET expires_at = now() + (renewal.ttl_seconds * interval '1 second')
FROM unnest($1::uuid[], $2::int[]) AS renewal(id, ttl_seconds)
WHERE lease.id = renewal.id
  AND lease.released_at IS NULL
  AND lease.expires_at > now()`, leaseIDs, ttlSeconds)
	if err != nil {
		return err
	}
	if tag.RowsAffected() != int64(len(leaseIDs)) {
		return fmt.Errorf("%w: renewed %d of %d leases", ErrConcurrencyLeaseLost, tag.RowsAffected(), len(leaseIDs))
	}
	return nil
}

func (s *Store) AttachRateLimitResultToAttempt(ctx context.Context, attemptID string, result RateLimitResult) error {
	if attemptID == "" || (len(result.Reservations) == 0 && len(result.Leases) == 0) {
		return nil
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollbackTransaction(tx)
	// The subsequent durable `submitting` transition is the safety barrier for
	// these task-scoped associations. If execution stops before that point,
	// the task admission owns and reaps the leases by task_id.
	if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil {
		return err
	}

	for _, reservation := range result.Reservations {
		if reservation.ReservationID == "" {
			continue
		}
		if _, err := tx.Exec(ctx, `
UPDATE gateway_rate_limit_reservations
SET attempt_id = $2::uuid,
    updated_at = now()
WHERE id = $1::uuid`, reservation.ReservationID, attemptID); err != nil {
			return err
		}
	}
	for _, lease := range result.Leases {
		leaseID := lease.ID
		if leaseID == "" {
			continue
		}
		if _, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET attempt_id = $2::uuid
WHERE id = $1::uuid`, leaseID, attemptID); err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}

func concurrencyLeaseIDs(leases []ConcurrencyLease) []string {
	ids := make([]string, 0, len(leases))
	for _, lease := range leases {
		if lease.ID != "" {
			ids = append(ids, lease.ID)
		}
	}
	return ids
}

func (s *Store) RecoverInterruptedRuntimeState(ctx context.Context) (RuntimeRecoveryResult, error) {
	var result RuntimeRecoveryResult
	if err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
		recovered, err := recoverExpiredRuntimeTasksTx(ctx, tx)
		if err != nil {
			return err
		}
		result = recovered
		return nil
	}); err != nil {
		return RuntimeRecoveryResult{}, err
	}
	if err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
		cleaned, err := cleanupTerminalRuntimeStateTx(ctx, tx)
		if err != nil {
			return err
		}
		result.ReleasedConcurrencyLeases += cleaned.ReleasedConcurrencyLeases
		result.ReleasedRateReservations += cleaned.ReleasedRateReservations
		result.CleanedTaskAdmissions += cleaned.CleanedTaskAdmissions
		return nil
	}); err != nil {
		return RuntimeRecoveryResult{}, err
	}
	if err := s.repairRuntimeClientCounters(ctx); err != nil {
		return RuntimeRecoveryResult{}, err
	}
	if result.RequeuedAsyncTasks > 0 || result.FailedTasks > 0 || result.CleanedTaskAdmissions > 0 {
		s.notifyTaskAdmissionBestEffort(ctx, "*")
	}
	return result, nil
}

type runtimeRecoveryTask struct {
	ID               string
	Async            bool
	Production       bool
	HasGatewayUser   bool
	SubmissionUnsafe bool
}

func recoverExpiredRuntimeTasksTx(ctx context.Context, tx pgx.Tx) (RuntimeRecoveryResult, error) {
	rows, err := tx.Query(ctx, `
SELECT task.id::text,
       task.async_mode,
       task.run_mode IN ('production', 'acceptance', 'acceptance_canary'),
       task.gateway_user_id IS NOT NULL,
       COALESCE(task.remote_task_id, '') = ''
       AND COALESCE((
         SELECT (
           (attempt.status = 'running' AND attempt.upstream_submission_status IN ('submitting', 'response_received'))
           OR (attempt.status = 'failed' AND attempt.upstream_submission_status = 'submitting')
           OR (
             attempt.status = 'failed'
             AND attempt.upstream_submission_status = 'response_received'
             AND attempt.error_code = 'execution_lease_expired'
           )
         )
         FROM gateway_task_attempts attempt
         WHERE attempt.task_id = task.id
         ORDER BY attempt.attempt_no DESC, attempt.started_at DESC
         LIMIT 1
       ), false)
FROM gateway_tasks task
WHERE task.status = 'running'
  AND task.execution_lease_expires_at IS NOT NULL
  AND task.execution_lease_expires_at <= now()
  AND (
    task.locked_by IS NULL
    OR NOT EXISTS (
      SELECT 1
      FROM gateway_worker_instances worker
      WHERE worker.instance_id = task.locked_by
        AND worker.status = 'active'
        AND worker.heartbeat_at > now() - $1::interval
    )
  )
ORDER BY task.execution_lease_expires_at ASC, task.id ASC
LIMIT $2
FOR UPDATE OF task SKIP LOCKED`, runtimeWorkerStaleAfter.String(), runtimeRecoveryBatchSize)
	if err != nil {
		return RuntimeRecoveryResult{}, err
	}
	tasks := make([]runtimeRecoveryTask, 0, runtimeRecoveryBatchSize)
	for rows.Next() {
		var task runtimeRecoveryTask
		if err := rows.Scan(&task.ID, &task.Async, &task.Production, &task.HasGatewayUser, &task.SubmissionUnsafe); err != nil {
			rows.Close()
			return RuntimeRecoveryResult{}, err
		}
		tasks = append(tasks, task)
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return RuntimeRecoveryResult{}, err
	}
	rows.Close()
	if len(tasks) == 0 {
		return RuntimeRecoveryResult{}, nil
	}

	result := RuntimeRecoveryResult{}
	taskIDs := make([]string, 0, len(tasks))
	releaseBillingTaskIDs := make([]string, 0, len(tasks))
	for _, task := range tasks {
		taskIDs = append(taskIDs, task.ID)
		eventType := "task.failed"
		eventStatus := "failed"
		if task.Production && task.SubmissionUnsafe {
			if err := markTaskExecutionManualReviewTx(ctx, tx, task.ID, task.HasGatewayUser); err != nil {
				return RuntimeRecoveryResult{}, err
			}
			result.FailedTasks++
		} else if task.Async {
			if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'queued',
    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,
    next_run_at = now(),
    finished_at = NULL,
    updated_at = now()
WHERE id = $1::uuid`, task.ID); err != nil {
				return RuntimeRecoveryResult{}, err
			}
			eventType = "task.queued"
			eventStatus = "queued"
			result.RequeuedAsyncTasks++
		} else {
			if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
    billing_status = CASE
      WHEN gateway_user_id IS NULL THEN 'not_required'
      WHEN reservation_amount > 0 THEN 'pending'
      ELSE 'released'
    END,
    billing_updated_at = now(),
    error = NULL,
    error_code = 'server_restarted',
    error_message = 'task interrupted after its owner and execution lease expired',
    remote_task_payload = '{}'::jsonb,
    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`, task.ID); err != nil {
				return RuntimeRecoveryResult{}, err
			}
			if task.Production && task.HasGatewayUser {
				releaseBillingTaskIDs = append(releaseBillingTaskIDs, task.ID)
			}
			result.FailedTasks++
		}
		if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
SELECT $1::uuid,
       COALESCE(MAX(event.seq), 0) + 1,
       $2,
       $3,
       'runtime_recovery',
       0,
       NULL,
       '{}'::jsonb,
       false
FROM gateway_task_events event
WHERE event.task_id = $1::uuid
HAVING NOT EXISTS (
  SELECT 1
  FROM gateway_task_events latest
  WHERE latest.task_id = $1::uuid
    AND latest.seq = (SELECT MAX(last_event.seq) FROM gateway_task_events last_event WHERE last_event.task_id = $1::uuid)
    AND latest.event_type = $2
    AND latest.status = $3
)`, task.ID, eventType, eventStatus); err != nil {
			return RuntimeRecoveryResult{}, err
		}
	}
	if len(releaseBillingTaskIDs) > 0 {
		if _, err := tx.Exec(ctx, `
INSERT INTO settlement_outbox (
  task_id, event_type, action, amount, currency, pricing_snapshot, payload,
  status, next_attempt_at
)
SELECT task.id, 'task.billing.release', 'release', task.reservation_amount,
       task.billing_currency, task.pricing_snapshot,
       jsonb_build_object('taskId', task.id, 'reason', 'execution_lease_expired'),
       'pending', now()
FROM gateway_tasks task
WHERE task.id = ANY($1::uuid[])
  AND task.reservation_amount > 0
ON CONFLICT (task_id, event_type) DO NOTHING`, releaseBillingTaskIDs); err != nil {
			return RuntimeRecoveryResult{}, err
		}
	}
	tag, err := tx.Exec(ctx, `
UPDATE gateway_task_attempts attempt
SET status = 'failed',
    retryable = task.error_code IS DISTINCT FROM 'upstream_submission_unknown',
    error_code = CASE
      WHEN task.error_code = 'upstream_submission_unknown' THEN 'upstream_submission_unknown'
      ELSE 'execution_lease_expired'
    END,
    error_message = CASE
      WHEN task.error_code = 'upstream_submission_unknown' THEN 'upstream submission result is unknown'
      ELSE 'attempt execution lease expired after worker heartbeat became stale'
    END,
    finished_at = now()
FROM gateway_tasks task
WHERE attempt.task_id = task.id
  AND attempt.task_id = ANY($1::uuid[])
  AND attempt.status = 'running'`, taskIDs)
	if err != nil {
		return RuntimeRecoveryResult{}, err
	}
	result.FailedAttempts = tag.RowsAffected()
	tag, err = tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = ANY($1::uuid[])
  AND released_at IS NULL`, taskIDs)
	if err != nil {
		return RuntimeRecoveryResult{}, err
	}
	result.ReleasedConcurrencyLeases = tag.RowsAffected()
	tag, err = tx.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = ANY($1::uuid[])`, taskIDs)
	if err != nil {
		return RuntimeRecoveryResult{}, err
	}
	result.CleanedTaskAdmissions = tag.RowsAffected()
	released, err := releaseRuntimeRateReservationsTx(ctx, tx, taskIDs)
	if err != nil {
		return RuntimeRecoveryResult{}, err
	}
	result.ReleasedRateReservations = released
	return result, nil
}

func cleanupTerminalRuntimeStateTx(ctx context.Context, tx pgx.Tx) (RuntimeRecoveryResult, error) {
	rows, err := tx.Query(ctx, `
SELECT task.id::text
FROM gateway_tasks task
WHERE task.status IN ('succeeded', 'failed', 'cancelled')
  AND (
    EXISTS (
      SELECT 1 FROM gateway_concurrency_leases lease
      WHERE lease.task_id = task.id AND lease.released_at IS NULL
    )
    OR EXISTS (
      SELECT 1 FROM gateway_task_admissions admission
      WHERE admission.task_id = task.id
    )
    OR EXISTS (
      SELECT 1 FROM gateway_rate_limit_reservations reservation
      WHERE reservation.task_id = task.id AND reservation.status = 'reserved'
    )
  )
ORDER BY task.updated_at ASC, task.id ASC
LIMIT $1
FOR UPDATE OF task SKIP LOCKED`, runtimeRecoveryBatchSize)
	if err != nil {
		return RuntimeRecoveryResult{}, err
	}
	taskIDs := make([]string, 0, runtimeRecoveryBatchSize)
	for rows.Next() {
		var taskID string
		if err := rows.Scan(&taskID); err != nil {
			rows.Close()
			return RuntimeRecoveryResult{}, err
		}
		taskIDs = append(taskIDs, taskID)
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return RuntimeRecoveryResult{}, err
	}
	rows.Close()
	if len(taskIDs) == 0 {
		return RuntimeRecoveryResult{}, nil
	}
	var result RuntimeRecoveryResult
	tag, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = ANY($1::uuid[])
  AND released_at IS NULL`, taskIDs)
	if err != nil {
		return RuntimeRecoveryResult{}, err
	}
	result.ReleasedConcurrencyLeases = tag.RowsAffected()
	tag, err = tx.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = ANY($1::uuid[])`, taskIDs)
	if err != nil {
		return RuntimeRecoveryResult{}, err
	}
	result.CleanedTaskAdmissions = tag.RowsAffected()
	result.ReleasedRateReservations, err = releaseRuntimeRateReservationsTx(ctx, tx, taskIDs)
	return result, err
}

func releaseRuntimeRateReservationsTx(ctx context.Context, tx pgx.Tx, taskIDs []string) (int64, error) {
	rows, err := tx.Query(ctx, `
UPDATE gateway_rate_limit_reservations reservation
SET status = 'released',
    reason = 'runtime_state_recovered',
    finalized_at = now(),
    updated_at = now()
WHERE reservation.task_id = ANY($1::uuid[])
  AND reservation.status = 'reserved'
RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`, taskIDs)
	if err != nil {
		return 0, err
	}
	reservations := make([]RateLimitReservation, 0, len(taskIDs))
	for rows.Next() {
		var reservation RateLimitReservation
		if err := rows.Scan(
			&reservation.ScopeType,
			&reservation.ScopeKey,
			&reservation.Metric,
			&reservation.WindowStart,
			&reservation.Amount,
		); err != nil {
			rows.Close()
			return 0, err
		}
		reservations = append(reservations, reservation)
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return 0, err
	}
	rows.Close()
	for _, reservation := range reservations {
		if err := releaseCounterReservation(
			ctx,
			tx,
			reservation.ScopeType,
			reservation.ScopeKey,
			reservation.Metric,
			reservation.WindowStart,
			reservation.Amount,
		); err != nil {
			return 0, err
		}
	}
	return int64(len(reservations)), nil
}

func (s *Store) repairRuntimeClientCounters(ctx context.Context) error {
	_, err := s.pool.Exec(ctx, `
WITH candidates AS MATERIALIZED (
  SELECT state.client_id
  FROM runtime_client_states state
  WHERE state.running_count IS DISTINCT FROM (
    SELECT count(*)
    FROM gateway_task_attempts attempt
    WHERE attempt.client_id = state.client_id
      AND attempt.status = 'running'
  )
  ORDER BY state.updated_at ASC, state.client_id ASC
  LIMIT $1
  FOR UPDATE SKIP LOCKED
)
UPDATE runtime_client_states state
SET running_count = (
      SELECT count(*)
      FROM gateway_task_attempts attempt
      WHERE attempt.client_id = state.client_id
        AND attempt.status = 'running'
    ),
    updated_at = now()
FROM candidates
WHERE state.client_id = candidates.client_id`, runtimeCounterRepairBatch)
	return err
}

func (s *Store) finishRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64, status string, reason string) error {
	if len(reservations) == 0 {
		return nil
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollbackTransaction(tx)

	for _, reservation := range reservations {
		if reservation.ReservationID == "" {
			continue
		}
		actualAmount := actualByMetric[reservation.Metric]
		if status == "committed" && actualAmount <= 0 {
			actualAmount = reservation.Amount
		}
		var stored RateLimitReservation
		err := tx.QueryRow(ctx, `
UPDATE gateway_rate_limit_reservations
SET status = $2::text,
    reason = NULLIF($3::text, ''),
    actual_amount = CASE WHEN $2::text = 'committed' THEN $4 ELSE actual_amount END,
    finalized_at = now(),
    updated_at = now()
WHERE id = $1::uuid
  AND status = 'reserved'
RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`,
			reservation.ReservationID,
			status,
			reason,
			actualAmount,
		).Scan(&stored.ScopeType, &stored.ScopeKey, &stored.Metric, &stored.WindowStart, &stored.Amount)
		if err != nil {
			if errors.Is(err, pgx.ErrNoRows) {
				continue
			}
			return err
		}
		if status == "committed" {
			if err := commitCounterReservation(ctx, tx, stored.ScopeType, stored.ScopeKey, stored.Metric, stored.WindowStart, stored.Amount, actualAmount); err != nil {
				return err
			}
			continue
		}
		if err := releaseCounterReservation(ctx, tx, stored.ScopeType, stored.ScopeKey, stored.Metric, stored.WindowStart, stored.Amount); err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}

func commitCounterReservation(ctx context.Context, tx pgx.Tx, scopeType string, scopeKey string, metric string, windowStart time.Time, reservedAmount float64, actualAmount float64) error {
	_, err := tx.Exec(ctx, `
UPDATE gateway_rate_limit_counters
SET reserved_value = GREATEST(reserved_value - $5, 0),
    used_value = used_value + $6,
    updated_at = now()
WHERE scope_type = $1
  AND scope_key = $2
  AND metric = $3
  AND window_start = $4`,
		scopeType, scopeKey, metric, windowStart, reservedAmount, actualAmount)
	return err
}

func releaseCounterReservation(ctx context.Context, tx pgx.Tx, scopeType string, scopeKey string, metric string, windowStart time.Time, reservedAmount float64) error {
	_, err := tx.Exec(ctx, `
UPDATE gateway_rate_limit_counters
SET reserved_value = GREATEST(reserved_value - $5, 0),
    updated_at = now()
WHERE scope_type = $1
  AND scope_key = $2
  AND metric = $3
  AND window_start = $4`,
		scopeType, scopeKey, metric, windowStart, reservedAmount)
	return err
}