notarius/internal/framework/pipeline/handoff_test.go

package pipeline

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

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

func handoffFixture(t *testing.T, value codecNotes) (RunInput, PreparedPipelineStep, contracts.SerializedOutput) {
	t.Helper()
	prepared := preparedOrderedPipeline(t, 1,
		orderedLaneSpec{id: "notes", profile: "notes"},
		orderedLaneSpec{id: "score", profile: "score"},
	)
	consumer := &prepared.Steps[1].lanes[0]
	installGeneratedTarget := func(target *ResolvedReferenceTarget, stage ModuleStage, module string) {
		*target = ResolvedReferenceTarget{
			Stage:  stage,
			StepID: consumer.resolved.StepID,
			LaneID: consumer.resolved.ID,
			Module: module,
			Bindings: []ReferenceBinding{{
				Stage:    stage,
				LaneID:   consumer.resolved.ID,
				SlotName: "producer-output",
				Artifact: &ArtifactReference{Step: "step-1", Lane: "notes"},
			}},
			ReferenceSet: contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
				"producer-output": {Slot: contracts.ReferenceSlot{
					Name:                  "producer-output",
					AcceptedArtifactKinds: []contracts.ArtifactKind{"test/notes"},
					AcceptedMediaTypes:    []string{"application/json"},
				}},
			}},
		}
	}
	installGeneratedTarget(&consumer.resolved.ExtractReferences, StageExtract, consumer.resolved.Extract.Module)
	installGeneratedTarget(&consumer.resolved.MergeReferences, StageMerge, consumer.resolved.Merge.Module)
	installGeneratedTarget(&consumer.resolved.NormalizeReferences, StageNormalize, consumer.resolved.Normalize.Module)

	artifact, err := checkpointArtifact(prepared.Steps[0].lanes[0].typed.codec, "notes", "typed/normalize", "source", value)
	if err != nil {
		t.Fatalf("checkpointArtifact() error = %v", err)
	}
	input := RunInput{Prepared: prepared}
	input.pipeline = prepared.resolved
	return input, prepared.Steps[1], contracts.SerializedOutput{
		StepID:        "step-1",
		LaneID:        "notes",
		NormalizerKey: "typed/normalize",
		SourceID:      "source",
		Artifact:      contracts.CloneSerializedArtifact(artifact.Artifact),
	}
}

func TestBuildStepReferenceSetsCanonicalizesAndClonesFanout(t *testing.T) {
	input, step, producerOutput := handoffFixture(t, codecNotes{Items: []string{"first"}})
	sets, provenance, err := buildStepReferenceSets(input, step, []contracts.SerializedOutput{producerOutput})
	if err != nil {
		t.Fatalf("buildStepReferenceSets() error = %v", err)
	}
	if got, want := len(sets), 3; got != want {
		t.Fatalf("reference target set count = %d, want %d", got, want)
	}
	if got, want := len(provenance), 3; got != want {
		t.Fatalf("generated provenance count = %d, want %d", got, want)
	}

	var first []byte
	for key, set := range sets {
		item := set.Slots["producer-output"].Items[0]
		if item.Origin.Type != "generated" || item.Origin.URI != "" {
			t.Fatalf("target %v origin = %#v, want generated origin without URI", key, item.Origin)
		}
		if item.Producer.StepID != "step-1" || item.Producer.LaneID != "notes" || item.Producer.ModuleKey != "typed/normalize" {
			t.Fatalf("target %v producer = %#v, want producer provenance", key, item.Producer)
		}
		if first == nil {
			first = item.Content
			continue
		}
		if &first[0] == &item.Content[0] {
			t.Fatalf("target %v shares generated content backing storage", key)
		}
	}
	if got := len(step.lanes[0].resolved.ExtractReferences.ReferenceSet.Slots["producer-output"].Items); got != 0 {
		t.Fatalf("prepared extract reference items = %d, want zero", got)
	}

	encoded, err := json.Marshal(provenance)
	if err != nil {
		t.Fatalf("marshal generated provenance: %v", err)
	}
	text := string(encoded)
	for _, forbidden := range []string{"first", "file://", "content_base64"} {
		if strings.Contains(text, forbidden) {
			t.Fatalf("generated provenance contains forbidden %q: %s", forbidden, text)
		}
	}
	if !strings.Contains(text, "producer_pipeline_id") || !strings.Contains(text, "schema_digest") {
		t.Fatalf("generated provenance = %s, want bounded producer and schema identity", text)
	}
}

