notarius/internal/modules/dnd/extract/locationoccurrences/extractor_test.go

package locationoccurrences

import (
	"context"
	"encoding/json"
	"errors"
	"reflect"
	"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/modules/dnd"
	locationcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locationregistry"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/identity"
	locationregistry "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/registry"
)

func TestExtractMapsKindsOrdersOccurrencesAndPreservesIndependentFacts(t *testing.T) {
	locations := locationRegistry(t, "The Tavern", "The Tavern")
	first, second := locations.Locations[0], locations.Locations[1]
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{
		{Name: second.Name, RegistryRefs: registryRefs(second), Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30)},
		{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
		{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "recalled", SourceRefs: occurrenceRefs(20, 10)},
		{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "planned", SourceRefs: occurrenceRefs(10, 10)},
		{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "visited", SourceRefs: append(occurrenceRefs(10, 10), occurrenceRefs(10, 10)...)},
		{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "visited", SourceRefs: occurrenceRefs(20, 20)},
		{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "visited", SourceRefs: occurrenceRefs(10, 10)},
	}}}
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	result, err := newExtractor(t, client, references).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	if len(result.Value.Occurrences) != 6 {
		t.Fatalf("occurrences = %#v, want exact duplicate removed", result.Value.Occurrences)
	}
	got := result.Value.Occurrences
	if kinds := []dnd.LocationOccurrenceKind{got[0].Kind, got[1].Kind, got[2].Kind, got[3].Kind}; !reflect.DeepEqual(kinds, []dnd.LocationOccurrenceKind{dnd.LocationOccurrenceKindVisited, dnd.LocationOccurrenceKindPlanned, dnd.LocationOccurrenceKindRecalled, dnd.LocationOccurrenceKindMentioned}) {
		t.Fatalf("same-evidence kind order = %#v", kinds)
	}
	if got[4].Kind != dnd.LocationOccurrenceKindVisited || got[4].SourceRefs[0].StartUnitID != 20 || got[5].LocationID != second.ID || got[5].SourceRefs[0].StartUnitID != 30 {
		t.Fatalf("occurrence order = %#v", got)
	}
	if !reflect.DeepEqual(got[0].SourceRefs, []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 10, EndUnitID: 10}}) {
		t.Fatalf("canonical evidence = %#v", got[0].SourceRefs)
	}
	if !reflect.DeepEqual(got[2].SourceRefs, []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 10, EndUnitID: 20}}) {
		t.Fatalf("reversed evidence = %#v", got[2].SourceRefs)
	}
}

func TestExtractResolvesContextualSelectorsAndUsesCurrentTranscriptEvidenceOnly(t *testing.T) {
	locations := locationRegistry(t, "The Tavern", "The Tavern")
	first, second := locations.Locations[0], locations.Locations[1]
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
		Name: second.Name, RegistryRefs: registryRefs(second), Kind: "visited", SourceRefs: occurrenceRefs(10, 10),
	}}}}
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	result, err := newExtractor(t, client, references).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	if occurrence := result.Value.Occurrences[0]; occurrence.LocationID != second.ID || occurrence.Name != second.Name || occurrence.SourceRefs[0].SourceID != req.Source.ID {
		t.Fatalf("occurrence = %#v", occurrence)
	}
	input := client.requests[0].Inputs[LocationRegistryReferenceSlot]
	if input.Name != LocationRegistryReferenceSlot || !strings.Contains(string(input.Content), `"registry_refs":[{"start_unit_id":20,"end_unit_id":20}]`) {
		t.Fatalf("location prompt input = %#v", input)
	}
	for _, forbidden := range []string{"source_refs", "source_id", first.ID, second.ID} {
		if strings.Contains(string(input.Content), forbidden) {
			t.Fatalf("location prompt leaked %q: %s", forbidden, input.Content)
		}
	}
	if strings.Contains(string(client.requests[0].Inputs["transcript"].Content), "other-session") {
		t.Fatal("transcript input contains registry evidence")
	}
	metadata, err := json.Marshal(newExtractor(t, &fakeOccurrencesLLMClient{}, references).ManifestMetadata())
	if err != nil || strings.Contains(string(metadata), "other-session") || strings.Contains(string(metadata), first.ID) {
		t.Fatalf("manifest metadata = %s, %v", metadata, err)
	}
}

