ModelRouter/internal/gateway/token_units_test.go

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"}},
	})

	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"}},
	})

	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"}},
	})

	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", 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)
	}
}