Compare commits

...

4 Commits

Author SHA1 Message Date
cd82835f25 chore(version): 1.7.2 -> 1.7.3
启动重复播种 admin key 与配置封存非幂等的修复,走 patch 发布。
2026-09-28 22:53:19 +08:00
26ea782350 fix(core): 修复启动重复播种 admin key + 配置封存非幂等
根因是 unseal 时序:NewFromConfig 把解密放在最后,而之前几步已经在读凭据。

1. seedKeys 重复播种(生产已累积 4 个同名 admin key)
   seedKeys 用 cfg.Keys[i].Key 与明文 gateway_keys 比对去重,但此时内存里的
   key 还是密文 enc:v1:…,比对永不命中 ⇒ 每次重启追加一个同值 admin key。
   实测:core.New(path) 连续重启,seeded key 数 2→3→4 递增。
   (旧测试用 NewFromConfig 构造全新内存对象,没有「盘上已有密文」这个前提,
    复现不出 —— 必须走 core.New 这条读盘的生产路径。)

2. 启动恒重写 config.yaml
   migratePlaintextSecrets 按内存状态判断,而 Save() 末尾会把内存恢复为明文,
   于是每次调用都判定「还有明文」并重写;注释却自称幂等。
   改为 UnsealSecrets 在解密前记录「盘上是否明文」,SealIfNeeded 据此决定
   是否写回 ⇒ 已封存的配置启动不再落盘。

原测试 TestMigratePlaintextSecretsIsIdempotent 用 ModTime 比较,两次写落在同一
时间戳刻度内就看不出来,所以表现为 ~1/6 概率的 flake 而非稳定失败。已改为比较
文件内容并走真实启动路径(UnsealSecrets + SealIfNeeded),并顺带消除该 flake。

附带更正:先前判断「rebuildRegistry 也会拿到密文 API key」不成立 ——
mergedSources → resolveSourceKey 对每个 source 独立解密(belt-and-braces),
provider 始终拿到明文。unseal 前置仍予保留,以消除对该兜底路径的隐性依赖、
并让 seedKeys 在明文下比较。

判据:
- TestRestartDoesNotDuplicateSeededKeys(敏感:回退顺序必红)
- TestSealingIsIdempotentAcrossStarts(12/12 稳定,原先 1/6 flake)
- TestProvidersGetPlaintextCredentials(钉 provider 必须拿到明文这一不变量)
2026-09-28 22:52:51 +08:00
5c58244781 chore(version): 1.7.1 -> 1.7.2
token 统计单位修复(流式改用上游真实 usage、图片不再记 token),
影响 per-model 配额计费口径,走 patch 发布。
2026-09-28 22:19:24 +08:00
0121d23f91 fix(tokens): 流式统计改用上游真实 usage,图片不再记 token
两处 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。
2026-09-28 22:19:12 +08:00
11 changed files with 621 additions and 29 deletions

View File

