notarius/internal/framework/pipeline/runner_chunk_plan.go

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
}