diff --git a/limits.go b/limits.go index fee20e1..6e9f63d 100644 --- a/limits.go +++ b/limits.go @@ -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 diff --git a/limits_test.go b/limits_test.go index 222d2d1..cd8ee9f 100644 --- a/limits_test.go +++ b/limits_test.go @@ -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]) + } +} diff --git a/mcp_a2a_test.go b/mcp_a2a_test.go index 5b05d93..2636453 100644 --- a/mcp_a2a_test.go +++ b/mcp_a2a_test.go @@ -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 --- diff --git a/providers.go b/providers.go index 3c5471b..af7ee87 100644 --- a/providers.go +++ b/providers.go @@ -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 diff --git a/providers_test.go b/providers_test.go new file mode 100644 index 0000000..ba7beec --- /dev/null +++ b/providers_test.go @@ -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) + } +} diff --git a/retry.go b/retry.go index d9abdd3..9cfce2f 100644 --- a/retry.go +++ b/retry.go @@ -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 } diff --git a/retry_test.go b/retry_test.go index 689bea7..bbf51b8 100644 --- a/retry_test.go +++ b/retry_test.go @@ -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) + } +} diff --git a/routes_media.go b/routes_media.go index eafdeb0..c05cb2a 100644 --- a/routes_media.go +++ b/routes_media.go @@ -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()) } } diff --git a/routes_media_test.go b/routes_media_test.go index 23a81a8..4d19e89 100644 --- a/routes_media_test.go +++ b/routes_media_test.go @@ -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 diff --git a/routes_passthrough.go b/routes_passthrough.go index ef149a4..1ed3c6e 100644 --- a/routes_passthrough.go +++ b/routes_passthrough.go @@ -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) diff --git a/routes_passthrough_test.go b/routes_passthrough_test.go index 04c570a..da716c9 100644 --- a/routes_passthrough_test.go +++ b/routes_passthrough_test.go @@ -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 diff --git a/routes_unified.go b/routes_unified.go index 99fe9f6..32a9e43 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -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