add feature for weekly and monthly quotas

Author: thedavidwhiteside

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+*.pyo

README.md

@@ -1,17 +1,17 @@
 [![Tests](https://github.com/thedavidwhiteside/claude-code-tokenbudget/actions/workflows/test.yml/badge.svg?branch=main)](https://github.com/thedavidwhiteside/claude-code-tokenbudget/actions/workflows/test.yml)
 
-# Claude Code Daily Token Quota Plugin
+# Claude Code Token Quota Plugin
 
-Claude Code has no built-in spending guardrails. This plugin tracks your daily token usage and hard stops new prompts once you hit your limit. It works with **any backend**: Bedrock, Vertex, direct API, or subscription.
+Claude Code has no built-in spending guardrails. This plugin tracks your token usage and hard stops new prompts once you hit your limit. Supports daily, weekly (rolling 7-day), and monthly (calendar month) quotas. It works with **any backend**: Bedrock, Vertex, direct API, or subscription.
 
 ## How it works
 
 | Hook | Event | Action |
 |------|-------|--------|
-| `enforce_quota.py` | `UserPromptSubmit` | Blocks the prompt if today's usage ≥ limit |
+| `enforce_quota.py` | `UserPromptSubmit` | Blocks the prompt if any quota (daily/weekly/monthly) is exceeded |
 | `track_tokens.py` | `Stop` | Records token usage after each turn |
 
-Usage is stored in `~/.claude-token-quota/YYYY-MM-DD.json` and resets automatically each day.
+Usage is stored in `~/.claude-token-quota/YYYY-MM-DD.json` per day. Daily quotas reset at midnight; the weekly quota uses a rolling 7-day window; the monthly quota resets on the 1st.
 
 ---
 
@@ -50,6 +50,8 @@ Override any of these in your `~/.claude/settings.json`:
 {
   "env": {
     "TOKEN_QUOTA_DAILY": "1000000",
+    "TOKEN_QUOTA_WEEKLY": "5000000",
+    "TOKEN_QUOTA_MONTHLY": "15000000",
     "TOKEN_QUOTA_DIR": "~/.claude-token-quota",
     "TOKEN_QUOTA_RETAIN_DAYS": "30"
   }
@@ -59,9 +61,13 @@ Override any of these in your `~/.claude/settings.json`:
 | Variable | Default | Description |
 |---|---|---|
 | `TOKEN_QUOTA_DAILY` | `1000000` | Daily token limit |
+| `TOKEN_QUOTA_WEEKLY` | _(unset)_ | Rolling 7-day token limit (optional) |
+| `TOKEN_QUOTA_MONTHLY` | _(unset)_ | Calendar-month token limit (optional) |
 | `TOKEN_QUOTA_DIR` | `~/.claude-token-quota` | Where ledger files are stored |
 | `TOKEN_QUOTA_RETAIN_DAYS` | `30` | How many days of usage history to keep |
 
+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.
+
 **Rough token budgets by spend goal — AWS Bedrock example (Claude Sonnet 4.6):**
 
 > **Note:** Prices below are AWS Bedrock examples only and will change. For current rates check the [AWS Bedrock pricing page](https://aws.amazon.com/bedrock/pricing/). Direct API users: see the [Anthropic pricing page](https://www.anthropic.com/pricing) for your model's rates, then apply the same blended-cost formula below.

hooks/enforce_quota.py

@@ -1,27 +1,33 @@
 #!/usr/bin/env python3
 """
 enforce_quota.py — UserPromptSubmit hook
-Reads today's token ledger and blocks the prompt if the daily limit is exceeded.
+Reads token ledgers and blocks the prompt if any quota is exceeded.
 
-Set your daily limit via environment variable:
-  TOKEN_QUOTA_DAILY=500000   (default: 500,000 tokens)
+Set limits via environment variables:
+  TOKEN_QUOTA_DAILY=1000000    (default: 1,000,000 tokens)
+  TOKEN_QUOTA_WEEKLY=5000000   (optional; rolling 7-day window)
+  TOKEN_QUOTA_MONTHLY=15000000 (optional; current calendar month)
   TOKEN_QUOTA_DIR=~/.claude-token-quota  (default ledger location)
 """
 
 import json
 import sys
 import os
-from datetime import date
+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))
 
-def today_ledger() -> Path:
-    return LEDGER_DIR / f"{date.today().isoformat()}.json"
+_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 = today_ledger()
+    p = LEDGER_DIR / f"{date.today().isoformat()}.json"
     if not p.exists():
         return 0
     try:
@@ -30,39 +36,86 @@ def get_today_total() -> int:
     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():
-    # Read the hook input (we don't need it, but consume stdin cleanly)
     try:
         sys.stdin.read()
     except Exception:
         pass
 
-    used = get_today_total()
-    remaining = DAILY_LIMIT - used
-    pct = (used / DAILY_LIMIT) * 100 if DAILY_LIMIT > 0 else 0
+    today = date.today()
+    used_daily = get_today_total()
+    pct_daily = (used_daily / DAILY_LIMIT) * 100 if DAILY_LIMIT > 0 else 0
 
-    if used >= DAILY_LIMIT:
-        # Block the prompt by returning a decision=block JSON
+    if used_daily >= DAILY_LIMIT:
         result = {
             "decision": "block",
             "reason": (
                 f"Daily token quota exceeded.\n"
-                f"Used:  {used:,} / {DAILY_LIMIT:,} tokens ({pct:.1f}%)\n"
+                f"Used:  {used_daily:,} / {DAILY_LIMIT:,} 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)
 
-    # Warn at 80% and 95%
-    if pct >= 95:
-        warning = f"Token quota at {pct:.1f}% ({used:,} / {DAILY_LIMIT:,}). Nearly exhausted."
+    if 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
+            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"Edit TOKEN_QUOTA_WEEKLY to change the limit."
+                )
+            }
+            print(json.dumps(result))
+            sys.exit(0)
+
+    if 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
+            result = {
+                "decision": "block",
+                "reason": (
+                    f"Monthly token quota exceeded.\n"
+                    f"Used:  {used_monthly:,} / {MONTHLY_LIMIT:,} 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
+    if pct_daily >= 95:
+        warning = f"Token quota at {pct_daily:.1f}% ({used_daily:,} / {DAILY_LIMIT:,}). Nearly exhausted."
         result = {"decision": "allow", "reason": warning}
         print(json.dumps(result))
-    elif pct >= 85:
-        print(f"[token-quota] {pct:.1f}% of daily quota used ({remaining:,} tokens remaining)", file=sys.stderr)
+    elif pct_daily >= 85:
+        print(f"[token-quota] {pct_daily:.1f}% of daily quota used ({remaining:,} tokens remaining)", file=sys.stderr)
 
     sys.exit(0)
 
+
 if __name__ == "__main__":
     main()

hooks/quota_status.py

@@ -2,49 +2,103 @@
 """
 quota_status.py — run manually to check today's token usage
 Usage:  python3 quota_status.py
-        TOKEN_QUOTA_DAILY=1000000 python3 quota_status.py
+        TOKEN_QUOTA_DAILY=1000000 TOKEN_QUOTA_WEEKLY=5000000 python3 quota_status.py
 """
 
 import json
 import os
-from datetime import date
+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))
 
-def main():
-    ledger_file = LEDGER_DIR / f"{date.today().isoformat()}.json"
+_weekly_raw = os.environ.get("TOKEN_QUOTA_WEEKLY")
+WEEKLY_LIMIT = int(_weekly_raw) if _weekly_raw else None
 
-    if not ledger_file.exists():
-        print(f"No usage recorded today ({date.today().isoformat()}).")
-        print(f"Daily limit: {DAILY_LIMIT:,} tokens")
-        return
+_monthly_raw = os.environ.get("TOKEN_QUOTA_MONTHLY")
+MONTHLY_LIMIT = int(_monthly_raw) if _monthly_raw else None
 
-    data = json.loads(ledger_file.read_text())
-    used = data.get("total_tokens", 0)
-    remaining = max(0, DAILY_LIMIT - used)
-    pct = (used / DAILY_LIMIT) * 100 if DAILY_LIMIT > 0 else 0
-    sessions = data.get("sessions", [])
 
-    bar_width = 30
-    filled = int(bar_width * pct / 100)
-    bar = "█" * filled + "░" * (bar_width - filled)
+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 _render_bar(used: int, limit: int) -> tuple[str, float, str]:
+    pct = (used / limit) * 100 if limit > 0 else 0
+    filled = min(30, int(30 * pct / 100))
+    bar = "█" * filled + "░" * (30 - filled)
     status = "✅ OK" if pct < 80 else ("⚠️  WARNING" if pct < 100 else "🚫 EXCEEDED")
+    return bar, pct, status
+
+
+def main():
+    today = date.today()
 
     print(f"\n{'─'*50}")
-    print(f"  Claude Code Token Quota — {date.today().isoformat()}")
+    print(f"  Claude Code Token Quota — {today.isoformat()}")
     print(f"{'─'*50}")
-    print(f"  [{bar}] {pct:.1f}%")
-    print(f"  Used:      {used:>12,} tokens")
-    print(f"  Remaining: {remaining:>12,} tokens")
-    print(f"  Limit:     {DAILY_LIMIT:>12,} tokens")
-    print(f"  Status:    {status}")
-    print(f"  Turns:     {len(sessions)}")
-    if sessions:
-        print(f"  Last turn: {sessions[-1]['timestamp']}")
-    print(f"{'─'*50}\n")
+
+    ledger_file = LEDGER_DIR / f"{today.isoformat()}.json"
+    if not ledger_file.exists():
+        print(f"  No usage recorded today ({today.isoformat()}).")
+        print(f"  Daily limit: {DAILY_LIMIT:,} tokens")
+        sessions = []
+    else:
+        data = json.loads(ledger_file.read_text())
+        used = data.get("total_tokens", 0)
+        remaining = max(0, DAILY_LIMIT - used)
+        sessions = data.get("sessions", [])
+        bar, pct, status = _render_bar(used, DAILY_LIMIT)
+
+        print(f"  [{bar}] {pct:.1f}%")
+        print(f"  Used:      {used:>12,} tokens")
+        print(f"  Remaining: {remaining:>12,} tokens")
+        print(f"  Limit:     {DAILY_LIMIT:>12,} tokens")
+        print(f"  Status:    {status}")
+        print(f"  Turns:     {len(sessions)}")
+        if sessions:
+            print(f"  Last turn: {sessions[-1]['timestamp']}")
+
+    if 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)
+
+        print(f"\n  Weekly (rolling 7-day: {week_start.isoformat()} – {today.isoformat()})")
+        print(f"  [{bar}] {pct:.1f}%")
+        print(f"  Used:      {used_weekly:>12,} tokens")
+        print(f"  Remaining: {remaining_weekly:>12,} tokens")
+        print(f"  Limit:     {WEEKLY_LIMIT:>12,} tokens")
+        print(f"  Status:    {status}")
+
+    if 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)
+
+        print(f"\n  Monthly ({today.strftime('%B %Y')})")
+        print(f"  [{bar}] {pct:.1f}%")
+        print(f"  Used:      {used_monthly:>12,} tokens")
+        print(f"  Remaining: {remaining_monthly:>12,} tokens")
+        print(f"  Limit:     {MONTHLY_LIMIT:>12,} tokens")
+        print(f"  Status:    {status}")
+
+    print(f"\n{'─'*50}\n")
+
 
 if __name__ == "__main__":
     main()