notarius/internal/cli/assembled_spell_pipeline_contract_test.go

package cli

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"sync"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	spellnormalize "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/normalize/spells"
	spellcatalog "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/spells/catalog"
)

const assembledSpellExtractorKey = "test/dnd/spell-casts"

const assembledCorrectingSpellExtractorKey = "test/dnd/correcting-spell-casts"

const assembledDirectSpellValidatorKey = "test/dnd/direct-spell-correction"

func TestAssembledSpellPipelineNormalizesMergedCasts(t *testing.T) {
	registries, resolved, extractor := assembledSpellPipeline(t, assembledSpellPipelineOptions{})
	prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{})
	if err != nil {
		t.Fatalf("Prepare() error = %v, want nil", err)
	}

	output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{
		Prepared:       prepared,
		RawInput:       readRepositoryFile(t, "examples", "seriatim-minimal-transcript.json"),
		ChunkCacheMode: pipeline.ChunkCacheBypass,
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}

	chunkIndexes := extractor.chunkIndexesSnapshot()
	sort.Ints(chunkIndexes)
	if !reflect.DeepEqual(chunkIndexes, []int{0, 1}) {
		t.Fatalf("extractor chunk indexes = %#v, want two chunk-boundary calls", chunkIndexes)
	}
	if output.Manifest.ValidationStatus != "approved" || len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 1 {
		t.Fatalf("run output = %#v, want approved normalized output without rejections", output)
	}
	if output.NormalizeOutputs[0].NormalizerKey != spellnormalize.Key {
		t.Fatalf("normalized output module = %q, want %q", output.NormalizeOutputs[0].NormalizerKey, spellnormalize.Key)
	}

	var normalized dnd.SpellList
	if err := json.Unmarshal(output.NormalizeOutputs[0].Artifact.Content, &normalized); err != nil {
		t.Fatalf("decode normalized output: %v", err)
	}
	if len(normalized.SpellCasts) != 2 {
		t.Fatalf("normalized casts = %#v, want collapsed duplicate plus distinct evidence", normalized.SpellCasts)
	}
	first, distinct := normalized.SpellCasts[0], normalized.SpellCasts[1]
	if first.Spell != "Cure Wounds" || first.Caster != "  Aria \t" {
		t.Fatalf("retained cast = %#v, want canonical spell with first occurrence caster", first)
	}
	if !reflect.DeepEqual(first.SourceRefs, []source.SourceRef{{SourceID: "session-alpha", StartUnitID: 1, EndUnitID: 1}, {SourceID: "session-alpha", StartUnitID: 2, EndUnitID: 2}}) {
		t.Fatalf("retained refs = %#v, want sorted complete evidence", first.SourceRefs)
	}
	if distinct.Spell != "Cure Wounds" || distinct.Caster != "aria" || !reflect.DeepEqual(distinct.SourceRefs, []source.SourceRef{{SourceID: "session-alpha", StartUnitID: 2, EndUnitID: 2}}) {
		t.Fatalf("distinct cast = %#v, want separate evidence event", distinct)
	}

	wantWarningReasons := []string{
		spellnormalize.ReasonCodeSpellNameCanonicalized,
		spellnormalize.ReasonCodeSourceReferencesNormalized,
		spellnormalize.ReasonCodeDuplicateSpellCastCollapsed,
		"spell_not_near_source",
	}
	gotWarningReasons := make([]string, len(output.Diagnostics.Groups))
	for index, group := range output.Diagnostics.Groups {
		gotWarningReasons[index] = group.ReasonCode
	}
	if !reflect.DeepEqual(gotWarningReasons, wantWarningReasons) {
		t.Fatalf("diagnostics = %#v, want deterministic normalize and validation diagnostics", output.Diagnostics)
	}
	if output.Diagnostics.Groups[2].Samples[0].Scope != "spell_casts[0]" || !strings.Contains(output.Diagnostics.Groups[2].Samples[0].Message, "retained input index 0") || !strings.Contains(output.Diagnostics.Groups[2].Samples[0].Message, "removed input indices [1]") {
		t.Fatalf("duplicate diagnostic = %#v, want retained and removed merged indices", output.Diagnostics.Groups[2])
	}

	warningsFile := decodeAssembledOutput[struct {
		Groups []contracts.DiagnosticGroup `json:"groups"`
	}](t, output.OutputFiles, "warnings.json")
	if len(warningsFile.Groups) != 0 {
		t.Fatalf("warnings file = %#v, want no process warnings for advisory-only diagnostics", warningsFile.Groups)
	}
	diagnosticsFile := decodeAssembledOutput[struct {
		SchemaVersion                string                      `json:"schema_version"`
		GroupCount                   int                         `json:"group_count"`
		OccurrenceCount              int                         `json:"occurrence_count"`
		Truncated                    bool                        `json:"truncated"`
		UnrepresentedOccurrenceCount int                         `json:"unrepresented_occurrence_count"`
		Groups                       []contracts.DiagnosticGroup `json:"groups"`
	}](t, output.OutputFiles, "diagnostics.json")
	if diagnosticsFile.SchemaVersion != "notarius.diagnostics.v1" || diagnosticsFile.GroupCount != len(output.Diagnostics.Groups) || !reflect.DeepEqual(diagnosticsFile.Groups, output.Diagnostics.Groups) || diagnosticsFile.OccurrenceCount != diagnosticOccurrenceCount(output.Diagnostics.Groups)+output.Diagnostics.UnrepresentedOccurrenceCount || diagnosticsFile.Truncated != output.Diagnostics.Truncated || diagnosticsFile.UnrepresentedOccurrenceCount != output.Diagnostics.UnrepresentedOccurrenceCount {
		t.Fatalf("diagnostics file = %#v, run diagnostics = %#v", diagnosticsFile, output.Diagnostics)
	}
	manifest := decodeAssembledOutput[artifacts.RunManifest](t, output.OutputFiles, "manifest.json")
	if len(manifest.ArtifactLanes) != 1 || manifest.ArtifactLanes[0].Normalizer != spellnormalize.Key {
		t.Fatalf("manifest lanes = %#v, want assembled spell normalizer", manifest.ArtifactLanes)
	}
	normalizerMetadata, ok := manifest.ArtifactLanes[0].Metadata["normalizer"].(map[string]any)
	_, hasOverlayIDs := normalizerMetadata["catalog_overlay_ids"]
	if !ok || normalizerMetadata["catalog_base_id"] != spellcatalog.SRD5E2014ID || !strings.HasPrefix(stringValue(normalizerMetadata["catalog_digest"]), "sha256:") || !hasOverlayIDs {
		t.Fatalf("normalizer manifest metadata = %#v, want base ID, digest, and overlay IDs", manifest.ArtifactLanes[0].Metadata)
	}
}

