notarius/internal/framework/pipeline/runner_concurrent.go

package pipeline

import (
	"context"
	"errors"
	"fmt"
	"path"
	"sort"
	"strings"
	"sync"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

type laneExtractState struct {
	index      int
	prepared   preparedLaneExecutor
	deps       []CheckpointFingerprint
	decision   CheckpointDecision
	values     []erasedExtractArtifact
	serialized []CheckpointArtifact
	warnings   []contracts.Warning
	rejected   []contracts.RejectedOutput
	results    map[int]extractJobResult
	remaining  int
	failed     bool
}

type finalizedExtractResults struct {
	accepted   []erasedExtractArtifact
	serialized []CheckpointArtifact
	warnings   []contracts.Warning
	rejected   []contracts.RejectedOutput
	decision   CheckpointDecision
}

func loadExtract(loader CheckpointLoader, laneID, moduleKey string, deps []CheckpointFingerprint) (ExtractCheckpoint, CheckpointDecision) {
	return loader.Extract(laneID, moduleKey, deps)
}

func recordExtract(recorder CheckpointRecorder, laneID, moduleKey string, deps []CheckpointFingerprint, outputs []CheckpointArtifact, rejected []contracts.RejectedOutput, warnings []contracts.Warning) error {
	return recorder.ExtractSucceeded(laneID, moduleKey, deps, outputs, rejected, warnings)
}

type extractJob struct {
	lane  *laneExtractState
	chunk source.Chunk
}

type extractJobResult struct {
	laneIndex  int
	chunkIndex int
	value      erasedExtractArtifact
	serialized CheckpointArtifact
	warnings   []contracts.Warning
	rejected   *contracts.RejectedOutput
	err        error
}

type laneCompletion struct {
	index  int
	output RunOutput
	err    error
}

type orderedRunError struct {
	stage int
	lane  int
	chunk int
	err   error
}

func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints CheckpointRecorder, loader CheckpointLoader, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []source.Chunk) (RunOutput, error) {
	output := RunOutput{Manifest: manifestFromPipeline(input)}
	states := make([]*laneExtractState, len(input.Prepared.lanes))
	for i, prepared := range input.Prepared.lanes {
		if prepared.typed == nil {
			return output, fmt.Errorf("typed lane %q executor is not prepared", prepared.resolved.ID)
		}
		setTypedLaneManifestMetadata(&output, prepared.resolved.ID, prepared.typed.extractor, prepared.typed.merger, prepared.typed.normalizer)
		state, err := prepareLaneExtract(input, loader, doc, chunks, i, prepared)
		if err != nil {
			return output, err
		}
		if !state.decision.Reused {
			if err := checkpoints.ExtractRunning(prepared.resolved.ID, prepared.resolved.Extract.Module, state.deps); err != nil {
				return output, fmt.Errorf("write extract checkpoint for lane %q: %w", prepared.resolved.ID, err)
			}
		} else if err := finalizeLaneExtract(checkpoints, state); err != nil {
			return output, err
		}
		states[i] = state
	}

	workerCount := input.ExtractWorkers
	if workerCount < 1 {
		workerCount = 1
	}
	ctx, cancel := context.WithCancel(parent)
	defer cancel()
	jobs := make(chan extractJob, workerCount)
	results := make(chan extractJobResult, workerCount)
	completions := make(chan laneCompletion, len(states))
	continuations := make(chan *laneExtractState, workerCount)

	var workers sync.WaitGroup
	for i := 0; i < workerCount; i++ {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for job := range jobs {
				if ctx.Err() != nil {
					continue
				}
				result := r.runExtractJob(ctx, input, doc, sourceInput, sessionID, job)
				results <- result
			}
		}()
	}
	go func() {
		defer close(jobs)
		for chunkIndex := range chunks {
			for laneIndex := range states {
				state := states[laneIndex]
				if state.decision.Reused {
					continue
				}
				select {
				case jobs <- extractJob{lane: state, chunk: chunks[chunkIndex]}:
				case <-ctx.Done():
					return
				}
			}
		}
	}()
	go func() { workers.Wait(); close(results) }()
	var continuationWorkers sync.WaitGroup
	for i := 0; i < workerCount; i++ {
		continuationWorkers.Add(1)
		go func() {
			defer continuationWorkers.Done()
			for state := range continuations {
				if err := ctx.Err(); err != nil {
					completions <- laneCompletion{index: state.index, err: err}
					continue
				}
				laneOutput, err := r.continueLane(ctx, input, checkpoints, loader, doc, sourceInput, sessionID, chunks, state)
				completions <- laneCompletion{index: state.index, output: laneOutput, err: err}
			}
		}()
	}

	completedOutputs := make([]RunOutput, len(states))
	var runErrors []orderedRunError
	var pendingContinuations []*laneExtractState
	launched, completed := 0, 0
	for _, state := range states {
		if state.decision.Reused {
			pendingContinuations = append(pendingContinuations, state)
		}
	}

	resultChannel := results
	for resultChannel != nil || len(pendingContinuations) > 0 || completed < launched {
		var continuationChannel chan<- *laneExtractState
		var nextContinuation *laneExtractState
		if len(pendingContinuations) > 0 && ctx.Err() == nil {
			continuationChannel = continuations
			nextContinuation = pendingContinuations[0]
		} else if ctx.Err() != nil {
			pendingContinuations = nil
		}
		select {
		case continuationChannel <- nextContinuation:
			pendingContinuations = pendingContinuations[1:]
			launched++
		case result, ok := <-resultChannel:
			if !ok {
				resultChannel = nil
				continue
			}
			state := states[result.laneIndex]
			state.remaining--
			if result.err != nil {
				state.failed = true
				runErrors = append(runErrors, orderedRunError{stage: 0, lane: state.index, chunk: result.chunkIndex, err: result.err})
				_ = checkpoints.ExtractFailed(state.prepared.resolved.ID, state.prepared.resolved.Extract.Module, state.deps, result.err)
				cancel()
			} else {
				state.results[result.chunkIndex] = result
			}
			if state.remaining == 0 && !state.failed && ctx.Err() == nil {
				if err := finalizeLaneExtract(checkpoints, state); err != nil {
					state.failed = true
					runErrors = append(runErrors, orderedRunError{stage: 0, lane: state.index, chunk: len(chunks), err: err})
					cancel()
				} else {
					pendingContinuations = append(pendingContinuations, state)
				}
			}
		case completion := <-completions:
			completed++
			completedOutputs[completion.index] = completion.output
			if completion.err != nil {
				runErrors = append(runErrors, classifyLaneError(completion.index, len(chunks), completion.err))
				cancel()
			}
		}
	}
	close(continuations)
	continuationWorkers.Wait()
	for i := range completedOutputs {
		mergeLaneOutput(&output, completedOutputs[i])
	}
	if err := selectRunError(parent, runErrors); err != nil {
		return output, err
	}
	return output, nil
}

