Background Quota Refresh & Automated Token Management

config.example.yaml

@@ -52,6 +52,12 @@ logs-max-total-size-mb: 0
 # When false, disable in-memory usage statistics aggregation
 usage-statistics-enabled: false
 
+# Quota refresh interval in seconds. When set to a positive value, the server
+# will periodically fetch quota data for all configured providers in the background
+# and cache it in memory. This eliminates the need to manually refresh quota data.
+# Set to 0 to disable background refresh (fetch on-demand only). Default: 0
+# quota-refresh-interval: 300  # 5 minutes
+
 # Proxy URL. Supports socks5/http/https protocols. Example: socks5://user:[email protected]:1080/
 proxy-url: ""
 
@@ -75,7 +81,6 @@ routing:
 
 # When true, enable authentication for the WebSocket API (/v1/ws).
 ws-auth: false
-
 # Streaming behavior (SSE keep-alives + safe bootstrap retries).
 # streaming:
 #   keepalive-seconds: 15   # Default: 0 (disabled). <= 0 disables keep-alives.

internal/api/handlers/management/handler.go

@@ -3,6 +3,7 @@
 package management
 
 import (
+	"context"
 	"crypto/subtle"
 	"fmt"
 	"net/http"
@@ -15,6 +16,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/buildinfo"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/quota"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
 	sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
 	coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
@@ -36,6 +38,7 @@ type Handler struct {
 	authManager         *coreauth.Manager
 	usageStats          *usage.RequestStatistics
 	tokenStore          coreauth.Store
+	quotaManager        *quota.Manager
 	localPassword       string
 	allowRemoteOverride bool
 	envSecret           string
@@ -47,13 +50,23 @@ func NewHandler(cfg *config.Config, configFilePath string, manager *coreauth.Man
 	envSecret, _ := os.LookupEnv("MANAGEMENT_PASSWORD")
 	envSecret = strings.TrimSpace(envSecret)
 
+	tokenStore := sdkAuth.GetTokenStore()
+	quotaManager := quota.NewManager(tokenStore, nil)
+
+	// Configure quota refresh interval if set in config
+	if cfg != nil && cfg.QuotaRefreshInterval > 0 {
+		interval := time.Duration(cfg.QuotaRefreshInterval) * time.Second
+		quotaManager.SetRefreshInterval(interval)
+	}
+
 	return &Handler{
 		cfg:                 cfg,
 		configFilePath:      configFilePath,
 		failedAttempts:      make(map[string]*attemptInfo),
 		authManager:         manager,
 		usageStats:          usage.GetRequestStatistics(),
-		tokenStore:          sdkAuth.GetTokenStore(),
+		tokenStore:          tokenStore,
+		quotaManager:        quotaManager,
 		allowRemoteOverride: envSecret != "",
 		envSecret:           envSecret,
 	}
@@ -89,6 +102,22 @@ func (h *Handler) SetLogDirectory(dir string) {
 	h.logDir = dir
 }
 
+// StartBackgroundWorkers starts background workers (quota refresh, etc.)
+// This should be called when the server is ready to handle requests.
+func (h *Handler) StartBackgroundWorkers(ctx context.Context) {
+	if h.quotaManager != nil {
+		h.quotaManager.StartWorker(ctx)
+	}
+}
+
+// StopBackgroundWorkers stops all background workers.
+// This should be called during server shutdown.
+func (h *Handler) StopBackgroundWorkers() {
+	if h.quotaManager != nil {
+		h.quotaManager.StopWorker()
+	}
+}
+
 // Middleware enforces access control for management endpoints.
 // All requests (local and remote) require a valid management key.
 // Additionally, remote access requires allow-remote-management=true.

internal/api/handlers/management/quota_fetchers.go

@@ -0,0 +1,132 @@
+package management
+
+import (
+	"errors"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/quota"
+)
+
+// GetAllQuotas returns quota for all connected accounts.
+// GET /v0/management/quotas
+func (h *Handler) GetAllQuotas(c *gin.Context) {
+	if h.quotaManager == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "quota manager not initialized"})
+		return
+	}
+
+	ctx := c.Request.Context()
+	quotas, err := h.quotaManager.FetchAllQuotas(ctx)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(http.StatusOK, quotas)
+}
+
+// GetProviderQuotas returns quota for a specific provider.
+// GET /v0/management/quotas/:provider
+func (h *Handler) GetProviderQuotas(c *gin.Context) {
+	if h.quotaManager == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "quota manager not initialized"})
+		return
+	}
+
+	provider := c.Param("provider")
+	if provider == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "provider is required"})
+		return
+	}
+
+	ctx := c.Request.Context()
+	quotas, err := h.quotaManager.FetchProviderQuotas(ctx, provider)
+	if err != nil {
+		if errors.Is(err, quota.ErrUnknownProvider) {
+			c.JSON(http.StatusNotFound, gin.H{
+				"error":           err.Error(),
+				"known_providers": h.quotaManager.GetKnownProviders(),
+			})
+			return
+		}
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(http.StatusOK, quotas)
+}
+
+// GetAccountQuota returns quota for a specific account.
+// GET /v0/management/quotas/:provider/:account
+func (h *Handler) GetAccountQuota(c *gin.Context) {
+	if h.quotaManager == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "quota manager not initialized"})
+		return
+	}
+
+	provider := c.Param("provider")
+	account := c.Param("account")
+	if provider == "" || account == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "provider and account are required"})
+		return
+	}
+
+	ctx := c.Request.Context()
+	quotaResp, err := h.quotaManager.FetchAccountQuota(ctx, provider, account)
+	if err != nil {
+		if errors.Is(err, quota.ErrUnknownProvider) {
+			c.JSON(http.StatusNotFound, gin.H{
+				"error":           err.Error(),
+				"known_providers": h.quotaManager.GetKnownProviders(),
+			})
+			return
+		}
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(http.StatusOK, quotaResp)
+}
+
+// RefreshQuotas forces a quota refresh for all or specific providers.
+// POST /v0/management/quotas/refresh
+func (h *Handler) RefreshQuotas(c *gin.Context) {
+	if h.quotaManager == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "quota manager not initialized"})
+		return
+	}
+
+	var req quota.RefreshRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		// Allow empty body - refresh all
+		req.Providers = nil
+	}
+
+	ctx := c.Request.Context()
+	quotas, err := h.quotaManager.RefreshQuotas(ctx, req.Providers)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(http.StatusOK, quotas)
+}
+
+// GetSubscriptionInfo returns subscription/tier info for Antigravity accounts.
+// GET /v0/management/subscription-info
+func (h *Handler) GetSubscriptionInfo(c *gin.Context) {
+	if h.quotaManager == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "quota manager not initialized"})
+		return
+	}
+
+	ctx := c.Request.Context()
+	info, err := h.quotaManager.GetSubscriptionInfo(ctx)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(http.StatusOK, info)
+}