func TestAssembledSpellPipelineCorrectsRejectedDirectExtraction(t *testing.T) {
	registries, resolved, extractor := assembledCorrectingSpellPipeline(t)
	prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{})
	if err != nil {
		t.Fatalf("Prepare() error = %v, want nil", err)
	}

	output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{
		Prepared:       prepared,
		RawInput:       readRepositoryFile(t, "examples", "seriatim-minimal-transcript.json"),
		ChunkCacheMode: pipeline.ChunkCacheBypass,
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if len(output.Rejected) != 0 || output.Manifest.ValidationStatus != "approved" || len(output.NormalizeOutputs) != 1 {
		t.Fatalf("run output = %#v, want corrected accepted spell output", output)
	}
	correction := extractor.correctionSnapshot()
	if correction == nil || string(correction.AssistantResponse) != `{"spell":"Mysterious Burst"}` || !strings.Contains(correction.UserGuidance, "use a known spell name") || !strings.Contains(correction.UserGuidance, "complete corrected replacement") || strings.Contains(correction.UserGuidance, "unknown_spell") || strings.Contains(correction.UserGuidance, "spell is not in the catalog") {
		t.Fatalf("extract correction = %#v, want exact rejected model response and semantic replacement guidance only", correction)
	}
}

func TestAssembledSpellPipelineHonorsNormalizeValidatorOverride(t *testing.T) {
	registries, resolved, _ := assembledSpellPipeline(t, assembledSpellPipelineOptions{normalizeValidatorOverride: true})
	var normalizeChain *pipeline.ResolvedValidatorChain
	for index := range resolved.ValidatorChains {
		chain := &resolved.ValidatorChains[index]
		if chain.Stage == pipeline.StageNormalize && chain.ModuleKey == spellnormalize.Key && chain.LaneID == "spells" {
			normalizeChain = chain
			break
		}
	}
	if normalizeChain == nil || len(normalizeChain.Validators) != 1 || normalizeChain.Validators[0].Binding.Module != "generic/always_accept" {
		t.Fatalf("normalize validator chain = %#v, want explicit always-accept override", normalizeChain)
	}

	prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{})
	if err != nil {
		t.Fatalf("Prepare() error = %v, want nil", err)
	}
	output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{
		Prepared:       prepared,
		RawInput:       readRepositoryFile(t, "examples", "seriatim-minimal-transcript.json"),
		ChunkCacheMode: pipeline.ChunkCacheBypass,
	})
	if err != nil || output.Manifest.ValidationStatus != "approved" || len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 1 {
		t.Fatalf("Run() error = %v output = %#v, want approved override run", err, output)
	}
	for _, group := range output.Diagnostics.Groups {
		if group.ReasonCode == "spell_not_near_source" {
			t.Fatalf("diagnostics = %#v, want explicit validator override to replace default relatedness chain", output.Diagnostics)
		}
	}
}

