easyai-ai-gateway/apps/api/internal/runner/queue_client_test.go

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 COALESCE(cardinality(attempted_by), 0)
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)
	}
}

func TestRecoverOrphanedAsyncRiverJobRequiresInactiveWorker(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 orphaned River job 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)
	service := New(config.Config{AppEnv: "test"}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
	service.StartAsyncQueueClient(ctx)

	suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
	createQueuedTask := func(label string) store.GatewayTask {
		t.Helper()
		task, createErr := db.CreateTask(ctx, store.CreateTaskInput{
			Kind:    "images.edits",
			Model:   "orphan-recovery-" + label + "-" + suffix,
			Request: map[string]any{"prompt": "recover orphaned River job"},
			Async:   true,
			RunMode: "simulation",
		}, &auth.User{ID: "orphan-recovery-" + suffix, Source: "gateway"})
		if createErr != nil {
			t.Fatalf("create %s task: %v", label, createErr)
		}
		if enqueueErr := service.EnqueueAsyncTask(ctx, task); enqueueErr != nil {
			t.Fatalf("enqueue %s task: %v", label, enqueueErr)
		}
		queued, getErr := db.GetTask(ctx, task.ID)
		if getErr != nil {
			t.Fatalf("load %s task: %v", label, getErr)
		}
		return queued
	}

	orphaned := createQueuedTask("orphaned")
	protected := createQueuedTask("protected")
	staleWorkerID := "orphan-recovery-stale-" + suffix
	activeWorkerID := "orphan-recovery-active-" + suffix
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cleanupCancel()
		_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_worker_instances WHERE instance_id = ANY($1::text[])`, []string{staleWorkerID, activeWorkerID})
		_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM river_job WHERE id = ANY($1::bigint[])`, []int64{orphaned.RiverJobID, protected.RiverJobID})
		_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`, []string{orphaned.ID, protected.ID})
	})

	if _, err := db.Pool().Exec(ctx, `
INSERT INTO gateway_worker_instances (
  instance_id, status, desired_capacity, allocated_capacity, heartbeat_at, updated_at
)
VALUES
  ($1, 'active', 1, 0, now() - interval '2 minutes', now() - interval '2 minutes'),
  ($2, 'active', 1, 1, now(), now())`,
		staleWorkerID,
		activeWorkerID,
	); err != nil {
		t.Fatalf("seed worker heartbeats: %v", err)
	}
	if _, err := db.Pool().Exec(ctx, `
UPDATE river_job
SET state = 'running',
    attempt = 1,
    attempted_at = now() - interval '1 minute',
    attempted_by = ARRAY[
      CASE id
        WHEN $1::bigint THEN $3
        WHEN $2::bigint THEN $4
      END
    ]::text[]
WHERE id = ANY($5::bigint[])`,
		orphaned.RiverJobID,
		protected.RiverJobID,
		staleWorkerID+"-exec-1-test",
		activeWorkerID+"-exec-1-test",
		[]int64{orphaned.RiverJobID, protected.RiverJobID},
	); err != nil {
		t.Fatalf("mark River jobs running: %v", err)
	}

	recovered, err := db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10)
	if err != nil {
		t.Fatalf("recover orphaned River jobs: %v", err)
	}
	if recovered != 1 {
		t.Fatalf("recovered jobs=%d, want 1", recovered)
	}
	var orphanedState, protectedState string
	if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, orphaned.RiverJobID).Scan(&orphanedState); err != nil {
		t.Fatalf("read orphaned River job state: %v", err)
	}
	if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, protected.RiverJobID).Scan(&protectedState); err != nil {
		t.Fatalf("read protected River job state: %v", err)
	}
	if orphanedState != "retryable" || protectedState != "running" {
		t.Fatalf("River states orphaned=%s protected=%s, want retryable/running", orphanedState, protectedState)
	}
	recovered, err = db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10)
	if err != nil || recovered != 0 {
		t.Fatalf("second orphan recovery count=%d err=%v, want idempotent zero", recovered, err)
	}
}