Preserve structured repair settings across pipeline boundaries

internal/framework/llm/promptkit_client.go

@@ -144,6 +144,9 @@ func (c *PromptKitClient) CompleteStructured(ctx context.Context, req contracts.
 	if req.StructuredOutputRepairAttempts != nil {
 		inspection, err := c.engine.InspectPrompt(ctx, promptID, strings.TrimSpace(req.PromptVersion))
 		if err != nil {
+			if ctxErr := ctx.Err(); ctxErr != nil {
+				return contracts.StructuredCompletionResponse{}, ctxErr
+			}
 			return contracts.StructuredCompletionResponse{}, fmt.Errorf("inspect PromptKit prompt %q: %v", promptID, redactPromptKitError(err))
 		}
 		contract := inspection.OutputContract

internal/framework/llm/promptkit_client_test.go

@@ -1147,19 +1147,31 @@ func TestPromptKitClientTranslatesBackendCapacityExhaustion(t *testing.T) {
 }
 
 func TestPromptKitClientContextCancellationIsRespected(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
+	attempts := 1
-	cancel()
+	for _, test := range []struct {
-	client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"ok":true}`})
+		name     string
+		attempts *int
+	}{
+		{name: "during preparation"},
+		{name: "during override inspection", attempts: &attempts},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel()
+			client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"ok":true}`})
 
-	var out map[string]any
+			var out map[string]any
-	_, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
+			_, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
-		PromptID: "adapter.test",
+				PromptID:                       "adapter.test",
-		Inputs: contracts.LLMInputSet{
+				StructuredOutputRepairAttempts: test.attempts,
-			"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
+				Inputs: contracts.LLMInputSet{
-		},
+					"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
-	}, &out)
+				},
-	if !errors.Is(err, context.Canceled) || errors.Is(err, contracts.ErrInvalidStructuredOutput) {
+			}, &out)
-		t.Fatalf("CompleteStructured() error = %v, want context canceled", err)
+			if !errors.Is(err, context.Canceled) || errors.Is(err, contracts.ErrInvalidStructuredOutput) {
+				t.Fatalf("CompleteStructured() error = %v, want context canceled", err)
+			}
+		})
 	}
 }
 

internal/framework/pipeline/merger_registry.go

@@ -85,7 +85,7 @@ func RegisterMergerBuilder[T any](registry *MergerRegistry, spec ModuleSpec, val
 				}
 				outputs[i] = contracts.ExtractArtifact[T]{LaneID: output.LaneID, ExtractorKey: output.ExtractorKey, SourceID: output.SourceID, ChunkID: output.ChunkID, ChunkIndex: output.ChunkIndex, ChunkRef: output.ChunkRef, Value: value}
 			}
-			result, err := merger.Merge(ctx, contracts.TypedMergeRequest[T]{Source: request.Source, LaneID: request.LaneID, ExtractOutputs: outputs, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, Metadata: request.Metadata})
+			result, err := merger.Merge(ctx, contracts.TypedMergeRequest[T]{Source: request.Source, LaneID: request.LaneID, ExtractOutputs: outputs, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, StructuredOutputRepairAttempts: request.StructuredOutputRepairAttempts, Metadata: request.Metadata})
 			if err != nil {
 				return erasedTypedResult{}, err
 			}

internal/framework/pipeline/normalizer_registry.go

