ModelRouter/internal/gateway/stats_test.go

package gateway

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestStatsByStatus(t *testing.T) {
	s := NewStats(100)
	s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 200, OK: true})
	s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 402, OK: false})
	s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 400, OK: false})
	snap := s.Snapshot(0, "")
	bs, ok := snap["by_status"].([]agrRow)
	if !ok {
		t.Fatalf("by_status missing: %#v", snap["by_status"])
	}
	if len(bs) != 3 {
		t.Fatalf("want 3 status buckets, got %d: %#v", len(bs), bs)
	}
	if bs[0].Name != "200" || bs[0].OK != 1 || bs[0].Err != 0 {
		t.Fatalf("bucket 200 wrong: %#v", bs[0])
	}
	if bs[1].Name != "400" || bs[1].Err != 1 {
		t.Fatalf("bucket 400 wrong: %#v", bs[1])
	}
	if bs[2].Name != "402" || bs[2].Err != 1 {
		t.Fatalf("bucket 402 wrong: %#v", bs[2])
	}
}

func TestAuditRotation(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	s := NewStats(10)
	s.LoadAudit(path)

	oldRotate, oldKeep := auditRotateBytes, auditKeepOld
	auditRotateBytes, auditKeepOld = 64, 10
	defer func() { auditRotateBytes, auditKeepOld = oldRotate, oldKeep }()

	oldFiles := func() []string {
		matches, _ := filepath.Glob(path + ".*.old")
		return matches
	}

	for i := 0; i < 3; i++ {
		s.AppendAudit("ev", map[string]interface{}{"i": i})
	}
	if got := len(oldFiles()); got != 1 {
		t.Fatalf("want 1 rotated file after first overflow, got %d", got)
	}
	if b, err := os.ReadFile(path); err != nil || len(b) == 0 {
		t.Fatalf("active audit file must continue appending: %v %d bytes", err, len(b))
	}

	// seed 12 fake old files; the next rotation must prune back to keep=10
	for i := 1; i <= 12; i++ {
		name := fmt.Sprintf("%s.%010d.old", path, i)
		_ = os.WriteFile(name, []byte("x\n"), 0644)
	}
	s.AppendAudit("ev", map[string]interface{}{"i": 98})
	s.AppendAudit("ev", map[string]interface{}{"i": 99})
	if got := len(oldFiles()); got != auditKeepOld {
		t.Fatalf("want keeper %d old files, got %d", auditKeepOld, got)
	}
}

func TestAuditRotationRecords(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	s := NewStats(10)
	s.LoadAudit(path)

	oldRotate := auditRotateBytes
	auditRotateBytes = 64
	defer func() { auditRotateBytes = oldRotate }()

	for i := 0; i < 5; i++ {
		s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 200, OK: true})
	}
	matches, _ := filepath.Glob(path + ".*.old")
	if len(matches) != 1 {
		t.Fatalf("Record must rotate too: got %d old files", len(matches))
	}
}

