package traefikllmgateway import ( "context" "net/http" "net/http/httptest" "sync" "testing" ) // This file measures counterStore round trips per request — the concrete // evidence for admin-redesign WP-A's own hot-path promise (COMMON3.md, // plan §2's command table): "round trips per request unchanged (2 sync + // 1 async); cache hit +1 async RT; defaults (umodel/latency off) add <=3 // INCRBY to account and <=2 SET/min to admission." Every test below drives // a real Gateway (New(), the same constructor production uses) through a // single-group chat request and counts calls to the limiter's own // counterStore — never entries within a call, since a redisStore call // pipelines its whole batch into ONE network round trip regardless of how // many entries it carries (counterStore's own doc comment, limits.go), so // counting CALLS is exactly counting round trips. // roundTripCountingStore wraps a real counterStore (a fresh memoryStore, // so behavior is otherwise identical to a Redis-less deployment) and // counts every call to its three methods. type roundTripCountingStore struct { delegate counterStore mu sync.Mutex calls int } func (s *roundTripCountingStore) getMulti(keys []string) ([]int64, error) { s.mu.Lock() s.calls++ s.mu.Unlock() return s.delegate.getMulti(keys) } func (s *roundTripCountingStore) incrMulti(entries []counterIncr) ([]int64, error) { s.mu.Lock() s.calls++ s.mu.Unlock() return s.delegate.incrMulti(entries) } func (s *roundTripCountingStore) incrAndGetMulti(entries []counterIncr, reads []string) ([]int64, []int64, error) { s.mu.Lock() s.calls++ s.mu.Unlock() return s.delegate.incrAndGetMulti(entries, reads) } func (s *roundTripCountingStore) total() int { s.mu.Lock() defer s.mu.Unlock() return s.calls } // buildRoundTripTestGateway constructs a real Gateway for a single-group, // single-user, single-provider chat deployment — cfgFn, when non-nil, may // layer on admin/stats/cache/metrics config before construction — then // swaps the limiter's store for a call-counting wrapper and forces spawn // synchronous, so an async write (recordProviderAttempt, recordCacheHit) // lands deterministically before the test inspects the count. The swap // happens strictly AFTER New() returns: construction itself never calls a // counterStore method, so it is never counted either way. func buildRoundTripTestGateway(t *testing.T, srv *httptest.Server, cfgFn func(*Config)) (*Gateway, *roundTripCountingStore) { t.Helper() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{ {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, }} if cfgFn != nil { cfgFn(cfg) } next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } counting := &roundTripCountingStore{delegate: newMemoryStore()} gw.limiter.store = counting gw.limiter.spawn = func(f func()) { f() } // land every async write synchronously before this test inspects counts return gw, counting } const roundTripTestRespBody = `{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}` func roundTripTestUpstream(t *testing.T) *httptest.Server { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(roundTripTestRespBody)) })) t.Cleanup(srv.Close) return srv } func roundTripChatRequest(t *testing.T) *http.Request { t.Helper() body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} return newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) } // TestRoundTrips_SingleGroupChat_Defaults_TwoSyncOneAsync is the // command-table baseline: a single-group chat request with every admin // stat off pays exactly 2 SYNCHRONOUS counterStore round trips // (checkAndCount's storeIncrAndGetMulti, account's storeIncrMulti) plus 1 // ASYNCHRONOUS one (recordProviderAttempt's spawned storeIncrMulti, // forced synchronous here for a deterministic count) — 3 counterStore // calls total. This is BEFORE, in the same sense the task's own // instructions ask for: the number every "all stats on" comparison below // must not exceed. func TestRoundTrips_SingleGroupChat_Defaults_TwoSyncOneAsync(t *testing.T) { srv := roundTripTestUpstream(t) gw, counting := buildRoundTripTestGateway(t, srv, nil) rec := httptest.NewRecorder() gw.ServeHTTP(rec, roundTripChatRequest(t)) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } got := counting.total() t.Logf("BEFORE (defaults, no admin stats): %d counterStore round trips for one single-group chat request", got) if got != 3 { t.Errorf("round trips = %d, want 3 (2 sync + 1 async, command table)", got) } } // TestRoundTrips_SingleGroupChat_AllStatsOn_NoNewSyncRoundTrips is this // task's own deliverable: turning EVERY new admin stat on (userModel, // latency; admin enabled; metrics enabled to also exercise the latency // recorder's metrics-side wiring) for a single-group chat request that is // neither a cache hit nor a failover adds ZERO new SYNCHRONOUS counterStore // round trips — every new counter family this task adds to the request's // own two synchronous calls (checkAndCount, accountWith) rides an EXISTING // batch (accountWith's own umodel/latency extras blocks), never a new one. // // It DOES add exactly one new ASYNCHRONOUS round trip, recordLastSeen's own // fire-and-forget write (P9 fix, admin dashboard redesign verify round): // last-seen used to ride checkAndCount's own synchronous batch // unconditionally, which meant a REJECTED request still moved that scope's // last-seen timestamp forward. It is now recorded only from the "admission // confirmed" branch, off the request's hot path via countAsync/l.spawn — // exactly like recordProviderAttempt's own pre-existing async write — so // this is a SEPARATE async call, not a batch this one rides. 2 sync + 2 // async = 4 total, once per admitted request from a scope l.lastSeen.due // has not already throttled within lastSeenGateThrottle (60s) — a fresh // limiter's first request for a scope, as this test drives. func TestRoundTrips_SingleGroupChat_AllStatsOn_NoNewSyncRoundTrips(t *testing.T) { srv := roundTripTestUpstream(t) gw, counting := buildRoundTripTestGateway(t, srv, func(cfg *Config) { cfg.Admin = &AdminConfig{Enabled: true, Stats: &AdminStatsConfig{UserModel: true, Latency: true}} cfg.Metrics = &MetricsConfig{Enabled: true} }) rec := httptest.NewRecorder() gw.ServeHTTP(rec, roundTripChatRequest(t)) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } got := counting.total() t.Logf("AFTER (admin.stats.userModel=true, admin.stats.latency=true, metrics on, no cache): %d counterStore round trips for the SAME single-group chat request", got) if got != 4 { t.Errorf("round trips with all stats on = %d, want 4 (2 sync + 2 async: checkAndCount, accountWith, recordProviderAttempt, recordLastSeen) — every new counter family rides an EXISTING sync batch except last-seen, which is its own async write by design (P9 fix)", got) } } // TestRoundTrips_SingleGroupChat_AdminEnabledStatsOff_FourRoundTripsWhenLastSeenDue // is N4 (verify-redesign-final.md): README's own round-trip table used to // fold "admin.enabled: false" and "admin.enabled: true with both // admin.stats opt-ins off" into one "Defaults" row claiming 3 round trips // for both. That is only true for the first: l.statsAdmin is // adminEnabled(cfg) alone (llmgateway.go, newConfiguredLimiter) — // independent of admin.stats.userModel/latency — so admin.enabled: true // with BOTH opt-ins off still gates every "always-on-with-admin" family // on, last-seen included. For a scope whose last-seen is DUE (this // test's case: a fresh limiter, first request for the scope), that is // recordLastSeen's own async write same as // TestRoundTrips_SingleGroupChat_AllStatsOn_NoNewSyncRoundTrips — 4 round // trips, not 3, even though neither opt-in stat is on and there is no // umodel/latency entry in any batch. func TestRoundTrips_SingleGroupChat_AdminEnabledStatsOff_FourRoundTripsWhenLastSeenDue(t *testing.T) { srv := roundTripTestUpstream(t) gw, counting := buildRoundTripTestGateway(t, srv, func(cfg *Config) { cfg.Admin = &AdminConfig{Enabled: true} // Stats left nil: userModel/latency both default off }) rec := httptest.NewRecorder() gw.ServeHTTP(rec, roundTripChatRequest(t)) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } got := counting.total() t.Logf("admin.enabled=true, stats.userModel/latency both off, last-seen due: %d counterStore round trips", got) if got != 4 { t.Errorf("round trips = %d, want 4 (2 sync + 2 async: checkAndCount, accountWith, recordProviderAttempt, recordLastSeen) — admin.enabled alone gates last-seen, not the opt-in stats", got) } } // TestRoundTrips_SingleGroupChat_CacheHit_AddsExactlyOneAsyncRoundTrip is // the command table's one documented exception ("cache hit +1 async RT"): // a response-cache HIT costs checkAndCount's own admission round trip // (every metered route enforces request-rate limits before it even knows // whether the response will be a hit) PLUS recordCacheHit's own async // write — 2 counterStore calls, never account()'s or // recordProviderAttempt's own (the hit branch returns before ever // reaching either, since no upstream call was made). This is +1 over // this feature's own pre-task baseline of "a hit costs only // checkAndCount, 1 call" — cache stats did not exist before this task. func TestRoundTrips_SingleGroupChat_CacheHit_AddsExactlyOneAsyncRoundTrip(t *testing.T) { srv := roundTripTestUpstream(t) redisLn := newBehavioralRedisServer(t) gw, counting := buildRoundTripTestGateway(t, srv, func(cfg *Config) { cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} cfg.Admin = &AdminConfig{Enabled: true} }) // buildRoundTripTestGateway already swapped gw.limiter.store for a // counting wrapper over a fresh memoryStore — the response CACHE // itself is a *responseCache (cache.go), not a counterStore: its own // Redis calls (via cfg.Redis's real behavioral fake) are never // counted here, matching the command table's own "counterStore round // trips" scope exactly. req1 := roundTripChatRequest(t) rec1 := httptest.NewRecorder() gw.ServeHTTP(rec1, req1) if rec1.Code != http.StatusOK || rec1.Header().Get("X-Llmgw-Cache") != "miss" { t.Fatalf("first request: status=%d cache=%q, want 200/miss", rec1.Code, rec1.Header().Get("X-Llmgw-Cache")) } baseline := counting.total() t.Logf("miss (cache+admin on): %d cumulative counterStore round trips", baseline) req2 := roundTripChatRequest(t) rec2 := httptest.NewRecorder() gw.ServeHTTP(rec2, req2) if rec2.Code != http.StatusOK || rec2.Header().Get("X-Llmgw-Cache") != "hit" { t.Fatalf("second request: status=%d cache=%q, want 200/hit", rec2.Code, rec2.Header().Get("X-Llmgw-Cache")) } hitOwnCalls := counting.total() - baseline t.Logf("hit (cache+admin on): %d counterStore round trips for this ONE request (checkAndCount + recordCacheHit)", hitOwnCalls) if hitOwnCalls != 2 { t.Errorf("cache-hit round trips = %d, want 2 (checkAndCount + recordCacheHit; no account(), no recordProviderAttempt)", hitOwnCalls) } }