feat(queue): 增加非文本模型分布式准入队列

使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配，并将生产 API 与独立 Worker 角色拆分。补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本；未默认启用任何生产 queue_size 策略。已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收；合入最新主干后将重新执行发布门禁。
2026-07-29 16:15:43 +08:00
parent 2e5a90731b
commit 9e4fc7362d
55 changed files with 4565 additions and 237 deletions
@@ -0,0 +1,499 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"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)
+	}
+	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")
+	}
+	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)
+	}
+
+	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)
+	}
+	var riverJobID int64
+	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)
+	}
+}
+
+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() - interval '16 seconds'
+WHERE instance_id = $1`, secondID); 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)
+	}
+}
@@ -7,6 +7,7 @@ import (

 type AsyncWorkerCapacitySnapshot struct {
 	Capacity        int
+	GlobalCapacity  int
 	Desired         int
 	HardLimit       int
 	Capped          bool
@@ -16,6 +17,8 @@ type AsyncWorkerCapacitySnapshot struct {
 	UnlimitedGroups int
 	ModelDesired    int
 	GroupDesired    int
+	ActiveInstances int
+	InstanceID      string
 }

 func (s *Store) AsyncWorkerCapacity(ctx context.Context, hardLimit int) (AsyncWorkerCapacitySnapshot, error) {
@@ -131,6 +134,7 @@ func asyncWorkerCapacityFromPolicySets(modelPolicies []map[string]any, groupPoli
 		snapshot.Capacity = hardLimit
 		snapshot.Capped = true
 	}
+	snapshot.GlobalCapacity = snapshot.Capacity
 	return snapshot
 }

@@ -212,6 +212,11 @@ WHERE id = $1::uuid AND deleted_at IS NULL`, id)

 func (s *Store) CreateUserGroup(ctx context.Context, input UserGroupInput) (UserGroup, error) {
 	input = normalizeUserGroupInput(input)
+	normalizedPolicy, err := NormalizeAndValidateRateLimitPolicy(input.RateLimitPolicy)
+	if err != nil {
+		return UserGroup{}, err
+	}
+	input.RateLimitPolicy = normalizedPolicy
 	rechargeDiscountPolicy, _ := json.Marshal(emptyObjectIfNil(input.RechargeDiscountPolicy))
 	billingDiscountPolicy, _ := json.Marshal(emptyObjectIfNil(input.BillingDiscountPolicy))
 	rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.RateLimitPolicy))
