promptkit/json_contract_test.go

package promptkit_test

import (
	"encoding/json"
	"math"
	"reflect"
	"strconv"
	"strings"
	"testing"
	"time"

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

func stringPointerForTest(value string) *string {
	return &value
}

func TestGenerateResponseUsagePresenceJSON(t *testing.T) {
	value := promptkit.GenerateResponse{
		Content: "ok",
		Usage:   promptkit.TokenUsage{TotalTokens: 5},
		UsagePresence: &promptkit.TokenUsagePresence{
			PromptTokens: true, CompletionTokens: true, TotalTokens: true,
			CachedTokens: true, CacheWriteTokens: true,
		},
	}
	payload, err := json.Marshal(value)
	if err != nil {
		t.Fatal(err)
	}
	object := decodeJSONObject(t, payload)
	if _, ok := object["usage_presence"]; !ok {
		t.Fatalf("missing usage_presence: %s", payload)
	}
	var decoded promptkit.GenerateResponse
	if err := json.Unmarshal(payload, &decoded); err != nil || !reflect.DeepEqual(decoded, value) {
		t.Fatalf("round trip = %+v, %v", decoded, err)
	}
	payload, err = json.Marshal(promptkit.GenerateResponse{})
	if err != nil {
		t.Fatal(err)
	}
	if _, ok := decodeJSONObject(t, payload)["usage_presence"]; ok {
		t.Fatalf("nil usage_presence was serialized: %s", payload)
	}
}

func TestPromptVersionJSONDistinguishesAbsenceAndDeclaredEmpty(t *testing.T) {
	tests := []struct {
		name        string
		value       *string
		wantPresent bool
		wantValue   string
	}{
		{name: "absent"},
		{name: "declared empty", value: stringPointerForTest(""), wantPresent: true},
		{name: "declared text", value: stringPointerForTest("release/v1"), wantPresent: true, wantValue: "release/v1"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			check := func(label string, marshal func() ([]byte, error)) {
				t.Helper()
				payload, err := marshal()
				if err != nil {
					t.Fatal(err)
				}
				object := decodeJSONObject(t, payload)
				value, present := object["prompt_version"]
				if present != test.wantPresent {
					t.Errorf("%s prompt_version present=%t, want %t: %s", label, present, test.wantPresent, payload)
				}
				var decoded string
				if test.wantPresent && (json.Unmarshal(value, &decoded) != nil || decoded != test.wantValue) {
					t.Errorf("%s prompt_version=%#v, want %q", label, value, test.wantValue)
				}
			}
			check("prepared", func() ([]byte, error) {
				return json.Marshal(promptkit.PreparedRun{PromptVersion: test.value})
			})
			check("result", func() ([]byte, error) {
				return json.Marshal(promptkit.RunResult{PromptVersion: test.value})
			})
		})
	}
}

func TestPreparedRunJSONContractRoundTripsAllFields(t *testing.T) {
	value := fullyPopulatedPreparedRun()
	payload, err := json.Marshal(value)
	if err != nil {
		t.Fatalf("marshal PreparedRun: %v", err)
	}

	object := decodeJSONObject(t, payload)
	assertJSONFields(t, object,
		"prompt_id",
		"prompt_version",
		"prompt_hash",
		"selected_profile_id",
		"selected_backend_id",
		"effective_model_params",
		"output_contract",
		"structured_output",
		"input_hashes",
		"session_id",
		"rendered_prompt_hash",
		"messages",
		"start_time",
		"end_time",
		"duration_ms",
	)

	var messages []map[string]json.RawMessage
	if err := json.Unmarshal(object["messages"], &messages); err != nil {
		t.Fatalf("decode PreparedRun messages: %v", err)
	}
	if len(messages) != 2 {
		t.Fatalf("message count = %d, want 2", len(messages))
	}
	if _, ok := messages[0]["cache_control"]; !ok {
		t.Fatalf("first message omitted cache_control: %s", object["messages"])
	}
	if _, ok := messages[1]["cache_control"]; ok {
		t.Fatalf("second message included empty cache_control: %s", object["messages"])
	}

	var decoded promptkit.PreparedRun
	if err := json.Unmarshal(payload, &decoded); err != nil {
		t.Fatalf("unmarshal PreparedRun: %v", err)
	}
	if !reflect.DeepEqual(decoded, value) {
		t.Fatalf("PreparedRun did not round trip:\ngot  %#v\nwant %#v", decoded, value)
	}

	zeroPayload, err := json.Marshal(promptkit.PreparedRun{})
	if err != nil {
		t.Fatalf("marshal zero PreparedRun: %v", err)
	}
	zeroObject := decodeJSONObject(t, zeroPayload)
	for _, field := range []string{
		"prompt_version",
		"prompt_hash",
		"selected_backend_id",
		"structured_output",
		"input_hashes",
		"session_id",
		"start_time",
		"end_time",
		"duration_ms",
	} {
		if _, ok := zeroObject[field]; ok {
			t.Fatalf("zero PreparedRun included %q: %s", field, zeroPayload)
		}
	}
}

