feat(usage): add import/export functionality for usage statistics and enhance deduplication logic

This commit is contained in:
Luis Pater
2025-12-26 11:49:51 +08:00
parent 33e53a2a56
commit 3ce0d76aa4
7 changed files with 225 additions and 20 deletions

View File

@@ -1,12 +1,25 @@
package management package management
import ( import (
"encoding/json"
"net/http" "net/http"
"time"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage" "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
) )
type usageExportPayload struct {
Version int `json:"version"`
ExportedAt time.Time `json:"exported_at"`
Usage usage.StatisticsSnapshot `json:"usage"`
}
type usageImportPayload struct {
Version int `json:"version"`
Usage usage.StatisticsSnapshot `json:"usage"`
}
// GetUsageStatistics returns the in-memory request statistics snapshot. // GetUsageStatistics returns the in-memory request statistics snapshot.
func (h *Handler) GetUsageStatistics(c *gin.Context) { func (h *Handler) GetUsageStatistics(c *gin.Context) {
var snapshot usage.StatisticsSnapshot var snapshot usage.StatisticsSnapshot
@@ -18,3 +31,49 @@ func (h *Handler) GetUsageStatistics(c *gin.Context) {
"failed_requests": snapshot.FailureCount, "failed_requests": snapshot.FailureCount,
}) })
} }
// ExportUsageStatistics returns a complete usage snapshot for backup/migration.
func (h *Handler) ExportUsageStatistics(c *gin.Context) {
var snapshot usage.StatisticsSnapshot
if h != nil && h.usageStats != nil {
snapshot = h.usageStats.Snapshot()
}
c.JSON(http.StatusOK, usageExportPayload{
Version: 1,
ExportedAt: time.Now().UTC(),
Usage: snapshot,
})
}
// ImportUsageStatistics merges a previously exported usage snapshot into memory.
func (h *Handler) ImportUsageStatistics(c *gin.Context) {
if h == nil || h.usageStats == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "usage statistics unavailable"})
return
}
data, err := c.GetRawData()
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "failed to read request body"})
return
}
var payload usageImportPayload
if err := json.Unmarshal(data, &payload); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json"})
return
}
if payload.Version != 0 && payload.Version != 1 {
c.JSON(http.StatusBadRequest, gin.H{"error": "unsupported version"})
return
}
result := h.usageStats.MergeSnapshot(payload.Usage)
snapshot := h.usageStats.Snapshot()
c.JSON(http.StatusOK, gin.H{
"added": result.Added,
"skipped": result.Skipped,
"total_requests": snapshot.TotalRequests,
"failed_requests": snapshot.FailureCount,
})
}

View File

@@ -476,6 +476,8 @@ func (s *Server) registerManagementRoutes() {
mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware()) mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware())
{ {
mgmt.GET("/usage", s.mgmt.GetUsageStatistics) mgmt.GET("/usage", s.mgmt.GetUsageStatistics)
mgmt.GET("/usage/export", s.mgmt.ExportUsageStatistics)
mgmt.POST("/usage/import", s.mgmt.ImportUsageStatistics)
mgmt.GET("/config", s.mgmt.GetConfig) mgmt.GET("/config", s.mgmt.GetConfig)
mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML) mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML)
mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML) mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML)

View File

