Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 53 additions & 0 deletions internal/api/handlers/management/config_basic.go
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,59 @@ func (h *Handler) PutUsageStatisticsEnabled(c *gin.Context) {
h.updateBoolField(c, func(v bool) { h.cfg.UsageStatisticsEnabled = v })
}

// UsagePersistence
func (h *Handler) GetUsagePersistence(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(200, gin.H{"usage-persistence": gin.H{}})
return
}
status := gin.H{}
if h.usagePersistence != nil {
persistStatus := h.usagePersistence.Status()
status["runtime"] = persistStatus
}
c.JSON(200, gin.H{
"usage-persistence": h.cfg.UsagePersistence,
"status": status,
})
}

func (h *Handler) PutUsagePersistence(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid config state"})
return
}
var body struct {
Enabled *bool `json:"enabled"`
FilePath *string `json:"file-path"`
IntervalSeconds *int `json:"interval-seconds"`
}
if errBindJSON := c.ShouldBindJSON(&body); errBindJSON != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
return
}
if body.Enabled != nil {
h.cfg.UsagePersistence.Enabled = *body.Enabled
}
if body.FilePath != nil {
h.cfg.UsagePersistence.FilePath = strings.TrimSpace(*body.FilePath)
if h.cfg.UsagePersistence.FilePath == "" {
h.cfg.UsagePersistence.FilePath = "usage-statistics.json"
}
}
if body.IntervalSeconds != nil {
if *body.IntervalSeconds <= 0 {
c.JSON(http.StatusBadRequest, gin.H{"error": "interval-seconds must be greater than 0"})
return
}
h.cfg.UsagePersistence.IntervalSeconds = *body.IntervalSeconds
}
if h.usagePersistence != nil {
h.usagePersistence.ApplyConfig(h.cfg.UsagePersistence)
}
h.persist(c)
}

// UsageStatisticsEnabled
func (h *Handler) GetLoggingToFile(c *gin.Context) {
c.JSON(200, gin.H{"logging-to-file": h.cfg.LoggingToFile})
Expand Down
25 changes: 24 additions & 1 deletion internal/api/handlers/management/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ type Handler struct {
failedAttempts map[string]*attemptInfo // keyed by client IP
authManager *coreauth.Manager
usageStats *usage.RequestStatistics
usagePersistence *usage.PersistenceManager
tokenStore coreauth.Store
localPassword string
allowRemoteOverride bool
Expand All @@ -61,10 +62,14 @@ func NewHandler(cfg *config.Config, configFilePath string, manager *coreauth.Man
failedAttempts: make(map[string]*attemptInfo),
authManager: manager,
usageStats: usage.GetRequestStatistics(),
usagePersistence: usage.NewPersistenceManager(usage.GetRequestStatistics(), filepath.Dir(configFilePath)),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The PersistenceManager is instantiated here within NewHandler, but it is immediately replaced by the Server via SetUsagePersistenceManager. This creates a temporary, unused PersistenceManager instance on every NewHandler call, which is inefficient and confusing. The management.Handler should receive its dependencies rather than creating them.

tokenStore: sdkAuth.GetTokenStore(),
allowRemoteOverride: envSecret != "",
envSecret: envSecret,
}
if cfg != nil {
h.usagePersistence.ApplyConfig(cfg.UsagePersistence)
}
h.startAttemptCleanup()
return h
}
Expand Down Expand Up @@ -105,14 +110,22 @@ func NewHandlerWithoutConfigFilePath(cfg *config.Config, manager *coreauth.Manag
}

// SetConfig updates the in-memory config reference when the server hot-reloads.
func (h *Handler) SetConfig(cfg *config.Config) { h.cfg = cfg }
func (h *Handler) SetConfig(cfg *config.Config) {
h.cfg = cfg
if h != nil && h.usagePersistence != nil && cfg != nil {
h.usagePersistence.ApplyConfig(cfg.UsagePersistence)
}
}
Comment on lines +113 to +118
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The ApplyConfig method for usagePersistence is called here within SetConfig, and also in Server.UpdateClients. This results in a redundant call during a hot reload. The Server should be the single source of truth for orchestrating configuration updates to its components.

func (h *Handler) SetConfig(cfg *config.Config) {
	h.cfg = cfg
}


// SetAuthManager updates the auth manager reference used by management endpoints.
func (h *Handler) SetAuthManager(manager *coreauth.Manager) { h.authManager = manager }

// SetUsageStatistics allows replacing the usage statistics reference.
func (h *Handler) SetUsageStatistics(stats *usage.RequestStatistics) { h.usageStats = stats }

// SetUsagePersistenceManager replaces the usage persistence manager.
func (h *Handler) SetUsagePersistenceManager(manager *usage.PersistenceManager) { h.usagePersistence = manager }

// SetLocalPassword configures the runtime-local password accepted for localhost requests.
func (h *Handler) SetLocalPassword(password string) { h.localPassword = password }

Expand All @@ -134,6 +147,16 @@ func (h *Handler) SetPostAuthHook(hook coreauth.PostAuthHook) {
h.postAuthHook = hook
}

// Stop releases background resources owned by management handler.
func (h *Handler) Stop() {
if h == nil {
return
}
if h.usagePersistence != nil {
h.usagePersistence.Stop(true)
}
}

// 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.
Expand Down
42 changes: 42 additions & 0 deletions internal/api/handlers/management/usage.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,3 +77,45 @@ func (h *Handler) ImportUsageStatistics(c *gin.Context) {
"failed_requests": snapshot.FailureCount,
})
}

