promptkit/internal/usecase/prepared_execution_test.go

package usecase

import (
	"context"
	"errors"
	"os"
	"reflect"
	"testing"

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

type recordingPreparedValidation struct {
	contract       domain.OutputContract
	schemaDocument any
	results        []domain.ValidationResult
	errs           []error
	artifacts      []string
}

func (p *recordingPreparedValidation) Validate(
	_ context.Context,
	artifact *domain.Artifact,
) (domain.ValidationResult, error) {
	p.artifacts = append(p.artifacts, string(artifact.Body))
	index := len(p.artifacts) - 1
	if index < len(p.errs) && p.errs[index] != nil {
		return domain.ValidationResult{}, p.errs[index]
	}
	if len(p.results) == 0 {
		return domain.ValidationResult{
			Status:  domain.ValidationPassed,
			Mode:    p.contract.ValidationMode,
			IsValid: true,
		}, nil
	}
	if index >= len(p.results) {
		index = len(p.results) - 1
	}
	return p.results[index], nil
}

func (p *recordingPreparedValidation) SchemaDocument() any {
	return p.schemaDocument
}

type recordingValidationPreparer struct {
	plan                *recordingPreparedValidation
	prepareErr          error
	prepareCalls        int
	directValidateCalls int
}

func (v *recordingValidationPreparer) Validate(
	context.Context,
	*domain.Artifact,
	domain.OutputContract,
) (domain.ValidationResult, error) {
	v.directValidateCalls++
	return domain.ValidationResult{}, errors.New("live validation must not be used")
}

func (v *recordingValidationPreparer) PrepareValidation(
	_ context.Context,
	contract domain.OutputContract,
) (validate.PreparedValidation, error) {
	v.prepareCalls++
	if v.prepareErr != nil {
		return nil, v.prepareErr
	}
	v.plan.contract = contract
	return v.plan, nil
}

func TestRunnerPrepareExecutionCompletesWithoutAdmissionOrGeneration(t *testing.T) {
	schemaDocument := map[string]any{
		"type": "object",
		"properties": map[string]any{
			"value": map[string]any{"type": "string"},
			"":      map[string]any{"type": "boolean"},
		},
	}
	def := promptDef(domain.FormatJSON, domain.ValidationJSONSchema, 0)
	def.Validation.SchemaPath = "schema.json"
	reader := defaultArtifactReader()
	renderer := &fakeRenderer{rendered: &domain.RenderedPrompt{
		SessionID: "prepared-session",
		Messages: []domain.RenderedMessage{{
			Role:    "user",
			Content: "original message",
		}},
	}}
	llmClient := &fakeLLM{forbid: true}
	validator := &recordingValidationPreparer{
		plan: &recordingPreparedValidation{schemaDocument: schemaDocument},
	}
	admitter := &fakeRunAdmitter{}
	profile := defaultExecutionProfile()
	profile.ExtraParams = map[string]any{
		"metadata": map[string]any{"source": "original"},
	}
	runner := NewRunner(
		&fakePromptRepo{def: def},
		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}},
		nil,
		reader,
		renderer,
		llmClient,
		validator,
		admitter,
	)

	prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
		PromptID:  "p",
		ProfileID: "exec",
		APIKey:    "direct-test-key",
		Inputs:    singleInputRef(),
	})
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}
	defer prepared.Discard()

	if validator.prepareCalls != 1 || validator.directValidateCalls != 0 {
		t.Fatalf(
			"validation calls=(prepare=%d direct=%d), want (1, 0)",
			validator.prepareCalls,
			validator.directValidateCalls,
		)
	}
	if reader.calls != 1 || renderer.calls != 1 {
		t.Fatalf("completion calls=(artifact=%d render=%d), want (1, 1)", reader.calls, renderer.calls)
	}
	if len(admitter.backendIDs) != 0 || llmClient.calls != 0 {
		t.Fatalf("prepare invoked execution collaborators: admission=%v generation=%d", admitter.backendIDs, llmClient.calls)
	}

	first := prepared.Details()
	if first == nil {
		t.Fatal("prepared details are nil")
	}
	if first.EffectiveModelParams.APIKey != "" {
		t.Fatal("prepared details retained the direct API key")
	}
	if first.StructuredOutput == nil ||
		first.StructuredOutput.JSONSchema == nil ||
		!reflect.DeepEqual(first.StructuredOutput.JSONSchema.Schema, schemaDocument) {
		t.Fatalf("prepared details have unexpected structured output: %#v", first.StructuredOutput)
	}

	first.Messages[0].Content = "caller mutation"
	first.InputHashes["input"] = "caller mutation"
	first.EffectiveModelParams.ExtraParams["metadata"].(map[string]any)["source"] = "caller mutation"
	first.StructuredOutput.JSONSchema.Schema.(map[string]any)["type"] = "string"
	renderer.rendered.Messages[0].Content = "source mutation"

	second := prepared.Details()
	if second.Messages[0].Content != "original message" ||
		second.InputHashes["input"] == "caller mutation" ||
		second.EffectiveModelParams.ExtraParams["metadata"].(map[string]any)["source"] != "original" ||
		second.StructuredOutput.JSONSchema.Schema.(map[string]any)["type"] != "object" {
		t.Fatalf("details did not preserve an independent snapshot: %#v", second)
	}
}