func TestBuildStepReferenceSetsIndexesManyOutputsForFanout(t *testing.T) {
	input, step, producerOutput := handoffFixture(t, codecNotes{Items: []string{"first"}})
	outputs := make([]contracts.SerializedOutput, 0, 65)
	for index := range 64 {
		outputs = append(outputs, contracts.SerializedOutput{StepID: fmt.Sprintf("unrelated-step-%d", index), LaneID: "unrelated-lane"})
	}
	outputs = append(outputs, producerOutput)

	sets, provenance, err := buildStepReferenceSets(input, step, outputs)
	if err != nil {
		t.Fatalf("buildStepReferenceSets() error = %v", err)
	}
	if len(sets) != 3 || len(provenance) != 3 {
		t.Fatalf("reference sets/provenance = %d/%d, want 3/3", len(sets), len(provenance))
	}
	for target, set := range sets {
		item := set.Slots["producer-output"].Items[0]
		if item.Producer.StepID != "step-1" || item.Producer.LaneID != "notes" || string(item.Content) == "" {
			t.Fatalf("target %v item = %#v, want canonical producer artifact", target, item)
		}
	}
}

func TestGeneratedReferenceAcceptsTypedEmptyCollection(t *testing.T) {
	input, step, producerOutput := handoffFixture(t, codecNotes{})
	sets, _, err := buildStepReferenceSets(input, step, []contracts.SerializedOutput{producerOutput})
	if err != nil {
		t.Fatalf("buildStepReferenceSets() error = %v, want nil for accepted empty collection", err)
	}
	if got := len(sets[keyForReferenceTarget(step.lanes[0].resolved.ExtractReferences)].Slots["producer-output"].Items); got != 1 {
		t.Fatalf("generated empty collection item count = %d, want one artifact", got)
	}
}

func TestGeneratedReferenceRejectsInvalidProducerOutputsDeterministically(t *testing.T) {
	tests := []struct {
		name    string
		outputs func(contracts.SerializedOutput) []contracts.SerializedOutput
		mutate  func(*ResolvedReferenceTarget, *contracts.SerializedOutput)
		want    string
	}{
		{name: "missing", outputs: func(contracts.SerializedOutput) []contracts.SerializedOutput { return nil }, want: "no accepted normalized output"},
		{name: "multiple", outputs: func(output contracts.SerializedOutput) []contracts.SerializedOutput {
			return []contracts.SerializedOutput{output, output}
		}, want: "exactly one is required"},
		{name: "multiple after identifier normalization", outputs: func(output contracts.SerializedOutput) []contracts.SerializedOutput {
			duplicate := contracts.CloneSerializedOutput(output)
			duplicate.StepID = " step-1 "
			duplicate.LaneID = " notes "
			return []contracts.SerializedOutput{output, duplicate}
		}, want: "producer has 2 accepted normalized outputs; exactly one is required"},
		{name: "kind mismatch", outputs: func(output contracts.SerializedOutput) []contracts.SerializedOutput {
			return []contracts.SerializedOutput{output}
		}, mutate: func(target *ResolvedReferenceTarget, _ *contracts.SerializedOutput) {
			target.ReferenceSet.Slots["producer-output"] = contracts.ResolvedReferenceSlot{Slot: contracts.ReferenceSlot{Name: "producer-output", AcceptedArtifactKinds: []contracts.ArtifactKind{"test/score"}}}
		}, want: "does not accept artifact kind"},
		{name: "schema mismatch", outputs: func(output contracts.SerializedOutput) []contracts.SerializedOutput {
			return []contracts.SerializedOutput{output}
		}, mutate: func(_ *ResolvedReferenceTarget, output *contracts.SerializedOutput) {
			output.Artifact.Schema.ID = "wrong.schema"
		}, want: "schema identity"},
		{name: "media mismatch", outputs: func(output contracts.SerializedOutput) []contracts.SerializedOutput {
			return []contracts.SerializedOutput{output}
		}, mutate: func(_ *ResolvedReferenceTarget, output *contracts.SerializedOutput) {
			output.Artifact.MediaType = "text/plain"
		}, want: "media type"},
		{name: "size mismatch", outputs: func(output contracts.SerializedOutput) []contracts.SerializedOutput {
			return []contracts.SerializedOutput{output}
		}, mutate: func(target *ResolvedReferenceTarget, _ *contracts.SerializedOutput) {
			slot := target.ReferenceSet.Slots["producer-output"]
			slot.Slot.MaxBytes = 1
			target.ReferenceSet.Slots["producer-output"] = slot
		}, want: "limit"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			input, step, output := handoffFixture(t, codecNotes{Items: []string{"first"}})
			output = contracts.CloneSerializedOutput(output)
			if test.mutate != nil {
				target := &step.lanes[0].resolved.ExtractReferences
				test.mutate(target, &output)
			}
			_, _, err := buildStepReferenceSets(input, step, test.outputs(output))
			if err == nil || !strings.Contains(err.Error(), test.want) {
				t.Fatalf("buildStepReferenceSets() error = %v, want %q", err, test.want)
			}
		})
	}
}

