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

Write workspace checkpoints during runs

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-08 02:46:38 +00:00
parent f044c00a7c
commit 1d3a444df8
11 changed files with 1258 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

161
internal/framework/pipeline/checkpoint.go Normal file
View File

@@ -0,0 +1,161 @@
package pipeline
import (
"crypto/sha256"
"encoding/hex"
"fmt"
"sort"
"strings"
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)
type CheckpointFingerprint struct {
Name string `json:"name"`
Value string `json:"value"`
}
type CheckpointRecorder interface {
SourceRunning(moduleKey string) error
SourceSucceeded(moduleKey string, doc *source.SourceDocument) error
SourceFailed(moduleKey string, err error) error
ChunkRunning(moduleKey string, sourceDigest string) error
ChunkSucceeded(moduleKey string, sourceDigest string, chunks []contracts.SourceChunk, warnings []contracts.Warning) error
ChunkRejected(moduleKey string, sourceDigest string, rejected contracts.RejectedOutput) error
ChunkFailed(moduleKey string, sourceDigest string, err error) error
ExtractRunning(laneID string, moduleKey string, dependencies []CheckpointFingerprint) error
ExtractSucceeded(laneID string, moduleKey string, dependencies []CheckpointFingerprint, outputs []contracts.ExtractOutput, rejected []contracts.RejectedOutput, warnings []contracts.Warning) error
ExtractFailed(laneID string, moduleKey string, dependencies []CheckpointFingerprint, err error) error
MergeRunning(laneID string, moduleKey string, dependencies []CheckpointFingerprint) error
MergeSucceeded(laneID string, moduleKey string, dependencies []CheckpointFingerprint, output contracts.MergeOutput, warnings []contracts.Warning) error
MergeRejected(laneID string, moduleKey string, dependencies []CheckpointFingerprint, rejected contracts.RejectedOutput) error
MergeFailed(laneID string, moduleKey string, dependencies []CheckpointFingerprint, err error) error
NormalizeRunning(laneID string, moduleKey string, dependencies []CheckpointFingerprint) error
NormalizeSucceeded(laneID string, moduleKey string, dependencies []CheckpointFingerprint, output contracts.NormalizeOutput, warnings []contracts.Warning) error
NormalizeRejected(laneID string, moduleKey string, dependencies []CheckpointFingerprint, rejected contracts.RejectedOutput) error
NormalizeFailed(laneID string, moduleKey string, dependencies []CheckpointFingerprint, err error) error
}
type noopCheckpointRecorder struct{}
func NoopCheckpointRecorder() CheckpointRecorder { return noopCheckpointRecorder{} }
func (noopCheckpointRecorder) SourceRunning(string) error { return nil }
func (noopCheckpointRecorder) SourceSucceeded(string, *source.SourceDocument) error { return nil }
func (noopCheckpointRecorder) SourceFailed(string, error) error { return nil }
func (noopCheckpointRecorder) ChunkRunning(string, string) error { return nil }
func (noopCheckpointRecorder) ChunkSucceeded(string, string, []contracts.SourceChunk, []contracts.Warning) error {
return nil
}
func (noopCheckpointRecorder) ChunkRejected(string, string, contracts.RejectedOutput) error {
return nil
}
func (noopCheckpointRecorder) ChunkFailed(string, string, error) error { return nil }
func (noopCheckpointRecorder) ExtractRunning(string, string, []CheckpointFingerprint) error {
return nil
}
func (noopCheckpointRecorder) ExtractSucceeded(string, string, []CheckpointFingerprint, []contracts.ExtractOutput, []contracts.RejectedOutput, []contracts.Warning) error {
return nil
}
func (noopCheckpointRecorder) ExtractFailed(string, string, []CheckpointFingerprint, error) error {
return nil
}
func (noopCheckpointRecorder) MergeRunning(string, string, []CheckpointFingerprint) error { return nil }
func (noopCheckpointRecorder) MergeSucceeded(string, string, []CheckpointFingerprint, contracts.MergeOutput, []contracts.Warning) error {
return nil
}
func (noopCheckpointRecorder) MergeRejected(string, string, []CheckpointFingerprint, contracts.RejectedOutput) error {
return nil
}
func (noopCheckpointRecorder) MergeFailed(string, string, []CheckpointFingerprint, error) error {
return nil
}
func (noopCheckpointRecorder) NormalizeRunning(string, string, []CheckpointFingerprint) error {
return nil
}
func (noopCheckpointRecorder) NormalizeSucceeded(string, string, []CheckpointFingerprint, contracts.NormalizeOutput, []contracts.Warning) error {
return nil
}
func (noopCheckpointRecorder) NormalizeRejected(string, string, []CheckpointFingerprint, contracts.RejectedOutput) error {
return nil
}
func (noopCheckpointRecorder) NormalizeFailed(string, string, []CheckpointFingerprint, error) error {
return nil
}
func rawOutputDigests(payloads []contracts.RawPayload) []CheckpointFingerprint {
values := make([]CheckpointFingerprint, 0, len(payloads))
for i, payload := range payloads {
values = append(values, CheckpointFingerprint{
Name: fmt.Sprintf("payload[%d]", i),
Value: checkpointContentDigest(payload.Content),
})
}
return normalizeCheckpointFingerprints(values)
}
func extractPayloads(outputs []contracts.ExtractOutput) []contracts.RawPayload {
if len(outputs) == 0 {
return nil
}
payloads := make([]contracts.RawPayload, 0, len(outputs))
for _, output := range outputs {
payloads = append(payloads, output.Payload)
}
return payloads
}
func digestFingerprints(name string, digest string) []CheckpointFingerprint {
digest = strings.TrimSpace(digest)
if digest == "" {
return nil
}
return []CheckpointFingerprint{{Name: name, Value: digest}}
}
func joinedChunkDigest(chunks []contracts.SourceChunk) string {
if len(chunks) == 0 {
return ""
}
values := make([]string, 0, len(chunks))
for _, chunk := range chunks {
values = append(values, chunk.ID+"="+checkpointContentDigest(chunk.Content))
}
sort.Strings(values)
sum := sha256.Sum256([]byte(strings.Join(values, "\n")))
return "sha256:" + hex.EncodeToString(sum[:])
}
func normalizeCheckpointFingerprints(values []CheckpointFingerprint) []CheckpointFingerprint {
if len(values) == 0 {
return nil
}
byName := make(map[string]string, len(values))
for _, value := range values {
name := strings.TrimSpace(value.Name)
fingerprint := strings.TrimSpace(value.Value)
if name == "" || fingerprint == "" {
continue
}
byName[name] = fingerprint
}
if len(byName) == 0 {
return nil
}
names := make([]string, 0, len(byName))
for name := range byName {
names = append(names, name)
}
sort.Strings(names)
out := make([]CheckpointFingerprint, 0, len(names))
for _, name := range names {
out = append(out, CheckpointFingerprint{Name: name, Value: byName[name]})
}
return out
}
func checkpointContentDigest(content []byte) string {
sum := sha256.Sum256(content)
return "sha256:" + hex.EncodeToString(sum[:])
}

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

