Organize D&D extensions by domain

internal/modules/dnd/shared/assets.go

@@ -0,0 +1,47 @@
+package shared
+
+import (
+	"embed"
+	"io/fs"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/promptfs"
+)
+
+//go:embed assets/prompts/*.md
+var embeddedAssets embed.FS
+
+var sharedPromptFiles = []string{
+	"common-dnd-system.md",
+	"common-dnd-transcript.md",
+	"common-dnd-references.md",
+}
+
+func SharedPromptFiles() []promptfs.SharedPromptFile {
+	files := make([]promptfs.SharedPromptFile, 0, len(sharedPromptFiles))
+	for _, name := range sharedPromptFiles {
+		files = append(files, promptfs.SharedPromptFile{
+			Name: name,
+			FS:   embeddedAssets,
+			Path: "assets/prompts/" + name,
+		})
+	}
+	return files
+}
+
+func CommonHashParts() []llm.AssetHashPart {
+	return []llm.AssetHashPart{
+		{FS: embeddedAssets, Path: "assets/prompts/common-dnd-system.md"},
+		{FS: embeddedAssets, Path: "assets/prompts/common-dnd-transcript.md"},
+	}
+}
+
+func ReferenceHashParts() []llm.AssetHashPart {
+	return []llm.AssetHashPart{
+		{FS: embeddedAssets, Path: "assets/prompts/common-dnd-references.md"},
+	}
+}
+
+func ModulePromptFS(moduleDir string, moduleFS fs.FS, files []promptfs.ModulePromptFile) (fs.FS, error) {
+	return promptfs.ModulePromptFS(moduleDir, moduleFS, files, SharedPromptFiles()...)
+}

internal/modules/dnd/shared/assets/prompts/common-dnd-references.md

@@ -0,0 +1,12 @@
+Optional reference material for this Dungeons & Dragons campaign is provided
+below. Use it only to disambiguate names, aliases, speakers, campaign terms, or
+spell names already present in the transcript.
+
+Player list reference:
+{{ input "players" }}
+
+Party roster reference:
+{{ input "party" }}
+
+Glossary reference:
+{{ input "glossary" }}

internal/modules/dnd/shared/assets/prompts/common-dnd-system.md

@@ -0,0 +1,8 @@
+You work with Dungeons & Dragons gameplay transcripts.
+
+Use only the provided transcript and reference material. Source text may contain
+transcription errors, repeated lines, incomplete sentences, and misheard proper
+nouns. Reference material, when present, is supporting context only and must not
+be treated as a source of extracted events by itself.
+
+Return only valid JSON matching the configured response schema.

internal/modules/dnd/shared/assets/prompts/common-dnd-transcript.md

@@ -0,0 +1,3 @@
+A transcript of a Dungeons & Dragons gameplay session is provided below.
+
+{{ input "transcript" }}

internal/modules/dnd/shared/assets_test.go

