ModelRouter/internal/gateway/stats.go

package gateway

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"sync"
	"time"
)

// Req is one recorded gateway request (audit trail + per-key/per-model stats).
type Req struct {
	Time  int64  `json:"time"`  // unix milliseconds
	Key   string `json:"key"`   // gateway key id
	Type  string `json:"type"`  // chat | stream | image
	Model string `json:"model"` // effective model used upstream
	// Source provider/source name
	Source string `json:"source"`
	// Prompt prompt tokens
	Prompt int64 `json:"prompt_tokens"`
	// Compl completion tokens
	Compl int64 `json:"completion_tokens"`
	// LatMs total handling time ms
	LatMs int64 `json:"latency_ms"`
	// FirstByteMs time-to-first-byte for streaming (ms from request start
	// to the first SSE chunk sent to the client); for non-streaming it
	// equals LatMs. 0 when unmeasured (legacy records).
	FirstByteMs int64 `json:"first_byte_ms,omitempty"`
	// CacheHit / CacheMiss carry the upstream prompt-cache accounting
	// (DeepSeek-style hit/miss tokens) when the upstream reports it.
	// Both 0 = upstream gave no cache data.
	CacheHit  int64 `json:"cache_hit_tokens,omitempty"`
	CacheMiss int64 `json:"cache_miss_tokens,omitempty"`
	// CacheReported marks that the upstream usage reported cache
	// accounting at all (even when the hit count is 0). The WebUI shows
	// "0%" instead of "—" for such rows.
	CacheReported bool `json:"cache_reported,omitempty"`
	OK    bool  `json:"ok"`
	// Status http status code
	Status int `json:"status"`
	// Err short error message
	Err string `json:"error,omitempty"`
}

// Stat aggregates counters for one dimension row.
type Stat struct {
	Reqs        int64 `json:"reqs"`
	OK          int64 `json:"ok"`
	Err         int64 `json:"err"`
	Tokens      int64 `json:"tokens"`
	Prompt      int64 `json:"prompt_tokens"`
	Compl       int64 `json:"completion_tokens"`
	LatSum      int64 `json:"latency_sum_ms"`
	LatMax      int64 `json:"latency_max_ms"`
	FirstByteSum int64 `json:"first_byte_sum_ms,omitempty"`
}

type agrRow struct {
	Name string `json:"name"`
	Stat
}

// Stats collects per-key / per-model / per-source aggregates plus a bounded
// ring of raw request records, all guarded by one mutex.
type Stats struct {
	mu         sync.Mutex
	active     int64
	byKey      map[string]*Stat
	byModel    map[string]*Stat
	bySrc      map[string]*Stat
	byKeyModel map[string]map[string]*Stat
	byKeySrc   map[string]map[string]*Stat
	byStatus   map[int]*Stat // per http status code aggregates (incl. 402/400)
	recs       []Req
	maxRecs    int
	auditPath  string
	modelHour  map[string]map[int64]int64 // model -> unix-hour bucket -> tokens
}

const hourSec = 3600

// auditRotateBytes rotates the audit file once it grows past this size (the
// file is renamed to <path>.<unix>.old and a fresh one is started); pruning
// keeps at most auditKeepOld rotated files. Both are vars so tests can shrink
// the threshold.
var (
	auditRotateBytes int64 = 64 << 20
	auditKeepOld           = 10
)

func NewStats(maxRecords int) *Stats {
	if maxRecords <= 0 {
		maxRecords = 3000
	}
	return &Stats{
		byKey:      map[string]*Stat{},
		byModel:    map[string]*Stat{},
		bySrc:      map[string]*Stat{},
		byKeyModel: map[string]map[string]*Stat{},
		byKeySrc:   map[string]map[string]*Stat{},
		byStatus:   map[int]*Stat{},
		modelHour:  map[string]map[int64]int64{},
		maxRecs:    maxRecords,
	}
}

// Begin accounts an in-flight request; the returned func must be called once
// the request finished (defer ok).
func (s *Stats) Begin() func() {
	s.mu.Lock()
	s.active++
	s.mu.Unlock()
	return func() {
		s.mu.Lock()
		s.active--
		s.mu.Unlock()
	}
}

func inc(m map[string]*Stat, name string, r Req) {
	a := m[name]
	if a == nil {
		a = &Stat{}
		m[name] = a
	}
	incStatus(a, name, r)
}

