package traefikllmgateway import ( "bufio" "bytes" "context" "encoding/json" "errors" "fmt" "io" "math" "net" "net/http" "net/http/httptest" "os" "strconv" "strings" "sync" "testing" "time" "github.com/stretchr/testify/assert" ) // newUnifiedRequest builds a POST request against the unified routes, // JSON-encoding body and attaching apiKey as a Bearer token when non-empty. func newUnifiedRequest(t *testing.T, method, path, apiKey string, body map[string]any) *http.Request { t.Helper() b, err := json.Marshal(body) if err != nil { t.Fatalf("marshal request body: %v", err) } req := httptest.NewRequest(method, path, bytes.NewReader(b)) if apiKey != "" { req.Header.Set("Authorization", "Bearer "+apiKey) } return req } // recordingWriter is a minimal http.ResponseWriter/http.Flusher that logs // the sequence of Write and Flush calls it receives, so a streaming test // can assert flushes happen incrementally rather than once at the end. type recordingWriter struct { hdr http.Header events []string buf bytes.Buffer status int } func newRecordingWriter() *recordingWriter { return &recordingWriter{hdr: http.Header{}} } func (r *recordingWriter) Header() http.Header { return r.hdr } func (r *recordingWriter) WriteHeader(status int) { r.status = status } func (r *recordingWriter) Write(b []byte) (int, error) { if r.status == 0 { r.status = http.StatusOK // mirrors net/http's implicit-200 default when Write precedes any explicit WriteHeader } n, _ := r.buf.Write(b) r.events = append(r.events, "write") return n, nil } func (r *recordingWriter) Flush() { r.events = append(r.events, "flush") } // alwaysErrStore is a counterStore stub whose every operation fails, used // to drive the limiter's fail-closed (storeDown) path deterministically. type alwaysErrStore struct{} var errStoreDownStub = errors.New("stub: store down") func (alwaysErrStore) getMulti([]string) ([]int64, error) { return nil, errStoreDownStub } func (alwaysErrStore) incrMulti([]counterIncr) ([]int64, error) { return nil, errStoreDownStub } func (alwaysErrStore) incrAndGetMulti([]counterIncr, []string) ([]int64, []int64, error) { return nil, nil, errStoreDownStub } // TestHandleChat_HappyPath_NonStreaming_AccountsUsage drives a full // chat-completion request through Gateway.ServeHTTP against a fake // openai-shaped upstream, and asserts the response is forwarded verbatim // and the reported usage lands in the limiter's own counters. func TestHandleChat_HappyPath_NonStreaming_AccountsUsage(t *testing.T) { const respBody = `{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-real","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}` 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(respBody)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{ {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, }} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } if rec.Body.String() != respBody { t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody) } 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 // completion is forwarded to the client chunk by chunk, flushing after // each one, rather than buffered until the response completes. func TestHandleChat_Streaming_FlushesIncrementally(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) fl, ok := w.(http.Flusher) if !ok { t.Fatal("fake upstream ResponseWriter does not support Flush") } for _, c := range []string{ `{"id":"c1","object":"chat.completion.chunk","choices":[{"delta":{"content":"Hel"}}]}`, `{"id":"c1","object":"chat.completion.chunk","choices":[{"delta":{"content":"lo"}}]}`, } { _, _ = w.Write([]byte("data: " + c + "\n\n")) fl.Flush() } _, _ = w.Write([]byte("data: [DONE]\n\n")) fl.Flush() })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rw := newRecordingWriter() h.ServeHTTP(rw, req) if rw.status != http.StatusOK { t.Fatalf("status = %d, want 200", rw.status) } writes, flushes := 0, 0 for _, e := range rw.events { switch e { case "write": writes++ case "flush": flushes++ } } if writes < 3 { t.Fatalf("want >=3 writes (2 content chunks + [DONE]), got %d (events=%v)", writes, rw.events) } if flushes < writes { t.Fatalf("want a flush per write (incremental delivery), got %d writes and %d flushes", writes, flushes) } if !strings.Contains(rw.buf.String(), "data: [DONE]") { t.Errorf("body missing terminal [DONE], got %q", rw.buf.String()) } } // --- usage provenance (feat: expose token-accounting provenance) --- // // Token accounting has three provenances (usage.estimated's own doc // comment, limits.go): reported (real provider usage), estimated (a // non-streaming response with no usage, prompt substituted from body // size), and unbilled (a streaming response with no usage, request // counted but zero tokens billed). These three tests drive each real // HTTP shape through Gateway.ServeHTTP and assert both the provenance // classification AND that billing itself is byte-identical to the // pre-feature contract — this feature must never change a billed number. // TestUsageProvenance_ReportedUsage_ClassifiedReported proves the normal, // accurate case (a provider that returns real usage) is classified // "reported", not merely "not estimated" by omission. // // MUTATION VERIFIED: changing routes_unified.go's provenance switch // default from `provenance := provenanceReported` to `provenance := // provenanceEstimated` made this test fail — the seen provenance became // "estimated" instead of "reported". Reverted before committing. func TestUsageProvenance_ReportedUsage_ClassifiedReported(t *testing.T) { t.Parallel() 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(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`)) })) defer srv.Close() cfg := CreateConfig() cfg.Metrics = &MetricsConfig{Enabled: true} cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} rec := httptest.NewRecorder() h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } snaps := gw.provenance.snapshot() if len(snaps) != 1 { t.Fatalf("g.provenance has %d entries, want 1: %+v", len(snaps), snaps) } s := snaps[0] if s.key.provider != "openai" || s.key.provenance != provenanceReported { t.Errorf("provenance key = %+v, want provider=openai provenance=reported", s.key) } if s.requests != 1 { t.Errorf("requests = %d, want 1", s.requests) } if s.tokens != 15 { t.Errorf("tokens = %d, want 15 (10 prompt + 5 completion, the real reported usage)", s.tokens) } // Billing itself must be exactly what TestHandleChat_HappyPath_ // NonStreaming_AccountsUsage already proves for this identical // response shape — this feature must never change a billed number. tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now()) tokOut, _ := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now()) if tokIn != 10 || tokOut != 5 { t.Errorf("billed tokin/tokout = %d/%d, want 10/5 — provenance recording must never change what gets billed", tokIn, tokOut) } } // TestUsageProvenance_NonStreamingZeroUsage_ClassifiedEstimated proves a // non-streaming response that carries no usage at all is classified // "estimated", with the SAME prompt-estimate/zero-completion billing // routes_unified.go already applied before this feature existed — // ceil(request body size / 4) prompt, zero completion. // // MUTATION VERIFIED: changing the recording switch's `case // result.estimated:` guard (routes_unified.go) to `case false:` made // this test fail — the seen provenance stayed "reported" instead of // becoming "estimated". Reverted before committing. func TestUsageProvenance_NonStreamingZeroUsage_ClassifiedEstimated(t *testing.T) { t.Parallel() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) // No "usage" field at all — decodes to chatUsagePayload's zero // value (provider_openai.go's forwardJSON), the identical shape a // provider that omits usage from a non-streaming response sends. _, _ = w.Write([]byte(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}`)) })) defer srv.Close() cfg := CreateConfig() cfg.Metrics = &MetricsConfig{Enabled: true} cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } reqBody := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} // The exact bytes newUnifiedRequest below marshals reqBody into — // routes_unified.go's own estimate is ceil(len(body) / 4) over these // SAME bytes, so this reproduces the source's formula against a // black-box-observable input (the request the test itself sends), // not a re-implementation of internal logic. marshaled, err := json.Marshal(reqBody) if err != nil { t.Fatalf("marshal request body: %v", err) } wantPrompt := int64(math.Ceil(float64(len(marshaled)) / 4)) rec := httptest.NewRecorder() h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", reqBody)) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } snaps := gw.provenance.snapshot() if len(snaps) != 1 { t.Fatalf("g.provenance has %d entries, want 1: %+v", len(snaps), snaps) } s := snaps[0] if s.key.provider != "openai" || s.key.provenance != provenanceEstimated { t.Errorf("provenance key = %+v, want provider=openai provenance=estimated", s.key) } if s.tokens != wantPrompt { t.Errorf("provenance tokens = %d, want %d (ceil(request body size / 4), the estimate)", s.tokens, wantPrompt) } // Billing unchanged from the pre-feature contract: prompt substituted // with the estimate, completion billed as zero. tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now()) tokOut, _ := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now()) if tokIn != wantPrompt { t.Errorf("billed tokin = %d, want %d (ceil(request body size / 4))", tokIn, wantPrompt) } if tokOut != 0 { t.Errorf("billed tokout = %d, want 0 — completion is billed as zero for an estimated response", tokOut) } } // TestUsageProvenance_StreamingZeroUsage_ClassifiedUnbilled_DistinctFromEstimated // proves a streaming response that carries no usage at all is classified // "unbilled" — DISTINCT from "estimated" (this feature's own // correctness requirement: collapsing the two into one "not reported" // signal would hide that a provider serving completions entirely free // is a materially different failure than an approximated bill) — and // that billing stays exactly what routes_unified.go already applied // before this feature existed: the request counted, zero tokens billed. // // MUTATION VERIFIED: changing the recording switch's `case streaming && // result.total() == 0:` guard (routes_unified.go) to `case // result.estimated:` (merging unbilled into the estimated branch) made // this test fail — the seen provenance became "estimated" instead of // "unbilled", and the distinctness assertion below failed too. Reverted // before committing. func TestUsageProvenance_StreamingZeroUsage_ClassifiedUnbilled_DistinctFromEstimated(t *testing.T) { t.Parallel() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) fl, ok := w.(http.Flusher) if !ok { t.Fatal("fake upstream ResponseWriter does not support Flush") } _, _ = w.Write([]byte("data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n")) fl.Flush() // No usage-only chunk — the provider never honored // stream_options.include_usage, the exact shape this feature // exists to surface. _, _ = w.Write([]byte("data: [DONE]\n\n")) fl.Flush() })) defer srv.Close() cfg := CreateConfig() cfg.Metrics = &MetricsConfig{Enabled: true} cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} rw := newRecordingWriter() h.ServeHTTP(rw, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)) if rw.status != http.StatusOK { t.Fatalf("status = %d, want 200", rw.status) } snaps := gw.provenance.snapshot() if len(snaps) != 1 { t.Fatalf("g.provenance has %d entries, want 1: %+v", len(snaps), snaps) } s := snaps[0] if s.key.provider != "openai" || s.key.provenance != provenanceUnbilled { t.Errorf("provenance key = %+v, want provider=openai provenance=unbilled", s.key) } if s.key.provenance == provenanceEstimated { t.Errorf("provenance = estimated, want unbilled — a streaming zero-usage response must never be classified the same as a non-streaming one") } if s.tokens != 0 { t.Errorf("provenance tokens = %d, want 0 — unbilled charges nothing at all", s.tokens) } // Billing unchanged from the pre-feature contract: the request is // counted (req/min bumped) but zero tokens are billed. tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now()) tokOut, _ := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now()) if tokIn != 0 || tokOut != 0 { t.Errorf("billed tokin/tokout = %d/%d, want 0/0 — a streaming zero-usage response bills nothing", tokIn, tokOut) } reqCount, ok2 := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now()) if !ok2 || reqCount != 1 { t.Errorf("request/min counter = %d (ok=%v), want 1 — the request must still be COUNTED even though zero tokens are billed", reqCount, ok2) } } // TestHandleChat_NoAPIKey_Returns401 asserts the unified routes require // authentication before any resolve/limit work happens. func TestHandleChat_NoAPIKey_Returns401(t *testing.T) { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "", map[string]any{"model": "gpt-test"}) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusUnauthorized { t.Fatalf("status = %d, want 401, body=%s", rec.Code, rec.Body.String()) } } // TestHandleChat_DeniedModel_Returns403 covers a model that a configured // provider knows, but the caller's group is not authorized to use. func TestHandleChat_DeniedModel_Returns403(t *testing.T) { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {Models: []string{"gpt-allowed-only"}}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "gpt-test", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusForbidden { t.Fatalf("status = %d, want 403, body=%s", rec.Code, rec.Body.String()) } } // TestHandleChat_UnknownModel_Returns404 covers a model no configured // provider knows at all. func TestHandleChat_UnknownModel_Returns404(t *testing.T) { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}} cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "no-such-model", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusNotFound { t.Fatalf("status = %d, want 404, body=%s", rec.Code, rec.Body.String()) } } // TestHandleChat_RequestLimitExceeded_Returns429WithRetryAfter drives two // requests against a user limited to one per minute, and asserts the // second is refused with a Retry-After header. func TestHandleChat_RequestLimitExceeded_Returns429WithRetryAfter(t *testing.T) { const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` 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(respBody)) })) defer srv.Close() 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{RequestsPerMinute: 1}}, }} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "gpt-test", "messages": []any{}} req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec1 := httptest.NewRecorder() h.ServeHTTP(rec1, req1) if rec1.Code != http.StatusOK { t.Fatalf("first request status = %d, want 200, body=%s", rec1.Code, rec1.Body.String()) } req2 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec2 := httptest.NewRecorder() h.ServeHTTP(rec2, req2) if rec2.Code != http.StatusTooManyRequests { t.Fatalf("second request status = %d, want 429, body=%s", rec2.Code, rec2.Body.String()) } if ra := rec2.Header().Get("Retry-After"); ra == "" { t.Error("want a Retry-After header on 429, got none") } } // TestHandleChat_StoreDown_FailClosed_Returns503 wires a counterStore // stub whose every operation fails, with failOpen=false, and asserts the // request is refused with 503 rather than silently enforcing limits // against an unreachable backend. func TestHandleChat_StoreDown_FailClosed_Returns503(t *testing.T) { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", 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{RequestsPerMinute: 100}}, }} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } gw.limiter = newLimiter(alwaysErrStore{}, false) body := map[string]any{"model": "gpt-test", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusServiceUnavailable { t.Fatalf("status = %d, want 503, body=%s", rec.Code, rec.Body.String()) } } // TestHandleEmbeddings_HappyPath mirrors the chat-completion happy path // against POST /v1/embeddings. func TestHandleEmbeddings_HappyPath(t *testing.T) { const respBody = `{"object":"list","data":[{"embedding":[0.1,0.2]}],"model":"text-embedding-3","usage":{"prompt_tokens":5,"total_tokens":5}}` srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/embeddings" { t.Errorf("path = %q, want /v1/embeddings", r.URL.Path) } w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"text-embedding-3"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "text-embedding-3", "input": "hello"} req := newUnifiedRequest(t, http.MethodPost, "/v1/embeddings", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } if rec.Body.String() != respBody { t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody) } } // TestHandleChat_OpenAI_AliasKeyStrippedFromUpstreamBody asserts the // gatewayAliasKey convention field never reaches a real openai-type // provider on POST /v1/chat/completions: unlike anthropic/gemini, this // adapter marshals req verbatim as the upstream wire body, so a leftover // alias entry would arrive as an unrecognized request field. This test // fails if provider_openai.go's chatCompletion ever drops its // delete(req, gatewayAliasKey) call. func TestHandleChat_OpenAI_AliasKeyStrippedFromUpstreamBody(t *testing.T) { const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` var gotBody []byte srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotBody, _ = io.ReadAll(r.Body) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() 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"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } const requestedID = "openai/gpt-test" body := map[string]any{"model": requestedID, "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } sentUpstream := decodeJSONBody(t, gotBody) if _, present := sentUpstream[gatewayAliasKey]; present { t.Errorf("gatewayAliasKey %q leaked into the openai chat upstream request body: %v", gatewayAliasKey, sentUpstream) } if sentUpstream["model"] != "gpt-test" { t.Errorf("upstream request model = %v, want bare upstream id %q", sentUpstream["model"], "gpt-test") } } // TestHandleEmbeddings_OpenAI_AliasKeyStrippedFromUpstreamBody mirrors // TestHandleChat_OpenAI_AliasKeyStrippedFromUpstreamBody for POST // /v1/embeddings: this test fails if provider_openai.go's embeddings ever // drops its delete(req, gatewayAliasKey) call. func TestHandleEmbeddings_OpenAI_AliasKeyStrippedFromUpstreamBody(t *testing.T) { const respBody = `{"object":"list","data":[{"embedding":[0.1]}],"usage":{"prompt_tokens":1,"total_tokens":1}}` var gotBody []byte srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotBody, _ = io.ReadAll(r.Body) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"text-embedding-3"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } const requestedID = "openai/text-embedding-3" body := map[string]any{"model": requestedID, "input": "hello"} req := newUnifiedRequest(t, http.MethodPost, "/v1/embeddings", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } sentUpstream := decodeJSONBody(t, gotBody) if _, present := sentUpstream[gatewayAliasKey]; present { t.Errorf("gatewayAliasKey %q leaked into the openai embeddings upstream request body: %v", gatewayAliasKey, sentUpstream) } if sentUpstream["model"] != "text-embedding-3" { t.Errorf("upstream request model = %v, want bare upstream id %q", sentUpstream["model"], "text-embedding-3") } } // TestHandleChat_AnthropicProviderPrefixedModel_EchoesAliasInResponse // requests a provider-prefixed model id ("anthropic/claude-x"): the // upstream call must use the bare upstream id, but the translated // response's "model" field must echo back the exact id the client // requested (ruling a, ALIAS ECHO), not the bare upstream form. func TestHandleChat_AnthropicProviderPrefixedModel_EchoesAliasInResponse(t *testing.T) { const anthResp = `{"id":"msg_01ABC","type":"message","role":"assistant","content":[{"type":"text","text":"hi there"}],"model":"claude-x","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":4}}` var gotBody []byte srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotBody, _ = io.ReadAll(r.Body) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(anthResp)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-ant", Models: []string{"claude-x"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } const requestedID = "anthropic/claude-x" body := map[string]any{"model": requestedID, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } var out map[string]any if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { t.Fatalf("decode response: %v", err) } if out["model"] != requestedID { t.Errorf("response model = %v, want the client-requested alias %q echoed back", out["model"], requestedID) } sentUpstream := decodeJSONBody(t, gotBody) if sentUpstream["model"] != "claude-x" { t.Errorf("upstream request model = %v, want bare upstream id %q", sentUpstream["model"], "claude-x") } if _, present := sentUpstream[gatewayAliasKey]; present { t.Errorf("gatewayAliasKey %q leaked into the upstream request body: %v", gatewayAliasKey, sentUpstream) } } // --- model aliases end-to-end (spec §5, v0.2) --- // TestHandleChat_ModelAlias_AnthropicTarget_EchoesAliasInResponse mirrors // TestHandleChat_AnthropicProviderPrefixedModel_EchoesAliasInResponse // above, but through an operator-defined modelAliases entry rather than a // bare provider-prefixed id: the upstream call uses the target's own // upstream model id, and the translated response's "model" field echoes // the client's exact alias string back (ruling a, ALIAS ECHO — aliases // inherit the existing __alias machinery for free, per registry.go's // resolve doc comment). func TestHandleChat_ModelAlias_AnthropicTarget_EchoesAliasInResponse(t *testing.T) { const anthResp = `{"id":"msg_01ABC","type":"message","role":"assistant","content":[{"type":"text","text":"hi there"}],"model":"claude-x","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":4}}` var gotBody []byte srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotBody, _ = io.ReadAll(r.Body) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(anthResp)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-ant", Models: []string{"claude-x"}}, } cfg.ModelAliases = map[string]string{"aliased/coding": "anthropic/claude-x"} cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } const requestedID = "aliased/coding" body := map[string]any{"model": requestedID, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } var out map[string]any if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { t.Fatalf("decode response: %v", err) } if out["model"] != requestedID { t.Errorf("response model = %v, want the client-requested alias %q echoed back", out["model"], requestedID) } sentUpstream := decodeJSONBody(t, gotBody) if sentUpstream["model"] != "claude-x" { t.Errorf("upstream request model = %v, want the alias target's bare upstream id %q", sentUpstream["model"], "claude-x") } if _, present := sentUpstream[gatewayAliasKey]; present { t.Errorf("gatewayAliasKey %q leaked into the upstream request body: %v", gatewayAliasKey, sentUpstream) } } // TestHandleChat_ModelAlias_OpenAITarget_UpstreamEchoNotAlias documents the // v0.1 passthrough asymmetry (routes_unified.go's gatewayAliasKey doc // comment) as it applies to aliases: an openai-type target's response is // forwarded verbatim, so its own "model" field carries whatever the mock // upstream itself returned — the target's bare upstream id, NOT the // alias. The gateway makes no attempt, and is not expected, to rewrite an // openai-type body in flight. func TestHandleChat_ModelAlias_OpenAITarget_UpstreamEchoNotAlias(t *testing.T) { const respBody = `{"id":"c1","model":"gpt-test","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` var gotBody []byte srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotBody, _ = io.ReadAll(r.Body) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, } cfg.ModelAliases = map[string]string{"aliased/fast": "openai/gpt-test"} cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "aliased/fast", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } var out map[string]any if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { t.Fatalf("decode response: %v", err) } if out["model"] != "gpt-test" { t.Errorf("response model = %v, want the openai-type passthrough's own upstream id %q, not the alias", out["model"], "gpt-test") } sentUpstream := decodeJSONBody(t, gotBody) if sentUpstream["model"] != "gpt-test" { t.Errorf("upstream request model = %v, want bare upstream id %q", sentUpstream["model"], "gpt-test") } } // TestHandleEmbeddings_ModelAlias_ResolvesToTargetUpstreamModel proves // aliasing applies to /v1/embeddings exactly like /v1/chat/completions — // the upstream call uses the target's bare upstream model id. func TestHandleEmbeddings_ModelAlias_ResolvesToTargetUpstreamModel(t *testing.T) { const respBody = `{"object":"list","data":[{"embedding":[0.1]}],"usage":{"prompt_tokens":1,"total_tokens":1}}` var gotBody []byte srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotBody, _ = io.ReadAll(r.Body) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"text-embedding-3"}}, } cfg.ModelAliases = map[string]string{"aliased/embed": "openai/text-embedding-3"} cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "aliased/embed", "input": "hello"} req := newUnifiedRequest(t, http.MethodPost, "/v1/embeddings", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } sentUpstream := decodeJSONBody(t, gotBody) if sentUpstream["model"] != "text-embedding-3" { t.Errorf("upstream request model = %v, want the alias target's bare upstream id %q", sentUpstream["model"], "text-embedding-3") } } // TestHandleChat_ModelAlias_AuthorizedViaAliasNameAlone_EndToEnd proves // spec §5's authorization rule end-to-end: a group whose models glob // matches only the ALIAS name (not the target, in either form) can still // use it through the real ServeHTTP pipeline, not just modelRegistry.resolve // in isolation. func TestHandleChat_ModelAlias_AuthorizedViaAliasNameAlone_EndToEnd(t *testing.T) { const respBody = `{"id":"c1","model":"gpt-test","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` 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(respBody)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, } cfg.ModelAliases = map[string]string{"aliased/coding": "openai/gpt-test"} // Matches only the alias name, not "openai/gpt-test" nor "gpt-test". cfg.Groups = map[string]*GroupConfig{"narrow": {Models: []string{"aliased/coding"}}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "narrow", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "aliased/coding", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200 (authorized via the alias name alone), body=%s", rec.Code, rec.Body.String()) } } // TestHandleChat_ModelAlias_UnresolvedTarget_Returns404NamingAliasAndTarget // covers spec §5's lazy-resolution 404 end-to-end: an alias whose target // names no currently-known model returns 404 with a message naming both // the alias and the missing target — not the generic "unknown model" text // a plain unresolved bare/prefixed id gets. func TestHandleChat_ModelAlias_UnresolvedTarget_Returns404NamingAliasAndTarget(t *testing.T) { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}, } cfg.ModelAliases = map[string]string{"aliased/future": "openai/not-yet-discovered"} cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "aliased/future", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusNotFound { t.Fatalf("status = %d, want 404, body=%s", rec.Code, rec.Body.String()) } var out struct { Error struct { Message string `json:"message"` } `json:"error"` } if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { t.Fatalf("decode response: %v", err) } if !strings.Contains(out.Error.Message, "aliased/future") { t.Errorf("error message = %q, want it to name the alias %q", out.Error.Message, "aliased/future") } if !strings.Contains(out.Error.Message, "openai/not-yet-discovered") { t.Errorf("error message = %q, want it to name the target %q", out.Error.Message, "openai/not-yet-discovered") } } // TestHandleChat_UpstreamNon2xx_WrapsProviderErrorEnvelope asserts a // non-2xx upstream response is passed through with its status code, but // wrapped in the gateway's own envelope with the raw upstream body // embedded under error.upstream (ruling c). func TestHandleChat_UpstreamNon2xx_WrapsProviderErrorEnvelope(t *testing.T) { const upstreamErrBody = `{"error":{"message":"rate limited upstream","type":"rate_limit_error"}}` srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(upstreamErrBody)) })) defer srv.Close() 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"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "gpt-test", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusTooManyRequests { t.Fatalf("status = %d, want 429, body=%s", rec.Code, rec.Body.String()) } var out struct { Error struct { Upstream map[string]any `json:"upstream"` Message string `json:"message"` Type string `json:"type"` Code string `json:"code"` } `json:"error"` } if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { t.Fatalf("decode response: %v", err) } if out.Error.Type != "upstream_error" { t.Errorf("error.type = %q, want upstream_error", out.Error.Type) } if out.Error.Code != "429" { t.Errorf("error.code = %q, want %q", out.Error.Code, "429") } if out.Error.Upstream == nil { t.Fatal("error.upstream is nil, want the parsed upstream body") } errObj, _ := out.Error.Upstream["error"].(map[string]any) if errObj["message"] != "rate limited upstream" { t.Errorf("error.upstream.error.message = %v, want %q", errObj["message"], "rate limited upstream") } } // TestHandleChat_MidStreamUpstreamDrop_EmitsErrorEventNoDone proves an // adapter error arriving after an SSE response has already started // streaming is logged AND turned into exactly one trailing OpenAI-style // `data: {"error":{...}}` event (finding 15 fix, review-routes.md — // supersedes this test's own former "no trailing envelope" ruling c: a // client stream that just ends, with no terminal signal at all, is // indistinguishable from a silently truncated success to any SDK that // does not strictly require "data: [DONE]"). It is never a second, // conflicting HTTP status/envelope — no second WriteHeader is possible // once headers are already committed — and it is never followed by // "data: [DONE]\n\n", since this is not a clean completion. func TestHandleChat_MidStreamUpstreamDrop_EmitsErrorEventNoDone(t *testing.T) { const firstChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n" srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) fl, ok := w.(http.Flusher) if !ok { t.Fatal("fake upstream ResponseWriter does not support Flush") } _, _ = w.Write([]byte(firstChunk)) fl.Flush() hj, ok := w.(http.Hijacker) if !ok { t.Fatal("fake upstream ResponseWriter does not support Hijack") } conn, _, hjErr := hj.Hijack() if hjErr == nil { _ = conn.Close() } })) defer srv.Close() 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"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } origStderr := os.Stderr pr, pw, _ := os.Pipe() os.Stderr = pw defer func() { os.Stderr = origStderr }() body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) _ = pw.Close() // closing the pipe write end to unblock the read; error not actionable in a test os.Stderr = origStderr var logBuf bytes.Buffer if _, err := io.Copy(&logBuf, pr); err != nil { t.Fatalf("io.Copy: %v", err) } if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200 (headers already committed before the drop)", rec.Code) } const wantErrEvent = "data: {\"error\":{\"message\":\"upstream connection error\",\"type\":\"server_error\"}}\n\n" if got := rec.Body.String(); got != firstChunk+wantErrEvent { t.Errorf("body = %q, want the first chunk plus one trailing error event %q", got, wantErrEvent) } if strings.Contains(rec.Body.String(), "[DONE]") { t.Errorf("body = %q, must not contain [DONE] — a mid-stream failure is not a clean completion", rec.Body.String()) } if !strings.Contains(logBuf.String(), "llmgw[llmgw] ERROR") { t.Errorf("want the mid-stream error logged, got %q", logBuf.String()) } } // TestMidStreamErrorMessage covers every branch of midStreamErrorMessage // (finding 15 fix, review-routes.md): the message/type pair handed to a // mid-stream SSE error event mirrors exactly what the equivalent // non-streamed envelope(...) call would have shown for the same error // type, and a context.Canceled produces ok=false — nothing to write to a // client that is already gone. func TestMidStreamErrorMessage(t *testing.T) { tests := []struct { err error name string wantMsg string wantErrType string wantOK bool }{ { name: "providerHTTPError", err: &providerHTTPError{status: 503, body: []byte(`{"error":"down"}`)}, wantMsg: "openai upstream error", wantErrType: "upstream_error", wantOK: true, }, { name: "translateError", err: &translateError{msg: "n>1 is not supported"}, wantMsg: "n>1 is not supported", wantErrType: "invalid_request_error", wantOK: true, }, { name: "responseTranslationError", err: &responseTranslationError{err: errors.New("boom")}, wantMsg: "failed to translate provider response", wantErrType: "server_error", wantOK: true, }, { name: "generic upstream error", err: fmt.Errorf("%w: connection reset", errUpstream), wantMsg: "upstream connection error", wantErrType: "server_error", wantOK: true, }, { name: "context.Canceled: nothing to write", err: context.Canceled, wantOK: false, }, { name: "wrapped context.Canceled: nothing to write", err: fmt.Errorf("%w: %w", errUpstream, context.Canceled), wantOK: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { msg, errType, ok := midStreamErrorMessage(tt.err, "openai") if ok != tt.wantOK { t.Fatalf("ok = %v, want %v", ok, tt.wantOK) } if !ok { return } if msg != tt.wantMsg { t.Errorf("msg = %q, want %q", msg, tt.wantMsg) } if errType != tt.wantErrType { t.Errorf("errType = %q, want %q", errType, tt.wantErrType) } }) } } // TestHandleChat_MidStreamDropAfterUsageChunk_AccountsPartialUsage is the // regression test for the critical accounting gap: an adapter that // captures a usage frame before the connection drops must still have that // usage billed, even though the request as a whole fails. The fake // upstream sends one content chunk, then the terminal usage-only chunk // (openaiAdapter.forwardStream always captures this into its returned // usage, whether or not the client asked to see it), then hijacks and // closes the connection without a clean chunked terminator — forcing // readSSE to return a genuine error after the usage frame was already // parsed. Without ruling (the CRITICAL fix), that captured usage would be // silently discarded instead of reaching the limiter, letting a client // dodge every token/cost budget by aborting right after the usage chunk // arrives. func TestHandleChat_MidStreamDropAfterUsageChunk_AccountsPartialUsage(t *testing.T) { const contentChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n" const usageChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3}}\n\n" srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) fl, ok := w.(http.Flusher) if !ok { t.Fatal("fake upstream ResponseWriter does not support Flush") } _, _ = w.Write([]byte(contentChunk)) fl.Flush() _, _ = w.Write([]byte(usageChunk)) fl.Flush() hj, ok := w.(http.Hijacker) if !ok { t.Fatal("fake upstream ResponseWriter does not support Hijack") } conn, _, hjErr := hj.Hijack() if hjErr == nil { _ = conn.Close() } })) defer srv.Close() 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{}}, }} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } origStderr := os.Stderr pr, pw, _ := os.Pipe() os.Stderr = pw defer func() { os.Stderr = origStderr }() body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) _ = pw.Close() // closing the pipe write end to unblock the read; error not actionable in a test os.Stderr = origStderr var logBuf bytes.Buffer if _, err := io.Copy(&logBuf, pr); err != nil { t.Fatalf("io.Copy: %v", err) } if !strings.Contains(logBuf.String(), "llmgw[llmgw] ERROR") { t.Errorf("want the mid-stream error logged, got %q", logBuf.String()) } 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) } } // --- response cache: end-to-end through Gateway.ServeHTTP (task 2) --- // behavioralRedisServer is a minimal in-memory RESP2 server implementing // just enough of the protocol (SELECT/AUTH, GET, SET ... EX, INCRBY, // EXPIRE) to drive a real Gateway end to end. Unlike resp_test.go's // respStep-scripted fake server — exact command count and order, used for // resp.go's own unit tests — this one behaves like a tiny real store: the // limiter's counters and the response cache share one respClient // (spec §2), so a scripted sequence would have to predict every INCRBY/ // EXPIRE/GET/SET this test's whole request pipeline issues, in order. // Behaving like real Redis instead of asserting on the wire trace is more // robust for a multi-request end-to-end test, and no less faithful: the // wire-level framing (encodeCommand, RESP replies) is exactly what // resp_test.go's scripted tests already pin. type behavioralRedisServer struct { data map[string]string // lastSetEX is the EX seconds argument from the most recent SET...EX // command this server has handled — item B's per-group cache TTL // override needs a way to observe the actual wire value a real // Gateway.ServeHTTP request produced, not just the reply data's // content. Empty until the first such SET arrives. lastSetEX string mu sync.Mutex } // newBehavioralRedisServer starts the server on an OS-assigned port, // closed automatically at test cleanup (via newFakeListener), and returns // its listener. func newBehavioralRedisServer(t *testing.T) net.Listener { t.Helper() ln, _ := newBehavioralRedisServerAndHandle(t) return ln } // newBehavioralRedisServerAndHandle is newBehavioralRedisServer, additionally // returning the *behavioralRedisServer itself so a caller can inspect state // the wire protocol alone does not expose — namely lastSetEX (see setEX // below), used by TestHandleChat_CachePerGroupTTL_EndToEnd_DifferentEXPerGroup // to prove effectiveTTL's per-group override (cache.go) actually reaches // the SET...EX command a real Gateway.ServeHTTP request issues. func newBehavioralRedisServerAndHandle(t *testing.T) (net.Listener, *behavioralRedisServer) { t.Helper() ln := newFakeListener(t) srv := &behavioralRedisServer{data: make(map[string]string)} go srv.acceptLoop(ln) return ln, srv } // setEX returns the EX seconds argument from the most recent SET...EX // command this server has handled, and whether any such SET has arrived // yet. func (s *behavioralRedisServer) setEX() (string, bool) { s.mu.Lock() defer s.mu.Unlock() return s.lastSetEX, s.lastSetEX != "" } func (s *behavioralRedisServer) acceptLoop(ln net.Listener) { for { conn, err := ln.Accept() if err != nil { return // listener closed by t.Cleanup; test is finishing } go s.serve(conn) } } func (s *behavioralRedisServer) serve(conn net.Conn) { defer func() { _ = conn.Close() }() r := bufio.NewReader(conn) for { args, err := readRESPCommand(r) if err != nil { return } if _, err := conn.Write(s.handle(args)); err != nil { return } } } // handle dispatches one command to its RESP2 reply. SELECT and AUTH // always succeed (this fake never enforces auth or multiple databases); // EXPIRE always succeeds without tracking any real expiry — every test // using this server runs well within any TTL it configures. func (s *behavioralRedisServer) handle(args []string) []byte { if len(args) == 0 { return []byte("-ERR empty command\r\n") } switch strings.ToUpper(args[0]) { case "SELECT", "AUTH", "EXPIRE": return []byte("+OK\r\n") case "GET": s.mu.Lock() v, ok := s.data[args[1]] s.mu.Unlock() if !ok { return []byte("$-1\r\n") } return []byte(fmt.Sprintf("$%d\r\n%s\r\n", len(v), v)) case "SET": s.mu.Lock() s.data[args[1]] = args[2] if len(args) >= 5 && strings.EqualFold(args[3], "EX") { s.lastSetEX = args[4] } s.mu.Unlock() return []byte("+OK\r\n") case "INCRBY": n, _ := strconv.ParseInt(args[2], 10, 64) s.mu.Lock() cur, _ := strconv.ParseInt(s.data[args[1]], 10, 64) cur += n s.data[args[1]] = strconv.FormatInt(cur, 10) s.mu.Unlock() return []byte(fmt.Sprintf(":%d\r\n", cur)) default: return []byte("-ERR unknown command\r\n") } } // newCacheTestGateway builds a *Gateway with cfg.Redis pointed at a // behavioralRedisServer and cfg.Cache enabled, an openai provider pointed // at srv, group "default" (its Cache override left to groupCache, nil // unless the caller sets it), and one user ("alice") with an (empty, // unlimited) LimitsConfig — kept for parity with the other unified-route // fixtures in this file, though since the v0.21 accounting fix // buildLimitScopes always produces alice's own scope regardless of whether // Limits is set at all. func newCacheTestGateway(t *testing.T, srv *httptest.Server, groupCache *bool, maxBodyBytes int) *Gateway { t.Helper() redisLn := newBehavioralRedisServer(t) 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": {Cache: groupCache}} cfg.Users = &UsersConfig{Inline: []*UserConfig{ {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, }} cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m", MaxBodyBytes: maxBodyBytes} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } return gw } // TestHandleChat_CacheHitMiss_EndToEnd_CountersAndHeaders drives two // identical chat-completion requests through a real Gateway with caching // enabled: the first is a miss (upstream called, response stored, // X-Llmgw-Cache: miss), the second is a hit served straight from the // cache (upstream NOT called again, X-Llmgw-Cache: hit) with the exact // same body. Counters prove the operator's "count requests, free tokens" // decision (spec §2): both requests increment req:day by 1, but only the // first (the real upstream call) moves tok:day — the hit adds zero. func TestHandleChat_CacheHitMiss_EndToEnd_CountersAndHeaders(t *testing.T) { const respBody = `{"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}}` var upstreamCalls int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamCalls++ w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() gw := newCacheTestGateway(t, srv, nil, defaultCacheMaxBodyBytes) body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} // First request: miss. req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec1 := httptest.NewRecorder() gw.ServeHTTP(rec1, req1) if rec1.Code != http.StatusOK || rec1.Body.String() != respBody { t.Fatalf("first request: status=%d body=%q, want 200 and the upstream body verbatim", rec1.Code, rec1.Body.String()) } if got := rec1.Header().Get("X-Llmgw-Cache"); got != "miss" { t.Errorf("first request X-Llmgw-Cache = %q, want %q", got, "miss") } if upstreamCalls != 1 { t.Fatalf("upstreamCalls after first request = %d, want 1", upstreamCalls) } reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowDay, time.Now()) if !ok || reqCount != 1 { t.Errorf("req:day after first request = %d (ok=%v), want 1", reqCount, 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. req2 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec2 := httptest.NewRecorder() gw.ServeHTTP(rec2, req2) if rec2.Code != http.StatusOK || rec2.Body.String() != respBody { t.Fatalf("second request: status=%d body=%q, want 200 and the identical cached body", rec2.Code, rec2.Body.String()) } if got := rec2.Header().Get("X-Llmgw-Cache"); got != "hit" { t.Errorf("second request X-Llmgw-Cache = %q, want %q", got, "hit") } if upstreamCalls != 1 { t.Errorf("upstreamCalls after second (cache-hit) request = %d, want still 1 (upstream must not be called again)", upstreamCalls) } reqCount, ok = gw.limiter.getCounter("user", "alice", metricReq, windowDay, time.Now()) 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) } 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) } } // TestHandleChat_CacheGroupOptOut_NeverServesFromCache is spec §2's // per-group override: GroupConfig.Cache=false must opt the group out even // though the global cache is enabled — every request goes to the // upstream, and no X-Llmgw-Cache header is ever set. func TestHandleChat_CacheGroupOptOut_NeverServesFromCache(t *testing.T) { const respBody = `{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` var upstreamCalls int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamCalls++ w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() optOut := false gw := newCacheTestGateway(t, srv, &optOut, defaultCacheMaxBodyBytes) body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} for i := 0; i < 2; i++ { req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() gw.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("request %d: status = %d, want 200", i, rec.Code) } if got := rec.Header().Get("X-Llmgw-Cache"); got != "" { t.Errorf("request %d: X-Llmgw-Cache = %q, want unset (group opted out)", i, got) } } if upstreamCalls != 2 { t.Errorf("upstreamCalls = %d, want 2 (an opted-out group must never be served from cache)", upstreamCalls) } } // TestHandleChat_CacheOversizeBody_SkipsStore_StillServesEveryRequest // proves an upstream body larger than cache.maxBodyBytes is served // normally but never cached: a second identical request still goes to // the upstream, since store() silently skipped the first one. func TestHandleChat_CacheOversizeBody_SkipsStore_StillServesEveryRequest(t *testing.T) { respBody := `{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"` + strings.Repeat("x", 200) + `"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` var upstreamCalls int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamCalls++ w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() gw := newCacheTestGateway(t, srv, nil, 32) // maxBodyBytes far smaller than respBody body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} for i := 0; i < 2; i++ { req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() gw.ServeHTTP(rec, req) if rec.Code != http.StatusOK || rec.Body.String() != respBody { t.Fatalf("request %d: status=%d body=%q, want 200 and the full upstream body regardless of the cache's maxBodyBytes", i, rec.Code, rec.Body.String()) } if got := rec.Header().Get("X-Llmgw-Cache"); got != "miss" { t.Errorf("request %d: X-Llmgw-Cache = %q, want %q (oversize is a store-time skip, not a lookup-time change)", i, got, "miss") } } if upstreamCalls != 2 { t.Errorf("upstreamCalls = %d, want 2 (an oversize response must never be served from cache on a later request)", upstreamCalls) } } // TestHandleChat_StreamingRequest_NeverCachedOrServedFromCache is the // direct regression for runMeteredCall's own `cacheable := !streaming && // ...` guard (routes_unified.go): a streaming chat-completion request // must never get an X-Llmgw-Cache header at all (it is never looked up // nor stored — the whole cacheable block is skipped), and nothing the // streaming call did may have populated the cache under that body's own // key. // // MUTATION PROVEN: changing `cacheable := !streaming && g.cache != nil && // groupCacheEnabled(grp)` to drop the `!streaming` term made this test // fail — the streaming response picked up an X-Llmgw-Cache: miss header // — confirmed via `go test -run // TestHandleChat_StreamingRequest_NeverCachedOrServedFromCache`, then // reverted with `git checkout -- routes_unified.go`. func TestHandleChat_StreamingRequest_NeverCachedOrServedFromCache(t *testing.T) { var upstreamCalls int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamCalls++ w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) fl, ok := w.(http.Flusher) if !ok { t.Fatal("fake upstream ResponseWriter does not support Flush") } _, _ = w.Write([]byte("data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n")) fl.Flush() _, _ = w.Write([]byte("data: [DONE]\n\n")) fl.Flush() })) defer srv.Close() gw := newCacheTestGateway(t, srv, nil, defaultCacheMaxBodyBytes) body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rw := newRecordingWriter() gw.ServeHTTP(rw, req) if rw.status != http.StatusOK { t.Fatalf("status = %d, want 200", rw.status) } if got := rw.Header().Get("X-Llmgw-Cache"); got != "" { t.Errorf(`X-Llmgw-Cache = %q for a streaming request, want no header at all (never cacheable)`, got) } if upstreamCalls != 1 { t.Fatalf("upstreamCalls = %d, want 1", upstreamCalls) } // Nothing the streaming call did may have populated the cache under // the key a non-streaming request with the same provider/model/body // would have used. key := cacheKey("openai", "gpt-test", "gpt-test", cacheEndpointChat, body) if _, hit := gw.cache.lookup(key); hit { t.Error("cache has an entry for the streaming request's key, want none stored") } } // TestHandleChat_CacheRedisDown_MissesAndServesNormallyWithLog proves // spec §2's "redis errors during cache ops → treat as miss ... never fail // the request because the cache is down": cfg.Redis points at an address // nothing listens on (a respClient dials lazily, so construction still // succeeds), so every lookup/store call fails, yet the request is served // normally every time, with a "miss" header and a rate-limited error log // line. The rate-limit behavior itself (once per storeErrorLogEvery) is // already pinned precisely at the responseCache level // (TestResponseCache_Lookup_RedisDown_MissWithRateLimitedLog, // cache_test.go); this test only proves the failure never surfaces to // the client end to end. func TestHandleChat_CacheRedisDown_MissesAndServesNormallyWithLog(t *testing.T) { deadLn := newFakeListener(t) deadAddr := deadLn.Addr().String() if err := deadLn.Close(); err != nil { t.Fatalf("Close: %v", err) } const respBody = `{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` 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(respBody)) })) defer srv.Close() 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"}}} cfg.Redis = &RedisConfig{Address: deadAddr} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } origStderr := os.Stderr pr, pw, _ := os.Pipe() os.Stderr = pw body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) _ = pw.Close() os.Stderr = origStderr var logBuf bytes.Buffer if _, err := io.Copy(&logBuf, pr); err != nil { t.Fatalf("io.Copy: %v", err) } if rec.Code != http.StatusOK || rec.Body.String() != respBody { t.Fatalf("status=%d body=%q, want 200 and the upstream body served despite redis being down", rec.Code, rec.Body.String()) } if got := rec.Header().Get("X-Llmgw-Cache"); got != "miss" { t.Errorf("X-Llmgw-Cache = %q, want %q (redis down is treated as a miss)", got, "miss") } if !strings.Contains(logBuf.String(), "response cache: redis error") { t.Errorf("want a logged cache redis error, got %q", logBuf.String()) } } // TestHandleChat_AnthropicTranslatedResponse_CachedAboveAdapter proves // the cache sits above the adapter layer: an anthropic-type provider's // translated, OpenAI-shaped response is what gets cached, and a second // identical request is served that exact translated body from the cache // — with the anthropic upstream called only once. func TestHandleChat_AnthropicTranslatedResponse_CachedAboveAdapter(t *testing.T) { const anthResp = `{"id":"msg_01ABC","type":"message","role":"assistant","content":[{"type":"text","text":"hi there"}],"model":"claude-x","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":4}}` var upstreamCalls int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamCalls++ w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(anthResp)) })) defer srv.Close() redisLn := newBehavioralRedisServer(t) cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-ant", Models: []string{"claude-x"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "claude-x", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec1 := httptest.NewRecorder() h.ServeHTTP(rec1, req1) if rec1.Code != http.StatusOK { t.Fatalf("first request: status = %d, body=%s", rec1.Code, rec1.Body.String()) } if got := rec1.Header().Get("X-Llmgw-Cache"); got != "miss" { t.Errorf("first request X-Llmgw-Cache = %q, want %q", got, "miss") } var out1 map[string]any if err := json.Unmarshal(rec1.Body.Bytes(), &out1); err != nil { t.Fatalf("decode first response: %v", err) } if _, ok := out1["choices"]; !ok { t.Fatalf("first response is not OpenAI-shaped (no \"choices\"): %s", rec1.Body.String()) } req2 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec2 := httptest.NewRecorder() h.ServeHTTP(rec2, req2) if rec2.Code != http.StatusOK { t.Fatalf("second request: status = %d, body=%s", rec2.Code, rec2.Body.String()) } if got := rec2.Header().Get("X-Llmgw-Cache"); got != "hit" { t.Errorf("second request X-Llmgw-Cache = %q, want %q", got, "hit") } if rec2.Body.String() != rec1.Body.String() { t.Errorf("second (cached) response body = %q, want identical to the first translated response %q", rec2.Body.String(), rec1.Body.String()) } if upstreamCalls != 1 { t.Errorf("upstreamCalls = %d, want 1 (the anthropic upstream must not be called again on a cache hit)", upstreamCalls) } } // TestHandleChat_CacheMixedAliasForm_NeverHits_SameFormHits is the // review-fix (finding 1, CRITICAL) end-to-end regression: a bare alias // ("claude-x") and its provider-prefixed form ("anthropic/claude-x") // resolve to the identical provider+upstreamModel, but a translating // adapter bakes the client's own requested alias into the cached // response body's "model" field — so a request using one form must never // be served from an entry cached under the other form, or the client // would receive a "model" value it never sent. The same form repeated // must still hit. func TestHandleChat_CacheMixedAliasForm_NeverHits_SameFormHits(t *testing.T) { const anthResp = `{"id":"msg_01ABC","type":"message","role":"assistant","content":[{"type":"text","text":"hi there"}],"model":"claude-x","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":4}}` var upstreamCalls int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamCalls++ w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(anthResp)) })) defer srv.Close() redisLn := newBehavioralRedisServer(t) cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-ant", Models: []string{"claude-x"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } bareBody := map[string]any{"model": "claude-x", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} prefixedBody := map[string]any{"model": "anthropic/claude-x", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} // A: bare alias form. Miss, cached under the bare-alias key. reqA := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", bareBody) recA := httptest.NewRecorder() h.ServeHTTP(recA, reqA) if recA.Code != http.StatusOK || recA.Header().Get("X-Llmgw-Cache") != "miss" { t.Fatalf("request A: status=%d cache=%q, want 200/miss", recA.Code, recA.Header().Get("X-Llmgw-Cache")) } var outA map[string]any if err := json.Unmarshal(recA.Body.Bytes(), &outA); err != nil { t.Fatalf("decode A: %v", err) } if outA["model"] != "claude-x" { t.Fatalf("request A model = %v, want %q", outA["model"], "claude-x") } if upstreamCalls != 1 { t.Fatalf("upstreamCalls after A = %d, want 1", upstreamCalls) } // B: same upstream model, PREFIXED alias form. Must be a genuine // miss, never served A's cached "claude-x" body. reqB := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", prefixedBody) recB := httptest.NewRecorder() h.ServeHTTP(recB, reqB) if recB.Code != http.StatusOK { t.Fatalf("request B: status = %d, body=%s", recB.Code, recB.Body.String()) } if got := recB.Header().Get("X-Llmgw-Cache"); got != "miss" { t.Errorf("request B X-Llmgw-Cache = %q, want %q (a different alias form must never hit A's entry)", got, "miss") } var outB map[string]any if err := json.Unmarshal(recB.Body.Bytes(), &outB); err != nil { t.Fatalf("decode B: %v", err) } if outB["model"] != "anthropic/claude-x" { t.Fatalf("request B model = %v, want its own requested alias %q, never A's cached %q", outB["model"], "anthropic/claude-x", "claude-x") } if upstreamCalls != 2 { t.Fatalf("upstreamCalls after B = %d, want 2 (a different alias form is a genuine miss)", upstreamCalls) } // C: bare form again, matching A exactly. Must hit A's entry. reqC := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", bareBody) recC := httptest.NewRecorder() h.ServeHTTP(recC, reqC) if got := recC.Header().Get("X-Llmgw-Cache"); got != "hit" { t.Errorf("request C X-Llmgw-Cache = %q, want %q (same alias form as A must hit)", got, "hit") } if recC.Body.String() != recA.Body.String() { t.Errorf("request C body = %q, want identical to A's cached body %q", recC.Body.String(), recA.Body.String()) } if upstreamCalls != 2 { t.Errorf("upstreamCalls after C = %d, want still 2 (C is a real hit, no upstream call)", upstreamCalls) } } // TestHandleChat_CacheableRequest_UpstreamDown_502HasNoCacheHeader is the // review-fix (finding 4, folded) regression: a cacheable request whose // upstream is unreachable gets a 502 with NO X-Llmgw-Cache header — the // pre-set "miss" header (set before call() runs, so headers precede a // successful body) must be cleared once handleAdapterError decides the // request actually failed, since nothing was served from or stored to // the cache. func TestHandleChat_CacheableRequest_UpstreamDown_502HasNoCacheHeader(t *testing.T) { deadLn := newFakeListener(t) deadAddr := deadLn.Addr().String() if err := deadLn.Close(); err != nil { t.Fatalf("Close: %v", err) } redisLn := newBehavioralRedisServer(t) cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: "http://" + deadAddr, APIKey: "k", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusBadGateway { t.Fatalf("status = %d, want 502 (connection refused), body=%s", rec.Code, rec.Body.String()) } if got := rec.Header().Get("X-Llmgw-Cache"); got != "" { t.Errorf("X-Llmgw-Cache = %q, want unset on a 502 (nothing was served from or stored to the cache)", got) } } // --- per-group cache TTL: end-to-end through a real Gateway (item B) --- // TestHandleChat_CachePerGroupTTL_EndToEnd_DifferentEXPerGroup is item B's // e2e proof: two groups share one global cache.ttl ("1m" = EX 60), one // ("override") overrides it to "30s", the other ("inherit") leaves it // unset. A cache-miss request from each group's own user must produce a // SET...EX command carrying that group's own resolved TTL — 30 for // "override", 60 for "inherit" — proving effectiveTTL's per-group // resolution (cache.go) actually reaches the wire through the full // ServeHTTP path, not just the unit level // (TestResponseCache_Store_EffectiveTTLDiffersPerGroup_SETCarriesGroupTTL, // cache_test.go). The two requests carry deliberately different bodies — // cacheKey has no group/user component (routes_unified.go), so identical // bodies from different groups would collide into one cache entry and the // second request would be a hit, never issuing a second SET at all. func TestHandleChat_CachePerGroupTTL_EndToEnd_DifferentEXPerGroup(t *testing.T) { const respBody = `{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` 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(respBody)) })) defer srv.Close() redisLn, redisSrv := newBehavioralRedisServerAndHandle(t) cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{ "override": {CacheTTL: "30s"}, "inherit": {}, } cfg.Users = &UsersConfig{Inline: []*UserConfig{ {Name: "bob", Group: "override", APIKey: "sk-bob", Limits: &LimitsConfig{}}, {Name: "alice", Group: "inherit", APIKey: "sk-alice", Limits: &LimitsConfig{}}, }} cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m", MaxBodyBytes: defaultCacheMaxBodyBytes} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } // "override" group's user (bob): expect EX 30. overrideBody := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "override-group"}}} reqOverride := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-bob", overrideBody) recOverride := httptest.NewRecorder() h.ServeHTTP(recOverride, reqOverride) if recOverride.Code != http.StatusOK { t.Fatalf("override-group request: status = %d, want 200", recOverride.Code) } if ex, ok := redisSrv.setEX(); !ok || ex != "30" { t.Errorf("override-group SET...EX = %q (ok=%v), want %q", ex, ok, "30") } // "inherit" group's user (alice), a genuinely different body so it // misses (and issues its own SET) rather than hitting bob's entry: // expect EX 60 (the global 1m). inheritBody := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "inherit-group"}}} reqInherit := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", inheritBody) recInherit := httptest.NewRecorder() h.ServeHTTP(recInherit, reqInherit) if recInherit.Code != http.StatusOK { t.Fatalf("inherit-group request: status = %d, want 200", recInherit.Code) } if ex, ok := redisSrv.setEX(); !ok || ex != "60" { t.Errorf("inherit-group SET...EX = %q (ok=%v), want %q", ex, ok, "60") } } // --- cacheCaptureWriter: bounded capture (review finding 2) --- // TestCacheCaptureWriter_BoundsBufferAtMaxBodyBytes_StillWritesFullBodyThrough // is the review-fix (finding 2, Important) regression: a single Write of // a body larger than maxBodyBytes must not buffer the whole thing — // capture stops at the cap — while the real client-visible write is // never truncated. func TestCacheCaptureWriter_BoundsBufferAtMaxBodyBytes_StillWritesFullBodyThrough(t *testing.T) { rec := httptest.NewRecorder() sw := &statusTrackingWriter{ResponseWriter: rec} const maxBody = 16 cw := newCacheCaptureWriter(sw, maxBody) fullBody := []byte(strings.Repeat("x", maxBody+50)) // one write, well over the cap cw.Header().Set("Content-Type", "application/json") cw.WriteHeader(http.StatusOK) n, err := cw.Write(fullBody) if err != nil { t.Fatalf("Write: %v", err) } if n != len(fullBody) { t.Errorf("Write returned n=%d, want %d (the real client write must never be truncated)", n, len(fullBody)) } if !cw.oversize { t.Error("oversize = false, want true") } // "buffer <= cap + one write": this implementation holds the capture // at exactly the cap, a tighter bound the looser one still accepts. if cw.buf.Len() > maxBody { t.Errorf("captured buffer = %d bytes, want <= maxBodyBytes (%d)", cw.buf.Len(), maxBody) } if rec.Body.String() != string(fullBody) { t.Errorf("client-visible body = %d bytes, want the full %d-byte body (capture must never truncate the real response)", rec.Body.Len(), len(fullBody)) } } // TestCacheCaptureWriter_MultipleWrites_StopsCapturingOnceCapReached // covers the multi-write case: capture stops appending once the cap is // crossed on a later write, not just a single oversize one, while every // write still reaches the real client in full. func TestCacheCaptureWriter_MultipleWrites_StopsCapturingOnceCapReached(t *testing.T) { rec := httptest.NewRecorder() sw := &statusTrackingWriter{ResponseWriter: rec} const maxBody = 10 cw := newCacheCaptureWriter(sw, maxBody) cw.WriteHeader(http.StatusOK) if _, err := cw.Write([]byte("12345")); err != nil { // 5 bytes, under the cap t.Fatalf("Write: %v", err) } if cw.oversize { t.Fatal("oversize = true after a write under the cap") } if _, err := cw.Write([]byte("67890ABCDE")); err != nil { // 10 more bytes, crosses the 10-byte cap t.Fatalf("Write: %v", err) } if !cw.oversize { t.Error("oversize = false, want true once the cap is crossed") } if cw.buf.Len() != maxBody { t.Errorf("captured buffer = %d bytes, want exactly maxBodyBytes (%d)", cw.buf.Len(), maxBody) } if rec.Body.String() != "1234567890ABCDE" { t.Errorf("client-visible body = %q, want the full, unsplit content across both writes", rec.Body.String()) } } // TestCacheCaptureWriter_ImplicitStatus_CapturesContentType is the // review-fix (finding 4, folded) regression: the implicit-200 Write path // (no explicit WriteHeader call) must still snapshot Content-Type, not // just status — mirroring WriteHeader's own capture. func TestCacheCaptureWriter_ImplicitStatus_CapturesContentType(t *testing.T) { rec := httptest.NewRecorder() sw := &statusTrackingWriter{ResponseWriter: rec} cw := newCacheCaptureWriter(sw, defaultCacheMaxBodyBytes) cw.Header().Set("Content-Type", "application/json") // set, but WriteHeader never explicitly called if _, err := cw.Write([]byte(`{"ok":true}`)); err != nil { t.Fatalf("Write: %v", err) } if cw.status != http.StatusOK { t.Errorf("status = %d, want 200 (implicit)", cw.status) } if cw.contentType != "application/json" { t.Errorf("contentType = %q, want %q (must be captured on the implicit-200 path too)", cw.contentType, "application/json") } } // TestUnifiedCostMicros covers ruling (d)'s two-id price lookup order: // canonical ("provider/model") first, falling back to bare — the upstream // model id alone — only when canonical has no configured price at all, in // neither overrides nor the built-in table. // TestBuildLimitScopes_AlwaysBuildsBothScopes proves buildLimitScopes // builds a user AND a group limitScope unconditionally (v0.21 accounting // fix): whether either entity has Limits configured or not, both scopes // must always be present — a nil Limits field turns off ENFORCEMENT // (checkAndCount's own nil check, limits.go) for that scope, never // accounting. Table-driven over every nil/non-nil combination of user and // group limits. func TestBuildLimitScopes_AlwaysBuildsBothScopes(t *testing.T) { cases := []struct { userLimits *LimitsConfig grpLimits *LimitsConfig name string }{ {name: "both nil (the exact bug scenario)"}, {name: "user limited, group unlimited", userLimits: &LimitsConfig{RequestsPerMinute: 5}}, {name: "user unlimited, group limited", grpLimits: &LimitsConfig{RequestsPerMinute: 5}}, {name: "both limited", userLimits: &LimitsConfig{RequestsPerMinute: 5}, grpLimits: &LimitsConfig{RequestsPerMinute: 9}}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { u := &user{name: "u1", limits: tc.userLimits} grp := &group{name: "g1", limits: tc.grpLimits} scopes := buildLimitScopes(u, grp) if len(scopes) != 2 { t.Fatalf("len(scopes) = %d, want 2 (unconditional user+group accounting)", len(scopes)) } if scopes[0].kind != "user" || scopes[0].id != "u1" || scopes[0].limits != tc.userLimits { t.Errorf("scopes[0] = %+v, want user scope for u1 with limits %+v", scopes[0], tc.userLimits) } if scopes[1].kind != "group" || scopes[1].id != "g1" || scopes[1].limits != tc.grpLimits { t.Errorf("scopes[1] = %+v, want group scope for g1 with limits %+v", scopes[1], tc.grpLimits) } }) } } // TestBuildLimitScopes_MultiGroup_OneScopePerMemberGroup proves a // multi-group principal's own limit scopes are one PER MEMBER GROUP, // each under its own name and its own limits — never merged into one // combined scope, and never carrying the synthetic principal's own "+" // -joined name. func TestBuildLimitScopes_MultiGroup_OneScopePerMemberGroup(t *testing.T) { a := &group{name: "eng", limits: &LimitsConfig{RequestsPerMinute: 5}} b := &group{name: "ops", limits: &LimitsConfig{RequestsPerMinute: 9}} grp := effectiveGroup([]*group{a, b}, nil) u := &user{name: "u1"} scopes := buildLimitScopes(u, grp) if len(scopes) != 3 { t.Fatalf("len(scopes) = %d, want 3 (1 user + 2 member groups)", len(scopes)) } if scopes[0].kind != "user" || scopes[0].id != "u1" { t.Errorf("scopes[0] = %+v, want the user scope first", scopes[0]) } got := map[string]*LimitsConfig{} for _, s := range scopes[1:] { if s.kind != "group" { t.Errorf("scope kind = %q, want %q", s.kind, "group") } got[s.id] = s.limits } if got["eng"] != a.limits || got["ops"] != b.limits { t.Errorf("group scopes = %v, want each member group's own limits under its own name", got) } } func TestUnifiedCostMicros(t *testing.T) { u := usage{prompt: 1_000_000, completion: 1_000_000} cases := []struct { overrides map[string]*ModelPricing meta map[string]*ModelMetaConfig name string canonical string bare string want int64 }{ { name: "canonical override price wins over bare", canonical: "openai/gpt-5", bare: "gpt-5", overrides: map[string]*ModelPricing{ "openai/gpt-5": {InputPerM: 1, OutputPerM: 2}, }, want: 3_000_000, // (1M prompt * $1/M) + (1M completion * $2/M), in micros }, { name: "canonical unpriced falls back to bare's built-in price", canonical: "unknown-provider/gpt-5", bare: "gpt-5", overrides: nil, // builtinPricing["gpt-5"] = {InputPerM: 1.25, OutputPerM: 10}: // (1M prompt * $1.25/M) + (1M completion * $10/M), in micros. want: 11_250_000, }, { name: "neither id priced returns zero", canonical: "unknown-provider/unknown-model", bare: "unknown-model", overrides: nil, want: 0, }, { // verify-core round-4 regression: gpt-oss-120b is a real bare // id in the generated builtinModelMetaTable (a common self- // hosted model name). An operator declaring it free:true under // a self-hosted canonical id must bill 0, never the table's // real price — priceKnown already treats this model as priced // (modelMetaFree, modelmeta.go); billing must agree. name: "modelMeta free:true bills zero even though bare id has a built-in table price", canonical: "gx10/gpt-oss-120b", bare: "gpt-oss-120b", overrides: nil, meta: map[string]*ModelMetaConfig{"gx10/gpt-oss-120b": {Free: true}}, want: 0, }, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { assert.Equal(t, c.want, unifiedCostMicros(c.canonical, c.bare, u, c.overrides, c.meta, func(string) {})) }) } } // TestHandleChat_ContextDeadlineExceeded_RecordsProviderFailure is the // route-level regression test the review's blocker (v0.22, round 2) // specifically demanded: TestIsDeadlineExceeded_YaegiSafeUnwrap's // hand-built error shapes (limits_test.go) proved matchesSentinel's own // unwrap logic works in isolation, but that alone is exactly what let the // real bug ship — every hand-built shape there was single-%w, while the // live chain through upstreamBytes' client.Do call is providers.go's // own double-%w `fmt.Errorf("%w: %w", errUpstream, err)`, one layer // short of what a table test alone would ever exercise. This drives a // REAL request whose context deadline expires while a REAL httptest // upstream is still sleeping — client.Do, *url.Error, and the wrap chain // providers.go actually builds, all for real — and asserts the resulting // provider counters directly, not a synthetic error value. func TestHandleChat_ContextDeadlineExceeded_RecordsProviderFailure(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(600 * time.Millisecond) // well past the request's own 60ms deadline below w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[],"usage":{}}`)) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{ {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, }} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } // SHOULD-5: run recordProviderAttempt's store write synchronously so // the counter read below is deterministic (see limiter.spawn's own // doc comment, limits.go). FOLDED-2's own test (limits_test.go) // covers the production async default separately. gw.limiter.spawn = func(f func()) { f() } ctx, cancel := context.WithTimeout(context.Background(), 60*time.Millisecond) defer cancel() req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-test","messages":[{"role":"user","content":"hi"}]}`)).WithContext(ctx) req.Header.Set("Authorization", "Bearer sk-alice") req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() h.ServeHTTP(rec, req) now := time.Now() attempts, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, now) if attempts != 1 { t.Errorf("provider attempts/day = %d, want 1", attempts) } fails, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvFail, windowDay, now) if fails != 1 { t.Errorf("provider fails/day = %d, want 1 — a real context-deadline timeout must count as a provider-health failure (SHOULD-1)", fails) } } // TestHandleChat_UnpricedModelCostBudget_Returns402AndRecordsEvent is the // unified route's own analog of routes_passthrough_test.go's // TestHandlePassthrough_UnpricedModelCostBudget_Returns402: it drives a // REAL request for an unpriced model against a caller with a configured // cost budget through the security-audit run-1 finding F-1 gate // (routes_unified.go, just before call()), and additionally proves item 5 // (this round): the refusal now reaches GET /admin/api/events as an // eventKindUnpriced/402 entry, not just a logged line — the ONLY route- // level coverage this specific 402 gate had before this test was the // upstream body never being reached (implicitly, via other tests); this // asserts the response, the event, and that the upstream was never // called, together. func TestHandleChat_UnpricedModelCostBudget_Returns402AndRecordsEvent(t *testing.T) { var upstreamCalled bool srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamCalled = true w.WriteHeader(http.StatusOK) })) defer srv.Close() cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"totally-unpriced-model"}}, } cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{ {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{CostPerDayUSD: 5}}, }} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } rec := httptest.NewRecorder() h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", map[string]any{"model": "totally-unpriced-model", "messages": []any{map[string]any{"role": "user", "content": "hi"}}})) if rec.Code != http.StatusPaymentRequired { t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusPaymentRequired, rec.Body.String()) } if upstreamCalled { t.Error("upstream must never be called: the refusal fires before call()") } events := gw.events.ring.snapshot(eventRingCap) if len(events) == 0 { t.Fatal("want an event recorded for the 402 refusal, got none") } ev := events[0] // newest first if ev.Kind != eventKindUnpriced { t.Errorf("Kind = %q, want %q", ev.Kind, eventKindUnpriced) } if ev.Status != http.StatusPaymentRequired { t.Errorf("Status = %d, want %d", ev.Status, http.StatusPaymentRequired) } if ev.Route != routeChatCompletions { t.Errorf("Route = %q, want %q", ev.Route, routeChatCompletions) } if ev.User != "alice" { t.Errorf("User = %q, want %q", ev.User, "alice") } } // --- security review finding 1, 2026-08-22: admission before decode, // body-admission semaphore, narrower unified body cap --- // readTrackingBody is an io.ReadCloser that records whether Read was ever // invoked, wrapping a real reader so a request that DOES get read still // decodes normally. Used to prove a request rejected before the body is // read (finding 1a) never actually touches it. type readTrackingBody struct { r io.Reader called bool } func (b *readTrackingBody) Read(p []byte) (int, error) { b.called = true return b.r.Read(p) } func (b *readTrackingBody) Close() error { return nil } // TestHandleChat_RequestLimitExceeded_BodyNeverRead is the GATE's own // explicit requirement: a 429'd user's body must never be read at all, // proving admitRequestForRoute now runs strictly before io.ReadAll in runUnified // (finding 1a) rather than merely returning the same status code for a // different reason. func TestHandleChat_RequestLimitExceeded_BodyNeverRead(t *testing.T) { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{ "openai": {Type: "openai", 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{RequestsPerMinute: 1}}, }} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } // First request consumes the one-per-minute budget; its body is a // real, valid body so the request completes normally end to end // (upstream is unreachable — irrelevant here, only the request // counter matters). req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", map[string]any{"model": "gpt-test", "messages": []any{}}) rec1 := httptest.NewRecorder() h.ServeHTTP(rec1, req1) // Second request: now over budget. Its body is wrapped in // readTrackingBody — if runUnified reads it before checking the // limit, called flips to true. tracked := &readTrackingBody{r: strings.NewReader(`{"model":"gpt-test","messages":[]}`)} req2 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", tracked) req2.Header.Set("Authorization", "Bearer sk-alice") rec2 := httptest.NewRecorder() h.ServeHTTP(rec2, req2) if rec2.Code != http.StatusTooManyRequests { t.Fatalf("status = %d, want 429, body=%s", rec2.Code, rec2.Body.String()) } if tracked.called { t.Error("a 429'd request's body must never be read — admitRequestForRoute must run before io.ReadAll") } } // TestHandleImagesGenerations_RequestLimitExceeded_BodyNeverRead mirrors // the unified-route test above for the media JSON routes, which share the // identical admission-before-decode ordering (finding 1a). func TestHandleImagesGenerations_RequestLimitExceeded_BodyNeverRead(t *testing.T) { cfg := newMediaTestConfig("http://127.0.0.1:1", "img-test") cfg.Users = &UsersConfig{Inline: []*UserConfig{ {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerMinute: 1}}, }} gw := newMediaTestGateway(t, cfg) req1 := newUnifiedRequest(t, http.MethodPost, imagesGenerationsPath, "sk-alice", map[string]any{"model": "img-test", "prompt": "a cat"}) rec1 := httptest.NewRecorder() gw.ServeHTTP(rec1, req1) tracked := &readTrackingBody{r: strings.NewReader(`{"model":"img-test","prompt":"a cat"}`)} req2 := httptest.NewRequest(http.MethodPost, imagesGenerationsPath, tracked) req2.Header.Set("Authorization", "Bearer sk-alice") rec2 := httptest.NewRecorder() gw.ServeHTTP(rec2, req2) if rec2.Code != http.StatusTooManyRequests { t.Fatalf("status = %d, want 429, body=%s", rec2.Code, rec2.Body.String()) } if tracked.called { t.Error("a 429'd media request's body must never be read — admitRequestForRoute must run before decode") } } // TestAcquireBodyAdmission_ExhaustionReturns503WithRetryAfter is the // GATE's own explicit requirement: on semaphore exhaustion, // acquireBodyAdmission must refuse immediately with 503 and a // Retry-After header, never block. gw.bodyAdmission is filled directly // (same package — this test manipulates the exact channel newGateway // constructed) to make exhaustion deterministic without needing real // concurrent goroutines. func TestAcquireBodyAdmission_ExhaustionReturns503WithRetryAfter(t *testing.T) { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}} cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} cfg.MaxInFlightBodyRequests = 1 next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } if cap(gw.bodyAdmission) != 1 { t.Fatalf("bodyAdmission cap = %d, want 1 (explicit override)", cap(gw.bodyAdmission)) } // Occupy the only slot directly, simulating another in-flight request. gw.bodyAdmission <- struct{}{} defer func() { <-gw.bodyAdmission }() req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", map[string]any{"model": "gpt-test", "messages": []any{}}) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusServiceUnavailable { t.Fatalf("status = %d, want 503, body=%s", rec.Code, rec.Body.String()) } if ra := rec.Header().Get("Retry-After"); ra == "" { t.Error("want a Retry-After header on 503, got none") } } // TestNewGateway_BodyAdmissionCap_SelfTunesFromGOMAXPROCS_ExplicitOverrideWins // pins both halves of the house rule (self-tuning over operator knobs, // but an explicit override always wins): a zero-config MaxInFlightBodyRequests // self-tunes to defaultBodyAdmissionCap(), and a positive value overrides // it exactly, regardless of what GOMAXPROCS would otherwise infer. func TestNewGateway_BodyAdmissionCap_SelfTunesFromGOMAXPROCS_ExplicitOverrideWins(t *testing.T) { baseCfg := func() *Config { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} cfg.Groups = map[string]*GroupConfig{"default": {}} return cfg } t.Run("zero value self-tunes", func(t *testing.T) { cfg := baseCfg() h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw := h.(*Gateway) if got, want := cap(gw.bodyAdmission), defaultBodyAdmissionCap(); got != want { t.Errorf("bodyAdmission cap = %d, want %d (defaultBodyAdmissionCap)", got, want) } }) t.Run("explicit override wins", func(t *testing.T) { cfg := baseCfg() cfg.MaxInFlightBodyRequests = 7 h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw := h.(*Gateway) if got := cap(gw.bodyAdmission); got != 7 { t.Errorf("bodyAdmission cap = %d, want 7 (explicit override)", got) } }) } // TestDefaultBodyAdmissionCap_ClampedToBounds pins defaultBodyAdmissionCap's // own contract directly: its result always lies within // [minBodyAdmissionCap, maxBodyAdmissionCap], regardless of the host's own // GOMAXPROCS. func TestDefaultBodyAdmissionCap_ClampedToBounds(t *testing.T) { got := defaultBodyAdmissionCap() if got < minBodyAdmissionCap || got > maxBodyAdmissionCap { t.Errorf("defaultBodyAdmissionCap() = %d, want within [%d, %d]", got, minBodyAdmissionCap, maxBodyAdmissionCap) } } // TestHandleChat_ZeroConfig_ConcurrentAboveDefaultCap_NoThrottling_SlotsReleasedBeforeUpstream // is the DEFAULT-PRESERVING gate's own explicit requirement, replacing a // SEQUENTIAL version that could not actually exercise it (MUST-2, round // 3, 2026-08-22 review: the reviewer re-ran the sequential test's exact // shape with MaxInFlightBodyRequests=1 and it still passed — a // capacity-one semaphore is never contended by requests that arrive one // at a time, so it was not evidence the semaphore leaves ordinary // traffic unthrottled). // // This drives n = defaultBodyAdmissionCap()+20 requests genuinely // CONCURRENTLY (real goroutines, real network round trips to an // httptest.Server upstream — n is deliberately larger than the // self-tuned default cap) against an upstream gated to block until every // one of them has arrived, and asserts two things: // // 1. While all n requests are parked waiting on the gated upstream — // more in-flight requests than admission slots exist — the // semaphore (gw.bodyAdmission) holds ZERO slots. This is the // regression's own regression test (MUST-1, round 3): before that // fix, a slot was held for the WHOLE handler including this exact // upstream wait, so n > cap concurrent parked requests would have // left the semaphore fully occupied (and, with a large enough n // relative to cap, some later arrivals would have found it // exhausted). With the fix, a slot is released the instant // read+decode finishes, strictly before the upstream call ever // starts — by the time a request reaches the gated upstream, it is // provably holding no slot at all. // 2. Once every parked request completes, all n get 200 — none was // ever refused with 503 merely because n other requests happened to // be concurrently in flight, which is exactly what "zero-config // ordinary traffic is not throttled" has to mean on an I/O-bound // gateway where many requests are routinely in flight for seconds // at a time. func TestHandleChat_ZeroConfig_ConcurrentAboveDefaultCap_NoThrottling_SlotsReleasedBeforeUpstream(t *testing.T) { const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` // gate blocks every upstream response until the test explicitly // closes it — simulating an LLM completion parked on a slow // upstream for a while, the exact scenario the MUST-1 regression // measured (a real completion is I/O-bound for seconds to minutes). gate := make(chan struct{}) var arrived sync.WaitGroup srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { arrived.Done() // this request's admission slot is already released by now — see below <-gate w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) })) defer srv.Close() 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"}}} // Deliberately no MaxInFlightBodyRequests set — every config that // predates this finding. next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } n := cap(gw.bodyAdmission) + 20 arrived.Add(n) statuses := make([]int, n) var wg sync.WaitGroup wg.Add(n) for i := 0; i < n; i++ { go func(i int) { defer wg.Done() body := map[string]any{"model": "gpt-test", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) statuses[i] = rec.Code }(i) } // Blocks until every one of the n requests has reached the gated // upstream — which, by this call chain's own synchronous ordering // (runUnified -> readAndDecodeUnifiedBody's deferred release -> // registry.resolve -> the adapter's own client.Do), can only happen // AFTER that request's admission slot has already been released. // Bounded by a timeout, not a bare Wait(): if the regression this // test guards against reappeared, fewer than n requests would ever // reach the upstream at all (the rest would be refused with an // immediate 503 — acquireBodyAdmission never blocks, see its own doc // comment — and never call arrived.Done()), which would hang a bare // Wait() forever instead of failing the test. arrivedCh := make(chan struct{}) go func() { arrived.Wait() close(arrivedCh) }() select { case <-arrivedCh: case <-time.After(10 * time.Second): close(gate) // unstick any goroutine that did make it through, so the test can still finish reporting statuses t.Fatal("timed out waiting for all n requests to reach the gated upstream — some are likely being refused with 503 before ever reaching it, which is the regression this test guards against") } if held := len(gw.bodyAdmission); held != 0 { t.Errorf("body-admission slots held while %d requests (cap=%d) are parked in the upstream = %d, want 0 — a request parked past decode must never still hold a slot", n, cap(gw.bodyAdmission), held) } close(gate) wg.Wait() for i, status := range statuses { if status != http.StatusOK { t.Errorf("request %d: status = %d, want 200 (cap=%d, n=%d concurrent) — ordinary traffic must never see 503 from the body-admission semaphore merely for being concurrent", i, status, cap(gw.bodyAdmission), n) } } } // TestHandleChat_OversizedBody_TruncatedByNarrowerUnifiedCap proves the // unified routes now enforce maxUnifiedRequestBytes (4MiB), narrower than // the general maxRequestBytes (10MiB) media/MCP-federation routes still // use (finding 1c): a body whose valid JSON only closes past 4MiB is // caught by readCapped and reported as 413, not the misleading "invalid // JSON body" 400 an oversized body used to get (finding 14 fix, review- // routes.md — see readAndDecodeUnifiedBody's own doc comment): a caller // whose body is genuinely too large has no way to tell that apart from a // genuinely malformed one under the old status code. func TestHandleChat_OversizedBody_TruncatedByNarrowerUnifiedCap(t *testing.T) { cfg := CreateConfig() cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}} cfg.Groups = map[string]*GroupConfig{"default": {}} cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } // A padding field alone bigger than maxUnifiedRequestBytes, so // readCapped's own cap lands before the object closes. padding := strings.Repeat("A", maxUnifiedRequestBytes+1024) body := `{"model":"gpt-test","padding":"` + padding + `","messages":[]}` req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body)) req.Header.Set("Authorization", "Bearer sk-alice") req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusRequestEntityTooLarge { t.Fatalf("status = %d, want 413 (oversize body reported explicitly, not misreported as invalid JSON), body=%s", rec.Code, rec.Body.String()) } } // TestHandleChat_UnderNewUnifiedCap_InlineBase64Image_StillWorks proves // the chosen 4MiB unified cap (finding 1c) preserves the documented // single-inline-base64-image use case maxUnifiedRequestBytes' own doc // comment names: a chat request carrying a realistic ~1.5MiB base64 // image_url content part, comfortably under 4MiB, must decode and reach // the adapter exactly as before this finding — proving 1MiB (the audit's // first suggestion, rejected in the doc comment) would have been too // tight for this real feature, while the chosen number is not. func TestHandleChat_UnderNewUnifiedCap_InlineBase64Image_StillWorks(t *testing.T) { var gotBody []byte srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotBody, _ = io.ReadAll(r.Body) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{"id":"c1","choices":[{"message":{"role":"assistant","content":"ok"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)) })) defer srv.Close() 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"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } // ~1.5MiB of base64 payload — a realistic single inline photo, well // under maxUnifiedRequestBytes (4MiB) and well OVER the audit's // rejected 1MiB suggestion, proving that number would have broken // this exact case. imageB64 := strings.Repeat("A", 1_500_000) body := map[string]any{ "model": "gpt-test", "messages": []any{ map[string]any{ "role": "user", "content": []any{ map[string]any{"type": "text", "text": "what is this?"}, map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/jpeg;base64," + imageB64}}, }, }, }, } req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200 (a ~1.5MiB inline image must fit under the new 4MiB cap), body=%s", rec.Code, rec.Body.String()) } if !strings.Contains(string(gotBody), imageB64) { t.Error("upstream did not receive the full inline image payload — the cap silently truncated a request under its own limit") } } // --- admin-redesign WP-A step 7: kindModel req must not count a failed // failover candidate (Q1, DECISIONS; pre-existing defect fixed here) --- // TestHandleChat_Failover_ModelReqCounter_NotChargedToFailedCandidate is // the regression test for the pre-existing defect DECISIONS Q1 fixes: // before this fix, account() ran once per candidate — including a // candidate that failed before writing anything — and kindModel's own req // counter was written UNCONDITIONALLY for every one of them. alpha fails // with a pre-send 500 (zero usage, callErr != nil): its own "alpha/shared" // model-scope req counter must stay at 0. beta succeeds: its own // "beta/shared" counter must read exactly 1. func TestHandleChat_Failover_ModelReqCounter_NotChargedToFailedCandidate(t *testing.T) { alphaSrv := jsonServer(http.StatusInternalServerError, `{"error":{"message":"down","type":"server_error"}}`) defer alphaSrv.Close() betaSrv := jsonServer(http.StatusOK, successRespBody) defer betaSrv.Close() cfg := twoProviderFailoverConfig(alphaSrv, betaSrv) 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") } body := map[string]any{"model": "shared", "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } now := time.Now() alphaModelReq, ok := gw.limiter.getCounter(kindModel, "alpha/shared", metricReq, windowDay, now) if !ok { t.Fatal("getCounter(alpha/shared) ok = false") } if alphaModelReq != 0 { t.Errorf("alpha/shared model req = %d, want 0 (a candidate that failed with zero usage must never touch kindModel's req counter)", alphaModelReq) } betaModelReq, ok := gw.limiter.getCounter(kindModel, "beta/shared", metricReq, windowDay, now) if !ok { t.Fatal("getCounter(beta/shared) ok = false") } if betaModelReq != 1 { t.Errorf("beta/shared model req = %d, want 1", betaModelReq) } } // TestHandleChat_ZeroUsageStreamingSuccess_StillCountsModelReq proves the // step-7 fix's narrow exclusion is exactly callErr!=nil&&total()==0 — a // SUCCESSFUL streaming response with zero reported usage (callErr == nil) // must still count kindModel's req, since it was genuinely served. func TestHandleChat_ZeroUsageStreamingSuccess_StillCountsModelReq(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("data: {\"id\":\"x\",\"object\":\"chat.completion.chunk\",\"model\":\"gpt-test\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"},\"finish_reason\":null}]}\n\n")) _, _ = w.Write([]byte("data: [DONE]\n\n")) })) defer srv.Close() 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"}}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } modelReq, ok := gw.limiter.getCounter(kindModel, "openai/gpt-test", metricReq, windowDay, time.Now()) if !ok || modelReq != 1 { t.Errorf("openai/gpt-test model req = %d (ok=%v), want 1 (a genuinely served, zero-usage stream still counts)", modelReq, ok) } } // --- admin-redesign WP-A step 6: cache-hit/miss counters end to end --- // TestHandleChat_CacheHitMiss_RecordsCmissOnMissAndChitCsaveOnHit drives // two identical chat-completion requests through a real Gateway with // caching AND admin stats enabled (newBehavioralRedisServer, the same // fake-but-behaviorally-correct Redis TestHandleChat_CacheHitMiss_ // EndToEnd_CountersAndHeaders already uses): the first (miss) must record // total/model cmiss in the SAME synchronous account batch; the second // (hit) must record total/model chit and csave via the async // recordCacheHit path — spawn is forced synchronous so both land // deterministically before this test reads them back. func TestHandleChat_CacheHitMiss_RecordsCmissOnMissAndChitCsaveOnHit(t *testing.T) { const respBody = `{"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}}` 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(respBody)) })) defer srv.Close() redisLn := newBehavioralRedisServer(t) 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{}}}} cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} cfg.Admin = &AdminConfig{Enabled: true} // Explicit price override — "gpt-test" is not a real model id, so // leaving this unset would price it 0 (unpriced) and make the csave // assertion below meaningless. cfg.Pricing = map[string]*ModelPricing{"gpt-test": {InputPerM: 1, OutputPerM: 2}} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } gw.limiter.spawn = func(f func()) { f() } // land recordProviderAttempt/recordCacheHit synchronously body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec1 := httptest.NewRecorder() h.ServeHTTP(rec1, req1) if rec1.Code != http.StatusOK || rec1.Header().Get("X-Llmgw-Cache") != "miss" { t.Fatalf("first request: status=%d cache-header=%q, want 200/miss, body=%s", rec1.Code, rec1.Header().Get("X-Llmgw-Cache"), rec1.Body.String()) } now := time.Now() totalCmiss, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricCmiss, windowDay, now) if !ok || totalCmiss != 1 { t.Errorf("after miss: total cmiss/day = %d (ok=%v), want 1", totalCmiss, ok) } modelCmiss, ok := gw.limiter.getCounter(kindModel, "openai/gpt-test", metricCmiss, windowDay, now) if !ok || modelCmiss != 1 { t.Errorf("after miss: model cmiss/day = %d (ok=%v), want 1", modelCmiss, ok) } req2 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec2 := httptest.NewRecorder() h.ServeHTTP(rec2, req2) if rec2.Code != http.StatusOK || rec2.Header().Get("X-Llmgw-Cache") != "hit" { t.Fatalf("second request: status=%d cache-header=%q, want 200/hit, body=%s", rec2.Code, rec2.Header().Get("X-Llmgw-Cache"), rec2.Body.String()) } if rec2.Body.String() != rec1.Body.String() { t.Errorf("hit body = %q, want identical to the miss's own stored body %q", rec2.Body.String(), rec1.Body.String()) } totalChit, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricChit, windowDay, now) if !ok || totalChit != 1 { t.Errorf("after hit: total chit/day = %d (ok=%v), want 1", totalChit, ok) } modelChit, ok := gw.limiter.getCounter(kindModel, "openai/gpt-test", metricChit, windowDay, now) if !ok || modelChit != 1 { t.Errorf("after hit: model chit/day = %d (ok=%v), want 1", modelChit, ok) } totalCsave, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricCsave, windowDay, now) if !ok || totalCsave <= 0 { t.Errorf("after hit: total csave/day = %d (ok=%v), want > 0 (gpt-test has a built-in price)", totalCsave, ok) } } // TestHandleChat_CacheMissFailover_CountsExactlyOneCmiss is P3 (admin // dashboard redesign verify round): cacheMiss extras used to be set from // `cacheable` alone — true for EVERY candidate that attempted its own // cache lookup and missed, including one that went on to fail over. A // 2-candidate failover chain where alpha misses its own cache lookup and // then fails, and beta then serves, used to count 1 end-user request as 2 // total/day cmiss increments. Gated on the SAME "served or reported // usage" condition as the pre-existing Q1 model-scope fix (routes_ // unified.go), only the candidate whose result actually gets served/ // billed reports a miss now. func TestHandleChat_CacheMissFailover_CountsExactlyOneCmiss(t *testing.T) { alphaSrv := jsonServer(http.StatusInternalServerError, `{"error":{"message":"down","type":"server_error"}}`) defer alphaSrv.Close() betaSrv := jsonServer(http.StatusOK, successRespBody) defer betaSrv.Close() cfg := twoProviderFailoverConfig(alphaSrv, betaSrv) redisLn := newBehavioralRedisServer(t) cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} cfg.Admin = &AdminConfig{Enabled: true} next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } gw.limiter.spawn = func(f func()) { f() } body := map[string]any{"model": "shared", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200 (beta must serve after alpha fails), body=%s", rec.Code, rec.Body.String()) } now := time.Now() cmiss, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricCmiss, windowDay, now) if !ok || cmiss != 1 { t.Errorf("total cmiss/day = %d (ok=%v), want exactly 1 for one request with failover (never one per candidate that missed its own cache lookup)", cmiss, ok) } fover, ok := gw.limiter.getCounter(kindProvider, "alpha", metricFover, windowDay, now) if !ok || fover != 1 { t.Errorf("fover(alpha)/day = %d (ok=%v), want 1", fover, ok) } } // TestCacheHit_BodyParseSkippedWhenStatsAdminOff is P4 (admin dashboard // redesign verify round, hot path): recordCacheHit itself already no-ops // when !statsAdmin, but cachedUsage(cached.Body) — a full JSON unmarshal // of the cached response body — used to run UNCONDITIONALLY before that // check, so a cache-enabled, admin-DISABLED deployment paid the parse // cost on every single hit for nothing (measured elsewhere at 639µs for a // 1 MiB body). Measured here via testing.AllocsPerRun against two real // Gateways (cache configured, admin on vs admin off) serving repeated // hits for the SAME large (256 KiB) cached body: admin off must allocate // dramatically less per hit than admin on, whose own cachedUsage+ // unifiedCostMicros parse of that body — plus recordCacheHit's own // synchronous (spawn forced sync for determinism) counterStore write — // dominates its count. func TestCacheHit_BodyParseSkippedWhenStatsAdminOff(t *testing.T) { const bodySize = 2 << 20 // 2 MiB content, comfortably past BenchmarkV6_CachedUsage's own 1 MiB/639µs figure buildHitFunc := func(t *testing.T, adminOn bool) func() { t.Helper() content := strings.Repeat("a", bodySize) respBody := `{"id":"x","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"` + content + `"}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}` 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(respBody)) })) t.Cleanup(srv.Close) redisLn := newBehavioralRedisServer(t) 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{}}}} cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} // MaxBodyBytes raised to maxCacheMaxBodyBytes: the default // (defaultCacheMaxBodyBytes, cache.go, exactly 1 MiB) would refuse // to cache this test's own >1 MiB body at all, turning every // "hit" attempt below back into a miss. cfg.Cache = CacheConfig{Enabled: true, TTL: "1m", MaxBodyBytes: maxCacheMaxBodyBytes} if adminOn { cfg.Admin = &AdminConfig{Enabled: true} } next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) h, err := New(context.Background(), next, cfg, "llmgw") if err != nil { t.Fatalf("New: %v", err) } gw, ok := h.(*Gateway) if !ok { t.Fatal("handler is not *Gateway") } gw.limiter.spawn = func(f func()) { f() } body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} rec := httptest.NewRecorder() h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)) if rec.Code != http.StatusOK || rec.Header().Get("X-Llmgw-Cache") != "miss" { t.Fatalf("warm-up: status=%d cache=%q, want 200/miss", rec.Code, rec.Header().Get("X-Llmgw-Cache")) } return func() { hrec := httptest.NewRecorder() h.ServeHTTP(hrec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)) if hrec.Code != http.StatusOK || hrec.Header().Get("X-Llmgw-Cache") != "hit" { t.Fatalf("status=%d cache=%q, want 200/hit", hrec.Code, hrec.Header().Get("X-Llmgw-Cache")) } } } hitOn := buildHitFunc(t, true) hitOff := buildHitFunc(t, false) // encoding/json's decoder skips a field cachedUsage's target struct // does not name (the cached body's own "choices[0].message.content") // WITHOUT materializing it into a Go value — so the parse's own cost // shows up as CPU time spent scanning every byte, not as a // proportionally large allocation count or byte total (the actual // perf finding this fix addresses was itself time-based: "1 MiB 639 // µs per hit"). NsPerOp is therefore the metric that actually reflects // this fix, not AllocsPerRun/AllocedBytesPerOp. resultOn := testing.Benchmark(func(b *testing.B) { for i := 0; i < b.N; i++ { hitOn() } }) resultOff := testing.Benchmark(func(b *testing.B) { for i := 0; i < b.N; i++ { hitOff() } }) nsOn := resultOn.NsPerOp() nsOff := resultOff.NsPerOp() t.Logf("cache-hit ns/op for a %d MiB cached body: admin on = %d, admin off = %d", bodySize/(1<<20), nsOn, nsOff) if nsOn < nsOff+nsOff/3 { t.Errorf("ns/op with admin on (%d) is not meaningfully slower than admin off (%d) — the cached body's own JSON parse (cachedUsage) must be skipped entirely when statsAdmin is off, not merely cheap", nsOn, nsOff) } }