feat(api): refactor provider key logic for API key usage and add test for compatibility grouping

- Extracted provider key determination into `apiKeyUsageProviderKey` for reuse and better readability.
- Updated `GetAPIKeyUsage` to utilize the new function.
- Added a new test case to validate compatibility grouping logic via `compat_name` attribute.

Closes: #3940 #3941
This commit is contained in:
Luis Pater
2026-06-21 11:37:52 +08:00
parent babef2a1ef
commit 369e560f1d
2 changed files with 62 additions and 4 deletions
@@ -40,6 +40,19 @@ func mergeRecentRequestBuckets(dst, src []coreauth.RecentRequestBucket) []coreau
return dst
}
func apiKeyUsageProviderKey(auth *coreauth.Auth) string {
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
if auth.Attributes != nil {
if compatName := strings.TrimSpace(auth.Attributes["compat_name"]); compatName != "" {
provider = strings.ToLower(compatName)
}
}
if provider == "" {
return "unknown"
}
return provider
}
// GetAPIKeyUsage returns recent request buckets for all in-memory api_key auths,
// grouped by provider and keyed by "base_url|api_key".
func (h *Handler) GetAPIKeyUsage(c *gin.Context) {
@@ -78,10 +91,7 @@ func (h *Handler) GetAPIKeyUsage(c *gin.Context) {
}
}
compositeKey := baseURL + "|" + apiKey
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
if provider == "" {
provider = "unknown"
}
provider := apiKeyUsageProviderKey(auth)
recent := auth.RecentRequestsSnapshot(now)
providerBucket, ok := out[provider]
@@ -92,3 +92,51 @@ func TestGetAPIKeyUsage_GroupsByProviderAndAPIKey(t *testing.T) {
t.Fatalf("claude totals = %d/%d, want 1/0", claudeSuccess, claudeFailed)
}
}
func TestGetAPIKeyUsage_GroupsOpenAICompatibleByCompatName(t *testing.T) {
t.Setenv("MANAGEMENT_PASSWORD", "")
manager := coreauth.NewManager(nil, nil, nil)
if _, err := manager.Register(context.Background(), &coreauth.Auth{
ID: "vast-auth",
Provider: "openai-compatible-vast",
Attributes: map[string]string{
"api_key": "vast-key",
"base_url": "https://www.vastnum.com/v1",
"compat_name": "VAST",
},
}); err != nil {
t.Fatalf("register vast auth: %v", err)
}
manager.MarkResult(context.Background(), coreauth.Result{AuthID: "vast-auth", Provider: "openai-compatible-vast", Model: "gpt-5", 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]apiKeyUsageEntry
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
t.Fatalf("decode payload: %v", err)
}
if _, exists := payload["openai-compatible-vast"]; exists {
t.Fatalf("unexpected namespaced provider bucket in payload: %#v", payload)
}
vastBucket, exists := payload["vast"]
if !exists {
t.Fatalf("missing compat provider bucket in payload: %#v", payload)
}
vastEntry := vastBucket["https://www.vastnum.com/v1|vast-key"]
if vastEntry.Success != 1 || vastEntry.Failed != 0 {
t.Fatalf("vast totals = %d/%d, want 1/0", vastEntry.Success, vastEntry.Failed)
}
}