func TestLoadAuditFullReplay(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	// timestamps anchored to whole unix hours (see aggregateLocked): window
	// assertions must hold no matter which minute-of-hour the suite runs at,
	// so rows are placed relative to hour boundaries, not raw offsets from now.
	now := time.Now()
	thisHour := now.Truncate(time.Hour)
	twoBucketsAgo := thisHour.Add(-90 * time.Minute) // bucket H-2
	prevBucketMid := thisHour.Add(-30 * time.Minute) // bucket H-1
	lines := []string{
		`{"obj":"access","time":1699999999000,"key":"k","method":"GET","path":"/api/stats","status":200}`,
		fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":100,"completion_tokens":50,"latency_ms":10,"ok":true,"status":200}`, twoBucketsAgo.UnixMilli()),
		`{this is not valid json`,
		fmt.Sprintf(`{"time":%d,"key":"k","type":"stream","model":"m","source":"s","prompt_tokens":200,"completion_tokens":20,"latency_ms":20,"ok":false,"status":503}`, prevBucketMid.UnixMilli()),
		"garbage-not-json\n",
		fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m2","source":"s","prompt_tokens":7,"completion_tokens":3,"latency_ms":5,"ok":true,"status":200}`, prevBucketMid.UnixMilli()),
	}
	loaded := strings.Join(lines, "\n") + "\n" + strings.Repeat("x", 1<<18) + "\n"
	// oversized row at the END proves the scanner tolerates >64KB lines and
	// still finishes the replay instead of truncating silently.
	loaded += fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":1,"ok":true,"status":200}`, now.UnixMilli()) + "\n"
	if err := os.WriteFile(path, []byte(loaded), 0644); err != nil {
		t.Fatal(err)
	}

	s := NewStats(1000)
	s.LoadAudit(path)

	got := s.byModel["m"]
	if got == nil || got.Tokens != 100+50+200+20+1+1 {
		t.Fatalf("aggregates must replay EVERY request row, got %#v", got)
	}
	if s.byModel["m2"] == nil || s.byModel["m2"].Tokens != 10 {
		t.Fatalf("m2 must be replayed too, got %#v", s.byModel["m2"])
	}
	// access rows are not requests: 4 real rows, junk skipped
	if len(s.recs) != 4 {
		t.Fatalf("ring must hold only real requests, got %d rows: %#v", len(s.recs), s.recs)
	}
	// quota window rebuilt from full history. All-time and multi-hour windows
	// must see everything. A 1h window must NOT return all records (the old
	// ms/seconds unit bug made any sec>0 window return everything): buckets
	// are whole unix hours, so only rows in the current hour bucket qualify
	// — exactly the just-now row (2 tokens).
	if w := s.WindowTokens("m", "s", 0); w != 372 {
		t.Fatalf("all-time window want 372, got %d", w)
	}
	if w := s.WindowTokens("m", "s", 3*hourSec); w != 372 {
		t.Fatalf("3h window want 372, got %d", w)
	}
	if w := s.WindowTokens("m", "s", 24*hourSec); w != 372 {
		t.Fatalf("24h window want 372, got %d", w)
	}
	if w := s.WindowTokens("m", "s", hourSec); w != 2 {
		t.Fatalf("1h window want 2 (current-bucket row only), got %d", w)
	}
	if w := s.WindowTokens("m2", "s", 24*hourSec); w != 10 {
		t.Fatalf("m2 24h window want 10, got %d", w)
	}
	// by_status only from requests (200 x3, 503 x1) — access line must not count
	if st := s.byStatus[200]; st == nil || st.Reqs != 3 {
		t.Fatalf("by_status 200 want 3 reqs, got %#v", st)
	}
	if s.byStatus[200].Err != 0 || s.byStatus[503] == nil || s.byStatus[503].Reqs != 1 {
		t.Fatalf("by_status wrong: %#v", s.byStatus)
	}
}

// ---- on-demand log loading (plan 阶段 3) ----

// writeAuditFile writes n synthetic request rows (plus some non-request event
// rows, which must be skipped) and returns the path.
func writeAuditFile(t *testing.T, path string, n int, key string, startMs int64) {
	t.Helper()
	f, err := os.Create(path)
	if err != nil {
		t.Fatalf("create audit: %v", err)
	}
	defer f.Close()
	for i := 0; i < n; i++ {
		fmt.Fprintf(f, `{"time":%d,"key":%q,"type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":5,"ok":true,"status":200}`+"\n",
			startMs+int64(i)*1000, key)
		if i%10 == 0 {
			// an access-log/event row: must never appear as a record
			fmt.Fprintf(f, `{"obj":"access","time":%d,"path":"/v1/models"}`+"\n", startMs+int64(i)*1000)
		}
	}
}

