notarius/internal/modules/dnd/shared/entityreconcile/entityreconcile_test.go

package entityreconcile

import (
	"bytes"
	"encoding/json"
	"io/fs"
	"reflect"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
	"github.com/santhosh-tekuri/jsonschema/v6"
)

func TestBuildContextUsesContextualSelectorsSourceOrderAndOwnedData(t *testing.T) {
	doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{
		{ID: 40, Kind: "narration", Text: "zero"},
		{ID: 10, Kind: "speech", Text: "one", Metadata: map[string]any{"speaker": map[string]any{"name": "Mira"}}},
		{ID: 70, Kind: "speech", Text: "two"},
		{ID: 20, Kind: "narration", Text: "three"},
		{ID: 90, Kind: "speech", Text: "four"},
	}}
	candidates := []Candidate{
		{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 20}}},
		{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 90, EndUnitID: 90}}},
		{Name: "Broken", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 20, EndUnitID: 10}}},
	}
	before := cloneCandidates(candidates)

	materials, ready, err := BuildContext(doc, candidates, 1)
	if err != nil || !ready {
		t.Fatalf("BuildContext() = %#v, %t, %v; want ready materials", materials, ready, err)
	}
	if !reflect.DeepEqual(candidates, before) {
		t.Fatalf("BuildContext() mutated candidates: %#v", candidates)
	}
	if got, want := materials.CandidateKeys(), []string{"candidate-000001", "candidate-000002", "candidate-000003"}; !reflect.DeepEqual(got, want) {
		t.Fatalf("CandidateKeys() = %#v, want %#v", got, want)
	}
	if got, want := materials.EligibleCandidateKeys(), []string{"candidate-000001", "candidate-000002"}; !reflect.DeepEqual(got, want) {
		t.Fatalf("EligibleCandidateKeys() = %#v, want %#v", got, want)
	}
	if !json.Valid(materials.Candidates.Content) || !json.Valid(materials.Transcript.Content) {
		t.Fatalf("prompt materials are not JSON: %#v", materials)
	}
	if strings.Contains(string(materials.Candidates.Content), doc.ID) {
		t.Fatalf("candidate material leaked source identity: %s", materials.Candidates.Content)
	}

	var candidatePayload candidateInput
	if err := json.Unmarshal(materials.Candidates.Content, &candidatePayload); err != nil {
		t.Fatal(err)
	}
	if len(candidatePayload.Candidates) != 2 || candidatePayload.Candidates[0].Name != candidatePayload.Candidates[1].Name || strings.Contains(string(materials.Candidates.Content), "candidate-") {
		t.Fatalf("candidate payload = %#v, want contextual descriptors without keys", candidatePayload)
	}
	if got := candidatePayload.Candidates[0].SourceRefs[0]; got != (SourceRange{StartUnitID: 10, EndUnitID: 20}) {
		t.Fatalf("candidate reference = %#v", got)
	}

	var transcript transcriptInput
	if err := json.Unmarshal(materials.Transcript.Content, &transcript); err != nil {
		t.Fatal(err)
	}
	if len(transcript.Windows) != 1 || len(transcript.Windows[0].Units) != 5 {
		t.Fatalf("windows = %#v, want one bounded coalesced window", transcript.Windows)
	}
	units := transcript.Windows[0].Units
	for index, wantID := range []int{40, 10, 70, 20, 90} {
		if units[index].ID != wantID {
			t.Fatalf("window unit %d = %d, want source-order %d", index, units[index].ID, wantID)
		}
	}
	if units[0].Cited || !units[1].Cited || !units[2].Cited || !units[3].Cited || !units[4].Cited {
		t.Fatalf("citation flags = %#v", units)
	}

	windows, err := contextWindows(doc.Units, []sourceInterval{{start: 1, end: 1}}, make([]bool, len(doc.Units)))
	if err != nil {
		t.Fatal(err)
	}
	windows[0].Units[0].Metadata["speaker"].(map[string]any)["name"] = "changed"
	if doc.Units[1].Metadata["speaker"].(map[string]any)["name"] != "Mira" {
		t.Fatal("context metadata aliases source document")
	}
}

