From 980408d045950ed43fc2a4c275bd82d5eec934f5 Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Mon, 28 Sep 2026 22:19:12 +0800 Subject: [PATCH] =?UTF-8?q?fix(tokens):=20=E6=B5=81=E5=BC=8F=E7=BB=9F?= =?UTF-8?q?=E8=AE=A1=E6=94=B9=E7=94=A8=E4=B8=8A=E6=B8=B8=E7=9C=9F=E5=AE=9E?= =?UTF-8?q?=20usage=EF=BC=8C=E5=9B=BE=E7=89=87=E4=B8=8D=E5=86=8D=E8=AE=B0?= =?UTF-8?q?=20token?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 两处 token 单位错误，均影响 per-model 配额计费： 1. 流式路径的 prompt/completion 只是「字节÷3」估算。 pumpStream 明明收到了上游最后一帧的真实 usage，却只发给客户端、 从不回写审计记录，于是配额按估算值扣。生产实测同一请求： 上游 prompt=37/completion=179 → 记账 27/262，prompt 低估 1.4x、 completion 高估 1.5x（双向失真）。同模型流式 completion 中位数 是非流式的 4-27 倍。非流式路径本就用真实值，两路不一致。 修法：lastUsage 非零时写回 rec.Prompt/rec.Compl，估算降为兜底 （上游不报 usage 时仍保留原估算行为）。 2. 图片请求把「图片张数」记成 completion_tokens。 rec.Compl = int64(len(resp.ImageData))，len 是切片长度即张数 （生产 38 条 image 记录全是 1），且被计入 token 总量。 图片生成无 token 概念 ⇒ 新增 Req.ImageCount 独立字段， Prompt/Compl 归 0；UI 记录表 image 行改显示张数（新增 i18n thImgs）。 顺带补 TestUILocaleKeyParity：此前无人校验 zh/en 键集合一致， 单边加键不会报错，只会显示原始键名。 新增 token_units_test.go（定值上游 6 项），做过变异验证： 回退修复实测复现 stream=16/173 vs chat=44/100、image completion=3。 --- internal/gateway/chat.go | 21 ++- internal/gateway/stats.go | 4 + internal/gateway/token_units_test.go | 231 +++++++++++++++++++++++++++ internal/gateway/ui/index.html | 4 +- internal/gateway/ui_contract_test.go | 57 +++++++ 5 files changed, 314 insertions(+), 3 deletions(-) create mode 100644 internal/gateway/token_units_test.go 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 @@ ${fmtTime(r.time)} ${r.ok ? `${r.status || 200}` : `${r.status || 500}`} ${esc(keyNames[r.key] ? keyNames[r.key] + " · " + r.key : r.key)}${esc(r.type)}${esc(r.model)}${esc(r.source || "")} - ${fmtTok(r.prompt_tokens)}${fmtTok(r.completion_tokens)} + ${fmtTok(r.prompt_tokens)}${r.type === "image" ? fmtN(r.image_count || 0) + " " + t("thImgs") : fmtTok(r.completion_tokens)} ${cacheCell(r)} ${fmtMs(r.latency_ms)}`; } diff --git a/internal/gateway/ui_contract_test.go b/internal/gateway/ui_contract_test.go index 73edcac..cfa5d50 100644 --- a/internal/gateway/ui_contract_test.go +++ b/internal/gateway/ui_contract_test.go @@ -347,3 +347,60 @@ func cssBlock(src, selector string) (string, bool) { } return src[i : i+j+1], true } + +// The WebUI ships two locale objects (`zh` and `en`) that every rendered string +// goes through. A key added to only one of them does not error: t() falls back +// to `t("key") || "literal"` in some call sites and to the bare key string in +// others, so the user sees either a hardcoded local string or a raw key name. +// Nothing in the Go test suite noticed — this was found by hand while adding +// thImgs. Pin the key sets so a one-sided edit fails here instead of shipping. +func TestUILocaleKeyParity(t *testing.T) { + src := uiSource(t) + iZh := strings.Index(src, "zh: {") + iEn := strings.Index(src, "en: {") + if iZh < 0 || iEn < 0 || iEn < iZh { + t.Fatalf("locale blocks not found (zh=%d en=%d)", iZh, iEn) + } + zh := src[iZh:iEn] + en := src[iEn:] + // End the EN block at its closing brace, so keys from later objects + // (config templates etc.) do not pollute the comparison. + if end := strings.Index(en, "\n },"); end > 0 { + en = en[:end] + } + + keyRe := regexp.MustCompile(`(?m)^\s*([A-Za-z_][A-Za-z0-9_]*)\s*:`) + collect := func(seg string) map[string]bool { + m := map[string]bool{} + for _, g := range keyRe.FindAllStringSubmatch(seg, -1) { + m[g[1]] = true + } + return m + } + zk, ek := collect(zh), collect(en) + // The locale markers themselves are the two block headers. + delete(zk, "zh") + delete(ek, "en") + + var missingEN, missingZH []string + for k := range zk { + if !ek[k] { + missingEN = append(missingEN, k) + } + } + for k := range ek { + if !zk[k] { + missingZH = append(missingZH, k) + } + } + if len(missingEN) > 0 { + t.Errorf("translation keys present in zh but missing in en: %v", missingEN) + } + if len(missingZH) > 0 { + t.Errorf("translation keys present in en but missing in zh: %v", missingZH) + } + if len(zk) < 100 { + t.Errorf("only %d zh keys parsed — the block boundary regexp has drifted "+ + "and this test is no longer checking anything", len(zk)) + } +}