Code MVP: hooks and budgets

"""
chapter 12: hooks and budgets.
HookManager fires deterministic handlers on lifecycle events; a PreToolUse
hook can veto a call. Budget enforces max_turns and max_cost (the real
version of Chapter 2's crude limit). Reuses Chapter 4's token estimate.
"""
from dataclasses import dataclass, field

def estimate_tokens(text: str) -> int:
    return max(1, len(str(text)) // 4)

class HookManager:
    def __init__(self):
        self.hooks = {"PreToolUse": [], "PostToolUse": [], "Stop": [], "PreCompact": []}

    def register(self, event: str, handler):
        self.hooks.setdefault(event, []).append(handler)

    def fire(self, event: str, payload: dict) -> dict:
        """Run all handlers for an event. A PreToolUse handler may return
        {'block': True, 'reason': ...} to veto the action."""
        for handler in self.hooks.get(event, []):
            result = handler(payload) or {}
            if result.get("block"):
                return result
        return {"block": False}

class BudgetExceeded(Exception):
    pass

@dataclass
class Budget:
    max_turns: int = 30
    max_cost: float = 5.0           # arbitrary cost units
    input_rate: float = 1.0
    output_rate: float = 4.0
    turns: int = 0
    cost: float = 0.0

    def charge(self, prompt_text: str, output_text: str):
        self.turns += 1
        self.cost += estimate_tokens(prompt_text) / 1000 * self.input_rate
        self.cost += estimate_tokens(output_text) / 1000 * self.output_rate
        if self.turns > self.max_turns:
            raise BudgetExceeded(f"error_max_turns ({self.max_turns})")
        if self.cost > self.max_cost:
            raise BudgetExceeded(f"error_max_budget ({self.max_cost})")

# --- example deterministic hooks ---
def block_pushes(payload):
    cmd = payload.get("args", {}).get("command", "")
    if "git push" in cmd:
        return {"block": True, "reason": "pushes are blocked by a PreToolUse hook"}

def autoformat(payload):
    if payload.get("name") == "edit":
        print(f"[PostToolUse] would auto-format {payload['args'].get('path')}")

if __name__ == "__main__":
    hm = HookManager()
    hm.register("PreToolUse", block_pushes)
    hm.register("PostToolUse", autoformat)
    budget = Budget(max_turns=3, max_cost=100)

    # A blocked call:
    print(hm.fire("PreToolUse", {"name": "shell", "args": {"command": "git push"}}))
    # An allowed call, then a post hook:
    print(hm.fire("PreToolUse", {"name": "shell", "args": {"command": "ls"}}))
    hm.fire("PostToolUse", {"name": "edit", "args": {"path": "main.py"}})

    # Budget stops a runaway loop:
    try:
        for i in range(10):
            budget.charge("prompt " * 100, "output")
            print("turn", budget.turns, "cost", round(budget.cost, 2))
    except BudgetExceeded as e:
        print("STOPPED:", e)

Run it: the git push is vetoed by the hook, ls passes, the edit triggers a post-hook, and the loop halts at the turn cap with a clear error_max_turns reason, just like the SDK's ResultMessage.subtype. Hooks and budgets are small, deterministic, and exactly the controls that make unattended runs safe.