promptkit/schema_file_contract_test.go

package promptkit_test

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/santhosh-tekuri/jsonschema/v6"
	"go.maximumdirect.net/configkit"
	"go.maximumdirect.net/promptkit"
)

type rejectWireSchemaLoad struct{}

func (rejectWireSchemaLoad) Load(uri string) (any, error) {
	return nil, fmt.Errorf("unexpected external resource %s", uri)
}

func TestSchemaFilePreparationAndNativeSubmission(t *testing.T) {
	firstDir, secondDir := t.TempDir(), t.TempDir()
	first, second := filepath.Join(firstDir, "schema.json"), filepath.Join(secondDir, "schema.json")
	if err := os.WriteFile(first, []byte(`{"type":"string"}`), 0600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(second, []byte(`{"$defs":{"value":{"type":"integer"}},"$ref":"#/$defs/value"}`), 0600); err != nil {
		t.Fatal(err)
	}
	input := configkit.Input{Source: configkit.Mapping("schema test", map[string]any{
		"profiles": []any{map[string]any{"id": "profile", "endpoint": "https://example.invalid/v1", "model": "model"}},
		"prompts":  []any{map[string]any{"id": "prompt", "default_profile": "profile", "messages": []any{map[string]any{"role": "user", "content": "hello"}}, "output": map[string]any{"format": "json", "validation_mode": "json_schema", "schema_path": "schema.json"}}},
	})}
	var payload map[string]any
	client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
		if err := json.NewDecoder(request.Body).Decode(&payload); err != nil {
			t.Fatal(err)
		}
		return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"3"}}]}`))}, nil
	})}
	engine, err := promptkit.NewEngine(promptkit.Config{HTTPClient: client, Sources: promptkit.ConfigurationSources{Main: &input}}, promptkit.WithSchemaFile(first), promptkit.WithSchemaFile(second))
	if err != nil {
		t.Fatal(err)
	}
	if err := engine.ValidateRequest(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); err != nil {
		t.Fatalf("validate request: %v", err)
	}
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if err != nil || prepared.Details().StructuredOutput == nil {
		t.Fatalf("prepared = %+v, %v", prepared, err)
	}
	if err := os.Remove(second); err != nil {
		t.Fatal(err)
	}
	result, err := engine.RunPrepared(context.Background(), prepared)
	if err != nil || result.Validation.Status != promptkit.ValidationPassed {
		t.Fatalf("run = %+v, %v", result, err)
	}
	encoded, err := json.Marshal(payload["response_format"])
	if err != nil || !strings.Contains(string(encoded), `"$defs"`) || !strings.Contains(string(encoded), `#/$defs/value`) {
		t.Fatalf("wire schema = %s, %v", encoded, err)
	}
	compiler := jsonschema.NewCompiler()
	compiler.UseLoader(rejectWireSchemaLoad{})
	const submission = "https://provider.invalid/request-schema"
	document := payload["response_format"].(map[string]any)["json_schema"].(map[string]any)["schema"]
	if err := compiler.AddResource(submission, document); err != nil {
		t.Fatal(err)
	}
	if _, err := compiler.Compile(submission); err != nil {
		t.Fatalf("outbound schema needs local source context: %v", err)
	}
	if _, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); !errors.Is(err, promptkit.ErrValidation) {
		t.Fatalf("missing root error = %v", err)
	}
}

func TestSchemaFilePreparationDetachesLargeAnnotation(t *testing.T) {
	// Annotation values are not constraints; this array exceeds the former
	// provider-value copier node limit without requiring a large fixture file.
	const entries = 100_001
	path := filepath.Join(t.TempDir(), "schema.json")
	schema := `{"type":"string","examples":[` + strings.Repeat(`0,`, entries-1) + `0]}`
	if err := os.WriteFile(path, []byte(schema), 0o600); err != nil {
		t.Fatal(err)
	}
	input := configkit.Input{Source: configkit.Mapping("schema test", map[string]any{
		"profiles": []any{map[string]any{"id": "profile", "endpoint": "https://example.invalid/v1", "model": "model"}},
		"prompts":  []any{map[string]any{"id": "prompt", "default_profile": "profile", "messages": []any{map[string]any{"role": "user", "content": "hello"}}, "output": map[string]any{"format": "json", "validation_mode": "json_schema", "schema_path": "schema.json"}}},
	})}
	engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &input}}, promptkit.WithSchemaFile(path))
	if err != nil {
		t.Fatal(err)
	}
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if err != nil {
		t.Fatalf("prepare large schema: %v", err)
	}
	first := prepared.Details().StructuredOutput.JSONSchema.Schema.(map[string]any)
	examples := first["examples"].([]any)
	if len(examples) != entries {
		t.Fatalf("examples = %d, want %d", len(examples), entries)
	}
	examples[0] = "changed"
	second := prepared.Details().StructuredOutput.JSONSchema.Schema.(map[string]any)
	if _, ok := second["examples"].([]any)[0].(json.Number); !ok {
		t.Fatal("schema details share mutable annotation values")
	}
}

func TestSchemaFileRejectsInvalidSourcesAndExternalReferences(t *testing.T) {
	dir := t.TempDir()
	for _, name := range []string{filepath.Join(dir, "missing.json"), dir} {
		if _, err := promptkit.NewEngine(promptkit.Config{}, promptkit.WithSchemaFile(name)); !errors.Is(err, promptkit.ErrInvalidConfig) {
			t.Fatalf("source %q = %v", name, err)
		}
	}
	root := filepath.Join(dir, "schema.json")
	if err := os.WriteFile(root, []byte(`{"$defs":{"unused":{"$ref":"other.json"}}}`), 0600); err != nil {
		t.Fatal(err)
	}
	input := configkit.Input{Source: configkit.Mapping("schema test", map[string]any{
		"profiles": []any{map[string]any{"id": "profile", "endpoint": "https://example.invalid/v1", "model": "model"}},
		"prompts":  []any{map[string]any{"id": "prompt", "default_profile": "profile", "messages": []any{map[string]any{"role": "user", "content": "hello"}}, "output": map[string]any{"format": "json", "validation_mode": "json_schema", "schema_path": "schema.json"}}},
	})}
	calls := 0
	client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { calls++; return nil, errors.New("unexpected dispatch") })}
	engine, err := promptkit.NewEngine(promptkit.Config{HTTPClient: client, Sources: promptkit.ConfigurationSources{Main: &input}}, promptkit.WithSchemaFile(root))
	if err != nil {
		t.Fatal(err)
	}
	_, err = engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if !errors.Is(err, promptkit.ErrValidation) || calls != 0 {
		t.Fatalf("external reference = %v, calls = %d", err, calls)
	}
}