// GetUsagePersistenceStatus returns runtime usage persistence status.
func (h *Handler) GetUsagePersistenceStatus(c *gin.Context) {
if h == nil || h.usagePersistence == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "usage persistence unavailable"})
return
}
c.JSON(http.StatusOK, gin.H{"status": h.usagePersistence.Status()})
}

// SaveUsageStatistics persists current usage snapshot immediately.
func (h *Handler) SaveUsageStatistics(c *gin.Context) {
if h == nil || h.usagePersistence == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "usage persistence unavailable"})
return
}
status, err := h.usagePersistence.SaveNow()
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error(), "status": status})
return
}
c.JSON(http.StatusOK, gin.H{"status": status})
}

// LoadUsageStatistics loads usage snapshot from persistence and merges into memory.
func (h *Handler) LoadUsageStatistics(c *gin.Context) {
if h == nil || h.usagePersistence == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "usage persistence unavailable"})
return
}
result, err := h.usagePersistence.LoadNow()
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error(), "result": result})
return
}
snapshot := h.usageStats.Snapshot()
c.JSON(http.StatusOK, gin.H{
"result": result,
"total_requests": snapshot.TotalRequests,
"failed_requests": snapshot.FailureCount,
})
}
23 changes: 23 additions & 0 deletions internal/api/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,8 @@ type Server struct {
keepAliveOnTimeout func()
keepAliveHeartbeat chan struct{}
keepAliveStop chan struct{}

usagePersistence *usage.PersistenceManager
}

// NewServer creates and initializes a new API server instance.
Expand Down Expand Up @@ -251,6 +253,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
currentPath: wd,
envManagementSecret: envManagementSecret,
wsRoutes: make(map[string]struct{}),
usagePersistence: usage.NewPersistenceManager(usage.GetRequestStatistics(), filepath.Dir(configFilePath)),
}
s.wsAuthEnabled.Store(cfg.WebsocketAuth)
// Save initial YAML snapshot
Expand All @@ -263,6 +266,10 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
// Initialize management handler
s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager)
if s.usagePersistence != nil {
s.usagePersistence.ApplyConfig(cfg.UsagePersistence)
s.mgmt.SetUsagePersistenceManager(s.usagePersistence)
Comment on lines +269 to +271

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Stop old persistence manager before swapping handler manager

When usage-persistence.enabled is true at startup, management.NewHandler has already created and configured its own PersistenceManager (which can start autosave), and this assignment then swaps in s.usagePersistence without stopping the original one. That leaves an orphaned goroutine writing with stale settings, so you can get duplicate/background saves and persistence may continue even after later disabling the active manager.

Useful? React with 👍 / 👎.

}
if optionState.localPassword != "" {
s.mgmt.SetLocalPassword(optionState.localPassword)
}
Expand Down Expand Up @@ -489,6 +496,9 @@ func (s *Server) registerManagementRoutes() {
mgmt.GET("/usage", s.mgmt.GetUsageStatistics)
mgmt.GET("/usage/export", s.mgmt.ExportUsageStatistics)
mgmt.POST("/usage/import", s.mgmt.ImportUsageStatistics)
mgmt.GET("/usage/persistence-status", s.mgmt.GetUsagePersistenceStatus)
mgmt.POST("/usage/save", s.mgmt.SaveUsageStatistics)
mgmt.POST("/usage/load", s.mgmt.LoadUsageStatistics)
mgmt.GET("/config", s.mgmt.GetConfig)
mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML)
mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML)
Expand All @@ -514,6 +524,10 @@ func (s *Server) registerManagementRoutes() {
mgmt.PUT("/usage-statistics-enabled", s.mgmt.PutUsageStatisticsEnabled)
mgmt.PATCH("/usage-statistics-enabled", s.mgmt.PutUsageStatisticsEnabled)

mgmt.GET("/usage-persistence", s.mgmt.GetUsagePersistence)
mgmt.PUT("/usage-persistence", s.mgmt.PutUsagePersistence)
mgmt.PATCH("/usage-persistence", s.mgmt.PutUsagePersistence)

mgmt.GET("/proxy-url", s.mgmt.GetProxyURL)
mgmt.PUT("/proxy-url", s.mgmt.PutProxyURL)
mgmt.PATCH("/proxy-url", s.mgmt.PutProxyURL)
Expand Down Expand Up @@ -830,6 +844,10 @@ func (s *Server) Stop(ctx context.Context) error {
}
}

