notarius/internal/framework/pipeline/runner_concurrency_test.go

package pipeline

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	frameworkllm "gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
)

type countingProvider struct {
	active  atomic.Int32
	maximum atomic.Int32
	calls   atomic.Int32
}

func (p *countingProvider) CompleteStructured(ctx context.Context, _ contracts.StructuredCompletionRequest, _ any) (contracts.StructuredCompletionResponse, error) {
	p.calls.Add(1)
	current := p.active.Add(1)
	defer p.active.Add(-1)
	for {
		seen := p.maximum.Load()
		if current <= seen || p.maximum.CompareAndSwap(seen, current) {
			break
		}
	}
	select {
	case <-time.After(5 * time.Millisecond):
		return contracts.StructuredCompletionResponse{}, nil
	case <-ctx.Done():
		return contracts.StructuredCompletionResponse{}, ctx.Err()
	}
}

func preparedConcurrentPipeline(t *testing.T, chunkCount int) *PreparedPipeline {
	t.Helper()
	catalog := typedResolutionCatalog(t, completeTypedCatalogOptions())
	resolved, err := ResolvePipeline(typedResolutionProfile(), ResolveOptions{}, catalog)
	if err != nil {
		t.Fatalf("ResolvePipeline() error = %v", err)
	}
	prepared, err := Prepare(resolved, registriesFromModuleCatalog(catalog), ModuleDependencies{})
	if err != nil {
		t.Fatalf("Prepare() error = %v", err)
	}
	doc := typedTestDocumentWithUnits(chunkCount)
	if adapter, ok := prepared.input.(*typedTestInput); ok {
		adapter.doc = doc
	} else {
		t.Fatalf("prepared input = %T, want *typedTestInput", prepared.input)
	}
	prepared.chunker = &typedTestChunker{key: "typed/chunk", plan: typedTestPlan(doc)}
	return prepared
}

type orderedLaneSpec struct {
	id      string
	profile string
}

func preparedOrderedPipeline(t *testing.T, chunkCount int, specs ...orderedLaneSpec) *PreparedPipeline {
	t.Helper()
	catalog := typedResolutionCatalog(t, completeTypedCatalogOptions())
	profile := typedResolutionProfile()
	profile.Artifacts = nil
	for index, spec := range specs {
		lane, ok := typedResolutionProfile().Artifacts[spec.profile]
		if !ok {
			t.Fatalf("typed lane profile %q is not defined", spec.profile)
		}
		profile.Steps = append(profile.Steps, PipelineStepProfile{
			ID:        fmt.Sprintf("step-%d", index+1),
			Artifacts: map[string]ArtifactLaneProfile{spec.id: lane},
		})
	}
	resolved, err := ResolvePipeline(profile, ResolveOptions{}, catalog)
	if err != nil {
		t.Fatalf("ResolvePipeline() error = %v", err)
	}
	prepared, err := Prepare(resolved, registriesFromModuleCatalog(catalog), ModuleDependencies{})
	if err != nil {
		t.Fatalf("Prepare() error = %v", err)
	}
	doc := typedTestDocumentWithUnits(chunkCount)
	if adapter, ok := prepared.input.(*typedTestInput); ok {
		adapter.doc = doc
	} else {
		t.Fatalf("prepared input = %T, want *typedTestInput", prepared.input)
	}
	prepared.chunker = &typedTestChunker{key: "typed/chunk", plan: typedTestPlan(doc)}
	return prepared
}

type countingOrderedOutput struct {
	calls atomic.Int32
}

func (o *countingOrderedOutput) Key() string { return "typed/output" }
func (o *countingOrderedOutput) Encode(context.Context, contracts.OutputRequest) (contracts.OutputResult, error) {
	o.calls.Add(1)
	return contracts.OutputResult{}, nil
}

type countingInputAdapter struct {
	contracts.InputAdapter
	calls atomic.Int32
}

func (a *countingInputAdapter) Parse(ctx context.Context, request contracts.ParseRequest) (*source.SourceDocument, error) {
	a.calls.Add(1)
	return a.InputAdapter.Parse(ctx, request)
}

