promptkit/generation_error_contract_test.go

package promptkit_test

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"testing"

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

func TestBuiltInGenerationError(t *testing.T) {
	const (
		codeMarker    = "provider-code-marker"
		typeMarker    = "provider-type-marker"
		messageMarker = "provider-message-marker"
	)
	engine := newBuiltInGenerationErrorEngine(t, http.StatusUnprocessableEntity,
		`{"error":{"code":"`+codeMarker+`","type":"`+typeMarker+`","message":"`+messageMarker+`"}}`)

	result, err := engine.Run(context.Background(), generationErrorRunRequest())
	if result != nil {
		t.Fatalf("Run result = %#v, want nil", result)
	}
	assertGenerationError(t, err, http.StatusUnprocessableEntity, codeMarker, typeMarker, messageMarker)

	preparedEngine := newBuiltInGenerationErrorEngine(t, http.StatusServiceUnavailable, `{"error":{}}`)
	prepared, err := preparedEngine.PrepareExecution(context.Background(), generationErrorRunRequest())
	if err != nil {
		t.Fatalf("PrepareExecution: %v", err)
	}
	result, err = preparedEngine.RunPrepared(context.Background(), prepared)
	if result != nil {
		t.Fatalf("RunPrepared result = %#v, want nil", result)
	}
	assertGenerationError(t, err, http.StatusServiceUnavailable, "", "", "")
}

func TestBuiltInGenerationErrorWithAppendedMessages(t *testing.T) {
	const messageMarker = "combined-request-provider-marker"
	engine := newBuiltInGenerationErrorEngine(t, http.StatusUnprocessableEntity,
		`{"error":{"message":"`+messageMarker+`"}}`)
	request := generationErrorRunRequest()
	request.AppendedMessages = []promptkit.RenderedMessage{
		{Role: promptkit.RoleAssistant, Content: "previous response"},
		{Role: promptkit.RoleUser, Content: "consumer correction"},
	}

	result, err := engine.Run(context.Background(), request)
	if result != nil {
		t.Fatalf("Run result = %#v, want nil", result)
	}
	assertGenerationError(t, err, http.StatusUnprocessableEntity, "", "", messageMarker)
}

func TestBuiltInRepairGenerationError(t *testing.T) {
	const (
		codeMarker    = "repair-code-marker"
		typeMarker    = "repair-type-marker"
		messageMarker = "repair-message-marker"
	)
	calls := 0
	config := contractConfig(frameworkSchemaDir)
	config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		calls++
		if calls == 1 {
			body := `{"choices":[{"message":{"content":"not-json"}}]}`
			return &http.Response{StatusCode: http.StatusOK, ContentLength: int64(len(body)), Body: io.NopCloser(strings.NewReader(body))}, nil
		}
		body := `{"error":{"code":"` + codeMarker + `","type":"` + typeMarker + `","message":"` + messageMarker + `"}}`
		return &http.Response{StatusCode: http.StatusUnprocessableEntity, ContentLength: int64(len(body)), Body: io.NopCloser(strings.NewReader(body))}, nil
	})}
	engine, err := promptkit.NewEngine(config)
	if err != nil {
		t.Fatalf("NewEngine: %v", err)
	}
	req := generationErrorRunRequest()
	req.Validation = &promptkit.OutputContract{
		Format:         promptkit.FormatJSON,
		ValidationMode: promptkit.ValidationJSON,
		RepairAttempts: 1,
	}

	result, err := engine.Run(context.Background(), req)
	if result != nil {
		t.Fatalf("Run result = %#v, want nil", result)
	}
	if calls != 2 {
		t.Fatalf("provider calls = %d, want 2", calls)
	}
	assertGenerationError(t, err, http.StatusUnprocessableEntity, codeMarker, typeMarker, messageMarker)
}

func TestBuiltInGenerationFailureCategories(t *testing.T) {
	transportCause := errors.New("private transport diagnostic")
	config := contractConfig(frameworkSchemaDir)
	config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		return nil, transportCause
	})}
	engine, err := promptkit.NewEngine(config)
	if err != nil {
		t.Fatal(err)
	}
	result, err := engine.Run(context.Background(), generationErrorRunRequest())
	if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, promptkit.ErrTransport) || !errors.Is(err, transportCause) {
		t.Fatalf("transport result=%#v, error=%v", result, err)
	}
	if errors.Is(err, promptkit.ErrMalformedResponse) || strings.Contains(fmt.Sprintf("%#v", err), transportCause.Error()) {
		t.Fatalf("transport category or formatting = %v", err)
	}

	for _, tc := range []struct {
		name string
		body string
	}{
		{name: "invalid framing", body: `{"choices":[]}{}`},
		{name: "missing choices", body: `{"choices":[]}`},
		{name: "missing content", body: `{"choices":[{"message":{}}]}`},
	} {
		t.Run(tc.name, func(t *testing.T) {
			engine := newBuiltInGenerationErrorEngine(t, http.StatusOK, tc.body)
			result, err := engine.Run(context.Background(), generationErrorRunRequest())
			if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, promptkit.ErrMalformedResponse) || errors.Is(err, promptkit.ErrTransport) {
				t.Fatalf("malformed result=%#v, error=%v", result, err)
			}
		})
	}
}

