refactor for DRY, setup a quora_store class to be shared across tools

Author: thedavidwhiteside

README.md

@@ -55,7 +55,8 @@ Override any of these in your `~/.claude/settings.json`:
     "TOKEN_QUOTA_DIR": "~/.claude-token-quota",
     "TOKEN_QUOTA_RETAIN_DAYS": "31",
     "TOKEN_QUOTA_WARN_CRITICAL": "95",
-    "TOKEN_QUOTA_WARN": "85"
+    "TOKEN_QUOTA_WARN": "85",
+    "TOKEN_QUOTA_SNOOZE_TOKENS": "1000000"
   }
 }
 ```
@@ -70,6 +71,7 @@ Override any of these in your `~/.claude/settings.json`:
 | `TOKEN_QUOTA_WARN_CRITICAL` | `95` | % of daily limit at which a visible warning is shown |
 | `TOKEN_QUOTA_WARN` | `85` | % of daily limit at which a stderr warning is shown |
 | `TOKEN_QUOTA_COST_PER_M` | _(unset)_ | Blended cost per 1M tokens — enables `~$X.XX` estimates in status output |
+| `TOKEN_QUOTA_SNOOZE_TOKENS` | `1000000` | Extra tokens added to all limits when `/tokenbudget:snooze` is run |
 
 Weekly and monthly limits are opt-in — omit them to enforce only the daily limit. When multiple limits are set, any one being exceeded blocks new prompts.
 

hooks/enforce_quota.py

@@ -10,6 +10,8 @@
   TOKEN_QUOTA_DIR=~/.claude-token-quota  (default ledger location)
   TOKEN_QUOTA_WARN_CRITICAL=95 (default: warn loudly at 95% of daily limit)
   TOKEN_QUOTA_WARN=85          (default: warn quietly at 85% of daily limit)
+  TOKEN_QUOTA_COST_PER_M=5.40  (optional; blended cost per 1M tokens for dollar estimates)
+  TOKEN_QUOTA_SNOOZE_TOKENS=1000000 (tokens added when snoozed; default 1,000,000)
 """
 
 import json
@@ -18,101 +20,76 @@
 from datetime import date, timedelta
 from pathlib import Path
 
-LEDGER_DIR = Path(os.environ.get("TOKEN_QUOTA_DIR", Path.home() / ".claude-token-quota"))
-DAILY_LIMIT = int(os.environ.get("TOKEN_QUOTA_DAILY", 1_000_000))
+sys.path.insert(0, str(Path(__file__).parent))
+from quota_store import QuotaStore
+
 WARN_CRITICAL = int(os.environ.get("TOKEN_QUOTA_WARN_CRITICAL", 95))
 WARN = int(os.environ.get("TOKEN_QUOTA_WARN", 85))
 
-_weekly_raw = os.environ.get("TOKEN_QUOTA_WEEKLY")
-WEEKLY_LIMIT = int(_weekly_raw) if _weekly_raw else None
-
-_monthly_raw = os.environ.get("TOKEN_QUOTA_MONTHLY")
-MONTHLY_LIMIT = int(_monthly_raw) if _monthly_raw else None
-
-
-def get_today_total() -> int:
-    p = LEDGER_DIR / f"{date.today().isoformat()}.json"
-    if not p.exists():
-        return 0
-    try:
-        data = json.loads(p.read_text())
-        return data.get("total_tokens", 0)
-    except Exception:
-        return 0
-
-
-def get_period_total(start: date, end: date) -> int:
-    total = 0
-    current = start
-    while current <= end:
-        p = LEDGER_DIR / f"{current.isoformat()}.json"
-        if p.exists():
-            try:
-                data = json.loads(p.read_text())
-                total += data.get("total_tokens", 0)
-            except Exception:
-                pass
-        current += timedelta(days=1)
-    return total
-
 
 def main():
     try:
         sys.stdin.read()
     except Exception:
         pass
 
+    store = QuotaStore()
     today = date.today()
-    used_daily = get_today_total()
-    pct_daily = (used_daily / DAILY_LIMIT) * 100 if DAILY_LIMIT > 0 else 0
+    snooze = store.get_snooze()
+
+    used_daily = store.get_today_total()
+    effective_daily = store.daily_limit + snooze
+    pct_daily = (used_daily / effective_daily) * 100 if effective_daily > 0 else 0
 
-    if used_daily >= DAILY_LIMIT:
+    if used_daily >= effective_daily:
         result = {
             "decision": "block",
             "reason": (
                 f"Daily token quota exceeded.\n"
-                f"Used:  {used_daily:,} / {DAILY_LIMIT:,} tokens ({pct_daily:.1f}%)\n"
+                f"Used:  {used_daily:,} / {effective_daily:,} tokens ({pct_daily:.1f}%)\n"
                 f"Quota resets at midnight. Edit TOKEN_QUOTA_DAILY to change the limit."
             )
         }
         print(json.dumps(result))
         sys.exit(0)
 
-    if WEEKLY_LIMIT is not None:
+    if store.weekly_limit is not None:
         week_start = today - timedelta(days=6)
-        used_weekly = get_period_total(week_start, today)
-        if used_weekly >= WEEKLY_LIMIT:
-            pct = (used_weekly / WEEKLY_LIMIT) * 100
+        used_weekly = store.get_period_total(week_start, today)
+        effective_weekly = store.weekly_limit + snooze
+        if used_weekly >= effective_weekly:
+            pct = (used_weekly / effective_weekly) * 100
             result = {
                 "decision": "block",
                 "reason": (
                     f"Weekly token quota exceeded (rolling 7-day window).\n"
-                    f"Used:  {used_weekly:,} / {WEEKLY_LIMIT:,} tokens ({pct:.1f}%)\n"
+                    f"Used:  {used_weekly:,} / {effective_weekly:,} tokens ({pct:.1f}%)\n"
                     f"Edit TOKEN_QUOTA_WEEKLY to change the limit."
                 )
             }
             print(json.dumps(result))
             sys.exit(0)
 
-    if MONTHLY_LIMIT is not None:
+    if store.monthly_limit is not None:
         month_start = today.replace(day=1)
-        used_monthly = get_period_total(month_start, today)
-        if used_monthly >= MONTHLY_LIMIT:
-            pct = (used_monthly / MONTHLY_LIMIT) * 100
+        used_monthly = store.get_period_total(month_start, today)
+        effective_monthly = store.monthly_limit + snooze
+        if used_monthly >= effective_monthly:
+            pct = (used_monthly / effective_monthly) * 100
             result = {
                 "decision": "block",
                 "reason": (
                     f"Monthly token quota exceeded.\n"
-                    f"Used:  {used_monthly:,} / {MONTHLY_LIMIT:,} tokens ({pct:.1f}%)\n"
+                    f"Used:  {used_monthly:,} / {effective_monthly:,} tokens ({pct:.1f}%)\n"
                     f"Quota resets on the 1st. Edit TOKEN_QUOTA_MONTHLY to change the limit."
                 )
             }
             print(json.dumps(result))
             sys.exit(0)
 
-    remaining = DAILY_LIMIT - used_daily
+    remaining = effective_daily - used_daily
     if pct_daily >= WARN_CRITICAL:
-        warning = f"Token quota at {pct_daily:.1f}% ({used_daily:,} / {DAILY_LIMIT:,}). Nearly exhausted."
+        warning = f"Token quota at {pct_daily:.1f}% ({used_daily:,} / {effective_daily:,}). Nearly exhausted."
         result = {"decision": "allow", "reason": warning}
         print(json.dumps(result))
     elif pct_daily >= WARN:

hooks/quota_status.py

@@ -12,39 +12,16 @@
 import os
 from datetime import date, timedelta
 from pathlib import Path
+import sys
 
-LEDGER_DIR = Path(os.environ.get("TOKEN_QUOTA_DIR", Path.home() / ".claude-token-quota"))
-DAILY_LIMIT = int(os.environ.get("TOKEN_QUOTA_DAILY", 1_000_000))
-
-_weekly_raw = os.environ.get("TOKEN_QUOTA_WEEKLY")
-WEEKLY_LIMIT = int(_weekly_raw) if _weekly_raw else None
-
-_monthly_raw = os.environ.get("TOKEN_QUOTA_MONTHLY")
-MONTHLY_LIMIT = int(_monthly_raw) if _monthly_raw else None
+sys.path.insert(0, str(Path(__file__).parent))
+from quota_store import QuotaStore
 
 _cost_raw = os.environ.get("TOKEN_QUOTA_COST_PER_M")
 COST_PER_M = float(_cost_raw) if _cost_raw else None
 
-RETAIN_DAYS = int(os.environ.get("TOKEN_QUOTA_RETAIN_DAYS", 31))
-
-
-def get_period_total(start: date, end: date) -> int:
-    total = 0
-    current = start
-    while current <= end:
-        p = LEDGER_DIR / f"{current.isoformat()}.json"
-        if p.exists():
-            try:
-                data = json.loads(p.read_text())
-                total += data.get("total_tokens", 0)
-            except Exception:
-                pass
-        current += timedelta(days=1)
-    return total
-
 
 def _cost_str(tokens: int) -> str:
-    """Return a formatted cost string, or empty string if cost tracking is off."""
     if COST_PER_M is None:
         return ""
     return f"  (~${tokens / 1_000_000 * COST_PER_M:.2f})"
@@ -59,62 +36,74 @@ def _render_bar(used: int, limit: int) -> tuple[str, float, str]:
 
 
 def main():
+    store = QuotaStore()
     today = date.today()
 
-    if MONTHLY_LIMIT is not None and RETAIN_DAYS < 31:
-        print(f"  ⚠️  WARNING: TOKEN_QUOTA_RETAIN_DAYS={RETAIN_DAYS} is below 31.")
+    if store.monthly_limit is not None and store.retain_days < 31:
+        print(f"  ⚠️  WARNING: TOKEN_QUOTA_RETAIN_DAYS={store.retain_days} is below 31.")
         print(f"  Monthly totals may be incomplete for 31-day months.")
         print(f"  Set TOKEN_QUOTA_RETAIN_DAYS=31 or higher to fix this.")
 
     print(f"\n{'─'*50}")
     print(f"  Claude Code Token Quota — {today.isoformat()}")
     print(f"{'─'*50}")
 
-    ledger_file = LEDGER_DIR / f"{today.isoformat()}.json"
+    snooze = store.get_snooze()
+
+    ledger_file = store.ledger_path()
     if not ledger_file.exists():
         print(f"  No usage recorded today ({today.isoformat()}).")
-        print(f"  Daily limit: {DAILY_LIMIT:,} tokens")
+        print(f"  Daily limit: {store.daily_limit:,} tokens")
         sessions = []
     else:
         data = json.loads(ledger_file.read_text())
         used = data.get("total_tokens", 0)
-        remaining = max(0, DAILY_LIMIT - used)
+        effective_daily = store.daily_limit + snooze
+        remaining = max(0, effective_daily - used)
         sessions = data.get("sessions", [])
-        bar, pct, status = _render_bar(used, DAILY_LIMIT)
+        bar, pct, status = _render_bar(used, effective_daily)
 
         print(f"  [{bar}] {pct:.1f}%")
         print(f"  Used:      {used:>12,} tokens{_cost_str(used)}")
         print(f"  Remaining: {remaining:>12,} tokens{_cost_str(remaining)}")
-        print(f"  Limit:     {DAILY_LIMIT:>12,} tokens{_cost_str(DAILY_LIMIT)}")
+        print(f"  Limit:     {store.daily_limit:>12,} tokens{_cost_str(store.daily_limit)}")
+        if snooze:
+            print(f"  Snooze:    {snooze:>12,} tokens  (expires midnight)")
         print(f"  Status:    {status}")
         print(f"  Turns:     {len(sessions)}")
         if sessions:
             print(f"  Last turn: {sessions[-1]['timestamp']}")
 
-    if WEEKLY_LIMIT is not None:
+    if store.weekly_limit is not None:
         week_start = today - timedelta(days=6)
-        used_weekly = get_period_total(week_start, today)
-        remaining_weekly = max(0, WEEKLY_LIMIT - used_weekly)
-        bar, pct, status = _render_bar(used_weekly, WEEKLY_LIMIT)
+        used_weekly = store.get_period_total(week_start, today)
+        effective_weekly = store.weekly_limit + snooze
+        remaining_weekly = max(0, effective_weekly - used_weekly)
+        bar, pct, status = _render_bar(used_weekly, effective_weekly)
 
         print(f"\n  Weekly (rolling 7-day: {week_start.isoformat()} – {today.isoformat()})")
         print(f"  [{bar}] {pct:.1f}%")
         print(f"  Used:      {used_weekly:>12,} tokens{_cost_str(used_weekly)}")
         print(f"  Remaining: {remaining_weekly:>12,} tokens{_cost_str(remaining_weekly)}")
-        print(f"  Limit:     {WEEKLY_LIMIT:>12,} tokens{_cost_str(WEEKLY_LIMIT)}")
+        print(f"  Limit:     {store.weekly_limit:>12,} tokens{_cost_str(store.weekly_limit)}")
+        if snooze:
+            print(f"  Snooze:    {snooze:>12,} tokens  (expires midnight)")
         print(f"  Status:    {status}")
 
-    if MONTHLY_LIMIT is not None:
+    if store.monthly_limit is not None:
         month_start = today.replace(day=1)
-        used_monthly = get_period_total(month_start, today)
-        remaining_monthly = max(0, MONTHLY_LIMIT - used_monthly)
-        bar, pct, status = _render_bar(used_monthly, MONTHLY_LIMIT)
+        used_monthly = store.get_period_total(month_start, today)
+        effective_monthly = store.monthly_limit + snooze
+        remaining_monthly = max(0, effective_monthly - used_monthly)
+        bar, pct, status = _render_bar(used_monthly, effective_monthly)
 
         print(f"\n  Monthly ({today.strftime('%B %Y')})")
         print(f"  [{bar}] {pct:.1f}%")
         print(f"  Used:      {used_monthly:>12,} tokens{_cost_str(used_monthly)}")
         print(f"  Remaining: {remaining_monthly:>12,} tokens{_cost_str(remaining_monthly)}")
-        print(f"  Limit:     {MONTHLY_LIMIT:>12,} tokens{_cost_str(MONTHLY_LIMIT)}")
+        print(f"  Limit:     {store.monthly_limit:>12,} tokens{_cost_str(store.monthly_limit)}")
+        if snooze:
+            print(f"  Snooze:    {snooze:>12,} tokens  (expires midnight)")
         print(f"  Status:    {status}")
 
     print(f"\n{'─'*50}\n")

hooks/quota_store.py

@@ -0,0 +1,105 @@
+"""
+quota_store.py — shared data access layer for all token-quota hooks.
+
+Encapsulates config parsing, ledger reads/writes, snooze state, and cleanup.
+Import via:
+    import sys
+    from pathlib import Path
+    sys.path.insert(0, str(Path(__file__).parent))
+    from quota_store import QuotaStore
+"""
+
+import json
+import os
+from datetime import date, timedelta
+from pathlib import Path
+
+
+class QuotaStore:
+    def __init__(self):
+        self.ledger_dir = Path(os.environ.get("TOKEN_QUOTA_DIR", Path.home() / ".claude-token-quota"))
+        self.daily_limit = int(os.environ.get("TOKEN_QUOTA_DAILY", 1_000_000))
+        self.retain_days = int(os.environ.get("TOKEN_QUOTA_RETAIN_DAYS", 31))
+        self.snooze_tokens = int(os.environ.get("TOKEN_QUOTA_SNOOZE_TOKENS", 1_000_000))
+
+        _weekly = os.environ.get("TOKEN_QUOTA_WEEKLY")
+        self.weekly_limit = int(_weekly) if _weekly else None
+
+        _monthly = os.environ.get("TOKEN_QUOTA_MONTHLY")
+        self.monthly_limit = int(_monthly) if _monthly else None
+
+    # ── Ledger paths ──────────────────────────────────────────────────────────
+
+    def ledger_path(self, for_date: date | None = None) -> Path:
+        return self.ledger_dir / f"{(for_date or date.today()).isoformat()}.json"
+
+    # ── Token reads ───────────────────────────────────────────────────────────
+
+    def get_today_total(self) -> int:
+        p = self.ledger_path()
+        if not p.exists():
+            return 0
+        try:
+            return json.loads(p.read_text()).get("total_tokens", 0)
+        except Exception:
+            return 0
+
+    def get_period_total(self, start: date, end: date) -> int:
+        total = 0
+        current = start
+        while current <= end:
+            p = self.ledger_path(current)
+            if p.exists():
+                try:
+                    total += json.loads(p.read_text()).get("total_tokens", 0)
+                except Exception:
+                    pass
+            current += timedelta(days=1)
+        return total
+
+    # ── Ledger read/write ─────────────────────────────────────────────────────
+
+    def load_ledger(self) -> dict:
+        p = self.ledger_path()
+        if p.exists():
+            try:
+                return json.loads(p.read_text())
+            except Exception:
+                pass
+        return {"date": date.today().isoformat(), "total_tokens": 0, "sessions": []}
+
+    def save_ledger(self, ledger: dict):
+        self.ledger_dir.mkdir(parents=True, exist_ok=True)
+        self.ledger_path().write_text(json.dumps(ledger, indent=2))
+
+    def cleanup_old_ledgers(self):
+        cutoff = date.today() - timedelta(days=self.retain_days)
+        for f in self.ledger_dir.glob("????-??-??.json"):
+            try:
+                if date.fromisoformat(f.stem) < cutoff:
+                    f.unlink()
+            except ValueError:
+                pass
+
+    # ── Snooze ────────────────────────────────────────────────────────────────
+
+    def get_snooze(self) -> int:
+        """Return extra tokens from an active snooze, or 0 if none is active."""
+        p = self.ledger_dir / "snooze.json"
+        if not p.exists():
+            return 0
+        try:
+            data = json.loads(p.read_text())
+            if data.get("expires") == date.today().isoformat():
+                return int(data.get("extra_tokens", 0))
+        except Exception:
+            pass
+        return 0
+
+    def write_snooze(self, extra_tokens: int):
+        self.ledger_dir.mkdir(parents=True, exist_ok=True)
+        snooze_file = self.ledger_dir / "snooze.json"
+        snooze_file.write_text(json.dumps({
+            "extra_tokens": extra_tokens,
+            "expires": date.today().isoformat(),
+        }, indent=2))