// TestLoadAuditBoundedTail is the startup-memory fix: only the tail of a large
// audit file is replayed, and the ring stays capped.
func TestLoadAuditBoundedTail(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	writeAuditFile(t, path, 20000, "k1", 1_700_000_000_000)

	fi, err := os.Stat(path)
	if err != nil {
		t.Fatal(err)
	}
	orig := auditReplayBytes
	auditReplayBytes = 64 << 10 // 64 KB tail out of a much larger file
	defer func() { auditReplayBytes = orig }()
	if fi.Size() <= auditReplayBytes {
		t.Fatalf("fixture too small (%d bytes) to exercise bounded replay", fi.Size())
	}

	s := NewStats(0)
	s.LoadAudit(path)
	if !s.ReplayPartial() {
		t.Fatal("a truncated replay must be reported as partial")
	}
	snap := s.Snapshot(0, "")
	recs, _ := snap["records"].([]Req)
	if len(recs) == 0 {
		t.Fatal("bounded replay must still load the recent window")
	}
	if len(recs) > firstScreenRecords {
		t.Fatalf("snapshot shipped %d records, must be capped at the first screen (%d)",
			len(recs), firstScreenRecords)
	}
	// the aggregates must come from the tail only, not the whole file
	total, _ := snap["total"].(Stat)
	if total.Reqs == 0 {
		t.Fatal("tail replay must feed the aggregates")
	}
	if total.Reqs >= 20000 {
		t.Fatalf("aggregates replayed the whole file (%d reqs); replay must be bounded", total.Reqs)
	}
}

// TestLoadAuditFullFileWhenSmall: a file inside the replay window is loaded
// completely and is not flagged partial.
func TestLoadAuditFullFileWhenSmall(t *testing.T) {
	path := filepath.Join(t.TempDir(), "audit.jsonl")
	writeAuditFile(t, path, 50, "k1", 1_700_000_000_000)
	s := NewStats(0)
	s.LoadAudit(path)
	if s.ReplayPartial() {
		t.Fatal("a small file must be replayed in full, not flagged partial")
	}
	snap := s.Snapshot(0, "")
	total, _ := snap["total"].(Stat)
	if total.Reqs != 50 {
		t.Fatalf("total reqs = %d, want 50", total.Reqs)
	}
}

// TestQuotaWindowsSurviveBoundedReplay is the regression guard the plan calls
// out: shrinking the ring must not break token-quota accounting, which reads the
// hour buckets rather than the ring.
func TestQuotaWindowsSurviveBoundedReplay(t *testing.T) {
	s := NewStats(10) // deliberately tiny ring
	now := time.Now().UnixMilli()
	for i := 0; i < 100; i++ {
		s.Record(Req{Time: now, Key: "k", Model: "m", Source: "s", Prompt: 2, Compl: 3, OK: true, Status: 200})
	}
	if got := s.WindowTokens("m", "s", 24*hourSec); got != 500 {
		t.Fatalf("window tokens = %d, want 500 (quota accounting must not depend on the ring)", got)
	}
	if got := s.ModelTokens("k")["m"]; got != 500 {
		t.Fatalf("model tokens = %d, want 500", got)
	}
	if got := s.KeyTokens("k"); got != 500 {
		t.Fatalf("key tokens = %d, want 500", got)
	}
}

// TestAuditPagePagesBackwards walks the whole history one page at a time and
// asserts strict newest-first order with no gaps or repeats.
func TestAuditPagePagesBackwards(t *testing.T) {
	path := filepath.Join(t.TempDir(), "audit.jsonl")
	const n = 250
	writeAuditFile(t, path, n, "k1", 1_700_000_000_000)
	s := NewStats(0)
	s.LoadAudit(path)

	var seen []Req
	cursor := ""
	for i := 0; i < 20; i++ { // bounded loop guards against a cursor that never advances
		page := s.AuditPage(cursor, 40, "")
		if len(page.Records) == 0 {
			break
		}
		seen = append(seen, page.Records...)
		if !page.HasMore {
			break
		}
		if page.Next == cursor {
			t.Fatal("cursor did not advance")
		}
		cursor = page.Next
	}
	if len(seen) != n {
		t.Fatalf("paged %d records, want %d", len(seen), n)
	}
	for i := 1; i < len(seen); i++ {
		if seen[i-1].Time < seen[i].Time {
			t.Fatalf("records not newest-first at %d: %d then %d", i, seen[i-1].Time, seen[i].Time)
		}
	}
	// newest record first, oldest last
	if seen[0].Time != 1_700_000_000_000+int64(n-1)*1000 {
		t.Fatalf("first record time = %d, want the newest", seen[0].Time)
	}
	if seen[len(seen)-1].Time != 1_700_000_000_000 {
		t.Fatalf("last record time = %d, want the oldest", seen[len(seen)-1].Time)
	}
}

