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:
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user