feat: improve model rate limit tracking

2026-05-12 03:22:29 +08:00
parent 05632172d0
commit ba850a06c6
25 changed files with 1223 additions and 96 deletions
@@ -3,6 +3,7 @@ package store
 import (
 	"context"
 	"fmt"
+	"strings"

 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
 )
@@ -38,6 +39,7 @@ WHERE p.status = 'enabled'
  AND m.enabled = true
  AND m.model_type @> jsonb_build_array($2)
  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)
    OR (
@@ -151,6 +153,11 @@ ORDER BY effective_priority ASC,
 		return nil, err
 	}
 	if len(items) == 0 {
+		if unavailableErr, err := s.modelCandidateCooldownError(ctx, model, modelType); err != nil {
+			return nil, err
+		} else if unavailableErr != nil {
+			return nil, unavailableErr
+		}
 		return nil, ErrNoModelCandidate
 	}
 	items, err = s.filterCandidatesByAccessRules(ctx, user, items)
@@ -162,3 +169,89 @@ ORDER BY effective_priority ASC,
 	}
 	return items, nil
 }
+
+func (s *Store) modelCandidateCooldownError(ctx context.Context, model string, modelType string) (error, error) {
+	rows, err := s.pool.Query(ctx, `
+SELECT p.name,
+       COALESCE(NULLIF(m.display_name, ''), NULLIF(m.model_alias, ''), m.model_name),
+       COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
+       GREATEST(COALESCE(EXTRACT(EPOCH FROM p.cooldown_until - now()), 0), 0)::float8,
+       COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
+       GREATEST(COALESCE(EXTRACT(EPOCH FROM m.cooldown_until - now()), 0), 0)::float8
+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
+WHERE p.status = 'enabled'
+  AND p.deleted_at IS NULL
+  AND m.enabled = true
+  AND m.model_type @> jsonb_build_array($2)
+  AND (
+    (COALESCE(m.model_alias, '') <> '' AND m.model_alias = $1)
+    OR (
+      COALESCE(m.model_alias, '') = ''
+      AND (
+        m.model_name = $1
+        OR b.canonical_model_key = $1
+        OR b.provider_model_name = $1
+      )
+    )
+  )
+ORDER BY GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) DESC,
+         p.priority ASC,
+         m.created_at ASC`, model, modelType)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var platformName string
+		var displayName string
+		var platformCooldownUntil string
+		var platformRemainingSeconds float64
+		var modelCooldownUntil string
+		var modelRemainingSeconds float64
+		if err := rows.Scan(
+			&platformName,
+			&displayName,
+			&platformCooldownUntil,
+			&platformRemainingSeconds,
+			&modelCooldownUntil,
+			&modelRemainingSeconds,
+		); err != nil {
+			return nil, err
+		}
+		if modelRemainingSeconds > 0 {
+			return &ModelCandidateUnavailableError{
+				Code:    "model_cooling_down",
+				Message: cooldownErrorMessage("模型", displayName, modelRemainingSeconds, modelCooldownUntil),
+			}, nil
+		}
+		if platformRemainingSeconds > 0 {
+			return &ModelCandidateUnavailableError{
+				Code:    "platform_cooling_down",
+				Message: cooldownErrorMessage("平台", platformName, platformRemainingSeconds, platformCooldownUntil),
+			}, nil
+		}
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return nil, nil
+}
+
+func cooldownErrorMessage(scope string, name string, remainingSeconds float64, cooldownUntil string) string {
+	name = strings.TrimSpace(name)
+	if name == "" {
+		name = "候选"
+	}
+	remainingMinutes := remainingSeconds / 60
+	if remainingMinutes < 0.1 {
+		remainingMinutes = 0.1
+	}
+	message := fmt.Sprintf("%s %s 冷却中，剩余 %.1f 分钟", scope, name, remainingMinutes)
+	if strings.TrimSpace(cooldownUntil) != "" {
+		message += "，预计恢复时间 " + cooldownUntil
+	}
+	return message
+}
@@ -193,6 +193,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
          capabilities, pricing_mode, COALESCE(discount_factor, 0)::float8,
          COALESCE(pricing_rule_set_id::text, ''), billing_config_override, billing_config,
          permission_config, retry_policy, rate_limit_policy, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override,
