Escape schema resources and reuse compiled plans

internal/usecase/prepared_execution.go

@@ -42,25 +42,12 @@ func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*
 		return nil, err
 	}
 
-	validationPlan, err := r.prepareValidation(ctx, state.effectiveContract)
+	operation, err := r.completePreparation(ctx, req, state)
 	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)
+	executionSnapshot, err := clonePreparedRun(operation.run)
 	if err != nil {
 		return nil, fmt.Errorf("%w: failed to copy prepared execution: %v", ErrInvalidRequest, err)
 	}
@@ -76,52 +63,12 @@ func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*
 		details: details,
 		payload: &preparedExecutionPayload{
 			prepared:   executionSnapshot,
-			validation: validationPlan,
+			validation: operation.validation,
 			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 {

internal/usecase/prepared_execution_test.go

@@ -52,6 +52,12 @@ type recordingValidationPreparer struct {
 	directValidateCalls int
 }
 
+type validationOnly struct{}
+
+func (validationOnly) Validate(context.Context, *domain.Artifact, domain.OutputContract) (domain.ValidationResult, error) {
+	return domain.ValidationResult{}, nil
+}
+
 func (v *recordingValidationPreparer) Validate(
 	context.Context,
 	*domain.Artifact,
@@ -532,7 +538,7 @@ func TestRunnerPrepareExecutionRequiresValidationPreparer(t *testing.T) {
 		reader,
 		defaultRenderer(),
 		&fakeLLM{forbid: true},
-		&fakeValidator{},
+		validationOnly{},
 		nil,
 	)
 

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

internal/usecase/runner_test.go

@@ -191,16 +191,34 @@ func (f *fakeValidator) Validate(ctx context.Context, artifact *domain.Artifact,
 	return f.result, nil
 }
 
-func (f *fakeValidator) LoadSchemaDocument(ctx context.Context, schemaPath string) (any, error) {
-	f.schemaLoads++
-	f.schemaLoadPath = schemaPath
-	if f.schemaErr != nil {
-		return nil, f.schemaErr
+func (f *fakeValidator) PrepareValidation(_ context.Context, contract domain.OutputContract) (validate.PreparedValidation, error) {
+	var schemaDocument any
+	if contract.ValidationMode == domain.ValidationJSONSchema {
+		f.schemaLoads++
+		f.schemaLoadPath = contract.SchemaPath
+		if f.schemaErr != nil {
+			return nil, f.schemaErr
+		}
+		schemaDocument = f.schemaDoc
+		if schemaDocument == nil {
+			schemaDocument = map[string]any{"type": "object"}
+		}
 	}
-	if f.schemaDoc != nil {
-		return f.schemaDoc, nil
-	}
-	return map[string]any{"type": "object"}, nil
+	return &fakePreparedValidator{validator: f, contract: contract, schemaDocument: schemaDocument}, nil
+}
+
+type fakePreparedValidator struct {
+	validator      *fakeValidator
+	contract       domain.OutputContract
+	schemaDocument any
+}
+
+func (p *fakePreparedValidator) Validate(ctx context.Context, artifact *domain.Artifact) (domain.ValidationResult, error) {
+	return p.validator.Validate(ctx, artifact, p.contract)
+}
+
+func (p *fakePreparedValidator) SchemaDocument() any {
+	return p.schemaDocument
 }
 
 type fakeRepairer struct {
@@ -2278,6 +2296,9 @@ func TestRunnerRunJSONSchemaRepairCarriesStructuredOutputSpec(t *testing.T) {
 	if repairer.reqs[0].StructuredOutput.JSONSchema.Name != "p_1" {
 		t.Fatalf("expected derived schema name p_1, got %q", repairer.reqs[0].StructuredOutput.JSONSchema.Name)
 	}
+	if validator.schemaLoads != 1 || validator.validateCalls != 2 {
+		t.Fatalf("schema preparation/validation calls = (%d, %d), want (1, 2)", validator.schemaLoads, validator.validateCalls)
+	}
 }
 
 func TestExecutionProfileToTargetPopulatesAllFieldsAndCopiesExtraParams(t *testing.T) {