Escape schema resources and reuse compiled plans

internal/usecase/runner.go

@@ -71,6 +71,11 @@ type preparationState struct {
 	start                time.Time
 }
 
+type preparedOperation struct {
+	run        *domain.PreparedRun
+	validation validate.PreparedValidation
+}
+
 func NewRunner(
 	promptDefs promptdef.Repository,
 	profiles profile.Repository,
@@ -137,18 +142,20 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
 	}
 	defer release()
 
-	prepared, err := r.completePreparation(ctx, req, state)
+	operation, err := r.completePreparation(ctx, req, state)
 	if err != nil {
 		return nil, err
 	}
 	directAPIKey := state.effectiveModel.APIKey
 
-	return r.executePreparedRun(ctx, prepared, directAPIKey, runID, start, func(
+	return r.executePreparedRun(ctx, operation.run, directAPIKey, runID, start, func(
 		ctx context.Context,
 		artifact *domain.Artifact,
 		attemptsUsed int,
 	) (domain.ValidationResult, error) {
-		return r.validateOutput(ctx, artifact, prepared.OutputContract, attemptsUsed)
+		result, err := operation.validation.Validate(ctx, artifact)
+		result.RepairAttempts = attemptsUsed
+		return result, err
 	})
 }
 
@@ -250,7 +257,11 @@ func (r *Runner) Prepare(ctx context.Context, req domain.RunRequest) (*domain.Pr
 	if err != nil {
 		return nil, err
 	}
-	return r.completePreparation(ctx, req, state)
+	operation, err := r.completePreparation(ctx, req, state)
+	if err != nil {
+		return nil, err
+	}
+	return operation.run, nil
 }
 
 func (r *Runner) resolvePreparation(
@@ -315,16 +326,64 @@ func (r *Runner) completePreparation(
 	ctx context.Context,
 	req domain.RunRequest,
 	state *preparationState,
-) (*domain.PreparedRun, error) {
-	structuredOutput, err := r.resolveStructuredOutput(
-		ctx,
+) (*preparedOperation, error) {
+	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
 	}
-	return r.completePreparationWithStructuredOutput(ctx, req, state, structuredOutput)
+	prepared, err := r.completePreparationWithStructuredOutput(ctx, req, state, structuredOutput)
+	if err != nil {
+		return nil, err
+	}
+	return &preparedOperation{run: prepared, validation: validationPlan}, 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
 }
 
 func (r *Runner) completePreparationWithStructuredOutput(
@@ -398,24 +457,6 @@ func (r *Runner) admitRun(ctx context.Context, backendID string) (func(), error)
 	return release, nil
 }
 
-func (r *Runner) resolveStructuredOutput(ctx context.Context, def *domain.PromptDefinition, contract domain.OutputContract) (*domain.StructuredOutputSpec, error) {
-	if contract.ValidationMode != domain.ValidationJSONSchema {
-		return nil, nil
-	}
-
-	loader, ok := r.validator.(validate.SchemaDocumentLoader)
-	if !ok || loader == nil {
-		return nil, fmt.Errorf("%w: json_schema output requires schema document loader", ErrValidation)
-	}
-
-	schemaDoc, err := loader.LoadSchemaDocument(ctx, contract.SchemaPath)
-	if err != nil {
-		return nil, fmt.Errorf("%w: failed to load json schema for structured output: %v", ErrValidation, err)
-	}
-
-	return structuredOutputSpec(def, schemaDoc), nil
-}
-
 func structuredOutputSpec(def *domain.PromptDefinition, schemaDocument any) *domain.StructuredOutputSpec {
 	return &domain.StructuredOutputSpec{
 		Type: domain.StructuredOutputJSONSchema,
@@ -453,25 +494,6 @@ func deriveStructuredSchemaName(promptID string, promptVersion string) string {
 	return name
 }
 
-func (r *Runner) validateOutput(ctx context.Context, artifact *domain.Artifact, contract domain.OutputContract, attemptsUsed int) (domain.ValidationResult, error) {
-	if r.validator == nil || contract.ValidationMode == domain.ValidationNone {
-		return domain.ValidationResult{
-			Status:         domain.ValidationSkipped,
-			Mode:           contract.ValidationMode,
-			SchemaPath:     contract.SchemaPath,
-			RepairAttempts: attemptsUsed,
-			IsValid:        true,
-		}, nil
-	}
-
-	res, err := r.validator.Validate(ctx, artifact, contract)
-	if err != nil {
-		return domain.ValidationResult{}, err
-	}
-	res.RepairAttempts = attemptsUsed
-	return res, nil
-}
-
 func (r *Runner) shouldAttemptRepair(contract domain.OutputContract, validationResult domain.ValidationResult) bool {
 	if r.repairer == nil {
 		return false