func TestAssembledSpellPipelinePromotesTerminalUnknownSpellWarning(t *testing.T) {
	registries, resolved, _ := assembledSpellPipeline(t, assembledSpellPipelineOptions{unknownSpell: true})
	resolved.Steps[0].ArtifactLanes[0].NormalizeValidationPolicy.SemanticRejection = pipeline.SemanticRejectionRejectOutput
	prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{})
	if err != nil {
		t.Fatalf("Prepare() error = %v, want nil", err)
	}
	output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{
		Prepared:       prepared,
		RawInput:       readRepositoryFile(t, "examples", "seriatim-minimal-transcript.json"),
		ChunkCacheMode: pipeline.ChunkCacheBypass,
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if output.Manifest.ValidationStatus != "rejected" || len(output.NormalizeOutputs) != 0 || len(output.Rejected) != 1 {
		t.Fatalf("run output = %#v, want one rejected normalize candidate and no normalized output", output)
	}
	rejection := output.Rejected[0]
	if rejection.Stage != string(pipeline.StageNormalize) || rejection.LaneID != "spells" || rejection.ModuleKey != spellnormalize.Key || rejection.ValidatorName != "extract/dnd/spells/catalog" || rejection.ReasonCode != "unknown_spell" {
		t.Fatalf("rejection = %#v, want durable normalize catalog rejection", rejection)
	}
	rejectedFile := decodeAssembledOutput[struct {
		Rejected []contracts.RejectedOutput `json:"rejected"`
	}](t, output.OutputFiles, "rejected.json")
	if !reflect.DeepEqual(rejectedFile.Rejected, output.Rejected) {
		t.Fatalf("rejected file = %#v, run rejections = %#v, want durable rejection diagnostic", rejectedFile.Rejected, output.Rejected)
	}
	if len(output.Diagnostics.Groups) != 2 || output.Diagnostics.Groups[0].ReasonCode != spellnormalize.ReasonCodeSpellNameUnresolved || output.Diagnostics.Groups[0].Samples[0].Scope != "spell_casts[0]" || output.Diagnostics.Groups[1].ReasonCode != "spell_not_near_source" {
		t.Fatalf("diagnostics = %#v, want complete terminal normalize validation diagnostics", output.Diagnostics)
	}
}

func TestAssembledSpellPipelinePromotesUnknownSpellWarningWhenOverrideAccepts(t *testing.T) {
	registries, resolved, _ := assembledSpellPipeline(t, assembledSpellPipelineOptions{normalizeValidatorOverride: true, unknownSpell: true})
	prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{})
	if err != nil {
		t.Fatalf("Prepare() error = %v, want nil", err)
	}
	output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{
		Prepared:       prepared,
		RawInput:       readRepositoryFile(t, "examples", "seriatim-minimal-transcript.json"),
		ChunkCacheMode: pipeline.ChunkCacheBypass,
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if output.Manifest.ValidationStatus != "approved" || len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 1 {
		t.Fatalf("run output = %#v, want accepted unknown spell with explicit validator override", output)
	}
	var normalized dnd.SpellList
	if err := json.Unmarshal(output.NormalizeOutputs[0].Artifact.Content, &normalized); err != nil {
		t.Fatalf("decode normalized output: %v", err)
	}
	if len(normalized.SpellCasts) != 1 || normalized.SpellCasts[0].Spell != "Mysterious Burst" {
		t.Fatalf("normalized casts = %#v, want unresolved name preserved", normalized.SpellCasts)
	}
	if len(output.Diagnostics.Groups) != 1 || output.Diagnostics.Groups[0].ReasonCode != spellnormalize.ReasonCodeSpellNameUnresolved || output.Diagnostics.Groups[0].Samples[0].Scope != "spell_casts[0]" {
		t.Fatalf("diagnostics = %#v, want promoted scoped unresolved-name diagnostic", output.Diagnostics)
	}
	warningsFile := decodeAssembledOutput[struct {
		Groups []contracts.DiagnosticGroup `json:"groups"`
	}](t, output.OutputFiles, "warnings.json")
	if len(warningsFile.Groups) != 0 {
		t.Fatalf("warnings file = %#v, want no process warnings for an advisory diagnostic", warningsFile.Groups)
	}
	diagnosticsFile := decodeAssembledOutput[struct {
		Groups []contracts.DiagnosticGroup `json:"groups"`
	}](t, output.OutputFiles, "diagnostics.json")
	if !reflect.DeepEqual(diagnosticsFile.Groups, output.Diagnostics.Groups) {
		t.Fatalf("diagnostics file = %#v, run diagnostics = %#v", diagnosticsFile.Groups, output.Diagnostics)
	}
}

type assembledSpellPipelineOptions struct {
	normalizeValidatorOverride bool
	unknownSpell               bool
}

func assembledCorrectingSpellPipeline(t *testing.T) (pipeline.Registries, pipeline.ResolvedPipeline, *assembledCorrectingSpellExtractor) {
	t.Helper()
	components := productionTestComponents(t)
	extractor := &assembledCorrectingSpellExtractor{}
	if err := pipeline.RegisterExtractor[dnd.SpellList](components.registries.Extractors, pipeline.ModuleSpec{
		Key:                assembledCorrectingSpellExtractorKey,
		Stage:              pipeline.StageExtract,
		ExecutionClass:     contracts.ExecutionClassLLMBacked,
		CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
		Requires:           []string{"chunks", "source.transcript"},
		Provides:           []string{"dnd.spell_casts"},
		ArtifactKind:       dnd.SpellListKind,
	}, func() (contracts.Extractor[dnd.SpellList], error) {
		return extractor, nil
	}); err != nil {
		t.Fatalf("register correcting extractor: %v", err)
	}
	if err := pipeline.RegisterTypedValidator[dnd.SpellList](components.registries.Validators, dnd.SpellListKind, pipeline.ValidatorSpec{Key: assembledDirectSpellValidatorKey, ExecutionClass: contracts.ExecutionClassDeterministic}, func() (contracts.TypedValidator[dnd.SpellList], error) {
		return assembledDirectSpellValidator{}, nil
	}); err != nil {
		t.Fatalf("register direct spell validator: %v", err)
	}

	extract := pipeline.Binding(assembledCorrectingSpellExtractorKey)
	extract.Retries = 1
	extract.Validators = pipeline.ValidatorOverride{Set: true, Validators: []pipeline.ModuleBinding{{Module: assembledDirectSpellValidatorKey}}}
	resolved, err := pipeline.ResolvePipeline(pipeline.PipelineProfile{
		ID:    "assembled-dnd-correcting-spells",
		Input: pipeline.Binding("seriatim"),
		Chunk: pipeline.ModuleBinding{Module: "generic", Options: map[string]any{"max_units": 1}},
		Artifacts: map[string]pipeline.ArtifactLaneProfile{
			"spells": {Extract: extract, Normalize: pipeline.Binding(spellnormalize.Key)},
		},
		Output: pipeline.Binding("json"),
	}, pipeline.ResolveOptions{}, catalogFromRegistries(components.registries))
	if err != nil {
		t.Fatalf("ResolvePipeline() error = %v, want nil", err)
	}
	return components.registries, resolved, extractor
}

