notarius/internal/cli/spell_catalog_retry_contract_test.go

package cli

import (
	"context"
	"encoding/json"
	"fmt"
	"path/filepath"
	"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"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/spells"
)

func TestProductionSpellCatalogValidationRetries(t *testing.T) {
	const retries = 2
	tests := []struct {
		name            string
		responses       []string
		wantCalls       int
		wantRejected    bool
		wantSpell       string
		wantWarningCode string
	}{
		{
			name: "unknown spell remains rejected after exhaustion",
			responses: []string{
				productionSpellResponse("Unknown Spell"),
				productionSpellResponse("Unknown Spell"),
				productionSpellResponse("Unknown Spell"),
			},
			wantCalls:    retries + 1,
			wantRejected: true,
		},
		{
			name: "overlay spell becomes valid on retry",
			responses: []string{
				productionSpellResponse("Unknown Spell"),
				productionSpellResponse("Aegis of Emberfall"),
			},
			wantCalls:       2,
			wantSpell:       "Aegis of Emberfall",
			wantWarningCode: "spell_not_near_source",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			components := productionTestComponents(t)
			configPath := writeProductionSpellCatalogContractConfig(t)
			cfg := loadMaintainedExample(t, configPath)
			effective, err := cfg.Resolve(config.ResolveInput{PipelineID: "dnd-session", Catalog: catalogFromRegistries(components.registries)})
			if err != nil {
				t.Fatalf("resolve production configuration: %v", err)
			}
			materialized, _, err := pipeline.MaterializeReferences(effective.ResolvedPipeline, catalogFromRegistries(components.registries), pipeline.ReferenceMaterializationOptions{
				ConfigPath: configPath,
				WorkingDir: filepath.Dir(configPath),
			})
			if err != nil {
				t.Fatalf("materialize production references: %v", err)
			}
			materialized.Steps[0].ArtifactLanes[0].Extract.Retries = retries

			llmClient := &catalogRetryLLMClient{responses: tt.responses}
			prepared, err := pipeline.Prepare(materialized, components.registries, pipeline.ModuleDependencies{LLM: llmClient})
			if err != nil {
				t.Fatalf("prepare production pipeline: %v", err)
			}
			output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{
				Prepared:       prepared,
				RawInput:       readRepositoryFile(t, "examples", "seriatim-minimal-transcript.json"),
				ChunkCacheMode: pipeline.ChunkCacheBypass,
			})
			if err != nil {
				t.Fatalf("Run() error = %v, want nil", err)
			}
			if calls := llmClient.CallCount(); calls > retries+1 || calls != tt.wantCalls {
				t.Fatalf("LLM calls = %d, want %d and no more than %d", calls, tt.wantCalls, retries+1)
			}

			if tt.wantRejected {
				if len(output.Rejected) != 1 || len(output.NormalizeOutputs) != 0 {
					t.Fatalf("rejected = %#v normalized = %#v, want one nonfatal rejection and no merge output", output.Rejected, output.NormalizeOutputs)
				}
				rejection := output.Rejected[0]
				if rejection.ReasonCode != "unknown_spell" || rejection.AttemptCount != retries+1 {
					t.Fatalf("rejection = %#v, want exhausted unknown-spell rejection", rejection)
				}
				if len(output.Warnings) != 0 {
					t.Fatalf("warnings = %#v, want no warnings from rejected attempts", output.Warnings)
				}
				return
			}

			if len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 1 {
				t.Fatalf("rejected = %#v normalized = %#v, want only accepted output", output.Rejected, output.NormalizeOutputs)
			}
			var value dnd.SpellList
			if err := json.Unmarshal(output.NormalizeOutputs[0].Artifact.Content, &value); err != nil {
				t.Fatalf("decode normalized spell list: %v", err)
			}
			if len(value.SpellCasts) != 1 || value.SpellCasts[0].Spell != tt.wantSpell {
				t.Fatalf("normalized spell list = %#v, want accepted overlay spell", value)
			}
			if len(output.Warnings) != 2 || output.Warnings[0].ReasonCode != tt.wantWarningCode || output.Warnings[1].ReasonCode != tt.wantWarningCode {
				t.Fatalf("warnings = %#v, want accepted-attempt warnings from extract and normalize validation", output.Warnings)
			}
		})
	}
}

type catalogRetryLLMClient struct {
	mu        sync.Mutex
	responses []string
	calls     int
}

func (client *catalogRetryLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	if err := ctx.Err(); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	if req.PromptID != spells.PromptID {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("unexpected prompt %q", req.PromptID)
	}
	client.mu.Lock()
	index := client.calls
	client.calls++
	client.mu.Unlock()
	if index >= len(client.responses) {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("missing fake response %d", index)
	}
	content := []byte(client.responses[index])
	if err := json.Unmarshal(content, out); err != nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("populate fake structured target: %w", err)
	}
	return contracts.StructuredCompletionResponse{Content: content, Provider: "test", Model: "deterministic", ProfileID: req.ProfileID}, nil
}

func (client *catalogRetryLLMClient) CallCount() int {
	client.mu.Lock()
	defer client.mu.Unlock()
	return client.calls
}

func productionSpellResponse(name string) string {
	content, err := json.Marshal(dnd.SpellList{SpellCasts: []dnd.SpellCast{{
		Caster:     "Aria",
		Spell:      name,
		SourceRefs: []source.SourceRef{{SourceID: "session-alpha", StartUnitID: 1, EndUnitID: 1}},
	}}})
	if err != nil {
		panic(err)
	}
	return string(content)
}