func incStatus(a *Stat, name string, r Req) {
	a.Reqs++
	if r.OK {
		a.OK++
	} else {
		a.Err++
	}
	a.Tokens += r.Prompt + r.Compl
	a.Prompt += r.Prompt
	a.Compl += r.Compl
	a.LatSum += r.LatMs
	if r.LatMs > a.LatMax {
		a.LatMax = r.LatMs
	}
	if r.FirstByteMs > 0 {
		a.FirstByteSum += r.FirstByteMs
	}
}

func (s *Stats) LoadAudit(path string) {
	f, err := os.Open(path)
	if err == nil {
		defer f.Close()
		s.mu.Lock()
		defer s.mu.Unlock()
		var recs []Req
		// tolerate long error summaries / oversized junk lines
		sc := bufio.NewScanner(f)
		sc.Buffer(make([]byte, 64*1024), 16*1024*1024)
		for sc.Scan() {
			var r Req
			if json.Unmarshal(sc.Bytes(), &r) != nil || r.Type == "" {
				continue // access/event rows (obj) and malformed lines are no requests
			}
			// replay EVERY request row into the aggregates so totals and
			// quota windows survive restarts; only the ring buffer view stays
			// capped at maxRecs.
			s.aggregateLocked(r)
			recs = append(recs, r)
		}
		if n := len(recs); n > s.maxRecs {
			recs = append([]Req(nil), recs[n-s.maxRecs:]...)
		}
		s.recs = recs
		s.auditPath = path
		return
	}
	s.mu.Lock()
	s.auditPath = path
	s.mu.Unlock()
}

// Record appends a finished request to the aggregates and ring buffer.
func (s *Stats) Record(r Req) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.aggregateLocked(r)
	s.recs = append(s.recs, r)
	if len(s.recs) > s.maxRecs {
		s.recs = s.recs[len(s.recs)-s.maxRecs:]
	}
	if s.auditPath != "" {
		s.rotateAuditLocked()
		appendAuditLine(s.auditPath, r)
	}
}

// aggregateLocked folds r into every aggregate row and the quota window
// bucket. Caller must hold s.mu.
func (s *Stats) aggregateLocked(r Req) {
	inc(s.byKey, r.Key, r)
	if r.Model != "" {
		inc(s.byModel, r.Model, r)
	}
	inc(s.bySrc, r.Source, r)
	km := s.byKeyModel[r.Key]
	if km == nil {
		km = map[string]*Stat{}
		s.byKeyModel[r.Key] = km
	}
	if r.Model != "" {
		inc(km, r.Model, r)
	}
	ks := s.byKeySrc[r.Key]
	if ks == nil {
		ks = map[string]*Stat{}
		s.byKeySrc[r.Key] = ks
	}
	inc(ks, r.Source, r)
	if r.Status > 0 {
		a := s.byStatus[r.Status]
		if a == nil {
			a = &Stat{}
			s.byStatus[r.Status] = a
		}
		incStatus(a, strconv.Itoa(r.Status), r)
	}
	// window bucket for quota enforcement (per source-model pair, per unix
	// hour). r.Time is unix MILLISECONDS (audit format); hourSec is seconds,
	// so convert before bucketing — otherwise the bucket width would be
	// 3.6s and every WindowTokens cutoff comparison would be off by ~1000x.
	tok := r.Prompt + r.Compl
	if tok > 0 && r.Model != "" {
		key := r.Model
		if r.Source != "" {
			key = r.Source + "::" + r.Model
		}
		h := (r.Time / 1000) / hourSec
		hm := s.modelHour[key]
		if hm == nil {
			hm = map[int64]int64{}
			s.modelHour[key] = hm
		}
		hm[h] += tok
		// retention: 24*40 = 960 hourly buckets ≈ 40 days of history (covers
		// the longest "month" quota window)
		if len(hm) > 24*40 {
			for k := range hm {
				if k < h-24*40 {
					delete(hm, k)
				}
			}
		}
	}
}

// rotateAuditLocked renames the audit file to <path>.<unix>.old once it
// exceeds auditRotateBytes and prunes old files beyond auditKeepOld, keeping
// the newest ones. Caller must hold s.mu.
func (s *Stats) rotateAuditLocked() {
	if s.auditPath == "" || auditRotateBytes <= 0 {
		return
	}
	if fi, err := os.Stat(s.auditPath); err == nil && fi.Size() < auditRotateBytes {
		return
	}
	ts := time.Now().Unix()
	if os.Rename(s.auditPath, fmt.Sprintf("%s.%d.old", s.auditPath, ts)) == nil {
		old, _ := filepath.Glob(s.auditPath + ".*.old")
		sort.Sort(sort.Reverse(sort.StringSlice(old)))
		for i := auditKeepOld; i < len(old); i++ {
			_ = os.Remove(old[i])
		}
	}
}

