feat: add river-backed async task queue

2026-05-12 10:11:54 +08:00
parent d69aaed444
commit 7e220b7477
30 changed files with 1342 additions and 200 deletions
@@ -19,7 +19,7 @@ SELECT p.id::text, p.platform_key, p.name, p.provider,
       COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
       m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
       COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), m.model_name, COALESCE(m.model_alias, ''),
-       $2 AS requested_model_type, m.display_name, m.capabilities, m.capability_override,
+       $2::text AS requested_model_type, m.display_name, m.capabilities, m.capability_override,
       COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override,
       m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''),
       COALESCE(b.pricing_rule_set_id::text, ''),
@@ -33,21 +33,21 @@ LEFT JOIN model_catalog_providers cp ON cp.provider_key = p.provider OR cp.provi
 LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
 LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
 LEFT JOIN runtime_client_states s
-  ON s.client_id = p.platform_key || ':' || $2 || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
+  ON s.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
 WHERE p.status = 'enabled'
  AND p.deleted_at IS NULL
  AND m.enabled = true
-  AND m.model_type @> jsonb_build_array($2)
+  AND m.model_type @> jsonb_build_array($2::text)
  AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
  AND (m.cooldown_until IS NULL OR m.cooldown_until <= now())
  AND (
-    (COALESCE(m.model_alias, '') <> '' AND m.model_alias = $1)
+    (COALESCE(m.model_alias, '') <> '' AND m.model_alias = $1::text)
    OR (
      COALESCE(m.model_alias, '') = ''
      AND (
-        m.model_name = $1
-        OR b.canonical_model_key = $1
-        OR b.provider_model_name = $1
+        m.model_name = $1::text
+        OR b.canonical_model_key = $1::text
+        OR b.provider_model_name = $1::text
      )
    )
  )
@@ -184,15 +184,15 @@ LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
 WHERE p.status = 'enabled'
  AND p.deleted_at IS NULL
  AND m.enabled = true
-  AND m.model_type @> jsonb_build_array($2)
+  AND m.model_type @> jsonb_build_array($2::text)
  AND (
-    (COALESCE(m.model_alias, '') <> '' AND m.model_alias = $1)
+    (COALESCE(m.model_alias, '') <> '' AND m.model_alias = $1::text)
    OR (
      COALESCE(m.model_alias, '') = ''
      AND (
-        m.model_name = $1
-        OR b.canonical_model_key = $1
-        OR b.provider_model_name = $1
+        m.model_name = $1::text
+        OR b.canonical_model_key = $1::text
+        OR b.provider_model_name = $1::text
      )
    )
  )
@@ -53,6 +53,10 @@ func (s *Store) Ping(ctx context.Context) error {
 	return s.pool.Ping(ctx)
 }

+func (s *Store) Pool() *pgxpool.Pool {
+	return s.pool
+}
+
 type Platform struct {
 	ID                    string         `json:"id"`
 	Provider              string         `json:"provider"`
@@ -374,6 +378,7 @@ type CreateTaskInput struct {
 	Kind    string         `json:"kind"`
 	Model   string         `json:"model"`
 	RunMode string         `json:"runMode"`
+	Async   bool           `json:"async"`
 	Request map[string]any `json:"request"`
 }

@@ -398,7 +403,12 @@ type GatewayTask struct {
 	ResolvedModel      string         `json:"resolvedModel,omitempty"`
 	RequestID          string         `json:"requestId,omitempty"`
 	Request            map[string]any `json:"request,omitempty"`
+	AsyncMode          bool           `json:"asyncMode"`
+	RiverJobID         int64          `json:"riverJobId,omitempty"`
 	Status             string         `json:"status"`
+	AttemptCount       int            `json:"attemptCount"`
+	RemoteTaskID       string         `json:"remoteTaskId,omitempty"`
+	RemoteTaskPayload  map[string]any `json:"remoteTaskPayload,omitempty"`
 	Result             map[string]any `json:"result,omitempty"`
 	Billings           []any          `json:"billings,omitempty"`
 	Usage              map[string]any `json:"usage"`
@@ -423,7 +433,9 @@ COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_
 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, ''),
-request, status, COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
+request, COALESCE(async_mode, false), COALESCE(river_job_id, 0), status, 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),
 COALESCE(final_charge_amount, 0)::float8, COALESCE(response_started_at::text, ''),
 COALESCE(response_finished_at::text, ''), COALESCE(response_duration_ms, 0), COALESCE(error, ''),
