Add generic AI job queue lifecycle

internal/store/store.go

@@ -19,6 +19,13 @@ type Store struct {
 	pool *pgxpool.Pool
 }
 
+const jobSelectColumns = `
+id, owner_service, owner_ref, task_type, model_profile, priority, status,
+attempts, max_attempts, input, result, error_code, error_message,
+scheduled_at, started_at, completed_at, worker_id, heartbeat_at,
+created_at, updated_at, idempotency_key
+`
+
 func Open(ctx context.Context, databaseURL string) (*Store, error) {
 	if strings.TrimSpace(databaseURL) == "" {
 		return nil, errors.New("DATABASE_URL is required")
@@ -74,9 +81,7 @@ INSERT INTO ai_jobs (
 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)
 ON CONFLICT (idempotency_key) WHERE idempotency_key IS NOT NULL
 DO UPDATE SET updated_at = ai_jobs.updated_at
-RETURNING id, owner_service, owner_ref, task_type, model_profile, priority, status,
-          attempts, max_attempts, input, result, error_code, error_message,
-          scheduled_at, started_at, completed_at, created_at, updated_at, idempotency_key
+RETURNING ` + jobSelectColumns + `
 `
 	row := s.pool.QueryRow(ctx, q,
 		in.OwnerService,
@@ -111,7 +116,8 @@ func (s *Store) GetJob(ctx context.Context, id uuid.UUID) (*model.Job, error) {
 	const q = `
 SELECT id, owner_service, owner_ref, task_type, model_profile, priority, status,
        attempts, max_attempts, input, result, error_code, error_message,
-       scheduled_at, started_at, completed_at, created_at, updated_at, idempotency_key
+       scheduled_at, started_at, completed_at, worker_id, heartbeat_at,
+       created_at, updated_at, idempotency_key
 FROM ai_jobs
 WHERE id = $1
 `
@@ -131,13 +137,16 @@ SET status = 'pending',
     completed_at = NULL,
     error_code = NULL,
     error_message = NULL,
+    worker_id = NULL,
+    heartbeat_at = NULL,
     scheduled_at = NOW(),
     updated_at = NOW()
 WHERE id = $1
   AND status IN ('failed', 'running')
 RETURNING id, owner_service, owner_ref, task_type, model_profile, priority, status,
           attempts, max_attempts, input, result, error_code, error_message,
-          scheduled_at, started_at, completed_at, created_at, updated_at, idempotency_key
+          scheduled_at, started_at, completed_at, worker_id, heartbeat_at,
+          created_at, updated_at, idempotency_key
 `
 	job, err := scanJob(s.pool.QueryRow(ctx, q, id))
 	if errors.Is(err, pgx.ErrNoRows) {
@@ -146,6 +155,115 @@ RETURNING id, owner_service, owner_ref, task_type, model_profile, priority, stat
 	return job, err
 }
 
+func (s *Store) ClaimJobs(ctx context.Context, in model.ClaimJobs) ([]*model.Job, error) {
+	if in.Limit <= 0 {
+		in.Limit = 1
+	}
+	if in.Limit > 100 {
+		in.Limit = 100
+	}
+	workerID := strings.TrimSpace(in.WorkerID)
+	if workerID == "" {
+		workerID = "unknown"
+	}
+	if in.TaskTypes == nil {
+		in.TaskTypes = []string{}
+	}
+	if in.ModelProfiles == nil {
+		in.ModelProfiles = []string{}
+	}
+	const q = `
+WITH picked AS (
+  SELECT id
+  FROM ai_jobs
+  WHERE status = 'pending'
+    AND attempts < max_attempts
+    AND scheduled_at <= NOW()
+    AND (cardinality($1::text[]) = 0 OR task_type = ANY($1::text[]))
+    AND (cardinality($2::text[]) = 0 OR model_profile = ANY($2::text[]))
+  ORDER BY priority DESC, scheduled_at ASC, created_at ASC
+  LIMIT $3
+  FOR UPDATE SKIP LOCKED
+)
+UPDATE ai_jobs j
+SET status = 'running',
+    attempts = j.attempts + 1,
+    started_at = NOW(),
+    completed_at = NULL,
+    error_code = NULL,
+    error_message = NULL,
+    worker_id = $4,
+    heartbeat_at = NOW(),
+    updated_at = NOW()
+FROM picked
+WHERE j.id = picked.id
+RETURNING ` + jobSelectColumns + `
+`
+	rows, err := s.pool.Query(ctx, q, in.TaskTypes, in.ModelProfiles, in.Limit, workerID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var out []*model.Job
+	for rows.Next() {
+		job, err := scanJob(rows)
+		if err != nil {
+			return nil, err
+		}
+		out = append(out, job)
+	}
+	return out, rows.Err()
+}
+
+func (s *Store) CompleteJob(ctx context.Context, id uuid.UUID, in model.CompleteJob) (*model.Job, error) {
+	if len(in.Result) == 0 {
+		in.Result = json.RawMessage(`{}`)
+	}
+	const q = `
+UPDATE ai_jobs
+SET status = 'done',
+    result = $2,
+    error_code = NULL,
+    error_message = NULL,
+    completed_at = NOW(),
+    heartbeat_at = NOW(),
+    updated_at = NOW()
+WHERE id = $1
+  AND status = 'running'
+RETURNING ` + jobSelectColumns + `
+`
+	job, err := scanJob(s.pool.QueryRow(ctx, q, id, in.Result))
+	if errors.Is(err, pgx.ErrNoRows) {
+		return nil, nil
+	}
+	return job, err
+}
+
+func (s *Store) FailJob(ctx context.Context, id uuid.UUID, in model.FailJob) (*model.Job, error) {
+	errorCode := strings.TrimSpace(in.ErrorCode)
+	if errorCode == "" {
+		errorCode = "unknown"
+	}
+	errorMessage := strings.TrimSpace(in.ErrorMessage)
+	const q = `
+UPDATE ai_jobs
+SET status = 'failed',
+    error_code = $2,
+    error_message = $3,
+    completed_at = NOW(),
+    heartbeat_at = NOW(),
+    updated_at = NOW()
+WHERE id = $1
+  AND status = 'running'
+RETURNING ` + jobSelectColumns + `
+`
+	job, err := scanJob(s.pool.QueryRow(ctx, q, id, errorCode, errorMessage))
+	if errors.Is(err, pgx.ErrNoRows) {
+		return nil, nil
+	}
+	return job, err
+}
+
 func (s *Store) Stats(ctx context.Context) (*model.Stats, error) {
 	out := &model.Stats{At: time.Now().UTC()}
 
@@ -214,6 +332,8 @@ func scanJob(row pgx.Row) (*model.Job, error) {
 		&job.ScheduledAt,
 		&job.StartedAt,
 		&job.CompletedAt,
+		&job.WorkerID,
+		&job.HeartbeatAt,
 		&job.CreatedAt,
 		&job.UpdatedAt,
 		&job.IdempotencyKey,