@@ -231,6 +236,11 @@ RETURNING `+userGroupColumns,

 func (s *Store) UpdateUserGroup(ctx context.Context, id string, input UserGroupInput) (UserGroup, error) {
 	input = normalizeUserGroupInput(input)
+	normalizedPolicy, err := NormalizeAndValidateRateLimitPolicy(input.RateLimitPolicy)
+	if err != nil {
+		return UserGroup{}, err
+	}
+	input.RateLimitPolicy = normalizedPolicy
 	rechargeDiscountPolicy, _ := json.Marshal(emptyObjectIfNil(input.RechargeDiscountPolicy))
 	billingDiscountPolicy, _ := json.Marshal(emptyObjectIfNil(input.BillingDiscountPolicy))
 	rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.RateLimitPolicy))
@@ -93,6 +93,19 @@ WHERE resource_type = 'platform_model'
 func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier, input CreatePlatformModelInput) (PlatformModel, error) {
 	input.ModelName = strings.TrimSpace(input.ModelName)
 	input.ProviderModelName = strings.TrimSpace(input.ProviderModelName)
+	normalizedRateLimitPolicy, err := NormalizeAndValidateRateLimitPolicy(input.RateLimitPolicy)
+	if err != nil {
+		return PlatformModel{}, err
+	}
+	input.RateLimitPolicy = normalizedRateLimitPolicy
+	if rawPolicy, ok := input.RuntimePolicyOverride["rateLimitPolicy"].(map[string]any); ok {
+		normalizedRuntimeOverridePolicy, normalizeErr := NormalizeAndValidateRateLimitPolicy(rawPolicy)
+		if normalizeErr != nil {
+			return PlatformModel{}, normalizeErr
+		}
+		input.RuntimePolicyOverride = cloneObject(input.RuntimePolicyOverride)
+		input.RuntimePolicyOverride["rateLimitPolicy"] = normalizedRuntimeOverridePolicy
+	}
 	base, err := s.lookupBaseModel(ctx, q, input.BaseModelID, input.CanonicalModelKey, input.ModelName)
 	if err != nil && !IsNotFound(err) {
 		return PlatformModel{}, err
@@ -76,7 +76,11 @@ var (
 )

 func Connect(ctx context.Context, databaseURL string) (*Store, error) {
-	config, err := postgresPoolConfig(databaseURL)
+	return ConnectWithMaxConns(ctx, databaseURL, 0)
+}
+
+func ConnectWithMaxConns(ctx context.Context, databaseURL string, maxConns int) (*Store, error) {
+	config, err := postgresPoolConfig(databaseURL, maxConns)
 	if err != nil {
 		return nil, err
 	}
@@ -91,13 +95,16 @@ func Connect(ctx context.Context, databaseURL string) (*Store, error) {
 	return &Store{pool: pool}, nil
 }

-func postgresPoolConfig(databaseURL string) (*pgxpool.Config, error) {
+func postgresPoolConfig(databaseURL string, maxConns ...int) (*pgxpool.Config, error) {
 	config, err := pgxpool.ParseConfig(databaseURL)
 	if err != nil {
 		return nil, err
 	}
 	config.ConnConfig.ConnectTimeout = postgresConnectTimeout
 	config.ConnConfig.RuntimeParams["application_name"] = postgresApplicationName
+	if len(maxConns) > 0 && maxConns[0] > 0 {
+		config.MaxConns = int32(maxConns[0])
+	}
 	return config, nil
 }

@@ -523,6 +530,8 @@ type GatewayTask struct {
 	AsyncMode                     bool           `json:"asyncMode"`
 	RiverJobID                    int64          `json:"riverJobId,omitempty"`
 	Status                        string         `json:"status"`
+	QueueKey                      string         `json:"-"`
+	Priority                      int            `json:"-"`
 	Cancellable                   *bool          `json:"cancellable,omitempty"`
 	Submitted                     *bool          `json:"submitted,omitempty"`
 	Message                       string         `json:"message,omitempty"`
@@ -570,7 +579,8 @@ COALESCE(api_key_id, ''), COALESCE(api_key_name, ''), COALESCE(api_key_prefix, '
 COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model,
 COALESCE(model_type, ''), COALESCE(requested_model, ''), COALESCE(resolved_model, ''), COALESCE(request_id, ''),
 COALESCE(conversation_id::text, ''), COALESCE(new_message_count, 0),
-request, COALESCE(async_mode, false), COALESCE(river_job_id, 0), status, COALESCE(attempt_count, 0),
+request, COALESCE(async_mode, false), COALESCE(river_job_id, 0), status,
+COALESCE(queue_key, 'default'), COALESCE(priority, 100), COALESCE(attempt_count, 0),
 COALESCE(remote_task_id, ''), COALESCE(remote_task_payload, '{}'::jsonb),
 COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
 COALESCE(usage, '{}'::jsonb), COALESCE(metrics, '{}'::jsonb), COALESCE(billing_summary, '{}'::jsonb),
@@ -715,6 +725,11 @@ ORDER BY COALESCE(dynamic_priority, priority) ASC, priority ASC, created_at DESC
 }

 func (s *Store) CreatePlatform(ctx context.Context, input CreatePlatformInput) (Platform, error) {
+	normalizedPolicy, err := NormalizeAndValidateRateLimitPolicy(input.RateLimitPolicy)
+	if err != nil {
+		return Platform{}, err
+	}
+	input.RateLimitPolicy = normalizedPolicy
 	credentials, _ := json.Marshal(emptyObjectIfNil(input.Credentials))
 	config, _ := json.Marshal(emptyObjectIfNil(input.Config))
 	retryPolicy, _ := json.Marshal(emptyObjectIfNil(input.RetryPolicy))
@@ -734,7 +749,7 @@ func (s *Store) CreatePlatform(ctx context.Context, input CreatePlatformInput) (
 	var retryPolicyBytes []byte
 	var rateLimitPolicyBytes []byte
 	var dynamicPriority sql.NullInt64
-	err := s.pool.QueryRow(ctx, `
+	err = s.pool.QueryRow(ctx, `
 INSERT INTO integration_platforms (
  provider, platform_key, name, internal_name, base_url, auth_type, credentials, config,
  default_pricing_mode, default_discount_factor, pricing_rule_set_id,
@@ -789,6 +804,11 @@ RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), C
 }

 func (s *Store) UpdatePlatform(ctx context.Context, id string, input CreatePlatformInput) (Platform, error) {
+	normalizedPolicy, err := NormalizeAndValidateRateLimitPolicy(input.RateLimitPolicy)
+	if err != nil {
+		return Platform{}, err
+	}
+	input.RateLimitPolicy = normalizedPolicy
 	var credentials any
 	if input.Credentials != nil {
 		credentialsBytes, _ := json.Marshal(emptyObjectIfNil(input.Credentials))
@@ -812,7 +832,7 @@ func (s *Store) UpdatePlatform(ctx context.Context, id string, input CreatePlatf
 	var retryPolicyBytes []byte
 	var rateLimitPolicyBytes []byte
 	var dynamicPriority sql.NullInt64
-	err := s.pool.QueryRow(ctx, `
+	err = s.pool.QueryRow(ctx, `
 UPDATE integration_platforms
 SET provider = $2,
    platform_key = COALESCE(NULLIF($3, ''), platform_key),
@@ -2138,6 +2158,8 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
 		&task.AsyncMode,
 		&task.RiverJobID,
 		&task.Status,
+		&task.QueueKey,
+		&task.Priority,
 		&task.AttemptCount,
 		&task.RemoteTaskID,
 		&remoteTaskPayloadBytes,
@@ -1,13 +1,24 @@
 package store

 import (
+	"errors"
+	"fmt"
 	"math"
 	"strings"
 )

+var ErrInvalidRateLimitPolicy = errors.New("invalid rate limit policy")
+
+func IsInvalidRateLimitPolicy(err error) bool {
+	return errors.Is(err, ErrInvalidRateLimitPolicy)
+}
+
 const (
 	RateLimitPolicyModeInherit  = "inherit"
 	RateLimitPolicyModeOverride = "override"
+	DefaultQueueMaxWaitSeconds  = 600
+	MaxQueueSize                = 10000
+	MaxQueueWaitSeconds         = 3600
 )

 type EffectiveRateLimitPolicyInput struct {
@@ -78,46 +89,150 @@ func NormalizeRateLimitPolicy(policy map[string]any) map[string]any {
 	}
 	out := clonePolicy(policy)
 	rules, _ := out["rules"].([]any)
-	if len(rules) > 0 {
-		return out
-	}
-
-	legacyScopes := []map[string]any{policy}
-	for _, key := range []string{"platformLimits", "modelLimits", "platform_limits", "model_limits"} {
-		if nested, ok := policy[key].(map[string]any); ok {
-			legacyScopes = append(legacyScopes, nested)
+	if len(rules) == 0 {
+		legacyScopes := []map[string]any{policy}
+		for _, key := range []string{"platformLimits", "modelLimits", "platform_limits", "model_limits"} {
+			if nested, ok := policy[key].(map[string]any); ok {
+				legacyScopes = append(legacyScopes, nested)
+			}
+		}
+		if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
+			"max_concurrent_requests", "maxConcurrentRequests", "concurrent"); ok {
+			rules = append(rules, map[string]any{
+				"metric":          "concurrent",
+				"limit":           limit,
+				"leaseTtlSeconds": 120,
+			})
+		}
+		if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
+			"max_request_per_minute", "maxRequestsPerMinute", "rpm"); ok {
+			rules = append(rules, map[string]any{
+				"metric":        "rpm",
+				"limit":         limit,
+				"windowSeconds": 60,
+			})
+		}
+		if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
+			"max_tokens_per_minute", "maxTokensPerMinute", "tpm_total"); ok {
+			rules = append(rules, map[string]any{
+				"metric":        "tpm_total",
+				"limit":         limit,
+				"windowSeconds": 60,
+			})
 		}
 	}
-	if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
-		"max_concurrent_requests", "maxConcurrentRequests", "concurrent"); ok {
-		rules = append(rules, map[string]any{
-			"metric":          "concurrent",
-			"limit":           limit,
-			"leaseTtlSeconds": 120,
-		})
-	}
-	if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
-		"max_request_per_minute", "maxRequestsPerMinute", "rpm"); ok {
-		rules = append(rules, map[string]any{
-			"metric":        "rpm",
-			"limit":         limit,
-			"windowSeconds": 60,
-		})
-	}
-	if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
-		"max_tokens_per_minute", "maxTokensPerMinute", "tpm_total"); ok {
-		rules = append(rules, map[string]any{
-			"metric":        "tpm_total",
-			"limit":         limit,
-			"windowSeconds": 60,
-		})
-	}
-	if len(rules) > 0 {
-		out["rules"] = rules
+	normalizedRules := make([]any, 0, len(rules))
+	for _, rawRule := range rules {
+		rule, ok := rawRule.(map[string]any)
+		if !ok {
+			normalizedRules = append(normalizedRules, rawRule)
+			continue
+		}
+		normalizedRule := clonePolicy(rule)
+		if strings.TrimSpace(stringValue(normalizedRule["metric"])) == "queue_size" {
+			if floatValue(normalizedRule["limit"]) <= 0 {
+				continue
+			}
+			if floatValue(normalizedRule["maxWaitSeconds"]) <= 0 {
+				normalizedRule["maxWaitSeconds"] = DefaultQueueMaxWaitSeconds
+			}
+		} else {
+			delete(normalizedRule, "maxWaitSeconds")
+		}
+		normalizedRules = append(normalizedRules, normalizedRule)
 	}
+	out["rules"] = normalizedRules
 	return out
 }

+type QueuePolicyRule struct {
+	Limit          int
+	MaxWaitSeconds int
+	Policy         map[string]any
+}
+
+func QueueRuleFromPolicy(policy map[string]any) (QueuePolicyRule, bool) {
+	normalized := NormalizeRateLimitPolicy(policy)
+	rules, _ := normalized["rules"].([]any)
+	for _, rawRule := range rules {
+		rule, _ := rawRule.(map[string]any)
+		if strings.TrimSpace(stringValue(rule["metric"])) != "queue_size" {
+			continue
+		}
+		limit := int(math.Floor(floatValue(rule["limit"])))
+		if limit <= 0 {
+			return QueuePolicyRule{}, false
+		}
+		maxWaitSeconds := int(math.Floor(floatValue(rule["maxWaitSeconds"])))
+		if maxWaitSeconds <= 0 {
+			maxWaitSeconds = DefaultQueueMaxWaitSeconds
+		}
+		return QueuePolicyRule{
+			Limit:          limit,
+			MaxWaitSeconds: maxWaitSeconds,
+			Policy:         normalized,
+		}, true
+	}
+	return QueuePolicyRule{}, false
+}
+
+// NormalizeAndValidateRateLimitPolicy validates only the canonical queue
+// extension. Existing rate-limit metrics intentionally keep their historical
+// permissive parsing contract.
+func NormalizeAndValidateRateLimitPolicy(policy map[string]any) (map[string]any, error) {
+	if policy == nil {
+		return nil, nil
+	}
+	rulesValue, rulesPresent := policy["rules"]
+	if rulesPresent {
+		if _, ok := rulesValue.([]any); !ok {
+			return nil, fmt.Errorf("%w: rateLimitPolicy.rules must be an array", ErrInvalidRateLimitPolicy)
+		}
+	}
+	rules, _ := rulesValue.([]any)
+	for index, rawRule := range rules {
+		rule, ok := rawRule.(map[string]any)
+		if !ok {
+			return nil, fmt.Errorf("%w: rateLimitPolicy.rules[%d] must be an object", ErrInvalidRateLimitPolicy, index)
+		}
+		metric := strings.TrimSpace(stringValue(rule["metric"]))
+		maxWait, hasMaxWait := numericRuleValue(rule, "maxWaitSeconds")
+		if metric != "queue_size" {
+			if hasMaxWait {
+				return nil, fmt.Errorf("%w: rateLimitPolicy.rules[%d].maxWaitSeconds is only valid for queue_size", ErrInvalidRateLimitPolicy, index)
+			}
+			continue
+		}
+		limit, hasLimit := numericRuleValue(rule, "limit")
+		if !hasLimit {
+			return nil, fmt.Errorf("%w: rateLimitPolicy.rules[%d].limit must be an integer", ErrInvalidRateLimitPolicy, index)
+		}
+		if limit != math.Trunc(limit) || limit < 0 || limit > MaxQueueSize {
+			return nil, fmt.Errorf("%w: queue_size limit must be 0 or an integer between 1 and %d", ErrInvalidRateLimitPolicy, MaxQueueSize)
+		}
+		if limit == 0 {
+			continue
+		}
+		if hasMaxWait && (maxWait != math.Trunc(maxWait) || maxWait < 1 || maxWait > MaxQueueWaitSeconds) {
+			return nil, fmt.Errorf("%w: queue_size maxWaitSeconds must be an integer between 1 and %d", ErrInvalidRateLimitPolicy, MaxQueueWaitSeconds)
+		}
+	}
+	return NormalizeRateLimitPolicy(policy), nil
+}
+
+func numericRuleValue(rule map[string]any, key string) (float64, bool) {
+	raw, ok := rule[key]
+	if !ok || raw == nil {
+		return 0, false
+	}
+	switch raw.(type) {
+	case float64, float32, int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
+		return floatValue(raw), true
+	default:
+		return 0, false
+	}
+}
+
 func ConcurrentPolicyCapacity(policy map[string]any) (int, bool) {
 	limit, ok := RateLimitPolicyMetric(policy, "concurrent")
 	if !ok || limit <= 0 {
@@ -2,6 +2,66 @@ package store

 import "testing"

+func TestNormalizeAndValidateQueuePolicy(t *testing.T) {
+	t.Run("defaults maximum wait and preserves extensions", func(t *testing.T) {
+		policy, err := NormalizeAndValidateRateLimitPolicy(map[string]any{
+			"strategy": "strict",
+			"rules": []any{map[string]any{
+				"metric": "queue_size",
+				"limit":  12,
+				"source": "admin",
+			}},
+		})
+		if err != nil {
+			t.Fatalf("NormalizeAndValidateRateLimitPolicy() error = %v", err)
+		}
+		queue, ok := QueueRuleFromPolicy(policy)
+		if !ok || queue.Limit != 12 || queue.MaxWaitSeconds != 600 {
+			t.Fatalf("queue rule = %+v, %v", queue, ok)
+		}
+		rules := policy["rules"].([]any)
+		rule := rules[0].(map[string]any)
+		if rule["source"] != "admin" {
+			t.Fatalf("unknown extension was lost: %+v", rule)
+		}
+	})
+
+	t.Run("normalizes zero to a removed queue rule", func(t *testing.T) {
+		policy, err := NormalizeAndValidateRateLimitPolicy(map[string]any{
+			"rules": []any{
+				map[string]any{"metric": "rpm", "limit": 30},
+				map[string]any{"metric": "queue_size", "limit": 0},
+			},
+		})
+		if err != nil {
+			t.Fatalf("NormalizeAndValidateRateLimitPolicy() error = %v", err)
+		}
+		if _, ok := QueueRuleFromPolicy(policy); ok {
+			t.Fatalf("zero queue rule remained enabled: %+v", policy)
+		}
+		rules := policy["rules"].([]any)
+		if len(rules) != 1 {
+			t.Fatalf("rules length = %d, want 1", len(rules))
+		}
+	})
+
+	for _, test := range []struct {
+		name   string
+		policy map[string]any
+	}{
+		{name: "fractional queue", policy: map[string]any{"rules": []any{map[string]any{"metric": "queue_size", "limit": 1.5}}}},
+		{name: "queue too large", policy: map[string]any{"rules": []any{map[string]any{"metric": "queue_size", "limit": 10001}}}},
+		{name: "wait too large", policy: map[string]any{"rules": []any{map[string]any{"metric": "queue_size", "limit": 1, "maxWaitSeconds": 3601}}}},
+		{name: "wait on rpm", policy: map[string]any{"rules": []any{map[string]any{"metric": "rpm", "limit": 1, "maxWaitSeconds": 60}}}},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			if _, err := NormalizeAndValidateRateLimitPolicy(test.policy); err == nil {
+				t.Fatalf("invalid policy was accepted: %+v", test.policy)
+			}
+		})
+	}
+}
+
 func TestEffectiveRateLimitPolicyPrecedence(t *testing.T) {
 	policy := func(limit float64) map[string]any {
 		return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
@@ -42,6 +42,10 @@ type ModelRateLimitStatus struct {
 	ModelCooldownUntil        string                   `json:"modelCooldownUntil,omitempty"`
 	Concurrent                RateLimitMetricStatus    `json:"concurrent"`
 	QueuedTasks               float64                  `json:"queuedTasks"`
+	WaitingSyncTasks          int                      `json:"waitingSyncTasks,omitempty"`
+	WaitingAsyncTasks         int                      `json:"waitingAsyncTasks,omitempty"`
+	OldestQueueWaitSeconds    float64                  `json:"oldestQueueWaitSeconds,omitempty"`
+	QueueLimit                int                      `json:"queueLimit,omitempty"`
 	RPM                       RateLimitMetricStatus    `json:"rpm"`
 	TPM                       RateLimitMetricStatus    `json:"tpm"`
 	LoadRatio                 float64                  `json:"loadRatio"`
@@ -151,6 +155,9 @@ func (s *Store) ListModelRateLimitStatuses(ctx context.Context) ([]ModelRateLimi
 	       COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
 	       COALESCE(con.active, 0)::float8,
 	       COALESCE(queued.waiting, 0)::float8,
+	       COALESCE(admission.waiting_sync, 0)::int,
+	       COALESCE(admission.waiting_async, 0)::int,
+	       COALESCE(admission.oldest_wait_seconds, 0)::float8,
 	       COALESCE(rpm.used_value, 0)::float8, COALESCE(rpm.reserved_value, 0)::float8, COALESCE(rpm.reset_at::text, ''),
 	       COALESCE(tpm.used_value, 0)::float8, COALESCE(tpm.reserved_value, 0)::float8, COALESCE(tpm.reset_at::text, '')
 FROM platform_models m
@@ -192,6 +199,15 @@ LEFT JOIN (
  ) queued_sources
  GROUP BY queued_sources.platform_model_id
 ) queued ON queued.platform_model_id = m.id::text
+LEFT JOIN (
+  SELECT platform_model_id::text AS platform_model_id,
+         COUNT(*) FILTER (WHERE mode = 'sync') AS waiting_sync,
+         COUNT(*) FILTER (WHERE mode = 'async') AS waiting_async,
+         EXTRACT(EPOCH FROM now() - MIN(enqueued_at)) AS oldest_wait_seconds
+  FROM gateway_task_admissions
+  WHERE status = 'waiting'
+  GROUP BY platform_model_id
+) admission ON admission.platform_model_id = m.id::text
 LEFT JOIN (
  SELECT DISTINCT ON (scope_key) scope_key, used_value, reserved_value, reset_at
  FROM gateway_rate_limit_counters
@@ -231,6 +247,9 @@ ORDER BY p.priority ASC, m.model_name ASC`)
 		var modelCooldownUntil string
 		var concurrentCurrent float64
 		var queuedTasks float64
+		var waitingSyncTasks int
+		var waitingAsyncTasks int
+		var oldestQueueWaitSeconds float64
 		var rpmUsed float64
 		var rpmReserved float64
 		var rpmResetAt string
@@ -263,6 +282,9 @@ ORDER BY p.priority ASC, m.model_name ASC`)
 			&modelCooldownUntil,
 			&concurrentCurrent,
 			&queuedTasks,
+			&waitingSyncTasks,
+			&waitingAsyncTasks,
+			&oldestQueueWaitSeconds,
 			&rpmUsed,
 			&rpmReserved,
 			&rpmResetAt,
@@ -287,6 +309,12 @@ ORDER BY p.priority ASC, m.model_name ASC`)
 		item.ModelCooldownUntil = modelCooldownUntil
 		item.RateLimitPolicy = policy
 		item.QueuedTasks = queuedTasks
