mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
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:
@@ -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
@@ -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])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 ---
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
@@ -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())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user