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

Reuse valid workspace checkpoints on request

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-08 03:02:50 +00:00
parent 1d3a444df8
commit ae9c2e1d5e
14 changed files with 1406 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1023,6 +1023,127 @@ func TestRunDoesNotPassCheckpointPathsToModules(t *testing.T) {
}
}
func TestRunReusesCheckpointedWorkflowOutputs(t *testing.T) {
modules := defaultRunnerModules()
doc := validSourceDocument()
chunks := []contracts.SourceChunk{sourceChunkWithID("chunk-0", 0)}
extractOutput := contracts.ExtractOutput{
LaneID: "alpha",
ExtractorKey: "extract-alpha",
SourceID: doc.ID,
ChunkID: "chunk-0",
ChunkIndex: 0,
Payload: contracts.RawPayload{
Content: []byte(`{"cached_extract":true}`),
MediaType: "application/json",
},
}
mergeOutput := contracts.MergeOutput{
LaneID: "alpha",
MergerKey: "merge",
SourceID: doc.ID,
Payload: contracts.RawPayload{
Content: []byte(`{"cached_merge":true}`),
MediaType: "application/json",
},
}
normalizeOutput := contracts.NormalizeOutput{
LaneID: "alpha",
NormalizerKey: "normalize",
SourceID: doc.ID,
Payload: contracts.RawPayload{
Content: []byte(`{"cached_normalize":true}`),
MediaType: "application/json",
},
}
loader := &runnerCheckpointLoader{
source: SourceCheckpoint{Document: doc},
chunk: ChunkCheckpoint{Chunks: chunks},
extract: ExtractCheckpoint{Outputs: []contracts.ExtractOutput{extractOutput}},
merge: MergeCheckpoint{Output: mergeOutput},
normalize: NormalizeCheckpoint{Output: normalizeOutput},
reuse: map[string]bool{
"source": true,
"chunk": true,
"extract": true,
"merge": true,
"normalize": true,
},
}
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{
Pipeline: resolvedPipeline(),
Checkpoint: loader,
})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if len(modules.input.requests) != 0 || len(modules.chunker.requests) != 0 || len(modules.extractors["extract-alpha"].requests) != 0 || len(modules.mergers["merge"].requests) != 0 || len(modules.normalizers["normalize"].requests) != 0 {
t.Fatalf("module requests = input:%d chunk:%d extract:%d merge:%d normalize:%d, want all skipped", len(modules.input.requests), len(modules.chunker.requests), len(modules.extractors["extract-alpha"].requests), len(modules.mergers["merge"].requests), len(modules.normalizers["normalize"].requests))
}
if len(output.NormalizeOutputs) != 1 || string(output.NormalizeOutputs[0].Payload.Content) != `{"cached_normalize":true}` {
t.Fatalf("NormalizeOutputs = %#v, want cached normalize output", output.NormalizeOutputs)
}
if len(output.CheckpointEvents) != 5 {
t.Fatalf("checkpoint events = %#v, want one per reusable workflow step", output.CheckpointEvents)
}
for _, event := range output.CheckpointEvents {
if event.Action != "reused" {
t.Fatalf("checkpoint event = %#v, want reused", event)
}
}
}
func TestRunPreservesCheckpointedExtractRejections(t *testing.T) {
modules := defaultRunnerModules()
extractOutput := contracts.ExtractOutput{
LaneID: "alpha",
ExtractorKey: "extract-alpha",
SourceID: "source-1",
ChunkID: "chunk-1",
ChunkIndex: 1,
Payload: contracts.RawPayload{
Content: []byte(`{"cached_extract":true}`),
MediaType: "application/json",
},
}
rejected := contracts.RejectedOutput{
Stage: string(StageExtract),
LaneID: "alpha",
ModuleKey: "extract-alpha",
ChunkID: "chunk-0",
ReasonCode: "invalid_shape",
Message: "invalid extract",
}
loader := &runnerCheckpointLoader{
extract: ExtractCheckpoint{
Outputs: []contracts.ExtractOutput{extractOutput},
Rejected: []contracts.RejectedOutput{rejected},
},
reuse: map[string]bool{"extract": true},
}
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{
Pipeline: resolvedPipeline(),
Checkpoint: loader,
})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if len(modules.extractors["extract-alpha"].requests) != 0 {
t.Fatalf("extract requests = %d, want reused checkpoint", len(modules.extractors["extract-alpha"].requests))
}
if len(output.Rejected) != 1 || output.Rejected[0].ChunkID != "chunk-0" {
t.Fatalf("rejected outputs = %#v, want checkpointed extract rejection", output.Rejected)
}
mergeRequests := modules.mergers["merge"].requests
if len(mergeRequests) != 1 || len(mergeRequests[0].ExtractOutputs) != 1 || mergeRequests[0].ExtractOutputs[0].ChunkID != "chunk-1" {
t.Fatalf("merge extract outputs = %#v, want only checkpointed accepted extract", mergeRequests)
}
}
func TestRunOmitsRejectedExtractOutputsFromMerge(t *testing.T) {
modules := defaultRunnerModules()
validator := &runnerChainValidator{name: "chain-extract", approved: []bool{false, true}, reason: "bad_extract", message: "extract rejected"}
@@ -2089,6 +2210,54 @@ func (encoder *runnerOutputEncoder) ManifestMetadata() map[string]any {
return encoder.manifestMetadata
}
type runnerCheckpointLoader struct {
source SourceCheckpoint
chunk ChunkCheckpoint
extract ExtractCheckpoint
merge MergeCheckpoint
normalize NormalizeCheckpoint
reuse map[string]bool
}
func (loader *runnerCheckpointLoader) Enabled() bool {
return true
}
func (loader *runnerCheckpointLoader) Source(string) (SourceCheckpoint, CheckpointDecision) {
if loader.reuse["source"] {
return loader.source, CheckpointDecision{Reused: true, Reason: "test checkpoint"}
}
return SourceCheckpoint{}, CheckpointDecision{Reason: "test checkpoint missing"}
}
func (loader *runnerCheckpointLoader) Chunk(string, string) (ChunkCheckpoint, CheckpointDecision) {
if loader.reuse["chunk"] {
return loader.chunk, CheckpointDecision{Reused: true, Reason: "test checkpoint"}
}
return ChunkCheckpoint{}, CheckpointDecision{Reason: "test checkpoint missing"}
}
func (loader *runnerCheckpointLoader) Extract(string, string, []CheckpointFingerprint) (ExtractCheckpoint, CheckpointDecision) {
if loader.reuse["extract"] {
return loader.extract, CheckpointDecision{Reused: true, Reason: "test checkpoint"}
}
return ExtractCheckpoint{}, CheckpointDecision{Reason: "test checkpoint missing"}
}
func (loader *runnerCheckpointLoader) Merge(string, string, []CheckpointFingerprint) (MergeCheckpoint, CheckpointDecision) {
if loader.reuse["merge"] {
return loader.merge, CheckpointDecision{Reused: true, Reason: "test checkpoint"}
}
return MergeCheckpoint{}, CheckpointDecision{Reason: "test checkpoint missing"}
}
func (loader *runnerCheckpointLoader) Normalize(string, string, []CheckpointFingerprint) (NormalizeCheckpoint, CheckpointDecision) {
if loader.reuse["normalize"] {
return loader.normalize, CheckpointDecision{Reused: true, Reason: "test checkpoint"}
}
return NormalizeCheckpoint{}, CheckpointDecision{Reason: "test checkpoint missing"}
}
type fakeLLMClient struct{}
func (client fakeLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {