promptkit/transient_retries_contract_test.go

package promptkit_test

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

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

func retryHTTPResponse(status int, content string) *http.Response {
	body := `{"error":{"message":"private-provider-marker"}}`
	if status == http.StatusOK {
		body = `{"choices":[{"message":{"content":` + strconv.Quote(content) + `}}]}`
	}
	return &http.Response{StatusCode: status, Body: io.NopCloser(strings.NewReader(body))}
}

func retryTestEngine(t *testing.T, budget *int, next func(int) *http.Response) (*promptkit.Engine, *int) {
	t.Helper()
	calls := 0
	config := contractConfig(frameworkSchemaDir)
	config.TransientRetries = budget
	config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		calls++
		return next(calls), nil
	})}
	engine, err := promptkit.NewEngine(config)
	if err != nil {
		t.Fatal(err)
	}
	return engine, &calls
}

func TestBuiltInRetryDefaultsAndRequestOverride(t *testing.T) {
	engine, calls := retryTestEngine(t, nil, func(call int) *http.Response {
		if call < 3 {
			return retryHTTPResponse(503, "")
		}
		return retryHTTPResponse(200, "ok")
	})
	result, err := engine.Run(context.Background(), generationErrorRunRequest())
	if err != nil || result == nil || result.RawOutput != "ok" || *calls != 3 {
		t.Fatalf("default result=%+v, error=%v, calls=%d", result, err, *calls)
	}

	engineBudget := 1
	engine, calls = retryTestEngine(t, &engineBudget, func(int) *http.Response { return retryHTTPResponse(503, "") })
	engineBudget = 5 // Construction copied the configured value.
	zero := 0
	req := generationErrorRunRequest()
	req.TransientRetries = &zero
	_, err = engine.Run(context.Background(), req)
	if !errors.Is(err, promptkit.ErrLLMGenerate) || errors.Is(err, promptkit.ErrRetriesExhausted) || *calls != 1 {
		t.Fatalf("zero override: error=%v, calls=%d", err, *calls)
	}
	req.TransientRetries = nil
	_, err = engine.Run(context.Background(), req)
	var exhausted *promptkit.RetryExhaustedError
	if !errors.As(err, &exhausted) || exhausted.Attempts() != 2 || *calls != 3 {
		t.Fatalf("inherited budget: error=%v, calls=%d", err, *calls)
	}
}

func TestRetryValidationAndPreparedFreezing(t *testing.T) {
	negative := -1
	config := contractConfig(frameworkSchemaDir)
	config.TransientRetries = &negative
	if _, err := promptkit.NewEngine(config); !errors.Is(err, promptkit.ErrInvalidConfig) {
		t.Fatalf("negative engine budget = %v", err)
	}

	engine, calls := retryTestEngine(t, nil, func(int) *http.Response { return retryHTTPResponse(503, "") })
	req := generationErrorRunRequest()
	req.TransientRetries = &negative
	if _, err := engine.Prepare(context.Background(), req); !errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("Prepare negative budget = %v", err)
	}
	if _, err := engine.PrepareExecution(context.Background(), req); !errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("PrepareExecution negative budget = %v", err)
	}
	if _, err := engine.Run(context.Background(), req); !errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("Run negative budget = %v", err)
	}
	if *calls != 0 {
		t.Fatalf("invalid requests dispatched %d calls", *calls)
	}
	maximum := int(^uint(0) >> 1)
	largeBudgetEngine, largeCalls := retryTestEngine(t, &maximum, func(int) *http.Response {
		return retryHTTPResponse(200, "ok")
	})
	if _, err := largeBudgetEngine.Run(context.Background(), generationErrorRunRequest()); err != nil || *largeCalls != 1 {
		t.Fatalf("large valid budget: error=%v, calls=%d", err, *largeCalls)
	}

	zero := 0
	req.TransientRetries = &zero
	prepared, err := engine.PrepareExecution(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	zero = 2 // The handle retains the selected zero.
	_, err = engine.RunPrepared(context.Background(), prepared)
	if !errors.Is(err, promptkit.ErrLLMGenerate) || errors.Is(err, promptkit.ErrRetriesExhausted) || *calls != 1 {
		t.Fatalf("frozen prepared budget: error=%v, calls=%d", err, *calls)
	}
}

func TestBuiltInRetryBudgetAppliesToRepairGeneration(t *testing.T) {
	one := 1
	engine, calls := retryTestEngine(t, &one, func(call int) *http.Response {
		switch call {
		case 1, 3:
			return retryHTTPResponse(503, "")
		case 2:
			return retryHTTPResponse(200, "not-json")
		default:
			return retryHTTPResponse(200, `{"ok":true}`)
		}
	})
	req := generationErrorRunRequest()
	req.Validation = &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}
	result, err := engine.Run(context.Background(), req)
	if err != nil || result == nil || !result.Validation.IsValid || *calls != 4 {
		t.Fatalf("repair result=%+v, error=%v, calls=%d", result, err, *calls)
	}
}

func TestRetryExhaustionRetainsPublicProviderCause(t *testing.T) {
	one := 1
	engine, calls := retryTestEngine(t, &one, func(int) *http.Response { return retryHTTPResponse(503, "") })
	result, err := engine.Run(context.Background(), generationErrorRunRequest())
	if result != nil || *calls != 2 || !errors.Is(err, promptkit.ErrRetriesExhausted) || !errors.Is(err, promptkit.ErrLLMGenerate) {
		t.Fatalf("exhaustion result=%+v, error=%v, calls=%d", result, err, *calls)
	}
	var exhausted *promptkit.RetryExhaustedError
	var provider *promptkit.GenerationError
	if !errors.As(err, &exhausted) || exhausted.Attempts() != 2 || !errors.As(err, &provider) || provider.StatusCode() != 503 || provider.ProviderMessage() != "private-provider-marker" {
		t.Fatalf("exhaustion cause=%v", err)
	}
	for _, rendered := range []string{fmt.Sprintf("%v", err), fmt.Sprintf("%#v", err)} {
		if strings.Contains(rendered, "private-provider-marker") {
			t.Fatalf("provider detail leaked: %s", rendered)
		}
	}
	var zero *promptkit.RetryExhaustedError
	if zero.Attempts() != 0 || !errors.Is(zero, promptkit.ErrRetriesExhausted) || zero.Error() == "" || zero.GoString() == "" || zero.Unwrap() != nil {
		t.Fatalf("nil error behavior = %v", zero)
	}
	zeroValue := &promptkit.RetryExhaustedError{}
	if zeroValue.Attempts() != 0 || !errors.Is(zeroValue, promptkit.ErrRetriesExhausted) || zeroValue.Error() == "" || zeroValue.GoString() == "" || zeroValue.Unwrap() != nil {
		t.Fatalf("zero error behavior = %v", zeroValue)
	}
}

func TestInjectedClientIsNotAutomaticallyRetried(t *testing.T) {
	client := &fakeLLMClient{err: errors.New("injected failure")}
	engine := newContractEngineWithOptions(t, frameworkSchemaDir, promptkit.WithLLMClient(client))
	count := 3
	req := generationErrorRunRequest()
	req.TransientRetries = &count
	_, err := engine.Run(context.Background(), req)
	if len(client.requests) != 1 || !errors.Is(err, client.err) || errors.Is(err, promptkit.ErrRetriesExhausted) {
		t.Fatalf("injected calls=%d, error=%v", len(client.requests), err)
	}
}