+          COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
          enabled, created_at, updated_at`,
 		input.PlatformID,
 		baseID,
@@ -234,6 +235,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
 		&rateLimitPolicyBytes,
 		&model.RuntimePolicySetID,
 		&runtimePolicyOverrideBytes,
+		&model.CooldownUntil,
 		&model.Enabled,
 		&model.CreatedAt,
 		&model.UpdatedAt,
@@ -70,6 +70,7 @@ type Platform struct {
 	CredentialsPreview    map[string]any `json:"credentialsPreview,omitempty"`
 	RetryPolicy           map[string]any `json:"retryPolicy,omitempty"`
 	RateLimitPolicy       map[string]any `json:"rateLimitPolicy,omitempty"`
+	CooldownUntil         string         `json:"cooldownUntil,omitempty"`
 	CreatedAt             time.Time      `json:"createdAt"`
 	UpdatedAt             time.Time      `json:"updatedAt"`
 }
@@ -151,6 +152,7 @@ type PlatformModel struct {
 	RateLimitPolicy       map[string]any `json:"rateLimitPolicy,omitempty"`
 	RuntimePolicySetID    string         `json:"runtimePolicySetId,omitempty"`
 	RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride,omitempty"`
+	CooldownUntil         string         `json:"cooldownUntil,omitempty"`
 	Enabled               bool           `json:"enabled"`
 	CreatedAt             time.Time      `json:"createdAt"`
 	UpdatedAt             time.Time      `json:"updatedAt"`