+		item.WaitingSyncTasks = waitingSyncTasks
+		item.WaitingAsyncTasks = waitingAsyncTasks
+		item.OldestQueueWaitSeconds = oldestQueueWaitSeconds
+		if queueRule, ok := QueueRuleFromPolicy(policy); ok {
+			item.QueueLimit = queueRule.Limit
+		}
 		item.Concurrent = metricStatus(concurrentCurrent, concurrentCurrent, 0, rateLimitForMetric(policy, "concurrent"), "")
 		item.RPM = metricStatus(rpmUsed+rpmReserved, rpmUsed, rpmReserved, rateLimitForMetric(policy, "rpm"), rpmResetAt)
 		item.TPM = metricStatus(tpmUsed+tpmReserved, tpmUsed, tpmReserved, tpmLimit(policy), tpmResetAt)
@@ -20,6 +20,84 @@ type RuntimeRecoveryResult struct {

 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 {
@@ -30,6 +108,9 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
 	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
 		}
@@ -377,15 +458,28 @@ func (s *Store) RecoverInterruptedRuntimeState(ctx context.Context) (RuntimeReco
 		return RuntimeRecoveryResult{}, err
 	}
 	defer tx.Rollback(ctx)
+	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-runtime-recovery', 0))`); err != nil {
+		return RuntimeRecoveryResult{}, err
+	}

 	result := RuntimeRecoveryResult{}
 	rows, err := tx.Query(ctx, `
-UPDATE gateway_rate_limit_reservations
+UPDATE gateway_rate_limit_reservations reservation
 SET status = 'released',
-    reason = 'server_restarted',
+    reason = 'execution_lease_expired',
    finalized_at = now(),
    updated_at = now()
-WHERE status = 'reserved'
+FROM gateway_tasks task
+WHERE reservation.task_id = task.id
+  AND reservation.status = 'reserved'
+  AND (
+    task.status IN ('succeeded', 'failed', 'cancelled')
+    OR (
+      task.status IN ('queued', 'running')
+      AND task.execution_lease_expires_at IS NOT NULL
+      AND task.execution_lease_expires_at <= now()
+    )
+  )
 RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`)
 	if err != nil {
 		return RuntimeRecoveryResult{}, err
@@ -415,7 +509,7 @@ RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`)
 UPDATE gateway_concurrency_leases
 SET released_at = now()
 WHERE released_at IS NULL
-  AND expires_at > now()`)
+  AND expires_at <= now()`)
 	if err != nil {
 		return RuntimeRecoveryResult{}, err
 	}
@@ -425,10 +519,22 @@ WHERE released_at IS NULL
 UPDATE gateway_task_attempts
 SET status = 'failed',
    retryable = true,
-    error_code = 'server_restarted',
-    error_message = 'attempt interrupted by service restart',
+    error_code = 'execution_lease_expired',
+    error_message = 'attempt execution lease expired',
    finished_at = now()
-WHERE status = 'running'`)
+WHERE status = 'running'
+  AND EXISTS (
+    SELECT 1
+    FROM gateway_tasks task
+    WHERE task.id = gateway_task_attempts.task_id
+      AND (
+        task.status IN ('succeeded', 'failed', 'cancelled')
+        OR (
+          task.execution_lease_expires_at IS NOT NULL
+          AND task.execution_lease_expires_at <= now()
+        )
+      )
+  )`)
 	if err != nil {
 		return RuntimeRecoveryResult{}, err
 	}
@@ -448,6 +554,8 @@ SET status = 'queued',
    updated_at = now()
 WHERE async_mode = true
  AND status = 'running'
+  AND execution_lease_expires_at IS NOT NULL
+  AND execution_lease_expires_at <= now()
 RETURNING id::text`)
 	if err != nil {
 		return RuntimeRecoveryResult{}, err
@@ -516,6 +624,8 @@ SET status = 'failed',
    updated_at = now()
 WHERE async_mode = false
  AND status = 'running'
+  AND execution_lease_expires_at IS NOT NULL
+  AND execution_lease_expires_at <= now()
 RETURNING id::text`)
 	if err != nil {
 		return RuntimeRecoveryResult{}, err
@@ -76,6 +76,11 @@ func (s *Store) ListRuntimePolicySets(ctx context.Context) ([]RuntimePolicySet,

 func (s *Store) CreateRuntimePolicySet(ctx context.Context, input RuntimePolicySetInput) (RuntimePolicySet, error) {
 	input = normalizeRuntimePolicyInput(input)
+	normalizedPolicy, err := NormalizeAndValidateRateLimitPolicy(input.RateLimitPolicy)
+	if err != nil {
+		return RuntimePolicySet{}, err
+	}
+	input.RateLimitPolicy = normalizedPolicy
 	rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.RateLimitPolicy))
 	retryPolicy, _ := json.Marshal(emptyObjectIfNil(input.RetryPolicy))
 	autoDisablePolicy, _ := json.Marshal(emptyObjectIfNil(input.AutoDisablePolicy))
@@ -94,6 +99,11 @@ RETURNING `+runtimePolicyColumns,

 func (s *Store) UpdateRuntimePolicySet(ctx context.Context, id string, input RuntimePolicySetInput) (RuntimePolicySet, error) {
 	input = normalizeRuntimePolicyInput(input)
+	normalizedPolicy, err := NormalizeAndValidateRateLimitPolicy(input.RateLimitPolicy)
+	if err != nil {
+		return RuntimePolicySet{}, err
+	}
+	input.RateLimitPolicy = normalizedPolicy
 	rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.RateLimitPolicy))
 	retryPolicy, _ := json.Marshal(emptyObjectIfNil(input.RetryPolicy))
 	autoDisablePolicy, _ := json.Marshal(emptyObjectIfNil(input.AutoDisablePolicy))
