notarius/internal/modules/dnd/normalize/locationregistry/normalizer_test.go

package locationregistry

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

	"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/framework/semanticreconcile"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/identity"
)

func TestModuleContractAndMetadata(t *testing.T) {
	want := pipeline.ModuleSpec{Key: Key, Stage: pipeline.StageNormalize, ExecutionClass: contracts.ExecutionClassLLMBacked, CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1, Requires: []string{"merged"}, Provides: []string{"normalized"}, ArtifactKind: dnd.LocationRegistryKind}
	if got := ModuleSpec(); !reflect.DeepEqual(got, want) {
		t.Fatalf("ModuleSpec() = %#v, want %#v", got, want)
	}
	registry := pipeline.NewNormalizerRegistry()
	if err := Register(registry); err != nil {
		t.Fatal(err)
	}
	if _, err := New(nil, Options{}); err == nil {
		t.Fatal("New() accepted nil client")
	}
	normalizer := newNormalizer(t, &recordingLocationNormalizerClient{})
	metadata := normalizer.ManifestMetadata()
	limits, ok := metadata["semantic_reconciliation_limits"].(map[string]any)
	if !ok || metadata["identity_policy"] != identity.Policy || metadata["normalization_policy"] != normalizationPolicy || metadata["prompt_id"] != PromptID || metadata["prompt_version"] != PromptVersion || metadata["response_schema_key"] != string(semanticreconcile.ResponseSchemaKey) || metadata["response_schema_id"] != semanticreconcile.ResponseSchemaID || metadata["response_schema_name"] != semanticreconcile.ResponseSchemaName || metadata["semantic_reconciliation_policy"] != semanticreconcile.Policy || len(limits) != 3 {
		t.Fatalf("metadata = %#v", metadata)
	}
	for _, name := range []string{"prompt", "response_schema", "semantic_reconciliation_policy", "semantic_reconciliation_limits", "identity_policy", "normalization_policy"} {
		if !hasFingerprint(normalizer.CheckpointFingerprints(), name) {
			t.Fatalf("fingerprints missing %q", name)
		}
	}
}

func TestNormalizeRejectsNilSourceDocument(t *testing.T) {
	_, err := newNormalizer(t, &recordingLocationNormalizerClient{}).Normalize(
		context.Background(),
		contracts.TypedNormalizeRequest[dnd.LocationRegistry]{},
	)
	if err == nil || !strings.Contains(err.Error(), "source document must not be nil") {
		t.Fatalf("Normalize() error = %v, want nil source rejection", err)
	}
}

func TestNormalizePreparesOnlyExactDuplicatesAndRetainsSameNameAndNestedPlaces(t *testing.T) {
	input := dnd.LocationRegistry{Locations: []dnd.Location{
		{Name: "  The Tavern  ", SourceRefs: []source.SourceRef{{SourceID: "session", StartUnitID: 1, EndUnitID: 1}}},
		{Name: "the tavern", SourceRefs: []source.SourceRef{{SourceID: "session", StartUnitID: 1, EndUnitID: 1}}},
		{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: "session", StartUnitID: 2, EndUnitID: 2}}},
		{Name: "The Tavern Cellar", SourceRefs: []source.SourceRef{{SourceID: "session", StartUnitID: 2, EndUnitID: 2}}},
	}}
	before := dnd.LocationRegistry{Locations: append([]dnd.Location(nil), input.Locations...)}
	client := &recordingLocationNormalizerClient{}
	result, err := newNormalizer(t, client).Normalize(context.Background(), normalizeRequest(input))
	if err != nil || len(result.Value.Locations) != 3 {
		t.Fatalf("Normalize() = %#v, %v; want one exact duplicate removed", result, err)
	}
	if !reflect.DeepEqual(input, before) {
		t.Fatalf("Normalize() mutated input: %#v", input)
	}
	if got := []string{result.Value.Locations[0].Name, result.Value.Locations[1].Name, result.Value.Locations[2].Name}; !reflect.DeepEqual(got, []string{"The Tavern", "The Tavern", "The Tavern Cellar"}) {
		t.Fatalf("locations = %#v, want same names and nested place retained", got)
	}
	if result.Value.Locations[0].ID == result.Value.Locations[1].ID || !hasDiagnostic(result.Diagnostics, ReasonCodeDuplicateLocationCollapsed, contracts.DiagnosticDispositionObservation) || len(client.requests) != 0 {
		t.Fatalf("result = %#v, want evidence-anchored IDs and exact duplicate finding", result)
	}
}

func BenchmarkExactDuplicateGroupsManyDistinct(b *testing.B) {
	records := make([]normalizedRecord, 1_000)
	for index := range records {
		records[index] = normalizedRecord{location: dnd.Location{
			Name:       "Location " + strconv.Itoa(index),
			SourceRefs: []source.SourceRef{{SourceID: "session", StartUnitID: index + 1, EndUnitID: index + 1}},
		}}
	}
	b.ReportAllocs()
	b.ResetTimer()
	for iteration := 0; iteration < b.N; iteration++ {
		if groups := exactDuplicateGroups(records); len(groups) != len(records) {
			b.Fatalf("group count = %d, want %d", len(groups), len(records))
		}
	}
}