func appendAuditLine(path string, row interface{}) {
	f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer f.Close()
	if b, err := json.Marshal(row); err == nil {
		_, _ = f.Write(append(b, '\n'))
	}
}

// AppendAudit writes a generic event line (access log entry, login event,
// config change, …) to the same audit file without touching the aggregates.
func (s *Stats) AppendAudit(obj string, data map[string]interface{}) {
	row := map[string]interface{}{"obj": obj, "time": time.Now().UnixMilli()}
	for k, v := range data {
		row[k] = v
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.auditPath == "" {
		return
	}
	s.rotateAuditLocked()
	appendAuditLine(s.auditPath, row)
}

// ModelTokens returns the tokens consumed per model for one gateway key id
// (used for per-model token quota enforcement).
func (s *Stats) ModelTokens(key string) map[string]int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := map[string]int64{}
	for k, v := range s.byKeyModel[key] {
		out[k] = v.Tokens
	}
	return out
}

// KeyTokens returns the total tokens consumed by one gateway key id.
func (s *Stats) KeyTokens(key string) int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	a := s.byKey[key]
	if a == nil {
		return 0
	}
	return a.Tokens
}

// AutoPeriodSeconds maps a quota reset period to its window length in
// seconds. "" → 0 (never resets); "hour" → 1h; "week" → 7d; "month" → 30d;
// "nhour" → Hours (>=1) hours.
func AutoPeriodSeconds(period string, hours int64) int64 {
	switch period {
	case "hour":
		return hourSec
	case "week":
		return 7 * 24 * hourSec
	case "month":
		return 30 * 24 * hourSec
	case "nhour":
		if hours < 1 {
			hours = 1
		}
		return hours * hourSec
	}
	return 0
}

// WindowTokens returns the tokens consumed for one model (optionally pinned
// to a single source) within the window; sec <= 0 means all time. Buckets are
// whole unix hours, so a sliding window overcounts by up to one hour — an
// accepted truncation for quota enforcement.
func (s *Stats) WindowTokens(model, source string, sec int64) int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	key := model
	if source != "" {
		key = source + "::" + model
	}
	now := time.Now().Unix()
	hm := s.modelHour[key]
	if len(hm) == 0 {
		return 0
	}
	var total int64
	if sec <= 0 {
		for _, v := range hm {
			total += v
		}
		return total
	}
	cut := now - sec
	for h, v := range hm {
		if h*hourSec >= cut {
			total += v
		}
	}
	return total
}

func rows(m map[string]*Stat) []StatsRow {
	out := make([]StatsRow, 0, len(m))
	for k, v := range m {
		out = append(out, StatsRow{Name: k, Stat: *v})
	}
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j].Reqs > out[j-1].Reqs; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	return out
}

// StatsRow is one aggregated row for the dashboard.
type StatsRow struct {
	Name string `json:"name"`
	Stat
}

// AuditRecords returns every request row from the audit file plus its rotated
// .old files within the [from,to] unix-millisecond window, optionally for one
// masked key id. Unlike Records it is not bounded by the in-memory ring, so it
// yields a complete history for CSV export regardless of restarts or churn.
func (s *Stats) AuditRecords(from, to int64, key string) []Req {
	s.mu.Lock()
	path := s.auditPath
	s.mu.Unlock()
	if path == "" {
		return nil
	}
	files := []string{path}
	if olds, err := filepath.Glob(path + ".*.old"); err == nil {
		files = append(files, olds...)
	}
	var out []Req
	for _, p := range files {
		f, err := os.Open(p)
		if err != nil {
			continue
		}
		sc := bufio.NewScanner(f)
		sc.Buffer(make([]byte, 64*1024), 16*1024*1024)
		for sc.Scan() {
			var r Req
			if json.Unmarshal(sc.Bytes(), &r) != nil || r.Type == "" {
				continue
			}
			if key != "" && r.Key != key {
				continue
			}
			if from > 0 && r.Time < from {
				continue
			}
			if to > 0 && r.Time > to {
				continue
			}
			out = append(out, r)
		}
		f.Close()
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Time < out[j].Time })
	return out
}

