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

Implement runner retries and raw validation

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-07 19:19:04 +00:00
parent bcedf19a08
commit cc6b050367
8 changed files with 805 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
internal/framework/pipeline/raw_validation_registry.go Normal file
View File

@@ -0,0 +1,73 @@
package pipeline
import (
"fmt"
"strings"
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)
type rawValidationKey struct {
stage ModuleStage
module string
}
type RawValidationRegistry struct {
chains map[rawValidationKey][]contracts.RawValidator
}
func NewRawValidationRegistry() *RawValidationRegistry {
return &RawValidationRegistry{
chains: make(map[rawValidationKey][]contracts.RawValidator),
}
}
func (r *RawValidationRegistry) Register(stage ModuleStage, module string, validators ...contracts.RawValidator) error {
if r == nil {
return fmt.Errorf("raw validation registry must not be nil")
}
normalizedModule := strings.TrimSpace(module)
if normalizedModule == "" {
return fmt.Errorf("raw validation module key must not be empty")
}
switch stage {
case StageChunk, StageExtract, StageMerge, StageNormalize:
default:
return fmt.Errorf("raw validation stage %q is not supported", stage)
}
if len(validators) == 0 {
return fmt.Errorf("raw validation chain for %q %q must not be empty", stage, normalizedModule)
}
chain := make([]contracts.RawValidator, 0, len(validators))
for i, validator := range validators {
if validator == nil {
return fmt.Errorf("raw validator %d for %q %q must not be nil", i, stage, normalizedModule)
}
if strings.TrimSpace(validator.Name()) == "" {
return fmt.Errorf("raw validator %d for %q %q must not have an empty name", i, stage, normalizedModule)
}
chain = append(chain, validator)
}
if r.chains == nil {
r.chains = make(map[rawValidationKey][]contracts.RawValidator)
}
key := rawValidationKey{stage: stage, module: normalizedModule}
if _, exists := r.chains[key]; exists {
return fmt.Errorf("raw validation chain for %q %q is already registered", stage, normalizedModule)
}
r.chains[key] = append([]contracts.RawValidator(nil), chain...)
return nil
}
func (r *RawValidationRegistry) Validators(stage ModuleStage, module string) []contracts.RawValidator {
if r == nil {
return nil
}
chain := r.chains[rawValidationKey{stage: stage, module: strings.TrimSpace(module)}]
if len(chain) == 0 {
return nil
}
return append([]contracts.RawValidator(nil), chain...)
}

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