func assembledSpellPipeline(t *testing.T, options assembledSpellPipelineOptions) (pipeline.Registries, pipeline.ResolvedPipeline, *assembledSpellExtractor) {
	t.Helper()
	components := productionTestComponents(t)
	extractor := &assembledSpellExtractor{unknownSpell: options.unknownSpell}
	if err := pipeline.RegisterExtractor[dnd.SpellList](components.registries.Extractors, pipeline.ModuleSpec{
		Key:            assembledSpellExtractorKey,
		Stage:          pipeline.StageExtract,
		ExecutionClass: contracts.ExecutionClassDeterministic,
		Requires:       []string{"chunks", "source.transcript"},
		Provides:       []string{"dnd.spell_casts"},
		ArtifactKind:   dnd.SpellListKind,
	}, func() (contracts.Extractor[dnd.SpellList], error) {
		return extractor, nil
	}); err != nil {
		t.Fatalf("register assembled extractor: %v", err)
	}

	normalize := pipeline.Binding(spellnormalize.Key)
	if options.normalizeValidatorOverride {
		normalize.Validators = pipeline.ValidatorOverride{
			Set:        true,
			Validators: []pipeline.ModuleBinding{pipeline.Binding("generic/always_accept")},
		}
	}
	resolved, err := pipeline.ResolvePipeline(pipeline.PipelineProfile{
		ID:    "assembled-dnd-spells",
		Input: pipeline.Binding("seriatim"),
		Chunk: pipeline.ModuleBinding{Module: "generic", Options: map[string]any{"max_units": 1}},
		Artifacts: map[string]pipeline.ArtifactLaneProfile{
			"spells": {Extract: pipeline.Binding(assembledSpellExtractorKey), Normalize: normalize},
		},
		Output: pipeline.Binding("json"),
	}, pipeline.ResolveOptions{}, catalogFromRegistries(components.registries))
	if err != nil {
		t.Fatalf("ResolvePipeline() error = %v, want nil", err)
	}
	return components.registries, resolved, extractor
}

type assembledSpellExtractor struct {
	mu           sync.Mutex
	chunkIndexes []int
	unknownSpell bool
}

type assembledCorrectingSpellExtractor struct {
	mu         sync.Mutex
	correction *contracts.SemanticCorrection
}

func (*assembledCorrectingSpellExtractor) Key() string { return assembledCorrectingSpellExtractorKey }

func (*assembledCorrectingSpellExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil }

func (e *assembledCorrectingSpellExtractor) Extract(_ context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SpellList], error) {
	if req.Source == nil || req.Chunk == nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, fmt.Errorf("correcting assembled extractor requires source and chunk")
	}
	response := `{"spell":"accepted"}`
	value := dnd.SpellList{SpellCasts: []dnd.SpellCast{}}
	if req.Chunk.Index == 0 && req.Correction == nil {
		response = `{"spell":"Mysterious Burst"}`
		value.SpellCasts = []dnd.SpellCast{{Caster: "Aria", Spell: "Mysterious Burst", SourceRefs: []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 1, EndUnitID: 1}}}}
	}
	if req.Chunk.Index == 0 && req.Correction != nil {
		correction, err := contracts.CloneSemanticCorrection(req.Correction)
		if err != nil {
			return contracts.TypedExtractionResult[dnd.SpellList]{}, err
		}
		e.mu.Lock()
		e.correction = correction
		e.mu.Unlock()
		value.SpellCasts = []dnd.SpellCast{{Caster: "Aria", Spell: "Cure Wounds", SourceRefs: []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 1, EndUnitID: 1}}}}
	}
	candidate, err := contracts.NewModelCandidate([]byte(response), contracts.CorrectionProtocolSingleResponseV1)
	if err != nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, err
	}
	return contracts.TypedExtractionResult[dnd.SpellList]{Value: value, ModelCandidate: candidate}, nil
}