// TestAuditPageFirstScreenIsSmall: the first page must be exactly one screen,
// independent of how large the audit file is.
func TestAuditPageFirstScreenIsSmall(t *testing.T) {
	path := filepath.Join(t.TempDir(), "audit.jsonl")
	writeAuditFile(t, path, 5000, "k1", 1_700_000_000_000)
	s := NewStats(0)
	s.LoadAudit(path)
	page := s.AuditPage("", 100, "")
	if len(page.Records) != 100 {
		t.Fatalf("first page = %d records, want 100", len(page.Records))
	}
	if !page.HasMore || page.Next == "" {
		t.Fatal("a large file must report more pages plus a cursor")
	}
}

// TestAuditPageFiltersByKey: a user key may only page its own records.
func TestAuditPageFiltersByKey(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	writeAuditFile(t, path, 30, "mine", 1_700_000_000_000)
	f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 30; i++ {
		fmt.Fprintf(f, `{"time":%d,"key":"other","type":"chat","model":"m","source":"s","ok":true,"status":200}`+"\n",
			1_700_000_100_000+int64(i)*1000)
	}
	f.Close()

	s := NewStats(0)
	s.LoadAudit(path)
	page := s.AuditPage("", 1000, "mine")
	if len(page.Records) != 30 {
		t.Fatalf("filtered page = %d records, want 30", len(page.Records))
	}
	for _, r := range page.Records {
		if r.Key != "mine" {
			t.Fatalf("key filter leaked record for %q", r.Key)
		}
	}
}

// TestAuditPageSpansRotatedFiles: paging must continue into .old files instead
// of stopping at the live file's start.
func TestAuditPageSpansRotatedFiles(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	writeAuditFile(t, path, 40, "k1", 1_700_000_500_000)            // newest
	writeAuditFile(t, path+".100.old", 40, "k1", 1_700_000_300_000) // older
	writeAuditFile(t, path+".050.old", 40, "k1", 1_700_000_100_000) // oldest

	s := NewStats(0)
	s.LoadAudit(path)
	var seen []Req
	cursor := ""
	for i := 0; i < 30; i++ {
		page := s.AuditPage(cursor, 25, "")
		seen = append(seen, page.Records...)
		if !page.HasMore {
			break
		}
		cursor = page.Next
	}
	if len(seen) != 120 {
		t.Fatalf("paged %d records across rotated files, want 120", len(seen))
	}
	for i := 1; i < len(seen); i++ {
		if seen[i-1].Time < seen[i].Time {
			t.Fatalf("rotation boundary broke the ordering at %d", i)
		}
	}
}

// TestAuditPageRotatedCursor: a cursor whose file rotated away must restart from
// the newest page and say so, rather than silently returning nothing.
func TestAuditPageRotatedCursor(t *testing.T) {
	path := filepath.Join(t.TempDir(), "audit.jsonl")
	writeAuditFile(t, path, 20, "k1", 1_700_000_000_000)
	s := NewStats(0)
	s.LoadAudit(path)
	page := s.AuditPage("/gone/audit.jsonl:12345", 10, "")
	if !page.Rotated {
		t.Fatal("a stale cursor must be reported as rotated")
	}
	if len(page.Records) == 0 {
		t.Fatal("a stale cursor must fall back to the newest page")
	}
}

// TestAuditPageEmptyAndNoFile cover the degenerate inputs.
func TestAuditPageEmptyAndNoFile(t *testing.T) {
	s := NewStats(0)
	if page := s.AuditPage("", 10, ""); len(page.Records) != 0 || page.HasMore {
		t.Fatalf("no audit path must yield an empty page, got %+v", page)
	}
	path := filepath.Join(t.TempDir(), "audit.jsonl")
	if err := os.WriteFile(path, nil, 0644); err != nil {
		t.Fatal(err)
	}
	s2 := NewStats(0)
	s2.LoadAudit(path)
	if page := s2.AuditPage("", 10, ""); len(page.Records) != 0 || page.HasMore {
		t.Fatalf("empty audit file must yield an empty page, got %+v", page)
	}
}