@@ -0,0 +1,87 @@
+package shared
+
+import (
+	"io/fs"
+	"testing"
+	"testing/fstest"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/promptfs"
+)
+
+func TestSharedPromptFilesReturnsNewSlice(t *testing.T) {
+	first := SharedPromptFiles()
+	second := SharedPromptFiles()
+
+	if len(first) != 3 || len(second) != 3 {
+		t.Fatalf("SharedPromptFiles() lengths = %d and %d, want 3", len(first), len(second))
+	}
+	first[0].Name = "changed.md"
+	if second[0].Name != "common-dnd-system.md" {
+		t.Fatalf("SharedPromptFiles() reused descriptor slice: %#v", second)
+	}
+}
+
+func TestSharedPromptFilesReferenceEmbeddedAssets(t *testing.T) {
+	for _, file := range SharedPromptFiles() {
+		if file.FS == nil {
+			t.Fatalf("SharedPromptFiles() descriptor %q has nil FS", file.Name)
+		}
+		if _, err := fs.ReadFile(file.FS, file.Path); err != nil {
+			t.Fatalf("ReadFile(%q) error = %v, want nil", file.Path, err)
+		}
+	}
+}
+
+func TestHashPartsReferenceSharedPrompts(t *testing.T) {
+	assertHashParts(t, "common", CommonHashParts(), []string{
+		"assets/prompts/common-dnd-system.md",
+		"assets/prompts/common-dnd-transcript.md",
+	})
+	assertHashParts(t, "reference", ReferenceHashParts(), []string{
+		"assets/prompts/common-dnd-references.md",
+	})
+
+	for _, part := range append(CommonHashParts(), ReferenceHashParts()...) {
+		if _, err := fs.ReadFile(part.FS, part.Path); err != nil {
+			t.Fatalf("ReadFile(%q) error = %v, want nil", part.Path, err)
+		}
+	}
+}
+
+func assertHashParts(t *testing.T, name string, parts []llm.AssetHashPart, want []string) {
+	t.Helper()
+	if len(parts) != len(want) {
+		t.Fatalf("%s hash parts length = %d, want %d", name, len(parts), len(want))
+	}
+	for i, part := range parts {
+		if part.Path != want[i] {
+			t.Fatalf("%s hash part %d path = %q, want %q", name, i, part.Path, want[i])
+		}
+		if part.FS == nil {
+			t.Fatalf("%s hash part %d has nil FS", name, i)
+		}
+	}
+}
+
+func TestModulePromptFSMountsDNDSharedPrompts(t *testing.T) {
+	fsys, err := ModulePromptFS("dnd.test", fstest.MapFS{
+		"assets/prompts/dnd.test.yaml": {Data: []byte("id: dnd.test")},
+	}, []promptfs.ModulePromptFile{
+		{Name: "dnd.test.yaml", Path: "assets/prompts/dnd.test.yaml"},
+	})
+	if err != nil {
+		t.Fatalf("ModulePromptFS() error = %v, want nil", err)
+	}
+
+	for _, path := range []string{
+		"assets/prompts/dnd.test/dnd.test.yaml",
+		"assets/prompts/dnd.test/sharedassets/common-dnd-system.md",
+		"assets/prompts/dnd.test/sharedassets/common-dnd-transcript.md",
+		"assets/prompts/dnd.test/sharedassets/common-dnd-references.md",
+	} {
+		if _, err := fs.ReadFile(fsys, path); err != nil {
+			t.Fatalf("ReadFile(%q) error = %v, want nil", path, err)
+		}
+	}
+}

internal/modules/dnd/shared/prompt_inputs.go

@@ -0,0 +1,84 @@
+package shared
+
+import (
+	"bytes"
+	"fmt"
+	"sort"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+)
+
+func PromptInputs(sourceInput contracts.LLMInputMaterial, references contracts.ReferenceSet) contracts.LLMInputSet {
+	partySlot := references.Slots["party"]
+	if len(partySlot.Items) == 0 {
+		partySlot = references.Slots["roster"]
+	}
+	return contracts.LLMInputSet{
+		"transcript": TranscriptPromptMaterial(sourceInput),
+		"players":    ReferencePromptMaterial("players", references.Slots["players"]),
+		"party":      ReferencePromptMaterial("party", partySlot),
+		"glossary":   ReferencePromptMaterial("glossary", references.Slots["glossary"]),
+	}
+}
+
+func TranscriptPromptMaterial(material contracts.LLMInputMaterial) contracts.LLMInputMaterial {
+	out := material.Clone()
+	out.Name = "transcript"
+	return out
+}
+
+func ReferencePromptMaterial(name string, slot contracts.ResolvedReferenceSlot) contracts.LLMInputMaterial {
+	body := ReferencePromptInput(slot)
+	digest := ""
+	originURI := ""
+	if len(slot.Items) == 1 {
+		digest = slot.Items[0].Digest
+		originURI = slot.Items[0].Origin.URI
+	}
+	return contracts.NewLLMInputMaterial(name, "text/plain", body, digest, originURI)
+}
+
+func ReferencePromptInput(slot contracts.ResolvedReferenceSlot) []byte {
+	if len(slot.Items) == 0 {
+		return []byte(" ")
+	}
+	items := append([]contracts.ReferenceItem(nil), slot.Items...)
+	sort.SliceStable(items, func(i, j int) bool {
+		if items[i].Origin.URI != items[j].Origin.URI {
+			return items[i].Origin.URI < items[j].Origin.URI
+		}
+		if items[i].Digest != items[j].Digest {
+			return items[i].Digest < items[j].Digest
+		}
+		return string(items[i].Content) < string(items[j].Content)
+	})
+	if len(items) == 1 {
+		return append([]byte(nil), items[0].Content...)
+	}
+
+	var b bytes.Buffer
+	for i, item := range items {
+		if i > 0 {
+			b.WriteString("\n\n")
+		}
+		fmt.Fprintf(&b, "Reference %d\n", i+1)
+		if item.Origin.Type != "" {
+			fmt.Fprintf(&b, "Origin-Type: %s\n", item.Origin.Type)
+		}
+		if item.Origin.URI != "" {
+			fmt.Fprintf(&b, "Origin-URI: %s\n", item.Origin.URI)
+		}
+		if item.Digest != "" {
+			fmt.Fprintf(&b, "Digest: %s\n", item.Digest)
+		}
+		if item.MediaType != "" {
+			fmt.Fprintf(&b, "Media-Type: %s\n", item.MediaType)
+		}
+		if item.SizeBytes > 0 {
+			fmt.Fprintf(&b, "Size-Bytes: %d\n", item.SizeBytes)
+		}
+		b.WriteString("\n")
+		b.Write(item.Content)
+	}
+	return b.Bytes()
+}

