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

Integrate chunk validation retries

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-27 00:35:37 +00:00
parent b487d93186
commit 29872b2e28
7 changed files with 221 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

102
internal/framework/pipeline/runner_chunk_plan_test.go
View File

@@ -84,6 +84,35 @@ type retryingChunker struct {
calls int
}
type correctionAwareChunker struct {
key string
plan source.ChunkPlan
calls int
correction *contracts.SemanticCorrection
}
func (c *correctionAwareChunker) Key() string { return c.key }
func (*correctionAwareChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
func (c *correctionAwareChunker) Plan(_ context.Context, request contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
c.calls++
if request.Correction != nil {
correction, err := contracts.CloneSemanticCorrection(request.Correction)
if err != nil {
return contracts.ChunkPlanResult{}, err
}
c.correction = correction
}
response := "initial chunk response"
if c.calls > 1 {
response = "corrected chunk response"
}
candidate, err := contracts.NewModelCandidate([]byte(response), contracts.CorrectionProtocolSingleResponseV1)
if err != nil {
return contracts.ChunkPlanResult{}, err
}
return contracts.ChunkPlanResult{Plan: source.CloneChunkPlan(c.plan), ModelCandidate: candidate}, nil
}
func (c *retryingChunker) Key() string { return c.key }
func (*retryingChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
func (c *retryingChunker) Plan(context.Context, contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
@@ -275,6 +304,53 @@ func TestRunnerOmitsChunkMapForRejectedPlan(t *testing.T) {
}
}
func TestRunnerCorrectsRejectedGeneratedChunkPlan(t *testing.T) {
prepared, plan := preparedTerminalDebugPipeline(t)
prepared.resolved.ChunkExecutionClass = contracts.ExecutionClassLLMBacked
prepared.resolved.Chunk.Retries = 1
chunker := &correctionAwareChunker{key: prepared.resolved.Chunk.Module, plan: plan}
prepared.chunker = chunker
validatorCalls := 0
prepared.chunkValidators.validators = []preparedValidator{{
resolved: ResolvedValidator{Binding: Binding("reject-then-approve"), Target: ValidatorTargetChunk},
chunk: chunkValidationFunc{name: "reject-then-approve", validate: func(context.Context, contracts.ChunkValidationRequest) (contracts.ValidationResult, error) {
validatorCalls++
if validatorCalls == 1 {
return contracts.ValidationResult{ReasonCode: "chunk_scope", Message: "scope needs correction", CorrectionGuidance: "correct the chunk scope"}, nil
}
return contracts.ValidationResult{Approved: true}, nil
}},
}}
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")})
if err != nil {
t.Fatalf("Run() error = %v", err)
}
if len(output.Rejected) != 0 || chunker.calls != 2 || validatorCalls != 2 {
t.Fatalf("output = %#v chunker calls = %d validator calls = %d", output.Rejected, chunker.calls, validatorCalls)
}
if got := string(chunker.correction.AssistantResponse); got != "initial chunk response" {
t.Fatalf("correction response = %q", got)
}
}
func TestRunnerDoesNotPublishValidationIncompleteChunkPlan(t *testing.T) {
prepared, plan := preparedTerminalDebugPipeline(t)
prepared.resolved.ChunkValidationPolicy.ValidatorFailure = ValidatorFailureWarnContinue
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan}
prepared.chunkValidators.validators = []preparedValidator{{
resolved: ResolvedValidator{Binding: Binding("unavailable"), Target: ValidatorTargetChunk},
chunk: &countingChunkValidator{err: errors.New("validator unavailable")},
}}
store := &recordingChunkPlanStore{decision: ChunkPlanDecision{Status: ChunkPlanMissing}}
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ChunkCacheMode: ChunkCacheAuto, ChunkPlans: store})
if err != nil {
t.Fatalf("Run() error = %v", err)
}
if output.ChunkPlan == nil || output.ChunkPlan.ValidationStatus != "incomplete" || output.ChunkPlan.PublicationStatus != "not_published" || store.saves != 0 {
t.Fatalf("chunk summary = %#v saves = %d", output.ChunkPlan, store.saves)
}
}
func TestRunnerRetainsChunkMapAfterLaneRejection(t *testing.T) {
prepared, _ := preparedTerminalDebugPipeline(t)
encoder := &capturingChunkMapOutput{}
@@ -322,8 +398,8 @@ func TestRunnerChunkPlanManifestRetainsCandidateOnRejection(t *testing.T) {
if err != nil {
t.Fatal(err)
}
if output.Manifest.ChunkPlan.Action != "refreshed" || output.Manifest.ChunkPlan.PlanDigest == "" || output.ChunkPlan.ValidationStatus != "rejected" || output.ChunkPlan.PublicationStatus != "not_published" {
t.Fatalf("manifest = %#v summary = %#v", output.Manifest.ChunkPlan, output.ChunkPlan)
if output.Manifest.ChunkPlan.Action != "" || output.Manifest.ChunkPlan.PlanDigest != "" || output.ChunkPlan.ValidationStatus != "rejected" || output.ChunkPlan.PublicationStatus != "not_published" {
t.Fatalf("manifest = %#v summary = %#v; rejected output must not retain a usable chunk candidate", output.Manifest.ChunkPlan, output.ChunkPlan)
}
}
@@ -455,20 +531,32 @@ func TestRunnerAutoHitValidatesOnceWithoutRegenerationOrMutation(t *testing.T) {
} else if err != nil {
t.Fatalf("Run() error = %v", err)
}
if validator.calls != 1 || calls != 0 || llmCalls != 0 || store.loads != 1 || store.saves != 0 {
wantProducerCalls := 0
if tc.wantReject {
wantProducerCalls = 1
}
if validator.calls != 1+wantProducerCalls || calls != wantProducerCalls || llmCalls != wantProducerCalls || store.loads != 1 || store.saves != 0 {
t.Fatalf("calls = validator %d module %d llm %d load %d save %d", validator.calls, calls, llmCalls, store.loads, store.saves)
}
assertAttemptEnvelopeSequence(t, debug, "chunk")
if tc.wantReject {
assertAttemptEnvelopeSequence(t, debug, "chunk", 1)
} else {
assertAttemptEnvelopeSequence(t, debug, "chunk")
}
if tc.wantError == "" {
encoded := string(debug.json["chunk/output.json"])
if !strings.Contains(encoded, `"status":"hit"`) || !strings.Contains(encoded, `"plan":`) || !strings.Contains(encoded, `"materialized_chunks":`) {
if tc.wantReject {
if !strings.Contains(encoded, `"status":"invalid"`) || strings.Contains(encoded, `"plan":`) {
t.Fatalf("rejected cache debug = %s", encoded)
}
} else if !strings.Contains(encoded, `"status":"hit"`) || !strings.Contains(encoded, `"plan":`) || !strings.Contains(encoded, `"materialized_chunks":`) {
t.Fatalf("chunk hit debug = %s", encoded)
}
}
if tc.wantReject && (len(output.Rejected) != 1 || output.Rejected[0].ReasonCode != "rejected") {
t.Fatalf("rejected = %#v", output.Rejected)
}
if tc.wantError == "" && len(output.Warnings) != 1+len(tc.result.Warnings) {
if tc.wantError == "" && !tc.wantReject && len(output.Warnings) != 1+len(tc.result.Warnings) {
t.Fatalf("warnings = %#v, want stored warning once plus current warnings", output.Warnings)
}
})
@@ -485,7 +573,7 @@ func TestRunnerPublishesOnlyAcceptedGeneratedPlans(t *testing.T) {
wantReject bool
}{
{name: "module error", moduleErr: errors.New("generation failed"), wantCalls: 2},
{name: "validator rejection", validator: contracts.ValidationResult{Approved: false, ReasonCode: "rejected", Message: "rejected generated plan"}, wantCalls: 2, wantReject: true},
{name: "validator rejection", validator: contracts.ValidationResult{Approved: false, ReasonCode: "rejected", Message: "rejected generated plan"}, wantCalls: 1, wantReject: true},
{name: "cancellation", cancel: true},
} {
t.Run(tc.name, func(t *testing.T) {