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.
This commit is contained in:
Luis Pater
2026-05-05 02:53:04 +08:00
parent ba5d8ca733
commit 61b39d49bd
4 changed files with 209 additions and 0 deletions

View File

@@ -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
}

View File

@@ -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)
}
}

View File

@@ -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)

View File

@@ -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