notarius/internal/framework/llm/asset_registry_test.go

package llm

import (
	"context"
	"io/fs"
	"strings"
	"testing"
	"testing/fstest"
	"time"

	"gitea.maximumdirect.net/eric/promptkit"
)

func TestAssetRegistryCombinesPromptAndSchemaSources(t *testing.T) {
	registry := NewAssetRegistry()
	mustRegisterPromptFS(t, registry, fstest.MapFS{
		"prompts/test.yaml":           {Data: []byte(validPromptYAML("schemas/out.json"))},
		"prompts/messages/user.tmpl":  {Data: []byte(`Input: {{ input "transcript" }}`)},
		"prompts/messages/task.tmpl":  {Data: []byte("Return JSON.")},
		"schemas/ignored/schema.json": {Data: []byte(`{"type":"object"}`)},
	}, "prompts")
	mustRegisterSchemaFS(t, registry, fstest.MapFS{
		"root/schemas/out.json": {Data: []byte(`{"type":"object"}`)},
	}, "root")

	engine := newPromptKitAssetTestEngine(t, registry)
	prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
		PromptID:  "asset.test",
		ProfileID: "asset-test-profile",
		Inputs: map[string]promptkit.ArtifactRef{
			"transcript": promptkit.Inline(`{"ok":true}`),
		},
	})
	if err != nil {
		t.Fatalf("Prepare() error = %v, want nil", err)
	}
	if got := len(prepared.Messages); got != 2 {
		t.Fatalf("message count = %d, want 2", got)
	}
	if prepared.OutputContract.SchemaPath != "schemas/out.json" {
		t.Fatalf("schema path = %q, want schemas/out.json", prepared.OutputContract.SchemaPath)
	}
}

func TestAssetRegistryPrepareFailsForMissingPromptAsset(t *testing.T) {
	registry := NewAssetRegistry()
	mustRegisterPromptFS(t, registry, fstest.MapFS{
		"test.yaml": {Data: []byte(validPromptYAML("out.json"))},
	}, ".")
	mustRegisterSchemaFS(t, registry, fstest.MapFS{
		"out.json": {Data: []byte(`{"type":"object"}`)},
	}, ".")

	engine := newPromptKitAssetTestEngine(t, registry)
	_, err := engine.Prepare(context.Background(), promptkit.RunRequest{
		PromptID:  "asset.test",
		ProfileID: "asset-test-profile",
		Inputs: map[string]promptkit.ArtifactRef{
			"transcript": promptkit.Inline(`{"ok":true}`),
		},
	})
	if err == nil || !strings.Contains(err.Error(), "content_file") {
		t.Fatalf("Prepare() error = %v, want missing content_file error", err)
	}
}

func TestAssetRegistryPrepareFailsForMissingSchemaAsset(t *testing.T) {
	registry := NewAssetRegistry()
	mustRegisterPromptFS(t, registry, fstest.MapFS{
		"test.yaml":          {Data: []byte(validPromptYAML("missing.json"))},
		"messages/user.tmpl": {Data: []byte(`Input: {{ input "transcript" }}`)},
		"messages/task.tmpl": {Data: []byte("Return JSON.")},
	}, ".")
	mustRegisterSchemaFS(t, registry, fstest.MapFS{
		"present.json": {Data: []byte(`{"type":"object"}`)},
	}, ".")

	engine := newPromptKitAssetTestEngine(t, registry)
	_, err := engine.Prepare(context.Background(), promptkit.RunRequest{
		PromptID:  "asset.test",
		ProfileID: "asset-test-profile",
		Inputs: map[string]promptkit.ArtifactRef{
			"transcript": promptkit.Inline(`{"ok":true}`),
		},
	})
	if err == nil || !strings.Contains(err.Error(), "missing.json") {
		t.Fatalf("Prepare() error = %v, want missing schema error", err)
	}
}