func TestNormalizeAppliesSafeAliasGroupAndUsesContextualInputs(t *testing.T) {
	client := &recordingLocationNormalizerClient{response: `{"duplicate_groups":[{"candidate_numbers":[1,2],"canonical_candidate_number":2}]}`}
	doc := semanticDocument()
	input := dnd.LocationRegistry{Locations: []dnd.Location{
		{Name: "Old Mill", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}}},
		{Name: "the Greencloak's refuge", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 20, EndUnitID: 20}}},
		{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 30, EndUnitID: 30}}},
	}}
	result, err := newNormalizer(t, client).Normalize(context.Background(), normalizeRequestWithSource(input, doc))
	if err != nil || result.Retry != nil || len(result.Value.Locations) != 2 {
		t.Fatalf("Normalize() = %#v, %v", result, err)
	}
	merged := result.Value.Locations[0]
	if merged.Name != "the Greencloak's refuge" || merged.ID != identity.DeriveID(merged.Name, merged.SourceRefs) || len(merged.SourceRefs) != 2 || !hasDiagnostic(result.Diagnostics, ReasonCodeDuplicateLocationCollapsed, contracts.DiagnosticDispositionObservation) {
		t.Fatalf("merged location = %#v, diagnostics = %#v", merged, result.Diagnostics)
	}
	encoded := string(client.requests[0].Inputs["candidates"].Content) + string(client.requests[0].Inputs["transcript"].Content)
	if strings.Contains(encoded, doc.ID) || strings.Contains(encoded, "candidate-") || strings.Contains(encoded, merged.ID) || !strings.Contains(encoded, `"source_refs"`) {
		t.Fatalf("private inputs = %s", encoded)
	}
}

func TestNormalizeRejectsUnsafeAndOverlappingGroupsWithoutLosingCandidates(t *testing.T) {
	doc := semanticDocument()
	input := dnd.LocationRegistry{Locations: []dnd.Location{
		{Name: "Old Mill", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}}},
		{Name: "Mill", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 20, EndUnitID: 20}}},
		{Name: "Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 30, EndUnitID: 30}}},
	}}
	client := &recordingLocationNormalizerClient{response: `{"duplicate_groups":[{"candidate_numbers":[1,2],"canonical_candidate_number":1},{"candidate_numbers":[2,3],"canonical_candidate_number":3}]}`}
	result, err := newNormalizer(t, client).Normalize(context.Background(), normalizeRequestWithSource(input, doc))
	if err != nil || result.Retry == nil || len(result.Value.Locations) != 3 || !strings.Contains(result.Retry.Message, "overlapping_member") {
		t.Fatalf("Normalize() = %#v, %v; want safe retry fallback", result, err)
	}
	if !strings.Contains(result.Retry.CorrectionGuidance, "Duplicate groups 1, 2") || !strings.Contains(result.Retry.CorrectionGuidance, "at most one duplicate group") || strings.Contains(result.Retry.CorrectionGuidance, "overlapping_member") {
		t.Fatalf("correction guidance = %q, want model-facing overlap explanation", result.Retry.CorrectionGuidance)
	}
	if len(result.Retry.FallbackDiagnostics) != 1 || !strings.Contains(result.Retry.FallbackDiagnostics[0].Samples[0].Message, "2 proposal group") || !strings.Contains(result.Retry.FallbackDiagnostics[0].Samples[0].Message, "preserved separately") {
		t.Fatalf("fallback diagnostics = %#v", result.Retry.FallbackDiagnostics)
	}
}

func TestReconciliationCandidatesKeepSameNameEvidenceDistinct(t *testing.T) {
	doc := semanticDocument()
	records := []normalizedRecord{
		{location: dnd.Location{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}}}, inputIndexes: []int{0}, earliest: 0},
		{location: dnd.Location{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 30, EndUnitID: 30}}}, inputIndexes: []int{1}, earliest: 1},
	}
	candidates, _, err := reconciliationInputs(records)
	if err != nil {
		t.Fatalf("reconciliationInputs() error = %v", err)
	}
	preparation, err := semanticreconcile.Prepare(doc, candidates, semanticreconcile.DefaultLimits())
	if err != nil || preparation.Disposition() != semanticreconcile.Ready {
		t.Fatalf("Prepare() = %#v, %v; want ready candidates", preparation, err)
	}
	var candidateInput struct {
		Candidates []struct {
			CandidateNumber int    `json:"candidate_number"`
			Label           string `json:"label"`
		} `json:"candidates"`
	}
	if err := json.Unmarshal(preparation.Materials()["candidates"].Content, &candidateInput); err != nil {
		t.Fatal(err)
	}
	if len(candidateInput.Candidates) != 2 || candidateInput.Candidates[0].CandidateNumber != 1 || candidateInput.Candidates[1].CandidateNumber != 2 || candidateInput.Candidates[0].Label != "The Tavern" || candidateInput.Candidates[1].Label != "The Tavern" {
		t.Fatalf("candidate input = %#v, want distinct integer handles for equal names", candidateInput)
	}
}

