notarius/internal/framework/pipeline/runner_extract_handoff_test.go

package pipeline

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

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

type extractCaptureRecorder struct {
	CheckpointRecorder
	checkpoint ExtractCheckpoint
}

type pipelineTypedMetadata map[string]any

type mutatingMetadataValidator struct {
	seen string
}

func (*mutatingMetadataValidator) Name() string { return "test/mutating-metadata" }
func (*mutatingMetadataValidator) ExecutionClass() contracts.ExecutionClass {
	return contracts.ExecutionClassDeterministic
}
func (v *mutatingMetadataValidator) Validate(_ context.Context, request contracts.ChunkValidationRequest) (contracts.ValidationResult, error) {
	nested := request.Chunks[0].Units[0].Metadata["typed"].(pipelineTypedMetadata)
	v.seen = string(nested["raw"].(json.RawMessage))
	nested["raw"].(json.RawMessage)[0] = '['
	nested["changed"] = true
	return contracts.ValidationResult{Approved: true}, nil
}

func (r *extractCaptureRecorder) ExtractSucceeded(_ string, _ string, _ []CheckpointFingerprint, outputs []CheckpointArtifact, rejected []contracts.RejectedOutput) error {
	r.checkpoint = ExtractCheckpoint{
		Outputs:  cloneCheckpointArtifacts(outputs),
		Rejected: cloneRejectedOutputs(rejected),
	}
	return nil
}

type extractResultLoader struct {
	CheckpointLoader
	checkpoint ExtractCheckpoint
	decision   CheckpointDecision
}

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

func (l *extractResultLoader) Extract(string, string, []CheckpointFingerprint) (ExtractCheckpoint, CheckpointDecision) {
	return ExtractCheckpoint{
		Outputs:  cloneCheckpointArtifacts(l.checkpoint.Outputs),
		Rejected: cloneRejectedOutputs(l.checkpoint.Rejected),
	}, l.decision
}

func cloneCheckpointArtifacts(values []CheckpointArtifact) []CheckpointArtifact {
	if len(values) == 0 {
		return nil
	}
	cloned := make([]CheckpointArtifact, len(values))
	for i := range values {
		cloned[i] = cloneCheckpointArtifact(values[i])
	}
	return cloned
}

func TestRunnerPassesIndependentAnnotationScopesToExtractors(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	chunker := prepared.chunker.(*typedTestChunker)
	chunker.plan.Ranges[0].Annotations = source.ChunkAnnotations{"same": json.RawMessage(`{"range":1}`)}
	chunker.plan.Annotations = source.ChunkAnnotations{"same": json.RawMessage(`{"plan":2}`)}

	var gotRange, gotPlan string
	installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		gotRange = string(request.Chunk.Annotations["same"])
		gotPlan = string(request.Chunk.PlanAnnotations["same"])
		request.Chunk.Annotations["same"][0] = '['
		request.Chunk.PlanAnnotations["same"][0] = '['
		return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil
	})
	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if gotRange != `{"range":1}` || gotPlan != `{"plan":2}` {
		t.Fatalf("extract annotations = %q / %q", gotRange, gotPlan)
	}
	if string(chunker.plan.Ranges[0].Annotations["same"]) != `{"range":1}` || string(chunker.plan.Annotations["same"]) != `{"plan":2}` {
		t.Fatal("extract request shares annotation bytes with chunker output")
	}
}