@@ -478,7 +480,9 @@ func (s *Store) ListPlatforms(ctx context.Context) ([]Platform, error) {
 	rows, err := s.pool.Query(ctx, `
 SELECT id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status, priority,
       default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''),
-       config, credentials, retry_policy, rate_limit_policy, created_at, updated_at
+       config, credentials, retry_policy, rate_limit_policy,
+       COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
+       created_at, updated_at
 FROM integration_platforms
 ORDER BY priority ASC, created_at DESC`)
 	if err != nil {
@@ -510,6 +514,7 @@ ORDER BY priority ASC, created_at DESC`)
 			&credentialsBytes,
 			&retryPolicyBytes,
 			&rateLimitPolicyBytes,
+			&platform.CooldownUntil,
 			&platform.CreatedAt,
 			&platform.UpdatedAt,
 		); err != nil {
@@ -555,7 +560,9 @@ VALUES (
 )
 RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status, priority,
          default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''),
-          config, credentials, retry_policy, rate_limit_policy, created_at, updated_at`,
+          config, credentials, retry_policy, rate_limit_policy,
+          COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
+          created_at, updated_at`,
 		input.Provider, input.PlatformKey, input.Name, strings.TrimSpace(input.InternalName), input.BaseURL, input.AuthType, credentials, config,
 		input.DefaultPricingMode, input.DefaultDiscountFactor, input.PricingRuleSetID, input.Priority,
 		string(retryPolicy), string(rateLimitPolicy),
@@ -576,6 +583,7 @@ RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), C
 		&credentialsResultBytes,
 		&retryPolicyBytes,
 		&rateLimitPolicyBytes,
+		&platform.CooldownUntil,
 		&platform.CreatedAt,
 		&platform.UpdatedAt,
 	)
@@ -636,7 +644,9 @@ SET provider = $2,
 WHERE id = $1::uuid
 RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status, priority,
          default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''),
-          config, credentials, retry_policy, rate_limit_policy, created_at, updated_at`,
+          config, credentials, retry_policy, rate_limit_policy,
+          COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
+          created_at, updated_at`,
 		id,
 		input.Provider,
 		input.PlatformKey,
@@ -669,6 +679,7 @@ RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), C
 		&credentialsResultBytes,
 		&retryPolicyBytes,
 		&rateLimitPolicyBytes,
+		&platform.CooldownUntil,
 		&platform.CreatedAt,
 		&platform.UpdatedAt,
 	)
@@ -746,6 +757,7 @@ SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.p
       m.capability_override, m.capabilities, m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
       COALESCE(m.pricing_rule_set_id::text, ''), m.billing_config_override, m.billing_config,
       m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override,
+       COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
       m.enabled, m.created_at, m.updated_at
 FROM platform_models m
 JOIN integration_platforms p ON p.id = m.platform_id
@@ -791,6 +803,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
 			&rateLimitPolicy,
 			&model.RuntimePolicySetID,
 			&runtimePolicyOverride,
+			&model.CooldownUntil,
 			&model.Enabled,
 			&model.CreatedAt,
 			&model.UpdatedAt,
@@ -0,0 +1,249 @@
+package store
+
+import (
+	"context"
+	"sort"
+	"strings"
+)
+
+type RateLimitMetricStatus struct {
+	CurrentValue  float64 `json:"currentValue"`
+	UsedValue     float64 `json:"usedValue"`
+	ReservedValue float64 `json:"reservedValue"`
+	LimitValue    float64 `json:"limitValue"`
+	Limited       bool    `json:"limited"`
+	Ratio         float64 `json:"ratio"`
+	ResetAt       string  `json:"resetAt,omitempty"`
+}
+
+type ModelRateLimitStatus struct {
+	PlatformModelID       string                `json:"platformModelId"`
+	PlatformID            string                `json:"platformId"`
+	PlatformName          string                `json:"platformName"`
+	Provider              string                `json:"provider"`
+	ModelName             string                `json:"modelName"`
+	ProviderModelName     string                `json:"providerModelName,omitempty"`
+	ModelAlias            string                `json:"modelAlias,omitempty"`
+	DisplayName           string                `json:"displayName"`
+	ModelType             []string              `json:"modelType"`
+	Enabled               bool                  `json:"enabled"`
+	RateLimitPolicy       map[string]any        `json:"rateLimitPolicy,omitempty"`
+	PlatformCooldownUntil string                `json:"platformCooldownUntil,omitempty"`
+	ModelCooldownUntil    string                `json:"modelCooldownUntil,omitempty"`
+	Concurrent            RateLimitMetricStatus `json:"concurrent"`
+	RPM                   RateLimitMetricStatus `json:"rpm"`
+	TPM                   RateLimitMetricStatus `json:"tpm"`
+	LoadRatio             float64               `json:"loadRatio"`
+}
+
+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, '')
+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
+LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
+LEFT JOIN (
+  SELECT scope_key, SUM(lease_value) AS active
+  FROM gateway_concurrency_leases
+  WHERE scope_type = 'platform_model'
+    AND released_at IS NULL
+    AND expires_at > now()
+  GROUP BY scope_key
+) con ON con.scope_key = m.id::text
+LEFT JOIN (
+  SELECT DISTINCT ON (scope_key) scope_key, used_value, reserved_value, reset_at
+  FROM gateway_rate_limit_counters
+  WHERE scope_type = 'platform_model'
+    AND metric = 'rpm'
+    AND reset_at > now()
+  ORDER BY scope_key, window_start DESC
+) rpm ON rpm.scope_key = m.id::text
+LEFT JOIN (
+  SELECT scope_key, SUM(used_value) AS used_value, SUM(reserved_value) AS reserved_value, MAX(reset_at) AS reset_at
+  FROM gateway_rate_limit_counters
+  WHERE scope_type = 'platform_model'
+    AND metric LIKE 'tpm%'
+    AND reset_at > now()
+  GROUP BY scope_key
+) tpm ON tpm.scope_key = m.id::text
+WHERE p.deleted_at IS NULL
+ORDER BY p.priority ASC, m.model_name ASC`)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	items := make([]ModelRateLimitStatus, 0)
+	for rows.Next() {
+		var item ModelRateLimitStatus
+		var modelTypeBytes []byte
+		var platformPolicyBytes []byte
+		var runtimePolicyBytes []byte
+		var runtimeOverrideBytes []byte
+		var modelPolicyBytes []byte
+		var platformCooldownUntil string
+		var modelCooldownUntil string
+		var concurrentCurrent float64
+		var rpmUsed float64
+		var rpmReserved float64
+		var rpmResetAt string
+		var tpmUsed float64
+		var tpmReserved float64
+		var tpmResetAt string
+		if err := rows.Scan(
+			&item.PlatformModelID,
+			&item.PlatformID,
+			&item.PlatformName,
+			&item.Provider,
+			&item.ModelName,
+			&item.ProviderModelName,
+			&item.ModelAlias,
+			&modelTypeBytes,
+			&item.DisplayName,
+			&item.Enabled,
+			&platformPolicyBytes,
+			&runtimePolicyBytes,
+			&runtimeOverrideBytes,
+			&modelPolicyBytes,
+			&platformCooldownUntil,
+			&modelCooldownUntil,
+			&concurrentCurrent,
+			&rpmUsed,
+			&rpmReserved,
+			&rpmResetAt,
+			&tpmUsed,
+			&tpmReserved,
+			&tpmResetAt,
+		); err != nil {
+			return nil, err
+		}
+		item.ModelType = decodeStringArray(modelTypeBytes)
+		policy := effectiveModelRateLimitPolicy(
+			decodeObject(platformPolicyBytes),
+			decodeObject(runtimePolicyBytes),
+			decodeObject(runtimeOverrideBytes),
+			decodeObject(modelPolicyBytes),
+		)
+		item.PlatformCooldownUntil = platformCooldownUntil
+		item.ModelCooldownUntil = modelCooldownUntil
+		item.RateLimitPolicy = policy
+		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)
+		item.LoadRatio = maxFloat(item.Concurrent.Ratio, item.RPM.Ratio, item.TPM.Ratio)
+		items = append(items, item)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	sort.SliceStable(items, func(i, j int) bool {
+		if items[i].LoadRatio == items[j].LoadRatio {
+			return strings.ToLower(items[i].DisplayName) < strings.ToLower(items[j].DisplayName)
+		}
+		return items[i].LoadRatio > items[j].LoadRatio
+	})
+	return items, nil
+}
+
+func effectiveModelRateLimitPolicy(platformPolicy map[string]any, runtimePolicy map[string]any, runtimeOverride map[string]any, modelPolicy map[string]any) map[string]any {
+	policy := platformPolicy
+	if hasRateLimitRules(runtimePolicy) {
+		policy = shallowMergeMap(policy, runtimePolicy)
+	}
+	if nested, ok := runtimeOverride["rateLimitPolicy"].(map[string]any); ok && len(nested) > 0 {
+		policy = shallowMergeMap(policy, nested)
+	}
+	if hasRateLimitRules(modelPolicy) {
+		policy = shallowMergeMap(policy, modelPolicy)
+	}
+	if hasRateLimitRules(policy) {
+		return policy
+	}
+	return nil
+}
+
+func hasRateLimitRules(policy map[string]any) bool {
+	rules, _ := policy["rules"].([]any)
+	return len(rules) > 0
+}
+
+func shallowMergeMap(base map[string]any, override map[string]any) map[string]any {
+	out := map[string]any{}
+	for key, value := range base {
+		out[key] = value
+	}
+	for key, value := range override {
+		out[key] = value
+	}
+	return out
+}
+
+func rateLimitForMetric(policy map[string]any, metric string) float64 {
+	rules, _ := policy["rules"].([]any)
+	for _, rawRule := range rules {
+		rule, _ := rawRule.(map[string]any)
+		if strings.TrimSpace(stringValue(rule["metric"])) == metric {
+			return floatValue(rule["limit"])
+		}
+	}
+	return 0
+}
+
+func tpmLimit(policy map[string]any) float64 {
+	if limit := rateLimitForMetric(policy, "tpm_total"); limit > 0 {
+		return limit
+	}
+	return rateLimitForMetric(policy, "tpm_input") + rateLimitForMetric(policy, "tpm_output")
+}
+
+func metricStatus(current float64, used float64, reserved float64, limit float64, resetAt string) RateLimitMetricStatus {
+	status := RateLimitMetricStatus{
+		CurrentValue:  current,
+		UsedValue:     used,
+		ReservedValue: reserved,
+		LimitValue:    limit,
+		Limited:       limit > 0,
+		ResetAt:       resetAt,
+	}
+	if status.Limited {
+		status.Ratio = current / limit
+	}
+	return status
+}
+
+func maxFloat(values ...float64) float64 {
+	out := 0.0
+	for _, value := range values {
+		if value > out {
+			out = value
+		}
+	}
+	return out
+}
+
+func stringValue(value any) string {
+	text, _ := value.(string)
+	return strings.TrimSpace(text)
+}
+
+func floatValue(value any) float64 {
+	switch typed := value.(type) {
+	case int:
+		return float64(typed)
+	case int64:
+		return float64(typed)
+	case float64:
+		return typed
+	default:
+		return 0
+	}
+}
@@ -0,0 +1,32 @@
+package store
+
+import "testing"
+
+func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing.T) {
+	policy := effectiveModelRateLimitPolicy(
+		map[string]any{"rules": []any{
+			map[string]any{"metric": "rpm", "limit": 500},
+			map[string]any{"metric": "tpm_total", "limit": 100000},
+		}},
+		map[string]any{"rules": []any{
+			map[string]any{"metric": "rpm", "limit": 120},
+			map[string]any{"metric": "tpm_total", "limit": 240000},
+			map[string]any{"metric": "concurrent", "limit": 6},
+		}},
+		map[string]any{},
+		map[string]any{"rules": []any{
+			map[string]any{"metric": "rpm", "limit": 30},
+			map[string]any{"metric": "concurrent", "limit": 2},
+		}},
+	)
+
+	if got := rateLimitForMetric(policy, "rpm"); got != 30 {
+		t.Fatalf("expected model rpm limit to win, got %v", got)
+	}
+	if got := rateLimitForMetric(policy, "concurrent"); got != 2 {
+		t.Fatalf("expected model concurrent limit to win, got %v", got)
+	}
+	if got := rateLimitForMetric(policy, "tpm_total"); got != 0 {
+		t.Fatalf("expected missing model tpm limit to mean unlimited, got %v", got)
+	}
+}
@@ -3,9 +3,19 @@ package store
 import (
 	"context"
 	"errors"
+	"time"
+
+	"github.com/jackc/pgx/v5"
 )