@@ -18,13 +18,14 @@ import (
)
type Registries struct {
Inputs *InputAdapterRegistry
Chunkers *ChunkerRegistry
Extractors *ExtractorRegistry
Mergers *MergerRegistry
Normalizers *NormalizerRegistry
Validators *ValidatorRegistry
Outputs *OutputEncoderRegistry
Inputs *InputAdapterRegistry
Chunkers *ChunkerRegistry
Extractors *ExtractorRegistry
Mergers *MergerRegistry
Normalizers *NormalizerRegistry
Validators *ValidatorRegistry
RawValidators *RawValidationRegistry
Outputs *OutputEncoderRegistry
}
type Runner struct {
@@ -103,31 +104,51 @@ 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)
chunkResult, err := chunker.Chunk(ctx, contracts.ChunkRequest{
Source: doc,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(input.Pipeline.ChunkReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: input.Pipeline.Chunk.LLMProfile,
Options: cloneOptions(input.Pipeline.Chunk.Options),
Metadata: input.Metadata,
var canonicalChunks []contracts.SourceChunk
var chunkWarnings []contracts.Warning
chunksAccepted, chunkRejection, err := runWithRetry(ctx, input.Pipeline.Chunk.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
chunkResult, err := chunker.Chunk(ctx, contracts.ChunkRequest{
Source: doc,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(input.Pipeline.ChunkReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: input.Pipeline.Chunk.LLMProfile,
Options: cloneOptions(input.Pipeline.Chunk.Options),
Metadata: input.Metadata,
})
if err != nil {
return false, nil, fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), err)
}
if len(chunkResult.Chunks) == 0 {
return false, nil, fmt.Errorf("chunker %q returned no chunks", chunker.Key())
}
chunks, err := validateAndCanonicalizeChunkResult(doc, chunkResult.Chunks)
if err != nil {
return false, nil, fmt.Errorf("validate chunks from chunker %q: %w", chunker.Key(), err)
}
rejection, err := r.validateChunksRaw(ctx, doc, chunker.Key(), chunks, input.Metadata, attempt)
if err != nil || rejection != nil {
return false, rejection, err
}
canonicalChunks = chunks
chunkWarnings = cloneWarnings(chunkResult.Warnings)
return true, nil, nil
})
output.Warnings = append(output.Warnings, chunkResult.Warnings...)
if err != nil {
return failOutput(output), fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), err)
return failOutput(output), err
}
if len(chunkResult.Chunks) == 0 {
return failOutput(output), fmt.Errorf("chunker %q returned no chunks", chunker.Key())
}
canonicalChunks, err := validateAndCanonicalizeChunkResult(doc, chunkResult.Chunks)
if err != nil {
return failOutput(output), fmt.Errorf("validate chunks from chunker %q: %w", chunker.Key(), err)
if !chunksAccepted {
output.Rejected = append(output.Rejected, *chunkRejection)
} else {
output.Warnings = append(output.Warnings, chunkWarnings...)
}
for _, lane := range input.Pipeline.ArtifactLanes {
if err := r.runLane(ctx, input, doc, sourceInput, sessionID, canonicalChunks, lane, &output); err != nil {
return failOutput(output), err
if chunksAccepted {
for _, lane := range input.Pipeline.ArtifactLanes {
if err := r.runLane(ctx, input, doc, sourceInput, sessionID, canonicalChunks, lane, &output); err != nil {
return failOutput(output), err
}
}
}
@@ -183,70 +204,165 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
extractOutputs := make([]contracts.ExtractOutput, 0, len(chunks))
for index := range chunks {
chunk := chunks[index]
result, err := extractor.Extract(ctx, contracts.ExtractionRequest{
var acceptedOutput contracts.ExtractOutput
var acceptedWarnings []contracts.Warning
accepted, rejection, err := runWithRetry(ctx, lane.Extract.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
result, err := extractor.Extract(ctx, contracts.ExtractionRequest{
Source: doc,
Chunk: &chunk,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.ExtractReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Extract.LLMProfile,
Options: cloneOptions(lane.Extract.Options),
Metadata: input.Metadata,
})
if err != nil {
return false, nil, fmt.Errorf("extract lane %q chunk %q with extractor %q: %w", lane.ID, chunk.ID, extractor.Key(), err)
}
extractOutput := result.Output
extractOutput.LaneID = lane.ID
extractOutput.ExtractorKey = extractor.Key()
extractOutput.SourceID = doc.ID
extractOutput.ChunkID = chunk.ID
extractOutput.ChunkIndex = chunk.Index
extractOutput.Payload.Warnings = append(extractOutput.Payload.Warnings, result.Warnings...)
validationWarnings, rejection, err := r.validateRaw(ctx, rawValidationTarget{
stage: StageExtract,
laneID: lane.ID,
moduleKey: extractor.Key(),
source: doc,
sourceID: doc.ID,
chunkID: chunk.ID,
chunkIndex: chunk.Index,
schema: extractOutput.Schema,
payload: extractOutput.Payload,
metadata: input.Metadata,
attempt: attempt,
})
if err != nil || rejection != nil {
return false, rejection, err
}
acceptedOutput = cloneExtractOutput(extractOutput)
acceptedWarnings = append(cloneWarnings(result.Warnings), validationWarnings...)
return true, nil, nil
})
if err != nil {
return err
}
if !accepted {
output.Rejected = append(output.Rejected, *rejection)
continue
}
output.Warnings = append(output.Warnings, acceptedWarnings...)
extractOutputs = append(extractOutputs, acceptedOutput)
}
if len(extractOutputs) == 0 {
return nil
}
var acceptedMerge contracts.MergeOutput
var mergeWarnings []contracts.Warning
mergeAccepted, mergeRejection, err := runWithRetry(ctx, lane.Merge.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
mergeResult, err := merger.Merge(ctx, contracts.MergeRequest{
Source: doc,
LaneID: lane.ID,
ExtractOutputs: cloneExtractOutputs(extractOutputs),
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.MergeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Merge.LLMProfile,
Options: cloneOptions(lane.Merge.Options),
Metadata: input.Metadata,
})
if err != nil {
return false, nil, fmt.Errorf("merge lane %q with merger %q: %w", lane.ID, merger.Key(), err)
}
mergeOutput := mergeResult.Output
mergeOutput.LaneID = lane.ID
mergeOutput.MergerKey = merger.Key()
mergeOutput.SourceID = doc.ID
mergeOutput.Payload.Warnings = append(mergeOutput.Payload.Warnings, mergeResult.Warnings...)
validationWarnings, rejection, err := r.validateRaw(ctx, rawValidationTarget{
stage: StageMerge,
laneID: lane.ID,
moduleKey: merger.Key(),
source: doc,
sourceID: doc.ID,
schema: mergeOutput.Schema,
payload: mergeOutput.Payload,
metadata: input.Metadata,
attempt: attempt,
})
if err != nil || rejection != nil {
return false, rejection, err
}
acceptedMerge = cloneMergeOutput(mergeOutput)
mergeWarnings = append(cloneWarnings(mergeResult.Warnings), validationWarnings...)
return true, nil, nil
})
if err != nil {
return err
}
if !mergeAccepted {
output.Rejected = append(output.Rejected, *mergeRejection)
return nil
}
output.Warnings = append(output.Warnings, mergeWarnings...)
var acceptedNormalize contracts.NormalizeOutput
var normalizeWarnings []contracts.Warning
normalizeAccepted, normalizeRejection, err := runWithRetry(ctx, lane.Normalize.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
normalizeResult, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
Source: doc,
Chunk: &chunk,
LaneID: lane.ID,
MergeOutput: cloneMergeOutput(acceptedMerge),
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.ExtractReferences.ReferenceSet),
References: CloneReferenceSet(lane.NormalizeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Extract.LLMProfile,
Options: cloneOptions(lane.Extract.Options),
LLMProfile: lane.Normalize.LLMProfile,
Options: cloneOptions(lane.Normalize.Options),
Metadata: input.Metadata,
})
output.Warnings = append(output.Warnings, result.Warnings...)
if err != nil {
return fmt.Errorf("extract lane %q chunk %q with extractor %q: %w", lane.ID, chunk.ID, extractor.Key(), err)
return false, nil, fmt.Errorf("normalize lane %q with normalizer %q: %w", lane.ID, normalizer.Key(), err)
}
extractOutput := result.Output
extractOutput.LaneID = lane.ID
extractOutput.ExtractorKey = extractor.Key()
extractOutput.SourceID = doc.ID
extractOutput.ChunkID = chunk.ID
extractOutput.ChunkIndex = chunk.Index
extractOutput.Payload.Warnings = append(extractOutput.Payload.Warnings, result.Warnings...)
extractOutputs = append(extractOutputs, cloneExtractOutput(extractOutput))
}
mergeResult, err := merger.Merge(ctx, contracts.MergeRequest{
Source: doc,
LaneID: lane.ID,
ExtractOutputs: cloneExtractOutputs(extractOutputs),
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.MergeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Merge.LLMProfile,
Options: cloneOptions(lane.Merge.Options),
Metadata: input.Metadata,
normalizeOutput := normalizeResult.Output
normalizeOutput.LaneID = lane.ID
normalizeOutput.NormalizerKey = normalizer.Key()
normalizeOutput.SourceID = doc.ID
normalizeOutput.Payload.Warnings = append(normalizeOutput.Payload.Warnings, normalizeResult.Warnings...)
validationWarnings, rejection, err := r.validateRaw(ctx, rawValidationTarget{
stage: StageNormalize,
laneID: lane.ID,
moduleKey: normalizer.Key(),
source: doc,
sourceID: doc.ID,
schema: normalizeOutput.Schema,
payload: normalizeOutput.Payload,
metadata: input.Metadata,
attempt: attempt,
})
if err != nil || rejection != nil {
return false, rejection, err
}
acceptedNormalize = cloneNormalizeOutput(normalizeOutput)
normalizeWarnings = append(cloneWarnings(normalizeResult.Warnings), validationWarnings...)
return true, nil, nil
})
output.Warnings = append(output.Warnings, mergeResult.Warnings...)
if err != nil {
return fmt.Errorf("merge lane %q with merger %q: %w", lane.ID, merger.Key(), err)
return err
}
normalizeResult, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
Source: doc,
LaneID: lane.ID,
MergeOutput: cloneMergeOutput(mergeResult.Output),
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.NormalizeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Normalize.LLMProfile,
Options: cloneOptions(lane.Normalize.Options),
Metadata: input.Metadata,
})
output.Warnings = append(output.Warnings, normalizeResult.Warnings...)
if err != nil {
return fmt.Errorf("normalize lane %q with normalizer %q: %w", lane.ID, normalizer.Key(), err)
if !normalizeAccepted {
output.Rejected = append(output.Rejected, *normalizeRejection)
return nil
}
normalizeOutput := normalizeResult.Output
normalizeOutput.LaneID = lane.ID
normalizeOutput.NormalizerKey = normalizer.Key()
normalizeOutput.SourceID = doc.ID
normalizeOutput.Payload.Warnings = append(normalizeOutput.Payload.Warnings, normalizeResult.Warnings...)
output.NormalizeOutputs = append(output.NormalizeOutputs, cloneNormalizeOutput(normalizeOutput))
output.Warnings = append(output.Warnings, normalizeWarnings...)
output.NormalizeOutputs = append(output.NormalizeOutputs, acceptedNormalize)
return nil
}
@@ -255,6 +371,146 @@ type validatorExecution struct {
binding ModuleBinding
}
type rawValidationTarget struct {
stage ModuleStage
laneID string
moduleKey string
source *source.SourceDocument
sourceID string
chunkID string
chunkIndex int
schema contracts.ResponseSchema
payload contracts.RawPayload
metadata map[string]any
attempt int
}
func runWithRetry(ctx context.Context, retries int, run func(attempt int) (bool, *contracts.RejectedOutput, error)) (bool, *contracts.RejectedOutput, error) {
attempts := 1
if retries > 0 {
attempts += retries
}
var lastRejection *contracts.RejectedOutput
for attempt := 1; attempt <= attempts; attempt++ {
if err := ctx.Err(); err != nil {
return false, nil, err
}
accepted, rejection, err := run(attempt)
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return false, nil, ctxErr
}
if attempt == attempts {
return false, nil, fmt.Errorf("failed after %d attempt(s): %w", attempt, err)
}
continue
}
if accepted {
return true, nil, nil
}
if rejection != nil {
rejection.AttemptCount = attempt
lastRejection = rejection
}
if ctxErr := ctx.Err(); ctxErr != nil {
return false, nil, ctxErr
}
if attempt == attempts {
if lastRejection == nil {
lastRejection = &contracts.RejectedOutput{
ReasonCode: "raw_output_rejected",
Message: "raw output rejected",
AttemptCount: attempt,
}
}
return false, lastRejection, nil
}
}
return false, lastRejection, nil
}
func (r *Runner) validateChunksRaw(ctx context.Context, doc *source.SourceDocument, moduleKey string, chunks []contracts.SourceChunk, metadata map[string]any, attempt int) (*contracts.RejectedOutput, error) {
for _, chunk := range chunks {
_, rejection, err := r.validateRaw(ctx, rawValidationTarget{
stage: StageChunk,
moduleKey: moduleKey,
source: doc,
sourceID: doc.ID,
chunkID: chunk.ID,
chunkIndex: chunk.Index,
payload: contracts.RawPayload{
Content: append([]byte(nil), chunk.Content...),
MediaType: chunk.MediaType,
Metadata: cloneMetadata(chunk.Metadata),
},
metadata: metadata,
attempt: attempt,
})
if err != nil {
return nil, err
}
if rejection != nil {
return rejection, nil
}
}
return nil, nil
}
func (r *Runner) validateRaw(ctx context.Context, target rawValidationTarget) ([]contracts.Warning, *contracts.RejectedOutput, error) {
validators := r.registries.RawValidators.Validators(target.stage, target.moduleKey)
if len(validators) == 0 {
return nil, nil, nil
}
request := contracts.RawValidationRequest{
Stage: string(target.stage),
LaneID: target.laneID,
ModuleKey: target.moduleKey,
Source: target.source,
SourceID: target.sourceID,
ChunkID: target.chunkID,
ChunkIndex: target.chunkIndex,
Schema: target.schema,
Payload: cloneRawPayload(target.payload),
Metadata: cloneMetadata(target.metadata),
}
var warnings []contracts.Warning
for _, validator := range validators {
result, err := validator.ValidateRaw(ctx, request)
if err != nil {
return nil, nil, fmt.Errorf("validate raw %s output with validator %q: %w", target.stage, validator.Name(), err)
}
if !result.Approved {
reasonCode := strings.TrimSpace(result.ReasonCode)
if reasonCode == "" {
reasonCode = "raw_output_rejected"
}
message := strings.TrimSpace(result.Message)
if message == "" {
message = "raw output rejected"
}
return nil, &contracts.RejectedOutput{
Stage: string(target.stage),
LaneID: target.laneID,
ModuleKey: target.moduleKey,
ChunkID: target.chunkID,
ChunkIndex: target.chunkIndex,
ValidatorName: validator.Name(),
ReasonCode: reasonCode,
Message: message,
AttemptCount: target.attempt,
DiagnosticArtifactPath: result.DiagnosticArtifactPath,
}, nil
}
warnings = append(warnings, result.Warnings...)
}
return warnings, nil, nil
}
func (r *Runner) buildConfiguredValidators(lane ResolvedArtifactLane) ([]validatorExecution, error) {
validators := make([]validatorExecution, 0, len(lane.Validators))
for _, binding := range lane.Validators {

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

@@ -854,6 +854,209 @@ func TestRunPassesPerChunkRawOutputsToMergeAndNormalize(t *testing.T) {
}
}
func TestRunPassesChunkContentAndMediaTypeToExtractors(t *testing.T) {
modules := defaultRunnerModules()
modules.chunker.chunks = []contracts.SourceChunk{
sourceChunkWithContent("chunk-0", 0, []byte(`{"chunk":0}`), "application/vnd.test+json"),
}
_, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: resolvedPipeline()})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
req := modules.extractors["extract-alpha"].requests[0]
if req.Chunk == nil {
t.Fatal("extractor chunk = nil, want chunk")
}
if got := string(req.Chunk.Content); got != `{"chunk":0}` {
t.Fatalf("chunk content = %q, want raw chunk content", got)
}
if req.Chunk.MediaType != "application/vnd.test+json" {
t.Fatalf("chunk media type = %q, want application/vnd.test+json", req.Chunk.MediaType)
}
}
func TestRunOmitsRejectedExtractOutputsFromMerge(t *testing.T) {
modules := defaultRunnerModules()
validator := &runnerRawValidator{name: "raw-extract", approved: []bool{false, true}, reason: "bad_extract", message: "extract rejected"}
modules.rawValidators = rawValidationRegistry(t, StageExtract, "extract-alpha", validator)
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: resolvedPipeline()})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if len(output.Rejected) != 1 || output.Rejected[0].Stage != string(StageExtract) || output.Rejected[0].ChunkID != "chunk-0" {
t.Fatalf("rejected outputs = %#v, want rejected first extract", output.Rejected)
}
extractOutputs := modules.mergers["merge"].requests[0].ExtractOutputs
if len(extractOutputs) != 1 || extractOutputs[0].ChunkID != "chunk-1" {
t.Fatalf("merge extract outputs = %#v, want only accepted second chunk", extractOutputs)
}
if output.Manifest.ValidationStatus != "rejected" {
t.Fatalf("ValidationStatus = %q, want rejected", output.Manifest.ValidationStatus)
}
}
func TestRunOmitsLaneWithNoAcceptedExtractOutputs(t *testing.T) {
modules := defaultRunnerModules()
validator := &runnerRawValidator{name: "raw-extract", approved: []bool{false}, reason: "bad_extract", message: "extract rejected"}
modules.rawValidators = rawValidationRegistry(t, StageExtract, "extract-alpha", validator)
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: resolvedPipeline()})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if len(output.Rejected) != 2 {
t.Fatalf("len(Rejected) = %d, want one rejected record per chunk", len(output.Rejected))
}
if len(modules.mergers["merge"].requests) != 0 {
t.Fatalf("merge requests = %d, want none", len(modules.mergers["merge"].requests))
}
if len(modules.normalizers["normalize"].requests) != 0 {
t.Fatalf("normalize requests = %d, want none", len(modules.normalizers["normalize"].requests))
}
if len(modules.output.requests) != 1 || len(modules.output.requests[0].NormalizeOutputs) != 0 {
t.Fatalf("output normalize outputs = %#v, want none", modules.output.requests)
}
}
func TestRunRejectedMergePreventsNormalizeForLane(t *testing.T) {
modules := defaultRunnerModules()
validator := &runnerRawValidator{name: "raw-merge", approved: []bool{false}, reason: "bad_merge", message: "merge rejected"}
modules.rawValidators = rawValidationRegistry(t, StageMerge, "merge", validator)
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: resolvedPipeline()})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if len(output.Rejected) != 1 || output.Rejected[0].Stage != string(StageMerge) {
t.Fatalf("rejected outputs = %#v, want rejected merge", output.Rejected)
}
if len(modules.normalizers["normalize"].requests) != 0 {
t.Fatalf("normalize requests = %d, want none", len(modules.normalizers["normalize"].requests))
}
if len(modules.output.requests[0].NormalizeOutputs) != 0 {
t.Fatalf("output normalize outputs = %#v, want none", modules.output.requests[0].NormalizeOutputs)
}
}
func TestRunRejectedNormalizePreventsOutputForLane(t *testing.T) {
modules := defaultRunnerModules()
validator := &runnerRawValidator{name: "raw-normalize", approved: []bool{false}, reason: "bad_normalize", message: "normalize rejected"}
modules.rawValidators = rawValidationRegistry(t, StageNormalize, "normalize", validator)
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: resolvedPipeline()})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if len(output.Rejected) != 1 || output.Rejected[0].Stage != string(StageNormalize) {
t.Fatalf("rejected outputs = %#v, want rejected normalize", output.Rejected)
}
if len(output.NormalizeOutputs) != 0 {
t.Fatalf("NormalizeOutputs = %#v, want none", output.NormalizeOutputs)
}
if len(modules.output.requests[0].NormalizeOutputs) != 0 {
t.Fatalf("output normalize outputs = %#v, want none", modules.output.requests[0].NormalizeOutputs)
}
}
func TestRunRetriesSameModuleInputAfterFrameworkError(t *testing.T) {
modules := defaultRunnerModules()
modules.extractors["extract-alpha"].failuresBeforeSuccess = 1
modules.extractors["extract-alpha"].failureErr = errors.New("transient extract failure")
pipeline := resolvedPipeline()
pipeline.ArtifactLanes[0].Extract.Retries = 1
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: pipeline})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
requests := modules.extractors["extract-alpha"].requests
if len(requests) != 3 {
t.Fatalf("extract requests = %d, want retry plus remaining chunk", len(requests))
}
if requests[0].Chunk.ID != "chunk-0" || requests[1].Chunk.ID != "chunk-0" {
t.Fatalf("retried chunks = %q, %q; want same first chunk input", requests[0].Chunk.ID, requests[1].Chunk.ID)
}
if output.Manifest.ValidationStatus != "approved" {
t.Fatalf("ValidationStatus = %q, want approved", output.Manifest.ValidationStatus)
}
}
func TestRunRetriesSameModuleInputAfterValidatorRejection(t *testing.T) {
modules := defaultRunnerModules()
validator := &runnerRawValidator{name: "raw-extract", approved: []bool{false, true, true}, reason: "bad_extract", message: "extract rejected"}
modules.rawValidators = rawValidationRegistry(t, StageExtract, "extract-alpha", validator)
pipeline := resolvedPipeline()
pipeline.ArtifactLanes[0].Extract.Retries = 1
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: pipeline})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
requests := modules.extractors["extract-alpha"].requests
if len(requests) != 3 {
t.Fatalf("extract requests = %d, want retry plus remaining chunk", len(requests))
}
if requests[0].Chunk.ID != "chunk-0" || requests[1].Chunk.ID != "chunk-0" {
t.Fatalf("retried chunks = %q, %q; want same first chunk input", requests[0].Chunk.ID, requests[1].Chunk.ID)
}
if len(output.Rejected) != 0 {
t.Fatalf("Rejected = %#v, want transient rejection omitted after retry approval", output.Rejected)
}
}
func TestRunStopsRetryAfterConfiguredAttemptsAndRecordsAttemptCount(t *testing.T) {
modules := defaultRunnerModules()
validator := &runnerRawValidator{name: "raw-extract", approved: []bool{false}, reason: "bad_extract", message: "extract rejected"}
modules.rawValidators = rawValidationRegistry(t, StageExtract, "extract-alpha", validator)
pipeline := resolvedPipeline()
pipeline.ArtifactLanes[0].Extract.Retries = 1
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: pipeline})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if len(output.Rejected) != 2 {
t.Fatalf("len(Rejected) = %d, want rejected record per chunk", len(output.Rejected))
}
if output.Rejected[0].AttemptCount != 2 || output.Rejected[1].AttemptCount != 2 {
t.Fatalf("attempt counts = %#v, want final attempt count 2", output.Rejected)
}
if len(modules.extractors["extract-alpha"].requests) != 4 {
t.Fatalf("extract requests = %d, want two attempts per chunk", len(modules.extractors["extract-alpha"].requests))
}
}
func TestRunContextCancellationStopsRetries(t *testing.T) {
modules := defaultRunnerModules()
modules.extractors["extract-alpha"].err = errors.New("extract failed")
pipeline := resolvedPipeline()
pipeline.ArtifactLanes[0].Extract.Retries = 2
ctx, cancel := context.WithCancel(context.Background())
cancel()
output, err := New(newRunnerRegistries(t, modules)).Run(ctx, RunInput{Pipeline: pipeline})
if !errors.Is(err, context.Canceled) {
t.Fatalf("Run() error = %v, want context.Canceled", err)
}
if len(modules.extractors["extract-alpha"].requests) != 0 {
t.Fatalf("extract requests = %d, want none after cancellation", len(modules.extractors["extract-alpha"].requests))
}
if output.Manifest.ValidationStatus != "failed" {
t.Fatalf("ValidationStatus = %q, want failed", output.Manifest.ValidationStatus)
}
}
func TestRunRecordsConfiguredValidatorsInManifest(t *testing.T) {
modules := defaultRunnerModules()
@@ -1278,6 +1481,7 @@ type runnerModules struct {
mergers map[string]*runnerMerger
normalizers map[string]*runnerNormalizer
validators map[string]*runnerValidator
rawValidators *RawValidationRegistry
output *runnerOutputEncoder
inputBuildErr error
chunkerBuildErr error
@@ -1316,13 +1520,14 @@ func newRunnerRegistries(t *testing.T, modules *runnerModules) Registries {
}
registries := Registries{
Inputs: NewInputAdapterRegistry(),
Chunkers: NewChunkerRegistry(),
Extractors: NewExtractorRegistry(),
Mergers: NewMergerRegistry(),
Normalizers: NewNormalizerRegistry(),
Validators: NewValidatorRegistry(),
Outputs: NewOutputEncoderRegistry(),
Inputs: NewInputAdapterRegistry(),
Chunkers: NewChunkerRegistry(),
Extractors: NewExtractorRegistry(),
Mergers: NewMergerRegistry(),
Normalizers: NewNormalizerRegistry(),
Validators: NewValidatorRegistry(),
RawValidators: modules.rawValidators,
Outputs: NewOutputEncoderRegistry(),
}
if err := registries.Inputs.Register("input", func() (contracts.InputAdapter, error) {
if modules.inputBuildErr != nil {
@@ -1392,12 +1597,14 @@ func (adapter *runnerInputAdapter) ManifestMetadata() map[string]any {
}
type runnerChunker struct {
key string
chunks []contracts.SourceChunk
warnings []contracts.Warning
err error
manifestMetadata map[string]any
requests []contracts.ChunkRequest
key string
chunks []contracts.SourceChunk
warnings []contracts.Warning
err error
failureErr error
failuresBeforeSuccess int
manifestMetadata map[string]any
requests []contracts.ChunkRequest
}
func (chunker *runnerChunker) Key() string {
@@ -1410,6 +1617,14 @@ func (chunker *runnerChunker) ReferenceSlots() []contracts.ReferenceSlot {
func (chunker *runnerChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
chunker.requests = append(chunker.requests, req)
if chunker.failuresBeforeSuccess > 0 {
chunker.failuresBeforeSuccess--
err := chunker.failureErr
if err == nil {
err = errors.New("transient chunk failure")
}
return contracts.ChunkResult{}, err
}
return contracts.ChunkResult{
Chunks: chunker.chunks,
Warnings: chunker.warnings,
@@ -1421,15 +1636,17 @@ func (chunker *runnerChunker) ManifestMetadata() map[string]any {
}
type runnerExtractor struct {
key string
manifestMetadata map[string]any
output *contracts.ExtractOutput
warnings []contracts.Warning
err error
requests []contracts.ExtractionRequest
seenChunkIDs []string
seenLLMClients []contracts.StructuredLLMClient
seenMetadata []map[string]any
key string
manifestMetadata map[string]any
output *contracts.ExtractOutput
warnings []contracts.Warning
err error
failureErr error
failuresBeforeSuccess int
requests []contracts.ExtractionRequest
seenChunkIDs []string
seenLLMClients []contracts.StructuredLLMClient
seenMetadata []map[string]any
}
func (extractor *runnerExtractor) Key() string {
@@ -1452,6 +1669,15 @@ func (extractor *runnerExtractor) Extract(ctx context.Context, req contracts.Ext
extractor.seenLLMClients = append(extractor.seenLLMClients, req.LLMClient)
extractor.seenMetadata = append(extractor.seenMetadata, req.Metadata)
if extractor.failuresBeforeSuccess > 0 {
extractor.failuresBeforeSuccess--
err := extractor.failureErr
if err == nil {
err = errors.New("transient extract failure")
}
return contracts.ExtractionResult{}, err
}
output := contracts.ExtractOutput{
Schema: contracts.ResponseSchema{ID: "runner.raw", Name: "runner_raw", Version: "v1"},
Payload: contracts.RawPayload{
@@ -1472,11 +1698,13 @@ func (extractor *runnerExtractor) Extract(ctx context.Context, req contracts.Ext
}
type runnerMerger struct {
key string
result *contracts.MergeOutput
warnings []contracts.Warning
err error
requests []contracts.MergeRequest
key string
result *contracts.MergeOutput
warnings []contracts.Warning
err error
failureErr error
failuresBeforeSuccess int
requests []contracts.MergeRequest
}
func (merger *runnerMerger) Key() string {
@@ -1485,6 +1713,14 @@ func (merger *runnerMerger) Key() string {
func (merger *runnerMerger) Merge(ctx context.Context, req contracts.MergeRequest) (contracts.MergeResult, error) {
merger.requests = append(merger.requests, req)
if merger.failuresBeforeSuccess > 0 {
merger.failuresBeforeSuccess--
err := merger.failureErr
if err == nil {
err = errors.New("transient merge failure")
}
return contracts.MergeResult{}, err
}
output := contracts.MergeOutput{
LaneID: req.LaneID,
SourceID: req.Source.ID,
@@ -1504,11 +1740,13 @@ func (merger *runnerMerger) Merge(ctx context.Context, req contracts.MergeReques
}
type runnerNormalizer struct {
key string
result *contracts.NormalizeOutput
warnings []contracts.Warning
err error
requests []contracts.NormalizeRequest
key string
result *contracts.NormalizeOutput
warnings []contracts.Warning
err error
failureErr error
failuresBeforeSuccess int
requests []contracts.NormalizeRequest
}
func (normalizer *runnerNormalizer) Key() string {
@@ -1521,6 +1759,14 @@ func (normalizer *runnerNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
func (normalizer *runnerNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
normalizer.requests = append(normalizer.requests, req)
if normalizer.failuresBeforeSuccess > 0 {
normalizer.failuresBeforeSuccess--
err := normalizer.failureErr
if err == nil {
err = errors.New("transient normalize failure")
}
return contracts.NormalizeResult{}, err
}
output := contracts.NormalizeOutput{
LaneID: req.LaneID,
SourceID: req.MergeOutput.SourceID,
@@ -1547,6 +1793,43 @@ type runnerValidator struct {
requests []contracts.ValidationRequest
}
type runnerRawValidator struct {
name string
approved []bool
reason string
message string
warnings []contracts.Warning
err error
calls int
requests []contracts.RawValidationRequest
}
func (validator *runnerRawValidator) Name() string {
return validator.name
}
func (validator *runnerRawValidator) ValidateRaw(ctx context.Context, req contracts.RawValidationRequest) (contracts.RawValidationResult, error) {
validator.calls++
validator.requests = append(validator.requests, req)
if validator.err != nil {
return contracts.RawValidationResult{}, validator.err
}
approved := true
if len(validator.approved) > 0 {
index := validator.calls - 1
if index >= len(validator.approved) {
index = len(validator.approved) - 1
}
approved = validator.approved[index]
}
return contracts.RawValidationResult{
Approved: approved,
ReasonCode: validator.reason,
Message: validator.message,
Warnings: validator.warnings,
}, nil
}
func (validator *runnerValidator) Name() string {
return validator.name
}
@@ -1668,6 +1951,13 @@ func sourceChunkWithID(id string, index int) contracts.SourceChunk {
}
}
func sourceChunkWithContent(id string, index int, content []byte, mediaType string) contracts.SourceChunk {
chunk := sourceChunkWithID(id, index)
chunk.Content = append([]byte(nil), content...)
chunk.MediaType = mediaType
return chunk
}
func unitWithID(id string) source.SourceUnit {
switch id {
case "u1":
@@ -1723,3 +2013,13 @@ func assertRunError(t *testing.T, err error, want string) {
t.Fatalf("Run() error = %q, want substring %q", err.Error(), want)
}
}
func rawValidationRegistry(t *testing.T, stage ModuleStage, module string, validators ...contracts.RawValidator) *RawValidationRegistry {
t.Helper()
registry := NewRawValidationRegistry()
if err := registry.Register(stage, module, validators...); err != nil {
t.Fatalf("register raw validators: %v", err)
}
return registry
}