func TestBuiltInGenerationPreservesContextIdentity(t *testing.T) {
	for _, tc := range []struct {
		name  string
		cause error
	}{
		{name: "canceled", cause: context.Canceled},
		{name: "deadline", cause: context.DeadlineExceeded},
	} {
		t.Run(tc.name, func(t *testing.T) {
			config := contractConfig(frameworkSchemaDir)
			config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
				return nil, tc.cause
			})}
			engine, err := promptkit.NewEngine(config)
			if err != nil {
				t.Fatal(err)
			}
			result, err := engine.Run(context.Background(), generationErrorRunRequest())
			if result != nil || !errors.Is(err, tc.cause) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, promptkit.ErrTransport) {
				t.Fatalf("context result=%#v, error=%v", result, err)
			}
		})
	}
}

func TestInjectedGenerationFailureCategories(t *testing.T) {
	for _, tc := range []struct {
		name      string
		client    *fakeLLMClient
		malformed bool
	}{
		{name: "nil response", client: &fakeLLMClient{}, malformed: true},
		{name: "arbitrary error", client: &fakeLLMClient{err: errors.New("injected failure")}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			engine := newContractEngineWithOptions(t, frameworkSchemaDir, promptkit.WithLLMClient(tc.client))
			result, err := engine.Run(context.Background(), generationErrorRunRequest())
			if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || errors.Is(err, promptkit.ErrTransport) || errors.Is(err, promptkit.ErrMalformedResponse) != tc.malformed {
				t.Fatalf("injected result=%#v, error=%v", result, err)
			}
			if tc.client.err != nil && !errors.Is(err, tc.client.err) {
				t.Fatalf("injected cause lost: %v", err)
			}
		})
	}
}

func assertGenerationError(t *testing.T, err error, statusCode int, code, providerType, message string) {
	t.Helper()

	if !errors.Is(err, promptkit.ErrLLMGenerate) {
		t.Fatalf("errors.Is(%v, ErrLLMGenerate) = false", err)
	}
	var generationErr *promptkit.GenerationError
	if !errors.As(err, &generationErr) || generationErr == nil {
		t.Fatalf("error = %T, want *GenerationError", err)
	}
	if generationErr.StatusCode() != statusCode || generationErr.ProviderCode() != code || generationErr.ProviderType() != providerType || generationErr.ProviderMessage() != message {
		t.Fatalf("GenerationError = %#v", generationErr)
	}

	wantFormatted := fmt.Sprintf("failed to generate output: provider returned HTTP status %d", statusCode)
	for _, rendered := range []string{fmt.Sprintf("%v", generationErr), fmt.Sprintf("%+v", generationErr), fmt.Sprintf("%#v", generationErr)} {
		if rendered != wantFormatted {
			t.Fatalf("formatted error = %q, want %q", rendered, wantFormatted)
		}
		for _, marker := range []string{code, providerType, message} {
			if marker != "" && strings.Contains(rendered, marker) {
				t.Fatalf("formatted error exposed provider marker %q: %q", marker, rendered)
			}
		}
	}
}

func newBuiltInGenerationErrorEngine(t *testing.T, statusCode int, body string) *promptkit.Engine {
	t.Helper()

	config := contractConfig(frameworkSchemaDir)
	noRetries := 0
	config.TransientRetries = &noRetries
	config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		return &http.Response{
			StatusCode:    statusCode,
			ContentLength: int64(len(body)),
			Body:          io.NopCloser(strings.NewReader(body)),
		}, nil
	})}
	engine, err := promptkit.NewEngine(config)
	if err != nil {
		t.Fatalf("NewEngine: %v", err)
	}
	return engine
}

func generationErrorRunRequest() promptkit.RunRequest {
	return promptkit.RunRequest{
		PromptID: frameworkMarkdownSummaryPromptID,
		Inputs: map[string]promptkit.ArtifactRef{
			"transcript": promptkit.Inline("Rin opens the gate."),
			"glossary":   promptkit.Inline("gate: A guarded passage."),
		},
	}
}