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
	record    *ChunkPlanRecord
	action    string
	summary   artifacts.ChunkPlanSummary
}

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"},
		summary: artifacts.ChunkPlanSummary{
			Mode: string(mode), SourceDigest: doc.Digest, RequestedModule: chunker.Key(), LookupStatus: "skipped",
			LookupReason: "chunk plan lookup skipped", ValidationStatus: "not_run", PublicationStatus: "not_requested",
		},
	}

	if mode == ChunkCacheAuto {
		record, decision, err := input.ChunkPlans.Load(doc.Digest)
		result.lookup = decision
		result.summary.LookupStatus = chunkPlanLookupStatus(decision.Status)
		result.summary.LookupReason = chunkPlanLookupReason(decision.Status)
		if err != nil {
			result.summary.LookupStatus = "skipped"
			result.summary.LookupReason = chunkPlanLookupReason("")
			return result, fmt.Errorf("load chunk plan: %w", err)
		}
		switch decision.Status {
		case ChunkPlanHit:
			plan, chunks, validationErr := validateAndMaterializeChunkPlan(doc, record.Plan)
			if validationErr == nil {
				if err := result.setCandidate(record, "reused"); err != nil {
					return result, fmt.Errorf("clone reused chunk plan record: %w", err)
				}
				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
				result.setValidation(validationWarnings, rejection, err)
				return result, err
			}
			result.lookup = ChunkPlanDecision{Status: ChunkPlanInvalid, Reason: chunkPlanLookupReason(ChunkPlanInvalid)}
			result.summary.LookupStatus = "invalid"
			result.summary.LookupReason = chunkPlanLookupReason(ChunkPlanInvalid)
		case ChunkPlanMissing, ChunkPlanInvalid:
			// Generate below.
		default:
			return result, fmt.Errorf("load chunk plan returned unsupported decision status")
		}
	}
	if mode == ChunkCacheAuto || mode == ChunkCacheRefresh {
		result.summary.PublicationStatus = "not_published"
	}

	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})
		requestMetadata, metadataErr := cloneMetadata(input.Metadata)
		if metadataErr != nil {
			return false, nil, terminal.record(nil, fmt.Errorf("clone chunk request metadata: %w", metadataErr))
		}
		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: requestMetadata,
		})
		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)
		}
		planDigest, digestErr := source.DigestChunkPlan(plan)
		if digestErr != nil {
			return false, nil, terminal.record(nil, fmt.Errorf("digest generated chunk plan: %w", digestErr))
		}
		producerMetadata, _, metadataErr := moduleManifestMetadata(chunker)
		if metadataErr != nil {
			return false, nil, terminal.record(nil, fmt.Errorf("clone chunker manifest metadata: %w", metadataErr))
		}
		profile := ""
		if input.pipeline.ChunkExecutionClass == contracts.ExecutionClassLLMBacked {
			profile = input.pipeline.Chunk.LLMProfile
		}
		candidate := ChunkPlanRecord{
			SchemaVersion: ChunkPlanSchemaVersion, SourceDigest: doc.Digest, PlanDigest: planDigest,
			Plan: source.CloneChunkPlan(plan),
			Producer: ChunkPlanProducer{
				InputModule: input.Prepared.input.Key(), ChunkModule: chunker.Key(), LLMProfile: profile,
				References: append([]artifacts.ReferenceProvenance(nil), referenceTargetProvenance(input.pipeline.ChunkReferences)...),
				Metadata:   producerMetadata,
			},
			Warnings: cloneWarnings(chunkResult.Warnings), CreatedAt: time.Now().UTC(),
		}
		action := "generated"
		if mode == ChunkCacheRefresh {
			action = "refreshed"
		}
		if mode == ChunkCacheBypass {
			action = "bypassed"
		}
		if candidateErr := result.setCandidate(candidate, action); candidateErr != nil {
			return false, nil, terminal.record(nil, fmt.Errorf("clone generated chunk plan record: %w", candidateErr))
		}
		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 {
			result.setValidation(validationWarnings, rejected, validationErr)
			return false, rejected, terminal.record(payload, validationErr)
		}
		result.chunks = chunks
		result.plan = &plan
		result.warnings = attemptWarnings
		producerWarnings = cloneWarnings(chunkResult.Warnings)
		result.setValidation(validationWarnings, nil, nil)
		if debugErr := terminal.record(payload, nil); debugErr != nil {
			return false, nil, debugErr
		}
		return true, nil, nil
	})
	if err != nil {
		if result.summary.ValidationStatus == "not_run" {
			result.summary.ValidationStatus = "error"
		}
		return result, err
	}
	result.accepted = accepted
	result.rejection = rejection
	if !accepted {
		return result, nil
	}

	if mode == ChunkCacheAuto || mode == ChunkCacheRefresh {
		record, cloneErr := cloneChunkPlanRecord(*result.record)
		if cloneErr != nil {
			return result, fmt.Errorf("clone chunk plan record for publication: %w", cloneErr)
		}
		record.Warnings = cloneWarnings(producerWarnings)
		if err := input.ChunkPlans.Save(record); err != nil {
			return result, fmt.Errorf("save chunk plan: %w", err)
		}
		result.summary.PublicationStatus = "published"
	}
	return result, nil
}