-func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, reservations []RateLimitReservation) (RateLimitResult, error) {
+type RuntimeRecoveryResult struct {
+	ReleasedConcurrencyLeases int64 `json:"releasedConcurrencyLeases"`
+	ReleasedRateReservations  int64 `json:"releasedRateReservations"`
+	FailedAttempts            int64 `json:"failedAttempts"`
+	FailedTasks               int64 `json:"failedTasks"`
+}
+
+func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID string, reservations []RateLimitReservation) (RateLimitResult, error) {
 	tx, err := s.pool.Begin(ctx)
 	if err != nil {
 		return RateLimitResult{}, err
@@ -24,70 +34,122 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, reservatio
 			reservation.WindowSeconds = 60
 		}
 		if reservation.Metric == "concurrent" {
-			if reservation.LeaseTTLSeconds <= 0 {
-				reservation.LeaseTTLSeconds = 120
+			leaseID, err := reserveConcurrencyLease(ctx, tx, taskID, attemptID, reservation)
+			if err != nil {
+				return RateLimitResult{}, err
 			}
-			var active float64
-			if err := tx.QueryRow(ctx, `
+			result.LeaseIDs = append(result.LeaseIDs, leaseID)
+			continue
+		}
+		normalized, err := reserveCounterWindow(ctx, tx, taskID, attemptID, reservation)
+		if err != nil {
+			return RateLimitResult{}, err
+		}
+		result.Reservations = append(result.Reservations, normalized)
+	}
+	return result, tx.Commit(ctx)
+}
+
+func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (string, error) {
+	if reservation.LeaseTTLSeconds <= 0 {
+		reservation.LeaseTTLSeconds = 120
+	}
+	var active float64
+	if err := tx.QueryRow(ctx, `
 SELECT COALESCE(SUM(lease_value), 0)::float8
 FROM gateway_concurrency_leases
 WHERE scope_type = $1
  AND scope_key = $2
  AND released_at IS NULL
  AND expires_at > now()`,
-				reservation.ScopeType,
-				reservation.ScopeKey,
-			).Scan(&active); err != nil {
-				return RateLimitResult{}, err
-			}
-			if active+reservation.Amount > reservation.Limit {
-				return RateLimitResult{}, ErrRateLimited
-			}
-			var leaseID string
-			if err := tx.QueryRow(ctx, `
-INSERT INTO gateway_concurrency_leases (task_id, scope_type, scope_key, lease_value, expires_at)
-VALUES ($1::uuid, $2, $3, $4, now() + ($5::int * interval '1 second'))
+		reservation.ScopeType,
+		reservation.ScopeKey,
+	).Scan(&active); err != nil {
+		return "", err
+	}
+	if active+reservation.Amount > reservation.Limit {
+		return "", ErrRateLimited
+	}
+	var leaseID string
+	if err := tx.QueryRow(ctx, `
+INSERT INTO gateway_concurrency_leases (task_id, attempt_id, scope_type, scope_key, lease_value, expires_at)
+VALUES ($1::uuid, NULLIF($2, '')::uuid, $3, $4, $5, now() + ($6::int * interval '1 second'))
 RETURNING id::text`,
-				taskID,
-				reservation.ScopeType,
-				reservation.ScopeKey,
-				reservation.Amount,
-				reservation.LeaseTTLSeconds,
-			).Scan(&leaseID); err != nil {
-				return RateLimitResult{}, err
-			}
-			result.LeaseIDs = append(result.LeaseIDs, leaseID)
-			continue
-		}
-		tag, err := tx.Exec(ctx, `
+		taskID,
+		attemptID,
+		reservation.ScopeType,
+		reservation.ScopeKey,
+		reservation.Amount,
+		reservation.LeaseTTLSeconds,
+	).Scan(&leaseID); err != nil {
+		return "", err
+	}
+	return leaseID, nil
+}
+
+func reserveCounterWindow(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (RateLimitReservation, error) {
+	usedAmount := 0.0
+	reservedAmount := reservation.Amount
+	var windowStart time.Time
+	err := tx.QueryRow(ctx, `
 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, 0,
-  date_trunc('minute', now()) + ($6::int * interval '1 second')
+  $1, $2, $3, date_trunc('minute', now()), $4, $5, $6,
+  date_trunc('minute', now()) + ($7::int * interval '1 second')
 )
 ON CONFLICT (scope_type, scope_key, metric, window_start) DO UPDATE
 SET limit_value = EXCLUDED.limit_value,
    used_value = gateway_rate_limit_counters.used_value + EXCLUDED.used_value,
+    reserved_value = gateway_rate_limit_counters.reserved_value + EXCLUDED.reserved_value,
    reset_at = EXCLUDED.reset_at,
    updated_at = now()
-WHERE gateway_rate_limit_counters.used_value + EXCLUDED.used_value <= EXCLUDED.limit_value`,
-			reservation.ScopeType,
-			reservation.ScopeKey,
-			reservation.Metric,
-			reservation.Limit,
-			reservation.Amount,
-			reservation.WindowSeconds,
-		)
-		if err != nil {
-			return RateLimitResult{}, err
-		}
-		if tag.RowsAffected() == 0 {
-			return RateLimitResult{}, ErrRateLimited
+WHERE gateway_rate_limit_counters.used_value + gateway_rate_limit_counters.reserved_value + EXCLUDED.used_value + EXCLUDED.reserved_value <= EXCLUDED.limit_value
+RETURNING window_start`,
+		reservation.ScopeType,
+		reservation.ScopeKey,
+		reservation.Metric,
+		reservation.Limit,
+		usedAmount,
+		reservedAmount,
+		reservation.WindowSeconds,
+	).Scan(&windowStart)
+	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return RateLimitReservation{}, ErrRateLimited
 		}