func prepareLaneExtract(input RunInput, loader CheckpointLoader, doc *source.SourceDocument, chunks []source.Chunk, index int, prepared preparedLaneExecutor) (*laneExtractState, error) {
	lane, typed := prepared.resolved, prepared.typed
	digest, err := joinedChunkDigest(chunks)
	if err != nil {
		return nil, fmt.Errorf("digest chunks for lane %q: %w", lane.ID, err)
	}
	state := &laneExtractState{index: index, prepared: prepared, deps: digestFingerprints("chunks", digest), remaining: len(chunks), results: make(map[int]extractJobResult, len(chunks))}
	cp, decision := loadExtract(loader, lane.ID, lane.Extract.Module, state.deps)
	if decision.Reused {
		for _, stored := range cp.Outputs {
			if _, decodeErr := decodeCheckpointArtifact(typed.codec, stored); decodeErr != nil {
				decision = CheckpointDecision{Reason: "extract artifact checkpoint codec is incompatible: " + decodeErr.Error()}
				break
			}
		}
	}
	state.decision = decision
	if decision.Reused {
		state.remaining = 0
		for _, stored := range cp.Outputs {
			value, decodeErr := decodeCheckpointArtifact(typed.codec, stored)
			if decodeErr != nil {
				return nil, fmt.Errorf("decode extract checkpoint for lane %q: %w", lane.ID, decodeErr)
			}
			stored = hydrateCheckpointArtifact(typed.codec, stored, value)
			artifact := erasedExtractArtifact{LaneID: lane.ID, ExtractorKey: lane.Extract.Module, SourceID: doc.ID, ChunkID: stored.ChunkID, ChunkIndex: stored.ChunkIndex, ChunkRef: stored.ChunkRef, Value: value}
			if stored.ChunkIndex >= 0 && stored.ChunkIndex < len(chunks) && artifact.ChunkRef == (source.SourceRef{}) {
				artifact.ChunkRef = chunks[stored.ChunkIndex].Ref
			}
			state.values = append(state.values, artifact)
			state.serialized = append(state.serialized, cloneCheckpointArtifact(stored))
		}
		state.warnings, state.rejected = cloneWarnings(cp.Warnings), cloneRejectedOutputs(cp.Rejected)
	}
	return state, nil
}

