mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
feat: split token counters and hour-resolution usage history
Data layer for the upcoming Vue admin charts (v0.2 data-layer task).
Tokens-in/out split (limits.go): metricTok replaced by metricTokIn
("tokin") and metricTokOut ("tokout"). account() records usage.prompt
and usage.completion independently, skipping a direction only when its
own value is zero. A TokensPerDay/TokensPerMonth limit still enforces
one TOTAL budget: tokenBudgetViolation reads both counters in one
storeGetMulti batch and sums them; the violation message still says
"tokens". Estimated usage (the request-body-size fallback) lands
entirely in tokin, documented at its call site.
Hour window (limits.go): windowHour ("hour"), bucket format
"2006010215" UTC, hourWindowTTL (48h). checkAndCount/account both write
it on every metered route; no LimitsConfig field ever names it, so
evaluateScope never evaluates it. dayWindowTTL/monthWindowTTL bumped to
35d/400d so the usage-history API's largest span always finds a live
key.
Total scope (limits.go, routes_unified.go): a synthetic
{kind:"total", id:"all"} scope, limits always nil. withTotalScope
wraps buildLimitScopes' result at every metered route (unified,
media, passthrough, mcp/a2a) so none can forget it; handleAdminAPI's
own accounting stays on buildLimitScopes alone, so admin traffic never
contributes to it.
History API (admin.go): GET /admin/api/usage/history?scope=&metric=&
window=&span= — admin-gated like the other two JSON routes, but
skips checkAndCount ahead of that same change landing for every
/admin/api/* route in a later task. Validates scope kind/metric/window
(400), span bounds per window (400), and unknown user/group id (404).
Computes every bucket key up front and reads them in ONE storeGetMulti
batch; a storeDown read is 503, never a silently-zero series. Works
against memoryStore too.
Admin usage API (admin.go): scopeUsage/adminUsageEntryView now expose
tokensInPerDay/tokensOutPerDay and their per-month counterparts in
place of the combined tokensPerDay/tokensPerMonth pair; buildAdminUsage
adds a "total" row to the response, still one storeGetMulti round trip
overall. usageKeysPerScope names the per-scope key stride (was a bare
6, now 8), closing a deferred review follow-up while the count changed
anyway.
Tests: TDD throughout — bucket/window-key math incl. UTC and hour
boundaries, the month step-back day-31 AddDate trap, split accounting,
token budget summing both directions, total-scope counting-without-
evaluating on every metered route, and the full history API (shape,
oldest-first, span bounds, validation table, 404, 503, one-batch
assertion via a counting stub). Integration: TestAdminDashboard now
fetches an hour history series for user alice and for total and
asserts a non-zero current-hour bucket.
Migration: none — old "tok" counter keys simply expire on their
existing TTL; nothing is deployed yet.
This commit is contained in:
@@ -2,20 +2,23 @@ package traefikllmgateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// The three routes the read-only admin dashboard registers (spec §4,
|
||||
// v0.2), matched only when adminEnabled(g.cfg) — see llmgateway.go's
|
||||
// ServeHTTP dispatch.
|
||||
// The four routes the read-only admin dashboard registers (spec §4, v0.2;
|
||||
// adminUsageHistoryPath added by the v0.2 data-layer task), matched only
|
||||
// when adminEnabled(g.cfg) — see llmgateway.go's ServeHTTP dispatch.
|
||||
const (
|
||||
adminPagePath = "/admin"
|
||||
adminOverviewPath = "/admin/api/overview"
|
||||
adminUsagePath = "/admin/api/usage"
|
||||
adminPagePath = "/admin"
|
||||
adminOverviewPath = "/admin/api/overview"
|
||||
adminUsagePath = "/admin/api/usage"
|
||||
adminUsageHistoryPath = "/admin/api/usage/history"
|
||||
)
|
||||
|
||||
// adminCSP is the Content-Security-Policy header served with all three
|
||||
@@ -43,9 +46,9 @@ func adminEnabled(cfg *Config) bool {
|
||||
return cfg.Admin != nil && cfg.Admin.Enabled
|
||||
}
|
||||
|
||||
// isAdminPath reports whether path is one of the three admin routes.
|
||||
// isAdminPath reports whether path is one of the four admin routes.
|
||||
func isAdminPath(path string) bool {
|
||||
return path == adminPagePath || path == adminOverviewPath || path == adminUsagePath
|
||||
return path == adminPagePath || path == adminOverviewPath || path == adminUsagePath || path == adminUsageHistoryPath
|
||||
}
|
||||
|
||||
// handleAdmin is ServeHTTP's single entry point for all three /admin*
|
||||
@@ -70,13 +73,21 @@ func (g *Gateway) handleAdmin(sw *statusTrackingWriter, r *http.Request) {
|
||||
g.handleAdminAPI(sw, r)
|
||||
}
|
||||
|
||||
// handleAdminAPI is the gate for the two /admin/api/* JSON routes,
|
||||
// handleAdminAPI is the gate for the three /admin/api/* JSON routes,
|
||||
// applying spec §4's gate order: unauthenticated → 401, authenticated
|
||||
// non-admin → 403, admin → serve. Like every other authenticated route
|
||||
// it counts request counters via checkAndCount before serving: an admin
|
||||
// over their own req/min limit gets a 429 here exactly as they would on
|
||||
// any other route (spec §4's "Accounting: admin routes count request
|
||||
// counters like any authed route").
|
||||
// non-admin → 403, admin → serve. GET /admin/api/overview and GET
|
||||
// /admin/api/usage count request counters via checkAndCount before
|
||||
// serving, like every other authenticated route: an admin over their own
|
||||
// req/min limit gets a 429 here exactly as they would on any other route
|
||||
// (spec §4's "Accounting: admin routes count request counters like any
|
||||
// authed route").
|
||||
//
|
||||
// GET /admin/api/usage/history (v0.2 data-layer task) is the one
|
||||
// exception: it skips checkAndCount entirely. The operator-directed
|
||||
// removal of admin-route stat counting (progress ledger, 2026-08-20)
|
||||
// lands for every /admin/api/* route in a later task; this new route is
|
||||
// built consistent with that direction from the start, rather than
|
||||
// counted now and un-counted later.
|
||||
func (g *Gateway) handleAdminAPI(sw *statusTrackingWriter, r *http.Request) {
|
||||
u, grp, ok := g.auth.identify(r)
|
||||
g.logAuthEvent(ok, authEventUserName(u), r)
|
||||
@@ -88,6 +99,12 @@ func (g *Gateway) handleAdminAPI(sw *statusTrackingWriter, r *http.Request) {
|
||||
writeOAIError(sw, http.StatusForbidden, "invalid_request_error", "admin access required")
|
||||
return
|
||||
}
|
||||
|
||||
if r.URL.Path == adminUsageHistoryPath {
|
||||
g.serveAdminUsageHistory(sw, r)
|
||||
return
|
||||
}
|
||||
|
||||
if violation := g.limiter.checkAndCount(buildLimitScopes(u, grp)); violation != nil {
|
||||
writeLimitViolation(sw, violation)
|
||||
return
|
||||
@@ -322,12 +339,16 @@ func (g *Gateway) serveAdminOverview(w http.ResponseWriter) {
|
||||
_ = json.NewEncoder(w).Encode(g.buildAdminOverview())
|
||||
}
|
||||
|
||||
// adminUsageEntryView is one user's or group's row in GET
|
||||
// /admin/api/usage (spec §4, v0.2): its current-window counter values,
|
||||
// with its configured limit echoed beside each. GroupName is set only for
|
||||
// a user entry (empty for a group entry) — the group a user currently
|
||||
// belongs to, for display; it carries no access-control meaning of its
|
||||
// own.
|
||||
// adminUsageEntryView is one user's, group's, or the total scope's row in
|
||||
// GET /admin/api/usage (spec §4, v0.2): its current-window counter
|
||||
// values, with its configured limit echoed beside each. GroupName is set
|
||||
// only for a user entry (empty for a group or the total entry) — the
|
||||
// group a user currently belongs to, for display; it carries no
|
||||
// access-control meaning of its own. TokensInPerDay/TokensOutPerDay and
|
||||
// their per-month counterparts replace a single combined
|
||||
// tokensPerDay/tokensPerMonth pair (v0.2 data-layer task, tokens-in/
|
||||
// tokens-out split) — a client wanting the combined total sums the two
|
||||
// itself.
|
||||
type adminUsageEntryView struct {
|
||||
Limits *LimitsConfig `json:"limits,omitempty"`
|
||||
Kind string `json:"kind"`
|
||||
@@ -335,8 +356,10 @@ type adminUsageEntryView struct {
|
||||
GroupName string `json:"groupName,omitempty"`
|
||||
RequestsPerMinute int64 `json:"requestsPerMinute"`
|
||||
RequestsPerDay int64 `json:"requestsPerDay"`
|
||||
TokensPerDay int64 `json:"tokensPerDay"`
|
||||
TokensPerMonth int64 `json:"tokensPerMonth"`
|
||||
TokensInPerDay int64 `json:"tokensInPerDay"`
|
||||
TokensOutPerDay int64 `json:"tokensOutPerDay"`
|
||||
TokensInPerMonth int64 `json:"tokensInPerMonth"`
|
||||
TokensOutPerMonth int64 `json:"tokensOutPerMonth"`
|
||||
CostPerDayMicroUSD int64 `json:"costPerDayMicroUsd"`
|
||||
CostPerMonthMicroUSD int64 `json:"costPerMonthMicroUsd"`
|
||||
StoreDown bool `json:"storeDown,omitempty"`
|
||||
@@ -344,11 +367,14 @@ type adminUsageEntryView struct {
|
||||
|
||||
// adminUsageResponse is the full body of GET /admin/api/usage: every
|
||||
// currently active user, and every configured group, each with its own
|
||||
// usage+limits row. Both slices are sorted (authStore.snapshot's own
|
||||
// contract) and built fresh per request.
|
||||
// usage+limits row, plus the synthetic total scope's own row (v0.2
|
||||
// data-layer task) — usage summed across every user and group combined,
|
||||
// limits always nil (see totalScopeKind). Users and Groups are both
|
||||
// sorted (authStore.snapshot's own contract) and built fresh per request.
|
||||
type adminUsageResponse struct {
|
||||
Users []adminUsageEntryView `json:"users"`
|
||||
Groups []adminUsageEntryView `json:"groups"`
|
||||
Total adminUsageEntryView `json:"total"`
|
||||
}
|
||||
|
||||
// usageEntryView converts one limiter.currentUsage result plus its
|
||||
@@ -360,8 +386,10 @@ func usageEntryView(su scopeUsage, limits *LimitsConfig) adminUsageEntryView {
|
||||
Limits: limits,
|
||||
RequestsPerMinute: su.requestsPerMinute,
|
||||
RequestsPerDay: su.requestsPerDay,
|
||||
TokensPerDay: su.tokensPerDay,
|
||||
TokensPerMonth: su.tokensPerMonth,
|
||||
TokensInPerDay: su.tokensInPerDay,
|
||||
TokensOutPerDay: su.tokensOutPerDay,
|
||||
TokensInPerMonth: su.tokensInPerMonth,
|
||||
TokensOutPerMonth: su.tokensOutPerMonth,
|
||||
CostPerDayMicroUSD: su.costPerDayMicros,
|
||||
CostPerMonthMicroUSD: su.costPerMonthMicros,
|
||||
StoreDown: su.storeDown,
|
||||
@@ -371,28 +399,32 @@ func usageEntryView(su scopeUsage, limits *LimitsConfig) adminUsageEntryView {
|
||||
// buildAdminUsage assembles adminUsageResponse: authStore.snapshot lists
|
||||
// every currently active user and every configured group (names, group
|
||||
// membership, and limits only — never a key or its digest), and
|
||||
// limiter.currentUsage reads every one's six current-window counters in
|
||||
// ONE storeGetMulti round trip total (v0.2 final review wave, 2026-08-20;
|
||||
// see limiter.currentUsage's own doc comment for the tradeoff this
|
||||
// supersedes) — users and groups are concatenated into a single scopes
|
||||
// slice before that one call, then the flat result is sliced back into
|
||||
// the two response sections at the same split point, order preserved.
|
||||
// Unlike buildLimitScopes (routes_unified.go), this never omits an
|
||||
// entity for having nil limits — the dashboard shows usage for every
|
||||
// user and group, limited or not.
|
||||
// limiter.currentUsage reads every one's current-window counters — plus
|
||||
// the synthetic total scope's own (v0.2 data-layer task) — in ONE
|
||||
// storeGetMulti round trip total (v0.2 final review wave, 2026-08-20; see
|
||||
// limiter.currentUsage's own doc comment for the tradeoff this
|
||||
// supersedes) — users, groups, and the total scope are concatenated into
|
||||
// a single scopes slice before that one call, then the flat result is
|
||||
// sliced back into the three response sections at the same split points,
|
||||
// order preserved. Unlike buildLimitScopes (routes_unified.go), this
|
||||
// never omits an entity for having nil limits — the dashboard shows usage
|
||||
// for every user and group, limited or not.
|
||||
func (g *Gateway) buildAdminUsage() adminUsageResponse {
|
||||
userSummaries, groupSummaries := g.auth.snapshot()
|
||||
|
||||
scopes := make([]limitScope, 0, len(userSummaries)+len(groupSummaries))
|
||||
scopes := make([]limitScope, 0, len(userSummaries)+len(groupSummaries)+1)
|
||||
for _, us := range userSummaries {
|
||||
scopes = append(scopes, limitScope{kind: "user", id: us.name, limits: us.limits})
|
||||
}
|
||||
for _, gs := range groupSummaries {
|
||||
scopes = append(scopes, limitScope{kind: "group", id: gs.name, limits: gs.limits})
|
||||
}
|
||||
scopes = append(scopes, limitScope{kind: totalScopeKind, id: totalScopeID, limits: nil})
|
||||
|
||||
allUsage := g.limiter.currentUsage(scopes)
|
||||
userUsage, groupUsage := allUsage[:len(userSummaries)], allUsage[len(userSummaries):]
|
||||
userUsage := allUsage[:len(userSummaries)]
|
||||
groupUsage := allUsage[len(userSummaries) : len(userSummaries)+len(groupSummaries)]
|
||||
totalUsage := allUsage[len(userSummaries)+len(groupSummaries)]
|
||||
|
||||
users := make([]adminUsageEntryView, len(userUsage))
|
||||
for i, su := range userUsage {
|
||||
@@ -403,8 +435,9 @@ func (g *Gateway) buildAdminUsage() adminUsageResponse {
|
||||
for i, su := range groupUsage {
|
||||
groups[i] = usageEntryView(su, groupSummaries[i].limits)
|
||||
}
|
||||
total := usageEntryView(totalUsage, nil)
|
||||
|
||||
return adminUsageResponse{Users: users, Groups: groups}
|
||||
return adminUsageResponse{Users: users, Groups: groups, Total: total}
|
||||
}
|
||||
|
||||
// serveAdminUsage writes buildAdminUsage's result as JSON.
|
||||
@@ -413,6 +446,189 @@ func (g *Gateway) serveAdminUsage(w http.ResponseWriter) {
|
||||
_ = json.NewEncoder(w).Encode(g.buildAdminUsage())
|
||||
}
|
||||
|
||||
// historyMaxSpan caps GET /admin/api/usage/history's span query parameter
|
||||
// per window (v0.2 data-layer task), matching each window's own retention
|
||||
// (limits.go: hourWindowTTL/dayWindowTTL/monthWindowTTL) rounded down to
|
||||
// whole buckets: 48 hourly buckets, 35 daily buckets, 13 monthly buckets
|
||||
// (a one-month margin over monthWindowTTL's 400*24h ≈ 13.3 months). A
|
||||
// span beyond what its window's TTL could ever have kept alive would only
|
||||
// echo zeros for the missing tail, so the cap keeps every requested
|
||||
// bucket meaningful. window is assumed already validated by
|
||||
// validHistoryWindow — the default case is a programming error, not a
|
||||
// client input, mirroring windowKey/windowEnd's own panic convention.
|
||||
func historyMaxSpan(window string) int {
|
||||
switch window {
|
||||
case windowHour:
|
||||
return 48
|
||||
case windowDay:
|
||||
return 35
|
||||
case windowMonth:
|
||||
return 13
|
||||
default:
|
||||
panic(fmt.Sprintf("llmgateway: historyMaxSpan: unknown window %q", window))
|
||||
}
|
||||
}
|
||||
|
||||
// validHistoryMetric reports whether metric is one of GET
|
||||
// /admin/api/usage/history's four accepted metric values.
|
||||
func validHistoryMetric(metric string) bool {
|
||||
switch metric {
|
||||
case metricReq, metricTokIn, metricTokOut, metricCost:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// validHistoryWindow reports whether window is one of GET
|
||||
// /admin/api/usage/history's three accepted window values.
|
||||
func validHistoryWindow(window string) bool {
|
||||
switch window {
|
||||
case windowHour, windowDay, windowMonth:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// parseHistoryScope parses GET /admin/api/usage/history's "scope" query
|
||||
// parameter: "user:{id}", "group:{id}", or the literal "total". ok is
|
||||
// false for anything else, including a bare "total:{id}" form — the total
|
||||
// scope carries no id component of its own, it is always totalScopeID
|
||||
// (limits.go).
|
||||
func parseHistoryScope(raw string) (kind, id string, ok bool) {
|
||||
if raw == totalScopeKind {
|
||||
return totalScopeKind, totalScopeID, true
|
||||
}
|
||||
k, i, found := strings.Cut(raw, ":")
|
||||
if !found || i == "" {
|
||||
return "", "", false
|
||||
}
|
||||
if k != "user" && k != "group" {
|
||||
return "", "", false
|
||||
}
|
||||
return k, i, true
|
||||
}
|
||||
|
||||
// parseHistorySpan parses GET /admin/api/usage/history's "span" query
|
||||
// parameter for window: an empty raw value defaults to window's own
|
||||
// historyMaxSpan; otherwise raw must parse as an integer between 1 and
|
||||
// that max inclusive. ok is false for anything else — a non-integer, a
|
||||
// value below 1, or a value beyond window's max.
|
||||
func parseHistorySpan(raw, window string) (span int, ok bool) {
|
||||
max := historyMaxSpan(window)
|
||||
if raw == "" {
|
||||
return max, true
|
||||
}
|
||||
n, err := strconv.Atoi(raw)
|
||||
if err != nil || n < 1 || n > max {
|
||||
return 0, false
|
||||
}
|
||||
return n, true
|
||||
}
|
||||
|
||||
// scopeExists reports whether kind/id names a currently active user or a
|
||||
// configured group — GET /admin/api/usage/history's 404 check for an
|
||||
// unknown scope id. authStore.snapshot's own listing is the same
|
||||
// membership buildAdminUsage already trusts for "every currently active
|
||||
// user and every configured group"; kind is assumed already restricted to
|
||||
// "user" or "group" by parseHistoryScope (the total scope's id is never
|
||||
// checked against it — totalScopeID always exists).
|
||||
func (g *Gateway) scopeExists(kind, id string) bool {
|
||||
users, groups := g.auth.snapshot()
|
||||
switch kind {
|
||||
case "user":
|
||||
for _, u := range users {
|
||||
if u.name == id {
|
||||
return true
|
||||
}
|
||||
}
|
||||
case "group":
|
||||
for _, gr := range groups {
|
||||
if gr.name == id {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// usageHistoryPointView is one bucket in GET /admin/api/usage/history's
|
||||
// points array.
|
||||
type usageHistoryPointView struct {
|
||||
Bucket string `json:"bucket"`
|
||||
Value int64 `json:"value"`
|
||||
}
|
||||
|
||||
// usageHistoryResponse is the full body of GET /admin/api/usage/history
|
||||
// (v0.2 data-layer task): scope/metric/window echoed back exactly as
|
||||
// validated (scope in its raw "user:{id}"/"group:{id}"/"total" query
|
||||
// form), plus span buckets oldest-first — Points[len-1] is the current,
|
||||
// possibly partial, bucket.
|
||||
type usageHistoryResponse struct {
|
||||
Scope string `json:"scope"`
|
||||
Metric string `json:"metric"`
|
||||
Window string `json:"window"`
|
||||
Points []usageHistoryPointView `json:"points"`
|
||||
}
|
||||
|
||||
// serveAdminUsageHistory writes GET /admin/api/usage/history's bucketed
|
||||
// series for one scope/metric/window (v0.2 data-layer task, the Vue admin
|
||||
// panel's chart data source): "scope"=user:{id}|group:{id}|total,
|
||||
// "metric"=req|tokin|tokout|cost, "window"=hour|day|month, optional
|
||||
// "span"=N (default and max per historyMaxSpan). Validates every
|
||||
// parameter before touching the store — 400 for an unrecognized scope
|
||||
// kind/metric/window or an out-of-range span, 404 for a user/group id
|
||||
// that names no configured entity — then reads the series via
|
||||
// limiter.history, ONE storeGetMulti round trip for the whole span. A
|
||||
// storeDown read answers 503, not a silently-zero series: a chart built
|
||||
// on this API must never mistake "the store was unreachable" for "usage
|
||||
// was genuinely zero". checkAndCount is deliberately never called here —
|
||||
// see handleAdminAPI's own doc comment.
|
||||
func (g *Gateway) serveAdminUsageHistory(sw *statusTrackingWriter, r *http.Request) {
|
||||
q := r.URL.Query()
|
||||
|
||||
rawScope := q.Get("scope")
|
||||
kind, id, ok := parseHistoryScope(rawScope)
|
||||
if !ok {
|
||||
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "unknown scope")
|
||||
return
|
||||
}
|
||||
metric := q.Get("metric")
|
||||
if !validHistoryMetric(metric) {
|
||||
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "unknown metric")
|
||||
return
|
||||
}
|
||||
window := q.Get("window")
|
||||
if !validHistoryWindow(window) {
|
||||
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "unknown window")
|
||||
return
|
||||
}
|
||||
span, ok := parseHistorySpan(q.Get("span"), window)
|
||||
if !ok {
|
||||
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid span")
|
||||
return
|
||||
}
|
||||
if kind != totalScopeKind && !g.scopeExists(kind, id) {
|
||||
writeOAIError(sw, http.StatusNotFound, "invalid_request_error", "unknown "+kind+" id")
|
||||
return
|
||||
}
|
||||
|
||||
points, storeOK := g.limiter.history(kind, id, metric, window, g.limiter.now(), span)
|
||||
if !storeOK {
|
||||
writeOAIError(sw, http.StatusServiceUnavailable, "server_error", "usage history store unavailable")
|
||||
return
|
||||
}
|
||||
|
||||
view := make([]usageHistoryPointView, len(points))
|
||||
for i, p := range points {
|
||||
view[i] = usageHistoryPointView{Bucket: p.bucket, Value: p.value}
|
||||
}
|
||||
|
||||
setAdminJSONHeaders(sw)
|
||||
_ = json.NewEncoder(sw).Encode(usageHistoryResponse{Scope: rawScope, Metric: metric, Window: window, Points: view})
|
||||
}
|
||||
|
||||
// adminPageHTML is the entire /admin single-page dashboard: markup, CSS,
|
||||
// and vanilla JS in one Go raw-string const. Yaegi interpretation forbids
|
||||
// the embed directive (no filesystem access from an interpreted plugin),
|
||||
@@ -531,14 +747,14 @@ const adminPageHTML = `<!doctype html>
|
||||
<section>
|
||||
<h2>Groups</h2>
|
||||
<table id="groups"><thead><tr>
|
||||
<th>Name</th><th>Members</th><th>Limits</th><th>req/min</th><th>req/day</th><th>tok/day</th><th>tok/month</th><th>cost/day</th><th>cost/month</th>
|
||||
<th>Name</th><th>Members</th><th>Limits</th><th>req/min</th><th>req/day</th><th>tokIn/day</th><th>tokOut/day</th><th>tokIn/month</th><th>tokOut/month</th><th>cost/day</th><th>cost/month</th>
|
||||
</tr></thead><tbody></tbody></table>
|
||||
</section>
|
||||
|
||||
<section>
|
||||
<h2>Users</h2>
|
||||
<table id="users"><thead><tr>
|
||||
<th>Name</th><th>Group</th><th>Limits</th><th>req/min</th><th>req/day</th><th>tok/day</th><th>tok/month</th><th>cost/day</th><th>cost/month</th>
|
||||
<th>Name</th><th>Group</th><th>Limits</th><th>req/min</th><th>req/day</th><th>tokIn/day</th><th>tokOut/day</th><th>tokIn/month</th><th>tokOut/month</th><th>cost/day</th><th>cost/month</th>
|
||||
</tr></thead><tbody></tbody></table>
|
||||
</section>
|
||||
|
||||
@@ -660,8 +876,10 @@ const adminPageHTML = `<!doctype html>
|
||||
tr.appendChild(el("td", fmtLimits(entry.limits)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : String(entry.requestsPerMinute)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : String(entry.requestsPerDay)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : String(entry.tokensPerDay)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : String(entry.tokensPerMonth)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : String(entry.tokensInPerDay)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : String(entry.tokensOutPerDay)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : String(entry.tokensInPerMonth)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : String(entry.tokensOutPerMonth)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : fmtCost(entry.costPerDayMicroUsd)));
|
||||
tr.appendChild(el("td", entry.storeDown ? "?" : fmtCost(entry.costPerMonthMicroUsd)));
|
||||
if (entry.storeDown) tr.className = "err";
|
||||
|
||||
+291
-20
@@ -381,6 +381,12 @@ func TestAdminUsage_MathAgainstSeededCounters(t *testing.T) {
|
||||
gw.limiter.incrCounter("group", "agroup", metricReq, windowDay, fixedNow, 9, dayWindowTTL)
|
||||
gw.limiter.account([]limitScope{{kind: "group", id: "agroup"}}, usage{prompt: 100, completion: 50}, 9_000_000)
|
||||
|
||||
// Seed the synthetic total scope directly too (v0.2 data-layer task):
|
||||
// buildAdminUsage always includes it, independent of whether any
|
||||
// route's withTotalScope call ever ran in this test.
|
||||
gw.limiter.incrCounter(totalScopeKind, totalScopeID, metricReq, windowDay, fixedNow, 16, dayWindowTTL)
|
||||
gw.limiter.account([]limitScope{{kind: totalScopeKind, id: totalScopeID}}, usage{prompt: 140, completion: 60}, 11_500_000)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, adminUsagePath, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
@@ -396,8 +402,9 @@ func TestAdminUsage_MathAgainstSeededCounters(t *testing.T) {
|
||||
if alice.RequestsPerMinute != 3 || alice.RequestsPerDay != 7 {
|
||||
t.Errorf("alice requests = min:%d day:%d, want min:3 day:7", alice.RequestsPerMinute, alice.RequestsPerDay)
|
||||
}
|
||||
if alice.TokensPerDay != 50 || alice.TokensPerMonth != 50 {
|
||||
t.Errorf("alice tokens = day:%d month:%d, want 50/50", alice.TokensPerDay, alice.TokensPerMonth)
|
||||
if alice.TokensInPerDay != 40 || alice.TokensOutPerDay != 10 || alice.TokensInPerMonth != 40 || alice.TokensOutPerMonth != 10 {
|
||||
t.Errorf("alice tokens = in/day:%d out/day:%d in/month:%d out/month:%d, want 40/10/40/10",
|
||||
alice.TokensInPerDay, alice.TokensOutPerDay, alice.TokensInPerMonth, alice.TokensOutPerMonth)
|
||||
}
|
||||
if alice.CostPerDayMicroUSD != 2_500_000 || alice.CostPerMonthMicroUSD != 2_500_000 {
|
||||
t.Errorf("alice cost = day:%d month:%d, want 2500000/2500000", alice.CostPerDayMicroUSD, alice.CostPerMonthMicroUSD)
|
||||
@@ -416,8 +423,9 @@ func TestAdminUsage_MathAgainstSeededCounters(t *testing.T) {
|
||||
if agroup.RequestsPerMinute != 5 || agroup.RequestsPerDay != 9 {
|
||||
t.Errorf("agroup requests = min:%d day:%d, want min:5 day:9", agroup.RequestsPerMinute, agroup.RequestsPerDay)
|
||||
}
|
||||
if agroup.TokensPerDay != 150 || agroup.TokensPerMonth != 150 {
|
||||
t.Errorf("agroup tokens = day:%d month:%d, want 150/150", agroup.TokensPerDay, agroup.TokensPerMonth)
|
||||
if agroup.TokensInPerDay != 100 || agroup.TokensOutPerDay != 50 || agroup.TokensInPerMonth != 100 || agroup.TokensOutPerMonth != 50 {
|
||||
t.Errorf("agroup tokens = in/day:%d out/day:%d in/month:%d out/month:%d, want 100/50/100/50",
|
||||
agroup.TokensInPerDay, agroup.TokensOutPerDay, agroup.TokensInPerMonth, agroup.TokensOutPerMonth)
|
||||
}
|
||||
if agroup.CostPerDayMicroUSD != 9_000_000 {
|
||||
t.Errorf("agroup cost/day = %d, want 9000000", agroup.CostPerDayMicroUSD)
|
||||
@@ -435,6 +443,28 @@ func TestAdminUsage_MathAgainstSeededCounters(t *testing.T) {
|
||||
if admin1.StoreDown {
|
||||
t.Error("admin1.storeDown must be false")
|
||||
}
|
||||
|
||||
// The synthetic total scope row (v0.2 data-layer task): present,
|
||||
// carries no limit (limits.go's totalScopeKind is never evaluated),
|
||||
// and reflects exactly the counters seeded above.
|
||||
if got.Total.Kind != totalScopeKind || got.Total.ID != totalScopeID {
|
||||
t.Errorf("total kind/id = %q/%q, want %q/%q", got.Total.Kind, got.Total.ID, totalScopeKind, totalScopeID)
|
||||
}
|
||||
if got.Total.Limits != nil {
|
||||
t.Errorf("total limits = %+v, want nil (the total scope is never evaluated)", got.Total.Limits)
|
||||
}
|
||||
if got.Total.RequestsPerDay != 16 {
|
||||
t.Errorf("total requestsPerDay = %d, want 16", got.Total.RequestsPerDay)
|
||||
}
|
||||
if got.Total.TokensInPerDay != 140 || got.Total.TokensOutPerDay != 60 {
|
||||
t.Errorf("total tokens = in/day:%d out/day:%d, want 140/60", got.Total.TokensInPerDay, got.Total.TokensOutPerDay)
|
||||
}
|
||||
if got.Total.CostPerDayMicroUSD != 11_500_000 {
|
||||
t.Errorf("total cost/day = %d, want 11500000", got.Total.CostPerDayMicroUSD)
|
||||
}
|
||||
if got.Total.StoreDown {
|
||||
t.Error("total.storeDown must be false")
|
||||
}
|
||||
}
|
||||
|
||||
func findUsageEntry(t *testing.T, entries []adminUsageEntryView, id string) adminUsageEntryView {
|
||||
@@ -466,8 +496,9 @@ func TestLimiterCurrentUsage_StoreDown(t *testing.T) {
|
||||
if !su.storeDown {
|
||||
t.Fatal("want storeDown true when the configured store errors and failOpen is false")
|
||||
}
|
||||
if su.requestsPerMinute != 0 || su.requestsPerDay != 0 || su.tokensPerDay != 0 ||
|
||||
su.tokensPerMonth != 0 || su.costPerDayMicros != 0 || su.costPerMonthMicros != 0 {
|
||||
if su.requestsPerMinute != 0 || su.requestsPerDay != 0 || su.tokensInPerDay != 0 ||
|
||||
su.tokensInPerMonth != 0 || su.tokensOutPerDay != 0 || su.tokensOutPerMonth != 0 ||
|
||||
su.costPerDayMicros != 0 || su.costPerMonthMicros != 0 {
|
||||
t.Errorf("storeDown entry must report zero values, got %+v", su)
|
||||
}
|
||||
}
|
||||
@@ -698,11 +729,11 @@ func TestAdminOverview_BaseURLStripsCredentials(t *testing.T) {
|
||||
|
||||
// countingMultiStore is a counterStore stub whose getMulti serves a fixed
|
||||
// value per key and counts how many times it was called — it proves
|
||||
// limiter.currentUsage flattens every scope's six counter reads into ONE
|
||||
// getMulti call total (v0.2 final review wave, 2026-08-20), rather than
|
||||
// six separate get calls or even one getMulti call per scope.
|
||||
// incrBy/get are never exercised by currentUsage and just return zero
|
||||
// values.
|
||||
// limiter.currentUsage flattens every scope's usageKeysPerScope counter
|
||||
// reads into ONE getMulti call total (v0.2 final review wave,
|
||||
// 2026-08-20), rather than usageKeysPerScope separate get calls or even
|
||||
// one getMulti call per scope. incrBy/get are never exercised by
|
||||
// currentUsage and just return zero values.
|
||||
type countingMultiStore struct {
|
||||
values map[string]int64
|
||||
getMultiCalls int
|
||||
@@ -725,13 +756,15 @@ func TestLimiterCurrentUsage_BatchesOneGetMultiCallForAllScopes(t *testing.T) {
|
||||
t.Parallel()
|
||||
fixedNow := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
||||
store := &countingMultiStore{values: map[string]int64{
|
||||
windowKey("user", "alice", metricReq, windowMin, fixedNow): 3,
|
||||
windowKey("user", "alice", metricReq, windowDay, fixedNow): 7,
|
||||
windowKey("user", "alice", metricTok, windowDay, fixedNow): 50,
|
||||
windowKey("user", "alice", metricTok, windowMonth, fixedNow): 60,
|
||||
windowKey("user", "alice", metricCost, windowDay, fixedNow): 1_000,
|
||||
windowKey("user", "alice", metricCost, windowMonth, fixedNow): 2_000,
|
||||
windowKey("group", "g1", metricReq, windowMin, fixedNow): 1,
|
||||
windowKey("user", "alice", metricReq, windowMin, fixedNow): 3,
|
||||
windowKey("user", "alice", metricReq, windowDay, fixedNow): 7,
|
||||
windowKey("user", "alice", metricTokIn, windowDay, fixedNow): 50,
|
||||
windowKey("user", "alice", metricTokIn, windowMonth, fixedNow): 60,
|
||||
windowKey("user", "alice", metricTokOut, windowDay, fixedNow): 20,
|
||||
windowKey("user", "alice", metricTokOut, windowMonth, fixedNow): 25,
|
||||
windowKey("user", "alice", metricCost, windowDay, fixedNow): 1_000,
|
||||
windowKey("user", "alice", metricCost, windowMonth, fixedNow): 2_000,
|
||||
windowKey("group", "g1", metricReq, windowMin, fixedNow): 1,
|
||||
}}
|
||||
l := newLimiter(store, true)
|
||||
l.nowFn = func() time.Time { return fixedNow }
|
||||
@@ -752,8 +785,9 @@ func TestLimiterCurrentUsage_BatchesOneGetMultiCallForAllScopes(t *testing.T) {
|
||||
if alice.storeDown {
|
||||
t.Fatal("alice.storeDown must be false")
|
||||
}
|
||||
if alice.requestsPerMinute != 3 || alice.requestsPerDay != 7 || alice.tokensPerDay != 50 ||
|
||||
alice.tokensPerMonth != 60 || alice.costPerDayMicros != 1_000 || alice.costPerMonthMicros != 2_000 {
|
||||
if alice.requestsPerMinute != 3 || alice.requestsPerDay != 7 || alice.tokensInPerDay != 50 ||
|
||||
alice.tokensInPerMonth != 60 || alice.tokensOutPerDay != 20 || alice.tokensOutPerMonth != 25 ||
|
||||
alice.costPerDayMicros != 1_000 || alice.costPerMonthMicros != 2_000 {
|
||||
t.Errorf("alice usage = %+v, want the seeded values", alice)
|
||||
}
|
||||
g1 := got[1]
|
||||
@@ -839,3 +873,240 @@ func TestAdmin_FileUserAdminFlag_GrantsAccess(t *testing.T) {
|
||||
t.Fatalf("status = %d, want 200 for a file-sourced admin user, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// --- usage/history: gate, validation, shape, batching, storeDown (v0.2 data-layer task) ---
|
||||
|
||||
// TestAdminUsageHistory_GateMatrix mirrors TestAdmin_GateMatrix for the
|
||||
// third /admin/api/* route: unauthenticated → 401, authenticated
|
||||
// non-admin → 403, admin → 200.
|
||||
func TestAdminUsageHistory_GateMatrix(t *testing.T) {
|
||||
cfg := newAdminTestConfig()
|
||||
h, _ := newAdminGatewayHandle(t, cfg)
|
||||
query := adminUsageHistoryPath + "?scope=total&metric=req&window=hour"
|
||||
|
||||
t.Run("unauthenticated", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, query, ""))
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status = %d, want 401", rec.Code)
|
||||
}
|
||||
assertOAIErrorEnvelope(t, rec)
|
||||
})
|
||||
t.Run("authenticated_non_admin", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, query, "sk-alice"))
|
||||
if rec.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want 403", rec.Code)
|
||||
}
|
||||
assertOAIErrorEnvelope(t, rec)
|
||||
})
|
||||
t.Run("admin", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, query, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestAdminUsageHistory_DoesNotCountStats proves GET
|
||||
// /admin/api/usage/history skips checkAndCount entirely (unlike
|
||||
// overview/usage — handleAdminAPI's own doc comment, admin.go): repeated
|
||||
// requests never move admin1's own req/min counter, and the route stays
|
||||
// reachable even once admin1 is already pinned at a requestsPerMinute
|
||||
// limit tight enough to 429 on the other two admin routes.
|
||||
func TestAdminUsageHistory_DoesNotCountStats(t *testing.T) {
|
||||
cfg := newAdminTestConfig()
|
||||
cfg.Users.Inline[1].Limits = &LimitsConfig{RequestsPerMinute: 1} // admin1
|
||||
h, gw := newAdminGatewayHandle(t, cfg)
|
||||
|
||||
query := adminUsageHistoryPath + "?scope=total&metric=req&window=hour"
|
||||
for i := 0; i < 5; i++ {
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, query, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("request %d: status = %d, want 200 (never 429), body=%s", i, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
reqCount, ok := gw.limiter.getCounter("user", "admin1", metricReq, windowMin, time.Now())
|
||||
if !ok || reqCount != 0 {
|
||||
t.Errorf("admin1 req/min counter = %d (ok=%v), want 0 (history route never calls checkAndCount)", reqCount, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdminUsageHistory_ValidationTable drives GET
|
||||
// /admin/api/usage/history's query-parameter validation end to end: an
|
||||
// unrecognized scope kind/metric/window or an out-of-range span is 400,
|
||||
// an unknown user/group id is 404, and every well-formed combination
|
||||
// (including span pinned exactly at its window's max) is 200.
|
||||
func TestAdminUsageHistory_ValidationTable(t *testing.T) {
|
||||
cfg := newAdminTestConfig()
|
||||
h, _ := newAdminGatewayHandle(t, cfg)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
query string
|
||||
want int
|
||||
}{
|
||||
{"unknown scope kind", "scope=nope&metric=req&window=hour", http.StatusBadRequest},
|
||||
{"total scope rejects an id suffix", "scope=total:all&metric=req&window=hour", http.StatusBadRequest},
|
||||
{"scope missing a colon", "scope=user&metric=req&window=hour", http.StatusBadRequest},
|
||||
{"scope with an empty id", "scope=user:&metric=req&window=hour", http.StatusBadRequest},
|
||||
{"unknown metric", "scope=total&metric=nope&window=hour", http.StatusBadRequest},
|
||||
{"unknown window", "scope=total&metric=req&window=nope", http.StatusBadRequest},
|
||||
{"span not an integer", "scope=total&metric=req&window=hour&span=abc", http.StatusBadRequest},
|
||||
{"span zero", "scope=total&metric=req&window=hour&span=0", http.StatusBadRequest},
|
||||
{"span negative", "scope=total&metric=req&window=hour&span=-1", http.StatusBadRequest},
|
||||
{"span beyond hour max (48)", "scope=total&metric=req&window=hour&span=49", http.StatusBadRequest},
|
||||
{"span beyond day max (35)", "scope=total&metric=req&window=day&span=36", http.StatusBadRequest},
|
||||
{"span beyond month max (13)", "scope=total&metric=req&window=month&span=14", http.StatusBadRequest},
|
||||
{"span pinned exactly at the hour max is valid", "scope=total&metric=req&window=hour&span=48", http.StatusOK},
|
||||
{"unknown user id", "scope=user:ghost&metric=req&window=hour", http.StatusNotFound},
|
||||
{"unknown group id", "scope=group:ghost&metric=req&window=hour", http.StatusNotFound},
|
||||
{"valid total/req/hour, default span", "scope=total&metric=req&window=hour", http.StatusOK},
|
||||
{"valid user/tokin/day", "scope=user:alice&metric=tokin&window=day", http.StatusOK},
|
||||
{"valid group/tokout/month", "scope=group:agroup&metric=tokout&window=month", http.StatusOK},
|
||||
{"valid cost metric", "scope=total&metric=cost&window=day", http.StatusOK},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, adminUsageHistoryPath+"?"+c.query, "sk-admin1"))
|
||||
if rec.Code != c.want {
|
||||
t.Errorf("status = %d, want %d, body=%s", rec.Code, c.want, rec.Body.String())
|
||||
}
|
||||
if c.want >= 400 {
|
||||
assertOAIErrorEnvelope(t, rec)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdminUsageHistory_DefaultSpanIsMax asserts an omitted span defaults
|
||||
// to each window's own historyMaxSpan.
|
||||
func TestAdminUsageHistory_DefaultSpanIsMax(t *testing.T) {
|
||||
cfg := newAdminTestConfig()
|
||||
h, _ := newAdminGatewayHandle(t, cfg)
|
||||
|
||||
cases := []struct {
|
||||
window string
|
||||
want int
|
||||
}{
|
||||
{windowHour, 48},
|
||||
{windowDay, 35},
|
||||
{windowMonth, 13},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.window, func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, adminUsageHistoryPath+"?scope=total&metric=req&window="+c.window, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var got usageHistoryResponse
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if len(got.Points) != c.want {
|
||||
t.Errorf("len(points) = %d, want %d (default span = max)", len(got.Points), c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdminUsageHistory_ShapeOldestFirstOneBatchCall drives the route
|
||||
// against a limiter whose store is swapped for a counting stub: it must
|
||||
// decode to the documented {"scope","metric","window","points":[...]}
|
||||
// shape, points must be oldest-first with the current (possibly partial)
|
||||
// bucket last, and the whole span must cost exactly ONE getMulti call —
|
||||
// never one call per bucket.
|
||||
func TestAdminUsageHistory_ShapeOldestFirstOneBatchCall(t *testing.T) {
|
||||
cfg := newAdminTestConfig()
|
||||
h, gw := newAdminGatewayHandle(t, cfg)
|
||||
|
||||
fixedNow := time.Date(2026, 8, 20, 14, 0, 0, 0, time.UTC)
|
||||
countStore := &countingMultiStore{values: map[string]int64{
|
||||
windowKey(totalScopeKind, totalScopeID, metricReq, windowHour, fixedNow.Add(-2*time.Hour)): 5,
|
||||
windowKey(totalScopeKind, totalScopeID, metricReq, windowHour, fixedNow.Add(-1*time.Hour)): 7,
|
||||
windowKey(totalScopeKind, totalScopeID, metricReq, windowHour, fixedNow): 9,
|
||||
}}
|
||||
gw.limiter = newLimiter(countStore, true)
|
||||
gw.limiter.nowFn = func() time.Time { return fixedNow }
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
query := adminUsageHistoryPath + "?scope=total&metric=req&window=hour&span=3"
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, query, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if countStore.getMultiCalls != 1 {
|
||||
t.Errorf("getMultiCalls = %d, want 1 (one batch for the whole span)", countStore.getMultiCalls)
|
||||
}
|
||||
|
||||
var got usageHistoryResponse
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if got.Scope != "total" || got.Metric != "req" || got.Window != "hour" {
|
||||
t.Errorf("echoed scope/metric/window = %q/%q/%q, want total/req/hour", got.Scope, got.Metric, got.Window)
|
||||
}
|
||||
if len(got.Points) != 3 {
|
||||
t.Fatalf("len(points) = %d, want 3", len(got.Points))
|
||||
}
|
||||
wantValues := []int64{5, 7, 9}
|
||||
for i, want := range wantValues {
|
||||
if got.Points[i].Value != want {
|
||||
t.Errorf("points[%d].value = %d, want %d (oldest-first)", i, got.Points[i].Value, want)
|
||||
}
|
||||
}
|
||||
if got.Points[2].Bucket != bucketFor(fixedNow, windowHour) {
|
||||
t.Errorf("points[2].bucket = %q, want the current bucket %q", got.Points[2].Bucket, bucketFor(fixedNow, windowHour))
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdminUsageHistory_StoreDown503 asserts a fail-closed store error
|
||||
// answers 503 with the standard OAI-shaped error envelope, never a
|
||||
// silently-zero series — data unavailable beats a chart that lies about
|
||||
// zero usage.
|
||||
func TestAdminUsageHistory_StoreDown503(t *testing.T) {
|
||||
cfg := newAdminTestConfig()
|
||||
failOpen := false
|
||||
cfg.Redis = &RedisConfig{Address: "127.0.0.1:1", FailOpen: &failOpen}
|
||||
h, _ := newAdminGatewayHandle(t, cfg)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
query := adminUsageHistoryPath + "?scope=total&metric=req&window=hour"
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, query, "sk-admin1"))
|
||||
if rec.Code != http.StatusServiceUnavailable {
|
||||
t.Fatalf("status = %d, want 503, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
assertOAIErrorEnvelope(t, rec)
|
||||
}
|
||||
|
||||
// TestAdminUsageHistory_MemoryStoreFallbackWorks proves the history API
|
||||
// serves real data through the in-process memoryStore fallback too — a
|
||||
// deployment with no Redis configured at all still gets charts of
|
||||
// whatever the process itself accumulated (spec's "memoryStore works
|
||||
// too" requirement).
|
||||
func TestAdminUsageHistory_MemoryStoreFallbackWorks(t *testing.T) {
|
||||
cfg := newAdminTestConfig() // cfg.Redis left nil: limiter.store is nil, every op uses the fallback
|
||||
h, gw := newAdminGatewayHandle(t, cfg)
|
||||
|
||||
fixedNow := time.Date(2026, 8, 20, 9, 0, 0, 0, time.UTC)
|
||||
gw.limiter.nowFn = func() time.Time { return fixedNow }
|
||||
gw.limiter.account([]limitScope{{kind: totalScopeKind, id: totalScopeID}}, usage{prompt: 12, completion: 4}, 0)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
query := adminUsageHistoryPath + "?scope=total&metric=tokin&window=hour&span=1"
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, query, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var got usageHistoryResponse
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if len(got.Points) != 1 || got.Points[0].Value != 12 {
|
||||
t.Errorf("points = %+v, want one point with value 12", got.Points)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -757,6 +757,29 @@ func adminUsageEntry(t *testing.T, kind, id string) map[string]any {
|
||||
return nil
|
||||
}
|
||||
|
||||
// adminTotalTokensPerDay sums an admin usage entry's tokensInPerDay and
|
||||
// tokensOutPerDay (v0.2 data-layer task: the JSON API split what was
|
||||
// previously a single combined tokensPerDay field into in/out — see
|
||||
// adminUsageEntryView, admin.go).
|
||||
func adminTotalTokensPerDay(entry map[string]any) float64 {
|
||||
in, _ := entry["tokensInPerDay"].(float64)
|
||||
out, _ := entry["tokensOutPerDay"].(float64)
|
||||
return in + out
|
||||
}
|
||||
|
||||
// adminUsageHistory issues GET /admin/api/usage/history as adminKey for
|
||||
// scope/metric/window and returns the decoded response body — the v0.2
|
||||
// data-layer task's third admin API.
|
||||
func adminUsageHistory(t *testing.T, scope, metric, window string) map[string]any {
|
||||
t.Helper()
|
||||
url := traefik1URL + "/admin/api/usage/history?scope=" + scope + "&metric=" + metric + "&window=" + window
|
||||
resp, body := doJSON(t, http.MethodGet, url, adminKey, nil)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("GET /admin/api/usage/history (scope=%s metric=%s window=%s): status = %d, body=%#v", scope, metric, window, resp.StatusCode, body)
|
||||
}
|
||||
return body
|
||||
}
|
||||
|
||||
// TestResponseCache covers v0.2 integration coverage: two identical chat
|
||||
// completions get X-Llmgw-Cache: miss then hit, and the admin usage API
|
||||
// (spec §4) proves the hit added zero tokens — only the miss's real
|
||||
@@ -784,9 +807,9 @@ func TestResponseCache(t *testing.T) {
|
||||
t.Error("redis holds no llmgw:cache:* key after a cache miss, want the entry the miss's SET (cache.go) should have written")
|
||||
}
|
||||
|
||||
tokensAfterMiss, _ := adminUsageEntry(t, "users", "alice")["tokensPerDay"].(float64)
|
||||
tokensAfterMiss := adminTotalTokensPerDay(adminUsageEntry(t, "users", "alice"))
|
||||
if tokensAfterMiss <= 0 {
|
||||
t.Fatalf("tokensPerDay after the miss = %v, want > 0 (the mock's real usage should have been accounted)", tokensAfterMiss)
|
||||
t.Fatalf("tokensInPerDay+tokensOutPerDay after the miss = %v, want > 0 (the mock's real usage should have been accounted)", tokensAfterMiss)
|
||||
}
|
||||
|
||||
resp2, body2 := doJSON(t, http.MethodPost, traefik1URL+"/v1/chat/completions", aliceKey, reqBody)
|
||||
@@ -800,9 +823,9 @@ func TestResponseCache(t *testing.T) {
|
||||
t.Errorf("second request: content = %q, want the cached %q", content, "mock openai response")
|
||||
}
|
||||
|
||||
tokensAfterHit, _ := adminUsageEntry(t, "users", "alice")["tokensPerDay"].(float64)
|
||||
tokensAfterHit := adminTotalTokensPerDay(adminUsageEntry(t, "users", "alice"))
|
||||
if tokensAfterHit != tokensAfterMiss {
|
||||
t.Errorf("tokensPerDay changed across the cache hit: after miss = %v, after hit = %v, want unchanged", tokensAfterMiss, tokensAfterHit)
|
||||
t.Errorf("tokensInPerDay+tokensOutPerDay changed across the cache hit: after miss = %v, after hit = %v, want unchanged", tokensAfterMiss, tokensAfterHit)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -978,4 +1001,27 @@ func TestAdminDashboard(t *testing.T) {
|
||||
if reqDay, _ := adminUsage["requestsPerDay"].(float64); reqDay <= 0 {
|
||||
t.Errorf("admin requestsPerDay = %v, want > 0 (this test's own admin API calls count too)", reqDay)
|
||||
}
|
||||
|
||||
// v0.2 data-layer task: GET /admin/api/usage/history's hour series for
|
||||
// alice and for the synthetic total scope both must show a non-zero
|
||||
// current-hour bucket, reflecting the real LLM traffic this suite
|
||||
// already generated (TestRetryFlakyRecovers, TestResponseCache, ...)
|
||||
// against a real Redis-backed limiter, not just a unit-test stub.
|
||||
for _, scope := range []string{"user:alice", "total"} {
|
||||
hist := adminUsageHistory(t, scope, "req", "hour")
|
||||
if got, _ := hist["scope"].(string); got != scope {
|
||||
t.Errorf("history scope=%q: echoed scope = %q, want %q", scope, got, scope)
|
||||
}
|
||||
points, ok := hist["points"].([]any)
|
||||
if !ok || len(points) == 0 {
|
||||
t.Fatalf("history scope=%q: points = %#v, want a non-empty array", scope, hist["points"])
|
||||
}
|
||||
current, ok := points[len(points)-1].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("history scope=%q: last point is not an object, got %#v", scope, points[len(points)-1])
|
||||
}
|
||||
if val, _ := current["value"].(float64); val <= 0 {
|
||||
t.Errorf("history scope=%q: current-hour bucket value = %v, want > 0 (this suite's own traffic)", scope, val)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -27,18 +27,30 @@ const storeErrorLogEvery = 30 * time.Second
|
||||
const storeDownLatchFor = 5 * time.Second
|
||||
|
||||
// Window identifiers used throughout windowKey, counter metrics, and
|
||||
// retry-after calculation.
|
||||
// retry-after calculation. windowHour is stats-only (v0.2 data-layer
|
||||
// task, see hourWindowTTL): checkAndCount/account both write it, but no
|
||||
// LimitsConfig field ever names it, so evaluateScope has nothing to
|
||||
// evaluate it against — GET /admin/api/usage/history (admin.go) is its
|
||||
// only reader.
|
||||
const (
|
||||
windowHour = "hour"
|
||||
windowMin = "min"
|
||||
windowDay = "day"
|
||||
windowMonth = "month"
|
||||
)
|
||||
|
||||
// Counter metric names embedded in windowKey.
|
||||
// Counter metric names embedded in windowKey. metricTokIn/metricTokOut
|
||||
// replace a single combined "tok" metric (v0.2 data-layer task, folding
|
||||
// the operator's tokens-in/tokens-out split directive): account records
|
||||
// usage.prompt under metricTokIn and usage.completion under metricTokOut
|
||||
// separately, so the admin usage/history APIs can chart each direction on
|
||||
// its own. A TokensPerDay/TokensPerMonth limit still enforces a TOTAL
|
||||
// budget across both — see tokenBudgetViolation.
|
||||
const (
|
||||
metricReq = "req"
|
||||
metricTok = "tok"
|
||||
metricCost = "cost"
|
||||
metricReq = "req"
|
||||
metricTokIn = "tokin"
|
||||
metricTokOut = "tokout"
|
||||
metricCost = "cost"
|
||||
)
|
||||
|
||||
// TTLs applied to counter keys. They exceed their window's natural length
|
||||
@@ -46,10 +58,18 @@ const (
|
||||
// embedded in the key (see windowKey) already makes a rolled-over window
|
||||
// use a different key, so these TTLs only bound how long a stale key
|
||||
// lingers in memoryStore before an opportunistic sweep reclaims it.
|
||||
// dayWindowTTL/monthWindowTTL were bumped from their original 25h/32d
|
||||
// (v0.2 data-layer task): the admin usage-history API's day/month charts
|
||||
// need a bucket to stay readable for the whole retention span a chart
|
||||
// can request (historyMaxSpan, admin.go — 35 daily, 13 monthly buckets),
|
||||
// not just long enough to survive its own single rollover. The TTL >
|
||||
// window-length invariant this comment already documented still holds at
|
||||
// the new values, and at hourWindowTTL for the new hour window.
|
||||
const (
|
||||
minWindowTTL = 2 * time.Minute
|
||||
dayWindowTTL = 25 * time.Hour
|
||||
monthWindowTTL = 32 * 24 * time.Hour
|
||||
hourWindowTTL = 48 * time.Hour
|
||||
dayWindowTTL = 35 * 24 * time.Hour
|
||||
monthWindowTTL = 400 * 24 * time.Hour
|
||||
)
|
||||
|
||||
// usdToMicroFactor scales a USD amount to micro-USD (1e6 micro-USD per
|
||||
@@ -248,16 +268,31 @@ func (u usage) total() int64 {
|
||||
return u.prompt + u.completion
|
||||
}
|
||||
|
||||
// limitScope is one entity (a user or their group) whose limits apply to a
|
||||
// request. checkAndCount and account evaluate every scope in the slice
|
||||
// they are given, so a request is counted and checked against both its
|
||||
// user's and its group's limits in one call.
|
||||
// limitScope is one entity (a user, their group, or the synthetic total
|
||||
// scope — see totalScopeKind) whose limits apply to a request.
|
||||
// checkAndCount and account evaluate every scope in the slice they are
|
||||
// given, so a request is counted and checked against both its user's and
|
||||
// its group's limits in one call.
|
||||
type limitScope struct {
|
||||
limits *LimitsConfig
|
||||
kind string // "user" or "group"
|
||||
kind string // "user", "group", or "total"
|
||||
id string
|
||||
}
|
||||
|
||||
// totalScopeKind and totalScopeID name the synthetic "all LLM traffic"
|
||||
// scope withTotalScope (routes_unified.go) appends to every metered
|
||||
// route's scopes slice (v0.2 data-layer task): {kind: totalScopeKind, id:
|
||||
// totalScopeID}, limits always nil. evaluateScope's nil-limits check
|
||||
// (below) skips it during evaluation, so checkAndCount still counts it
|
||||
// like any other scope — its req/tokin/tokout/cost counters accumulate
|
||||
// every LLM request across all users and groups combined, the admin
|
||||
// usage-history "total" series — but no configuration can ever throttle
|
||||
// traffic in its name.
|
||||
const (
|
||||
totalScopeKind = "total"
|
||||
totalScopeID = "all"
|
||||
)
|
||||
|
||||
// limitViolation describes the first limit a request breached.
|
||||
type limitViolation struct {
|
||||
message string
|
||||
@@ -479,36 +514,46 @@ func storeDownViolation() *limitViolation {
|
||||
return &limitViolation{message: "limit store unavailable", storeDown: true}
|
||||
}
|
||||
|
||||
// windowKey builds the counterStore key for one (kind, id, metric, window)
|
||||
// counter at time t: llmgw:{kind}:{id}:{metric}:{window}:{bucket}. bucket
|
||||
// is t.UTC() formatted to the window's granularity, so a counter's key
|
||||
// changes automatically when its window rolls over.
|
||||
func windowKey(kind, id, metric, window string, t time.Time) string {
|
||||
// bucketFor formats t (in UTC) to window's bucket granularity: the string
|
||||
// windowKey embeds as a key's final component, and the same string GET
|
||||
// /admin/api/usage/history echoes as each point's "bucket" field
|
||||
// (historyBucketKeys, below).
|
||||
func bucketFor(t time.Time, window string) string {
|
||||
u := t.UTC()
|
||||
var bucket string
|
||||
switch window {
|
||||
case windowHour:
|
||||
return u.Format("2006010215")
|
||||
case windowMin:
|
||||
bucket = u.Format("200601021504")
|
||||
return u.Format("200601021504")
|
||||
case windowDay:
|
||||
bucket = u.Format("20060102")
|
||||
return u.Format("20060102")
|
||||
case windowMonth:
|
||||
bucket = u.Format("200601")
|
||||
return u.Format("200601")
|
||||
default:
|
||||
// window is always one of the three constants above, set by this
|
||||
// window is always one of the four constants above, set by this
|
||||
// file's own callers, never by request input — an unknown value
|
||||
// here is a programming error. ServeHTTP's recoverPanic turns
|
||||
// this into a logged 500 instead of crashing the process.
|
||||
panic(fmt.Sprintf("llmgateway: windowKey: unknown window %q", window))
|
||||
panic(fmt.Sprintf("llmgateway: bucketFor: unknown window %q", window))
|
||||
}
|
||||
return fmt.Sprintf("llmgw:%s:%s:%s:%s:%s", kind, id, metric, window, bucket)
|
||||
}
|
||||
|
||||
// windowKey builds the counterStore key for one (kind, id, metric, window)
|
||||
// counter at time t: llmgw:{kind}:{id}:{metric}:{window}:{bucket}. bucket
|
||||
// is t.UTC() formatted to the window's granularity (bucketFor), so a
|
||||
// counter's key changes automatically when its window rolls over.
|
||||
func windowKey(kind, id, metric, window string, t time.Time) string {
|
||||
return fmt.Sprintf("llmgw:%s:%s:%s:%s:%s", kind, id, metric, window, bucketFor(t, window))
|
||||
}
|
||||
|
||||
// windowEnd returns the UTC instant at which window's bucket containing t
|
||||
// ends: the next minute boundary, the next UTC midnight, or the first of
|
||||
// next month UTC.
|
||||
// ends: the next hour or minute boundary, the next UTC midnight, or the
|
||||
// first of next month UTC.
|
||||
func windowEnd(t time.Time, window string) time.Time {
|
||||
u := t.UTC()
|
||||
switch window {
|
||||
case windowHour:
|
||||
return time.Date(u.Year(), u.Month(), u.Day(), u.Hour(), 0, 0, 0, time.UTC).Add(time.Hour)
|
||||
case windowMin:
|
||||
return time.Date(u.Year(), u.Month(), u.Day(), u.Hour(), u.Minute(), 0, 0, time.UTC).Add(time.Minute)
|
||||
case windowDay:
|
||||
@@ -575,6 +620,14 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
|
||||
return storeDownViolation()
|
||||
}
|
||||
dayCounts[i] = v
|
||||
|
||||
// Hour is stats-only (windowHour's own doc comment): its value is
|
||||
// never read back here, only written, so the admin usage-history
|
||||
// API (admin.go) has an hour-granularity req series to read
|
||||
// later.
|
||||
if _, ok = l.incrCounter(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL); !ok {
|
||||
return storeDownViolation()
|
||||
}
|
||||
}
|
||||
|
||||
for i, sc := range scopes {
|
||||
@@ -607,10 +660,10 @@ func (l *limiter) evaluateScope(sc limitScope, minCount, dayCount int64, now tim
|
||||
if v := requestLimitViolation(sc, "requests-per-day", lim.RequestsPerDay, dayCount, windowDay, now); v != nil {
|
||||
return v
|
||||
}
|
||||
if v := l.budgetViolation(sc, metricTok, "tokens-per-day", lim.TokensPerDay, windowDay, now); v != nil {
|
||||
if v := l.tokenBudgetViolation(sc, "tokens-per-day", lim.TokensPerDay, windowDay, now); v != nil {
|
||||
return v
|
||||
}
|
||||
if v := l.budgetViolation(sc, metricTok, "tokens-per-month", lim.TokensPerMonth, windowMonth, now); v != nil {
|
||||
if v := l.tokenBudgetViolation(sc, "tokens-per-month", lim.TokensPerMonth, windowMonth, now); v != nil {
|
||||
return v
|
||||
}
|
||||
if v := l.budgetViolation(sc, metricCost, "cost-per-day", usdToMicros(lim.CostPerDayUSD), windowDay, now); v != nil {
|
||||
@@ -657,84 +710,145 @@ func (l *limiter) budgetViolation(sc limitScope, metric, name string, limit int6
|
||||
}
|
||||
}
|
||||
|
||||
// account records u's total tokens and costMicros against every scope's
|
||||
// day and month counters. A zero-valued metric is skipped entirely — no
|
||||
// store write for a metric this call has nothing to report. A sample that
|
||||
// tokenBudgetViolation mirrors budgetViolation for a token limit, which —
|
||||
// unlike a request or cost limit — enforces a TOTAL budget over two split
|
||||
// counters (metricTokIn, metricTokOut; v0.2 data-layer task): it reads
|
||||
// both of the scope's accumulated tok-in/tok-out counters for window in
|
||||
// ONE storeGetMulti round trip and compares their sum against limit.
|
||||
// limit<=0 means unlimited. A fail-closed read returns
|
||||
// storeDownViolation, matching budgetViolation's own contract; name is
|
||||
// still e.g. "tokens-per-day", so the violation message still says
|
||||
// "tokens" regardless of which direction pushed the total over.
|
||||
func (l *limiter) tokenBudgetViolation(sc limitScope, name string, limit int64, window string, now time.Time) *limitViolation {
|
||||
if limit <= 0 {
|
||||
return nil
|
||||
}
|
||||
keys := []string{
|
||||
windowKey(sc.kind, sc.id, metricTokIn, window, now),
|
||||
windowKey(sc.kind, sc.id, metricTokOut, window, now),
|
||||
}
|
||||
vals, ok := l.storeGetMulti(keys)
|
||||
if !ok || len(vals) != 2 {
|
||||
return storeDownViolation()
|
||||
}
|
||||
used := vals[0] + vals[1]
|
||||
if used < limit {
|
||||
return nil
|
||||
}
|
||||
return &limitViolation{
|
||||
message: fmt.Sprintf("%s %q exceeded %s budget", sc.kind, sc.id, name),
|
||||
retryAfter: retryAfterSeconds(now, window),
|
||||
}
|
||||
}
|
||||
|
||||
// account records u's prompt/completion tokens — split into metricTokIn/
|
||||
// metricTokOut (v0.2 data-layer task) — and costMicros against every
|
||||
// scope's hour, day, and month counters. Each metric/direction is skipped
|
||||
// independently when its own value is 0 — no store write for a direction
|
||||
// this call has nothing to report, so a streaming response cut off before
|
||||
// any completion tokens arrived still writes tokin alone. A sample that
|
||||
// hits a fail-closed store error (see storeIncrBy) is dropped silently:
|
||||
// account has no error return to signal it, and dropping — rather than
|
||||
// counting it in the fallback — avoids double-counting once the store
|
||||
// recovers.
|
||||
func (l *limiter) account(scopes []limitScope, u usage, costMicros int64) {
|
||||
now := l.now()
|
||||
total := u.total()
|
||||
|
||||
for _, sc := range scopes {
|
||||
if total != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricTok, windowDay, now, total, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTok, windowMonth, now, total, monthWindowTTL)
|
||||
if u.prompt != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricTokIn, windowHour, now, u.prompt, hourWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTokIn, windowDay, now, u.prompt, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTokIn, windowMonth, now, u.prompt, monthWindowTTL)
|
||||
}
|
||||
if u.completion != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricTokOut, windowHour, now, u.completion, hourWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTokOut, windowDay, now, u.completion, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTokOut, windowMonth, now, u.completion, monthWindowTTL)
|
||||
}
|
||||
if costMicros != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricCost, windowHour, now, costMicros, hourWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricCost, windowDay, now, costMicros, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricCost, windowMonth, now, costMicros, monthWindowTTL)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// scopeUsage is one scope's (a user's or a group's) current-window
|
||||
// counter values — the admin dashboard's usage table (spec §4, v0.2). The
|
||||
// six value fields mirror LimitsConfig's six limit kinds exactly, so the
|
||||
// admin API can echo a usage value beside its matching limit.
|
||||
// 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
|
||||
// v0.2 final review wave follow-up (2026-08-20) that flagged the previous
|
||||
// literal "6" as a magic-number stride: this task grows the per-scope key
|
||||
// count from 6 to 8 for the tokens-in/tokens-out split, which is exactly
|
||||
// the moment a silent off-by-N here would have gone unnoticed.
|
||||
const usageKeysPerScope = 8
|
||||
|
||||
// scopeUsage is one scope's (a user's, a group's, or the synthetic total
|
||||
// scope's — totalScopeKind) current-window counter values — the admin
|
||||
// dashboard's usage table (spec §4, v0.2) and GET /admin/api/usage's
|
||||
// total row (v0.2 data-layer task). tokensInPerDay/tokensOutPerDay and
|
||||
// their per-month counterparts replace a single combined
|
||||
// tokensPerDay/tokensPerMonth pair (v0.2 data-layer task, tokens-in/
|
||||
// tokens-out split): a caller wanting the combined total a
|
||||
// LimitsConfig.TokensPerDay/TokensPerMonth limit is actually evaluated
|
||||
// against (tokenBudgetViolation) sums the two itself.
|
||||
type scopeUsage struct {
|
||||
kind string // "user" or "group", mirroring limitScope.kind
|
||||
kind string // "user", "group", or "total", mirroring limitScope.kind
|
||||
id string
|
||||
requestsPerMinute int64
|
||||
requestsPerDay int64
|
||||
tokensPerDay int64
|
||||
tokensPerMonth int64
|
||||
tokensInPerDay int64
|
||||
tokensInPerMonth int64
|
||||
tokensOutPerDay int64
|
||||
tokensOutPerMonth int64
|
||||
costPerDayMicros int64
|
||||
costPerMonthMicros int64
|
||||
// storeDown reports whether reading any of the six counters above
|
||||
// failed closed (a configured store errored and failOpen is false) —
|
||||
// mirrors limitViolation.storeDown. Every value field is 0 in that
|
||||
// case; a caller must not present them as "confirmed zero usage".
|
||||
// storeDown reports whether reading any of the usageKeysPerScope
|
||||
// counters above failed closed (a configured store errored and
|
||||
// failOpen is false) — mirrors limitViolation.storeDown. Every value
|
||||
// field is 0 in that case; a caller must not present them as
|
||||
// "confirmed zero usage".
|
||||
storeDown bool
|
||||
}
|
||||
|
||||
// usageWindowKeys returns the six windowKey strings currentUsage reads
|
||||
// for sc at time now, in a fixed order — req/min, req/day, tok/day,
|
||||
// tok/month, cost/day, cost/month — matching scopeUsage's field order
|
||||
// exactly, so currentUsage can map storeGetMulti's result slice back to
|
||||
// named fields by plain index.
|
||||
// usageWindowKeys returns the usageKeysPerScope windowKey strings
|
||||
// currentUsage reads for sc at time now, in a fixed order — req/min,
|
||||
// req/day, tokin/day, tokin/month, tokout/day, tokout/month, cost/day,
|
||||
// cost/month — matching scopeUsage's field order exactly, so currentUsage
|
||||
// can map storeGetMulti's result slice back to named fields by plain
|
||||
// index.
|
||||
func usageWindowKeys(sc limitScope, now time.Time) []string {
|
||||
return []string{
|
||||
windowKey(sc.kind, sc.id, metricReq, windowMin, now),
|
||||
windowKey(sc.kind, sc.id, metricReq, windowDay, now),
|
||||
windowKey(sc.kind, sc.id, metricTok, windowDay, now),
|
||||
windowKey(sc.kind, sc.id, metricTok, windowMonth, now),
|
||||
windowKey(sc.kind, sc.id, metricTokIn, windowDay, now),
|
||||
windowKey(sc.kind, sc.id, metricTokIn, windowMonth, now),
|
||||
windowKey(sc.kind, sc.id, metricTokOut, windowDay, now),
|
||||
windowKey(sc.kind, sc.id, metricTokOut, windowMonth, now),
|
||||
windowKey(sc.kind, sc.id, metricCost, windowDay, now),
|
||||
windowKey(sc.kind, sc.id, metricCost, windowMonth, now),
|
||||
}
|
||||
}
|
||||
|
||||
// currentUsage reads every scope's six current-window counters — the
|
||||
// same (kind, id, metric, window) combinations checkAndCount/account
|
||||
// already write — via the limiter's own storeGetMulti, so it applies the
|
||||
// identical fail-open/fail-closed policy every enforcement read already
|
||||
// does, and it is read-only: unlike checkAndCount, it never increments
|
||||
// anything.
|
||||
// currentUsage reads every scope's usageKeysPerScope current-window
|
||||
// counters — the same (kind, id, metric, window) combinations
|
||||
// checkAndCount/account already write — via the limiter's own
|
||||
// storeGetMulti, so it applies the identical fail-open/fail-closed policy
|
||||
// every enforcement read already does, and it is read-only: unlike
|
||||
// checkAndCount, it never increments anything.
|
||||
//
|
||||
// ONE storeGetMulti call for the whole scopes slice (v0.2 final review
|
||||
// wave, 2026-08-20; supersedes the "one call per scope" amendment this
|
||||
// comment previously described, 2026-08-20 review): every scope's six
|
||||
// keys are flattened into a single round trip against the shared store,
|
||||
// comment previously described, 2026-08-20 review): every scope's keys
|
||||
// are flattened into a single round trip against the shared store,
|
||||
// rather than N separate ones. The reviewer-confirmed reason for this
|
||||
// second amendment: GET /admin/api/usage's per-scope loop was contending
|
||||
// with request admission on the same store connection as live traffic,
|
||||
// and N round trips (even pipelined 6-at-a-time) scales with the number
|
||||
// of configured users and groups in a way one round trip does not.
|
||||
// buildAdminUsage (admin.go) drives this by concatenating its users and
|
||||
// groups into one scopes slice before calling in, then slicing the flat
|
||||
// result back into its two response sections at the same split point.
|
||||
// and N round trips (even pipelined usageKeysPerScope-at-a-time) scales
|
||||
// with the number of configured users and groups in a way one round trip
|
||||
// does not. buildAdminUsage (admin.go) drives this by concatenating its
|
||||
// users, groups, and the total scope into one scopes slice before calling
|
||||
// in, then slicing the flat result back into its response sections at the
|
||||
// same split points.
|
||||
//
|
||||
// Tradeoff given up by this second amendment: a transient store error
|
||||
// now marks every scope in the call storeDown together, where the
|
||||
@@ -751,7 +865,7 @@ func (l *limiter) currentUsage(scopes []limitScope) []scopeUsage {
|
||||
}
|
||||
|
||||
now := l.now()
|
||||
allKeys := make([]string, 0, len(scopes)*6)
|
||||
allKeys := make([]string, 0, len(scopes)*usageKeysPerScope)
|
||||
for _, sc := range scopes {
|
||||
allKeys = append(allKeys, usageWindowKeys(sc, now)...)
|
||||
}
|
||||
@@ -771,13 +885,86 @@ func (l *limiter) currentUsage(scopes []limitScope) []scopeUsage {
|
||||
}
|
||||
|
||||
for i, sc := range scopes {
|
||||
v := vals[i*6 : i*6+6]
|
||||
v := vals[i*usageKeysPerScope : i*usageKeysPerScope+usageKeysPerScope]
|
||||
out[i] = scopeUsage{
|
||||
kind: sc.kind, id: sc.id,
|
||||
requestsPerMinute: v[0], requestsPerDay: v[1],
|
||||
tokensPerDay: v[2], tokensPerMonth: v[3],
|
||||
costPerDayMicros: v[4], costPerMonthMicros: v[5],
|
||||
tokensInPerDay: v[2], tokensInPerMonth: v[3],
|
||||
tokensOutPerDay: v[4], tokensOutPerMonth: v[5],
|
||||
costPerDayMicros: v[6], costPerMonthMicros: v[7],
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// historyPoint is one bucket's counter value in a limiter.history result:
|
||||
// bucket is the same string bucketFor produces for that instant (e.g.
|
||||
// "2026082114" for an hour bucket), value is the counter's reading — 0
|
||||
// for a bucket that was never written, matching get/getMulti's existing
|
||||
// "0 for a missing key" contract.
|
||||
type historyPoint struct {
|
||||
bucket string
|
||||
value int64
|
||||
}
|
||||
|
||||
// historyStepBack returns the instant window's bucket was current i steps
|
||||
// before now: i=0 is now's own (current, possibly partial) bucket, i=1
|
||||
// the previous one, and so on. hour and day step by a fixed duration/day
|
||||
// count from now itself; month anchors on the 1st of now's own month
|
||||
// first, then steps whole calendar months from there —
|
||||
// time.Time.AddDate normalizes an overflowing day-of-month (e.g. "Mar 31"
|
||||
// minus one month becomes "Mar 3", not "Feb 28/29"), so stepping from day
|
||||
// 1 (which exists in every month) is the only way whole-month arithmetic
|
||||
// stays exact regardless of what day of the month now is.
|
||||
func historyStepBack(now time.Time, window string, i int) time.Time {
|
||||
u := now.UTC()
|
||||
switch window {
|
||||
case windowHour:
|
||||
return u.Add(-time.Duration(i) * time.Hour)
|
||||
case windowDay:
|
||||
return u.AddDate(0, 0, -i)
|
||||
case windowMonth:
|
||||
anchor := time.Date(u.Year(), u.Month(), 1, 0, 0, 0, 0, time.UTC)
|
||||
return anchor.AddDate(0, -i, 0)
|
||||
default:
|
||||
panic(fmt.Sprintf("llmgateway: historyStepBack: unknown window %q", window))
|
||||
}
|
||||
}
|
||||
|
||||
// historyBucketKeys returns the span counterStore keys and their matching
|
||||
// bucket label strings for (kind, id, metric, window), stepping back from
|
||||
// now one whole window unit at a time — oldest first, the current
|
||||
// (possibly partial) bucket last, matching GET /admin/api/usage/history's
|
||||
// "inclusive of current bucket, oldest-first" contract.
|
||||
func historyBucketKeys(kind, id, metric, window string, now time.Time, span int) (keys, buckets []string) {
|
||||
keys = make([]string, span)
|
||||
buckets = make([]string, span)
|
||||
for i := 0; i < span; i++ {
|
||||
t := historyStepBack(now, window, span-1-i)
|
||||
keys[i] = windowKey(kind, id, metric, window, t)
|
||||
buckets[i] = bucketFor(t, window)
|
||||
}
|
||||
return keys, buckets
|
||||
}
|
||||
|
||||
// history returns span counter values for (kind, id, metric, window)
|
||||
// ending at now, oldest-first — GET /admin/api/usage/history's data
|
||||
// source (v0.2 data-layer task). It computes every bucket key up front
|
||||
// and reads them in ONE storeGetMulti round trip, the same single-batch
|
||||
// discipline currentUsage already applies to a scope's current-window
|
||||
// keys, scaled here to a whole span of one metric/window instead of a
|
||||
// fixed set of usageKeysPerScope. ok is false only in the fail-closed
|
||||
// case (storeGetMulti's own contract): a caller must not present the
|
||||
// returned points as real data then.
|
||||
func (l *limiter) history(kind, id, metric, window string, now time.Time, span int) ([]historyPoint, bool) {
|
||||
keys, buckets := historyBucketKeys(kind, id, metric, window, now, span)
|
||||
vals, ok := l.storeGetMulti(keys)
|
||||
if !ok || len(vals) != len(keys) {
|
||||
return nil, false
|
||||
}
|
||||
points := make([]historyPoint, span)
|
||||
for i := range points {
|
||||
points[i] = historyPoint{bucket: buckets[i], value: vals[i]}
|
||||
}
|
||||
return points, true
|
||||
}
|
||||
|
||||
+336
-10
@@ -168,13 +168,15 @@ func TestLimiterRequestLimitRetryAfter(t *testing.T) {
|
||||
|
||||
// TestWindowKeyFormat pins the exact key layout every store operation
|
||||
// depends on: llmgw:{kind}:{id}:{metric}:{window}:{bucket}, with bucket
|
||||
// formatted from t.UTC() at each window's granularity.
|
||||
// formatted from t.UTC() at each window's granularity — including hour
|
||||
// (v0.2 data-layer task).
|
||||
func TestWindowKeyFormat(t *testing.T) {
|
||||
tm := time.Date(2026, 8, 20, 10, 4, 59, 0, time.FixedZone("UTC+1", 3600)) // 09:04:59 UTC
|
||||
cases := []struct {
|
||||
window string
|
||||
want string
|
||||
}{
|
||||
{windowHour, "llmgw:user:a:req:hour:2026082009"},
|
||||
{windowMin, "llmgw:user:a:req:min:202608200904"},
|
||||
{windowDay, "llmgw:user:a:req:day:20260820"},
|
||||
{windowMonth, "llmgw:user:a:req:month:202608"},
|
||||
@@ -186,6 +188,270 @@ func TestWindowKeyFormat(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestWindowKeyFormat_HourUTCBoundary pins the hour bucket rolling over
|
||||
// exactly at the UTC hour boundary, and proves a non-UTC input is
|
||||
// converted to UTC first (an input at 00:04:59 in a UTC+1 zone is 23:04:59
|
||||
// UTC the PREVIOUS day — a bug converting the hour field alone without
|
||||
// re-deriving the date would produce "...23" glued onto the wrong day).
|
||||
func TestWindowKeyFormat_HourUTCBoundary(t *testing.T) {
|
||||
justBefore := time.Date(2026, 8, 20, 8, 59, 59, 0, time.UTC)
|
||||
justAfter := time.Date(2026, 8, 20, 9, 0, 0, 0, time.UTC)
|
||||
if got := windowKey("user", "a", "req", windowHour, justBefore); got != "llmgw:user:a:req:hour:2026082008" {
|
||||
t.Errorf("just before the hour boundary: got %q", got)
|
||||
}
|
||||
if got := windowKey("user", "a", "req", windowHour, justAfter); got != "llmgw:user:a:req:hour:2026082009" {
|
||||
t.Errorf("just after the hour boundary: got %q", got)
|
||||
}
|
||||
|
||||
// 00:04:59 in UTC+1 is 23:04:59 UTC the previous day.
|
||||
crossMidnight := time.Date(2026, 8, 21, 0, 4, 59, 0, time.FixedZone("UTC+1", 3600))
|
||||
if got := windowKey("user", "a", "req", windowHour, crossMidnight); got != "llmgw:user:a:req:hour:2026082023" {
|
||||
t.Errorf("cross-midnight non-UTC input: got %q, want the previous UTC day's hour 23 bucket", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWindowEnd_HourBoundary asserts windowEnd's new hour case (v0.2
|
||||
// data-layer task) returns the next UTC hour boundary, and that a value
|
||||
// exactly on the boundary reports zero seconds remaining in that boundary
|
||||
// (retryAfterSeconds rounds a zero duration up to 0, distinct from
|
||||
// rounding a fractional second up to 1 — see TestLimiterRetryAfter for
|
||||
// the day/month/min cases this mirrors).
|
||||
func TestWindowEnd_HourBoundary(t *testing.T) {
|
||||
t.Run("mid-hour", func(t *testing.T) {
|
||||
got := windowEnd(time.Date(2026, 8, 20, 9, 30, 0, 0, time.UTC), windowHour)
|
||||
want := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
||||
if !got.Equal(want) {
|
||||
t.Errorf("windowEnd = %v, want %v", got, want)
|
||||
}
|
||||
})
|
||||
t.Run("exactly on the boundary", func(t *testing.T) {
|
||||
got := windowEnd(time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC), windowHour)
|
||||
want := time.Date(2026, 8, 20, 11, 0, 0, 0, time.UTC)
|
||||
if !got.Equal(want) {
|
||||
t.Errorf("windowEnd = %v, want %v", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestHistoryStepBack_HourAndDay asserts the fixed-duration step-back
|
||||
// cases (hour, day) land i units before now exactly, including across a
|
||||
// UTC day boundary.
|
||||
func TestHistoryStepBack_HourAndDay(t *testing.T) {
|
||||
now := time.Date(2026, 8, 20, 1, 30, 0, 0, time.UTC)
|
||||
|
||||
if got := historyStepBack(now, windowHour, 0); !got.Equal(now) {
|
||||
t.Errorf("i=0 must be now itself, got %v", got)
|
||||
}
|
||||
wantHour := time.Date(2026, 8, 19, 23, 30, 0, 0, time.UTC) // 2 hours back, crossing midnight
|
||||
if got := historyStepBack(now, windowHour, 2); !got.Equal(wantHour) {
|
||||
t.Errorf("i=2 hours back = %v, want %v", got, wantHour)
|
||||
}
|
||||
|
||||
wantDay := time.Date(2026, 8, 17, 1, 30, 0, 0, time.UTC)
|
||||
if got := historyStepBack(now, windowDay, 3); !got.Equal(wantDay) {
|
||||
t.Errorf("i=3 days back = %v, want %v", got, wantDay)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHistoryStepBack_MonthAnchorsOnDayOne is the month-arithmetic trap
|
||||
// TestLimiterRetryAfter's own month case does not exercise: stepping back
|
||||
// from a high day-of-month (the 31st) with a naive
|
||||
// now.AddDate(0, -i, 0) overflows a shorter target month (e.g. "Mar 31"
|
||||
// minus one month normalizes to "Mar 3", not "Feb 28") and would land in
|
||||
// the WRONG calendar month's bucket. historyStepBack avoids this by
|
||||
// anchoring on day 1 first — this test drives it from Mar 31 and Jan 31
|
||||
// specifically because those are the inputs that would expose a
|
||||
// regression back to the naive form.
|
||||
func TestHistoryStepBack_MonthAnchorsOnDayOne(t *testing.T) {
|
||||
cases := []struct {
|
||||
now time.Time
|
||||
name string
|
||||
want string
|
||||
i int
|
||||
}{
|
||||
{name: "Mar 31 minus 1 month = Feb", now: time.Date(2026, 3, 31, 12, 0, 0, 0, time.UTC), i: 1, want: "202602"},
|
||||
{name: "Jan 31 minus 1 month = Dec (year rollover)", now: time.Date(2026, 1, 31, 0, 0, 0, 0, time.UTC), i: 1, want: "202512"},
|
||||
{name: "May 31 minus 3 months = Feb", now: time.Date(2026, 5, 31, 0, 0, 0, 0, time.UTC), i: 3, want: "202602"},
|
||||
{name: "i=0 is now's own month", now: time.Date(2026, 3, 31, 0, 0, 0, 0, time.UTC), i: 0, want: "202603"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
got := bucketFor(historyStepBack(c.now, windowMonth, c.i), windowMonth)
|
||||
if got != c.want {
|
||||
t.Errorf("bucketFor(historyStepBack(%v, month, %d)) = %q, want %q", c.now, c.i, got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestHistoryBucketKeys_OldestFirstInclusiveOfCurrent asserts
|
||||
// historyBucketKeys returns span buckets oldest-first, with the LAST
|
||||
// entry being now's own (current, possibly partial) bucket — the
|
||||
// "inclusive of current bucket" contract GET /admin/api/usage/history
|
||||
// promises.
|
||||
func TestHistoryBucketKeys_OldestFirstInclusiveOfCurrent(t *testing.T) {
|
||||
now := time.Date(2026, 8, 20, 14, 0, 0, 0, time.UTC)
|
||||
keys, buckets := historyBucketKeys("total", "all", "req", windowHour, now, 4)
|
||||
|
||||
wantBuckets := []string{"2026082011", "2026082012", "2026082013", "2026082014"}
|
||||
if len(buckets) != len(wantBuckets) {
|
||||
t.Fatalf("len(buckets) = %d, want %d", len(buckets), len(wantBuckets))
|
||||
}
|
||||
for i, want := range wantBuckets {
|
||||
if buckets[i] != want {
|
||||
t.Errorf("buckets[%d] = %q, want %q", i, buckets[i], want)
|
||||
}
|
||||
wantKey := "llmgw:total:all:req:hour:" + want
|
||||
if keys[i] != wantKey {
|
||||
t.Errorf("keys[%d] = %q, want %q", i, keys[i], wantKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestLimiterHistory_OneBatchCallOldestFirst drives limiter.history
|
||||
// against a counting stub store: it must issue exactly ONE getMulti call
|
||||
// for the whole span, and the returned points must be oldest-first with
|
||||
// the seeded values in the right buckets.
|
||||
func TestLimiterHistory_OneBatchCallOldestFirst(t *testing.T) {
|
||||
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
||||
countStore := &historyCountingStore{values: map[string]int64{
|
||||
windowKey("total", "all", metricReq, windowHour, now.Add(-2*time.Hour)): 5,
|
||||
windowKey("total", "all", metricReq, windowHour, now): 9,
|
||||
}}
|
||||
l := newLimiter(countStore, true)
|
||||
l.nowFn = func() time.Time { return now }
|
||||
|
||||
points, ok := l.history("total", "all", "req", windowHour, now, 3)
|
||||
if !ok {
|
||||
t.Fatal("want ok=true")
|
||||
}
|
||||
if countStore.getMultiCalls != 1 {
|
||||
t.Errorf("getMultiCalls = %d, want 1", countStore.getMultiCalls)
|
||||
}
|
||||
if len(points) != 3 {
|
||||
t.Fatalf("len(points) = %d, want 3", len(points))
|
||||
}
|
||||
if points[0].value != 5 || points[0].bucket != bucketFor(now.Add(-2*time.Hour), windowHour) {
|
||||
t.Errorf("points[0] (oldest) = %+v, want value 5 at the -2h bucket", points[0])
|
||||
}
|
||||
if points[1].value != 0 {
|
||||
t.Errorf("points[1] (unseeded -1h bucket) = %+v, want value 0", points[1])
|
||||
}
|
||||
if points[2].value != 9 || points[2].bucket != bucketFor(now, windowHour) {
|
||||
t.Errorf("points[2] (newest, current bucket) = %+v, want value 9", points[2])
|
||||
}
|
||||
}
|
||||
|
||||
// historyCountingStore is a counterStore stub whose getMulti serves fixed
|
||||
// values and counts its own calls — mirrors admin_test.go's
|
||||
// countingMultiStore, defined separately here since limits_test.go must
|
||||
// not depend on admin_test.go's test-only types.
|
||||
type historyCountingStore struct {
|
||||
values map[string]int64
|
||||
getMultiCalls int
|
||||
}
|
||||
|
||||
func (s *historyCountingStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, nil }
|
||||
func (s *historyCountingStore) get(string) (int64, error) { return 0, nil }
|
||||
func (s *historyCountingStore) getMulti(keys []string) ([]int64, error) {
|
||||
s.getMultiCalls++
|
||||
out := make([]int64, len(keys))
|
||||
for i, k := range keys {
|
||||
out[i] = s.values[k]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// TestLimiterHistory_StoreDown asserts limiter.history reports ok=false
|
||||
// when the configured store errors and failOpen is false — the caller
|
||||
// (GET /admin/api/usage/history) must answer 503, never a silently-zero
|
||||
// series.
|
||||
func TestLimiterHistory_StoreDown(t *testing.T) {
|
||||
l := newLimiter(alwaysErrStore{}, false)
|
||||
points, ok := l.history("total", "all", "req", windowDay, time.Now(), 5)
|
||||
if ok {
|
||||
t.Fatal("want ok=false when the configured store errors and failOpen is false")
|
||||
}
|
||||
if points != nil {
|
||||
t.Errorf("points = %+v, want nil on a storeDown read", points)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWithTotalScope_AppendsUnlimitedTotalScope asserts withTotalScope
|
||||
// appends exactly one {kind: totalScopeKind, id: totalScopeID, limits:
|
||||
// nil} scope after whatever buildLimitScopes already returned.
|
||||
func TestWithTotalScope_AppendsUnlimitedTotalScope(t *testing.T) {
|
||||
base := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 1}}}
|
||||
got := withTotalScope(base)
|
||||
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("len(got) = %d, want 2", len(got))
|
||||
}
|
||||
if got[0].kind != "user" || got[0].id != "u" {
|
||||
t.Errorf("got[0] = %+v, want the original user scope untouched", got[0])
|
||||
}
|
||||
total := got[1]
|
||||
if total.kind != totalScopeKind || total.id != totalScopeID || total.limits != nil {
|
||||
t.Errorf("got[1] = %+v, want kind=%q id=%q limits=nil", total, totalScopeKind, totalScopeID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWithTotalScope_CountedButNeverEvaluated drives checkAndCount many
|
||||
// times against a scopes slice built via withTotalScope: the total
|
||||
// scope's own limits are always nil, so no volume of traffic can ever
|
||||
// make it violate — evaluateScope's nil-limits check skips it — while its
|
||||
// req counter still climbs by exactly one per call, proving
|
||||
// checkAndCount's counting loop does not skip a nil-limits scope the way
|
||||
// evaluation does.
|
||||
func TestWithTotalScope_CountedButNeverEvaluated(t *testing.T) {
|
||||
l := newLimiter(nil, true)
|
||||
now := time.Date(2026, 8, 20, 12, 0, 0, 0, time.UTC)
|
||||
l.nowFn = func() time.Time { return now }
|
||||
|
||||
scopes := withTotalScope([]limitScope{{kind: "user", id: "u", limits: nil}})
|
||||
const n = 250
|
||||
for i := 0; i < n; i++ {
|
||||
if v := l.checkAndCount(scopes); v != nil {
|
||||
t.Fatalf("iteration %d: total scope must never violate (limits always nil), got %+v", i, v)
|
||||
}
|
||||
}
|
||||
|
||||
total, ok := l.getCounter(totalScopeKind, totalScopeID, metricReq, windowDay, now)
|
||||
if !ok || total != n {
|
||||
t.Errorf("total req:day = %d, ok=%v, want %d (counted every call despite never evaluating)", total, ok, n)
|
||||
}
|
||||
totalHour, ok := l.getCounter(totalScopeKind, totalScopeID, metricReq, windowHour, now)
|
||||
if !ok || totalHour != n {
|
||||
t.Errorf("total req:hour = %d, ok=%v, want %d", totalHour, ok, n)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWithTotalScope_AccountWritesHourBuckets asserts account (v0.2
|
||||
// data-layer task) writes the total scope's hour-window tokin/tokout/cost
|
||||
// counters, not just day/month — checkAndCount's own hour write is
|
||||
// covered by TestWithTotalScope_CountedButNeverEvaluated above.
|
||||
func TestWithTotalScope_AccountWritesHourBuckets(t *testing.T) {
|
||||
l := newLimiter(nil, true)
|
||||
now := time.Date(2026, 8, 20, 12, 0, 0, 0, time.UTC)
|
||||
l.nowFn = func() time.Time { return now }
|
||||
|
||||
scopes := withTotalScope([]limitScope{{kind: "user", id: "u", limits: nil}})
|
||||
l.account(scopes, usage{prompt: 4, completion: 6}, 700)
|
||||
|
||||
tokIn, ok := l.getCounter(totalScopeKind, totalScopeID, metricTokIn, windowHour, now)
|
||||
if !ok || tokIn != 4 {
|
||||
t.Errorf("total tokin:hour = %d, ok=%v, want 4", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := l.getCounter(totalScopeKind, totalScopeID, metricTokOut, windowHour, now)
|
||||
if !ok || tokOut != 6 {
|
||||
t.Errorf("total tokout:hour = %d, ok=%v, want 6", tokOut, ok)
|
||||
}
|
||||
cost, ok := l.getCounter(totalScopeKind, totalScopeID, metricCost, windowHour, now)
|
||||
if !ok || cost != 700 {
|
||||
t.Errorf("total cost:hour = %d, ok=%v, want 700", cost, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMemoryStoreIncrAndGet exercises the fallback store directly: fresh
|
||||
// keys start at n, repeated incrBy accumulates, and get on an unknown key
|
||||
// returns 0 with no error.
|
||||
@@ -303,9 +569,10 @@ func TestLimiterCostBudgetRoundsUSDConversion(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccountSkipsZeroMetrics asserts account writes no tokens counter
|
||||
// when usage totals 0, and no cost counter when costMicros is 0 — each
|
||||
// metric's absence is independently observable via get.
|
||||
// TestAccountSkipsZeroMetrics asserts account writes no tokin/tokout
|
||||
// counter when the matching usage direction is 0, and no cost counter
|
||||
// when costMicros is 0 — each metric's absence is independently
|
||||
// observable via get.
|
||||
func TestAccountSkipsZeroMetrics(t *testing.T) {
|
||||
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
||||
l := newLimiter(nil, true)
|
||||
@@ -314,9 +581,13 @@ func TestAccountSkipsZeroMetrics(t *testing.T) {
|
||||
|
||||
l.account(scopes, usage{}, 0)
|
||||
|
||||
tok, ok := l.getCounter("user", "z", metricTok, windowDay, now)
|
||||
if !ok || tok != 0 {
|
||||
t.Errorf("tok counter = %d, ok=%v, want 0, true", tok, ok)
|
||||
tokIn, ok := l.getCounter("user", "z", metricTokIn, windowDay, now)
|
||||
if !ok || tokIn != 0 {
|
||||
t.Errorf("tokin counter = %d, ok=%v, want 0, true", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := l.getCounter("user", "z", metricTokOut, windowDay, now)
|
||||
if !ok || tokOut != 0 {
|
||||
t.Errorf("tokout counter = %d, ok=%v, want 0, true", tokOut, ok)
|
||||
}
|
||||
cost, ok := l.getCounter("user", "z", metricCost, windowDay, now)
|
||||
if !ok || cost != 0 {
|
||||
@@ -324,6 +595,54 @@ func TestAccountSkipsZeroMetrics(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccountSplitsTokensIndependently asserts account writes prompt
|
||||
// tokens to metricTokIn and completion tokens to metricTokOut
|
||||
// independently — including the asymmetric case (prompt only, no
|
||||
// completion yet) the estimation fallback and a dropped stream both rely
|
||||
// on: a nonzero prompt with a zero completion must still write tokin
|
||||
// alone, never skip it because completion is 0.
|
||||
func TestAccountSplitsTokensIndependently(t *testing.T) {
|
||||
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
||||
l := newLimiter(nil, true)
|
||||
l.nowFn = func() time.Time { return now }
|
||||
scopes := []limitScope{{kind: "user", id: "split", limits: nil}}
|
||||
|
||||
l.account(scopes, usage{prompt: 25}, 0) // completion left at zero
|
||||
|
||||
tokIn, ok := l.getCounter("user", "split", metricTokIn, windowDay, now)
|
||||
if !ok || tokIn != 25 {
|
||||
t.Errorf("tokin counter = %d, ok=%v, want 25, true", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := l.getCounter("user", "split", metricTokOut, windowDay, now)
|
||||
if !ok || tokOut != 0 {
|
||||
t.Errorf("tokout counter = %d, ok=%v, want 0, true (no completion tokens reported)", tokOut, ok)
|
||||
}
|
||||
|
||||
l.account(scopes, usage{completion: 9}, 0) // prompt left at zero this time
|
||||
tokIn, ok = l.getCounter("user", "split", metricTokIn, windowDay, now)
|
||||
if !ok || tokIn != 25 {
|
||||
t.Errorf("tokin counter = %d, ok=%v, want still 25 (this call reported no prompt tokens)", tokIn, ok)
|
||||
}
|
||||
tokOut, ok = l.getCounter("user", "split", metricTokOut, windowDay, now)
|
||||
if !ok || tokOut != 9 {
|
||||
t.Errorf("tokout counter = %d, ok=%v, want 9", tokOut, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenBudgetSumsInAndOut asserts a TokensPerDay/TokensPerMonth limit
|
||||
// enforces a TOTAL budget across metricTokIn and metricTokOut combined —
|
||||
// the operator directive's "limits stay total-token (limiter sums both at
|
||||
// read)": neither direction alone crosses the budget, but their sum does.
|
||||
func TestTokenBudgetSumsInAndOut(t *testing.T) {
|
||||
l := newLimiter(nil, true)
|
||||
scopes := []limitScope{{kind: "group", id: "g", limits: &LimitsConfig{TokensPerDay: 100}}}
|
||||
|
||||
l.account(scopes, usage{prompt: 60, completion: 45}, 0) // 105 > 100, split 60/45
|
||||
if v := l.checkAndCount(scopes); v == nil {
|
||||
t.Fatal("60 tokin + 45 tokout = 105 > 100 must refuse, even though neither direction alone exceeds 100")
|
||||
}
|
||||
}
|
||||
|
||||
// TestLimiterRequestPerDayLimit is carried-item (b): a per-day request
|
||||
// limit blocks once the day's count exceeds it, independent of the
|
||||
// (unset, unlimited) per-minute limit.
|
||||
@@ -491,12 +810,19 @@ func TestLimiter_FailClosed_AccountDropsSampleOnStoreError(t *testing.T) {
|
||||
|
||||
l.account(scopes, usage{prompt: 10, completion: 10}, 500)
|
||||
|
||||
got, err := l.fallback.get(windowKey("user", "u", metricTok, windowDay, l.now()))
|
||||
gotIn, err := l.fallback.get(windowKey("user", "u", metricTokIn, windowDay, l.now()))
|
||||
if err != nil {
|
||||
t.Fatalf("fallback.get: %v", err)
|
||||
}
|
||||
if got != 0 {
|
||||
t.Errorf("fallback tok:day counter = %d, want 0 (sample must be dropped, not counted locally)", got)
|
||||
if gotIn != 0 {
|
||||
t.Errorf("fallback tokin:day counter = %d, want 0 (sample must be dropped, not counted locally)", gotIn)
|
||||
}
|
||||
gotOut, err := l.fallback.get(windowKey("user", "u", metricTokOut, windowDay, l.now()))
|
||||
if err != nil {
|
||||
t.Fatalf("fallback.get: %v", err)
|
||||
}
|
||||
if gotOut != 0 {
|
||||
t.Errorf("fallback tokout:day counter = %d, want 0 (sample must be dropped, not counted locally)", gotOut)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+1
-1
@@ -166,7 +166,7 @@ func (g *Gateway) handleTargetProxy(w http.ResponseWriter, r *http.Request, u *u
|
||||
return
|
||||
}
|
||||
|
||||
scopes := buildLimitScopes(u, grp)
|
||||
scopes := withTotalScope(buildLimitScopes(u, grp))
|
||||
if violation := g.limiter.checkAndCount(scopes); violation != nil {
|
||||
writeLimitViolation(w, violation)
|
||||
return
|
||||
|
||||
+15
-3
@@ -660,9 +660,13 @@ func TestHandleTargetProxy_JSONResponse_OnlyRequestCounterMoves(t *testing.T) {
|
||||
t.Errorf("body = %q, want verbatim upstream body", rec.Body.String())
|
||||
}
|
||||
|
||||
tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||
if !ok || tok != 0 {
|
||||
t.Errorf("user token/day counter = %d (ok=%v), want 0 — target proxy must never account response usage", tok, ok)
|
||||
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
if !ok || tokIn != 0 {
|
||||
t.Errorf("user tokin/day counter = %d (ok=%v), want 0 — target proxy must never account response usage", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
||||
if !ok || tokOut != 0 {
|
||||
t.Errorf("user tokout/day counter = %d (ok=%v), want 0 — target proxy must never account response usage", tokOut, ok)
|
||||
}
|
||||
cost, ok := gw.limiter.getCounter("user", "alice", metricCost, windowDay, time.Now())
|
||||
if !ok || cost != 0 {
|
||||
@@ -672,6 +676,14 @@ func TestHandleTargetProxy_JSONResponse_OnlyRequestCounterMoves(t *testing.T) {
|
||||
if !ok || reqCount != 1 {
|
||||
t.Errorf("user request/min counter = %d (ok=%v), want 1 — the request itself is still accounted", reqCount, ok)
|
||||
}
|
||||
|
||||
// withTotalScope (handleTargetProxy, mcp_a2a.go) must have appended
|
||||
// the synthetic total scope too (v0.2 data-layer task): counted like
|
||||
// any other scope, still never token/cost-accounted here.
|
||||
totalReq, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricReq, windowMin, time.Now())
|
||||
if !ok || totalReq != 1 {
|
||||
t.Errorf("total request/min counter = %d (ok=%v), want 1", totalReq, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// --- config validation at construction ---
|
||||
|
||||
+1
-1
@@ -69,7 +69,7 @@ func (g *Gateway) resolveMediaRequest(sw *statusTrackingWriter, u *user, grp *gr
|
||||
writeModelResolveError(sw, err)
|
||||
return nil, "", false
|
||||
}
|
||||
if violation := g.limiter.checkAndCount(buildLimitScopes(u, grp)); violation != nil {
|
||||
if violation := g.limiter.checkAndCount(withTotalScope(buildLimitScopes(u, grp))); violation != nil {
|
||||
writeLimitViolation(sw, violation)
|
||||
return nil, "", false
|
||||
}
|
||||
|
||||
+12
-2
@@ -164,8 +164,18 @@ func TestHandleImagesGenerations_OpenAI_HappyPath_NativeForward(t *testing.T) {
|
||||
if !ok || reqCount != 1 {
|
||||
t.Errorf("user request/min counter = %d (ok=%v), want 1", reqCount, ok)
|
||||
}
|
||||
if tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now()); ok && tok != 0 {
|
||||
t.Errorf("user token/day counter = %d, want 0 (images are never token-accounted)", tok)
|
||||
if tokIn, okIn := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now()); okIn && tokIn != 0 {
|
||||
t.Errorf("user tokin/day counter = %d, want 0 (images are never token-accounted)", tokIn)
|
||||
}
|
||||
if tokOut, okOut := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now()); okOut && tokOut != 0 {
|
||||
t.Errorf("user tokout/day counter = %d, want 0 (images are never token-accounted)", tokOut)
|
||||
}
|
||||
|
||||
// withTotalScope (resolveMediaRequest, routes_media.go) must have
|
||||
// appended the synthetic total scope too (v0.2 data-layer task).
|
||||
totalReq, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricReq, windowMin, time.Now())
|
||||
if !ok || totalReq != 1 {
|
||||
t.Errorf("total request/min counter = %d (ok=%v), want 1", totalReq, ok)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -266,7 +266,7 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
|
||||
return
|
||||
}
|
||||
|
||||
scopes := buildLimitScopes(u, grp)
|
||||
scopes := withTotalScope(buildLimitScopes(u, grp))
|
||||
if violation := g.limiter.checkAndCount(scopes); violation != nil {
|
||||
writeLimitViolation(w, violation)
|
||||
return
|
||||
|
||||
@@ -461,9 +461,24 @@ func TestHandlePassthrough_NonStreamJSON_AccountsUsage(t *testing.T) {
|
||||
t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody)
|
||||
}
|
||||
|
||||
tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||
if !ok || tok != 10 {
|
||||
t.Errorf("user token/day counter = %d (ok=%v), want 10", tok, ok)
|
||||
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
if !ok || tokIn != 7 {
|
||||
t.Errorf("user tokin/day counter = %d (ok=%v), want 7", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
||||
if !ok || tokOut != 3 {
|
||||
t.Errorf("user tokout/day counter = %d (ok=%v), want 3", tokOut, ok)
|
||||
}
|
||||
|
||||
// withTotalScope (handlePassthrough, routes_passthrough.go) must have
|
||||
// appended the synthetic total scope too (v0.2 data-layer task).
|
||||
totalTokIn, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricTokIn, windowDay, time.Now())
|
||||
if !ok || totalTokIn != 7 {
|
||||
t.Errorf("total tokin/day counter = %d (ok=%v), want 7", totalTokIn, ok)
|
||||
}
|
||||
totalTokOut, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricTokOut, windowDay, time.Now())
|
||||
if !ok || totalTokOut != 3 {
|
||||
t.Errorf("total tokout/day counter = %d (ok=%v), want 3", totalTokOut, ok)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -623,9 +638,13 @@ func TestHandlePassthrough_AcceptEncodingStripped_AccountsUsageThroughGzip(t *te
|
||||
t.Errorf("body = %q, want decompressed upstream body %q", rec.Body.String(), respBody)
|
||||
}
|
||||
|
||||
tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||
if !ok || tok != 10 {
|
||||
t.Errorf("user token/day counter = %d (ok=%v), want 10 — usage must still be accounted through gzip", tok, ok)
|
||||
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
if !ok || tokIn != 7 {
|
||||
t.Errorf("user tokin/day counter = %d (ok=%v), want 7 — usage must still be accounted through gzip", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
||||
if !ok || tokOut != 3 {
|
||||
t.Errorf("user tokout/day counter = %d (ok=%v), want 3 — usage must still be accounted through gzip", tokOut, ok)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -757,9 +776,13 @@ func TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_RequestOnlyAccou
|
||||
t.Errorf("client body = %d bytes, want the full %d bytes (client copy must never truncate)", rec.Body.Len(), len(respBody))
|
||||
}
|
||||
|
||||
tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||
if !ok || tok != 0 {
|
||||
t.Errorf("user token/day counter = %d (ok=%v), want 0 — usage accounting must be skipped over the tee cap", tok, ok)
|
||||
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
if !ok || tokIn != 0 {
|
||||
t.Errorf("user tokin/day counter = %d (ok=%v), want 0 — usage accounting must be skipped over the tee cap", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
||||
if !ok || tokOut != 0 {
|
||||
t.Errorf("user tokout/day counter = %d (ok=%v), want 0 — usage accounting must be skipped over the tee cap", tokOut, ok)
|
||||
}
|
||||
reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now())
|
||||
if !ok || reqCount != 1 {
|
||||
|
||||
+21
-1
@@ -113,7 +113,7 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr
|
||||
return
|
||||
}
|
||||
|
||||
scopes := buildLimitScopes(u, grp)
|
||||
scopes := withTotalScope(buildLimitScopes(u, grp))
|
||||
if violation := g.limiter.checkAndCount(scopes); violation != nil {
|
||||
writeLimitViolation(sw, violation)
|
||||
return
|
||||
@@ -225,6 +225,12 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr
|
||||
// apply to a given request. The user scope is listed first, so
|
||||
// checkAndCount reports a user's own violation ahead of their group's when
|
||||
// both are breached by the same request.
|
||||
//
|
||||
// Callers metering actual LLM traffic wrap this result in withTotalScope
|
||||
// before passing it to checkAndCount/account; handleAdminAPI (admin.go)
|
||||
// calls this directly, without withTotalScope, so an admin request's own
|
||||
// req/min-req/day accounting never contributes to the total-scope
|
||||
// LLM-traffic series (v0.2 data-layer task).
|
||||
func buildLimitScopes(u *user, grp *group) []limitScope {
|
||||
scopes := make([]limitScope, 0, 2)
|
||||
if u.limits != nil {
|
||||
@@ -236,6 +242,20 @@ func buildLimitScopes(u *user, grp *group) []limitScope {
|
||||
return scopes
|
||||
}
|
||||
|
||||
// withTotalScope returns scopes with the synthetic total scope
|
||||
// (totalScopeKind/totalScopeID, limits.go) appended — the single helper
|
||||
// every metered route (runUnified above; resolveMediaRequest,
|
||||
// routes_media.go; handlePassthrough, routes_passthrough.go;
|
||||
// handleTargetProxy, mcp_a2a.go) calls around its own buildLimitScopes
|
||||
// result, so none of them can forget it and none of them duplicate the
|
||||
// scope literal. Deliberately not folded into buildLimitScopes itself:
|
||||
// that helper is also called directly by handleAdminAPI (admin.go) for
|
||||
// its own req/min-req/day accounting, and admin traffic must never
|
||||
// contribute to the total LLM-traffic scope.
|
||||
func withTotalScope(scopes []limitScope) []limitScope {
|
||||
return append(scopes, limitScope{kind: totalScopeKind, id: totalScopeID, limits: nil})
|
||||
}
|
||||
|
||||
// unifiedCostMicros resolves the price to charge one request's usage
|
||||
// against: canonical ("provider/model") first, falling back to bare — the
|
||||
// upstream model id alone — only when canonical has no configured price at
|
||||
|
||||
+40
-12
@@ -127,14 +127,34 @@ func TestHandleChat_HappyPath_NonStreaming_AccountsUsage(t *testing.T) {
|
||||
t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody)
|
||||
}
|
||||
|
||||
tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||
if !ok || tok != 15 {
|
||||
t.Errorf("user token/day counter = %d (ok=%v), want 15", tok, ok)
|
||||
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
if !ok || tokIn != 10 {
|
||||
t.Errorf("user tokin/day counter = %d (ok=%v), want 10", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
||||
if !ok || tokOut != 5 {
|
||||
t.Errorf("user tokout/day counter = %d (ok=%v), want 5", tokOut, ok)
|
||||
}
|
||||
reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now())
|
||||
if !ok || reqCount != 1 {
|
||||
t.Errorf("user request/min counter = %d (ok=%v), want 1", reqCount, ok)
|
||||
}
|
||||
|
||||
// withTotalScope (routes_unified.go) must have appended the synthetic
|
||||
// total scope to runUnified's own checkAndCount/account calls: its
|
||||
// req/tokin/tokout counters mirror alice's own (v0.2 data-layer task).
|
||||
totalReq, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricReq, windowMin, time.Now())
|
||||
if !ok || totalReq != 1 {
|
||||
t.Errorf("total request/min counter = %d (ok=%v), want 1", totalReq, ok)
|
||||
}
|
||||
totalTokIn, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricTokIn, windowDay, time.Now())
|
||||
if !ok || totalTokIn != 10 {
|
||||
t.Errorf("total tokin/day counter = %d (ok=%v), want 10", totalTokIn, ok)
|
||||
}
|
||||
totalTokOut, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricTokOut, windowDay, time.Now())
|
||||
if !ok || totalTokOut != 5 {
|
||||
t.Errorf("total tokout/day counter = %d (ok=%v), want 5", totalTokOut, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleChat_Streaming_FlushesIncrementally proves a streaming chat
|
||||
@@ -980,9 +1000,13 @@ func TestHandleChat_MidStreamDropAfterUsageChunk_AccountsPartialUsage(t *testing
|
||||
t.Errorf("want the mid-stream error logged, got %q", logBuf.String())
|
||||
}
|
||||
|
||||
tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||
if !ok || tok != 10 {
|
||||
t.Errorf("user token/day counter = %d (ok=%v), want 10 (the usage chunk captured before the drop, 7 prompt + 3 completion)", tok, ok)
|
||||
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
if !ok || tokIn != 7 {
|
||||
t.Errorf("user tokin/day counter = %d (ok=%v), want 7 (the usage chunk captured before the drop)", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
||||
if !ok || tokOut != 3 {
|
||||
t.Errorf("user tokout/day counter = %d (ok=%v), want 3 (the usage chunk captured before the drop)", tokOut, ok)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1183,9 +1207,13 @@ func TestHandleChat_CacheHitMiss_EndToEnd_CountersAndHeaders(t *testing.T) {
|
||||
if !ok || reqCount != 1 {
|
||||
t.Errorf("req:day after first request = %d (ok=%v), want 1", reqCount, ok)
|
||||
}
|
||||
tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||
if !ok || tok != 15 {
|
||||
t.Errorf("tok:day after first request = %d (ok=%v), want 15 (10 prompt + 5 completion)", tok, ok)
|
||||
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
if !ok || tokIn != 10 {
|
||||
t.Errorf("tokin:day after first request = %d (ok=%v), want 10", tokIn, ok)
|
||||
}
|
||||
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
||||
if !ok || tokOut != 5 {
|
||||
t.Errorf("tokout:day after first request = %d (ok=%v), want 5", tokOut, ok)
|
||||
}
|
||||
|
||||
// Second, identical request: hit.
|
||||
@@ -1207,9 +1235,9 @@ func TestHandleChat_CacheHitMiss_EndToEnd_CountersAndHeaders(t *testing.T) {
|
||||
if !ok || reqCount != 2 {
|
||||
t.Errorf("req:day after the cache hit = %d (ok=%v), want 2 (hit still counts as one more request)", reqCount, ok)
|
||||
}
|
||||
tok, ok = gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now())
|
||||
if !ok || tok != 15 {
|
||||
t.Errorf("tok:day after the cache hit = %d (ok=%v), want still 15 (a hit accounts zero tokens, no estimation)", tok, ok)
|
||||
tokIn, ok = gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
if !ok || tokIn != 10 {
|
||||
t.Errorf("tokin:day after the cache hit = %d (ok=%v), want still 10 (a hit accounts zero tokens, no estimation)", tokIn, ok)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user