+		return RateLimitReservation{}, err
 	}
-	return result, tx.Commit(ctx)
+	reservation.WindowStart = windowStart
+	if err := tx.QueryRow(ctx, `
+INSERT INTO gateway_rate_limit_reservations (
+  task_id, attempt_id, scope_type, scope_key, metric, window_start, limit_value, reserved_amount, status
+)
+VALUES (
+  $1::uuid, NULLIF($2, '')::uuid, $3, $4, $5, $6, $7, $8, 'reserved'
+)
+RETURNING id::text`,
+		taskID,
+		attemptID,
+		reservation.ScopeType,
+		reservation.ScopeKey,
+		reservation.Metric,
+		windowStart,
+		reservation.Limit,
+		reservedAmount,
+	).Scan(&reservation.ReservationID); err != nil {
+		return RateLimitReservation{}, err
+	}
+	return reservation, nil
+}
+
+func (s *Store) CommitRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64) error {
+	return s.finishRateLimitReservations(ctx, reservations, actualByMetric, "committed", "success")
+}
+
+func (s *Store) ReleaseRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, reason string) error {
+	return s.finishRateLimitReservations(ctx, reservations, nil, "released", reason)
 }

 func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leaseIDs []string) error {
@@ -107,3 +169,193 @@ WHERE id = $1::uuid AND released_at IS NULL`, leaseID); err != nil && !errors.Is
 	}
 	return nil
 }
+
+func (s *Store) RecoverInterruptedRuntimeState(ctx context.Context) (RuntimeRecoveryResult, error) {
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return RuntimeRecoveryResult{}, err
+	}
+	defer tx.Rollback(ctx)
+
+	result := RuntimeRecoveryResult{}
+	rows, err := tx.Query(ctx, `
+UPDATE gateway_rate_limit_reservations
+SET status = 'released',
+    reason = 'server_restarted',
+    finalized_at = now(),
+    updated_at = now()
+WHERE status = 'reserved'
+RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`)
+	if err != nil {
+		return RuntimeRecoveryResult{}, err
+	}
+	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++
+	}
+	if err := rows.Err(); err != nil {
+		rows.Close()
+		return RuntimeRecoveryResult{}, err
+	}
+	rows.Close()
+
+	tag, err := tx.Exec(ctx, `
+UPDATE gateway_concurrency_leases
+SET released_at = now()
+WHERE released_at IS NULL
+  AND expires_at > now()`)
+	if err != nil {
+		return RuntimeRecoveryResult{}, err
+	}
+	result.ReleasedConcurrencyLeases = tag.RowsAffected()
+
+	tag, err = tx.Exec(ctx, `
+UPDATE gateway_task_attempts
+SET status = 'failed',
+    retryable = false,
+    error_code = 'server_restarted',
+    error_message = 'attempt interrupted by service restart',
+    finished_at = now()
+WHERE status = 'running'`)
+	if err != nil {
+		return RuntimeRecoveryResult{}, err
+	}
+	result.FailedAttempts = tag.RowsAffected()
+
+	taskRows, err := tx.Query(ctx, `
+UPDATE gateway_tasks
+SET status = 'failed',
+    error = 'task interrupted by service restart',
+    error_code = 'server_restarted',
+    error_message = 'task interrupted by service restart',
+    finished_at = now(),
+    updated_at = now()
+WHERE status IN ('queued', 'running')
+RETURNING id::text`)
+	if err != nil {
+		return RuntimeRecoveryResult{}, err
+	}
+	taskIDs := make([]string, 0)
+	for taskRows.Next() {
+		var taskID string
+		if err := taskRows.Scan(&taskID); err != nil {
+			taskRows.Close()
+			return RuntimeRecoveryResult{}, err
+		}
+		taskIDs = append(taskIDs, taskID)
+	}
+	if err := taskRows.Err(); err != nil {
+		taskRows.Close()
+		return RuntimeRecoveryResult{}, err
+	}
+	taskRows.Close()
+
+	for _, taskID := range taskIDs {
+		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',
+  'failed',
+  'recovered',
+  1,
+  'task interrupted by service restart',
+  '{"code":"server_restarted"}'::jsonb,
+  false
+)`, taskID); err != nil {
+			return RuntimeRecoveryResult{}, err
+		}
+	}
+	result.FailedTasks = int64(len(taskIDs))
+
+	return result, tx.Commit(ctx)
+}
+
+func (s *Store) finishRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64, status string, reason string) error {
+	if len(reservations) == 0 {
+		return nil
+	}
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback(ctx)
+
+	for _, reservation := range reservations {
+		if reservation.ReservationID == "" {
+			continue
+		}
+		actualAmount := actualByMetric[reservation.Metric]
+		if status == "committed" && actualAmount <= 0 {
+			actualAmount = reservation.Amount
+		}
+		var stored RateLimitReservation
+		err := tx.QueryRow(ctx, `
+UPDATE gateway_rate_limit_reservations
+SET status = $2,
+    reason = NULLIF($3, ''),
+    actual_amount = CASE WHEN $2 = 'committed' THEN $4 ELSE actual_amount END,
+    finalized_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid
+  AND status = 'reserved'
+RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`,
+			reservation.ReservationID,
+			status,
+			reason,
+			actualAmount,
+		).Scan(&stored.ScopeType, &stored.ScopeKey, &stored.Metric, &stored.WindowStart, &stored.Amount)
+		if err != nil {
+			if errors.Is(err, pgx.ErrNoRows) {
+				continue
+			}
+			return err
+		}
+		if status == "committed" {
+			if err := commitCounterReservation(ctx, tx, stored.ScopeType, stored.ScopeKey, stored.Metric, stored.WindowStart, stored.Amount, actualAmount); err != nil {
+				return err
+			}
+			continue
+		}
+		if err := releaseCounterReservation(ctx, tx, stored.ScopeType, stored.ScopeKey, stored.Metric, stored.WindowStart, stored.Amount); err != nil {
+			return err
+		}
+	}
+	return tx.Commit(ctx)
+}
+
+func commitCounterReservation(ctx context.Context, tx pgx.Tx, scopeType string, scopeKey string, metric string, windowStart time.Time, reservedAmount float64, actualAmount float64) error {
+	_, err := tx.Exec(ctx, `
+UPDATE gateway_rate_limit_counters
+SET reserved_value = GREATEST(reserved_value - $5, 0),
+    used_value = used_value + $6,
+    updated_at = now()
+WHERE scope_type = $1
+  AND scope_key = $2
+  AND metric = $3
+  AND window_start = $4`,
+		scopeType, scopeKey, metric, windowStart, reservedAmount, actualAmount)
+	return err
+}
+
+func releaseCounterReservation(ctx context.Context, tx pgx.Tx, scopeType string, scopeKey string, metric string, windowStart time.Time, reservedAmount float64) error {
+	_, err := tx.Exec(ctx, `
+UPDATE gateway_rate_limit_counters
+SET reserved_value = GREATEST(reserved_value - $5, 0),
+    updated_at = now()
+WHERE scope_type = $1
+  AND scope_key = $2
+  AND metric = $3
+  AND window_start = $4`,
+		scopeType, scopeKey, metric, windowStart, reservedAmount)
+	return err
+}
@@ -134,6 +134,21 @@ WHERE id = $1::uuid`, platformID, cooldownSeconds)
 	return err
 }

+func (s *Store) CooldownCandidatePlatformModel(ctx context.Context, platformModelID string, cooldownSeconds int) error {
+	if strings.TrimSpace(platformModelID) == "" {
+		return nil
+	}
+	if cooldownSeconds <= 0 {
+		cooldownSeconds = 300
+	}
+	_, err := s.pool.Exec(ctx, `
+UPDATE platform_models
+SET cooldown_until = now() + ($2::int * interval '1 second'),
+    updated_at = now()
+WHERE id = $1::uuid`, platformModelID, cooldownSeconds)
+	return err
+}
+
 func (s *Store) DemoteCandidatePlatformPriority(ctx context.Context, platformID string, demoteStep int) error {
 	if strings.TrimSpace(platformID) == "" {
 		return nil
@@ -2,6 +2,7 @@ package store

 import (
 	"errors"
+	"strings"
 	"time"
 )

@@ -10,6 +11,27 @@ var (
 	ErrRateLimited      = errors.New("rate limit exceeded")
 )

+type ModelCandidateUnavailableError struct {
+	Code    string
+	Message string
+}
+
+func (e *ModelCandidateUnavailableError) Error() string {
+	return e.Message
+}
+
+func (e *ModelCandidateUnavailableError) Unwrap() error {
+	return ErrNoModelCandidate
+}
+
+func ModelCandidateErrorCode(err error) string {
+	var candidateErr *ModelCandidateUnavailableError
+	if errors.As(err, &candidateErr) && strings.TrimSpace(candidateErr.Code) != "" {
+		return candidateErr.Code
+	}
+	return "no_model_candidate"
+}
+
 type CreatePlatformModelInput struct {
 	PlatformID            string         `json:"platformId"`
 	BaseModelID           string         `json:"baseModelId"`
@@ -81,6 +103,7 @@ type RuntimeModelCandidate struct {
 }

 type RateLimitReservation struct {
+	ReservationID   string
 	ScopeType       string
 	ScopeKey        string
 	Metric          string
@@ -88,10 +111,12 @@ type RateLimitReservation struct {
 	Amount          float64
 	WindowSeconds   int
 	LeaseTTLSeconds int
+	WindowStart     time.Time
 }

 type RateLimitResult struct {
-	LeaseIDs []string
+	LeaseIDs     []string
+	Reservations []RateLimitReservation
 }

 type CreateTaskAttemptInput struct {