func TestRunnerPrepareExecutionRejectsExcessivelyDeepPreparedSchema(t *testing.T) {
	def := promptDef(domain.FormatJSON, domain.ValidationJSONSchema, 0)
	def.Validation.SchemaPath = "schema.json"
	llmClient := &fakeLLM{forbid: true}
	validator := &recordingValidationPreparer{
		plan: &recordingPreparedValidation{schemaDocument: excessivelyDeepPreparedJSONValue()},
	}
	runner := NewRunner(
		&fakePromptRepo{def: def},
		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
		nil,
		defaultArtifactReader(),
		defaultRenderer(),
		llmClient,
		validator,
		nil,
	)

	_, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
		PromptID:  "p",
		ProfileID: "exec",
		Inputs:    singleInputRef(),
	})
	if !errors.Is(err, ErrInvalidRequest) {
		t.Fatalf("expected ErrInvalidRequest, got %v", err)
	}
	if llmClient.calls != 0 {
		t.Fatalf("invalid prepared schema reached generation: %d calls", llmClient.calls)
	}
}

func excessivelyDeepPreparedJSONValue() any {
	const clearlyUnsafeContainerDepth = 1_000
	var value any = true
	for level := 0; level < clearlyUnsafeContainerDepth; level++ {
		value = map[string]any{"child": value}
	}
	return value
}

func TestRunnerRunPreparedRechecksEnvironmentCredentialBeforeAdmission(t *testing.T) {
	const environmentName = "PROMPTKIT_PREPARED_EXECUTION_TEST_KEY"
	t.Setenv(environmentName, "available-during-preparation")

	profile := defaultExecutionProfile()
	profile.APIKeyEnv = environmentName
	validator := &recordingValidationPreparer{plan: &recordingPreparedValidation{}}
	admitter := &fakeRunAdmitter{}
	llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: "unexpected"}}
	runner := NewRunner(
		&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)},
		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}},
		nil,
		defaultArtifactReader(),
		defaultRenderer(),
		llmClient,
		validator,
		admitter,
	)
	prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
		PromptID:  "p",
		ProfileID: "exec",
		Inputs:    singleInputRef(),
	})
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}
	if err := os.Unsetenv(environmentName); err != nil {
		t.Fatalf("unset credential environment: %v", err)
	}

	result, err := runner.RunPrepared(context.Background(), prepared)
	if result != nil {
		t.Fatalf("credential failure returned partial result: %+v", result)
	}
	if !errors.Is(err, ErrInvalidRequest) || !errors.Is(err, ErrAPIKeyEnvMissing) {
		t.Fatalf("credential error identities are missing: %v", err)
	}
	if len(admitter.backendIDs) != 0 || llmClient.calls != 0 {
		t.Fatalf("credential failure reached admission or generation: admission=%v generation=%d", admitter.backendIDs, llmClient.calls)
	}
	if _, err := runner.RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
		t.Fatalf("credential failure did not consume execution: %v", err)
	}
}