func TestExtractRejectsUnknownMalformedOrMismatchedSelectorsAtomically(t *testing.T) {
	locations := locationRegistry(t, "The Mill")
	known := locations.Locations[0]
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	for _, occurrences := range [][]occurrenceResponse{
		{{Name: "Unknown", RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
		{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 0}}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
		{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 10}}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
		{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}, {Name: "Unknown", RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(20, 20)}},
	} {
		client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: occurrences}}
		if result, err := newExtractor(t, client, references).Extract(context.Background(), req); err == nil || result.Value.Occurrences != nil || !strings.Contains(err.Error(), "location selector") {
			t.Fatalf("Extract() = %#v, %v", result, err)
		}
	}

	sharedName := "The Tavern"
	firstRefs := []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 10, EndUnitID: 10}, {SourceID: req.Source.ID, StartUnitID: 20, EndUnitID: 20}}
	secondRefs := []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 30, EndUnitID: 30}}
	duplicateLocations := dnd.LocationRegistry{Locations: []dnd.Location{
		{ID: identity.DeriveID(sharedName, firstRefs), Name: sharedName, SourceRefs: firstRefs},
		{ID: identity.DeriveID(sharedName, secondRefs), Name: sharedName, SourceRefs: secondRefs},
	}}
	duplicateReferences := registryReferences(t, duplicateLocations)
	for _, selector := range []struct {
		name string
		refs []locationregistry.RegistryRef
	}{
		{name: sharedName, refs: []locationregistry.RegistryRef{}},
		{name: sharedName, refs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 10}}},
		{name: sharedName, refs: []locationregistry.RegistryRef{{StartUnitID: 20, EndUnitID: 20}, {StartUnitID: 10, EndUnitID: 10}}},
	} {
		client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{Name: selector.name, RegistryRefs: selector.refs, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}}}}
		duplicateRequest := extractionRequest()
		duplicateRequest.References = duplicateReferences
		if _, err := newExtractor(t, client, duplicateReferences).Extract(context.Background(), duplicateRequest); err == nil || !strings.Contains(err.Error(), "location selector") {
			t.Fatalf("Extract(%#v) error = %v", selector, err)
		}
	}
}

func TestExtractRequiresRegistryAndAcceptsEmptyRegistryWithNoOccurrences(t *testing.T) {
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{}}}
	if _, err := newExtractor(t, client).Extract(context.Background(), extractionRequest()); err == nil || !strings.Contains(err.Error(), "location registry reference is required") {
		t.Fatalf("Extract() error = %v", err)
	}
	if len(client.requests) != 0 {
		t.Fatalf("LLM calls = %d", len(client.requests))
	}
	empty := dnd.LocationRegistry{Locations: []dnd.Location{}}
	references := registryReferences(t, empty)
	req := extractionRequest()
	req.References = references
	result, err := newExtractor(t, client, references).Extract(context.Background(), req)
	if err != nil || result.Value.Occurrences == nil || len(result.Value.Occurrences) != 0 {
		t.Fatalf("empty registry result = %#v, %v", result, err)
	}
}

func TestExtractResolvesGeneratedRegistryAtOperationTimeAndDoesNotMutateResponse(t *testing.T) {
	locations := locationRegistry(t, "The Mill")
	location := locations.Locations[0]
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
		Name: location.Name, RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30),
	}}}}
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	extractor := newExtractor(t, client)
	result, err := extractor.Extract(context.Background(), req)
	if err != nil || result.Value.Occurrences[0].Name != "The Mill" {
		t.Fatalf("Extract() = %#v, %v", result, err)
	}
	if input := client.requests[0].Inputs[LocationRegistryReferenceSlot]; strings.Contains(string(input.Content), location.ID) || input.OriginURI != "" {
		t.Fatalf("generated registry prompt input = %#v", input)
	}
	if _, ok := extractor.ManifestMetadata()["location_registry_digest"]; ok {
		t.Fatalf("operation registry leaked into static metadata: %#v", extractor.ManifestMetadata())
	}
	if client.response.Occurrences[0].SourceRefs[0].StartUnitID != 30 {
		t.Fatalf("model response mutated: %#v", client.response)
	}
}

