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

使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配，并将生产 API 与独立 Worker 角色拆分。

补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本；未默认启用任何生产 queue_size 策略。

已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收；合入最新主干后将重新执行发布门禁。

apps/api/internal/runner/admission.go

@@ -0,0 +1,432 @@
+package runner
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"time"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+)
+
+type taskAdmissionPlan struct {
+	Candidate store.RuntimeModelCandidate
+	Body      map[string]any
+	ModelType string
+	GroupID   string
+	Scopes    []store.AdmissionScope
+	Eligible  bool
+}
+
+func distributedAdmissionModelType(modelType string) bool {
+	switch strings.ToLower(strings.TrimSpace(modelType)) {
+	case "image_generate",
+		"image_edit",
+		"image_analysis",
+		"image_vectorize",
+		"video_generate",
+		"video_enhance",
+		"image_to_video",
+		"text_to_video",
+		"video_edit",
+		"video_reference",
+		"video_first_last_frame",
+		"video_understanding",
+		"omni_video",
+		"omni",
+		"audio_generate",
+		"audio_understanding",
+		"text_to_speech",
+		"voice_clone":
+		return true
+	default:
+		return false
+	}
+}
+
+func (s *Service) buildTaskAdmissionPlan(ctx context.Context, task store.GatewayTask, user *auth.User) (taskAdmissionPlan, error) {
+	restoredRequest, err := s.restoreTaskRequestReferences(ctx, task)
+	if err != nil {
+		return taskAdmissionPlan{}, err
+	}
+	body := normalizeRequest(task.Kind, restoredRequest)
+	modelType := modelTypeFromKind(task.Kind, body)
+	if !distributedAdmissionModelType(modelType) {
+		return taskAdmissionPlan{Body: body, ModelType: modelType}, nil
+	}
+	if err := validateRequest(task.Kind, body); err != nil {
+		return taskAdmissionPlan{}, parameterPreprocessClientError(err)
+	}
+	body, clonedVoice, err := s.resolveClonedVoiceBinding(ctx, user, task.Kind, body)
+	if err != nil {
+		return taskAdmissionPlan{}, err
+	}
+	runnerPolicy, err := s.store.GetActiveRunnerPolicy(ctx)
+	if err != nil {
+		return taskAdmissionPlan{}, err
+	}
+	cacheAffinityKeys := buildCacheAffinityKeys(task.Kind, modelType, body)
+	candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user, store.ListModelCandidatesOptions{
+		CacheAffinityKey:    cacheAffinityKeys.Primary,
+		CacheAffinityKeys:   cacheAffinityKeys.Lookup,
+		CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy,
+	})
+	if err == nil {
+		candidates, err = filterCandidatesByRequestedPlatform(candidates, body)
+	}
+	if err == nil {
+		candidates, err = filterCandidatesByClonedVoiceBinding(candidates, clonedVoice)
+	}
+	if err == nil {
+		candidates, _, err = filterRuntimeCandidatesByRequest(task.Kind, task.Model, modelType, body, candidates)
+	}
+	if err == nil {
+		candidates, _, err = filterRuntimeCandidatesByOutputTokens(task.Kind, task.Model, modelType, body, candidates)
+	}
+	if err != nil {
+		return taskAdmissionPlan{}, err
+	}
+	for _, candidate := range candidates {
+		available, availabilityErr := s.store.RuntimeCandidateAvailable(ctx, candidate.PlatformID, candidate.PlatformModelID)
+		if availabilityErr != nil {
+			return taskAdmissionPlan{}, availabilityErr
+		}
+		if !available {
+			continue
+		}
+		scopes, groupID := s.admissionScopes(ctx, user, candidate)
+		hasConcurrentLimit := false
+		for _, scope := range scopes {
+			if scope.ConcurrentLimit > 0 {
+				hasConcurrentLimit = true
+				break
+			}
+		}
+		if !hasConcurrentLimit {
+			return taskAdmissionPlan{
+				Candidate: candidate,
+				Body:      body,
+				ModelType: modelType,
+			}, nil
+		}
+		if err := s.store.CheckRateLimits(ctx, s.rateLimitReservations(ctx, user, candidate, body)); err != nil {
+			return taskAdmissionPlan{}, err
+		}
+		return taskAdmissionPlan{
+			Candidate: candidate,
+			Body:      body,
+			ModelType: modelType,
+			GroupID:   groupID,
+			Scopes:    scopes,
+			Eligible:  true,
+		}, nil
+	}
+	return taskAdmissionPlan{}, store.ErrNoModelCandidate
+}
+
+func (s *Service) tryTaskAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan, waiterID string) (store.TaskAdmissionResult, error) {
+	return s.tryTaskAdmissionWithAdmittedHook(ctx, task, plan, waiterID, nil)
+}
+
+func (s *Service) tryTaskAdmissionWithAdmittedHook(
+	ctx context.Context,
+	task store.GatewayTask,
+	plan taskAdmissionPlan,
+	waiterID string,
+	onAdmitted func(pgx.Tx) error,
+) (store.TaskAdmissionResult, error) {
+	mode := "sync"
+	if task.AsyncMode {
+		mode = "async"
+	}
+	input := store.TaskAdmissionInput{
+		TaskID:          task.ID,
+		PlatformID:      plan.Candidate.PlatformID,
+		PlatformModelID: plan.Candidate.PlatformModelID,
+		UserGroupID:     plan.GroupID,
+		QueueKey:        plan.Candidate.QueueKey,
+		Mode:            mode,
+		Priority:        task.Priority,
+		WaiterID:        waiterID,
+		Scopes:          plan.Scopes,
+	}
+	current, currentErr := s.store.GetTaskAdmission(ctx, task.ID)
+	if currentErr != nil && !errors.Is(currentErr, pgx.ErrNoRows) {
+		return store.TaskAdmissionResult{}, currentErr
+	}
+	if currentErr == nil && current.Status == "waiting" &&
+		(current.PlatformID != input.PlatformID || current.PlatformModelID != input.PlatformModelID || current.UserGroupID != input.UserGroupID) {
+		if _, rebindErr := s.store.RebindWaitingTaskAdmission(ctx, input); rebindErr != nil {
+			return store.TaskAdmissionResult{}, rebindErr
+		}
+		s.observeTaskAdmission("candidate_migrated")
+	}
+	var result store.TaskAdmissionResult
+	var err error
+	if onAdmitted == nil {
+		result, err = s.store.TryTaskAdmission(ctx, input)
+	} else {
+		result, err = s.store.TryTaskAdmissionWithAdmittedHook(ctx, input, onAdmitted)
+	}
+	if result.NewlyAdmitted {
+		s.observeTaskAdmission("admitted")
+		s.observeTaskAdmissionWait(time.Since(result.Admission.EnqueuedAt))
+	}
+	var limitErr *store.RateLimitExceededError
+	switch {
+	case errors.As(err, &limitErr) && limitErr.Reason == "queue_full":
+		s.observeTaskAdmission("queue_full")
+	case errors.Is(err, store.ErrQueueTimeout):
+		s.observeTaskAdmission("timeout")
+	}
+	return result, err
+}
+
+func (s *Service) waitForSynchronousAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan) (store.TaskAdmissionResult, error) {
+	waiterID := uuid.NewString()
+	return s.waitForTaskAdmission(ctx, task, plan, waiterID)
+}
+
+func (s *Service) waitForTaskAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan, waiterID string) (store.TaskAdmissionResult, error) {
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+	for {
+		result, err := s.tryTaskAdmission(ctx, task, plan, waiterID)
+		if err != nil || result.Admitted {
+			return result, err
+		}
+		select {
+		case <-ctx.Done():
+			_ = s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID)
+			s.observeTaskAdmission("cancelled")
+			return store.TaskAdmissionResult{}, ctx.Err()
+		case <-s.currentAdmissionWake():
+		case <-ticker.C:
+		}
+	}
+}
+
+func (s *Service) ensureCandidateAdmission(
+	ctx context.Context,
+	task store.GatewayTask,
+	user *auth.User,
+	body map[string]any,
+	candidate store.RuntimeModelCandidate,
+) (store.TaskAdmissionResult, bool, error) {
+	scopes, groupID := s.admissionScopes(ctx, user, candidate)
+	hasConcurrentLimit := false
+	for _, scope := range scopes {
+		if scope.ConcurrentLimit > 0 {
+			hasConcurrentLimit = true
+			break
+		}
+	}
+	if !hasConcurrentLimit {
+		if err := s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
+			return store.TaskAdmissionResult{}, false, err
+		}
+		return store.TaskAdmissionResult{}, false, nil
+	}
+	if err := s.store.CheckRateLimits(ctx, s.rateLimitReservations(ctx, user, candidate, body)); err != nil {
+		return store.TaskAdmissionResult{}, true, err
+	}
+	plan := taskAdmissionPlan{
+		Candidate: candidate,
+		Body:      body,
+		ModelType: candidate.ModelType,
+		GroupID:   groupID,
+		Scopes:    scopes,
+		Eligible:  true,
+	}
+	waiterID := ""
+	if !task.AsyncMode {
+		waiterID = uuid.NewString()
+	}
+	result, err := s.waitForTaskAdmission(ctx, task, plan, waiterID)
+	return result, true, err
+}
+
+func (s *Service) observeTaskAdmission(event string) {
+	observer, ok := s.billingMetrics.(interface {
+		ObserveTaskAdmission(string)
+	})
+	if ok {
+		observer.ObserveTaskAdmission(event)
+	}
+}
+
+func (s *Service) observeTaskAdmissionWait(wait time.Duration) {
+	observer, ok := s.billingMetrics.(interface {
+		ObserveTaskAdmissionWait(time.Duration)
+	})
+	if ok {
+		observer.ObserveTaskAdmissionWait(wait)
+	}
+}
+
+func (s *Service) StartAdmissionNotifier(ctx context.Context) {
+	s.admissionListener.Do(func() {
+		go func() {
+			for ctx.Err() == nil {
+				err := s.store.ListenTaskAdmissionNotifications(ctx, func(string) {
+					s.broadcastAdmissionWake()
+				})
+				if ctx.Err() != nil {
+					return
+				}
+				if s.logger != nil {
+					s.logger.Warn("task admission LISTEN connection interrupted; periodic polling remains active", "error", err)
+				}
+				timer := time.NewTimer(time.Second)
+				select {
+				case <-ctx.Done():
+					timer.Stop()
+					return
+				case <-timer.C:
+				}
+			}
+		}()
+	})
+}
+
+func (s *Service) currentAdmissionWake() <-chan struct{} {
+	s.admissionWakeMu.Lock()
+	defer s.admissionWakeMu.Unlock()
+	return s.admissionWake
+}
+
+func (s *Service) broadcastAdmissionWake() {
+	s.admissionWakeMu.Lock()
+	close(s.admissionWake)
+	s.admissionWake = make(chan struct{})
+	s.admissionWakeMu.Unlock()
+}
+
+func (s *Service) SubmitAsyncTask(ctx context.Context, task store.GatewayTask) error {
+	if !task.AsyncMode {
+		return errors.New("only asynchronous tasks can be submitted to the async queue")
+	}
+	if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
+		return ctx.Err()
+	}
+	user := authUserFromTask(task)
+	plan, err := s.buildTaskAdmissionPlan(ctx, task, user)
+	if err != nil {
+		if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
+			return ctx.Err()
+		}
+		_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error())
+		return err
+	}
+	if !plan.Eligible {
+		if err := s.EnqueueAsyncTask(ctx, task); err != nil {
+			if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
+				return ctx.Err()
+			}
+			_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, "enqueue_failed", err.Error())
+			return err
+		}
+		if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
+			return ctx.Err()
+		}
+		return nil
+	}
+	result, err := s.tryTaskAdmissionWithAdmittedHook(ctx, task, plan, "", func(tx pgx.Tx) error {
+		return s.enqueueAsyncTaskTx(ctx, tx, task.ID, asyncTaskInsertOpts(task))
+	})
+	if err != nil {
+		if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
+			return ctx.Err()
+		}
+		_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error())
+		return err
+	}
+	if !result.Admitted {
+		if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
+			return ctx.Err()
+		}
+		return nil
+	}
+	if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
+		return ctx.Err()
+	}
+	return nil
+}
+
+func (s *Service) cancelAsyncSubmissionIfDisconnected(ctx context.Context, taskID string) bool {
+	if ctx.Err() == nil {
+		return false
+	}
+	cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
+	defer cancel()
+	if _, changed, err := s.store.CancelQueuedTask(cleanupCtx, taskID, "client disconnected before upstream submission"); err != nil {
+		if s.logger != nil {
+			s.logger.Warn("cancel disconnected asynchronous task failed", "taskID", taskID, "error", err)
+		}
+	} else if changed {
+		s.observeTaskAdmission("cancelled")
+	}
+	return true
+}
+
+func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) {
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-s.currentAdmissionWake():
+		case <-ticker.C:
+		}
+		taskIDs, err := s.store.ListWaitingAsyncAdmissionTaskIDs(ctx, 100)
+		if err != nil {
+			s.logger.Warn("list waiting async admissions failed", "error", err)
+			continue
+		}
+		for _, taskID := range taskIDs {
+			task, err := s.store.GetTask(ctx, taskID)
+			if err != nil {
+				if !errors.Is(err, pgx.ErrNoRows) {
+					s.logger.Warn("load waiting async task failed", "taskID", taskID, "error", err)
+				}
+				continue
+			}
+			if task.Status != "queued" || task.RiverJobID > 0 {
+				continue
+			}
+			if err := s.SubmitAsyncTask(ctx, task); err != nil {
+				s.logger.Warn("dispatch waiting async admission failed", "taskID", taskID, "error", err)
+			}
+		}
+	}
+}
+
+func (s *Service) reapExpiredTaskAdmissions(ctx context.Context) {
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+		}
+		result, err := s.store.ReapExpiredTaskAdmissions(ctx, 500)
+		if err != nil {
+			if s.logger != nil {
+				s.logger.Warn("reap expired task admissions failed", "error", err)
+			}
+			continue
+		}
+		for range result.ExpiredWaiters {
+			s.observeTaskAdmission("expired")
+		}
+		for range result.ExpiredDeadlines {
+			s.observeTaskAdmission("timeout")
+		}
+	}
+}

