package store

import (
	"context"
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

func TestDistributedTaskAdmissionQueueAcrossStores(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 task admission PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)

	first, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect first store: %v", err)
	}
	defer first.Close()
	second, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect second store: %v", err)
	}
	defer second.Close()
	var databaseName string
	if err := first.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		t.Fatalf("read test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
	platform, err := first.CreatePlatform(ctx, CreatePlatformInput{
		Provider:    "queue-test",
		PlatformKey: "queue-test-" + suffix,
		Name:        "Queue Test " + suffix,
		AuthType:    "none",
		Status:      "enabled",
	})
	if err != nil {
		t.Fatalf("create platform: %v", err)
	}
	var platformModelID string
	if err := first.pool.QueryRow(ctx, `
INSERT INTO platform_models (
  platform_id, model_name, provider_model_name, model_alias, model_type,
  display_name, pricing_mode, enabled
)
VALUES ($1::uuid, $2, $2, $2, '["image_generate"]'::jsonb, $2, 'inherit_discount', true)
RETURNING id::text`, platform.ID, "queue-model-"+suffix).Scan(&platformModelID); err != nil {
		t.Fatalf("create platform model: %v", err)
	}
	taskIDs := make([]string, 0, 4)
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cleanupCancel()
		for _, taskID := range taskIDs {
			_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, taskID)
		}
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM platform_models WHERE id = $1::uuid`, platformModelID)
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_worker_instances WHERE instance_id LIKE $1`, "queue-test-"+suffix+"%")
	})

	createTask := func(async bool) GatewayTask {
		t.Helper()
		task, createErr := first.CreateTask(ctx, CreateTaskInput{
			Kind:    "images.generate",
			Model:   "queue-model-" + suffix,
			RunMode: "production",
			Async:   async,
			Request: map[string]any{"prompt": "queue admission test"},
		}, &auth.User{ID: "queue-test-user-" + suffix, Source: "gateway"})
		if createErr != nil {
			t.Fatalf("create task: %v", createErr)
		}
		taskIDs = append(taskIDs, task.ID)
		return task
	}
	scope := AdmissionScope{
		ScopeType:       "platform_model",
		ScopeKey:        platformModelID,
		ScopeName:       "queue-model",
		ConcurrentLimit: 1,
		Amount:          1,
		LeaseTTLSeconds: 137,
		QueueLimit:      2,
		MaxWaitSeconds:  600,
	}
	inputFor := func(task GatewayTask, priority int, waiterID string) TaskAdmissionInput {
		mode := "sync"
		if task.AsyncMode {
			mode = "async"
		}
		return TaskAdmissionInput{
			TaskID:          task.ID,
			PlatformID:      platform.ID,
			PlatformModelID: platformModelID,
			QueueKey:        "queue-test:" + suffix,
			Mode:            mode,
			Priority:        priority,
			WaiterID:        waiterID,
			Scopes:          []AdmissionScope{scope},
		}
	}

	running := createTask(false)
	result, err := first.TryTaskAdmission(ctx, inputFor(running, 100, "waiter-running"))
	if err != nil || !result.Admitted || len(result.Leases) != 1 {
		t.Fatalf("first task admission = %+v, err=%v", result, err)
	}
	reloadedLeases, err := second.ActiveTaskAdmissionLeases(ctx, running.ID)
	if err != nil || len(reloadedLeases) != 1 || reloadedLeases[0].TTL != 137*time.Second {
		t.Fatalf("reloaded admission leases = %+v, err=%v, want one 137s lease", reloadedLeases, err)
	}
	lowerPriority := createTask(false)
	result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
	if err != nil || result.Admitted {
		t.Fatalf("second task should wait: result=%+v err=%v", result, err)
	}
	if len(result.Blockers) != 1 || result.Blockers[0].Reason != "saturated" || result.Blockers[0].ScopeType != "platform_model" {
		t.Fatalf("second task blockers=%+v, want saturated platform_model", result.Blockers)
	}
	higherPriorityAsync := createTask(true)
	result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
	if err != nil || result.Admitted {
		t.Fatalf("third task should wait: result=%+v err=%v", result, err)
	}
	rejected := createTask(true)
	_, err = second.TryTaskAdmission(ctx, inputFor(rejected, 30, ""))
	var limitErr *RateLimitExceededError
	if !errors.As(err, &limitErr) || limitErr.Metric != "queue_size" || limitErr.Reason != "queue_full" ||
		limitErr.QueueDepth != 2 || limitErr.QueueLimit != 2 {
		t.Fatalf("fourth task error = %#v, want queue_full depth 2/2", err)
	}

	var queueCounters int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_rate_limit_counters
