promptkit/public_contract_test.go

package promptkit_test

import (
	"context"
	"encoding/json"
	"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 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")},
		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)
	}
	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.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{"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 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.WithProfiles(promptkit.Profile{
			ID:       "profile",
			Endpoint: "http://example.test/v1",
			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))},
	}
}