// Package usage provides usage tracking and logging functionality for the CLI Proxy API server.// It includes plugins for monitoring API usage, token consumption, and other metrics// to help with observability and billing purposes.package usage import ( "context" "fmt" "sync" "sync/atomic" "time" "github.com/gin-gonic/gin" coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage") var statisticsEnabled atomic.Bool func init() { statisticsEnabled.Store(true) coreusage.RegisterPlugin(NewLoggerPlugin())} // LoggerPlugin collects in-memory request statistics for usage analysis.// It implements coreusage.Plugin to receive usage records emitted by the runtime.type LoggerPlugin struct { stats *RequestStatistics} // NewLoggerPlugin constructs a new logger plugin instance.//// Returns:// - *LoggerPlugin: A new logger plugin instance wired to the shared statistics store.func NewLoggerPlugin() *LoggerPlugin { return &LoggerPlugin{stats: defaultRequestStatistics} } // HandleUsage implements coreusage.Plugin.// It updates the in-memory statistics store whenever a usage record is received.//// Parameters:// - ctx: The context for the usage record// - record: The usage record to aggregatefunc (p *LoggerPlugin) HandleUsage(ctx context.Context, record coreusage.Record) { if !statisticsEnabled.Load() { return } if p == nil || p.stats == nil { return } p.stats.Record(ctx, record)} // SetStatisticsEnabled toggles whether in-memory statistics are recorded.func SetStatisticsEnabled(enabled bool) { statisticsEnabled.Store(enabled) } // StatisticsEnabled reports the current recording state.func StatisticsEnabled() bool { return statisticsEnabled.Load() } // RequestStatistics maintains aggregated request metrics in memory.type RequestStatistics struct { mu sync.RWMutex totalRequests int64 successCount int64 failureCount int64 totalTokens int64 apis map[string]*apiStats requestsByDay map[string]int64 requestsByHour map[int]int64 tokensByDay map[string]int64 tokensByHour map[int]int64} // apiStats holds aggregated metrics for a single API key.type apiStats struct { TotalRequests int64 TotalTokens int64 Models map[string]*modelStats} // modelStats holds aggregated metrics for a specific model within an API.type modelStats struct { TotalRequests int64 TotalTokens int64 Details []RequestDetail} // RequestDetail stores the timestamp and token usage for a single request.type RequestDetail struct { Timestamp time.Time `json:"timestamp"` Source string `json:"source"` Tokens TokenStats `json:"tokens"`} // TokenStats captures the token usage breakdown for a request.type TokenStats struct { InputTokens int64 `json:"input_tokens"` OutputTokens int64 `json:"output_tokens"` ReasoningTokens int64 `json:"reasoning_tokens"` CachedTokens int64 `json:"cached_tokens"` TotalTokens int64 `json:"total_tokens"`} // StatisticsSnapshot represents an immutable view of the aggregated metrics.type StatisticsSnapshot struct { TotalRequests int64 `json:"total_requests"` SuccessCount int64 `json:"success_count"` FailureCount int64 `json:"failure_count"` TotalTokens int64 `json:"total_tokens"` APIs map[string]APISnapshot `json:"apis"` RequestsByDay map[string]int64 `json:"requests_by_day"` RequestsByHour map[string]int64 `json:"requests_by_hour"` TokensByDay map[string]int64 `json:"tokens_by_day"` TokensByHour map[string]int64 `json:"tokens_by_hour"`} // APISnapshot summarises metrics for a single API key.type APISnapshot struct { TotalRequests int64 `json:"total_requests"` TotalTokens int64 `json:"total_tokens"` Models map[string]ModelSnapshot `json:"models"`} // ModelSnapshot summarises metrics for a specific model.type ModelSnapshot struct { TotalRequests int64 `json:"total_requests"` TotalTokens int64 `json:"total_tokens"` Details []RequestDetail `json:"details"`} var defaultRequestStatistics = NewRequestStatistics() // GetRequestStatistics returns the shared statistics store.func GetRequestStatistics() *RequestStatistics { return defaultRequestStatistics } // NewRequestStatistics constructs an empty statistics store.func NewRequestStatistics() *RequestStatistics { return &RequestStatistics{ apis: make(map[string]*apiStats), requestsByDay: make(map[string]int64), requestsByHour: make(map[int]int64), tokensByDay: make(map[string]int64), tokensByHour: make(map[int]int64), }} // Record ingests a new usage record and updates the aggregates.func (s *RequestStatistics) Record(ctx context.Context, record coreusage.Record) { if s == nil { return } if !statisticsEnabled.Load() { return } timestamp := record.RequestedAt if timestamp.IsZero() { timestamp = time.Now() } detail := normaliseDetail(record.Detail) totalTokens := detail.TotalTokens statsKey := record.APIKey if statsKey == "" { statsKey = resolveAPIIdentifier(ctx, record) } success := resolveSuccess(ctx) modelName := record.Model if modelName == "" { modelName = "unknown" } dayKey := timestamp.Format("2006-01-02") hourKey := timestamp.Hour() s.mu.Lock() defer s.mu.Unlock() s.totalRequests++ if success { s.successCount++ } else { s.failureCount++ } s.totalTokens += totalTokens stats, ok := s.apis[statsKey] if !ok { stats = &apiStats{Models: make(map[string]*modelStats)} s.apis[statsKey] = stats } s.updateAPIStats(stats, modelName, RequestDetail{ Timestamp: timestamp, Source: record.Source, Tokens: detail, }) s.requestsByDay[dayKey]++ s.requestsByHour[hourKey]++ s.tokensByDay[dayKey] += totalTokens s.tokensByHour[hourKey] += totalTokens} func (s *RequestStatistics) updateAPIStats(stats *apiStats, model string, detail RequestDetail) { stats.TotalRequests++ stats.TotalTokens += detail.Tokens.TotalTokens modelStatsValue, ok := stats.Models[model] if !ok { modelStatsValue = &modelStats{} stats.Models[model] = modelStatsValue } modelStatsValue.TotalRequests++ modelStatsValue.TotalTokens += detail.Tokens.TotalTokens modelStatsValue.Details = append(modelStatsValue.Details, detail)} // Snapshot returns a copy of the aggregated metrics for external consumption.func (s *RequestStatistics) Snapshot() StatisticsSnapshot { result := StatisticsSnapshot{} if s == nil { return result } s.mu.RLock() defer s.mu.RUnlock() result.TotalRequests = s.totalRequests result.SuccessCount = s.successCount result.FailureCount = s.failureCount result.TotalTokens = s.totalTokens result.APIs = make(map[string]APISnapshot, len(s.apis)) for apiName, stats := range s.apis { apiSnapshot := APISnapshot{ TotalRequests: stats.TotalRequests, TotalTokens: stats.TotalTokens, Models: make(map[string]ModelSnapshot, len(stats.Models)), } for modelName, modelStatsValue := range stats.Models { requestDetails := make([]RequestDetail, len(modelStatsValue.Details)) copy(requestDetails, modelStatsValue.Details) apiSnapshot.Models[modelName] = ModelSnapshot{ TotalRequests: modelStatsValue.TotalRequests, TotalTokens: modelStatsValue.TotalTokens, Details: requestDetails, } } result.APIs[apiName] = apiSnapshot } result.RequestsByDay = make(map[string]int64, len(s.requestsByDay)) for k, v := range s.requestsByDay { result.RequestsByDay[k] = v } result.RequestsByHour = make(map[string]int64, len(s.requestsByHour)) for hour, v := range s.requestsByHour { key := formatHour(hour) result.RequestsByHour[key] = v } result.TokensByDay = make(map[string]int64, len(s.tokensByDay)) for k, v := range s.tokensByDay { result.TokensByDay[k] = v } result.TokensByHour = make(map[string]int64, len(s.tokensByHour)) for hour, v := range s.tokensByHour { key := formatHour(hour) result.TokensByHour[key] = v } return result} func resolveAPIIdentifier(ctx context.Context, record coreusage.Record) string { if ctx != nil { if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil { path := ginCtx.FullPath() if path == "" && ginCtx.Request != nil { path = ginCtx.Request.URL.Path } method := "" if ginCtx.Request != nil { method = ginCtx.Request.Method } if path != "" { if method != "" { return method + " " + path } return path } } } if record.Provider != "" { return record.Provider } return "unknown"} func resolveSuccess(ctx context.Context) bool { if ctx == nil { return true } ginCtx, ok := ctx.Value("gin").(*gin.Context) if !ok || ginCtx == nil { return true } status := ginCtx.Writer.Status() if status == 0 { return true } return status < httpStatusBadRequest} const httpStatusBadRequest = 400 func normaliseDetail(detail coreusage.Detail) TokenStats { tokens := TokenStats{ InputTokens: detail.InputTokens, OutputTokens: detail.OutputTokens, ReasoningTokens: detail.ReasoningTokens, CachedTokens: detail.CachedTokens, TotalTokens: detail.TotalTokens, } if tokens.TotalTokens == 0 { tokens.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens } if tokens.TotalTokens == 0 { tokens.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens + detail.CachedTokens } return tokens} func formatHour(hour int) string { if hour < 0 { hour = 0 } hour = hour % 24 return fmt.Sprintf("%02d", hour)}
// Package usage provides usage tracking and logging functionality for the CLI Proxy API server.
// It includes plugins for monitoring API usage, token consumption, and other metrics
// to help with observability and billing purposes.
package usage
import (
"context"
"fmt"
"sync"
"sync/atomic"
"time"
"github.com/gin-gonic/gin"
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
)
var statisticsEnabled atomic.Bool
func init() {
statisticsEnabled.Store(true)
coreusage.RegisterPlugin(NewLoggerPlugin())
}
// LoggerPlugin collects in-memory request statistics for usage analysis.
// It implements coreusage.Plugin to receive usage records emitted by the runtime.
type LoggerPlugin struct {
stats *RequestStatistics
}
// NewLoggerPlugin constructs a new logger plugin instance.
//
// Returns:
// - *LoggerPlugin: A new logger plugin instance wired to the shared statistics store.
func NewLoggerPlugin() *LoggerPlugin { return &LoggerPlugin{stats: defaultRequestStatistics} }
// HandleUsage implements coreusage.Plugin.
// It updates the in-memory statistics store whenever a usage record is received.
//
// Parameters:
// - ctx: The context for the usage record
// - record: The usage record to aggregate
func (p *LoggerPlugin) HandleUsage(ctx context.Context, record coreusage.Record) {
if !statisticsEnabled.Load() {
return
}
if p == nil || p.stats == nil {
return
}
p.stats.Record(ctx, record)
}
// SetStatisticsEnabled toggles whether in-memory statistics are recorded.
func SetStatisticsEnabled(enabled bool) { statisticsEnabled.Store(enabled) }
// StatisticsEnabled reports the current recording state.
func StatisticsEnabled() bool { return statisticsEnabled.Load() }
// RequestStatistics maintains aggregated request metrics in memory.
type RequestStatistics struct {
mu sync.RWMutex
totalRequests int64
successCount int64
failureCount int64
totalTokens int64
apis map[string]*apiStats
requestsByDay map[string]int64
requestsByHour map[int]int64
tokensByDay map[string]int64
tokensByHour map[int]int64
}
// apiStats holds aggregated metrics for a single API key.
type apiStats struct {
TotalRequests int64
TotalTokens int64
Models map[string]*modelStats
}
// modelStats holds aggregated metrics for a specific model within an API.
type modelStats struct {
TotalRequests int64
TotalTokens int64
Details []RequestDetail
}
// RequestDetail stores the timestamp and token usage for a single request.
type RequestDetail struct {
Timestamp time.Time `json:"timestamp"`
Source string `json:"source"`
Tokens TokenStats `json:"tokens"`
}
// TokenStats captures the token usage breakdown for a request.
type TokenStats struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
ReasoningTokens int64 `json:"reasoning_tokens"`
CachedTokens int64 `json:"cached_tokens"`
TotalTokens int64 `json:"total_tokens"`
}
// StatisticsSnapshot represents an immutable view of the aggregated metrics.
type StatisticsSnapshot struct {
TotalRequests int64 `json:"total_requests"`
SuccessCount int64 `json:"success_count"`
FailureCount int64 `json:"failure_count"`
TotalTokens int64 `json:"total_tokens"`
APIs map[string]APISnapshot `json:"apis"`
RequestsByDay map[string]int64 `json:"requests_by_day"`
RequestsByHour map[string]int64 `json:"requests_by_hour"`
TokensByDay map[string]int64 `json:"tokens_by_day"`
TokensByHour map[string]int64 `json:"tokens_by_hour"`
}
// APISnapshot summarises metrics for a single API key.
type APISnapshot struct {
TotalRequests int64 `json:"total_requests"`
TotalTokens int64 `json:"total_tokens"`
Models map[string]ModelSnapshot `json:"models"`
}
// ModelSnapshot summarises metrics for a specific model.
type ModelSnapshot struct {
TotalRequests int64 `json:"total_requests"`
TotalTokens int64 `json:"total_tokens"`
Details []RequestDetail `json:"details"`
}
var defaultRequestStatistics = NewRequestStatistics()
// GetRequestStatistics returns the shared statistics store.
func GetRequestStatistics() *RequestStatistics { return defaultRequestStatistics }
// NewRequestStatistics constructs an empty statistics store.
func NewRequestStatistics() *RequestStatistics {
return &RequestStatistics{
apis: make(map[string]*apiStats),
requestsByDay: make(map[string]int64),
requestsByHour: make(map[int]int64),
tokensByDay: make(map[string]int64),
tokensByHour: make(map[int]int64),
}
}
// Record ingests a new usage record and updates the aggregates.
func (s *RequestStatistics) Record(ctx context.Context, record coreusage.Record) {
if s == nil {
return
}
if !statisticsEnabled.Load() {
return
}
timestamp := record.RequestedAt
if timestamp.IsZero() {
timestamp = time.Now()
}
detail := normaliseDetail(record.Detail)
totalTokens := detail.TotalTokens
statsKey := record.APIKey
if statsKey == "" {
statsKey = resolveAPIIdentifier(ctx, record)
}
success := resolveSuccess(ctx)
modelName := record.Model
if modelName == "" {
modelName = "unknown"
}
dayKey := timestamp.Format("2006-01-02")
hourKey := timestamp.Hour()
s.mu.Lock()
defer s.mu.Unlock()
s.totalRequests++
if success {
s.successCount++
} else {
s.failureCount++
}
s.totalTokens += totalTokens
stats, ok := s.apis[statsKey]
if !ok {
stats = &apiStats{Models: make(map[string]*modelStats)}
s.apis[statsKey] = stats
}
s.updateAPIStats(stats, modelName, RequestDetail{
Timestamp: timestamp,
Source: record.Source,
Tokens: detail,
})
s.requestsByDay[dayKey]++
s.requestsByHour[hourKey]++
s.tokensByDay[dayKey] += totalTokens
s.tokensByHour[hourKey] += totalTokens
}
func (s *RequestStatistics) updateAPIStats(stats *apiStats, model string, detail RequestDetail) {
stats.TotalRequests++
stats.TotalTokens += detail.Tokens.TotalTokens
modelStatsValue, ok := stats.Models[model]
if !ok {
modelStatsValue = &modelStats{}
stats.Models[model] = modelStatsValue
}
modelStatsValue.TotalRequests++
modelStatsValue.TotalTokens += detail.Tokens.TotalTokens
modelStatsValue.Details = append(modelStatsValue.Details, detail)
}
// Snapshot returns a copy of the aggregated metrics for external consumption.
func (s *RequestStatistics) Snapshot() StatisticsSnapshot {
result := StatisticsSnapshot{}
if s == nil {
return result
}
s.mu.RLock()
defer s.mu.RUnlock()
result.TotalRequests = s.totalRequests
result.SuccessCount = s.successCount
result.FailureCount = s.failureCount
result.TotalTokens = s.totalTokens
result.APIs = make(map[string]APISnapshot, len(s.apis))
for apiName, stats := range s.apis {
apiSnapshot := APISnapshot{
TotalRequests: stats.TotalRequests,
TotalTokens: stats.TotalTokens,
Models: make(map[string]ModelSnapshot, len(stats.Models)),
}
for modelName, modelStatsValue := range stats.Models {
requestDetails := make([]RequestDetail, len(modelStatsValue.Details))
copy(requestDetails, modelStatsValue.Details)
apiSnapshot.Models[modelName] = ModelSnapshot{
TotalRequests: modelStatsValue.TotalRequests,
TotalTokens: modelStatsValue.TotalTokens,
Details: requestDetails,
}
}
result.APIs[apiName] = apiSnapshot
}
result.RequestsByDay = make(map[string]int64, len(s.requestsByDay))
for k, v := range s.requestsByDay {
result.RequestsByDay[k] = v
}
result.RequestsByHour = make(map[string]int64, len(s.requestsByHour))
for hour, v := range s.requestsByHour {
key := formatHour(hour)
result.RequestsByHour[key] = v
}
result.TokensByDay = make(map[string]int64, len(s.tokensByDay))
for k, v := range s.tokensByDay {
result.TokensByDay[k] = v
}
result.TokensByHour = make(map[string]int64, len(s.tokensByHour))
for hour, v := range s.tokensByHour {
key := formatHour(hour)
result.TokensByHour[key] = v
}
return result
}
func resolveAPIIdentifier(ctx context.Context, record coreusage.Record) string {
if ctx != nil {
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil {
path := ginCtx.FullPath()
if path == "" && ginCtx.Request != nil {
path = ginCtx.Request.URL.Path
}
method := ""
if ginCtx.Request != nil {
method = ginCtx.Request.Method
}
if path != "" {
if method != "" {
return method + " " + path
}
return path
}
}
}
if record.Provider != "" {
return record.Provider
}
return "unknown"
}
func resolveSuccess(ctx context.Context) bool {
if ctx == nil {
return true
}
ginCtx, ok := ctx.Value("gin").(*gin.Context)
if !ok || ginCtx == nil {
return true
}
status := ginCtx.Writer.Status()
if status == 0 {
return true
}
return status < httpStatusBadRequest
}
const httpStatusBadRequest = 400
func normaliseDetail(detail coreusage.Detail) TokenStats {
tokens := TokenStats{
InputTokens: detail.InputTokens,
OutputTokens: detail.OutputTokens,
ReasoningTokens: detail.ReasoningTokens,
CachedTokens: detail.CachedTokens,
TotalTokens: detail.TotalTokens,
}
if tokens.TotalTokens == 0 {
tokens.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
if tokens.TotalTokens == 0 {
tokens.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens + detail.CachedTokens
}
return tokens
}
func formatHour(hour int) string {
if hour < 0 {
hour = 0
}
hour = hour % 24
return fmt.Sprintf("%02d", hour)
}