promptkit/internal/usecase/prepared_execution.go

package usecase

import (
	"context"
	"fmt"
	"sync"
	"time"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/jsonvalue"
	"gitea.maximumdirect.net/eric/promptkit/internal/validate"
)

type preparedExecutionState uint8

const (
	preparedExecutionReady preparedExecutionState = iota
	preparedExecutionClaimed
	preparedExecutionDiscarded
)

// PreparedExecution owns one frozen, single-use runner execution.
type PreparedExecution struct {
	owner   *Runner
	mu      sync.Mutex
	state   preparedExecutionState
	details *domain.PreparedRun
	payload *preparedExecutionPayload
}

type preparedExecutionPayload struct {
	prepared   *domain.PreparedRun
	validation validate.PreparedValidation
	directKey  string
}

// PrepareExecution completes preparation without generation or admission and
// returns a runner-bound, single-use execution.
func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*PreparedExecution, error) {
	state, err := r.resolvePreparation(ctx, req, time.Now().UTC())
	if err != nil {
		return nil, err
	}

	validationPlan, err := r.prepareValidation(ctx, state.effectiveContract)
	if err != nil {
		return nil, err
	}

	structuredOutput, err := r.structuredOutputFromValidationPlan(
		state.definition,
		state.effectiveContract,
		validationPlan,
	)
	if err != nil {
		return nil, err
	}
	prepared, err := r.completePreparationWithStructuredOutput(ctx, req, state, structuredOutput)
	if err != nil {
		return nil, err
	}

	executionSnapshot, err := clonePreparedRun(prepared)
	if err != nil {
		return nil, fmt.Errorf("%w: failed to copy prepared execution: %v", ErrInvalidRequest, err)
	}

	details, err := clonePreparedRun(executionSnapshot)
	if err != nil {
		return nil, fmt.Errorf("%w: failed to copy prepared execution details: %v", ErrInvalidRequest, err)
	}

	return &PreparedExecution{
		owner:   r,
		state:   preparedExecutionReady,
		details: details,
		payload: &preparedExecutionPayload{
			prepared:   executionSnapshot,
			validation: validationPlan,
			directKey:  state.effectiveModel.APIKey,
		},
	}, nil
}

func (r *Runner) prepareValidation(
	ctx context.Context,
	contract domain.OutputContract,
) (validate.PreparedValidation, error) {
	if r.validator == nil {
		return noOpPreparedValidation{contract: contract}, nil
	}
	preparer, ok := r.validator.(validate.ValidationPreparer)
	if !ok {
		return nil, fmt.Errorf("%w: validator does not support prepared validation", ErrValidation)
	}
	plan, err := preparer.PrepareValidation(ctx, contract)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrValidation, err)
	}
	if plan == nil {
		return nil, fmt.Errorf("%w: validator returned nil prepared validation", ErrValidation)
	}
	return plan, nil
}

func (r *Runner) structuredOutputFromValidationPlan(
	def *domain.PromptDefinition,
	contract domain.OutputContract,
	plan validate.PreparedValidation,
) (*domain.StructuredOutputSpec, error) {
	if contract.ValidationMode != domain.ValidationJSONSchema {
		return nil, nil
	}

	schemaDocument := plan.SchemaDocument()
	if schemaDocument == nil {
		if r.validator == nil {
			return nil, nil
		}
		return nil, fmt.Errorf("%w: prepared json_schema validation has no schema document", ErrValidation)
	}
	return structuredOutputSpec(def, schemaDocument), nil
}

// Details returns a fresh credential-redacted copy of the prepared run.
func (p *PreparedExecution) Details() *domain.PreparedRun {
	if p == nil {
		return nil
	}
	p.mu.Lock()
	detailsSnapshot := p.details
	p.mu.Unlock()

	details, err := clonePreparedRun(detailsSnapshot)
	if err != nil {
		return nil
	}
	return details
}

// Discard invalidates an unclaimed execution and drops its private payload.
func (p *PreparedExecution) Discard() {
	if p == nil {
		return
	}

	p.mu.Lock()
	if p.state != preparedExecutionReady {
		p.mu.Unlock()
		return
	}
	p.state = preparedExecutionDiscarded
	payload := p.payload
	p.payload = nil
	p.mu.Unlock()

	payload.clear()
}

