@0dai-dev/cli 4.3.6 → 4.3.8

package/lib/python/usage_ledger.py

@@ -0,0 +1,569 @@
+#!/usr/bin/env python3
+"""0dai Usage Ledger — per-request cost accounting tied to billing plans.
+
+The ledger bridges three existing systems:
+1. model_usage_log.py — raw token counts per agent invocation
+2. swarm_cost.py — USD cost estimates per model
+3. tier_gate.py — daily swarm task quotas per billing plan
+4. billing_state.py — payment records and plan entitlements
+
+It produces:
+- ai/telemetry/usage-ledger.jsonl — append-only row per billing event
+- ai/telemetry/usage-ledger-daily.json — daily rollup (cost, tasks, by-model)
+- ai/telemetry/usage-ledger-monthly.json — monthly rollup for billing cycles
+
+CLI:
+    python3 usage_ledger.py record --task-id <id> --agent <agent> --model <model> \
+        --tokens-input <n> --tokens-output <n> [--plan <plan>]
+    python3 usage_ledger.py daily [--date YYYY-MM-DD] [--json]
+    python3 usage_ledger.py monthly [--month YYYY-MM] [--json]
+    python3 usage_ledger.py status [--plan <plan>] [--json]
+"""
+from __future__ import annotations
+
+import argparse
+import datetime as _dt
+import hashlib
+import json
+import pathlib
+import sys
+import time
+from dataclasses import dataclass, asdict
+
+SCRIPTS_DIR = pathlib.Path(__file__).resolve().parent
+sys.path.insert(0, str(SCRIPTS_DIR))
+
+import swarm_cost as sc  # noqa: E402
+
+LEDGER_DIR = pathlib.Path("ai/telemetry")
+LEDGER_FILE = LEDGER_DIR / "usage-ledger.jsonl"
+LEDGER_DAILY = LEDGER_DIR / "usage-ledger-daily.json"
+LEDGER_MONTHLY = LEDGER_DIR / "usage-ledger-monthly.json"
+
+# Free tier defaults from ai/work/strategy/free-tier-spec.md.
+FREE_TIER_LIMITS = {
+    "monthly_merges": 10,
+    "monthly_tokens": 100_000,
+    "concurrent_agents": 2,
+    "api_requests_per_hour": 60,
+    "storage_kg_entries": 500,
+    "team_members": 1,
+    "reset_policy": "calendar_month_utc",
+}
+
+# Plan entitlements (mirrors tier_gate.py PLAN_LIMITS for task budgets).
+PLAN_ENTITLEMENTS = {
+    "free": {"daily_tasks": 0, "monthly_budget_usd": 0.0, "max_model_tier": "fast", "free_tier": FREE_TIER_LIMITS},
+    "pro": {"daily_tasks": 50, "monthly_budget_usd": 15.0, "max_model_tier": "balanced"},
+    "team": {"daily_tasks": 200, "monthly_budget_usd": 49.0, "max_model_tier": "deep"},
+    "enterprise": {"daily_tasks": 999999, "monthly_budget_usd": 999999.0, "max_model_tier": "deep"},
+}
+
+# Model tier mapping
+MODEL_TIERS = {
+    "fast": {"haiku", "gpt-5.4-mini", "gemini-1.5-flash"},
+    "balanced": {"sonnet", "gpt-5.3", "gpt-5.4", "gemini-3", "gemini-2.5-pro"},
+    "deep": {"opus", "gpt-5.5", "gpt-5.3-codex", "claude-opus-4.7", "gemini-3.1-pro"},
+}
+
+# Reverse lookup: model name -> tier
+_MODEL_TO_TIER: dict[str, str] = {}
+for tier, models in MODEL_TIERS.items():
+    for m in models:
+        _MODEL_TO_TIER[m] = tier
+
+
+def _get_model_tier(model: str) -> str:
+    """Return the cost tier for a model name."""
+    # Normalize
+    normalized = model.lower().strip()
+    if normalized in _MODEL_TO_TIER:
+        return _MODEL_TO_TIER[normalized]
+    # Fallback: check rate aliases
+    for tier, models in MODEL_TIERS.items():
+        for m in models:
+            if m in normalized or normalized in m:
+                return tier
+    return "balanced"  # default assumption
+
+
+@dataclass
+class LedgerEntry:
+    """A single billing event."""
+    ts: str                     # ISO-8601 UTC
+    task_id: str                # Swarm task ID or loop run ID
+    agent: str                  # claude, codex, gemini, loop, etc.
+    model: str                  # Model name
+    model_tier: str             # fast, balanced, deep
+    tokens_input: int
+    tokens_output: int
+    tokens_total: int
+    cost_usd: float             # Computed from swarm_cost rates
+    plan: str                   # free, pro, team, enterprise
+    auth_mode: str              # api-key, subscription
+    source: str                 # swarm, loop, manual, adhoc
+    source_hash: str            # Dedup hash
+
+
+def _load_ledger() -> list[dict]:
+    if not LEDGER_FILE.is_file():
+        return []
+    entries = []
+    with LEDGER_FILE.open() as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                try:
+                    entries.append(json.loads(line))
+                except json.JSONDecodeError:
+                    continue
+    return entries
+
+
+def _append_entry(entry: dict) -> None:
+    LEDGER_DIR.mkdir(parents=True, exist_ok=True)
+    with LEDGER_FILE.open("a") as f:
+        f.write(json.dumps(entry) + "\n")
+
+
+def _source_hash(task_id: str, agent: str, ts: str) -> str:
+    """Short SHA-256 for dedup."""
+    return hashlib.sha256(f"{task_id}:{agent}:{ts}".encode()).hexdigest()[:12]
+
+
+def compute_cost(model: str, tokens_input: int, tokens_output: int) -> float:
+    """Compute USD cost using swarm_cost rates."""
+    try:
+        rate = sc.DEFAULT_RATES.get(model)
+        if rate:
+            input_rate, output_rate = rate
+            cost = (tokens_input / 1_000_000) * input_rate + (tokens_output / 1_000_000) * output_rate
+            return round(cost, 6)
+        # Fallback: estimate using sonnet rates
+        input_rate, output_rate = sc.DEFAULT_RATES.get("sonnet", (3.0, 15.0))
+        cost = (tokens_input / 1_000_000) * input_rate + (tokens_output / 1_000_000) * output_rate
+        return round(cost, 6)
+    except Exception:  # noqa: BLE001 — cost estimation failed
+        return 0.0
+
+
+def record_event(
+    task_id: str,
+    agent: str,
+    model: str,
+    tokens_input: int,
+    tokens_output: int,
+    plan: str = "pro",
+    auth_mode: str = "subscription",
+    source: str = "swarm",
+) -> LedgerEntry:
+    """Record a usage event in the ledger."""
+    ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+    model_tier = _get_model_tier(model)
+    cost = compute_cost(model, tokens_input, tokens_output)
+    tokens_total = tokens_input + tokens_output
+    shash = _source_hash(task_id, agent, ts)
+
+    entry = LedgerEntry(
+        ts=ts,
+        task_id=task_id,
+        agent=agent,
+        model=model,
+        model_tier=model_tier,
+        tokens_input=tokens_input,
+        tokens_output=tokens_output,
+        tokens_total=tokens_total,
+        cost_usd=cost,
+        plan=plan,
+        auth_mode=auth_mode,
+        source=source,
+        source_hash=shash,
+    )
+    entry_dict = asdict(entry)
+    _append_entry(entry_dict)
+    _update_daily_rollup(entry_dict)
+    _update_monthly_rollup(entry_dict)
+    return entry
+
+
+def _update_daily_rollup(entry: dict) -> None:
+    """Update the daily rollup file."""
+    daily = _load_daily_rollup()
+    date = entry["ts"][:10]  # YYYY-MM-DD
+    day = daily.setdefault(date, {
+        "date": date,
+        "total_cost_usd": 0.0,
+        "total_tasks": 0,
+        "total_tokens": 0,
+        "by_model": {},
+        "by_agent": {},
+        "by_plan": {},
+        "by_source": {},
+    })
+    day["total_cost_usd"] = round(day["total_cost_usd"] + entry["cost_usd"], 6)
+    day["total_tasks"] += 1
+    day["total_tokens"] += entry["tokens_total"]
+
+    # By model
+    model_stats = day["by_model"].setdefault(entry["model"], {"tasks": 0, "cost_usd": 0.0, "tokens": 0})
+    model_stats["tasks"] += 1
+    model_stats["cost_usd"] = round(model_stats["cost_usd"] + entry["cost_usd"], 6)
+    model_stats["tokens"] += entry["tokens_total"]
+
+    # By agent
+    agent_stats = day["by_agent"].setdefault(entry["agent"], {"tasks": 0, "cost_usd": 0.0})
+    agent_stats["tasks"] += 1
+    agent_stats["cost_usd"] = round(agent_stats["cost_usd"] + entry["cost_usd"], 6)
+
+    # By plan
+    plan_stats = day["by_plan"].setdefault(entry["plan"], {"tasks": 0, "cost_usd": 0.0})
+    plan_stats["tasks"] += 1
+    plan_stats["cost_usd"] = round(plan_stats["cost_usd"] + entry["cost_usd"], 6)
+
+    # By source
+    source_stats = day["by_source"].setdefault(entry["source"], {"tasks": 0, "cost_usd": 0.0, "tokens": 0})
+    source_stats["tasks"] += 1
+    source_stats["cost_usd"] = round(source_stats["cost_usd"] + entry["cost_usd"], 6)
+    source_stats["tokens"] += entry["tokens_total"]
+
+    _save_daily_rollup(daily)
+
+
+def _update_monthly_rollup(entry: dict) -> None:
+    """Update the monthly rollup file."""
+    monthly = _load_monthly_rollup()
+    month = entry["ts"][:7]  # YYYY-MM
+    mon = monthly.setdefault(month, {
+        "month": month,
+        "total_cost_usd": 0.0,
+        "total_tasks": 0,
+        "total_tokens": 0,
+        "daily_count": 0,
+        "avg_daily_cost_usd": 0.0,
+        "by_model": {},
+        "by_plan": {},
+        "by_source": {},
+    })
+    mon["total_cost_usd"] = round(mon["total_cost_usd"] + entry["cost_usd"], 6)
+    mon["total_tasks"] += 1
+    mon["total_tokens"] += entry["tokens_total"]
+    mon["daily_count"] = len(set(e[:10] for e in _get_daily_keys()))
+
+    # Recompute averages
+    days_in_month = max(1, mon["daily_count"])
+    mon["avg_daily_cost_usd"] = round(mon["total_cost_usd"] / days_in_month, 6)
+
+    # By model
+    model_stats = mon["by_model"].setdefault(entry["model"], {"tasks": 0, "cost_usd": 0.0, "tokens": 0})
+    model_stats["tasks"] += 1
+    model_stats["cost_usd"] = round(model_stats["cost_usd"] + entry["cost_usd"], 6)
+    model_stats["tokens"] += entry["tokens_total"]
+
+    # By plan
+    plan_stats = mon["by_plan"].setdefault(entry["plan"], {"tasks": 0, "cost_usd": 0.0})
+    plan_stats["tasks"] += 1
+    plan_stats["cost_usd"] = round(plan_stats["cost_usd"] + entry["cost_usd"], 6)
+
+    # By source
+    source_stats = mon["by_source"].setdefault(entry["source"], {"tasks": 0, "cost_usd": 0.0, "tokens": 0})
+    source_stats["tasks"] += 1
+    source_stats["cost_usd"] = round(source_stats["cost_usd"] + entry["cost_usd"], 6)
+    source_stats["tokens"] += entry["tokens_total"]
+
+    _save_monthly_rollup(monthly)
+
+
+def _load_daily_rollup() -> dict:
+    if not LEDGER_DAILY.is_file():
+        return {}
+    try:
+        return json.loads(LEDGER_DAILY.read_text("utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def _save_daily_rollup(data: dict) -> None:
+    LEDGER_DIR.mkdir(parents=True, exist_ok=True)
+    LEDGER_DAILY.write_text(json.dumps(data, indent=2), "utf-8")
+
+
+def _load_monthly_rollup() -> dict:
+    if not LEDGER_MONTHLY.is_file():
+        return {}
+    try:
+        return json.loads(LEDGER_MONTHLY.read_text("utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def _save_monthly_rollup(data: dict) -> None:
+    LEDGER_DIR.mkdir(parents=True, exist_ok=True)
+    LEDGER_MONTHLY.write_text(json.dumps(data, indent=2), "utf-8")
+
+
+def _get_daily_keys() -> list[str]:
+    daily = _load_daily_rollup()
+    return list(daily.keys())
+
+
+def get_daily_summary(date: str | None = None) -> dict:
+    """Get daily usage summary."""
+    if date is None:
+        date = _dt.date.today().isoformat()
+    daily = _load_daily_rollup()
+    return daily.get(date, {"date": date, "total_cost_usd": 0.0, "total_tasks": 0, "total_tokens": 0, "by_model": {}, "by_agent": {}, "by_plan": {}, "by_source": {}})
+
+
+def get_monthly_summary(month: str | None = None) -> dict:
+    """Get monthly usage summary."""
+    if month is None:
+        month = _dt.date.today().strftime("%Y-%m")
+    monthly = _load_monthly_rollup()
+    return monthly.get(month, {"month": month, "total_cost_usd": 0.0, "total_tasks": 0, "total_tokens": 0, "daily_count": 0, "avg_daily_cost_usd": 0.0, "by_model": {}, "by_plan": {}, "by_source": {}})
+
+
+def _next_month_reset(today: _dt.date) -> _dt.date:
+    """Return the next calendar-month UTC reset date."""
+    if today.month == 12:
+        return _dt.date(today.year + 1, 1, 1)
+    return _dt.date(today.year, today.month + 1, 1)
+
+
+def _free_tier_meter(monthly: dict, today: _dt.date) -> dict:
+    """Build the free-tier meter from monthly rollup data."""
+    limits = FREE_TIER_LIMITS.copy()
+    by_source = monthly.get("by_source", {})
+    merge_stats = by_source.get("merge", {})
+    merges_used = int(merge_stats.get("tasks", 0))
+    tokens_used = int(monthly.get("total_tokens", 0))
+    reset_at = _next_month_reset(today)
+    return {
+        "limits": limits,
+        "usage": {
+            "merges_used": merges_used,
+            "tokens_used": tokens_used,
+        },
+        "remaining": {
+            "merges": max(0, limits["monthly_merges"] - merges_used),
+            "tokens": max(0, limits["monthly_tokens"] - tokens_used),
+        },
+        "reset": {
+            "policy": limits["reset_policy"],
+            "reset_at": reset_at.isoformat(),
+            "days_remaining": max(0, (reset_at - today).days),
+        },
+        "cap_reached": merges_used >= limits["monthly_merges"] or tokens_used >= limits["monthly_tokens"],
+    }
+
+
+def get_plan_status(plan: str) -> dict:
+    """Get current plan usage status with entitlements."""
+    entitlements = PLAN_ENTITLEMENTS.get(plan, PLAN_ENTITLEMENTS["pro"])
+    today = _dt.date.today().isoformat()
+    month = today[:7]
+    daily = get_daily_summary(today)
+    monthly = get_monthly_summary(month)
+    status = {
+        "plan": plan,
+        "entitlements": entitlements,
+        "today": {
+            "date": today,
+            "tasks_used": daily.get("total_tasks", 0),
+            "tasks_remaining": max(0, entitlements["daily_tasks"] - daily.get("total_tasks", 0)),
+            "cost_usd": daily.get("total_cost_usd", 0.0),
+        },
+        "this_month": {
+            "month": month,
+            "tasks_used": monthly.get("total_tasks", 0),
+            "cost_usd": monthly.get("total_cost_usd", 0.0),
+            "budget_remaining": round(max(0, entitlements["monthly_budget_usd"] - monthly.get("total_cost_usd", 0.0)), 2),
+        },
+    }
+    if plan == "free":
+        status["free_tier"] = _free_tier_meter(monthly, _dt.date.fromisoformat(today))
+    return status
+
+
+def check_budget_alerts(plan: str) -> list[dict]:
+    """Check if usage is approaching budget limits. Returns alerts if any."""
+    status = get_plan_status(plan)
+    alerts = []
+
+    # Daily task usage
+    daily_limit = status["entitlements"]["daily_tasks"]
+    daily_pct = status["today"]["tasks_used"] / max(1, daily_limit)
+    if daily_limit > 0 and daily_pct >= 0.9:
+        alerts.append({
+            "severity": "warning" if daily_pct < 1.0 else "critical",
+            "type": "daily_tasks_near_limit",
+            "message": f"Daily tasks at {daily_pct * 100:.0f}% ({status['today']['tasks_used']}/{status['entitlements']['daily_tasks']})",
+        })
+
+    # Monthly budget
+    monthly_budget = status["entitlements"]["monthly_budget_usd"]
+    monthly_pct = status["this_month"]["cost_usd"] / max(0.01, monthly_budget)
+    if monthly_budget > 0 and monthly_pct >= 0.8:
+        alerts.append({
+            "severity": "warning" if monthly_pct < 1.0 else "critical",
+            "type": "monthly_budget_near_limit",
+            "message": f"Monthly budget at {monthly_pct * 100:.0f}% (${status['this_month']['cost_usd']:.2f}/${status['entitlements']['monthly_budget_usd']:.2f})",
+        })
+
+    free_tier = status.get("free_tier")
+    if free_tier:
+        limits = free_tier["limits"]
+        usage = free_tier["usage"]
+        merge_pct = usage["merges_used"] / max(1, limits["monthly_merges"])
+        token_pct = usage["tokens_used"] / max(1, limits["monthly_tokens"])
+        if merge_pct >= 0.8:
+            alerts.append({
+                "severity": "warning" if merge_pct < 1.0 else "critical",
+                "type": "free_merges_near_limit",
+                "message": f"Free tier merges at {merge_pct * 100:.0f}% ({usage['merges_used']}/{limits['monthly_merges']})",
+            })
+        if token_pct >= 0.8:
+            alerts.append({
+                "severity": "warning" if token_pct < 1.0 else "critical",
+                "type": "free_tokens_near_limit",
+                "message": f"Free tier tokens at {token_pct * 100:.0f}% ({usage['tokens_used']}/{limits['monthly_tokens']})",
+            })
+
+    return alerts
+
+
+def check_free_tier_merge_allowed(plan: str) -> tuple[bool, str]:
+    """Return (allowed, message). Blocks when free tier merge/token cap reached."""
+    if plan != "free":
+        return True, ""
+    status = get_plan_status("free")
+    free_tier = status.get("free_tier") or {}
+    if free_tier.get("cap_reached"):
+        usage = free_tier.get("usage", {})
+        limits = free_tier.get("limits", {})
+        reset = free_tier.get("reset", {})
+        merges_used = usage.get("merges_used", "?")
+        merge_limit = limits.get("monthly_merges", 10)
+        days = reset.get("days_remaining", "?")
+        return False, (
+            f"BLOCKED: free tier cap reached ({merges_used}/{merge_limit} merges). "
+            f"Reset in {days}d. Upgrade: 0dai upgrade pro"
+        )
+    remaining = free_tier.get("remaining", {}).get("merges", 0)
+    if remaining == 1:
+        return True, "WARNING: 1 merge left this month on free tier"
+    return True, ""
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = parser.add_subparsers(dest="command")
+
+    # record
+    p_record = sub.add_parser("record", help="Record a usage event")
+    p_record.add_argument("--task-id", required=True)
+    p_record.add_argument("--agent", required=True)
+    p_record.add_argument("--model", required=True)
+    p_record.add_argument("--tokens-input", type=int, required=True)
+    p_record.add_argument("--tokens-output", type=int, required=True)
+    p_record.add_argument("--plan", default="pro")
+    p_record.add_argument("--auth-mode", default="subscription")
+    p_record.add_argument("--source", default="swarm")
+    p_record.add_argument("--json", dest="as_json", action="store_true")
+
+    # daily
+    p_daily = sub.add_parser("daily", help="Show daily usage summary")
+    p_daily.add_argument("--date", help="YYYY-MM-DD (default today)")
+    p_daily.add_argument("--json", dest="as_json", action="store_true")
+
+    # monthly
+    p_monthly = sub.add_parser("monthly", help="Show monthly usage summary")
+    p_monthly.add_argument("--month", help="YYYY-MM (default current month)")
+    p_monthly.add_argument("--json", dest="as_json", action="store_true")
+
+    # status
+    p_status = sub.add_parser("status", help="Show plan status and budget alerts")
+    p_status.add_argument("--plan", default="pro")
+    p_status.add_argument("--json", dest="as_json", action="store_true")
+
+    p_check = sub.add_parser("check-merge", help="Exit 0 if plan may perform a merge")
+    p_check.add_argument("--plan", default="pro")
+
+    args = parser.parse_args()
+
+    if args.command == "check-merge":
+        allowed, message = check_free_tier_merge_allowed(args.plan)
+        if message:
+            print(message, file=sys.stderr if not allowed else sys.stdout)
+        return 0 if allowed else 2
+
+    if args.command == "record":
+        entry = record_event(
+            task_id=args.task_id,
+            agent=args.agent,
+            model=args.model,
+            tokens_input=args.tokens_input,
+            tokens_output=args.tokens_output,
+            plan=args.plan,
+            auth_mode=args.auth_mode,
+            source=args.source,
+        )
+        if args.as_json:
+            print(json.dumps(asdict(entry), indent=2))
+        else:
+            print(f"[recorded] {entry.task_id} — {entry.model} ({entry.model_tier}) — {entry.tokens_total:,} tokens — ${entry.cost_usd:.6f}")
+
+    elif args.command == "daily":
+        summary = get_daily_summary(getattr(args, "date", None))
+        if args.as_json:
+            print(json.dumps(summary, indent=2))
+        else:
+            print(f"📅 {summary['date']} — {summary.get('total_tasks', 0)} tasks — ${summary.get('total_cost_usd', 0):.6f} — {summary.get('total_tokens', 0):,} tokens")
+            for model, stats in summary.get("by_model", {}).items():
+                print(f"  {model}: {stats['tasks']} tasks — ${stats['cost_usd']:.6f}")
+
+    elif args.command == "monthly":
+        summary = get_monthly_summary(getattr(args, "month", None))
+        if args.as_json:
+            print(json.dumps(summary, indent=2))
+        else:
+            print(f"📆 {summary['month']} — {summary.get('total_tasks', 0)} tasks — ${summary.get('total_cost_usd', 0):.6f} — avg ${summary.get('avg_daily_cost_usd', 0):.6f}/day")
+            for model, stats in summary.get("by_model", {}).items():
+                print(f"  {model}: {stats['tasks']} tasks — ${stats['cost_usd']:.6f}")
+
+    elif args.command == "status":
+        status = get_plan_status(args.plan)
+        alerts = check_budget_alerts(args.plan)
+        output = {"status": status, "alerts": alerts}
+        if args.as_json:
+            print(json.dumps(output, indent=2))
+        else:
+            print(f"Plan: {status['plan']}")
+            if status["entitlements"]["daily_tasks"] > 0:
+                print(f"  Today: {status['today']['tasks_used']}/{status['entitlements']['daily_tasks']} tasks — ${status['today']['cost_usd']:.6f}")
+            else:
+                print(f"  Today: {status['today']['tasks_used']} tasks — ${status['today']['cost_usd']:.6f}")
+            print(f"  This month: {status['this_month']['tasks_used']} tasks — ${status['this_month']['cost_usd']:.2f}/${status['entitlements']['monthly_budget_usd']:.2f}")
+            if "free_tier" in status:
+                ft = status["free_tier"]
+                print(
+                    "  Free tier: "
+                    f"{ft['remaining']['merges']}/{ft['limits']['monthly_merges']} merges remaining, "
+                    f"{ft['remaining']['tokens']:,}/{ft['limits']['monthly_tokens']:,} tokens remaining"
+                )
+                print(f"  Reset: {ft['reset']['reset_at']} UTC ({ft['reset']['days_remaining']} days)")
+            if alerts:
+                print("  ⚠️  Alerts:")
+                for a in alerts:
+                    print(f"    [{a['severity']}] {a['message']}")
+            else:
+                print("  ✅ Within budget")
+
+    else:
+        parser.print_help()
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())