Merge pull request #54 from kaitranntt/kai/fix/52-upstream-sync-usage-stats

fix: restore usage statistics support

Author: kaitranntt

.ccs-fork-upstream.env

@@ -1,2 +1,2 @@
-UPSTREAM_TAG=v6.9.45
-UPSTREAM_COMMIT=8b286e8fb39e1cc95dd86d8923e8baa83dec8722
+UPSTREAM_TAG=v6.10.0
+UPSTREAM_COMMIT=18bb9c315fced2c428f57b4a0e66b06183c46c06

.goreleaser.yml

@@ -19,6 +19,8 @@ builds:
 archives:
   - id: "cli-proxy-api-plus"
     format: tar.gz
+    name_template: >-
+      {{ .ProjectName }}_{{ .Version }}_{{ .Os }}_{{- if eq .Arch "arm64" -}}aarch64{{- else -}}{{ .Arch }}{{- end -}}
     format_overrides:
       - goos: windows
         format: zip

cmd/server/main.go

@@ -25,6 +25,7 @@ import (
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/store"
 	_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator"
@@ -470,6 +471,7 @@ func main() {
 		}
 	}
 	usage.SetStatisticsEnabled(cfg.UsageStatisticsEnabled)
+	redisqueue.SetUsageStatisticsEnabled(cfg.UsageStatisticsEnabled)
 	coreauth.SetQuotaCooldownDisabled(cfg.DisableCooling)
 
 	if err = logging.ConfigureLogOutput(cfg); err != nil {

docker-build.sh

@@ -5,123 +5,13 @@
 # This script automates the process of building and running the Docker container
 # with version information dynamically injected at build time.
 
-# Hidden feature: Preserve usage statistics across rebuilds
-# Usage: ./docker-build.sh --with-usage
-# First run prompts for management API key, saved to temp/stats/.api_secret
-
 set -euo pipefail
 
-STATS_DIR="temp/stats"
-STATS_FILE="${STATS_DIR}/.usage_backup.json"
-SECRET_FILE="${STATS_DIR}/.api_secret"
-WITH_USAGE=false
-
-get_port() {
-  if [[ -f "config.yaml" ]]; then
-    grep -E "^port:" config.yaml | sed -E 's/^port: *["'"'"']?([0-9]+)["'"'"']?.*$/\1/'
-  else
-    echo "8317"
-  fi
-}
-
-export_stats_api_secret() {
-  if [[ -f "${SECRET_FILE}" ]]; then
-    API_SECRET=$(cat "${SECRET_FILE}")
-  else
-    if [[ ! -d "${STATS_DIR}" ]]; then
-      mkdir -p "${STATS_DIR}"
-    fi
-    echo "First time using --with-usage. Management API key required."
-    read -r -p "Enter management key: " -s API_SECRET
-    echo
-    echo "${API_SECRET}" > "${SECRET_FILE}"
-    chmod 600 "${SECRET_FILE}"
-  fi
-}
-
-check_container_running() {
-  local port
-  port=$(get_port)
-
-  if ! curl -s -o /dev/null -w "%{http_code}" "http://localhost:${port}/" | grep -q "200"; then
-    echo "Error: cli-proxy-api service is not responding at localhost:${port}"
-    echo "Please start the container first or use without --with-usage flag."
-    exit 1
-  fi
-}
-
-export_stats() {
-  local port
-  port=$(get_port)
-
-  if [[ ! -d "${STATS_DIR}" ]]; then
-    mkdir -p "${STATS_DIR}"
-  fi
-  check_container_running
-  echo "Exporting usage statistics..."
-  EXPORT_RESPONSE=$(curl -s -w "\n%{http_code}" -H "X-Management-Key: ${API_SECRET}" \
-    "http://localhost:${port}/v0/management/usage/export")
-  HTTP_CODE=$(echo "${EXPORT_RESPONSE}" | tail -n1)
-  RESPONSE_BODY=$(echo "${EXPORT_RESPONSE}" | sed '$d')
-
-  if [[ "${HTTP_CODE}" != "200" ]]; then
-    echo "Export failed (HTTP ${HTTP_CODE}): ${RESPONSE_BODY}"
-    exit 1
-  fi
-
-  echo "${RESPONSE_BODY}" > "${STATS_FILE}"
-  echo "Statistics exported to ${STATS_FILE}"
-}
-
-import_stats() {
-  local port
-  port=$(get_port)
-
-  echo "Importing usage statistics..."
-  IMPORT_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST \
-    -H "X-Management-Key: ${API_SECRET}" \
-    -H "Content-Type: application/json" \
-    -d @"${STATS_FILE}" \
-    "http://localhost:${port}/v0/management/usage/import")
-  IMPORT_CODE=$(echo "${IMPORT_RESPONSE}" | tail -n1)
-  IMPORT_BODY=$(echo "${IMPORT_RESPONSE}" | sed '$d')
-
-  if [[ "${IMPORT_CODE}" == "200" ]]; then
-    echo "Statistics imported successfully"
-  else
-    echo "Import failed (HTTP ${IMPORT_CODE}): ${IMPORT_BODY}"
-  fi
-
-  rm -f "${STATS_FILE}"
-}
-
-wait_for_service() {
-  local port
-  port=$(get_port)
-
-  echo "Waiting for service to be ready..."
-  for i in {1..30}; do
-    if curl -s -o /dev/null -w "%{http_code}" "http://localhost:${port}/" | grep -q "200"; then
-      break
-    fi
-    sleep 1
-  done
-  sleep 2
-}
-
-case "${1:-}" in
-  "")
-    ;;
-  "--with-usage")
-    WITH_USAGE=true
-    export_stats_api_secret
-    ;;
-  *)
-    echo "Error: unknown option '${1}'. Did you mean '--with-usage'?"
-    echo "Usage: ./docker-build.sh [--with-usage]"
-    exit 1
-    ;;
-esac
+if [[ "${1:-}" != "" ]]; then
+  echo "Error: unknown option '${1}'."
+  echo "Usage: ./docker-build.sh"
+  exit 1
+fi
 
 # --- Step 1: Choose Environment ---
 echo "Please select an option:"