WHERE metric = 'queue_size'
  AND scope_key = $1`, platformModelID).Scan(&queueCounters); err != nil {
		t.Fatalf("count queue counters: %v", err)
	}
	if queueCounters != 0 {
		t.Fatalf("queue_size wrote %d fixed-window counters", queueCounters)
	}
	var attempts int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_task_attempts
WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&attempts); err != nil {
		t.Fatalf("count attempts: %v", err)
	}
	if attempts != 0 {
		t.Fatalf("waiting admissions created %d task attempts", attempts)
	}

	if err := first.DeleteTaskAdmission(ctx, running.ID); err != nil {
		t.Fatalf("release running admission: %v", err)
	}
	result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
	if err != nil {
		t.Fatalf("try lower-priority task after release: %v", err)
	}
	if result.Admitted {
		t.Fatal("lower-priority task bypassed higher-priority asynchronous waiter")
	}
	if len(result.Blockers) != 1 || result.Blockers[0].Reason != "not_queue_head" || result.Blockers[0].ScopeType != "platform_model" {
		t.Fatalf("lower-priority blockers=%+v, want platform FIFO head blocker", result.Blockers)
	}
	result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
	if err != nil || !result.Admitted {
		t.Fatalf("higher-priority async task was not admitted first: result=%+v err=%v", result, err)
	}
	for _, taskID := range taskIDs {
		_ = first.DeleteTaskAdmission(ctx, taskID)
	}

	concurrentTasks := []GatewayTask{createTask(false), createTask(true), createTask(false), createTask(true)}
	type concurrentResult struct {
		result TaskAdmissionResult
		err    error
	}
	start := make(chan struct{})
	results := make(chan concurrentResult, len(concurrentTasks))
	var waitGroup sync.WaitGroup
	for index, task := range concurrentTasks {
		waitGroup.Add(1)
		go func(index int, task GatewayTask) {
			defer waitGroup.Done()
			<-start
			target := first
			if index%2 == 1 {
				target = second
			}
			waiterID := ""
			if !task.AsyncMode {
				waiterID = "concurrent-waiter-" + task.ID
			}
			admission, admissionErr := target.TryTaskAdmission(ctx, inputFor(task, 100, waiterID))
			results <- concurrentResult{result: admission, err: admissionErr}
		}(index, task)
	}
	close(start)
	waitGroup.Wait()
	close(results)
	admittedCount := 0
	waitingCount := 0
	queueFullCount := 0
	for item := range results {
		if item.err == nil && item.result.Admitted {
			admittedCount++
			continue
		}
		if item.err == nil {
			waitingCount++
			continue
		}
		if errors.As(item.err, &limitErr) && limitErr.Metric == "queue_size" {
			queueFullCount++
			continue
		}
		t.Fatalf("unexpected concurrent admission result: %+v err=%v", item.result, item.err)
	}
	if admittedCount != 1 || waitingCount != 2 || queueFullCount != 1 {
		t.Fatalf("concurrent outcomes admitted=%d waiting=%d queue_full=%d, want 1/2/1", admittedCount, waitingCount, queueFullCount)
	}
	var activeLeases, waitingAdmissions int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
  AND scope_key = $1
  AND released_at IS NULL
  AND expires_at > now()`, platformModelID).Scan(&activeLeases); err != nil {
		t.Fatalf("count active leases: %v", err)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_task_admissions
WHERE platform_model_id = $1::uuid
  AND status = 'waiting'`, platformModelID).Scan(&waitingAdmissions); err != nil {
		t.Fatalf("count waiting admissions: %v", err)
	}
	if activeLeases != 1 || waitingAdmissions != 2 {
		t.Fatalf("database bounds active=%d waiting=%d, want 1/2", activeLeases, waitingAdmissions)
	}

	var deadlineTaskID string
	if err := first.pool.QueryRow(ctx, `
SELECT task_id::text
FROM gateway_task_admissions
WHERE platform_model_id = $1::uuid
  AND status = 'waiting'
ORDER BY task_id
LIMIT 1`, platformModelID).Scan(&deadlineTaskID); err != nil {
		t.Fatalf("select deadline waiter: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET enqueued_at = now() - interval '2 seconds',
    wait_deadline_at = now() - interval '1 second'
WHERE task_id = $1::uuid`, deadlineTaskID); err != nil {
		t.Fatalf("expire queue deadline: %v", err)
	}
	reaped, err := second.ReapExpiredTaskAdmissions(ctx, 10)
	if err != nil || reaped.ExpiredDeadlines < 1 {
		t.Fatalf("deadline reap = %+v, err=%v", reaped, err)
	}
	var taskStatus, taskErrorCode string
	if err := first.pool.QueryRow(ctx, `
SELECT status, COALESCE(error_code, '')
FROM gateway_tasks
WHERE id = $1::uuid`, deadlineTaskID).Scan(&taskStatus, &taskErrorCode); err != nil {
		t.Fatalf("read deadline task: %v", err)
	}
	if taskStatus != "failed" || taskErrorCode != "queue_timeout" {
		t.Fatalf("deadline task status=%s code=%s, want failed/queue_timeout", taskStatus, taskErrorCode)
	}

	for _, taskID := range taskIDs {
		_ = first.DeleteTaskAdmission(ctx, taskID)
	}
	leaseHolder := createTask(false)
	if result, err := first.TryTaskAdmission(ctx, inputFor(leaseHolder, 100, "lease-holder")); err != nil || !result.Admitted {
		t.Fatalf("admit lease holder: result=%+v err=%v", result, err)
	}
	disconnectedWaiter := createTask(false)
	if result, err := second.TryTaskAdmission(ctx, inputFor(disconnectedWaiter, 100, "disconnected-waiter")); err != nil || result.Admitted {
		t.Fatalf("queue disconnected waiter: result=%+v err=%v", result, err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET waiter_lease_expires_at = now() - interval '1 second'
WHERE task_id = $1::uuid`, disconnectedWaiter.ID); err != nil {
		t.Fatalf("expire waiter lease: %v", err)
	}
	reaped, err = second.ReapExpiredTaskAdmissions(ctx, 10)
	if err != nil || reaped.ExpiredWaiters < 1 {
		t.Fatalf("waiter reap = %+v, err=%v", reaped, err)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT status, COALESCE(error_code, '')
FROM gateway_tasks
WHERE id = $1::uuid`, disconnectedWaiter.ID).Scan(&taskStatus, &taskErrorCode); err != nil {
		t.Fatalf("read disconnected task: %v", err)
	}
	if taskStatus != "cancelled" || taskErrorCode != "client_disconnected" {
		t.Fatalf("disconnected task status=%s code=%s, want cancelled/client_disconnected", taskStatus, taskErrorCode)
	}
	if err := first.DeleteTaskAdmission(ctx, leaseHolder.ID); err != nil {
		t.Fatalf("release lease holder admission: %v", err)
	}

	preSubmission := createTask(false)
	claimed, err := first.ClaimTaskPreparation(ctx, preSubmission.ID, uuid.NewString(), time.Minute)
	if err != nil {
		t.Fatalf("claim pre-submission task: %v", err)
	}
	result, err = first.TryTaskAdmission(ctx, inputFor(claimed, 100, "pre-submission"))
	if err != nil || !result.Admitted {
		t.Fatalf("admit pre-submission task: result=%+v err=%v", result, err)
	}
	var attemptID string
	if err := first.pool.QueryRow(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
  upstream_submission_status
)
VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'not_submitted')
RETURNING id::text`,
		preSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
	).Scan(&attemptID); err != nil {
		t.Fatalf("create pre-submission attempt: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_param_preprocessing_logs (
  task_id, attempt_id, attempt_no, model_type
)
VALUES ($1::uuid, $2::uuid, 1, 'image_generate')`, preSubmission.ID, attemptID); err != nil {
		t.Fatalf("create pre-submission log: %v", err)
	}
	cancelled, changed, err := first.CancelTaskBeforeUpstreamSubmission(
		ctx,
		preSubmission.ID,
		claimed.ExecutionToken,
		"client disconnected before upstream submission",
	)
	if err != nil || !changed || cancelled.Status != "cancelled" || cancelled.ErrorCode != "client_disconnected" {
		t.Fatalf("cancel pre-submission task = %+v changed=%v err=%v", cancelled, changed, err)
	}
	var remainingAdmissions, remainingLeases, remainingAttempts, remainingPreprocessing int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL),
  (SELECT COUNT(*) FROM gateway_task_attempts WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_task_param_preprocessing_logs WHERE task_id = $1::uuid)`,
		preSubmission.ID,
	).Scan(&remainingAdmissions, &remainingLeases, &remainingAttempts, &remainingPreprocessing); err != nil {
		t.Fatalf("read pre-submission cleanup: %v", err)
	}
	if remainingAdmissions != 0 || remainingLeases != 0 || remainingAttempts != 0 || remainingPreprocessing != 0 {
		t.Fatalf(
			"pre-submission cleanup admissions=%d leases=%d attempts=%d preprocessing=%d",
			remainingAdmissions,
			remainingLeases,
			remainingAttempts,
			remainingPreprocessing,
		)
	}

	postSubmission := createTask(false)
	claimed, err = first.ClaimTaskPreparation(ctx, postSubmission.ID, uuid.NewString(), time.Minute)
	if err != nil {
		t.Fatalf("claim post-submission task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
  upstream_submission_status
)
VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'submitting')`,
		postSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
	); err != nil {
		t.Fatalf("create submitted attempt: %v", err)
	}
	_, changed, err = first.CancelTaskBeforeUpstreamSubmission(
		ctx,
		postSubmission.ID,
		claimed.ExecutionToken,
		"client disconnected after upstream submission",
	)
	if err != nil || changed {
		t.Fatalf("post-submission cancellation changed=%v err=%v, want unchanged", changed, err)
	}

	queuedAtomicTask := createTask(true)
	queuedHookFailure := errors.New("synthetic queued hook failure")
	_, err = first.QueueTaskAdmissionWithHook(ctx, inputFor(queuedAtomicTask, 100, ""), func(pgx.Tx) error {
		return queuedHookFailure
	})
	if !errors.Is(err, queuedHookFailure) {
		t.Fatalf("queued hook failure = %v, want synthetic failure", err)
	}
	var queuedAdmissions, queuedLeases int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		queuedAtomicTask.ID,
	).Scan(&queuedAdmissions, &queuedLeases); err != nil {
		t.Fatalf("read rolled back queued hook state: %v", err)
	}
	if queuedAdmissions != 0 || queuedLeases != 0 {
		t.Fatalf("failed queued hook left admissions=%d leases=%d", queuedAdmissions, queuedLeases)
	}
	const queuedSyntheticRiverJobID int64 = 987654320
	queuedAdmission, err := first.QueueTaskAdmissionWithHook(ctx, inputFor(queuedAtomicTask, 100, ""), func(tx pgx.Tx) error {
		_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2
WHERE id = $1::uuid`, queuedAtomicTask.ID, queuedSyntheticRiverJobID)
		return hookErr
	})
	if err != nil || queuedAdmission.Status != "waiting" {
		t.Fatalf("atomic queued admission=%+v err=%v, want waiting", queuedAdmission, err)
	}
	if len(queuedAdmission.Scopes) == 0 {
		t.Fatal("queued admission did not persist its scope snapshot")
	}
	waitingAdmissionSnapshots, err := first.ListWaitingAsyncAdmissions(ctx, 1000)
	if err != nil {
		t.Fatalf("list waiting async admission snapshots: %v", err)
	}
	var listedSnapshot *TaskAdmission
	for index := range waitingAdmissionSnapshots {
		if waitingAdmissionSnapshots[index].TaskID == queuedAtomicTask.ID {
			listedSnapshot = &waitingAdmissionSnapshots[index]
			break
		}
	}
	if listedSnapshot == nil || len(listedSnapshot.Scopes) != len(queuedAdmission.Scopes) {
		t.Fatalf("listed admission snapshot=%+v, want %d scopes", listedSnapshot, len(queuedAdmission.Scopes))
	}
	_, markedForReselect, err := first.RequestTaskAdmissionReselect(ctx, queuedAtomicTask.ID, platformModelID)
	if err != nil {
		t.Fatalf("request waiting admission reselection: %v", err)
	}
	if !markedForReselect {
		t.Fatal("queued task was not marked for reselection")
	}
	var platformReselectCount int
	if err := first.pool.QueryRow(ctx, `
SELECT count(*)
FROM gateway_task_admissions
WHERE platform_model_id = $1::uuid
  AND reselect_requested_at IS NOT NULL`, platformModelID).Scan(&platformReselectCount); err != nil {
		t.Fatalf("count task-scoped reselections: %v", err)
	}
	if platformReselectCount != 1 {
		t.Fatalf("platform reselection fanout=%d, want exactly one task", platformReselectCount)
	}
	reselectAdmission, err := first.GetTaskAdmission(ctx, queuedAtomicTask.ID)
	if err != nil {
		t.Fatalf("read admission reselection marker: %v", err)
	}
	if reselectAdmission.ReselectRequestedAt.IsZero() {
		t.Fatal("queued admission was not marked for candidate reselection")
	}
	var riverJobID int64
	if err := first.pool.QueryRow(ctx, `
SELECT
  river_job_id,
  (SELECT COUNT(*)
   FROM gateway_concurrency_leases
   WHERE task_id = $1::uuid AND released_at IS NULL)
FROM gateway_tasks
WHERE id = $1::uuid`, queuedAtomicTask.ID).Scan(&riverJobID, &queuedLeases); err != nil {
		t.Fatalf("read atomic queued River marker: %v", err)
	}
	if riverJobID != queuedSyntheticRiverJobID || queuedLeases != 0 {
		t.Fatalf("queued River marker=%d leases=%d, want %d/0", riverJobID, queuedLeases, queuedSyntheticRiverJobID)
	}
	waitingAsyncTaskIDs, err := first.ListWaitingAsyncAdmissionTaskIDs(ctx, 1000)
	if err != nil {
		t.Fatalf("list waiting async tasks with missing River job: %v", err)
	}
	waitingAsyncFound := false
	for _, taskID := range waitingAsyncTaskIDs {
		if taskID == queuedAtomicTask.ID {
			waitingAsyncFound = true
			break
		}
	}
	if !waitingAsyncFound {
		t.Fatal("waiting async task with missing River job was not dispatchable")
	}
	recoverableTasks, err := first.ListRecoverableAsyncTasks(ctx, 1000)
	if err != nil {
		t.Fatalf("list recoverable tasks while admission is waiting: %v", err)
	}
	for _, item := range recoverableTasks {
		if item.ID == queuedAtomicTask.ID {
			t.Fatal("generic River recovery claimed a task owned by the async admission dispatcher")
		}
	}
	recoveryGraceTask := createTask(true)
	recoverableTasks, err = first.ListRecoverableAsyncTasks(ctx, 1000)
	if err != nil {
		t.Fatalf("list fresh recoverable tasks: %v", err)
	}
	for _, item := range recoverableTasks {
		if item.ID == recoveryGraceTask.ID {
			t.Fatal("generic River recovery claimed a task still preparing admission")
		}
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET created_at = now() - interval '6 minutes'
WHERE id = $1::uuid`, recoveryGraceTask.ID); err != nil {
		t.Fatalf("age task beyond River recovery grace period: %v", err)
	}
	recoverableTasks, err = first.ListRecoverableAsyncTasks(ctx, 1000)
	if err != nil {
		t.Fatalf("list aged recoverable tasks: %v", err)
	}
	recoveryGraceFound := false
	for _, item := range recoverableTasks {
		if item.ID == recoveryGraceTask.ID {
			recoveryGraceFound = true
			break
		}
	}
	if !recoveryGraceFound {
		t.Fatal("generic River recovery did not claim a task beyond the recovery grace period")
	}

	batchScope := AdmissionScope{
		ScopeType:       "worker_capacity",
		ScopeKey:        "batch-" + suffix,
		ScopeName:       "batch capacity",
		ConcurrentLimit: 48,
		Amount:          1,
		LeaseTTLSeconds: 120,
		QueueLimit:      100,
		MaxWaitSeconds:  600,
	}
	batchTasks := make([]GatewayTask, 48)
	for index := range batchTasks {
		batchTasks[index] = createTask(true)
	}
	batchInputs := make([]TaskAdmissionInput, 0, len(batchTasks))
	batchJobIDs := make(map[string]int64, len(batchTasks))
	for index, task := range batchTasks {
		input := inputFor(task, 100+index, "")
		input.Scopes = []AdmissionScope{batchScope}
		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
			t.Fatalf("queue batch task %d: %v", index, err)
		}
		batchInputs = append(batchInputs, input)
		batchJobIDs[task.ID] = int64(987654400 + index)
	}
	batchOutcomes, err := first.TryTaskAdmissionBatchWithAdmittedHook(
		ctx,
		batchInputs,
		func(tx pgx.Tx, input TaskAdmissionInput) error {
			_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2
WHERE id = $1::uuid`, input.TaskID, batchJobIDs[input.TaskID])
			return hookErr
		},
	)
	if err != nil || len(batchOutcomes) != len(batchTasks) {
		t.Fatalf("batch admission outcomes=%+v err=%v", batchOutcomes, err)
	}
	for index, outcome := range batchOutcomes {
		if outcome.Err != nil || !outcome.Result.Admitted || !outcome.Result.NewlyAdmitted ||
			len(outcome.Result.Leases) != 1 {
			t.Fatalf("batch outcome %d=%+v", index, outcome)
		}
	}
	batchTaskIDs := make([]string, 0, len(batchTasks))
	for _, task := range batchTasks {
		batchTaskIDs = append(batchTaskIDs, task.ID)
	}
	var batchAdmissions, batchLeases, batchRiverJobs int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*)
   FROM gateway_task_admissions
   WHERE task_id = ANY($1::uuid[]) AND status = 'admitted'),
  (SELECT count(*)
   FROM gateway_concurrency_leases
   WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL AND expires_at > now()),
  (SELECT count(*)
   FROM gateway_tasks
   WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
		batchTaskIDs,
	).Scan(&batchAdmissions, &batchLeases, &batchRiverJobs); err != nil {
		t.Fatalf("read committed batch admission: %v", err)
	}
	if batchAdmissions != len(batchTasks) || batchLeases != len(batchTasks) || batchRiverJobs != len(batchTasks) {
		t.Fatalf(
			"batch admissions=%d leases=%d River jobs=%d, want %d/%d/%d",
			batchAdmissions,
			batchLeases,
			batchRiverJobs,
			len(batchTasks),
			len(batchTasks),
			len(batchTasks),
		)
	}
	for _, task := range batchTasks {
		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
			t.Fatalf("release batch task %s: %v", task.ID, err)
		}
	}

	rollbackScope := batchScope
	rollbackScope.ScopeKey = "batch-rollback-" + suffix
	rollbackScope.ConcurrentLimit = 2
	rollbackTasks := []GatewayTask{createTask(true), createTask(true)}
	rollbackInputs := make([]TaskAdmissionInput, 0, len(rollbackTasks))
	for index, task := range rollbackTasks {
		input := inputFor(task, 200+index, "")
		input.Scopes = []AdmissionScope{rollbackScope}
		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
			t.Fatalf("queue rollback batch task %d: %v", index, err)
		}
		rollbackInputs = append(rollbackInputs, input)
	}
	batchHookFailure := errors.New("synthetic batch hook failure")
	partialOutcomes, err := first.TryTaskAdmissionBatchWithAdmittedHook(
		ctx,
		rollbackInputs,
		func(tx pgx.Tx, input TaskAdmissionInput) error {
			if input.TaskID == rollbackTasks[1].ID {
				return batchHookFailure
			}
			_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = 987654499
WHERE id = $1::uuid`, input.TaskID)
			return hookErr
		},
	)
	if !errors.Is(err, batchHookFailure) {
		t.Fatalf("batch hook failure=%v, want synthetic failure", err)
	}
	if len(partialOutcomes) != 2 || !partialOutcomes[0].Result.Admitted ||
		!errors.Is(partialOutcomes[1].Err, batchHookFailure) {
		t.Fatalf("per-task batch outcomes=%+v, want first committed and second failed", partialOutcomes)
	}
	var rollbackWaiting, rollbackLeases, rollbackRiverJobs int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*)
   FROM gateway_task_admissions
   WHERE task_id = ANY($1::uuid[]) AND status = 'waiting'),
  (SELECT count(*)
   FROM gateway_concurrency_leases
   WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL),
  (SELECT count(*)
   FROM gateway_tasks
   WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
		[]string{rollbackTasks[0].ID, rollbackTasks[1].ID},
	).Scan(&rollbackWaiting, &rollbackLeases, &rollbackRiverJobs); err != nil {
		t.Fatalf("read rolled back batch admission: %v", err)
	}
	if rollbackWaiting != 1 || rollbackLeases != 1 || rollbackRiverJobs != 1 {
		t.Fatalf(
			"per-task rollback waiting=%d leases=%d River jobs=%d, want 1/1/1",
			rollbackWaiting,
			rollbackLeases,
			rollbackRiverJobs,
		)
	}
	for _, task := range rollbackTasks {
		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
			t.Fatalf("release rollback batch task %s: %v", task.ID, err)
		}
	}

	atomicScope := batchScope
	atomicScope.ScopeKey = "atomic-microbatch-" + suffix
	atomicScope.ConcurrentLimit = 4
	atomicTasks := []GatewayTask{createTask(true), createTask(true), createTask(true), createTask(true)}
	atomicInputs := make([]TaskAdmissionInput, 0, len(atomicTasks))
	for index, task := range atomicTasks {
		input := inputFor(task, 300+index, "")
		input.Scopes = []AdmissionScope{atomicScope}
		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
			t.Fatalf("queue atomic microbatch task %d: %v", index, err)
		}
		atomicInputs = append(atomicInputs, input)
	}
	atomicOutcomes, err := first.TryTaskAdmissionAtomicBatchWithAdmittedHook(
		ctx,
		atomicInputs,
		func(tx pgx.Tx, input TaskAdmissionInput) error {
			_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = 987654600
WHERE id = $1::uuid`, input.TaskID)
			return hookErr
		},
	)
	if err != nil || len(atomicOutcomes) != len(atomicTasks) {
		t.Fatalf("atomic microbatch outcomes=%+v err=%v", atomicOutcomes, err)
	}
	var microbatchAdmissions, microbatchLeases, microbatchRiverJobs int
	atomicTaskIDs := []string{atomicTasks[0].ID, atomicTasks[1].ID, atomicTasks[2].ID, atomicTasks[3].ID}
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*) FROM gateway_task_admissions WHERE task_id = ANY($1::uuid[]) AND status = 'admitted'),
  (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL),
  (SELECT count(*) FROM gateway_tasks WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
		atomicTaskIDs,
	).Scan(&microbatchAdmissions, &microbatchLeases, &microbatchRiverJobs); err != nil {
		t.Fatalf("read atomic microbatch: %v", err)
	}
	if microbatchAdmissions != 4 || microbatchLeases != 4 || microbatchRiverJobs != 4 {
		t.Fatalf("atomic microbatch admissions=%d leases=%d jobs=%d, want 4/4/4",
			microbatchAdmissions, microbatchLeases, microbatchRiverJobs)
	}
	for _, task := range atomicTasks {
		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
			t.Fatalf("release atomic microbatch task %s: %v", task.ID, err)
		}
	}

	atomicRollbackScope := batchScope
	atomicRollbackScope.ScopeKey = "atomic-rollback-" + suffix
	atomicRollbackScope.ConcurrentLimit = 2
	atomicRollbackTasks := []GatewayTask{createTask(true), createTask(true)}
	atomicRollbackInputs := make([]TaskAdmissionInput, 0, len(atomicRollbackTasks))
	for index, task := range atomicRollbackTasks {
		input := inputFor(task, 400+index, "")
		input.Scopes = []AdmissionScope{atomicRollbackScope}
		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
			t.Fatalf("queue atomic rollback task %d: %v", index, err)
		}
		atomicRollbackInputs = append(atomicRollbackInputs, input)
	}
	_, err = first.TryTaskAdmissionAtomicBatchWithAdmittedHook(
		ctx,
		atomicRollbackInputs,
		func(_ pgx.Tx, input TaskAdmissionInput) error {
			if input.TaskID == atomicRollbackTasks[1].ID {
				return batchHookFailure
			}
			return nil
		},
	)
	if !errors.Is(err, batchHookFailure) {
		t.Fatalf("atomic rollback error=%v, want synthetic failure", err)
	}
	var atomicRollbackWaiting, atomicRollbackLeases int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*) FROM gateway_task_admissions WHERE task_id = ANY($1::uuid[]) AND status = 'waiting'),
  (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL)`,
		[]string{atomicRollbackTasks[0].ID, atomicRollbackTasks[1].ID},
	).Scan(&atomicRollbackWaiting, &atomicRollbackLeases); err != nil {
		t.Fatalf("read atomic rollback: %v", err)
	}
	if atomicRollbackWaiting != 2 || atomicRollbackLeases != 0 {
		t.Fatalf("atomic rollback waiting=%d leases=%d, want 2/0",
			atomicRollbackWaiting, atomicRollbackLeases)
	}
	for _, task := range atomicRollbackTasks {
		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
			t.Fatalf("release atomic rollback task %s: %v", task.ID, err)
		}
	}

	result, err = second.TryTaskAdmission(ctx, inputFor(queuedAtomicTask, 100, ""))
	if err != nil || !result.Admitted || len(result.Leases) != 1 {
		t.Fatalf("worker-time queued admission result=%+v err=%v", result, err)
	}
	if err := first.DeleteTaskAdmission(ctx, queuedAtomicTask.ID); err != nil {
		t.Fatalf("release queued atomic task: %v", err)
	}

	atomicTask := createTask(true)
	hookFailure := errors.New("synthetic admitted hook failure")
	_, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(pgx.Tx) error {
		return hookFailure
	})
	if !errors.Is(err, hookFailure) {
		t.Fatalf("admitted hook failure = %v, want synthetic failure", err)
	}
	var atomicAdmissions, atomicLeases int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		atomicTask.ID,
	).Scan(&atomicAdmissions, &atomicLeases); err != nil {
		t.Fatalf("read rolled back admitted hook state: %v", err)
	}
	if atomicAdmissions != 0 || atomicLeases != 0 {
		t.Fatalf("failed admitted hook left admissions=%d leases=%d", atomicAdmissions, atomicLeases)
	}
	const syntheticRiverJobID int64 = 987654321
	result, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(tx pgx.Tx) error {
		_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2
WHERE id = $1::uuid`, atomicTask.ID, syntheticRiverJobID)
		return hookErr
	})
	if err != nil || !result.Admitted {
		t.Fatalf("atomic admitted hook result=%+v err=%v", result, err)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT river_job_id
FROM gateway_tasks
WHERE id = $1::uuid`, atomicTask.ID).Scan(&riverJobID); err != nil {
		t.Fatalf("read atomic River job marker: %v", err)
	}
	if riverJobID != syntheticRiverJobID {
		t.Fatalf("atomic River job marker=%d, want %d", riverJobID, syntheticRiverJobID)
	}
	if err := first.DeleteTaskAdmission(ctx, atomicTask.ID); err != nil {
		t.Fatalf("release atomic admitted hook task: %v", err)
	}

	recoveryTask := createTask(true)
	result, err = first.TryTaskAdmission(ctx, inputFor(recoveryTask, 100, ""))
	if err != nil || !result.Admitted {
		t.Fatalf("admit interrupted recovery task: result=%+v err=%v", result, err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    execution_token = $2::uuid,
    execution_lease_expires_at = now() - interval '1 second'
WHERE id = $1::uuid`, recoveryTask.ID, uuid.NewString()); err != nil {
		t.Fatalf("expire interrupted recovery task lease: %v", err)
	}
	recovery, err := first.RecoverInterruptedRuntimeState(ctx)
	if err != nil {
		t.Fatalf("recover interrupted admitted task: %v", err)
	}
	if recovery.RequeuedAsyncTasks < 1 || recovery.CleanedTaskAdmissions < 1 {
		t.Fatalf("interrupted task recovery = %+v, want requeue and admission cleanup", recovery)
	}
	var recoveredStatus string
	var recoveredExecutionToken *string
	if err := first.pool.QueryRow(ctx, `
SELECT status, execution_token::text
FROM gateway_tasks
WHERE id = $1::uuid`, recoveryTask.ID).Scan(&recoveredStatus, &recoveredExecutionToken); err != nil {
		t.Fatalf("read recovered task state: %v", err)
	}
	if recoveredStatus != "queued" || recoveredExecutionToken != nil {
		t.Fatalf("recovered task status=%s execution_token=%v, want queued without token", recoveredStatus, recoveredExecutionToken)
	}
	var recoveredAdmissions, recoveredLeases int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		recoveryTask.ID,
	).Scan(&recoveredAdmissions, &recoveredLeases); err != nil {
		t.Fatalf("read recovered task cleanup: %v", err)
	}
	if recoveredAdmissions != 0 || recoveredLeases != 0 {
		t.Fatalf("recovered task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases)
	}

	remoteRecoveryTask := createTask(true)
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    remote_task_id = 'remote-video-task',
    execution_token = gen_random_uuid(),
    execution_lease_expires_at = now() - interval '1 second'
WHERE id = $1::uuid`,
		remoteRecoveryTask.ID,
	); err != nil {
		t.Fatalf("prepare persisted remote task recovery task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, queue_key, status, upstream_submission_status
)
VALUES ($1::uuid, 1, 'integration-test', 'running', 'submitting')`,
		remoteRecoveryTask.ID,
	); err != nil {
		t.Fatalf("prepare persisted remote task recovery: %v", err)
	}
	recovery, err = first.RecoverInterruptedRuntimeState(ctx)
	if err != nil {
		t.Fatalf("recover persisted remote task: %v", err)
	}
	var remoteStatus, remoteErrorCode string
	var remoteAttemptRetryable bool
	if err := first.pool.QueryRow(ctx, `
SELECT task.status, COALESCE(task.error_code, ''), attempt.retryable
FROM gateway_tasks task
JOIN gateway_task_attempts attempt ON attempt.task_id = task.id
WHERE task.id = $1::uuid`, remoteRecoveryTask.ID).Scan(
		&remoteStatus,
		&remoteErrorCode,
		&remoteAttemptRetryable,
	); err != nil {
		t.Fatalf("read persisted remote task recovery: %v", err)
	}
	if remoteStatus != "queued" || remoteErrorCode != "" || !remoteAttemptRetryable {
		t.Fatalf(
			"persisted remote task recovery status=%s code=%s retryable=%t, want queued resume",
			remoteStatus,
			remoteErrorCode,
			remoteAttemptRetryable,
		)
	}

	lockedRecoveryTask := createTask(true)
	unlockedRecoveryTask := createTask(true)
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    execution_token = gen_random_uuid(),
    execution_lease_expires_at = now() - interval '1 second'
WHERE id = ANY($1::uuid[])`, []string{lockedRecoveryTask.ID, unlockedRecoveryTask.ID}); err != nil {
		t.Fatalf("prepare locked recovery tasks: %v", err)
	}
	lockTx, err := second.pool.Begin(ctx)
	if err != nil {
		t.Fatalf("begin task row lock: %v", err)
	}
	var lockedTaskID string
	if err := lockTx.QueryRow(ctx, `
SELECT id::text
FROM gateway_tasks
WHERE id = $1::uuid
FOR UPDATE`, lockedRecoveryTask.ID).Scan(&lockedTaskID); err != nil {
		rollbackTransaction(lockTx)
		t.Fatalf("lock interrupted task row: %v", err)
	}
	recoveryCtx, recoveryCancel := context.WithTimeout(ctx, 3*time.Second)
	recovery, err = first.RecoverInterruptedRuntimeState(recoveryCtx)
	recoveryCancel()
	if err != nil {
		rollbackTransaction(lockTx)
		t.Fatalf("recover unlocked task while another row is locked: %v", err)
	}
	if recovery.RequeuedAsyncTasks < 1 {
		rollbackTransaction(lockTx)
		t.Fatalf("SKIP LOCKED recovery = %+v, want at least one requeued task", recovery)
	}
	var lockedStatus, unlockedStatus string
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT status FROM gateway_tasks WHERE id = $1::uuid),
  (SELECT status FROM gateway_tasks WHERE id = $2::uuid)`,
		lockedRecoveryTask.ID,
		unlockedRecoveryTask.ID,
	).Scan(&lockedStatus, &unlockedStatus); err != nil {
		rollbackTransaction(lockTx)
		t.Fatalf("read SKIP LOCKED recovery states: %v", err)
	}
	if lockedStatus != "running" || unlockedStatus != "queued" {
		rollbackTransaction(lockTx)
		t.Fatalf("SKIP LOCKED recovery statuses = %s/%s, want running/queued", lockedStatus, unlockedStatus)
	}
	rollbackTransaction(lockTx)
	recovery, err = first.RecoverInterruptedRuntimeState(ctx)
	if err != nil {
		t.Fatalf("recover previously locked task: %v", err)
	}
	if recovery.RequeuedAsyncTasks < 1 {
		t.Fatalf("unlocked recovery = %+v, want requeued task", recovery)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT status
FROM gateway_tasks
WHERE id = $1::uuid`, lockedRecoveryTask.ID).Scan(&lockedStatus); err != nil {
		t.Fatalf("read unlocked recovery status: %v", err)
	}
	if lockedStatus != "queued" {
		t.Fatalf("unlocked recovery status = %s, want queued", lockedStatus)
	}

	preSubmissionTask := createTask(true)
	preSubmissionToken := uuid.NewString()
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    execution_token = $2::uuid,
    execution_lease_expires_at = now() + interval '5 minutes'