@@ -78,6 +78,9 @@ type RateLimitExceededError struct {
 	Message       string
 	RetryAfter    time.Duration
 	Retryable     bool
+	Reason        string
+	QueueDepth    int
+	QueueLimit    int
 }

 func (e *RateLimitExceededError) Error() string {
@@ -90,6 +93,10 @@ func (e *RateLimitExceededError) Error() string {
 	return ErrRateLimited.Error()
 }

+func (e *RateLimitExceededError) ErrorCode() string {
+	return "rate_limit"
+}
+
 func (e *RateLimitExceededError) Unwrap() error {
 	return ErrRateLimited
 }
@@ -389,6 +389,82 @@ RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second
 	return task, nil
 }

+func (s *Store) ClaimTaskPreparation(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) {
+	if leaseTTL <= 0 {
+		leaseTTL = 5 * time.Minute
+	}
+	task, err := scanGatewayTask(s.pool.QueryRow(ctx, `
+UPDATE gateway_tasks
+SET execution_token = $2::uuid,
+    execution_lease_expires_at = now() + ($3::int * interval '1 second'),
+    locked_at = now(),
+    heartbeat_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid
+  AND status = 'queued'
+  AND next_run_at <= now()
+  AND (
+    execution_token IS NULL
+    OR execution_lease_expires_at IS NULL
+    OR execution_lease_expires_at <= now()
+  )
+RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
+	if errors.Is(err, pgx.ErrNoRows) {
+		return GatewayTask{}, ErrTaskExecutionLeaseUnavailable
+	}
+	return task, err
+}
+
+func (s *Store) ReleaseTaskPreparation(ctx context.Context, taskID string, executionToken string) error {
+	tag, err := s.pool.Exec(ctx, `
+UPDATE gateway_tasks
+SET execution_token = NULL,
+    execution_lease_expires_at = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    updated_at = now()
+WHERE id = $1::uuid
+  AND status = 'queued'
+  AND execution_token = $2::uuid`, taskID, executionToken)
+	if err != nil {
+		return err
+	}
+	if tag.RowsAffected() != 1 {
+		return ErrTaskExecutionLeaseLost
+	}
+	return nil
+}
+
+func (s *Store) FailQueuedTask(ctx context.Context, taskID string, code string, message string) (GatewayTask, error) {
+	tag, err := s.pool.Exec(ctx, `
+UPDATE gateway_tasks
+SET status = 'failed',
+    error = NULL,
+    error_code = NULLIF($2, ''),
+    error_message = NULLIF($3, ''),
+    billing_status = CASE
+      WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
+      ELSE 'released'
+    END,
+    billing_updated_at = now(),
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    execution_token = NULL,
+    execution_lease_expires_at = NULL,
+    finished_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid
+  AND status = 'queued'`, taskID, strings.TrimSpace(code), truncateUTF8Bytes(message, 2048))
+	if err != nil {
+		return GatewayTask{}, err
+	}
+	if tag.RowsAffected() != 1 {
+		return GatewayTask{}, ErrTaskExecutionFinished
+	}
+	return s.GetTask(ctx, taskID)
+}
+
 func (s *Store) RenewTaskExecutionLease(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) error {
 	if leaseTTL <= 0 {
 		leaseTTL = 5 * time.Minute
@@ -399,7 +475,7 @@ SET execution_lease_expires_at = now() + ($3::int * interval '1 second'),
    heartbeat_at = now(),
    updated_at = now()
 WHERE id = $1::uuid
-  AND status = 'running'
+  AND status IN ('queued', 'running')
  AND execution_token = $2::uuid`, taskID, executionToken, int(leaseTTL/time.Second))
 	if err != nil {
 		return err
@@ -407,7 +483,7 @@ WHERE id = $1::uuid
 	if tag.RowsAffected() != 1 {
 		var stillRunning bool
 		if err := s.pool.QueryRow(ctx, `
-SELECT COALESCE((SELECT status = 'running' FROM gateway_tasks WHERE id = $1::uuid), false)`, taskID).Scan(&stillRunning); err != nil {
+SELECT COALESCE((SELECT status IN ('queued', 'running') FROM gateway_tasks WHERE id = $1::uuid), false)`, taskID).Scan(&stillRunning); err != nil {
 			return err
 		}
 		if !stillRunning {
@@ -428,7 +504,7 @@ SET status = 'running',
 	    heartbeat_at = now(),
 	    updated_at = now()
 WHERE id = $1::uuid
-  AND status = 'running'
+  AND status IN ('queued', 'running')
  AND execution_token = $2::uuid`, taskID, executionToken, modelType)
 	if err != nil {
 		return err
@@ -492,7 +568,7 @@ SET status = 'queued',
    error_message = NULL,
    updated_at = now()
 WHERE id = $1::uuid
-  AND status = 'running'
+  AND status IN ('queued', 'running')
  AND execution_token = $2::uuid
 RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt, strings.TrimSpace(queueKey)))
 }
@@ -566,7 +642,14 @@ func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message str
 		message = "任务已取消"
 	}
 	message = truncateUTF8Bytes(message, 2048)
-	tag, err := s.pool.Exec(ctx, `
+	var task GatewayTask
+	changed := false
+	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
+		if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil {
+			return err
+		}
+		var err error
+		task, err = scanGatewayTask(tx.QueryRow(ctx, `
 UPDATE gateway_tasks
 SET status = 'cancelled',
    error = NULL,
@@ -576,22 +659,123 @@ SET status = 'cancelled',
    locked_by = NULL,
    locked_at = NULL,
    heartbeat_at = NULL,
+    execution_token = NULL,
+    execution_lease_expires_at = NULL,
    finished_at = now(),
    updated_at = now()
 WHERE id = $1::uuid
  AND status = 'queued'
-  AND COALESCE(remote_task_id, '') = ''`, taskID, message)
+  AND COALESCE(remote_task_id, '') = ''
+RETURNING `+gatewayTaskColumns, taskID, message))
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		changed = true
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_concurrency_leases
+SET released_at = now()
+WHERE task_id = $1::uuid
+  AND released_at IS NULL`, taskID); err != nil {
+			return err
+		}
+		if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
+			return err
+		}
+		return notifyTaskAdmissionTx(ctx, tx, taskID)
+	})
 	if err != nil {
 		return GatewayTask{}, false, err
 	}
-	if tag.RowsAffected() == 0 {
-		return GatewayTask{}, false, nil
+	return task, changed, nil
+}
+
+func (s *Store) CancelTaskBeforeUpstreamSubmission(
+	ctx context.Context,
+	taskID string,
+	executionToken string,
+	message string,
+) (GatewayTask, bool, error) {
+	message = strings.TrimSpace(message)
+	if message == "" {
+		message = "client disconnected before upstream submission"
 	}
-	task, err := s.GetTask(ctx, taskID)
+	message = truncateUTF8Bytes(message, 2048)
+	var task GatewayTask
+	changed := false
+	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
+		if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil {
+			return err
+		}
+		var err error
+		task, err = scanGatewayTask(tx.QueryRow(ctx, `
+UPDATE gateway_tasks task
+SET status = 'cancelled',
+    error = NULL,
+    error_code = 'client_disconnected',
+    error_message = NULLIF($3::text, ''),
+    billing_status = CASE
+      WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
+      ELSE 'released'
+    END,
+    billing_updated_at = now(),
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    execution_token = NULL,
+    execution_lease_expires_at = NULL,
+    finished_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid
+  AND status IN ('queued', 'running')
+  AND execution_token = NULLIF($2, '')::uuid
+  AND COALESCE(remote_task_id, '') = ''
+  AND NOT EXISTS (
+    SELECT 1
+    FROM gateway_task_attempts attempt
+    WHERE attempt.task_id = task.id
+      AND attempt.upstream_submission_status <> 'not_submitted'
+  )
+RETURNING `+gatewayTaskColumns, taskID, strings.TrimSpace(executionToken), message))
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		changed = true
+		if _, err := tx.Exec(ctx, `
+DELETE FROM gateway_task_param_preprocessing_logs log
+USING gateway_task_attempts attempt
+WHERE log.attempt_id = attempt.id
+  AND attempt.task_id = $1::uuid
+  AND attempt.upstream_submission_status = 'not_submitted'`, taskID); err != nil {
+			return err
+		}
+		if _, err := tx.Exec(ctx, `
+DELETE FROM gateway_task_attempts
+WHERE task_id = $1::uuid
+  AND upstream_submission_status = 'not_submitted'`, taskID); err != nil {
+			return err
+		}
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_concurrency_leases
+SET released_at = now()
+WHERE task_id = $1::uuid
+  AND released_at IS NULL`, taskID); err != nil {
+			return err
+		}
+		if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
+			return err
+		}
+		return notifyTaskAdmissionTx(ctx, tx, taskID)
+	})
 	if err != nil {
-		return GatewayTask{}, true, err
+		return GatewayTask{}, false, err
 	}
