Enforce USD budget at each agent invocation (incl. pruning): skip pruning when tripped, fail immediately with no retries at zero/negative balance #76

New issue

Open

opened 2026-07-07 10:06:46 +00:00 by Graa · 0 comments

Graa commented

2026-07-07 10:06:46 +00:00

Metadata

Branch name: bugfix/m2-76-budget-pre-flight
Commit description: fix(graph): enforce USD budget at each agent invocation with pre-flight gate

Summary

Cost/credit enforcement (ADR-2029, issue #15) is currently a post-node accumulate-then-check guard living entirely in langgraph/pure_graph.py. On each agent node the executor prices the node's reported token usage, adds it to self._accumulated_cost, and only then compares against max_cost_usd (main path pure_graph.py:916-960, streaming path :1722-1750, plus the interactive path). The ExecutionError(kind="cost", reason="budget_exhausted") is re-raised and never swallowed, so the graph halts before the next node.

Two gaps follow from checking only at the node boundary:

Pruning sub-calls are never budget-gated. Tool-output pruning makes an extra LLM call inside the node's tool loop (agents/llm.py:636 _run_pruning_pass → :704-708 _prune_invoke → prune_model.ainvoke). This spend happens inside the node and is only reconciled after the node returns. So once the budget is already tripped we can still burn credit on pruning passes for the remainder of that node.
No pre-flight gate before (re)invoking the main agent. If the tenant balance is already at zero or negative when a node is about to run, the executor still starts the node (and its tool loop / retries via agents/retry.py) and only fails after incurring a full node's cost. There is no "we're already broke, don't start" check.

(Note: the token_budget_percent logic in agents/llm.py:281-294, 833-860 is a context-window cap, not a dollar cap — unrelated to credit enforcement.)

Requested behavior

Enforce the USD budget at each agent invocation, including pruning calls, not just at the node boundary:

Gate every LLM invocation inside a node against the budget. Before each main-agent round and before each pruning pass, check whether max_cost_usd has already been tripped by _accumulated_cost.
If the budget is tripped, stop pruning. Do not issue further pruning passes (_run_pruning_pass) once we are at/over budget — pruning is optimization spend and must not push us further past the ceiling. Fall back to the unpruned tool output.
Pre-flight the main agent. If the main agent is about to be invoked but we are already at zero or negative remaining budget, fail the graph execution immediately with ExecutionError(kind="cost", reason="budget_exhausted") and no further retries — do not start the node, do not enter the retry mechanism (agents/retry.py).

Acceptance criteria

A budget check runs before each main-agent ainvoke round and before each pruning pass, not only after the node completes.
When _accumulated_cost has already crossed max_cost_usd, no pruning pass is issued for the remainder of that node; unpruned output is used.
When remaining budget is <= 0 and a main-agent invocation would start, the graph fails immediately with kind="cost", reason="budget_exhausted", and the retry path is not entered (0 retries).
Behavior is consistent across all three executor paths: _execute_from_node, the streaming path (~:1670-1750), and the interactive path.
No silent $0 fallback is introduced (preserve the existing missing_pricing_entry hard-fail contract, pure_graph.py:864-915).
Tests: (a) budget already exhausted → pruning skipped, unpruned output returned; (b) zero/negative balance before a node → immediate budget_exhausted, zero retries, no node work started; (c) budget tripped mid tool-loop → no further pruning for that node.

Notes / open design points

This narrows the overshoot window but does not by itself make spend strictly <= balance within a single main-agent call — a single main invocation still runs to completion before its cost is reconciled. A true "never exceeds balance" guarantee still needs either a router-side safety margin on max_cost_usd sized to worst-case single-call cost, or a pre-flight per-call cost reservation. This ticket covers the pruning-gate + zero-balance-pre-flight asks; a strict per-call reservation can be a follow-up.

Refs: langgraph/pure_graph.py:916-960, :1670-1750; agents/llm.py:636, :704-708, :747-1019; agents/retry.py. Related: ADR-2029, issue #15.

## Metadata **Branch name**: `bugfix/m2-76-budget-pre-flight` **Commit description**: `fix(graph): enforce USD budget at each agent invocation with pre-flight gate` ## Summary Cost/credit enforcement (ADR-2029, issue #15) is currently a **post-node accumulate-then-check** guard living entirely in `langgraph/pure_graph.py`. On each agent node the executor prices the node's reported token usage, adds it to `self._accumulated_cost`, and only then compares against `max_cost_usd` (main path `pure_graph.py:916-960`, streaming path `:1722-1750`, plus the interactive path). The `ExecutionError(kind="cost", reason="budget_exhausted")` is re-raised and never swallowed, so the graph halts before the **next** node. Two gaps follow from checking only at the node boundary: 1. **Pruning sub-calls are never budget-gated.** Tool-output pruning makes an *extra* LLM call inside the node's tool loop (`agents/llm.py:636 _run_pruning_pass` → `:704-708 _prune_invoke` → `prune_model.ainvoke`). This spend happens *inside* the node and is only reconciled after the node returns. So once the budget is already tripped we can still burn credit on pruning passes for the remainder of that node. 2. **No pre-flight gate before (re)invoking the main agent.** If the tenant balance is already at zero or negative when a node is about to run, the executor still starts the node (and its tool loop / retries via `agents/retry.py`) and only fails *after* incurring a full node's cost. There is no "we're already broke, don't start" check. (Note: the `token_budget_percent` logic in `agents/llm.py:281-294, 833-860` is a **context-window** cap, not a dollar cap — unrelated to credit enforcement.) ## Requested behavior Enforce the USD budget at **each agent invocation**, including pruning calls, not just at the node boundary: 1. **Gate every LLM invocation inside a node against the budget.** Before each main-agent round and before each pruning pass, check whether `max_cost_usd` has already been tripped by `_accumulated_cost`. 2. **If the budget is tripped, stop pruning.** Do not issue further pruning passes (`_run_pruning_pass`) once we are at/over budget — pruning is optimization spend and must not push us further past the ceiling. Fall back to the unpruned tool output. 3. **Pre-flight the main agent.** If the main agent is about to be invoked but we are already at **zero or negative** remaining budget, **fail the graph execution immediately** with `ExecutionError(kind="cost", reason="budget_exhausted")` and **no further retries** — do not start the node, do not enter the retry mechanism (`agents/retry.py`). ## Acceptance criteria - [ ] A budget check runs *before* each main-agent `ainvoke` round and *before* each pruning pass, not only after the node completes. - [ ] When `_accumulated_cost` has already crossed `max_cost_usd`, no pruning pass is issued for the remainder of that node; unpruned output is used. - [ ] When remaining budget is `<= 0` and a main-agent invocation would start, the graph fails immediately with `kind="cost"`, `reason="budget_exhausted"`, and the retry path is **not** entered (0 retries). - [ ] Behavior is consistent across all three executor paths: `_execute_from_node`, the streaming path (`~:1670-1750`), and the interactive path. - [ ] No silent `$0` fallback is introduced (preserve the existing `missing_pricing_entry` hard-fail contract, `pure_graph.py:864-915`). - [ ] Tests: (a) budget already exhausted → pruning skipped, unpruned output returned; (b) zero/negative balance before a node → immediate `budget_exhausted`, zero retries, no node work started; (c) budget tripped mid tool-loop → no further pruning for that node. ## Notes / open design points This narrows the overshoot window but does **not** by itself make spend strictly `<= balance` within a single main-agent call — a single main invocation still runs to completion before its cost is reconciled. A true "never exceeds balance" guarantee still needs either a router-side safety margin on `max_cost_usd` sized to worst-case single-call cost, or a pre-flight per-call cost reservation. This ticket covers the pruning-gate + zero-balance-pre-flight asks; a strict per-call reservation can be a follow-up. Refs: `langgraph/pure_graph.py:916-960`, `:1670-1750`; `agents/llm.py:636`, `:704-708`, `:747-1019`; `agents/retry.py`. Related: ADR-2029, issue #15.

Rows
Columns