promptkit/public_contract_test.go

package promptkit_test

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

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

func TestPreparedRunJSONOmitsZeroTimingValues(t *testing.T) {
	payload, err := json.Marshal(promptkit.PreparedRun{})
	if err != nil {
		t.Fatalf("marshal prepared run: %v", err)
	}
	for _, field := range []string{"start_time", "end_time", "duration_ms"} {
		if strings.Contains(string(payload), `"`+field+`"`) {
			t.Fatalf("expected zero %s to be omitted, got %s", field, payload)
		}
	}
}

func TestBackendIdentityJSONNamesAndOmission(t *testing.T) {
	t.Run("execution target round trip", func(t *testing.T) {
		value := promptkit.ExecutionTarget{BackendID: promptkit.BackendOpenRouter}
		payload, err := json.Marshal(value)
		if err != nil {
			t.Fatalf("marshal execution target: %v", err)
		}
		var decoded promptkit.ExecutionTarget
		if err := json.Unmarshal(payload, &decoded); err != nil {
			t.Fatalf("unmarshal execution target: %v", err)
		}
		if decoded.BackendID != value.BackendID {
			t.Fatalf("backend identity did not round trip: got %q want %q", decoded.BackendID, value.BackendID)
		}
	})

	t.Run("prepared run round trip", func(t *testing.T) {
		value := promptkit.PreparedRun{SelectedBackendID: promptkit.BackendOpenRouter}
		payload, err := json.Marshal(value)
		if err != nil {
			t.Fatalf("marshal prepared run: %v", err)
		}
		var decoded promptkit.PreparedRun
		if err := json.Unmarshal(payload, &decoded); err != nil {
			t.Fatalf("unmarshal prepared run: %v", err)
		}
		if decoded.SelectedBackendID != value.SelectedBackendID {
			t.Fatalf("backend identity did not round trip: got %q want %q", decoded.SelectedBackendID, value.SelectedBackendID)
		}
	})

	t.Run("run result round trip", func(t *testing.T) {
		value := promptkit.RunResult{SelectedBackendID: promptkit.BackendOpenRouter}
		payload, err := json.Marshal(value)
		if err != nil {
			t.Fatalf("marshal run result: %v", err)
		}
		var decoded promptkit.RunResult
		if err := json.Unmarshal(payload, &decoded); err != nil {
			t.Fatalf("unmarshal run result: %v", err)
		}
		if decoded.SelectedBackendID != value.SelectedBackendID {
			t.Fatalf("backend identity did not round trip: got %q want %q", decoded.SelectedBackendID, value.SelectedBackendID)
		}
	})

	payload, err := json.Marshal(promptkit.ExecutionTarget{})
	if err != nil {
		t.Fatalf("marshal empty execution target: %v", err)
	}
	if strings.Contains(string(payload), `"backend_id"`) {
		t.Fatalf("empty backend identity was not omitted: %s", payload)
	}
}

func TestEndpointOnlyProfileOmitsBackendIdentityFromStableJSON(t *testing.T) {
	client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "ok"}}
	engine, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
		promptkit.WithProfiles(promptkit.Profile{
			ID: "profile", Endpoint: "http://example.test/v1", Model: "model",
		}),
		promptkit.WithLLMClient(client),
	)
	if err != nil {
		t.Fatalf("construct engine: %v", err)
	}

	prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if err != nil {
		t.Fatalf("prepare endpoint-only profile: %v", err)
	}
	result, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if err != nil {
		t.Fatalf("run endpoint-only profile: %v", err)
	}
	if prepared.SelectedBackendID != "" ||
		prepared.EffectiveModelParams.BackendID != "" ||
		result.SelectedBackendID != "" ||
		result.EffectiveModelParams.BackendID != "" {
		t.Fatalf("endpoint-only profile acquired backend identity: prepared=%+v result=%+v", prepared, result)
	}
	for _, value := range []any{prepared, result} {
		payload, err := json.Marshal(value)
		if err != nil {
			t.Fatalf("marshal endpoint-only value: %v", err)
		}
		if strings.Contains(string(payload), `"backend_id"`) || strings.Contains(string(payload), `"selected_backend_id"`) {
			t.Fatalf("endpoint-only backend identity was not omitted: %s", payload)
		}
	}
}