// TestAuditPageHandlesLongLines: a record larger than the reverse-read chunk
// must still be decoded (the reader carries partial lines across chunks).
func TestAuditPageHandlesLongLines(t *testing.T) {
	path := filepath.Join(t.TempDir(), "audit.jsonl")
	f, err := os.Create(path)
	if err != nil {
		t.Fatal(err)
	}
	big := strings.Repeat("e", revChunk+1024) // one record spanning >1 chunk
	fmt.Fprintf(f, `{"time":1,"key":"k","type":"chat","model":"m","source":"s","ok":false,"status":500,"error":%q}`+"\n", big)
	fmt.Fprintf(f, `{"time":2,"key":"k","type":"chat","model":"m","source":"s","ok":true,"status":200}`+"\n")
	f.Close()

	s := NewStats(0)
	s.LoadAudit(path)
	page := s.AuditPage("", 10, "")
	if len(page.Records) != 2 {
		t.Fatalf("got %d records, want 2 (a >chunk record must survive reverse reads)", len(page.Records))
	}
	if page.Records[0].Time != 2 || page.Records[1].Time != 1 {
		t.Fatalf("ordering broken: %d then %d", page.Records[0].Time, page.Records[1].Time)
	}
	if len(page.Records[1].Err) != len(big) {
		t.Fatalf("oversized field truncated: %d bytes, want %d", len(page.Records[1].Err), len(big))
	}
}

// TestStreamAuditRecordsIsWindowed: export streaming must honour the time window
// and the key filter, and must hand records over one at a time.
func TestStreamAuditRecordsIsWindowed(t *testing.T) {
	path := filepath.Join(t.TempDir(), "audit.jsonl")
	writeAuditFile(t, path, 100, "k1", 1_700_000_000_000)
	s := NewStats(0)
	s.LoadAudit(path)

	from := int64(1_700_000_010_000)
	to := int64(1_700_000_019_000)
	var got []Req
	if err := s.StreamAuditRecords(from, to, "", func(r Req) error {
		got = append(got, r)
		return nil
	}); err != nil {
		t.Fatalf("stream: %v", err)
	}
	if len(got) != 10 {
		t.Fatalf("streamed %d records, want 10 in [%d,%d]", len(got), from, to)
	}
	for _, r := range got {
		if r.Time < from || r.Time > to {
			t.Fatalf("record outside the window: %d", r.Time)
		}
	}
}

// TestStreamAuditRecordsAbortStops: a callback error (client disconnect) must
// stop the walk immediately instead of reading the rest of the history.
func TestStreamAuditRecordsAbortStops(t *testing.T) {
	path := filepath.Join(t.TempDir(), "audit.jsonl")
	writeAuditFile(t, path, 500, "k1", 1_700_000_000_000)
	s := NewStats(0)
	s.LoadAudit(path)

	seen := 0
	stop := fmt.Errorf("client gone")
	err := s.StreamAuditRecords(0, 0, "", func(Req) error {
		seen++
		if seen == 5 {
			return stop
		}
		return nil
	})
	if err != stop {
		t.Fatalf("abort must surface the callback error, got %v", err)
	}
	if seen != 5 {
		t.Fatalf("walk continued after abort: %d records", seen)
	}
}

// TestAuditRotationThresholds pins the smaller-file/more-files trade-off that
// makes reverse paging cheap while keeping the same total budget.
func TestAuditRotationThresholds(t *testing.T) {
	if auditRotateBytes != 16<<20 {
		t.Fatalf("auditRotateBytes = %d, want 16 MiB", auditRotateBytes)
	}
	if auditKeepOld != 16 {
		t.Fatalf("auditKeepOld = %d, want 16", auditKeepOld)
	}
	if total := auditRotateBytes * int64(auditKeepOld); total != 256<<20 {
		t.Fatalf("total audit budget = %d, want 256 MiB", total)
	}
}