type cancelingDebugRecorder struct {
	cancel context.CancelFunc
	onName string
	once   sync.Once
}

func (r *cancelingDebugRecorder) Enabled() bool { return true }

func (r *cancelingDebugRecorder) WriteJSON(name string, _ any) error {
	if name == r.onName {
		r.once.Do(r.cancel)
	}
	return nil
}

func (*cancelingDebugRecorder) WriteBytes(string, []byte) error { return nil }

type cancelingOutputEncoder struct {
	contracts.OutputEncoder
	cancel context.CancelFunc
	calls  atomic.Int32
}

func (e *cancelingOutputEncoder) Encode(context.Context, contracts.OutputRequest) (contracts.OutputResult, error) {
	e.calls.Add(1)
	e.cancel()
	return contracts.OutputResult{
		Files: []contracts.OutputFile{{Name: "result.txt", ContentType: "text/plain", Bytes: []byte("result")}},
	}, nil
}

func TestRunnerSkipsInputAdapterWhenContextIsAlreadyCanceled(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 1)
	adapter := &countingInputAdapter{InputAdapter: prepared.input}
	prepared.input = adapter
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input")})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Run() error = %v, want context.Canceled", err)
	}
	if got := adapter.calls.Load(); got != 0 {
		t.Fatalf("input Parse calls = %d, want none", got)
	}
}

func TestRunnerSkipsOutputEncodingAfterDebugCancellation(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 1)
	encoder := &countingOrderedOutput{}
	prepared.output = encoder
	ctx, cancel := context.WithCancel(context.Background())

	_, err := New().Run(ctx, RunInput{
		Prepared: prepared,
		RawInput: []byte("input"),
		Debug:    &cancelingDebugRecorder{cancel: cancel, onName: "output/input.json"},
	})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Run() error = %v, want context.Canceled", err)
	}
	if got := encoder.calls.Load(); got != 0 {
		t.Fatalf("output Encode calls = %d, want none", got)
	}
}

func TestRunnerDiscardsOutputReturnedAfterCancellation(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 1)
	ctx, cancel := context.WithCancel(context.Background())
	encoder := &cancelingOutputEncoder{OutputEncoder: prepared.output, cancel: cancel}
	prepared.output = encoder

	output, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input")})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Run() error = %v, want context.Canceled", err)
	}
	if got := encoder.calls.Load(); got != 1 {
		t.Fatalf("output Encode calls = %d, want one", got)
	}
	if len(output.OutputFiles) != 0 {
		t.Fatalf("output files = %#v, want none", output.OutputFiles)
	}
	for _, group := range output.Diagnostics.Groups {
		if group.ReasonCode == "returned-after-cancel" {
			t.Fatalf("output diagnostics include encoder diagnostic after cancellation")
		}
	}
}

func TestRunnerFailsGeneratedHandoffBeforeConsumerExtraction(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 1)
	input := &countingInputAdapter{InputAdapter: prepared.input}
	prepared.input = input
	preparedLane := &prepared.Steps[0].lanes[0]
	lane := &preparedLane.resolved
	binding := ReferenceBinding{
		Stage:    StageExtract,
		LaneID:   lane.ID,
		SlotName: "generated",
		Artifact: &ArtifactReference{Step: "producer", Lane: "source"},
	}
	lane.ExtractReferences.Bindings = []ReferenceBinding{binding}
	lane.ExtractReferences.ReferenceSet = contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"generated": {Slot: contracts.ReferenceSlot{Name: "generated", AcceptedArtifactKinds: []contracts.ArtifactKind{"test/notes"}}},
	}}
	prepared.resolved.Steps[0].ArtifactLanes[0].ExtractReferences = CloneReferenceTarget(lane.ExtractReferences)
	var extractCalls atomic.Int32
	preparedLane.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		extractCalls.Add(1)
		return erasedTypedResult{Value: codecNotes{}}, nil
	}

	_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")})
	if err == nil || !strings.Contains(err.Error(), "generated dependency") {
		t.Fatalf("Run() error = %v, want generated dependency failure", err)
	}
	if got := input.calls.Load(); got != 1 {
		t.Fatalf("input Parse calls = %d, want one before handoff", got)
	}
	if got := extractCalls.Load(); got != 0 {
		t.Fatalf("consumer extract calls = %d, want zero", got)
	}
}