WHERE id = $1::uuid`,
		preSubmissionTask.ID,
		preSubmissionToken,
	); err != nil {
		t.Fatalf("prepare pre-submission lock timeout task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, queue_key, status, upstream_submission_status
)
VALUES ($1::uuid, 1, 'integration-test', 'running', 'not_submitted')`,
		preSubmissionTask.ID,
	); err != nil {
		t.Fatalf("create pre-submission attempt: %v", err)
	}
	requeuedTask, changed, err := first.RequeueTaskBeforeUpstreamSubmission(
		ctx,
		preSubmissionTask.ID,
		preSubmissionToken,
		time.Second,
	)
	if err != nil || !changed || requeuedTask.Status != "queued" {
		t.Fatalf("pre-submission lock timeout requeue task=%+v changed=%v err=%v", requeuedTask, changed, err)
	}
	var preSubmissionAttempts int
	if err := first.pool.QueryRow(ctx, `
SELECT count(*)
FROM gateway_task_attempts
WHERE task_id = $1::uuid`, preSubmissionTask.ID).Scan(&preSubmissionAttempts); err != nil {
		t.Fatalf("count cleaned pre-submission attempts: %v", err)
	}
	if preSubmissionAttempts != 0 {
		t.Fatalf("pre-submission lock timeout left %d attempts, want 0", preSubmissionAttempts)
	}

	ambiguousTask := createTask(true)
	ambiguousToken := uuid.NewString()
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    execution_token = $2::uuid,
    execution_lease_expires_at = now() + interval '5 minutes'
