ModelRouter/internal/gateway/stats_period_records_test.go

package gateway

import (
	"encoding/json"
	"net/url"
	"path/filepath"
	"testing"
	"time"
)

// The dashboard paints its records table from the `records` field of
// /api/stats (paintRecords(st.records)). The period branch of that handler
// used to hand-write a response map and simply omit the key, so switching the
// selector to day / week / month rendered an EMPTY table no matter how much
// traffic the window had — while the lifetime view worked, because it goes
// through Snapshot() which carries records.
//
// These tests pin the contract across every period: same records, same cursor.
func TestStatsPeriodViewsShipRecords(t *testing.T) {
	td := t.TempDir()
	g := newTestGateway(t)

	// Write audit records the period aggregator and the record pager both read.
	g.stats.auditPath = filepath.Join(td, "audit.jsonl")
	now := time.Now().UnixMilli()
	for i := 0; i < 5; i++ {
		g.stats.Record(Req{
			Time: now - int64(i)*1000, Type: "chat", OK: true,
			Model: "m1", Source: "s1", Status: 200,
			Prompt: 10, Compl: 2,
		})
	}

	for _, p := range []string{"day", "week", "month", "all"} {
		rr := doReq(t, g, "GET", "/api/stats?period="+p, "")
		if rr.Code != 200 {
			t.Fatalf("period=%s: HTTP %d: %s", p, rr.Code, rr.Body.String())
		}
		var st map[string]interface{}
		if err := json.Unmarshal(rr.Body.Bytes(), &st); err != nil {
			t.Fatalf("period=%s: %v", p, err)
		}
		recs, ok := st["records"].([]interface{})
		if !ok {
			t.Fatalf("period=%s: response has no `records` array — the dashboard "+
				"records table renders empty for this view (keys: %v)", p, keysOf(st))
		}
		if len(recs) != 5 {
			t.Fatalf("period=%s: got %d records, want 5", p, len(recs))
		}
		// The scroll handler pages from next_cursor; without it the table is
		// stuck at one screen no matter how far the user scrolls.
		if _, ok := st["next_cursor"]; !ok {
			t.Fatalf("period=%s: no next_cursor — paging can never continue", p)
		}
	}
}

// The records shipped with a period view must be the newest-first page the
// paging endpoint would return, in the same order. If /api/stats and
// /api/stats/records disagreed, scrolling would duplicate or skip rows.
func TestStatsPeriodRecordsMatchPager(t *testing.T) {
	td := t.TempDir()
	g := newTestGateway(t)
	g.stats.auditPath = filepath.Join(td, "audit.jsonl")
	now := time.Now().UnixMilli()
	for i := 0; i < 8; i++ {
		g.stats.Record(Req{
			Time: now - int64(i)*1000, Type: "chat", OK: true,
			Model: "m1", Status: 200, Prompt: 5, Compl: 1,
		})
	}

	rr := doReq(t, g, "GET", "/api/stats?period=day&limit=100", "")
	var st struct {
		Records []struct {
			Time  int64  `json:"time"`
			Model string `json:"model"`
		} `json:"records"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &st); err != nil {
		t.Fatal(err)
	}
	rr2 := doReq(t, g, "GET", "/api/stats/records?limit=100", "")
	var pg struct {
		Records []struct {
			Time  int64  `json:"time"`
			Model string `json:"model"`
		} `json:"records"`
	}
	if err := json.Unmarshal(rr2.Body.Bytes(), &pg); err != nil {
		t.Fatal(err)
	}
	if len(st.Records) != len(pg.Records) {
		t.Fatalf("period view shipped %d records, pager returned %d — the two "+
			"sources must agree or scrolling duplicates/skips rows",
			len(st.Records), len(pg.Records))
	}
	for i := range st.Records {
		if st.Records[i].Time != pg.Records[i].Time {
			t.Fatalf("record %d differs: period view t=%d, pager t=%d",
				i, st.Records[i].Time, pg.Records[i].Time)
		}
	}
	// The two sources must agree on DIRECTION as well as content. Which
	// direction that is depends on how many reverse-read chunks the audit file
	// spans (a small file is read whole, a large one chunk-by-chunk), so the
	// test asserts agreement rather than hard-coding newest-first — but they
	// can never disagree, because the table interleaves both streams.
	if len(st.Records) > 1 {
		periodAsc := st.Records[1].Time > st.Records[0].Time
		pagerAsc := pg.Records[1].Time > pg.Records[0].Time
		if periodAsc != pagerAsc {
			t.Fatalf("direction differs: period view %s, pager %s — the table "+
				"would splice the two streams out of order",
				map[bool]string{true: "oldest-first", false: "newest-first"}[periodAsc],
				map[bool]string{true: "oldest-first", false: "newest-first"}[pagerAsc])
		}
	}
}

// A non-admin caller cannot widen their view by passing someone else's key in
// the query string: exportKey overrides ?key= with the caller's own masked id.
// This is a data-leak boundary, and it applies to the records the period view
// ships just as much as to the aggregates.
func TestStatsPeriodRecordsIgnoreForeignKeyForUsers(t *testing.T) {
	td := t.TempDir()
	g := newTestGateway(t)
	g.stats.auditPath = filepath.Join(td, "audit.jsonl")

	other := keyID("sk-gw-cccccccccccccccc")
	own := keyID("sk-test")
	now := time.Now().UnixMilli()
	g.stats.Record(Req{Time: now, OK: true, Model: "m", Status: 200, Key: other})
	g.stats.Record(Req{Time: now - 1000, OK: true, Model: "m", Status: 200, Key: own})

	// sk-test authenticates as a non-admin (newTestGateway seeds it under
	// gateway_keys, not keys), so ?key=<other> must be ignored and the caller
	// sees only its own rows.
	rr := doReq(t, g, "GET",
		"/api/stats?period=day&key="+url.QueryEscape(other), "")
	if rr.Code != 200 {
		t.Fatalf("HTTP %d: %s", rr.Code, rr.Body.String())
	}
	var st struct {
		Records []struct {
			Key string `json:"key"`
		} `json:"records"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &st); err != nil {
		t.Fatal(err)
	}
	for _, r := range st.Records {
		if r.Key == other {
			t.Fatalf("another key's record leaked into a user view "+
				"(caller=%s, requested=%s)", own, other)
		}
	}
}