notarius/internal/modules/integration/dnd_npc_grounded_test.go

package integration_test

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"sync"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/config"
	"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"
	combatcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/combatturns"
	npccodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/npcs"
	combatextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/combatturns"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/npcs"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/spells"
)

func TestNPCOutputGroundsSpellAndCombatConsumersThroughOneOperation(t *testing.T) {
	registries := productionNPCRegistries(t)
	catalog := moduleCatalog(registries)
	configValue := loadGroundedPipelineConfig(t)
	effective, err := configValue.Resolve(config.ResolveInput{PipelineID: "dnd-npc-grounded", Catalog: catalog})
	if err != nil {
		t.Fatalf("Resolve() error = %v", err)
	}
	materialized, warnings, err := pipeline.MaterializeReferences(effective.ResolvedPipeline, catalog, pipeline.ReferenceMaterializationOptions{})
	if err != nil || len(warnings) != 0 {
		t.Fatalf("MaterializeReferences() error = %v warnings = %#v", err, warnings)
	}

	client := &groundedDNDLLMClient{}
	output, err := runPreparedPipeline(t, registries, materialized, client, pipeline.RunInput{
		RawInput:       readNPCFixture(t),
		ExtractWorkers: 1,
	})
	if err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	if len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 3 {
		t.Fatalf("run outputs = %#v rejected = %#v, want NPC, spell, and combat outputs", output.NormalizeOutputs, output.Rejected)
	}

	var npcPayload []byte
	for _, serialized := range output.NormalizeOutputs {
		if serialized.LaneID == "npcs" {
			npcPayload = append([]byte(nil), serialized.Artifact.Content...)
			break
		}
	}
	if len(npcPayload) == 0 {
		t.Fatal("NPC producer did not publish a canonical payload")
	}
	npcValue, err := npccodec.New().Decode(npcPayload)
	if err != nil {
		t.Fatalf("decode NPC producer payload: %v", err)
	}
	npcPayload, err = npccodec.New().Encode(npcValue)
	if err != nil {
		t.Fatalf("encode canonical NPC producer payload: %v", err)
	}
	canonicalDigest := digestBytes(npcPayload)
	projection := []byte(`{"npcs":[{"name":"Mira Thorn"},{"name":"Hooded Guard"}]}`)
	projectionDigest := digestBytes(projection)
	seenConsumers := map[string]bool{}
	for _, request := range client.requestsSnapshot() {
		if request.PromptID != spells.PromptID && request.PromptID != combatextract.PromptID {
			continue
		}
		input := request.Inputs["npcs"]
		if input.MediaType != npccodec.MediaType || input.Digest != projectionDigest || string(input.Content) != string(projection) || input.OriginURI != "" {
			t.Fatalf("%s NPC prompt input = %#v, want names-only generated registry", request.PromptID, input)
		}
		seenConsumers[request.PromptID] = true
	}
	if !seenConsumers[spells.PromptID] || !seenConsumers[combatextract.PromptID] {
		t.Fatalf("consumer prompt IDs = %#v, want spell and combat requests", seenConsumers)
	}

	provenanceCount := 0
	for _, reference := range output.Manifest.References {
		if reference.SlotName != "npcs" {
			continue
		}
		provenanceCount++
		if reference.Digest != canonicalDigest || reference.OriginType != "generated" {
			t.Fatalf("NPC generated provenance = %#v, want canonical digest and generated origin", reference)
		}
	}
	if provenanceCount != 3 {
		t.Fatalf("NPC generated provenance count = %d, want spell extract plus combat extract/normalize", provenanceCount)
	}
	for _, lane := range output.Manifest.ArtifactLanes {
		for _, component := range []string{"extractor", "normalizer"} {
			metadata, ok := lane.Metadata[component].(map[string]any)
			if ok && metadata["npc_registry_digest"] != nil {
				t.Fatalf("%s %s metadata = %#v, want generated identity only in framework provenance", lane.ID, component, metadata)
			}
		}
	}

	var combatValue dnd.CombatTurnList
	for _, serialized := range output.NormalizeOutputs {
		switch serialized.LaneID {
		case "spells":
			spellValue := decodeRunnerSpellResponse(t, serialized.Artifact.Content)
			if len(spellValue.SpellCasts) != 1 || spellValue.SpellCasts[0].Caster != "Mira Thorn" {
				t.Fatalf("spell output = %#v, want one registry-grounded-context spell", spellValue)
			}
			assertCurrentEvidence(t, spellValue.SpellCasts[0].SourceRefs)
		case "combat":
			decoded, decodeErr := combatcodec.New().Decode(serialized.Artifact.Content)
			if decodeErr != nil {
				t.Fatalf("decode combat output: %v", decodeErr)
			}
			combatValue = decoded
		}
	}
	if len(combatValue.CombatTurns) != 1 || combatValue.CombatTurns[0].Actor != "Mira Thorn" {
		t.Fatalf("combat output = %#v, want registry-normalized actor", combatValue)
	}
	assertCurrentEvidence(t, combatValue.CombatTurns[0].SourceRefs)
}

