Configure built-in retries across execution paths

transient_retries_contract_test.go

@@ -0,0 +1,171 @@
+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)
+	}
+}