Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 59 additions & 0 deletions internal/api/handlers/management/usage.go
Original file line number Diff line number Diff line change
@@ -1,12 +1,25 @@
package management

import (
"encoding/json"
"net/http"
"time"

"github.com/gin-gonic/gin"
"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.
func (h *Handler) GetUsageStatistics(c *gin.Context) {
var snapshot usage.StatisticsSnapshot
Expand All @@ -18,3 +31,49 @@ func (h *Handler) GetUsageStatistics(c *gin.Context) {
"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,
})
}
2 changes: 2 additions & 0 deletions internal/api/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -496,6 +496,8 @@ func (s *Server) registerManagementRoutes() {
mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware())
{
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.yaml", s.mgmt.GetConfigYAML)
mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML)
Expand Down
2 changes: 1 addition & 1 deletion internal/runtime/executor/usage_helpers.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ type usageReporter struct {
provider string
model string
authID string
authIndex uint64
authIndex string
apiKey string
source string
requestedAt time.Time
Expand Down
125 changes: 124 additions & 1 deletion internal/usage/logger_plugin.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ package usage
import (
"context"
"fmt"
"strings"
"sync"
"sync/atomic"
"time"
Expand Down Expand Up @@ -90,7 +91,7 @@ type modelStats struct {
type RequestDetail struct {
Timestamp time.Time `json:"timestamp"`
Source string `json:"source"`
AuthIndex uint64 `json:"auth_index"`
AuthIndex string `json:"auth_index"`
Tokens TokenStats `json:"tokens"`
Failed bool `json:"failed"`
}
Expand Down Expand Up @@ -281,6 +282,118 @@ func (s *RequestStatistics) Snapshot() StatisticsSnapshot {
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 {
if ctx != nil {
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil {
Expand Down Expand Up @@ -340,6 +453,16 @@ func normaliseDetail(detail coreusage.Detail) TokenStats {
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
Comment on lines +456 to +463

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic in this new normaliseTokenStats function is identical to the logic within the existing normaliseDetail function (lines 447-452). This introduces code duplication. To improve maintainability, you could refactor normaliseDetail to use this new normaliseTokenStats function, which would centralize the token normalization logic.

}

func formatHour(hour int) string {
if hour < 0 {
hour = 0
Expand Down
4 changes: 2 additions & 2 deletions sdk/cliproxy/auth/conductor.go
Original file line number Diff line number Diff line change
Expand Up @@ -203,10 +203,10 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
if auth == nil {
return nil, nil
}
auth.EnsureIndex()
if auth.ID == "" {
auth.ID = uuid.NewString()
}
auth.EnsureIndex()
m.mu.Lock()
m.auths[auth.ID] = auth.Clone()
m.mu.Unlock()
Expand All @@ -221,7 +221,7 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
return nil, nil
}
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.indexAssigned = existing.indexAssigned
}
Expand Down
49 changes: 35 additions & 14 deletions sdk/cliproxy/auth/types.go
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
package auth

import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"

baseauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth"
Expand All @@ -15,8 +16,8 @@ import (
type Auth struct {
// ID uniquely identifies the auth record across restarts.
ID string `json:"id"`
// Index is a monotonically increasing runtime identifier used for diagnostics.
Index uint64 `json:"-"`
// Index is a stable runtime identifier derived from auth metadata (not persisted).
Index string `json:"-"`
// Provider is the upstream provider key (e.g. "gemini", "claude").
Provider string `json:"provider"`
// Prefix optionally namespaces models for routing (e.g., "teamA/gemini-3-pro-preview").
Expand Down Expand Up @@ -94,12 +95,6 @@ type ModelState struct {
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.
func (a *Auth) Clone() *Auth {
if a == nil {
Expand Down Expand Up @@ -128,15 +123,41 @@ func (a *Auth) Clone() *Auth {
return &copyAuth
}

// EnsureIndex returns the global index, assigning one if it was not set yet.
func (a *Auth) EnsureIndex() uint64 {
func stableAuthIndex(seed string) string {
seed = strings.TrimSpace(seed)
if seed == "" {
return ""
}
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 0
return ""
}
if a.indexAssigned {
if a.indexAssigned && 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.indexAssigned = true
return idx
Expand Down
2 changes: 1 addition & 1 deletion sdk/cliproxy/usage/manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ type Record struct {
Model string
APIKey string
AuthID string
AuthIndex uint64
AuthIndex string
Source string
RequestedAt time.Time
Failed bool
Expand Down