func assertCurrentEvidence(t *testing.T, references []source.SourceRef) {
	t.Helper()
	for _, reference := range references {
		if reference.SourceID != "npc-session" {
			t.Fatalf("evidence reference = %#v, want current source only", reference)
		}
	}
}

func loadGroundedPipelineConfig(t *testing.T) config.Config {
	t.Helper()
	fileConfig, err := config.LoadFileConfig(repositoryPathForIntegration("internal", "modules", "integration", "testdata", "dnd_npc_grounded_pipeline.yml"))
	if err != nil {
		t.Fatalf("LoadFileConfig() error = %v", err)
	}
	cfg := config.Default()
	if err := cfg.ApplyFileConfig(fileConfig); err != nil {
		t.Fatalf("ApplyFileConfig() error = %v", err)
	}
	return cfg
}

type groundedDNDLLMClient struct {
	mu       sync.Mutex
	requests []contracts.StructuredCompletionRequest
}

func (client *groundedDNDLLMClient) CompleteStructured(ctx context.Context, request contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	if err := ctx.Err(); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	client.mu.Lock()
	client.requests = append(client.requests, cloneStructuredCompletionRequest(request))
	client.mu.Unlock()

	var payload any
	switch request.PromptID {
	case npcs.PromptID:
		payload = map[string]any{"npcs": []any{
			map[string]any{
				"name": "Mira Thorn", "source_refs": []any{map[string]int{"start_unit_id": 1, "end_unit_id": 1}},
			},
			map[string]any{
				"name": "Hooded Guard", "source_refs": []any{map[string]int{"start_unit_id": 3, "end_unit_id": 3}},
			},
		}}
	case spells.PromptID:
		payload = map[string]any{"spell_casts": []any{map[string]any{
			"caster": "Mira Thorn", "spell": "Cure Wounds",
			"source_refs": []any{map[string]int{"start_unit_id": 1, "end_unit_id": 1}},
		}}}
	case combatextract.PromptID:
		payload = map[string]any{"combat_turns": []any{map[string]any{
			"actor":       "Mira Thorn",
			"turn_kind":   "turn",
			"source_refs": []any{map[string]int{"start_unit_id": 1, "end_unit_id": 1}},
		}}}
	default:
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("unexpected grounded prompt %q", request.PromptID)
	}
	content, err := json.Marshal(payload)
	if err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	if err := json.Unmarshal(content, out); err != nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("populate grounded response: %w", err)
	}
	return contracts.StructuredCompletionResponse{Content: content, Provider: "test", Model: "grounded-fake"}, nil
}

func (client *groundedDNDLLMClient) requestsSnapshot() []contracts.StructuredCompletionRequest {
	client.mu.Lock()
	defer client.mu.Unlock()
	return append([]contracts.StructuredCompletionRequest(nil), client.requests...)
}

func digestBytes(content []byte) string {
	sum := sha256.Sum256(content)
	return fmt.Sprintf("sha256:%x", sum[:])
}

var _ contracts.StructuredLLMClient = (*groundedDNDLLMClient)(nil)