mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
lukaszraczylo
aafd2a0433
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.
65 lines
2.4 KiB
Go
65 lines
2.4 KiB
Go
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)
|
|
}
|
|
}
|