WHERE id = $1::uuid`,
		ambiguousTask.ID,
		ambiguousToken,
	); err != nil {
		t.Fatalf("prepare ambiguous submission task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, queue_key, status, upstream_submission_status
)
VALUES ($1::uuid, 1, 'integration-test', 'running', 'submitting')`,
		ambiguousTask.ID,
	); err != nil {
		t.Fatalf("create ambiguous submission attempt: %v", err)
	}
	_, changed, err = first.RequeueTaskBeforeUpstreamSubmission(
		ctx,
		ambiguousTask.ID,
		ambiguousToken,
		time.Second,
	)
	if err != nil {
		t.Fatalf("guard ambiguous submission task requeue: %v", err)
	}
	if changed {
		t.Fatal("ambiguous submission task was requeued")
	}
	var ambiguousStatus string
	if err := first.pool.QueryRow(ctx, `
SELECT status
FROM gateway_tasks
WHERE id = $1::uuid`, ambiguousTask.ID).Scan(&ambiguousStatus); err != nil {
		t.Fatalf("read ambiguous submission task: %v", err)
	}
	if ambiguousStatus != "running" {
		t.Fatalf("ambiguous submission task status = %s, want running", ambiguousStatus)
	}

	globalHolder := createTask(true)
	globalWaiter := createTask(true)
	globalScope := AdmissionScope{
		ScopeType:       "worker_capacity",
		ScopeKey:        "global-blocker-" + suffix,
		ScopeName:       "global blocker",
		ConcurrentLimit: 1,
		Amount:          1,
		LeaseTTLSeconds: 120,
		QueueLimit:      10,
		MaxWaitSeconds:  600,
	}
	globalHolderInput := inputFor(globalHolder, 100, "")
	globalHolderInput.Scopes = []AdmissionScope{globalScope}
	globalWaiterInput := inputFor(globalWaiter, 100, "")
	globalWaiterInput.Scopes = []AdmissionScope{globalScope}
	if result, err := first.TryTaskAdmission(ctx, globalHolderInput); err != nil || !result.Admitted {
		t.Fatalf("admit global capacity holder: result=%+v err=%v", result, err)
	}
	if _, err := first.QueueTaskAdmissionWithHook(ctx, globalWaiterInput, nil); err != nil {
		t.Fatalf("queue global capacity waiter: %v", err)
	}
	globalWaitResult, err := second.TryTaskAdmission(ctx, globalWaiterInput)
	if err != nil || globalWaitResult.Admitted || len(globalWaitResult.Blockers) != 1 ||
		globalWaitResult.Blockers[0].Reason != "saturated" ||
		globalWaitResult.Blockers[0].ScopeType != "worker_capacity" {
		t.Fatalf("global capacity waiter result=%+v err=%v", globalWaitResult, err)
	}
	if !globalWaitResult.Admission.ReselectRequestedAt.IsZero() {
		t.Fatal("global capacity wait unexpectedly requested candidate reselection")
	}
	if err := first.DeleteTaskAdmission(ctx, globalHolder.ID); err != nil {
		t.Fatalf("release global capacity holder: %v", err)
	}
	if err := first.DeleteTaskAdmission(ctx, globalWaiter.ID); err != nil {
		t.Fatalf("delete global capacity waiter: %v", err)
	}

	orphanedAdmittedTask := createTask(true)
	orphanedAdmittedInput := inputFor(orphanedAdmittedTask, 100, "")
	orphanedAdmittedInput.Scopes = []AdmissionScope{{
		ScopeType:       "worker_capacity",
		ScopeKey:        "orphaned-admitted-" + suffix,
		ConcurrentLimit: 1,
		Amount:          1,
		LeaseTTLSeconds: 120,
		QueueLimit:      10,
		MaxWaitSeconds:  600,
	}}
	if result, err := first.TryTaskAdmission(ctx, orphanedAdmittedInput); err != nil || !result.Admitted {
		t.Fatalf("admit orphan candidate: result=%+v err=%v", result, err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = $1::uuid
  AND released_at IS NULL`, orphanedAdmittedTask.ID); err != nil {
		t.Fatalf("orphan admitted task leases: %v", err)
	}
	if result, err := second.TryTaskAdmission(ctx, orphanedAdmittedInput); err != nil || !result.Admitted || !result.NewlyAdmitted {
		t.Fatalf("recover admitted row without lease: result=%+v err=%v", result, err)
	}
	if err := first.DeleteTaskAdmission(ctx, orphanedAdmittedTask.ID); err != nil {
		t.Fatalf("delete recovered admitted row: %v", err)
	}

	resilienceTask := createTask(true)
	resilienceInput := inputFor(resilienceTask, 100, "")
	resilienceInput.Scopes = []AdmissionScope{{
		ScopeType:       "worker_capacity",
		ScopeKey:        "resilience-" + suffix,
		ScopeName:       "resilience capacity",
		ConcurrentLimit: 5,
		Amount:          1,
		LeaseTTLSeconds: 120,
		QueueLimit:      10,
		MaxWaitSeconds:  600,
	}}
	if _, err := first.QueueTaskAdmissionWithHook(ctx, resilienceInput, nil); err != nil {
		t.Fatalf("queue resilience task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET scope_snapshot = '[{"ScopeType":123}]'::jsonb
WHERE task_id = $1::uuid`, resilienceTask.ID); err != nil {
		t.Fatalf("corrupt resilience scope snapshot: %v", err)
	}
	listedAdmissions, err := first.ListWaitingAsyncAdmissions(ctx, 1000)
	if err != nil {
		t.Fatalf("list admissions with one corrupt snapshot: %v", err)
	}
	corruptFound := false
	for _, listed := range listedAdmissions {
		if listed.TaskID == resilienceTask.ID {
			corruptFound = listed.SnapshotInvalid
			break
		}
	}
	if !corruptFound {
		t.Fatal("corrupt admission snapshot was not isolated for rebuilding")
	}
	repairedAdmission, err := first.RebindWaitingTaskAdmission(ctx, resilienceInput)
	if err != nil || repairedAdmission.SnapshotInvalid || len(repairedAdmission.Scopes) != 1 {
		t.Fatalf("repair corrupt admission snapshot=%+v err=%v", repairedAdmission, err)
	}
	deferred, changed, err := first.DeferWaitingTaskAdmission(
		ctx,
		resilienceTask.ID,
		platformModelID,
		time.Now().Add(time.Minute),
		"deferred",
		"snapshot_invalid",
	)
	if err != nil || !changed || deferred.DispatchFailureCount != 1 || deferred.DispatchLastErrorCode != "snapshot_invalid" {
		t.Fatalf("defer corrupt admission=%+v changed=%v err=%v", deferred, changed, err)
	}
	listedAdmissions, err = first.ListWaitingAsyncAdmissions(ctx, 1000)
	if err != nil {
		t.Fatalf("list admissions after defer: %v", err)
	}
	for _, listed := range listedAdmissions {
		if listed.TaskID == resilienceTask.ID {
			t.Fatal("deferred admission remained dispatchable")
		}
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET dispatch_next_at = now(), scope_snapshot = '[]'::jsonb
WHERE task_id = $1::uuid`, resilienceTask.ID); err != nil {
		t.Fatalf("make resilience admission dispatchable: %v", err)
	}
	reselected, changed, err := first.RequestTaskAdmissionReselect(ctx, resilienceTask.ID, platformModelID)
	if err != nil || !changed || reselected.ReselectRequestedAt.IsZero() {
		t.Fatalf("task-scoped reselection=%+v changed=%v err=%v", reselected, changed, err)
	}
	if _, changed, err := second.RequestTaskAdmissionReselect(ctx, resilienceTask.ID, platformModelID); err != nil || changed {
		t.Fatalf("duplicate task-scoped reselection changed=%v err=%v", changed, err)
	}
	if _, err := first.FailQueuedTaskWithCallback(
		ctx,
		resilienceTask.ID,
		"no_model_candidate",
		"no enabled platform model matches request",
		"https://callback.invalid/admission-resilience",
		false,
	); err != nil {
		t.Fatalf("terminally fail resilience task: %v", err)
	}
	if _, err := second.FailQueuedTaskWithCallback(
		ctx,
		resilienceTask.ID,
		"no_model_candidate",
		"duplicate terminalization",
		"https://callback.invalid/admission-resilience",
		false,
	); !errors.Is(err, ErrTaskExecutionFinished) {
		t.Fatalf("duplicate terminalization error=%v, want task finished", err)
	}
	var resilienceStatus, resilienceCode string
	var resilienceAdmissions, resilienceLeases, resilienceFailedEvents, resilienceCallbacks int
	if err := first.pool.QueryRow(ctx, `
SELECT task.status, COALESCE(task.error_code, ''),
       (SELECT count(*) FROM gateway_task_admissions admission WHERE admission.task_id = task.id),
       (SELECT count(*) FROM gateway_concurrency_leases lease WHERE lease.task_id = task.id AND lease.released_at IS NULL),
       (SELECT count(*) FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.failed'),
       (SELECT count(*) FROM gateway_task_callback_outbox callback WHERE callback.task_id = task.id)
FROM gateway_tasks task
WHERE task.id = $1::uuid`, resilienceTask.ID).Scan(
		&resilienceStatus,
		&resilienceCode,
		&resilienceAdmissions,
		&resilienceLeases,
		&resilienceFailedEvents,
		&resilienceCallbacks,
	); err != nil {
		t.Fatalf("read resilience terminal state: %v", err)
	}
	if resilienceStatus != "failed" || resilienceCode != "no_model_candidate" ||
		resilienceAdmissions != 0 || resilienceLeases != 0 ||
		resilienceFailedEvents != 1 || resilienceCallbacks != 1 {
		t.Fatalf(
			"resilience terminal state=%s/%s admissions=%d leases=%d events=%d callbacks=%d",
			resilienceStatus,
			resilienceCode,
			resilienceAdmissions,
			resilienceLeases,
			resilienceFailedEvents,
			resilienceCallbacks,
		)
	}
	retryExpiredTask := createTask(true)
	retryExpiredInput := resilienceInput
	retryExpiredInput.TaskID = retryExpiredTask.ID
	if _, err := first.QueueTaskAdmissionWithHook(ctx, retryExpiredInput, nil); err != nil {
		t.Fatalf("queue retry-expired task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET enqueued_at = now() - interval '2 seconds',
    wait_deadline_at = now() - interval '1 second',
    dispatch_failure_count = 3,
    dispatch_last_outcome = 'deferred',
    dispatch_last_error_code = 'client_error',
    dispatch_last_error_at = now() - interval '1 minute'
WHERE task_id = $1::uuid`, retryExpiredTask.ID); err != nil {
		t.Fatalf("expire deferred task: %v", err)
	}
	reaped, err = first.ReapExpiredTaskAdmissions(ctx, 10, "https://callback.invalid/admission-timeout")
	if err != nil || reaped.ExpiredDeadlines < 1 {
		t.Fatalf("reap retry-expired task=%+v err=%v", reaped, err)
	}
	var retryExpiredCode string
	var retryExpiredEvents, retryExpiredCallbacks int
	if err := first.pool.QueryRow(ctx, `
SELECT COALESCE(task.error_code, ''),
       (SELECT count(*) FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.failed'),
       (SELECT count(*) FROM gateway_task_callback_outbox callback WHERE callback.task_id = task.id)
FROM gateway_tasks task
WHERE task.id = $1::uuid`, retryExpiredTask.ID).Scan(
		&retryExpiredCode,
		&retryExpiredEvents,
		&retryExpiredCallbacks,
	); err != nil {
		t.Fatalf("read retry-expired task: %v", err)
	}
	if retryExpiredCode != "admission_dispatch_failed" || retryExpiredEvents != 1 || retryExpiredCallbacks != 1 {
		t.Fatalf(
			"retry-expired task code=%s events=%d callbacks=%d, want admission_dispatch_failed/1/1",
			retryExpiredCode,
			retryExpiredEvents,
			retryExpiredCallbacks,
		)
	}

	disconnectedTask := createTask(true)
	disconnectedInput := resilienceInput
	disconnectedInput.TaskID = disconnectedTask.ID
	if _, err := first.QueueTaskAdmissionWithHook(ctx, disconnectedInput, nil); err != nil {
		t.Fatalf("queue disconnected task: %v", err)
	}
	if _, changed, err := first.CancelQueuedTaskWithCallback(
		ctx,
		disconnectedTask.ID,
		"client disconnected before upstream submission",
		"https://callback.invalid/admission-disconnect",
	); err != nil || !changed {
		t.Fatalf("cancel disconnected task changed=%v err=%v", changed, err)
	}
	var disconnectedEvents, disconnectedCallbacks, disconnectedAdmissions int
	if err := first.pool.QueryRow(ctx, `
SELECT (SELECT count(*) FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.cancelled'),
       (SELECT count(*) FROM gateway_task_callback_outbox callback WHERE callback.task_id = task.id),
       (SELECT count(*) FROM gateway_task_admissions admission WHERE admission.task_id = task.id)
FROM gateway_tasks task
WHERE task.id = $1::uuid`, disconnectedTask.ID).Scan(
		&disconnectedEvents,
		&disconnectedCallbacks,
		&disconnectedAdmissions,
	); err != nil {
		t.Fatalf("read disconnected task cleanup: %v", err)
	}
	if disconnectedEvents != 1 || disconnectedCallbacks != 1 || disconnectedAdmissions != 0 {
		t.Fatalf(
			"disconnected task events=%d callbacks=%d admissions=%d, want 1/1/0",
			disconnectedEvents,
			disconnectedCallbacks,
			disconnectedAdmissions,
		)
	}

	terminalResidue := createTask(true)
	result, err = first.TryTaskAdmission(ctx, inputFor(terminalResidue, 100, ""))
	if err != nil || !result.Admitted {
		t.Fatalf("admit terminal residue task: result=%+v err=%v", result, err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
    error_code = 'integration_test',
    finished_at = now(),
    updated_at = now()
WHERE id = $1::uuid`, terminalResidue.ID); err != nil {
		t.Fatalf("mark terminal residue task failed: %v", err)
	}
	recovery, err = first.RecoverInterruptedRuntimeState(ctx)
	if err != nil {
		t.Fatalf("clean terminal admission residue: %v", err)
	}
	if recovery.CleanedTaskAdmissions < 1 {
		t.Fatalf("terminal recovery = %+v, want admission cleanup", recovery)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		terminalResidue.ID,
	).Scan(&recoveredAdmissions, &recoveredLeases); err != nil {
		t.Fatalf("count terminal admission residue: %v", err)
	}
	if recoveredAdmissions != 0 || recoveredLeases != 0 {
		t.Fatalf("terminal task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases)
	}
}

func TestWaitingAdmissionLeadersDoNotCrossBlockIndependentCandidates(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 task admission PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()
	var databaseName string
	if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		t.Fatalf("read test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
	platform, err := db.CreatePlatform(ctx, CreatePlatformInput{
		Provider:    "leader-test",
		PlatformKey: "leader-test-" + suffix,
		Name:        "Admission Leader Test " + suffix,
		AuthType:    "none",
		Status:      "enabled",
	})
	if err != nil {
		t.Fatalf("create platform: %v", err)
	}
	modelIDs := make([]string, 0, 2)
	for index := 0; index < 2; index++ {
		modelName := fmt.Sprintf("leader-model-%d-%s", index, suffix)
		var modelID string
		if err := db.pool.QueryRow(ctx, `
INSERT INTO platform_models (
  platform_id, model_name, provider_model_name, model_alias, model_type,
  display_name, pricing_mode, enabled
)
VALUES ($1::uuid, $2, $2, $2, '["image_generate"]'::jsonb, $2, 'inherit_discount', true)
RETURNING id::text`, platform.ID, modelName).Scan(&modelID); err != nil {
			t.Fatalf("create platform model %d: %v", index, err)
		}
		modelIDs = append(modelIDs, modelID)
	}
	group, err := db.CreateUserGroup(ctx, UserGroupInput{
		GroupKey: "leader-test-" + suffix,
		Name:     "Admission Leader Test " + suffix,
		Status:   "active",
	})
	if err != nil {
		t.Fatalf("create user group: %v", err)
	}
	taskIDs := make([]string, 0, 4)
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cleanupCancel()
		for _, taskID := range taskIDs {
			_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, taskID)
		}
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_user_groups WHERE id = $1::uuid`, group.ID)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM platform_models WHERE id = ANY($1::uuid[])`, modelIDs)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
	})
	createTask := func(modelIndex int) GatewayTask {
		t.Helper()
		task, createErr := db.CreateTask(ctx, CreateTaskInput{
			Kind:    "images.generate",
			Model:   fmt.Sprintf("leader-model-%d-%s", modelIndex, suffix),
			RunMode: "production",
			Request: map[string]any{"prompt": "admission leader test"},
		}, &auth.User{ID: "leader-test-user-" + suffix, Source: "gateway"})
		if createErr != nil {
			t.Fatalf("create task: %v", createErr)
		}
		taskIDs = append(taskIDs, task.ID)
		return task
	}
	tasks := []GatewayTask{createTask(0), createTask(0), createTask(1), createTask(1)}
	insertWaiting := func(task GatewayTask, modelID string, groupID string, enqueuedAt time.Time) {
		t.Helper()
		if _, insertErr := db.pool.Exec(ctx, `
INSERT INTO gateway_task_admissions (
  task_id, platform_id, platform_model_id, user_group_id, queue_key, mode,
  status, priority, enqueued_at, wait_deadline_at, waiter_id,
  waiter_lease_expires_at
)
VALUES (
  $1::uuid, $2::uuid, $3::uuid, NULLIF($4, '')::uuid, $5, 'sync',
  'waiting', 100, $6, now() + interval '10 minutes', $7,
  now() + interval '15 seconds'
)`,
			task.ID,
			platform.ID,
			modelID,
			groupID,
			"leader-test:"+modelID,
			enqueuedAt,
			"waiter-"+task.ID,
		); insertErr != nil {
			t.Fatalf("insert waiting admission: %v", insertErr)
		}
	}
	baseTime := time.Now().Add(-time.Minute)
	insertWaiting(tasks[0], modelIDs[0], "", baseTime)
	insertWaiting(tasks[1], modelIDs[0], "", baseTime.Add(time.Second))
	insertWaiting(tasks[2], modelIDs[1], "", baseTime.Add(2*time.Second))
	insertWaiting(tasks[3], modelIDs[1], "", baseTime.Add(3*time.Second))
	leaders, err := db.ListWaitingTaskAdmissionIDs(ctx, 256)
	if err != nil {
		t.Fatalf("list independent candidate leaders: %v", err)
	}
	leaderSet := make(map[string]bool, len(leaders))
	for _, taskID := range leaders {
		leaderSet[taskID] = true
	}
	if !leaderSet[tasks[0].ID] || !leaderSet[tasks[2].ID] {
		t.Fatalf("independent candidate leaders = %v, want %s and %s", leaders, tasks[0].ID, tasks[2].ID)
	}
	if leaderSet[tasks[1].ID] || leaderSet[tasks[3].ID] {
		t.Fatalf("non-head independent candidate task was selected: %v", leaders)
	}

	if _, err := db.pool.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = ANY($1::uuid[])`, taskIDs); err != nil {
		t.Fatalf("clear independent admissions: %v", err)
	}
	insertWaiting(tasks[0], modelIDs[0], group.ID, baseTime)
	insertWaiting(tasks[2], modelIDs[1], group.ID, baseTime.Add(time.Second))
	leaders, err = db.ListWaitingTaskAdmissionIDs(ctx, 256)
	if err != nil {
		t.Fatalf("list grouped candidate leaders: %v", err)
	}
	leaderSet = make(map[string]bool, len(leaders))
	for _, taskID := range leaders {
		leaderSet[taskID] = true
	}
	if !leaderSet[tasks[0].ID] || leaderSet[tasks[2].ID] {
		t.Fatalf("grouped candidate leaders = %v, want only group head %s", leaders, tasks[0].ID)
	}
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = ANY($1::uuid[])`, taskIDs); err != nil {
		t.Fatalf("clear grouped admission leaders: %v", err)
	}

	createAsyncTask := func(modelIndex int) GatewayTask {
		t.Helper()
		task, createErr := db.CreateTask(ctx, CreateTaskInput{
			Kind:    "images.generate",
			Model:   fmt.Sprintf("leader-model-%d-%s", modelIndex, suffix),
			RunMode: "production",
			Async:   true,
			Request: map[string]any{"prompt": "async admission independence test"},
		}, &auth.User{ID: "leader-test-user-" + suffix, Source: "gateway"})
		if createErr != nil {
			t.Fatalf("create async task: %v", createErr)
		}
		taskIDs = append(taskIDs, task.ID)
		return task
	}
	platformScope := func(modelID string) AdmissionScope {
		return AdmissionScope{
			ScopeType:       "platform_model",
			ScopeKey:        modelID,
			ConcurrentLimit: 1,
			Amount:          1,
			LeaseTTLSeconds: 120,
			QueueLimit:      10,
			MaxWaitSeconds:  600,
		}
	}
	groupScope := AdmissionScope{
		ScopeType:       "user_group",
		ScopeKey:        group.ID,
		ConcurrentLimit: 1,
		Amount:          1,
		LeaseTTLSeconds: 120,
		QueueLimit:      10,
		MaxWaitSeconds:  600,
	}
	holder := createAsyncTask(0)
	blocked := createAsyncTask(0)
	independent := createAsyncTask(1)
	holderInput := TaskAdmissionInput{
		TaskID: holder.ID, PlatformID: platform.ID, PlatformModelID: modelIDs[0],
		UserGroupID: group.ID, QueueKey: "leader-test:" + modelIDs[0], Mode: "async", Priority: 100,
		Scopes: []AdmissionScope{platformScope(modelIDs[0]), groupScope},
	}
	blockedInput := holderInput
	blockedInput.TaskID = blocked.ID
	independentInput := TaskAdmissionInput{
		TaskID: independent.ID, PlatformID: platform.ID, PlatformModelID: modelIDs[1],
		QueueKey: "leader-test:" + modelIDs[1], Mode: "async", Priority: 100,
		Scopes: []AdmissionScope{platformScope(modelIDs[1])},
	}
	if result, err := db.TryTaskAdmission(ctx, holderInput); err != nil || !result.Admitted {
		t.Fatalf("admit local blocker holder: result=%+v err=%v", result, err)
	}
	if _, err := db.QueueTaskAdmissionWithHook(ctx, blockedInput, nil); err != nil {
		t.Fatalf("queue locally blocked task: %v", err)
	}
	if _, err := db.QueueTaskAdmissionWithHook(ctx, independentInput, nil); err != nil {
		t.Fatalf("queue independent task: %v", err)
	}
	outcomes, err := db.TryTaskAdmissionAtomicBatchWithAdmittedHook(
		ctx,
		[]TaskAdmissionInput{blockedInput, independentInput},
		nil,
	)
	if err != nil || len(outcomes) != 2 {
		t.Fatalf("independent atomic outcomes=%+v err=%v", outcomes, err)
	}
	if outcomes[0].Result.Admitted || len(outcomes[0].Result.Blockers) == 0 {
		t.Fatalf("locally blocked outcome=%+v, want explicit blockers", outcomes[0])
	}
	if !outcomes[1].Result.Admitted {
		t.Fatalf("independent candidate was cross-blocked: %+v", outcomes[1])
	}
	for _, taskID := range []string{holder.ID, blocked.ID, independent.ID} {
		if err := db.DeleteTaskAdmission(ctx, taskID); err != nil {
			t.Fatalf("clear independent atomic admission %s: %v", taskID, err)
		}
	}

	hookTasks := []GatewayTask{createAsyncTask(0), createAsyncTask(0), createAsyncTask(1)}
	hookInputs := make([]TaskAdmissionInput, 0, len(hookTasks))
	for index, task := range hookTasks {
		input := TaskAdmissionInput{
			TaskID: task.ID, PlatformID: platform.ID, PlatformModelID: modelIDs[index%len(modelIDs)],
			QueueKey: "hook-isolation:" + task.ID, Mode: "async", Priority: 100,
			Scopes: []AdmissionScope{{
				ScopeType:       "worker_capacity",
				ScopeKey:        "hook-isolation-" + task.ID,
				ConcurrentLimit: 1,
				Amount:          1,
				LeaseTTLSeconds: 120,
				QueueLimit:      10,
				MaxWaitSeconds:  600,
			}},
		}
		if _, err := db.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
			t.Fatalf("queue hook isolation task %d: %v", index, err)
		}
		hookInputs = append(hookInputs, input)
	}
	hookFailure := errors.New("synthetic per-task hook failure")
	outcomes, err = db.TryTaskAdmissionAtomicBatchWithAdmittedHook(
		ctx,
		hookInputs,
		func(_ pgx.Tx, input TaskAdmissionInput) error {
			if input.TaskID == hookTasks[1].ID {
				return hookFailure
			}
			return nil
		},
	)
	if !errors.Is(err, hookFailure) || len(outcomes) != 2 || outcomes[1].TaskID != hookTasks[1].ID {
		t.Fatalf("hook failure outcomes=%+v err=%v", outcomes, err)
	}
	var hookAdmitted, hookLeases int
	if err := db.pool.QueryRow(ctx, `
SELECT (SELECT count(*) FROM gateway_task_admissions WHERE task_id = ANY($1::uuid[]) AND status = 'admitted'),
       (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL)`,
		[]string{hookTasks[0].ID, hookTasks[1].ID, hookTasks[2].ID},
	).Scan(&hookAdmitted, &hookLeases); err != nil {
		t.Fatalf("read rolled back hook batch: %v", err)
	}
	if hookAdmitted != 0 || hookLeases != 0 {
		t.Fatalf("failed hook batch left admitted=%d leases=%d", hookAdmitted, hookLeases)
	}
	if _, changed, err := db.DeferWaitingTaskAdmission(
		ctx,
		hookTasks[1].ID,
		hookInputs[1].PlatformModelID,
		time.Now().Add(time.Minute),
		"deferred",
		"synthetic_hook_failure",
	); err != nil || !changed {
		t.Fatalf("isolate hook failure changed=%v err=%v", changed, err)
	}
	for _, index := range []int{0, 2} {
		if result, err := db.TryTaskAdmission(ctx, hookInputs[index]); err != nil || !result.Admitted {
			t.Fatalf("admit hook survivor %d result=%+v err=%v", index, result, err)
		}
	}
}

func TestWorkerCapacityAllocationAndFailover(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 worker registry PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()
	var databaseName string
	if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		t.Fatalf("read test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	prefix := "queue-test-" + strings.ReplaceAll(uuid.NewString(), "-", "")
	firstID := prefix + "-a"
	secondID := prefix + "-b"
	defer db.pool.Exec(context.Background(), `DELETE FROM gateway_worker_instances WHERE instance_id = ANY($1::text[])`, []string{firstID, secondID})

	first, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
		t.Fatalf("first allocation = %+v, err=%v", first, err)
	}
	second, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: secondID, DesiredCapacity: 5})
	if err != nil || second.Allocated != 2 || second.ActiveInstances != 2 {
		t.Fatalf("second allocation = %+v, err=%v", second, err)
	}
	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 3 || first.ActiveInstances != 2 {
		t.Fatalf("rebalanced first allocation = %+v, err=%v", first, err)
	}
	if _, err := db.pool.Exec(ctx, `
UPDATE gateway_worker_instances
SET heartbeat_at = now() - $2::interval
WHERE instance_id = $1`, secondID, (workerHeartbeatStaleAfter + time.Second).String()); err != nil {
		t.Fatalf("expire second heartbeat: %v", err)
	}
	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
		t.Fatalf("single-worker failover allocation = %+v, err=%v", first, err)
	}

	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
		InstanceID: firstID, DesiredCapacity: 100, CapacityLimit: 3,
	})
	if err != nil || first.Allocated != 3 || first.GlobalAllocated != 3 {
		t.Fatalf("bounded single-worker allocation = %+v, err=%v", first, err)
	}
	second, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
		InstanceID: secondID, DesiredCapacity: 100, CapacityLimit: 2,
	})
	if err != nil || second.Allocated != 2 || second.GlobalAllocated != 5 || second.ActiveInstances != 2 {
		t.Fatalf("bounded two-worker allocation = %+v, err=%v", second, err)
	}
	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
		InstanceID: firstID, DesiredCapacity: 100, CapacityLimit: 10,
		LoadMode: "adaptive", SafeCapacity: 0, HeavyCapacity: 1,
		ActiveTasks: 2, WaitingUpstreamTasks: 2,
		PressureState: "critical", PressureReason: "memory",
	})
	if err != nil || first.Allocated != 0 {
		t.Fatalf("critical worker allocation = %+v, err=%v", first, err)
	}
	second, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
		InstanceID: secondID, DesiredCapacity: 100, CapacityLimit: 10,
		LoadMode: "adaptive", SafeCapacity: 5, HeavyCapacity: 2,
		ActiveTasks: 3, PreparingTasks: 1, WaitingUpstreamTasks: 1, FinalizingTasks: 1,
		PressureState: "normal",
	})
	if err != nil || second.Allocated != 5 || second.GlobalAllocated != 5 {
		t.Fatalf("adaptive redistribution allocation = %+v, err=%v", second, err)
	}
	instances, err := db.ListWorkerInstanceRuntime(ctx)
	if err != nil {
		t.Fatalf("list adaptive worker runtime: %v", err)
	}
	foundCritical := false
	for _, instance := range instances {
		if instance.InstanceID == firstID {
			foundCritical = instance.SafeCapacity == 0 && instance.ReportedActiveTasks == 2 && instance.WaitingUpstreamTasks == 2 && instance.PressureState == "critical"
		}
	}
	if !foundCritical {
		t.Fatalf("adaptive runtime did not expose the critical worker: %+v", instances)
	}
}