internal/modules/dnd/shared/prompt_inputs_test.go

@@ -0,0 +1,151 @@
+package shared
+
+import (
+	"strings"
+	"testing"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+)
+
+func TestPromptInputsBuildExpectedInputs(t *testing.T) {
+	source := contracts.NewLLMInputMaterial("source", "application/json", []byte("source text"), "sha256:source", "file:///source.json")
+	references := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
+		"players":  slotWithContent("players", "Alice: Aria"),
+		"party":    slotWithContent("party", "Aria: cleric"),
+		"glossary": slotWithContent("glossary", "Brightmantle: temple"),
+	}}
+
+	inputs := PromptInputs(source, references)
+	for _, name := range []string{"transcript", "players", "party", "glossary"} {
+		if _, ok := inputs[name]; !ok {
+			t.Fatalf("PromptInputs() missing %q: %#v", name, inputs)
+		}
+	}
+	if _, ok := inputs["roster"]; ok {
+		t.Fatalf("PromptInputs() included roster input: %#v", inputs)
+	}
+	if got := inputs["transcript"].Name; got != "transcript" {
+		t.Fatalf("transcript name = %q, want transcript", got)
+	}
+	if got := string(inputs["transcript"].Content); got != "source text" {
+		t.Fatalf("transcript content = %q, want source text", got)
+	}
+	if got := string(inputs["players"].Content); got != "Alice: Aria" {
+		t.Fatalf("players content = %q, want player reference", got)
+	}
+	if got := string(inputs["party"].Content); got != "Aria: cleric" {
+		t.Fatalf("party content = %q, want party reference", got)
+	}
+	if got := string(inputs["glossary"].Content); got != "Brightmantle: temple" {
+		t.Fatalf("glossary content = %q, want glossary reference", got)
+	}
+}
+
+func TestPromptInputsUseRosterWhenPartyIsEmpty(t *testing.T) {
+	inputs := PromptInputs(contracts.LLMInputMaterial{}, contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
+		"party":  {},
+		"roster": slotWithContent("roster", "Legacy roster text"),
+	}})
+
+	if got := string(inputs["party"].Content); got != "Legacy roster text" {
+		t.Fatalf("party content = %q, want roster fallback content", got)
+	}
+	if _, ok := inputs["roster"]; ok {
+		t.Fatalf("PromptInputs() included roster input: %#v", inputs)
+	}
+}
+
+func TestTranscriptPromptMaterialClonesSource(t *testing.T) {
+	source := contracts.NewLLMInputMaterial("source", "text/plain", []byte("source text"), "sha256:source", "file:///source.txt")
+	got := TranscriptPromptMaterial(source)
+
+	if got.Name != "transcript" {
+		t.Fatalf("Name = %q, want transcript", got.Name)
+	}
+	if got.MediaType != source.MediaType || got.Digest != source.Digest || got.OriginURI != source.OriginURI || got.SizeBytes != source.SizeBytes {
+		t.Fatalf("TranscriptPromptMaterial() = %#v, want cloned metadata from %#v", got, source)
+	}
+	source.Content[0] = 'X'
+	if string(got.Content) != "source text" {
+		t.Fatalf("TranscriptPromptMaterial() reused content slice: %q", got.Content)
+	}
+}
+
+func TestReferencePromptMaterialUsesTextPlainAndSingleReferenceMetadata(t *testing.T) {
+	slot := contracts.ResolvedReferenceSlot{Items: []contracts.ReferenceItem{{
+		Content: []byte("single reference"),
+		Digest:  "sha256:reference",
+		Origin:  contracts.ReferenceOrigin{Type: "file", URI: "file:///reference.md"},
+	}}}
+
+	got := ReferencePromptMaterial("party", slot)
+	if got.Name != "party" || got.MediaType != "text/plain" {
+		t.Fatalf("ReferencePromptMaterial() name/media = %q/%q, want party/text/plain", got.Name, got.MediaType)
+	}
+	if got.Digest != "sha256:reference" || got.OriginURI != "file:///reference.md" {
+		t.Fatalf("ReferencePromptMaterial() digest/origin = %q/%q, want single reference metadata", got.Digest, got.OriginURI)
+	}
+	if string(got.Content) != "single reference" {
+		t.Fatalf("ReferencePromptMaterial() content = %q, want raw single reference", got.Content)
+	}
+}
+
+func TestReferencePromptMaterialOmitsAggregateMetadata(t *testing.T) {
+	got := ReferencePromptMaterial("party", contracts.ResolvedReferenceSlot{Items: []contracts.ReferenceItem{
+		{Content: []byte("one"), Digest: "sha256:one", Origin: contracts.ReferenceOrigin{URI: "file:///one.md"}},
+		{Content: []byte("two"), Digest: "sha256:two", Origin: contracts.ReferenceOrigin{URI: "file:///two.md"}},
+	}})
+
+	if got.Digest != "" || got.OriginURI != "" {
+		t.Fatalf("ReferencePromptMaterial() digest/origin = %q/%q, want empty aggregate metadata", got.Digest, got.OriginURI)
+	}
+}
+
+func TestReferencePromptInputRendering(t *testing.T) {
+	if got := string(ReferencePromptInput(contracts.ResolvedReferenceSlot{})); got != " " {
+		t.Fatalf("empty rendering = %q, want single space", got)
+	}
+	if got := string(ReferencePromptInput(slotWithContent("party", "single reference"))); got != "single reference" {
+		t.Fatalf("single rendering = %q, want raw content", got)
+	}
+
+	slot := contracts.ResolvedReferenceSlot{Items: []contracts.ReferenceItem{
+		{
+			SlotName:  "party",
+			MediaType: "text/plain",
+			Content:   []byte("second"),
+			Digest:    "sha256:bbb",
+			Origin:    contracts.ReferenceOrigin{Type: "file", URI: "file:///b.txt"},
+			SizeBytes: 6,
+		},
+		{
+			SlotName:  "party",
+			MediaType: "text/plain",
+			Content:   []byte("first"),
+			Digest:    "sha256:aaa",
+			Origin:    contracts.ReferenceOrigin{Type: "file", URI: "file:///a.txt"},
+			SizeBytes: 5,
+		},
+	}}
+	first := string(ReferencePromptInput(slot))
+	second := string(ReferencePromptInput(slot))
+	if first != second {
+		t.Fatalf("ReferencePromptInput() was not deterministic:\nfirst=%q\nsecond=%q", first, second)
+	}
+	if !strings.Contains(first, "Reference 1\nOrigin-Type: file\nOrigin-URI: file:///a.txt\nDigest: sha256:aaa\nMedia-Type: text/plain\nSize-Bytes: 5\n\nfirst") {
+		t.Fatalf("first reference block = %q, want sorted first reference metadata", first)
+	}
+	if strings.Index(first, "first") > strings.Index(first, "second") {
+		t.Fatalf("references were not sorted deterministically: %q", first)
+	}
+}
+
+func slotWithContent(name string, content string) contracts.ResolvedReferenceSlot {
+	return contracts.ResolvedReferenceSlot{
+		Slot: contracts.ReferenceSlot{Name: name},
+		Items: []contracts.ReferenceItem{{
+			SlotName: name,
+			Content:  []byte(content),
+		}},
+	}
+}