@@ -133,14 +23,7 @@ read -r -p "Enter choice [1-2]: " choice
 case "$choice" in
   1)
     echo "--- Running with Pre-built Image ---"
-    if [[ "${WITH_USAGE}" == "true" ]]; then
-      export_stats
-    fi
     docker compose up -d --remove-orphans --no-build
-    if [[ "${WITH_USAGE}" == "true" ]]; then
-      wait_for_service
-      import_stats
-    fi
     echo "Services are starting from remote image."
     echo "Run 'docker compose logs -f' to see the logs."
     ;;
@@ -167,18 +50,9 @@ case "$choice" in
       --build-arg COMMIT="${COMMIT}" \
       --build-arg BUILD_DATE="${BUILD_DATE}"
 
-    if [[ "${WITH_USAGE}" == "true" ]]; then
-      export_stats
-    fi
-
     echo "Starting the services..."
     docker compose up -d --remove-orphans --pull never
 
-    if [[ "${WITH_USAGE}" == "true" ]]; then
-      wait_for_service
-      import_stats
-    fi
-
     echo "Build complete. Services are starting."
     echo "Run 'docker compose logs -f' to see the logs."
     ;;

internal/api/handlers/management/api_key_usage.go

@@ -0,0 +1,107 @@
+package management
+
+import (
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+)
+
+type apiKeyUsageEntry struct {
+	Success        int64                          `json:"success"`
+	Failed         int64                          `json:"failed"`
+	RecentRequests []coreauth.RecentRequestBucket `json:"recent_requests"`
+}
+
+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 "base_url|api_key".
+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]apiKeyUsageEntry)
+	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
+		}
+		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"
+		}
+
+		recent := auth.RecentRequestsSnapshot(now)
+		providerBucket, ok := out[provider]
+		if !ok {
+			providerBucket = make(map[string]apiKeyUsageEntry)
+			out[provider] = providerBucket
+		}
+		if existing, exists := providerBucket[compositeKey]; exists {
+			existing.Success += auth.Success
+			existing.Failed += auth.Failed
+			existing.RecentRequests = mergeRecentRequestBuckets(existing.RecentRequests, recent)
+			providerBucket[compositeKey] = existing
+			continue
+		}
+		providerBucket[compositeKey] = apiKeyUsageEntry{
+			Success:        auth.Success,
+			Failed:         auth.Failed,
+			RecentRequests: recent,
+		}
+	}
+
+	c.JSON(http.StatusOK, out)
+}

internal/api/handlers/management/api_key_usage_test.go

@@ -0,0 +1,95 @@
+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",
+			"base_url": "https://codex.example.com",
+		},
+	}); 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",
+			"base_url": "https://claude.example.com",
+		},
+	}); 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]apiKeyUsageEntry
+	if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
+		t.Fatalf("decode payload: %v", err)
+	}
+
+	codexEntry := payload["codex"]["https://codex.example.com|codex-key"]
+	if codexEntry.Success != 1 || codexEntry.Failed != 1 {
+		t.Fatalf("codex totals = %d/%d, want 1/1", codexEntry.Success, codexEntry.Failed)
+	}
+	if len(codexEntry.RecentRequests) != 20 {
+		t.Fatalf("codex buckets len = %d, want 20", len(codexEntry.RecentRequests))
+	}
+	codexSuccess, codexFailed := sumRecentRequestBuckets(codexEntry.RecentRequests)
+	if codexSuccess != 1 || codexFailed != 1 {
+		t.Fatalf("codex totals = %d/%d, want 1/1", codexSuccess, codexFailed)
+	}
+
+	claudeEntry := payload["claude"]["https://claude.example.com|claude-key"]
+	if claudeEntry.Success != 1 || claudeEntry.Failed != 0 {
+		t.Fatalf("claude totals = %d/%d, want 1/0", claudeEntry.Success, claudeEntry.Failed)
+	}
+	if len(claudeEntry.RecentRequests) != 20 {
+		t.Fatalf("claude buckets len = %d, want 20", len(claudeEntry.RecentRequests))
+	}
+	claudeSuccess, claudeFailed := sumRecentRequestBuckets(claudeEntry.RecentRequests)
+	if claudeSuccess != 1 || claudeFailed != 0 {
+		t.Fatalf("claude totals = %d/%d, want 1/0", claudeSuccess, claudeFailed)
+	}
+}