Integrate chunk plan caching into the runner

internal/framework/pipeline/runner_chunk_plan.go

@@ -0,0 +1,137 @@
+package pipeline
+
+import (
+	"context"
+	"fmt"
+	"path"
+	"time"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
+	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+)
+
+type chunkPlanExecution struct {
+	accepted  bool
+	chunks    []source.Chunk
+	plan      *source.ChunkPlan
+	warnings  []contracts.Warning
+	rejection *contracts.RejectedOutput
+	lookup    ChunkPlanDecision
+}
+
+func effectiveChunkCacheMode(mode ChunkCacheMode) ChunkCacheMode {
+	if mode == "" {
+		return ChunkCacheBypass
+	}
+	parsed, _ := ParseChunkCacheMode(string(mode))
+	return parsed
+}
+
+// Chunk-plan lookup and publication happen serially before lane workers start.
+// The runner therefore does not add synchronization around the store.
+func (r *Runner) runChunkPlan(ctx context.Context, input RunInput, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string) (chunkPlanExecution, error) {
+	mode := effectiveChunkCacheMode(input.ChunkCacheMode)
+	chunker := input.Prepared.chunker
+	result := chunkPlanExecution{lookup: ChunkPlanDecision{Reason: "chunk plan lookup skipped"}}
+
+	if mode == ChunkCacheAuto {
+		record, decision, err := input.ChunkPlans.Load(doc.Digest)
+		result.lookup = decision
+		if err != nil {
+			return result, fmt.Errorf("load chunk plan: %w", err)
+		}
+		switch decision.Status {
+		case ChunkPlanHit:
+			plan, chunks, validationErr := validateAndMaterializeChunkPlan(doc, record.Plan)
+			if validationErr == nil {
+				validationWarnings, rejection, err := r.validateChunks(ctx, doc, chunker.Key(), chunks, sourceInput, sessionID, input.pipeline.ChunkReferences.ReferenceSet, input.Metadata, input.Prepared.chunkValidators, 1, input.Debug)
+				result.plan = &plan
+				result.chunks = chunks
+				result.warnings = append(cloneWarnings(record.Warnings), validationWarnings...)
+				result.rejection = rejection
+				result.accepted = rejection == nil && err == nil
+				return result, err
+			}
+			result.lookup = ChunkPlanDecision{Status: ChunkPlanInvalid, Reason: "stored chunk plan cannot be materialized against the current source"}
+		case ChunkPlanMissing, ChunkPlanInvalid:
+			// Generate below.
+		default:
+			return result, fmt.Errorf("load chunk plan returned unsupported status %q", decision.Status)
+		}
+	}
+
+	var producerWarnings []contracts.Warning
+	accepted, rejection, err := runWithRetry(ctx, input.pipeline.Chunk.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
+		attemptStarted := time.Now().UTC()
+		attemptPath := path.Join("chunk", fmt.Sprintf("attempt-%02d", attempt))
+		attemptCtx, llmScope := withDebugLLMScope(ctx, attemptPath)
+		terminal := newAttemptTerminalRecorder(input.Debug, attemptPath, "chunk", llmScope, debugTimedEnvelope{Stage: string(StageChunk), ModuleKey: chunker.Key(), Attempt: attempt, StartedAt: attemptStarted})
+		chunkResult, callErr := chunker.Plan(attemptCtx, contracts.ChunkRequest{
+			Source: doc, SourceInput: sourceInput.Clone(), SessionID: sessionID,
+			References: CloneReferenceSet(input.pipeline.ChunkReferences.ReferenceSet),
+			LLMProfile: input.pipeline.Chunk.LLMProfile, Metadata: input.Metadata,
+		})
+		if callErr != nil {
+			return false, nil, terminal.record(nil, fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), callErr))
+		}
+		plan, chunks, validationErr := validateAndMaterializeChunkPlan(doc, chunkResult.Plan)
+		if validationErr != nil {
+			attemptErr := fmt.Errorf("validate chunk plan from chunker %q: %w", chunker.Key(), validationErr)
+			payload := map[string]any{"plan": debugChunkPlanEnvelope(chunkResult.Plan), "warnings": debugWarningEnvelopes(chunkResult.Warnings)}
+			return false, nil, terminal.record(payload, attemptErr)
+		}
+		validationWarnings, rejected, validationErr := r.validateChunks(attemptCtx, doc, chunker.Key(), chunks, sourceInput, sessionID, input.pipeline.ChunkReferences.ReferenceSet, input.Metadata, input.Prepared.chunkValidators, attempt, input.Debug)
+		attemptWarnings := append(cloneWarnings(chunkResult.Warnings), validationWarnings...)
+		payload := map[string]any{
+			"plan": debugChunkPlanEnvelope(plan), "materialized_chunks": debugSourceChunkEnvelopes(chunks),
+			"warnings": debugWarningEnvelopes(attemptWarnings), "rejection": debugRejectedOutputPtr(rejected),
+		}
+		if validationErr != nil || rejected != nil {
+			return false, rejected, terminal.record(payload, validationErr)
+		}
+		result.chunks = chunks
+		result.plan = &plan
+		result.warnings = attemptWarnings
+		producerWarnings = cloneWarnings(chunkResult.Warnings)
+		if debugErr := terminal.record(payload, nil); debugErr != nil {
+			return false, nil, debugErr
+		}
+		return true, nil, nil
+	})
+	if err != nil {
+		return result, err
+	}
+	result.accepted = accepted
+	result.rejection = rejection
+	if !accepted {
+		return result, nil
+	}
+
+	if mode == ChunkCacheAuto || mode == ChunkCacheRefresh {
+		planDigest, err := source.DigestChunkPlan(*result.plan)
+		if err != nil {
+			return result, fmt.Errorf("digest generated chunk plan: %w", err)
+		}
+		producerMetadata, _ := moduleManifestMetadata(chunker)
+		record := ChunkPlanRecord{
+			SchemaVersion: ChunkPlanSchemaVersion,
+			SourceDigest:  doc.Digest,
+			PlanDigest:    planDigest,
+			Plan:          source.CloneChunkPlan(*result.plan),
+			Producer: ChunkPlanProducer{
+				InputModule: input.Prepared.input.Key(),
+				ChunkModule: chunker.Key(),
+				LLMProfile:  input.pipeline.Chunk.LLMProfile,
+				References:  append([]artifacts.ReferenceProvenance(nil), referenceTargetProvenance(input.pipeline.ChunkReferences)...),
+				Metadata:    cloneMetadata(producerMetadata),
+			},
+			Warnings:  cloneWarnings(producerWarnings),
+			CreatedAt: time.Now().UTC(),
+		}
+		if err := input.ChunkPlans.Save(record); err != nil {
+			return result, fmt.Errorf("save chunk plan: %w", err)
+		}
+	}
+	return result, nil
+}