func (r *Runner) runExtractJob(ctx context.Context, input RunInput, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, job extractJob) extractJobResult {
	state, chunk := job.lane, cloneSourceChunk(job.chunk)
	lane, typed := state.prepared.resolved, state.prepared.typed
	result := extractJobResult{laneIndex: state.index, chunkIndex: chunk.Index}
	var accepted erasedExtractArtifact
	var serialized CheckpointArtifact
	var acceptedWarnings []contracts.Warning
	ok, rejection, err := runWithRetry(ctx, lane.Extract.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
		started := time.Now().UTC()
		attemptPath := path.Join("extract", debugPathComponent(lane.ID), fmt.Sprintf("chunk-%06d", chunk.Index+1), fmt.Sprintf("attempt-%02d", attempt))
		attemptCtx, llmScope := withDebugLLMScope(ctx, attemptPath)
		terminal := newAttemptTerminalRecorder(input.Debug, attemptPath, "extract", llmScope, debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, Attempt: attempt, StartedAt: started})
		extracted, callErr := typed.extract(attemptCtx, typed.extractor, contracts.TypedExtractionRequest{Source: doc, Chunk: &chunk, SourceInput: chunkInputMaterial(sourceInput, chunk), SessionID: sessionID, References: CloneReferenceSet(lane.ExtractReferences.ReferenceSet), LLMProfile: lane.Extract.LLMProfile, Metadata: cloneMetadata(input.Metadata)})
		if callErr != nil {
			attemptErr := fmt.Errorf("extract lane %q chunk %q with extractor %q: %w", lane.ID, chunk.ID, lane.Extract.Module, callErr)
			return false, nil, terminal.record(nil, attemptErr)
		}
		artifact := erasedExtractArtifact{LaneID: lane.ID, ExtractorKey: lane.Extract.Module, SourceID: doc.ID, ChunkID: chunk.ID, ChunkIndex: chunk.Index, ChunkRef: chunk.Ref, Value: extracted.Value}
		attemptWarnings := cloneWarnings(extracted.Warnings)
		serializedCandidate, encodeErr := serializeCandidateArtifact(typed.codec, artifact.LaneID, artifact.ExtractorKey, artifact.SourceID, artifact.Value)
		if encodeErr != nil {
			attemptErr := fmt.Errorf("serialize extract candidate for lane %q chunk %q: %w", lane.ID, chunk.ID, encodeErr)
			payload := map[string]any{"warnings": debugWarningEnvelopes(attemptWarnings)}
			return false, nil, terminal.record(payload, attemptErr)
		}
		serializedCandidate.ChunkID, serializedCandidate.ChunkIndex, serializedCandidate.ChunkRef = artifact.ChunkID, artifact.ChunkIndex, artifact.ChunkRef
		warnings, rejected, validateErr := r.validateTypedArtifact(attemptCtx, typed.codec, typedValidationTarget{stage: StageExtract, laneID: lane.ID, moduleKey: lane.Extract.Module, source: doc, sourceID: doc.ID, sourceInput: chunkInputMaterial(sourceInput, chunk), sessionID: sessionID, references: lane.ExtractReferences.ReferenceSet, metadata: input.Metadata, chunk: &chunk, ref: chunk.Ref, value: extracted.Value, candidate: &serializedCandidate}, state.prepared.extractValidators, attempt, input.Debug)
		attemptWarnings = append(attemptWarnings, warnings...)
		payload := map[string]any{"output": debugCheckpointArtifact(serializedCandidate), "warnings": debugWarningEnvelopes(attemptWarnings), "rejection": debugRejectedOutputPtr(rejected)}
		if validateErr != nil || rejected != nil {
			return false, rejected, terminal.record(payload, validateErr)
		}
		stored, encodeErr := checkpointArtifact(typed.codec, artifact.LaneID, artifact.ExtractorKey, artifact.SourceID, artifact.Value)
		if encodeErr != nil {
			attemptErr := fmt.Errorf("serialize accepted extract output for lane %q chunk %q: %w", lane.ID, chunk.ID, encodeErr)
			return false, nil, terminal.record(payload, attemptErr)
		}
		stored.ChunkID, stored.ChunkIndex, stored.ChunkRef = artifact.ChunkID, artifact.ChunkIndex, artifact.ChunkRef
		accepted, serialized = artifact, stored
		acceptedWarnings = attemptWarnings
		if debugErr := terminal.record(payload, nil); debugErr != nil {
			return false, nil, debugErr
		}
		return true, nil, nil
	})
	result.err = err
	if err == nil && !ok {
		result.rejected = rejection
		return result
	}
	result.value, result.serialized, result.warnings = accepted, serialized, acceptedWarnings
	return result
}

