From 61b39d49bd8cad26c8d74eb0bd0f6b8fda16ab2c Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Tue, 5 May 2026 02:53:04 +0800 Subject: [PATCH] feat(management): add usage record retrieval endpoint - Implemented `/v0/management/usage` endpoint for fetching queued usage records from Redis. - Included validation for `count` parameter to ensure positive integers. - Added unit tests for queue retrieval and validation, with authentication validation in integration tests. - Updated management routing to include the new endpoint. --- internal/api/handlers/management/usage.go | 55 +++++++++++ .../api/handlers/management/usage_test.go | 98 +++++++++++++++++++ internal/api/server.go | 1 + internal/api/server_test.go | 55 +++++++++++ 4 files changed, 209 insertions(+) create mode 100644 internal/api/handlers/management/usage.go create mode 100644 internal/api/handlers/management/usage_test.go diff --git a/internal/api/handlers/management/usage.go b/internal/api/handlers/management/usage.go new file mode 100644 index 000000000..8cb175eb6 --- /dev/null +++ b/internal/api/handlers/management/usage.go @@ -0,0 +1,55 @@ +package management + +import ( + "encoding/json" + "errors" + "net/http" + "strconv" + "strings" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue" +) + +type usageQueueRecord []byte + +func (r usageQueueRecord) MarshalJSON() ([]byte, error) { + if json.Valid(r) { + return append([]byte(nil), r...), nil + } + return json.Marshal(string(r)) +} + +// GetUsage pops queued usage records from the Redis-compatible usage queue. +func (h *Handler) GetUsage(c *gin.Context) { + if h == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "handler unavailable"}) + return + } + + count, errCount := parseUsageQueueCount(c.Query("count")) + if errCount != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": errCount.Error()}) + return + } + + items := redisqueue.PopOldest(count) + records := make([]usageQueueRecord, 0, len(items)) + for _, item := range items { + records = append(records, usageQueueRecord(append([]byte(nil), item...))) + } + + c.JSON(http.StatusOK, records) +} + +func parseUsageQueueCount(value string) (int, error) { + value = strings.TrimSpace(value) + if value == "" { + return 1, nil + } + count, errCount := strconv.Atoi(value) + if errCount != nil || count <= 0 { + return 0, errors.New("count must be a positive integer") + } + return count, nil +} diff --git a/internal/api/handlers/management/usage_test.go b/internal/api/handlers/management/usage_test.go new file mode 100644 index 000000000..5c5f5c69d --- /dev/null +++ b/internal/api/handlers/management/usage_test.go @@ -0,0 +1,98 @@ +package management + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue" +) + +func TestGetUsagePopsRequestedRecords(t *testing.T) { + gin.SetMode(gin.TestMode) + withManagementUsageQueue(t, func() { + redisqueue.Enqueue([]byte(`{"id":1}`)) + redisqueue.Enqueue([]byte(`{"id":2}`)) + redisqueue.Enqueue([]byte(`{"id":3}`)) + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/usage?count=2", nil) + + h := &Handler{} + h.GetUsage(ginCtx) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + var payload []json.RawMessage + if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &payload); errUnmarshal != nil { + t.Fatalf("unmarshal response: %v", errUnmarshal) + } + if len(payload) != 2 { + t.Fatalf("response records = %d, want 2", len(payload)) + } + requireRecordID(t, payload[0], 1) + requireRecordID(t, payload[1], 2) + + remaining := redisqueue.PopOldest(10) + if len(remaining) != 1 || string(remaining[0]) != `{"id":3}` { + t.Fatalf("remaining queue = %q, want third item only", remaining) + } + }) +} + +func TestGetUsageInvalidCountDoesNotPop(t *testing.T) { + gin.SetMode(gin.TestMode) + withManagementUsageQueue(t, func() { + redisqueue.Enqueue([]byte(`{"id":1}`)) + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/usage?count=0", nil) + + h := &Handler{} + h.GetUsage(ginCtx) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusBadRequest, rec.Body.String()) + } + + remaining := redisqueue.PopOldest(10) + if len(remaining) != 1 || string(remaining[0]) != `{"id":1}` { + t.Fatalf("remaining queue = %q, want original item", remaining) + } + }) +} + +func withManagementUsageQueue(t *testing.T, fn func()) { + t.Helper() + + prevQueueEnabled := redisqueue.Enabled() + redisqueue.SetEnabled(false) + redisqueue.SetEnabled(true) + + defer func() { + redisqueue.SetEnabled(false) + redisqueue.SetEnabled(prevQueueEnabled) + }() + + fn() +} + +func requireRecordID(t *testing.T, raw json.RawMessage, want int) { + t.Helper() + + var payload struct { + ID int `json:"id"` + } + if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil { + t.Fatalf("unmarshal record: %v", errUnmarshal) + } + if payload.ID != want { + t.Fatalf("record id = %d, want %d", payload.ID, want) + } +} diff --git a/internal/api/server.go b/internal/api/server.go index 2e89ac5a3..5c43db48c 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -551,6 +551,7 @@ func (s *Server) registerManagementRoutes() { mgmt.PATCH("/api-keys", s.mgmt.PatchAPIKeys) mgmt.DELETE("/api-keys", s.mgmt.DeleteAPIKeys) mgmt.GET("/api-key-usage", s.mgmt.GetAPIKeyUsage) + mgmt.GET("/usage", s.mgmt.GetUsage) mgmt.GET("/gemini-api-key", s.mgmt.GetGeminiKeys) mgmt.PUT("/gemini-api-key", s.mgmt.PutGeminiKeys) diff --git a/internal/api/server_test.go b/internal/api/server_test.go index db1ef27d1..d5718091a 100644 --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -13,6 +13,7 @@ import ( gin "github.com/gin-gonic/gin" proxyconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config" internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue" sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access" "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config" @@ -84,6 +85,60 @@ func TestHealthz(t *testing.T) { }) } +func TestManagementUsageRequiresManagementAuthAndPopsArray(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "test-management-key") + + prevQueueEnabled := redisqueue.Enabled() + redisqueue.SetEnabled(false) + t.Cleanup(func() { + redisqueue.SetEnabled(false) + redisqueue.SetEnabled(prevQueueEnabled) + }) + + server := newTestServer(t) + + redisqueue.Enqueue([]byte(`{"id":1}`)) + redisqueue.Enqueue([]byte(`{"id":2}`)) + + missingKeyReq := httptest.NewRequest(http.MethodGet, "/v0/management/usage?count=2", nil) + missingKeyRR := httptest.NewRecorder() + server.engine.ServeHTTP(missingKeyRR, missingKeyReq) + if missingKeyRR.Code != http.StatusUnauthorized { + t.Fatalf("missing key status = %d, want %d body=%s", missingKeyRR.Code, http.StatusUnauthorized, missingKeyRR.Body.String()) + } + + authReq := httptest.NewRequest(http.MethodGet, "/v0/management/usage?count=2", nil) + authReq.Header.Set("Authorization", "Bearer test-management-key") + authRR := httptest.NewRecorder() + server.engine.ServeHTTP(authRR, authReq) + if authRR.Code != http.StatusOK { + t.Fatalf("authenticated status = %d, want %d body=%s", authRR.Code, http.StatusOK, authRR.Body.String()) + } + + var payload []json.RawMessage + if errUnmarshal := json.Unmarshal(authRR.Body.Bytes(), &payload); errUnmarshal != nil { + t.Fatalf("unmarshal response: %v body=%s", errUnmarshal, authRR.Body.String()) + } + if len(payload) != 2 { + t.Fatalf("response records = %d, want 2", len(payload)) + } + for i, raw := range payload { + var record struct { + ID int `json:"id"` + } + if errUnmarshal := json.Unmarshal(raw, &record); errUnmarshal != nil { + t.Fatalf("unmarshal record %d: %v", i, errUnmarshal) + } + if record.ID != i+1 { + t.Fatalf("record %d id = %d, want %d", i, record.ID, i+1) + } + } + + if remaining := redisqueue.PopOldest(1); len(remaining) != 0 { + t.Fatalf("remaining queue = %q, want empty", remaining) + } +} + func TestAmpProviderModelRoutes(t *testing.T) { testCases := []struct { name string