func TestBuildContextExcludesCollidingDescriptors(t *testing.T) {
	doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{{ID: 1}, {ID: 2}, {ID: 3}}}
	candidates := []Candidate{
		{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 1, EndUnitID: 1}}},
		{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 1, EndUnitID: 1}}},
		{Name: "The Market", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 3, EndUnitID: 3}}},
	}
	materials, ready, err := BuildContext(doc, candidates, 0)
	if err != nil || ready || len(materials.EligibleCandidateKeys()) != 1 || strings.Contains(string(materials.Candidates.Content), "The Tavern") {
		t.Fatalf("BuildContext() = %#v, %t, %v", materials, ready, err)
	}
	assessment := materials.Assess(ProposalResponse{DuplicateGroups: []DuplicateGroup{{
		Members:   []Selector{{Name: "The Tavern", SourceRefs: []SourceRange{{StartUnitID: 1, EndUnitID: 1}}}, {Name: "The Market", SourceRefs: []SourceRange{{StartUnitID: 3, EndUnitID: 3}}}},
		Canonical: Selector{Name: "The Market", SourceRefs: []SourceRange{{StartUnitID: 3, EndUnitID: 3}}},
	}}})
	if !hasIssue(assessment.Issues(), "member_ineligible") {
		t.Fatalf("Assess() issues = %#v, want collided descriptor rejection", assessment.Issues())
	}
}

func TestAssessmentRejectsPartialAndReorderedDescriptors(t *testing.T) {
	doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{{ID: 10}, {ID: 20}, {ID: 30}}}
	materials, ready, err := BuildContext(doc, []Candidate{
		{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 20, EndUnitID: 20}, {SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}}},
		{Name: "The Market", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 30, EndUnitID: 30}}},
	}, 0)
	if err != nil || !ready {
		t.Fatalf("BuildContext() = %#v, %t, %v", materials, ready, err)
	}
	selectors := materialSelectors(t, materials)
	for _, refs := range [][]SourceRange{
		{{StartUnitID: 10, EndUnitID: 10}},
		{{StartUnitID: 20, EndUnitID: 20}, {StartUnitID: 10, EndUnitID: 10}},
	} {
		assessment := materials.Assess(ProposalResponse{DuplicateGroups: []DuplicateGroup{{
			Members:   []Selector{{Name: "The Tavern", SourceRefs: refs}, selectors[1]},
			Canonical: selectors[1],
		}}})
		if !hasIssue(assessment.Issues(), "member_unknown") {
			t.Fatalf("Assess(%#v) issues = %#v, want descriptor mismatch rejection", refs, assessment.Issues())
		}
	}
}

func TestBuildContextExcludesUnsafeReferencesAndCoalescesAdjacentWindows(t *testing.T) {
	doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{{ID: 9}, {ID: 3}, {ID: 8}, {ID: 1}, {ID: 7}}}
	candidates := []Candidate{
		{Name: "One", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 3, EndUnitID: 3}}},
		{Name: "Two", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 8, EndUnitID: 8}}},
		{Name: "Missing", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 99, EndUnitID: 99}}},
		{Name: "Foreign", SourceRefs: []source.SourceRef{{SourceID: "other", StartUnitID: 1, EndUnitID: 1}}},
		{Name: "Blank"},
	}
	materials, ready, err := BuildContext(doc, candidates, 0)
	if err != nil || !ready {
		t.Fatalf("BuildContext() error = %v, ready = %t", err, ready)
	}
	if got, want := materials.EligibleCandidateKeys(), []string{"candidate-000001", "candidate-000002"}; !reflect.DeepEqual(got, want) {
		t.Fatalf("EligibleCandidateKeys() = %#v, want %#v", got, want)
	}
	var transcript transcriptInput
	if err := json.Unmarshal(materials.Transcript.Content, &transcript); err != nil {
		t.Fatal(err)
	}
	if len(transcript.Windows) != 1 || len(transcript.Windows[0].Units) != 2 || transcript.Windows[0].Units[0].ID != 3 || transcript.Windows[0].Units[1].ID != 8 {
		t.Fatalf("windows = %#v, want adjacent document-order units coalesced", transcript.Windows)
	}
	if _, ready, err := BuildContext(doc, candidates, -1); err == nil || ready {
		t.Fatalf("BuildContext(radius=-1) = ready %t, err %v", ready, err)
	}
}

