feat: add river-backed async task queue

apps/api/internal/runner/service.go

@@ -3,6 +3,7 @@ package runner
 import (
 	"context"
 	"errors"
+	"fmt"
 	"log/slog"
 	"strconv"
 	"strings"
@@ -12,6 +13,8 @@ import (
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+	"github.com/jackc/pgx/v5"
+	"github.com/riverqueue/river"
 )
 
 type Service struct {
@@ -20,6 +23,7 @@ type Service struct {
 	logger      *slog.Logger
 	clients     map[string]clients.Client
 	httpClients *httpClientCache
+	riverClient *river.Client[pgx.Tx]
 }
 
 type Result struct {
@@ -27,6 +31,20 @@ type Result struct {
 	Output map[string]any
 }
 
+var ErrTaskQueued = errors.New("task queued")
+
+type TaskQueuedError struct {
+	Delay time.Duration
+}
+
+func (e *TaskQueuedError) Error() string {
+	return ErrTaskQueued.Error()
+}
+
+func (e *TaskQueuedError) Is(target error) bool {
+	return target == ErrTaskQueued
+}
+
 func New(cfg config.Config, db *store.Store, logger *slog.Logger) *Service {
 	httpClients := newHTTPClientCache()
 	return &Service{
@@ -55,6 +73,14 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 	executeStartedAt := time.Now()
 	body := normalizeRequest(task.Kind, task.Request)
 	modelType := modelTypeFromKind(task.Kind, body)
+	if err := s.store.MarkTaskRunning(ctx, task.ID, modelType, body); err != nil {
+		return Result{}, err
+	}
+	if task.Status != "running" {
+		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 {
 		failed, finishErr := s.failTask(ctx, task.ID, "bad_request", err.Error(), task.RunMode == "simulation", err)
 		if finishErr != nil {
@@ -83,9 +109,6 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 			return Result{}, err
 		}
 	}
-	if err := s.store.MarkTaskRunning(ctx, task.ID, modelType, body); err != nil {
-		return Result{}, err
-	}
 	if err := s.emit(ctx, task.ID, "task.progress", "running", "normalizing", 0.15, "request normalized", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
 		return Result{}, err
 	}
@@ -96,7 +119,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 	}
 	maxPlatforms := maxPlatformsForCandidates(candidates, runnerPolicy)
 	maxFailoverDuration := maxFailoverDurationForCandidates(candidates, runnerPolicy)
-	attemptNo := 0
+	attemptNo := task.AttemptCount
 	var lastErr error
 	for index, candidate := range candidates {
 		if index >= maxPlatforms {
@@ -251,6 +274,20 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 	if lastErr != nil {
 		message = lastErr.Error()
 	}
+	if task.AsyncMode && ctx.Err() != nil {
+		queued, queueErr := s.requeueInterruptedAsyncTask(context.WithoutCancel(ctx), task)
+		if queueErr != nil {
+			return Result{}, queueErr
+		}
+		return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: 0}
+	}
+	if task.AsyncMode && errors.Is(lastErr, store.ErrRateLimited) && store.RateLimitRetryable(lastErr) {
+		queued, delay, queueErr := s.requeueRateLimitedTask(ctx, task, lastErr)
+		if queueErr != nil {
+			return Result{}, queueErr
+		}
+		return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: delay}
+	}
 	failed, err := s.failTask(ctx, task.ID, code, message, task.RunMode == "simulation", lastErr)
 	if err != nil {
 		return Result{}, err
@@ -261,7 +298,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta) (clients.Response, error) {
 	simulated := isSimulation(task, candidate)
 	if err := s.emit(ctx, task.ID, "task.attempt.started", "running", "submitting", 0.25, "client attempt started", map[string]any{"attempt": attemptNo, "clientId": candidate.ClientID}, simulated); err != nil {
-		return clients.Response{}, err
+		return clients.Response{}, fmt.Errorf("emit attempt started: %w", err)
 	}
 	attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
 		TaskID:          task.ID,
@@ -276,21 +313,22 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 		Metrics:         attemptMetrics(candidate, attemptNo, simulated),
 	})
 	if err != nil {
-		return clients.Response{}, err
+		return clients.Response{}, fmt.Errorf("create task attempt: %w", err)
 	}
 	reservations := s.rateLimitReservations(ctx, user, candidate, body)
 	limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, attemptID, reservations)
 	if err != nil {
-		clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: false}
+		retryable := store.RateLimitRetryable(err)
+		clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable}
 		_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
 			AttemptID:    attemptID,
 			Status:       "failed",
-			Retryable:    false,
-			Metrics:      mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), map[string]any{"error": err.Error(), "retryable": false, "trace": []any{failureTraceEntry(clientErr, false)}}),
+			Retryable:    retryable,
+			Metrics:      mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), map[string]any{"error": err.Error(), "retryable": retryable, "retryAfterMs": localRateLimitRetryAfter(err).Milliseconds(), "trace": []any{failureTraceEntry(clientErr, retryable)}}),
 			ErrorCode:    "rate_limit",
 			ErrorMessage: err.Error(),
 		})
