diff --git a/CHANGELOG.md b/CHANGELOG.md index 1088eb3..2f0580b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,16 @@ and this project adheres to [Clever Semantic Versioning](https://www.w3.org/subm ### Added +- **Pre-Flight USD Budget Enforcement at Each Agent Invocation (issue #76)** (`pure_graph.py`, `llm.py`, `retry.py`): 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. + + **Pre-flight gate (`pure_graph.py`):** Before a node runs, `_check_budget_pre_flight()` raises `ExecutionError(kind="cost", reason="budget_exhausted")` immediately — with zero retries — when `_accumulated_cost >= max_cost_usd`. This prevents entering a node (and its retry mechanism) when the remaining budget is zero or negative. The check fires in both `_execute_from_node` (non-streaming) and all three branches of `_stream_from_node` (terminal AGENT, intermediate AGENT, non-AGENT). + + **In-node budget gating (`llm.py`):** `_check_budget_before_invoke()` gates every `_retry_ainvoke()` call inside `_execute_tool_loop` — covering the main response, synthesis, synthesis follow-up, stuck-model, and stuck-model follow-up rounds. `_should_skip_pruning()` returns `True` when the budget is exhausted, causing pruning passes to be silently skipped (raw tool output is used as-is) instead of raising a cost error. + + **ContextVar plumbing (`retry.py`, `pure_graph.py`):** `current_accumulated_cost` and `current_max_cost_usd` ContextVars carry the per-node accumulated cost and limit from `PureLangGraph` into `LLMAgent`. The ContextVars are set before `node.execute()` and always reset in `finally` blocks to prevent stale values leaking between nodes. + + **Module:** `src/cleveractors/langgraph/pure_graph.py`, `src/cleveractors/agents/llm.py`, `src/cleveractors/agents/retry.py`. BDD: scenarios in `features/execution_limits.feature`. + - **LLM Agent Token-Budget Awareness and Tool Output Pruning (issue #61, #65)** (`llm.py`): Two complementary mechanisms to prevent context-window exhaustion in the multi-turn tool-call loop, with accurate token tracking for billing. **Token-budget awareness** (`token_budget_percent` config, default off): tracks actual token consumption from LLM response metadata before each invocation. Emits a warning at 75% budget consumption. When the budget ceiling is exceeded, injects a synthesis prompt, permits one final tool-call round, then forces a text-only response. Token counts from all rounds (including budget-exhaustion synthesis, stuck-model synthesis, and pruning passes) are accumulated for accurate billing. diff --git a/features/execution_limits.feature b/features/execution_limits.feature index 22b145f..16b260e 100644 --- a/features/execution_limits.feature +++ b/features/execution_limits.feature @@ -205,6 +205,20 @@ Feature: Execution Limits Enforcement in PureLangGraph (ADR-2029) Then no ExecutionError should be raised (elim) And a warning should have been logged for non-dict token usage (elim) + # ── Budget pre-flight enforcement (issue #76) ────────────────────────────── + + Scenario: Zero max_cost_usd triggers pre-flight failure before first node + Given a graph with an LLM node using 100 prompt tokens and max_cost_usd 0.00 (elim) + When the graph is executed with pricing for "openai"/"gpt-4.1-mini" at 0.15 prompt rate (elim) + Then an ExecutionError should be raised with kind "cost" (elim) + And the reason should be "budget_exhausted" (elim) + + Scenario: Second LLM node fails pre-flight when first node already exhausted budget + Given a graph with 2 LLM nodes each using 3000000 prompt tokens and max_cost_usd 0.40 (elim) + When the graph is executed with pricing for "openai"/"gpt-4.1-mini" at 0.15 prompt rate (elim) + Then an ExecutionError should be raised with kind "cost" (elim) + And the reason should be "budget_exhausted" (elim) + # ── Happy-path result correctness ───────────────────────────────────────── Scenario: Graph within max_depth returns non-None result diff --git a/src/cleveractors/agents/llm.py b/src/cleveractors/agents/llm.py index d0fbd5c..f62eb71 100644 --- a/src/cleveractors/agents/llm.py +++ b/src/cleveractors/agents/llm.py @@ -39,7 +39,12 @@ from cleveractors.agents.base import AgentWithMemory from cleveractors.agents.llm_client import build_chat_model from cleveractors.agents.llm_imports import populate_langchain_globals from cleveractors.agents.llm_tools import normalize_tool_entry as _normalize_tool_entry -from cleveractors.agents.retry import _get_provider_url, call_with_retry +from cleveractors.agents.retry import ( + _get_provider_url, + call_with_retry, + current_accumulated_cost, + current_max_cost_usd, +) from cleveractors.core.exceptions import ( AgentCreationError, ConfigurationError, @@ -493,6 +498,52 @@ class LLMAgent(AgentWithMemory): provider_url=provider_url, ) + def _check_budget_before_invoke(self, where: str = "main") -> None: + """Raise ``budget_exhausted`` if accumulated cost already at or exceeds limit. + + Called *before* each main-agent ``ainvoke`` round inside the tool + loop so that zero/negative remaining budget fails fast (issue #76). + Does **not** apply to pruning passes — those are silently skipped + via :meth:`_should_skip_pruning` instead. + """ + acc = current_accumulated_cost.get() + limit = current_max_cost_usd.get() + if limit is not None and acc >= limit: + logger.error( + "Agent %s: budget exhausted before %s invoke ($%.4f >= $%.4f limit)", + self.name, + where, + acc, + limit, + ) + raise ExecutionError( + f"Cost limit exceeded pre-flight: " + f"{acc:.6f} USD >= {limit} USD at {where} invoke", + kind="cost", + reason="budget_exhausted", + ) + + def _should_skip_pruning(self, tool_name: str) -> bool: + """Return True when the budget is already exhausted and pruning should be skipped. + + When the accumulated cost has already crossed the limit, pruning + passes (extra LLM calls) are suppressed and the raw tool output is + used instead (issue #76). + """ + acc = current_accumulated_cost.get() + limit = current_max_cost_usd.get() + if limit is not None and acc >= limit: + logger.info( + "Agent %s: budget exhausted, skipping pruning for %r " + "($%.4f >= $%.4f limit)", + self.name, + tool_name, + acc, + limit, + ) + return True + return False + def _get_model_context_window(self) -> int: """Return the model's advertised context window size in tokens. @@ -867,6 +918,7 @@ class LLMAgent(AgentWithMemory): "call to complete your answer, do it now." ) messages.append(HumanMessage(content=_synthesis_text)) + self._check_budget_before_invoke("synthesis") response = await self._retry_ainvoke(messages, **invoke_kwargs) _any_invocation_made = True _bp, _bc, _ = self._extract_token_counts(response) @@ -952,6 +1004,7 @@ class LLMAgent(AgentWithMemory): and _bq_output_prune is not False and tool_name in self._pruning_tool_filter and len(raw_out) > self._pruning_threshold + and not self._should_skip_pruning(tool_name) ): try: ( @@ -1011,6 +1064,7 @@ class LLMAgent(AgentWithMemory): tool_call_id=call_id, ) ) + self._check_budget_before_invoke("synthesis_followup") response = await self._retry_ainvoke(messages) _any_invocation_made = True _bp3, _bc3, _ = self._extract_token_counts(response) @@ -1019,6 +1073,9 @@ class LLMAgent(AgentWithMemory): _budget_exhausted = True break + # ── Pre-invoke budget check (issue #76) ───────────────────── + self._check_budget_before_invoke("main") + # ── Regular ainvoke ──────────────────────────────────────── response = await self._retry_ainvoke(messages, **invoke_kwargs) _any_invocation_made = True @@ -1103,6 +1160,7 @@ class LLMAgent(AgentWithMemory): and _output_prune is not False and tool_name in self._pruning_tool_filter and len(raw_tool_output) > self._pruning_threshold + and not self._should_skip_pruning(tool_name) ): _task_context = list(messages) try: @@ -1182,6 +1240,7 @@ class LLMAgent(AgentWithMemory): ) ) ) + self._check_budget_before_invoke("stuck_model") response = await self._retry_ainvoke(messages, tools=self._lc_tools) _any_invocation_made = True _sp, _sc, _ = self._extract_token_counts(response) @@ -1274,6 +1333,7 @@ class LLMAgent(AgentWithMemory): tool_call_id=call_id, ) ) + self._check_budget_before_invoke("stuck_model_followup") response = await self._retry_ainvoke(messages) _any_invocation_made = True _sfp, _sfc, _ = self._extract_token_counts(response) diff --git a/src/cleveractors/agents/retry.py b/src/cleveractors/agents/retry.py index 8050c3f..9fd79f6 100644 --- a/src/cleveractors/agents/retry.py +++ b/src/cleveractors/agents/retry.py @@ -27,6 +27,18 @@ current_graph_name: contextvars.ContextVar[str] = contextvars.ContextVar( "current_graph_name", default="" ) +# Per-task budget state carried from PureLangGraph into LLMAgent so +# in-node LLM invocations (main-agent rounds and pruning passes) can +# check whether the USD budget has already been exhausted (issue #76). +# Set by PureLangGraph before each node execution; read by LLMAgent's +# _execute_tool_loop and _run_pruning_pass. +current_accumulated_cost: contextvars.ContextVar[float] = contextvars.ContextVar( + "current_accumulated_cost", default=0.0 +) +current_max_cost_usd: contextvars.ContextVar[float | None] = contextvars.ContextVar( + "current_max_cost_usd", default=None +) + _INITIAL_RETRY_DELAY: float = 0.5 diff --git a/src/cleveractors/langgraph/pure_graph.py b/src/cleveractors/langgraph/pure_graph.py index 282f078..7991ea4 100644 --- a/src/cleveractors/langgraph/pure_graph.py +++ b/src/cleveractors/langgraph/pure_graph.py @@ -8,6 +8,7 @@ allowing it to work properly in run mode without timeouts. from __future__ import annotations import asyncio +import contextvars import copy import logging from collections import defaultdict, deque @@ -17,6 +18,10 @@ from pathlib import Path from typing import Any, Deque, Final, List, Optional, Set from cleveractors.agents.base import Agent +from cleveractors.agents.retry import ( + current_accumulated_cost, + current_max_cost_usd, +) from cleveractors.context_manager import ContextManager from cleveractors.core.exceptions import ExecutionError from cleveractors.langgraph.dynamic_router import DynamicRouterNode @@ -535,6 +540,75 @@ class PureLangGraph: global_context.clear() global_context.update(final_metadata) + def _set_budget_context_for_node( + self, node_name: str + ) -> tuple[contextvars.Token[float] | None, contextvars.Token[float | None] | None]: + """Set budget ContextVars for in-node budget checks (issue #76). + + Must be called before each node execution. Returns a pair of + ``(cost_token, max_token)`` that the caller must reset in a + ``finally`` block after the node completes. + + When ``_pricing`` is empty or ``max_cost_usd`` is absent, the + ContextVars are set to zero / ``None`` respectively so that + in-node checks are no-ops. + """ + _cost = self._accumulated_cost + _max: float | None = None + if self._pricing: + _raw = self._limits.get("max_cost_usd") + if _raw is not None and not isinstance(_raw, bool): + try: + _max = float(_raw) + except (TypeError, ValueError): + pass + _cost_token = current_accumulated_cost.set(_cost) + _max_token: contextvars.Token[float | None] | None = None + if _max is not None: + _max_token = current_max_cost_usd.set(_max) + return (_cost_token, _max_token) + + @staticmethod + def _reset_budget_context( + cost_token: contextvars.Token[float] | None, + max_token: contextvars.Token[float | None] | None, + ) -> None: + """Reset budget ContextVars after node execution.""" + if cost_token is not None: + current_accumulated_cost.reset(cost_token) + if max_token is not None: + current_max_cost_usd.reset(max_token) + + def _check_budget_pre_flight(self, node_name: str) -> None: + """Raise ``budget_exhausted`` if accumulated cost already at or exceeds limit. + + Called *before* a node executes so that a zero/negative remaining + budget fails immediately without entering the node or its retry + mechanism (issue #76). No-op when pricing is empty or no limit is set. + """ + if self._pricing: + _raw = self._limits.get("max_cost_usd") + if _raw is not None and not isinstance(_raw, bool): + try: + _max = float(_raw) + except (TypeError, ValueError): + return + if self._accumulated_cost >= _max: + self.logger.warning( + "Budget pre-flight: accumulated $%.4f >= limit $%.4f " + "before node %s", + self._accumulated_cost, + _max, + node_name, + ) + raise ExecutionError( + f"Cost limit exceeded pre-flight: " + f"{self._accumulated_cost:.6f} USD >= {_max} USD" + f" before node '{node_name}'", + kind="cost", + reason="budget_exhausted", + ) + async def _execute_from_node( self, node_name: str, message: Any, depth: int = 0 ) -> Any: @@ -790,6 +864,18 @@ class PureLangGraph: # Pass graph name for retry error reporting (ADR-2032 D-7). state.metadata["graph_name"] = self.name + # ── Budget pre-flight check (issue #76) ──────────────────────────── + # If the accumulated cost has already reached or exceeded the limit + # before this node starts, fail immediately with zero retries. + # This prevents entering the node (and its retry mechanism) when + # there is no remaining budget. + self._check_budget_pre_flight(node_name) + + # Set budget ContextVars for in-node budget checks (issue #76). + _budget_cost_token, _budget_max_token = self._set_budget_context_for_node( + node_name + ) + # Execute node with current state try: self.logger.debug(f"About to execute node.execute() for {node_name}") @@ -981,13 +1067,14 @@ class PureLangGraph: break except ExecutionError: - # Limit-enforcement errors (depth, model_calls, tool_calls, timeout, - # cost) must propagate without suppression so the router can map - # them to the correct HTTP status code (ADR-2029, issue #15). raise except Exception as e: self.logger.error(f"Error executing node {node_name}: {e}", exc_info=True) output_message = message + finally: + # Reset budget ContextVars after node execution so that stale + # values from this node cannot leak into the next node (issue #76). + self._reset_budget_context(_budget_cost_token, _budget_max_token) # Determine next nodes self.logger.debug( @@ -1577,6 +1664,9 @@ class PureLangGraph: # they are also yielded immediately (fast path) or held back # until the stream completes (slow path). _collected_tokens: List[str] = [] + # ── Budget pre-flight check (issue #76) ─────────── + self._check_budget_pre_flight(node_name) + _bcost_tok_t, _bmax_tok_t = self._set_budget_context_for_node(node_name) try: async for token in node._stream_agent(state): _collected_tokens.append(token) # always accumulate @@ -1622,6 +1712,9 @@ class PureLangGraph: ) raise ExecutionError("Agent node streaming failed") from e + # ── Reset budget ContextVars (issue #76) ──────────── + self._reset_budget_context(_bcost_tok_t, _bmax_tok_t) + # Update last_output for routing (mirrors non-streaming path). # Note: this is only reached on the success path; the except blocks # above re-raise, so last_output retains its previous value on failure. @@ -1768,6 +1861,9 @@ class PureLangGraph: else: # Intermediate AGENT node: use ainvoke() (node.execute) to avoid # buffering tokens only to discard them (AC2). + # ── Budget pre-flight check (issue #76) ─────────── + self._check_budget_pre_flight(node_name) + _bcost_tok_i, _bmax_tok_i = self._set_budget_context_for_node(node_name) try: result = await node.execute(state) if isinstance(result, dict): @@ -1912,6 +2008,9 @@ class PureLangGraph: ) raise ExecutionError("Intermediate agent node failed") from e + # ── Reset budget ContextVars (issue #76) ──────────── + self._reset_budget_context(_bcost_tok_i, _bmax_tok_i) + self.state_manager.state.metadata["last_output"] = full_response # Parse CleverAgents v2.0 routing commands from agent output so @@ -2023,6 +2122,9 @@ class PureLangGraph: else: # Non-AGENT node: run with node.execute() (ainvoke path) + # ── Budget pre-flight check (issue #76) ─────────── + self._check_budget_pre_flight(node_name) + _bcost_tok_n, _bmax_tok_n = self._set_budget_context_for_node(node_name) output_message: Any = message try: result = await node.execute(state) @@ -2063,6 +2165,9 @@ class PureLangGraph: ) output_message = message + # ── Reset budget ContextVars (issue #76) ──────────── + self._reset_budget_context(_bcost_tok_n, _bmax_tok_n) + self.state_manager.state.metadata["last_output"] = output_message # Parse CleverAgents v2.0 routing commands from non-AGENT node output so