From 16a998055cab04843d424ac4de99d39cf88362ae Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Thu, 23 Jul 2026 15:02:01 +0000 Subject: [PATCH] Add opt-in JSON chunk map export --- .../modules/generic/output/json/encoder.go | 70 ++++++- .../generic/output/json/encoder_test.go | 196 ++++++++++++++++++ 2 files changed, 261 insertions(+), 5 deletions(-) diff --git a/internal/modules/generic/output/json/encoder.go b/internal/modules/generic/output/json/encoder.go index 3f9082f..1775b73 100644 --- a/internal/modules/generic/output/json/encoder.go +++ b/internal/modules/generic/output/json/encoder.go @@ -9,6 +9,7 @@ import ( "sort" "strings" + "gitea.maximumdirect.net/eric/notarius/internal/framework/chunkmap" "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" "gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline" ) @@ -17,18 +18,25 @@ const Key = "json" const contentTypeJSON = "application/json" +const chunkMapFileName = "chunk-map.json" + var safeOutputFileChar = regexp.MustCompile(`[^A-Za-z0-9._-]`) var _ contracts.OutputEncoder = (*Encoder)(nil) -type Options struct{} +type Options struct { + IncludeChunkMap bool +} type Encoder struct { options Options } func New() *Encoder { - options, _ := DecodeOptions(nil) + return NewWithOptions(Options{}) +} + +func NewWithOptions(options Options) *Encoder { return &Encoder{options: options} } @@ -47,7 +55,7 @@ func (e *Encoder) Encode(ctx context.Context, req contracts.OutputRequest) (cont return contracts.OutputResult{}, encoderErrorf("context error before encoding: %w", err) } - files, err := logicalFiles(req) + files, err := logicalFiles(req, e.options) if err != nil { return contracts.OutputResult{}, err } @@ -79,9 +87,16 @@ func validateOptions(options map[string]any) error { } func DecodeOptions(options map[string]any) (Options, error) { - if err := pipeline.RejectUnknownOptions(options); err != nil { + if err := pipeline.RejectUnknownOptions(options, "include_chunk_map"); err != nil { return Options{}, encoderErrorf("%w", err) } + if value, ok := options["include_chunk_map"]; ok { + enabled, ok := value.(bool) + if !ok { + return Options{}, encoderErrorf("option %q must be a boolean", "include_chunk_map") + } + return Options{IncludeChunkMap: enabled}, nil + } return Options{}, nil } @@ -90,6 +105,16 @@ type indexFile struct { OutputFiles []outputFileIndex `json:"output_files"` RejectedFile string `json:"rejected_file"` WarningsFile string `json:"warnings_file"` + ChunkMap *chunkMapIndex `json:"chunk_map,omitempty"` +} + +type chunkMapIndex struct { + ArtifactKind contracts.ArtifactKind `json:"artifact_kind"` + File string `json:"file"` + MediaType string `json:"media_type"` + SchemaID string `json:"schema_id"` + SchemaName string `json:"schema_name"` + SchemaVersion string `json:"schema_version"` } type outputFileIndex struct { @@ -110,7 +135,7 @@ type warningsFile struct { Warnings []contracts.Warning `json:"warnings"` } -func logicalFiles(req contracts.OutputRequest) ([]contracts.OutputFile, error) { +func logicalFiles(req contracts.OutputRequest, options Options) ([]contracts.OutputFile, error) { outputs := cloneNormalizeOutputs(req.NormalizeOutputs) sort.SliceStable(outputs, func(i, j int) bool { return outputs[i].LaneID < outputs[j].LaneID @@ -156,6 +181,14 @@ func logicalFiles(req contracts.OutputRequest) ([]contracts.OutputFile, error) { RejectedFile: "rejected.json", WarningsFile: "warnings.json", } + if options.IncludeChunkMap && req.ChunkMap != nil { + chunkMapOutput, chunkMapDescriptor, err := serializedChunkMapFile(*req.ChunkMap) + if err != nil { + return nil, err + } + files = append(files, chunkMapOutput) + index.ChunkMap = &chunkMapDescriptor + } indexOutput, err := jsonFile("index.json", index) if err != nil { return nil, err @@ -175,6 +208,33 @@ func logicalFiles(req contracts.OutputRequest) ([]contracts.OutputFile, error) { return files, nil } +func serializedChunkMapFile(artifact contracts.SerializedArtifact) (contracts.OutputFile, chunkMapIndex, error) { + if artifact.Kind != chunkmap.ArtifactKind { + return contracts.OutputFile{}, chunkMapIndex{}, 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") + } + if strings.TrimSpace(artifact.MediaType) != chunkmap.MediaType { + return contracts.OutputFile{}, chunkMapIndex{}, 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) + } + file, err := serializedOutputFile(chunkMapFileName, artifact) + if err != nil { + return contracts.OutputFile{}, chunkMapIndex{}, err + } + return file, chunkMapIndex{ + ArtifactKind: artifact.Kind, + File: chunkMapFileName, + MediaType: chunkmap.MediaType, + SchemaID: artifact.Schema.ID, + SchemaName: artifact.Schema.Name, + SchemaVersion: artifact.Schema.Version, + }, nil +} + func serializedOutputFile(name string, artifact contracts.SerializedArtifact) (contracts.OutputFile, error) { content := append([]byte(nil), artifact.Content...) if len(content) == 0 { diff --git a/internal/modules/generic/output/json/encoder_test.go b/internal/modules/generic/output/json/encoder_test.go index 08386b9..82a79b9 100644 --- a/internal/modules/generic/output/json/encoder_test.go +++ b/internal/modules/generic/output/json/encoder_test.go @@ -8,6 +8,8 @@ import ( "testing" "gitea.maximumdirect.net/eric/notarius/internal/core/artifacts" + "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/pipeline" ) @@ -37,6 +39,40 @@ func TestModuleSpecAndRegister(t *testing.T) { if err := registry.ValidateOptions(Key, map[string]any{"unexpected": true}); err == nil || !strings.Contains(err.Error(), "unknown option") { t.Fatalf("ValidateOptions() error = %v, want unknown option error", err) } + if err := registry.ValidateOptions(Key, map[string]any{"include_chunk_map": true}); err != nil { + t.Fatalf("ValidateOptions() error = %v, want nil", err) + } +} + +func TestDecodeOptions(t *testing.T) { + for _, test := range []struct { + name string + options map[string]any + want Options + wantErr string + }{ + {name: "omitted", want: Options{}}, + {name: "disabled", options: map[string]any{"include_chunk_map": false}, want: Options{}}, + {name: "enabled", options: map[string]any{"include_chunk_map": true}, want: Options{IncludeChunkMap: true}}, + {name: "wrong type", options: map[string]any{"include_chunk_map": "true"}, wantErr: "must be a boolean"}, + {name: "unknown", options: map[string]any{"unknown": true}, wantErr: "unknown option"}, + } { + t.Run(test.name, func(t *testing.T) { + got, err := DecodeOptions(test.options) + if test.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), test.wantErr) { + t.Fatalf("DecodeOptions() error = %v, want %q", err, test.wantErr) + } + return + } + if err != nil { + t.Fatalf("DecodeOptions() error = %v, want nil", err) + } + if got != test.want { + t.Fatalf("DecodeOptions() = %#v, want %#v", got, test.want) + } + }) + } } func TestEncodeReturnsLogicalFilesForNormalizedOutputs(t *testing.T) { @@ -129,6 +165,133 @@ func TestEncodeIncludesRejectedAndWarningsWhenEmpty(t *testing.T) { } } +func TestEncodeChunkMapExportIsOptIn(t *testing.T) { + artifact := acceptedChunkMapArtifact(t) + request := contracts.OutputRequest{ + Manifest: artifacts.RunManifest{RunID: "run-1"}, + ChunkMap: &artifact, + } + for _, test := range []struct { + name string + encoder *Encoder + }{ + {name: "default", encoder: New()}, + {name: "explicitly disabled", encoder: NewWithOptions(Options{IncludeChunkMap: false})}, + } { + t.Run(test.name, func(t *testing.T) { + result, err := test.encoder.Encode(context.Background(), request) + if err != nil { + t.Fatalf("Encode() error = %v, want nil", err) + } + if got, want := outputFileNames(result.Files), []string{"index.json", "manifest.json", "rejected.json", "warnings.json"}; !reflect.DeepEqual(got, want) { + t.Fatalf("file names = %#v, want %#v", got, want) + } + index := decodeObject(t, fileBytes(t, result.Files, "index.json")) + if _, ok := index["chunk_map"]; ok { + t.Fatalf("index = %#v, want no chunk map descriptor", index) + } + }) + } +} + +func TestEncodeIncludesValidatedChunkMap(t *testing.T) { + artifact := acceptedChunkMapArtifact(t) + result, err := NewWithOptions(Options{IncludeChunkMap: true}).Encode(context.Background(), contracts.OutputRequest{ + Manifest: artifacts.RunManifest{RunID: "run-1"}, + NormalizeOutputs: []contracts.SerializedOutput{ + normalizeOutput("spells", `{"spell_casts":[]}`), + }, + ChunkMap: &artifact, + }) + if err != nil { + t.Fatalf("Encode() error = %v, want nil", err) + } + if got, want := outputFileNames(result.Files), []string{"chunk-map.json", "index.json", "lanes/spells.json", "manifest.json", "rejected.json", "warnings.json"}; !reflect.DeepEqual(got, want) { + t.Fatalf("file names = %#v, want %#v", got, want) + } + chunkMapFile := fileBytes(t, result.Files, chunkMapFileName) + if !stdjson.Valid(chunkMapFile) || !strings.HasSuffix(string(chunkMapFile), "\n") || !strings.Contains(string(chunkMapFile), "\n \"source_id\"") { + t.Fatalf("chunk map file = %q, want pretty valid newline-terminated JSON", chunkMapFile) + } + if _, err := chunkmap.New().Decode(chunkMapFile); err != nil { + t.Fatalf("Decode(chunk map file) error = %v, want nil", err) + } + index := decodeObject(t, fileBytes(t, result.Files, "index.json")) + if got, want := index["chunk_map"], map[string]any{ + "artifact_kind": "source/chunk-map", + "file": chunkMapFileName, + "media_type": chunkmap.MediaType, + "schema_id": chunkmap.SchemaID, + "schema_name": chunkmap.SchemaName, + "schema_version": chunkmap.SchemaVersion, + }; !reflect.DeepEqual(got, want) { + t.Fatalf("chunk map descriptor = %#v, want %#v", got, want) + } + for _, entry := range index["output_files"].([]any) { + if entry.(map[string]any)["file"] == chunkMapFileName { + t.Fatalf("output_files = %#v, want no chunk map", index["output_files"]) + } + } +} + +func TestEncodeOmitsChunkMapWithoutAcceptedArtifact(t *testing.T) { + result, err := NewWithOptions(Options{IncludeChunkMap: true}).Encode(context.Background(), contracts.OutputRequest{ + Manifest: artifacts.RunManifest{RunID: "run-1"}, + }) + if err != nil { + t.Fatalf("Encode() error = %v, want nil", err) + } + if got := outputFileNames(result.Files); containsString(got, chunkMapFileName) { + t.Fatalf("file names = %#v, want no chunk map", got) + } + index := decodeObject(t, fileBytes(t, result.Files, "index.json")) + if _, ok := index["chunk_map"]; ok { + t.Fatalf("index = %#v, want no chunk map descriptor", index) + } +} + +func TestEncodeRejectsInvalidChunkMapArtifact(t *testing.T) { + artifact := acceptedChunkMapArtifact(t) + for _, test := range []struct { + name string + mutate func(*contracts.SerializedArtifact) + }{ + {name: "kind", mutate: func(artifact *contracts.SerializedArtifact) { artifact.Kind = "other/chunk-map" }}, + {name: "schema", mutate: func(artifact *contracts.SerializedArtifact) { artifact.Schema.Version = "v2" }}, + {name: "media type", mutate: func(artifact *contracts.SerializedArtifact) { artifact.MediaType = "text/plain" }}, + {name: "content", mutate: func(artifact *contracts.SerializedArtifact) { artifact.Content = []byte(`{}`) }}, + } { + t.Run(test.name, func(t *testing.T) { + candidate := contracts.CloneSerializedArtifact(artifact) + test.mutate(&candidate) + _, err := NewWithOptions(Options{IncludeChunkMap: true}).Encode(context.Background(), contracts.OutputRequest{ChunkMap: &candidate}) + if err == nil { + t.Fatal("Encode() error = nil, want invalid chunk map error") + } + }) + } +} + +func TestEncodeChunkMapDoesNotMutateRequest(t *testing.T) { + artifact := acceptedChunkMapArtifact(t) + artifact.Metadata = map[string]any{"owner": "caller"} + request := contracts.OutputRequest{ChunkMap: &artifact} + before := contracts.CloneSerializedArtifact(artifact) + result, err := NewWithOptions(Options{IncludeChunkMap: true}).Encode(context.Background(), request) + if err != nil { + t.Fatalf("Encode() error = %v, want nil", err) + } + if !reflect.DeepEqual(*request.ChunkMap, before) { + t.Fatalf("chunk map artifact mutated:\nbefore: %#v\nafter: %#v", before, *request.ChunkMap) + } + request.ChunkMap.Content[0] = '[' + request.ChunkMap.Schema.JSONSchema[0] = '[' + request.ChunkMap.Metadata["owner"] = "changed" + if _, err := chunkmap.New().Decode(fileBytes(t, result.Files, chunkMapFileName)); err != nil { + t.Fatalf("chunk map output changed after request mutation: %v", err) + } +} + func TestEncodePrettyPrintsJSON(t *testing.T) { result, err := New().Encode(context.Background(), contracts.OutputRequest{ Manifest: artifacts.RunManifest{RunID: "run-1"}, @@ -380,6 +543,39 @@ func normalizeOutput(laneID string, content string) contracts.SerializedOutput { } } +func acceptedChunkMapArtifact(t *testing.T) contracts.SerializedArtifact { + t.Helper() + document := &source.SourceDocument{ + ID: "source-1", + Kind: "text", + Format: "text/plain", + Units: []source.SourceUnit{{ + ID: 1, + Kind: "text", + Text: "Accepted source content.", + Ref: source.SourceRef{SourceID: "source-1", StartUnitID: 1, EndUnitID: 1}, + }}, + } + digest, err := source.DigestDocument(document) + if err != nil { + t.Fatal(err) + } + document.Digest = digest + plan := source.ChunkPlan{SourceDigest: digest, Ranges: []source.ChunkRange{{StartUnitID: 1, EndUnitID: 1}}} + chunks, err := source.MaterializeChunkPlan(document, plan) + if err != nil { + t.Fatal(err) + } + artifact, err := chunkmap.Serialize(chunkmap.BuildRequest{ + Source: document, Plan: plan, Chunks: chunks, RequestedChunker: "generic/units", + Producer: chunkmap.Producer{InputModule: "test/input", ChunkModule: "generic/units"}, + }) + if err != nil { + t.Fatal(err) + } + return artifact +} + func outputFileNames(files []contracts.OutputFile) []string { names := make([]string, 0, len(files)) for _, file := range files {