promptkit/prepared_execution_contract_test.go

package promptkit_test

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"reflect"
	"strings"
	"sync"
	"testing"
	"testing/fstest"
	"time"

	"gitea.maximumdirect.net/eric/promptkit"
)

func TestPreparedExecutionFreezesSourcesAndReturnsIndependentDetails(t *testing.T) {
	promptSource := preparedPromptSource("original")
	profileSource := preparedProfileSource("original-model")
	schemaSource := preparedSchemaSource()
	reader := &mutablePreparedArtifactReader{
		body: "original artifact",
		hash: "original-input-hash",
	}
	client := &preparedRecordingClient{
		response: &promptkit.GenerateResponse{Content: `{"value":3}`},
	}
	engine, err := promptkit.NewEngine(
		promptkit.Config{},
		promptkit.WithPromptFS(promptSource, "."),
		promptkit.WithProfileFS(profileSource, "."),
		promptkit.WithSchemaFS(schemaSource, "."),
		promptkit.WithArtifactReader(reader),
		promptkit.WithLLMClient(client),
	)
	if err != nil {
		t.Fatalf("construct engine: %v", err)
	}

	temperature := 0.25
	extraParams := map[string]any{
		"nested": map[string]any{"source": "original"},
	}
	request := promptkit.RunRequest{
		PromptID: "prepared",
		Inputs: map[string]promptkit.ArtifactRef{
			"input": promptkit.Inline("original request input"),
		},
		Vars: map[string]string{"label": "original variable"},
		Execution: &promptkit.ExecutionTargetOverride{
			Temperature: &temperature,
			ExtraParams: extraParams,
		},
	}
	preparationContext, cancelPreparation := context.WithCancel(context.Background())
	prepared, err := engine.PrepareExecution(preparationContext, request)
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}
	cancelPreparation()

	request.PromptID = "changed"
	request.Inputs["input"] = promptkit.Inline("changed request input")
	request.Vars["label"] = "changed variable"
	temperature = 1.5
	extraParams["nested"].(map[string]any)["source"] = "changed"
	promptSource["prompt.yaml"] = &fstest.MapFile{Data: []byte(`id: changed`)}
	profileSource["profile.yaml"] = &fstest.MapFile{Data: []byte(`id: changed`)}
	schemaSource["schema.json"] = &fstest.MapFile{Data: []byte(`{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "title": "changed root",
  "type": "string"
}`)}
	schemaSource["value.json"] = &fstest.MapFile{Data: []byte(`{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "type": "string"
}`)}
	reader.set("changed artifact", "changed-input-hash")

	first := prepared.Details()
	first.Messages[0].Content = "changed details"
	first.InputHashes["input"] = "changed-details-hash"
	first.EffectiveModelParams.ExtraParams["nested"].(map[string]any)["source"] = "changed details"
	first.StructuredOutput.JSONSchema.Schema.(map[string]any)["title"] = "changed details"

	second := prepared.Details()
	if second.Messages[0].Content != "Input=original artifact Label=original variable" {
		t.Fatalf("details message changed: %q", second.Messages[0].Content)
	}
	if second.InputHashes["input"] != "original-input-hash" {
		t.Fatalf("details input hash changed: %q", second.InputHashes["input"])
	}
	if second.EffectiveModelParams.Model != "original-model" ||
		second.EffectiveModelParams.Temperature != 0.25 ||
		second.EffectiveModelParams.ExtraParams["nested"].(map[string]any)["source"] != "original" {
		t.Fatalf("details target changed: %+v", second.EffectiveModelParams)
	}
	schema := second.StructuredOutput.JSONSchema.Schema.(map[string]any)
	if schema["title"] != "original root" {
		t.Fatalf("details schema changed: %#v", schema)
	}

	result, err := engine.RunPrepared(context.Background(), prepared)
	if err != nil {
		t.Fatalf("run prepared after preparation-context cancellation: %v", err)
	}
	if result.Validation.Status != promptkit.ValidationPassed || !result.Validation.IsValid {
		t.Fatalf("frozen schema did not validate original output: %+v", result.Validation)
	}
	if reader.callCount() != 1 {
		t.Fatalf("execution reopened artifact source: calls=%d", reader.callCount())
	}

	requests := client.snapshot()
	if len(requests) != 1 {
		t.Fatalf("generation calls=%d, want 1", len(requests))
	}
	generated := requests[0]
	if generated.Prompt.Messages[0].Content != second.Messages[0].Content ||
		generated.Target.Model != second.EffectiveModelParams.Model ||
		!reflect.DeepEqual(generated.Target.ExtraParams, second.EffectiveModelParams.ExtraParams) ||
		!reflect.DeepEqual(generated.StructuredOutput, second.StructuredOutput) {
		t.Fatalf("generation did not use frozen details:\nrequest=%+v\ndetails=%+v", generated, second)
	}
	if result.PromptID != second.PromptID ||
		result.PromptVersion != second.PromptVersion ||
		result.PromptHash != second.PromptHash ||
		result.SessionID != second.SessionID ||
		result.RenderedPromptHash != second.RenderedPromptHash ||
		result.SelectedProfileID != second.SelectedProfileID ||
		result.SelectedBackendID != second.SelectedBackendID ||
		!reflect.DeepEqual(result.EffectiveModelParams, second.EffectiveModelParams) ||
		!reflect.DeepEqual(result.InputHashes, second.InputHashes) {
		t.Fatalf("result provenance does not match details:\nresult=%+v\ndetails=%+v", result, second)
	}
}

