chore(git): 合并远端主分支最新变更

apps/api/internal/httpapi/async_worker_acceptance_integration_test.go

@@ -54,6 +54,7 @@ func TestAsyncWorkerDynamicConcurrencyAcceptance(t *testing.T) {
 		JWTSecret:                         "test-secret",
 		BillingEngineMode:                 "observe",
 		CORSAllowedOrigin:                 "*",
+		AsyncQueueWorkerEnabled:           true,
 		AsyncWorkerHardLimit:              256,
 		AsyncWorkerRefreshIntervalSeconds: 1,
 	}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))

apps/api/internal/httpapi/server.go

@@ -122,8 +122,11 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
 	server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
 	if cfg.AsyncQueueWorkerEnabled {
 		server.runner.StartAsyncQueueWorker(ctx)
-	} else if logger != nil {
-		logger.Info("asynchronous queue worker disabled for this process")
+	} else {
+		server.runner.StartAsyncQueueClient(ctx)
+		if logger != nil {
+			logger.Info("asynchronous queue worker disabled for this process")
+		}
 	}
 	server.runner.StartBillingSettlementWorker(ctx)
 	server.runner.StartTaskHistoryWorkers(ctx)

apps/api/internal/runner/queue_client_test.go

@@ -0,0 +1,74 @@
+package runner
+
+import (
+	"context"
+	"io"
+	"log/slog"
+	"os"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+)
+
+func TestAsyncQueueClientEnqueuesWithoutExecutionWorker(t *testing.T) {
+	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
+	if databaseURL == "" {
+		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the queue client integration test")
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	db, err := store.Connect(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect store: %v", err)
+	}
+	t.Cleanup(db.Close)
+
+	suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
+	task, err := db.CreateTask(ctx, store.CreateTaskInput{
+		Kind:    "images.edits",
+		Model:   "queue-client-" + suffix,
+		Request: map[string]any{"prompt": "enqueue without execution worker"},
+		Async:   true,
+		RunMode: "simulation",
+	}, &auth.User{ID: "queue-client-" + suffix, Source: "gateway"})
+	if err != nil {
+		t.Fatalf("create task: %v", err)
+	}
+	t.Cleanup(func() {
+		_, _ = db.Pool().Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id = $1::uuid`, task.ID)
+	})
+
+	service := New(config.Config{AppEnv: "test"}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
+	service.StartAsyncQueueClient(ctx)
+	if err := service.EnqueueAsyncTask(ctx, task); err != nil {
+		t.Fatalf("enqueue task without execution worker: %v", err)
+	}
+
+	queued, err := db.GetTask(ctx, task.ID)
+	if err != nil {
+		t.Fatalf("load queued task: %v", err)
+	}
+	if queued.RiverJobID <= 0 {
+		t.Fatal("queue client did not persist a River job ID")
+	}
+	t.Cleanup(func() {
+		_, _ = db.Pool().Exec(context.Background(), `DELETE FROM river_job WHERE id = $1`, queued.RiverJobID)
+	})
+
+	var attemptedByCount int
+	if err := db.Pool().QueryRow(ctx, `
+SELECT cardinality(attempted_by)
+FROM river_job
+WHERE id = $1`, queued.RiverJobID).Scan(&attemptedByCount); err != nil {
+		t.Fatalf("load River job: %v", err)
+	}
+	if attemptedByCount != 0 {
+		t.Fatalf("control-only queue client executed the job: attempted_by=%d", attemptedByCount)
+	}
+}

apps/api/internal/runner/queue_worker.go

@@ -78,13 +78,20 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs
 }
 
 func (s *Service) StartAsyncQueueWorker(ctx context.Context) {
-	if err := s.startRiverQueue(ctx); err != nil {
+	if err := s.startRiverQueue(ctx, true); err != nil {
 		s.logger.Error("start river async queue failed", "error", err)
 		panic(err)
 	}
 }
 
-func (s *Service) startRiverQueue(ctx context.Context) error {
+func (s *Service) StartAsyncQueueClient(ctx context.Context) {
+	if err := s.startRiverQueue(ctx, false); err != nil {
+		s.logger.Error("start river async queue client failed", "error", err)
+		panic(err)
+	}
+}
+
+func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error {
 	driver := riverpgxv5.New(s.store.Pool())
 	migrator, err := rivermigrate.New(driver, nil)
 	if err != nil {
@@ -102,6 +109,23 @@ func (s *Service) startRiverQueue(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
+	if !workerEnabled {
+		s.riverMu.Lock()
+		s.riverControlClient = controlClient
+		s.riverExecutionClient = nil
+		s.riverWorkerCapacity = 0
+		s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
+		s.riverMu.Unlock()
+		if err := s.recoverAsyncRiverJobs(ctx); err != nil {
+			s.riverMu.Lock()
+			s.riverControlClient = nil
+			s.riverMu.Unlock()
+			return err
+		}
+		go s.stopAsyncWorkersOnShutdown(ctx)
+		s.logger.Info("river async queue client initialized without execution workers")
+		return nil
+	}
 	snapshot, err := s.loadAsyncWorkerCapacity(ctx)
 	if err != nil {
 		return fmt.Errorf("calculate initial async worker capacity: %w", err)
@@ -271,6 +295,7 @@ func (s *Service) drainAsyncWorkerClient(client asyncExecutionClient) {
 func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
 	<-ctx.Done()
 	s.riverMu.Lock()
+	s.riverControlClient = nil
 	clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients))
 	if s.riverExecutionClient != nil {
 		clients = append(clients, s.riverExecutionClient)