func TestRunResultJSONContractRoundTripsAllFields(t *testing.T) {
	value := fullyPopulatedRunResult()
	payload, err := json.Marshal(value)
	if err != nil {
		t.Fatalf("marshal RunResult: %v", err)
	}

	object := decodeJSONObject(t, payload)
	assertJSONFields(t, object,
		"run_id",
		"artifact",
		"raw_output",
		"validation",
		"prompt_id",
		"prompt_version",
		"prompt_hash",
		"session_id",
		"rendered_prompt_hash",
		"selected_profile_id",
		"selected_backend_id",
		"model_name",
		"endpoint",
		"effective_model_params",
		"input_hashes",
		"usage",
		"start_time",
		"end_time",
		"duration_ms",
	)
	if _, ok := object["duration"]; ok {
		t.Fatalf("RunResult included nanosecond duration field: %s", payload)
	}

	var decoded promptkit.RunResult
	if err := json.Unmarshal(payload, &decoded); err != nil {
		t.Fatalf("unmarshal RunResult: %v", err)
	}
	if !reflect.DeepEqual(decoded, value) {
		t.Fatalf("RunResult did not round trip:\ngot  %#v\nwant %#v", decoded, value)
	}

	zeroPayload, err := json.Marshal(promptkit.RunResult{})
	if err != nil {
		t.Fatalf("marshal zero RunResult: %v", err)
	}
	zeroObject := decodeJSONObject(t, zeroPayload)
	for _, field := range []string{
		"prompt_version",
		"prompt_hash",
		"session_id",
		"selected_backend_id",
		"input_hashes",
		"start_time",
		"end_time",
		"duration_ms",
		"duration",
	} {
		if _, ok := zeroObject[field]; ok {
			t.Fatalf("zero RunResult included %q: %s", field, zeroPayload)
		}
	}
}

func TestRunResultJSONDurationMillisecondBoundaries(t *testing.T) {
	maxMilliseconds := int64(time.Duration(math.MaxInt64) / time.Millisecond)
	minMilliseconds := int64(time.Duration(math.MinInt64) / time.Millisecond)

	for _, milliseconds := range []int64{minMilliseconds, maxMilliseconds} {
		t.Run(strconv.FormatInt(milliseconds, 10), func(t *testing.T) {
			var decoded promptkit.RunResult
			if err := json.Unmarshal(durationPayload(milliseconds), &decoded); err != nil {
				t.Fatalf("decode representable duration_ms: %v", err)
			}
			want := time.Duration(milliseconds) * time.Millisecond
			if decoded.Duration != want {
				t.Fatalf("Duration = %v, want %v", decoded.Duration, want)
			}
		})
	}

	for _, milliseconds := range []int64{
		minMilliseconds - 1,
		maxMilliseconds + 1,
		math.MinInt64,
		math.MaxInt64,
	} {
		t.Run(strconv.FormatInt(milliseconds, 10), func(t *testing.T) {
			original := fullyPopulatedRunResult()
			decoded := original
			err := json.Unmarshal(durationPayload(milliseconds), &decoded)
			if err == nil || !strings.Contains(err.Error(), "duration_ms") || !strings.Contains(err.Error(), "time.Duration") {
				t.Fatalf("overflow error = %v", err)
			}
			if !reflect.DeepEqual(decoded, original) {
				t.Fatalf("failed decode partially updated receiver:\ngot  %#v\nwant %#v", decoded, original)
			}
		})
	}
}