func TestPreparedExecutionLifecycleAndEngineBinding(t *testing.T) {
	ownerClient := &preparedRecordingClient{
		response: &promptkit.GenerateResponse{Content: "ok"},
	}
	owner := newPreparedContractEngine(t, ownerClient, "owner content")
	foreign := newPreparedContractEngine(t, &preparedRecordingClient{
		response: &promptkit.GenerateResponse{Content: "unexpected"},
	}, "foreign content")

	prepared, err := owner.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared"})
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}
	copied := *prepared

	var nilEngine *promptkit.Engine
	if result, err := nilEngine.RunPrepared(context.Background(), prepared); result != nil ||
		!errors.Is(err, promptkit.ErrInvalidConfig) {
		t.Fatalf("nil engine result=(%+v, %v), want ErrInvalidConfig", result, err)
	}
	if result, err := foreign.RunPrepared(context.Background(), prepared); result != nil ||
		!errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("foreign engine result=(%+v, %v), want ErrInvalidRequest", result, err)
	}
	if result, err := owner.RunPrepared(context.Background(), nil); result != nil ||
		!errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("nil handle result=(%+v, %v), want ErrInvalidRequest", result, err)
	}
	if result, err := owner.RunPrepared(context.Background(), &promptkit.PreparedExecution{}); result != nil ||
		!errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("zero handle result=(%+v, %v), want ErrInvalidRequest", result, err)
	}

	result, err := owner.RunPrepared(context.Background(), &copied)
	if err != nil || result == nil {
		t.Fatalf("owner run prepared=(%+v, %v), want success", result, err)
	}
	for name, handle := range map[string]*promptkit.PreparedExecution{
		"original": prepared,
		"copy":     &copied,
	} {
		if result, err := owner.RunPrepared(context.Background(), handle); result != nil ||
			!errors.Is(err, promptkit.ErrInvalidRequest) {
			t.Fatalf("%s reused handle result=(%+v, %v), want ErrInvalidRequest", name, result, err)
		}
		if handle.Details().PromptID != "prepared" {
			t.Fatalf("%s details unavailable after execution", name)
		}
	}
	if len(ownerClient.snapshot()) != 1 {
		t.Fatalf("owner generation calls=%d, want 1", len(ownerClient.snapshot()))
	}

	collaboratorFailure := errors.New("prepared collaborator failure")
	failingClient := &preparedRecordingClient{err: collaboratorFailure}
	failingEngine := newPreparedContractEngine(t, failingClient, "failure content")
	failing, err := failingEngine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared"})
	if err != nil {
		t.Fatalf("prepare failing execution: %v", err)
	}
	if result, err := failingEngine.RunPrepared(context.Background(), failing); result != nil ||
		!errors.Is(err, promptkit.ErrLLMGenerate) ||
		!errors.Is(err, collaboratorFailure) {
		t.Fatalf("generation failure result=(%+v, %v), want public and collaborator identities", result, err)
	}
	if result, err := failingEngine.RunPrepared(context.Background(), failing); result != nil ||
		!errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("failed execution was reusable: result=(%+v, %v)", result, err)
	}

	cancellationRelease := make(chan struct{})
	cancellationStarted := make(chan struct{}, 1)
	cancelingEngine := newPreparedContractEngine(t, &preparedRecordingClient{
		response: &promptkit.GenerateResponse{Content: "unexpected"},
		started:  cancellationStarted,
		release:  cancellationRelease,
	}, "cancellation content")
	canceling, err := cancelingEngine.PrepareExecution(
		context.Background(),
		promptkit.RunRequest{PromptID: "prepared"},
	)
	if err != nil {
		t.Fatalf("prepare canceled execution: %v", err)
	}
	executionContext, cancelExecution := context.WithCancel(context.Background())
	type canceledOutcome struct {
		result *promptkit.RunResult
		err    error
	}
	canceledResult := make(chan canceledOutcome, 1)
	go func() {
		result, runErr := cancelingEngine.RunPrepared(executionContext, canceling)
		canceledResult <- canceledOutcome{result: result, err: runErr}
	}()
	select {
	case <-cancellationStarted:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for cancelable generation")
	}
	cancelExecution()
	select {
	case outcome := <-canceledResult:
		if outcome.result != nil ||
			!errors.Is(outcome.err, promptkit.ErrLLMGenerate) ||
			!errors.Is(outcome.err, context.Canceled) {
			t.Fatalf(
				"canceled execution=(%+v, %v), want generation and context identities",
				outcome.result,
				outcome.err,
			)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for canceled execution")
	}
	if result, err := cancelingEngine.RunPrepared(context.Background(), canceling); result != nil ||
		!errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("canceled execution was reusable: result=(%+v, %v)", result, err)
	}
}

