feat(providers): success-rate accounting

Track per-provider and per-(provider,upstream-model) upstream-attempt
outcomes using the existing fixed-window counter machinery (limits.go),
following mcp_a2a.go's per-target counter pattern: kind-embedded keys
(kindProvider="prov", kindProviderModel="provmodel") keep a provider
scope collision-safe against a user/group/MCP-target/agent of the same
name.

limiter.recordProviderAttempt(provider, model, resp, err) increments an
"attempts" counter on every call, and a "failures" counter only for a
provider-fault outcome — reusing retry.go's own isTransient(resp, err)
classification exactly (transport/connection error, HTTP 429, or a 5xx),
not forking it. Both counters are written at minute and day granularity
only; provider health is a now-and-today question, not a billing one, so
there is no month window. limiter.providerUsage batches every scope's
read into one storeGetMulti round trip, mirroring targetUsage's own
single-batch discipline.

Wired at every chokepoint an upstream HTTP attempt resolves:
- providers.go's upstreamJSON/upstreamRawBytes path (chat, embeddings,
  images, audio) via retry.go's retryPolicy.do, which now looks up an
  optional attemptRecorder from ctx and invokes it once per retry
  attempt (attempt-accounting: a request retried twice before
  succeeding reports three attempts, not one).
- routes_unified.go's runUnified (chat/embeddings, streaming and
  non-streaming alike — both funnel through the same upstreamJSON call).
- routes_media.go's three media handlers (images.generations,
  audio.speech, audio.transcriptions).
- routes_passthrough.go's proxyUpstream (native provider passthrough) —
  a single un-retried attempt, provider-level only: the upstream model
  lives in the response body, read only after the attempt resolves.