func TestRunnerRunPreparedKeepsDirectCredentialOutOfMetadata(t *testing.T) {
	const directKey = "direct-prepared-test-key"

	profile := defaultExecutionProfile()
	profile.APIKeyRequired = true
	client := &fakeLLM{resp: &domain.GenerateResponse{Content: "ok"}}
	runner := NewRunner(
		&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)},
		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}},
		nil,
		defaultArtifactReader(),
		defaultRenderer(),
		client,
		&recordingValidationPreparer{plan: &recordingPreparedValidation{}},
		nil,
	)

	prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
		PromptID:  "p",
		ProfileID: "exec",
		APIKey:    directKey,
		Inputs:    singleInputRef(),
	})
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}
	if details := prepared.Details(); details.EffectiveModelParams.APIKey != "" {
		t.Fatal("prepared details retained direct credential")
	}

	result, err := runner.RunPrepared(context.Background(), prepared)
	if err != nil {
		t.Fatalf("run prepared: %v", err)
	}
	if client.lastReq.Target.APIKey != directKey {
		t.Fatal("generation did not receive direct credential")
	}
	if result.EffectiveModelParams.APIKey != "" {
		t.Fatal("run result retained direct credential")
	}
}

func TestRunnerRunPreparedUsesFrozenValidationForInitialAndRepairOutputs(t *testing.T) {
	validator := &recordingValidationPreparer{
		plan: &recordingPreparedValidation{
			results: []domain.ValidationResult{
				{
					Status:  domain.ValidationFailed,
					Mode:    domain.ValidationJSON,
					Errors:  []string{"invalid"},
					IsValid: false,
				},
				{
					Status:  domain.ValidationPassed,
					Mode:    domain.ValidationJSON,
					IsValid: true,
				},
			},
		},
	}
	repairer := &fakeRepairer{
		responses: []*domain.GenerateResponse{{Content: `{"repaired":true}`}},
	}
	admitter := &fakeRunAdmitter{}
	reader := defaultArtifactReader()
	renderer := defaultRenderer()
	runner := NewRunnerWithRepairer(
		&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 1)},
		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
		nil,
		reader,
		renderer,
		&fakeLLM{resp: &domain.GenerateResponse{Content: `{"broken":true}`}},
		validator,
		repairer,
		admitter,
	)
	prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
		PromptID:  "p",
		ProfileID: "exec",
		Inputs:    singleInputRef(),
	})
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}

	result, err := runner.RunPrepared(context.Background(), prepared)
	if err != nil {
		t.Fatalf("run prepared: %v", err)
	}
	if !reflect.DeepEqual(validator.plan.artifacts, []string{`{"broken":true}`, `{"repaired":true}`}) {
		t.Fatalf("prepared validation artifacts=%#v", validator.plan.artifacts)
	}
	if validator.directValidateCalls != 0 || repairer.calls != 1 {
		t.Fatalf("validation/repair calls=(direct=%d repair=%d), want (0, 1)", validator.directValidateCalls, repairer.calls)
	}
	if result.Validation.Status != domain.ValidationPassed || result.Validation.RepairAttempts != 1 {
		t.Fatalf("unexpected repaired validation result: %+v", result.Validation)
	}
	if admitter.releaseCalls != 1 {
		t.Fatalf("admission releases=%d, want 1", admitter.releaseCalls)
	}
	if reader.calls != 1 || renderer.calls != 1 {
		t.Fatalf("execution reopened preparation sources: artifact=%d render=%d", reader.calls, renderer.calls)
	}
}

func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
	generationFailure := errors.New("generation failed")
	validationFailure := errors.New("validation failed")
	repairFailure := errors.New("repair failed")

	tests := []struct {
		name          string
		generationErr error
		validation    *recordingPreparedValidation
		repairer      *fakeRepairer
		wantError     error
	}{
		{
			name:          "generation failure",
			generationErr: generationFailure,
			validation:    &recordingPreparedValidation{},
			wantError:     ErrLLMGenerate,
		},
		{
			name: "validation failure",
			validation: &recordingPreparedValidation{
				errs: []error{validationFailure},
			},
			wantError: ErrValidation,
		},
		{
			name: "repair failure",
			validation: &recordingPreparedValidation{
				results: []domain.ValidationResult{{
					Status:  domain.ValidationFailed,
					Mode:    domain.ValidationJSON,
					Errors:  []string{"invalid"},
					IsValid: false,
				}},
			},
			repairer:  &fakeRepairer{err: repairFailure},
			wantError: ErrValidation,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			def := promptDef(domain.FormatJSON, domain.ValidationJSON, 1)
			if test.repairer == nil {
				def.Validation.RepairAttempts = 0
			}
			validator := &recordingValidationPreparer{plan: test.validation}
			admitter := &fakeRunAdmitter{}
			runner := NewRunnerWithRepairer(
				&fakePromptRepo{def: def},
				&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
				nil,
				defaultArtifactReader(),
				defaultRenderer(),
				&fakeLLM{
					resp: &domain.GenerateResponse{Content: `{"value":true}`},
					err:  test.generationErr,
				},
				validator,
				test.repairer,
				admitter,
			)
			prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
				PromptID:  "p",
				ProfileID: "exec",
				Inputs:    singleInputRef(),
			})
			if err != nil {
				t.Fatalf("prepare execution: %v", err)
			}

			result, err := runner.RunPrepared(context.Background(), prepared)
			if result != nil || !errors.Is(err, test.wantError) {
				t.Fatalf("run prepared=(%+v, %v), want %v", result, err, test.wantError)
			}
			if len(admitter.backendIDs) != 1 || admitter.releaseCalls != 1 {
				t.Fatalf(
					"admission calls=%#v releases=%d, want one each",
					admitter.backendIDs,
					admitter.releaseCalls,
				)
			}
		})
	}
}

