Files
company-haness/.claude/hooks/token_ledger.py
T

192 lines
7.9 KiB
Python

#!/usr/bin/env python3
"""Token ledger + budget gate (권고 #1).
Claude Code에서 서브에이전트가 끝나면 정확한 토큰 수를 Orchestrator에 반환한다. 그 수치를
append-only 원장에 적고(log), 대표용 대시보드로 렌더(dashboard)하며, tier 예산 초과를
게이트(check)한다. 강제(초과 시 collapse 강등)의 주체는 Orchestrator이고, 이 도구는 계측·게이트다.
Usage:
token_ledger.py log --workflow WF --role ROLE --tokens N [--wave V] [--tier T]
token_ledger.py dashboard # -> reports/TOKENS.md
token_ledger.py check --workflow WF --tier T [--wave V] [--add N] # 예산 초과면 exit 2
예산(tier)은 per-wave 다. check/dashboard 는 워크플로 전체가 아니라 해당 wave 만 대조한다
(finding #19). --wave 미지정이면 '-' wave 로 묶여 단일-wave 워크플로 동작이 보존된다.
"""
import json
import os
import sys
from datetime import datetime, timezone
import yaml
ROOT = os.environ.get("CLAUDE_PROJECT_DIR") or os.path.dirname(
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import _workspace as W # noqa: E402
_ORGWORK = os.path.join(ROOT, "org-os", "06-agent-work") # KPI 예산 = SSOT, org-os 유지
LEDGER = os.path.join(W.state_dir(), "token-ledger.jsonl")
KPI = os.path.join(_ORGWORK, "agent-operating-kpi.yaml")
DASH = os.path.join(W.reports_dir(), "TOKENS.md")
def budgets():
try:
d = yaml.safe_load(open(KPI))["agent-operating-kpi"]["token-budgets"]
return d.get("per-wave", {}), float(d.get("cost-per-1k-tokens-usd", 0.015))
except Exception:
return {"light": 150000, "standard": 500000, "heavy": 2000000}, 0.015
def rows():
if not os.path.exists(LEDGER):
return []
out = []
for line in open(LEDGER):
line = line.strip()
if line:
try:
out.append(json.loads(line))
except json.JSONDecodeError:
pass
return out
def log(workflow, role, tokens, wave=None, tier=None, usage_source_id=None):
per_wave, _ = budgets()
if tier not in per_wave:
raise ValueError(f"미등록 tier: {tier!r}")
value = int(tokens)
if value < 0:
raise ValueError("tokens는 0 이상이어야 한다")
if not str(workflow or "").strip() or workflow == "-":
raise ValueError("canonical workflow id 필수")
if usage_source_id and any(r.get("usage-source-id") == usage_source_id for r in rows()):
raise ValueError(f"중복 usage-source-id: {usage_source_id}")
os.makedirs(os.path.dirname(LEDGER), exist_ok=True)
rec = {
"at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
"workflow": workflow, "wave": wave or "-", "role": role,
"tokens": value, "tier": tier,
}
if usage_source_id:
rec["usage-source-id"] = usage_source_id
with open(LEDGER, "a") as f:
try:
import fcntl
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
except Exception:
pass
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
f.flush()
os.fsync(f.fileno())
print(f"[token_ledger] +{tokens} tok · {workflow}/{role}")
def sum_workflow(workflow):
return sum(r["tokens"] for r in rows() if r.get("workflow") == workflow)
def sum_wave(workflow, wave):
"""단일 wave의 토큰 합. 예산(per-wave)은 wave 단위로 대조해야 하므로 이걸 쓴다.
finding #19: 예전 check()/dashboard()는 sum_workflow(워크플로 전체 합)를 per-wave
예산과 비교해, 정상 wave가 여러 번 쌓이면 뒤 wave에서 허위 초과가 났다. wave 미지정
로그는 '-' 버킷으로 묶이므로 단일-wave 워크플로의 기존 동작은 그대로 보존된다."""
w = wave or "-"
return sum(r["tokens"] for r in rows()
if r.get("workflow") == workflow and (r.get("wave") or "-") == w)
def status_icon(used, budget):
if not budget:
return "—"
r = used / budget
return "🚨" if r > 1 else ("⚠️" if r > 0.8 else "✅")
def dashboard():
# finding #19: 예산 대비는 per-wave 이므로 (워크플로, wave) 단위로 그룹핑해 각 wave를 대조한다.
per_wave, cost1k = budgets()
data = rows()
groups = {}
for r in data:
groups.setdefault((r.get("workflow", "-"), r.get("wave") or "-"), []).append(r)
workflows = {k[0] for k in groups}
ts = datetime.now().strftime("%Y-%m-%d %H:%M")
total = sum(r["tokens"] for r in data)
L = ["# 💰 토큰 대시보드 (대표용)", "",
f"생성: {ts} · 총 {total:,} tok · 추정 ${total/1000*cost1k:,.2f} · 워크플로 {len(workflows)}개 · wave {len(groups)}개",
f"> tier 예산(per-wave): light {per_wave.get('light',0):,} · standard {per_wave.get('standard',0):,} · heavy {per_wave.get('heavy',0):,} · 초과 시 Orchestrator가 collapse로 강등. (예산 대비는 wave 단위)",
"", "| 워크플로 | wave | 워커수 | 토큰 | 추정$ | tier | wave예산대비 | 상태 |", "|---|---|--:|--:|--:|---|---|:--:|"]
for wf, wv in sorted(groups):
rs = groups[(wf, wv)]
tok = sum(r["tokens"] for r in rs)
tier = next((r.get("tier") for r in rs if r.get("tier") and r.get("tier") != "-"), "-")
budget = per_wave.get(tier)
pct = f"{tok/budget*100:.0f}% of {tier}" if budget else "-"
L.append(f"| {wf} | {wv} | {len(rs)} | {tok:,} | ${tok/1000*cost1k:,.2f} | {tier} | {pct} | {status_icon(tok, budget)} |")
L += ["", "## 워커별 상세", "", "| at | 워크플로 | 역할 | 토큰 |", "|---|---|---|--:|"]
for r in sorted(data, key=lambda x: x.get("at", ""), reverse=True):
L.append(f"| {r.get('at','-')} | {r.get('workflow','-')} | {r.get('role','-')} | {r['tokens']:,} |")
os.makedirs(os.path.dirname(DASH), exist_ok=True)
with open(DASH, "w") as f:
f.write("\n".join(L) + "\n")
print(f"[token_ledger] dashboard -> {os.path.relpath(DASH, ROOT)} ({total:,} tok)")
def check(workflow, tier, add=0, wave=None):
# finding #19: 예산은 per-wave 이므로 현재 wave의 토큰만 대조한다(워크플로 전체 합 아님).
per_wave, _ = budgets()
budget = per_wave.get(tier)
if budget is None:
sys.stderr.write(f"[token_ledger] 미등록 tier: {tier!r}\n")
sys.exit(2)
addition = int(add or 0)
if addition < 0:
sys.stderr.write("[token_ledger] --add는 0 이상이어야 한다\n")
sys.exit(2)
w = wave or "-"
used = sum_wave(workflow, w) + addition
if budget and used > budget:
sys.stderr.write(
f"[token_ledger] BUDGET EXCEEDED {workflow}/wave {w}: {used:,} > {tier} per-wave 예산 {budget:,}. "
f"fan-out을 collapse(단일 종합)로 강등하거나 tier를 올려라.\n")
sys.exit(2)
print(f"[token_ledger] OK {workflow}/wave {w}: {used:,}/{budget or '∞'} ({tier})")
sys.exit(0)
def main():
a = sys.argv[1:]
if not a:
sys.stderr.write(__doc__)
sys.exit(1)
cmd, opt = a[0], {}
i = 1
while i < len(a):
if a[i].startswith("--"):
opt[a[i][2:]] = a[i + 1] if i + 1 < len(a) and not a[i + 1].startswith("--") else True
i += 2
else:
i += 1
if cmd == "log":
try:
log(opt.get("workflow", "-"), opt.get("role", "-"), opt.get("tokens", 0),
opt.get("wave"), opt.get("tier"), opt.get("usage-source-id"))
except (TypeError, ValueError) as exc:
sys.stderr.write(f"[token_ledger] log 거부: {exc}\n")
sys.exit(2)
elif cmd == "dashboard":
dashboard()
elif cmd == "check":
check(opt.get("workflow", "-"), opt.get("tier", "standard"), opt.get("add", 0), opt.get("wave"))
else:
sys.stderr.write(f"unknown command: {cmd}\n")
sys.exit(1)
if __name__ == "__main__":
main()