@@ -1675,11 +1687,11 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
 	INSERT INTO gateway_tasks (
 	  kind, run_mode, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
 	  api_key_id, api_key_name, api_key_prefix, user_group_id, user_group_key,
-	  model, requested_model, request, status, result, billings, finished_at
+	  model, requested_model, request, async_mode, status, result, billings, finished_at
 	)
-	VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17::jsonb, $18::jsonb, CASE WHEN $19 THEN now() ELSE NULL END)
+	VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17, $18::jsonb, $19::jsonb, CASE WHEN $20 THEN now() ELSE NULL END)
 	RETURNING `+gatewayTaskColumns,
-		input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, status, resultBody, billingsBody, false,
+		input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, false,
 	))
 	if err != nil {
 		return GatewayTask{}, err
@@ -1689,7 +1701,7 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
 		payload, _ := json.Marshal(event.Payload)
 		if _, err := tx.Exec(ctx, `
 INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
-VALUES ($1::uuid, $2, $3, NULLIF($4, ''), NULLIF($5, ''), $6, NULLIF($7, ''), $8::jsonb, $9)`,
+VALUES ($1::uuid, $2, $3::text, NULLIF($4::text, ''), NULLIF($5::text, ''), $6, NULLIF($7::text, ''), $8::jsonb, $9)`,
 			task.ID, event.Seq, event.EventType, event.Status, event.Phase, event.Progress, event.Message, string(payload), event.Simulated,
 		); err != nil {
 			return GatewayTask{}, err
@@ -1730,6 +1742,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
 	var usageBytes []byte
 	var metricsBytes []byte
 	var billingSummaryBytes []byte
+	var remoteTaskPayloadBytes []byte
 	if err := scanner.Scan(
 		&task.ID,
 		&task.Kind,
@@ -1751,7 +1764,12 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
 		&task.ResolvedModel,
 		&task.RequestID,
 		&requestBytes,
+		&task.AsyncMode,
+		&task.RiverJobID,
 		&task.Status,
+		&task.AttemptCount,
+		&task.RemoteTaskID,
+		&remoteTaskPayloadBytes,
 		&resultBytes,
 		&billingsBytes,
 		&usageBytes,
@@ -1771,6 +1789,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
 		return GatewayTask{}, err
 	}
 	task.Request = decodeObject(requestBytes)
+	task.RemoteTaskPayload = decodeObject(remoteTaskPayloadBytes)
 	task.Result = decodeObject(resultBytes)
 	task.Billings = decodeArray(billingsBytes)
 	task.Usage = decodeObject(usageBytes)
@@ -31,6 +31,7 @@ type ModelRateLimitStatus struct {
 	PlatformCooldownUntil string                `json:"platformCooldownUntil,omitempty"`
 	ModelCooldownUntil    string                `json:"modelCooldownUntil,omitempty"`
 	Concurrent            RateLimitMetricStatus `json:"concurrent"`
+	QueuedTasks           float64               `json:"queuedTasks"`
 	RPM                   RateLimitMetricStatus `json:"rpm"`
 	TPM                   RateLimitMetricStatus `json:"tpm"`
 	LoadRatio             float64               `json:"loadRatio"`
@@ -38,15 +39,16 @@ type ModelRateLimitStatus struct {

 func (s *Store) ListModelRateLimitStatuses(ctx context.Context) ([]ModelRateLimitStatus, error) {
 	rows, err := s.pool.Query(ctx, `
-SELECT m.id::text, m.platform_id::text, p.name, p.provider,
-       m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''),
-       m.model_type, m.display_name, m.enabled,
-       p.rate_limit_policy, COALESCE(rp.rate_limit_policy, '{}'::jsonb), COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy,
-       COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
-       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(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, '')
+	SELECT m.id::text, m.platform_id::text, p.name, p.provider,
+	       m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''),
+	       m.model_type, m.display_name, m.enabled,
+	       p.rate_limit_policy, COALESCE(rp.rate_limit_policy, '{}'::jsonb), COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy,
+	       COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
+	       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(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
 JOIN integration_platforms p ON p.id = m.platform_id
 LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
@@ -59,6 +61,19 @@ LEFT JOIN (
    AND expires_at > now()
  GROUP BY scope_key
 ) con ON con.scope_key = m.id::text
+LEFT JOIN (
+  SELECT latest.platform_model_id, COUNT(*) AS waiting
+  FROM (
+    SELECT DISTINCT ON (a.task_id) a.task_id, a.platform_model_id::text AS platform_model_id
+    FROM gateway_tasks t
+    JOIN gateway_task_attempts a ON a.task_id = t.id
+    WHERE t.async_mode = true
+      AND t.status = 'queued'
+      AND a.platform_model_id IS NOT NULL
+    ORDER BY a.task_id, a.attempt_no DESC, a.started_at DESC
+  ) latest
+  GROUP BY latest.platform_model_id
+) queued ON queued.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
@@ -93,6 +108,7 @@ ORDER BY p.priority ASC, m.model_name ASC`)
 		var platformCooldownUntil string
 		var modelCooldownUntil string
 		var concurrentCurrent float64
+		var queuedTasks float64
 		var rpmUsed float64
 		var rpmReserved float64
 		var rpmResetAt string
@@ -117,6 +133,7 @@ ORDER BY p.priority ASC, m.model_name ASC`)
 			&platformCooldownUntil,
 			&modelCooldownUntil,
 			&concurrentCurrent,
+			&queuedTasks,
 			&rpmUsed,
 			&rpmReserved,
 			&rpmResetAt,
@@ -136,6 +153,7 @@ ORDER BY p.priority ASC, m.model_name ASC`)
 		item.PlatformCooldownUntil = platformCooldownUntil
 		item.ModelCooldownUntil = modelCooldownUntil
 		item.RateLimitPolicy = policy
+		item.QueuedTasks = queuedTasks
 		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)
@@ -3,6 +3,7 @@ package store
 import (
 	"context"
 	"errors"
+	"fmt"
 	"time"

 	"github.com/jackc/pgx/v5"
@@ -13,6 +14,7 @@ type RuntimeRecoveryResult struct {
 	ReleasedRateReservations  int64 `json:"releasedRateReservations"`
 	FailedAttempts            int64 `json:"failedAttempts"`
 	FailedTasks               int64 `json:"failedTasks"`
+	RequeuedAsyncTasks        int64 `json:"requeuedAsyncTasks"`
 }

 func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID string, reservations []RateLimitReservation) (RateLimitResult, error) {
@@ -28,7 +30,11 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
 			continue
 		}
 		if reservation.Metric == "" || reservation.Amount > reservation.Limit {
-			return RateLimitResult{}, ErrRateLimited
+			return RateLimitResult{}, &RateLimitExceededError{
+				Metric:    reservation.Metric,
+				Message:   fmt.Sprintf("rate limit exceeded: %s request amount %.0f is greater than limit %.0f", reservation.Metric, reservation.Amount, reservation.Limit),
+				Retryable: false,
+			}
 		}
 		if reservation.WindowSeconds <= 0 {
 			reservation.WindowSeconds = 60
@@ -55,8 +61,10 @@ func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, atte
 		reservation.LeaseTTLSeconds = 120
 	}
 	var active float64
+	var nextAvailableAt time.Time
 	if err := tx.QueryRow(ctx, `
-SELECT COALESCE(SUM(lease_value), 0)::float8
+SELECT COALESCE(SUM(lease_value), 0)::float8,
+       COALESCE(MIN(expires_at), now() + ($3::int * interval '1 second'))
 FROM gateway_concurrency_leases
 WHERE scope_type = $1
  AND scope_key = $2
@@ -64,11 +72,17 @@ WHERE scope_type = $1
  AND expires_at > now()`,
 		reservation.ScopeType,
 		reservation.ScopeKey,
-	).Scan(&active); err != nil {
+		reservation.LeaseTTLSeconds,
+	).Scan(&active, &nextAvailableAt); err != nil {
 		return "", err
 	}
 	if active+reservation.Amount > reservation.Limit {
-		return "", ErrRateLimited
+		return "", &RateLimitExceededError{
+			Metric:     reservation.Metric,
+			Message:    fmt.Sprintf("rate limit exceeded: concurrent active %.0f plus request %.0f is greater than limit %.0f", active, reservation.Amount, reservation.Limit),
+			RetryAfter: concurrencyRetryAfter(nextAvailableAt),
+			Retryable:  true,
+		}
 	}
 	var leaseID string
 	if err := tx.QueryRow(ctx, `
@@ -92,13 +106,16 @@ func reserveCounterWindow(ctx context.Context, tx pgx.Tx, taskID string, attempt
 	reservedAmount := reservation.Amount
 	var windowStart time.Time
 	err := tx.QueryRow(ctx, `
+WITH bounds AS (
+  SELECT
+    to_timestamp(floor(extract(epoch FROM now()) / $7::int) * $7::int) AS window_start,
+    to_timestamp(floor(extract(epoch FROM now()) / $7::int) * $7::int) + ($7::int * interval '1 second') AS reset_at
+)
 INSERT INTO gateway_rate_limit_counters (
  scope_type, scope_key, metric, window_start, limit_value, used_value, reserved_value, reset_at
 )
-VALUES (
-  $1, $2, $3, date_trunc('minute', now()), $4, $5, $6,
-  date_trunc('minute', now()) + ($7::int * interval '1 second')
-)
+SELECT $1, $2, $3, bounds.window_start, $4, $5, $6, bounds.reset_at
+FROM bounds
 ON CONFLICT (scope_type, scope_key, metric, window_start) DO UPDATE
 SET limit_value = EXCLUDED.limit_value,
    used_value = gateway_rate_limit_counters.used_value + EXCLUDED.used_value,
@@ -117,7 +134,28 @@ RETURNING window_start`,
 	).Scan(&windowStart)
 	if err != nil {
 		if errors.Is(err, pgx.ErrNoRows) {
-			return RateLimitReservation{}, ErrRateLimited
+			resetAt := time.Now().Add(time.Duration(reservation.WindowSeconds) * time.Second)
+			_ = tx.QueryRow(ctx, `
+WITH bounds AS (
+  SELECT to_timestamp(floor(extract(epoch FROM now()) / $4::int) * $4::int) AS window_start
+)
+SELECT counters.reset_at
+FROM gateway_rate_limit_counters counters
+JOIN bounds ON counters.window_start = bounds.window_start
+WHERE scope_type = $1
+  AND scope_key = $2
+  AND metric = $3`,
+				reservation.ScopeType,
+				reservation.ScopeKey,
+				reservation.Metric,
+				reservation.WindowSeconds,
+			).Scan(&resetAt)
+			return RateLimitReservation{}, &RateLimitExceededError{
+				Metric:     reservation.Metric,
+				Message:    fmt.Sprintf("rate limit exceeded: %s window has no remaining capacity", reservation.Metric),
+				RetryAfter: retryAfterUntil(resetAt),
+				Retryable:  true,
+			}
 		}
 		return RateLimitReservation{}, err
 	}
@@ -144,6 +182,28 @@ RETURNING id::text`,
 	return reservation, nil
 }

+func retryAfterUntil(when time.Time) time.Duration {
+	if when.IsZero() {
+		return 0
+	}
+	duration := time.Until(when)
+	if duration < time.Second {
+		return time.Second
+	}
+	return duration
+}
+
+func concurrencyRetryAfter(leaseExpiresAt time.Time) time.Duration {
+	if leaseExpiresAt.IsZero() {
+		return time.Second
+	}
+	duration := time.Until(leaseExpiresAt)
+	if duration <= time.Second {
+		return time.Second
+	}
+	return time.Second
+}
+
 func (s *Store) CommitRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64) error {
 	return s.finishRateLimitReservations(ctx, reservations, actualByMetric, "committed", "success")
 }
@@ -189,23 +249,26 @@ RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`)
 	if err != nil {
 		return RuntimeRecoveryResult{}, err
 	}
+	reservations := make([]RateLimitReservation, 0)
 	for rows.Next() {
 		var reservation RateLimitReservation
 		if err := rows.Scan(&reservation.ScopeType, &reservation.ScopeKey, &reservation.Metric, &reservation.WindowStart, &reservation.Amount); err != nil {
 			rows.Close()
 			return RuntimeRecoveryResult{}, err
 		}
-		if err := releaseCounterReservation(ctx, tx, reservation.ScopeType, reservation.ScopeKey, reservation.Metric, reservation.WindowStart, reservation.Amount); err != nil {
-			rows.Close()
-			return RuntimeRecoveryResult{}, err
-		}
-		result.ReleasedRateReservations++
+		reservations = append(reservations, reservation)
 	}
 	if err := rows.Err(); err != nil {
 		rows.Close()
 		return RuntimeRecoveryResult{}, err
 	}
 	rows.Close()
+	for _, reservation := range reservations {
+		if err := releaseCounterReservation(ctx, tx, reservation.ScopeType, reservation.ScopeKey, reservation.Metric, reservation.WindowStart, reservation.Amount); err != nil {
+			return RuntimeRecoveryResult{}, err
+		}
+	}
+	result.ReleasedRateReservations = int64(len(reservations))

 	tag, err := tx.Exec(ctx, `
 UPDATE gateway_concurrency_leases
@@ -220,7 +283,7 @@ WHERE released_at IS NULL
 	tag, err = tx.Exec(ctx, `
 UPDATE gateway_task_attempts
 SET status = 'failed',
-    retryable = false,
+    retryable = true,
    error_code = 'server_restarted',
    error_message = 'attempt interrupted by service restart',
    finished_at = now()
@@ -230,6 +293,57 @@ WHERE status = 'running'`)
 	}
 	result.FailedAttempts = tag.RowsAffected()

