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

Add canonical chunk plan source model

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-17 23:47:22 +00:00
parent 1c13e1d64a
commit 3bfac14397
16 changed files with 755 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
internal/framework/checkpoint/loader.go
View File

@@ -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

32
internal/framework/checkpoint/recorder.go
View File

@@ -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

33
internal/framework/checkpoint/recorder_test.go
View File

@@ -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) {

70
internal/framework/pipeline/chunk_annotations_test.go Normal file
View File

@@ -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")
}
}

26
internal/framework/pipeline/chunk_validation.go
View File

@@ -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,
})
}

32
internal/framework/pipeline/debug.go
View File

@@ -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),
}
}

2
internal/framework/pipeline/runner.go
View File

@@ -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
}

2
internal/framework/pipeline/runner_concurrent.go
View File

@@ -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

27
internal/framework/pipeline/runner_extract_handoff_test.go
View File

@@ -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