func TestUnknownProfileBackendHasProfileLoadIdentity(t *testing.T) {
	engine, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
		promptkit.WithProfiles(promptkit.Profile{ID: "profile", BackendID: "unknown", Model: "model"}),
	)
	if err != nil {
		t.Fatalf("construct engine: %v", err)
	}

	_, err = engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if !errors.Is(err, promptkit.ErrProfileLoad) {
		t.Fatalf("expected ErrProfileLoad, got %v", err)
	}
	if errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("unknown backend should not have invalid-request identity: %v", err)
	}
}

func TestCustomBackendFlowsThroughProfilesOverridesAndInjectedClient(t *testing.T) {
	t.Setenv("CUSTOM_LLM_KEY", "test-key")
	client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "ok"}}
	engine, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prompt", "backend-profile", "message"), "."),
		promptkit.WithBackend(promptkit.Backend{
			ID:        "  custom  ",
			Endpoint:  "  http://backend.example/v1  ",
			APIKeyEnv: "  CUSTOM_LLM_KEY  ",
			ExtraParams: map[string]any{
				"provider": "custom",
			},
		}),
		promptkit.WithProfiles(
			promptkit.Profile{ID: "backend-profile", BackendID: "custom", Model: "backend-model"},
			promptkit.Profile{ID: "profile-endpoint", BackendID: "custom", Endpoint: "http://profile.example/v1", Model: "profile-model"},
			promptkit.Profile{ID: "blank-profile-endpoint", BackendID: "custom", Endpoint: " \t ", Model: "profile-model"},
		),
		promptkit.WithLLMClient(client),
	)
	if err != nil {
		t.Fatalf("construct engine: %v", err)
	}

	result, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if err != nil {
		t.Fatalf("run with custom backend: %v", err)
	}
	if len(client.requests) != 1 {
		t.Fatalf("expected one injected-client request, got %d", len(client.requests))
	}
	target := client.requests[0].Target
	if target.BackendID != "custom" ||
		target.Endpoint != "http://backend.example/v1" ||
		target.APIKeyEnv != "CUSTOM_LLM_KEY" ||
		target.Model != "backend-model" ||
		target.ExtraParams["provider"] != "custom" ||
		result.SelectedBackendID != "custom" {
		t.Fatalf("unexpected custom backend settings: target=%+v result_backend=%q", target, result.SelectedBackendID)
	}

	prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
		PromptID: "prompt", ProfileID: "profile-endpoint",
	})
	if err != nil {
		t.Fatalf("prepare profile endpoint override: %v", err)
	}
	if prepared.SelectedBackendID != "custom" || prepared.EffectiveModelParams.Endpoint != "http://profile.example/v1" {
		t.Fatalf("profile endpoint override changed backend identity: %+v", prepared)
	}

	prepared, err = engine.Prepare(context.Background(), promptkit.RunRequest{
		PromptID: "prompt", ProfileID: "blank-profile-endpoint",
	})
	if err != nil {
		t.Fatalf("prepare blank profile endpoint: %v", err)
	}
	if prepared.SelectedBackendID != "custom" || prepared.EffectiveModelParams.Endpoint != "http://backend.example/v1" {
		t.Fatalf("blank profile endpoint did not inherit backend endpoint: %+v", prepared)
	}

	prepared, err = engine.Prepare(context.Background(), promptkit.RunRequest{
		PromptID: "prompt",
		Execution: &promptkit.ExecutionTargetOverride{
			Endpoint: "http://request.example/v1",
		},
	})
	if err != nil {
		t.Fatalf("prepare request endpoint override: %v", err)
	}
	if prepared.SelectedBackendID != "custom" || prepared.EffectiveModelParams.Endpoint != "http://request.example/v1" {
		t.Fatalf("request endpoint override changed backend identity: %+v", prepared)
	}
}