-	return task, true, nil
+	return task, changed, nil
 }

 // CancelSubmittedTask records a confirmed upstream cancellation. Callers must
@@ -1506,7 +1690,7 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
 	    finished_at = now(),
 	    updated_at = now()
 WHERE id = $1::uuid
-  AND status = 'running'
+  AND status IN ('queued', 'running')
  AND execution_token = $10::uuid`,
 			input.TaskID,
 			message,
@@ -0,0 +1,162 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+)
+
+const workerHeartbeatStaleAfter = 15 * time.Second
+
+type WorkerRegistrationInput struct {
+	InstanceID      string
+	PodUID          string
+	PodName         string
+	Site            string
+	Revision        string
+	DesiredCapacity int
+}
+
+type WorkerAllocation struct {
+	InstanceID      string
+	DesiredCapacity int
+	Allocated       int
+	ActiveInstances int
+	HeartbeatAt     time.Time
+}
+
+func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistrationInput) (WorkerAllocation, error) {
+	input.InstanceID = strings.TrimSpace(input.InstanceID)
+	if input.InstanceID == "" {
+		return WorkerAllocation{}, errors.New("worker instance ID is required")
+	}
+	if input.DesiredCapacity < 0 {
+		return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative")
+	}
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return WorkerAllocation{}, err
+	}
+	defer tx.Rollback(ctx)
+	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-worker-capacity', 0))`); err != nil {
+		return WorkerAllocation{}, err
+	}
+	if _, err := tx.Exec(ctx, `
+INSERT INTO gateway_worker_instances (
+  instance_id, pod_uid, pod_name, site, revision, status,
+  desired_capacity, allocated_capacity, started_at, heartbeat_at, updated_at
+)
+VALUES ($1, $2, $3, $4, $5, 'active', $6, 0, now(), now(), now())
+ON CONFLICT (instance_id) DO UPDATE
+SET pod_uid = EXCLUDED.pod_uid,
+    pod_name = EXCLUDED.pod_name,
+    site = EXCLUDED.site,
+    revision = EXCLUDED.revision,
+    status = 'active',
+    desired_capacity = EXCLUDED.desired_capacity,
+    heartbeat_at = now(),
+    updated_at = now()`,
+		input.InstanceID,
+		strings.TrimSpace(input.PodUID),
+		strings.TrimSpace(input.PodName),
+		strings.TrimSpace(input.Site),
+		strings.TrimSpace(input.Revision),
+		input.DesiredCapacity,
+	); err != nil {
+		return WorkerAllocation{}, err
+	}
+	if _, err := tx.Exec(ctx, `
+UPDATE gateway_worker_instances
+SET allocated_capacity = 0,
+    updated_at = now()
+WHERE status <> 'active'
+   OR heartbeat_at <= now() - $1::interval`, workerHeartbeatStaleAfter.String()); err != nil {
+		return WorkerAllocation{}, err
+	}
+
+	rows, err := tx.Query(ctx, `
+SELECT instance_id
+FROM gateway_worker_instances
+WHERE status = 'active'
+  AND heartbeat_at > now() - $1::interval
+ORDER BY instance_id ASC
+FOR UPDATE`, workerHeartbeatStaleAfter.String())
+	if err != nil {
+		return WorkerAllocation{}, err
+	}
+	activeIDs := make([]string, 0)
+	for rows.Next() {
+		var instanceID string
+		if err := rows.Scan(&instanceID); err != nil {
+			rows.Close()
+			return WorkerAllocation{}, err
+		}
+		activeIDs = append(activeIDs, instanceID)
+	}
+	if err := rows.Err(); err != nil {
+		rows.Close()
+		return WorkerAllocation{}, err
+	}
+	rows.Close()
+	if len(activeIDs) == 0 {
+		return WorkerAllocation{}, errors.New("worker registration was not active after heartbeat")
+	}
+
+	base := input.DesiredCapacity / len(activeIDs)
+	remainder := input.DesiredCapacity % len(activeIDs)
+	allocated := 0
+	for index, instanceID := range activeIDs {
+		capacity := base
+		if index < remainder {
+			capacity++
+		}
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_worker_instances
+SET desired_capacity = $2,
+    allocated_capacity = $3,
+    updated_at = now()
+WHERE instance_id = $1`, instanceID, input.DesiredCapacity, capacity); err != nil {
+			return WorkerAllocation{}, err
+		}
+		if instanceID == input.InstanceID {
+			allocated = capacity
+		}
+	}
+	var heartbeatAt time.Time
+	if err := tx.QueryRow(ctx, `
+SELECT heartbeat_at
+FROM gateway_worker_instances
+WHERE instance_id = $1`, input.InstanceID).Scan(&heartbeatAt); err != nil {
+		return WorkerAllocation{}, err
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return WorkerAllocation{}, err
+	}
+	return WorkerAllocation{
+		InstanceID:      input.InstanceID,
+		DesiredCapacity: input.DesiredCapacity,
+		Allocated:       allocated,
+		ActiveInstances: len(activeIDs),
+		HeartbeatAt:     heartbeatAt,
+	}, nil
+}
+
+func (s *Store) MarkWorkerDraining(ctx context.Context, instanceID string) error {
+	result, err := s.pool.Exec(ctx, `
+UPDATE gateway_worker_instances
+SET status = 'draining',
+    allocated_capacity = 0,
+    heartbeat_at = now(),
+    updated_at = now()
+WHERE instance_id = $1`, strings.TrimSpace(instanceID))
+	if err != nil {
+		return err
+	}
+	if result.RowsAffected() == 0 {
+		return pgx.ErrNoRows
+	}
+	return nil
+}