func BenchmarkIndexGeneratedOutputs(b *testing.B) {
	outputs := make([]contracts.SerializedOutput, 0, 1024)
	for index := range 1024 {
		outputs = append(outputs, contracts.SerializedOutput{StepID: fmt.Sprintf("step-%d", index), LaneID: fmt.Sprintf("lane-%d", index)})
	}
	b.ResetTimer()
	for range b.N {
		_ = indexGeneratedOutputs(outputs)
	}
}

func TestGeneratedReferenceFingerprintChangesWithCanonicalContent(t *testing.T) {
	input, step, first := handoffFixture(t, codecNotes{Items: []string{"first"}})
	firstSets, _, err := buildStepReferenceSets(input, step, []contracts.SerializedOutput{first})
	if err != nil {
		t.Fatalf("build first reference set: %v", err)
	}
	_, _, second := handoffFixture(t, codecNotes{Items: []string{"second"}})
	secondSets, _, err := buildStepReferenceSets(input, step, []contracts.SerializedOutput{second})
	if err != nil {
		t.Fatalf("build second reference set: %v", err)
	}
	firstDeps := generatedReferenceDependencies(firstSets[keyForReferenceTarget(step.lanes[0].resolved.ExtractReferences)])
	secondDeps := generatedReferenceDependencies(secondSets[keyForReferenceTarget(step.lanes[0].resolved.ExtractReferences)])
	if reflect.DeepEqual(firstDeps, secondDeps) {
		t.Fatalf("generated dependencies = %#v, want content-sensitive fingerprint", firstDeps)
	}
}

func TestGeneratedReferenceFingerprintChangesWithCanonicalSize(t *testing.T) {
	input, step, producerOutput := handoffFixture(t, codecNotes{Items: []string{"first"}})
	sets, _, err := buildStepReferenceSets(input, step, []contracts.SerializedOutput{producerOutput})
	if err != nil {
		t.Fatalf("build reference set: %v", err)
	}
	key := keyForReferenceTarget(step.lanes[0].resolved.ExtractReferences)
	firstDeps := generatedReferenceDependencies(sets[key])
	changed := CloneReferenceSet(sets[key])
	item := changed.Slots["producer-output"].Items[0]
	item.SizeBytes++
	changed.Slots["producer-output"] = contracts.ResolvedReferenceSlot{
		Slot:  changed.Slots["producer-output"].Slot,
		Items: []contracts.ReferenceItem{item},
	}
	secondDeps := generatedReferenceDependencies(changed)
	if reflect.DeepEqual(firstDeps, secondDeps) {
		t.Fatalf("generated dependencies = %#v, want size-sensitive fingerprint", firstDeps)
	}
}