@@ -76,7 +76,7 @@ func RegisterNormalizerBuilder[T any](registry *NormalizerRegistry, spec ModuleS
 			if err != nil {
 				return erasedTypedResult{}, err
 			}
-			result, err := normalizer.Normalize(ctx, contracts.TypedNormalizeRequest[T]{Source: request.Source, LaneID: request.LaneID, MergeOutput: contracts.MergeArtifact[T]{LaneID: request.MergeOutput.LaneID, MergerKey: request.MergeOutput.MergerKey, SourceID: request.MergeOutput.SourceID, Value: value}, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, Metadata: request.Metadata})
+			result, err := normalizer.Normalize(ctx, contracts.TypedNormalizeRequest[T]{Source: request.Source, LaneID: request.LaneID, MergeOutput: contracts.MergeArtifact[T]{LaneID: request.MergeOutput.LaneID, MergerKey: request.MergeOutput.MergerKey, SourceID: request.MergeOutput.SourceID, Value: value}, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, StructuredOutputRepairAttempts: request.StructuredOutputRepairAttempts, Metadata: request.Metadata})
 			if err != nil {
 				return erasedTypedResult{}, err
 			}

internal/framework/pipeline/runner.go

@@ -493,9 +493,9 @@ func (r *Runner) validateChunks(ctx context.Context, doc *source.SourceDocument,
 		}
 		switch item.resolved.Target {
 		case ValidatorTargetChunk:
-			result, err = item.chunk.Validate(validatorCtx, contracts.ChunkValidationRequest{ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, Metadata: requestMetadata, Chunks: requestChunks})
+			result, err = item.chunk.Validate(validatorCtx, contracts.ChunkValidationRequest{ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(binding.StructuredOutputRepairAttempts), Metadata: requestMetadata, Chunks: requestChunks})
 		case ValidatorTargetSerialized:
-			result, err = item.serialized.Validate(validatorCtx, contracts.SerializedValidationRequest{Stage: string(StageChunk), ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, Metadata: requestMetadata, Chunks: requestChunks, Schema: contracts.CloneArtifactSchema(schema), MediaType: "application/json", Content: append([]byte(nil), content...)})
+			result, err = item.serialized.Validate(validatorCtx, contracts.SerializedValidationRequest{Stage: string(StageChunk), ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(binding.StructuredOutputRepairAttempts), Metadata: requestMetadata, Chunks: requestChunks, Schema: contracts.CloneArtifactSchema(schema), MediaType: "application/json", Content: append([]byte(nil), content...)})
 		default:
 			return nil, nil, fmt.Errorf("validator %q is incompatible with chunk validation", binding.Module)
 		}

internal/framework/pipeline/runner_terminal_debug_test.go

@@ -41,6 +41,10 @@ type observingChunkValidator struct {
 	request contracts.ChunkValidationRequest
 }
 
+type observingSerializedValidator struct {
+	request contracts.SerializedValidationRequest
+}
+
 func (*observingChunkValidator) Name() string { return "observing/chunk-validator" }
 func (*observingChunkValidator) ExecutionClass() contracts.ExecutionClass {
 	return contracts.ExecutionClassDeterministic
@@ -50,6 +54,15 @@ func (v *observingChunkValidator) Validate(_ context.Context, request contracts.
 	return contracts.ValidationResult{Approved: true}, nil
 }
 
+func (*observingSerializedValidator) Name() string { return "observing/serialized-validator" }
+func (*observingSerializedValidator) ExecutionClass() contracts.ExecutionClass {
+	return contracts.ExecutionClassDeterministic
+}
+func (v *observingSerializedValidator) Validate(_ context.Context, request contracts.SerializedValidationRequest) (contracts.ValidationResult, error) {
+	v.request = request
+	return contracts.ValidationResult{Approved: true}, nil
+}
+
 func (terminalChunkValidator) Name() string { return "terminal/chunk-validator" }
 
 func (terminalChunkValidator) ExecutionClass() contracts.ExecutionClass {
@@ -168,6 +181,33 @@ func TestRunnerMaterializesAnnotatedPlanBeforeChunkValidation(t *testing.T) {
 	}
 }
 
+func TestRunnerForwardsDetachedRepairAttemptsToChunkValidators(t *testing.T) {
+	prepared, plan := preparedTerminalDebugPipeline(t)
+	prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan}
+	chunkValidator := &observingChunkValidator{}
+	serializedValidator := &observingSerializedValidator{}
+	chunkAttempts := 2
+	serializedAttempts := 3
+	chunkBinding := Binding(chunkValidator.Name())
+	chunkBinding.StructuredOutputRepairAttempts = &chunkAttempts
+	serializedBinding := Binding(serializedValidator.Name())
+	serializedBinding.StructuredOutputRepairAttempts = &serializedAttempts
+	prepared.chunkValidators.validators = []preparedValidator{
+		{resolved: ResolvedValidator{Binding: chunkBinding, Target: ValidatorTargetChunk}, chunk: chunkValidator},
+		{resolved: ResolvedValidator{Binding: serializedBinding, Target: ValidatorTargetSerialized}, serialized: serializedValidator},
+	}
+
+	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
+		t.Fatalf("Run() error = %v, want nil", err)
+	}
+	if got := chunkValidator.request.StructuredOutputRepairAttempts; got == nil || *got != chunkAttempts || got == chunkBinding.StructuredOutputRepairAttempts {
+		t.Fatalf("chunk validator repair attempts = %v, want detached value %d", got, chunkAttempts)
+	}
+	if got := serializedValidator.request.StructuredOutputRepairAttempts; got == nil || *got != serializedAttempts || got == serializedBinding.StructuredOutputRepairAttempts {
+		t.Fatalf("serialized validator repair attempts = %v, want detached value %d", got, serializedAttempts)
+	}
+}
+
 func TestRunnerRetriesMalformedPlanWithDebugEnabled(t *testing.T) {
 	prepared, plan := preparedTerminalDebugPipeline(t)
 	plan.Ranges[0].Annotations = source.ChunkAnnotations{"broken": []byte(`{"value":`)}

internal/framework/pipeline/typed_registry_request_test.go

@@ -0,0 +1,95 @@
+package pipeline
+
+import (
+	"context"
+	"testing"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+)
+
+type repairObservingNotesMerger struct {
+	attempts *int
+}
+
+func (*repairObservingNotesMerger) Key() string { return "test/repair-observing-merge" }
+
+func (m *repairObservingNotesMerger) Merge(_ context.Context, request contracts.TypedMergeRequest[codecNotes]) (contracts.TypedMergeResult[codecNotes], error) {
+	m.attempts = request.StructuredOutputRepairAttempts
+	return contracts.TypedMergeResult[codecNotes]{Value: codecNotes{Items: []string{"merged"}}}, nil
+}
+
+type repairObservingNotesNormalizer struct {
+	attempts *int
+}
+
+func (*repairObservingNotesNormalizer) Key() string { return "test/repair-observing-normalize" }
+func (*repairObservingNotesNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
+	return nil
+}
+
+func (n *repairObservingNotesNormalizer) Normalize(_ context.Context, request contracts.TypedNormalizeRequest[codecNotes]) (contracts.TypedNormalizeResult[codecNotes], error) {
+	n.attempts = request.StructuredOutputRepairAttempts
+	return contracts.TypedNormalizeResult[codecNotes]{Value: request.MergeOutput.Value}, nil
+}
+
+func TestTypedRegistryErasurePreservesStructuredOutputRepairAttempts(t *testing.T) {
+	const artifactKind contracts.ArtifactKind = "test/notes"
+
+	t.Run("merge", func(t *testing.T) {
+		implementation := &repairObservingNotesMerger{}
+		registry := NewMergerRegistry()
+		if err := RegisterMerger(registry, ModuleSpec{Key: implementation.Key(), Stage: StageMerge, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Merger[codecNotes], error) {
+			return implementation, nil
+		}); err != nil {
+			t.Fatalf("RegisterMerger() error = %v", err)
+		}
+		entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
+		if !ok {
+			t.Fatal("typed merger entry missing")
+		}
+		built, err := entry.builder(BuildRequest{})
+		if err != nil {
+			t.Fatalf("builder() error = %v", err)
+		}
+		attempts := 2
+		_, err = entry.merge(context.Background(), built, contracts.TypedMergeRequest[any]{
+			ExtractOutputs:                 []contracts.ExtractArtifact[any]{{Value: codecNotes{Items: []string{"extracted"}}}},
+			StructuredOutputRepairAttempts: &attempts,
+		})
+		if err != nil {
+			t.Fatalf("merge() error = %v", err)
+		}
+		if implementation.attempts == nil || *implementation.attempts != attempts {
+			t.Fatalf("repair attempts = %v, want %d", implementation.attempts, attempts)
+		}
+	})
+
+	t.Run("normalize", func(t *testing.T) {
+		implementation := &repairObservingNotesNormalizer{}
+		registry := NewNormalizerRegistry()
+		if err := RegisterNormalizer(registry, ModuleSpec{Key: implementation.Key(), Stage: StageNormalize, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Normalizer[codecNotes], error) {
+			return implementation, nil
+		}); err != nil {
+			t.Fatalf("RegisterNormalizer() error = %v", err)
+		}
+		entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
+		if !ok {
+			t.Fatal("typed normalizer entry missing")
+		}
+		built, err := entry.builder(BuildRequest{})
+		if err != nil {
+			t.Fatalf("builder() error = %v", err)
+		}
+		attempts := 3
+		_, err = entry.normalize(context.Background(), built, contracts.TypedNormalizeRequest[any]{
+			MergeOutput:                    contracts.MergeArtifact[any]{Value: codecNotes{Items: []string{"merged"}}},
+			StructuredOutputRepairAttempts: &attempts,
+		})
+		if err != nil {
+			t.Fatalf("normalize() error = %v", err)
+		}
+		if implementation.attempts == nil || *implementation.attempts != attempts {
+			t.Fatalf("repair attempts = %v, want %d", implementation.attempts, attempts)
+		}
+	})
+}