func TestPreparedExecutionConcurrentClaimAllowsOneGeneration(t *testing.T) {
	release := make(chan struct{})
	client := &preparedRecordingClient{
		response: &promptkit.GenerateResponse{Content: "ok"},
		started:  make(chan struct{}, 1),
		release:  release,
	}
	engine := newPreparedContractEngine(t, client, "concurrent content")
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared"})
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}

	type outcome struct {
		result *promptkit.RunResult
		err    error
	}
	outcomes := make(chan outcome, 2)
	for i := 0; i < 2; i++ {
		go func() {
			result, runErr := engine.RunPrepared(context.Background(), prepared)
			outcomes <- outcome{result: result, err: runErr}
		}()
	}

	select {
	case <-client.started:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for generation")
	}
	select {
	case loser := <-outcomes:
		if loser.result != nil || !errors.Is(loser.err, promptkit.ErrInvalidRequest) {
			t.Fatalf("concurrent loser=(%+v, %v), want ErrInvalidRequest", loser.result, loser.err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for rejected concurrent claim")
	}

	close(release)
	select {
	case winner := <-outcomes:
		if winner.err != nil || winner.result == nil {
			t.Fatalf("concurrent winner=(%+v, %v), want success", winner.result, winner.err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for successful concurrent claim")
	}
	if len(client.snapshot()) != 1 {
		t.Fatalf("generation calls=%d, want 1", len(client.snapshot()))
	}
}

func TestPreparedExecutionRunAndDiscardRaceHasOneWinner(t *testing.T) {
	const attempts = 32

	for i := 0; i < attempts; i++ {
		client := &preparedRecordingClient{
			response: &promptkit.GenerateResponse{Content: "ok"},
		}
		engine := newPreparedContractEngine(t, client, "race content")
		prepared, err := engine.PrepareExecution(
			context.Background(),
			promptkit.RunRequest{PromptID: "prepared"},
		)
		if err != nil {
			t.Fatalf("attempt %d prepare execution: %v", i, err)
		}

		start := make(chan struct{})
		type outcome struct {
			result *promptkit.RunResult
			err    error
		}
		runOutcome := make(chan outcome, 1)
		discardDone := make(chan struct{})

		go func() {
			<-start
			result, runErr := engine.RunPrepared(context.Background(), prepared)
			runOutcome <- outcome{result: result, err: runErr}
		}()
		go func() {
			<-start
			prepared.Discard()
			close(discardDone)
		}()

		close(start)
		runResult := <-runOutcome
		<-discardDone
		calls := len(client.snapshot())
		switch {
		case runResult.err == nil:
			if runResult.result == nil || calls != 1 {
				t.Fatalf(
					"attempt %d run won with outcome=(%+v, %v), generation calls=%d",
					i,
					runResult.result,
					runResult.err,
					calls,
				)
			}
		case errors.Is(runResult.err, promptkit.ErrInvalidRequest):
			if runResult.result != nil || calls != 0 {
				t.Fatalf(
					"attempt %d discard won with outcome=(%+v, %v), generation calls=%d",
					i,
					runResult.result,
					runResult.err,
					calls,
				)
			}
		default:
			t.Fatalf("attempt %d unexpected run outcome=(%+v, %v)", i, runResult.result, runResult.err)
		}
		if prepared.Details().PromptID != "prepared" {
			t.Fatalf("attempt %d details unavailable after race", i)
		}
	}
}

func TestPreparedExecutionDiscardAndFormattingDoNotExposePrivateState(t *testing.T) {
	const (
		directCredential = "pk-test-direct-credential-41f7"
		renderedContent  = "rendered-content-sentinel-98d2"
	)
	client := &preparedRecordingClient{
		response: &promptkit.GenerateResponse{Content: "generated output"},
	}
	engine, err := promptkit.NewEngine(
		promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prepared", "profile", renderedContent), "."),
		promptkit.WithProfiles(promptkit.Profile{
			ID:             "profile",
			Endpoint:       "http://example.test/v1",
			Model:          "model",
			APIKeyRequired: true,
		}),
		promptkit.WithLLMClient(client),
	)
	if err != nil {
		t.Fatalf("construct engine: %v", err)
	}
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{
		PromptID: "prepared",
		APIKey:   directCredential,
	})
	if err != nil {
		t.Fatalf("prepare execution: %v", err)
	}

	copied := *prepared
	zeroValue := promptkit.PreparedExecution{}
	var nilHandle *promptkit.PreparedExecution
	for name, value := range map[string]any{
		"original pointer": prepared,
		"copied value":     copied,
		"zero value":       zeroValue,
		"zero pointer":     &zeroValue,
	} {
		for format, formatted := range map[string]string{
			"String":   fmt.Sprintf("%s", value),
			"GoString": fmt.Sprintf("%#v", value),
			"v":        fmt.Sprintf("%v", value),
			"+v":       fmt.Sprintf("%+v", value),
		} {
			if formatted != "promptkit.PreparedExecution{opaque}" {
				t.Fatalf("%s %s formatting = %q, want opaque representation", name, format, formatted)
			}
			assertPreparedPrivateValuesAbsent(t, formatted, directCredential, renderedContent)
		}
	}
	for format, formatted := range map[string]string{
		"String":   fmt.Sprintf("%s", nilHandle),
		"GoString": fmt.Sprintf("%#v", nilHandle),
		"v":        fmt.Sprintf("%v", nilHandle),
		"+v":       fmt.Sprintf("%+v", nilHandle),
	} {
		if formatted != "<nil>" {
			t.Fatalf("nil pointer %s formatting = %q, want <nil>", format, formatted)
		}
		assertPreparedPrivateValuesAbsent(t, formatted, directCredential, renderedContent)
	}
	payload, err := json.Marshal(prepared)
	if err != nil {
		t.Fatalf("marshal opaque handle: %v", err)
	}
	assertPreparedPrivateValuesAbsent(t, string(payload), directCredential, renderedContent)

	detailsBefore := prepared.Details()
	detailsJSON, err := json.Marshal(detailsBefore)
	if err != nil {
		t.Fatalf("marshal prepared details: %v", err)
	}
	assertPreparedPrivateValuesAbsent(t, string(detailsJSON), directCredential)

	executionResult, err := engine.RunPrepared(context.Background(), &copied)
	if err != nil {
		t.Fatalf("run copied execution after formatting: %v", err)
	}
	requests := client.snapshot()
	if len(requests) != 1 || requests[0].APIKey != directCredential {
		t.Fatalf("direct credential did not reach only the client credential field: %#v", requests)
	}
	requestJSON, err := json.Marshal(requests[0])
	if err != nil {
		t.Fatalf("marshal captured generate request: %v", err)
	}
	for _, value := range []string{
		fmt.Sprint(requests[0]),
		fmt.Sprintf("%+v", requests[0]),
		fmt.Sprintf("%#v", requests[0]),
		string(requestJSON),
		fmt.Sprint(executionResult),
	} {
		assertPreparedPrivateValuesAbsent(t, value, directCredential)
	}
	resultJSON, err := json.Marshal(executionResult)
	if err != nil {
		t.Fatalf("marshal execution result: %v", err)
	}
	assertPreparedPrivateValuesAbsent(t, string(resultJSON), directCredential)

	nilHandle.Discard()
	if !reflect.DeepEqual(nilHandle.Details(), promptkit.PreparedRun{}) {
		t.Fatalf("nil handle details=%+v, want zero value", nilHandle.Details())
	}
	zeroHandle := &zeroValue
	zeroHandle.Discard()
	if !reflect.DeepEqual(zeroHandle.Details(), promptkit.PreparedRun{}) {
		t.Fatalf("zero handle details=%+v, want zero value", zeroHandle.Details())
	}

	discarded, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{
		PromptID: "prepared",
		APIKey:   directCredential,
	})
	if err != nil {
		t.Fatalf("prepare execution for discard: %v", err)
	}
	discardedDetails := discarded.Details()
	discarded.Discard()
	discarded.Discard()
	result, lifecycleErr := engine.RunPrepared(context.Background(), discarded)
	if result != nil || !errors.Is(lifecycleErr, promptkit.ErrInvalidRequest) {
		t.Fatalf("discarded execution result=(%+v, %v), want ErrInvalidRequest", result, lifecycleErr)
	}
	assertPreparedPrivateValuesAbsent(t, lifecycleErr.Error(), directCredential, renderedContent)
	if !reflect.DeepEqual(discarded.Details(), discardedDetails) {
		t.Fatal("details changed after discard")
	}
}

func TestPreparedExecutionCredentialCapacityAndTimingBoundaries(t *testing.T) {
	t.Run("credential is rechecked before generation", func(t *testing.T) {
		const (
			environmentName = "PROMPTKIT_PREPARED_CONTRACT_KEY"
			environmentKey  = "environment-credential-sentinel"
		)
		t.Setenv(environmentName, environmentKey)
		client := &preparedRecordingClient{
			response: &promptkit.GenerateResponse{Content: "unexpected"},
		}
		engine, err := promptkit.NewEngine(
			promptkit.Config{},
			promptkit.WithPromptFS(contractPromptFS("prepared", "profile", "content"), "."),
			promptkit.WithProfileFS(preparedCredentialProfileSource(environmentName), "."),
			promptkit.WithLLMClient(client),
		)
		if err != nil {
			t.Fatalf("construct credential engine: %v", err)
		}
		prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared"})
		if err != nil {
			t.Fatalf("prepare credential execution: %v", err)
		}
		if err := os.Unsetenv(environmentName); err != nil {
			t.Fatalf("unset credential environment: %v", err)
		}

		result, err := engine.RunPrepared(context.Background(), prepared)
		if result != nil ||
			!errors.Is(err, promptkit.ErrInvalidRequest) ||
			!errors.Is(err, promptkit.ErrAPIKeyEnvMissing) {
			t.Fatalf("credential execution=(%+v, %v), want credential identities", result, err)
		}
		if len(client.snapshot()) != 0 {
			t.Fatalf("credential failure reached generation: %d calls", len(client.snapshot()))
		}
		assertPreparedPrivateValuesAbsent(t, err.Error(), environmentKey)
		if result, err := engine.RunPrepared(context.Background(), prepared); result != nil ||
			!errors.Is(err, promptkit.ErrInvalidRequest) {
			t.Fatalf("credential failure did not consume handle: result=(%+v, %v)", result, err)
		}
	})

	t.Run("preparation does not admit and execution timing starts after retention", func(t *testing.T) {
		release := make(chan struct{})
		client := newCapacityGateClient(release, 4)
		engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), nil)

		activeRun := make(chan capacityRunResult, 1)
		go runCapacityRequest(
			engine,
			context.Background(),
			promptkit.RunRequest{PromptID: "prompt"},
			activeRun,
		)
		awaitCapacityRequest(t, client.started)

		prepared, err := engine.PrepareExecution(
			context.Background(),
			promptkit.RunRequest{PromptID: "prompt"},
		)
		if err != nil {
			t.Fatalf("prepare while capacity is full: %v", err)
		}
		if _, _, calls := client.snapshot(); calls != 1 {
			t.Fatalf("preparation invoked generation: calls=%d", calls)
		}
		if result, err := engine.RunPrepared(context.Background(), prepared); result != nil ||
			!errors.Is(err, promptkit.ErrCapacityExceeded) {
			t.Fatalf("capacity execution=(%+v, %v), want ErrCapacityExceeded", result, err)
		} else {
			var capacityErr *promptkit.CapacityError
			if !errors.As(err, &capacityErr) || capacityErr == nil || capacityErr.BackendID != "limited" {
				t.Fatalf("capacity execution=%v, want limited CapacityError", err)
			}
		}
		if result, err := engine.RunPrepared(context.Background(), prepared); result != nil ||
			!errors.Is(err, promptkit.ErrInvalidRequest) {
			t.Fatalf("capacity rejection did not consume handle: result=(%+v, %v)", result, err)
		}
		if prepared.Details().PromptID != "prompt" {
			t.Fatal("details unavailable after capacity rejection")
		}

		close(release)
		activeOutcome := awaitCapacityRun(t, activeRun)
		if activeOutcome.err != nil || activeOutcome.result == nil {
			t.Fatalf("active run outcome=(%+v, %v), want success", activeOutcome.result, activeOutcome.err)
		}

		timed, err := engine.PrepareExecution(
			context.Background(),
			promptkit.RunRequest{PromptID: "prompt"},
		)
		if err != nil {
			t.Fatalf("prepare timed execution: %v", err)
		}
		details := timed.Details()
		time.Sleep(25 * time.Millisecond)
		executionFloor := time.Now().UTC()
		result, err := engine.RunPrepared(context.Background(), timed)
		if err != nil {
			t.Fatalf("run timed execution: %v", err)
		}
		if result.StartTime.Before(executionFloor) ||
			!result.StartTime.After(details.EndTime) ||
			result.EndTime.Before(result.StartTime) ||
			result.Duration != result.EndTime.Sub(result.StartTime) {
			t.Fatalf(
				"execution timing includes preparation or retention: details_end=%s floor=%s result=%+v",
				details.EndTime,
				executionFloor,
				result,
			)
		}
		if _, _, calls := client.snapshot(); calls != 2 {
			t.Fatalf("generation calls=%d, want active and timed executions only", calls)
		}
	})
}

