eric/notarius
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Harden chunk map export and retire completed plans

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-23 17:38:04 +00:00
parent 06148074a2
commit b490297cde
6 changed files with 166 additions and 733 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
internal/framework/chunkmap/codec.go
View File

@@ -8,9 +8,11 @@ import (
"io"
"regexp"
"strings"
"sync"
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
"github.com/santhosh-tekuri/jsonschema/v6"
)
//go:embed assets/schemas/source_chunk_map.v1.json
@@ -18,6 +20,13 @@ var schemaAssets embed.FS
var digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
var (
loadSchemaOnce sync.Once
loadedSchema []byte
compiledSchema *jsonschema.Schema
loadSchemaErr error
)
// Codec owns strict serialization for the durable chunk-map contract.
type Codec struct{}
@@ -123,7 +132,7 @@ func (c *Codec) Encode(value ChunkMap) ([]byte, error) {
if _, err := c.schemaBytes(); err != nil {
return nil, err
}
canonical, err := canonicalize(value)
canonical, err := canonicalize(clone(value))
if err != nil {
return nil, fmt.Errorf("encode source chunk map: %w", err)
}
@@ -131,6 +140,9 @@ func (c *Codec) Encode(value ChunkMap) ([]byte, error) {
if err != nil {
return nil, fmt.Errorf("encode source chunk map: %w", err)
}
if err := validateSchemaInstance(content); err != nil {
return nil, fmt.Errorf("encode source chunk map: %w", err)
}
return content, nil
}
@@ -138,6 +150,9 @@ func (c *Codec) Decode(content []byte) (ChunkMap, error) {
if _, err := c.schemaBytes(); err != nil {
return ChunkMap{}, err
}
if err := validateSchemaInstance(content); err != nil {
return ChunkMap{}, fmt.Errorf("decode source chunk map: %w", err)
}
decoder := json.NewDecoder(bytes.NewReader(content))
decoder.DisallowUnknownFields()
var value ChunkMap
@@ -156,23 +171,61 @@ func (c *Codec) Decode(content []byte) (ChunkMap, error) {
}
func (c *Codec) schemaBytes() ([]byte, error) {
loadSchemaOnce.Do(loadAndCompileSchema)
if loadSchemaErr != nil {
return nil, loadSchemaErr
}
return append([]byte(nil), loadedSchema...), nil
}
func loadAndCompileSchema() {
raw, err := schemaAssets.ReadFile("assets/schemas/source_chunk_map.v1.json")
if err != nil {
return nil, fmt.Errorf("read source chunk map schema: %w", err)
loadSchemaErr = fmt.Errorf("read source chunk map schema: %w", err)
return
}
var schema struct {
var identity struct {
ID string `json:"$id"`
Title string `json:"title"`
Type string `json:"type"`
Required []string `json:"required"`
}
if err := json.Unmarshal(raw, &schema); err != nil {
return nil, fmt.Errorf("decode source chunk map schema: %w", err)
if err := json.Unmarshal(raw, &identity); err != nil {
loadSchemaErr = fmt.Errorf("decode source chunk map schema: %w", err)
return
}
if schema.ID != SchemaID || schema.Title != SchemaName || schema.Type != "object" || !hasRequiredFields(schema.Required) {
return nil, fmt.Errorf("source chunk map schema identity or required fields are invalid")
if identity.ID != SchemaID || identity.Title != SchemaName || identity.Type != "object" || !hasRequiredFields(identity.Required) {
loadSchemaErr = fmt.Errorf("source chunk map schema identity or required fields are invalid")
return
}
return append([]byte(nil), raw...), nil
schemaDocument, err := jsonschema.UnmarshalJSON(bytes.NewReader(raw))
if err != nil {
loadSchemaErr = fmt.Errorf("parse source chunk map schema: %w", err)
return
}
compiler := jsonschema.NewCompiler()
if err := compiler.AddResource("source-chunk-map-schema.json", schemaDocument); err != nil {
loadSchemaErr = fmt.Errorf("load source chunk map schema: %w", err)
return
}
compiled, err := compiler.Compile("source-chunk-map-schema.json")
if err != nil {
loadSchemaErr = fmt.Errorf("compile source chunk map schema: %w", err)
return
}
loadedSchema = append([]byte(nil), raw...)
compiledSchema = compiled
}
func validateSchemaInstance(content []byte) error {
instance, err := jsonschema.UnmarshalJSON(bytes.NewReader(content))
if err != nil {
return fmt.Errorf("payload is not valid JSON: %w", err)
}
if err := compiledSchema.Validate(instance); err != nil {
return fmt.Errorf("payload does not conform to source chunk map schema: %w", err)
}
return nil
}
func hasRequiredFields(required []string) bool {

72
internal/framework/chunkmap/codec_test.go
View File

@@ -4,6 +4,7 @@ import (
"bytes"
"encoding/json"
"os"
"reflect"
"strings"
"testing"
@@ -140,6 +141,62 @@ func TestCodecRejectsInvalidDurableBoundaries(t *testing.T) {
}
}
func TestDecodeEnforcesRequiredSchemaFieldsAndTypes(t *testing.T) {
request := acceptedBuildRequest(t)
request.Plan.Annotations = nil
var err error
request.Chunks, err = source.MaterializeChunkPlan(request.Source, request.Plan)
if err != nil {
t.Fatal(err)
}
artifact, err := Serialize(request)
if err != nil {
t.Fatal(err)
}
for _, test := range []struct {
name string
mutate func(map[string]any)
}{
{name: "missing plan annotations", mutate: func(value map[string]any) { delete(value, "plan_annotations") }},
{name: "null plan annotations", mutate: func(value map[string]any) { value["plan_annotations"] = nil }},
{name: "missing first index", mutate: func(value map[string]any) { delete(chunkDocument(value, 0), "index") }},
{name: "null first index", mutate: func(value map[string]any) { chunkDocument(value, 0)["index"] = nil }},
{name: "missing empty chunk annotations", mutate: func(value map[string]any) { delete(chunkDocument(value, 1), "annotations") }},
{name: "null empty chunk annotations", mutate: func(value map[string]any) { chunkDocument(value, 1)["annotations"] = nil }},
{name: "explicit empty llm profile", mutate: func(value map[string]any) {
value["producer"].(map[string]any)["llm_profile"] = ""
}},
} {
t.Run(test.name, func(t *testing.T) {
value := decodeJSONDocument(t, artifact.Content)
test.mutate(value)
content, err := json.Marshal(value)
if err != nil {
t.Fatal(err)
}
if _, err := New().Decode(content); err == nil {
t.Fatalf("Decode(%s) error = nil, want schema rejection", content)
}
})
}
}
func TestEncodeDoesNotMutateValue(t *testing.T) {
value, err := Build(acceptedBuildRequest(t))
if err != nil {
t.Fatal(err)
}
value.Chunks[0].Annotations["dnd/scenes"] = json.RawMessage(" { \n \"kind\" : \"narrative\" \n } ")
before := clone(value)
if _, err := New().Encode(value); err != nil {
t.Fatalf("Encode() error = %v", err)
}
if !reflect.DeepEqual(value, before) {
t.Fatalf("Encode() mutated value:\nbefore: %#v\nafter: %#v", before, value)
}
}
func TestChunkMapOwnershipIsIndependent(t *testing.T) {
request := acceptedBuildRequest(t)
first, err := Build(request)
@@ -157,6 +214,21 @@ func TestChunkMapOwnershipIsIndependent(t *testing.T) {
}
}
func decodeJSONDocument(t *testing.T, content []byte) map[string]any {
t.Helper()
decoder := json.NewDecoder(bytes.NewReader(content))
decoder.UseNumber()
var value map[string]any
if err := decoder.Decode(&value); err != nil {
t.Fatal(err)
}
return value
}
func chunkDocument(value map[string]any, index int) map[string]any {
return value["chunks"].([]any)[index].(map[string]any)
}
func acceptedBuildRequest(t *testing.T) BuildRequest {
t.Helper()
document := &source.SourceDocument{

10
internal/modules/generic/output/json/encoder.go
View File

@@ -1,9 +1,11 @@
package json
import (
"bytes"
"context"
stdjson "encoding/json"
"fmt"
"io"
"mime"
"regexp"
"sort"
@@ -247,10 +249,16 @@ func serializedOutputFile(name string, artifact contracts.SerializedArtifact) (c
if !isJSONMediaType(mediaType) {
return contracts.OutputFile{}, encoderErrorf("normalized output %q has unsupported media type %q", name, mediaType)
}
decoder := stdjson.NewDecoder(bytes.NewReader(content))
decoder.UseNumber()
var decoded any
if err := stdjson.Unmarshal(content, &decoded); err != nil {
if err := decoder.Decode(&decoded); err != nil {
return contracts.OutputFile{}, encoderErrorf("normalized output %q contains invalid JSON: %w", name, err)
}
var trailing any
if err := decoder.Decode(&trailing); err != io.EOF {
return contracts.OutputFile{}, encoderErrorf("normalized output %q contains multiple JSON values", name)
}
pretty, err := marshalPretty(decoded)
if err != nil {
return contracts.OutputFile{}, err

24
internal/modules/generic/output/json/encoder_test.go
View File

@@ -234,6 +234,23 @@ func TestEncodeIncludesValidatedChunkMap(t *testing.T) {
}
}
func TestEncodePreservesChunkMapAnnotationNumbers(t *testing.T) {
artifact := acceptedChunkMapArtifactWithPlanAnnotation(t, stdjson.RawMessage(`{"decimal":1.0,"large":9007199254740993}`))
result, err := NewWithOptions(Options{IncludeChunkMap: true}).Encode(context.Background(), contracts.OutputRequest{
ChunkMap: &artifact,
})
if err != nil {
t.Fatalf("Encode() error = %v", err)
}
value, err := chunkmap.New().Decode(fileBytes(t, result.Files, chunkMapFileName))
if err != nil {
t.Fatalf("Decode(emitted chunk map) error = %v", err)
}
if got, want := string(value.PlanAnnotations["test/numbers"]), `{"decimal":1.0,"large":9007199254740993}`; got != want {
t.Fatalf("numeric annotation = %s, want %s", got, want)
}
}
func TestEncodeOmitsChunkMapWithoutAcceptedArtifact(t *testing.T) {
result, err := NewWithOptions(Options{IncludeChunkMap: true}).Encode(context.Background(), contracts.OutputRequest{
Manifest: artifacts.RunManifest{RunID: "run-1"},
@@ -544,6 +561,10 @@ func normalizeOutput(laneID string, content string) contracts.SerializedOutput {
}
func acceptedChunkMapArtifact(t *testing.T) contracts.SerializedArtifact {
return acceptedChunkMapArtifactWithPlanAnnotation(t, nil)
}
func acceptedChunkMapArtifactWithPlanAnnotation(t *testing.T, annotation stdjson.RawMessage) contracts.SerializedArtifact {
t.Helper()
document := &source.SourceDocument{
ID: "source-1",
@@ -562,6 +583,9 @@ func acceptedChunkMapArtifact(t *testing.T) contracts.SerializedArtifact {
}
document.Digest = digest
plan := source.ChunkPlan{SourceDigest: digest, Ranges: []source.ChunkRange{{StartUnitID: 1, EndUnitID: 1}}}
if annotation != nil {
plan.Annotations = source.ChunkAnnotations{"test/numbers": append(stdjson.RawMessage(nil), annotation...)}
}
chunks, err := source.MaterializeChunkPlan(document, plan)
if err != nil {
t.Fatal(err)