From b0dc9df887ef9f8fd9fad5bd9e4ebd639d6de8f3 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Fri, 1 May 2026 23:34:18 +0800 Subject: [PATCH] feat: add API key usage endpoint with provider and key grouping - Implemented `GetAPIKeyUsage` to expose recent request data grouped by provider and API key. - Added supporting function `mergeRecentRequestBuckets` for bucket aggregation. - Registered new endpoint `/v0/management/api-key-usage` in the management API. - Included extensive unit tests for provider and key-based grouping validation. - Updated `formatRecentRequestBucketLabel` to support configurable bucket duration. --- .../api/handlers/management/api_key_usage.go | 86 ++++++++++++++++++ .../handlers/management/api_key_usage_test.go | 87 +++++++++++++++++++ internal/api/server.go | 1 + sdk/cliproxy/auth/types.go | 2 +- 4 files changed, 175 insertions(+), 1 deletion(-) create mode 100644 internal/api/handlers/management/api_key_usage.go create mode 100644 internal/api/handlers/management/api_key_usage_test.go diff --git a/internal/api/handlers/management/api_key_usage.go b/internal/api/handlers/management/api_key_usage.go new file mode 100644 index 00000000..599fbad9 --- /dev/null +++ b/internal/api/handlers/management/api_key_usage.go @@ -0,0 +1,86 @@ +package management + +import ( + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" +) + +func mergeRecentRequestBuckets(dst, src []coreauth.RecentRequestBucket) []coreauth.RecentRequestBucket { + if len(dst) == 0 { + return src + } + if len(src) == 0 { + return dst + } + if len(dst) != len(src) { + n := len(dst) + if len(src) < n { + n = len(src) + } + for i := 0; i < n; i++ { + dst[i].Success += src[i].Success + dst[i].Failed += src[i].Failed + } + return dst + } + for i := range dst { + dst[i].Success += src[i].Success + dst[i].Failed += src[i].Failed + } + return dst +} + +// GetAPIKeyUsage returns recent request buckets for all in-memory api_key auths, +// grouped by provider and keyed by the raw api-key value. +func (h *Handler) GetAPIKeyUsage(c *gin.Context) { + if h == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "handler not initialized"}) + return + } + + h.mu.Lock() + manager := h.authManager + h.mu.Unlock() + if manager == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) + return + } + + now := time.Now() + out := make(map[string]map[string][]coreauth.RecentRequestBucket) + for _, auth := range manager.List() { + if auth == nil { + continue + } + kind, apiKey := auth.AccountInfo() + if !strings.EqualFold(strings.TrimSpace(kind), "api_key") { + continue + } + apiKey = strings.TrimSpace(apiKey) + if apiKey == "" { + continue + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if provider == "" { + provider = "unknown" + } + + recent := auth.RecentRequestsSnapshot(now) + providerBucket, ok := out[provider] + if !ok { + providerBucket = make(map[string][]coreauth.RecentRequestBucket) + out[provider] = providerBucket + } + if existing, exists := providerBucket[apiKey]; exists { + providerBucket[apiKey] = mergeRecentRequestBuckets(existing, recent) + continue + } + providerBucket[apiKey] = recent + } + + c.JSON(http.StatusOK, out) +} diff --git a/internal/api/handlers/management/api_key_usage_test.go b/internal/api/handlers/management/api_key_usage_test.go new file mode 100644 index 00000000..230dca4a --- /dev/null +++ b/internal/api/handlers/management/api_key_usage_test.go @@ -0,0 +1,87 @@ +package management + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" +) + +func sumRecentRequestBuckets(buckets []coreauth.RecentRequestBucket) (int64, int64) { + var success int64 + var failed int64 + for _, bucket := range buckets { + success += bucket.Success + failed += bucket.Failed + } + return success, failed +} + +func TestGetAPIKeyUsage_GroupsByProviderAndAPIKey(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + gin.SetMode(gin.TestMode) + + manager := coreauth.NewManager(nil, nil, nil) + if _, err := manager.Register(context.Background(), &coreauth.Auth{ + ID: "codex-auth", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "codex-key", + }, + }); err != nil { + t.Fatalf("register codex auth: %v", err) + } + if _, err := manager.Register(context.Background(), &coreauth.Auth{ + ID: "claude-auth", + Provider: "claude", + Attributes: map[string]string{ + "api_key": "claude-key", + }, + }); err != nil { + t.Fatalf("register claude auth: %v", err) + } + + manager.MarkResult(context.Background(), coreauth.Result{AuthID: "codex-auth", Provider: "codex", Model: "gpt-5", Success: true}) + manager.MarkResult(context.Background(), coreauth.Result{AuthID: "codex-auth", Provider: "codex", Model: "gpt-5", Success: false}) + manager.MarkResult(context.Background(), coreauth.Result{AuthID: "claude-auth", Provider: "claude", Model: "claude-4", Success: true}) + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager) + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodGet, "/v0/management/api-key-usage", nil) + ginCtx.Request = req + h.GetAPIKeyUsage(ginCtx) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + var payload map[string]map[string][]coreauth.RecentRequestBucket + if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil { + t.Fatalf("decode payload: %v", err) + } + + codexBuckets := payload["codex"]["codex-key"] + if len(codexBuckets) != 20 { + t.Fatalf("codex buckets len = %d, want 20", len(codexBuckets)) + } + codexSuccess, codexFailed := sumRecentRequestBuckets(codexBuckets) + if codexSuccess != 1 || codexFailed != 1 { + t.Fatalf("codex totals = %d/%d, want 1/1", codexSuccess, codexFailed) + } + + claudeBuckets := payload["claude"]["claude-key"] + if len(claudeBuckets) != 20 { + t.Fatalf("claude buckets len = %d, want 20", len(claudeBuckets)) + } + claudeSuccess, claudeFailed := sumRecentRequestBuckets(claudeBuckets) + if claudeSuccess != 1 || claudeFailed != 0 { + t.Fatalf("claude totals = %d/%d, want 1/0", claudeSuccess, claudeFailed) + } +} diff --git a/internal/api/server.go b/internal/api/server.go index 8421357b..4d51460d 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -554,6 +554,7 @@ func (s *Server) registerManagementRoutes() { mgmt.PUT("/api-keys", s.mgmt.PutAPIKeys) mgmt.PATCH("/api-keys", s.mgmt.PatchAPIKeys) mgmt.DELETE("/api-keys", s.mgmt.DeleteAPIKeys) + mgmt.GET("/api-key-usage", s.mgmt.GetAPIKeyUsage) mgmt.GET("/gemini-api-key", s.mgmt.GetGeminiKeys) mgmt.PUT("/gemini-api-key", s.mgmt.PutGeminiKeys) diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go index 93dd3881..4a394ad4 100644 --- a/sdk/cliproxy/auth/types.go +++ b/sdk/cliproxy/auth/types.go @@ -164,7 +164,7 @@ func recentRequestBucketIndex(bucketID int64) int { func formatRecentRequestBucketLabel(bucketID int64) string { start := time.Unix(bucketID*recentRequestBucketSeconds, 0).In(time.Local) - end := start.Add(10 * time.Minute) + end := start.Add(time.Duration(recentRequestBucketSeconds) * time.Second) return start.Format("15:04") + "-" + end.Format("15:04") }