func finalizeLaneExtract(checkpoints CheckpointRecorder, state *laneExtractState) error {
	lane := state.prepared.resolved
	indexes := make([]int, 0, len(state.results))
	for index := range state.results {
		indexes = append(indexes, index)
	}
	sort.Ints(indexes)
	for _, index := range indexes {
		result := state.results[index]
		if result.rejected != nil {
			state.rejected = append(state.rejected, *result.rejected)
			continue
		}
		state.values = append(state.values, result.value)
		state.serialized = append(state.serialized, result.serialized)
		state.warnings = append(state.warnings, result.warnings...)
	}
	sort.SliceStable(state.values, func(i, j int) bool { return state.values[i].ChunkIndex < state.values[j].ChunkIndex })
	sort.SliceStable(state.serialized, func(i, j int) bool { return state.serialized[i].ChunkIndex < state.serialized[j].ChunkIndex })
	sort.SliceStable(state.rejected, func(i, j int) bool { return state.rejected[i].ChunkIndex < state.rejected[j].ChunkIndex })
	if !state.decision.Reused {
		if err := recordExtract(checkpoints, lane.ID, lane.Extract.Module, state.deps, state.serialized, state.rejected, state.warnings); err != nil {
			return fmt.Errorf("write extract checkpoint for lane %q: %w", lane.ID, err)
		}
	}
	return nil
}

