notarius/internal/framework/llm/asset_registry_test.go

package llm

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

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

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 := newAssetTestEngine(t, registry)
	prepared, err := engine.Prepare(context.Background(), scriptorium.RunRequest{
		PromptID:  "asset.test",
		ProfileID: "asset-test-profile",
		Inputs: map[string]scriptorium.ArtifactRef{
			"transcript": scriptorium.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 := newAssetTestEngine(t, registry)
	_, err := engine.Prepare(context.Background(), scriptorium.RunRequest{
		PromptID:  "asset.test",
		ProfileID: "asset-test-profile",
		Inputs: map[string]scriptorium.ArtifactRef{
			"transcript": scriptorium.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 := newAssetTestEngine(t, registry)
	_, err := engine.Prepare(context.Background(), scriptorium.RunRequest{
		PromptID:  "asset.test",
		ProfileID: "asset-test-profile",
		Inputs: map[string]scriptorium.ArtifactRef{
			"transcript": scriptorium.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 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 newAssetTestEngine(t *testing.T, registry *AssetRegistry) *scriptorium.Engine {
	t.Helper()
	options, err := registry.ScriptoriumOptions()
	if err != nil {
		t.Fatalf("ScriptoriumOptions() error = %v, want nil", err)
	}
	options = append(options, scriptorium.WithProfiles(scriptorium.OpenAICompatibleProfile(scriptorium.OpenAICompatibleProfileConfig{
		ID:       "asset-test-profile",
		Endpoint: "http://127.0.0.1:1/v1",
		Model:    "asset-test-model",
	})))
	engine, err := scriptorium.NewEngine(scriptorium.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
`
}