feat(graph): enforce USD budget at each agent invocation with pre-flight gate #79

2026-07-07T22:32:03Z

CoreRasurae commented

2026-07-07 22:32:03 +00:00

Summary

Enforces the USD budget (max_cost_usd) as a hard pre-flight gate before every node execution and at each main-agent ainvoke round inside the multi-turn tool loop, with pruning passes silently skipped when the budget is exhausted.

Key changes

Pre-flight gate (pure_graph.py): _check_budget_pre_flight() raises ExecutionError(kind="cost", reason="budget_exhausted") immediately before node execution in both _execute_from_node and all three branches of _stream_from_node.
In-node budget gating (llm.py): _check_budget_before_invoke() gates every _retry_ainvoke() call. Pruning passes are skipped via _should_skip_pruning() when the budget is exhausted.
ContextVar plumbing (retry.py, pure_graph.py): current_accumulated_cost and current_max_cost_usd ContextVars carry per-node budget into LLMAgent, always reset in finally blocks.
BDD tests: Two new scenarios in features/execution_limits.feature verify pre-flight failure on zero budget and on second-node exhaustion.

Quality gates

nox -s lint — pass
nox -s format -- --check — pass
nox -s typecheck — pass (0 errors)
nox -s security_scan — pass (0 findings)
nox -s dead_code — pass
nox -s unit_tests — 2748/2748 scenarios pass
nox -s coverage_report — 96.7% (threshold: 96.5%)

Closes #76

## Summary Enforces the USD budget (`max_cost_usd`) as a hard pre-flight gate before every node execution and at each main-agent `ainvoke` round inside the multi-turn tool loop, with pruning passes silently skipped when the budget is exhausted. ### Key changes - **Pre-flight gate** (`pure_graph.py`): `_check_budget_pre_flight()` raises `ExecutionError(kind="cost", reason="budget_exhausted")` immediately before node execution in both `_execute_from_node` and all three branches of `_stream_from_node`. - **In-node budget gating** (`llm.py`): `_check_budget_before_invoke()` gates every `_retry_ainvoke()` call. Pruning passes are skipped via `_should_skip_pruning()` when the budget is exhausted. - **ContextVar plumbing** (`retry.py`, `pure_graph.py`): `current_accumulated_cost` and `current_max_cost_usd` ContextVars carry per-node budget into `LLMAgent`, always reset in `finally` blocks. - **BDD tests**: Two new scenarios in `features/execution_limits.feature` verify pre-flight failure on zero budget and on second-node exhaustion. ### Quality gates - `nox -s lint` — pass - `nox -s format -- --check` — pass - `nox -s typecheck` — pass (0 errors) - `nox -s security_scan` — pass (0 findings) - `nox -s dead_code` — pass - `nox -s unit_tests` — 2748/2748 scenarios pass - `nox -s coverage_report` — 96.7% (threshold: 96.5%) Closes #76

CoreRasurae added 1 commit 2026-07-07 22:32:03 +00:00

feat(graph): enforce USD budget at each agent invocation with pre-flight gate

CI / lint (pull_request) Successful in 47s

Details

CI / typecheck (pull_request) Successful in 52s

Details

CI / security (pull_request) Successful in 50s

Details

CI / quality (pull_request) Successful in 34s

Details

CI / unit_tests (pull_request) Successful in 3m6s

Details

CI / integration_tests (pull_request) Successful in 1m8s

Details

CI / build (pull_request) Successful in 35s

Details

CI / coverage (pull_request) Successful in 3m9s

Details

CI / status-check (pull_request) Successful in 3s

Details

b85bc27e0b

Enforces max_cost_usd as a hard pre-flight gate before every node
execution in both _execute_from_node and _stream_from_node (all three
branches: terminal AGENT, intermediate AGENT, non-AGENT).

Inside _execute_tool_loop, each main-agent ainvoke round (main,
synthesis, synthesis_followup, stuck_model, stuck_model_followup) is
gated by _check_budget_before_invoke. Pruning passes are silently
skipped when the budget is exhausted via _should_skip_pruning.

current_accumulated_cost and current_max_cost_usd ContextVars carry
the per-node budget from PureLangGraph into LLMAgent, set before
node.execute() and reset in finally blocks.

BDD scenarios in features/execution_limits.feature.

ISSUES CLOSED: #76

CoreRasurae added a new dependency 2026-07-07 22:32:44 +00:00

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

CoreRasurae added the

Type

Feature

label 2026-07-07 22:33:19 +00:00

hurui200320 requested changes 2026-07-08 05:00:50 +00:00

hurui200320 left a comment