func TestNormalizeLimitSkipDoesNotCallLLMAndAddsBoundedFallbackWarning(t *testing.T) {
	client := &recordingLocationNormalizerClient{}
	doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{{ID: 1, Kind: "narration", Text: "A sprawling city"}}}
	limit := semanticreconcile.DefaultLimits().MaximumCandidates
	input := dnd.LocationRegistry{Locations: make([]dnd.Location, limit+1)}
	for index := range input.Locations {
		input.Locations[index] = dnd.Location{Name: fmt.Sprintf("Place %d", index), SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 1, EndUnitID: 1}}}
	}
	result, err := newNormalizer(t, client).Normalize(context.Background(), normalizeRequestWithSource(input, doc))
	if err != nil || result.Retry != nil {
		t.Fatalf("Normalize() = %#v, %v; want deterministic limit fallback", result, err)
	}
	if len(client.requests) != 0 || len(result.Value.Locations) != limit+1 {
		t.Fatalf("completion calls = %d, locations = %d; want no call and all records", len(client.requests), len(result.Value.Locations))
	}
	if !hasDiagnostic(result.Diagnostics, ReasonCodeLocationSemanticReconciliationExhausted, contracts.DiagnosticDispositionWarning) {
		t.Fatalf("diagnostics = %#v, want reconciliation fallback warning", result.Diagnostics)
	}
}

func TestNormalizeHandlesRetryFallbackAndErrors(t *testing.T) {
	doc := semanticDocument()
	input := dnd.LocationRegistry{Locations: []dnd.Location{{Name: "Old Mill", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}}}, {Name: "Mill", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 20, EndUnitID: 20}}}}}
	invalid, err := newNormalizer(t, &recordingLocationNormalizerClient{err: contracts.ErrInvalidStructuredOutput}).Normalize(context.Background(), normalizeRequestWithSource(input, doc))
	if err != nil || invalid.Retry == nil || invalid.Retry.ReasonCode != ReasonCodeLocationSemanticProposalInvalid {
		t.Fatalf("invalid result = %#v, %v", invalid, err)
	}
	if invalid.Retry.CorrectionGuidance != "" || !strings.Contains(invalid.Retry.FallbackDiagnostics[0].Samples[0].Message, "structured-output contract") || !strings.Contains(invalid.Retry.FallbackDiagnostics[0].Samples[0].Message, "preserved separately") {
		t.Fatalf("invalid structured retry = %#v, want feedback-free retry and informative fallback", invalid.Retry)
	}
	_, err = newNormalizer(t, &recordingLocationNormalizerClient{err: errors.New("provider unavailable")}).Normalize(context.Background(), normalizeRequestWithSource(input, doc))
	if err == nil || !strings.Contains(err.Error(), "provider unavailable") {
		t.Fatalf("provider error = %v", err)
	}
}

func TestNormalizeOrdersEvidenceAndIsIdempotent(t *testing.T) {
	doc := &source.SourceDocument{ID: "ordered", Units: []source.SourceUnit{{ID: 30}, {ID: 10}}}
	input := dnd.LocationRegistry{Locations: []dnd.Location{{Name: "Old Mill", SourceRefs: []source.SourceRef{
		{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10},
		{SourceID: doc.ID, StartUnitID: 30, EndUnitID: 30},
		{SourceID: doc.ID, StartUnitID: 30, EndUnitID: 30},
	}}}}
	normalizer := newNormalizer(t, &recordingLocationNormalizerClient{})
	first, err := normalizer.Normalize(context.Background(), normalizeRequestWithSource(input, doc))
	if err != nil || !reflect.DeepEqual([]int{first.Value.Locations[0].SourceRefs[0].StartUnitID, first.Value.Locations[0].SourceRefs[1].StartUnitID}, []int{30, 10}) {
		t.Fatalf("first Normalize() = %#v, %v; want document-ordered evidence", first, err)
	}
	second, err := normalizer.Normalize(context.Background(), normalizeRequestWithSource(first.Value, doc))
	if err != nil || !reflect.DeepEqual(second.Value, first.Value) {
		t.Fatalf("second Normalize() = %#v, %v; want idempotent value %#v", second, err, first.Value)
	}
}

func hasFingerprint(fingerprints []pipeline.CheckpointFingerprint, name string) bool {
	for _, fingerprint := range fingerprints {
		if fingerprint.Name == name && fingerprint.Value != "" {
			return true
		}
	}
	return false
}