internal/modules/dnd/shared/references.go

@@ -0,0 +1,47 @@
+package shared
+
+import "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+
+type ReferenceSlotDescriptions struct {
+	Glossary string
+	Party    string
+	Players  string
+	Roster   string
+}
+
+var referenceMediaTypes = []string{
+	"application/json",
+	"application/x-yaml",
+	"application/yaml",
+	"text/markdown",
+	"text/plain",
+}
+
+func ReferenceMediaTypes() []string {
+	return append([]string(nil), referenceMediaTypes...)
+}
+
+func ReferenceSlots(descriptions ReferenceSlotDescriptions) []contracts.ReferenceSlot {
+	return contracts.CloneReferenceSlots([]contracts.ReferenceSlot{
+		{
+			Name:               "glossary",
+			Description:        descriptions.Glossary,
+			AcceptedMediaTypes: ReferenceMediaTypes(),
+		},
+		{
+			Name:               "party",
+			Description:        descriptions.Party,
+			AcceptedMediaTypes: ReferenceMediaTypes(),
+		},
+		{
+			Name:               "players",
+			Description:        descriptions.Players,
+			AcceptedMediaTypes: ReferenceMediaTypes(),
+		},
+		{
+			Name:               "roster",
+			Description:        descriptions.Roster,
+			AcceptedMediaTypes: ReferenceMediaTypes(),
+		},
+	})
+}