type mutablePreparedArtifactReader struct {
	mu    sync.Mutex
	body  string
	hash  string
	calls int
}

func (r *mutablePreparedArtifactReader) Read(
	_ context.Context,
	_ promptkit.ArtifactRef,
) (*promptkit.Artifact, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.calls++
	return &promptkit.Artifact{
		Body: []byte(r.body),
		Hash: r.hash,
	}, nil
}

func (r *mutablePreparedArtifactReader) set(body, hash string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.body = body
	r.hash = hash
}

func (r *mutablePreparedArtifactReader) callCount() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.calls
}

type preparedRecordingClient struct {
	mu       sync.Mutex
	response *promptkit.GenerateResponse
	err      error
	requests []promptkit.GenerateRequest
	started  chan struct{}
	release  <-chan struct{}
}

func (c *preparedRecordingClient) Generate(
	ctx context.Context,
	request promptkit.GenerateRequest,
) (*promptkit.GenerateResponse, error) {
	c.mu.Lock()
	c.requests = append(c.requests, request)
	c.mu.Unlock()

	if c.started != nil {
		c.started <- struct{}{}
	}
	if c.release != nil {
		select {
		case <-c.release:
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	if c.err != nil {
		return nil, c.err
	}
	return c.response, nil
}

func (c *preparedRecordingClient) snapshot() []promptkit.GenerateRequest {
	c.mu.Lock()
	defer c.mu.Unlock()
	return append([]promptkit.GenerateRequest(nil), c.requests...)
}

func newPreparedContractEngine(
	t *testing.T,
	client promptkit.LLMClient,
	message string,
) *promptkit.Engine {
	t.Helper()
	engine, err := promptkit.NewEngine(
		promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prepared", "profile", message), "."),
		promptkit.WithProfiles(promptkit.Profile{
			ID:       "profile",
			Endpoint: "http://example.test/v1",
			Model:    "model",
		}),
		promptkit.WithLLMClient(client),
	)
	if err != nil {
		t.Fatalf("construct prepared execution engine: %v", err)
	}
	return engine
}

func preparedPromptSource(label string) fstest.MapFS {
	return fstest.MapFS{
		"prompt.yaml": &fstest.MapFile{Data: []byte(`id: prepared
version: "1"
default_profile: profile
inputs:
  - name: input
    required: true
messages:
  - role: user
    content: 'Input={{input "input"}} Label={{.label}}'
description: ` + label + `
output:
  format: json
  validation_mode: json_schema
  schema_path: schema.json
`)},
	}
}

func preparedProfileSource(model string) fstest.MapFS {
	return fstest.MapFS{
		"profile.yaml": &fstest.MapFile{Data: []byte(`id: profile
endpoint: http://example.test/v1
model: ` + model + `
`)},
	}
}

func preparedCredentialProfileSource(environmentName string) fstest.MapFS {
	return fstest.MapFS{
		"profile.yaml": &fstest.MapFile{Data: []byte(`id: profile
endpoint: http://example.test/v1
model: model
api_key_env: ` + environmentName + `
`)},
	}
}

func preparedSchemaSource() fstest.MapFS {
	return fstest.MapFS{
		"schema.json": &fstest.MapFile{Data: []byte(`{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "title": "original root",
  "type": "object",
  "required": ["value"],
  "properties": {
    "value": {"$ref": "value.json"}
  }
}`)},
		"value.json": &fstest.MapFile{Data: []byte(`{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "type": "integer",
  "minimum": 2
}`)},
	}
}

func assertPreparedPrivateValuesAbsent(t *testing.T, value string, privateValues ...string) {
	t.Helper()
	for _, privateValue := range privateValues {
		if strings.Contains(value, privateValue) {
			t.Fatalf("value exposed private data %q: %s", privateValue, value)
		}
	}
}