feat(postgres): GetStepProgress, EnqueueTx, UpdateRunSpec

experimental/store/postgres/queue.go

@@ -12,8 +12,37 @@ import (
 	"github.com/deepnoodle-ai/workflow/experimental/worker"
 )
 
-// Enqueue implements worker.QueueStore.
+// Enqueue implements worker.QueueStore. The insert runs in its own
+// connection. When the insert must be atomic with writes to adjacent
+// tables (credit ledger, idempotency keys, audit records, …), use
+// EnqueueTx inside a caller-owned transaction instead.
 func (s *Store) Enqueue(ctx context.Context, run worker.NewRun) error {
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return fmt.Errorf("postgres: begin enqueue tx: %w", err)
+	}
+	defer func() { _ = tx.Rollback(ctx) }()
+	if err := s.EnqueueTx(ctx, tx, run); err != nil {
+		return err
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return fmt.Errorf("postgres: commit enqueue %s: %w", run.ID, err)
+	}
+	return nil
+}
+
+// EnqueueTx inserts a queued run inside a caller-provided pgx
+// transaction. The caller owns the tx lifecycle (Begin, Commit,
+// Rollback). Use this when the run insert must be atomic with writes
+// to tables outside the store's schema — e.g., debiting a credit
+// ledger and creating the run in one commit.
+//
+// The tx must be against the same database as the Store's pool; the
+// library does not verify this.
+func (s *Store) EnqueueTx(ctx context.Context, tx pgx.Tx, run worker.NewRun) error {
+	if tx == nil {
+		return fmt.Errorf("postgres: EnqueueTx requires a non-nil tx")
+	}
 	if run.ID == "" {
 		return fmt.Errorf("postgres: NewRun.ID is required")
 	}
@@ -29,7 +58,7 @@ func (s *Store) Enqueue(ctx context.Context, run worker.NewRun) error {
 		)
 		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
 	`, s.t("workflow_runs"))
-	if _, err := s.pool.Exec(ctx, query,
+	if _, err := tx.Exec(ctx, query,
 		run.ID, run.Spec, string(worker.StatusQueued),
 		nullableString(run.OrgID),
 		nullableString(run.ProjectID),
@@ -188,6 +217,39 @@ func (s *Store) Complete(ctx context.Context, claim *worker.Claim, outcome worke
 	return nil
 }
 
+// UpdateRunSpec replaces the spec on a running claim. It fences on
+// (claim_id, worker_id, attempt) and status = running, and returns
+// ErrLeaseLost if the fence fails — matching Heartbeat and Complete.
+//
+// Use this during long-running activities that mutate the run spec
+// incrementally (e.g., a KB-apply loop persisting progress between
+// steps) and need the update durable without waiting for the next
+// checkpoint. The caller retains responsibility for producing a
+// valid spec; the store does not inspect it.
+func (s *Store) UpdateRunSpec(ctx context.Context, claim *worker.Claim, spec []byte) error {
+	if claim == nil {
+		return fmt.Errorf("postgres: UpdateRunSpec requires a non-nil claim")
+	}
+	query := fmt.Sprintf(`
+		UPDATE %s
+		SET spec = $1
+		WHERE id         = $2
+		  AND claimed_by = $3
+		  AND attempt    = $4
+		  AND status     = $5
+	`, s.t("workflow_runs"))
+	tag, err := s.pool.Exec(ctx, query,
+		spec, claim.ID, claim.WorkerID, claim.Attempt, string(worker.StatusRunning),
+	)
+	if err != nil {
+		return fmt.Errorf("postgres: update run spec %s: %w", claim.ID, err)
+	}
+	if tag.RowsAffected() == 0 {
+		return worker.ErrLeaseLost
+	}
+	return nil
+}
+
 // ReclaimStale implements worker.QueueStore.
 func (s *Store) ReclaimStale(ctx context.Context, staleBefore time.Time, maxAttempts int, excludeIDs []string) (int, error) {
 	query := fmt.Sprintf(`

experimental/store/postgres/step_progress.go

@@ -57,6 +57,64 @@ func (s *Store) UpdateStepProgress(ctx context.Context, executionID string, p wo
 	return nil
 }
 
+// GetStepProgress returns every step progress row recorded for an
+// execution, ordered by started_at (NULLS LAST) then step_name. One
+// row per (step_name, branch_id). Returns an empty slice if no rows
+// exist. Use this on the read side to render per-step status for a
+// run whose identity came back from runquery.Store.GetRun, which
+// intentionally does not carry step progress.
+func (s *Store) GetStepProgress(ctx context.Context, executionID string) ([]workflow.StepProgress, error) {
+	query := fmt.Sprintf(`
+		SELECT step_name, branch_id, status, activity, attempt,
+		       detail, started_at, finished_at, error
+		FROM %s
+		WHERE execution_id = $1
+		ORDER BY started_at ASC NULLS LAST, step_name ASC, branch_id ASC
+	`, s.t("workflow_step_progress"))
+
+	rows, err := s.pool.Query(ctx, query, executionID)
+	if err != nil {
+		return nil, fmt.Errorf("postgres: query step progress %s: %w", executionID, err)
+	}
+	defer rows.Close()
+
+	var out []workflow.StepProgress
+	for rows.Next() {
+		var (
+			p          workflow.StepProgress
+			status     string
+			detail     []byte
+			startedAt  *time.Time
+			finishedAt *time.Time
+		)
+		if err := rows.Scan(
+			&p.StepName, &p.BranchID, &status, &p.ActivityName, &p.Attempt,
+			&detail, &startedAt, &finishedAt, &p.Error,
+		); err != nil {
+			return nil, fmt.Errorf("postgres: scan step progress: %w", err)
+		}
+		p.Status = workflow.StepStatus(status)
+		if len(detail) > 0 {
+			pd := &workflow.ProgressDetail{}
+			if err := json.Unmarshal(detail, pd); err != nil {
+				return nil, fmt.Errorf("postgres: unmarshal progress detail for %s/%s: %w", executionID, p.StepName, err)
+			}
+			p.Detail = pd
+		}
+		if startedAt != nil {
+			p.StartedAt = *startedAt
+		}
+		if finishedAt != nil {
+			p.FinishedAt = *finishedAt
+		}
+		out = append(out, p)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("postgres: iterate step progress %s: %w", executionID, err)
+	}
+	return out, nil
+}
+
 func nullTime(t time.Time) any {
 	if t.IsZero() {
 		return nil

experimental/store/postgres/store_test.go

@@ -2,6 +2,7 @@ package postgres_test
 
 import (
 	"context"
+	"errors"
 	"os"
 	"testing"
 	"time"
@@ -11,6 +12,7 @@ import (
 	"github.com/deepnoodle-ai/workflow"
 	"github.com/deepnoodle-ai/workflow/experimental/store/postgres"
 	"github.com/deepnoodle-ai/workflow/experimental/worker"
+	"github.com/deepnoodle-ai/workflow/experimental/worker/runquery"
 )
 
 // These tests require a real Postgres instance. Set WORKFLOW_PG_DSN
@@ -269,3 +271,209 @@ func TestStore_StepProgressAndActivityLog(t *testing.T) {
 		t.Fatalf("unexpected history: %+v", hist)
 	}
 }
+
+func TestStore_GetStepProgress(t *testing.T) {
+	store, _ := openTestStore(t)
+	ctx := context.Background()
+
+	// Empty executions return an empty slice, not an error.
+	got, err := store.GetStepProgress(ctx, "nobody")
+	if err != nil {
+		t.Fatalf("empty get: %v", err)
+	}
+	if len(got) != 0 {
+		t.Fatalf("expected empty slice, got %+v", got)
+	}
+
+	t0 := time.Date(2026, 4, 12, 10, 0, 0, 0, time.UTC)
+	rows := []workflow.StepProgress{
+		{
+			StepName:     "second",
+			BranchID:     "main",
+			Status:       workflow.StepStatusCompleted,
+			ActivityName: "print",
+			Attempt:      1,
+			StartedAt:    t0.Add(5 * time.Second),
+			FinishedAt:   t0.Add(6 * time.Second),
+		},
+		{
+			StepName:     "first",
+			BranchID:     "main",
+			Status:       workflow.StepStatusCompleted,
+			ActivityName: "print",
+			Attempt:      1,
+			StartedAt:    t0,
+			FinishedAt:   t0.Add(1 * time.Second),
+			Detail: &workflow.ProgressDetail{
+				Message: "halfway",
+				Data:    map[string]any{"pct": float64(50)},
+			},
+		},
+		{
+			StepName:     "pending-step",
+			BranchID:     "main",
+			Status:       workflow.StepStatusPending,
+			ActivityName: "print",
+			Attempt:      0,
+		},
+	}
+	for _, p := range rows {
+		if err := store.UpdateStepProgress(ctx, "run-sp", p); err != nil {
+			t.Fatalf("update: %v", err)
+		}
+	}
+
+	// Progress for a different execution must not leak.
+	if err := store.UpdateStepProgress(ctx, "other-run", workflow.StepProgress{
+		StepName: "x", BranchID: "main", Status: workflow.StepStatusRunning, ActivityName: "print", Attempt: 1,
+		StartedAt: t0,
+	}); err != nil {
+		t.Fatalf("update other: %v", err)
+	}
+
+	got, err = store.GetStepProgress(ctx, "run-sp")
+	if err != nil {
+		t.Fatalf("get: %v", err)
+	}
+	if len(got) != 3 {
+		t.Fatalf("expected 3 rows, got %d: %+v", len(got), got)
+	}
+	// Ordered by started_at NULLS LAST, step_name, branch_id:
+	// "first" (t0) → "second" (t0+5s) → "pending-step" (NULL).
+	if got[0].StepName != "first" || got[1].StepName != "second" || got[2].StepName != "pending-step" {
+		t.Fatalf("unexpected order: [%s, %s, %s]", got[0].StepName, got[1].StepName, got[2].StepName)
+	}
+	// Detail round-trips.
+	if got[0].Detail == nil || got[0].Detail.Message != "halfway" || got[0].Detail.Data["pct"] != float64(50) {
+		t.Fatalf("detail round-trip mismatch: %+v", got[0].Detail)
+	}
+	// Pending row has zero started_at.
+	if !got[2].StartedAt.IsZero() || !got[2].FinishedAt.IsZero() {
+		t.Fatalf("expected zero times on pending row: %+v", got[2])
+	}
+	// Status parses back correctly.
+	if got[0].Status != workflow.StepStatusCompleted || got[2].Status != workflow.StepStatusPending {
+		t.Fatalf("status mismatch: %+v", got)
+	}
+}
+
+func TestStore_EnqueueTxRollbackAndCommit(t *testing.T) {
+	store, pool := openTestStore(t)
+	ctx := context.Background()
+
+	// Rollback path: an EnqueueTx inside a tx that the caller
+	// aborts must leave no row behind.
+	tx, err := pool.Begin(ctx)
+	if err != nil {
+		t.Fatalf("begin: %v", err)
+	}
+	if err := store.EnqueueTx(ctx, tx, worker.NewRun{ID: "rb-run", Spec: []byte(`{}`)}); err != nil {
+		_ = tx.Rollback(ctx)
+		t.Fatalf("enqueue tx: %v", err)
+	}
+	if err := tx.Rollback(ctx); err != nil {
+		t.Fatalf("rollback: %v", err)
+	}
+	if _, err := store.GetRun(ctx, "", "rb-run"); !errors.Is(err, runquery.ErrRunNotFound) {
+		t.Fatalf("expected not found after rollback, got err=%v", err)
+	}
+
+	// Commit path: EnqueueTx + another write in the same tx are
+	// both visible after commit.
+	tx2, err := pool.Begin(ctx)
+	if err != nil {
+		t.Fatalf("begin 2: %v", err)
+	}
+	if err := store.EnqueueTx(ctx, tx2, worker.NewRun{
+		ID: "cm-run", Spec: []byte(`{}`), OrgID: "org-x",
+	}); err != nil {
+		_ = tx2.Rollback(ctx)
+		t.Fatalf("enqueue tx 2: %v", err)
+	}
+	// Write to an adjacent row inside the same tx to simulate a
+	// credit ledger debit: both rows must commit atomically.
+	if _, err := tx2.Exec(ctx, `
+		INSERT INTO workflow_credit_ledger (id, org_id, run_id, amount, reason, created_at)
+		VALUES ('led-1', 'org-x', 'cm-run', -5, 'debit', NOW())
+	`); err != nil {
+		_ = tx2.Rollback(ctx)
+		t.Fatalf("ledger insert: %v", err)
+	}
+	if err := tx2.Commit(ctx); err != nil {
+		t.Fatalf("commit: %v", err)
+	}
+
+	run, err := store.GetRun(ctx, "org-x", "cm-run")
+	if err != nil {
+		t.Fatalf("GetRun: %v", err)
+	}
+	if run.ID != "cm-run" || run.OrgID != "org-x" {
+		t.Fatalf("unexpected run: %+v", run)
+	}
+	var ledgerCount int
+	if err := pool.QueryRow(ctx, `SELECT count(*) FROM workflow_credit_ledger WHERE run_id = $1`, "cm-run").Scan(&ledgerCount); err != nil {
+		t.Fatalf("ledger count: %v", err)
+	}
+	if ledgerCount != 1 {
+		t.Fatalf("expected 1 ledger row, got %d", ledgerCount)
+	}
+
+	// Nil tx is rejected.
+	if err := store.EnqueueTx(ctx, nil, worker.NewRun{ID: "nil-tx", Spec: []byte(`{}`)}); err == nil {
+		t.Fatalf("expected error for nil tx")
+	}
+}
+
+func TestStore_UpdateRunSpecFencing(t *testing.T) {
+	store, pool := openTestStore(t)
+	ctx := context.Background()
+
+	if err := store.Enqueue(ctx, worker.NewRun{ID: "spec-run", Spec: []byte(`{"v":1}`)}); err != nil {
+		t.Fatalf("enqueue: %v", err)
+	}
+	claim, err := store.ClaimQueued(ctx, "w1")
+	if err != nil || claim == nil {
+		t.Fatalf("claim: %v / %+v", err, claim)
+	}
+
+	// Happy path: same lease updates the spec.
+	newSpec := []byte(`{"v":2}`)
+	if err := store.UpdateRunSpec(ctx, claim, newSpec); err != nil {
+		t.Fatalf("update spec: %v", err)
+	}
+	var stored []byte
+	if err := pool.QueryRow(ctx, `SELECT spec FROM workflow_runs WHERE id = $1`, "spec-run").Scan(&stored); err != nil {
+		t.Fatalf("read spec: %v", err)
+	}
+	if string(stored) != `{"v":2}` {
+		t.Fatalf("spec not updated: %s", stored)
+	}
+
+	// Wrong worker rejected.
+	wrong := *claim
+	wrong.WorkerID = "w2"
+	if err := store.UpdateRunSpec(ctx, &wrong, []byte(`{"v":3}`)); err != worker.ErrLeaseLost {
+		t.Fatalf("wrong worker: expected ErrLeaseLost, got %v", err)
+	}
+
+	// Wrong attempt rejected.
+	badAttempt := *claim
+	badAttempt.Attempt = 99
+	if err := store.UpdateRunSpec(ctx, &badAttempt, []byte(`{"v":3}`)); err != worker.ErrLeaseLost {
+		t.Fatalf("wrong attempt: expected ErrLeaseLost, got %v", err)
+	}
+
+	// After completion, the run is no longer running, so updates
+	// must fail even with the right lease.
+	if err := store.Complete(ctx, claim, worker.Outcome{Status: worker.StatusCompleted}); err != nil {
+		t.Fatalf("complete: %v", err)
+	}
+	if err := store.UpdateRunSpec(ctx, claim, []byte(`{"v":4}`)); err != worker.ErrLeaseLost {
+		t.Fatalf("after complete: expected ErrLeaseLost, got %v", err)
+	}
+
+	// Nil claim rejected.
+	if err := store.UpdateRunSpec(ctx, nil, []byte(`{"v":5}`)); err == nil {
+		t.Fatalf("expected error for nil claim")
+	}
+}