func TestCustomBackendSupportsFileProfileAndBothSelectionPaths(t *testing.T) {
	engine, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prompt", "file-profile", "message"), "."),
		promptkit.WithProfileFS(fstest.MapFS{
			"profile.yaml": &fstest.MapFile{Data: []byte(`id: file-profile
backend: file-backend
endpoint: "  "
model: file-model
`)},
		}, "."),
		promptkit.WithBackend(promptkit.Backend{
			ID:       "file-backend",
			Endpoint: "http://file-backend.example/v1",
		}),
	)
	if err != nil {
		t.Fatalf("construct engine: %v", err)
	}

	for _, request := range []promptkit.RunRequest{
		{PromptID: "prompt"},
		{PromptID: "prompt", ProfileID: "file-profile"},
	} {
		prepared, err := engine.Prepare(context.Background(), request)
		if err != nil {
			t.Fatalf("prepare file profile: %v", err)
		}
		if prepared.SelectedBackendID != "file-backend" ||
			prepared.EffectiveModelParams.Endpoint != "http://file-backend.example/v1" {
			t.Fatalf("unexpected file-profile backend resolution: %+v", prepared)
		}
	}
}

func TestBackendOptionsAccumulateAndRegistrationsAreEngineLocal(t *testing.T) {
	engine, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prompt", "first-profile", "message"), "."),
		promptkit.WithBackend(promptkit.Backend{ID: "first", Endpoint: "http://first.example/v1"}),
		promptkit.WithBackend(promptkit.Backend{ID: "second", Endpoint: "http://second.example/v1"}),
		promptkit.WithProfiles(
			promptkit.Profile{ID: "first-profile", BackendID: "first", Model: "model"},
			promptkit.Profile{ID: "second-profile", BackendID: "second", Model: "model"},
		),
	)
	if err != nil {
		t.Fatalf("construct engine with accumulated registrations: %v", err)
	}
	for profileID, wantEndpoint := range map[string]string{
		"first-profile":  "http://first.example/v1",
		"second-profile": "http://second.example/v1",
	} {
		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
			PromptID: "prompt", ProfileID: profileID,
		})
		if err != nil {
			t.Fatalf("prepare %s: %v", profileID, err)
		}
		if prepared.EffectiveModelParams.Endpoint != wantEndpoint {
			t.Fatalf("profile %s endpoint=%q, want %q", profileID, prepared.EffectiveModelParams.Endpoint, wantEndpoint)
		}
	}

	newEngine := func(endpoint string) *promptkit.Engine {
		t.Helper()
		value, err := promptkit.NewEngine(promptkit.Config{},
			promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
			promptkit.WithBackend(promptkit.Backend{ID: "same-id", Endpoint: endpoint}),
			promptkit.WithProfiles(promptkit.Profile{ID: "profile", BackendID: "same-id", Model: "model"}),
		)
		if err != nil {
			t.Fatalf("construct isolated engine: %v", err)
		}
		return value
	}
	firstEngine := newEngine("http://one.example/v1")
	secondEngine := newEngine("http://two.example/v1")
	for engine, wantEndpoint := range map[*promptkit.Engine]string{
		firstEngine:  "http://one.example/v1",
		secondEngine: "http://two.example/v1",
	} {
		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
		if err != nil {
			t.Fatalf("prepare isolated engine: %v", err)
		}
		if prepared.EffectiveModelParams.Endpoint != wantEndpoint {
			t.Fatalf("isolated engine endpoint=%q, want %q", prepared.EffectiveModelParams.Endpoint, wantEndpoint)
		}
	}
}

func TestWithBackendCopiesQueueCapacity(t *testing.T) {
	queueCapacity := 4
	option := promptkit.WithBackend(promptkit.Backend{
		ID:               "custom",
		Endpoint:         "http://custom.example/v1",
		ConcurrencyLimit: 1,
		QueueCapacity:    &queueCapacity,
	})
	queueCapacity = -1

	_, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
		option,
		promptkit.WithProfiles(promptkit.Profile{
			ID: "profile", BackendID: "custom", Model: "model",
		}),
	)
	if err != nil {
		t.Fatalf("construct engine after mutating queue pointer: %v", err)
	}
}