func TestRunnerExecutesOrderedStepsWithHardBarriers(t *testing.T) {
	prepared := preparedOrderedPipeline(t, 1,
		orderedLaneSpec{id: "notes", profile: "notes"},
		orderedLaneSpec{id: "score", profile: "score"},
	)
	var mu sync.Mutex
	var events []string
	record := func(event string) {
		mu.Lock()
		events = append(events, event)
		mu.Unlock()
	}
	first := &prepared.Steps[0].lanes[0]
	second := &prepared.Steps[1].lanes[0]
	first.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return erasedTypedResult{Value: codecNotes{Items: []string{"first"}}}, nil
	}
	originalFirstNormalize := first.typed.normalize
	first.typed.normalize = func(ctx context.Context, implementation any, request contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
		result, err := originalFirstNormalize(ctx, implementation, request)
		if err == nil {
			record("first-normalize-done")
		}
		return result, err
	}
	second.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		record("second-extract-started")
		return erasedTypedResult{Value: codecScore{Value: 2}}, nil
	}

	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
	if err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	mu.Lock()
	gotEvents := append([]string(nil), events...)
	mu.Unlock()
	if want := []string{"first-normalize-done", "second-extract-started"}; !reflect.DeepEqual(gotEvents, want) {
		t.Fatalf("ordered events = %#v, want %#v", gotEvents, want)
	}
	if got := []string{output.NormalizeOutputs[0].LaneID, output.NormalizeOutputs[1].LaneID}; !reflect.DeepEqual(got, []string{"notes", "score"}) {
		t.Fatalf("normalized lane order = %#v, want notes then score", got)
	}
}

func TestRunnerStopsLaterOrderedStepsAfterFailure(t *testing.T) {
	prepared := preparedOrderedPipeline(t, 1,
		orderedLaneSpec{id: "first", profile: "notes"},
		orderedLaneSpec{id: "failure", profile: "score"},
		orderedLaneSpec{id: "later", profile: "notes"},
	)
	first := &prepared.Steps[0].lanes[0]
	second := &prepared.Steps[1].lanes[0]
	third := &prepared.Steps[2].lanes[0]
	first.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return erasedTypedResult{Value: codecNotes{Items: []string{"first"}}}, nil
	}
	second.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return erasedTypedResult{}, errors.New("ordered step extraction failed")
	}
	var thirdCalls atomic.Int32
	third.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		thirdCalls.Add(1)
		return erasedTypedResult{Value: codecNotes{Items: []string{"later"}}}, nil
	}
	encoder := &countingOrderedOutput{}
	prepared.output = encoder

	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
	if err == nil || !strings.Contains(err.Error(), `execute pipeline step "step-2"`) {
		t.Fatalf("Run() error = %v, want step-scoped failure", err)
	}
	if len(output.NormalizeOutputs) != 1 || output.NormalizeOutputs[0].LaneID != "first" {
		t.Fatalf("completed outputs = %#v, want only the first step output", output.NormalizeOutputs)
	}
	if got := thirdCalls.Load(); got != 0 {
		t.Fatalf("later step extract calls = %d, want zero", got)
	}
	if got := encoder.calls.Load(); got != 0 {
		t.Fatalf("output encoder calls = %d, want zero after failure", got)
	}
}

func installExtractOperation(prepared *PreparedPipeline, laneIndex int, operation func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error)) {
	prepared.Steps[0].lanes[laneIndex].typed.extract = func(ctx context.Context, _ any, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return operation(ctx, request)
	}
}

func typedValueForLane(laneIndex int, chunkIndex int) any {
	if laneIndex == 0 {
		return codecNotes{Items: []string{fmt.Sprintf("chunk-%d", chunkIndex)}}
	}
	return codecScore{Value: chunkIndex}
}