-		return clients.Response{}, clientErr
+		return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)}
 	}
 	rateReservationsFinalized := false
 	defer func() {
@@ -301,7 +339,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.LeaseIDs)
 
 	if err := s.store.RecordClientAssignment(ctx, candidate); err != nil {
-		return clients.Response{}, err
+		return clients.Response{}, fmt.Errorf("record client assignment: %w", err)
 	}
 	defer s.store.RecordClientRelease(context.WithoutCancel(ctx), candidate.ClientID, "")
 
@@ -315,17 +353,25 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 			ErrorCode:    clients.ErrorCode(err),
 			ErrorMessage: err.Error(),
 		})
-		return clients.Response{}, err
+		return clients.Response{}, fmt.Errorf("prepare http client: %w", err)
 	}
 	client := s.clientFor(candidate, simulated)
 	callStartedAt := time.Now()
 	response, err := client.Run(ctx, clients.Request{
-		Kind:        task.Kind,
-		ModelType:   candidate.ModelType,
-		Model:       task.Model,
-		Body:        body,
-		Candidate:   candidate,
-		HTTPClient:  requestHTTPClient,
+		Kind:              task.Kind,
+		ModelType:         candidate.ModelType,
+		Model:             task.Model,
+		Body:              body,
+		Candidate:         candidate,
+		HTTPClient:        requestHTTPClient,
+		RemoteTaskID:      task.RemoteTaskID,
+		RemoteTaskPayload: task.RemoteTaskPayload,
+		OnRemoteTaskSubmitted: func(remoteTaskID string, payload map[string]any) error {
+			if strings.TrimSpace(remoteTaskID) == "" {
+				return nil
+			}
+			return s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, attemptID, remoteTaskID, payload)
+		},
 		Stream:      boolFromMap(body, "stream"),
 		StreamDelta: onDelta,
 	})
@@ -400,11 +446,11 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	response.Result = uploadedResult
 	for _, progress := range response.Progress {
 		if err := s.emit(ctx, task.ID, "task.progress", "running", progress.Phase, progress.Progress, progress.Message, progress.Payload, simulated); err != nil {
-			return clients.Response{}, err
+			return clients.Response{}, fmt.Errorf("emit task progress: %w", err)
 		}
 	}
 	if err := s.store.CommitRateLimitReservations(ctx, limitResult.Reservations, tokenUsageAmounts(response.Usage)); err != nil {
-		return clients.Response{}, err
+		return clients.Response{}, fmt.Errorf("commit rate limit reservations: %w", err)
 	}
 	rateReservationsFinalized = true
 	if err := s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
@@ -418,7 +464,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 		ResponseFinishedAt: response.ResponseFinishedAt,
 		ResponseDurationMS: response.ResponseDurationMS,
 	}); err != nil {
-		return clients.Response{}, err
+		return clients.Response{}, fmt.Errorf("finish task attempt: %w", err)
 	}
 	return response, nil
 }
@@ -459,6 +505,41 @@ func (s *Service) failTask(ctx context.Context, taskID string, code string, mess
 	return failed, nil
 }
 
+func (s *Service) requeueRateLimitedTask(ctx context.Context, task store.GatewayTask, cause error) (store.GatewayTask, time.Duration, error) {
+	delay := localRateLimitRetryAfter(cause)
+	if delay <= 0 {
+		delay = 5 * time.Second
+	}
+	queued, err := s.store.RequeueTask(ctx, task.ID, delay)
+	if err != nil {
+		return store.GatewayTask{}, 0, err
+	}
+	payload := map[string]any{
+		"code":         "rate_limit",
+		"message":      cause.Error(),
+		"retryAfterMs": delay.Milliseconds(),
+	}
+	if eventErr := s.emit(ctx, task.ID, "task.queued", "queued", "rate_limited", 0.2, "task queued by local rate limit", payload, task.RunMode == "simulation"); eventErr != nil {
+		return store.GatewayTask{}, 0, eventErr
+	}
+	return queued, delay, nil
+}
+
+func (s *Service) requeueInterruptedAsyncTask(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) {
+	queued, err := s.store.RequeueTask(ctx, task.ID, 0)
+	if err != nil {
+		return store.GatewayTask{}, err
+	}
+	payload := map[string]any{"code": "worker_interrupted"}
+	if task.RemoteTaskID != "" {
+		payload["remoteTaskId"] = task.RemoteTaskID
+	}
+	if eventErr := s.emit(ctx, task.ID, "task.queued", "queued", "worker_interrupted", 0.2, "async task queued after worker interruption", payload, task.RunMode == "simulation"); eventErr != nil {
+		return store.GatewayTask{}, eventErr
+	}
+	return queued, nil
+}
+
 func (s *Service) withAttemptHistory(ctx context.Context, taskID string, metrics map[string]any) map[string]any {
 	attempts, err := s.store.ListTaskAttempts(ctx, taskID)
 	if err != nil {