internal/api/server.go

@@ -595,6 +595,13 @@ func (s *Server) registerManagementRoutes() {
 		mgmt.POST("/iflow-auth-url", s.mgmt.RequestIFlowCookieToken)
 		mgmt.POST("/oauth-callback", s.mgmt.PostOAuthCallback)
 		mgmt.GET("/get-auth-status", s.mgmt.GetAuthStatus)
+
+		// Quota fetching endpoints
+		mgmt.GET("/quotas", s.mgmt.GetAllQuotas)
+		mgmt.GET("/quotas/:provider", s.mgmt.GetProviderQuotas)
+		mgmt.GET("/quotas/:provider/:account", s.mgmt.GetAccountQuota)
+		mgmt.POST("/quotas/refresh", s.mgmt.RefreshQuotas)
+		mgmt.GET("/subscription-info", s.mgmt.GetSubscriptionInfo)
 	}
 }
 
@@ -764,6 +771,22 @@ func (s *Server) Start() error {
 	return nil
 }
 
+// StartBackgroundWorkers starts background workers such as quota refresh.
+// This should be called before Start() with the run context.
+func (s *Server) StartBackgroundWorkers(ctx context.Context) {
+	if s.mgmt != nil {
+		s.mgmt.StartBackgroundWorkers(ctx)
+	}
+}
+
+// StopBackgroundWorkers stops all background workers.
+// This is called automatically by Stop() but can be called manually if needed.
+func (s *Server) StopBackgroundWorkers() {
+	if s.mgmt != nil {
+		s.mgmt.StopBackgroundWorkers()
+	}
+}
+
 // Stop gracefully shuts down the API server without interrupting any
 // active connections.
 //
@@ -775,6 +798,9 @@ func (s *Server) Start() error {
 func (s *Server) Stop(ctx context.Context) error {
 	log.Debug("Stopping API server...")
 
+	// Stop background workers first
+	s.StopBackgroundWorkers()
+
 	if s.keepAliveEnabled {
 		select {
 		case s.keepAliveStop <- struct{}{}:

internal/config/config.go

@@ -55,6 +55,12 @@ type Config struct {
 	// DisableCooling disables quota cooldown scheduling when true.
 	DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`
 
+	// QuotaRefreshInterval is the interval in seconds for background quota refresh.
+	// When set to a positive value, the server will periodically fetch quota data
+	// for all configured providers in the background and cache it in memory.
+	// Set to 0 to disable background refresh (fetch on-demand only). Default: 0
+	QuotaRefreshInterval int `yaml:"quota-refresh-interval" json:"quota-refresh-interval"`
+
 	// RequestRetry defines the retry times when the request failed.
 	RequestRetry int `yaml:"request-retry" json:"request-retry"`
 	// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.