ModelRouter/internal/gateway/stats.go

package gateway

import (
	"bufio"
	"encoding/json"
	"os"
	"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"`
	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"`
}

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
	recs       []Req
	maxRecs    int
	auditPath  string
	modelHour  map[string]map[int64]int64 // model -> unix-hour bucket -> tokens
}

const hourSec = 3600

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{},
		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
	}
	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
	}
}

func (s *Stats) LoadAudit(path string) {
	f, err := os.Open(path)
	if err == nil {
		var recs []Req
		sc := bufio.NewScanner(f)
		for sc.Scan() {
			var r Req
			if json.Unmarshal(sc.Bytes(), &r) == nil {
				recs = append(recs, r)
			}
		}
		_ = f.Close()
		if len(recs) > s.maxRecs {
			recs = recs[len(recs)-s.maxRecs:]
		}
		for _, r := range recs {
			s.Record(r)
		}
	}
	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()
	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)
	// window bucket for quota enforcement (per source-model pair, per unix hour)
	tok := r.Prompt + r.Compl
	if tok > 0 && r.Model != "" {
		key := r.Model
		if r.Source != "" {
			key = r.Source + "::" + r.Model
		}
		h := r.Time / hourSec
		hm := s.modelHour[key]
		if hm == nil {
			hm = map[int64]int64{}
			s.modelHour[key] = hm
		}
		hm[h] += tok
		if len(hm) > 24*40 {
			for k := range hm {
				if k < h-24*40 {
					delete(hm, k)
				}
			}
		}
	}
	s.recs = append(s.recs, r)
	if len(s.recs) > s.maxRecs {
		s.recs = s.recs[len(s.recs)-s.maxRecs:]
	}
	if s.auditPath != "" {
		if f, err := os.OpenFile(s.auditPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); err == nil {
			if b, err := json.Marshal(r); err == nil {
				_, _ = f.Write(append(b, '\n'))
			}
			_ = f.Close()
		}
	}
}

// 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 billed for the model within the last `sec`
// seconds (0 = since forever).
// WindowTokens returns the tokens consumed for one model (optionally pinned
// to a single source) within the window; sec <= 0 means all time.
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
}

// Records returns request records filtered by unix-millisecond time range and key.
func (s *Stats) Records(from, to int64, key string) []Req {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]Req, 0, len(s.recs))
	for _, r := range s.recs {
		if key != "" && r.Key != key {
			continue
		}
		if from > 0 && r.Time < from {
			continue
		}
		if to > 0 && r.Time > to {
			continue
		}
		out = append(out, r)
	}
	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
		}
	}
	return map[string]interface{}{
		"active":    s.active,
		"total":     total,
		"by_key":    rows(byKey),
		"by_model":  rows(byModel),
		"by_source": rows(bySrc),
		"records":   append([]Req(nil), recs...),
	}
}