@@ -19,7 +19,7 @@ type usageReporter struct {
provider string provider string
model string model string
authID string authID string
authIndex uint64 authIndex string
apiKey string apiKey string
source string source string
requestedAt time.Time requestedAt time.Time

View File

@@ -6,6 +6,7 @@ package usage
import ( import (
"context" "context"
"fmt" "fmt"
"strings"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "time"
@@ -90,7 +91,7 @@ type modelStats struct {
type RequestDetail struct { type RequestDetail struct {
Timestamp time.Time `json:"timestamp"` Timestamp time.Time `json:"timestamp"`
Source string `json:"source"` Source string `json:"source"`
AuthIndex uint64 `json:"auth_index"` AuthIndex string `json:"auth_index"`
Tokens TokenStats `json:"tokens"` Tokens TokenStats `json:"tokens"`
Failed bool `json:"failed"` Failed bool `json:"failed"`
} }
@@ -281,6 +282,118 @@ func (s *RequestStatistics) Snapshot() StatisticsSnapshot {
return result return result
} }
type MergeResult struct {
Added int64 `json:"added"`
Skipped int64 `json:"skipped"`
}
// MergeSnapshot merges an exported statistics snapshot into the current store.
// Existing data is preserved and duplicate request details are skipped.
func (s *RequestStatistics) MergeSnapshot(snapshot StatisticsSnapshot) MergeResult {
result := MergeResult{}
if s == nil {
return result
}
s.mu.Lock()
defer s.mu.Unlock()
seen := make(map[string]struct{})
for apiName, stats := range s.apis {
if stats == nil {
continue
}
for modelName, modelStatsValue := range stats.Models {
if modelStatsValue == nil {
continue
}
for _, detail := range modelStatsValue.Details {
seen[dedupKey(apiName, modelName, detail)] = struct{}{}
}
}
}
for apiName, apiSnapshot := range snapshot.APIs {
apiName = strings.TrimSpace(apiName)
if apiName == "" {
continue
}
stats, ok := s.apis[apiName]
if !ok || stats == nil {
stats = &apiStats{Models: make(map[string]*modelStats)}
s.apis[apiName] = stats
} else if stats.Models == nil {
stats.Models = make(map[string]*modelStats)
}
for modelName, modelSnapshot := range apiSnapshot.Models {
modelName = strings.TrimSpace(modelName)
if modelName == "" {
modelName = "unknown"
}
for _, detail := range modelSnapshot.Details {
detail.Tokens = normaliseTokenStats(detail.Tokens)
if detail.Timestamp.IsZero() {
detail.Timestamp = time.Now()
}
key := dedupKey(apiName, modelName, detail)
if _, exists := seen[key]; exists {
result.Skipped++
continue
}
seen[key] = struct{}{}
s.recordImported(apiName, modelName, stats, detail)
result.Added++
}
}
}
return result
}
func (s *RequestStatistics) recordImported(apiName, modelName string, stats *apiStats, detail RequestDetail) {
totalTokens := detail.Tokens.TotalTokens
if totalTokens < 0 {
totalTokens = 0
}
s.totalRequests++
if detail.Failed {
s.failureCount++
} else {
s.successCount++
}
s.totalTokens += totalTokens
s.updateAPIStats(stats, modelName, detail)
dayKey := detail.Timestamp.Format("2006-01-02")
hourKey := detail.Timestamp.Hour()
s.requestsByDay[dayKey]++
s.requestsByHour[hourKey]++
s.tokensByDay[dayKey] += totalTokens
s.tokensByHour[hourKey] += totalTokens
}
func dedupKey(apiName, modelName string, detail RequestDetail) string {
timestamp := detail.Timestamp.UTC().Format(time.RFC3339Nano)
tokens := normaliseTokenStats(detail.Tokens)
return fmt.Sprintf(
"%s|%s|%s|%s|%s|%t|%d|%d|%d|%d|%d",
apiName,
modelName,
timestamp,
detail.Source,
detail.AuthIndex,
detail.Failed,
tokens.InputTokens,
tokens.OutputTokens,
tokens.ReasoningTokens,
tokens.CachedTokens,
tokens.TotalTokens,
)
}
func resolveAPIIdentifier(ctx context.Context, record coreusage.Record) string { func resolveAPIIdentifier(ctx context.Context, record coreusage.Record) string {
if ctx != nil { if ctx != nil {
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil { if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil {
@@ -340,6 +453,16 @@ func normaliseDetail(detail coreusage.Detail) TokenStats {
return tokens return tokens
} }
func normaliseTokenStats(tokens TokenStats) TokenStats {
if tokens.TotalTokens == 0 {
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
}
if tokens.TotalTokens == 0 {
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens + tokens.CachedTokens
}
return tokens
}
func formatHour(hour int) string { func formatHour(hour int) string {
if hour < 0 { if hour < 0 {
hour = 0 hour = 0

View File

@@ -203,10 +203,10 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
if auth == nil { if auth == nil {
return nil, nil return nil, nil
} }
auth.EnsureIndex()
if auth.ID == "" { if auth.ID == "" {
auth.ID = uuid.NewString() auth.ID = uuid.NewString()
} }
auth.EnsureIndex()
m.mu.Lock() m.mu.Lock()
m.auths[auth.ID] = auth.Clone() m.auths[auth.ID] = auth.Clone()
m.mu.Unlock() m.mu.Unlock()
@@ -221,7 +221,7 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
return nil, nil return nil, nil
} }
m.mu.Lock() m.mu.Lock()
if existing, ok := m.auths[auth.ID]; ok && existing != nil && !auth.indexAssigned && auth.Index == 0 { if existing, ok := m.auths[auth.ID]; ok && existing != nil && !auth.indexAssigned && auth.Index == "" {
auth.Index = existing.Index auth.Index = existing.Index
auth.indexAssigned = existing.indexAssigned auth.indexAssigned = existing.indexAssigned
} }

View File

@@ -1,11 +1,12 @@
package auth package auth
import ( import (
"crypto/sha256"
"encoding/hex"
"encoding/json" "encoding/json"
"strconv" "strconv"
"strings" "strings"
"sync" "sync"
"sync/atomic"
"time" "time"
baseauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth" baseauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth"
@@ -15,8 +16,8 @@ import (
type Auth struct { type Auth struct {
// ID uniquely identifies the auth record across restarts. // ID uniquely identifies the auth record across restarts.
ID string `json:"id"` ID string `json:"id"`
// Index is a monotonically increasing runtime identifier used for diagnostics. // Index is a stable runtime identifier derived from auth metadata (not persisted).
Index uint64 `json:"-"` Index string `json:"-"`
// Provider is the upstream provider key (e.g. "gemini", "claude"). // Provider is the upstream provider key (e.g. "gemini", "claude").
Provider string `json:"provider"` Provider string `json:"provider"`
// Prefix optionally namespaces models for routing (e.g., "teamA/gemini-3-pro-preview"). // Prefix optionally namespaces models for routing (e.g., "teamA/gemini-3-pro-preview").
@@ -94,12 +95,6 @@ type ModelState struct {
UpdatedAt time.Time `json:"updated_at"` UpdatedAt time.Time `json:"updated_at"`
} }
var authIndexCounter atomic.Uint64
func nextAuthIndex() uint64 {
return authIndexCounter.Add(1) - 1
}
// Clone shallow copies the Auth structure, duplicating maps to avoid accidental mutation. // Clone shallow copies the Auth structure, duplicating maps to avoid accidental mutation.
func (a *Auth) Clone() *Auth { func (a *Auth) Clone() *Auth {
if a == nil { if a == nil {
@@ -128,15 +123,41 @@ func (a *Auth) Clone() *Auth {
return &copyAuth return &copyAuth
} }
// EnsureIndex returns the global index, assigning one if it was not set yet. func stableAuthIndex(seed string) string {
func (a *Auth) EnsureIndex() uint64 { seed = strings.TrimSpace(seed)
if a == nil { if seed == "" {
return 0 return ""
} }
if a.indexAssigned { sum := sha256.Sum256([]byte(seed))
return hex.EncodeToString(sum[:8])
}
// EnsureIndex returns a stable index derived from the auth file name or API key.
func (a *Auth) EnsureIndex() string {
if a == nil {
return ""
}
if a.indexAssigned && a.Index != "" {
return a.Index return a.Index
} }
idx := nextAuthIndex()
seed := strings.TrimSpace(a.FileName)
if seed != "" {
seed = "file:" + seed
} else if a.Attributes != nil {
if apiKey := strings.TrimSpace(a.Attributes["api_key"]); apiKey != "" {
seed = "api_key:" + apiKey
}
}
if seed == "" {
if id := strings.TrimSpace(a.ID); id != "" {
seed = "id:" + id
} else {
return ""
}
}
idx := stableAuthIndex(seed)
a.Index = idx a.Index = idx
a.indexAssigned = true a.indexAssigned = true
return idx return idx

View File

@@ -14,7 +14,7 @@ type Record struct {
Model string Model string
APIKey string APIKey string
AuthID string AuthID string
AuthIndex uint64 AuthIndex string
Source string Source string
RequestedAt time.Time RequestedAt time.Time
Failed bool Failed bool