Alexander
273213cbed
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.
90 lines
2.2 KiB
Go
90 lines
2.2 KiB
Go
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")
|
|
}
|
|
}
|