internal/modules/dnd/shared/references_test.go

@@ -0,0 +1,60 @@
+package shared
+
+import (
+	"reflect"
+	"testing"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+)
+
+func TestReferenceMediaTypesReturnsDefensiveCopy(t *testing.T) {
+	want := []string{
+		"application/json",
+		"application/x-yaml",
+		"application/yaml",
+		"text/markdown",
+		"text/plain",
+	}
+	got := ReferenceMediaTypes()
+	if !reflect.DeepEqual(got, want) {
+		t.Fatalf("ReferenceMediaTypes() = %#v, want %#v", got, want)
+	}
+	got[0] = "changed"
+	if again := ReferenceMediaTypes(); again[0] != "application/json" {
+		t.Fatalf("ReferenceMediaTypes() reused backing storage: %#v", again)
+	}
+}
+
+func TestReferenceSlotsUseDescriptionsAndExpectedOrder(t *testing.T) {
+	descriptions := ReferenceSlotDescriptions{
+		Glossary: "Glossary reference",
+		Party:    "Party reference",
+		Players:  "Players reference",
+		Roster:   "Roster reference",
+	}
+	got := ReferenceSlots(descriptions)
+	want := []contracts.ReferenceSlot{
+		{Name: "glossary", Description: descriptions.Glossary, AcceptedMediaTypes: ReferenceMediaTypes()},
+		{Name: "party", Description: descriptions.Party, AcceptedMediaTypes: ReferenceMediaTypes()},
+		{Name: "players", Description: descriptions.Players, AcceptedMediaTypes: ReferenceMediaTypes()},
+		{Name: "roster", Description: descriptions.Roster, AcceptedMediaTypes: ReferenceMediaTypes()},
+	}
+	if !reflect.DeepEqual(got, want) {
+		t.Fatalf("ReferenceSlots() = %#v, want %#v", got, want)
+	}
+}
+
+func TestReferenceSlotsReturnDefensiveCopies(t *testing.T) {
+	first := ReferenceSlots(ReferenceSlotDescriptions{})
+	second := ReferenceSlots(ReferenceSlotDescriptions{})
+
+	first[0].Name = "changed"
+	first[0].AcceptedMediaTypes[0] = "changed"
+
+	if second[0].Name != "glossary" {
+		t.Fatalf("ReferenceSlots() reused slot slice: %#v", second)
+	}
+	if second[0].AcceptedMediaTypes[0] != "application/json" {
+		t.Fatalf("ReferenceSlots() reused media type slice: %#v", second)
+	}
+}

