mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-18 05:47:38 +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:
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user