@ -1,12 +1,12 @@
{
"name": "modelrouter-gui",
"version": "1.7.1",
"version": "1.7.3",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "modelrouter-gui",
"version": "1.7.1",
"version": "1.7.3",
"license": "Proprietary",
"devDependencies": {
"electron": "^33.0.0",
@ -582,7 +582,7 @@
}
},
"node_modules/@peculiar/webcrypto": {
"version": "1.7.1",
"version": "1.7.3",
"resolved": "https://registry.npmjs.org/@peculiar/webcrypto/-/webcrypto-1.7.1.tgz",
"integrity": "sha512-ODOov0sGMJMf3jPonOkgGqPknTsu+DdQ7kD++gz8aI+aFMOMHFbWAA2taqXXVTdP+OTOQR/znGvSpmkeI0WTYQ==",
"dev": true,
@ -3335,7 +3335,7 @@
}
},
"node_modules/resedit": {
"version": "1.7.2",
"version": "1.7.3",
"resolved": "https://registry.npmjs.org/resedit/-/resedit-1.7.2.tgz",
"integrity": "sha512-vHjcY2MlAITJhC0eRD/Vv8Vlgmu9Sd3LX9zZvtGzU5ZImdTN3+d6e/4mnTyV8vEbyf1sgNIrWxhWlrys52OkEA==",
"dev": true,

View File

@ -1,6 +1,6 @@
{
"name": "modelrouter-gui",
"version": "1.7.1",
"version": "1.7.3",
"description": "ModelRouter Desktop — embedded ModelRouter core with tray",
"author": "ModelRouter",
"main": "main.js",

View File

@ -193,11 +193,55 @@ func (c *Config) countPlaintextSecrets() int {
// with ciphertext still in place, so the unseal has to happen before the
// registry (and any Save the startup path performs) sees the values.
func (c *Config) NormalizeSecretsForRun(box *SecretBox) error {
c.AttachSecretBox(box)
if err := c.normalizeSecrets(box); err != nil {
hadPlaintext, err := c.UnsealSecrets(box)
if err != nil {
return err
}
return c.migratePlaintextSecrets()
return c.SealIfNeeded(hadPlaintext)
}
// UnsealSecrets decrypts every sealed credential in memory and reports whether
// the config ON DISK still held plaintext (i.e. whether a sealing write is
// needed). It never writes; a failed decrypt (wrong master key) is returned so
// the process refuses to start instead of running with unusable credentials.
//
// The return value must be computed BEFORE unsealing and from the disk state,
// not from memory: after a Save the in-memory values are always plaintext, so a
// "is anything plaintext?" test run afterwards is unconditionally true and a
// caller would rewrite the file on every start. That was the actual behaviour -
// migratePlaintextSecrets() claimed to be idempotent in a comment but rewrote
// config.yaml on every boot.
//
// Callers that consume credentials (the provider registry, key seeding) must
// unseal FIRST. Seeding compares cfg.Keys[i].Key against the plaintext
// gateway_keys entries; running it while keys are still ciphertext made the
// dedupe never match, so every restart appended another copy of the same admin
// key (production accumulated four).
func (c *Config) UnsealSecrets(box *SecretBox) (bool, error) {
if box == nil {
return false, nil
}
c.AttachSecretBox(box)
hadPlaintext := c.hasPlaintextSecrets()
if err := c.normalizeSecrets(box); err != nil {
return hadPlaintext, err
}
return hadPlaintext, nil
}
// SealIfNeeded writes the config back once if hadPlaintext reported that the
// file still held clear-text credentials. When it is false the file is left
// untouched, which is what makes startup a no-op for an already-sealed config.
func (c *Config) SealIfNeeded(hadPlaintext bool) error {
if !hadPlaintext || c.box == nil || c.Path == "" {
return nil
}
n := c.countPlaintextSecrets()
if err := c.Save(); err != nil {
return err
}
log.Printf("[config] sealed %d plaintext credential(s) in %s", n, c.Path)
return nil
}
// AttachSecretBox wires the encryption box into the config so Save can seal
@ -210,8 +254,9 @@ func (c *Config) AttachSecretBox(box *SecretBox) { c.box = box }
func (c *Config) SecretBox() *SecretBox { return c.box }
// migratePlaintextSecrets seals any credential still in the clear and writes the
// file once. It is idempotent: a config that is already sealed (or has no
// secrets) is left alone and nothing is written.
// file once. Kept for callers that attach the box themselves; it decides from
// the in-memory state, which is why UnsealSecrets + SealIfNeeded (which decide
// from the on-disk state) are preferred on the startup path.
func (c *Config) migratePlaintextSecrets() error {
if c.box == nil || c.Path == "" {
return nil

View File

@ -119,7 +119,24 @@ sources:
}
}
func TestMigratePlaintextSecretsIsIdempotent(t *testing.T) {
// Sealing a config is a one-time migration: the FIRST run writes, every later
// start must leave the file alone.
//
// This test previously called migratePlaintextSecrets() twice and compared
// ModTime. That assertion was both weak and wrong:
//
// - migratePlaintextSecrets decides from IN-MEMORY state, and Save() ends by
// unsealing memory so the running process keeps working. So every call sees
// "plaintext present" and rewrites the file. Comparing ModTime hid this
// because two writes inside one filesystem timestamp tick look identical —
// the test flaked (~1 in 6) instead of failing.
// - it also exercised a function production no longer calls on the startup
// path, which now uses UnsealSecrets + SealIfNeeded (they decide from the
// ON-DISK state).
//
// So the test targets that real path and compares file CONTENT, which cannot be
// fooled by timestamp granularity.
func TestSealingIsIdempotentAcrossStarts(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "config.yaml")
runtime := filepath.Join(dir, "runtime.json")
@ -132,29 +149,52 @@ func TestMigratePlaintextSecretsIsIdempotent(t *testing.T) {
}
cfg.AttachSecretBox(box)
if err := cfg.migratePlaintextSecrets(); err != nil {
t.Fatalf("first migrate: %v", err)
hadPlaintext, err := cfg.UnsealSecrets(box)
if err != nil {
t.Fatalf("unseal: %v", err)
}
if !hadPlaintext {
t.Fatal("a config written with a clear-text credential must report plaintext")
}
if err := cfg.SealIfNeeded(hadPlaintext); err != nil {
t.Fatalf("seal: %v", err)
}
first, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(first), encPrefix) {
t.Fatal("migration did not seal the value")
t.Fatal("sealing did not encrypt the value on disk")
}
// In-memory must be plaintext so the running process keeps working.
if cfg.Sources[0].APIKey != "sk-clear" {
t.Errorf("in-memory api_key = %q, want plaintext", cfg.Sources[0].APIKey)
}
// Second run: already sealed => no write.
before, _ := os.Stat(path)
if err := cfg.migratePlaintextSecrets(); err != nil {
t.Fatalf("second migrate: %v", err)
// Simulate the next start: load from disk, unseal, seal-if-needed. The file
// was already sealed, so nothing may be written.
cfg2, err := Load(path)
if err != nil {
t.Fatal(err)
}
after, _ := os.Stat(path)
if before.ModTime() != after.ModTime() {
t.Error("second migrate rewrote an already-sealed config")
cfg2.AttachSecretBox(box)
hadPlaintext2, err := cfg2.UnsealSecrets(box)
if err != nil {
t.Fatalf("second unseal: %v", err)
}
if hadPlaintext2 {
t.Error("a sealed config reported plaintext — SealIfNeeded would rewrite " +
"the file on every start")
}
if err := cfg2.SealIfNeeded(hadPlaintext2); err != nil {
t.Fatalf("second seal: %v", err)
}
second, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(first) != string(second) {
t.Error("the second start rewrote an already-sealed config; startup must be a no-op")
}
}

View File

@ -56,8 +56,26 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
c.store = config.NewStore(cfg.RuntimeFile)
// Share one box between the runtime store and config.yaml so a single
// master.key seals both files. config.Load left the config holding
// ciphertext (if it was sealed); unseal it now that the box exists.
// ciphertext (if it was sealed); unseal it NOW, before anything reads a
// credential.
//
// ORDER IS LOAD-BEARING. These steps each consume secrets and must run
// after the unseal:
// - seedKeys compares cfg.Keys[i].Key against the plaintext
// gateway_keys entries; with ciphertext keys the comparison never
// matched and every restart appended another duplicate admin key
// (production had four copies of the same admin key).
// - rebuildRegistry hands cfg.Sources[i].APIKey to the providers; with
// ciphertext it built every upstream client with "enc:v1:..." as its
// bearer token.
// Previously the unseal happened at the END of this function, and things
// only appeared to work because seedKeys' Save() unsealed memory as a side
// effect. Removing the redundant saves exposed the real ordering bug.
cfg.AttachSecretBox(c.store.SecretBox())
hadPlaintextSecrets, err := cfg.UnsealSecrets(c.store.SecretBox())
if err != nil {
return nil, fmt.Errorf("unseal secrets: %w", err)
}
if err := c.store.Load(); err != nil {
return nil, fmt.Errorf("runtime store: %w", err)
}
@ -79,10 +97,10 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
if err := c.seedPresetTemplates(); err != nil {
return nil, fmt.Errorf("seed preset templates: %w", err)
}
// Seal any credential still in the clear in config.yaml. Idempotent: an
// already-sealed config is not rewritten, so a normal restart writes
// nothing. This is the only place that rewrites the file on startup.
if err := cfg.NormalizeSecretsForRun(c.store.SecretBox()); err != nil {
// Seal any credential still in the clear in config.yaml. Startup writes
// nothing when the file was already sealed (hadPlaintextSecrets is decided
// from the on-disk state, before the unseal above).
if err := cfg.SealIfNeeded(hadPlaintextSecrets); err != nil {
return nil, fmt.Errorf("normalize secrets: %w", err)
}
return c, nil

View File

@ -0,0 +1,178 @@
package core
import (
"os"
"path/filepath"
"strings"
"testing"
"llmsproxy/internal/config"
)
// Startup must be idempotent, and every credential consumer must see PLAINTEXT.
//
// Two bugs shared one root cause: NewFromConfig unsealed the config at the END,
// after the steps that read credentials had already run.
//
// 1. seedKeys() deduped cfg.Keys[i].Key against the plaintext gateway_keys
// entries. With ciphertext keys the comparison never matched, so every
// restart appended another copy of the same admin key — production reached
// four identical admin keys.
// 2. rebuildRegistry() handed cfg.Sources[i].APIKey to provider.New, so with
// ciphertext it built every upstream client with "enc:v1:..." as its bearer
// token and every upstream call would 401.
//
// They masked each other: seedKeys' Save() unsealed memory as a side effect, so
// bug 2 was invisible until the redundant saves were removed. These tests drive
// the real entry point (core.New -> config.Load off disk), because building a
// fresh in-memory config per attempt hides both bugs — that is exactly how the
// first draft of this test failed to reproduce anything.
// writeSealedInstall creates a config on disk the way a real install looks
// after its first run: credentials already sealed, gateway_keys still plaintext.
func writeSealedInstall(t *testing.T, dir string) string {
t.Helper()
path := filepath.Join(dir, "config.yaml")
cfg := &config.Config{
Path: path,
AdapterDir: filepath.Join(dir, "adapters"),
RuntimeFile: filepath.Join(dir, "runtime.json"),
Listen: "127.0.0.1:0",
DefaultModel: "AUTO",
GatewayKeys: []string{"sk-gw-seed-me"},
Sources: []config.Source{{
Name: "up", BaseURL: "http://127.0.0.1:1/v1", APIKey: "sk-upstream-secret",
Adapter: "openai", Models: []config.Model{{ID: "gpt-4o", Priority: 10}},
}},
}
c, err := NewFromConfig(cfg)
if err != nil {
t.Fatalf("initial install: %v", err)
}
c.Close()
raw, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(raw), "enc:v1:") {
t.Fatalf("fixture is wrong: on-disk config holds no sealed credential, "+
"so this test cannot detect ciphertext leaking into consumers:\n%s", raw)
}
return path
}
// writeSealedInstallNoSeed is writeSealedInstall without gateway_keys, so no
// key is ever seeded and no startup Save() can unseal the config as a side
// effect. That isolation is what makes the ciphertext leak observable.
func writeSealedInstallNoSeed(t *testing.T, dir string) string {
t.Helper()
path := filepath.Join(dir, "config.yaml")
cfg := &config.Config{
Path: path,
AdapterDir: filepath.Join(dir, "adapters"),
RuntimeFile: filepath.Join(dir, "runtime.json"),
Listen: "127.0.0.1:0",
DefaultModel: "AUTO",
Sources: []config.Source{{
Name: "up", BaseURL: "http://127.0.0.1:1/v1", APIKey: "sk-upstream-secret",
Adapter: "openai", Models: []config.Model{{ID: "gpt-4o", Priority: 10}},
}},
}
c, err := NewFromConfig(cfg)
if err != nil {
t.Fatalf("initial install: %v", err)
}
c.Close()
raw, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(raw), "enc:v1:") {
t.Fatalf("fixture is wrong: no sealed credential on disk:\n%s", raw)
}
return path
}
// countKeysWithSecret loads the config, unseals it, and counts keys whose
// plaintext equals one of the gateway_keys entries.
func countKeysWithSecret(t *testing.T, path string) (total, matching int) {
t.Helper()
cfg, err := config.Load(path)
if err != nil {
t.Fatal(err)
}
box, err := config.NewSecretBox(cfg.RuntimeFile)
if err != nil {
t.Fatal(err)
}
if _, err := cfg.UnsealSecrets(box); err != nil {
t.Fatal(err)
}
want := map[string]bool{}
for _, g := range cfg.GatewayKeys {
want[g] = true
}
for _, k := range cfg.Keys {
total++
if want[k.Key] {
matching++
}
}
return total, matching
}
// Restarting must not grow the key list. Before the fix this went 2 -> 3 -> 4.
func TestRestartDoesNotDuplicateSeededKeys(t *testing.T) {
dir := t.TempDir()
path := writeSealedInstall(t, dir)
if _, n := countKeysWithSecret(t, path); n != 1 {
t.Fatalf("after install: %d keys carry the seeded secret, want 1", n)
}
for i := 2; i <= 4; i++ {
c, err := New(path) // production path: reads the file
if err != nil {
t.Fatalf("restart %d: %v", i, err)
}
c.Close()
total, n := countKeysWithSecret(t, path)
if n != 1 {
t.Fatalf("restart %d: %d keys carry the seeded secret (total %d), want "+
"exactly 1 — seedKeys is comparing ciphertext against plaintext and "+
"appending a duplicate on every start", i, n, total)
}
}
}
// Every provider must hold the real upstream credential, not its ciphertext.
//
// This deliberately runs with gateway_keys EMPTY. The buggy order was masked
// whenever seedKeys had work to do, because its Save() unsealed memory as a
// side effect — providers then happened to see plaintext. With no key to seed,
// no Save runs, and the ciphertext reaches the registry directly. An earlier
// version of this test kept gateway_keys populated and was insensitive: it
// passed even with the bug present.
func TestProvidersGetPlaintextCredentials(t *testing.T) {
dir := t.TempDir()
path := writeSealedInstallNoSeed(t, dir)
c, err := New(path)
if err != nil {
t.Fatalf("restart: %v", err)
}
defer c.Close()
for _, p := range c.registry.Providers() {
key := p.Config().APIKey
if strings.HasPrefix(key, "enc:v1:") {
t.Errorf("provider %q holds CIPHERTEXT api_key %q — every upstream call "+
"would 401", p.Name(), key)
}
if key != "sk-upstream-secret" {
t.Errorf("provider %q api_key = %q, want the plaintext upstream secret",
p.Name(), key)
}
}
}

View File

@ -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(),

View File

@ -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.

View File

@ -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)
}
}

View File

@ -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 @@
<td class="t-tag">${fmtTime(r.time)}</td>
<td class="num">${r.ok ? `<span class="tag tag-green">${r.status || 200}</span>` : `<span class="tag tag-red" title="${esc(r.error || "")}">${r.status || 500}</span>`}</td>
<td>${esc(keyNames[r.key] ? keyNames[r.key] + " · " + r.key : r.key)}</td><td class="t-tag">${esc(r.type)}</td><td>${esc(r.model)}</td><td>${esc(r.source || "")}</td>
<td class="num">${fmtTok(r.prompt_tokens)}</td><td class="num">${fmtTok(r.completion_tokens)}</td>
<td class="num">${fmtTok(r.prompt_tokens)}</td><td class="num">${r.type === "image" ? fmtN(r.image_count || 0) + " " + t("thImgs") : fmtTok(r.completion_tokens)}</td>
<td class="num">${cacheCell(r)}</td>
<td class="num">${fmtMs(r.latency_ms)}</td></tr>`;
}

View File

@ -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))
}
}