internal/modules/dnd/shared/unit_refs.go

@@ -0,0 +1,118 @@
+package shared
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
+)
+
+type UnitRef struct {
+	value      int
+	fromNumber bool
+}
+
+type SourceRefResponse struct {
+	SourceID    string  `json:"source_id"`
+	StartUnitID UnitRef `json:"start_unit_id"`
+	EndUnitID   UnitRef `json:"end_unit_id"`
+}
+
+func UnitRefFromString(value string) UnitRef {
+	parsed, _ := parseUnitRefNumber(value)
+	return UnitRef{value: parsed}
+}
+
+func UnitRefFromInt(value int) UnitRef {
+	return UnitRef{
+		value:      value,
+		fromNumber: true,
+	}
+}
+
+func (ref UnitRef) String() string {
+	if ref.value == 0 {
+		return ""
+	}
+	return strconv.Itoa(ref.value)
+}
+
+func (ref UnitRef) Int() int {
+	return ref.value
+}
+
+func (ref *UnitRef) UnmarshalJSON(raw []byte) error {
+	raw = bytes.TrimSpace(raw)
+	if len(raw) == 0 {
+		return fmt.Errorf("unit ref must be a string or integer")
+	}
+	if raw[0] == '"' {
+		var value string
+		if err := json.Unmarshal(raw, &value); err != nil {
+			return err
+		}
+		number, err := parseUnitRefNumber(value)
+		if err != nil {
+			return err
+		}
+		*ref = UnitRef{value: number}
+		return nil
+	}
+
+	number, err := parseUnitRefNumber(string(raw))
+	if err != nil {
+		return err
+	}
+	*ref = UnitRefFromInt(number)
+	return nil
+}
+
+func (ref UnitRef) MarshalJSON() ([]byte, error) {
+	if ref.fromNumber {
+		return []byte(strconv.Itoa(ref.value)), nil
+	}
+	return json.Marshal(ref.String())
+}
+
+func ResolveUnitID(doc *source.SourceDocument, field string, ref UnitRef) (int, error) {
+	if ref.value <= 0 {
+		return 0, fmt.Errorf("%s must be positive", field)
+	}
+	if _, ok := source.UnitIndex(doc, ref.value); !ok {
+		return 0, fmt.Errorf("%s %d was not found", field, ref.value)
+	}
+	return ref.value, nil
+}
+
+func SourceRefCandidate(doc *source.SourceDocument, ref SourceRefResponse) source.SourceRef {
+	return source.SourceRef{
+		SourceID:    strings.TrimSpace(ref.SourceID),
+		StartUnitID: unitIDCandidate(ref.StartUnitID),
+		EndUnitID:   unitIDCandidate(ref.EndUnitID),
+	}
+}
+
+func unitIDCandidate(ref UnitRef) int {
+	return ref.value
+}
+
+func parseUnitRefNumber(value string) (int, error) {
+	trimmed := strings.TrimSpace(value)
+	if trimmed == "" {
+		return 0, fmt.Errorf("unit ref must not be empty")
+	}
+	if trimmed != value {
+		return 0, fmt.Errorf("unit ref must not contain leading or trailing whitespace")
+	}
+	number, err := strconv.Atoi(value)
+	if err != nil {
+		return 0, fmt.Errorf("unit ref must be an integer")
+	}
+	if number <= 0 {
+		return 0, fmt.Errorf("unit ref must be positive")
+	}
+	return number, nil
+}

internal/modules/dnd/shared/unit_refs_test.go