func TestBackendRegistrationRejectsInvalidAndDuplicateDefinitions(t *testing.T) {
	cycle := map[string]any{}
	cycle["self"] = cycle
	tests := []struct {
		name     string
		backends []promptkit.Backend
	}{
		{name: "blank id", backends: []promptkit.Backend{{Endpoint: "http://example.test/v1"}}},
		{name: "invalid endpoint", backends: []promptkit.Backend{{ID: "custom", Endpoint: "ftp://example.test/v1"}}},
		{name: "invalid environment", backends: []promptkit.Backend{{ID: "custom", Endpoint: "http://example.test/v1", APIKeyEnv: "BAD-NAME"}}},
		{name: "reserved extra parameter", backends: []promptkit.Backend{{ID: "custom", Endpoint: "http://example.test/v1", ExtraParams: map[string]any{"model": "override"}}}},
		{name: "cyclic extra parameter", backends: []promptkit.Backend{{ID: "custom", Endpoint: "http://example.test/v1", ExtraParams: cycle}}},
		{name: "malformed JSON number", backends: []promptkit.Backend{{ID: "custom", Endpoint: "http://example.test/v1", ExtraParams: map[string]any{"value": json.Number("01")}}}},
		{name: "duplicate consumer id", backends: []promptkit.Backend{
			{ID: " custom ", Endpoint: "http://one.example/v1"},
			{ID: "custom", Endpoint: "http://two.example/v1"},
		}},
		{name: "reserved built-in id", backends: []promptkit.Backend{{
			ID: promptkit.BackendOpenRouter, Endpoint: "http://replacement.example/v1",
		}}},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			options := []promptkit.Option{
				promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
			}
			for _, backend := range tt.backends {
				options = append(options, promptkit.WithBackend(backend))
			}
			_, err := promptkit.NewEngine(promptkit.Config{}, options...)
			if !errors.Is(err, promptkit.ErrInvalidConfig) {
				t.Fatalf("expected ErrInvalidConfig, got %v", err)
			}
		})
	}
}

func TestBackendExtraParamsAreDeeplyCopiedAtConstructionAndLookup(t *testing.T) {
	nested := map[string]any{"value": "original"}
	extraParams := map[string]any{"nested": nested}
	engine, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
		promptkit.WithBackend(promptkit.Backend{
			ID: "custom", Endpoint: "http://example.test/v1", ExtraParams: extraParams,
		}),
		promptkit.WithProfiles(promptkit.Profile{ID: "profile", BackendID: "custom", Model: "model"}),
	)
	if err != nil {
		t.Fatalf("construct engine: %v", err)
	}
	nested["value"] = "mutated input"
	extraParams["later"] = true

	prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if err != nil {
		t.Fatalf("first prepare: %v", err)
	}
	gotNested := prepared.EffectiveModelParams.ExtraParams["nested"].(map[string]any)
	if gotNested["value"] != "original" || prepared.EffectiveModelParams.ExtraParams["later"] != nil {
		t.Fatalf("backend retained caller mutations: %#v", prepared.EffectiveModelParams.ExtraParams)
	}
	gotNested["value"] = "mutated lookup"

	prepared, err = engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if err != nil {
		t.Fatalf("second prepare: %v", err)
	}
	gotNested = prepared.EffectiveModelParams.ExtraParams["nested"].(map[string]any)
	if gotNested["value"] != "original" {
		t.Fatalf("backend retained lookup mutation: %#v", prepared.EffectiveModelParams.ExtraParams)
	}
}