func TestAssetRegistryRejectsDuplicateAssetPaths(t *testing.T) {
	registry := NewAssetRegistry()
	mustRegisterPromptFS(t, registry, fstest.MapFS{"one/prompt.yaml": {Data: []byte("id: one")}}, "one")
	mustRegisterPromptFS(t, registry, fstest.MapFS{"two/prompt.yaml": {Data: []byte("id: two")}}, "two")

	_, err := registry.PromptFS()
	if err == nil || !strings.Contains(err.Error(), "duplicate asset path") {
		t.Fatalf("PromptFS() error = %v, want duplicate path error", err)
	}
}

func TestAssetRegistryCombinesFallbackProfileSources(t *testing.T) {
	registry := NewAssetRegistry()
	if err := registry.RegisterFallbackProfileFS(fstest.MapFS{
		"first/profiles/one.yaml": {Data: []byte("id: one\nmodel: first\n")},
	}, "first/profiles"); err != nil {
		t.Fatalf("RegisterFallbackProfileFS() error = %v, want nil", err)
	}
	if err := registry.RegisterFallbackProfileFS(fstest.MapFS{
		"second/two.yaml": {Data: []byte("id: two\nmodel: second\n")},
	}, "second"); err != nil {
		t.Fatalf("RegisterFallbackProfileFS() error = %v, want nil", err)
	}

	fallbackFS, err := registry.FallbackProfileFS()
	if err != nil {
		t.Fatalf("FallbackProfileFS() error = %v, want nil", err)
	}
	for _, name := range []string{"one.yaml", "two.yaml"} {
		if _, err := fs.ReadFile(fallbackFS, name); err != nil {
			t.Fatalf("FallbackProfileFS().ReadFile(%q) error = %v, want nil", name, err)
		}
	}
}

func TestAssetRegistryRejectsInvalidFallbackProfileRoot(t *testing.T) {
	registry := NewAssetRegistry()
	err := registry.RegisterFallbackProfileFS(fstest.MapFS{}, "../profiles")
	if err == nil || !strings.Contains(err.Error(), "invalid path") {
		t.Fatalf("RegisterFallbackProfileFS() error = %v, want invalid root error", err)
	}
}

func TestAssetRegistryRejectsUnreadableFallbackProfileAssets(t *testing.T) {
	registry := NewAssetRegistry()
	if err := registry.RegisterFallbackProfileFS(unreadableAssetFS{}, "."); err != nil {
		t.Fatalf("RegisterFallbackProfileFS() error = %v, want nil", err)
	}
	_, err := registry.FallbackProfileFS()
	if err == nil || !strings.Contains(err.Error(), "permission denied") {
		t.Fatalf("FallbackProfileFS() error = %v, want unreadable asset error", err)
	}
}

func TestAssetRegistryRejectsDuplicateFallbackProfilePaths(t *testing.T) {
	registry := NewAssetRegistry()
	if err := registry.RegisterFallbackProfileFS(fstest.MapFS{"first/profile.yaml": {Data: []byte("id: first\n")}}, "first"); err != nil {
		t.Fatal(err)
	}
	if err := registry.RegisterFallbackProfileFS(fstest.MapFS{"second/profile.yaml": {Data: []byte("id: second\n")}}, "second"); err != nil {
		t.Fatal(err)
	}
	_, err := registry.FallbackProfileFS()
	if err == nil || !strings.Contains(err.Error(), "duplicate asset path") {
		t.Fatalf("FallbackProfileFS() error = %v, want duplicate path error", err)
	}
}

func TestAssetRegistryFallbackProfileDigestTracksContentWithoutLeakingIt(t *testing.T) {
	digestFor := func(content string) string {
		t.Helper()
		registry := NewAssetRegistry()
		if err := registry.RegisterFallbackProfileFS(fstest.MapFS{
			"profiles/profile.yaml": {Data: []byte(content)},
		}, "profiles"); err != nil {
			t.Fatal(err)
		}
		digest, err := registry.FallbackProfileDigest()
		if err != nil {
			t.Fatal(err)
		}
		return digest
	}

	first := digestFor("id: fallback\nmodel: model-one\n")
	second := digestFor("id: fallback\nmodel: model-two\n")
	if first == second {
		t.Fatalf("fallback profile digests = %q and %q, want content change", first, second)
	}
	if !strings.HasPrefix(first, "sha256:") || strings.Contains(first, "model-one") || strings.Contains(first, "profile.yaml") {
		t.Fatalf("fallback profile digest leaked source details: %q", first)
	}
}

