Job lifecycle

kind   text NOT NULL CHECK (kind IN ('test','discover','preview','extract','load')),
status text NOT NULL DEFAULT 'pending'
       CHECK (status IN ('pending','leased','running','succeeded','failed','expired')),
attempt      int NOT NULL DEFAULT 0,
max_attempts int NOT NULL DEFAULT 3,
lease_expires_at timestamptz,
result jsonb,
error  text
-- run_id / connection_id are nullable FKs, ON DELETE CASCADE.
-- CREATE INDEX agent_job_lease_idx ON agent_job (agent_id, status);

-- name: LeaseNextJob :one
-- Atomically claim the oldest pending job for this agent. FOR UPDATE SKIP LOCKED
-- lets concurrent leasers each grab a distinct row without blocking: the inner
-- SELECT locks one pending row and skips rows already locked by a peer, so the
-- same job is never handed out twice.
UPDATE agent_job SET
  status = 'leased',
  attempt = attempt + 1,
  lease_expires_at = now() + make_interval(secs => sqlc.arg('lease_seconds')::int),
  updated_at = now()
WHERE id = (
  SELECT j.id FROM agent_job j
  WHERE j.agent_id = sqlc.arg('agent_id') AND j.status = 'pending'
  ORDER BY j.created_at
  FOR UPDATE SKIP LOCKED
  LIMIT 1
)
RETURNING *;

// heartbeatLoop renews a job's lease on a fixed interval until ctx is cancelled.
func heartbeatLoop(ctx context.Context, c *client.Client, jobID string) {
	t := time.NewTicker(heartbeatInterval) // 20s, well under the 60s lease TTL
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			if err := c.Heartbeat(ctx, jobID); err != nil && ctx.Err() == nil {
				fmt.Fprintf(os.Stderr, "heartbeat job %s: %v\n", jobID, err)
			}
		}
	}
}

-- name: HeartbeatJob :exec
UPDATE agent_job SET
  lease_expires_at = now() + make_interval(secs => sqlc.arg('lease_seconds')::int),
  updated_at = now()
WHERE id = sqlc.arg('id')
  AND agent_id = sqlc.arg('agent_id')
  AND status = 'leased';

-- name: RequeueExpiredJobs :many
-- Lease reaper: reclaim jobs whose lease elapsed while still 'leased'. Jobs with
-- attempts to spare go back to 'pending'; those that exhausted max_attempts become
-- terminal 'expired'. Returns the affected rows.
UPDATE agent_job SET
  status = CASE WHEN attempt >= max_attempts THEN 'expired' ELSE 'pending' END,
  lease_expires_at = NULL,
  updated_at = now()
WHERE status = 'leased' AND lease_expires_at < now()
RETURNING *;

func (s *Server) StartAgentJobReaper(ctx context.Context, every time.Duration) {
	t := time.NewTicker(every)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			rows, err := s.q.RequeueExpiredJobs(ctx)
			if err != nil {
				s.log.Error("agent-job reaper failed", "err", err)
				continue
			}
			if len(rows) > 0 {
				s.log.Info("reaped expired agent-job leases", "count", len(rows))
			}
		}
	}
}

spec, err := c.ExtractSpec(ctx, job.ID)     // fetch secretless extract-spec
// ...
hbCtx, stopHB := context.WithCancel(ctx)
go heartbeatLoop(hbCtx, c, job.ID)          // renew the lease while copying
res, err := exec.Extract(ctx, spec, cfg.Secrets)
stopHB()                                    // stop the heartbeat goroutine
if err != nil {
	reportFailure(ctx, c, job.ID, err.Error()) // POST result status=failed
	return
}
c.Result(ctx, job.ID, "succeeded",
	map[string]any{"blobPrefix": res.BlobPrefix, "rows": res.Rows}, "")

s.q.CompleteAgentJob(ctx, sqlc.CompleteAgentJobParams{
	ID: id, AgentID: agentID, Status: req.Status, Result: resultBytes, Error: req.Error,
})
job.Status = req.Status
job.Result = resultBytes
if req.Status == "succeeded" {
	s.onAgentJobResult(ctx, job) // extract → dispatch load; load → complete run; discover → catalog
} else {
	s.onAgentJobFailed(ctx, job, req.Error) // extract/load → fail the parent run
}

func (s *Server) onAgentJobResult(ctx context.Context, job sqlc.AgentJob) {
	if job.Status != "succeeded" {
		return
	}
	// Never resurrect a run that was cancelled or already finished.
	if job.RunID.Valid && (job.Kind == "extract" || job.Kind == "load") &&
		!s.runIsActive(ctx, uuid.UUID(job.RunID.Bytes)) {
		return
	}
	switch job.Kind {
	case "extract":
		s.onAgentExtractResult(ctx, job) // record artifact, phase=load, dispatch load
	case "load":
		s.onAgentLoadResult(ctx, job)    // MarkRunSucceeded with row counts
	case "discover":
		s.onAgentDiscoverResult(ctx, job) // register streams into the catalog
	}
}

// ... CreateRunArtifact (upsert on run_id) + SetRunPhase("load") + SetRunTotal(rows)
load, dstID, ok := s.runLoadPlacement(ctx, job.WorkspaceID, runID)
if !ok {
	return
}
if !load.Cloud {
	// Dest is agent-placed: hand off to that agent's load job.
	s.q.CreateAgentJob(ctx, sqlc.CreateAgentJobParams{
		AgentID: load.AgentID, WorkspaceID: job.WorkspaceID, Kind: "load",
		RunID: pgtype.UUID{Bytes: runID, Valid: true},
		ConnectionID: pgtype.UUID{Bytes: dstID, Valid: true},
		Payload: []byte("{}"),
	})
	s.runEvent(ctx, runID, "info", "", "Loading staged data on the destination runtime")
	return
}
// Dest is cloud: the worker loads from the agent-staged blob.
s.enq.EnqueueMigrate(ctx, runID)
s.runEvent(ctx, runID, "info", "", "Loading staged data on Setu cloud")

-- name: CancelRunAgentJobs :execrows
UPDATE agent_job SET status = 'failed', error = 'cancelled', updated_at = now()
WHERE run_id = $1 AND status IN ('pending', 'leased');

// runEvent appends one row to a run's activity log (best-effort). Agent-executed
// runs bypass the worker's event stream, so the apiserver emits these itself at
// each dispatch/result transition — otherwise the run monitor shows nothing.
func (s *Server) runEvent(ctx context.Context, runID uuid.UUID, level, stream, msg string) {
	if _, err := s.q.AppendRunEvent(ctx, sqlc.AppendRunEventParams{
		RunID: runID, Level: level, Stream: stream, Message: msg,
	}); err != nil {
		s.log.Warn("append run event failed", "run", runID, "err", err)
	}
}