// SourceRecent counts real gateway requests per source within the last
// window (unix seconds). It reads only the in-memory ring, so it is cheap and
// reflects live traffic — used by /api/status so the source status column is
// driven by what actually happens, not just a probe.
func (s *Stats) SourceRecent(windowSec int64) map[string][2]int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	cut := time.Now().Unix() - windowSec
	out := map[string][2]int64{}
	for _, r := range s.recs {
		if r.Time/1000 < cut {
			continue
		}
		if r.Source == "" {
			continue
		}
		v := out[r.Source]
		if r.OK {
			v[0]++
		} else {
			v[1]++
		}
		out[r.Source] = v
	}
	return out
}

// SourceAvg carries per-source performance averages for the status page.
type SourceAvg struct {
	// AvgFirstByteMs is the mean time-to-first-byte over successful
	// requests in the window (0 when no measured samples).
	AvgFirstByteMs int64 `json:"avg_first_byte_ms"`
	// AvgTokPerS is the aggregate completion throughput:
	// sum(completion_tokens) / sum(latency_seconds) (0 when no samples).
	AvgTokPerS int64 `json:"avg_tok_per_s"`
	// Samples is the number of successful requests the averages cover.
	Samples int64 `json:"samples"`
}

// SourceAverages computes TTFB and tokens/s averages per source from the
// in-memory ring within the window (unix seconds). Only successful chat/
// stream rows count; image and failed rows are skipped.
func (s *Stats) SourceAverages(windowSec int64) map[string]SourceAvg {
	s.mu.Lock()
	defer s.mu.Unlock()
	cut := time.Now().Unix() - windowSec
	type acc struct {
		fbSum, latSum, complSum, n int64
	}
	accs := map[string]*acc{}
	for _, r := range s.recs {
		if !r.OK || r.Source == "" || r.Time/1000 < cut {
			continue
		}
		if r.Type != "chat" && r.Type != "stream" {
			continue
		}
		if r.LatMs <= 0 {
			continue
		}
		a := accs[r.Source]
		if a == nil {
			a = &acc{}
			accs[r.Source] = a
		}
		a.latSum += r.LatMs
		a.complSum += r.Compl
		a.fbSum += r.FirstByteMs
		if r.FirstByteMs > 0 {
			a.n++
		}
	}
	out := make(map[string]SourceAvg, len(accs))
	for src, a := range accs {
		avg := SourceAvg{Samples: a.n}
		if a.n > 0 {
			avg.AvgFirstByteMs = a.fbSum / a.n
		}
		if a.latSum > 0 && a.complSum > 0 {
			avg.AvgTokPerS = a.complSum * 1000 / a.latSum
		}
		out[src] = avg
	}
	return out
}

// Snapshot returns the whole dashboard payload; when key != "" the records
// and aggregate views are restricted to that gateway key.
func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
	s.mu.Lock()
	defer s.mu.Unlock()
	if limit <= 0 {
		limit = s.maxRecs
	}
	start := 0
	if len(s.recs) > limit {
		start = len(s.recs) - limit
	}
	recs := s.recs[start:]
	if key != "" {
		filt := recs[:0]
		for _, r := range recs {
			if r.Key == key {
				filt = append(filt, r)
			}
		}
		recs = filt
	}
	var total Stat
	var byModel, byKey, bySrc map[string]*Stat
	if key == "" {
		byKey, byModel, bySrc = s.byKey, s.byModel, s.bySrc
	} else {
		byKey = map[string]*Stat{key: s.byKey[key]}
		byModel = s.byKeyModel[key]
		bySrc = s.byKeySrc[key]
	}
	for _, a := range byModel {
		total.Reqs += a.Reqs
		total.OK += a.OK
		total.Err += a.Err
		total.Tokens += a.Tokens
		total.Prompt += a.Prompt
		total.Compl += a.Compl
		total.LatSum += a.LatSum
		if a.LatMax > total.LatMax {
			total.LatMax = a.LatMax
		}
	}
	bs := make([]agrRow, 0, len(s.byStatus))
	for code := range s.byStatus {
		bs = append(bs, agrRow{Name: strconv.Itoa(code), Stat: *s.byStatus[code]})
	}
	sort.Slice(bs, func(i, j int) bool {
		ci, _ := strconv.Atoi(bs[i].Name)
		cj, _ := strconv.Atoi(bs[j].Name)
		return ci < cj
	})
	return map[string]interface{}{
		"active":    s.active,
		"total":     total,
		"by_key":    rows(byKey),
		"by_model":  rows(byModel),
		"by_source": rows(bySrc),
		"by_status": bs,
		"records":   append([]Req(nil), recs...),
	}
}