// RunPrepared claims and executes one prepared execution owned by this runner.
func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (*domain.RunResult, error) {
	payload, err := prepared.claim(r)
	if err != nil {
		return nil, err
	}
	defer payload.clear()

	runID, err := newRunID()
	if err != nil {
		return nil, fmt.Errorf("failed to create run id: %w", err)
	}
	start := time.Now().UTC()

	target := payload.prepared.EffectiveModelParams
	if err := validateAPIKey(target.APIKeyEnv, payload.directKey, target.APIKeyRequired); err != nil {
		return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
	}

	release, err := r.admitRun(ctx, target.BackendID)
	if err != nil {
		return nil, err
	}
	defer release()

	return r.executePreparedRun(ctx, payload.prepared, payload.directKey, runID, start, func(
		ctx context.Context,
		artifact *domain.Artifact,
		attemptsUsed int,
	) (domain.ValidationResult, error) {
		result, validationErr := payload.validation.Validate(ctx, artifact)
		if validationErr != nil {
			return domain.ValidationResult{}, validationErr
		}
		result.RepairAttempts = attemptsUsed
		return result, nil
	})
}

func (p *PreparedExecution) claim(owner *Runner) (*preparedExecutionPayload, error) {
	if p == nil || owner == nil || p.owner != owner {
		return nil, fmt.Errorf("%w: prepared execution does not belong to this runner", ErrInvalidRequest)
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if p.state != preparedExecutionReady || p.payload == nil {
		return nil, fmt.Errorf("%w: prepared execution is not ready", ErrInvalidRequest)
	}
	p.state = preparedExecutionClaimed
	payload := p.payload
	p.payload = nil
	return payload, nil
}

func (p *preparedExecutionPayload) clear() {
	if p == nil {
		return
	}
	if p.prepared != nil {
		p.prepared.EffectiveModelParams.APIKey = ""
	}
	p.prepared = nil
	p.validation = nil
	p.directKey = ""
}

type noOpPreparedValidation struct {
	contract domain.OutputContract
}

func (p noOpPreparedValidation) Validate(
	ctx context.Context,
	_ *domain.Artifact,
) (domain.ValidationResult, error) {
	if err := ctx.Err(); err != nil {
		return domain.ValidationResult{}, err
	}
	return domain.ValidationResult{
		Status:         domain.ValidationSkipped,
		Mode:           p.contract.ValidationMode,
		SchemaPath:     p.contract.SchemaPath,
		RepairAttempts: p.contract.RepairAttempts,
		IsValid:        true,
	}, nil
}

func (noOpPreparedValidation) SchemaDocument() any {
	return nil
}

func clonePreparedRun(source *domain.PreparedRun) (*domain.PreparedRun, error) {
	if source == nil {
		return nil, nil
	}

	copied := *source
	extraParams, err := jsonvalue.CopyMap(source.EffectiveModelParams.ExtraParams)
	if err != nil {
		return nil, err
	}
	copied.EffectiveModelParams.ExtraParams = extraParams

	if source.InputHashes != nil {
		copied.InputHashes = make(map[string]string, len(source.InputHashes))
		for name, hash := range source.InputHashes {
			copied.InputHashes[name] = hash
		}
	}
	if source.Messages != nil {
		copied.Messages = make([]domain.RenderedMessage, len(source.Messages))
		for i, message := range source.Messages {
			copied.Messages[i] = message
			if message.CacheControl != nil {
				cacheControl := *message.CacheControl
				copied.Messages[i].CacheControl = &cacheControl
			}
		}
	}

	if source.StructuredOutput != nil {
		structuredOutput := *source.StructuredOutput
		copied.StructuredOutput = &structuredOutput
		if source.StructuredOutput.JSONSchema != nil {
			jsonSchema := *source.StructuredOutput.JSONSchema
			copied.StructuredOutput.JSONSchema = &jsonSchema
			schema, copyErr := cloneJSONValue(source.StructuredOutput.JSONSchema.Schema)
			if copyErr != nil {
				return nil, copyErr
			}
			copied.StructuredOutput.JSONSchema.Schema = schema
		}
	}
	return &copied, nil
}

func cloneJSONValue(source any) (any, error) {
	return jsonvalue.Copy(source)
}