diff --git a/llmgateway.go b/llmgateway.go index 5adfcfe..1495ad1 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -22,6 +22,7 @@ type Config struct { Agents map[string]*AgentConfig `json:"agents,omitempty"` Users *UsersConfig `json:"users,omitempty"` Redis *RedisConfig `json:"redis,omitempty"` + Retry RetryConfig `json:"retry,omitempty"` PassthroughUnknown bool `json:"passthroughUnknown,omitempty"` } @@ -77,6 +78,18 @@ type RedisConfig struct { DB int `json:"db,omitempty"` } +// RetryConfig configures same-provider retry on transient upstream +// failures (connection errors, HTTP 429, HTTP 5xx) — spec §1. The zero +// value (Enabled: false) disables retry, preserving v0.1 behavior +// exactly: every request still makes exactly one upstream attempt. +// Attempts and Backoff are validated (and defaulted, when left zero) by +// newRetryPolicy (retry.go) only when Enabled is true. +type RetryConfig struct { + Backoff string `json:"backoff,omitempty"` + Attempts int `json:"attempts,omitempty"` + Enabled bool `json:"enabled,omitempty"` +} + // ModelPricing overrides the built-in per-model price table. type ModelPricing struct { InputPerM float64 `json:"inputPerM"` diff --git a/provider_anthropic.go b/provider_anthropic.go index 50f6184..c3571a2 100644 --- a/provider_anthropic.go +++ b/provider_anthropic.go @@ -19,7 +19,13 @@ const ( // anthropicAdapter is the providerAdapter for provider type "anthropic". type anthropicAdapter struct { - client *http.Client + client *http.Client + // retry is nil for every adapter built by a bare newAnthropicAdapter + // call (every v0.1 test, and any caller that never wires one) — + // nil-safe, per retryPolicy.do's own contract, so leaving it unset is + // exactly v0.1's one-attempt behavior. buildAdapters (providers.go) + // sets it from cfg.Retry for production use. + retry *retryPolicy adapterName string baseURL string apiKey string @@ -106,7 +112,7 @@ func (a *anthropicAdapter) chatCompletion(ctx context.Context, w http.ResponseWr hdr.Set("Accept", "text/event-stream") } - resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+anthropicMessagesPath, hdr, body) + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+anthropicMessagesPath, hdr, body, a.retry) if err != nil { return usage{}, err } @@ -200,7 +206,7 @@ func (a *anthropicAdapter) embeddings(_ context.Context, _ http.ResponseWriter, // falls back to a provider's explicitly configured models on error rather // than this adapter guessing at a fallback itself. func (a *anthropicAdapter) listModels(ctx context.Context) ([]string, error) { - resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+anthropicModelsPath, a.requestHeaders(false), nil) + resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+anthropicModelsPath, a.requestHeaders(false), nil, a.retry) if err != nil { return nil, err } diff --git a/provider_gemini.go b/provider_gemini.go index 21b4a71..ad8a80d 100644 --- a/provider_gemini.go +++ b/provider_gemini.go @@ -13,7 +13,13 @@ import ( // geminiAdapter is the providerAdapter for provider type "gemini". type geminiAdapter struct { - client *http.Client + client *http.Client + // retry is nil for every adapter built by a bare newGeminiAdapter call + // (every v0.1 test, and any caller that never wires one) — nil-safe, + // per retryPolicy.do's own contract, so leaving it unset is exactly + // v0.1's one-attempt behavior. buildAdapters (providers.go) sets it + // from cfg.Retry for production use. + retry *retryPolicy adapterName string baseURL string apiKey string @@ -119,7 +125,7 @@ func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite endpoint = a.baseURL + geminiAPIPrefix + escapedModel + ":streamGenerateContent?alt=sse" } - resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, hdr, body) + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, hdr, body, a.retry) if err != nil { return usage{}, err } @@ -228,7 +234,7 @@ func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, r endpoint = a.baseURL + geminiAPIPrefix + escapedModel + ":batchEmbedContents" } - resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, a.requestHeaders(true), body) + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, a.requestHeaders(true), body, a.retry) if err != nil { return usage{}, err } @@ -273,7 +279,7 @@ type geminiModelsListPayload struct { // model ids, matching the shape openai-type and anthropic-type adapters // both return. func (a *geminiAdapter) listModels(ctx context.Context) ([]string, error) { - resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+"/v1beta/models", a.requestHeaders(false), nil) + resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+"/v1beta/models", a.requestHeaders(false), nil, a.retry) if err != nil { return nil, err } diff --git a/provider_openai.go b/provider_openai.go index 7e19923..7062baa 100644 --- a/provider_openai.go +++ b/provider_openai.go @@ -13,7 +13,13 @@ import ( // itself, and any OpenAI-compatible upstream (the operator's own gateway // included — a real deployment target for this adapter is keyless). type openaiAdapter struct { - client *http.Client + client *http.Client + // retry is nil for every adapter built by a bare newOpenAIAdapter call + // (every v0.1 test, and any caller that never wires one) — nil-safe, + // per retryPolicy.do's own contract, so leaving it unset is exactly + // v0.1's one-attempt behavior. buildAdapters (providers.go) sets it + // from cfg.Retry for production use. + retry *retryPolicy adapterName string baseURL string apiKey string @@ -131,7 +137,7 @@ func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite hdr.Set("Accept", "text/event-stream") } - resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/chat/completions", hdr, req) + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/chat/completions", hdr, req, a.retry) if err != nil { return usage{}, err } @@ -156,7 +162,7 @@ func (a *openaiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite // non-streaming path. func (a *openaiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { delete(req, gatewayAliasKey) // see chatCompletion's identical delete for why - resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/embeddings", a.requestHeaders(true), req) + resp, err := upstreamJSON(ctx, a.client, http.MethodPost, a.baseURL+"/v1/embeddings", a.requestHeaders(true), req, a.retry) if err != nil { return usage{}, err } @@ -258,7 +264,7 @@ type modelsPayload struct { // listModels implements providerAdapter. func (a *openaiAdapter) listModels(ctx context.Context) ([]string, error) { - resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+"/v1/models", a.requestHeaders(false), nil) + resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+"/v1/models", a.requestHeaders(false), nil, a.retry) if err != nil { return nil, err } diff --git a/providers.go b/providers.go index f23d918..b78bc37 100644 --- a/providers.go +++ b/providers.go @@ -159,34 +159,54 @@ func newAdapterHTTPClient() *http.Client { // a non-2xx status is not itself an error here, since some callers (e.g. // listModels) have nothing to write through and every current caller // checks resp.StatusCode itself. -func upstreamJSON(ctx context.Context, client *http.Client, method, url string, hdr http.Header, body any) (*http.Response, error) { - var r io.Reader +// +// policy (retry.go) wraps the whole request-and-receive-status exchange: +// a nil policy, or a disabled one, makes exactly one attempt (v0.1 +// behavior, byte-identical). Every caller of upstreamJSON — chatCompletion, +// embeddings, listModels, across all three adapter types — sits upstream +// of any response-body forwarding (forwardJSON/forwardStream), so a retry +// here can only ever replace an attempt whose body no caller has read or +// forwarded yet. That is what makes "zero bytes reached the client" hold +// by construction: forwardStream is invoked only after upstreamJSON has +// already returned its final, retried-or-not response, so a stream that +// dies while forwardStream is mid-flight never re-enters this function +// and is therefore never retried. +func upstreamJSON(ctx context.Context, client *http.Client, method, url string, hdr http.Header, body any, policy *retryPolicy) (*http.Response, error) { + var bodyBytes []byte if body != nil { b, err := json.Marshal(body) if err != nil { return nil, fmt.Errorf("%w: encode request body: %w", errUpstream, err) } - r = bytes.NewReader(b) + bodyBytes = b } - req, err := http.NewRequestWithContext(ctx, method, url, r) - if err != nil { - return nil, fmt.Errorf("%w: build request: %w", errUpstream, err) - } - for k, vs := range hdr { - for _, v := range vs { - req.Header.Add(k, v) + call := func() (*http.Response, error) { + var r io.Reader + if bodyBytes != nil { + r = bytes.NewReader(bodyBytes) } - } - if body != nil && req.Header.Get("Content-Type") == "" { - req.Header.Set("Content-Type", "application/json") + req, err := http.NewRequestWithContext(ctx, method, url, r) + if err != nil { + return nil, fmt.Errorf("%w: build request: %w", errUpstream, err) + } + for k, vs := range hdr { + for _, v := range vs { + req.Header.Add(k, v) + } + } + if bodyBytes != nil && req.Header.Get("Content-Type") == "" { + req.Header.Set("Content-Type", "application/json") + } + + resp, err := client.Do(req) //nolint:bodyclose // caller closes resp.Body; upstreamJSON hands the response, not its lifecycle, back + if err != nil { + return nil, fmt.Errorf("%w: %w", errUpstream, err) + } + return resp, nil } - resp, err := client.Do(req) //nolint:bodyclose // caller closes resp.Body; upstreamJSON hands the response, not its lifecycle, back - if err != nil { - return nil, fmt.Errorf("%w: %w", errUpstream, err) - } - return resp, nil + return policy.do(ctx, call) } // configNamePattern is the character set allowed in a provider, mcpServers, @@ -247,7 +267,21 @@ func validateConfigName(kind, name string) error { // that, per ruling (c); this function just propagates whatever error it // returns. A configured gemini provider follows the same constructor-error // ruling — newGeminiAdapter enforces it too. +// +// buildAdapters also derives cfg.Retry into a single *retryPolicy +// (newRetryPolicy, retry.go — an invalid Retry config fails construction +// the same way any other bad config value does) and attaches it to every +// adapter it constructs. Deriving it here, from cfg, rather than taking it +// as a parameter built upstream in newGateway, keeps buildAdapters' +// signature at its v0.1 shape — buildAdapters(cfg *Config) — so every +// existing direct call to it (production and test) keeps compiling and +// behaving identically when cfg.Retry is left at its zero value. func buildAdapters(cfg *Config) (map[string]providerAdapter, error) { + policy, err := newRetryPolicy(cfg.Retry) + if err != nil { + return nil, err + } + adapters := make(map[string]providerAdapter, len(cfg.Providers)) for name, pc := range cfg.Providers { if pc == nil { @@ -268,18 +302,22 @@ func buildAdapters(cfg *Config) (map[string]providerAdapter, error) { switch pc.Type { case providerTypeOpenAI: - adapters[name] = newOpenAIAdapter(name, base, apiKey) + a := newOpenAIAdapter(name, base, apiKey) + a.retry = policy + adapters[name] = a case providerTypeAnthropic: a, err := newAnthropicAdapter(name, base, apiKey) if err != nil { return nil, err } + a.retry = policy adapters[name] = a case providerTypeGemini: g, err := newGeminiAdapter(name, base, apiKey) if err != nil { return nil, err } + g.retry = policy adapters[name] = g default: return nil, fmt.Errorf("llmgateway: provider %q: unknown type %q", name, pc.Type) diff --git a/retry.go b/retry.go new file mode 100644 index 0000000..ea4b1cf --- /dev/null +++ b/retry.go @@ -0,0 +1,165 @@ +package traefikllmgateway + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "strconv" + "time" +) + +// defaultRetryAttempts and defaultRetryBackoff are the values +// newRetryPolicy applies when a Retry config block is Enabled but leaves +// Attempts or Backoff at its zero value — spec §1's stated defaults. +const ( + defaultRetryAttempts = 1 + defaultRetryBackoff = "250ms" +) + +// maxRetryAttempts is the highest RetryConfig.Attempts newRetryPolicy +// accepts — spec §1's "attempts: int (default 1, max 3)". +const maxRetryAttempts = 3 + +// maxRetryWait caps every wait retryPolicy.do performs between attempts, +// whether computed from exponential backoff or read from an upstream's +// Retry-After header — spec §1's "capped at 2s per wait". +const maxRetryWait = 2 * time.Second + +// retryPolicy implements spec §1's same-provider retry: at most +// p.attempts total tries per upstream call, exponential backoff between +// them (capped at maxRetryWait), honoring a qualifying 429 Retry-After +// header outright. A nil *retryPolicy, or one with enabled=false, makes +// do try call exactly once — the same shape every v0.1 adapter +// constructor leaves its adapter's retry field in, since none of them set +// it. sleepFn is a field (not a direct time.Sleep call) so tests can +// inject a non-blocking spy instead of paying real wall-clock time. +type retryPolicy struct { + sleepFn func(time.Duration) + backoff time.Duration + attempts int + enabled bool +} + +// newRetryPolicy validates rc and returns the retryPolicy buildAdapters +// attaches to every adapter it constructs. Disabled config (the v0.1 +// default) skips validation and defaulting entirely — an unset +// Attempts/Backoff on a disabled block is never an error, since neither +// value is ever read. Enabled config defaults an unset Attempts to +// defaultRetryAttempts and an unset Backoff to defaultRetryBackoff before +// validating: Attempts must fall within 1..maxRetryAttempts, and Backoff +// must parse via time.ParseDuration. +func newRetryPolicy(rc RetryConfig) (*retryPolicy, error) { + if !rc.Enabled { + return &retryPolicy{enabled: false, sleepFn: time.Sleep}, nil + } + + attempts := rc.Attempts + if attempts == 0 { + attempts = defaultRetryAttempts + } + if attempts < 1 || attempts > maxRetryAttempts { + return nil, fmt.Errorf("llmgateway: retry.attempts must be between 1 and %d, got %d", maxRetryAttempts, attempts) + } + + backoffStr := rc.Backoff + if backoffStr == "" { + backoffStr = defaultRetryBackoff + } + backoff, err := time.ParseDuration(backoffStr) + if err != nil { + return nil, fmt.Errorf("llmgateway: retry.backoff: %w", err) + } + + return &retryPolicy{enabled: true, attempts: attempts, backoff: backoff, sleepFn: time.Sleep}, nil +} + +// isTransient reports whether resp/err — the result of one upstream +// attempt, mutually exclusive per http.Client's own contract (a non-nil +// err always means resp is nil) — is a failure retryPolicy.do should +// retry: any error (a dial/EOF/reset failure reaching upstreamJSON's +// client.Do) other than a context cancellation or deadline, or a +// response carrying HTTP 429 or a 5xx status. Every other outcome — +// success, or a non-429 4xx — is not transient, per spec §1's "Never on +// 4xx≠429, translate errors, or context cancellation". +func isTransient(resp *http.Response, err error) bool { + if err != nil { + return !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) + } + if resp == nil { + return false + } + return resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= http.StatusInternalServerError +} + +// waitBefore returns how long do should wait before the attempt after +// attempt (1-based) — the one that just produced resp (nil on a network +// error). A 429 carrying a Retry-After header parseable as whole seconds, +// and no greater than maxRetryWait, wins outright: the upstream said +// exactly how long to back off. Every other transient failure falls back +// to exponential backoff (p.backoff * 2^(attempt-1)), capped at +// maxRetryWait either way. +func (p *retryPolicy) waitBefore(attempt int, resp *http.Response) time.Duration { + if resp != nil && resp.StatusCode == http.StatusTooManyRequests { + if ra := resp.Header.Get("Retry-After"); ra != "" { + if secs, err := strconv.Atoi(ra); err == nil && secs >= 0 { + if wait := time.Duration(secs) * time.Second; wait <= maxRetryWait { + return wait + } + } + } + } + + wait := p.backoff * time.Duration(int64(1)< maxRetryWait { + wait = maxRetryWait + } + return wait +} + +// drainAndClose discards a retried-away response's remaining body +// (capped at maxResponseBytes, defensively) and closes it, so the +// underlying connection is eligible for reuse on the next attempt instead +// of being forced closed by an un-drained body. +func drainAndClose(resp *http.Response) { + if resp == nil || resp.Body == nil { + return + } + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxResponseBytes)) + _ = resp.Body.Close() +} + +// do executes call, retrying while isTransient(resp, err) reports a +// transient failure, up to p.attempts total tries. A nil p, or one with +// enabled=false, tries call exactly once — the disabled/unwired shape. +// Between attempts it drains and closes the failed attempt's response +// body, then waits per waitBefore via p.sleepFn. do never starts a wait +// once ctx is already done: at that point it returns the last attempt's +// (resp, err) as-is, the same value it would return had attempts run out +// naturally. do returns the final attempt's result either way — success, +// a still-transient failure once attempts are exhausted, or a +// non-transient failure on any attempt. +func (p *retryPolicy) do(ctx context.Context, call func() (*http.Response, error)) (*http.Response, error) { + attempts := 1 + if p != nil && p.enabled { + attempts = p.attempts + } + + var resp *http.Response + var err error + for attempt := 1; attempt <= attempts; attempt++ { + resp, err = call() + if !isTransient(resp, err) || attempt == attempts { + return resp, err + } + if ctx.Err() != nil { + return resp, err + } + + wait := p.waitBefore(attempt, resp) + drainAndClose(resp) + p.sleepFn(wait) + } + return resp, err +} diff --git a/retry_test.go b/retry_test.go new file mode 100644 index 0000000..11e5987 --- /dev/null +++ b/retry_test.go @@ -0,0 +1,583 @@ +package traefikllmgateway + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +// fakeResp builds a canned *http.Response for retryPolicy.do tests that +// exercise the transience decision without a real network round trip. +// body is always readable and closable, matching a real *http.Response. +func fakeResp(status int, hdr http.Header) *http.Response { + if hdr == nil { + hdr = http.Header{} + } + return &http.Response{StatusCode: status, Header: hdr, Body: io.NopCloser(strings.NewReader("body"))} +} + +// TestIsTransient covers the transience classification retryPolicy.do +// relies on: net errors (wrapping errUpstream) and HTTP 429/5xx are +// transient; a successful or non-429 4xx status, and a context +// cancellation/deadline error, are not. +func TestIsTransient(t *testing.T) { + tests := []struct { + resp *http.Response + err error + name string + want bool + }{ + {name: "503", resp: fakeResp(http.StatusServiceUnavailable, nil), want: true}, + {name: "500", resp: fakeResp(http.StatusInternalServerError, nil), want: true}, + {name: "429", resp: fakeResp(http.StatusTooManyRequests, nil), want: true}, + {name: "200", resp: fakeResp(http.StatusOK, nil), want: false}, + {name: "400", resp: fakeResp(http.StatusBadRequest, nil), want: false}, + {name: "404", resp: fakeResp(http.StatusNotFound, nil), want: false}, + {name: "network error", err: fmt.Errorf("%w: dial tcp: connection reset", errUpstream), want: true}, + {name: "context canceled", err: fmt.Errorf("%w: %w", errUpstream, context.Canceled), want: false}, + {name: "context deadline exceeded", err: fmt.Errorf("%w: %w", errUpstream, context.DeadlineExceeded), want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isTransient(tt.resp, tt.err); got != tt.want { + t.Errorf("isTransient(%v, %v) = %v, want %v", tt.resp, tt.err, got, tt.want) + } + }) + } +} + +// TestNewRetryPolicy_Disabled asserts the zero-value RetryConfig (the +// v0.1-preserving default) produces a policy with enabled=false and no +// validation error, even though Attempts/Backoff are left at their zero +// values. +func TestNewRetryPolicy_Disabled(t *testing.T) { + p, err := newRetryPolicy(RetryConfig{}) + if err != nil { + t.Fatalf("newRetryPolicy: %v", err) + } + if p == nil { + t.Fatal("newRetryPolicy returned a nil policy for a valid (disabled) config") + } + if p.enabled { + t.Error("p.enabled = true, want false for RetryConfig{}") + } +} + +// TestNewRetryPolicy_Defaults asserts Enabled=true with Attempts/Backoff +// left unset applies the spec defaults: 1 attempt, 250ms backoff. +func TestNewRetryPolicy_Defaults(t *testing.T) { + p, err := newRetryPolicy(RetryConfig{Enabled: true}) + if err != nil { + t.Fatalf("newRetryPolicy: %v", err) + } + if !p.enabled { + t.Error("p.enabled = false, want true") + } + if p.attempts != 1 { + t.Errorf("p.attempts = %d, want 1", p.attempts) + } + if p.backoff != 250*time.Millisecond { + t.Errorf("p.backoff = %v, want 250ms", p.backoff) + } +} + +// TestNewRetryPolicy_ValidatesAttempts covers the 1..3 attempts range. +func TestNewRetryPolicy_ValidatesAttempts(t *testing.T) { + tests := []struct { + name string + wantErr bool + attempt int + }{ + {name: "min", attempt: 1, wantErr: false}, + {name: "max", attempt: 3, wantErr: false}, + {name: "mid", attempt: 2, wantErr: false}, + {name: "zero uses default", attempt: 0, wantErr: false}, + {name: "too high", attempt: 4, wantErr: true}, + {name: "negative", attempt: -1, wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := newRetryPolicy(RetryConfig{Enabled: true, Attempts: tt.attempt}) + if (err != nil) != tt.wantErr { + t.Errorf("newRetryPolicy(Attempts:%d): err = %v, wantErr %v", tt.attempt, err, tt.wantErr) + } + }) + } +} + +// TestNewRetryPolicy_ValidatesBackoff covers backoff parsing: an +// unparseable duration string is a config error; a valid one is used +// verbatim (not defaulted). +func TestNewRetryPolicy_ValidatesBackoff(t *testing.T) { + if _, err := newRetryPolicy(RetryConfig{Enabled: true, Backoff: "not-a-duration"}); err == nil { + t.Fatal("newRetryPolicy: want error for unparseable backoff, got nil") + } + p, err := newRetryPolicy(RetryConfig{Enabled: true, Backoff: "10ms"}) + if err != nil { + t.Fatalf("newRetryPolicy: %v", err) + } + if p.backoff != 10*time.Millisecond { + t.Errorf("p.backoff = %v, want 10ms", p.backoff) + } +} + +// spySleep returns a sleepFn that records every requested duration +// instead of actually sleeping, so retry-wait tests run instantly and can +// assert on the exact durations retryPolicy.do computed. +func spySleep(waits *[]time.Duration) func(time.Duration) { + return func(d time.Duration) { + *waits = append(*waits, d) + } +} + +// callSequence returns a call func for retryPolicy.do that returns +// results[i] on its i-th invocation (0-indexed), and counts invocations +// into calls. +func callSequence(calls *int32, results ...func() (*http.Response, error)) func() (*http.Response, error) { + return func() (*http.Response, error) { + n := int(atomic.AddInt32(calls, 1)) - 1 + if n >= len(results) { + n = len(results) - 1 + } + return results[n]() + } +} + +// TestRetryPolicy_Do_RetriesOnTransientThenSucceeds covers the brief's +// three canonical transient-then-success cases: 503, a raw network error +// (connection reset), and 429 with a qualifying Retry-After header. Each +// retries exactly once, waits the expected duration, and returns the +// second attempt's success. +func TestRetryPolicy_Do_RetriesOnTransientThenSucceeds(t *testing.T) { + tests := []struct { + name string + first func() (*http.Response, error) + wantWaits []time.Duration + }{ + { + name: "503 then success", + first: func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil }, + wantWaits: []time.Duration{100 * time.Millisecond}, // backoff * 2^0 + }, + { + name: "connection reset then success", + first: func() (*http.Response, error) { + return nil, fmt.Errorf("%w: dial tcp: connection reset by peer", errUpstream) + }, + wantWaits: []time.Duration{100 * time.Millisecond}, + }, + { + name: "429 with Retry-After then success", + first: func() (*http.Response, error) { + h := http.Header{} + h.Set("Retry-After", "1") + return fakeResp(http.StatusTooManyRequests, h), nil + }, + wantWaits: []time.Duration{1 * time.Second}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var calls int32 + var waits []time.Duration + p := &retryPolicy{enabled: true, attempts: 2, backoff: 100 * time.Millisecond, sleepFn: spySleep(&waits)} + call := callSequence(&calls, tt.first, func() (*http.Response, error) { + return fakeResp(http.StatusOK, nil), nil + }) + + resp, err := p.do(context.Background(), call) + if err != nil { + t.Fatalf("do: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Errorf("resp.StatusCode = %d, want 200", resp.StatusCode) + } + if calls != 2 { + t.Errorf("calls = %d, want 2 (one retry)", calls) + } + if len(waits) != len(tt.wantWaits) || (len(waits) > 0 && waits[0] != tt.wantWaits[0]) { + t.Errorf("waits = %v, want %v", waits, tt.wantWaits) + } + }) + } +} + +// TestRetryPolicy_Do_NonTransientStatus_NoRetry asserts a 400 response +// (not 429, not 5xx) is returned as-is on the first attempt, with no +// retry and no wait. +func TestRetryPolicy_Do_NonTransientStatus_NoRetry(t *testing.T) { + var calls int32 + var waits []time.Duration + p := &retryPolicy{enabled: true, attempts: 3, backoff: time.Millisecond, sleepFn: spySleep(&waits)} + call := callSequence(&calls, func() (*http.Response, error) { return fakeResp(http.StatusBadRequest, nil), nil }) + + resp, err := p.do(context.Background(), call) + if err != nil { + t.Fatalf("do: %v", err) + } + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("resp.StatusCode = %d, want 400", resp.StatusCode) + } + if calls != 1 { + t.Errorf("calls = %d, want 1 (no retry on non-transient status)", calls) + } + if len(waits) != 0 { + t.Errorf("waits = %v, want none", waits) + } +} + +// TestRetryPolicy_Do_DisabledPolicy_NoRetry asserts an enabled=false +// policy (the v0.1-preserving default shape) tries exactly once even on +// a transient failure. +func TestRetryPolicy_Do_DisabledPolicy_NoRetry(t *testing.T) { + var calls int32 + p := &retryPolicy{enabled: false, sleepFn: func(time.Duration) { t.Error("sleepFn called on a disabled policy") }} + call := callSequence(&calls, func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil }) + + resp, err := p.do(context.Background(), call) + if err != nil { + t.Fatalf("do: %v", err) + } + if resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("resp.StatusCode = %d, want 503", resp.StatusCode) + } + if calls != 1 { + t.Errorf("calls = %d, want 1 (retry disabled)", calls) + } +} + +// TestRetryPolicy_Do_NilPolicy_NoRetry asserts a nil *retryPolicy — the +// shape every adapter constructed by v0.1's plain constructors has, since +// they never set the retry field — behaves identically to an explicitly +// disabled one: exactly one attempt, no panic. +func TestRetryPolicy_Do_NilPolicy_NoRetry(t *testing.T) { + var calls int32 + var p *retryPolicy + call := callSequence(&calls, func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil }) + + resp, err := p.do(context.Background(), call) + if err != nil { + t.Fatalf("do: %v", err) + } + if resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("resp.StatusCode = %d, want 503", resp.StatusCode) + } + if calls != 1 { + t.Errorf("calls = %d, want 1 (nil policy = disabled)", calls) + } +} + +// TestRetryPolicy_Do_AttemptsCap asserts a policy that never sees success +// stops after exactly p.attempts tries, waiting p.attempts-1 times with +// exponential backoff, capped at maxRetryWait. +func TestRetryPolicy_Do_AttemptsCap(t *testing.T) { + var calls int32 + var waits []time.Duration + p := &retryPolicy{enabled: true, attempts: 3, backoff: time.Second, sleepFn: spySleep(&waits)} + call := callSequence(&calls, func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil }) + + resp, err := p.do(context.Background(), call) + if err != nil { + t.Fatalf("do: %v", err) + } + if calls != 3 { + t.Errorf("calls = %d, want 3 (attempts cap)", calls) + } + if resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("resp.StatusCode = %d, want 503 (last attempt's result)", resp.StatusCode) + } + wantWaits := []time.Duration{maxRetryWait, maxRetryWait} // 1s*2^0=1s→wait capped? no: 1s, 2s both <= cap(2s) + _ = wantWaits + if len(waits) != 2 { + t.Fatalf("waits = %v, want 2 entries (between 3 attempts)", waits) + } + if waits[0] != time.Second { + t.Errorf("waits[0] = %v, want 1s (backoff * 2^0)", waits[0]) + } + if waits[1] != 2*time.Second { + t.Errorf("waits[1] = %v, want 2s (backoff * 2^1, at the cap)", waits[1]) + } +} + +// TestRetryPolicy_Do_BackoffCappedAtMaxWait asserts a backoff whose +// exponential growth would exceed maxRetryWait is clamped to it. +func TestRetryPolicy_Do_BackoffCappedAtMaxWait(t *testing.T) { + var calls int32 + var waits []time.Duration + p := &retryPolicy{enabled: true, attempts: 3, backoff: 2 * time.Second, sleepFn: spySleep(&waits)} + call := callSequence(&calls, func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil }) + + if _, err := p.do(context.Background(), call); err != nil { + t.Fatalf("do: %v", err) + } + for i, w := range waits { + if w > maxRetryWait { + t.Errorf("waits[%d] = %v, exceeds cap %v", i, w, maxRetryWait) + } + } +} + +// TestRetryPolicy_Do_CtxCanceledAbortsRetry asserts a context canceled +// before a retry wait would begin stops the loop immediately: no further +// call, no sleep, and the already-failed result is returned as-is. +func TestRetryPolicy_Do_CtxCanceledAbortsRetry(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() // canceled before do() ever runs + + var calls int32 + p := &retryPolicy{enabled: true, attempts: 3, backoff: time.Millisecond, sleepFn: func(time.Duration) { + t.Error("sleepFn called after ctx was already canceled") + }} + call := callSequence(&calls, func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil }) + + resp, err := p.do(ctx, call) + if err != nil { + t.Fatalf("do: %v", err) + } + if calls != 1 { + t.Errorf("calls = %d, want 1 (ctx already done, no retry)", calls) + } + if resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("resp.StatusCode = %d, want 503", resp.StatusCode) + } +} + +// TestRetryPolicy_Do_ClosesFailedAttemptBody asserts a retried-away +// response's body is drained and closed before the next attempt, so the +// underlying connection is eligible for reuse instead of being forced +// closed by an un-drained body. +func TestRetryPolicy_Do_ClosesFailedAttemptBody(t *testing.T) { + var calls int32 + var waits []time.Duration + closed := false + firstBody := &trackingCloser{Reader: strings.NewReader("stale body"), onClose: func() { closed = true }} + p := &retryPolicy{enabled: true, attempts: 2, backoff: time.Millisecond, sleepFn: spySleep(&waits)} + call := callSequence(&calls, + func() (*http.Response, error) { + return &http.Response{StatusCode: http.StatusServiceUnavailable, Header: http.Header{}, Body: firstBody}, nil + }, + func() (*http.Response, error) { return fakeResp(http.StatusOK, nil), nil }, + ) + + if _, err := p.do(context.Background(), call); err != nil { + t.Fatalf("do: %v", err) + } + if !closed { + t.Error("failed attempt's response body was never closed") + } +} + +// trackingCloser wraps an io.Reader with a Close method that records the +// call, for asserting a retried-away response body was actually closed. +type trackingCloser struct { + io.Reader + onClose func() +} + +func (c *trackingCloser) Close() error { + c.onClose() + return nil +} + +// TestBuildAdapters_WiresRetryPolicyFromConfig asserts buildAdapters +// derives a *retryPolicy from cfg.Retry and attaches it to every +// constructed adapter, and that an invalid Retry config fails +// construction the same way any other bad ProviderConfig does. +func TestBuildAdapters_WiresRetryPolicyFromConfig(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"p1": {Type: "openai", APIKey: "k"}} + cfg.Retry = RetryConfig{Enabled: true, Attempts: 2, Backoff: "50ms"} + + adapters, err := buildAdapters(cfg) + if err != nil { + t.Fatalf("buildAdapters: %v", err) + } + oa, ok := adapters["p1"].(*openaiAdapter) + if !ok { + t.Fatalf("adapters[%q] = %T, want *openaiAdapter", "p1", adapters["p1"]) + } + if oa.retry == nil { + t.Fatal("oa.retry is nil, want a wired retryPolicy") + } + if !oa.retry.enabled || oa.retry.attempts != 2 || oa.retry.backoff != 50*time.Millisecond { + t.Errorf("oa.retry = %+v, want enabled=true attempts=2 backoff=50ms", oa.retry) + } +} + +// TestBuildAdapters_InvalidRetryConfig_ReturnsError asserts an +// out-of-range Retry.Attempts fails buildAdapters the same way any other +// invalid config value does. +func TestBuildAdapters_InvalidRetryConfig_ReturnsError(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"p1": {Type: "openai", APIKey: "k"}} + cfg.Retry = RetryConfig{Enabled: true, Attempts: 9} + + if _, err := buildAdapters(cfg); err == nil { + t.Fatal("buildAdapters: want error for Attempts=9, got nil") + } +} + +// TestNewGateway_RetryDisabledByDefault_Unwired asserts a Config that +// never sets Retry at all — every v0.1 config — produces adapters whose +// retry field is a disabled (not necessarily nil, but enabled=false) +// policy, matching v0.1's always-one-attempt behavior. +func TestNewGateway_RetryDisabledByDefault_Unwired(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw := h.(*Gateway) + oa, ok := gw.adapters["openai"].(*openaiAdapter) + if !ok { + t.Fatalf("adapters[openai] = %T, want *openaiAdapter", gw.adapters["openai"]) + } + if oa.retry != nil && oa.retry.enabled { + t.Errorf("oa.retry = %+v, want disabled (Retry never configured)", oa.retry) + } +} + +// ---- adapter-level wiring tests (stream-initiation path) ---- + +// TestOpenAIAdapter_ChatCompletion_NonStream_RetriedOnTransientFailure +// proves upstreamJSON's retry covers a non-streaming chatCompletion call +// end to end: the client sees only the second (successful) response, and +// the upstream received exactly two requests. +func TestOpenAIAdapter_ChatCompletion_NonStream_RetriedOnTransientFailure(t *testing.T) { + const okBody = `{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":2}}` + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&calls, 1) == 1 { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(okBody)) + })) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + a.retry = &retryPolicy{enabled: true, attempts: 2, backoff: time.Millisecond, sleepFn: func(time.Duration) {}} + rec := httptest.NewRecorder() + + u, err := a.chatCompletion(context.Background(), rec, map[string]any{"model": "gpt-5", "messages": []any{}}) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + if calls != 2 { + t.Errorf("upstream calls = %d, want 2", calls) + } + if rec.Code != http.StatusOK { + t.Errorf("rec.Code = %d, want 200", rec.Code) + } + if rec.Body.String() != okBody { + t.Errorf("rec.Body = %q, want %q", rec.Body.String(), okBody) + } + if u.prompt != 1 || u.completion != 2 { + t.Errorf("usage = %+v, want {1 2}", u) + } +} + +// TestOpenAIAdapter_ChatCompletion_Stream_RetriedBeforeFirstChunk proves +// a streaming request whose first upstream exchange fails (before any +// SSE chunk exists) is retried, and the client receives only the +// successful stream — no trace of the failed first attempt. +func TestOpenAIAdapter_ChatCompletion_Stream_RetriedBeforeFirstChunk(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&calls, 1) == 1 { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl := w.(http.Flusher) + for _, f := range streamFrames { + _, _ = fmt.Fprintf(w, "data: %s\n\n", f) + fl.Flush() + } + _, _ = fmt.Fprint(w, "data: [DONE]\n\n") + fl.Flush() + })) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + a.retry = &retryPolicy{enabled: true, attempts: 2, backoff: time.Millisecond, sleepFn: func(time.Duration) {}} + rec := httptest.NewRecorder() + + u, err := a.chatCompletion(context.Background(), rec, map[string]any{"model": "gpt-5", "messages": []any{}, "stream": true}) + if err != nil { + t.Fatalf("chatCompletion: %v", err) + } + if calls != 2 { + t.Errorf("upstream calls = %d, want 2", calls) + } + if u.prompt != 7 || u.completion != 9 { + t.Errorf("usage = %+v, want {7 9}", u) + } + body := rec.Body.String() + if !strings.Contains(body, "Hel") || !strings.Contains(body, "[DONE]") { + t.Errorf("rec.Body = %q, want the full successful stream", body) + } + if strings.Count(body, "[DONE]") != 1 { + t.Errorf("rec.Body contains %d [DONE] markers, want exactly 1 (no leaked failed-attempt output)", strings.Count(body, "[DONE]")) + } +} + +// TestOpenAIAdapter_ChatCompletion_Stream_MidDeathNotRetried proves a +// stream that dies after already forwarding a chunk is NOT retried: the +// client keeps the partial data it already received (proving those bytes +// were forwarded), chatCompletion returns an error, and the upstream saw +// exactly one request. +func TestOpenAIAdapter_ChatCompletion_Stream_MidDeathNotRetried(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&calls, 1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl := w.(http.Flusher) + _, _ = fmt.Fprintf(w, "data: %s\n\n", streamFrames[0]) + fl.Flush() + + hj, ok := w.(http.Hijacker) + if !ok { + t.Fatal("ResponseWriter does not support Hijacker") + } + conn, _, err := hj.Hijack() + if err != nil { + t.Fatalf("Hijack: %v", err) + } + _ = conn.Close() // abrupt drop mid-chunked-body: client sees a read error, not a clean EOF + })) + defer srv.Close() + + a := newOpenAIAdapter("p1", srv.URL, "sk-test") + a.retry = &retryPolicy{enabled: true, attempts: 3, backoff: time.Millisecond, sleepFn: func(time.Duration) { + t.Error("sleepFn called: a mid-stream death must not trigger a retry") + }} + rec := httptest.NewRecorder() + + _, err := a.chatCompletion(context.Background(), rec, map[string]any{"model": "gpt-5", "messages": []any{}, "stream": true}) + if err == nil { + t.Fatal("chatCompletion: want an error for a connection dropped mid-stream, got nil") + } + if !errors.Is(err, errUpstream) { + t.Errorf("chatCompletion err = %v, want it to wrap errUpstream", err) + } + if calls != 1 { + t.Errorf("upstream calls = %d, want 1 (mid-stream death is not retried)", calls) + } + if !strings.Contains(rec.Body.String(), "Hel") { + t.Errorf("rec.Body = %q, want it to contain the chunk already forwarded before the drop", rec.Body.String()) + } +}