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

package itemoccurrences

import (
	"context"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	itemidentity "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/items/identity"
	itemoccurrenceshape "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/validate/itemoccurrences/shape"
)

func TestExtractGroundsOccurrencesInRequiredRegistry(t *testing.T) {
	id := itemidentity.DeriveID("Torch")
	rawResponse := []byte(`{"occurrences":[{"name":"Torch","kind":"lost","from":"party","source_refs":[{"start_unit_id":2,"end_unit_id":1}]}]}`)
	client := &fakeItemOccurrencesLLMClient{content: append([]byte(nil), rawResponse...)}
	req := extractionRequest()
	req.References = itemRegistryReferences(t)
	correction, err := contracts.NewSemanticCorrection([]byte(`{"occurrences":[]}`), "Keep the transcript-grounded item occurrence.")
	if err != nil {
		t.Fatalf("NewSemanticCorrection() error = %v", err)
	}
	req.Correction = correction
	result, err := newExtractor(t, client, req.References).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	if len(result.Value.Occurrences) != 1 || result.Value.Occurrences[0].ItemID != id || result.Value.Occurrences[0].Name != "Torch" {
		t.Fatalf("occurrences = %#v", result.Value.Occurrences)
	}
	if refs := result.Value.Occurrences[0].SourceRefs; len(refs) != 1 || refs[0].SourceID != req.Source.ID || refs[0].StartUnitID != 1 || refs[0].EndUnitID != 2 {
		t.Fatalf("occurrence evidence = %#v, want current-source unit range", refs)
	}
	request := client.requests[0]
	input := request.Inputs[ItemRegistryReferenceSlot]
	if input.Name != ItemRegistryReferenceSlot || string(input.Content) != `{"items":[{"name":"Torch"}]}` || strings.Contains(string(input.Content), "item:sha256:") {
		t.Fatalf("registry prompt input = %#v, want names-only projection", input)
	}
	if request.Correction == nil || string(request.Correction.AssistantResponse) != `{"occurrences":[]}` || request.Correction.UserGuidance != "Keep the transcript-grounded item occurrence." {
		t.Fatalf("correction = %#v, want exact request correction", request.Correction)
	}
	for _, material := range []string{string(request.Correction.AssistantResponse), request.Correction.UserGuidance} {
		if strings.Contains(material, id) || strings.Contains(material, "item:sha256:") {
			t.Fatalf("correction material leaked opaque item identity: %q", material)
		}
	}
	correction.AssistantResponse[0] = '['
	if got := string(request.Correction.AssistantResponse); got != `{"occurrences":[]}` {
		t.Fatalf("captured correction changed after caller mutation: %q", got)
	}
	if result.ModelCandidate == nil || result.ModelCandidate.Protocol != contracts.CorrectionProtocolSingleResponseV1 || string(result.ModelCandidate.Response) != string(rawResponse) {
		t.Fatalf("model candidate = %#v, want exact validated response", result.ModelCandidate)
	}
	client.content[0] = '['
	if got := string(result.ModelCandidate.Response); got != string(rawResponse) {
		t.Fatalf("model candidate changed after provider buffer mutation: %q", got)
	}
}

func TestExtractCanonicalizesComparisonEquivalentNames(t *testing.T) {
	id := itemidentity.DeriveID("Torch")
	client := &fakeItemOccurrencesLLMClient{response: extractionResponse{Occurrences: []itemOccurrenceResponse{
		{Name: "  tORCH ", Kind: "lost", From: "party", SourceRefs: responseRefs(1, 1)},
	}}}
	req := extractionRequest()
	req.References = itemRegistryReferences(t)
	result, err := newExtractor(t, client, req.References).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	occurrence := result.Value.Occurrences[0]
	if occurrence.ItemID != id || occurrence.Name != "Torch" {
		t.Fatalf("canonical occurrence = %#v", occurrence)
	}
}

func TestExtractRejectsUnknownNameAfterValidOccurrence(t *testing.T) {
	client := &fakeItemOccurrencesLLMClient{response: extractionResponse{Occurrences: []itemOccurrenceResponse{
		{Name: "Torch", Kind: "lost", From: "party", SourceRefs: responseRefs(1, 1)},
		{Name: "Unknown", Kind: "lost", From: "party", SourceRefs: responseRefs(2, 2)},
	}}}
	req := extractionRequest()
	req.References = itemRegistryReferences(t)
	result, err := newExtractor(t, client, req.References).Extract(context.Background(), req)
	if err == nil || !strings.Contains(err.Error(), "occurrences[1].name is not in the item registry") {
		t.Fatalf("Extract() error = %v, want occurrence index and unknown name", err)
	}
	if len(result.Value.Occurrences) != 0 {
		t.Fatalf("Extract() returned partial occurrences: %#v", result.Value.Occurrences)
	}
}

func TestExtractRequiresItemRegistry(t *testing.T) {
	_, err := newExtractor(t, &fakeItemOccurrencesLLMClient{}).Extract(context.Background(), extractionRequest())
	if err == nil {
		t.Fatal("Extract() error = nil, want required registry error")
	}
}

