From c06164ce6a97c62ce737a6a264f7caa513586c25 Mon Sep 17 00:00:00 2001 From: moneytools Date: Thu, 8 Oct 2026 21:44:36 -0500 Subject: [PATCH 1/3] feat(budget): enforce the budget cap in aegis hook and aegis install MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - hook: after the policy decision (BLOCK stands), budget gate for every covered tool call and for prompt events (Claude/Codex UserPromptSubmit, Gemini BeforeAgent); warn once per session per threshold via systemMessage / OpenCode toast / Copilot stderr; an unusable budget.yaml fails closed - install: budget hooks when a budget.yaml applies (--budget/--no-budget): all-tools matcher plus the prompt hook; reinstall and remove clean every aegis-owned event - Claude plugin and Gemini extension: budget-only entries whose wrapper exits before starting aegis when no budget.yaml applies - OpenCode plugin: every tool call with sessionID where a budget applies - design §6 table re-checked against each agent's hook docs --- CHANGELOG.md | 11 + docs/dev/DESIGN-v1.0-budget-cap.md | 18 +- hooks/aegis-hook.sh | 29 ++- hooks/hooks.json | 23 ++ plugins/claude/hooks/aegis-hook.sh | 29 ++- plugins/claude/hooks/hooks.json | 27 +- plugins/copilot/aegis-hook.sh | 29 ++- src/aegis_core/budget/gate.py | 120 +++++++++ src/aegis_core/cli.py | 40 ++- src/aegis_core/hook.py | 222 ++++++++++++++--- .../integrations/opencode-plugin.js | 60 ++++- tests/test_budget_hook.py | 235 ++++++++++++++++++ tests/test_plugins.py | 63 ++++- 13 files changed, 839 insertions(+), 67 deletions(-) create mode 100644 src/aegis_core/budget/gate.py create mode 100644 tests/test_budget_hook.py diff --git a/CHANGELOG.md b/CHANGELOG.md index dce7943..b19dd17 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,17 @@ test on each agent are in. rebuilds that report whether the day is complete, partial or of unknown completeness, and usage bucketed by the time the log recorded, never the time it was read. - Evaluation to allow / warn / deny with the measured usage, limits and observed overshoot. +- The budget cap in `aegis hook` (design §6): in a project whose policy directory holds a + `budget.yaml`, after the policy decision (a BLOCK always stands), every covered tool call is + denied once the usage available at the hook reaches a limit, and new prompts where the agent + can block them (Claude Code and Codex `UserPromptSubmit`, Gemini CLI `BeforeAgent`); a warning + is shown once per session per threshold (`systemMessage`; an OpenCode toast; stderr for Copilot + CLI). A `budget.yaml` that cannot be used (bad signature, shape, authority) fails closed. +- `aegis install` adds the budget hooks when the policy directory that applies holds a + `budget.yaml` (`--budget` / `--no-budget` to choose): the tool hook sees every tool call and + the prompt hook is added where the agent has one. The Claude Code plugin and the Gemini + extension carry budget-only entries that start no process in projects without a budget; the + OpenCode plugin sends every tool call, with its session id, only where a budget applies. - The example `authority.yaml` gives `admin` the `budget` class. - `agents.yaml`: a trusted AWS `source-identity` now warns (`source-identity-inherited`): a source identity survives role chaining, so an agent started in a session that carries it is diff --git a/docs/dev/DESIGN-v1.0-budget-cap.md b/docs/dev/DESIGN-v1.0-budget-cap.md index dd27c56..7d8c770 100644 --- a/docs/dev/DESIGN-v1.0-budget-cap.md +++ b/docs/dev/DESIGN-v1.0-budget-cap.md @@ -255,13 +255,17 @@ commands to **all tool calls**, and adds a prompt hook where the agent has one. | Agent | Before each tool call | Before a new prompt | How "warn" shows | |---|---|---|---| | Claude Code | `PreToolUse` (matcher `*`) | `UserPromptSubmit` | `systemMessage` to the user | -| Codex CLI | `PreToolUse` | — | stderr message | -| Gemini CLI | `BeforeTool` | `BeforeAgent` | `systemMessage` | -| OpenCode | plugin `tool.execute.before` | — | toast | -| Copilot CLI | `preToolUse` (premium-request cap) | — | stderr message | - -Hook and event names in this table come from the 2026-10-06 research and must be re-checked -against each agent's current version before implementation (step 4 of §10). +| Codex CLI | `PreToolUse` (matcher `.*`) | `UserPromptSubmit` | `systemMessage` (a UI warning) | +| Gemini CLI | `BeforeTool` (matcher `.*`) | `BeforeAgent` | `systemMessage` | +| OpenCode | plugin `tool.execute.before` (every tool, where a budget applies) | — | toast (`client.tui.showToast`) | +| Copilot CLI | `preToolUse` (no matcher: every tool; premium-request cap) | — (`userPromptSubmitted` cannot block) | stderr only: Copilot CLI has no message field | + +Re-checked against each agent's hook documentation on 2026-10-08 (implementation PR 2): Codex +now has `UserPromptSubmit` and `systemMessage`; Claude Code and Codex sub-agent hooks carry the +parent's `session_id`; OpenCode's `tool.execute.before` input carries `sessionID`. The Claude +Code plugin and the Gemini extension ship fixed hook files, so they carry extra budget-only +entries (every other tool, and the prompt hook) whose wrapper exits at once unless the policy +directory holds a `budget.yaml`: projects without a budget start no extra process. - **Below `warn_at`:** nothing changes; budget adds no output. - **Crossing `warn_at`:** one warning per session per threshold (not per call), with usage, the diff --git a/hooks/aegis-hook.sh b/hooks/aegis-hook.sh index 6b18db7..7c9c508 100755 --- a/hooks/aegis-hook.sh +++ b/hooks/aegis-hook.sh @@ -1,7 +1,11 @@ #!/usr/bin/env bash # Plugin entry point: find the aegis binary and hand it the hook payload. # -# aegis-hook.sh (the payload arrives on stdin) +# aegis-hook.sh (the payload arrives on stdin) +# aegis-hook.sh --budget-only (non-shell tools and prompts: only the +# budget cap applies, so this exits at +# once unless the policy directory that +# applies holds a budget.yaml) # # aegis itself decides everything (see aegis_core/hook.py): outside a # project with a policy it allows every command, inside one it blocks only @@ -12,6 +16,26 @@ # names an infrastructure binary, and allow the rest. set -u agent=${1:-claude} +mode=${2:-} + +# The policy directory aegis would use: $AEGIS_CONFIG_DIR, the nearest .aegis/ +# at or above the project, or ~/.config/aegis (as aegis_core/hook.py does). +policy_dir() { + if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi + d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + while [ -n "$d" ]; do + if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi + [ "$d" = "/" ] && break + d=$(dirname "$d") + done + if [ -d "$HOME/.config/aegis" ]; then printf '%s' "$HOME/.config/aegis"; return 0; fi + return 1 +} + +if [ "$mode" = "--budget-only" ]; then + dir=$(policy_dir) || exit 0 + [ -f "$dir/budget.yaml" ] || exit 0 +fi payload=$(cat) aegis=${AEGIS_BIN:-} @@ -27,6 +51,9 @@ if command -v uvx >/dev/null 2>&1; then case "$rc" in 0|2) exit "$rc" ;; esac fi +# Without aegis nothing can be measured; the budget-only entries allow. +[ "$mode" = "--budget-only" ] && exit 0 + opted_in() { [ -n "${AEGIS_CONFIG_DIR:-}" ] && return 0 [ -d "$HOME/.config/aegis" ] && return 0 diff --git a/hooks/hooks.json b/hooks/hooks.json index dd615f9..d6a2d05 100644 --- a/hooks/hooks.json +++ b/hooks/hooks.json @@ -11,6 +11,29 @@ "timeout": 30000 } ] + }, + { + "matcher": "^(?!run_shell_command$).*", + "hooks": [ + { + "name": "aegis-devops-budget", + "type": "command", + "command": "\"${extensionPath}${/}hooks${/}aegis-hook.sh\" gemini --budget-only", + "timeout": 30000 + } + ] + } + ], + "BeforeAgent": [ + { + "hooks": [ + { + "name": "aegis-devops-budget", + "type": "command", + "command": "\"${extensionPath}${/}hooks${/}aegis-hook.sh\" gemini --budget-only", + "timeout": 30000 + } + ] } ] } diff --git a/plugins/claude/hooks/aegis-hook.sh b/plugins/claude/hooks/aegis-hook.sh index 6b18db7..7c9c508 100755 --- a/plugins/claude/hooks/aegis-hook.sh +++ b/plugins/claude/hooks/aegis-hook.sh @@ -1,7 +1,11 @@ #!/usr/bin/env bash # Plugin entry point: find the aegis binary and hand it the hook payload. # -# aegis-hook.sh (the payload arrives on stdin) +# aegis-hook.sh (the payload arrives on stdin) +# aegis-hook.sh --budget-only (non-shell tools and prompts: only the +# budget cap applies, so this exits at +# once unless the policy directory that +# applies holds a budget.yaml) # # aegis itself decides everything (see aegis_core/hook.py): outside a # project with a policy it allows every command, inside one it blocks only @@ -12,6 +16,26 @@ # names an infrastructure binary, and allow the rest. set -u agent=${1:-claude} +mode=${2:-} + +# The policy directory aegis would use: $AEGIS_CONFIG_DIR, the nearest .aegis/ +# at or above the project, or ~/.config/aegis (as aegis_core/hook.py does). +policy_dir() { + if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi + d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + while [ -n "$d" ]; do + if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi + [ "$d" = "/" ] && break + d=$(dirname "$d") + done + if [ -d "$HOME/.config/aegis" ]; then printf '%s' "$HOME/.config/aegis"; return 0; fi + return 1 +} + +if [ "$mode" = "--budget-only" ]; then + dir=$(policy_dir) || exit 0 + [ -f "$dir/budget.yaml" ] || exit 0 +fi payload=$(cat) aegis=${AEGIS_BIN:-} @@ -27,6 +51,9 @@ if command -v uvx >/dev/null 2>&1; then case "$rc" in 0|2) exit "$rc" ;; esac fi +# Without aegis nothing can be measured; the budget-only entries allow. +[ "$mode" = "--budget-only" ] && exit 0 + opted_in() { [ -n "${AEGIS_CONFIG_DIR:-}" ] && return 0 [ -d "$HOME/.config/aegis" ] && return 0 diff --git a/plugins/claude/hooks/hooks.json b/plugins/claude/hooks/hooks.json index 4203326..b149463 100644 --- a/plugins/claude/hooks/hooks.json +++ b/plugins/claude/hooks/hooks.json @@ -4,7 +4,32 @@ { "matcher": "Bash", "hooks": [ - { "type": "command", "command": "\"${CLAUDE_PLUGIN_ROOT}\"/hooks/aegis-hook.sh claude", "timeout": 30 } + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}\"/hooks/aegis-hook.sh claude", + "timeout": 30 + } + ] + }, + { + "matcher": "^(?!Bash$).*", + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}\"/hooks/aegis-hook.sh claude --budget-only", + "timeout": 30 + } + ] + } + ], + "UserPromptSubmit": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}\"/hooks/aegis-hook.sh claude --budget-only", + "timeout": 30 + } ] } ] diff --git a/plugins/copilot/aegis-hook.sh b/plugins/copilot/aegis-hook.sh index 6b18db7..7c9c508 100755 --- a/plugins/copilot/aegis-hook.sh +++ b/plugins/copilot/aegis-hook.sh @@ -1,7 +1,11 @@ #!/usr/bin/env bash # Plugin entry point: find the aegis binary and hand it the hook payload. # -# aegis-hook.sh (the payload arrives on stdin) +# aegis-hook.sh (the payload arrives on stdin) +# aegis-hook.sh --budget-only (non-shell tools and prompts: only the +# budget cap applies, so this exits at +# once unless the policy directory that +# applies holds a budget.yaml) # # aegis itself decides everything (see aegis_core/hook.py): outside a # project with a policy it allows every command, inside one it blocks only @@ -12,6 +16,26 @@ # names an infrastructure binary, and allow the rest. set -u agent=${1:-claude} +mode=${2:-} + +# The policy directory aegis would use: $AEGIS_CONFIG_DIR, the nearest .aegis/ +# at or above the project, or ~/.config/aegis (as aegis_core/hook.py does). +policy_dir() { + if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi + d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + while [ -n "$d" ]; do + if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi + [ "$d" = "/" ] && break + d=$(dirname "$d") + done + if [ -d "$HOME/.config/aegis" ]; then printf '%s' "$HOME/.config/aegis"; return 0; fi + return 1 +} + +if [ "$mode" = "--budget-only" ]; then + dir=$(policy_dir) || exit 0 + [ -f "$dir/budget.yaml" ] || exit 0 +fi payload=$(cat) aegis=${AEGIS_BIN:-} @@ -27,6 +51,9 @@ if command -v uvx >/dev/null 2>&1; then case "$rc" in 0|2) exit "$rc" ;; esac fi +# Without aegis nothing can be measured; the budget-only entries allow. +[ "$mode" = "--budget-only" ] && exit 0 + opted_in() { [ -n "${AEGIS_CONFIG_DIR:-}" ] && return 0 [ -d "$HOME/.config/aegis" ] && return 0 diff --git a/src/aegis_core/budget/gate.py b/src/aegis_core/budget/gate.py new file mode 100644 index 0000000..83e71cf --- /dev/null +++ b/src/aegis_core/budget/gate.py @@ -0,0 +1,120 @@ +"""The budget check inside ``aegis hook`` (design v1.0, §6). + +``aegis hook `` first decides the command against the infrastructure +policy, as it always has; a BLOCK stands whatever the budget says. Then, +in a project whose policy directory holds a ``budget.yaml``, this module +brings the session's usage up to date from the agent's logs and decides: + +* **deny** once the usage available at this hook reaches a limit -- every + covered tool call, and new prompts where the agent has a prompt hook; +* **warn** once per session per threshold (not on every call); +* **allow** otherwise, silently. +""" + +from __future__ import annotations + +import contextlib +import json +import time +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from aegis_core.budget.accounting import BudgetStore, refresh +from aegis_core.budget.evaluate import ALLOW, DENY, WARN, evaluate +from aegis_core.budget.logs import AgentLog, default_logs +from aegis_core.budget.policy import BudgetPolicy +from aegis_core.budget.pricing import PriceTable, load_builtin_table + +# Hook events that come before a new prompt rather than a tool call. +PROMPT_EVENTS = frozenset({"UserPromptSubmit", "BeforeAgent"}) + +_HOW_TO_CONTINUE = { + "session": "start a new session", + "project_day": "wait for the next day", + "copilot_premium_requests": "start a new session", +} + + +@dataclass +class Outcome: + decision: str # allow | warn | deny + message: str = "" # the deny reason, or the warning to show ("" = nothing to show) + + +def session_ref(payload: Mapping[str, Any]) -> tuple[str | None, str | None]: + """``(session id, transcript path)`` from any agent's hook payload.""" + sid = payload.get("session_id") or payload.get("sessionId") or payload.get("sessionID") + transcript = payload.get("transcript_path") + return (sid if isinstance(sid, str) and sid else None, + transcript if isinstance(transcript, str) and transcript else None) + + +def _warned_path(store: BudgetStore, agent: str, sid: str) -> Path: + return store.record_path(agent, sid).with_suffix(".warned.json") + + +def _first_time(store: BudgetStore, agent: str, sid: str, keys: list[str]) -> list[str]: + """The warning keys not yet shown in this session; records them as + shown. Best effort: a lost file means a warning is shown again.""" + path = _warned_path(store, agent, sid) + try: + shown = set(json.loads(path.read_text())) + except (OSError, ValueError, TypeError): + shown = set() + new = [k for k in keys if k not in shown] + if new: + with contextlib.suppress(OSError): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(sorted(shown | set(new)))) + return new + + +def check(agent: str, payload: Mapping[str, Any], cwd: str | None, policy: BudgetPolicy, *, + store: BudgetStore | None = None, logs: Mapping[str, AgentLog] | None = None, + now: float | None = None) -> Outcome: + """The budget decision for one hook call. Raises on internal errors; + the caller fails closed inside an opted-in project.""" + if not policy.measures(agent): + return Outcome(ALLOW) + store = store or BudgetStore() + all_logs = logs or default_logs() + logs = {a: all_logs[a] for a in policy.agents if a in all_logs and policy.measures(a)} + sid, transcript = session_ref(payload) + loc = logs[agent].locate(sid, transcript) if sid else None + if loc is None: + what = (f"cannot find the log of {agent} session {sid}" if sid + else "the hook payload names no session") + if policy.on_unknown_log == "deny": + return Outcome(DENY, f"aegis budget: {what}, so its usage cannot be measured " + "(on_unknown_log: deny)") + if sid and not _first_time(store, agent, sid, ["log-not-found"]): + return Outcome(WARN) + return Outcome(WARN, f"aegis budget: {what}; this session is not being measured") + + snap = refresh(store, logs, loc, policy.zone, now=time.time() if now is None else now, + fallback_cwd=cwd) + prices = PriceTable(load_builtin_table(), policy.pricing) + verdict = evaluate(policy, prices, snap) + if verdict.decision == DENY: + over = [m.name for m in verdict.measures if m.level == DENY] + ways = sorted({_HOW_TO_CONTINUE[n] for n in over}) + ways.append("raise the limit in budget.yaml (it must be re-signed)") + reason = "; ".join(m.removeprefix("budget: ") for m in verdict.messages + if not m.endswith("left")) + return Outcome(DENY, f"aegis budget: {reason}. To continue: {', or '.join(ways)}.") + if verdict.decision == WARN: + keys = [] + for reason in verdict.reasons: + key = reason + if reason.startswith("project_day"): + key += f":{snap.day}" + elif reason == "estimated-price": + key += ":" + ",".join(sorted(verdict.estimated)) + keys.append(key) + new = set(_first_time(store, snap.record.agent, snap.record.session_id, keys)) + shown = [msg for key, msg in zip(keys, verdict.messages, strict=False) if key in new] + return Outcome(WARN, "aegis " + "; ".join(m.removeprefix("budget: ") for m in shown) + if shown else "") + return Outcome(ALLOW) diff --git a/src/aegis_core/cli.py b/src/aegis_core/cli.py index 018bf40..3d54f71 100644 --- a/src/aegis_core/cli.py +++ b/src/aegis_core/cli.py @@ -627,6 +627,17 @@ def _build_parser() -> argparse.ArgumentParser: install_parser.add_argument( "--remove", action="store_true", help="remove the aegis hook instead" ) + budget_flag = install_parser.add_mutually_exclusive_group() + budget_flag.add_argument( + "--budget", dest="budget", action="store_true", default=None, + help="install the budget-cap hooks (every tool call, and the prompt hook where the " + "agent has one) even if no budget.yaml is found now", + ) + budget_flag.add_argument( + "--no-budget", dest="budget", action="store_false", + help="install only the shell-command hook (default: budget hooks when the policy " + "directory that applies holds a budget.yaml)", + ) return parser @@ -1337,6 +1348,32 @@ def _load_identity(args: argparse.Namespace): return model, [*key_warnings, *authority_map.warnings, *model.warnings] +def load_budget_for_hook(config_dir: Path, extra_args: list[str] | None = None): + """``budget.yaml`` in an opted-in policy directory, verified like every + other policy file (signature, and its principal must hold ``budget``), + for ``aegis hook``. ``None`` when the directory has no budget. The key + is resolved as for ``aegis check`` (``--key``/``--insecure`` among the + hook's extra arguments, ``$AEGIS_SIGNING_KEY``, the example key next to + the constraints). Raises when the budget cannot be used; the hook fails + closed on that.""" + from aegis_core.budget.policy import BUDGET_FILE, load_budget_policy + + path = Path(config_dir) / BUDGET_FILE + if not path.exists(): + return None + opts = argparse.ArgumentParser(add_help=False) + opts.add_argument("--key") + opts.add_argument("--insecure", action="store_true") + known, _ = opts.parse_known_args(extra_args or []) + args = argparse.Namespace(config_dir=str(config_dir), constraints=None, authority=None, + environments=None, plan_constraints=None, key=known.key, + insecure=known.insecure) + _resolve_config_paths(args) + key, insecure = _resolve_key(args, []) + authority_map = load_authority_map(args.authority, key=key, insecure=insecure) + return load_budget_policy(path, authority_map=authority_map, key=key, insecure=insecure) + + def _run_audit_aws(args: argparse.Namespace) -> int: """``aegis audit-identity aws``: read-only; exit 1 on a problem.""" from aegis_core.audit_aws import InventoryError, audit, collect_inventory @@ -1640,7 +1677,8 @@ def _run(argv: list[str]) -> int: ) if args.command == "install": return hook_module.main_install( - args.agent, user=args.user, project=args.project, remove=args.remove + args.agent, user=args.user, project=args.project, remove=args.remove, + budget=args.budget, ) if args.command != "check": raise UsageError("unknown command") diff --git a/src/aegis_core/hook.py b/src/aegis_core/hook.py index c9414db..36ccd9d 100644 --- a/src/aegis_core/hook.py +++ b/src/aegis_core/hook.py @@ -56,6 +56,17 @@ AGENTS = ("claude", "codex", "copilot", "cursor", "gemini", "opencode", "vscode") +# Agents whose usage the budget cap can measure (design v1.0 §4). Cursor and +# VS Code expose no token data. +BUDGET_AGENTS = frozenset({"claude", "codex", "copilot", "gemini", "opencode"}) +BUDGET_FILE = "budget.yaml" +# Hook events that come before a new prompt rather than a tool call. +PROMPT_EVENTS = frozenset({"UserPromptSubmit", "BeforeAgent"}) +# Agents that show a hook's JSON "systemMessage" to the user (a warning that +# does not block). OpenCode's plugin shows "warning" as a toast; Copilot CLI +# has no such channel, so its warning goes to stderr. +_SYSTEM_MESSAGE_AGENTS = frozenset({"claude", "codex", "gemini"}) + # Binaries whose commands are gated even when the policy is unusable or the # command cannot be parsed. terraform/tofu have no argv parser (their plans # are checked as JSON) but are infrastructure tools all the same; k/tf are @@ -320,25 +331,102 @@ def render(agent: str, verdict: HookVerdict, *, escalate_as: str = "ask") -> tup return EXIT_ALLOW, json.dumps(body), "" +def _with_warning(agent: str, reply: tuple[int, str, str], warning: str) -> tuple[int, str, str]: + """Adds a budget warning to a reply that lets the call through, in the + agent's own channel.""" + code, out, err = reply + if not warning or code != EXIT_ALLOW: + return reply + if agent in _SYSTEM_MESSAGE_AGENTS or agent == "opencode": + try: + body = json.loads(out) if out else {} + except json.JSONDecodeError: + body = {} + body["systemMessage" if agent != "opencode" else "warning"] = warning + return code, json.dumps(body), err + return code, out, f"{err}\n{warning}".strip() + + +def render_prompt(agent: str, verdict: HookVerdict, warning: str = "") -> tuple[int, str, str]: + """The reply to a prompt hook (Claude Code / Codex ``UserPromptSubmit``, + Gemini CLI ``BeforeAgent``): block the prompt, or let it through, + optionally with a warning.""" + if verdict.decision == "deny": + reason = verdict.reason or "aegis: blocked" + body = ({"decision": "deny", "reason": reason} if agent == "gemini" + else {"decision": "block", "reason": reason}) + return EXIT_BLOCK, json.dumps(body), reason + if warning: + return EXIT_ALLOW, json.dumps({"systemMessage": warning}), "" + return EXIT_ALLOW, "", "" + + +def _budget_outcome(agent: str, payload: dict, request: HookRequest, config_dir: Path, + extra_args: list[str] | None): + """The budget gate's outcome for this call. A budget.yaml that cannot be + used denies (the project opted in to a budget; fail closed).""" + from aegis_core import cli # local: cli imports this module + from aegis_core.budget.gate import Outcome, check + + try: + policy = cli.load_budget_for_hook(config_dir, extra_args) + except Exception as exc: # noqa: BLE001 - any failure to load is reported, not raised + return Outcome("deny", f"aegis budget: {config_dir / BUDGET_FILE} cannot be used " + f"({exc.__class__.__name__}: {exc}); fix and re-sign it, or " + "remove it") + if policy is None: + return Outcome("allow") + return check(agent, payload, request.cwd, policy) + + def run_hook(agent: str, stdin_text: str, *, escalate_as: str = "ask", extra_args: list[str] | None = None) -> tuple[int, str, str]: """The whole hook: payload text in, ``(exit, stdout, stderr)`` out. Never raises. Outside an opted-in project every failure allows; inside - one, a failure blocks (fail closed where the user asked to be gated).""" + one, a failure blocks (fail closed where the user asked to be gated). + + Where the project's policy directory also holds a ``budget.yaml``, the + budget cap is checked after the policy decision (a BLOCK stands whatever + the budget says): for every tool call the hook sees, and for prompt + events.""" config_dir: Path | None = None try: try: payload = json.loads(stdin_text or "") except json.JSONDecodeError as exc: raise HookInputError(f"hook payload is not JSON: {exc}") from None - request = parse_request(agent, payload) - if request.command is None or not request.command.strip(): + if not isinstance(payload, dict): + raise HookInputError("hook payload is not a JSON object") + prompt = payload.get("hook_event_name") in PROMPT_EVENTS + if prompt: + cwd = payload.get("cwd") if isinstance(payload.get("cwd"), str) else None + request = HookRequest(None, cwd) + else: + request = parse_request(agent, payload) + has_command = bool(request.command and request.command.strip()) + if not has_command and agent not in BUDGET_AGENTS: return _allow(agent) config_dir = find_opt_in_config(_project_dir(request)) if config_dir is None: return _allow(agent) - verdict = decide(request.command, config_dir, extra_args) - return render(agent, verdict, escalate_as=escalate_as) + has_budget = agent in BUDGET_AGENTS and (config_dir / BUDGET_FILE).is_file() + if not has_command and not has_budget: + return _allow(agent) + + verdict = decide(request.command, config_dir, extra_args) if has_command \ + else HookVerdict("allow") + if verdict.decision == "deny": + return render(agent, verdict, escalate_as=escalate_as) + warning = "" + if has_budget: + outcome = _budget_outcome(agent, payload, request, config_dir, extra_args) + if outcome.decision == "deny": + verdict = HookVerdict("deny", outcome.message) + elif outcome.message: + warning = outcome.message + if prompt: + return render_prompt(agent, verdict, warning) + return _with_warning(agent, render(agent, verdict, escalate_as=escalate_as), warning) except Exception as exc: # noqa: BLE001 - a hook must answer, never crash if config_dir is None and not isinstance(exc, HookInputError): return _allow(agent) @@ -427,45 +515,80 @@ def _strip_ours(entries: list) -> list: return kept -def merge_config(agent: str, existing: dict, command: str | None) -> dict: - """``existing`` with the aegis hook added (``command``) or removed - (``None``), every other setting and hook left as it was.""" - doc = dict(existing) - hooks = dict(doc.get("hooks") or {}) +# Events aegis may own per agent; reinstalling or removing cleans all of them. +_EVENTS = { + "claude": ("PreToolUse", "UserPromptSubmit"), + "codex": ("PreToolUse", "UserPromptSubmit"), + "gemini": ("BeforeTool", "BeforeAgent"), + "copilot": ("preToolUse",), + "vscode": ("PreToolUse",), + "cursor": ("beforeShellExecution",), +} + + +def _entries(agent: str, command: str, budget: bool) -> list[tuple[str, dict]]: + """``(event, entry)`` pairs for this agent. With a budget, the tool hook + sees every tool call (the budget counts all of them, not only shell + commands) and a prompt hook is added where the agent has one that can + block (design v1.0 §6).""" if agent in ("claude", "codex"): - event, entry = "PreToolUse", { - "matcher": "Bash" if agent == "claude" else "^Bash$", - "hooks": [{"type": "command", "command": command, "timeout": 30}], - } - elif agent == "gemini": - event, entry = "BeforeTool", { - "matcher": "run_shell_command", - "hooks": [{"name": "aegis-devops", "type": "command", "command": command, - "timeout": 30000}], - } - elif agent == "copilot": - doc.setdefault("version", 1) - event, entry = "preToolUse", {"type": "command", "bash": command, "timeoutSec": 30} - elif agent == "vscode": + matcher = ("*" if agent == "claude" else ".*") if budget else ( + "Bash" if agent == "claude" else "^Bash$") + hook = {"type": "command", "command": command, "timeout": 30} + pairs = [("PreToolUse", {"matcher": matcher, "hooks": [hook]})] + if budget: + pairs.append(("UserPromptSubmit", {"hooks": [dict(hook)]})) + return pairs + if agent == "gemini": + hook = {"name": "aegis-devops", "type": "command", "command": command, "timeout": 30000} + pairs = [("BeforeTool", {"matcher": ".*" if budget else "run_shell_command", + "hooks": [hook]})] + if budget: + pairs.append(("BeforeAgent", {"hooks": [dict(hook)]})) + return pairs + if agent == "copilot": + # no matcher: Copilot CLI already sends every tool call; its prompt hook + # cannot block, so a budget adds nothing here + return [("preToolUse", {"type": "command", "bash": command, "timeoutSec": 30})] + if agent == "vscode": # VS Code's native format: PascalCase events and no "version" (a # numeric version marks the Copilot CLI format instead). VS Code # ignores matchers here, so the hook sees every tool call and # parse_request lets the non-terminal ones through. - event, entry = "PreToolUse", {"type": "command", "command": command, "timeout": 30} - else: # cursor + return [("PreToolUse", {"type": "command", "command": command, "timeout": 30})] + return [("beforeShellExecution", {"command": command, "failClosed": True})] # cursor + + +def merge_config(agent: str, existing: dict, command: str | None, *, + budget: bool = False) -> dict: + """``existing`` with the aegis hooks added (``command``) or removed + (``None``), every other setting and hook left as it was.""" + doc = dict(existing) + hooks = dict(doc.get("hooks") or {}) + if agent in ("copilot", "cursor") and command is not None: doc.setdefault("version", 1) - event, entry = "beforeShellExecution", {"command": command, "failClosed": True} - entries = _strip_ours(list(hooks.get(event) or [])) - if command is not None: - entries.append(entry) - if entries: - hooks[event] = entries - else: - hooks.pop(event, None) + wanted = _entries(agent, command, budget) if command is not None else [] + for event in _EVENTS[agent]: + entries = _strip_ours(list(hooks.get(event) or [])) + entries.extend(entry for ev, entry in wanted if ev == event) + if entries: + hooks[event] = entries + else: + hooks.pop(event, None) doc["hooks"] = hooks return doc +def budget_configured(*, user: bool, project: Path) -> bool: + """Whether the policy that would apply has a ``budget.yaml``: the + project's ``.aegis/`` (or ``$AEGIS_CONFIG_DIR``), or ``~/.config/aegis`` + for a user-level install.""" + env = os.environ.get(config_module.CONFIG_ENV_VAR) + dirs = [Path(env)] if env else [] + dirs.append(Path.home() / ".config" / "aegis" if user else project / ".aegis") + return any((d / BUDGET_FILE).is_file() for d in dirs) + + _AFTER_INSTALL = { "claude": "Restart Claude Code (or run /hooks) to load it.", "codex": "Codex runs a new hook only after you trust it: open codex and run /hooks.", @@ -497,7 +620,8 @@ def _install_opencode_plugin(path: Path, *, remove: bool) -> Path: return path -def install(agent: str, *, user: bool, project: Path, remove: bool = False) -> Path: +def install(agent: str, *, user: bool, project: Path, remove: bool = False, + budget: bool = False) -> Path: path = config_path(agent, user=user, project=project) if agent == "opencode": return _install_opencode_plugin(path, remove=remove) @@ -508,7 +632,8 @@ def install(agent: str, *, user: bool, project: Path, remove: bool = False) -> P existing = json.loads(text) if not isinstance(existing, dict): raise ValueError(f"{path} is not a JSON object; not touching it") - updated = merge_config(agent, existing, None if remove else hook_command(agent)) + updated = merge_config(agent, existing, None if remove else hook_command(agent), + budget=budget and agent in BUDGET_AGENTS) if remove and not updated["hooks"] and set(updated) <= {"version", "hooks"}: # nothing but our (now removed) hook was ever in it if path.exists(): @@ -519,12 +644,33 @@ def install(agent: str, *, user: bool, project: Path, remove: bool = False) -> P return path -def main_install(agent: str, *, user: bool, project: str, remove: bool) -> int: - path = install(agent, user=user, project=Path(project), remove=remove) +_BUDGET_INSTALL_NOTE = { + "claude": "budget: the hook sees every tool call, and UserPromptSubmit stops new prompts " + "over the limit", + "codex": "budget: the hook sees every tool call, and UserPromptSubmit stops new prompts " + "over the limit", + "gemini": "budget: the hook sees every tool call, and BeforeAgent stops new prompts over " + "the limit", + "copilot": "budget: the hook already sees every tool call (premium requests are " + "counted); Copilot CLI's prompt hook cannot block, so prompts are not stopped", + "opencode": "budget: the plugin checks every tool call in projects with a budget.yaml; " + "OpenCode has no prompt hook, so prompts are not stopped", +} + + +def main_install(agent: str, *, user: bool, project: str, remove: bool, + budget: bool | None = None) -> int: + if budget is None: + budget = budget_configured(user=user, project=Path(project)) + path = install(agent, user=user, project=Path(project), remove=remove, budget=budget) if remove: print(f"aegis: removed the {agent} hook from {path}") return 0 print(f"aegis: {agent} hook written to {path}") + if budget and agent in _BUDGET_INSTALL_NOTE: + print(f" {_BUDGET_INSTALL_NOTE[agent]}") + elif budget: + print(f" budget: not measured for {agent} (no token data); only the policy applies") print(f" {_AFTER_INSTALL[agent]}") if not user and not (Path(project) / ".aegis").is_dir(): print(f" No policy in {Path(project) / '.aegis'} yet, so every command is allowed " diff --git a/src/aegis_core/integrations/opencode-plugin.js b/src/aegis_core/integrations/opencode-plugin.js index d7dcc81..ab33a93 100644 --- a/src/aegis_core/integrations/opencode-plugin.js +++ b/src/aegis_core/integrations/opencode-plugin.js @@ -10,6 +10,10 @@ // policy blocks. If aegis cannot be started, the same rule applies in its // crudest form: in an opted-in project block commands naming an // infrastructure binary, allow the rest. +// +// Budget cap (aegis 1.0): where the policy directory holds a budget.yaml, +// every tool call goes to aegis, not only bash, with OpenCode's session id so +// aegis can read this session's usage; a warning comes back as a toast. import { spawnSync } from "node:child_process" import { existsSync } from "node:fs" import { homedir } from "node:os" @@ -22,25 +26,51 @@ const AEGIS = process.env.AEGIS_BIN ? [process.env.AEGIS_BIN] : INSTALLED || ["a const INFRA = /(^|[^\w.-])(kubectl|k|terraform|tf|tofu|aws|az|gcloud|gsutil|helm|argocd|flux|git|gh|psql|mysql|sqlite3|mongosh|pulumi|alembic|flyway|rails|prisma)([^\w.-]|$)/ -function optedIn(start) { - if (process.env.AEGIS_CONFIG_DIR) return true - if (existsSync(join(homedir(), ".config", "aegis"))) return true +// The policy directory aegis would use: $AEGIS_CONFIG_DIR, the nearest .aegis/ +// at or above the project, or ~/.config/aegis. +function policyDir(start) { + if (process.env.AEGIS_CONFIG_DIR) return process.env.AEGIS_CONFIG_DIR let dir = resolve(start || process.cwd()) for (;;) { - if (existsSync(join(dir, ".aegis"))) return true + if (existsSync(join(dir, ".aegis"))) return join(dir, ".aegis") const parent = dirname(dir) - if (parent === dir) return false + if (parent === dir) break dir = parent } + const user = join(homedir(), ".config", "aegis") + return existsSync(user) ? user : null +} + +function optedIn(start) { + return policyDir(start) !== null +} + +function hasBudget(start) { + const dir = policyDir(start) + return dir !== null && existsSync(join(dir, "budget.yaml")) +} + +async function toast(client, message) { + try { + await client.tui.showToast({ body: { message, variant: "warning" } }) + } catch { + console.warn(message) + } } -export const AegisDevOps = async ({ directory }) => ({ +export const AegisDevOps = async ({ directory, client }) => ({ "tool.execute.before": async (input, output) => { - if (input?.tool !== "bash") return - const command = output?.args?.command - if (typeof command !== "string" || !command.trim()) return const cwd = directory || process.cwd() - const payload = JSON.stringify({ tool_name: "bash", tool_input: { command }, cwd }) + const isBash = input?.tool === "bash" + const command = output?.args?.command + if (isBash && (typeof command !== "string" || !command.trim())) return + if (!isBash && !hasBudget(cwd)) return + const payload = JSON.stringify({ + tool_name: input?.tool, + tool_input: isBash ? { command } : output?.args || {}, + cwd, + session_id: input?.sessionID, + }) const run = spawnSync(AEGIS[0], [...AEGIS.slice(1), "hook", "opencode"], { input: payload, encoding: "utf8", @@ -48,7 +78,7 @@ export const AegisDevOps = async ({ directory }) => ({ cwd, }) if (run.error || run.status === null) { - if (optedIn(cwd) && INFRA.test(command)) { + if (isBash && optedIn(cwd) && INFRA.test(command)) { throw new Error( "aegis-devops could not run, so this infrastructure command cannot be checked: " + "run 'pip install aegis-devops' (blocking to fail closed)", @@ -56,7 +86,13 @@ export const AegisDevOps = async ({ directory }) => ({ } return } - if (run.status === 0) return + if (run.status === 0) { + try { + const reply = JSON.parse(run.stdout || "{}") + if (reply.warning) await toast(client, reply.warning) + } catch {} + return + } throw new Error((run.stderr || "").trim() || `aegis: blocked (exit ${run.status})`) }, }) diff --git a/tests/test_budget_hook.py b/tests/test_budget_hook.py new file mode 100644 index 0000000..8ea7c9e --- /dev/null +++ b/tests/test_budget_hook.py @@ -0,0 +1,235 @@ +"""The budget cap inside ``aegis hook`` and ``aegis install`` (design v1.0 §6). + +Each test runs the real hook against an ``aegis init`` policy directory with +a ``budget.yaml`` signed by the example key, and a synthetic Claude Code +transcript as the usage source (claude-opus-5-5: $4 per million input +tokens, so 250,000 input tokens is $1).""" + +import json +from pathlib import Path + +import pytest + +from aegis_core.config import init_config_dir +from aegis_core.hook import install, main_install, run_hook +from aegis_core.signing import load_key, sign_file + +SID = "sess-1" +MODEL = "claude-opus-5-5" +DAY = "2026-10-08" + + +@pytest.fixture +def env(tmp_path, monkeypatch): + home = tmp_path / "home" + project = tmp_path / "project" + home.mkdir() + project.mkdir() + for var in ("XDG_CACHE_HOME", "XDG_STATE_HOME", "XDG_DATA_HOME", "CLAUDE_CONFIG_DIR", + "CODEX_HOME", "AEGIS_CONFIG_DIR", "AEGIS_SIGNING_KEY"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("HOME", str(home)) + monkeypatch.chdir(project) + init_config_dir(project / ".aegis") + return {"home": home, "project": project, "tmp": tmp_path} + + +def write_budget(env, text: str) -> Path: + path = env["project"] / ".aegis" / "budget.yaml" + path.write_text(text) + sign_file(path, load_key(f"file:{env['project'] / '.aegis' / 'example-signing.key'}")) + return path + + +BUDGET = "principal: admin\nunit: usd\nsession: {limit: 1.0, warn_at: 0.5}\n" + + +def transcript(env, input_tokens: int, sid: str = SID) -> Path: + path = env["home"] / ".claude" / "projects" / "-project" / f"{sid}.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + recs = [{"type": "user", "sessionId": sid, "cwd": str(env["project"]), + "timestamp": f"{DAY}T09:00:00Z"}] + if input_tokens: + recs.append({"type": "assistant", "sessionId": sid, "requestId": "r1", + "timestamp": f"{DAY}T09:01:00Z", + "message": {"id": "m1", "model": MODEL, + "usage": {"input_tokens": input_tokens, "output_tokens": 0}}}) + path.write_text("".join(json.dumps(r) + "\n" for r in recs)) + return path + + +def call(env, tool="Edit", command=None, agent="claude", event=None, sid=SID, **extra): + payload = {"session_id": sid, "cwd": str(env["project"]), + "transcript_path": str(env["home"] / ".claude" / "projects" / "-project" + / f"{sid}.jsonl"), **extra} + if event: + payload["hook_event_name"] = event + payload["prompt"] = "hello" + else: + payload["tool_name"] = tool + payload["tool_input"] = {"command": command} if command else {"file_path": "x"} + return run_hook(agent, json.dumps(payload)) + + +# --- no budget: unchanged behaviour ------------------------------------------------------- + + +def test_without_a_budget_non_shell_tools_are_not_looked_at(env): + transcript(env, 10_000_000) + assert call(env) == (0, "", "") + assert not (env["home"] / ".cache" / "aegis").exists() + + +# --- under, near and over the limit --------------------------------------------------------- + + +def test_under_the_warning_threshold_is_silent(env): + write_budget(env, BUDGET) + transcript(env, 25_000) # $0.10 + assert call(env) == (0, "", "") + + +def test_crossing_the_warning_threshold_warns_once(env): + write_budget(env, BUDGET) + transcript(env, 150_000) # $0.60 of $1.00 + code, out, _ = call(env) + assert code == 0 + message = json.loads(out)["systemMessage"] + assert "60%" in message and "$0.40 left" in message + assert call(env) == (0, "", "") # once per session + + +def test_at_the_limit_every_tool_call_is_denied(env): + write_budget(env, BUDGET) + transcript(env, 300_000) # $1.20 + for tool, command in (("Edit", None), ("Read", None), ("Bash", "ls -la")): + code, out, err = call(env, tool=tool, command=command) + assert code == 2 + decision = json.loads(out)["hookSpecificOutput"] + assert decision["permissionDecision"] == "deny" + assert "session limit reached" in err and "start a new session" in err + assert "overshoot $0.20" in err + + +def test_a_policy_block_stands_whatever_the_budget(env): + write_budget(env, BUDGET) + transcript(env, 1_000) + code, _, err = call(env, tool="Bash", command="kubectl delete namespace prod") + assert code == 2 and "budget" not in err and "BLOCK" in err + + +# --- prompts -------------------------------------------------------------------------------- + + +def test_a_new_prompt_over_the_limit_is_stopped(env): + write_budget(env, BUDGET) + transcript(env, 300_000) + code, out, err = call(env, event="UserPromptSubmit") + assert code == 2 and json.loads(out) == {"decision": "block", "reason": err} + assert "session limit reached" in err + + +@pytest.mark.parametrize("agent,decision", [("claude", "block"), ("codex", "block"), + ("gemini", "deny")]) +def test_prompt_replies_in_each_agents_format(agent, decision): + from aegis_core.hook import HookVerdict, render_prompt + + code, out, err = render_prompt(agent, HookVerdict("deny", "over")) + assert (code, json.loads(out), err) == (2, {"decision": decision, "reason": "over"}, "over") + assert render_prompt(agent, HookVerdict("allow"), "near") == ( + 0, json.dumps({"systemMessage": "near"}), "") + assert render_prompt(agent, HookVerdict("allow")) == (0, "", "") + + +def test_a_prompt_under_the_limit_passes_silently(env): + write_budget(env, BUDGET) + transcript(env, 1_000) + assert call(env, event="UserPromptSubmit") == (0, "", "") + + +# --- failures ------------------------------------------------------------------------------- + + +def test_an_unusable_budget_file_fails_closed(env): + path = write_budget(env, BUDGET) + path.write_text(BUDGET.replace("1.0", "1000")) # edited after signing + transcript(env, 1_000) + code, _, err = call(env) + assert code == 2 and "budget.yaml cannot be used" in err and "re-sign" in err + + +def test_an_unlocatable_session_warns_once_or_denies(env): + write_budget(env, BUDGET) + code, out, _ = call(env, sid="nobody") + assert code == 0 and "not being measured" in json.loads(out)["systemMessage"] + assert call(env, sid="nobody") == (0, "", "") + write_budget(env, BUDGET + "on_unknown_log: deny\n") + code, _, err = call(env, sid="nobody") + assert code == 2 and "on_unknown_log: deny" in err + + +def test_unlisted_agents_are_not_measured(env): + write_budget(env, BUDGET + "agents: [codex]\n") + transcript(env, 300_000) + assert call(env) == (0, "", "") + + +# --- warning channels ----------------------------------------------------------------------- + + +def test_copilot_has_no_message_channel_so_its_warning_goes_to_stderr(env): + from aegis_core.hook import _with_warning + + assert _with_warning("copilot", (0, "", ""), "aegis budget: 80%") == ( + 0, "", "aegis budget: 80%") + code, out, _ = _with_warning("opencode", (0, "", ""), "aegis budget: 80%") + assert json.loads(out) == {"warning": "aegis budget: 80%"} + code, out, _ = _with_warning("claude", (0, '{"hookSpecificOutput": {}}', ""), "w") + assert json.loads(out) == {"hookSpecificOutput": {}, "systemMessage": "w"} + assert _with_warning("claude", (2, "x", "y"), "w") == (2, "x", "y") + + +# --- install ---------------------------------------------------------------------------------- + + +def _settings(env) -> dict: + return json.loads((env["project"] / ".claude" / "settings.json").read_text()) + + +def test_install_with_a_budget_widens_the_hook_and_adds_the_prompt_hook(env, capsys): + write_budget(env, BUDGET) + main_install("claude", user=False, project=str(env["project"]), remove=False) + hooks = _settings(env)["hooks"] + assert [e["matcher"] for e in hooks["PreToolUse"]] == ["*"] + assert len(hooks["UserPromptSubmit"]) == 1 + assert "UserPromptSubmit stops new prompts" in capsys.readouterr().out + # reinstalling without the budget narrows it again and drops the prompt hook + main_install("claude", user=False, project=str(env["project"]), remove=False, + budget=False) + hooks = _settings(env)["hooks"] + assert [e["matcher"] for e in hooks["PreToolUse"]] == ["Bash"] + assert "UserPromptSubmit" not in hooks + install("claude", user=False, project=env["project"], remove=True) + assert not (env["project"] / ".claude" / "settings.json").exists() + + +@pytest.mark.parametrize("agent,tool_event,matcher,prompt_event", [ + ("codex", "PreToolUse", ".*", "UserPromptSubmit"), + ("gemini", "BeforeTool", ".*", "BeforeAgent"), +]) +def test_install_budget_shapes_for_codex_and_gemini(env, agent, tool_event, matcher, + prompt_event): + path = install(agent, user=False, project=env["project"], budget=True) + hooks = json.loads(path.read_text())["hooks"] + assert [e["matcher"] for e in hooks[tool_event]] == [matcher] + assert len(hooks[prompt_event]) == 1 + + +def test_install_keeps_other_hooks_when_switching_budget_on_and_off(env): + settings = env["project"] / ".claude" / "settings.json" + settings.parent.mkdir() + mine = {"type": "command", "command": "echo mine"} + settings.write_text(json.dumps({"hooks": {"UserPromptSubmit": [{"hooks": [mine]}]}})) + install("claude", user=False, project=env["project"], budget=True) + install("claude", user=False, project=env["project"], budget=False) + assert _settings(env)["hooks"]["UserPromptSubmit"] == [{"hooks": [mine]}] diff --git a/tests/test_plugins.py b/tests/test_plugins.py index be5a502..58f9def 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -4,6 +4,7 @@ import json import os +import re import subprocess from pathlib import Path @@ -55,9 +56,16 @@ def test_claude_marketplace_points_at_the_plugin_folder(): def test_claude_plugin_hook_runs_the_wrapper(): hooks = json.loads((CLAUDE / "hooks" / "hooks.json").read_text()) - [entry] = hooks["hooks"]["PreToolUse"] - assert entry["matcher"] == "Bash" - assert entry["hooks"][0]["command"] == '"${CLAUDE_PLUGIN_ROOT}"/hooks/aegis-hook.sh claude' + shell, others = hooks["hooks"]["PreToolUse"] + assert shell["matcher"] == "Bash" + assert shell["hooks"][0]["command"] == '"${CLAUDE_PLUGIN_ROOT}"/hooks/aegis-hook.sh claude' + # budget cap: every other tool, and new prompts, through the budget-only fast path + budget_only = '"${CLAUDE_PLUGIN_ROOT}"/hooks/aegis-hook.sh claude --budget-only' + assert others["matcher"] == "^(?!Bash$).*" + assert others["hooks"][0]["command"] == budget_only + assert not re.search(others["matcher"], "Bash") and re.search(others["matcher"], "Edit") + [prompt] = hooks["hooks"]["UserPromptSubmit"] + assert prompt["hooks"][0]["command"] == budget_only def test_gemini_extension_manifest(): @@ -70,12 +78,18 @@ def test_gemini_extension_hook_runs_the_wrapper(): # Gemini CLI reads an extension's hooks from hooks/hooks.json at its root; # nothing else may live there (Gemini warns about unknown events) hooks = json.loads((ROOT / "hooks" / "hooks.json").read_text()) - assert set(hooks["hooks"]) == {"BeforeTool"} - [entry] = hooks["hooks"]["BeforeTool"] + assert set(hooks["hooks"]) == {"BeforeTool", "BeforeAgent"} + entry, others = hooks["hooks"]["BeforeTool"] assert entry["matcher"] == "run_shell_command" [h] = entry["hooks"] assert h["command"] == '"${extensionPath}${/}hooks${/}aegis-hook.sh" gemini' assert h["timeout"] == 30000 + budget_only = '"${extensionPath}${/}hooks${/}aegis-hook.sh" gemini --budget-only' + assert others["matcher"] == "^(?!run_shell_command$).*" + assert others["hooks"][0]["command"] == budget_only + assert not re.search(others["matcher"], "run_shell_command") + [prompt] = hooks["hooks"]["BeforeAgent"] + assert prompt["hooks"][0]["command"] == budget_only # --- the wrapper with no aegis installed -------------------------------------------- @@ -110,3 +124,42 @@ def test_wrapper_without_aegis_blocks_only_infra_in_opted_in_projects(tmp_path): (tmp_path / "home").rename(tmp_path / "old-home") (tmp_path / "project").rename(tmp_path / "old-project") assert _run_wrapper(tmp_path, "ls -la", opted_in=True) == (0, "") + + +# --- the wrapper's budget-only fast path -------------------------------------------------- + + +def _run_budget_only(tmp_path, *, budget: bool, aegis_ok: bool = True): + home, project = tmp_path / "home", tmp_path / "project" + home.mkdir() + (project / ".aegis").mkdir(parents=True) + if budget: + (project / ".aegis" / "budget.yaml").write_text("unit: usd\n") + marker = tmp_path / "aegis-ran" + fake = tmp_path / "aegis" + fake.write_text(f"#!/bin/sh\ncat >/dev/null\ntouch {marker}\nexit 0\n") + fake.chmod(0o755) + env = {"HOME": str(home), "PATH": "/usr/bin:/bin", + **({"AEGIS_BIN": str(fake)} if aegis_ok else {})} + payload = json.dumps({"tool_name": "Edit", "tool_input": {}, "session_id": "s"}) + proc = subprocess.run(["bash", str(WRAPPER), "claude", "--budget-only"], input=payload, + capture_output=True, text=True, cwd=project, env=env, check=False) + return proc.returncode, marker.exists() + + +@pytest.mark.skipif(os.name == "nt", reason="bash wrapper") +def test_budget_only_entries_cost_nothing_without_a_budget(tmp_path): + """No budget.yaml: the wrapper exits before starting aegis at all.""" + assert _run_budget_only(tmp_path, budget=False) == (0, False) + + +@pytest.mark.skipif(os.name == "nt", reason="bash wrapper") +def test_budget_only_entries_run_aegis_where_a_budget_applies(tmp_path): + assert _run_budget_only(tmp_path, budget=True) == (0, True) + + +@pytest.mark.skipif(os.name == "nt", reason="bash wrapper") +def test_budget_only_entries_allow_when_aegis_is_missing(tmp_path): + """Without aegis nothing can be measured; only shell commands are gated + by the wrapper's own fallback.""" + assert _run_budget_only(tmp_path, budget=True, aegis_ok=False) == (0, False) From cb43401c2fbb71e3bdde015c38919d41f348c4fe Mon Sep 17 00:00:00 2001 From: moneytools Date: Thu, 8 Oct 2026 21:55:34 -0500 Subject: [PATCH 2/3] fix(budget): discover budgets the way the hook does; honour --no-budget for OpenCode - wrapper --budget-only: decide from every cwd in the payload, the process directory and the project variables, starting aegis if any has a budget.yaml (or a cwd cannot be read); a cd into a worktree with a budget no longer bypasses it - install: detect the budget in the policy directory the hook would use (AEGIS_CONFIG_DIR, nearest .aegis/ upwards, ~/.config/aegis), never past a directory that shadows another - OpenCode plugin: BUDGET_MODE written at install (auto, or off with --no-budget); off sends only bash to aegis Review of #39. --- hooks/aegis-hook.sh | 38 +++++++++++--- plugins/claude/hooks/aegis-hook.sh | 38 +++++++++++--- plugins/copilot/aegis-hook.sh | 38 +++++++++++--- src/aegis_core/hook.py | 52 +++++++++++++------ .../integrations/opencode-plugin.js | 5 +- tests/test_budget_hook.py | 39 ++++++++++++++ tests/test_opencode_plugin.py | 37 +++++++++++++ tests/test_plugins.py | 50 ++++++++++++++++++ 8 files changed, 262 insertions(+), 35 deletions(-) diff --git a/hooks/aegis-hook.sh b/hooks/aegis-hook.sh index 7c9c508..84016e1 100755 --- a/hooks/aegis-hook.sh +++ b/hooks/aegis-hook.sh @@ -18,11 +18,11 @@ set -u agent=${1:-claude} mode=${2:-} -# The policy directory aegis would use: $AEGIS_CONFIG_DIR, the nearest .aegis/ -# at or above the project, or ~/.config/aegis (as aegis_core/hook.py does). +# The policy directory aegis would use from directory $1: $AEGIS_CONFIG_DIR, +# the nearest .aegis/ at or above $1, or ~/.config/aegis (as hook.py does). policy_dir() { if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi - d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + d=$1 while [ -n "$d" ]; do if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi [ "$d" = "/" ] && break @@ -32,11 +32,37 @@ policy_dir() { return 1 } +budget_applies_from() { + [ -n "$1" ] || return 1 + dir=$(policy_dir "$1") || return 1 + [ -f "$dir/budget.yaml" ] +} + +payload=$(cat) if [ "$mode" = "--budget-only" ]; then - dir=$(policy_dir) || exit 0 - [ -f "$dir/budget.yaml" ] || exit 0 + # aegis decides from the payload's cwd first (the agent's directory may have + # changed since the session started, e.g. into a worktree), then the + # project variables. Skip starting aegis only when none of those has a + # budget.yaml; if the payload names a cwd this cannot read, start it. + found=1 + cwds=$(printf '%s' "$payload" | grep -o '"cwd"[[:space:]]*:[[:space:]]*"[^"\\]*"' \ + | sed 's/.*:[[:space:]]*"\(.*\)"$/\1/') + ncwd=$(printf '%s' "$payload" | grep -o '"cwd"' | wc -l | tr -d ' ') + if [ "$ncwd" != "$(printf '%s' "$cwds" | grep -c .)" ]; then + found=0 # a cwd this cannot read (escaped characters): let aegis decide + else + while IFS= read -r d; do + if budget_applies_from "$d"; then found=0; break; fi + done </dev/null 2>&1; then aegis=aegis; fi diff --git a/plugins/claude/hooks/aegis-hook.sh b/plugins/claude/hooks/aegis-hook.sh index 7c9c508..84016e1 100755 --- a/plugins/claude/hooks/aegis-hook.sh +++ b/plugins/claude/hooks/aegis-hook.sh @@ -18,11 +18,11 @@ set -u agent=${1:-claude} mode=${2:-} -# The policy directory aegis would use: $AEGIS_CONFIG_DIR, the nearest .aegis/ -# at or above the project, or ~/.config/aegis (as aegis_core/hook.py does). +# The policy directory aegis would use from directory $1: $AEGIS_CONFIG_DIR, +# the nearest .aegis/ at or above $1, or ~/.config/aegis (as hook.py does). policy_dir() { if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi - d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + d=$1 while [ -n "$d" ]; do if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi [ "$d" = "/" ] && break @@ -32,11 +32,37 @@ policy_dir() { return 1 } +budget_applies_from() { + [ -n "$1" ] || return 1 + dir=$(policy_dir "$1") || return 1 + [ -f "$dir/budget.yaml" ] +} + +payload=$(cat) if [ "$mode" = "--budget-only" ]; then - dir=$(policy_dir) || exit 0 - [ -f "$dir/budget.yaml" ] || exit 0 + # aegis decides from the payload's cwd first (the agent's directory may have + # changed since the session started, e.g. into a worktree), then the + # project variables. Skip starting aegis only when none of those has a + # budget.yaml; if the payload names a cwd this cannot read, start it. + found=1 + cwds=$(printf '%s' "$payload" | grep -o '"cwd"[[:space:]]*:[[:space:]]*"[^"\\]*"' \ + | sed 's/.*:[[:space:]]*"\(.*\)"$/\1/') + ncwd=$(printf '%s' "$payload" | grep -o '"cwd"' | wc -l | tr -d ' ') + if [ "$ncwd" != "$(printf '%s' "$cwds" | grep -c .)" ]; then + found=0 # a cwd this cannot read (escaped characters): let aegis decide + else + while IFS= read -r d; do + if budget_applies_from "$d"; then found=0; break; fi + done </dev/null 2>&1; then aegis=aegis; fi diff --git a/plugins/copilot/aegis-hook.sh b/plugins/copilot/aegis-hook.sh index 7c9c508..84016e1 100755 --- a/plugins/copilot/aegis-hook.sh +++ b/plugins/copilot/aegis-hook.sh @@ -18,11 +18,11 @@ set -u agent=${1:-claude} mode=${2:-} -# The policy directory aegis would use: $AEGIS_CONFIG_DIR, the nearest .aegis/ -# at or above the project, or ~/.config/aegis (as aegis_core/hook.py does). +# The policy directory aegis would use from directory $1: $AEGIS_CONFIG_DIR, +# the nearest .aegis/ at or above $1, or ~/.config/aegis (as hook.py does). policy_dir() { if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi - d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + d=$1 while [ -n "$d" ]; do if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi [ "$d" = "/" ] && break @@ -32,11 +32,37 @@ policy_dir() { return 1 } +budget_applies_from() { + [ -n "$1" ] || return 1 + dir=$(policy_dir "$1") || return 1 + [ -f "$dir/budget.yaml" ] +} + +payload=$(cat) if [ "$mode" = "--budget-only" ]; then - dir=$(policy_dir) || exit 0 - [ -f "$dir/budget.yaml" ] || exit 0 + # aegis decides from the payload's cwd first (the agent's directory may have + # changed since the session started, e.g. into a worktree), then the + # project variables. Skip starting aegis only when none of those has a + # budget.yaml; if the payload names a cwd this cannot read, start it. + found=1 + cwds=$(printf '%s' "$payload" | grep -o '"cwd"[[:space:]]*:[[:space:]]*"[^"\\]*"' \ + | sed 's/.*:[[:space:]]*"\(.*\)"$/\1/') + ncwd=$(printf '%s' "$payload" | grep -o '"cwd"' | wc -l | tr -d ' ') + if [ "$ncwd" != "$(printf '%s' "$cwds" | grep -c .)" ]; then + found=0 # a cwd this cannot read (escaped characters): let aegis decide + else + while IFS= read -r d; do + if budget_applies_from "$d"; then found=0; break; fi + done </dev/null 2>&1; then aegis=aegis; fi diff --git a/src/aegis_core/hook.py b/src/aegis_core/hook.py index 36ccd9d..0cc2d7e 100644 --- a/src/aegis_core/hook.py +++ b/src/aegis_core/hook.py @@ -580,13 +580,18 @@ def merge_config(agent: str, existing: dict, command: str | None, *, def budget_configured(*, user: bool, project: Path) -> bool: - """Whether the policy that would apply has a ``budget.yaml``: the - project's ``.aegis/`` (or ``$AEGIS_CONFIG_DIR``), or ``~/.config/aegis`` - for a user-level install.""" - env = os.environ.get(config_module.CONFIG_ENV_VAR) - dirs = [Path(env)] if env else [] - dirs.append(Path.home() / ".config" / "aegis" if user else project / ".aegis") - return any((d / BUDGET_FILE).is_file() for d in dirs) + """Whether the policy directory that would apply holds a ``budget.yaml``, + resolved with the hook's own precedence and never by looking past a + directory that shadows another: for a project install, + ``$AEGIS_CONFIG_DIR``, else the nearest ``.aegis/`` at or above the + project, else ``~/.config/aegis`` (:func:`find_opt_in_config`); for a + user install, ``$AEGIS_CONFIG_DIR``, else ``~/.config/aegis``.""" + if user: + env = os.environ.get(config_module.CONFIG_ENV_VAR) + directory: Path | None = Path(env) if env else Path.home() / ".config" / "aegis" + else: + directory = find_opt_in_config(project) + return directory is not None and (directory / BUDGET_FILE).is_file() _AFTER_INSTALL = { @@ -602,9 +607,11 @@ def budget_configured(*, user: bool, project: Path) -> bool: } -def _install_opencode_plugin(path: Path, *, remove: bool) -> Path: +def _install_opencode_plugin(path: Path, *, remove: bool, budget: bool | None = None) -> Path: """OpenCode hooks are JS plugins: write (or delete) our plugin file, with - the argv of this aegis filled in.""" + the argv of this aegis and the budget coverage filled in. The plugin + detects a ``budget.yaml`` at run time, so the default is ``"auto"``; + ``--no-budget`` (``budget=False``) writes ``"off"``: shell commands only.""" if remove: if path.exists(): path.unlink() @@ -615,16 +622,26 @@ def _install_opencode_plugin(path: Path, *, remove: bool) -> Path: marker = "const INSTALLED = null" if marker not in template: raise ValueError("opencode plugin template has no INSTALLED marker") + mode_marker = 'const BUDGET_MODE = "auto"' + if mode_marker not in template: + raise ValueError("opencode plugin template has no BUDGET_MODE marker") + mode = "off" if budget is False else "auto" path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(template.replace(marker, f"const INSTALLED = {json.dumps(_aegis_command())}")) + path.write_text(template.replace(marker, f"const INSTALLED = {json.dumps(_aegis_command())}") + .replace(mode_marker, f"const BUDGET_MODE = {json.dumps(mode)}")) return path def install(agent: str, *, user: bool, project: Path, remove: bool = False, - budget: bool = False) -> Path: + budget: bool | None = None) -> Path: + """Writes (or removes) the agent's hook. ``budget``: ``True``/``False`` + install or leave out the budget hooks; ``None`` detects whether the + policy directory that applies holds a ``budget.yaml``.""" path = config_path(agent, user=user, project=project) if agent == "opencode": - return _install_opencode_plugin(path, remove=remove) + return _install_opencode_plugin(path, remove=remove, budget=budget) + if budget is None: + budget = budget_configured(user=user, project=project) existing: dict = {} if path.exists(): text = path.read_text() @@ -653,16 +670,19 @@ def install(agent: str, *, user: bool, project: Path, remove: bool = False, "the limit", "copilot": "budget: the hook already sees every tool call (premium requests are " "counted); Copilot CLI's prompt hook cannot block, so prompts are not stopped", - "opencode": "budget: the plugin checks every tool call in projects with a budget.yaml; " - "OpenCode has no prompt hook, so prompts are not stopped", + "opencode": "budget: the plugin checks every tool call wherever a budget.yaml applies " + "(--no-budget: shell commands only); OpenCode has no prompt hook, so prompts " + "are not stopped", } def main_install(agent: str, *, user: bool, project: str, remove: bool, budget: bool | None = None) -> int: - if budget is None: - budget = budget_configured(user=user, project=Path(project)) path = install(agent, user=user, project=Path(project), remove=remove, budget=budget) + if agent == "opencode": + budget = budget is not False # the plugin detects a budget.yaml at run time + elif budget is None: + budget = budget_configured(user=user, project=Path(project)) if remove: print(f"aegis: removed the {agent} hook from {path}") return 0 diff --git a/src/aegis_core/integrations/opencode-plugin.js b/src/aegis_core/integrations/opencode-plugin.js index ab33a93..bd112f3 100644 --- a/src/aegis_core/integrations/opencode-plugin.js +++ b/src/aegis_core/integrations/opencode-plugin.js @@ -22,6 +22,9 @@ import { dirname, join, resolve } from "node:path" // `aegis install opencode` replaces null with the argv that starts the aegis // it was run from (GUI apps do not see your shell's PATH). const INSTALLED = null +// `aegis install opencode --no-budget` replaces "auto" with "off": then only bash +// commands go to aegis, whatever budget.yaml says. +const BUDGET_MODE = "auto" const AEGIS = process.env.AEGIS_BIN ? [process.env.AEGIS_BIN] : INSTALLED || ["aegis"] const INFRA = /(^|[^\w.-])(kubectl|k|terraform|tf|tofu|aws|az|gcloud|gsutil|helm|argocd|flux|git|gh|psql|mysql|sqlite3|mongosh|pulumi|alembic|flyway|rails|prisma)([^\w.-]|$)/ @@ -64,7 +67,7 @@ export const AegisDevOps = async ({ directory, client }) => ({ const isBash = input?.tool === "bash" const command = output?.args?.command if (isBash && (typeof command !== "string" || !command.trim())) return - if (!isBash && !hasBudget(cwd)) return + if (!isBash && (BUDGET_MODE === "off" || !hasBudget(cwd))) return const payload = JSON.stringify({ tool_name: input?.tool, tool_input: isBash ? { command } : output?.args || {}, diff --git a/tests/test_budget_hook.py b/tests/test_budget_hook.py index 8ea7c9e..1588baf 100644 --- a/tests/test_budget_hook.py +++ b/tests/test_budget_hook.py @@ -233,3 +233,42 @@ def test_install_keeps_other_hooks_when_switching_budget_on_and_off(env): install("claude", user=False, project=env["project"], budget=True) install("claude", user=False, project=env["project"], budget=False) assert _settings(env)["hooks"]["UserPromptSubmit"] == [{"hooks": [mine]}] + + +# --- install-time detection follows the hook's precedence (review of #39) -------------------- + + +def _matchers(env, project: Path) -> list[str]: + path = install("claude", user=False, project=project) + return [e["matcher"] for e in json.loads(path.read_text())["hooks"]["PreToolUse"]] + + +def test_a_nested_project_inherits_its_parents_budget(env): + write_budget(env, BUDGET) + nested = env["project"] / "services" / "api" + nested.mkdir(parents=True) + assert _matchers(env, nested) == ["*"] + + +def test_a_project_without_aegis_inherits_the_user_budget(env, tmp_path): + user_dir = env["home"] / ".config" / "aegis" + init_config_dir(user_dir) + (user_dir / "budget.yaml").write_text(BUDGET) + other = tmp_path / "elsewhere" + other.mkdir() + assert _matchers(env, other) == ["*"] + + +def test_aegis_config_dir_shadows_a_local_budget(env, tmp_path, monkeypatch): + write_budget(env, BUDGET) + shadow = tmp_path / "central" + init_config_dir(shadow) # no budget.yaml there + monkeypatch.setenv("AEGIS_CONFIG_DIR", str(shadow)) + assert _matchers(env, env["project"]) == ["Bash"] + + +def test_a_nearer_aegis_without_a_budget_shadows_the_parents(env): + write_budget(env, BUDGET) + nested = env["project"] / "services" / "api" + init_config_dir(nested / ".aegis") + assert _matchers(env, nested) == ["Bash"] diff --git a/tests/test_opencode_plugin.py b/tests/test_opencode_plugin.py index 7a89d09..b819c69 100644 --- a/tests/test_opencode_plugin.py +++ b/tests/test_opencode_plugin.py @@ -85,3 +85,40 @@ def test_missing_aegis_blocks_only_infra_in_opted_in_projects(env): assert blocked["allowed"] is False assert "pip install aegis-devops" in blocked["message"] assert _call(env, "ls -la", missing) == {"allowed": True} + + +TOOL_HARNESS = """ +import { pathToFileURL } from "node:url" +const [pluginPath, directory, tool] = process.argv.slice(2) +const { AegisDevOps } = await import(pathToFileURL(pluginPath).href) +const hooks = await AegisDevOps({ directory, client: {} }) +try { + await hooks["tool.execute.before"]({ tool, sessionID: "ses_1" }, { args: { filePath: "x" } }) + console.log(JSON.stringify({ allowed: true })) +} catch (e) { + console.log(JSON.stringify({ allowed: false, message: e.message })) +} +""" + + +@pytest.mark.parametrize("budget,expect_aegis", [(None, True), (True, True), (False, False)]) +def test_no_budget_keeps_non_bash_tools_away_from_aegis(env, tmp_path, budget, expect_aegis): + """Review of #39: --no-budget must hold for OpenCode too. A fake aegis + that denies everything shows whether a read call reached it.""" + import os + + main(["init", str(env["project"] / ".aegis")]) + (env["project"] / ".aegis" / "budget.yaml").write_text("unit: usd\n") + plugin = hook.install("opencode", user=False, project=env["project"], budget=budget) + mjs = tmp_path / f"plugin-{budget}.mjs" + mjs.write_text(plugin.read_text()) + harness = tmp_path / "tool-harness.mjs" + harness.write_text(TOOL_HARNESS) + fake = tmp_path / "deny-all" + fake.write_text("#!/bin/sh\ncat >/dev/null\necho denied >&2\nexit 2\n") + fake.chmod(0o755) + out = subprocess.run([NODE, str(harness), str(mjs), str(env["project"]), "read"], + capture_output=True, text=True, check=True, timeout=60, + env={**os.environ, "AEGIS_BIN": str(fake)}) + result = json.loads(out.stdout.strip().splitlines()[-1]) + assert result["allowed"] is (not expect_aegis) diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 58f9def..d03868d 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -163,3 +163,53 @@ def test_budget_only_entries_allow_when_aegis_is_missing(tmp_path): """Without aegis nothing can be measured; only shell commands are gated by the wrapper's own fallback.""" assert _run_budget_only(tmp_path, budget=True, aegis_ok=False) == (0, False) + + +def _wrapper_starts_aegis(tmp_path, *, budget_in: str, cwd: str, payload_cwd: str | None, + project_env: str | None, raw_payload: str | None = None) -> bool: + """Runs the budget-only wrapper with a fake aegis; True if it was started.""" + dirs = {"orig": tmp_path / "orig", "wt": tmp_path / "work tree"} + for d in dirs.values(): + (d / ".aegis").mkdir(parents=True, exist_ok=True) + (dirs[budget_in] / ".aegis" / "budget.yaml").write_text("unit: usd\n") if budget_in \ + else None + home = tmp_path / "home" + home.mkdir(exist_ok=True) + marker = tmp_path / "aegis-ran" + marker.unlink(missing_ok=True) + fake = tmp_path / "aegis" + fake.write_text(f"#!/bin/sh\ncat >/dev/null\ntouch '{marker}'\nexit 0\n") + fake.chmod(0o755) + env = {"HOME": str(home), "PATH": "/usr/bin:/bin", "AEGIS_BIN": str(fake)} + if project_env: + env["CLAUDE_PROJECT_DIR"] = str(dirs[project_env]) + payload = raw_payload or json.dumps( + {"tool_name": "Edit", "tool_input": {}, "session_id": "s", + **({"cwd": str(dirs[payload_cwd])} if payload_cwd else {})}) + subprocess.run(["bash", str(WRAPPER), "claude", "--budget-only"], input=payload, + capture_output=True, text=True, cwd=dirs[cwd], env=env, check=False) + return marker.exists() + + +@pytest.mark.skipif(os.name == "nt", reason="bash wrapper") +def test_budget_only_follows_the_hooks_cwd_not_the_original_project(tmp_path): + """Review of #39: CLAUDE_PROJECT_DIR stays at the session's original root + while the payload's cwd follows a cd into a worktree with a budget (whose + path here has a space).""" + assert _wrapper_starts_aegis(tmp_path, budget_in="wt", cwd="wt", payload_cwd="wt", + project_env="orig") + assert _wrapper_starts_aegis(tmp_path, budget_in="wt", cwd="orig", payload_cwd="wt", + project_env="orig") + + +@pytest.mark.skipif(os.name == "nt", reason="bash wrapper") +def test_budget_only_skips_aegis_when_no_candidate_directory_has_a_budget(tmp_path): + assert not _wrapper_starts_aegis(tmp_path, budget_in="", cwd="wt", payload_cwd="wt", + project_env="orig") + + +@pytest.mark.skipif(os.name == "nt", reason="bash wrapper") +def test_budget_only_starts_aegis_when_it_cannot_read_the_cwd(tmp_path): + raw = '{"tool_name": "Edit", "cwd": "/tmp/a\\"b", "session_id": "s"}' + assert _wrapper_starts_aegis(tmp_path, budget_in="", cwd="orig", payload_cwd=None, + project_env=None, raw_payload=raw) From c896de2c173715d5c9d8c4795f74bfee5cb777d9 Mon Sep 17 00:00:00 2001 From: moneytools Date: Thu, 8 Oct 2026 22:00:56 -0500 Subject: [PATCH 3/3] fix(budget): wrapper resolves relative cwds and always ends its ancestor walk A nested {"cwd": "."} reached the walk unresolved and dirname(".") never reached "/". Relative candidates are resolved from the hook's directory, and the walk (and the older opted_in walk) stops when dirname stops changing. Review of #39. --- hooks/aegis-hook.sh | 12 ++++++---- plugins/claude/hooks/aegis-hook.sh | 12 ++++++---- plugins/copilot/aegis-hook.sh | 12 ++++++---- tests/test_plugins.py | 38 ++++++++++++++++++++++++++++++ 4 files changed, 62 insertions(+), 12 deletions(-) diff --git a/hooks/aegis-hook.sh b/hooks/aegis-hook.sh index 84016e1..67f983b 100755 --- a/hooks/aegis-hook.sh +++ b/hooks/aegis-hook.sh @@ -23,10 +23,12 @@ mode=${2:-} policy_dir() { if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi d=$1 + case "$d" in /*) ;; *) d="$PWD/$d" ;; esac # relative: from the hook's directory while [ -n "$d" ]; do if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi - [ "$d" = "/" ] && break - d=$(dirname "$d") + parent=$(dirname "$d") + [ "$parent" = "$d" ] && break # reached the root (or "." / "//") + d=$parent done if [ -d "$HOME/.config/aegis" ]; then printf '%s' "$HOME/.config/aegis"; return 0; fi return 1 @@ -84,10 +86,12 @@ opted_in() { [ -n "${AEGIS_CONFIG_DIR:-}" ] && return 0 [ -d "$HOME/.config/aegis" ] && return 0 d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + case "$d" in /*) ;; *) d="$PWD/$d" ;; esac while [ -n "$d" ]; do [ -d "$d/.aegis" ] && return 0 - [ "$d" = "/" ] && break - d=$(dirname "$d") + parent=$(dirname "$d") + [ "$parent" = "$d" ] && break + d=$parent done return 1 } diff --git a/plugins/claude/hooks/aegis-hook.sh b/plugins/claude/hooks/aegis-hook.sh index 84016e1..67f983b 100755 --- a/plugins/claude/hooks/aegis-hook.sh +++ b/plugins/claude/hooks/aegis-hook.sh @@ -23,10 +23,12 @@ mode=${2:-} policy_dir() { if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi d=$1 + case "$d" in /*) ;; *) d="$PWD/$d" ;; esac # relative: from the hook's directory while [ -n "$d" ]; do if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi - [ "$d" = "/" ] && break - d=$(dirname "$d") + parent=$(dirname "$d") + [ "$parent" = "$d" ] && break # reached the root (or "." / "//") + d=$parent done if [ -d "$HOME/.config/aegis" ]; then printf '%s' "$HOME/.config/aegis"; return 0; fi return 1 @@ -84,10 +86,12 @@ opted_in() { [ -n "${AEGIS_CONFIG_DIR:-}" ] && return 0 [ -d "$HOME/.config/aegis" ] && return 0 d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + case "$d" in /*) ;; *) d="$PWD/$d" ;; esac while [ -n "$d" ]; do [ -d "$d/.aegis" ] && return 0 - [ "$d" = "/" ] && break - d=$(dirname "$d") + parent=$(dirname "$d") + [ "$parent" = "$d" ] && break + d=$parent done return 1 } diff --git a/plugins/copilot/aegis-hook.sh b/plugins/copilot/aegis-hook.sh index 84016e1..67f983b 100755 --- a/plugins/copilot/aegis-hook.sh +++ b/plugins/copilot/aegis-hook.sh @@ -23,10 +23,12 @@ mode=${2:-} policy_dir() { if [ -n "${AEGIS_CONFIG_DIR:-}" ]; then printf '%s' "$AEGIS_CONFIG_DIR"; return 0; fi d=$1 + case "$d" in /*) ;; *) d="$PWD/$d" ;; esac # relative: from the hook's directory while [ -n "$d" ]; do if [ -d "$d/.aegis" ]; then printf '%s' "$d/.aegis"; return 0; fi - [ "$d" = "/" ] && break - d=$(dirname "$d") + parent=$(dirname "$d") + [ "$parent" = "$d" ] && break # reached the root (or "." / "//") + d=$parent done if [ -d "$HOME/.config/aegis" ]; then printf '%s' "$HOME/.config/aegis"; return 0; fi return 1 @@ -84,10 +86,12 @@ opted_in() { [ -n "${AEGIS_CONFIG_DIR:-}" ] && return 0 [ -d "$HOME/.config/aegis" ] && return 0 d=${CLAUDE_PROJECT_DIR:-${COPILOT_PROJECT_DIR:-${GEMINI_PROJECT_DIR:-$PWD}}} + case "$d" in /*) ;; *) d="$PWD/$d" ;; esac while [ -n "$d" ]; do [ -d "$d/.aegis" ] && return 0 - [ "$d" = "/" ] && break - d=$(dirname "$d") + parent=$(dirname "$d") + [ "$parent" = "$d" ] && break + d=$parent done return 1 } diff --git a/tests/test_plugins.py b/tests/test_plugins.py index d03868d..4110b30 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -213,3 +213,41 @@ def test_budget_only_starts_aegis_when_it_cannot_read_the_cwd(tmp_path): raw = '{"tool_name": "Edit", "cwd": "/tmp/a\\"b", "session_id": "s"}' assert _wrapper_starts_aegis(tmp_path, budget_in="", cwd="orig", payload_cwd=None, project_env=None, raw_payload=raw) + + +@pytest.mark.skipif(os.name == "nt", reason="bash wrapper") +@pytest.mark.parametrize("rel", [".", "..", "sub/dir", "./x/../y"]) +def test_budget_only_handles_relative_cwds_in_bounded_time(tmp_path, rel): + """Review of #39: a nested tool argument like {"cwd": "."} reached the + ancestor walk unresolved, and dirname(".") == "." never ended. Relative + candidates are resolved from the hook's directory; the walk stops when + dirname stops changing.""" + project = tmp_path / "project" + project.mkdir() + home = tmp_path / "home" + home.mkdir() + payload = json.dumps({"tool_name": "Edit", "cwd": str(project), "session_id": "s", + "tool_input": {"cwd": rel}}) + proc = subprocess.run(["bash", str(WRAPPER), "claude", "--budget-only"], input=payload, + capture_output=True, text=True, cwd=project, timeout=10, + env={"HOME": str(home), "PATH": "/usr/bin:/bin"}, check=False) + assert proc.returncode == 0 + + +@pytest.mark.skipif(os.name == "nt", reason="bash wrapper") +def test_a_relative_cwd_still_finds_a_budget_above_it(tmp_path): + project = tmp_path / "project" + (project / ".aegis").mkdir(parents=True) + (project / ".aegis" / "budget.yaml").write_text("unit: usd\n") + (project / "sub").mkdir() + home = tmp_path / "home" + home.mkdir() + marker = tmp_path / "ran" + fake = tmp_path / "aegis" + fake.write_text(f"#!/bin/sh\ncat >/dev/null\ntouch '{marker}'\n") + fake.chmod(0o755) + payload = json.dumps({"tool_name": "Edit", "session_id": "s", "cwd": "sub"}) + subprocess.run(["bash", str(WRAPPER), "claude", "--budget-only"], input=payload, + capture_output=True, text=True, cwd=project, timeout=10, check=False, + env={"HOME": str(home), "PATH": "/usr/bin:/bin", "AEGIS_BIN": str(fake)}) + assert marker.exists()