func TestRunnerPassesIndependentConcreteMetadataToValidatorsAndExtractors(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	input := prepared.input.(*typedTestInput)
	input.doc.Units[0].Metadata = map[string]any{
		"typed": pipelineTypedMetadata{"raw": json.RawMessage(`{"value":true}`)},
	}
	var err error
	input.doc.Digest, err = source.DigestDocument(input.doc)
	if err != nil {
		t.Fatal(err)
	}
	chunker := prepared.chunker.(*typedTestChunker)
	chunker.plan = typedTestPlan(input.doc)

	validator := &mutatingMetadataValidator{}
	prepared.chunkValidators.validators = []preparedValidator{{
		resolved: ResolvedValidator{Binding: Binding(validator.Name()), Target: ValidatorTargetChunk},
		chunk:    validator,
	}}
	var extractorSeen string
	installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		nested := request.Chunk.Units[0].Metadata["typed"].(pipelineTypedMetadata)
		extractorSeen = string(nested["raw"].(json.RawMessage))
		nested["raw"].(json.RawMessage)[0] = '['
		nested["changed"] = true
		return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil
	})

	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	if validator.seen != `{"value":true}` || extractorSeen != `{"value":true}` {
		t.Fatalf("validator/extractor metadata = %q / %q", validator.seen, extractorSeen)
	}
	original := input.doc.Units[0].Metadata["typed"].(pipelineTypedMetadata)
	if _, changed := original["changed"]; changed || string(original["raw"].(json.RawMessage)) != `{"value":true}` {
		t.Fatalf("source metadata changed through runner handoff: %#v", original)
	}
}

func TestRunnerRejectsMetadataThatCannotBeCloned(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	cyclic := make(map[string]any)
	cyclic["self"] = cyclic
	_, err := New().Run(context.Background(), RunInput{
		Prepared: prepared,
		RawInput: []byte("input"),
		Metadata: map[string]any{"cyclic": cyclic},
	})
	if err == nil || !strings.Contains(err.Error(), "clone run metadata: metadata.cyclic.self contains a cycle") {
		t.Fatalf("Run() error = %v, want contextual metadata clone failure", err)
	}
}

func TestRunnerContinuesFromFreshAndReusedExtractResults(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	extractCalls := 0
	installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		extractCalls++
		return erasedTypedResult{
			Value: typedValueForLane(0, request.Chunk.Index),
			Diagnostics: []contracts.ProducerDiagnostic{{
				Disposition:     contracts.DiagnosticDispositionObservation,
				Category:        contracts.DiagnosticCategoryNormalization,
				ReasonCode:      "accepted_extract_normalized",
				OccurrenceCount: 1,
				Samples:         []contracts.DiagnosticSample{{Scope: "extract", Message: "accepted extract"}},
			}},
		}, nil
	})

	freshDebug := newCapturedDebugRecorder()
	recorder := &extractCaptureRecorder{CheckpointRecorder: NoopCheckpointRecorder()}
	freshLoader := &extractResultLoader{
		CheckpointLoader: NoopCheckpointLoader(),
		decision:         CheckpointDecision{Reason: "extract checkpoint not found"},
	}
	fresh, err := New().Run(context.Background(), RunInput{
		Prepared:    prepared,
		RawInput:    []byte("input"),
		Checkpoints: recorder,
		Checkpoint:  freshLoader,
		Debug:       freshDebug,
	})
	if err != nil {
		t.Fatalf("fresh Run() error = %v, want nil", err)
	}
	if extractCalls != 1 {
		t.Fatalf("fresh extract calls = %d, want 1", extractCalls)
	}
	assertExtractDecision(t, fresh.CheckpointEvents, "executed", "extract checkpoint not found")
	assertExtractDebugPaths(t, freshDebug, true)

	reusedDebug := newCapturedDebugRecorder()
	reusedLoader := &extractResultLoader{
		CheckpointLoader: NoopCheckpointLoader(),
		checkpoint:       recorder.checkpoint,
		decision:         CheckpointDecision{Reused: true, Reason: "extract checkpoint matched"},
	}
	reused, err := New().Run(context.Background(), RunInput{
		Prepared:   prepared,
		RawInput:   []byte("input"),
		Checkpoint: reusedLoader,
		Debug:      reusedDebug,
	})
	if err != nil {
		t.Fatalf("reused Run() error = %v, want nil", err)
	}
	if extractCalls != 1 {
		t.Fatalf("extract calls after reuse = %d, want 1", extractCalls)
	}
	assertExtractDecision(t, reused.CheckpointEvents, "reused", "extract checkpoint matched")
	assertExtractDebugPaths(t, reusedDebug, false)
	if !reflect.DeepEqual(reused.NormalizeOutputs, fresh.NormalizeOutputs) {
		t.Fatalf("reused normalize outputs = %#v, want fresh outputs %#v", reused.NormalizeOutputs, fresh.NormalizeOutputs)
	}
	if !reflect.DeepEqual(reused.Diagnostics, fresh.Diagnostics) {
		t.Fatalf("reused diagnostics = %#v, want fresh diagnostics %#v", reused.Diagnostics, fresh.Diagnostics)
	}
}