func TestRunnerBoundsExtractJobsAndStabilizesReverseCompletion(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 3)
	started := make(chan int, 6)
	release := make([]chan struct{}, 6)
	for i := range release {
		release[i] = make(chan struct{})
	}
	var active atomic.Int32
	var maximum atomic.Int32
	for laneIndex := range prepared.Steps[0].lanes {
		lane := laneIndex
		installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
			jobIndex := request.Chunk.Index*2 + lane
			current := active.Add(1)
			defer active.Add(-1)
			for {
				seen := maximum.Load()
				if current <= seen || maximum.CompareAndSwap(seen, current) {
					break
				}
			}
			started <- jobIndex
			select {
			case <-release[jobIndex]:
			case <-ctx.Done():
				return erasedTypedResult{}, ctx.Err()
			}
			scope := fmt.Sprintf("lane-%d/chunk-%d", lane, request.Chunk.Index)
			return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index), Diagnostics: []contracts.ProducerDiagnostic{{Disposition: contracts.DiagnosticDispositionObservation, Category: contracts.DiagnosticCategoryNormalization, ReasonCode: "observed", OccurrenceCount: 1, Samples: []contracts.DiagnosticSample{{Scope: scope, Message: "ordered"}}}}}, nil
		})
	}

	type outcome struct {
		output RunOutput
		err    error
	}
	done := make(chan outcome, 1)
	go func() {
		output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
		done <- outcome{output: output, err: err}
	}()
	for pair := 0; pair < 3; pair++ {
		first, second := <-started, <-started
		want := map[int]bool{pair * 2: true, pair*2 + 1: true}
		if !want[first] || !want[second] || first == second {
			t.Fatalf("started jobs = %d, %d; want round-robin pair %d", first, second, pair)
		}
		close(release[second])
		close(release[first])
	}
	result := <-done
	if result.err != nil {
		t.Fatalf("Run() error = %v", result.err)
	}
	if got := maximum.Load(); got != 2 {
		t.Fatalf("maximum concurrent extract jobs = %d, want 2", got)
	}
	if len(result.output.Diagnostics.Groups) != 2 {
		t.Fatalf("diagnostic groups = %#v, want one deterministic group per lane", result.output.Diagnostics.Groups)
	}
	for lane, group := range result.output.Diagnostics.Groups {
		if group.Origin.LaneID != prepared.Steps[0].lanes[lane].resolved.ID {
			t.Fatalf("group origin = %#v, want configured lane %q", group.Origin, prepared.Steps[0].lanes[lane].resolved.ID)
		}
		indexes := make([]int, len(group.Samples))
		for index, sample := range group.Samples {
			if sample.ChunkIndex == nil {
				t.Fatalf("sample = %#v, want chunk index", sample)
			}
			indexes[index] = *sample.ChunkIndex
		}
		if !reflect.DeepEqual(indexes, []int{0, 1, 2}) {
			t.Fatalf("group sample chunk order = %#v, want canonical chunk order", indexes)
		}
	}
}

func TestRunnerStartsLaneContinuationWhileOtherLaneExtractsRemain(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 2)
	releaseOtherLane := make(chan struct{})
	mergeStarted := make(chan struct{}, 1)
	installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil
	})
	installExtractOperation(prepared, 1, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		select {
		case <-releaseOtherLane:
			return erasedTypedResult{Value: typedValueForLane(1, request.Chunk.Index)}, nil
		case <-ctx.Done():
			return erasedTypedResult{}, ctx.Err()
		}
	})
	originalMerge := prepared.Steps[0].lanes[0].typed.merge
	prepared.Steps[0].lanes[0].typed.merge = func(ctx context.Context, implementation any, request contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
		select {
		case mergeStarted <- struct{}{}:
		default:
		}
		return originalMerge(ctx, implementation, request)
	}
	done := make(chan error, 1)
	go func() {
		_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
		done <- err
	}()
	select {
	case <-mergeStarted:
		close(releaseOtherLane)
	case <-time.After(2 * time.Second):
		t.Fatal("first lane merge did not start while the second lane remained in extract")
	}
	if err := <-done; err != nil {
		t.Fatalf("Run() error = %v", err)
	}
}