+	asyncTaskRows, err := tx.Query(ctx, `
+UPDATE gateway_tasks
+SET status = 'queued',
+    error = NULL,
+    error_code = NULL,
+    error_message = NULL,
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    next_run_at = now(),
+    finished_at = NULL,
+    updated_at = now()
+WHERE async_mode = true
+  AND status = 'running'
+RETURNING id::text`)
+	if err != nil {
+		return RuntimeRecoveryResult{}, err
+	}
+	asyncTaskIDs := make([]string, 0)
+	for asyncTaskRows.Next() {
+		var taskID string
+		if err := asyncTaskRows.Scan(&taskID); err != nil {
+			asyncTaskRows.Close()
+			return RuntimeRecoveryResult{}, err
+		}
+		asyncTaskIDs = append(asyncTaskIDs, taskID)
+	}
+	if err := asyncTaskRows.Err(); err != nil {
+		asyncTaskRows.Close()
+		return RuntimeRecoveryResult{}, err
+	}
+	asyncTaskRows.Close()
+	for _, taskID := range asyncTaskIDs {
+		if _, err := tx.Exec(ctx, `
+INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
+VALUES (
+  $1::uuid,
+  COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1),
+  'task.recovered',
+  'queued',
+  'recovered',
+  0.2,
+  'async task recovered after service restart',
+  '{"code":"server_restarted"}'::jsonb,
+  false
+)`, taskID); err != nil {
+			return RuntimeRecoveryResult{}, err
+		}
+	}
+	result.RequeuedAsyncTasks = int64(len(asyncTaskIDs))
+
 	taskRows, err := tx.Query(ctx, `
 UPDATE gateway_tasks
 SET status = 'failed',
@@ -238,7 +352,8 @@ SET status = 'failed',
    error_message = 'task interrupted by service restart',
    finished_at = now(),
    updated_at = now()
-WHERE status IN ('queued', 'running')
+WHERE async_mode = false
+  AND status = 'running'
 RETURNING id::text`)
 	if err != nil {
 		return RuntimeRecoveryResult{}, err
@@ -301,9 +416,9 @@ func (s *Store) finishRateLimitReservations(ctx context.Context, reservations []
 		var stored RateLimitReservation
 		err := tx.QueryRow(ctx, `
 UPDATE gateway_rate_limit_reservations
-SET status = $2,
-    reason = NULLIF($3, ''),
-    actual_amount = CASE WHEN $2 = 'committed' THEN $4 ELSE actual_amount END,
+SET status = $2::text,
+    reason = NULLIF($3::text, ''),
+    actual_amount = CASE WHEN $2::text = 'committed' THEN $4 ELSE actual_amount END,
    finalized_at = now(),
    updated_at = now()
 WHERE id = $1::uuid
@@ -32,6 +32,43 @@ func ModelCandidateErrorCode(err error) string {
 	return "no_model_candidate"
 }

+type RateLimitExceededError struct {
+	Metric     string
+	Message    string
+	RetryAfter time.Duration
+	Retryable  bool
+}
+
+func (e *RateLimitExceededError) Error() string {
+	if strings.TrimSpace(e.Message) != "" {
+		return e.Message
+	}
+	if strings.TrimSpace(e.Metric) != "" {
+		return "rate limit exceeded for " + e.Metric
+	}
+	return ErrRateLimited.Error()
+}
+
+func (e *RateLimitExceededError) Unwrap() error {
+	return ErrRateLimited
+}
+
+func RateLimitRetryAfter(err error) time.Duration {
+	var limitErr *RateLimitExceededError
+	if errors.As(err, &limitErr) && limitErr.RetryAfter > 0 {
+		return limitErr.RetryAfter
+	}
+	return 0
+}
+
+func RateLimitRetryable(err error) bool {
+	var limitErr *RateLimitExceededError
+	if errors.As(err, &limitErr) {
+		return limitErr.Retryable
+	}
+	return errors.Is(err, ErrRateLimited)
+}
+
 type CreatePlatformModelInput struct {
 	PlatformID            string         `json:"platformId"`
 	BaseModelID           string         `json:"baseModelId"`
@@ -132,6 +169,11 @@ type CreateTaskAttemptInput struct {
 	Metrics         map[string]any
 }

+type AsyncTaskQueueItem struct {
+	ID       string
+	Priority int
+}
+
 type FinishTaskAttemptInput struct {
 	AttemptID          string
 	Status             string
@@ -157,7 +157,7 @@ func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType st
 	_, err := s.pool.Exec(ctx, `
 UPDATE gateway_tasks
 SET status = 'running',
-    model_type = NULLIF($2, ''),
+    model_type = NULLIF($2::text, ''),
    normalized_request = $3::jsonb,
    locked_at = now(),
    heartbeat_at = now(),
@@ -166,6 +166,124 @@ WHERE id = $1::uuid`, taskID, modelType, string(normalizedJSON))
 	return err
 }

+func (s *Store) ClaimAsyncQueuedTask(ctx context.Context, workerID string) (GatewayTask, error) {
+	return scanGatewayTask(s.pool.QueryRow(ctx, `
+WITH picked AS (
+  SELECT id AS task_id
+  FROM gateway_tasks
+  WHERE async_mode = true
+    AND status = 'queued'
+    AND next_run_at <= now()
+  ORDER BY priority ASC, created_at ASC
+  LIMIT 1
+  FOR UPDATE SKIP LOCKED
+)
+UPDATE gateway_tasks t
+SET status = 'running',
+    locked_by = NULLIF($1::text, ''),
+    locked_at = now(),
+    heartbeat_at = now(),
+    updated_at = now()
+FROM picked
+WHERE t.id = picked.task_id
+RETURNING `+gatewayTaskColumns, workerID))
+}
+
+func (s *Store) RequeueTask(ctx context.Context, taskID string, delay time.Duration) (GatewayTask, error) {
+	if delay < time.Second {
+		delay = time.Second
+	}
+	if delay > 10*time.Minute {
+		delay = 10 * time.Minute
+	}
+	nextRunAt := time.Now().Add(delay)
+	return scanGatewayTask(s.pool.QueryRow(ctx, `
+UPDATE gateway_tasks
+SET status = 'queued',
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    next_run_at = $2::timestamptz,
+    error = NULL,
+    error_code = NULL,
+    error_message = NULL,
+    updated_at = now()
+WHERE id = $1::uuid
+RETURNING `+gatewayTaskColumns, taskID, nextRunAt))
+}
+
+func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error {
+	if riverJobID <= 0 {
+		return nil
+	}
+	_, err := s.pool.Exec(ctx, `
+UPDATE gateway_tasks
+SET river_job_id = $2,
+    updated_at = now()
+WHERE id = $1::uuid`, taskID, riverJobID)
+	return err
+}
+
+func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, attemptID string, remoteTaskID string, payload map[string]any) error {
+	payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload))
+	return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_tasks
+SET remote_task_id = NULLIF($2::text, ''),
+    remote_task_payload = $3::jsonb,
+    updated_at = now()
+WHERE id = $1::uuid`,
+			taskID,
+			remoteTaskID,
+			string(payloadJSON),
+		); err != nil {
+			return err
+		}
+		if strings.TrimSpace(attemptID) == "" {
+			return nil
+		}
+		_, err := tx.Exec(ctx, `
+UPDATE gateway_task_attempts
+SET remote_task_id = NULLIF($2::text, ''),
+    response_snapshot = COALESCE(response_snapshot, '{}'::jsonb) || jsonb_build_object('remote_task_payload', $3::jsonb)
+WHERE id = $1::uuid`,
+			attemptID,
+			remoteTaskID,
+			string(payloadJSON),
+		)
+		return err
+	})
+}
+
+func (s *Store) ListRecoverableAsyncTasks(ctx context.Context, limit int) ([]AsyncTaskQueueItem, error) {
+	if limit <= 0 {
+		limit = 500
+	}
+	rows, err := s.pool.Query(ctx, `
+SELECT id::text, priority
+FROM gateway_tasks
+WHERE async_mode = true
+  AND status IN ('queued', 'running')
+ORDER BY priority ASC, created_at ASC
+LIMIT $1`, limit)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	items := make([]AsyncTaskQueueItem, 0)
+	for rows.Next() {
+		var item AsyncTaskQueueItem
+		if err := rows.Scan(&item.ID, &item.Priority); err != nil {
+			return nil, err
+		}
+		items = append(items, item)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
 func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) {
 	requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot))
 	metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