@@ -48,6 +48,7 @@ type RunInput struct {
LLMProfiles []artifacts.LLMProfileManifest
Metadata map[string]any
Warnings []contracts.Warning
Checkpoints CheckpointRecorder
}
type RunOutput struct {
@@ -70,6 +71,10 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
}
output.Manifest = manifestFromPipeline(input)
checkpoints := input.Checkpoints
if checkpoints == nil {
checkpoints = NoopCheckpointRecorder()
}
defer func() {
output.Manifest.LLMProfiles = mergeLLMProfileManifests(input.LLMProfiles, llmProfileManifests(input.LLMClient))
}()
@@ -80,6 +85,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
return failOutput(output), fmt.Errorf("build input adapter %q: %w", input.Pipeline.Input.Module, err)
}
attachModuleManifestMetadata(&output, "input", adapter)
if err := checkpoints.SourceRunning(adapter.Key()); err != nil {
return failOutput(output), fmt.Errorf("write source checkpoint: %w", err)
}
doc, err := adapter.Parse(ctx, contracts.ParseRequest{
SourceID: input.SourceID,
Path: input.Path,
@@ -89,11 +97,16 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
Metadata: input.Metadata,
})
if err != nil {
_ = checkpoints.SourceFailed(adapter.Key(), err)
return failOutput(output), fmt.Errorf("parse input with adapter %q: %w", adapter.Key(), err)
}
if err := source.ValidateDocument(doc); err != nil {
_ = checkpoints.SourceFailed(adapter.Key(), err)
return failOutput(output), fmt.Errorf("validate source document: %w", err)
}
if err := checkpoints.SourceSucceeded(adapter.Key(), doc); err != nil {
return failOutput(output), fmt.Errorf("write source checkpoint: %w", err)
}
sourceInput := sourceInputMaterial(input.Path, input.RawInput)
sessionID := resolvedSessionID(input.SessionID, doc.ID)
output.Manifest.Metadata = manifestMetadataWithSessionID(output.Manifest.Metadata, sessionID)
@@ -104,6 +117,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
return failOutput(output), fmt.Errorf("build chunker %q: %w", input.Pipeline.Chunk.Module, err)
}
attachModuleManifestMetadata(&output, "chunker", chunker)
if err := checkpoints.ChunkRunning(chunker.Key(), doc.Digest); err != nil {
return failOutput(output), fmt.Errorf("write chunk checkpoint: %w", err)
}
var canonicalChunks []contracts.SourceChunk
var chunkWarnings []contracts.Warning
chunksAccepted, chunkRejection, err := runWithRetry(ctx, input.Pipeline.Chunk.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
@@ -136,17 +152,24 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
return true, nil, nil
})
if err != nil {
_ = checkpoints.ChunkFailed(chunker.Key(), doc.Digest, err)
return failOutput(output), err
}
if !chunksAccepted {
output.Rejected = append(output.Rejected, *chunkRejection)
if err := checkpoints.ChunkRejected(chunker.Key(), doc.Digest, *chunkRejection); err != nil {
return failOutput(output), fmt.Errorf("write chunk checkpoint: %w", err)
}
} else {
output.Warnings = append(output.Warnings, chunkWarnings...)
if err := checkpoints.ChunkSucceeded(chunker.Key(), doc.Digest, canonicalChunks, chunkWarnings); err != nil {
return failOutput(output), fmt.Errorf("write chunk checkpoint: %w", err)
}
}
if chunksAccepted {
for _, lane := range input.Pipeline.ArtifactLanes {
if err := r.runLane(ctx, input, doc, sourceInput, sessionID, canonicalChunks, lane, &output); err != nil {
if err := r.runLane(ctx, input, checkpoints, doc, sourceInput, sessionID, canonicalChunks, lane, &output); err != nil {
return failOutput(output), err
}
}
@@ -187,7 +210,7 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
return output, nil
}
func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []contracts.SourceChunk, lane ResolvedArtifactLane, output *RunOutput) error {
func (r *Runner) runLane(ctx context.Context, input RunInput, checkpoints CheckpointRecorder, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []contracts.SourceChunk, lane ResolvedArtifactLane, output *RunOutput) error {
extractor, err := r.registries.Extractors.Build(lane.Extract.Module)
if err != nil {
return fmt.Errorf("build extractor %q for lane %q: %w", lane.Extract.Module, lane.ID, err)
@@ -203,6 +226,12 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
setLaneManifestMetadata(output, lane.ID, extractor, merger, normalizer)
extractOutputs := make([]contracts.ExtractOutput, 0, len(chunks))
extractWarnings := []contracts.Warning{}
extractRejectedStart := len(output.Rejected)
extractDependencies := digestFingerprints("chunks", joinedChunkDigest(chunks))
if err := checkpoints.ExtractRunning(lane.ID, extractor.Key(), extractDependencies); err != nil {
return fmt.Errorf("write extract checkpoint for lane %q: %w", lane.ID, err)
}
for index := range chunks {
chunk := chunks[index]
var acceptedOutput contracts.ExtractOutput
@@ -256,6 +285,7 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
return true, nil, nil
})
if err != nil {
_ = checkpoints.ExtractFailed(lane.ID, extractor.Key(), extractDependencies, err)
return err
}
if !accepted {
@@ -263,8 +293,13 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
continue
}
output.Warnings = append(output.Warnings, acceptedWarnings...)
extractWarnings = append(extractWarnings, acceptedWarnings...)
extractOutputs = append(extractOutputs, acceptedOutput)
}
extractRejected := cloneRejectedOutputs(output.Rejected[extractRejectedStart:])
if err := checkpoints.ExtractSucceeded(lane.ID, extractor.Key(), extractDependencies, extractOutputs, extractRejected, extractWarnings); err != nil {
return fmt.Errorf("write extract checkpoint for lane %q: %w", lane.ID, err)
}
if len(extractOutputs) == 0 {
return nil
@@ -272,6 +307,10 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
var acceptedMerge contracts.MergeOutput
var mergeWarnings []contracts.Warning
mergeDependencies := rawOutputDigests(extractPayloads(extractOutputs))
if err := checkpoints.MergeRunning(lane.ID, merger.Key(), mergeDependencies); err != nil {
return fmt.Errorf("write merge checkpoint for lane %q: %w", lane.ID, err)
}
mergeAccepted, mergeRejection, err := runWithRetry(ctx, lane.Merge.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
mergeResult, err := merger.Merge(ctx, contracts.MergeRequest{
Source: doc,
@@ -318,16 +357,27 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
return true, nil, nil
})
if err != nil {
_ = checkpoints.MergeFailed(lane.ID, merger.Key(), mergeDependencies, err)
return err
}
if !mergeAccepted {
output.Rejected = append(output.Rejected, *mergeRejection)
if err := checkpoints.MergeRejected(lane.ID, merger.Key(), mergeDependencies, *mergeRejection); err != nil {
return fmt.Errorf("write merge checkpoint for lane %q: %w", lane.ID, err)
}
return nil
}
output.Warnings = append(output.Warnings, mergeWarnings...)
if err := checkpoints.MergeSucceeded(lane.ID, merger.Key(), mergeDependencies, acceptedMerge, mergeWarnings); err != nil {
return fmt.Errorf("write merge checkpoint for lane %q: %w", lane.ID, err)
}
var acceptedNormalize contracts.NormalizeOutput
var normalizeWarnings []contracts.Warning
normalizeDependencies := rawOutputDigests([]contracts.RawPayload{acceptedMerge.Payload})
if err := checkpoints.NormalizeRunning(lane.ID, normalizer.Key(), normalizeDependencies); err != nil {
return fmt.Errorf("write normalize checkpoint for lane %q: %w", lane.ID, err)
}
normalizeAccepted, normalizeRejection, err := runWithRetry(ctx, lane.Normalize.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
normalizeResult, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
Source: doc,
@@ -374,13 +424,20 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
return true, nil, nil
})
if err != nil {
_ = checkpoints.NormalizeFailed(lane.ID, normalizer.Key(), normalizeDependencies, err)
return err
}
if !normalizeAccepted {
output.Rejected = append(output.Rejected, *normalizeRejection)
if err := checkpoints.NormalizeRejected(lane.ID, normalizer.Key(), normalizeDependencies, *normalizeRejection); err != nil {
return fmt.Errorf("write normalize checkpoint for lane %q: %w", lane.ID, err)
}
return nil
}
output.Warnings = append(output.Warnings, normalizeWarnings...)
if err := checkpoints.NormalizeSucceeded(lane.ID, normalizer.Key(), normalizeDependencies, acceptedNormalize, normalizeWarnings); err != nil {
return fmt.Errorf("write normalize checkpoint for lane %q: %w", lane.ID, err)
}
output.NormalizeOutputs = append(output.NormalizeOutputs, acceptedNormalize)
return nil
}

50
internal/framework/pipeline/runner_test.go
View File

@@ -992,6 +992,37 @@ func TestRunPassesChunkContentAndMediaTypeToExtractors(t *testing.T) {
}
}
func TestRunDoesNotPassCheckpointPathsToModules(t *testing.T) {
modules := defaultRunnerModules()
_, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{
Pipeline: resolvedPipeline(),
Checkpoints: NoopCheckpointRecorder(),
})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
for _, req := range modules.input.requests {
assertNoCheckpointMetadata(t, req.Metadata)
}
for _, req := range modules.chunker.requests {
assertNoCheckpointMetadata(t, req.Metadata)
}
for _, req := range modules.extractors["extract-alpha"].requests {
assertNoCheckpointMetadata(t, req.Metadata)
}
for _, req := range modules.mergers["merge"].requests {
assertNoCheckpointMetadata(t, req.Metadata)
}
for _, req := range modules.normalizers["normalize"].requests {
assertNoCheckpointMetadata(t, req.Metadata)
}
for _, req := range modules.output.requests {
assertNoCheckpointMetadata(t, req.Metadata)
}
}
func TestRunOmitsRejectedExtractOutputsFromMerge(t *testing.T) {
modules := defaultRunnerModules()
validator := &runnerChainValidator{name: "chain-extract", approved: []bool{false, true}, reason: "bad_extract", message: "extract rejected"}
@@ -2181,6 +2212,25 @@ func warningReasons(warnings []contracts.Warning) []string {
return reasons
}
func assertNoCheckpointMetadata(t *testing.T, metadata map[string]any) {
t.Helper()
for key, value := range metadata {
lowerKey := strings.ToLower(key)
if strings.Contains(lowerKey, "checkpoint") || strings.Contains(lowerKey, "workspace") {
t.Fatalf("metadata key %q exposes checkpoint/workspace state", key)
}
text, ok := value.(string)
if !ok {
continue
}
lowerValue := strings.ToLower(text)
if strings.Contains(lowerValue, "checkpoint") || strings.Contains(lowerValue, "workspace") {
t.Fatalf("metadata value for %q exposes checkpoint/workspace state: %q", key, text)
}
}
}
func assertRunError(t *testing.T, err error, want string) {
t.Helper()