func TestRunnerBoundsConcurrentLaneContinuations(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 1)
	base := prepared.Steps[0].lanes[0]
	lanes := make([]preparedLaneExecutor, 8)
	resolvedLanes := make([]ResolvedArtifactLane, len(lanes))
	publicLanes := make([]PreparedArtifactLane, len(lanes))
	var active atomic.Int32
	var maximum atomic.Int32
	for i := range lanes {
		lane := base
		lane.resolved.ID = fmt.Sprintf("lane-%02d", i)
		lane.typed = &preparedTypedLane{
			extractor:  base.typed.extractor,
			merger:     base.typed.merger,
			normalizer: base.typed.normalizer,
			extract: func(_ context.Context, _ any, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
				return erasedTypedResult{Value: codecNotes{Items: []string{request.Chunk.ID}}}, nil
			},
			merge: func(_ context.Context, _ any, _ contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
				current := active.Add(1)
				defer active.Add(-1)
				for {
					seen := maximum.Load()
					if current <= seen || maximum.CompareAndSwap(seen, current) {
						break
					}
				}
				time.Sleep(5 * time.Millisecond)
				return erasedTypedResult{Value: codecNotes{}}, nil
			},
			normalize: base.typed.normalize,
			codec:     base.typed.codec,
		}
		lanes[i] = lane
		resolvedLanes[i] = lane.resolved
		publicLanes[i] = PreparedArtifactLane{Resolved: lane.resolved}
	}
	prepared.Steps[0].lanes = lanes
	prepared.resolved.Steps[0].ArtifactLanes = resolvedLanes
	prepared.Steps[0].ArtifactLanes = publicLanes

	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
	if err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	if len(output.NormalizeOutputs) != len(lanes) {
		t.Fatalf("normalize outputs = %d, want %d", len(output.NormalizeOutputs), len(lanes))
	}
	if got := maximum.Load(); got != 2 {
		t.Fatalf("maximum concurrent lane continuations = %d, want 2", got)
	}
}

func TestRunnerSelectsFrameworkErrorByStableLaneOrder(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 1)
	ready := make(chan struct{}, 2)
	release := make(chan struct{})
	for laneIndex := range prepared.Steps[0].lanes {
		lane := laneIndex
		installExtractOperation(prepared, lane, func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
			ready <- struct{}{}
			<-release
			return erasedTypedResult{}, fmt.Errorf("lane-%d failure", lane)
		})
	}
	done := make(chan error, 1)
	go func() {
		_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
		done <- err
	}()
	<-ready
	<-ready
	close(release)
	if err := <-done; err == nil || !strings.Contains(err.Error(), "lane-0 failure") {
		t.Fatalf("Run() error = %v, want first resolved lane failure", err)
	}
}

func TestRunnerSelectsMergeErrorBeforeEarlierLaneNormalizeError(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 1)
	for laneIndex := range prepared.Steps[0].lanes {
		lane := laneIndex
		installExtractOperation(prepared, lane, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
			return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
		})
	}
	ready := make(chan struct{}, 2)
	release := make(chan struct{})
	prepared.Steps[0].lanes[0].typed.normalize = func(context.Context, any, contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
		ready <- struct{}{}
		<-release
		return erasedTypedResult{}, errors.New("earlier lane normalize failure")
	}
	prepared.Steps[0].lanes[1].typed.merge = func(context.Context, any, contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
		ready <- struct{}{}
		<-release
		return erasedTypedResult{}, errors.New("later lane merge failure")
	}
	done := make(chan error, 1)
	go func() {
		_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
		done <- err
	}()
	<-ready
	<-ready
	close(release)
	if err := <-done; err == nil || !strings.Contains(err.Error(), "later lane merge failure") {
		t.Fatalf("Run() error = %v, want merge failure before normalize failure", err)
	}
}

