fix: harden credential and input handling

2026-08-10 16:15:41 +01:00
parent 00ef697b03
commit ab181453f0
19 changed files with 6100 additions and 0 deletions
@@ -0,0 +1,407 @@
+// Package jobs runs Relata's background work: a bounded pool of workers
+// draining a persistent, leased queue stored in MongoDB.
+//
+// The design goals are recoverability and boundedness. A crashed worker's job
+// returns to the queue when its lease expires; no work is ever started without
+// a worker slot; AI calls are additionally limited by their own semaphore; and
+// shutdown drains in-flight jobs rather than abandoning them.
+package jobs
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math/rand"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/andydixon/relata/internal/analysis"
+	"github.com/andydixon/relata/internal/config"
+	"github.com/andydixon/relata/internal/domain"
+	"github.com/andydixon/relata/internal/embeddings"
+	"github.com/andydixon/relata/internal/ids"
+	"github.com/andydixon/relata/internal/observability"
+)
+
+// Handler executes one job. It must honour ctx cancellation and should be
+// idempotent: a job may run more than once after a crash.
+type Handler func(ctx context.Context, j *domain.Job) error
+
+// Pool is a set of workers draining the queue.
+type Pool struct {
+	repo     domain.JobRepository
+	cfg      config.WorkerConfig
+	metrics  *observability.Metrics
+	handlers map[domain.JobType]Handler
+
+	// aiSem bounds concurrent AI provider calls independently of worker count,
+	// so a small local model is not swamped by every worker at once.
+	aiSem chan struct{}
+
+	workerID string
+	wg       sync.WaitGroup
+	stopOnce sync.Once
+	stopped  chan struct{}
+}
+
+// NewPool builds a worker pool. Handlers are registered before Start.
+func NewPool(repo domain.JobRepository, cfg config.WorkerConfig, m *observability.Metrics) *Pool {
+	host, _ := os.Hostname()
+	if host == "" {
+		host = "relata"
+	}
+	aiConc := cfg.AIConcurrency
+	if aiConc <= 0 {
+		aiConc = 1
+	}
+	return &Pool{
+		repo:     repo,
+		cfg:      cfg,
+		metrics:  m,
+		handlers: make(map[domain.JobType]Handler),
+		aiSem:    make(chan struct{}, aiConc),
+		workerID: fmt.Sprintf("%s-%s", host, ids.New()[:8]),
+		stopped:  make(chan struct{}),
+	}
+}
+
+// Register attaches a handler to a job type. It panics on an unknown type,
+// which can only be a programming error and is better caught at startup.
+func (p *Pool) Register(t domain.JobType, h Handler) {
+	if !t.Valid() {
+		panic("jobs: cannot register handler for unknown job type " + string(t))
+	}
+	p.handlers[t] = h
+}
+
+// AcquireAI blocks until an AI slot is free, returning a release function.
+//
+// Handlers call this around provider requests so that the number of concurrent
+// model calls stays bounded regardless of how many workers are running.
+func (p *Pool) AcquireAI(ctx context.Context) (release func(), err error) {
+	select {
+	case p.aiSem <- struct{}{}:
+		var once sync.Once
+		return func() { once.Do(func() { <-p.aiSem }) }, nil
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
+// Start launches the workers and the lease reaper.
+//
+// The returned context governs the pool's lifetime: cancelling it begins a
+// graceful shutdown in which no new jobs are claimed and in-flight jobs are
+// given until Stop's deadline to finish.
+func (p *Pool) Start(ctx context.Context) {
+	if !p.cfg.Enabled {
+		observability.LoggerFrom(ctx).Info("background workers are disabled by configuration")
+		return
+	}
+	n := p.cfg.Count
+	if n <= 0 {
+		n = 1
+	}
+
+	types := make([]domain.JobType, 0, len(p.handlers))
+	for t := range p.handlers {
+		types = append(types, t)
+	}
+
+	for i := 0; i < n; i++ {
+		p.wg.Add(1)
+		go func(idx int) {
+			defer p.wg.Done()
+			p.run(ctx, fmt.Sprintf("%s-w%d", p.workerID, idx), types)
+		}(i)
+	}
+
+	p.wg.Add(1)
+	go func() {
+		defer p.wg.Done()
+		p.reapLeases(ctx)
+	}()
+
+	observability.LoggerFrom(ctx).Info("background workers started",
+		"workers", n, "ai_concurrency", cap(p.aiSem), "job_types", len(types))
+}
+
+// Stop waits for in-flight jobs to finish, up to the given timeout.
+func (p *Pool) Stop(timeout time.Duration) {
+	p.stopOnce.Do(func() { close(p.stopped) })
+
+	done := make(chan struct{})
+	go func() {
+		p.wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(timeout):
+		// In-flight jobs keep their leases, which expire and return them to
+		// the queue for another worker to pick up.
+	}
+}
+
+// run is one worker's claim-execute-report loop.
+func (p *Pool) run(ctx context.Context, workerID string, types []domain.JobType) {
+	log := observability.LoggerFrom(ctx).With("worker", workerID)
+
+	// A small random offset stops every worker polling on the same tick.
+	jitter := time.Duration(rand.Int63n(int64(p.cfg.PollInterval) + 1))
+	timer := time.NewTimer(jitter)
+	defer timer.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-p.stopped:
+			return
+		case <-timer.C:
+		}
+
+		job, err := p.repo.Claim(ctx, workerID, types, p.cfg.LeaseDuration)
+		if err != nil {
+			if ctx.Err() == nil {
+				log.Error("could not claim a job", "error", err.Error())
+			}
+			timer.Reset(p.cfg.PollInterval)
+			continue
+		}
+		if job == nil {
+			timer.Reset(p.cfg.PollInterval)
+			continue
+		}
+
+		p.execute(ctx, workerID, job)
+		// Work was found, so look again immediately rather than idling.
+		timer.Reset(time.Millisecond)
+	}
+}
+
+// execute runs one job with its own timeout, heartbeat and outcome reporting.
+func (p *Pool) execute(ctx context.Context, workerID string, job *domain.Job) {
+	log := observability.LoggerFrom(ctx).With(
+		"worker", workerID, "job_id", job.ID, "job_type", string(job.Type),
+		"attempt", job.Attempts, "item_id", job.ItemID)
+
+	handler, ok := p.handlers[job.Type]
+	if !ok {
+		// Nothing can run this job, and retrying will not change that.
+		p.finish(ctx, workerID, job, fmt.Errorf("no handler registered for job type %q", job.Type),
+			domain.ErrClassPermanent, time.Now())
+		return
+	}
+
+	jobCtx, cancel := context.WithTimeout(observability.WithJobID(ctx, job.ID), p.cfg.JobTimeout)
+	defer cancel()
+
+	// Renew the lease while the job runs, so a long job is not reclaimed from
+	// under this worker. Losing the lease cancels the work: another worker has
+	// taken over and this attempt's result is no longer wanted.
+	hbDone := make(chan struct{})
+	go func() {
+		defer close(hbDone)
+		t := time.NewTicker(p.cfg.HeartbeatPeriod)
+		defer t.Stop()
+		for {
+			select {
+			case <-jobCtx.Done():
+				return
+			case <-t.C:
+				if err := p.repo.Heartbeat(jobCtx, job.ID, workerID, p.cfg.LeaseDuration); err != nil {
+					if errors.Is(err, domain.ErrConflict) {
+						log.Warn("lease lost; abandoning this attempt")
+						cancel()
+						return
+					}
+					log.Debug("heartbeat failed", "error", err.Error())
+				}
+			}
+		}
+	}()
+
+	started := time.Now()
+	err := runGuarded(jobCtx, handler, job)
+
+	// Classify before cancelling: the check distinguishes a job that ran out of
+	// time from one cancelled here, and cancelling first would conflate them.
+	class := Classify(err, jobCtx)
+
+	// Stop the heartbeat and wait for it to exit. The cancel must come before
+	// the wait — the goroutine's only exit is jobCtx being done, so waiting
+	// first would block until the deferred cancel that never runs.
+	cancel()
+	<-hbDone
+	p.finish(ctx, workerID, job, err, class, started)
+
+	if p.metrics != nil {
+		outcome := "success"
+		if err != nil {
+			outcome = string(class)
+		}
+		p.metrics.JobsProcessed.WithLabelValues(string(job.Type), outcome).Inc()
+		p.metrics.JobDuration.WithLabelValues(string(job.Type)).Observe(time.Since(started).Seconds())
+	}
+
+	if err != nil {
+		log.Error("job failed", "error", err.Error(), "class", string(class),
+			"duration_ms", time.Since(started).Milliseconds())
+	} else {
+		log.Info("job completed", "duration_ms", time.Since(started).Milliseconds())
+	}
+}
+
+// runGuarded executes a handler, converting a panic into an error so that one
+// malformed document cannot take down a worker.
+func runGuarded(ctx context.Context, h Handler, job *domain.Job) (err error) {
+	defer func() {
+		if recover() != nil {
+			// Never persist the panic value: handler panics can include document
+			// content or credentials. The job and request IDs provide correlation.
+			err = errors.New("job handler panicked")
+		}
+	}()
+	return h(ctx, job)
+}
+
+// finish reports the attempt's outcome and decides whether to retry.
+func (p *Pool) finish(ctx context.Context, workerID string, job *domain.Job, err error, class domain.ErrorClass, started time.Time) {
+	// Report using a fresh context: the job's context may already be cancelled,
+	// and the outcome must still be recorded.
+	repCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 15*time.Second)
+	defer cancel()
+
+	att := domain.JobAttempt{
+		Number:     job.Attempts,
+		StartedAt:  started,
+		EndedAt:    time.Now().UTC(),
+		Worker:     workerID,
+		DurationMS: time.Since(started).Milliseconds(),
+	}
+
+	if err == nil {
+		if cerr := p.repo.Complete(repCtx, job.ID, workerID, att); cerr != nil {
+			observability.LoggerFrom(ctx).Error("could not mark job complete",
+				"job_id", job.ID, "error", cerr.Error())
+		}
+		return
+	}
+
+	att.Error = truncateError(err.Error())
+	att.ErrorClass = string(class)
+
+	maxAttempts := job.MaxAttempts
+	if maxAttempts <= 0 {
+		maxAttempts = p.cfg.MaxAttempts
+	}
+
+	// Only transient failures are retried, and only within the attempt budget.
+	// A permanent failure — invalid model output, an unparseable document, a
+	// missing handler — is terminal: repeating it would waste the queue.
+	var retryAfter *time.Time
+	if class == domain.ErrClassTransient && job.Attempts < maxAttempts {
+		delay := backoff(job.Attempts)
+		t := time.Now().UTC().Add(delay)
+		retryAfter = &t
+	}
+
+	if ferr := p.repo.Fail(repCtx, job.ID, workerID, att, retryAfter); ferr != nil {
+		observability.LoggerFrom(ctx).Error("could not record job failure",
+			"job_id", job.ID, "error", ferr.Error())
+	}
+}
+
+// backoff returns the delay before attempt n+1, with jitter.
+func backoff(attempt int) time.Duration {
+	base := 2 * time.Second
+	d := base << minInt(attempt, 6) // capped so the shift cannot overflow
+	if d > 5*time.Minute {
+		d = 5 * time.Minute
+	}
+	return d/2 + time.Duration(rand.Int63n(int64(d/2)+1))
+}
+
+func minInt(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+// Classify decides whether a failure is worth retrying.
+func Classify(err error, ctx context.Context) domain.ErrorClass {
+	if err == nil {
+		return ""
+	}
+	switch {
+	case errors.Is(err, context.Canceled):
+		return domain.ErrClassCancelled
+	case errors.Is(err, context.DeadlineExceeded):
+		// A job that ran out of time may succeed with a fresh budget.
+		return domain.ErrClassTimeout
+	case errors.Is(err, analysis.ErrInvalidOutput):
+		// Retrying a model that produced structurally invalid output tends to
+		// produce more of the same; the failure is recorded instead.
+		return domain.ErrClassInvalid
+	case errors.Is(err, analysis.ErrDisabled):
+		return domain.ErrClassPermanent
+	case embeddings.IsPermanent(err):
+		return domain.ErrClassPermanent
+	case errors.Is(err, domain.ErrNotFound),
+		errors.Is(err, domain.ErrInvalidInput),
+		errors.Is(err, domain.ErrUnsupported),
+		errors.Is(err, domain.ErrWorkspaceMiss):
+		return domain.ErrClassPermanent
+	case embeddings.Retryable(err):
+		return domain.ErrClassTransient
+	}
+	// An unrecognised error is treated as transient: the attempt budget bounds
+	// the cost, and a genuinely permanent failure will exhaust it quickly.
+	return domain.ErrClassTransient
+}
+
+// reapLeases periodically returns jobs from crashed workers to the queue and
+// refreshes the queue-depth gauges.
+func (p *Pool) reapLeases(ctx context.Context) {
+	interval := p.cfg.LeaseDuration / 2
+	if interval < 10*time.Second {
+		interval = 10 * time.Second
+	}
+	t := time.NewTicker(interval)
+	defer t.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-p.stopped:
+			return
+		case <-t.C:
+			released, err := p.repo.ReleaseExpired(ctx, time.Now().UTC())
+			if err != nil {
+				if ctx.Err() == nil {
+					observability.LoggerFrom(ctx).Error("could not release expired leases", "error", err.Error())
+				}
+				continue
+			}
+			if released > 0 {
+				observability.LoggerFrom(ctx).Warn(
+					"returned jobs from expired leases to the queue", "count", released)
+			}
+		}
+	}
+}
+
+// truncateError bounds an error string before it is stored. Error text can
+// include untrusted content, so its length is capped.
+func truncateError(s string) string {
+	const max = 2000
+	if len(s) <= max {
+		return s
+	}
+	return s[:max] + "…"
+}
@@ -0,0 +1,427 @@
+package jobs
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/andydixon/relata/internal/analysis"
+	"github.com/andydixon/relata/internal/config"
+	"github.com/andydixon/relata/internal/domain"
+	"github.com/andydixon/relata/internal/embeddings"
+)
+
+// fakeQueue is an in-memory job repository for exercising the worker loop.
+type fakeQueue struct {
+	mu         sync.Mutex
+	jobs       map[string]*domain.Job
+	order      []string
+	completed  []string
+	failed     []string
+	retried    []string
+	heartbeats int
+}
+
+func newFakeQueue() *fakeQueue {
+	return &fakeQueue{jobs: map[string]*domain.Job{}}
+}
+
+func (q *fakeQueue) Enqueue(_ context.Context, j *domain.Job) error {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	q.jobs[j.ID] = j
+	q.order = append(q.order, j.ID)
+	return nil
+}
+
+func (q *fakeQueue) Claim(_ context.Context, owner string, _ []domain.JobType, lease time.Duration) (*domain.Job, error) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	for _, id := range q.order {
+		j := q.jobs[id]
+		if j.Status != domain.JobPending {
+			continue
+		}
+		j.Status = domain.JobRunning
+		j.LeaseOwner = owner
+		j.Attempts++
+		exp := time.Now().Add(lease)
+		j.LeaseExpires = &exp
+		cp := *j
+		return &cp, nil
+	}
+	return nil, nil
+}
+
+func (q *fakeQueue) Heartbeat(_ context.Context, id, owner string, _ time.Duration) error {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	q.heartbeats++
+	j := q.jobs[id]
+	if j == nil || j.LeaseOwner != owner {
+		return domain.ErrConflict
+	}
+	return nil
+}
+
+func (q *fakeQueue) Complete(_ context.Context, id, owner string, _ domain.JobAttempt) error {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	j := q.jobs[id]
+	if j == nil || j.LeaseOwner != owner {
+		return domain.ErrConflict
+	}
+	j.Status = domain.JobSucceeded
+	q.completed = append(q.completed, id)
+	return nil
+}
+
+func (q *fakeQueue) Fail(_ context.Context, id, owner string, att domain.JobAttempt, retryAfter *time.Time) error {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	j := q.jobs[id]
+	if j == nil || j.LeaseOwner != owner {
+		return domain.ErrConflict
+	}
+	if retryAfter != nil {
+		j.Status = domain.JobPending
+		q.retried = append(q.retried, id)
+	} else {
+		j.Status = domain.JobFailed
+		q.failed = append(q.failed, id)
+	}
+	return nil
+}
+
+func (q *fakeQueue) Get(_ context.Context, id string) (*domain.Job, error) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	if j := q.jobs[id]; j != nil {
+		cp := *j
+		return &cp, nil
+	}
+	return nil, domain.ErrNotFound
+}
+
+func (q *fakeQueue) List(context.Context, string, []domain.JobStatus, []domain.JobType, domain.Page) ([]*domain.Job, int64, error) {
+	return nil, 0, nil
+}
+func (q *fakeQueue) Retry(context.Context, string, string) (*domain.Job, error) { return nil, nil }
+func (q *fakeQueue) CountByStatus(context.Context, string) (map[domain.JobStatus]int64, error) {
+	return nil, nil
+}
+func (q *fakeQueue) ReleaseExpired(context.Context, time.Time) (int64, error) { return 0, nil }
+func (q *fakeQueue) DeleteForItem(context.Context, string, string) (int64, error) {
+	return 0, nil
+}
+
+func (q *fakeQueue) counts() (done, failed, retried int) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	return len(q.completed), len(q.failed), len(q.retried)
+}
+
+func testConfig() config.WorkerConfig {
+	return config.WorkerConfig{
+		Count:           2,
+		AIConcurrency:   1,
+		PollInterval:    10 * time.Millisecond,
+		LeaseDuration:   2 * time.Second,
+		HeartbeatPeriod: 200 * time.Millisecond,
+		MaxAttempts:     3,
+		JobTimeout:      5 * time.Second,
+		Enabled:         true,
+	}
+}
+
+func queueJob(q *fakeQueue, id string, t domain.JobType) {
+	_ = q.Enqueue(context.Background(), &domain.Job{
+		ID: id, WorkspaceID: "ws", Type: t, Status: domain.JobPending,
+		MaxAttempts: 3, RunAfter: time.Now(),
+	})
+}
+
+// waitFor polls cond until it holds or the deadline passes.
+func waitFor(t *testing.T, timeout time.Duration, what string, cond func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		if cond() {
+			return
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	t.Fatalf("timed out after %s waiting for %s", timeout, what)
+}
+
+// This is the regression test for a deadlock in which the worker waited for its
+// heartbeat goroutine to exit, while that goroutine waited for a context that
+// was only cancelled once the worker returned. Every worker wedged after one
+// job and the queue stopped draining.
+func TestWorkerCompletesJobsAndDoesNotDeadlock(t *testing.T) {
+	q := newFakeQueue()
+	pool := NewPool(q, testConfig(), nil)
+
+	var ran int
+	var mu sync.Mutex
+	pool.Register(domain.JobParseItem, func(ctx context.Context, j *domain.Job) error {
+		mu.Lock()
+		ran++
+		mu.Unlock()
+		return nil
+	})
+
+	// More jobs than workers: if a worker wedges after its first job, the
+	// remainder never complete.
+	for i := 0; i < 8; i++ {
+		queueJob(q, string(rune('a'+i)), domain.JobParseItem)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	pool.Start(ctx)
+	defer pool.Stop(2 * time.Second)
+
+	waitFor(t, 5*time.Second, "all 8 jobs to complete", func() bool {
+		done, _, _ := q.counts()
+		return done == 8
+	})
+
+	mu.Lock()
+	defer mu.Unlock()
+	if ran != 8 {
+		t.Fatalf("handler ran %d times, want 8", ran)
+	}
+}
+
+func TestWorkerRetriesTransientFailures(t *testing.T) {
+	q := newFakeQueue()
+	pool := NewPool(q, testConfig(), nil)
+
+	pool.Register(domain.JobParseItem, func(context.Context, *domain.Job) error {
+		// A network-shaped error is transient and should be retried.
+		return &embeddings.StatusError{Status: 503, Op: "test", Body: "unavailable"}
+	})
+	queueJob(q, "j1", domain.JobParseItem)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	pool.Start(ctx)
+	defer pool.Stop(2 * time.Second)
+
+	waitFor(t, 5*time.Second, "the job to be retried", func() bool {
+		_, _, retried := q.counts()
+		return retried >= 1
+	})
+}
+
+func TestWorkerDoesNotRetryPermanentFailures(t *testing.T) {
+	q := newFakeQueue()
+	pool := NewPool(q, testConfig(), nil)
+
+	pool.Register(domain.JobParseItem, func(context.Context, *domain.Job) error {
+		// Invalid model output will never become valid by repeating the call.
+		return analysis.Reject("task", analysis.ReasonBadEnum, "nope", "raw")
+	})
+	queueJob(q, "j1", domain.JobParseItem)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	pool.Start(ctx)
+	defer pool.Stop(2 * time.Second)
+
+	waitFor(t, 5*time.Second, "the job to fail terminally", func() bool {
+		_, failed, _ := q.counts()
+		return failed == 1
+	})
+	if _, _, retried := q.counts(); retried != 0 {
+		t.Fatalf("a permanently invalid job was retried %d times", retried)
+	}
+}
+
+func TestWorkerSurvivesAPanickingHandler(t *testing.T) {
+	q := newFakeQueue()
+	pool := NewPool(q, testConfig(), nil)
+
+	pool.Register(domain.JobParseItem, func(context.Context, *domain.Job) error {
+		panic("a malformed document exploded")
+	})
+	pool.Register(domain.JobChunkItem, func(context.Context, *domain.Job) error { return nil })
+
+	queueJob(q, "boom", domain.JobParseItem)
+	queueJob(q, "fine", domain.JobChunkItem)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	pool.Start(ctx)
+	defer pool.Stop(2 * time.Second)
+
+	// The panicking job must not take the worker down: the other job still runs.
+	waitFor(t, 5*time.Second, "the healthy job to complete despite a panic", func() bool {
+		done, _, _ := q.counts()
+		return done == 1
+	})
+}
+
+func TestRunGuardedDoesNotPersistPanicValue(t *testing.T) {
+	err := runGuarded(context.Background(), func(context.Context, *domain.Job) error {
+		panic("bearer credential rlt_v1_secret")
+	}, &domain.Job{ID: "job-1"})
+	if err == nil {
+		t.Fatal("panic was not converted to an error")
+	}
+	if strings.Contains(err.Error(), "rlt_v1_secret") {
+		t.Fatalf("panic value reached the persistent job error: %q", err)
+	}
+}
+
+func TestUnregisteredJobTypeFailsPermanently(t *testing.T) {
+	q := newFakeQueue()
+	pool := NewPool(q, testConfig(), nil)
+	pool.Register(domain.JobParseItem, func(context.Context, *domain.Job) error { return nil })
+
+	// A job whose type nothing handles must not spin in the queue.
+	queueJob(q, "orphan", domain.JobExtractClaims)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	pool.Start(ctx)
+	defer pool.Stop(2 * time.Second)
+
+	waitFor(t, 5*time.Second, "the unhandled job to fail", func() bool {
+		_, failed, _ := q.counts()
+		return failed == 1
+	})
+}
+
+func TestAIConcurrencyIsBounded(t *testing.T) {
+	cfg := testConfig()
+	cfg.Count = 4
+	cfg.AIConcurrency = 2
+
+	q := newFakeQueue()
+	pool := NewPool(q, cfg, nil)
+
+	var mu sync.Mutex
+	var inFlight, peak int
+
+	pool.Register(domain.JobEmbedItem, func(ctx context.Context, j *domain.Job) error {
+		release, err := pool.AcquireAI(ctx)
+		if err != nil {
+			return err
+		}
+		defer release()
+
+		mu.Lock()
+		inFlight++
+		if inFlight > peak {
+			peak = inFlight
+		}
+		mu.Unlock()
+
+		time.Sleep(80 * time.Millisecond)
+
+		mu.Lock()
+		inFlight--
+		mu.Unlock()
+		return nil
+	})
+
+	for i := 0; i < 8; i++ {
+		queueJob(q, string(rune('a'+i)), domain.JobEmbedItem)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	pool.Start(ctx)
+	defer pool.Stop(3 * time.Second)
+
+	waitFor(t, 8*time.Second, "all embedding jobs to finish", func() bool {
+		done, _, _ := q.counts()
+		return done == 8
+	})
+
+	mu.Lock()
+	defer mu.Unlock()
+	if peak > cfg.AIConcurrency {
+		t.Fatalf("%d concurrent AI calls observed, limit is %d", peak, cfg.AIConcurrency)
+	}
+	if peak < 2 {
+		t.Fatalf("AI calls never overlapped (peak %d); the semaphore may be serialising everything", peak)
+	}
+}
+
+func TestStopDrainsInFlightWork(t *testing.T) {
+	q := newFakeQueue()
+	pool := NewPool(q, testConfig(), nil)
+
+	started := make(chan struct{})
+	var once sync.Once
+	pool.Register(domain.JobParseItem, func(ctx context.Context, j *domain.Job) error {
+		once.Do(func() { close(started) })
+		time.Sleep(150 * time.Millisecond)
+		return nil
+	})
+	queueJob(q, "j1", domain.JobParseItem)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	pool.Start(ctx)
+
+	<-started
+	pool.Stop(3 * time.Second)
+
+	// A job that was running when shutdown began should have been finished,
+	// not abandoned.
+	done, _, _ := q.counts()
+	if done != 1 {
+		t.Fatalf("in-flight job was not completed during shutdown (completed=%d)", done)
+	}
+}
+
+func TestClassifyDistinguishesFailureKinds(t *testing.T) {
+	ctx := context.Background()
+
+	cases := []struct {
+		name string
+		err  error
+		want domain.ErrorClass
+	}{
+		{"nil", nil, ""},
+		{"cancelled", context.Canceled, domain.ErrClassCancelled},
+		{"timeout", context.DeadlineExceeded, domain.ErrClassTimeout},
+		{"invalid model output", analysis.Reject("t", "r", "d", ""), domain.ErrClassInvalid},
+		{"analysis disabled", analysis.ErrDisabled, domain.ErrClassPermanent},
+		{"not found", domain.ErrNotFound, domain.ErrClassPermanent},
+		{"invalid input", domain.ErrInvalidInput, domain.ErrClassPermanent},
+		{"permanent provider error", embeddings.Permanent(errors.New("bad model")), domain.ErrClassPermanent},
+		{"503 from provider", &embeddings.StatusError{Status: 503}, domain.ErrClassTransient},
+		{"unknown", errors.New("something odd"), domain.ErrClassTransient},
+	}
+	for _, c := range cases {
+		if got := Classify(c.err, ctx); got != c.want {
+			t.Errorf("Classify(%s) = %q, want %q", c.name, got, c.want)
+		}
+	}
+}
+
+func TestBackoffIsBoundedAndJittered(t *testing.T) {
+	seen := map[time.Duration]bool{}
+	for attempt := 1; attempt <= 12; attempt++ {
+		for i := 0; i < 20; i++ {
+			d := backoff(attempt)
+			if d < 0 || d > 5*time.Minute {
+				t.Fatalf("backoff(%d) = %v, outside the expected bounds", attempt, d)
+			}
+			seen[d] = true
+		}
+	}
+	// Jitter means repeated calls must not all return the same value, or a
+	// provider restart would be met by every worker retrying in lockstep.
+	if len(seen) < 10 {
+		t.Fatalf("backoff produced only %d distinct delays; jitter is not being applied", len(seen))
+	}
+}