mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
GET /admin/api/events returns recent failures, newest first. It covers rate limits, budget and unpriced-model refusals, upstream errors, timeouts and capacity rejections. Events go into an in-process ring of 200 entries and are copied to a Redis list (llmgw:events, LPUSH plus LTRIM to 200). The copy happens only on the error path, is throttled to 50 per second per replica, and never blocks a request. When Redis is absent or down, the feed falls back to this replica's ring. Messages are fixed strings or scrubbed provider errors. Fields are assigned one at a time, because Yaegi drops values from a multi-value assignment that contains a call. Rejections per day are counted per scope and in total, in the same store call that already rolls back a rejected request. /admin/api/usage returns them as rejectionsPerDay. /admin/api/usage/models accepts span, which sums that many buckets, and prefix, which filters model ids before the limit applies. The overview reports the replica id, the instance name and the warnings logged while the middleware was built. The failover and retry fan-out warning is now logged at WARN. Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
2793 lines
117 KiB
Go
2793 lines
117 KiB
Go
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")
|
|
}
|
|
}
|