From 8c2f1a80d39e542d3d85d569d035d7df8d5c39f6 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 2 May 2026 02:20:49 +0800 Subject: [PATCH] feat: enhance API key usage grouping with base URL inclusion - Updated `GetAPIKeyUsage` to group API key usage by "base_url|api_key" composite keys. - Adjusted logic to handle `base_url` extraction from auth attributes. - Revised unit tests to validate "base_url|api_key" grouping behavior. --- .../api/handlers/management/api_key_usage.go | 16 ++++++++++++---- .../handlers/management/api_key_usage_test.go | 10 ++++++---- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/internal/api/handlers/management/api_key_usage.go b/internal/api/handlers/management/api_key_usage.go index 599fbad98..76b32bbb6 100644 --- a/internal/api/handlers/management/api_key_usage.go +++ b/internal/api/handlers/management/api_key_usage.go @@ -35,7 +35,7 @@ func mergeRecentRequestBuckets(dst, src []coreauth.RecentRequestBucket) []coreau } // GetAPIKeyUsage returns recent request buckets for all in-memory api_key auths, -// grouped by provider and keyed by the raw api-key value. +// grouped by provider and keyed by "base_url|api_key". func (h *Handler) GetAPIKeyUsage(c *gin.Context) { if h == nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "handler not initialized"}) @@ -64,6 +64,14 @@ func (h *Handler) GetAPIKeyUsage(c *gin.Context) { if apiKey == "" { continue } + baseURL := "" + if auth.Attributes != nil { + baseURL = strings.TrimSpace(auth.Attributes["base_url"]) + if baseURL == "" { + baseURL = strings.TrimSpace(auth.Attributes["base-url"]) + } + } + compositeKey := baseURL + "|" + apiKey provider := strings.ToLower(strings.TrimSpace(auth.Provider)) if provider == "" { provider = "unknown" @@ -75,11 +83,11 @@ func (h *Handler) GetAPIKeyUsage(c *gin.Context) { providerBucket = make(map[string][]coreauth.RecentRequestBucket) out[provider] = providerBucket } - if existing, exists := providerBucket[apiKey]; exists { - providerBucket[apiKey] = mergeRecentRequestBuckets(existing, recent) + if existing, exists := providerBucket[compositeKey]; exists { + providerBucket[compositeKey] = mergeRecentRequestBuckets(existing, recent) continue } - providerBucket[apiKey] = recent + providerBucket[compositeKey] = 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 index 230dca4a6..56617161c 100644 --- a/internal/api/handlers/management/api_key_usage_test.go +++ b/internal/api/handlers/management/api_key_usage_test.go @@ -31,7 +31,8 @@ func TestGetAPIKeyUsage_GroupsByProviderAndAPIKey(t *testing.T) { ID: "codex-auth", Provider: "codex", Attributes: map[string]string{ - "api_key": "codex-key", + "api_key": "codex-key", + "base_url": "https://codex.example.com", }, }); err != nil { t.Fatalf("register codex auth: %v", err) @@ -40,7 +41,8 @@ func TestGetAPIKeyUsage_GroupsByProviderAndAPIKey(t *testing.T) { ID: "claude-auth", Provider: "claude", Attributes: map[string]string{ - "api_key": "claude-key", + "api_key": "claude-key", + "base_url": "https://claude.example.com", }, }); err != nil { t.Fatalf("register claude auth: %v", err) @@ -67,7 +69,7 @@ func TestGetAPIKeyUsage_GroupsByProviderAndAPIKey(t *testing.T) { t.Fatalf("decode payload: %v", err) } - codexBuckets := payload["codex"]["codex-key"] + codexBuckets := payload["codex"]["https://codex.example.com|codex-key"] if len(codexBuckets) != 20 { t.Fatalf("codex buckets len = %d, want 20", len(codexBuckets)) } @@ -76,7 +78,7 @@ func TestGetAPIKeyUsage_GroupsByProviderAndAPIKey(t *testing.T) { t.Fatalf("codex totals = %d/%d, want 1/1", codexSuccess, codexFailed) } - claudeBuckets := payload["claude"]["claude-key"] + claudeBuckets := payload["claude"]["https://claude.example.com|claude-key"] if len(claudeBuckets) != 20 { t.Fatalf("claude buckets len = %d, want 20", len(claudeBuckets)) }