func TestAssessmentRejectsEveryUnsafeProposalCategory(t *testing.T) {
	materials := preparedMaterials(t, 4, true)
	selectors := materialSelectors(t, materials)
	unsafe := []struct {
		name     string
		response ProposalResponse
		category string
	}{
		{"blank member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{{}, selectors[1]}, Canonical: selectors[1]}}}, "member_blank"},
		{"unknown member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{{Name: "unknown", SourceRefs: []SourceRange{{StartUnitID: 99, EndUnitID: 99}}}, selectors[1]}, Canonical: selectors[1]}}}, "member_unknown"},
		{"repeated member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{selectors[0], selectors[0]}, Canonical: selectors[0]}}}, "repeated_member"},
		{"too small", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{selectors[0]}, Canonical: selectors[0]}}}, "fewer_than_two_members"},
		{"canonical blank", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{selectors[0], selectors[1]}, Canonical: Selector{}}}}, "canonical_blank"},
		{"canonical not member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{selectors[0], selectors[1]}, Canonical: selectors[2]}}}, "canonical_not_member"},
		{"overlapping", ProposalResponse{DuplicateGroups: []DuplicateGroup{
			{Members: []Selector{selectors[0], selectors[1]}, Canonical: selectors[0]},
			{Members: []Selector{selectors[1], selectors[2]}, Canonical: selectors[2]},
		}}, "overlapping_member"},
	}
	for _, test := range unsafe {
		t.Run(test.name, func(t *testing.T) {
			assessment := materials.Assess(test.response)
			if len(assessment.SafeGroups()) != 0 || assessment.DiscardedGroups() != len(test.response.DuplicateGroups) || !hasIssue(assessment.Issues(), test.category) {
				t.Fatalf("Assess() = groups %#v discarded %d issues %#v; want %q rejection", assessment.SafeGroups(), assessment.DiscardedGroups(), assessment.Issues(), test.category)
			}
		})
	}
}

func TestAssessmentReturnsNonOverlappingSafeGroupsAndDefensiveCopies(t *testing.T) {
	materials := preparedMaterials(t, 4, false)
	keys := materials.CandidateKeys()
	selectors := materialSelectors(t, materials)
	assessment := materials.Assess(ProposalResponse{DuplicateGroups: []DuplicateGroup{
		{Members: []Selector{selectors[1], selectors[0]}, Canonical: selectors[1]},
		{Members: []Selector{selectors[3], selectors[2]}, Canonical: selectors[2]},
	}})
	groups := assessment.SafeGroups()
	if assessment.DiscardedGroups() != 0 || len(assessment.Issues()) != 0 || len(groups) != 2 {
		t.Fatalf("assessment = %#v, %d, %#v", groups, assessment.DiscardedGroups(), assessment.Issues())
	}
	if got, want := groups[0].Members(), []string{keys[0], keys[1]}; !reflect.DeepEqual(got, want) || groups[0].Canonical() != keys[1] {
		t.Fatalf("first safe group = %#v / %q", got, groups[0].Canonical())
	}
	keys[0] = "changed"
	if materials.CandidateKeys()[0] == "changed" {
		t.Fatal("CandidateKeys() exposed retained keys")
	}
	members := groups[0].Members()
	members[0] = "changed"
	if groups[0].Members()[0] == "changed" || assessment.SafeGroups()[0].Members()[0] == "changed" {
		t.Fatal("SafeGroups() exposed retained members")
	}
}

