feat(ratelimit): persist window token counters across restarts

Save window state (resets_at + token counts) to ~/.claude/ on shutdown
and every poll cycle. On startup, restore counters if the window hasn't
rolled over. Fixes token counters resetting to zero on deploy.
This commit is contained in:
Alexander
2026-04-14 14:07:28 +02:00
parent b864092dad
commit 273213cbed
2 changed files with 92 additions and 1 deletions
+89
View File
@@ -0,0 +1,89 @@
package ratelimit
import (
"encoding/json"
"os"
"path/filepath"
"time"
"github.com/rs/zerolog/log"
)
const stateFile = ".anthropic-proxy-state.json"
type windowState struct {
ResetsAt time.Time `json:"resets_at"`
TokensIn int64 `json:"tokens_in"`
TokensOut int64 `json:"tokens_out"`
}
type persistedState struct {
FiveHour windowState `json:"five_hour"`
SevenDay windowState `json:"seven_day"`
}
func statePath() string {
home, err := os.UserHomeDir()
if err != nil {
return stateFile
}
return filepath.Join(home, ".claude", stateFile)
}
func (t *Tracker) Save() {
t.mu.RLock()
state := persistedState{
FiveHour: windowState{
ResetsAt: t.fiveHour.ResetsAt,
TokensIn: t.fiveHour.TokensIn.Load(),
TokensOut: t.fiveHour.TokensOut.Load(),
},
SevenDay: windowState{
ResetsAt: t.sevenDay.ResetsAt,
TokensIn: t.sevenDay.TokensIn.Load(),
TokensOut: t.sevenDay.TokensOut.Load(),
},
}
t.mu.RUnlock()
data, err := json.Marshal(state)
if err != nil {
log.Warn().Err(err).Msg("failed to marshal tracker state")
return
}
path := statePath()
if err := os.WriteFile(path, data, 0644); err != nil {
log.Warn().Err(err).Str("path", path).Msg("failed to save tracker state")
return
}
}
func (t *Tracker) Restore() {
path := statePath()
data, err := os.ReadFile(path)
if err != nil {
return
}
var state persistedState
if err := json.Unmarshal(data, &state); err != nil {
log.Warn().Err(err).Msg("failed to parse tracker state")
return
}
t.mu.Lock()
defer t.mu.Unlock()
if !state.FiveHour.ResetsAt.IsZero() && state.FiveHour.ResetsAt == t.fiveHour.ResetsAt {
t.fiveHour.TokensIn.Store(state.FiveHour.TokensIn)
t.fiveHour.TokensOut.Store(state.FiveHour.TokensOut)
log.Info().Int64("tokens_in", state.FiveHour.TokensIn).Int64("tokens_out", state.FiveHour.TokensOut).Msg("restored 5h window tokens")
}
if !state.SevenDay.ResetsAt.IsZero() && state.SevenDay.ResetsAt == t.sevenDay.ResetsAt {
t.sevenDay.TokensIn.Store(state.SevenDay.TokensIn)
t.sevenDay.TokensOut.Store(state.SevenDay.TokensOut)
log.Info().Int64("tokens_in", state.SevenDay.TokensIn).Int64("tokens_out", state.SevenDay.TokensOut).Msg("restored 7d window tokens")
}
}
+3 -1
View File
@@ -45,14 +45,16 @@ func NewTracker(tokenFn func() string) *Tracker {
// Start begins the background poll loop.
func (t *Tracker) Start(ctx context.Context) {
go func() {
// Poll immediately on start
t.poll(ctx)
t.Restore()
for {
select {
case <-ctx.Done():
t.Save()
return
case <-time.After(5 * time.Minute):
t.poll(ctx)
t.Save()
}
}
}()