@@ -182,7 +300,7 @@ INSERT INTO gateway_task_attempts (
  status, simulated, request_snapshot, metrics
 )
 VALUES (
-  $1::uuid, $2, NULLIF($3, '')::uuid, NULLIF($4, '')::uuid, NULLIF($5, ''), $6,
+  $1::uuid, $2::int, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, NULLIF($5::text, ''), $6,
  $7, $8, $9::jsonb, $10::jsonb
 )
 RETURNING id::text`,
@@ -202,7 +320,7 @@ RETURNING id::text`,
 	}
 	if _, err := tx.Exec(ctx, `
 UPDATE gateway_tasks
-SET attempt_count = GREATEST(attempt_count, $2), updated_at = now()
+SET attempt_count = GREATEST(attempt_count, $2::int), updated_at = now()
 WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil {
 		return "", err
 	}
@@ -252,7 +370,7 @@ SET metrics = jsonb_set(
      true
    )
 WHERE task_id = $1::uuid
-  AND attempt_no = $2`, taskID, attemptNo, string(entryJSON))
+  AND attempt_no = $2::int`, taskID, attemptNo, string(entryJSON))
 	return err
 }

@@ -386,17 +504,17 @@ func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptIn
 	metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
 	_, err := s.pool.Exec(ctx, `
 UPDATE gateway_task_attempts
