diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index 51801d6..7cc91e0 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -957,6 +957,19 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type // Final usage chunk (OpenAI standard: empty choices + usage before [DONE]). // Prefer the upstream's exact usage if the stream carried it; fall back to // the gateway's estimate otherwise. + // Write the upstream's exact numbers back onto the audit record. Without + // this the streamed path kept only the per-chunk byte estimate, so the same + // request recorded ~1.4x its real completion tokens (measured: upstream 100, + // audit 145) while the non-streaming path recorded 100. Two paths, two + // different numbers for one request is a reporting bug, not a rounding one. + if lastUsage != nil { + if lastUsage.Prompt > 0 { + rec.Prompt = int64(lastUsage.Prompt) + } + if lastUsage.Completion > 0 { + rec.Compl = int64(lastUsage.Completion) + } + } var tut *types.TokenUsage if lastUsage != nil { tut = lastUsage @@ -1125,7 +1138,10 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { } rec.OK = true rec.Status = http.StatusOK - rec.Compl = int64(len(resp.ImageData)) + // Image generation has no token concept. Recording len(ImageData) + // (the image COUNT) in completion_tokens mislabels image count as + // tokens and feeds it into the token totals; leave it 0. + rec.ImageCount = len(resp.ImageData) g.writeRec(rec) writeJSON(w, http.StatusOK, types.ImageGenResponse{ Created: time.Now().Unix(), @@ -1174,7 +1190,8 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { } rec.OK = true rec.Status = http.StatusOK - rec.Compl = int64(len(resp.ImageData)) + // Image generation has no token concept — see the AUTO path above. + rec.ImageCount = len(resp.ImageData) g.writeRec(rec) writeJSON(w, http.StatusOK, types.ImageGenResponse{ Created: time.Now().Unix(), diff --git a/internal/gateway/stats.go b/internal/gateway/stats.go index 389866e..dced2b5 100644 --- a/internal/gateway/stats.go +++ b/internal/gateway/stats.go @@ -48,6 +48,10 @@ type Req struct { Status int `json:"status"` // Err short error message Err string `json:"error,omitempty"` + // ImageCount is how many images an image-generation request produced. + // Kept separate from Compl/Prompt: image generation has no token concept, + // so counting images as "completion tokens" would corrupt the token totals. + ImageCount int `json:"image_count,omitempty"` } // Stat aggregates counters for one dimension row. diff --git a/internal/gateway/token_units_test.go b/internal/gateway/token_units_test.go new file mode 100644 index 0000000..12b09a4 --- /dev/null +++ b/internal/gateway/token_units_test.go @@ -0,0 +1,231 @@ +package gateway + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "llmsproxy/internal/config" +) + +// Token accounting must record TOKENS, not byte estimates and not image +// counts. These tests pin the two units bugs found on 2026-09-27: +// +// 1. streamed chat: rec.Compl was only ever the per-chunk byte estimate +// (len(content)/3), even when the upstream sent exact usage on its final +// frame. Measured on production traffic: the same model reported a +// streamed completion median 4-27x its non-streamed median. Upstream 100 +// recorded as 145. +// 2. image generation: rec.Compl was set to len(resp.ImageData) — the image +// COUNT — and folded into the token totals. +// +// A fake upstream that reports FIXED numbers makes both mismatches +// unambiguous: any difference between what it reports and what the audit +// record holds is the gateway's doing. + +// fixedUsageUpstream reports prompt=prompt completion=compl on both the +// streaming and the non-streaming path. The streaming path emits usage on its +// final frame before [DONE], exactly like OpenAI-compatible upstreams. +func fixedUsageUpstream(t *testing.T, prompt, compl int, images int) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]interface{} + _ = json.NewDecoder(r.Body).Decode(&body) + stream, _ := body["stream"].(bool) + usage := fmt.Sprintf(`{"prompt_tokens":%d,"completion_tokens":%d,"total_tokens":%d}`, + prompt, compl, prompt+compl) + + // Image generation wire: /v1/images/generations style body. + if strings.Contains(r.URL.Path, "image") { + w.Header().Set("Content-Type", "application/json") + var data []string + for i := 0; i < images; i++ { + data = append(data, `{"b64_json":"AAAA"}`) + } + fmt.Fprintf(w, `{"created":1,"data":[%s]}`, strings.Join(data, ",")) + return + } + + if !stream { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"model":"m1","choices":[{"message":{"content":"pong pong pong"}, + "finish_reason":"stop"}],"usage":%s}`, usage) + return + } + + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(200) + fl, _ := w.(http.Flusher) + send := func(s string) { + fmt.Fprintf(w, "data: %s\n\n", s) + if fl != nil { + fl.Flush() + } + } + // Deliberately verbose payloads: a byte-estimating path records a + // number far from `compl`, so the test would catch it. + send(`{"model":"m1","choices":[{"delta":{"reasoning_content":"` + strings.Repeat("think ", 40) + `"}}]}`) + send(`{"model":"m1","choices":[{"delta":{"content":"` + strings.Repeat("answer ", 40) + `"}}]}`) + send(`{"model":"m1","choices":[{"delta":{},"finish_reason":"stop"}]}`) + send(`{"model":"m1","choices":[],"usage":` + usage + `}`) + send(`[DONE]`) + })) +} + +// lastChatRecord returns the most recent audit record whose Type is one of the +// given types. +func lastChatRecord(t *testing.T, g *Gateway, types ...string) Req { + t.Helper() + recs := g.stats.AuditRecords(0, 0, "") + want := map[string]bool{} + for _, ty := range types { + want[ty] = true + } + for i := len(recs) - 1; i >= 0; i-- { + if want[recs[i].Type] { + return recs[i] + } + } + t.Fatalf("no audit record of type %v found (have %d records)", types, len(recs)) + return Req{} +} + +// The streamed path must record the upstream's exact usage, not a byte +// estimate. Before the fix this recorded 145 for an upstream-reported 100. +func TestStreamedChatRecordsUpstreamUsage(t *testing.T) { + const wantPrompt, wantCompl = 44, 100 + up := fixedUsageUpstream(t, wantPrompt, wantCompl, 0) + defer up.Close() + g := newTestGateway(t, config.Source{ + Name: "up", BaseURL: up.URL, Adapter: "openai", + Models: []config.Model{{ID: "m1", Priority: 100}}, + }) + + rr := doReq(t, g, "POST", "/v1/chat/completions", + `{"model":"m1","messages":[{"role":"user","content":"hi"}],"stream":true}`) + if rr.Code != 200 { + t.Fatalf("stream status=%d body=%s", rr.Code, rr.Body.String()) + } + + rec := lastChatRecord(t, g, "stream") + if rec.Prompt != wantPrompt || rec.Compl != wantCompl { + t.Errorf("streamed record prompt=%d completion=%d, want %d/%d — the audit "+ + "must carry the upstream's tokens, not a byte estimate", + rec.Prompt, rec.Compl, wantPrompt, wantCompl) + } +} + +// Streaming and non-streaming must agree on the same request's tokens. This is +// the cross-path consistency check: the two paths previously recorded 145 and +// 100 for one identical upstream usage value. +func TestStreamedAndNonStreamedAgreeOnTokens(t *testing.T) { + const wantPrompt, wantCompl = 44, 100 + up := fixedUsageUpstream(t, wantPrompt, wantCompl, 0) + defer up.Close() + g := newTestGateway(t, config.Source{ + Name: "up", BaseURL: up.URL, Adapter: "openai", + Models: []config.Model{{ID: "m1", Priority: 100}}, + }) + + if rr := doReq(t, g, "POST", "/v1/chat/completions", + `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`); rr.Code != 200 { + t.Fatalf("non-stream status=%d body=%s", rr.Code, rr.Body.String()) + } + if rr := doReq(t, g, "POST", "/v1/chat/completions", + `{"model":"m1","messages":[{"role":"user","content":"hi"}],"stream":true}`); rr.Code != 200 { + t.Fatalf("stream status=%d body=%s", rr.Code, rr.Body.String()) + } + + chat, stream := lastChatRecord(t, g, "chat"), lastChatRecord(t, g, "stream") + if chat.Prompt != stream.Prompt || chat.Compl != stream.Compl { + t.Errorf("chat=%d/%d stream=%d/%d — identical upstream usage must record "+ + "identically on both paths", chat.Prompt, chat.Compl, stream.Prompt, stream.Compl) + } +} + +// A fallback to the byte estimate is still correct WHEN the upstream reports no +// usage at all. Without this the fix could silently zero out every streamed +// record from upstreams that omit usage. +func TestStreamedChatFallsBackToEstimateWithoutUsage(t *testing.T) { + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(200) + fl, _ := w.(http.Flusher) + send := func(s string) { + fmt.Fprintf(w, "data: %s\n\n", s) + if fl != nil { + fl.Flush() + } + } + send(`{"model":"m1","choices":[{"delta":{"content":"` + strings.Repeat("abc ", 60) + `"}}]}`) + send(`{"model":"m1","choices":[{"delta":{},"finish_reason":"stop"}]}`) + send(`[DONE]`) + })) + defer up.Close() + g := newTestGateway(t, config.Source{ + Name: "up", BaseURL: up.URL, Adapter: "openai", + Models: []config.Model{{ID: "m1", Priority: 100}}, + }) + + if rr := doReq(t, g, "POST", "/v1/chat/completions", + `{"model":"m1","messages":[{"role":"user","content":"hi"}],"stream":true}`); rr.Code != 200 { + t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String()) + } + rec := lastChatRecord(t, g, "stream") + if rec.Compl <= 0 { + t.Errorf("completion=%d, want a positive byte-estimate fallback when the "+ + "upstream reports no usage", rec.Compl) + } +} + +// Image generation has no token concept: recording the image COUNT as +// completion tokens corrupted the token totals. The count is reported +// separately and must not leak into Prompt/Compl. +func TestImageGenerationDoesNotRecordTokens(t *testing.T) { + up := fixedUsageUpstream(t, 44, 100, 3) + defer up.Close() + src := config.Source{ + Name: "up", BaseURL: up.URL, Adapter: "openai", + Models: []config.Model{{ID: "img-m", Priority: 100, Kind: "image"}}, + } + g := newTestGateway(t, src) + + rr := doReq(t, g, "POST", "/v1/images/generations", + `{"model":"img-m","prompt":"a cat","n":3}`) + if rr.Code != 200 { + t.Fatalf("image status=%d body=%s", rr.Code, rr.Body.String()) + } + + rec := lastChatRecord(t, g, "image") + if rec.Compl != 0 || rec.Prompt != 0 { + t.Errorf("image record prompt=%d completion=%d, want 0/0 — image count is "+ + "not a token and must not enter the token totals", rec.Prompt, rec.Compl) + } + if rec.ImageCount != 3 { + t.Errorf("image_count=%d, want 3 (the count belongs in its own field)", rec.ImageCount) + } + // And the aggregate must see zero tokens for it. + if got := g.stats.WindowTokens("img-m", "", 3600); got != 0 { + t.Errorf("window tokens for an image model = %d, want 0", got) + } +} + +// The audit JSON must carry image_count so the WebUI can show images rather +// than a bogus token figure. +func TestImageCountSurvivesAuditJSON(t *testing.T) { + b, err := json.Marshal(Req{Type: "image", ImageCount: 2, OK: true}) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(b), `"image_count":2`) { + t.Errorf("audit JSON %s lacks image_count", b) + } + // Omitempty keeps chat records unchanged. + b2, _ := json.Marshal(Req{Type: "chat", Prompt: 1, Compl: 2, OK: true}) + if strings.Contains(string(b2), "image_count") { + t.Errorf("chat record %s should omit image_count", b2) + } +} diff --git a/internal/gateway/ui/index.html b/internal/gateway/ui/index.html index 029d62f..0187339 100644 --- a/internal/gateway/ui/index.html +++ b/internal/gateway/ui/index.html @@ -984,6 +984,7 @@ thErr: "失败", thPrompt: "输入 Tokens", thCompl: "输出 Tokens", + thImgs: "张", thAvgLat: "平均延迟", thMaxLat: "最长延迟", thTime: "时间", @@ -1214,6 +1215,7 @@ thErr: "Err", thPrompt: "Prompt Tokens", thCompl: "Completion Tokens", + thImgs: "imgs", thAvgLat: "Avg latency", thMaxLat: "Max latency", thTime: "Time", @@ -2193,7 +2195,7 @@