func TestSharedResponseSchemaIsPrivateStrictAndRegisterableOnce(t *testing.T) {
	schema, err := LoadResponseSchema()
	if err != nil {
		t.Fatal(err)
	}
	if schema.Key != ResponseSchemaKey || schema.ID != ResponseSchemaID || schema.Name != ResponseSchemaName || schema.Version != SchemaVersion || !strings.HasPrefix(schema.SHA256, "sha256:") || !json.Valid(schema.JSONSchema) {
		t.Fatalf("schema = %#v", schema)
	}
	for _, test := range []struct {
		name  string
		value any
		valid bool
	}{
		{"empty groups", map[string]any{"duplicate_groups": []any{}}, true},
		{"semantic proposal problem", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{map[string]any{"name": "Mira", "source_refs": []any{map[string]any{"start_unit_id": 1, "end_unit_id": 1}}}}, "canonical": map[string]any{"name": "Mira", "source_refs": []any{map[string]any{"start_unit_id": 1, "end_unit_id": 1}}}}}}, true},
		{"missing groups", map[string]any{}, false},
		{"unknown top level", map[string]any{"duplicate_groups": []any{}, "extra": true}, false},
		{"replacement name", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{}, "canonical": map[string]any{"name": "Mira", "source_refs": []any{}}, "name": "replacement"}}}, false},
		{"missing selector evidence", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{}, "canonical": map[string]any{"name": "Mira"}}}}, false},
		{"invalid range", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{map[string]any{"name": "Mira", "source_refs": []any{map[string]any{"start_unit_id": 0, "end_unit_id": 1}}}}, "canonical": map[string]any{"name": "Mira", "source_refs": []any{}}}}}, false},
	} {
		t.Run(test.name, func(t *testing.T) {
			content, err := json.Marshal(test.value)
			if err != nil {
				t.Fatal(err)
			}
			err = validateSchema(content, schema.JSONSchema)
			if (err == nil) != test.valid {
				t.Fatalf("validateSchema() error = %v, want valid=%t", err, test.valid)
			}
		})
	}
	first := schema.JSONSchema
	first[0] = '['
	second, err := LoadResponseSchema()
	if err != nil || !json.Valid(second.JSONSchema) || bytes.Equal(first, second.JSONSchema) {
		t.Fatalf("LoadResponseSchema() returned shared content: %s, %v", second.JSONSchema, err)
	}
	registry := llm.NewAssetRegistry()
	if err := RegisterSchemaAssets(registry); err != nil {
		t.Fatalf("RegisterSchemaAssets() error = %v", err)
	}
	if schemaFS, err := registry.SchemaFS(); err != nil {
		t.Fatalf("SchemaFS() error = %v", err)
	} else if content, err := fs.ReadFile(schemaFS, "dnd_entity_reconcile_llm.v1.json"); err != nil || !json.Valid(content) {
		t.Fatalf("shared schema asset = %s, %v", content, err)
	}
}

func preparedMaterials(t *testing.T, count int, includeIneligible bool) Materials {
	t.Helper()
	doc := &source.SourceDocument{ID: "session", Units: make([]source.SourceUnit, count)}
	candidates := make([]Candidate, count)
	for index := range candidates {
		doc.Units[index] = source.SourceUnit{ID: index + 1, Text: "unit"}
		candidates[index] = Candidate{Name: "same display name", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: index + 1, EndUnitID: index + 1}}}
	}
	if includeIneligible && count > 3 {
		candidates[3].SourceRefs = []source.SourceRef{{SourceID: "other", StartUnitID: 1, EndUnitID: 1}}
	}
	materials, ready, err := BuildContext(doc, candidates, 0)
	if err != nil || !ready {
		t.Fatalf("BuildContext() = %#v, %t, %v", materials, ready, err)
	}
	return materials
}

func materialSelectors(t *testing.T, materials Materials) []Selector {
	t.Helper()
	var input candidateInput
	if err := json.Unmarshal(materials.Candidates.Content, &input); err != nil {
		t.Fatal(err)
	}
	return input.Candidates
}

func cloneCandidates(input []Candidate) []Candidate {
	output := make([]Candidate, len(input))
	copy(output, input)
	for index := range output {
		output[index].SourceRefs = append([]source.SourceRef(nil), input[index].SourceRefs...)
	}
	return output
}

func hasIssue(issues []Issue, want string) bool {
	for _, issue := range issues {
		if issue.Category == want {
			return true
		}
	}
	return false
}

func validateSchema(instanceContent, schemaContent []byte) error {
	instance, err := jsonschema.UnmarshalJSON(bytes.NewReader(instanceContent))
	if err != nil {
		return err
	}
	document, err := jsonschema.UnmarshalJSON(bytes.NewReader(schemaContent))
	if err != nil {
		return err
	}
	compiler := jsonschema.NewCompiler()
	if err := compiler.AddResource("schema.json", document); err != nil {
		return err
	}
	compiled, err := compiler.Compile("schema.json")
	if err != nil {
		return err
	}
	return compiled.Validate(instance)
}