func TestRunnerPromotesOnlyAcceptedExtractRetryDiagnostics(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	prepared.Steps[0].lanes[0].resolved.Extract.Retries = 1
	attempts := 0
	installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		attempts++
		scope := "discarded"
		if attempts == 2 {
			scope = "accepted"
		}
		return erasedTypedResult{
			Value:          typedValueForLane(0, request.Chunk.Index),
			Diagnostics:    []contracts.ProducerDiagnostic{producerDiagnostic(scope, scope)},
			ModelCandidate: attemptCandidate(t, fmt.Sprintf(`{"items":["%s"]}`, scope)),
		}, nil
	})
	validatorCalls := 0
	prepared.Steps[0].lanes[0].extractValidators.validators[0].typedValidate = func(context.Context, any, typedValidationTarget) (contracts.ValidationResult, error) {
		validatorCalls++
		return contracts.ValidationResult{Approved: validatorCalls == 2, ReasonCode: "retry", Message: "retry extract", CorrectionGuidance: "return an acceptable extract"}, nil
	}
	debug := newCapturedDebugRecorder()
	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if attempts != 2 {
		t.Fatalf("extract attempts = %d, want 2", attempts)
	}
	if len(output.Diagnostics.Groups) != 1 || output.Diagnostics.Groups[0].Samples[0].Scope != "accepted" {
		t.Fatalf("promoted diagnostics = %#v, want accepted attempt only", output.Diagnostics)
	}
	assertAttemptEnvelopeSequence(t, debug, "extract/notes/chunk-000001", 1, 2)
	first := debug.envelope(t, "extract/notes/chunk-000001/attempt-01.json")
	if first.Error != "" || !strings.Contains(string(debug.json["extract/notes/chunk-000001/attempt-01.json"]), "rejection") {
		t.Fatalf("first extract attempt = %#v, want rejection without error", first)
	}
	name := "extract/notes/chunk-000001/attempt-02.json"
	if !debug.has(name) {
		t.Fatalf("debug artifact %q is missing; names = %#v", name, debug.names())
	}
}

func assertExtractDecision(t *testing.T, events []CheckpointEvent, action CheckpointDecisionCategory, reason string) {
	t.Helper()
	for _, event := range events {
		if event.Stage == string(StageExtract) {
			if event.Action != action || event.Reason != reason {
				t.Fatalf("extract checkpoint event = %#v, want action %q reason %q", event, action, reason)
			}
			return
		}
	}
	t.Fatalf("extract checkpoint event missing from %#v", events)
}

func assertExtractDebugPaths(t *testing.T, debug *capturedDebugRecorder, wantAttempt bool) {
	t.Helper()
	for _, name := range []string{"extract/notes/input.json", "extract/notes/output.json"} {
		if !debug.has(name) {
			t.Fatalf("debug artifact %q is missing; names = %#v", name, debug.names())
		}
	}
	attemptPath := "extract/notes/chunk-000001/attempt-01.json"
	if debug.has(attemptPath) != wantAttempt {
		t.Fatalf("attempt debug path present = %t, want %t; names = %#v", debug.has(attemptPath), wantAttempt, debug.names())
	}
	input := debug.envelope(t, "extract/notes/input.json")
	wantReuse := !wantAttempt
	if !strings.Contains(string(debug.json["extract/notes/input.json"]), `"reused":`+map[bool]string{true: "true", false: "false"}[wantReuse]) {
		t.Fatalf("extract input payload = %#v, want reused %t", input.Payload, wantReuse)
	}
}