From 273213cbed1565fbd7e4e33bd2ea3a224b1be06b Mon Sep 17 00:00:00 2001
From: Alexander <alexander0derevianko@gmail.com>
Date: Tue, 14 Apr 2026 14:07:28 +0200
Subject: [PATCH] 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.
---
 internal/ratelimit/persist.go | 89 +++++++++++++++++++++++++++++++++++
 internal/ratelimit/tracker.go |  4 +-
 2 files changed, 92 insertions(+), 1 deletion(-)
 create mode 100644 internal/ratelimit/persist.go

diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go
new file mode 100644
index 0000000..296bc68
--- /dev/null
+++ b/internal/ratelimit/persist.go
@@ -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")
+	}
+}
diff --git a/internal/ratelimit/tracker.go b/internal/ratelimit/tracker.go
index 42b5fdc..569b8e8 100644
--- a/internal/ratelimit/tracker.go
+++ b/internal/ratelimit/tracker.go
@@ -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()
 			}
 		}
 	}()