func TestExtractorContractsMetadataAndFailures(t *testing.T) {
	if _, err := New(nil, Options{}); err == nil || !strings.Contains(err.Error(), "LLM client") {
		t.Fatalf("New(nil) error = %v", err)
	}
	if _, err := New(&fakeOccurrencesLLMClient{}, Options{}, contracts.ReferenceSet{}, contracts.ReferenceSet{}); err == nil || !strings.Contains(err.Error(), "at most one reference set") {
		t.Fatalf("New() error = %v", err)
	}
	malformed := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{LocationRegistryReferenceSlot: {
		Items: []contracts.ReferenceItem{{SlotName: LocationRegistryReferenceSlot, MediaType: "application/json", Content: []byte(`{"secret":"registry evidence"}`)}},
	}}}
	if _, err := New(&fakeOccurrencesLLMClient{}, Options{}, malformed); err == nil || !strings.Contains(err.Error(), "prepare location registry") || strings.Contains(err.Error(), "registry evidence") {
		t.Fatalf("New() error = %v", err)
	}

	locations := locationRegistry(t, "The Mill")
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	extractor := newExtractor(t, &fakeOccurrencesLLMClient{}, references)
	var nilExtractor *Extractor
	for _, test := range []struct {
		name      string
		extractor *Extractor
		req       contracts.TypedExtractionRequest
		want      string
	}{
		{"nil extractor", nilExtractor, req, "extractor"},
		{"nil client", &Extractor{}, req, "LLM client"},
		{"invalid request", extractor, mismatchedSourceInputRequest(req), "must match chunk"},
	} {
		t.Run(test.name, func(t *testing.T) {
			if _, err := test.extractor.Extract(context.Background(), test.req); err == nil || !strings.Contains(err.Error(), test.want) {
				t.Fatalf("Extract() error = %v", err)
			}
		})
	}
	if _, err := newExtractor(t, &fakeOccurrencesLLMClient{err: errors.New("provider unavailable")}, references).Extract(context.Background(), req); err == nil || !strings.Contains(err.Error(), "provider unavailable") {
		t.Fatalf("provider error = %v", err)
	}

	spec := ModuleSpec()
	if spec.Key != Key || spec.Stage != pipeline.StageExtract || spec.ExecutionClass != contracts.ExecutionClassLLMBacked || spec.CorrectionProtocol != contracts.CorrectionProtocolSingleResponseV1 || spec.ArtifactKind != dnd.LocationOccurrenceListKind {
		t.Fatalf("ModuleSpec() = %#v", spec)
	}
	var slot contracts.ReferenceSlot
	for _, candidate := range spec.ReferenceSlots {
		if candidate.Name == LocationRegistryReferenceSlot {
			slot = candidate
		}
	}
	if !slot.Required || !reflect.DeepEqual(slot.AcceptedArtifactKinds, []contracts.ArtifactKind{dnd.LocationRegistryKind}) || slot.MaxBytes != LocationRegistryMaxBytes {
		t.Fatalf("location registry slot = %#v", slot)
	}
	registry := pipeline.NewExtractorRegistry()
	if err := Register(registry); err != nil {
		t.Fatal(err)
	}
	if _, ok := registry.Spec(Key); !ok {
		t.Fatalf("registration missing %q", Key)
	}
	if _, err := DecodeOptions(map[string]any{"unexpected": true}); err == nil {
		t.Fatal("DecodeOptions() accepted unknown options")
	}
	metadata := newExtractor(t, &fakeOccurrencesLLMClient{}, references).ManifestMetadata()
	for _, key := range []string{"prompt_sha256", "response_schema_sha256", "location_registry_digest"} {
		if value, ok := metadata[key].(string); !ok || !strings.HasPrefix(value, "sha256:") {
			t.Fatalf("metadata[%q] = %#v", key, metadata[key])
		}
	}
	if got := newExtractor(t, &fakeOccurrencesLLMClient{}, references).CheckpointFingerprints(); len(got) != 4 || got[3].Name != "location_registry" || got[3].Value != locationRegistryIdentityDigest(t, references) {
		t.Fatalf("fingerprints = %#v", got)
	}
}

func locationRegistryIdentityDigest(t *testing.T, references contracts.ReferenceSet) string {
	t.Helper()
	resolver, err := locationregistry.NewResolver(references)
	if err != nil {
		t.Fatal(err)
	}
	return resolver.Seeded().IdentityDigest()
}

func locationRegistry(t *testing.T, names ...string) dnd.LocationRegistry {
	t.Helper()
	locations := make([]dnd.Location, len(names))
	for index, name := range names {
		refs := []source.SourceRef{{SourceID: sourceDocument().ID, StartUnitID: (index + 1) * 10, EndUnitID: (index + 1) * 10}}
		locations[index] = dnd.Location{ID: identity.DeriveID(name, refs), Name: name, SourceRefs: refs}
	}
	return dnd.LocationRegistry{Locations: locations}
}

func registryRefs(location dnd.Location) []locationregistry.RegistryRef {
	refs := make([]locationregistry.RegistryRef, len(location.SourceRefs))
	for index, ref := range location.SourceRefs {
		refs[index] = locationregistry.RegistryRef{StartUnitID: ref.StartUnitID, EndUnitID: ref.EndUnitID}
	}
	return refs
}

func registryReferences(t *testing.T, locations dnd.LocationRegistry) contracts.ReferenceSet {
	t.Helper()
	content, err := locationcodec.New().Encode(locations)
	if err != nil {
		t.Fatal(err)
	}
	return contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{LocationRegistryReferenceSlot: {
		Slot:  contracts.ReferenceSlot{Name: LocationRegistryReferenceSlot},
		Items: []contracts.ReferenceItem{{SlotName: LocationRegistryReferenceSlot, MediaType: locationcodec.MediaType, Content: content, Origin: contracts.ReferenceOrigin{Type: "generated"}}},
	}}}
}