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 := simpleConfig("message") 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(), simpleGenerationRequest()) 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 := simpleGenerationRequest() 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 := simpleConfig("message") 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 := simpleGenerationRequest() 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(), simpleGenerationRequest()); 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 := simpleGenerationRequest() 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(), simpleGenerationRequest()) 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 := newSimpleEngine(t, "message", promptkit.WithLLMClient(client)) count := 3 req := simpleGenerationRequest() 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) } }