func TestPreparedRunJSONTimingRoundTrips(t *testing.T) {
	start := time.Date(2026, time.July, 29, 12, 0, 0, 0, time.UTC)
	prepared := promptkit.PreparedRun{
		PromptID:   "prompt",
		StartTime:  start,
		EndTime:    start.Add(1250 * time.Millisecond),
		DurationMS: 1250,
	}

	payload, err := json.Marshal(prepared)
	if err != nil {
		t.Fatalf("marshal prepared run: %v", err)
	}
	var decoded promptkit.PreparedRun
	if err := json.Unmarshal(payload, &decoded); err != nil {
		t.Fatalf("unmarshal prepared run: %v", err)
	}
	if decoded.DurationMS != prepared.DurationMS ||
		!decoded.StartTime.Equal(prepared.StartTime) ||
		!decoded.EndTime.Equal(prepared.EndTime) {
		t.Fatalf("timing values did not round trip: got %#v, want %#v", decoded, prepared)
	}
}

func TestRunResultJSONUsesMillisecondsAndRoundTrips(t *testing.T) {
	start := time.Date(2026, time.July, 29, 12, 0, 0, 0, time.UTC)
	result := promptkit.RunResult{
		RunID:     "opaque-run-id",
		Artifact:  promptkit.Artifact{Name: "output", ContentType: "text/plain", Body: []byte("ok")},
		SessionID: "session-123",
		StartTime: start,
		EndTime:   start.Add(1500 * time.Millisecond),
		Duration:  1500 * time.Millisecond,
	}

	payload, err := json.Marshal(result)
	if err != nil {
		t.Fatalf("marshal run result: %v", err)
	}
	var object map[string]any
	if err := json.Unmarshal(payload, &object); err != nil {
		t.Fatalf("decode run result JSON: %v", err)
	}
	if got := object["duration_ms"]; got != float64(1500) {
		t.Fatalf("expected duration_ms=1500, got %#v in %s", got, payload)
	}
	if _, exists := object["duration"]; exists {
		t.Fatalf("unexpected nanosecond duration field in %s", payload)
	}
	if got := object["session_id"]; got != result.SessionID {
		t.Fatalf("expected session_id=%q, got %#v in %s", result.SessionID, got, payload)
	}
	artifact, ok := object["artifact"].(map[string]any)
	if !ok || artifact["content_type"] != "text/plain" {
		t.Fatalf("expected stable artifact JSON fields, got %#v", object["artifact"])
	}

	var decoded promptkit.RunResult
	if err := json.Unmarshal(payload, &decoded); err != nil {
		t.Fatalf("unmarshal run result: %v", err)
	}
	if decoded.SessionID != result.SessionID ||
		decoded.Duration != result.Duration ||
		!decoded.StartTime.Equal(result.StartTime) ||
		!decoded.EndTime.Equal(result.EndTime) {
		t.Fatalf("timing values did not round trip: got %#v, want %#v", decoded, result)
	}

	payload, err = json.Marshal(promptkit.RunResult{})
	if err != nil {
		t.Fatalf("marshal zero run result: %v", err)
	}
	for _, field := range []string{"session_id", "start_time", "end_time", "duration_ms"} {
		if strings.Contains(string(payload), `"`+field+`"`) {
			t.Fatalf("expected zero %s to be omitted, got %s", field, payload)
		}
	}
	var decodedEmpty promptkit.RunResult
	if err := json.Unmarshal(payload, &decodedEmpty); err != nil {
		t.Fatalf("unmarshal run result without session_id: %v", err)
	}
	if decodedEmpty.SessionID != "" {
		t.Fatalf("expected absent session_id to decode empty, got %q", decodedEmpty.SessionID)
	}
}

func TestEngineValidationIsSinglePass(t *testing.T) {
	client := &fakeLLMClient{
		response: &promptkit.GenerateResponse{Content: "not-json"},
	}
	engine := newContractEngineWithOptions(t, frameworkSchemaDir, promptkit.WithLLMClient(client))

	result, err := engine.Run(context.Background(), promptkit.RunRequest{
		PromptID: frameworkMarkdownSummaryPromptID,
		Inputs: map[string]promptkit.ArtifactRef{
			"transcript": promptkit.Inline("Rin opens the gate."),
			"glossary":   promptkit.Inline("gate: A guarded passage."),
		},
		Validation: &promptkit.OutputContract{
			Format:         promptkit.FormatJSON,
			ValidationMode: promptkit.ValidationJSON,
			RepairAttempts: 3,
		},
	})
	if err != nil {
		t.Fatalf("run with failed content validation: %v", err)
	}
	if result.Validation.Status != promptkit.ValidationFailed ||
		result.Validation.RepairAttempts != 0 {
		t.Fatalf("expected failed single-pass validation, got %#v", result.Validation)
	}
	if len(client.requests) != 1 {
		t.Fatalf("expected one model generation, got %d", len(client.requests))
	}
}

