Publish evidence context in JSON output bundles

2026-07-27 18:23:48 +00:00
parent e61e522662
commit 256cc98ddb
6 changed files with 276 additions and 31 deletions
--- a/examples/dnd-complete.config.yml
+++ b/examples/dnd-complete.config.yml
@@ -28,6 +28,15 @@ pipelines:
      module: json
      options:
        include_chunk_map: true
+        evidence_context:
+          enabled: true
+          window_units: 3
+          lanes:
+            - item-events
+            - npcs
+            - spells
+            - combat-turns
+            - npc-interactions
    steps:
      # Establish session-wide reference artifacts alongside independent item events.
      - id: describe-session
--- a/internal/framework/pipeline/evidence_output.go
+++ b/internal/framework/pipeline/evidence_output.go
@@ -13,12 +13,16 @@ import (
 // identifiers and aggregate counts. The evidence document itself can include
 // source text and must never be written to this debug envelope.
 type debugEvidenceContextSummary struct {
-	SourceID      string   `json:"source_id"`
-	SelectedLanes []string `json:"selected_lanes"`
-	WindowUnits   int      `json:"window_units"`
-	ContextCount  int      `json:"context_count"`
-	UnitCount     int      `json:"unit_count"`
-	SourceDigest  string   `json:"source_digest"`
+	ArtifactKind  contracts.ArtifactKind `json:"artifact_kind"`
+	MediaType     string                 `json:"media_type"`
+	SchemaID      string                 `json:"schema_id"`
+	SchemaName    string                 `json:"schema_name"`
+	SchemaVersion string                 `json:"schema_version"`
+	SelectedLanes []string               `json:"selected_lanes"`
+	WindowUnits   int                    `json:"window_units"`
+	ContextCount  int                    `json:"context_count"`
+	UnitCount     int                    `json:"unit_count"`
+	SourceDigest  string                 `json:"source_digest"`
 }

 // buildOutputEvidenceContext projects the prepared output policy from accepted
@@ -90,7 +94,11 @@ func buildOutputEvidenceContext(prepared *PreparedPipeline, doc *source.SourceDo
 		Content:   content,
 	}
 	summary := debugEvidenceContextSummary{
-		SourceID:      document.SourceID,
+		ArtifactKind:  artifact.Kind,
+		MediaType:     artifact.MediaType,
+		SchemaID:      artifact.Schema.ID,
+		SchemaName:    artifact.Schema.Name,
+		SchemaVersion: artifact.Schema.Version,
 		SelectedLanes: append([]string(nil), document.SelectedLanes...),
 		WindowUnits:   document.WindowUnits,
 		ContextCount:  len(document.Contexts),
--- a/internal/framework/pipeline/evidence_output_test.go
+++ b/internal/framework/pipeline/evidence_output_test.go
@@ -104,7 +104,7 @@ func TestRunnerBuildsEvidenceContextFromSelectedNormalizedOutputs(t *testing.T)
 		t.Fatalf("evidence refs = %#v, want both selected lanes", got)
 	}
 	debugJSON := string(debug.json["output/evidence-context.json"])
-	if strings.Contains(debugJSON, "text-1") || strings.Contains(debugJSON, "metadata") || !strings.Contains(debugJSON, `"context_count":1`) || !strings.Contains(debugJSON, `"unit_count":3`) {
+	if strings.Contains(debugJSON, "text-1") || strings.Contains(debugJSON, "metadata") || !strings.Contains(debugJSON, `"artifact_kind":"source/evidence-context"`) || !strings.Contains(debugJSON, `"schema_id":"notarius.source.evidence_context"`) || !strings.Contains(debugJSON, `"context_count":1`) || !strings.Contains(debugJSON, `"unit_count":3`) {
 		t.Fatalf("evidence debug envelope = %s, want only allowlisted summary", debugJSON)
 	}
 }
--- a/internal/modules/generic/output/json/encoder.go
+++ b/internal/modules/generic/output/json/encoder.go
@@ -13,6 +13,7 @@ import (

 	"gitea.maximumdirect.net/eric/notarius/internal/framework/chunkmap"
 	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/evidencecontext"
 	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
 )

@@ -22,6 +23,8 @@ const contentTypeJSON = "application/json"

 const chunkMapFileName = "chunk-map.json"

+const evidenceContextFileName = "evidence-context.json"
+
 var safeOutputFileChar = regexp.MustCompile(`[^A-Za-z0-9._-]`)

 var _ contracts.OutputEncoder = (*Encoder)(nil)
@@ -223,14 +226,15 @@ func decodeEvidenceLaneIDs(value any) ([]string, error) {
 }

 type indexFile struct {
-	ManifestFile string            `json:"manifest_file"`
-	OutputFiles  []outputFileIndex `json:"output_files"`
-	RejectedFile string            `json:"rejected_file"`
-	WarningsFile string            `json:"warnings_file"`
-	ChunkMap     *chunkMapIndex    `json:"chunk_map,omitempty"`
+	ManifestFile    string            `json:"manifest_file"`
+	OutputFiles     []outputFileIndex `json:"output_files"`
+	RejectedFile    string            `json:"rejected_file"`
+	WarningsFile    string            `json:"warnings_file"`
+	ChunkMap        *artifactIndex    `json:"chunk_map,omitempty"`
+	EvidenceContext *artifactIndex    `json:"evidence_context,omitempty"`
 }

-type chunkMapIndex struct {
+type artifactIndex struct {
 	ArtifactKind  contracts.ArtifactKind `json:"artifact_kind"`
 	File          string                 `json:"file"`
 	MediaType     string                 `json:"media_type"`
@@ -264,7 +268,7 @@ func logicalFiles(req contracts.OutputRequest, options Options) ([]contracts.Out
 	})

 	outputIndexes := make([]outputFileIndex, 0, len(outputs))
-	files := make([]contracts.OutputFile, 0, len(outputs)+4)
+	files := make([]contracts.OutputFile, 0, len(outputs)+5)
 	manifestFile, err := jsonFile("manifest.json", req.Manifest)
 	if err != nil {
 		return nil, err
@@ -311,6 +315,14 @@ func logicalFiles(req contracts.OutputRequest, options Options) ([]contracts.Out
 		files = append(files, chunkMapOutput)
 		index.ChunkMap = &chunkMapDescriptor
 	}
+	if req.EvidenceContext != nil {
+		evidenceOutput, evidenceDescriptor, err := serializedEvidenceContextFile(*req.EvidenceContext)
+		if err != nil {
+			return nil, err
+		}
+		files = append(files, evidenceOutput)
+		index.EvidenceContext = &evidenceDescriptor
+	}
 	indexOutput, err := jsonFile("index.json", index)
 	if err != nil {
 		return nil, err
@@ -330,24 +342,24 @@ func logicalFiles(req contracts.OutputRequest, options Options) ([]contracts.Out
 	return files, nil
 }

-func serializedChunkMapFile(artifact contracts.SerializedArtifact) (contracts.OutputFile, chunkMapIndex, error) {
+func serializedChunkMapFile(artifact contracts.SerializedArtifact) (contracts.OutputFile, artifactIndex, error) {
 	if artifact.Kind != chunkmap.ArtifactKind {
-		return contracts.OutputFile{}, chunkMapIndex{}, encoderErrorf("chunk map has unexpected artifact kind %q", artifact.Kind)
+		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("chunk map has unexpected artifact kind %q", artifact.Kind)
 	}
 	if artifact.Schema.ID != chunkmap.SchemaID || artifact.Schema.Name != chunkmap.SchemaName || artifact.Schema.Version != chunkmap.SchemaVersion {
-		return contracts.OutputFile{}, chunkMapIndex{}, encoderErrorf("chunk map has unexpected schema identity")
+		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("chunk map has unexpected schema identity")
 	}
 	if strings.TrimSpace(artifact.MediaType) != chunkmap.MediaType {
-		return contracts.OutputFile{}, chunkMapIndex{}, encoderErrorf("chunk map has unsupported media type %q", artifact.MediaType)
+		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("chunk map has unsupported media type %q", artifact.MediaType)
 	}
 	if _, err := chunkmap.New().Decode(artifact.Content); err != nil {
-		return contracts.OutputFile{}, chunkMapIndex{}, encoderErrorf("decode chunk map: %w", err)
+		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("decode chunk map: %w", err)
 	}
 	file, err := serializedOutputFile(chunkMapFileName, artifact)
 	if err != nil {
-		return contracts.OutputFile{}, chunkMapIndex{}, err
+		return contracts.OutputFile{}, artifactIndex{}, err
 	}
-	return file, chunkMapIndex{
+	return file, artifactIndex{
 		ArtifactKind:  artifact.Kind,
 		File:          chunkMapFileName,
 		MediaType:     chunkmap.MediaType,
@@ -357,6 +369,32 @@ func serializedChunkMapFile(artifact contracts.SerializedArtifact) (contracts.Ou
 	}, nil
 }

+func serializedEvidenceContextFile(artifact contracts.SerializedArtifact) (contracts.OutputFile, artifactIndex, error) {
+	codec := evidencecontext.New()
+	expected := codec.Schema()
+	if artifact.Kind != evidencecontext.ArtifactKind ||
+		artifact.Schema.ID != expected.ID || artifact.Schema.Name != expected.Name || artifact.Schema.Version != expected.Version ||
+		contracts.DigestArtifactSchema(artifact.Schema) != contracts.DigestArtifactSchema(expected) ||
+		strings.TrimSpace(artifact.MediaType) != evidencecontext.MediaType {
+		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("evidence context artifact is invalid")
+	}
+	if _, err := codec.Decode(artifact.Content); err != nil {
+		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("evidence context artifact is invalid")
+	}
+	file, err := serializedOutputFile(evidenceContextFileName, artifact)
+	if err != nil {
+		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("evidence context artifact is invalid")
+	}
+	return file, artifactIndex{
+		ArtifactKind:  evidencecontext.ArtifactKind,
+		File:          evidenceContextFileName,
+		MediaType:     evidencecontext.MediaType,
+		SchemaID:      expected.ID,
+		SchemaName:    expected.Name,
+		SchemaVersion: expected.Version,
+	}, nil
+}
+
 func serializedOutputFile(name string, artifact contracts.SerializedArtifact) (contracts.OutputFile, error) {
 	content := append([]byte(nil), artifact.Content...)
 	if len(content) == 0 {
--- a/internal/modules/generic/output/json/encoder_test.go
+++ b/internal/modules/generic/output/json/encoder_test.go
@@ -11,6 +11,7 @@ import (
 	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
 	"gitea.maximumdirect.net/eric/notarius/internal/framework/chunkmap"
 	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/evidencecontext"
 	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
 )

@@ -338,6 +339,82 @@ func TestEncodeRejectsInvalidChunkMapArtifact(t *testing.T) {
 	}
 }

+func TestEncodeIncludesValidatedEvidenceContext(t *testing.T) {
+	artifact := acceptedEvidenceContextArtifact(t)
+	result, err := New().Encode(context.Background(), contracts.OutputRequest{EvidenceContext: &artifact})
+	if err != nil {
+		t.Fatalf("Encode() error = %v", err)
+	}
+	if got := string(fileBytes(t, result.Files, evidenceContextFileName)); !strings.HasSuffix(got, "\n") || !stdjson.Valid([]byte(got)) {
+		t.Fatalf("evidence context file = %q, want pretty valid newline-terminated JSON", got)
+	}
+	value, err := evidencecontext.New().Decode(fileBytes(t, result.Files, evidenceContextFileName))
+	if err != nil {
+		t.Fatalf("Decode(evidence context file) error = %v", err)
+	}
+	if len(value.Contexts) != 0 {
+		t.Fatalf("evidence context = %#v, want explicit empty contexts", value)
+	}
+	index := decodeObject(t, fileBytes(t, result.Files, "index.json"))
+	if got, want := index["evidence_context"], map[string]any{
+		"artifact_kind":  string(evidencecontext.ArtifactKind),
+		"file":           evidenceContextFileName,
+		"media_type":     evidencecontext.MediaType,
+		"schema_id":      evidencecontext.SchemaID,
+		"schema_name":    evidencecontext.SchemaName,
+		"schema_version": evidencecontext.SchemaVersion,
+	}; !reflect.DeepEqual(got, want) {
+		t.Fatalf("evidence context descriptor = %#v, want %#v", got, want)
+	}
+	for _, entry := range index["output_files"].([]any) {
+		if entry.(map[string]any)["file"] == evidenceContextFileName {
+			t.Fatalf("output_files = %#v, want no evidence context lane entry", index["output_files"])
+		}
+	}
+}
+
+func TestEncodeRejectsInvalidEvidenceContextArtifactWithoutContentLeakage(t *testing.T) {
+	artifact := acceptedEvidenceContextArtifact(t)
+	for _, test := range []struct {
+		name   string
+		mutate func(*contracts.SerializedArtifact)
+	}{
+		{name: "kind", mutate: func(artifact *contracts.SerializedArtifact) { artifact.Kind = "other/evidence" }},
+		{name: "schema identity", mutate: func(artifact *contracts.SerializedArtifact) { artifact.Schema.Version = "v2" }},
+		{name: "schema digest", mutate: func(artifact *contracts.SerializedArtifact) { artifact.Schema.JSONSchema[0] = '[' }},
+		{name: "media type", mutate: func(artifact *contracts.SerializedArtifact) { artifact.MediaType = "text/plain" }},
+		{name: "payload", mutate: func(artifact *contracts.SerializedArtifact) {
+			artifact.Content = []byte(`{"source_id":"secret transcript text"}`)
+		}},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			candidate := contracts.CloneSerializedArtifact(artifact)
+			test.mutate(&candidate)
+			_, err := New().Encode(context.Background(), contracts.OutputRequest{EvidenceContext: &candidate})
+			if err == nil || err.Error() != "json output encoder: evidence context artifact is invalid" {
+				t.Fatalf("Encode() error = %v, want fixed evidence artifact error", err)
+			}
+			if strings.Contains(err.Error(), "secret transcript text") {
+				t.Fatalf("Encode() leaked artifact content: %v", err)
+			}
+		})
+	}
+}
+
+func TestEncodeOmitsEvidenceContextWhenArtifactIsAbsent(t *testing.T) {
+	result, err := New().Encode(context.Background(), contracts.OutputRequest{Manifest: artifacts.RunManifest{RunID: "run-1"}})
+	if err != nil {
+		t.Fatalf("Encode() error = %v", err)
+	}
+	if got := outputFileNames(result.Files); containsString(got, evidenceContextFileName) {
+		t.Fatalf("file names = %#v, want no evidence context", got)
+	}
+	index := decodeObject(t, fileBytes(t, result.Files, "index.json"))
+	if _, ok := index["evidence_context"]; ok {
+		t.Fatalf("index = %#v, want no evidence context descriptor", index)
+	}
+}
+
 func TestEncodeChunkMapDoesNotMutateRequest(t *testing.T) {
 	artifact := acceptedChunkMapArtifact(t)
 	artifact.Metadata = map[string]any{"owner": "caller"}
@@ -649,6 +726,33 @@ func acceptedChunkMapArtifactWithPlanAnnotation(t *testing.T, annotation stdjson
 	return artifact
 }

+func acceptedEvidenceContextArtifact(t *testing.T) contracts.SerializedArtifact {
+	t.Helper()
+	document := &source.SourceDocument{
+		ID:     "source-1",
+		Kind:   "text",
+		Format: "text/plain",
+		Units: []source.SourceUnit{{
+			ID:   7,
+			Kind: "text",
+			Text: "Source content retained only in the evidence artifact.",
+			Ref:  source.SourceRef{SourceID: "source-1", StartUnitID: 7, EndUnitID: 7},
+		}},
+	}
+	digest, err := source.DigestDocument(document)
+	if err != nil {
+		t.Fatal(err)
+	}
+	document.Digest = digest
+	artifact, err := evidencecontext.Serialize(evidencecontext.BuildRequest{
+		Source: document, WindowUnits: 3, SelectedLanes: []string{"spells"},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	return artifact
+}
+
 func outputFileNames(files []contracts.OutputFile) []string {
 	names := make([]string, 0, len(files))
 	for _, file := range files {
--- a/internal/modules/integration/dnd_npc_grounded_test.go
+++ b/internal/modules/integration/dnd_npc_grounded_test.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"reflect"
 	"strings"
 	"sync"
 	"testing"
@@ -14,6 +15,7 @@ import (
 	"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/evidencecontext"
 	"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"
@@ -254,6 +256,65 @@ func TestNPCOutputGroundsSpellAndCombatConsumersThroughOneOperation(t *testing.T
 	assertCurrentEvidence(t, combatValue.CombatTurns[0].SourceRefs)
 }

+func TestProductionDNDOutputPublishesSelectedEvidenceContext(t *testing.T) {
+	registries := productionNPCRegistries(t)
+	configValue := loadGroundedPipelineConfig(t)
+	profile := configValue.Pipelines["dnd-npc-grounded"]
+	profile.Output.Options = map[string]any{"evidence_context": map[string]any{
+		"enabled": true, "window_units": 0, "lanes": []any{"npcs", "spells", "combat"},
+	}}
+	configValue.Pipelines["dnd-npc-grounded"] = profile
+	raw := strings.NewReplacer(
+		`"id": 1`, `"id": 10`,
+		`"id": 2`, `"id": 30`,
+		`"id": 3`, `"id": 20`,
+		`"id": 4`, `"id": 50`,
+		`"id": 5`, `"id": 40`,
+	).Replace(string(readNPCFixture(t)))
+	output := runGroundedPipelineWithRaw(t, configValue, registries, &groundedDNDLLMClient{firstUnitID: 10, thirdUnitID: 20}, nil, []byte(raw))
+
+	value, err := evidencecontext.New().Decode(outputFileContent(t, output.OutputFiles, "evidence-context.json"))
+	if err != nil {
+		t.Fatalf("Decode(evidence context) error = %v", err)
+	}
+	if !reflect.DeepEqual(value.SelectedLanes, []string{"combat", "npcs", "spells"}) {
+		t.Fatalf("selected lanes = %#v, want configured production lanes without scene descriptions", value.SelectedLanes)
+	}
+	if len(value.Contexts) != 2 || len(value.Contexts[0].Units) != 1 || len(value.Contexts[1].Units) != 1 || value.Contexts[0].Units[0].ID != 10 || value.Contexts[1].Units[0].ID != 20 {
+		t.Fatalf("evidence contexts = %#v, want source-position union with non-monotonic unit IDs", value.Contexts)
+	}
+	firstRefs := value.Contexts[0].EvidenceRefs
+	if len(firstRefs) != 3 || firstRefs[0].LaneID != "combat" || firstRefs[1].LaneID != "npcs" || firstRefs[2].LaneID != "spells" {
+		t.Fatalf("first context evidence = %#v, want overlapping selected lane references", firstRefs)
+	}
+	for _, context := range value.Contexts {
+		for _, reference := range context.EvidenceRefs {
+			if reference.LaneID == "scene-descriptions" {
+				t.Fatalf("evidence refs = %#v, want scene descriptions excluded by allowlist", value.Contexts)
+			}
+		}
+	}
+}
+
+func TestProductionDNDOutputCanExplicitlySelectSceneDescriptionEvidence(t *testing.T) {
+	registries := productionNPCRegistries(t)
+	configValue := loadGroundedPipelineConfig(t)
+	profile := configValue.Pipelines["dnd-npc-grounded"]
+	profile.Output.Options = map[string]any{"evidence_context": map[string]any{
+		"enabled": true, "lanes": []any{"scene-descriptions"},
+	}}
+	configValue.Pipelines["dnd-npc-grounded"] = profile
+	output := runGroundedPipeline(t, configValue, registries, &groundedDNDLLMClient{}, nil)
+
+	value, err := evidencecontext.New().Decode(outputFileContent(t, output.OutputFiles, "evidence-context.json"))
+	if err != nil {
+		t.Fatalf("Decode(evidence context) error = %v", err)
+	}
+	if !reflect.DeepEqual(value.SelectedLanes, []string{"scene-descriptions"}) || len(value.Contexts) == 0 || len(value.Contexts[0].EvidenceRefs) == 0 || value.Contexts[0].EvidenceRefs[0].LaneID != "scene-descriptions" {
+		t.Fatalf("evidence context = %#v, want explicitly selected scene-description evidence", value)
+	}
+}
+
 func TestGroundedPipelineSkipsCombatForExactNarrativeScene(t *testing.T) {
 	registries := productionNPCRegistries(t)
 	configValue := loadGroundedPipelineConfig(t)
@@ -338,6 +399,10 @@ func configureGroundedCampaignReferences(t *testing.T, cfg *config.Config) map[s
 }

 func runGroundedPipeline(t *testing.T, configValue config.Config, registries pipeline.Registries, client *groundedDNDLLMClient, checkpoint pipeline.CheckpointLoader) pipeline.RunOutput {
+	return runGroundedPipelineWithRaw(t, configValue, registries, client, checkpoint, readNPCFixture(t))
+}
+
+func runGroundedPipelineWithRaw(t *testing.T, configValue config.Config, registries pipeline.Registries, client *groundedDNDLLMClient, checkpoint pipeline.CheckpointLoader, raw []byte) pipeline.RunOutput {
 	t.Helper()
 	catalog := moduleCatalog(registries)
 	effective, err := configValue.Resolve(config.ResolveInput{PipelineID: "dnd-npc-grounded", Catalog: catalog})
@@ -354,7 +419,7 @@ func runGroundedPipeline(t *testing.T, configValue config.Config, registries pip
 	}
 	output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{
 		Prepared:       prepared,
-		RawInput:       readNPCFixture(t),
+		RawInput:       append([]byte(nil), raw...),
 		ExtractWorkers: 1,
 		Checkpoint:     checkpoint,
 	})
@@ -364,6 +429,17 @@ func runGroundedPipeline(t *testing.T, configValue config.Config, registries pip
 	return output
 }

+func outputFileContent(t *testing.T, files []contracts.OutputFile, name string) []byte {
+	t.Helper()
+	for _, file := range files {
+		if file.Name == name {
+			return append([]byte(nil), file.Bytes...)
+		}
+	}
+	t.Fatalf("output files = %#v, missing %q", files, name)
+	return nil
+}
+
 func normalizedCombatOutput(t *testing.T, output pipeline.RunOutput) dnd.CombatTurnList {
 	t.Helper()
 	for _, serialized := range output.NormalizeOutputs {
@@ -436,10 +512,12 @@ func (loader *generatedReferenceCheckpointLoader) extractDependencies(laneID str
 }

 type groundedDNDLLMClient struct {
-	mu         sync.Mutex
-	requests   []contracts.StructuredCompletionRequest
-	sceneKind  dnd.SceneKind
-	sceneTitle string
+	mu          sync.Mutex
+	requests    []contracts.StructuredCompletionRequest
+	sceneKind   dnd.SceneKind
+	sceneTitle  string
+	firstUnitID int
+	thirdUnitID int
 }

 func (client *groundedDNDLLMClient) CompleteStructured(ctx context.Context, request contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
@@ -451,14 +529,22 @@ func (client *groundedDNDLLMClient) CompleteStructured(ctx context.Context, requ
 	client.mu.Unlock()

 	var payload any
+	firstUnitID := client.firstUnitID
+	if firstUnitID == 0 {
+		firstUnitID = 1
+	}
+	thirdUnitID := client.thirdUnitID
+	if thirdUnitID == 0 {
+		thirdUnitID = 3
+	}
 	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}},
+				"name": "Mira Thorn", "source_refs": []any{map[string]int{"start_unit_id": firstUnitID, "end_unit_id": firstUnitID}},
 			},
 			map[string]any{
-				"name": "Hooded Guard", "source_refs": []any{map[string]int{"start_unit_id": 3, "end_unit_id": 3}},
+				"name": "Hooded Guard", "source_refs": []any{map[string]int{"start_unit_id": thirdUnitID, "end_unit_id": thirdUnitID}},
 			},
 		}}
 	case npcnormalize.PromptID:
@@ -476,13 +562,13 @@ func (client *groundedDNDLLMClient) CompleteStructured(ctx context.Context, requ
 	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}},
+			"source_refs": []any{map[string]int{"start_unit_id": firstUnitID, "end_unit_id": firstUnitID}},
 		}}}
 	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}},
+			"source_refs": []any{map[string]int{"start_unit_id": firstUnitID, "end_unit_id": firstUnitID}},
 		}}}
 	default:
 		return contracts.StructuredCompletionResponse{}, fmt.Errorf("unexpected grounded prompt %q", request.PromptID)