func TestExtractAcceptsOnlyEmptyResponseForEmptyRegistry(t *testing.T) {
	references := emptyItemRegistryReferences(t)
	for _, test := range []struct {
		name     string
		response extractionResponse
		wantErr  bool
	}{
		{name: "empty", response: extractionResponse{Occurrences: []itemOccurrenceResponse{}}},
		{name: "selection", response: extractionResponse{Occurrences: []itemOccurrenceResponse{{Name: "Torch", Kind: "lost", From: "party", SourceRefs: responseRefs(1, 1)}}}, wantErr: true},
	} {
		t.Run(test.name, func(t *testing.T) {
			client := &fakeItemOccurrencesLLMClient{response: test.response}
			req := extractionRequest()
			req.References = references
			result, err := newExtractor(t, client, references).Extract(context.Background(), req)
			if test.wantErr {
				if err == nil || !strings.Contains(err.Error(), "name is not in the item registry") || len(result.Value.Occurrences) != 0 {
					t.Fatalf("Extract() = %#v, %v; want no accepted occurrences", result, err)
				}
				return
			}
			if err != nil || result.Value.Occurrences == nil || len(result.Value.Occurrences) != 0 {
				t.Fatalf("Extract() = %#v, %v; want present empty occurrences", result, err)
			}
		})
	}
}

func TestExtractPreservesNullableFields(t *testing.T) {
	client := &fakeItemOccurrencesLLMClient{content: []byte(`{"occurrences":[{"name":"Torch","kind":"discovered","quantity":null,"from":null,"to":null,"source_refs":[{"start_unit_id":1,"end_unit_id":1}]}]}`)}
	req := extractionRequest()
	req.References = itemRegistryReferences(t)
	result, err := newExtractor(t, client, req.References).Extract(context.Background(), req)
	if err != nil || len(result.Value.Occurrences) != 1 || result.Value.Occurrences[0].Quantity != nil || result.Value.Occurrences[0].From != "" || result.Value.Occurrences[0].To != "" {
		t.Fatalf("result = %#v, %v", result, err)
	}
}

func TestExtractUsesOnlySupportedPromptInputs(t *testing.T) {
	client := &fakeItemOccurrencesLLMClient{response: extractionResponse{Occurrences: []itemOccurrenceResponse{}}}
	req := extractionRequest()
	req.References = itemRegistryReferences(t)
	req.References.Slots["unrelated"] = contracts.ResolvedReferenceSlot{Slot: contracts.ReferenceSlot{Name: "unrelated"}, Items: []contracts.ReferenceItem{{SlotName: "unrelated", Content: []byte("ignored")}}}
	if _, err := newExtractor(t, client, req.References).Extract(context.Background(), req); err != nil {
		t.Fatal(err)
	}
	if _, ok := client.requests[0].Inputs["unrelated"]; ok {
		t.Fatalf("unexpected prompt input: %#v", client.requests[0].Inputs)
	}
}

func TestExtractCarriesRejectedHolderGuidanceIntoCorrectedAttempt(t *testing.T) {
	references := itemRegistryReferences(t)
	req := extractionRequest()
	req.References = references
	defectiveResponse := []byte(`{"occurrences":[{"name":"Torch","kind":"acquired","quantity":null,"from":"Chest","to":"party","source_refs":[{"start_unit_id":2,"end_unit_id":2}]}]}`)
	defective, err := newExtractor(t, &fakeItemOccurrencesLLMClient{content: defectiveResponse}, references).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	validation, err := itemoccurrenceshape.New(itemoccurrenceshape.Options{}).Validate(context.Background(), contracts.TypedValidationRequest[dnd.ItemOccurrenceList]{Value: defective.Value})
	if err != nil || validation.Approved {
		t.Fatalf("holder validation = %#v, %v; want rejection", validation, err)
	}
	for _, fragment := range []string{"For `acquired` occurrences", "`from` to JSON null", "item \"Torch\"", "source unit 2"} {
		if !strings.Contains(validation.CorrectionGuidance, fragment) {
			t.Fatalf("CorrectionGuidance = %q, want %q", validation.CorrectionGuidance, fragment)
		}
	}
	for _, forbidden := range []string{defective.Value.Occurrences[0].ItemID, "item_id", "sha256", itemoccurrenceshape.Key, itemoccurrenceshape.ReasonCode} {
		if strings.Contains(validation.CorrectionGuidance, forbidden) {
			t.Fatalf("CorrectionGuidance leaked implementation identifier %q: %q", forbidden, validation.CorrectionGuidance)
		}
	}
	correction, err := contracts.NewSemanticCorrection(defective.ModelCandidate.Response, validation.CorrectionGuidance)
	if err != nil {
		t.Fatal(err)
	}
	correctedResponse := []byte(`{"occurrences":[{"name":"Torch","kind":"acquired","quantity":null,"from":null,"to":"party","source_refs":[{"start_unit_id":2,"end_unit_id":2}]}]}`)
	client := &fakeItemOccurrencesLLMClient{content: correctedResponse}
	req.Correction = correction
	corrected, err := newExtractor(t, client, references).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	accepted, err := itemoccurrenceshape.New(itemoccurrenceshape.Options{}).Validate(context.Background(), contracts.TypedValidationRequest[dnd.ItemOccurrenceList]{Value: corrected.Value})
	if err != nil || !accepted.Approved {
		t.Fatalf("corrected holder validation = %#v, %v; want approval", accepted, err)
	}
	if len(client.requests) != 1 || client.requests[0].Correction == nil || string(client.requests[0].Correction.AssistantResponse) != string(defectiveResponse) || client.requests[0].Correction.UserGuidance != validation.CorrectionGuidance {
		t.Fatalf("corrected request = %#v, want byte-faithful defective response and semantic guidance", client.requests)
	}
}