PR Review: !79 (Ticket #76)

Verdict: Request Changes

The pre-flight USD budget enforcement is wired consistently across all three executor paths and pruning is correctly guarded, but the budget ContextVars are not reset in finally blocks in the three streaming branches as the PR claims. On any ExecutionError (or wrapped exception re-raised as ExecutionError) the ContextVars leak for the remainder of the task, which contradicts the stated safety design and can leave stale budget state behind.

Critical Issues

None.

Major Issues

Budget ContextVars are not reset on exception paths in _stream_from_node
- Files: src/cleveractors/langgraph/pure_graph.py
- Lines: 1668–1716 (terminal AGENT branch), 1865–2012 (intermediate AGENT branch), 2126–2169 (non-AGENT branch)
- Problem: The reset calls are placed after the try/except blocks, not inside finally. When except ExecutionError: raise or except Exception: raise ExecutionError(...) fires, control propagates out of the method before _reset_budget_context runs. This leaves current_accumulated_cost and current_max_cost_usd set to the values from the failed node.
- Recommendation: Wrap each node execution block in a try/finally that resets the ContextVars, mirroring _execute_from_node (line 1077). For example, in the terminal AGENT branch move _reset_budget_context(_bcost_tok_t, _bmax_tok_t) into a finally clause attached to the existing try/except.

Minor Issues

No-tools LLM invocations are not gated by _check_budget_before_invoke
- Files: src/cleveractors/agents/llm.py
- Lines: 1580 (process_message no-tools ainvoke), 1914 (stream_message no-tools astream)
- Problem: The PR gates every _retry_ainvoke() inside _execute_tool_loop, but the plain single-shot paths used when self._lc_tools is None (or LangChain unavailable) bypass the in-node budget check. While the node-level pre-flight gate catches zero/negative balance before the node starts, this leaves an inconsistency with the stated goal of gating every main-agent LLM invocation.
- Recommendation: Add _check_budget_before_invoke("main") immediately before the no-tools ainvoke and astream calls, or document why these paths are intentionally exempt.
In-node checks cannot detect budget exhaustion caused by the current node
- Files: src/cleveractors/agents/retry.py, src/cleveractors/langgraph/pure_graph.py, src/cleveractors/agents/llm.py
- Lines: retry.py:35–40 (ContextVar defaults), pure_graph.py:543–572 (_set_budget_context_for_node), llm.py:501–548 (budget helpers)
- Problem: current_accumulated_cost is set once to the graph’s pre-node _accumulated_cost and is never updated as the node spends tokens. Consequently _check_budget_before_invoke and _should_skip_pruning only catch budget exhaustion from previous nodes, not from earlier invocations within the same node. This means acceptance criterion (c) — “budget tripped mid tool-loop → no further pruning for that node” — is not literally satisfied for intra-node spend.
- Recommendation: This is acknowledged in the ticket as a follow-up that requires per-call cost reservation or pricing-aware ContextVar updates. Ensure the follow-up ticket is created, or update the ContextVar during the node (e.g., by passing pricing into the agent or updating the ContextVar from PureLangGraph after each tracked sub-call).

Nits

Logging level inconsistency for the same condition
- Files: src/cleveractors/agents/llm.py:513, src/cleveractors/langgraph/pure_graph.py:594
- Problem: _check_budget_before_invoke logs at logger.error while _check_budget_pre_flight logs at logger.warning for the same “budget already exhausted” situation.
- Recommendation: Standardize on one level; warning is probably more appropriate since budget exhaustion is a state/limit condition rather than a code fault.
Duplicated max_cost_usd parsing logic
- Files: src/cleveractors/langgraph/pure_graph.py
- Lines: 553–560 and 586–592
- Problem: _set_budget_context_for_node and _check_budget_pre_flight both parse self._limits.get("max_cost_usd") with identical bool-guard and float-conversion logic.
- Recommendation: Extract a small helper such as _get_max_cost_usd_limit() -> float | None to remove the duplication and reduce the chance of the two sites diverging.

Summary

The PR correctly adds a hard pre-flight budget gate in PureLangGraph for all execution paths and wires pruning-skip logic into LLMAgent. The core logic is sound and matches the ticket’s explicit scope for pruning-gate + zero-balance pre-flight. However, the streaming branches fail to reset the budget ContextVars in finally blocks, which is a real state-hygiene bug that should be fixed before merge. Once the streaming reset is moved into finally and ideally the no-tools paths are also gated, the PR should be ready for approval.

## PR Review: !79 (Ticket #76) ### Verdict: Request Changes The pre-flight USD budget enforcement is wired consistently across all three executor paths and pruning is correctly guarded, but the budget ContextVars are **not** reset in `finally` blocks in the three streaming branches as the PR claims. On any `ExecutionError` (or wrapped exception re-raised as `ExecutionError`) the ContextVars leak for the remainder of the task, which contradicts the stated safety design and can leave stale budget state behind. ### Critical Issues None. ### Major Issues 1. **Budget ContextVars are not reset on exception paths in `_stream_from_node`** - **Files:** `src/cleveractors/langgraph/pure_graph.py` - **Lines:** 1668–1716 (terminal AGENT branch), 1865–2012 (intermediate AGENT branch), 2126–2169 (non-AGENT branch) - **Problem:** The reset calls are placed *after* the `try/except` blocks, not inside `finally`. When `except ExecutionError: raise` or `except Exception: raise ExecutionError(...)` fires, control propagates out of the method before `_reset_budget_context` runs. This leaves `current_accumulated_cost` and `current_max_cost_usd` set to the values from the failed node. - **Recommendation:** Wrap each node execution block in a `try/finally` that resets the ContextVars, mirroring `_execute_from_node` (line 1077). For example, in the terminal AGENT branch move `_reset_budget_context(_bcost_tok_t, _bmax_tok_t)` into a `finally` clause attached to the existing `try/except`. ### Minor Issues 1. **No-tools LLM invocations are not gated by `_check_budget_before_invoke`** - **Files:** `src/cleveractors/agents/llm.py` - **Lines:** 1580 (`process_message` no-tools `ainvoke`), 1914 (`stream_message` no-tools `astream`) - **Problem:** The PR gates every `_retry_ainvoke()` inside `_execute_tool_loop`, but the plain single-shot paths used when `self._lc_tools` is `None` (or LangChain unavailable) bypass the in-node budget check. While the node-level pre-flight gate catches zero/negative balance before the node starts, this leaves an inconsistency with the stated goal of gating every main-agent LLM invocation. - **Recommendation:** Add `_check_budget_before_invoke("main")` immediately before the no-tools `ainvoke` and `astream` calls, or document why these paths are intentionally exempt. 2. **In-node checks cannot detect budget exhaustion caused by the current node** - **Files:** `src/cleveractors/agents/retry.py`, `src/cleveractors/langgraph/pure_graph.py`, `src/cleveractors/agents/llm.py` - **Lines:** `retry.py:35–40` (ContextVar defaults), `pure_graph.py:543–572` (`_set_budget_context_for_node`), `llm.py:501–548` (budget helpers) - **Problem:** `current_accumulated_cost` is set once to the graph’s pre-node `_accumulated_cost` and is never updated as the node spends tokens. Consequently `_check_budget_before_invoke` and `_should_skip_pruning` only catch budget exhaustion from *previous* nodes, not from earlier invocations within the same node. This means acceptance criterion (c) — “budget tripped mid tool-loop → no further pruning for that node” — is not literally satisfied for intra-node spend. - **Recommendation:** This is acknowledged in the ticket as a follow-up that requires per-call cost reservation or pricing-aware ContextVar updates. Ensure the follow-up ticket is created, or update the ContextVar during the node (e.g., by passing pricing into the agent or updating the ContextVar from PureLangGraph after each tracked sub-call). ### Nits 1. **Logging level inconsistency for the same condition** - **Files:** `src/cleveractors/agents/llm.py:513`, `src/cleveractors/langgraph/pure_graph.py:594` - **Problem:** `_check_budget_before_invoke` logs at `logger.error` while `_check_budget_pre_flight` logs at `logger.warning` for the same “budget already exhausted” situation. - **Recommendation:** Standardize on one level; `warning` is probably more appropriate since budget exhaustion is a state/limit condition rather than a code fault. 2. **Duplicated `max_cost_usd` parsing logic** - **Files:** `src/cleveractors/langgraph/pure_graph.py` - **Lines:** 553–560 and 586–592 - **Problem:** `_set_budget_context_for_node` and `_check_budget_pre_flight` both parse `self._limits.get("max_cost_usd")` with identical bool-guard and float-conversion logic. - **Recommendation:** Extract a small helper such as `_get_max_cost_usd_limit() -> float | None` to remove the duplication and reduce the chance of the two sites diverging. ### Summary The PR correctly adds a hard pre-flight budget gate in `PureLangGraph` for all execution paths and wires pruning-skip logic into `LLMAgent`. The core logic is sound and matches the ticket’s explicit scope for pruning-gate + zero-balance pre-flight. However, the streaming branches fail to reset the budget ContextVars in `finally` blocks, which is a real state-hygiene bug that should be fixed before merge. Once the streaming reset is moved into `finally` and ideally the no-tools paths are also gated, the PR should be ready for approval.

CoreRasurae added

and removed

labels 2026-07-08 15:17:02 +00:00

CoreRasurae closed this pull request

2026-07-08 18:14:01 +00:00

CoreRasurae removed a dependency 2026-07-08 18:15:24 +00:00

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

CI / lint (pull_request) Successful in 47s

Details

CI / typecheck (pull_request) Successful in 52s

Details

CI / security (pull_request) Successful in 50s

Details

CI / quality (pull_request) Successful in 34s

Details

CI / unit_tests (pull_request) Successful in 3m6s

Details

CI / integration_tests (pull_request) Successful in 1m8s

Details

CI / build (pull_request) Successful in 35s

Details

CI / coverage (pull_request) Successful in 3m9s

Details

CI / status-check (pull_request) Successful in 3s

Details

Pull request closed

Please reopen this pull request to perform a merge.

Sign in to join this conversation.

2 Participants

Notifications

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: cleveragents/cleveractors-core#79