mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-20 05:56:56 +00:00
fix(gateway): close billing, limit, routing and translation gaps
Billing: - Passthrough SSE streams now bill the usage the upstream reports. The gateway falls back to a request-size estimate only on a 2xx stream or a truncated JSON body. - A unified stream that fails after it starts is estimated, not billed at zero. - Billing uses the built-in LiteLLM price table and honours modelMeta free: true. - Pricing entries are validated. Cost arithmetic rounds the total and saturates on overflow. - Gemini thinking tokens count as completion tokens. Limits and storage: - A rejected request rolls back its counts on the other scopes, so one user cannot lock out a group. - Redis EXPIRE replies are checked. Keys recreated without a TTL are repaired at once. - A mutating pipeline is not retried after reply bytes arrive. - History, model totals and target usage report a store outage. Routing and translation: - Discovered ids with a provider-like prefix are listed only in owner/id form, so listing and routing agree. - Anthropic and Gemini discovery follows pagination. - The developer role maps to system. - stream_options.include_usage is honoured on translated streams. - tool_choice none keeps tools when history holds tool blocks. - refusal maps to content_filter. Data URI parameters are stripped. Streaming, retry and failover: - SSE accepts CR, LF and CRLF line ends. This closes a CR injection path. - Retry-After is never jittered. A cancelled wait returns the context error. - The watchdog timer stops at end of body. - Oversized bodies return 413. - A failed stream ends with an error event. - Failover restores pre-existing headers. Client cancels do not count as provider successes. Also: "admin" is a reserved provider name. metrics.path must start with "/". Pricing warnings go to the right gateway instance. The admin targets view sends access as null or [] so the two states differ. Upstream organisation headers are stripped. The users-file reload error is logged once per error. Dead code is removed, and test-only helpers move to test files. Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
This commit is contained in:
@@ -181,6 +181,142 @@ func TestAnthropicAdapter_ChatCompletion_Streaming(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// sseDataChunks splits an SSE response body into its individual "data:
|
||||
// ..." payloads, in wire order, stripping the "data: " prefix — a small
|
||||
// shared parser for the M1 tests below, since forwardStream's per-event
|
||||
// framing (sseWriter.writeData, sse.go) always separates events with a
|
||||
// blank line.
|
||||
func sseDataChunks(t *testing.T, body string) []string {
|
||||
t.Helper()
|
||||
var out []string
|
||||
for _, frame := range strings.Split(body, "\n\n") {
|
||||
frame = strings.TrimSpace(frame)
|
||||
if frame == "" {
|
||||
continue
|
||||
}
|
||||
data, ok := strings.CutPrefix(frame, "data: ")
|
||||
if !ok {
|
||||
t.Fatalf("SSE frame missing %q prefix: %q", "data: ", frame)
|
||||
}
|
||||
out = append(out, data)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestAnthropicAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk
|
||||
// is the M1 regression test: a client that sets
|
||||
// stream_options.include_usage on a request routed to an anthropic-type
|
||||
// provider must still get a final usage-only chunk on the wire, exactly
|
||||
// as an OpenAI-type provider's own forwardStream relays from a real
|
||||
// upstream usage chunk (provider_openai.go) — Anthropic has no such
|
||||
// upstream chunk, so the adapter must synthesize one from the usage it
|
||||
// already accumulated translating message_start/message_delta.
|
||||
func TestAnthropicAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk(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)
|
||||
f, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
|
||||
}
|
||||
for _, frame := range anthropicStreamFrames {
|
||||
fmt.Fprintf(w, "event: %s\ndata: %s\n\n", frame.event, frame.data)
|
||||
f.Flush()
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
a, err := newAnthropicAdapter("p1", srv.URL, "sk-ant-test")
|
||||
if err != nil {
|
||||
t.Fatalf("newAnthropicAdapter: %v", err)
|
||||
}
|
||||
rec := newRecordingResponseWriter()
|
||||
|
||||
req := map[string]any{
|
||||
"model": "claude-opus-5",
|
||||
"stream": true,
|
||||
"stream_options": map[string]any{"include_usage": true},
|
||||
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
||||
}
|
||||
u, err := a.chatCompletion(context.Background(), rec, req)
|
||||
if err != nil {
|
||||
t.Fatalf("chatCompletion: %v", err)
|
||||
}
|
||||
if u.prompt != 6 || u.completion != 3 {
|
||||
t.Fatalf("usage = %+v, want {prompt:6 completion:3}", u)
|
||||
}
|
||||
|
||||
chunks := sseDataChunks(t, rec.body.String())
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("got %d SSE frames, want at least a usage chunk followed by [DONE]", len(chunks))
|
||||
}
|
||||
if chunks[len(chunks)-1] != "[DONE]" {
|
||||
t.Fatalf("last SSE frame = %q, want [DONE]", chunks[len(chunks)-1])
|
||||
}
|
||||
usageFrame := chunks[len(chunks)-2]
|
||||
var got struct {
|
||||
Usage *struct {
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
} `json:"usage"`
|
||||
Choices []any `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(usageFrame), &got); err != nil {
|
||||
t.Fatalf("decode usage chunk %q: %v", usageFrame, err)
|
||||
}
|
||||
if len(got.Choices) != 0 {
|
||||
t.Errorf("usage chunk choices = %v, want empty", got.Choices)
|
||||
}
|
||||
if got.Usage == nil {
|
||||
t.Fatalf("usage chunk has no usage object: %q", usageFrame)
|
||||
}
|
||||
if got.Usage.PromptTokens != 6 || got.Usage.CompletionTokens != 3 || got.Usage.TotalTokens != 9 {
|
||||
t.Errorf("usage chunk = %+v, want {prompt_tokens:6 completion_tokens:3 total_tokens:9}", *got.Usage)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk
|
||||
// proves the M1 fix is opt-in: a client that never sets
|
||||
// stream_options.include_usage gets the stream unchanged from before the
|
||||
// fix — no extra usage-only chunk on the wire.
|
||||
func TestAnthropicAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk(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)
|
||||
f, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
|
||||
}
|
||||
for _, frame := range anthropicStreamFrames {
|
||||
fmt.Fprintf(w, "event: %s\ndata: %s\n\n", frame.event, frame.data)
|
||||
f.Flush()
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
a, err := newAnthropicAdapter("p1", srv.URL, "sk-ant-test")
|
||||
if err != nil {
|
||||
t.Fatalf("newAnthropicAdapter: %v", err)
|
||||
}
|
||||
rec := newRecordingResponseWriter()
|
||||
|
||||
req := map[string]any{
|
||||
"model": "claude-opus-5",
|
||||
"stream": true,
|
||||
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
||||
}
|
||||
if _, err := a.chatCompletion(context.Background(), rec, req); err != nil {
|
||||
t.Fatalf("chatCompletion: %v", err)
|
||||
}
|
||||
|
||||
for _, chunk := range sseDataChunks(t, rec.body.String()) {
|
||||
if strings.Contains(chunk, `"usage"`) {
|
||||
t.Errorf("SSE frame %q carries a usage field though the client never set stream_options.include_usage", chunk)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType
|
||||
// proves the streaming-vs-JSON-fallback gate matches an upstream's
|
||||
// Content-Type case-insensitively — review fix, mirroring the same fix
|
||||
@@ -329,3 +465,76 @@ func TestAnthropicAdapter_ListModels(t *testing.T) {
|
||||
t.Errorf("got %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicAdapter_ListModels_FollowsPagination is the M7 regression
|
||||
// test: has_more:true plus a last_id must make listModels fetch a second
|
||||
// page via ?after_id=<last_id> and return the union of both pages, not
|
||||
// just the first page's 20 (Anthropic's own default page size).
|
||||
func TestAnthropicAdapter_ListModels_FollowsPagination(t *testing.T) {
|
||||
var gotAfterIDs []string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAfterIDs = append(gotAfterIDs, r.URL.Query().Get("after_id"))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.URL.Query().Get("after_id") == "" {
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"claude-opus-5"},{"id":"claude-sonnet-5"}],"has_more":true,"last_id":"claude-sonnet-5"}`))
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"claude-haiku-5"}],"has_more":false,"last_id":"claude-haiku-5"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
a, err := newAnthropicAdapter("p1", srv.URL, "sk-ant-test")
|
||||
if err != nil {
|
||||
t.Fatalf("newAnthropicAdapter: %v", err)
|
||||
}
|
||||
|
||||
got, err := a.listModels(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("listModels: %v", err)
|
||||
}
|
||||
want := []string{"claude-opus-5", "claude-sonnet-5", "claude-haiku-5"}
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("got %#v, want %#v", got, want)
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Errorf("got[%d] = %q, want %q", i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
if len(gotAfterIDs) != 2 || gotAfterIDs[0] != "" || gotAfterIDs[1] != "claude-sonnet-5" {
|
||||
t.Errorf("after_id sequence = %#v, want [\"\", \"claude-sonnet-5\"] (first page unqualified, second page's after_id = first page's last_id)", gotAfterIDs)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicAdapter_ListModels_PageCapStopsRunawayPagination proves
|
||||
// the M7 fix's other half: an upstream that always reports has_more:true
|
||||
// with an ever-advancing last_id does not page forever — listModels
|
||||
// stops after anthropicListModelsPageCap requests and returns
|
||||
// successfully with whatever it accumulated, rather than looping until
|
||||
// context cancellation or exhausting memory.
|
||||
func TestAnthropicAdapter_ListModels_PageCapStopsRunawayPagination(t *testing.T) {
|
||||
var requests int
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests++
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
id := fmt.Sprintf("model-%d", requests)
|
||||
fmt.Fprintf(w, `{"data":[{"id":%q}],"has_more":true,"last_id":%q}`, id, id)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
a, err := newAnthropicAdapter("p1", srv.URL, "sk-ant-test")
|
||||
if err != nil {
|
||||
t.Fatalf("newAnthropicAdapter: %v", err)
|
||||
}
|
||||
|
||||
got, err := a.listModels(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("listModels: %v", err)
|
||||
}
|
||||
if requests != anthropicListModelsPageCap {
|
||||
t.Errorf("upstream received %d requests, want exactly %d (the page cap)", requests, anthropicListModelsPageCap)
|
||||
}
|
||||
if len(got) != anthropicListModelsPageCap {
|
||||
t.Errorf("got %d ids, want %d (one per page up to the cap)", len(got), anthropicListModelsPageCap)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user