func TestRepeatedOptionsUseLastValueInEachCategory(t *testing.T) {
	profile := promptkit.Profile{ID: "profile", Endpoint: "http://example.test/v1", Model: "model"}

	t.Run("prompt source", func(t *testing.T) {
		engine, err := promptkit.NewEngine(promptkit.Config{},
			promptkit.WithPromptFS(contractPromptFS("first", "profile", "first"), "."),
			promptkit.WithPromptFS(contractPromptFS("second", "profile", "second"), "."),
			promptkit.WithProfiles(profile),
		)
		if err != nil {
			t.Fatalf("construct engine: %v", err)
		}
		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "second"})
		if err != nil {
			t.Fatalf("prepare from last prompt source: %v", err)
		}
		if prepared.Messages[0].Content != "second" {
			t.Fatalf("expected last prompt source, got %#v", prepared.Messages)
		}
	})

	t.Run("profile source", func(t *testing.T) {
		engine, err := promptkit.NewEngine(promptkit.Config{},
			promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
			promptkit.WithProfileFS(contractProfileFS("profile", "first-model"), "."),
			promptkit.WithProfileFS(contractProfileFS("profile", "second-model"), "."),
		)
		if err != nil {
			t.Fatalf("construct engine: %v", err)
		}
		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
		if err != nil {
			t.Fatalf("prepare from last profile source: %v", err)
		}
		if prepared.EffectiveModelParams.Model != "second-model" {
			t.Fatalf("expected last profile source, got %q", prepared.EffectiveModelParams.Model)
		}
	})

	t.Run("in-memory profiles", func(t *testing.T) {
		first := profile
		first.Model = "first-model"
		second := profile
		second.Model = "second-model"
		engine, err := promptkit.NewEngine(promptkit.Config{},
			promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
			promptkit.WithProfiles(first),
			promptkit.WithProfiles(second),
		)
		if err != nil {
			t.Fatalf("construct engine: %v", err)
		}
		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
		if err != nil {
			t.Fatalf("prepare from last in-memory profile option: %v", err)
		}
		if prepared.EffectiveModelParams.Model != "second-model" {
			t.Fatalf("expected last in-memory profiles, got %q", prepared.EffectiveModelParams.Model)
		}
	})

	t.Run("schema source", func(t *testing.T) {
		engine, err := promptkit.NewEngine(promptkit.Config{},
			promptkit.WithPromptFS(contractSchemaPromptFS(), "."),
			promptkit.WithProfiles(profile),
			promptkit.WithSchemaFS(contractSchemaFS("first"), "."),
			promptkit.WithSchemaFS(contractSchemaFS("second"), "."),
		)
		if err != nil {
			t.Fatalf("construct engine: %v", err)
		}
		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "schema-prompt"})
		if err != nil {
			t.Fatalf("prepare from last schema source: %v", err)
		}
		schema := prepared.StructuredOutput.JSONSchema.Schema.(map[string]any)
		if schema["title"] != "second" {
			t.Fatalf("expected last schema source, got %#v", schema)
		}
	})

	t.Run("model client", func(t *testing.T) {
		var firstCalls, secondCalls atomic.Int64
		engine, err := promptkit.NewEngine(promptkit.Config{},
			promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
			promptkit.WithProfiles(profile),
			promptkit.WithLLMClient(countingLLMClient{calls: &firstCalls}),
			promptkit.WithLLMClient(countingLLMClient{calls: &secondCalls}),
		)
		if err != nil {
			t.Fatalf("construct engine: %v", err)
		}
		if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); err != nil {
			t.Fatalf("run with last model client: %v", err)
		}
		if firstCalls.Load() != 0 || secondCalls.Load() != 1 {
			t.Fatalf("expected only last client call, got first=%d second=%d", firstCalls.Load(), secondCalls.Load())
		}
	})

	t.Run("artifact reader", func(t *testing.T) {
		engine, err := promptkit.NewEngine(promptkit.Config{},
			promptkit.WithPromptFS(contractInputPromptFS(), "."),
			promptkit.WithProfiles(profile),
			promptkit.WithArtifactReader(fixedArtifactReader("first")),
			promptkit.WithArtifactReader(fixedArtifactReader("second")),
		)
		if err != nil {
			t.Fatalf("construct engine: %v", err)
		}
		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
			PromptID: "input-prompt",
			Inputs:   map[string]promptkit.ArtifactRef{"input": promptkit.Inline("ignored")},
		})
		if err != nil {
			t.Fatalf("prepare with last artifact reader: %v", err)
		}
		if prepared.Messages[0].Content != "second" {
			t.Fatalf("expected last artifact reader, got %#v", prepared.Messages)
		}
	})
}