func fullyPopulatedPreparedRun() promptkit.PreparedRun {
	start := time.Date(2026, time.August, 11, 12, 13, 14, 150_000_000, time.UTC)
	return promptkit.PreparedRun{
		PromptID:             "prompt.prepared",
		PromptVersion:        stringPointerForTest("2.1.0"),
		PromptHash:           "prompt-hash",
		SelectedProfileID:    "profile-prepared",
		SelectedBackendID:    "backend-prepared",
		EffectiveModelParams: jsonContractExecutionTarget(),
		OutputContract: promptkit.OutputContract{
			Format:         promptkit.FormatJSON,
			ValidationMode: promptkit.ValidationJSONSchema,
			SchemaPath:     "schemas/prepared.json",
			RepairAttempts: 2,
		},
		StructuredOutput: &promptkit.StructuredOutputSpec{
			Type: promptkit.StructuredOutputJSONSchema,
			JSONSchema: &promptkit.StructuredOutputJSONSpec{
				Name:   "prepared_schema",
				Strict: true,
				Schema: map[string]any{
					"type":     "object",
					"required": []any{"value"},
					"properties": map[string]any{
						"value": map[string]any{"minimum": float64(1)},
					},
				},
			},
		},
		InputHashes:        map[string]string{"first": "hash-1", "second": "hash-2"},
		SessionID:          "session-prepared",
		RenderedPromptHash: "rendered-hash",
		Messages: []promptkit.RenderedMessage{
			{
				Role:    "system",
				Content: "System message",
				CacheControl: &promptkit.CacheControl{
					Type: promptkit.CacheControlEphemeral,
					TTL:  "1h",
				},
			},
			{Role: "user", Content: "User message"},
		},
		StartTime:  start,
		EndTime:    start.Add(1501 * time.Millisecond),
		DurationMS: 1501,
	}
}

func fullyPopulatedRunResult() promptkit.RunResult {
	start := time.Date(2026, time.August, 11, 15, 16, 17, 250_000_000, time.UTC)
	return promptkit.RunResult{
		RunID: "run-id",
		Artifact: promptkit.Artifact{
			Name:        "result.json",
			ContentType: "application/json",
			Body:        []byte(`{"value":2}`),
			URI:         "memory://result.json",
			Size:        11,
			Hash:        "artifact-hash",
		},
		RawOutput: `{"value":2}`,
		Validation: promptkit.ValidationResult{
			Status:         promptkit.ValidationFailed,
			Mode:           promptkit.ValidationJSONSchema,
			Errors:         []string{"first error", "second error"},
			SchemaPath:     "schemas/result.json",
			RepairAttempts: 2,
			IsValid:        false,
		},
		PromptID:             "prompt.result",
		PromptVersion:        stringPointerForTest("3.2.1"),
		PromptHash:           "result-prompt-hash",
		SessionID:            "session-result",
		RenderedPromptHash:   "result-rendered-hash",
		SelectedProfileID:    "profile-result",
		SelectedBackendID:    "backend-result",
		ModelName:            "model-result",
		Endpoint:             "https://result.example/v1",
		EffectiveModelParams: jsonContractExecutionTarget(),
		InputHashes:          map[string]string{"input": "input-hash"},
		Usage: promptkit.TokenUsage{
			PromptTokens:     101,
			CompletionTokens: 202,
			TotalTokens:      303,
			CachedTokens:     44,
			CacheWriteTokens: 55,
		},
		StartTime: start,
		EndTime:   start.Add(1750 * time.Millisecond),
		Duration:  1750 * time.Millisecond,
	}
}

func jsonContractExecutionTarget() promptkit.ExecutionTarget {
	return promptkit.ExecutionTarget{
		BackendID:       "backend-target",
		Endpoint:        "https://target.example/v1",
		Model:           "model-target",
		Temperature:     0.75,
		MaxTokens:       321,
		TopP:            0.875,
		TimeoutSeconds:  43,
		ServiceTier:     "priority",
		ReasoningEffort: "high",
		APIKeyEnv:       "PROMPTKIT_JSON_CONTRACT_KEY",
		ExtraParams: map[string]any{
			"enabled": true,
			"weight":  float64(1.25),
			"nested":  map[string]any{"name": "value"},
		},
	}
}

func durationPayload(milliseconds int64) []byte {
	return []byte(`{"duration_ms":` + strconv.FormatInt(milliseconds, 10) + `}`)
}

func decodeJSONObject(t *testing.T, payload []byte) map[string]json.RawMessage {
	t.Helper()
	var object map[string]json.RawMessage
	if err := json.Unmarshal(payload, &object); err != nil {
		t.Fatalf("decode JSON object: %v", err)
	}
	return object
}

func assertJSONFields(t *testing.T, object map[string]json.RawMessage, fields ...string) {
	t.Helper()
	want := make(map[string]struct{}, len(fields))
	for _, field := range fields {
		want[field] = struct{}{}
	}
	for field := range object {
		if _, ok := want[field]; !ok {
			t.Errorf("unexpected JSON field %q", field)
		}
	}
	for field := range want {
		if _, ok := object[field]; !ok {
			t.Errorf("missing JSON field %q", field)
		}
	}
}