if s.mgmt != nil {
s.mgmt.Stop()
}

// Shutdown the HTTP server.
if err := s.server.Shutdown(ctx); err != nil {
Comment on lines +848 to 852

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Flush persistence after HTTP shutdown to avoid lost requests

The stop sequence flushes usage persistence before s.server.Shutdown(ctx). During graceful shutdown, in-flight requests can still finish and update usage statistics after this early flush, but there is no second flush afterward, so tail requests are dropped from persisted stats under active traffic.

Useful? React with 👍 / 👎.

return fmt.Errorf("failed to shutdown HTTP server: %v", err)
Expand Down Expand Up @@ -904,6 +922,10 @@ func (s *Server) UpdateClients(cfg *config.Config) {
usage.SetStatisticsEnabled(cfg.UsageStatisticsEnabled)
}

if s.usagePersistence != nil && (oldCfg == nil || !reflect.DeepEqual(oldCfg.UsagePersistence, cfg.UsagePersistence)) {
s.usagePersistence.ApplyConfig(cfg.UsagePersistence)
}

if s.requestLogger != nil && (oldCfg == nil || oldCfg.ErrorLogsMaxFiles != cfg.ErrorLogsMaxFiles) {
if setter, ok := s.requestLogger.(interface{ SetErrorLogsMaxFiles(int) }); ok {
setter.SetErrorLogsMaxFiles(cfg.ErrorLogsMaxFiles)
Expand Down Expand Up @@ -970,6 +992,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
if s.mgmt != nil {
s.mgmt.SetConfig(cfg)
s.mgmt.SetAuthManager(s.handlers.AuthManager)
s.mgmt.SetUsagePersistenceManager(s.usagePersistence)
}

// Notify Amp module only when Amp config has changed.
Expand Down
24 changes: 24 additions & 0 deletions internal/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,9 @@ type Config struct {
// UsageStatisticsEnabled toggles in-memory usage aggregation; when false, usage data is discarded.
UsageStatisticsEnabled bool `yaml:"usage-statistics-enabled" json:"usage-statistics-enabled"`

// UsagePersistence controls optional periodic persistence of usage statistics.
UsagePersistence UsagePersistenceConfig `yaml:"usage-persistence" json:"usage-persistence"`

// DisableCooling disables quota cooldown scheduling when true.
DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`

Expand Down Expand Up @@ -128,6 +131,16 @@ type Config struct {
legacyMigrationPending bool `yaml:"-" json:"-"`
}

// UsagePersistenceConfig defines usage statistics persistence behavior.
type UsagePersistenceConfig struct {
// Enabled toggles automatic usage persistence.
Enabled bool `yaml:"enabled" json:"enabled"`
// FilePath is the output file path for usage snapshots. Relative paths are resolved from config directory.
FilePath string `yaml:"file-path" json:"file-path"`
// IntervalSeconds controls periodic flush interval in seconds.
IntervalSeconds int `yaml:"interval-seconds" json:"interval-seconds"`
}

// ClaudeHeaderDefaults configures default header values injected into Claude API requests
// when the client does not send them. Update these when Claude Code releases a new version.
type ClaudeHeaderDefaults struct {
Expand Down Expand Up @@ -553,6 +566,9 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
cfg.LogsMaxTotalSizeMB = 0
cfg.ErrorLogsMaxFiles = 10
cfg.UsageStatisticsEnabled = false
cfg.UsagePersistence.Enabled = false
cfg.UsagePersistence.FilePath = "usage-statistics.json"
cfg.UsagePersistence.IntervalSeconds = 30
Comment on lines +570 to +571
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The default file path "usage-statistics.json" and interval 30 for usage persistence are hardcoded as magic values in multiple files (internal/config/config.go, internal/api/handlers/management/config_basic.go, internal/usage/persistence.go). This makes the configuration harder to maintain and increases the risk of inconsistencies. These default values should be defined as exported constants in the config package and used throughout the codebase.

cfg.DisableCooling = false
cfg.Pprof.Enable = false
cfg.Pprof.Addr = DefaultPprofAddr
Expand Down Expand Up @@ -618,6 +634,14 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
cfg.MaxRetryCredentials = 0
}

cfg.UsagePersistence.FilePath = strings.TrimSpace(cfg.UsagePersistence.FilePath)
if cfg.UsagePersistence.FilePath == "" {
cfg.UsagePersistence.FilePath = "usage-statistics.json"
}
if cfg.UsagePersistence.IntervalSeconds <= 0 {
cfg.UsagePersistence.IntervalSeconds = 30
}

// Sanitize Gemini API key configuration and migrate legacy entries.
cfg.SanitizeGeminiKeys()

Expand Down
Loading
Loading