-SET status = $2,
+SET status = $2::text,
    retryable = $3,
-    request_id = NULLIF($4, ''),
+    request_id = NULLIF($4::text, ''),
    usage = $5::jsonb,
    metrics = $6::jsonb,
    response_snapshot = $7::jsonb,
    response_started_at = $8::timestamptz,
    response_finished_at = $9::timestamptz,
    response_duration_ms = $10,
-    error_code = NULLIF($11, ''),
-    error_message = NULLIF($12, ''),
+    error_code = NULLIF($11::text, ''),
+    error_message = NULLIF($12::text, ''),
    finished_at = now()
 WHERE id = $1::uuid`,
 		input.AttemptID,
@@ -438,6 +556,9 @@ SET status = 'succeeded',
    error = NULL,
    error_code = NULL,
    error_message = NULL,
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
    finished_at = now(),
    updated_at = now()
 WHERE id = $1::uuid`,
@@ -561,14 +682,17 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
 	if _, err := s.pool.Exec(ctx, `
 UPDATE gateway_tasks
 SET status = 'failed',
-    error = NULLIF($2, ''),
-    error_code = NULLIF($3, ''),
-    error_message = NULLIF($2, ''),
-    request_id = NULLIF($4, ''),
+    error = NULLIF($2::text, ''),
+    error_code = NULLIF($3::text, ''),
+    error_message = NULLIF($2::text, ''),
+    request_id = NULLIF($4::text, ''),
    metrics = $5::jsonb,
    response_started_at = $6::timestamptz,
    response_finished_at = $7::timestamptz,
    response_duration_ms = $8,
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
    finished_at = now(),
    updated_at = now()
 WHERE id = $1::uuid`,
@@ -604,7 +728,7 @@ WITH next_seq AS (
  WHERE task_id = $1::uuid
 )
 INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
-SELECT $1::uuid, next_seq.seq, $2, NULLIF($3, ''), NULLIF($4, ''), $5, NULLIF($6, ''), $7::jsonb, $8
+SELECT $1::uuid, next_seq.seq, $2::text, NULLIF($3::text, ''), NULLIF($4::text, ''), $5, NULLIF($6::text, ''), $7::jsonb, $8
 FROM next_seq
 RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''),
          COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at`,
@@ -688,7 +812,7 @@ func (s *Store) RecordClientRelease(ctx context.Context, clientID string, lastEr
 	_, err := s.pool.Exec(ctx, `
 UPDATE runtime_client_states
 SET running_count = GREATEST(running_count - 1, 0),
-    last_error = NULLIF($2, ''),
+    last_error = NULLIF($2::text, ''),
    updated_at = now()
 WHERE client_id = $1`, clientID, lastError)
 	return err