func (e *assembledCorrectingSpellExtractor) correctionSnapshot() *contracts.SemanticCorrection {
	e.mu.Lock()
	defer e.mu.Unlock()
	correction, err := contracts.CloneSemanticCorrection(e.correction)
	if err != nil {
		return nil
	}
	return correction
}

type assembledDirectSpellValidator struct{}

func (assembledDirectSpellValidator) Name() string { return assembledDirectSpellValidatorKey }

func (assembledDirectSpellValidator) ExecutionClass() contracts.ExecutionClass {
	return contracts.ExecutionClassDeterministic
}

func (assembledDirectSpellValidator) Validate(_ context.Context, req contracts.TypedValidationRequest[dnd.SpellList]) (contracts.ValidationResult, error) {
	for _, cast := range req.Value.SpellCasts {
		if cast.Spell == "Mysterious Burst" {
			return contracts.ValidationResult{Approved: false, ReasonCode: "unknown_spell", Message: "spell is not in the catalog", CorrectionGuidance: "use a known spell name"}, nil
		}
	}
	return contracts.ValidationResult{Approved: true}, nil
}

func (e *assembledSpellExtractor) Key() string { return assembledSpellExtractorKey }

func (*assembledSpellExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil }

func (e *assembledSpellExtractor) Extract(ctx context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SpellList], error) {
	if err := ctx.Err(); err != nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, err
	}
	if req.Chunk == nil || req.Source == nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, fmt.Errorf("assembled extractor requires source and chunk")
	}
	e.mu.Lock()
	e.chunkIndexes = append(e.chunkIndexes, req.Chunk.Index)
	e.mu.Unlock()
	refOne := source.SourceRef{SourceID: req.Source.ID, StartUnitID: 1, EndUnitID: 1}
	refTwo := source.SourceRef{SourceID: req.Source.ID, StartUnitID: 2, EndUnitID: 2}
	if e.unknownSpell {
		if req.Chunk.Index == 0 {
			return contracts.TypedExtractionResult[dnd.SpellList]{Value: dnd.SpellList{SpellCasts: []dnd.SpellCast{{
				Caster: "Aria", Spell: "Mysterious Burst", SourceRefs: []source.SourceRef{refOne},
			}}}}, nil
		}
		return contracts.TypedExtractionResult[dnd.SpellList]{Value: dnd.SpellList{SpellCasts: []dnd.SpellCast{}}}, nil
	}
	switch req.Chunk.Index {
	case 0:
		return contracts.TypedExtractionResult[dnd.SpellList]{Value: dnd.SpellList{SpellCasts: []dnd.SpellCast{{
			Caster: "  Aria \t", Spell: " cure   wounds ", SourceRefs: []source.SourceRef{refTwo, refOne},
		}}}}, nil
	case 1:
		return contracts.TypedExtractionResult[dnd.SpellList]{Value: dnd.SpellList{SpellCasts: []dnd.SpellCast{
			{Caster: "aria", Spell: "Cure Wounds", SourceRefs: []source.SourceRef{refOne, refTwo}},
			{Caster: "aria", Spell: "Cure Wounds", SourceRefs: []source.SourceRef{refTwo}},
		}}}, nil
	default:
		return contracts.TypedExtractionResult[dnd.SpellList]{}, fmt.Errorf("unexpected assembled chunk index %d", req.Chunk.Index)
	}
}

func (e *assembledSpellExtractor) chunkIndexesSnapshot() []int {
	e.mu.Lock()
	defer e.mu.Unlock()
	return append([]int(nil), e.chunkIndexes...)
}

func diagnosticOccurrenceCount(groups []contracts.DiagnosticGroup) int {
	count := 0
	for _, group := range groups {
		count += group.OccurrenceCount
	}
	return count
}

func decodeAssembledOutput[T any](t *testing.T, files []contracts.OutputFile, name string) T {
	t.Helper()
	for _, file := range files {
		if file.Name != name {
			continue
		}
		var value T
		if err := json.Unmarshal(file.Bytes, &value); err != nil {
			t.Fatalf("decode %s: %v", name, err)
		}
		return value
	}
	t.Fatalf("output files = %#v, want %q", files, name)
	return *new(T)
}