apps/api/internal/runner/admission_test.go

@@ -0,0 +1,21 @@
+package runner
+
+import "testing"
+
+func TestDistributedAdmissionModelTypeBoundary(t *testing.T) {
+	for _, modelType := range []string{
+		"image_generate", "image_edit", "image_analysis", "image_vectorize",
+		"video_generate", "video_enhance", "image_to_video", "text_to_video",
+		"video_edit", "video_reference", "video_first_last_frame", "video_understanding", "omni_video", "omni",
+		"audio_generate", "audio_understanding", "text_to_speech", "voice_clone",
+	} {
+		if !distributedAdmissionModelType(modelType) {
+			t.Errorf("%s should use distributed admission", modelType)
+		}
+	}
+	for _, modelType := range []string{"text_generate", "embedding", "rerank", "", "unknown"} {
+		if distributedAdmissionModelType(modelType) {
+			t.Errorf("%s should preserve immediate execution/rate-limit behavior", modelType)
+		}
+	}
+}

apps/api/internal/runner/limits.go

@@ -81,6 +81,74 @@ func (s *Service) rateLimitReservations(ctx context.Context, user *auth.User, ca
 	return out
 }
 
+func (s *Service) admissionScopes(ctx context.Context, user *auth.User, candidate store.RuntimeModelCandidate) ([]store.AdmissionScope, string) {
+	scopes := make([]store.AdmissionScope, 0, 2)
+	platformPolicy := effectiveRateLimitPolicy(candidate)
+	if scope, ok := admissionScopeFromPolicy(
+		"platform_model",
+		candidate.PlatformModelID,
+		firstNonEmptyString(candidate.DisplayName, candidate.ModelAlias, candidate.ModelName),
+		map[string]any{
+			"platformId":   candidate.PlatformID,
+			"platformName": candidate.PlatformName,
+			"modelAlias":   candidate.ModelAlias,
+			"modelName":    candidate.ModelName,
+		},
+		platformPolicy,
+	); ok {
+		scopes = append(scopes, scope)
+	}
+	groupID := ""
+	if group, err := s.store.ResolveUserGroupPolicy(ctx, user); err == nil && group.ID != "" {
+		groupPolicy := store.NormalizeRateLimitPolicy(group.RateLimitPolicy)
+		if scope, ok := admissionScopeFromPolicy(
+			"user_group",
+			group.ID,
+			firstNonEmptyString(group.Name, group.GroupKey),
+			map[string]any{"groupKey": group.GroupKey, "name": group.Name},
+			groupPolicy,
+		); ok {
+			groupID = group.ID
+			scopes = append(scopes, scope)
+		}
+	}
+	return scopes, groupID
+}
+
+func admissionScopeFromPolicy(scopeType string, scopeKey string, scopeName string, metadata map[string]any, policy map[string]any) (store.AdmissionScope, bool) {
+	concurrentLimit, hasConcurrent := store.RateLimitPolicyMetric(policy, "concurrent")
+	queueRule, hasQueue := store.QueueRuleFromPolicy(policy)
+	if !hasConcurrent && !hasQueue {
+		return store.AdmissionScope{}, false
+	}
+	scope := store.AdmissionScope{
+		ScopeType:       scopeType,
+		ScopeKey:        scopeKey,
+		ScopeName:       scopeName,
+		ScopeMetadata:   metadata,
+		ConcurrentLimit: concurrentLimit,
+		Amount:          1,
+		LeaseTTLSeconds: 120,
+		Policy:          policy,
+	}
+	if hasQueue {
+		scope.QueueLimit = queueRule.Limit
+		scope.MaxWaitSeconds = queueRule.MaxWaitSeconds
+	}
+	rules, _ := store.NormalizeRateLimitPolicy(policy)["rules"].([]any)
+	for _, rawRule := range rules {
+		rule, _ := rawRule.(map[string]any)
+		if strings.TrimSpace(stringFromMap(rule, "metric")) != "concurrent" {
+			continue
+		}
+		if ttl := int(floatFromAny(rule["leaseTtlSeconds"])); ttl > 0 {
+			scope.LeaseTTLSeconds = ttl
+		}
+		break
+	}
+	return scope, true
+}
+
 func effectiveRateLimitPolicy(candidate store.RuntimeModelCandidate) map[string]any {
 	return store.EffectiveRateLimitPolicy(store.EffectiveRateLimitPolicyInput{
 		BasePolicy:            candidate.BaseRateLimitPolicy,
@@ -118,7 +186,7 @@ func reservationsFromPolicy(scopeType string, scopeKey string, scopeName string,
 		rule, _ := rawRule.(map[string]any)
 		metric := strings.TrimSpace(stringFromMap(rule, "metric"))
 		limit := floatFromAny(rule["limit"])
-		if metric == "" || limit <= 0 {
+		if metric == "" || metric == "queue_size" || limit <= 0 {
 			continue
 		}
 		amount := 1.0

apps/api/internal/runner/queue_worker.go

@@ -102,7 +102,7 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
 	}
 
 	controlClient, err := river.NewClient(driver, &river.Config{
-		ID:       asyncWorkerID() + "-control",
+		ID:       s.workerInstanceID + "-control",
 		Logger:   s.logger,
 		TestOnly: s.cfg.AppEnv == "test",
 	})
@@ -130,15 +130,18 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
 	if err != nil {
 		return fmt.Errorf("calculate initial async worker capacity: %w", err)
 	}
-	executionClient, err := s.makeAsyncExecutionClient(snapshot.Capacity)
-	if err != nil {
-		return err
-	}
-	if err := executionClient.Start(ctx); err != nil {
-		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		_ = executionClient.StopAndCancel(cleanupCtx)
-		cancel()
-		return err
+	var executionClient asyncExecutionClient
+	if snapshot.Capacity > 0 {
+		executionClient, err = s.makeAsyncExecutionClient(snapshot.Capacity)
+		if err != nil {
+			return err
+		}
+		if err := executionClient.Start(ctx); err != nil {
+			cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			_ = executionClient.StopAndCancel(cleanupCtx)
+			cancel()
+			return err
+		}
 	}
 	s.riverMu.Lock()
 	s.riverControlClient = controlClient
@@ -161,11 +164,15 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
 	)
 	if err := s.recoverAsyncRiverJobs(ctx); err != nil {
 		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		_ = executionClient.StopAndCancel(cleanupCtx)
+		if executionClient != nil {
+			_ = executionClient.StopAndCancel(cleanupCtx)
+		}
 		cancel()
 		return err
 	}
 	go s.refreshAsyncWorkerCapacity(ctx)
+	go s.dispatchWaitingAsyncAdmissions(ctx)
+	go s.reapExpiredTaskAdmissions(ctx)
 	go s.stopAsyncWorkersOnShutdown(ctx)
 	return nil
 }
@@ -176,7 +183,7 @@ func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx.
 		return nil, err
 	}
 	return river.NewClient(riverpgxv5.New(s.store.Pool()), &river.Config{
-		ID:                          fmt.Sprintf("%s-exec-%d-%d", asyncWorkerID(), capacity, time.Now().UnixNano()),
+		ID:                          fmt.Sprintf("%s-exec-%d-%d", s.workerInstanceID, capacity, time.Now().UnixNano()),
 		JobTimeout:                  -1,
 		Logger:                      s.logger,
 		CompletedJobRetentionPeriod: 24 * time.Hour,
@@ -203,7 +210,26 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke
 	if s.asyncCapacityLoader != nil {
 		return s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit)
 	}
-	return s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
+	snapshot, err := s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
+	if err != nil {
+		return store.AsyncWorkerCapacitySnapshot{}, err
+	}
+	allocation, err := s.store.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{
+		InstanceID:      s.workerInstanceID,
+		PodUID:          strings.TrimSpace(os.Getenv("POD_UID")),
+		PodName:         strings.TrimSpace(os.Getenv("POD_NAME")),
+		Site:            strings.TrimSpace(os.Getenv("EASYAI_SITE")),
+		Revision:        strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
+		DesiredCapacity: snapshot.Capacity,
+	})
+	if err != nil {
+		return store.AsyncWorkerCapacitySnapshot{}, err
+	}
+	snapshot.Capacity = allocation.Allocated
+	snapshot.GlobalCapacity = allocation.DesiredCapacity
+	snapshot.ActiveInstances = allocation.ActiveInstances
+	snapshot.InstanceID = allocation.InstanceID
+	return snapshot, nil
 }
 
 func (s *Service) refreshAsyncWorkerCapacity(ctx context.Context) {
@@ -233,25 +259,30 @@ func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
 	if snapshot.Capacity == currentCapacity {
 		return
 	}
-	newClient, err := s.makeAsyncExecutionClient(snapshot.Capacity)
-	if err != nil {
-		s.observeAsyncWorkerResize("create_failed")
-		s.logger.Warn("create replacement async worker client failed; keeping current client",
-			"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
-		return
-	}
-	if err := newClient.Start(ctx); err != nil {
-		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		_ = newClient.StopAndCancel(cleanupCtx)
-		cancel()
-		s.observeAsyncWorkerResize("start_failed")
-		s.logger.Warn("start replacement async worker client failed; keeping current client",
-			"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
-		return
+	var newClient asyncExecutionClient
+	if snapshot.Capacity > 0 {
+		newClient, err = s.makeAsyncExecutionClient(snapshot.Capacity)
+		if err != nil {
+			s.observeAsyncWorkerResize("create_failed")
+			s.logger.Warn("create replacement async worker client failed; keeping current client",
+				"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
+			return
+		}
+		if err := newClient.Start(ctx); err != nil {
+			cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			_ = newClient.StopAndCancel(cleanupCtx)
+			cancel()
+			s.observeAsyncWorkerResize("start_failed")
+			s.logger.Warn("start replacement async worker client failed; keeping current client",
+				"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
+			return
+		}
 	}
 	if ctx.Err() != nil {
 		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		_ = newClient.StopAndCancel(cleanupCtx)
+		if newClient != nil {
+			_ = newClient.StopAndCancel(cleanupCtx)
+		}
 		cancel()
 		return
 	}
@@ -271,6 +302,8 @@ func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
 		"previousCapacity", currentCapacity,
 		"capacity", snapshot.Capacity,
 		"desiredCapacity", snapshot.Desired,
+		"activeInstances", snapshot.ActiveInstances,
+		"instanceID", snapshot.InstanceID,
 		"hardLimit", snapshot.HardLimit,
 		"modelDesired", snapshot.ModelDesired,
 		"groupDesired", snapshot.GroupDesired,
@@ -294,6 +327,13 @@ func (s *Service) drainAsyncWorkerClient(client asyncExecutionClient) {
 
 func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
 	<-ctx.Done()
+	if s.store != nil && strings.TrimSpace(s.workerInstanceID) != "" {
+		drainCtx, drainCancel := context.WithTimeout(context.Background(), 5*time.Second)
+		if err := s.store.MarkWorkerDraining(drainCtx, s.workerInstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
+			s.logger.Warn("mark async worker draining failed", "error", err)
+		}
+		drainCancel()
+	}
 	s.riverMu.Lock()
 	s.riverControlClient = nil
 	clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients))
@@ -328,6 +368,12 @@ func (s *Service) observeAsyncWorkerCapacity(snapshot store.AsyncWorkerCapacityS
 	if ok {
 		observer.SetAsyncWorkerCapacity(snapshot.Capacity, snapshot.Desired, snapshot.HardLimit, snapshot.Capped)
 	}
+	distributedObserver, ok := s.billingMetrics.(interface {
+		SetDistributedWorkerCapacity(activeInstances, globalCapacity, allocatedCapacity int)
+	})
+	if ok {
+		distributedObserver.SetDistributedWorkerCapacity(snapshot.ActiveInstances, snapshot.GlobalCapacity, snapshot.Capacity)
+	}
 }
 
 func (s *Service) observeAsyncWorkerResize(outcome string) {
@@ -340,16 +386,38 @@ func (s *Service) observeAsyncWorkerResize(outcome string) {
 }
 
 func (s *Service) EnqueueAsyncTask(ctx context.Context, task store.GatewayTask) error {
+	return s.enqueueAsyncTaskWithOptions(ctx, task.ID, asyncTaskInsertOpts(task))
+}
+
+func (s *Service) enqueueAsyncTaskWithOptions(ctx context.Context, taskID string, opts *river.InsertOpts) error {
+	tx, err := s.store.Pool().Begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback(ctx)
+	if err := s.enqueueAsyncTaskTx(ctx, tx, taskID, opts); err != nil {
+		return err
+	}
+	return tx.Commit(ctx)
+}
+
+func (s *Service) enqueueAsyncTaskTx(ctx context.Context, tx pgx.Tx, taskID string, opts *river.InsertOpts) error {
 	riverClient := s.asyncControlClient()
 	if riverClient == nil {
 		return errors.New("river async queue is not started")
 	}
-	result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: task.ID}, asyncTaskInsertOpts(task))
+	result, err := riverClient.InsertTx(ctx, tx, asyncTaskArgs{TaskID: taskID}, opts)
 	if err != nil {
 		return err
 	}
 	if result.Job != nil {
-		return s.store.SetTaskRiverJobID(ctx, task.ID, result.Job.ID)
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_tasks
+SET river_job_id = $2,
+    updated_at = now()
+WHERE id = $1::uuid`, taskID, result.Job.ID); err != nil {
+			return err
+		}
 	}
 	return nil
 }
@@ -373,14 +441,13 @@ func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
 		return err
 	}
 	for _, item := range items {
-		result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
-		if err != nil {
-			return err
+		if admission, admissionErr := s.store.GetTaskAdmission(ctx, item.ID); admissionErr == nil && admission.Status == "waiting" {
+			continue
+		} else if admissionErr != nil && !errors.Is(admissionErr, pgx.ErrNoRows) {
+			return admissionErr
 		}
-		if result.Job != nil {
-			if err := s.store.SetTaskRiverJobID(ctx, item.ID, result.Job.ID); err != nil {
-				return err
-			}
+		if err := s.enqueueAsyncTaskWithOptions(ctx, item.ID, asyncTaskRecoveryInsertOpts(item, time.Now())); err != nil {
+			return err
 		}
 	}
 	if len(items) > 0 {
@@ -418,10 +485,6 @@ func asyncTaskRecoveryInsertOpts(item store.AsyncTaskQueueItem, now time.Time) *
 	if item.NextRunAt.After(now) {
 		opts.ScheduledAt = item.NextRunAt
 	}
-	// A replacement process must not be blocked by a River row that the dead
-	// process left in running state. PostgreSQL execution leases still ensure
-	// that only one recovery job can call the upstream provider.
-	opts.UniqueOpts = river.UniqueOpts{}
 	return opts
 }
 

apps/api/internal/runner/request_context.go

@@ -0,0 +1,47 @@
+package runner
+
+import (
+	"context"
+	"sync/atomic"
+)
+
+type submissionCancellationGate struct {
+	submitted atomic.Bool
+}
+
+type submissionCancellationGateKey struct{}
+
+func NewRequestExecutionContext(requestContext context.Context, serverContext context.Context) (context.Context, context.CancelFunc) {
+	base, cancel := context.WithCancel(context.WithoutCancel(requestContext))
+	gate := &submissionCancellationGate{}
+	ctx := context.WithValue(base, submissionCancellationGateKey{}, gate)
+	go func() {
+		select {
+		case <-serverContext.Done():
+			cancel()
+		case <-requestContext.Done():
+			if !gate.submitted.Load() {
+				cancel()
+				return
+			}
+			select {
+			case <-serverContext.Done():
+				cancel()
+			case <-base.Done():
+			}
+		case <-base.Done():
+		}
+	}()
+	return ctx, cancel
+}
+
+func markUpstreamSubmissionStarted(ctx context.Context) {
+	if gate, ok := ctx.Value(submissionCancellationGateKey{}).(*submissionCancellationGate); ok {
+		gate.submitted.Store(true)
+	}
+}
+
+func requestCancelledBeforeSubmission(ctx context.Context) bool {
+	gate, ok := ctx.Value(submissionCancellationGateKey{}).(*submissionCancellationGate)
+	return ok && ctx.Err() != nil && !gate.submitted.Load()
+}

apps/api/internal/runner/request_context_test.go

@@ -0,0 +1,55 @@
+package runner
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+func TestRequestExecutionContextCancelsBeforeSubmission(t *testing.T) {
+	requestCtx, cancelRequest := context.WithCancel(context.Background())
+	serverCtx, cancelServer := context.WithCancel(context.Background())
+	defer cancelServer()
+	ctx, cancel := NewRequestExecutionContext(requestCtx, serverCtx)
+	defer cancel()
+
+	cancelRequest()
+	select {
+	case <-ctx.Done():
+	case <-time.After(time.Second):
+		t.Fatal("request context did not cancel before upstream submission")
+	}
+}
+
+func TestRequestExecutionContextSurvivesDisconnectAfterSubmission(t *testing.T) {
+	requestCtx, cancelRequest := context.WithCancel(context.Background())
+	serverCtx, cancelServer := context.WithCancel(context.Background())
+	ctx, cancel := NewRequestExecutionContext(requestCtx, serverCtx)
+	defer cancel()
+
+	markUpstreamSubmissionStarted(ctx)
+	cancelRequest()
+	select {
+	case <-ctx.Done():
+		t.Fatal("client disconnect propagated after upstream submission")
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	cancelServer()
+	select {
+	case <-ctx.Done():
+	case <-time.After(time.Second):
+		t.Fatal("server shutdown did not cancel detached execution")
+	}
+}
+
+func TestRequestExecutionContextPreservesRequestValues(t *testing.T) {
+	type contextKey struct{}
+	requestCtx := context.WithValue(context.Background(), contextKey{}, "trace-value")
+	ctx, cancel := NewRequestExecutionContext(requestCtx, context.Background())
+	defer cancel()
+
+	if got := ctx.Value(contextKey{}); got != "trace-value" {
+		t.Fatalf("request context value = %v, want trace-value", got)
+	}
+}

apps/api/internal/runner/service.go

@@ -36,7 +36,11 @@ type Service struct {
 	riverExecutionClient asyncExecutionClient
 	riverDrainingClients map[asyncExecutionClient]struct{}
 	riverWorkerCapacity  int
+	workerInstanceID     string
 	asyncCapacityLoader  func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error)
+	admissionWakeMu      sync.Mutex
+	admissionWake        chan struct{}
+	admissionListener    sync.Once
 	asyncClientFactory   func(int) (asyncExecutionClient, error)
 	billingMetrics       billingMetricsObserver
 }
@@ -135,7 +139,9 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
 			"universal":             clients.UniversalClient{HTTPClient: httpClients.none, ScriptExecutor: scriptExecutor},
 			"simulation":            clients.SimulationClient{},
 		},
-		httpClients: httpClients,
+		httpClients:      httpClients,
+		workerInstanceID: asyncWorkerID(),
+		admissionWake:    make(chan struct{}),
 	}
 	if len(observers) > 0 {
 		service.billingMetrics = observers[0]
@@ -172,7 +178,15 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 
 func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta, executionToken string) (Result, error) {
 	wasRunning := task.Status == "running"
-	claimed, err := s.store.ClaimTaskExecution(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
+	preliminaryModelType := modelTypeFromKind(task.Kind, normalizeRequest(task.Kind, task.Request))
+	distributedAdmission := distributedAdmissionModelType(preliminaryModelType)
+	var claimed store.GatewayTask
+	var err error
+	if distributedAdmission && !wasRunning {
+		claimed, err = s.store.ClaimTaskPreparation(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
+	} else {
+		claimed, err = s.store.ClaimTaskExecution(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
+	}
 	if err != nil {
 		return Result{}, err
 	}
@@ -181,6 +195,21 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
 	defer stopExecution()
 	go s.renewTaskExecutionLease(executionCtx, stopExecution, task.ID, task.ExecutionToken)
 	ctx = executionCtx
+	defer func() {
+		if !requestCancelledBeforeSubmission(ctx) {
+			return
+		}
+		cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
+		defer cancel()
+		if _, _, cleanupErr := s.store.CancelTaskBeforeUpstreamSubmission(
+			cleanupCtx,
+			task.ID,
+			task.ExecutionToken,
+			"client disconnected before upstream submission",
+		); cleanupErr != nil && s.logger != nil {
+			s.logger.Warn("cancel disconnected task before upstream submission failed", "taskID", task.ID, "error", cleanupErr)
+		}
+	}()
 	executeStartedAt := time.Now()
 	restoredRequest, err := s.restoreTaskRequestReferences(ctx, task)
 	if err != nil {
@@ -189,13 +218,16 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
 	body := normalizeRequest(task.Kind, restoredRequest)
 	responseExecution := responseExecutionContext{}
 	modelType := modelTypeFromKind(task.Kind, body)
-	if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
-		return Result{}, err
-	}
-	if !wasRunning {
-		if err := s.emit(ctx, task.ID, "task.running", "running", "starting", 0.12, "task pulled from queue and started", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
+	distributedAdmission = distributedAdmissionModelType(modelType)
+	if !distributedAdmission {
+		if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
 			return Result{}, err
 		}
+		if !wasRunning {
+			if err := s.emit(ctx, task.ID, "task.running", "running", "starting", 0.12, "task pulled from queue and started", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
+				return Result{}, err
+			}
+		}
 	}
 	if err := validateRequest(task.Kind, body); err != nil {
 		validationErr := parameterPreprocessClientError(err)
@@ -477,6 +509,90 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
 			}
 		}
 	}
+	var admittedLeases []store.ConcurrencyLease
+	admittedPlatformModelID := ""
+	if distributedAdmission && len(candidates) > 0 {
+		admissionScopes, groupID := s.admissionScopes(ctx, user, candidates[0])
+		hasConcurrentLimit := false
+		for _, scope := range admissionScopes {
+			if scope.ConcurrentLimit > 0 {
+				hasConcurrentLimit = true
+				break
+			}
+		}
+		if hasConcurrentLimit {
+			if err := s.store.CheckRateLimits(ctx, s.rateLimitReservations(ctx, user, candidates[0], body)); err != nil {
+				failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, clients.ErrorCode(err), err.Error(), task.RunMode == "simulation", err)
+				if finishErr != nil {
+					return Result{}, finishErr
+				}
+				return Result{Task: failed, Output: failed.Result}, err
+			}
+			plan := taskAdmissionPlan{
+				Candidate: candidates[0],
+				Body:      body,
+				ModelType: modelType,
+				GroupID:   groupID,
+				Scopes:    admissionScopes,
+				Eligible:  true,
+			}
+			var admissionResult store.TaskAdmissionResult
+			var admissionErr error
+			if task.AsyncMode {
+				admissionResult, admissionErr = s.tryTaskAdmission(ctx, task, plan, "")
+				if admissionErr == nil && !admissionResult.Admitted {
+					_ = s.store.ReleaseTaskPreparation(context.WithoutCancel(ctx), task.ID, task.ExecutionToken)
+					return Result{Task: task, Output: task.Result}, &TaskQueuedError{Delay: time.Second}
+				}
+			} else {
+				admissionResult, admissionErr = s.waitForSynchronousAdmission(ctx, task, plan)
+			}
+			if admissionErr != nil {
+				if requestCancelledBeforeSubmission(ctx) {
+					cancelled, changed, cancelErr := s.store.CancelTaskBeforeUpstreamSubmission(
+						context.WithoutCancel(ctx),
+						task.ID,
+						task.ExecutionToken,
+						"client disconnected before upstream submission",
+					)
+					if cancelErr != nil {
+						return Result{}, cancelErr
+					}
+					if changed {
+						return Result{Task: cancelled, Output: cancelled.Result}, admissionErr
+					}
+				}
+				code := clients.ErrorCode(admissionErr)
+				if errors.Is(admissionErr, context.Canceled) {
+					code = "client_disconnected"
+				}
+				failed, finishErr := s.failTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, code, admissionErr.Error(), task.RunMode == "simulation", admissionErr)
+				if finishErr != nil && !errors.Is(finishErr, store.ErrTaskExecutionLeaseLost) {
+					return Result{}, finishErr
+				}
+				return Result{Task: failed, Output: failed.Result}, admissionErr
+			}
+			admittedLeases = admissionResult.Leases
+			if len(admittedLeases) == 0 {
+				admittedLeases, err = s.store.ActiveTaskAdmissionLeases(ctx, task.ID)
+				if err != nil {
+					return Result{}, err
+				}
+			}
+			admittedPlatformModelID = admissionResult.Admission.PlatformModelID
+		}
+		if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
+			return Result{}, err
+		}
+		if !wasRunning {
+			if err := s.emit(ctx, task.ID, "task.running", "running", "starting", 0.12, "task admitted and started", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
+				return Result{}, err
+			}
+		}
+	}
+	if distributedAdmission {
+		defer s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID)
+	}
 	firstCandidateBody := body
 	normalizedModelType := modelType
 	attemptNo := task.AttemptCount
@@ -579,6 +695,33 @@ candidatesLoop:
 		if !available {
 			continue
 		}
+		if distributedAdmission && candidate.PlatformModelID != admittedPlatformModelID {
+			admissionResult, candidateLimited, admissionErr := s.ensureCandidateAdmission(ctx, task, user, body, candidate)
+			if admissionErr != nil {
+				lastErr = admissionErr
+				lastCandidate = candidate
+				if errors.Is(admissionErr, store.ErrQueueTimeout) {
+					break
+				}
+				if candidateLimited && errors.Is(admissionErr, store.ErrRateLimited) {
+					continue
+				}
+				return Result{}, admissionErr
+			}
+			if candidateLimited {
+				admittedPlatformModelID = admissionResult.Admission.PlatformModelID
+				admittedLeases = admissionResult.Leases
+				if len(admittedLeases) == 0 {
+					admittedLeases, err = s.store.ActiveTaskAdmissionLeases(ctx, task.ID)
+					if err != nil {
+						return Result{}, err
+					}
+				}
+			} else {
+				admittedPlatformModelID = ""
+				admittedLeases = nil
+			}
+		}
 		platformsVisited++
 		lastCandidate = candidate
 		clientAttempts := clientAttemptsForCandidate(candidate)
@@ -614,7 +757,21 @@ candidatesLoop:
 			}
 			candidateBody := preprocessing.Body
 			candidatePricing := pricingByCandidate[pricingCandidateKey(candidate)]
-			response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, candidatePricing, nextAttemptNo, candidateDelta, responseExecution, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record)
+			response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, candidatePricing, nextAttemptNo, candidateDelta, responseExecution, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record, admittedPlatformModelID, admittedLeases)
+			if err != nil && requestCancelledBeforeSubmission(ctx) {
+				cancelled, changed, cancelErr := s.store.CancelTaskBeforeUpstreamSubmission(
+					context.WithoutCancel(ctx),
+					task.ID,
+					task.ExecutionToken,
+					"client disconnected before upstream submission",
+				)
+				if cancelErr != nil {
+					return Result{}, cancelErr
+				}
+				if changed {
+					return Result{Task: cancelled, Output: cancelled.Result}, err
+				}
+			}
 			if err != nil && isVolcesRemoteTaskCancellation(candidate, err) {
 				cancelled, changed, cancelErr := s.store.CancelSubmittedTask(ctx, task.ID, task.ExecutionToken, "任务已由火山引擎取消")
 				if cancelErr != nil {
@@ -981,7 +1138,22 @@ func billingItemsFixedTotal(items []any) fixedAmount {
 	return total
 }
 
-func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, pricing resolvedPricing, attemptNo int, onDelta clients.StreamDelta, responseExecution responseExecutionContext, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (clients.Response, error) {
+func (s *Service) runCandidate(
+	ctx context.Context,
+	task store.GatewayTask,
+	user *auth.User,
+	body map[string]any,
+	preprocessing parameterPreprocessingLog,
+	candidate store.RuntimeModelCandidate,
+	pricing resolvedPricing,
+	attemptNo int,
+	onDelta clients.StreamDelta,
+	responseExecution responseExecutionContext,
+	cacheAffinityPolicy map[string]any,
+	cacheAffinityRecordKeys []string,
+	admittedPlatformModelID string,
+	admittedLeases []store.ConcurrencyLease,
+) (clients.Response, error) {
 	var err error
 	candidate, err = candidateWithEnvironmentCredentials(candidate)
 	if err != nil {
@@ -990,12 +1162,24 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	simulated := isSimulation(task, candidate)
 	baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
 	reservations := s.rateLimitReservations(ctx, user, candidate, body)
+	if admittedPlatformModelID == candidate.PlatformModelID && len(admittedLeases) > 0 {
+		filtered := make([]store.RateLimitReservation, 0, len(reservations))
+		for _, reservation := range reservations {
+			if reservation.Metric != "concurrent" {
+				filtered = append(filtered, reservation)
+			}
+		}
+		reservations = filtered
+	}
 	limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, "", reservations)
 	if err != nil {
 		retryable := store.RateLimitRetryable(err)
 		clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable}
 		return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)}
 	}
+	if admittedPlatformModelID == candidate.PlatformModelID {
+		limitResult.Leases = append(limitResult.Leases, admittedLeases...)
+	}
 	rateReservationsFinalized := false
 	defer func() {
 		if !rateReservationsFinalized {
@@ -1124,6 +1308,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	submissionStatus := "not_submitted"
 	if strings.TrimSpace(task.RemoteTaskID) != "" {
 		submissionStatus = "response_received"
+		markUpstreamSubmissionStarted(ctx)
 		if err := s.store.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, submissionStatus); err != nil {
 			return clients.Response{}, fmt.Errorf("restore upstream submission status: %w", err)
 		}
@@ -1181,7 +1366,11 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 			return setSubmissionStatus("response_received")
 		},
 		OnUpstreamSubmissionStarted: func() error {
-			return setSubmissionStatus("submitting")
+			if err := setSubmissionStatus("submitting"); err != nil {
+				return err
+			}
+			markUpstreamSubmissionStarted(ctx)
+			return nil
 		},
 		OnUpstreamResponseReceived: func() error {
 			return setSubmissionStatus("response_received")
@@ -1907,7 +2096,10 @@ func canonicalModelType(value string) string {
 
 func isKnownModelType(value string) bool {
 	switch value {
-	case "text_generate", "text_embedding", "text_rerank", "image_generate", "image_edit", "image_vectorize", "video_generate", "video_enhance", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "omni_video", "omni", "audio_generate", "text_to_speech", "voice_clone":
+	case "text_generate", "text_embedding", "text_rerank",
+		"image_generate", "image_edit", "image_analysis", "image_vectorize",
+		"video_generate", "video_enhance", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "video_understanding",
+		"omni_video", "omni", "audio_generate", "audio_understanding", "text_to_speech", "voice_clone":
 		return true
 	default:
 		return false