func TestRunnerKeepsSuccessfulChunksWhenAnotherExtractIsRejected(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 2)
	prepared.Steps[0].lanes[0].resolved.ExtractValidationPolicy.SemanticRejection = SemanticRejectionRejectOutput
	for laneIndex := range prepared.Steps[0].lanes {
		lane := laneIndex
		installExtractOperation(prepared, lane, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
			return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
		})
	}
	validator := &prepared.Steps[0].lanes[0].extractValidators.validators[0]
	validator.typedValidate = func(_ context.Context, _ any, target typedValidationTarget) (contracts.ValidationResult, error) {
		if target.chunk != nil && target.chunk.Index == 0 {
			return contracts.ValidationResult{Approved: false, ReasonCode: "expected_rejection", Message: "rejected by test", CorrectionGuidance: "return an acceptable candidate"}, nil
		}
		return contracts.ValidationResult{Approved: true}, nil
	}
	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
	if err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	if len(output.Rejected) != 1 || output.Rejected[0].LaneID != prepared.Steps[0].lanes[0].resolved.ID || output.Rejected[0].ChunkIndex != 0 {
		t.Fatalf("rejections = %#v, want the first lane's first chunk", output.Rejected)
	}
	if len(output.NormalizeOutputs) != 2 {
		t.Fatalf("normalize outputs = %d, want both lanes to continue", len(output.NormalizeOutputs))
	}
	if output.Manifest.ValidationStatus != "rejected" {
		t.Fatalf("validation status = %q, want rejected", output.Manifest.ValidationStatus)
	}
}

func TestRunnerSeparatelyBoundsExtractJobsAndSharedProviderCalls(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 3)
	provider := &countingProvider{}
	scheduler, err := frameworkllm.NewScheduler(2)
	if err != nil {
		t.Fatalf("NewScheduler() error = %v", err)
	}
	client := frameworkllm.NewScheduledClient(provider, scheduler)
	var jobActive atomic.Int32
	var jobMaximum atomic.Int32
	var attempts sync.Map
	for laneIndex := range prepared.Steps[0].lanes {
		lane := laneIndex
		prepared.Steps[0].lanes[lane].resolved.Extract.Retries = 1
		installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
			current := jobActive.Add(1)
			defer jobActive.Add(-1)
			for {
				seen := jobMaximum.Load()
				if current <= seen || jobMaximum.CompareAndSwap(seen, current) {
					break
				}
			}
			if _, callErr := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "extract"}, nil); callErr != nil {
				return erasedTypedResult{}, callErr
			}
			key := fmt.Sprintf("%d/%d", lane, request.Chunk.Index)
			countValue, _ := attempts.LoadOrStore(key, &atomic.Int32{})
			count := countValue.(*atomic.Int32).Add(1)
			if lane == 0 && request.Chunk.Index == 0 && count == 1 {
				return erasedTypedResult{}, errors.New("retry extract")
			}
			return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
		})
		validator := &prepared.Steps[0].lanes[lane].extractValidators.validators[0]
		validator.typedValidate = func(ctx context.Context, _ any, _ typedValidationTarget) (contracts.ValidationResult, error) {
			if _, callErr := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "validate"}, nil); callErr != nil {
				return contracts.ValidationResult{}, callErr
			}
			return contracts.ValidationResult{Approved: true}, nil
		}
	}
	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 3})
	if err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	if len(output.NormalizeOutputs) != 2 {
		t.Fatalf("normalize outputs = %d, want 2", len(output.NormalizeOutputs))
	}
	if got := jobMaximum.Load(); got != 3 {
		t.Fatalf("maximum extract jobs = %d, want 3", got)
	}
	if got := provider.maximum.Load(); got != 2 {
		t.Fatalf("maximum provider calls = %d, want 2", got)
	}
	if got := provider.calls.Load(); got != 13 {
		t.Fatalf("provider calls = %d, want 13 including retry and validators", got)
	}
}

func TestRunnerReturnsParentCancellationAndStopsQueuedExtracts(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 4)
	started := make(chan struct{}, 8)
	for laneIndex := range prepared.Steps[0].lanes {
		lane := laneIndex
		installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
			started <- struct{}{}
			<-ctx.Done()
			return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, ctx.Err()
		})
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		_, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
		done <- err
	}()
	<-started
	<-started
	cancel()
	if err := <-done; !errors.Is(err, context.Canceled) {
		t.Fatalf("Run() error = %v, want context.Canceled", err)
	}
	if got := len(started); got != 0 {
		t.Fatalf("additional started jobs = %d, want none after cancellation", got)
	}
}