func (r *Runner) continueLane(ctx context.Context, input RunInput, checkpoints CheckpointRecorder, loader CheckpointLoader, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []source.Chunk, state *laneExtractState) (RunOutput, error) {
	lane := state.prepared.resolved
	local := RunOutput{Manifest: manifestFromPipeline(input)}
	results := finalizedExtractResults{
		accepted:   state.values,
		serialized: state.serialized,
		warnings:   state.warnings,
		rejected:   state.rejected,
		decision:   state.decision,
	}
	local.Warnings = append(local.Warnings, cloneWarnings(results.warnings)...)
	local.Rejected = append(local.Rejected, cloneRejectedOutputs(results.rejected)...)
	recordCheckpointEvent(&local, loader, string(StageExtract), lane.ID, lane.Extract.Module, results.decision)
	if err := writeDebugTimed(input.Debug, path.Join("extract", debugPathComponent(lane.ID), "input.json"), debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": results.decision.Reused, "decision": results.decision, "source": debugSourceDocumentEnvelope(doc), "chunks": debugSourceChunkEnvelopes(chunks), "options": redactSensitiveMap(lane.Extract.Options), "metadata": redactSensitiveMap(input.Metadata)}}); err != nil {
		return local, &laneRunError{stage: StageExtract, err: err}
	}
	if err := writeDebugTimed(input.Debug, path.Join("extract", debugPathComponent(lane.ID), "output.json"), debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": results.decision.Reused, "outputs": debugCheckpointArtifacts(results.serialized), "rejected": debugRejectedOutputEnvelopes(results.rejected), "warnings": debugWarningEnvelopes(results.warnings)}}); err != nil {
		return local, &laneRunError{stage: StageExtract, err: err}
	}
	if len(results.accepted) == 0 {
		return local, nil
	}
	err := r.continueTypedLane(ctx, input, checkpoints, loader, doc, sourceInput, sessionID, state.prepared, results, &local)
	return local, err
}

func classifyLaneError(lane, sentinel int, err error) orderedRunError {
	stage := 1
	var laneErr *laneRunError
	if errors.As(err, &laneErr) {
		switch laneErr.stage {
		case StageExtract:
			stage = 0
		case StageNormalize:
			stage = 2
		}
	} else if strings.Contains(strings.ToLower(err.Error()), "normalize") {
		stage = 2
	}
	return orderedRunError{stage: stage, lane: lane, chunk: sentinel, err: err}
}

func selectRunError(parent context.Context, values []orderedRunError) error {
	if err := parent.Err(); err != nil {
		return err
	}
	if len(values) == 0 {
		return nil
	}
	hasReal := false
	for _, value := range values {
		if !errors.Is(value.err, context.Canceled) && !errors.Is(value.err, context.DeadlineExceeded) {
			hasReal = true
			break
		}
	}
	filtered := values[:0]
	for _, value := range values {
		if hasReal && (errors.Is(value.err, context.Canceled) || errors.Is(value.err, context.DeadlineExceeded)) {
			continue
		}
		filtered = append(filtered, value)
	}
	sort.SliceStable(filtered, func(i, j int) bool {
		if filtered[i].stage != filtered[j].stage {
			return filtered[i].stage < filtered[j].stage
		}
		if filtered[i].lane != filtered[j].lane {
			return filtered[i].lane < filtered[j].lane
		}
		return filtered[i].chunk < filtered[j].chunk
	})
	return filtered[0].err
}

func mergeLaneOutput(dst *RunOutput, src RunOutput) {
	if dst == nil {
		return
	}
	dst.NormalizeOutputs = append(dst.NormalizeOutputs, cloneSerializedOutputs(src.NormalizeOutputs)...)
	dst.Rejected = append(dst.Rejected, cloneRejectedOutputs(src.Rejected)...)
	dst.Warnings = append(dst.Warnings, cloneWarnings(src.Warnings)...)
	dst.CheckpointEvents = append(dst.CheckpointEvents, src.CheckpointEvents...)
	for i := range dst.Manifest.ArtifactLanes {
		for j := range src.Manifest.ArtifactLanes {
			if dst.Manifest.ArtifactLanes[i].ID == src.Manifest.ArtifactLanes[j].ID && src.Manifest.ArtifactLanes[j].Metadata != nil {
				dst.Manifest.ArtifactLanes[i].Metadata = cloneMetadata(src.Manifest.ArtifactLanes[j].Metadata)
			}
		}
	}
}