func TestEngineSupportsConcurrentPrepareAndRun(t *testing.T) {
	engine, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
		promptkit.WithBackend(promptkit.Backend{
			ID:       "concurrent",
			Endpoint: "http://example.test/v1",
		}),
		promptkit.WithProfiles(promptkit.Profile{
			ID:        "profile",
			BackendID: "concurrent",
			Model:     "model",
		}),
		promptkit.WithLLMClient(countingLLMClient{}),
	)
	if err != nil {
		t.Fatalf("construct engine: %v", err)
	}

	const calls = 40
	errs := make(chan error, calls)
	var wg sync.WaitGroup
	for i := 0; i < calls; i++ {
		wg.Add(1)
		go func(run bool) {
			defer wg.Done()
			request := promptkit.RunRequest{PromptID: "prompt"}
			if run {
				_, err := engine.Run(context.Background(), request)
				errs <- err
				return
			}
			_, err := engine.Prepare(context.Background(), request)
			errs <- err
		}(i%2 == 0)
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		if err != nil {
			t.Fatalf("concurrent call failed: %v", err)
		}
	}
}

type countingLLMClient struct {
	calls *atomic.Int64
}

func (c countingLLMClient) Generate(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
	if c.calls != nil {
		c.calls.Add(1)
	}
	return &promptkit.GenerateResponse{Content: "ok"}, nil
}

type fixedArtifactReader string

func (r fixedArtifactReader) Read(context.Context, promptkit.ArtifactRef) (*promptkit.Artifact, error) {
	return &promptkit.Artifact{Body: []byte(r)}, nil
}

func contractPromptFS(id, profileID, message string) fstest.MapFS {
	return fstest.MapFS{
		"prompt.yaml": &fstest.MapFile{Data: []byte(fmt.Sprintf(`id: %s
version: "1"
default_profile: %s
messages:
  - role: user
    content: %q
output:
  format: text
  validation_mode: none
`, id, profileID, message))},
	}
}

func contractInputPromptFS() fstest.MapFS {
	return fstest.MapFS{
		"prompt.yaml": &fstest.MapFile{Data: []byte(`id: input-prompt
version: "1"
default_profile: profile
inputs:
  - name: input
    required: true
messages:
  - role: user
    content: '{{input "input"}}'
output:
  format: text
  validation_mode: none
`)},
	}
}

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

func contractSchemaPromptFS() fstest.MapFS {
	return fstest.MapFS{
		"prompt.yaml": &fstest.MapFile{Data: []byte(`id: schema-prompt
version: "1"
default_profile: profile
messages:
  - role: user
    content: message
output:
  format: json
  validation_mode: json_schema
  schema_path: schema.json
`)},
	}
}

func contractSchemaFS(title string) fstest.MapFS {
	return fstest.MapFS{
		"schema.json": &fstest.MapFile{Data: []byte(fmt.Sprintf(`{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "title": %q,
  "type": "object"
}`, title))},
	}
}