The attemptRecorder is threaded via context (providers.go's
withAttemptRecorder/attemptRecorderFromContext), not a new function
parameter on every call site — proxyUpstream and retryPolicy.do are
both shared with paths that must NOT record provider accounting
(registry.go's discovery listModels calls, and the MCP/A2A target
proxy, mcp_a2a.go's handleTargetProxy), so scoping via context keeps
their signatures unchanged and those callers correctly opt out by
simply never wrapping their context.
This commit is contained in:
2026-08-21 23:22:16 +01:00
parent ad5adbcc1d
commit aafd2a0433
12 changed files with 653 additions and 4 deletions
+167
View File
@@ -3,6 +3,7 @@ package traefikllmgateway
import (
"fmt"
"math"
"net/http"
"sync"
"time"
)
@@ -1140,6 +1141,172 @@ func (l *limiter) targetUsage(scopes []limitScope) []targetCounters {
return out
}
// kindProvider and kindProviderModel are the limitScope.kind (and
// admin-API scope-kind) strings a provider's and a (provider, upstream
// model) pair's outcome counters use (Feature A, v0.22 — provider/model
// success-rate accounting) — collision-safety via kind embedding, mirroring
// scopeKindAgent's own rationale (mcp_a2a.go): windowKey embeds kind as the
// counter key's own leading segment, so a provider named identically to a
// user, group, or MCP/A2A target can never share a counter with it, and a
// provider-level scope can never collide with one of its own
// (provider, model) scopes either, even when a model happens to be named
// after its provider.
const (
kindProvider = "prov"
kindProviderModel = "provmodel"
)
// metricProvAttempt and metricProvFail are the counter metric names
// recordProviderAttempt writes (Feature A, v0.22): every upstream attempt
// increments metricProvAttempt; only a provider-fault outcome
// (isTransient's own classification, retry.go — reused, not forked, per
// spec ruling) additionally increments metricProvFail.
const (
metricProvAttempt = "attempt"
metricProvFail = "fail"
)
// providerModelScopeID joins provider and model into kindProviderModel's
// id convention: "provider/model" — the same canonical id shape
// pricing.go's unifiedCostMicros and the admin API already use elsewhere
// for a (provider, model) pair, so a reader who already knows that
// convention needs no new one here.
func providerModelScopeID(provider, model string) string {
return provider + "/" + model
}
// recordProviderAttempt accounts ONE upstream HTTP attempt against
// provider (Feature A, v0.22), and against the (provider, model) pair too
// when model is non-empty — a caller that cannot cheaply know the upstream
// model before the attempt resolves (native passthrough, routes_
// passthrough.go's handlePassthrough: the model lives in the response
// body, read only after the attempt already happened) passes "" and gets
// provider-level accounting only.
//
// Every attempt increments an "attempts" counter; only a provider-fault
// outcome — exactly isTransient(resp, err)'s own classification
// (retry.go: a transport/connection error, HTTP 429, or a 5xx status) —
// additionally increments a "failures" counter. Deliberately reused rather
// than forked: a 4xx-but-not-429 response means the provider answered
// correctly to a request it did not like, which is not a provider-health
// signal, exactly the line isTransient already draws for "worth retrying".
//
// attempt-accounting: this is invoked once per upstream attempt, not once
// per logical request — attemptRecorderFromContext (providers.go) is
// looked up fresh inside retryPolicy.do's own retry loop and proxyUpstream's
// single client.Do call, so a chat completion retried twice before
// succeeding counts as two attempts here (one failure, one success), the
// same "count what was actually attempted" discipline
// countTargetRequests/-Request already apply to target scopes (mcp_a2a.go
// callers), consistent with those counters' own attempt/request semantics.
//
// Both counters are written at minute AND day granularity only — no month
// window: provider health is a now-and-today question (the Providers tab's
// success-rate badge, webui), not a billing one, so there is nothing here
// for a month-long retention window to serve.
func (l *limiter) recordProviderAttempt(provider, model string, resp *http.Response, err error) {
if provider == "" {
return
}
now := l.now()
fail := isTransient(resp, err)
entries := make([]counterIncr, 0, 8)
entries = append(entries,
newCounterIncr(kindProvider, provider, metricProvAttempt, windowMin, now, 1, minWindowTTL),
newCounterIncr(kindProvider, provider, metricProvAttempt, windowDay, now, 1, dayWindowTTL),
)
if fail {
entries = append(entries,
newCounterIncr(kindProvider, provider, metricProvFail, windowMin, now, 1, minWindowTTL),
newCounterIncr(kindProvider, provider, metricProvFail, windowDay, now, 1, dayWindowTTL),
)
}
if model != "" {
id := providerModelScopeID(provider, model)
entries = append(entries,
newCounterIncr(kindProviderModel, id, metricProvAttempt, windowMin, now, 1, minWindowTTL),
newCounterIncr(kindProviderModel, id, metricProvAttempt, windowDay, now, 1, dayWindowTTL),
)
if fail {
entries = append(entries,
newCounterIncr(kindProviderModel, id, metricProvFail, windowMin, now, 1, minWindowTTL),
newCounterIncr(kindProviderModel, id, metricProvFail, windowDay, now, 1, dayWindowTTL),
)
}
}
l.storeIncrMulti(entries) // no error return by contract (doc comment above, mirroring account()); ok is intentionally discarded
}
// providerCounterKeysPerScope is the number of windowKey strings
// providerCounterKeys builds per scope, and the stride providerUsage's
// flat storeGetMulti result is sliced back into per-scope chunks by —
// mirrors targetCounterKeysPerScope's own convention for the identical
// reason.
const providerCounterKeysPerScope = 4
// providerCounterKeys returns the providerCounterKeysPerScope windowKey
// strings providerUsage reads for one (kind, id) scope at time now —
// attempt/fail at minute, then attempt/fail at day, matching
// providerCounters' own field order exactly so providerUsage can map
// storeGetMulti's result slice back by plain index.
func providerCounterKeys(kind, id string, now time.Time) []string {
return []string{
windowKey(kind, id, metricProvAttempt, windowMin, now),
windowKey(kind, id, metricProvFail, windowMin, now),
windowKey(kind, id, metricProvAttempt, windowDay, now),
windowKey(kind, id, metricProvFail, windowDay, now),
}
}
// providerCounters is one provider's or (provider, model) pair's current-
// window attempt/failure counters (Feature A, v0.22) — mirrors
// targetCounters' own shape and read-only, admin-view-only role.
type providerCounters struct {
attemptsMinute int64
failuresMinute int64
attemptsDay int64
failuresDay int64
}
// providerUsage reads every scope's current attempt/failure counters in
// ONE storeGetMulti round trip — mirrors targetUsage's own single-batch
// discipline (limits.go) for the identical reason: GET
// /admin/api/overview must not pay one round trip per configured provider
// and model. scopes is built by the caller (buildAdminOverview, admin.go)
// as kindProvider entries for the provider-level rows followed by
// kindProviderModel entries for the per-model breakdown, all in one
// slice, so both share this single round trip. Order is preserved:
// providerUsage(scopes)[i] corresponds to scopes[i]. A failed read
// (store down + fail-closed) returns every zero-value counters, exactly
// like targetUsage — see its own doc comment for why that is acceptable
// here: nothing downstream treats a provider's zero counters as
// "confirmed no traffic" the way scopeUsage.storeDown guards against for
// a limited scope.
func (l *limiter) providerUsage(scopes []limitScope) []providerCounters {
out := make([]providerCounters, len(scopes))
if len(scopes) == 0 {
return out
}
now := l.now()
allKeys := make([]string, 0, len(scopes)*providerCounterKeysPerScope)
for _, sc := range scopes {
allKeys = append(allKeys, providerCounterKeys(sc.kind, sc.id, now)...)
}
vals, ok := l.storeGetMulti(allKeys)
if !ok || len(vals) != len(allKeys) {
return out // zero-value counters; see doc comment above
}
for i := range scopes {
v := vals[i*providerCounterKeysPerScope : i*providerCounterKeysPerScope+providerCounterKeysPerScope]
out[i] = providerCounters{attemptsMinute: v[0], failuresMinute: v[1], attemptsDay: v[2], failuresDay: v[3]}
}
return out
}
// usageKeysPerScope is the number of windowKey strings usageWindowKeys
// builds per scope, and the stride currentUsage's flat storeGetMulti
// result is sliced back into per-scope chunks by. Named here, closing a
+172
View File
@@ -1,9 +1,11 @@
package traefikllmgateway
import (
"context"
"errors"
"fmt"
"math"
"net/http"
"testing"
"time"
@@ -1374,3 +1376,173 @@ func TestMemoryStore_IncrMulti_DayCounterStillClampsAt48h(t *testing.T) {
t.Errorf("day counter at 49h past write = %d, want 0 (expired at the 48h ceiling; day's own 25h floor never raises it)", got)
}
}
// --- Feature A (v0.22): per-provider/per-(provider,model) success-rate accounting ---
// TestRecordProviderAttempt_ClassifiesUsingIsTransient proves
// recordProviderAttempt's failure classification is exactly isTransient's
// own — reused, not forked (spec ruling) — across the full outcome table:
// success, a non-429 4xx (still a "success" for provider-health purposes:
// the provider answered correctly to a request it did not like), 429, a
// 5xx, a plain network error, and a context-canceled error (not a
// provider fault, per isTransient's own carve-out).
func TestRecordProviderAttempt_ClassifiesUsingIsTransient(t *testing.T) {
tests := []struct {
err error
resp *http.Response
name string
wantFail bool
}{
{name: "200 OK", resp: &http.Response{StatusCode: http.StatusOK}, wantFail: false},
{name: "404 not found (provider answered)", resp: &http.Response{StatusCode: http.StatusNotFound}, wantFail: false},
{name: "429 too many requests", resp: &http.Response{StatusCode: http.StatusTooManyRequests}, wantFail: true},
{name: "500 internal server error", resp: &http.Response{StatusCode: http.StatusInternalServerError}, wantFail: true},
{name: "503 service unavailable", resp: &http.Response{StatusCode: http.StatusServiceUnavailable}, wantFail: true},
{name: "network error", err: errors.New("dial tcp: connection refused"), wantFail: true},
{name: "context canceled (not a provider fault)", err: context.Canceled, wantFail: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
l := newLimiter(nil, true)
now := time.Date(2026, 8, 21, 12, 0, 0, 0, time.UTC)
l.nowFn = func() time.Time { return now }
l.recordProviderAttempt("openai", "gpt-4o", tt.resp, tt.err)
attempts, _ := l.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, now)
if attempts != 1 {
t.Errorf("provider attempts/day = %d, want 1", attempts)
}
fails, _ := l.getCounter(kindProvider, "openai", metricProvFail, windowDay, now)
wantFails := int64(0)
if tt.wantFail {
wantFails = 1
}
if fails != wantFails {
t.Errorf("provider fails/day = %d, want %d", fails, wantFails)
}
// Minute window and the (provider, model) scope both mirror the
// provider/day counters exactly.
attemptsMin, _ := l.getCounter(kindProvider, "openai", metricProvAttempt, windowMin, now)
if attemptsMin != 1 {
t.Errorf("provider attempts/min = %d, want 1", attemptsMin)
}
modelAttempts, _ := l.getCounter(kindProviderModel, "openai/gpt-4o", metricProvAttempt, windowDay, now)
if modelAttempts != 1 {
t.Errorf("model attempts/day = %d, want 1", modelAttempts)
}
modelFails, _ := l.getCounter(kindProviderModel, "openai/gpt-4o", metricProvFail, windowDay, now)
if modelFails != wantFails {
t.Errorf("model fails/day = %d, want %d", modelFails, wantFails)
}
})
}
}
// TestRecordProviderAttempt_EmptyModel_ProviderScopeOnly proves a caller
// that cannot cheaply know the upstream model before the attempt resolves
// (native passthrough, routes_passthrough.go's handlePassthrough: the
// model lives in the response body, read only afterward) gets
// provider-level accounting only when it passes model="".
func TestRecordProviderAttempt_EmptyModel_ProviderScopeOnly(t *testing.T) {
l := newLimiter(nil, true)
now := time.Date(2026, 8, 21, 12, 0, 0, 0, time.UTC)
l.nowFn = func() time.Time { return now }
l.recordProviderAttempt("openai", "", &http.Response{StatusCode: http.StatusInternalServerError}, nil)
attempts, _ := l.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, now)
if attempts != 1 {
t.Errorf("provider attempts/day = %d, want 1", attempts)
}
fails, _ := l.getCounter(kindProvider, "openai", metricProvFail, windowDay, now)
if fails != 1 {
t.Errorf("provider fails/day = %d, want 1", fails)
}
}
// TestRecordProviderAttempt_EmptyProvider_NoOp proves a caller with no
// resolved provider name (should never happen in production traffic, but
// must never panic) is a silent no-op.
func TestRecordProviderAttempt_EmptyProvider_NoOp(t *testing.T) {
l := newLimiter(nil, true)
l.recordProviderAttempt("", "some-model", &http.Response{StatusCode: http.StatusOK}, nil) // must not panic
}
// TestProviderModelScopeID proves the "provider/model" id convention
// providerUsage/recordProviderAttempt share with pricing.go's own
// unifiedCostMicros and the admin API.
func TestProviderModelScopeID(t *testing.T) {
if got := providerModelScopeID("openai", "gpt-4o"); got != "openai/gpt-4o" {
t.Errorf("providerModelScopeID = %q, want %q", got, "openai/gpt-4o")
}
}
// TestProviderUsage_BatchedSingleRoundTrip drives limiter.providerUsage
// against a counting stub store (historyCountingStore, shared with
// TestLimiterHistory_OneBatchCallOldestFirst above): it must issue
// exactly ONE getMulti call for both a provider-level and a
// (provider, model) scope together, and slice the flat result back to the
// right scope in order.
func TestProviderUsage_BatchedSingleRoundTrip(t *testing.T) {
now := time.Date(2026, 8, 21, 12, 0, 0, 0, time.UTC)
countStore := &historyCountingStore{values: map[string]int64{
windowKey(kindProvider, "openai", metricProvAttempt, windowMin, now): 3,
windowKey(kindProvider, "openai", metricProvFail, windowMin, now): 1,
windowKey(kindProvider, "openai", metricProvAttempt, windowDay, now): 30,
windowKey(kindProvider, "openai", metricProvFail, windowDay, now): 2,
windowKey(kindProviderModel, "openai/gpt-4o", metricProvAttempt, windowDay, now): 10,
}}
l := newLimiter(countStore, true)
l.nowFn = func() time.Time { return now }
scopes := []limitScope{
{kind: kindProvider, id: "openai"},
{kind: kindProviderModel, id: "openai/gpt-4o"},
}
got := l.providerUsage(scopes)
if countStore.getMultiCalls != 1 {
t.Errorf("getMultiCalls = %d, want 1 (both scopes' keys flattened into one call)", countStore.getMultiCalls)
}
if len(got) != 2 {
t.Fatalf("len(got) = %d, want 2", len(got))
}
if got[0].attemptsMinute != 3 || got[0].failuresMinute != 1 || got[0].attemptsDay != 30 || got[0].failuresDay != 2 {
t.Errorf("provider counters = %+v, want {3 1 30 2}", got[0])
}
if got[1].attemptsDay != 10 || got[1].failuresDay != 0 {
t.Errorf("model counters = %+v, want attemptsDay 10, failuresDay 0", got[1])
}
}
// TestProviderUsage_EmptyScopes_ReturnsEmptyWithoutTouchingStore mirrors
// targetUsage's own empty-input short circuit.
func TestProviderUsage_EmptyScopes_ReturnsEmptyWithoutTouchingStore(t *testing.T) {
countStore := &historyCountingStore{}
l := newLimiter(countStore, true)
got := l.providerUsage(nil)
if len(got) != 0 {
t.Errorf("len(got) = %d, want 0", len(got))
}
if countStore.getMultiCalls != 0 {
t.Errorf("getMultiCalls = %d, want 0 (empty scopes must not touch the store)", countStore.getMultiCalls)
}
}
// TestProviderUsage_StoreDown_ReturnsZeroCounters mirrors targetUsage's
// own fail-closed contract: a store error with failOpen=false must report
// every scope's counters as zero, not partial or stale values.
func TestProviderUsage_StoreDown_ReturnsZeroCounters(t *testing.T) {
l := newLimiter(alwaysErrStore{}, false)
got := l.providerUsage([]limitScope{{kind: kindProvider, id: "openai"}})
if len(got) != 1 {
t.Fatalf("len(got) = %d, want 1", len(got))
}
if got[0] != (providerCounters{}) {
t.Errorf("counters = %+v, want the zero value on a storeDown read", got[0])
}
}
+10
View File
@@ -684,6 +684,16 @@ func TestHandleTargetProxy_JSONResponse_OnlyRequestCounterMoves(t *testing.T) {
if !ok || totalReq != 1 {
t.Errorf("total request/min counter = %d (ok=%v), want 1", totalReq, ok)
}
// Feature A (v0.22): handleTargetProxy/proxyUpstream never wrap the
// request context with an attemptRecorder — an MCP/A2A target proxy
// attempt against the "openai" MCP server's own upstream must never be
// mistaken for provider traffic, even though a provider of the same
// name happens to be configured too.
provAttempts, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, time.Now())
if provAttempts != 0 {
t.Errorf("provider attempts/day = %d, want 0 — the MCP target proxy must never record provider accounting", provAttempts)
}
}
// --- config validation at construction ---
+44
View File
@@ -184,6 +184,50 @@ func newAdapterHTTPClient() *http.Client {
return &http.Client{Transport: tr}
}
// attemptRecorder is invoked once per upstream HTTP attempt a shared
// chokepoint makes — retryPolicy.do's own retry loop (retry.go), for every
// adapter method that goes through upstreamJSON/upstreamRawBytes, and
// proxyUpstream's single client.Do call (routes_passthrough.go) — with
// that attempt's raw resp/err, for Feature A's (v0.22) per-provider
// success-rate accounting (limiter.recordProviderAttempt, limits.go). A
// recorder must never touch resp.Body (still owned by its caller, read or
// forwarded afterward) and must never retain resp or err past the call.
type attemptRecorder func(resp *http.Response, err error)
// attemptRecorderCtxKey is the unexported context.Value key
// withAttemptRecorder/attemptRecorderFromContext share. An unexported
// struct type, not a string, so no other package's context.WithValue call
// can ever collide with it by accident.
type attemptRecorderCtxKey struct{}
// withAttemptRecorder returns a context carrying rec, so a chokepoint
// shared across call sites that must NOT all record provider accounting —
// retryPolicy.do and proxyUpstream are both also reached by traffic
// Feature A is deliberately out of scope for: registry.go's discovery
// listModels calls (no recorder ever set on their ctx) and the MCP/A2A
// target proxy (handleTargetProxy, mcp_a2a.go, which never wraps its
// request's context this way) — can look the recorder up without either
// function needing a dedicated parameter every call site would otherwise
// have to pass. A nil rec returns ctx unchanged, so a caller may call this
// unconditionally.
func withAttemptRecorder(ctx context.Context, rec attemptRecorder) context.Context {
if rec == nil {
return ctx
}
return context.WithValue(ctx, attemptRecorderCtxKey{}, rec)
}
// attemptRecorderFromContext returns the attemptRecorder ctx carries via
// withAttemptRecorder, or nil when none was set — the common case for any
// call path Feature A does not account (see withAttemptRecorder's doc
// comment). A bare comma-ok type assertion, not errors.As or any reflect-
// based check: see providerHTTPError's own doc comment for why a
// yaegi-interpreted plugin must avoid errors.As here.
func attemptRecorderFromContext(ctx context.Context) attemptRecorder {
rec, _ := ctx.Value(attemptRecorderCtxKey{}).(attemptRecorder)
return rec
}
// upstreamJSON issues an HTTP request to url: body, when non-nil, is
// marshaled as the JSON request body; hdr's values are added to the
// request (a caller builds this from its adapter's auth and content-type
+64
View File
@@ -0,0 +1,64 @@
package traefikllmgateway
import (
"context"
"net/http"
"testing"
)
// TestAttemptRecorderFromContext_RoundTrip proves withAttemptRecorder/
// attemptRecorderFromContext round-trip a non-nil recorder, and that the
// SAME func value comes back out (asserted indirectly here, by observing
// it fire when invoked) — Feature A's (v0.22) chokepoint plumbing between
// providers.go, retry.go's policy.do, and routes_passthrough.go's
// proxyUpstream.
func TestAttemptRecorderFromContext_RoundTrip(t *testing.T) {
var gotResp *http.Response
var gotErr error
rec := attemptRecorder(func(resp *http.Response, err error) {
gotResp = resp
gotErr = err
})
ctx := withAttemptRecorder(context.Background(), rec)
got := attemptRecorderFromContext(ctx)
if got == nil {
t.Fatal("attemptRecorderFromContext returned nil after withAttemptRecorder set one")
}
wantResp := &http.Response{StatusCode: http.StatusOK}
got(wantResp, nil)
if gotResp != wantResp {
t.Errorf("gotResp = %v, want %v (the recorder round-tripped through ctx must be the same func)", gotResp, wantResp)
}
if gotErr != nil {
t.Errorf("gotErr = %v, want nil", gotErr)
}
}
// TestAttemptRecorderFromContext_NoneSet proves a ctx that never called
// withAttemptRecorder — the common case: registry.go's discovery
// listModels calls, and handleTargetProxy's MCP/A2A target proxy
// (mcp_a2a.go), neither of which is Feature A's accounting scope —
// returns a nil recorder, not a panic or a zero-value func that panics
// when invoked.
func TestAttemptRecorderFromContext_NoneSet(t *testing.T) {
if got := attemptRecorderFromContext(context.Background()); got != nil {
t.Errorf("attemptRecorderFromContext(bare ctx) = %v, want nil", got)
}
}
// TestWithAttemptRecorder_NilRecorder_ReturnsSameContext proves passing a
// nil recorder is a safe no-op a caller may call unconditionally — it
// does not wrap ctx in a no-op value that would still satisfy a later
// "!= nil" nil-check with a non-functional func.
func TestWithAttemptRecorder_NilRecorder_ReturnsSameContext(t *testing.T) {
ctx := context.Background()
got := withAttemptRecorder(ctx, nil)
if got != ctx { //nolint:staticcheck // deliberately comparing context identity, not using it as a map key
t.Error("withAttemptRecorder(ctx, nil) returned a different context; want ctx unchanged")
}
if rec := attemptRecorderFromContext(got); rec != nil {
t.Errorf("attemptRecorderFromContext = %v, want nil after a nil recorder was passed in", rec)
}
}
+13
View File
@@ -215,13 +215,26 @@ func drainAndClose(resp *http.Response) {
// last attempt's (resp, err) as-is either way: on success, on a
// still-transient failure once tries are exhausted, on a non-transient
// failure at any attempt, or on a wait ctx aborted early.
//
// Feature A (v0.22) attempt-accounting: every call() invocation — not just
// the final one — is reported to ctx's attemptRecorder, if it carries one
// (attemptRecorderFromContext, providers.go), before the transient check
// above decides whether to retry. A request retried twice before
// succeeding therefore reports THREE attempts, not one — the caller
// (limiter.recordProviderAttempt, limits.go) classifies each with the
// identical isTransient(resp, err) this loop already uses, so "attempt"
// and "failure" here always agree with what actually got retried.
func (p *retryPolicy) do(ctx context.Context, call func() (*http.Response, error)) (*http.Response, error) {
tries := p.totalTries()
rec := attemptRecorderFromContext(ctx)
var resp *http.Response
var err error
for attempt := 1; attempt <= tries; attempt++ {
resp, err = call()
if rec != nil {
rec(resp, err)
}
if !isTransient(resp, err) || attempt == tries {
return resp, err
}
+58
View File
@@ -871,3 +871,61 @@ func TestOpenAIAdapter_ChatCompletion_Stream_MidDeathNotRetried(t *testing.T) {
t.Errorf("rec.Body = %q, want it to contain the chunk already forwarded before the drop", rec.Body.String())
}
}
// TestRetryPolicy_Do_AttemptRecorder_FiresOncePerAttempt proves
// attempt-accounting (Feature A, v0.22 — see do's own doc comment and
// limits.go's recordProviderAttempt): a request retried twice before
// succeeding reports THREE attempts to ctx's attemptRecorder, not one —
// each with that specific attempt's own resp, in call order.
func TestRetryPolicy_Do_AttemptRecorder_FiresOncePerAttempt(t *testing.T) {
var calls int32
p := &retryPolicy{enabled: true, attempts: 2, backoff: time.Millisecond, waitFn: spyWait(&[]time.Duration{})}
call := callSequence(&calls,
func() (*http.Response, error) { return fakeResp(http.StatusInternalServerError, nil), nil },
func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil },
func() (*http.Response, error) { return fakeResp(http.StatusOK, nil), nil },
)
var recorded []int
ctx := withAttemptRecorder(context.Background(), func(resp *http.Response, err error) {
if err != nil {
t.Fatalf("recorder got unexpected err: %v", err)
}
recorded = append(recorded, resp.StatusCode)
})
resp, err := p.do(ctx, call)
if err != nil {
t.Fatalf("do: %v", err)
}
if resp.StatusCode != http.StatusOK {
t.Errorf("resp.StatusCode = %d, want 200", resp.StatusCode)
}
want := []int{http.StatusInternalServerError, http.StatusServiceUnavailable, http.StatusOK}
if len(recorded) != len(want) {
t.Fatalf("recorded = %v, want %v", recorded, want)
}
for i, w := range want {
if recorded[i] != w {
t.Errorf("recorded[%d] = %d, want %d", i, recorded[i], w)
}
}
}
// TestRetryPolicy_Do_NoAttemptRecorderInContext_NoPanic proves the common
// case — a ctx that never called withAttemptRecorder, e.g. registry.go's
// discovery listModels calls or the MCP/A2A target proxy — is a silent
// no-op, not a nil-func-call panic.
func TestRetryPolicy_Do_NoAttemptRecorderInContext_NoPanic(t *testing.T) {
p := &retryPolicy{enabled: true, attempts: 1, backoff: time.Millisecond, waitFn: spyWait(&[]time.Duration{})}
call := func() (*http.Response, error) { return fakeResp(http.StatusOK, nil), nil }
resp, err := p.do(context.Background(), call) // must not panic
if err != nil {
t.Fatalf("do: %v", err)
}
if resp.StatusCode != http.StatusOK {
t.Errorf("resp.StatusCode = %d, want 200", resp.StatusCode)
}
}
+18 -3
View File
@@ -96,7 +96,12 @@ func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request
}
req["model"] = upstreamModel
if _, err := adapter.imagesGeneration(r.Context(), sw, req); err != nil {
// Feature A (v0.22): see routes_unified.go's runUnified for the
// identical pattern this mirrors.
ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
})
if _, err := adapter.imagesGeneration(ctx, sw, req); err != nil {
g.handleAdapterError(sw, err, adapter.name())
}
}
@@ -145,7 +150,12 @@ func (g *Gateway) handleAudioSpeech(w http.ResponseWriter, r *http.Request, u *u
return
}
if _, err := adapter.audioSpeech(r.Context(), sw, body, "application/json"); err != nil {
// Feature A (v0.22): see routes_unified.go's runUnified for the
// identical pattern this mirrors.
ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
})
if _, err := adapter.audioSpeech(ctx, sw, body, "application/json"); err != nil {
g.handleAdapterError(sw, err, adapter.name())
}
}
@@ -230,7 +240,12 @@ func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Reque
uploadContentType = rebuiltContentType
}
if _, err := adapter.audioTranscription(r.Context(), sw, uploadBody, uploadContentType); err != nil {
// Feature A (v0.22): see routes_unified.go's runUnified for the
// identical pattern this mirrors.
ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
})
if _, err := adapter.audioTranscription(ctx, sw, uploadBody, uploadContentType); err != nil {
g.handleAdapterError(sw, err, adapter.name())
}
}
+18
View File
@@ -179,6 +179,24 @@ func TestHandleImagesGenerations_OpenAI_HappyPath_NativeForward(t *testing.T) {
if !ok || totalReq != 1 {
t.Errorf("total request/min counter = %d (ok=%v), want 1", totalReq, ok)
}
// Feature A (v0.22): handleImagesGenerations wraps the request context
// with an attemptRecorder before calling adapter.imagesGeneration
// (routes_media.go) — a successful 200 upstream attempt must land as
// one attempt, zero failures, at both provider and (provider, model)
// scope.
attempts, ok := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, time.Now())
if !ok || attempts != 1 {
t.Errorf("provider attempts/day = %d (ok=%v), want 1", attempts, ok)
}
fails, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvFail, windowDay, time.Now())
if fails != 0 {
t.Errorf("provider fails/day = %d, want 0", fails)
}
modelAttempts, ok := gw.limiter.getCounter(kindProviderModel, "openai/img-test", metricProvAttempt, windowDay, time.Now())
if !ok || modelAttempts != 1 {
t.Errorf("model attempts/day = %d (ok=%v), want 1", modelAttempts, ok)
}
}
// TestHandleImagesGenerations_ModelAlias_ResolvesToTargetUpstreamModel
+19
View File
@@ -281,6 +281,15 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
upstreamURL += "?" + r.URL.RawQuery
}
// Feature A (v0.22): provider-level only — unlike runUnified/the media
// routes, the upstream model here lives in the response body
// (extractPassthroughUsage, below), read only after this attempt
// already resolved, so there is no model to attribute it to yet. See
// recordProviderAttempt's own doc comment (limits.go).
r = r.WithContext(withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
g.limiter.recordProviderAttempt(providerName, "", resp, attemptErr)
}))
result, ok := g.proxyUpstream(w, r, upstreamURL, adapter.httpClient(), adapter.injectAuth, true, "passthrough (provider "+providerName+")")
if !ok || !result.isJSON {
// A build/connection/copy failure already wrote its own response
@@ -384,6 +393,16 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
}
resp, err := client.Do(upstreamReq) //nolint:gosec // same upstreamReq built above; operator-fixed host, traversal-checked, see its construction comment
// Feature A (v0.22): proxyUpstream makes exactly one attempt (no
// retry.go policy wraps this path), so this fires once per call,
// whichever way it resolves. Only handlePassthrough's own context ever
// carries a recorder (attemptRecorderFromContext, providers.go) —
// handleTargetProxy (mcp_a2a.go), this function's other caller, never
// wraps r's context this way, so an MCP/A2A target proxy attempt is
// correctly never accounted as provider traffic.
if rec := attemptRecorderFromContext(r.Context()); rec != nil {
rec(resp, err)
}
if err != nil {
if errors.Is(err, context.Canceled) {
g.logf("%s: client canceled request: %v", logPrefix, err)
+62
View File
@@ -480,6 +480,68 @@ func TestHandlePassthrough_NonStreamJSON_AccountsUsage(t *testing.T) {
if !ok || totalTokOut != 3 {
t.Errorf("total tokout/day counter = %d (ok=%v), want 3", totalTokOut, ok)
}
// Feature A (v0.22): handlePassthrough wraps r's context with an
// attemptRecorder before calling g.proxyUpstream — provider-level
// only (model="": the upstream model here lives in the response body,
// read only after this attempt already resolved, per
// recordProviderAttempt's own doc comment, limits.go).
attempts, ok := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, time.Now())
if !ok || attempts != 1 {
t.Errorf("provider attempts/day = %d (ok=%v), want 1", attempts, ok)
}
fails, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvFail, windowDay, time.Now())
if fails != 0 {
t.Errorf("provider fails/day = %d, want 0", fails)
}
}
// TestHandlePassthrough_DeadUpstream_RecordsProviderFailure proves a
// connection-level failure (proxyUpstream's client.Do returning a network
// error) still reports one provider-level attempt AND one failure —
// Feature A (v0.22): passthrough makes no retry.go attempt at all, so
// this exercises proxyUpstream's own attemptRecorderFromContext call
// directly, not retryPolicy.do's.
func TestHandlePassthrough_DeadUpstream_RecordsProviderFailure(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
deadURL := srv.URL
srv.Close() // closed before any request: every dial fails
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{
"openai": {Type: "openai", BaseURL: deadURL, APIKey: "sk-up"},
}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{
{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}},
}}
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, ok := h.(*Gateway)
if !ok {
t.Fatal("handler is not *Gateway")
}
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(`{}`))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadGateway {
t.Fatalf("status = %d, want 502, body=%s", rec.Code, rec.Body.String())
}
attempts, ok := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, time.Now())
if !ok || attempts != 1 {
t.Errorf("provider attempts/day = %d (ok=%v), want 1", attempts, ok)
}
fails, ok := gw.limiter.getCounter(kindProvider, "openai", metricProvFail, windowDay, time.Now())
if !ok || fails != 1 {
t.Errorf("provider fails/day = %d (ok=%v), want 1", fails, ok)
}
}
// TestHandlePassthrough_GroupDeniesProvider_Returns403 covers a provider
+8 -1
View File
@@ -172,7 +172,14 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr
respWriter = capture
}
result, callErr := call(adapter, r.Context(), respWriter, req)
// Feature A (v0.22): every upstream attempt call() makes — via
// upstreamJSON's retryPolicy.do, providers.go — reports through to
// this recorder, closing over the adapter/model this request already
// resolved to. See attemptRecorder's own doc comment (providers.go).
ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
})
result, callErr := call(adapter, ctx, respWriter, req)
// Usage is accounted before the error branch below runs, not after:
// every adapter that can fail mid-stream (forwardStream in each of the