From 3bfac14397970f614ddca680a0b4baad2a861b11 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Fri, 17 Jul 2026 23:47:22 +0000 Subject: [PATCH] Add canonical chunk plan source model --- ...5-cache-canonical-chunk-plans-by-source.md | 5 +- docs/roadmap/implementation.md | 2 +- internal/core/source/chunk_plan.go | 213 ++++++++++++++++ internal/core/source/chunk_plan_test.go | 230 ++++++++++++++++++ internal/core/source/digest.go | 69 ++++-- internal/core/source/source.go | 36 ++- internal/core/source/source_test.go | 23 ++ internal/framework/checkpoint/loader.go | 26 +- internal/framework/checkpoint/recorder.go | 32 +-- .../framework/checkpoint/recorder_test.go | 33 ++- .../pipeline/chunk_annotations_test.go | 70 ++++++ .../framework/pipeline/chunk_validation.go | 26 +- internal/framework/pipeline/debug.go | 32 +-- internal/framework/pipeline/runner.go | 2 + .../framework/pipeline/runner_concurrent.go | 2 +- .../pipeline/runner_extract_handoff_test.go | 27 ++ 16 files changed, 755 insertions(+), 73 deletions(-) create mode 100644 internal/core/source/chunk_plan.go create mode 100644 internal/core/source/chunk_plan_test.go create mode 100644 internal/framework/pipeline/chunk_annotations_test.go diff --git a/docs/adr/0005-cache-canonical-chunk-plans-by-source.md b/docs/adr/0005-cache-canonical-chunk-plans-by-source.md index 3426325..bf98ec7 100644 --- a/docs/adr/0005-cache-canonical-chunk-plans-by-source.md +++ b/docs/adr/0005-cache-canonical-chunk-plans-by-source.md @@ -1,6 +1,6 @@ # ADR-0005: Cache one canonical chunk plan per source -**Status:** Proposed +**Status:** Accepted **Date:** 2026-07-17 ## Context @@ -70,6 +70,9 @@ Notarius state. One mutable active plan is stored under the canonical source identity and retains provenance for the module and relevant runtime inputs that produced it. +Refreshing the active plan atomically replaces that one mutable record; readers +must observe either the previous complete plan or the replacement complete +plan, never a partial update. The effective plan producer is reported separately from the chunk module requested by the current pipeline; reuse must not attribute cached boundaries or annotations to a module that did not produce them. diff --git a/docs/roadmap/implementation.md b/docs/roadmap/implementation.md index 0158b7e..66567a0 100644 --- a/docs/roadmap/implementation.md +++ b/docs/roadmap/implementation.md @@ -284,7 +284,7 @@ then fills candidate and producer fields when a plan reaches materialization. ## Stage 1: Finalize ADR and add the source-zone plan model -**Status:** Not started +**Status:** Complete ### Objective diff --git a/internal/core/source/chunk_plan.go b/internal/core/source/chunk_plan.go new file mode 100644 index 0000000..0d4a89d --- /dev/null +++ b/internal/core/source/chunk_plan.go @@ -0,0 +1,213 @@ +package source + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "strings" +) + +// CanonicalizeChunkAnnotations validates annotation namespaces and JSON values +// and returns an independently owned map whose values use canonical JSON bytes. +func CanonicalizeChunkAnnotations(annotations ChunkAnnotations) (ChunkAnnotations, error) { + if len(annotations) == 0 { + return nil, nil + } + canonical := make(ChunkAnnotations, len(annotations)) + for namespace, raw := range annotations { + if strings.TrimSpace(namespace) == "" { + return nil, fmt.Errorf("chunk annotation namespace must not be empty") + } + if strings.TrimSpace(namespace) != namespace { + return nil, fmt.Errorf("chunk annotation namespace %q must not contain leading or trailing whitespace", namespace) + } + value, err := decodeAnnotation(raw) + if err != nil { + return nil, fmt.Errorf("chunk annotation %q: %w", namespace, err) + } + encoded, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("chunk annotation %q contains an unsupported value: %w", namespace, err) + } + canonical[namespace] = encoded + } + return canonical, nil +} + +func decodeAnnotation(raw json.RawMessage) (any, error) { + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return nil, fmt.Errorf("must contain valid JSON: %w", err) + } + var trailing any + if err := decoder.Decode(&trailing); err != io.EOF { + if err == nil { + return nil, fmt.Errorf("must contain exactly one JSON value") + } + return nil, fmt.Errorf("must contain exactly one JSON value: %w", err) + } + return value, nil +} + +// ValidateChunkAnnotations requires annotations to already contain canonical +// JSON. CanonicalizeChunkAnnotations can be used at producer boundaries. +func ValidateChunkAnnotations(annotations ChunkAnnotations) error { + canonical, err := CanonicalizeChunkAnnotations(annotations) + if err != nil { + return err + } + for namespace, raw := range annotations { + if !bytes.Equal(raw, canonical[namespace]) { + return fmt.Errorf("chunk annotation %q must use canonical JSON", namespace) + } + } + return nil +} + +// CloneChunkAnnotations returns a deep clone, including every raw JSON value. +func CloneChunkAnnotations(annotations ChunkAnnotations) ChunkAnnotations { + if len(annotations) == 0 { + return nil + } + cloned := make(ChunkAnnotations, len(annotations)) + for namespace, raw := range annotations { + cloned[namespace] = append(json.RawMessage(nil), raw...) + } + return cloned +} + +// CloneChunkPlan returns a deep clone of a chunk plan. +func CloneChunkPlan(plan ChunkPlan) ChunkPlan { + cloned := ChunkPlan{ + SourceDigest: plan.SourceDigest, + Ranges: make([]ChunkRange, len(plan.Ranges)), + Annotations: CloneChunkAnnotations(plan.Annotations), + } + for i, chunkRange := range plan.Ranges { + cloned.Ranges[i] = ChunkRange{ + StartUnitID: chunkRange.StartUnitID, + EndUnitID: chunkRange.EndUnitID, + Annotations: CloneChunkAnnotations(chunkRange.Annotations), + } + } + return cloned +} + +// CanonicalizeChunkPlan returns a deep clone with canonical annotation bytes. +func CanonicalizeChunkPlan(plan ChunkPlan) (ChunkPlan, error) { + canonical := CloneChunkPlan(plan) + annotations, err := CanonicalizeChunkAnnotations(plan.Annotations) + if err != nil { + return ChunkPlan{}, fmt.Errorf("chunk plan annotations: %w", err) + } + canonical.Annotations = annotations + for i := range plan.Ranges { + annotations, err := CanonicalizeChunkAnnotations(plan.Ranges[i].Annotations) + if err != nil { + return ChunkPlan{}, fmt.Errorf("chunk plan range[%d] annotations: %w", i, err) + } + canonical.Ranges[i].Annotations = annotations + } + return canonical, nil +} + +// ValidateChunkPlan validates a canonical plan against the current source. +// Ranges may contain gaps or overlap, but their start positions must increase. +func ValidateChunkPlan(doc *SourceDocument, plan ChunkPlan) error { + if err := ValidateDocument(doc); err != nil { + return fmt.Errorf("source document: %w", err) + } + if plan.SourceDigest != doc.Digest { + return fmt.Errorf("chunk plan source_digest %q does not match source document digest %q", plan.SourceDigest, doc.Digest) + } + if len(plan.Ranges) == 0 { + return fmt.Errorf("chunk plan ranges must not be empty") + } + if err := ValidateChunkAnnotations(plan.Annotations); err != nil { + return fmt.Errorf("chunk plan annotations: %w", err) + } + previousStart := -1 + for i, chunkRange := range plan.Ranges { + start, ok := UnitIndex(doc, chunkRange.StartUnitID) + if !ok { + return fmt.Errorf("chunk plan range[%d] start_unit_id %d was not found", i, chunkRange.StartUnitID) + } + end, ok := UnitIndex(doc, chunkRange.EndUnitID) + if !ok { + return fmt.Errorf("chunk plan range[%d] end_unit_id %d was not found", i, chunkRange.EndUnitID) + } + if start > end { + return fmt.Errorf("chunk plan range[%d] start_unit_id %d appears after end_unit_id %d", i, chunkRange.StartUnitID, chunkRange.EndUnitID) + } + if start <= previousStart { + return fmt.Errorf("chunk plan range[%d] start_unit_id %d does not appear after the previous range start", i, chunkRange.StartUnitID) + } + if err := ValidateChunkAnnotations(chunkRange.Annotations); err != nil { + return fmt.Errorf("chunk plan range[%d] annotations: %w", i, err) + } + previousStart = start + } + return nil +} + +// MaterializeChunkPlan deterministically expands a validated plan into chunks. +func MaterializeChunkPlan(doc *SourceDocument, plan ChunkPlan) ([]Chunk, error) { + if err := ValidateChunkPlan(doc, plan); err != nil { + return nil, err + } + chunks := make([]Chunk, 0, len(plan.Ranges)) + for index, chunkRange := range plan.Ranges { + start, _ := UnitIndex(doc, chunkRange.StartUnitID) + end, _ := UnitIndex(doc, chunkRange.EndUnitID) + units := cloneSourceUnits(doc.Units[start : end+1]) + content, err := json.Marshal(struct { + Units []SourceUnit `json:"units"` + }{Units: units}) + if err != nil { + return nil, fmt.Errorf("encode chunk plan range[%d]: %w", index, err) + } + chunks = append(chunks, Chunk{ + ID: fmt.Sprintf("chunk-%06d", index+1), + SourceID: doc.ID, + Index: index, + Ref: SourceRef{SourceID: doc.ID, StartUnitID: chunkRange.StartUnitID, EndUnitID: chunkRange.EndUnitID}, + Content: content, + MediaType: "application/json", + Units: units, + Metadata: map[string]any{ + "start_unit_id": chunkRange.StartUnitID, + "end_unit_id": chunkRange.EndUnitID, + "unit_count": len(units), + }, + Annotations: CloneChunkAnnotations(chunkRange.Annotations), + PlanAnnotations: CloneChunkAnnotations(plan.Annotations), + }) + } + return chunks, nil +} + +func cloneSourceUnits(units []SourceUnit) []SourceUnit { + if len(units) == 0 { + return nil + } + cloned := make([]SourceUnit, len(units)) + for i, unit := range units { + cloned[i] = unit + cloned[i].Metadata = cloneJSONMap(unit.Metadata) + } + return cloned +} + +func cloneJSONMap(values map[string]any) map[string]any { + if len(values) == 0 { + return nil + } + cloned := make(map[string]any, len(values)) + for key, value := range values { + cloned[key] = value + } + return cloned +} diff --git a/internal/core/source/chunk_plan_test.go b/internal/core/source/chunk_plan_test.go new file mode 100644 index 0000000..8f0c9c7 --- /dev/null +++ b/internal/core/source/chunk_plan_test.go @@ -0,0 +1,230 @@ +package source + +import ( + "bytes" + "encoding/json" + "reflect" + "strings" + "testing" +) + +func TestCanonicalizeChunkAnnotations(t *testing.T) { + original := ChunkAnnotations{ + "domain/items": json.RawMessage(` { "z": [3, 2, 1], "a": 1.0 } `), + } + canonical, err := CanonicalizeChunkAnnotations(original) + if err != nil { + t.Fatalf("CanonicalizeChunkAnnotations() error = %v, want nil", err) + } + if got, want := string(canonical["domain/items"]), `{"a":1.0,"z":[3,2,1]}`; got != want { + t.Fatalf("canonical annotation = %q, want %q", got, want) + } + original["domain/items"][0] = '[' + if got := string(canonical["domain/items"]); got != `{"a":1.0,"z":[3,2,1]}` { + t.Fatalf("canonical annotation changed after input mutation: %q", got) + } + canonical["domain/items"][0] = '[' + if original["domain/items"][0] == '[' && bytes.Equal(original["domain/items"], canonical["domain/items"]) { + t.Fatal("input and canonical annotation share value storage") + } +} + +func TestCanonicalizeChunkAnnotationsRejectsInvalidValues(t *testing.T) { + tests := []struct { + name string + annotations ChunkAnnotations + want string + }{ + {name: "blank namespace", annotations: ChunkAnnotations{" \t": json.RawMessage(`true`)}, want: "namespace must not be empty"}, + {name: "untrimmed namespace", annotations: ChunkAnnotations{" items ": json.RawMessage(`true`)}, want: "leading or trailing whitespace"}, + {name: "invalid JSON", annotations: ChunkAnnotations{"items": json.RawMessage(`{"x":`)}, want: "valid JSON"}, + {name: "trailing JSON", annotations: ChunkAnnotations{"items": json.RawMessage(`true false`)}, want: "exactly one JSON value"}, + {name: "non-finite number", annotations: ChunkAnnotations{"items": json.RawMessage(`NaN`)}, want: "valid JSON"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := CanonicalizeChunkAnnotations(tt.annotations) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("CanonicalizeChunkAnnotations() error = %v, want containing %q", err, tt.want) + } + }) + } +} + +func TestValidateChunkAnnotationsRequiresCanonicalJSON(t *testing.T) { + if err := ValidateChunkAnnotations(ChunkAnnotations{"items": json.RawMessage(` {"b":2,"a":1}`)}); err == nil || !strings.Contains(err.Error(), "canonical JSON") { + t.Fatalf("ValidateChunkAnnotations() error = %v, want canonical JSON error", err) + } + if err := ValidateChunkAnnotations(ChunkAnnotations{"items": json.RawMessage(`{"a":1,"b":2}`)}); err != nil { + t.Fatalf("ValidateChunkAnnotations(canonical) error = %v, want nil", err) + } +} + +func TestCloneChunkPlanDoesNotShareAnnotationBytes(t *testing.T) { + plan := validChunkPlan(planDocument()) + cloned := CloneChunkPlan(plan) + cloned.Annotations["plan"][0] = '[' + cloned.Ranges[0].Annotations["range"][0] = '[' + if string(plan.Annotations["plan"]) != `{"value":1}` || string(plan.Ranges[0].Annotations["range"]) != `{"value":2}` { + t.Fatal("CloneChunkPlan() shares annotation value storage") + } +} + +func TestValidateChunkPlanRanges(t *testing.T) { + doc := planDocument() + tests := []struct { + name string + mutate func(*ChunkPlan) + want string + }{ + {name: "source mismatch", mutate: func(plan *ChunkPlan) { plan.SourceDigest = "sha256:other" }, want: "does not match"}, + {name: "missing ranges", mutate: func(plan *ChunkPlan) { plan.Ranges = nil }, want: "ranges must not be empty"}, + {name: "missing start", mutate: func(plan *ChunkPlan) { plan.Ranges[0].StartUnitID = 99 }, want: "start_unit_id 99 was not found"}, + {name: "missing end", mutate: func(plan *ChunkPlan) { plan.Ranges[0].EndUnitID = 99 }, want: "end_unit_id 99 was not found"}, + {name: "backward range", mutate: func(plan *ChunkPlan) { plan.Ranges[0] = ChunkRange{StartUnitID: 30, EndUnitID: 10} }, want: "appears after end_unit_id"}, + {name: "duplicate start", mutate: func(plan *ChunkPlan) { plan.Ranges[1].StartUnitID = plan.Ranges[0].StartUnitID }, want: "does not appear after"}, + {name: "backward starts", mutate: func(plan *ChunkPlan) { + plan.Ranges = []ChunkRange{{StartUnitID: 30, EndUnitID: 50}, {StartUnitID: 20, EndUnitID: 40}} + }, want: "does not appear after"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + plan := validChunkPlan(doc) + tt.mutate(&plan) + err := ValidateChunkPlan(doc, plan) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("ValidateChunkPlan() error = %v, want containing %q", err, tt.want) + } + }) + } + + for name, ranges := range map[string][]ChunkRange{ + "gap": {{StartUnitID: 10, EndUnitID: 20}, {StartUnitID: 40, EndUnitID: 50}}, + "overlap": {{StartUnitID: 10, EndUnitID: 30}, {StartUnitID: 20, EndUnitID: 50}}, + } { + t.Run(name, func(t *testing.T) { + plan := validChunkPlan(doc) + plan.Ranges = ranges + if err := ValidateChunkPlan(doc, plan); err != nil { + t.Fatalf("ValidateChunkPlan() error = %v, want nil", err) + } + }) + } +} + +func TestDigestChunkPlanIsStableAndCoversLogicalPlan(t *testing.T) { + doc := planDocument() + plan := validChunkPlan(doc) + first, err := DigestChunkPlan(plan) + if err != nil { + t.Fatalf("DigestChunkPlan() error = %v, want nil", err) + } + reformatted := CloneChunkPlan(plan) + reformatted.Annotations["plan"] = json.RawMessage(` { "value" : 1 } `) + second, err := DigestChunkPlan(reformatted) + if err != nil { + t.Fatalf("DigestChunkPlan(reformatted) error = %v, want nil", err) + } + if first != second { + t.Fatalf("digests = %q and %q, want stable canonical annotation digest", first, second) + } + + changes := []func(*ChunkPlan){ + func(value *ChunkPlan) { value.Ranges[0].EndUnitID = 30 }, + func(value *ChunkPlan) { value.Annotations["plan"] = json.RawMessage(`{"value":2}`) }, + func(value *ChunkPlan) { value.Ranges[0].Annotations["range"] = json.RawMessage(`{"value":3}`) }, + } + for i, change := range changes { + changed := CloneChunkPlan(plan) + change(&changed) + digest, err := DigestChunkPlan(changed) + if err != nil { + t.Fatalf("DigestChunkPlan(change %d) error = %v", i, err) + } + if digest == first { + t.Fatalf("DigestChunkPlan(change %d) = %q, want changed digest", i, digest) + } + } +} + +func TestMaterializeChunkPlanExactOutputAndMutationSafety(t *testing.T) { + doc := planDocument() + plan := validChunkPlan(doc) + plan.Ranges = []ChunkRange{ + {StartUnitID: 10, EndUnitID: 30, Annotations: ChunkAnnotations{"range": json.RawMessage(`{"value":2}`)}}, + {StartUnitID: 20, EndUnitID: 50, Annotations: ChunkAnnotations{"range": json.RawMessage(`{"value":3}`)}}, + } + chunks, err := MaterializeChunkPlan(doc, plan) + if err != nil { + t.Fatalf("MaterializeChunkPlan() error = %v, want nil", err) + } + if len(chunks) != 2 { + t.Fatalf("chunks = %d, want 2", len(chunks)) + } + first := chunks[0] + if first.ID != "chunk-000001" || first.SourceID != doc.ID || first.Index != 0 || first.MediaType != "application/json" { + t.Fatalf("first chunk identity = %#v", first) + } + if want := (SourceRef{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 30}); first.Ref != want { + t.Fatalf("first ref = %#v, want %#v", first.Ref, want) + } + if got, want := unitIDs(first.Units), []int{10, 20, 30}; !reflect.DeepEqual(got, want) { + t.Fatalf("first unit ids = %#v, want %#v", got, want) + } + wantContent, _ := json.Marshal(struct { + Units []SourceUnit `json:"units"` + }{Units: doc.Units[:3]}) + if !bytes.Equal(first.Content, wantContent) { + t.Fatalf("first content = %s, want %s", first.Content, wantContent) + } + if !reflect.DeepEqual(first.Metadata, map[string]any{"start_unit_id": 10, "end_unit_id": 30, "unit_count": 3}) { + t.Fatalf("first metadata = %#v", first.Metadata) + } + if string(first.Annotations["range"]) != `{"value":2}` || string(first.PlanAnnotations["plan"]) != `{"value":1}` { + t.Fatalf("first annotations = %#v / %#v", first.Annotations, first.PlanAnnotations) + } + if got, want := unitIDs(chunks[1].Units), []int{20, 30, 40, 50}; !reflect.DeepEqual(got, want) { + t.Fatalf("overlapping unit ids = %#v, want %#v", got, want) + } + + again, err := MaterializeChunkPlan(doc, plan) + if err != nil { + t.Fatalf("MaterializeChunkPlan(repeated) error = %v", err) + } + if !reflect.DeepEqual(chunks, again) { + t.Fatalf("repeated materialization differs:\nfirst: %#v\nagain: %#v", chunks, again) + } + chunks[0].Units[0].Text = "mutated" + chunks[0].Annotations["range"][0] = '[' + chunks[0].PlanAnnotations["plan"][0] = '[' + if doc.Units[0].Text == "mutated" || string(plan.Ranges[0].Annotations["range"]) != `{"value":2}` || string(plan.Annotations["plan"]) != `{"value":1}` { + t.Fatal("materialized chunk shares owned plan or source storage") + } +} + +func planDocument() *SourceDocument { + doc := &SourceDocument{ID: "source-plan", Kind: "test", Format: "application/test", Digest: "sha256:source-plan"} + for _, id := range []int{10, 20, 30, 40, 50} { + doc.Units = append(doc.Units, SourceUnit{ID: id, Kind: "line", Text: "unit", Ref: SourceRef{SourceID: doc.ID, StartUnitID: id, EndUnitID: id}}) + } + return doc +} + +func validChunkPlan(doc *SourceDocument) ChunkPlan { + return ChunkPlan{ + SourceDigest: doc.Digest, + Ranges: []ChunkRange{ + {StartUnitID: 10, EndUnitID: 20, Annotations: ChunkAnnotations{"range": json.RawMessage(`{"value":2}`)}}, + {StartUnitID: 30, EndUnitID: 50}, + }, + Annotations: ChunkAnnotations{"plan": json.RawMessage(`{"value":1}`)}, + } +} + +func unitIDs(units []SourceUnit) []int { + ids := make([]int, len(units)) + for i, unit := range units { + ids[i] = unit.ID + } + return ids +} diff --git a/internal/core/source/digest.go b/internal/core/source/digest.go index 80a637a..5e8e252 100644 --- a/internal/core/source/digest.go +++ b/internal/core/source/digest.go @@ -37,24 +37,36 @@ func DigestDocument(doc *SourceDocument) (string, error) { // DigestChunk returns a deterministic digest of a chunk, including its source // provenance, content, units, and metadata. func DigestChunk(chunk Chunk) (string, error) { + annotations, err := CanonicalizeChunkAnnotations(chunk.Annotations) + if err != nil { + return "", fmt.Errorf("canonicalize source chunk annotations: %w", err) + } + planAnnotations, err := CanonicalizeChunkAnnotations(chunk.PlanAnnotations) + if err != nil { + return "", fmt.Errorf("canonicalize source chunk plan annotations: %w", err) + } payload := struct { - ID string `json:"id"` - SourceID string `json:"source_id"` - Index int `json:"index"` - Ref SourceRef `json:"ref"` - Content []byte `json:"content"` - MediaType string `json:"media_type"` - Units []SourceUnit `json:"units"` - Metadata map[string]any `json:"metadata,omitempty"` + ID string `json:"id"` + SourceID string `json:"source_id"` + Index int `json:"index"` + Ref SourceRef `json:"ref"` + Content []byte `json:"content"` + MediaType string `json:"media_type"` + Units []SourceUnit `json:"units"` + Metadata map[string]any `json:"metadata,omitempty"` + Annotations ChunkAnnotations `json:"annotations,omitempty"` + PlanAnnotations ChunkAnnotations `json:"plan_annotations,omitempty"` }{ - ID: chunk.ID, - SourceID: chunk.SourceID, - Index: chunk.Index, - Ref: chunk.Ref, - Content: chunk.Content, - MediaType: chunk.MediaType, - Units: chunk.Units, - Metadata: chunk.Metadata, + ID: chunk.ID, + SourceID: chunk.SourceID, + Index: chunk.Index, + Ref: chunk.Ref, + Content: chunk.Content, + MediaType: chunk.MediaType, + Units: chunk.Units, + Metadata: chunk.Metadata, + Annotations: annotations, + PlanAnnotations: planAnnotations, } encoded, err := json.Marshal(payload) if err != nil { @@ -63,3 +75,28 @@ func DigestChunk(chunk Chunk) (string, error) { sum := sha256.Sum256(encoded) return "sha256:" + hex.EncodeToString(sum[:]), nil } + +// DigestChunkPlan returns a deterministic digest of the logical plan. Storage +// schema, producer provenance, warnings, and timestamps are intentionally not +// part of the digest. +func DigestChunkPlan(plan ChunkPlan) (string, error) { + canonical, err := CanonicalizeChunkPlan(plan) + if err != nil { + return "", err + } + if isBlank(canonical.SourceDigest) { + return "", fmt.Errorf("chunk plan source_digest must not be empty") + } + if hasSurroundingWhitespace(canonical.SourceDigest) { + return "", fmt.Errorf("chunk plan source_digest must not contain leading or trailing whitespace") + } + if len(canonical.Ranges) == 0 { + return "", fmt.Errorf("chunk plan ranges must not be empty") + } + encoded, err := json.Marshal(canonical) + if err != nil { + return "", fmt.Errorf("encode chunk plan for digest: %w", err) + } + sum := sha256.Sum256(encoded) + return "sha256:" + hex.EncodeToString(sum[:]), nil +} diff --git a/internal/core/source/source.go b/internal/core/source/source.go index e39cd87..6b426b4 100644 --- a/internal/core/source/source.go +++ b/internal/core/source/source.go @@ -1,5 +1,7 @@ package source +import "encoding/json" + type SourceDocument struct { ID string `json:"id"` Kind string `json:"kind"` @@ -23,13 +25,29 @@ type SourceRef struct { EndUnitID int `json:"end_unit_id"` } -type Chunk struct { - ID string `json:"id"` - SourceID string `json:"source_id"` - Index int `json:"index"` - Ref SourceRef `json:"ref"` - Content []byte `json:"-"` - MediaType string `json:"media_type"` - Units []SourceUnit `json:"units"` - Metadata map[string]any `json:"metadata,omitempty"` +type ChunkAnnotations map[string]json.RawMessage + +type ChunkPlan struct { + SourceDigest string `json:"source_digest"` + Ranges []ChunkRange `json:"ranges"` + Annotations ChunkAnnotations `json:"annotations,omitempty"` +} + +type ChunkRange struct { + StartUnitID int `json:"start_unit_id"` + EndUnitID int `json:"end_unit_id"` + Annotations ChunkAnnotations `json:"annotations,omitempty"` +} + +type Chunk struct { + ID string `json:"id"` + SourceID string `json:"source_id"` + Index int `json:"index"` + Ref SourceRef `json:"ref"` + Content []byte `json:"-"` + MediaType string `json:"media_type"` + Units []SourceUnit `json:"units"` + Metadata map[string]any `json:"metadata,omitempty"` + Annotations ChunkAnnotations `json:"annotations,omitempty"` + PlanAnnotations ChunkAnnotations `json:"plan_annotations,omitempty"` } diff --git a/internal/core/source/source_test.go b/internal/core/source/source_test.go index 8ee0d3f..0434826 100644 --- a/internal/core/source/source_test.go +++ b/internal/core/source/source_test.go @@ -1,6 +1,7 @@ package source import ( + "encoding/json" "strings" "testing" ) @@ -258,6 +259,28 @@ func TestDigestChunkIsDeterministicAndIncludesReference(t *testing.T) { } } +func TestDigestChunkIncludesAnnotationScopes(t *testing.T) { + doc := validDocument() + chunk := Chunk{ + ID: "chunk-1", SourceID: doc.ID, Ref: SourceRef{SourceID: doc.ID, StartUnitID: 1, EndUnitID: 2}, + Content: []byte("content"), MediaType: "text/plain", Units: doc.Units, + Annotations: ChunkAnnotations{"scope": json.RawMessage(`{"value":1}`)}, + PlanAnnotations: ChunkAnnotations{"scope": json.RawMessage(`{"value":2}`)}, + } + base, err := DigestChunk(chunk) + if err != nil { + t.Fatalf("DigestChunk() error = %v", err) + } + chunk.Annotations["scope"] = json.RawMessage(`{"value":3}`) + rangeChanged, _ := DigestChunk(chunk) + chunk.Annotations["scope"] = json.RawMessage(`{"value":1}`) + chunk.PlanAnnotations["scope"] = json.RawMessage(`{"value":3}`) + planChanged, _ := DigestChunk(chunk) + if base == rangeChanged || base == planChanged || rangeChanged == planChanged { + t.Fatalf("annotation scope digests did not change distinctly: %q %q %q", base, rangeChanged, planChanged) + } +} + func TestValidateRefValid(t *testing.T) { doc := validDocument() ref := SourceRef{ diff --git a/internal/framework/checkpoint/loader.go b/internal/framework/checkpoint/loader.go index e997e64..a11bbe4 100644 --- a/internal/framework/checkpoint/loader.go +++ b/internal/framework/checkpoint/loader.go @@ -242,15 +242,25 @@ func sourceChunksFromEnvelope(values []chunkEnvelope) ([]source.Chunk, error) { if err != nil { return nil, err } + annotations, err := source.CanonicalizeChunkAnnotations(value.Annotations) + if err != nil { + return nil, fmt.Errorf("canonicalize chunk %q annotations: %w", value.ID, err) + } + planAnnotations, err := source.CanonicalizeChunkAnnotations(value.PlanAnnotations) + if err != nil { + return nil, fmt.Errorf("canonicalize chunk %q plan_annotations: %w", value.ID, err) + } out = append(out, source.Chunk{ - ID: value.ID, - SourceID: value.SourceID, - Index: value.Index, - Ref: value.Ref, - Content: content, - MediaType: value.Content.MediaType, - Units: cloneSourceUnits(value.Units), - Metadata: cloneMetadata(value.Metadata), + ID: value.ID, + SourceID: value.SourceID, + Index: value.Index, + Ref: value.Ref, + Content: content, + MediaType: value.Content.MediaType, + Units: cloneSourceUnits(value.Units), + Metadata: cloneMetadata(value.Metadata), + Annotations: annotations, + PlanAnnotations: planAnnotations, }) } return out, nil diff --git a/internal/framework/checkpoint/recorder.go b/internal/framework/checkpoint/recorder.go index ebaa60a..1a12732 100644 --- a/internal/framework/checkpoint/recorder.go +++ b/internal/framework/checkpoint/recorder.go @@ -251,13 +251,15 @@ type chunksEnvelope struct { } type chunkEnvelope struct { - ID string `json:"id"` - SourceID string `json:"source_id"` - Index int `json:"index"` - Ref source.SourceRef `json:"ref"` - Content binaryEnvelope `json:"content"` - Units []source.SourceUnit `json:"units,omitempty"` - Metadata map[string]any `json:"metadata,omitempty"` + ID string `json:"id"` + SourceID string `json:"source_id"` + Index int `json:"index"` + Ref source.SourceRef `json:"ref"` + Content binaryEnvelope `json:"content"` + Units []source.SourceUnit `json:"units,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` + Annotations source.ChunkAnnotations `json:"annotations,omitempty"` + PlanAnnotations source.ChunkAnnotations `json:"plan_annotations,omitempty"` } type binaryEnvelope struct { @@ -320,13 +322,15 @@ func chunkEnvelopes(chunks []source.Chunk) []chunkEnvelope { out := make([]chunkEnvelope, 0, len(chunks)) for _, chunk := range chunks { out = append(out, chunkEnvelope{ - ID: chunk.ID, - SourceID: chunk.SourceID, - Index: chunk.Index, - Ref: chunk.Ref, - Content: binaryEnvelopeFromContent(chunk.Content, chunk.MediaType, chunk.Metadata, nil), - Units: cloneSourceUnits(chunk.Units), - Metadata: cloneMetadata(chunk.Metadata), + ID: chunk.ID, + SourceID: chunk.SourceID, + Index: chunk.Index, + Ref: chunk.Ref, + Content: binaryEnvelopeFromContent(chunk.Content, chunk.MediaType, chunk.Metadata, nil), + Units: cloneSourceUnits(chunk.Units), + Metadata: cloneMetadata(chunk.Metadata), + Annotations: source.CloneChunkAnnotations(chunk.Annotations), + PlanAnnotations: source.CloneChunkAnnotations(chunk.PlanAnnotations), }) } return out diff --git a/internal/framework/checkpoint/recorder_test.go b/internal/framework/checkpoint/recorder_test.go index cfe1794..3003aad 100644 --- a/internal/framework/checkpoint/recorder_test.go +++ b/internal/framework/checkpoint/recorder_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "os" "path/filepath" + "reflect" "strings" "testing" @@ -33,6 +34,12 @@ func TestWorkspaceRecorderWritesSuccessfulCheckpointFiles(t *testing.T) { Content: []byte("chunk content"), MediaType: "text/plain", Units: doc.Units, + Annotations: source.ChunkAnnotations{ + "same": json.RawMessage(`{"range":1}`), + }, + PlanAnnotations: source.ChunkAnnotations{ + "same": json.RawMessage(`{"plan":2}`), + }, }, } @@ -57,7 +64,9 @@ func TestWorkspaceRecorderWritesSuccessfulCheckpointFiles(t *testing.T) { assertManifestStatus(t, filepath.Join(root, "chunk", "manifest.json"), coreworkspace.StatusSucceeded) var chunkPayload struct { Chunks []struct { - Content struct { + Annotations source.ChunkAnnotations `json:"annotations"` + PlanAnnotations source.ChunkAnnotations `json:"plan_annotations"` + Content struct { ContentBase64 string `json:"content_base64"` ContentDigest string `json:"content_digest"` } `json:"content"` @@ -77,6 +86,28 @@ func TestWorkspaceRecorderWritesSuccessfulCheckpointFiles(t *testing.T) { if got, want := chunkPayload.Chunks[0].Content.ContentDigest, contentDigest([]byte("chunk content")); got != want { t.Fatalf("content digest = %q, want %q", got, want) } + if !jsonEqual(t, chunkPayload.Chunks[0].Annotations["same"], json.RawMessage(`{"range":1}`)) || !jsonEqual(t, chunkPayload.Chunks[0].PlanAnnotations["same"], json.RawMessage(`{"plan":2}`)) { + t.Fatalf("checkpoint annotations = %s / %s", chunkPayload.Chunks[0].Annotations["same"], chunkPayload.Chunks[0].PlanAnnotations["same"]) + } + loaded, decision := (&WorkspaceLoader{root: root}).Chunk("generic", doc.Digest) + if !decision.Reused || len(loaded.Chunks) != 1 { + t.Fatalf("loaded chunk checkpoint = %#v, decision = %#v", loaded, decision) + } + if string(loaded.Chunks[0].Annotations["same"]) != `{"range":1}` || string(loaded.Chunks[0].PlanAnnotations["same"]) != `{"plan":2}` { + t.Fatalf("loaded canonical annotations = %s / %s", loaded.Chunks[0].Annotations["same"], loaded.Chunks[0].PlanAnnotations["same"]) + } +} + +func jsonEqual(t *testing.T, left, right []byte) bool { + t.Helper() + var leftValue, rightValue any + if err := json.Unmarshal(left, &leftValue); err != nil { + t.Fatalf("decode left JSON: %v", err) + } + if err := json.Unmarshal(right, &rightValue); err != nil { + t.Fatalf("decode right JSON: %v", err) + } + return reflect.DeepEqual(leftValue, rightValue) } func TestWorkspaceArtifactCheckpointsRoundTripCodecIdentityAndBytes(t *testing.T) { diff --git a/internal/framework/pipeline/chunk_annotations_test.go b/internal/framework/pipeline/chunk_annotations_test.go new file mode 100644 index 0000000..f129526 --- /dev/null +++ b/internal/framework/pipeline/chunk_annotations_test.go @@ -0,0 +1,70 @@ +package pipeline + +import ( + "encoding/json" + "strings" + "testing" + + "gitea.maximumdirect.net/eric/notarius/internal/core/source" +) + +func TestChunkCanonicalizationAndClonePreserveAnnotationScopes(t *testing.T) { + doc := validSourceDocument() + chunk := source.Chunk{ + ID: "chunk-1", SourceID: doc.ID, Index: 0, + Ref: doc.Units[0].Ref, Content: []byte(`{"units":[1]}`), MediaType: "application/json", Units: doc.Units[:1], + Annotations: source.ChunkAnnotations{"shared": json.RawMessage(` {"range": true} `)}, + PlanAnnotations: source.ChunkAnnotations{"shared": json.RawMessage(` {"plan": true} `)}, + } + canonical, err := validateAndCanonicalizeChunkResult(doc, []source.Chunk{chunk}) + if err != nil { + t.Fatalf("validateAndCanonicalizeChunkResult() error = %v", err) + } + if got := string(canonical[0].Annotations["shared"]); got != `{"range":true}` { + t.Fatalf("range annotation = %q", got) + } + if got := string(canonical[0].PlanAnnotations["shared"]); got != `{"plan":true}` { + t.Fatalf("plan annotation = %q", got) + } + + cloned := cloneSourceChunk(canonical[0]) + cloned.Annotations["shared"][0] = '[' + cloned.PlanAnnotations["shared"][0] = '[' + if string(canonical[0].Annotations["shared"]) != `{"range":true}` || string(canonical[0].PlanAnnotations["shared"]) != `{"plan":true}` { + t.Fatal("cloneSourceChunk() shares annotation bytes") + } +} + +func TestSerializedChunkValidationRepresentationIncludesAnnotationScopes(t *testing.T) { + doc := validSourceDocument() + chunks := []source.Chunk{{ + ID: "chunk-1", SourceID: doc.ID, Ref: doc.Units[0].Ref, MediaType: "application/json", Units: doc.Units[:1], + Annotations: source.ChunkAnnotations{"same": json.RawMessage(`{"range":1}`)}, + PlanAnnotations: source.ChunkAnnotations{"same": json.RawMessage(`{"plan":2}`)}, + }} + encoded, err := json.Marshal(chunks) + if err != nil { + t.Fatalf("json.Marshal(chunks) error = %v", err) + } + got := string(encoded) + if !strings.Contains(got, `"annotations":{"same":{"range":1}}`) || !strings.Contains(got, `"plan_annotations":{"same":{"plan":2}}`) { + t.Fatalf("serialized chunks = %s, want both annotation scopes", got) + } +} + +func TestDebugChunkEnvelopeClonesAnnotationScopes(t *testing.T) { + chunk := source.Chunk{ + ID: "chunk-1", SourceID: "source-1", + Annotations: source.ChunkAnnotations{"same": json.RawMessage(`{"range":1}`)}, + PlanAnnotations: source.ChunkAnnotations{"same": json.RawMessage(`{"plan":2}`)}, + } + envelope := debugSourceChunkEnvelope(chunk) + if string(envelope.Annotations["same"]) != `{"range":1}` || string(envelope.PlanAnnotations["same"]) != `{"plan":2}` { + t.Fatalf("debug annotations = %#v / %#v", envelope.Annotations, envelope.PlanAnnotations) + } + envelope.Annotations["same"][0] = '[' + envelope.PlanAnnotations["same"][0] = '[' + if string(chunk.Annotations["same"]) != `{"range":1}` || string(chunk.PlanAnnotations["same"]) != `{"plan":2}` { + t.Fatal("debugSourceChunkEnvelope() shares annotation bytes") + } +} diff --git a/internal/framework/pipeline/chunk_validation.go b/internal/framework/pipeline/chunk_validation.go index 1baa952..0567090 100644 --- a/internal/framework/pipeline/chunk_validation.go +++ b/internal/framework/pipeline/chunk_validation.go @@ -48,6 +48,14 @@ func validateAndCanonicalizeChunkResult(doc *source.SourceDocument, chunks []sou if err := source.ValidateRef(doc, chunk.Ref); err != nil { return nil, fmt.Errorf("chunk %q ref: %w", chunk.ID, err) } + annotations, err := source.CanonicalizeChunkAnnotations(chunk.Annotations) + if err != nil { + return nil, fmt.Errorf("chunk %q annotations: %w", chunk.ID, err) + } + planAnnotations, err := source.CanonicalizeChunkAnnotations(chunk.PlanAnnotations) + if err != nil { + return nil, fmt.Errorf("chunk %q plan_annotations: %w", chunk.ID, err) + } seenUnitIDs := make(map[int]struct{}, len(chunk.Units)) previousSourceIndex := -1 @@ -84,14 +92,16 @@ func validateAndCanonicalizeChunkResult(doc *source.SourceDocument, chunks []sou } canonicalChunks = append(canonicalChunks, source.Chunk{ - ID: chunk.ID, - SourceID: chunk.SourceID, - Index: chunk.Index, - Ref: expectedRef, - Content: append([]byte(nil), chunk.Content...), - MediaType: chunk.MediaType, - Units: canonicalUnits, - Metadata: cloneMetadata(chunk.Metadata), + ID: chunk.ID, + SourceID: chunk.SourceID, + Index: chunk.Index, + Ref: expectedRef, + Content: append([]byte(nil), chunk.Content...), + MediaType: chunk.MediaType, + Units: canonicalUnits, + Metadata: cloneMetadata(chunk.Metadata), + Annotations: annotations, + PlanAnnotations: planAnnotations, }) } diff --git a/internal/framework/pipeline/debug.go b/internal/framework/pipeline/debug.go index 0d1a8fb..f15522c 100644 --- a/internal/framework/pipeline/debug.go +++ b/internal/framework/pipeline/debug.go @@ -100,13 +100,15 @@ type debugSourceDocument struct { } type debugSourceChunk struct { - ID string `json:"id"` - SourceID string `json:"source_id"` - Index int `json:"index"` - Ref source.SourceRef `json:"ref"` - Content debugBinaryEnvelope `json:"content"` - Units []source.SourceUnit `json:"units,omitempty"` - Metadata map[string]any `json:"metadata,omitempty"` + ID string `json:"id"` + SourceID string `json:"source_id"` + Index int `json:"index"` + Ref source.SourceRef `json:"ref"` + Content debugBinaryEnvelope `json:"content"` + Units []source.SourceUnit `json:"units,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` + Annotations source.ChunkAnnotations `json:"annotations,omitempty"` + PlanAnnotations source.ChunkAnnotations `json:"plan_annotations,omitempty"` } type debugSerializedOutput struct { @@ -464,13 +466,15 @@ func debugSourceDocumentEnvelope(doc *source.SourceDocument) *debugSourceDocumen func debugSourceChunkEnvelope(chunk source.Chunk) debugSourceChunk { return debugSourceChunk{ - ID: chunk.ID, - SourceID: chunk.SourceID, - Index: chunk.Index, - Ref: chunk.Ref, - Content: debugContentEnvelope(chunk.Content, chunk.MediaType, chunk.Metadata, nil), - Units: cloneSourceUnits(chunk.Units), - Metadata: redactSensitiveMap(chunk.Metadata), + ID: chunk.ID, + SourceID: chunk.SourceID, + Index: chunk.Index, + Ref: chunk.Ref, + Content: debugContentEnvelope(chunk.Content, chunk.MediaType, chunk.Metadata, nil), + Units: cloneSourceUnits(chunk.Units), + Metadata: redactSensitiveMap(chunk.Metadata), + Annotations: source.CloneChunkAnnotations(chunk.Annotations), + PlanAnnotations: source.CloneChunkAnnotations(chunk.PlanAnnotations), } } diff --git a/internal/framework/pipeline/runner.go b/internal/framework/pipeline/runner.go index c786faa..dac26a4 100644 --- a/internal/framework/pipeline/runner.go +++ b/internal/framework/pipeline/runner.go @@ -863,6 +863,8 @@ func cloneSourceChunk(chunk source.Chunk) source.Chunk { chunk.Content = append([]byte(nil), chunk.Content...) chunk.Units = cloneSourceUnits(chunk.Units) chunk.Metadata = cloneMetadata(chunk.Metadata) + chunk.Annotations = source.CloneChunkAnnotations(chunk.Annotations) + chunk.PlanAnnotations = source.CloneChunkAnnotations(chunk.PlanAnnotations) return chunk } diff --git a/internal/framework/pipeline/runner_concurrent.go b/internal/framework/pipeline/runner_concurrent.go index a8f0c19..ea5f2c9 100644 --- a/internal/framework/pipeline/runner_concurrent.go +++ b/internal/framework/pipeline/runner_concurrent.go @@ -258,7 +258,7 @@ func prepareLaneExtract(input RunInput, loader CheckpointLoader, doc *source.Sou } func (r *Runner) runExtractJob(ctx context.Context, input RunInput, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, job extractJob) extractJobResult { - state, chunk := job.lane, job.chunk + state, chunk := job.lane, cloneSourceChunk(job.chunk) lane, typed := state.prepared.resolved, state.prepared.typed result := extractJobResult{laneIndex: state.index, chunkIndex: chunk.Index} var accepted erasedExtractArtifact diff --git a/internal/framework/pipeline/runner_extract_handoff_test.go b/internal/framework/pipeline/runner_extract_handoff_test.go index 4b90fd3..eb19d45 100644 --- a/internal/framework/pipeline/runner_extract_handoff_test.go +++ b/internal/framework/pipeline/runner_extract_handoff_test.go @@ -2,10 +2,12 @@ package pipeline import ( "context" + "encoding/json" "reflect" "strings" "testing" + "gitea.maximumdirect.net/eric/notarius/internal/core/source" "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" ) @@ -50,6 +52,31 @@ func cloneCheckpointArtifacts(values []CheckpointArtifact) []CheckpointArtifact return cloned } +func TestRunnerPassesIndependentAnnotationScopesToExtractors(t *testing.T) { + prepared := preparedAttemptDebugPipeline(t) + chunker := prepared.chunker.(*typedTestChunker) + chunker.chunks[0].Annotations = source.ChunkAnnotations{"same": json.RawMessage(`{"range":1}`)} + chunker.chunks[0].PlanAnnotations = source.ChunkAnnotations{"same": json.RawMessage(`{"plan":2}`)} + + var gotRange, gotPlan string + installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) { + gotRange = string(request.Chunk.Annotations["same"]) + gotPlan = string(request.Chunk.PlanAnnotations["same"]) + request.Chunk.Annotations["same"][0] = '[' + request.Chunk.PlanAnnotations["same"][0] = '[' + return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil + }) + if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } + if gotRange != `{"range":1}` || gotPlan != `{"plan":2}` { + t.Fatalf("extract annotations = %q / %q", gotRange, gotPlan) + } + if string(chunker.chunks[0].Annotations["same"]) != `{"range":1}` || string(chunker.chunks[0].PlanAnnotations["same"]) != `{"plan":2}` { + t.Fatal("extract request shares annotation bytes with chunker output") + } +} + func TestRunnerContinuesFromFreshAndReusedExtractResults(t *testing.T) { prepared := preparedAttemptDebugPipeline(t) extractCalls := 0