func TestRunnerHandsOffAcceptedNormalizedOutputBeforeConsumerLanes(t *testing.T) {
	input, _, _ := handoffFixture(t, codecNotes{Items: []string{"first"}})
	prepared := input.Prepared
	var received contracts.ReferenceSet
	prepared.Steps[0].lanes[0].typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return erasedTypedResult{Value: codecNotes{Items: []string{"first"}}}, nil
	}
	prepared.Steps[1].lanes[0].typed.extract = func(_ context.Context, _ any, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		received = CloneReferenceSet(request.References)
		return erasedTypedResult{Value: codecScore{Value: 3}}, nil
	}

	output, err := New().Run(context.Background(), input)
	if err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	if got, want := len(output.NormalizeOutputs), 2; got != want {
		t.Fatalf("normalized output count = %d, want %d", got, want)
	}
	if output.NormalizeOutputs[0].StepID != "step-1" || output.NormalizeOutputs[1].StepID != "step-2" {
		t.Fatalf("normalized output step IDs = %#v, want step-1 and step-2", output.NormalizeOutputs)
	}
	item := received.Slots["producer-output"].Items[0]
	if item.Origin.Type != "generated" || item.Producer.StepID != "step-1" || item.Producer.LaneID != "notes" {
		t.Fatalf("consumer reference = %#v, want generated producer reference", item)
	}
	if len(output.Manifest.References) == 0 || output.Manifest.References[len(output.Manifest.References)-1].OriginType != "generated" {
		t.Fatalf("manifest references = %#v, want generated provenance", output.Manifest.References)
	}
}

type handoffDependencyLoader struct {
	CheckpointLoader
	extract   [][]CheckpointFingerprint
	merge     [][]CheckpointFingerprint
	normalize [][]CheckpointFingerprint
}

func (l *handoffDependencyLoader) Enabled() bool { return true }

func (l *handoffDependencyLoader) Extract(_ string, _ string, dependencies []CheckpointFingerprint) (ExtractCheckpoint, CheckpointDecision) {
	l.extract = append(l.extract, append([]CheckpointFingerprint(nil), dependencies...))
	return ExtractCheckpoint{}, CheckpointDecision{Reason: "not found"}
}

func (l *handoffDependencyLoader) Merge(_ string, _ string, dependencies []CheckpointFingerprint) (MergeCheckpoint, CheckpointDecision) {
	l.merge = append(l.merge, append([]CheckpointFingerprint(nil), dependencies...))
	return MergeCheckpoint{}, CheckpointDecision{Reason: "not found"}
}

func (l *handoffDependencyLoader) Normalize(_ string, _ string, dependencies []CheckpointFingerprint) (NormalizeCheckpoint, CheckpointDecision) {
	l.normalize = append(l.normalize, append([]CheckpointFingerprint(nil), dependencies...))
	return NormalizeCheckpoint{}, CheckpointDecision{Reason: "not found"}
}

func runHandoffWithProducerValue(t *testing.T, value codecNotes) *handoffDependencyLoader {
	t.Helper()
	input, _, _ := handoffFixture(t, value)
	prepared := input.Prepared
	prepared.Steps[0].lanes[0].typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return erasedTypedResult{Value: value}, nil
	}
	prepared.Steps[0].lanes[0].typed.normalize = func(context.Context, any, contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
		return erasedTypedResult{Value: value}, nil
	}
	prepared.Steps[1].lanes[0].typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return erasedTypedResult{Value: codecScore{Value: 3}}, nil
	}
	loader := &handoffDependencyLoader{CheckpointLoader: NoopCheckpointLoader()}
	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Checkpoint: loader}); err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	return loader
}

func TestRunnerAddsGeneratedDependencyToEveryReceivingCheckpoint(t *testing.T) {
	first := runHandoffWithProducerValue(t, codecNotes{Items: []string{"first"}})
	second := runHandoffWithProducerValue(t, codecNotes{Items: []string{"second"}})
	if len(first.extract) != 2 || len(first.merge) != 2 || len(first.normalize) != 2 {
		t.Fatalf("checkpoint load calls = extract %d merge %d normalize %d, want two each", len(first.extract), len(first.merge), len(first.normalize))
	}
	for name, calls := range map[string][][]CheckpointFingerprint{"extract": first.extract, "merge": first.merge, "normalize": first.normalize} {
		if got := generatedFingerprintCount(calls[1]); got != 1 {
			t.Fatalf("%s consumer dependency count = %d, want one", name, got)
		}
	}
	if reflect.DeepEqual(first.extract[1], second.extract[1]) || reflect.DeepEqual(first.merge[1], second.merge[1]) || reflect.DeepEqual(first.normalize[1], second.normalize[1]) {
		t.Fatalf("consumer checkpoint dependencies did not change with producer content")
	}
}

func generatedFingerprintCount(dependencies []CheckpointFingerprint) int {
	count := 0
	for _, dependency := range dependencies {
		if strings.HasPrefix(dependency.Name, "generated-reference:") {
			count++
		}
	}
	return count
}