func TestAssetRegistryCombinesNamespacedPromptSources(t *testing.T) {
	registry := NewAssetRegistry()
	mustRegisterPromptFS(t, registry, fstest.MapFS{
		"dnd.spells/dnd.spells.yaml": {Data: []byte(validPromptYAML("schema.json"))},
		"dnd.spells/task.md":         {Data: []byte("spell task")},
		"dnd.spells/instructions.md": {Data: []byte("spell instructions")},
	}, ".")
	mustRegisterPromptFS(t, registry, fstest.MapFS{
		"dnd.scenes/dnd.scenes.yaml": {Data: []byte(validPromptYAML("schema.json"))},
		"dnd.scenes/task.md":         {Data: []byte("scene task")},
		"dnd.scenes/instructions.md": {Data: []byte("scene instructions")},
	}, ".")

	fsys, err := registry.PromptFS()
	if err != nil {
		t.Fatalf("PromptFS() error = %v, want nil", err)
	}
	for _, name := range []string{
		"dnd.spells/dnd.spells.yaml",
		"dnd.spells/task.md",
		"dnd.spells/instructions.md",
		"dnd.scenes/dnd.scenes.yaml",
		"dnd.scenes/task.md",
		"dnd.scenes/instructions.md",
	} {
		if _, err := fsys.Open(name); err != nil {
			t.Fatalf("PromptFS().Open(%q) error = %v, want nil", name, err)
		}
	}
}

func TestHashAssetsOmitsRawAssetContent(t *testing.T) {
	hash, err := HashAssets([]AssetHashPart{{
		FS:   fstest.MapFS{"prompt.md": {Data: []byte("secret prompt text")}},
		Path: "prompt.md",
	}})
	if err != nil {
		t.Fatalf("HashAssets() error = %v, want nil", err)
	}
	if !strings.HasPrefix(hash, "sha256:") {
		t.Fatalf("hash = %q, want sha256-prefixed value", hash)
	}
	if strings.Contains(hash, "secret prompt text") {
		t.Fatalf("hash leaked asset content")
	}
}

func newPromptKitAssetTestEngine(t *testing.T, registry *AssetRegistry) *promptkit.Engine {
	t.Helper()
	options, err := registry.PromptKitOptions()
	if err != nil {
		t.Fatalf("PromptKitOptions() error = %v, want nil", err)
	}
	options = append(options, promptkit.WithProfiles(promptkit.OpenAICompatibleProfile(promptkit.OpenAICompatibleProfileConfig{
		ID:       "asset-test-profile",
		Endpoint: "http://127.0.0.1:1/v1",
		Model:    "asset-test-model",
	})))
	engine, err := promptkit.NewEngine(promptkit.Config{Timeout: time.Second}, options...)
	if err != nil {
		t.Fatalf("NewEngine() error = %v, want nil", err)
	}
	return engine
}

func mustRegisterPromptFS(t *testing.T, registry *AssetRegistry, fsys fstest.MapFS, root string) {
	t.Helper()
	if err := registry.RegisterPromptFS(fsys, root); err != nil {
		t.Fatalf("RegisterPromptFS() error = %v, want nil", err)
	}
}

func mustRegisterSchemaFS(t *testing.T, registry *AssetRegistry, fsys fstest.MapFS, root string) {
	t.Helper()
	if err := registry.RegisterSchemaFS(fsys, root); err != nil {
		t.Fatalf("RegisterSchemaFS() error = %v, want nil", err)
	}
}

func validPromptYAML(schemaPath string) string {
	return `id: asset.test
version: "v1"
inputs:
  - name: transcript
    required: true
    content_type: application/json
messages:
  - role: user
    content_file: ./messages/user.tmpl
  - role: user
    content_file: ./messages/task.tmpl
output:
  format: json
  validation_mode: json_schema
  schema_path: ` + schemaPath + `
  repair_attempts: 0
`
}

type unreadableAssetFS struct{}

func (unreadableAssetFS) Open(name string) (fs.File, error) {
	return nil, &fs.PathError{Op: "open", Path: name, Err: fs.ErrPermission}
}