func chunkPlanLookupStatus(status ChunkPlanStatus) string {
	switch status {
	case ChunkPlanHit:
		return "hit"
	case ChunkPlanMissing:
		return "missing"
	case ChunkPlanInvalid:
		return "invalid"
	default:
		return "skipped"
	}
}

func chunkPlanLookupReason(status ChunkPlanStatus) string {
	switch status {
	case ChunkPlanHit:
		return "stored chunk plan is valid"
	case ChunkPlanMissing:
		return "chunk plan not found"
	case ChunkPlanInvalid:
		return "stored chunk plan is invalid"
	default:
		return "chunk plan lookup skipped"
	}
}

func (result *chunkPlanExecution) setCandidate(record ChunkPlanRecord, action string) error {
	cloned, err := cloneChunkPlanRecord(record)
	if err != nil {
		return err
	}
	result.record = &cloned
	result.action = action
	result.summary.Action = action
	result.summary.SourceDigest = record.SourceDigest
	result.summary.CandidateDigest = record.PlanDigest
	return nil
}

func (result *chunkPlanExecution) setValidation(warnings []contracts.Warning, rejection *contracts.RejectedOutput, err error) {
	switch {
	case err != nil:
		result.summary.ValidationStatus = "error"
	case rejection != nil:
		result.summary.ValidationStatus = "rejected"
	case len(warnings) > 0:
		result.summary.ValidationStatus = "approved_with_warnings"
	default:
		result.summary.ValidationStatus = "approved"
	}
}

func cloneChunkPlanRecord(record ChunkPlanRecord) (ChunkPlanRecord, error) {
	record.Plan = source.CloneChunkPlan(record.Plan)
	record.Producer.References = append([]artifacts.ReferenceProvenance(nil), record.Producer.References...)
	metadata, err := cloneMetadata(record.Producer.Metadata)
	if err != nil {
		return ChunkPlanRecord{}, fmt.Errorf("clone chunk plan producer metadata: %w", err)
	}
	record.Producer.Metadata = metadata
	record.Warnings = cloneWarnings(record.Warnings)
	return record, nil
}

func applyChunkPlanExecution(output *RunOutput, result chunkPlanExecution) error {
	if output == nil {
		return nil
	}
	summary := result.summary
	output.ChunkPlan = &summary
	if output.Manifest.ChunkPlan == nil {
		return nil
	}
	manifest := output.Manifest.ChunkPlan
	manifest.Action = result.action
	if result.record == nil {
		return nil
	}
	record := result.record
	manifest.SourceDigest = record.SourceDigest
	manifest.PlanDigest = record.PlanDigest
	manifest.PlanSchemaVersion = record.SchemaVersion
	manifest.ProducerInputModule = record.Producer.InputModule
	manifest.ProducerModule = record.Producer.ChunkModule
	manifest.ProducerLLMProfile = record.Producer.LLMProfile
	manifest.ProducerReferences = append([]artifacts.ReferenceProvenance(nil), record.Producer.References...)
	metadata, err := cloneMetadata(record.Producer.Metadata)
	if err != nil {
		return fmt.Errorf("clone chunk plan manifest producer metadata: %w", err)
	}
	manifest.ProducerMetadata = metadata
	createdAt := record.CreatedAt
	manifest.CreatedAt = &createdAt
	return nil
}