func TestRunnerPreparedExecutionOwnershipUseAndDiscard(t *testing.T) {
	newRunner := func() *Runner {
		return NewRunner(
			&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)},
			&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
			nil,
			defaultArtifactReader(),
			defaultRenderer(),
			&fakeLLM{resp: &domain.GenerateResponse{Content: "ok"}},
			&recordingValidationPreparer{plan: &recordingPreparedValidation{}},
			nil,
		)
	}
	request := domain.RunRequest{
		PromptID:  "p",
		ProfileID: "exec",
		Inputs:    singleInputRef(),
	}

	owner := newRunner()
	prepared, err := owner.PrepareExecution(context.Background(), request)
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}
	if _, err := newRunner().RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
		t.Fatalf("foreign runner error=%v, want ErrInvalidRequest", err)
	}
	if _, err := owner.RunPrepared(context.Background(), prepared); err != nil {
		t.Fatalf("owner run prepared: %v", err)
	}
	if _, err := owner.RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
		t.Fatalf("second owner run error=%v, want ErrInvalidRequest", err)
	}
	if prepared.Details() == nil {
		t.Fatal("details unavailable after execution")
	}

	discarded, err := owner.PrepareExecution(context.Background(), request)
	if err != nil {
		t.Fatalf("prepare discarded execution: %v", err)
	}
	discarded.Discard()
	discarded.Discard()
	if _, err := owner.RunPrepared(context.Background(), discarded); !errors.Is(err, ErrInvalidRequest) {
		t.Fatalf("discarded execution error=%v, want ErrInvalidRequest", err)
	}
	if discarded.Details() == nil {
		t.Fatal("details unavailable after discard")
	}
}

func TestRunnerPreparedExecutionWithoutValidatorSkipsValidation(t *testing.T) {
	runner := NewRunner(
		&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 0)},
		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
		nil,
		defaultArtifactReader(),
		defaultRenderer(),
		&fakeLLM{resp: &domain.GenerateResponse{Content: `{}`}},
		nil,
		nil,
	)
	prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
		PromptID:  "p",
		ProfileID: "exec",
		Inputs:    singleInputRef(),
	})
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}

	result, err := runner.RunPrepared(context.Background(), prepared)
	if err != nil {
		t.Fatalf("run prepared: %v", err)
	}
	if result.Validation.Status != domain.ValidationSkipped || !result.Validation.IsValid {
		t.Fatalf("unexpected no-validator result: %+v", result.Validation)
	}
}

func TestRunnerPrepareExecutionRequiresValidationPreparer(t *testing.T) {
	reader := defaultArtifactReader()
	runner := NewRunner(
		&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationBasic, 0)},
		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
		nil,
		reader,
		defaultRenderer(),
		&fakeLLM{forbid: true},
		&fakeValidator{},
		nil,
	)

	prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
		PromptID:  "p",
		ProfileID: "exec",
		Inputs:    singleInputRef(),
	})
	if prepared != nil || !errors.Is(err, ErrValidation) {
		t.Fatalf("prepare execution=(%+v, %v), want ErrValidation", prepared, err)
	}
	if reader.calls != 0 {
		t.Fatalf("unsupported validator allowed completion, artifact calls=%d", reader.calls)
	}
}