@@ -0,0 +1,114 @@
+package shared
+
+import (
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
+)
+
+func TestUnitRefUnmarshalAcceptsIntegerAndNumericString(t *testing.T) {
+	var integerRef UnitRef
+	if err := json.Unmarshal([]byte(`12`), &integerRef); err != nil {
+		t.Fatalf("Unmarshal(integer) error = %v, want nil", err)
+	}
+	if got := integerRef.String(); got != "12" {
+		t.Fatalf("integer ref = %q, want 12", got)
+	}
+	if got := integerRef.Int(); got != 12 {
+		t.Fatalf("integer ref value = %d, want 12", got)
+	}
+
+	var stringRef UnitRef
+	if err := json.Unmarshal([]byte(`"12"`), &stringRef); err != nil {
+		t.Fatalf("Unmarshal(string) error = %v, want nil", err)
+	}
+	if got := stringRef.String(); got != "12" {
+		t.Fatalf("string ref = %q, want 12", got)
+	}
+}
+
+func TestUnitRefUnmarshalRejectsNonIntegerValues(t *testing.T) {
+	for _, raw := range []string{`true`, `null`, `1.5`, `{}`, `"seg-001"`, `" 1 "`, `0`, `-1`} {
+		t.Run(raw, func(t *testing.T) {
+			var ref UnitRef
+			err := json.Unmarshal([]byte(raw), &ref)
+			if err == nil {
+				t.Fatal("Unmarshal() error = nil, want error")
+			}
+		})
+	}
+}
+
+func TestResolveUnitIDReturnsExistingIntegerSourceUnitID(t *testing.T) {
+	doc := unitRefSourceDocument(2, 10)
+
+	got, err := ResolveUnitID(doc, "start_unit_id", UnitRefFromInt(2))
+	if err != nil {
+		t.Fatalf("ResolveUnitID() error = %v, want nil", err)
+	}
+	if got != 2 {
+		t.Fatalf("ResolveUnitID() = %d, want exact source unit ID", got)
+	}
+}
+
+func TestResolveUnitIDDoesNotFallbackToOneBasedUnitNumber(t *testing.T) {
+	doc := unitRefSourceDocument(10, 20)
+
+	_, err := ResolveUnitID(doc, "end_unit_id", UnitRefFromInt(2))
+	if err == nil {
+		t.Fatal("ResolveUnitID() error = nil, want missing source-unit ID")
+	}
+}
+
+func TestResolveUnitIDRejectsMissingUnit(t *testing.T) {
+	doc := unitRefSourceDocument(1)
+
+	_, err := ResolveUnitID(doc, "start_unit_id", UnitRefFromInt(9))
+	if err == nil {
+		t.Fatal("ResolveUnitID() error = nil, want error")
+	}
+	if !strings.Contains(err.Error(), "start_unit_id 9") {
+		t.Fatalf("ResolveUnitID() error = %q, want field and value context", err.Error())
+	}
+}
+
+func TestSourceRefCandidateCanonicalizesValidRefsAndPreservesInvalidRefs(t *testing.T) {
+	doc := unitRefSourceDocument(1, 2)
+
+	valid := SourceRefCandidate(doc, SourceRefResponse{
+		SourceID:    " session-alpha ",
+		StartUnitID: UnitRefFromInt(1),
+		EndUnitID:   UnitRefFromInt(2),
+	})
+	if valid != (source.SourceRef{SourceID: "session-alpha", StartUnitID: 1, EndUnitID: 2}) {
+		t.Fatalf("valid candidate = %#v, want canonical source ref", valid)
+	}
+
+	invalid := SourceRefCandidate(doc, SourceRefResponse{
+		SourceID:    "session-alpha",
+		StartUnitID: UnitRefFromInt(9),
+		EndUnitID:   UnitRefFromString("missing"),
+	})
+	if invalid != (source.SourceRef{SourceID: "session-alpha", StartUnitID: 9, EndUnitID: 0}) {
+		t.Fatalf("invalid candidate = %#v, want unresolved values for validator", invalid)
+	}
+}
+
+func unitRefSourceDocument(ids ...int) *source.SourceDocument {
+	doc := &source.SourceDocument{
+		ID:     "session-alpha",
+		Kind:   "transcript",
+		Format: "application/json",
+		Digest: "sha256:test",
+	}
+	for _, id := range ids {
+		doc.Units = append(doc.Units, source.SourceUnit{
+			ID:   id,
+			Kind: "transcript_segment",
+			Text: "text",
+		})
+	}
+	return doc
+}