diff --git a/CHANGELOG.md b/CHANGELOG.md index 64e3f37..e8dea5e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,7 @@ and this project adheres to [Clever Semantic Versioning](https://www.w3.org/subm **Module:** `src/cleveractors/agents/factory.py` (`AgentFactory.__init__`, `_create_agent_instance`), `src/cleveractors/agents/llm.py` (`LLMAgent.__init__`, `_execute_tool_loop`), `src/cleveractors/runtime.py` (`Executor.__init__`, `create_executor`), `src/cleveractors/runtime_dispatch.py` (`_execute_tool`, `_execute_graph`, `_execute_graph_stream`, `_execute_llm`, `_execute_llm_stream`, `_execute_multi_actor`), `src/cleveractors/core/application.py` (`load_configuration`, `_create_agents`). BDD: scenarios in `features/tool_agent_class_injection.feature`. +- **Pre-Flight USD Budget Enforcement at Each Agent Invocation (issue #76)** (`pure_graph.py`, `llm.py`, `retry.py`, `exceptions.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. Node-boundary pre-flight uses `>=` (blocks next node when at limit) while post-node checks use `>` (current node completes at limit). The `current_accumulated_cost`, `current_max_cost_usd`, and `current_pricing` ContextVars carry per-node budget state from `PureLangGraph` into `LLMAgent`, enabling in-node cost tracking after each `ainvoke` round so subsequent budget checks see the updated cost. ContextVars are reset in `finally` blocks across all four execution branches (`_execute_from_node`, terminal-stream, intermediate-stream, non-AGENT stream) to prevent leaks. Cost formula deduplication: `_compute_cost` in `exceptions.py` centralises the pricing-lookup + cost-computation logic formerly duplicated across 5 call sites (3 in `pure_graph.py`, 1 in `runtime_dispatch.py`, 1 in `llm.py`), eliminating 267 lines of redundancy. Single-LLM streaming path (`_execute_llm_stream`) now raises `missing_pricing_entry` when `max_cost_usd` is specified but the pricing table is empty, mirroring the graph-path behaviour. BDD: scenarios in `features/execution_limits.feature`, `features/llm_agent_tool_loop.feature`, and `features/execute_stream.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/environment.py b/features/environment.py index 754b9fa..5d168dc 100644 --- a/features/environment.py +++ b/features/environment.py @@ -227,6 +227,26 @@ def after_scenario(context, scenario): except Exception: pass + # Reset budget ContextVars set by test steps (issue #76 test hygiene). + # Without this, budget ContextVars set in llm_agent_tool_loop scenarios + # leak across scenarios and can cause order-dependent flakiness + # (Graa review PR #80). + if hasattr(context, "_budget_cost_token"): + from cleveractors.agents.retry import ( + current_accumulated_cost, + current_max_cost_usd, + current_pricing, + ) + + try: + current_accumulated_cost.reset(context._budget_cost_token) + if getattr(context, "_budget_max_token", None) is not None: + current_max_cost_usd.reset(context._budget_max_token) + if getattr(context, "_budget_pricing_token", None) is not None: + current_pricing.reset(context._budget_pricing_token) + except Exception: + pass + # Restore monkey-patched validation constants if hasattr(context, "_original_max_dfs"): import cleveractors.validation._limits as _limits diff --git a/features/execute_stream.feature b/features/execute_stream.feature index 61c09d5..040c2ff 100644 --- a/features/execute_stream.feature +++ b/features/execute_stream.feature @@ -866,6 +866,18 @@ Feature: Executor.execute_stream() - token-by-token streaming delivery # ── Major #1 fix: max_cost_usd enforced for intermediate AGENT nodes ────── + Scenario: Streaming graph exhausts budget via post-node check on first node + Given an Executor with a two-agent sequential graph and max_cost_usd 0.40 with pricing (stream) + And a mock astream yielding single token "BudgetNodeCost" with usage prompt=3000000 completion=0 (stream) + When I attempt execute_stream expecting an error with message "test" (stream) + Then an ExecutionError with kind "cost" and reason "budget_exhausted" should be raised (stream kind+reason) + + Scenario: Second streaming node blocked by pre-flight when first node exactly hits limit + Given an Executor with a two-agent sequential graph and max_cost_usd 0.30 at 0.15/M prompt pricing (stream) + And a mock astream yielding single token "ExactPreFlight" with usage prompt=2000000 completion=0 (stream) + When I attempt execute_stream expecting an error with message "test" (stream) + Then an ExecutionError with kind "cost" and reason "budget_exhausted" should be raised (stream kind+reason) + Scenario: execute_stream raises ExecutionError when intermediate AGENT node exceeds max_cost_usd Given an Executor with a two-agent sequential graph and max_cost_usd 0.0 with pricing (stream) And a mock astream yielding single token "IntCostToken" with usage prompt=1000 completion=1000 (stream) diff --git a/features/execution_limits.feature b/features/execution_limits.feature index 22b145f..30004e7 100644 --- a/features/execution_limits.feature +++ b/features/execution_limits.feature @@ -108,10 +108,11 @@ Feature: Execution Limits Enforcement in PureLangGraph (ADR-2029) Then an ExecutionError should be raised with kind "cost" (elim) And the reason should be "missing_pricing_entry" (elim) - Scenario: Empty pricing table skips cost calculation entirely + Scenario: Empty pricing table raises missing_pricing_entry when limit is specified Given a graph with an LLM node using 2000000 prompt tokens and max_cost_usd 0.001 (elim) When the graph is executed with an empty pricing table (elim) - Then no ExecutionError should be raised (elim) + Then an ExecutionError should be raised with kind "cost" (elim) + And the reason should be "missing_pricing_entry" (elim) Scenario: LLM node with incomplete model pricing entry missing prompt rate raises missing_pricing_entry Given a graph with an LLM node for provider "openai" model "gpt-4.1-mini" (elim) @@ -205,6 +206,26 @@ 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: Two LLM nodes exhaust budget via post-node check on first node + 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) + + Scenario: Second LLM node blocked by pre-flight when first node exactly hits limit + Given a graph with 2 LLM nodes each using 2000000 prompt tokens and max_cost_usd 0.30 (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/features/llm_agent_tool_loop.feature b/features/llm_agent_tool_loop.feature index cbe872b..41351eb 100644 --- a/features/llm_agent_tool_loop.feature +++ b/features/llm_agent_tool_loop.feature @@ -95,3 +95,22 @@ Feature: LLMAgent._execute_tool_loop() — shared tool-call orchestration helper Then a _ToolLoopError should be raised And the _ToolLoopError any_invocation_made should be false And the _ToolLoopError cause should be a ConfigurationError + + # ── USD cost budget exhaustion (in-node gating, issue #76) ──────────────── + + Scenario: _execute_tool_loop raises budget_exhausted when pricing context is set and cost exceeds limit + Given I have an LLMAgent with tool loop test config for tool "echo" and tool_max_rounds 5 + And I set a mock that makes one tool call then returns a final answer "cost triggered" + And the USD budget ContextVars are set for in-node gating + When I run _execute_tool_loop with an initial messages list + Then an ExecutionError with kind "cost" and reason "budget_exhausted" should be raised + + Scenario: _execute_tool_loop skips pruning pass when budget tripped mid-loop (issue #76) + Given I have an LLMAgent with tool loop test config for tool "file_read" and pruning enabled + And I set a mock that makes one file_read tool call then returns "Final text" + And I set a mock pruning model that adds extra tokens + And the USD budget ContextVars are set so first round exhausts limit + When I run _execute_tool_loop with an initial messages list + Then an ExecutionError with kind "cost" and reason "budget_exhausted" should be raised + And the accumulated prompt should NOT include pruning tokens + And the accumulated completion should NOT include pruning tokens diff --git a/features/steps/execute_stream_steps.py b/features/steps/execute_stream_steps.py index 9db7805..53c0850 100644 --- a/features/steps/execute_stream_steps.py +++ b/features/steps/execute_stream_steps.py @@ -4762,6 +4762,57 @@ def _two_agent_sequential_config() -> dict[str, Any]: } +@given( + "an Executor with a two-agent sequential graph and" + " max_cost_usd 0.40 with pricing (stream)" +) +def step_es_two_agent_max_cost_040(context: Any) -> None: + """Graph: START → agent_a (AGENT, intermediate) → agent_b (AGENT, terminal) → END. + + max_cost_usd=0.40 with prompt pricing at 0.50/M so that a first node + producing 3M prompt tokens costs $1.50 and exhausts the budget. The + post-node cost check catches the exhaustion; the pre-flight gate is also + exercised as a no-op before the first node. This exercises the streaming + path cost enforcement for a non-zero budget (issue #76, AC4). + """ + context.es_executor = create_executor( + config_dict=_two_agent_sequential_config(), + credentials={"openai": {"api_key": "test-key"}}, + limits={"max_cost_usd": 0.40}, + pricing={ + "openai": { + "gpt-3.5-turbo": {"prompt": 0.50, "completion": 1.50}, + } + }, + ) + + +@given( + "an Executor with a two-agent sequential graph and" + " max_cost_usd 0.30 at 0.15/M prompt pricing (stream)" +) +def step_es_two_agent_max_cost_030_exact(context: Any) -> None: + """Graph: START → agent_a (AGENT, intermediate) → agent_b (AGENT, terminal) → END. + + max_cost_usd=0.30 with prompt pricing at 0.15/M so that a first node + producing 2M prompt tokens costs exactly $0.30. The post-node check uses + ``>`` so the first node completes (0.30 is NOT > 0.30). The second node's + pre-flight gate uses ``>=`` so 0.30 >= 0.30 blocks it, exercising the + intentional asymmetry between post-node and pre-flight checks + (Graa review PR #80). + """ + context.es_executor = create_executor( + config_dict=_two_agent_sequential_config(), + credentials={"openai": {"api_key": "test-key"}}, + limits={"max_cost_usd": 0.30}, + pricing={ + "openai": { + "gpt-3.5-turbo": {"prompt": 0.15, "completion": 0.60}, + } + }, + ) + + @given( "an Executor with a two-agent sequential graph and max_cost_usd 1.0" " with missing provider pricing (stream)" diff --git a/features/steps/llm_agent_tool_loop_steps.py b/features/steps/llm_agent_tool_loop_steps.py index 7ce502c..7e395b4 100644 --- a/features/steps/llm_agent_tool_loop_steps.py +++ b/features/steps/llm_agent_tool_loop_steps.py @@ -11,7 +11,12 @@ from behave import given, then, when from langchain_core.messages import AIMessage, ToolMessage from cleveractors.agents.llm import LLMAgent, _ToolLoopError -from cleveractors.core.exceptions import ConfigurationError +from cleveractors.agents.retry import ( + current_accumulated_cost, + current_max_cost_usd, + current_pricing, +) +from cleveractors.core.exceptions import ConfigurationError, ExecutionError from cleveractors.templates.renderer import TemplateRenderer # --------------------------------------------------------------------------- @@ -383,6 +388,39 @@ def step_mock_one_round_then_error(context): # --------------------------------------------------------------------------- +@given("the USD budget ContextVars are set for in-node gating") +def step_set_usd_budget(context): + """Set budget ContextVars so the in-node cost gating fires. + + Uses a pricing table where a single 15/5 token invoke at + $1000/M prompt/$1000/M completion costs $0.02, exceeding $0.01 limit. + """ + # Cost per round: 15/1M * 1000 + 5/1M * 1000 = 0.015 + 0.005 = $0.02 + # With limit $0.01, the second main invoke check should fire. + # Save previous tokens so after_scenario can reset them (issue #76 test hygiene). + context._budget_cost_token = current_accumulated_cost.set(0.0) + context._budget_max_token = current_max_cost_usd.set(0.01) + context._budget_pricing_token = current_pricing.set( + {"openai": {"gpt-3.5-turbo": {"prompt": 1000.0, "completion": 1000.0}}} + ) + + +@given("the USD budget ContextVars are set so first round exhausts limit") +def step_set_usd_budget_exhausted_mid_round(context): + """Set budget ContextVars so the first round ainvoke exhausts the limit. + + The pre-flight check passes (0.0 < limit), but after the first ainvoke + round the accumulated cost exceeds the limit. This triggers pruning skip + and then ExecutionError on the next round pre-flight check. + """ + # Save previous tokens so after_scenario can reset them (issue #76 test hygiene). + context._budget_cost_token = current_accumulated_cost.set(0.0) + context._budget_max_token = current_max_cost_usd.set(0.015) + context._budget_pricing_token = current_pricing.set( + {"openai": {"gpt-3.5-turbo": {"prompt": 1000.0, "completion": 1000.0}}} + ) + + @when("I run _execute_tool_loop with an initial messages list") def step_run_execute_tool_loop(context): agent = _make_agent(context.tool_loop_config) @@ -549,6 +587,38 @@ def step_loop_result_includes_pruning_tokens(context): ) +@then("the accumulated prompt should NOT include pruning tokens") +def step_accumulated_prompt_no_pruning_tokens(context): + if hasattr(context, "tool_loop_extra_pruning_prompt"): + expected_min = context.tool_loop_extra_pruning_prompt + if context.tool_loop_result is not None: + actual = context.tool_loop_result.accumulated_prompt + elif isinstance(getattr(context, "tool_loop_error", None), _ToolLoopError): + actual = context.tool_loop_error.accumulated_prompt + else: + actual = 0 + assert actual < expected_min, ( + f"Expected accumulated_prompt to NOT include pruning tokens, " + f"but got {actual} >= pruning threshold {expected_min}" + ) + + +@then("the accumulated completion should NOT include pruning tokens") +def step_accumulated_completion_no_pruning_tokens(context): + if hasattr(context, "tool_loop_extra_pruning_completion"): + expected_min = context.tool_loop_extra_pruning_completion + if context.tool_loop_result is not None: + actual = context.tool_loop_result.accumulated_completion + elif isinstance(getattr(context, "tool_loop_error", None), _ToolLoopError): + actual = context.tool_loop_error.accumulated_completion + else: + actual = 0 + assert actual < expected_min, ( + f"Expected accumulated_completion to NOT include pruning tokens, " + f"but got {actual} >= pruning threshold {expected_min}" + ) + + @then("the messages list should contain a ToolMessage with not available content") def step_messages_contain_not_available(context): msgs = context.tool_loop_messages @@ -599,3 +669,19 @@ def step_tool_loop_error_cause_config_error(context): assert isinstance(context.tool_loop_error.cause, ConfigurationError), ( f"Expected cause=ConfigurationError, got {type(context.tool_loop_error.cause).__name__}" ) + + +@then('an ExecutionError with kind "{kind}" and reason "{reason}" should be raised') +def step_execution_error_kind_reason(context, kind, reason): + assert context.tool_loop_error is not None, "Expected an error but got none" + # The error may be a _ToolLoopError wrapping an ExecutionError + cause = context.tool_loop_error + if isinstance(cause, _ToolLoopError): + cause = cause.cause + assert isinstance(cause, ExecutionError), ( + f"Expected ExecutionError, got {type(cause).__name__}: {cause}" + ) + assert cause.kind == kind, f"Expected kind={kind!r}, got {cause.kind!r}: {cause}" + assert cause.reason == reason, ( + f"Expected reason={reason!r}, got {cause.reason!r}: {cause}" + ) diff --git a/src/cleveractors/agents/llm.py b/src/cleveractors/agents/llm.py index f6fc6c2..05467e6 100644 --- a/src/cleveractors/agents/llm.py +++ b/src/cleveractors/agents/llm.py @@ -39,13 +39,20 @@ 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, + current_pricing, +) from cleveractors.agents.tool import ToolAgent from cleveractors.core.exceptions import ( AgentCreationError, ConfigurationError, ExecutionError, MissingUsageMetadataError, + _compute_cost, ) from cleveractors.result import MAX_REASONABLE_TOKENS from cleveractors.templates.renderer import TemplateRenderer @@ -517,6 +524,114 @@ 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.warning( + "Agent %s: budget exhausted before %s invoke ($%.6f >= $%.6f 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 " + "($%.6f >= $%.6f limit)", + self.name, + tool_name, + acc, + limit, + ) + return True + return False + + def _accumulate_cost_after_invoke( + self, + prompt_tokens: int, + completion_tokens: int, + where: str = "main", + model: str | None = None, + ) -> None: + """Compute USD cost from token counts and advance ``current_accumulated_cost``. + + Uses the pricing table from ``current_pricing`` ContextVar (set by + ``PureLangGraph`` before each node execution). The cost is computed + using the same formula as ``PureLangGraph._execute_from_node``: + + cost = prompt_tokens / 1M * prompt_rate + + completion_tokens / 1M * completion_rate + + When *model* is provided (e.g. from a pruning pass that uses a + different model), that model is used for the pricing lookup instead + of ``self.model``. This ensures the in-node budget gate reflects + the actual spend even when ``pruning_model != self.model``, and works + correctly when the ``pruning_model != self.model`` guard is removed + in the future (Graa review PR #80). + + Fails fast with ``missing_pricing_entry`` when a budget limit + (``current_max_cost_usd``) is set but pricing is unavailable for + the agent's provider/model — silently skipping cost tracking would + allow uncontrolled spending against the defined limit. + """ + pricing = current_pricing.get() + limit = current_max_cost_usd.get() + provider = self.provider + model = model or self.model + if not pricing: + if limit is not None: + raise ExecutionError( + f"Cost limit (${limit}) specified but pricing table is empty " + f"or missing for agent '{self.name}' ({provider}/{model}). " + f"Provide pricing data or remove max_cost_usd.", + kind="cost", + reason="missing_pricing_entry", + ) + return + try: + _cost = _compute_cost( + pricing, provider, model, prompt_tokens, completion_tokens + ) + except ExecutionError: + if limit is not None: + raise + return + _current = current_accumulated_cost.get() + current_accumulated_cost.set(_current + _cost) + logger.debug( + "Agent %s: accumulated $%.6f after %s invoke " + "(%d prompt + %d completion tokens)", + self.name, + _current + _cost, + where, + prompt_tokens, + completion_tokens, + ) + def _get_model_context_window(self) -> int: """Return the model's advertised context window size in tokens. @@ -891,11 +1006,13 @@ 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) _accumulated_prompt += _bp _accumulated_completion += _bc + self._accumulate_cost_after_invoke(_bp, _bc, "synthesis") _synth_tool_calls: list[dict[str, Any]] = ( getattr(response, "tool_calls", None) or [] ) @@ -972,6 +1089,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: ( @@ -999,6 +1117,12 @@ class LLMAgent(AgentWithMemory): raise _accumulated_prompt += _bp_tok _accumulated_completion += _bc_tok + self._accumulate_cost_after_invoke( + _bp_tok, + _bc_tok, + "pruning", + model=(self._pruning_model or self.model), + ) logger.info( "Agent %s: pruning pass for %r: " "%d -> %d chars", @@ -1031,20 +1155,28 @@ 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) _accumulated_prompt += _bp3 _accumulated_completion += _bc3 + self._accumulate_cost_after_invoke( + _bp3, _bc3, "synthesis_followup" + ) _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 _mp, _mc, _ = self._extract_token_counts(response) _accumulated_prompt += _mp _accumulated_completion += _mc + self._accumulate_cost_after_invoke(_mp, _mc, "main") response_tool_calls: list[dict[str, Any]] = ( getattr(response, "tool_calls", None) or [] @@ -1121,6 +1253,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: @@ -1149,6 +1282,12 @@ class LLMAgent(AgentWithMemory): raise _accumulated_prompt += _mlp_tok _accumulated_completion += _mlc_tok + self._accumulate_cost_after_invoke( + _mlp_tok, + _mlc_tok, + "pruning", + model=(self._pruning_model or self.model), + ) logger.info( "Agent %s: pruning pass for %r: %d -> %d chars", self.name, @@ -1200,11 +1339,13 @@ 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) _accumulated_prompt += _sp _accumulated_completion += _sc + self._accumulate_cost_after_invoke(_sp, _sc, "stuck_model") # Allow *one* final tool-call round so the model can # write files or perform other last-minute operations. @@ -1288,11 +1429,15 @@ 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) _accumulated_prompt += _sfp _accumulated_completion += _sfc + self._accumulate_cost_after_invoke( + _sfp, _sfc, "stuck_model_followup" + ) # Log accumulated token usage across all rounds. if _accumulated_prompt > 0 or _accumulated_completion > 0: @@ -1531,10 +1676,12 @@ class LLMAgent(AgentWithMemory): _completion_tokens = _loop_result.accumulated_completion else: # No tools: single plain ainvoke() — no tool dispatch or synthesis. + self._check_budget_before_invoke("main") _no_tools_resp = await self._retry_ainvoke(messages) _pt, _ct, _ = self._extract_token_counts(_no_tools_resp) _prompt_tokens = _pt _completion_tokens = _ct + self._accumulate_cost_after_invoke(_pt, _ct, "main") response_text = str(_no_tools_resp.content) # Capture token counts immediately after ainvoke() succeeds. # Setting the sentinel variables here marks the "ainvoke succeeded" @@ -1851,6 +1998,9 @@ class LLMAgent(AgentWithMemory): return # generator done; finally block still runs # ── No-tools path: real token-by-token astream() ───────────── + # Budget gate before streaming starts (issue #76). + self._check_budget_before_invoke("main") + # Wrap the initial stream connection (first chunk) in retry so # transient LLM provider failures are retried at chunk granularity # (ADR-2032 D-5 / issue #70 review). Once the first chunk @@ -1952,6 +2102,9 @@ class LLMAgent(AgentWithMemory): _captured_completion = _completion_tokens self._last_token_usage = (_captured_prompt, _captured_completion) last_token_usage_var.set((_captured_prompt, _captured_completion)) + self._accumulate_cost_after_invoke( + _captured_prompt, _captured_completion, "stream" + ) # Update memory if enabled (spec §4.4.4). # Mirrors process_message() memory block. diff --git a/src/cleveractors/agents/retry.py b/src/cleveractors/agents/retry.py index 8050c3f..57e95c2 100644 --- a/src/cleveractors/agents/retry.py +++ b/src/cleveractors/agents/retry.py @@ -27,6 +27,21 @@ 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 +) +current_pricing: contextvars.ContextVar[dict[str, Any] | None] = contextvars.ContextVar( + "current_pricing", default=None +) + _INITIAL_RETRY_DELAY: float = 0.5 diff --git a/src/cleveractors/core/exceptions.py b/src/cleveractors/core/exceptions.py index 610ede6..1e38a71 100644 --- a/src/cleveractors/core/exceptions.py +++ b/src/cleveractors/core/exceptions.py @@ -61,6 +61,67 @@ class ExecutionError(CleverAgentsException): self.reason = reason +def _compute_cost( + pricing: dict, + provider: str, + model: str, + prompt_tokens: int, + completion_tokens: int, +) -> float: + """Compute USD cost from a pricing table. + + Looks up *provider → model → {prompt, completion}* rates from the + pricing dict, validates them, and returns the cost in USD. + + cost = prompt_tokens / 1M * prompt_rate + completion_tokens / 1M * completion_rate + + Raises ``ExecutionError(kind="cost", reason="missing_pricing_entry")`` + when the provider, model, or rate keys are missing or invalid. + + This exists in ``exceptions.py`` because it is the one module imported + by all three cost-enforcement consumers (``pure_graph.py``, + ``runtime_dispatch.py``, ``llm.py``) without circular-dependency risk. + """ + _provider_pricing = pricing.get(provider) + if not isinstance(_provider_pricing, dict): + raise ExecutionError( + f"Missing pricing entry for provider '{provider}'", + kind="cost", + reason="missing_pricing_entry", + ) + _model_pricing = _provider_pricing.get(model) + if not isinstance(_model_pricing, dict): + raise ExecutionError( + f"Missing pricing entry for model '{model}' under provider '{provider}'", + kind="cost", + reason="missing_pricing_entry", + ) + _prompt_rate_raw = _model_pricing.get("prompt") + _completion_rate_raw = _model_pricing.get("completion") + if _prompt_rate_raw is None or _completion_rate_raw is None: + raise ExecutionError( + f"Incomplete pricing entry for model '{model}' " + f"under provider '{provider}': missing 'prompt' " + f"or 'completion' rate key", + kind="cost", + reason="missing_pricing_entry", + ) + try: + _prompt_rate = float(_prompt_rate_raw) + _completion_rate = float(_completion_rate_raw) + except (TypeError, ValueError) as err: + raise ExecutionError( + f"Invalid pricing rate for model '{model}' " + f"under provider '{provider}': {err}", + kind="cost", + reason="missing_pricing_entry", + ) from err + return ( + prompt_tokens / 1_000_000.0 * _prompt_rate + + completion_tokens / 1_000_000.0 * _completion_rate + ) + + class ApplicationError(CleverAgentsException): """Exception raised for application-level errors.""" diff --git a/src/cleveractors/langgraph/pure_graph.py b/src/cleveractors/langgraph/pure_graph.py index 6d3da1d..b08fe6c 100644 --- a/src/cleveractors/langgraph/pure_graph.py +++ b/src/cleveractors/langgraph/pure_graph.py @@ -8,17 +8,24 @@ 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 from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager from dataclasses import dataclass, field 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, + current_pricing, +) from cleveractors.context_manager import ContextManager -from cleveractors.core.exceptions import ExecutionError +from cleveractors.core.exceptions import ExecutionError, _compute_cost from cleveractors.langgraph.dynamic_router import DynamicRouterNode from cleveractors.langgraph.nodes import Edge, Node, NodeConfig, NodeType from cleveractors.langgraph.state import GraphState, StateManager @@ -535,6 +542,152 @@ class PureLangGraph: global_context.clear() global_context.update(final_metadata) + def _get_max_cost_usd_limit(self) -> float | None: + """Parse and return the ``max_cost_usd`` limit, or ``None`` if absent. + + Raises ``ExecutionError(kind='cost', reason='missing_pricing_entry')`` + when ``max_cost_usd`` is specified in limits but pricing is empty — + a limit without pricing is a configuration error. + + Returns ``None`` (no-op) only when the limit key is absent, + the value is a bool (nonsensical), or the value cannot be coerced + to float. + """ + _raw = self._limits.get("max_cost_usd") + if _raw is None: + return None + if not self._pricing: + raise ExecutionError( + "max_cost_usd limit specified but pricing table is empty. " + "Provide pricing data or remove max_cost_usd.", + kind="cost", + reason="missing_pricing_entry", + ) + if isinstance(_raw, bool): + return None + try: + return float(_raw) + except (TypeError, ValueError): + return None + + def _set_budget_context_for_node( + self, + ) -> tuple[ + contextvars.Token[float] | None, + contextvars.Token[float | None] | None, + contextvars.Token[dict[str, Any] | None] | None, + ]: + """Set budget ContextVars for in-node budget checks (issue #76). + + Must be called before each node execution. Returns a tuple of + ``(cost_token, max_token, pricing_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 = self._get_max_cost_usd_limit() + _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) + _pricing_token: contextvars.Token[dict[str, Any] | None] | None = None + if self._pricing: + _pricing_token = current_pricing.set(self._pricing) + return (_cost_token, _max_token, _pricing_token) + + @staticmethod + def _reset_budget_context( + cost_token: contextvars.Token[float] | None, + max_token: contextvars.Token[float | None] | None, + pricing_token: contextvars.Token[dict[str, Any] | 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) + if pricing_token is not None: + current_pricing.reset(pricing_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. + + Uses ``>=`` (not ``>``) intentionally: exactly-at-limit blocks the + *next* node from entering but does not abort the current node's + post-node check (which uses ``>``). This guarantees that a node + that exactly hits the limit completes normally; only the next node + sees the pre-flight block (Graa review PR #80). + """ + _max = self._get_max_cost_usd_limit() + if _max is not None and self._accumulated_cost >= _max: + self.logger.warning( + "Budget pre-flight: accumulated $%.6f >= limit $%.6f 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", + ) + + @asynccontextmanager + async def _node_budget_context(self, node_name: str) -> AsyncGenerator[None, None]: + """Context manager wrapping the budget pre-flight, ContextVar setup, + and finally-reset that was copy-pasted across all four execution + branches. Using ``async with self._node_budget_context(node_name):`` + makes the reset un-forgettable. + """ + self._check_budget_pre_flight(node_name) + _cost_tok, _max_tok, _pricing_tok = self._set_budget_context_for_node() + try: + yield + finally: + self._reset_budget_context(_cost_tok, _max_tok, _pricing_tok) + + def _check_post_node_cost_limit(self) -> None: + """Check accumulated cost against ``max_cost_usd`` after a node completes. + + Raises ``ExecutionError(kind='cost', reason='budget_exhausted')`` when + the limit is exceeded. Uses ``>`` (not ``>=``) intentionally: a node + that exactly hits the limit completes normally; only the *next* node's + pre-flight check (``>=``) blocks it. + + This centralises the bool-guard + float() coercion that was previously + inlined in all three post-node cost blocks. + """ + _raw = self._limits.get("max_cost_usd") + if _raw is None: + return + if isinstance(_raw, bool): + raise ExecutionError( + f"Invalid max_cost_usd value {_raw!r}: bool is not a valid cost limit", + kind="cost", + ) + try: + _limit = float(_raw) + except (TypeError, ValueError) as _err: + raise ExecutionError( + f"Invalid max_cost_usd value {_raw!r}: {_err}", + kind="cost", + ) from _err + if self._accumulated_cost > _limit: + raise ExecutionError( + f"Cost limit exceeded: {self._accumulated_cost:.6f} USD > {_raw} USD", + kind="cost", + reason="budget_exhausted", + ) + async def _execute_from_node( self, node_name: str, message: Any, depth: int = 0 ) -> Any: @@ -790,204 +943,88 @@ class PureLangGraph: # Pass graph name for retry error reporting (ADR-2032 D-7). state.metadata["graph_name"] = self.name - # Execute node with current state - try: - self.logger.debug(f"About to execute node.execute() for {node_name}") - result = await node.execute(state) - result_str = ( - str(result)[:100] if isinstance(result, str) else str(result)[:100] - ) - self.logger.debug(f"Node {node_name} returned: {result_str}") + # ── Budget pre-flight + ContextVar setup (issue #76) ────────────── + async with self._node_budget_context(node_name): + try: + self.logger.debug(f"About to execute node.execute() for {node_name}") + result = await node.execute(state) + result_str = ( + str(result)[:100] if isinstance(result, str) else str(result)[:100] + ) + self.logger.debug(f"Node {node_name} returned: {result_str}") - # Update state with node results - if isinstance(result, dict): - # Extract token usage BEFORE passing to update_state so the - # internal side-channel key is never written into GraphState - # (m5 of review for issue #14). Popping here means a custom - # state_class that accidentally defines _node_token_usage as a - # real field will not receive the internal dict. - _tok_info = result.pop("_node_token_usage", None) + # Update state with node results + if isinstance(result, dict): + # Extract token usage BEFORE passing to update_state + _tok_info = result.pop("_node_token_usage", None) - self.state_manager.update_state(result, node_id=node_name) - # Extract the actual message content from the result - # Agents return messages in a 'messages' array with 'content' field - if "messages" in result and result["messages"]: - # Get the last message's content - last_message = result["messages"][-1] - if isinstance(last_message, dict) and "content" in last_message: - output_message = last_message["content"] + self.state_manager.update_state(result, node_id=node_name) + if "messages" in result and result["messages"]: + last_message = result["messages"][-1] + if isinstance(last_message, dict) and "content" in last_message: + output_message = last_message["content"] + else: + output_message = str(last_message) else: - output_message = str(last_message) + output_message = result.get( + "content", result.get("output", message) + ) + + if _tok_info is not None and not isinstance(_tok_info, dict): + self.logger.warning( + "Node %s returned non-dict _node_token_usage (%r); " + "token accounting skipped", + node_name, + type(_tok_info).__name__, + ) + if isinstance(_tok_info, dict): + _pt = _safe_token_int(_tok_info.get("prompt_tokens", 0)) + _ct = _safe_token_int(_tok_info.get("completion_tokens", 0)) + self._node_usages.append( + ( + str(_tok_info.get("node_id", node_name)), + str(_tok_info.get("provider", "unknown")), + str(_tok_info.get("model", "unknown")), + _pt, + _ct, + ) + ) + + if self._pricing: + _node_cost = _compute_cost( + self._pricing, + str(_tok_info.get("provider", "unknown")), + str(_tok_info.get("model", "unknown")), + _pt, + _ct, + ) + self._accumulated_cost += _node_cost + self._check_post_node_cost_limit() else: - # Fallback to direct content/output fields - output_message = result.get( - "content", result.get("output", message) - ) + output_message = result - # Collect per-node token usage if this agent node reported it - # (AC4, issue #14). Node._execute_agent() sets - # _node_token_usage in the result dict for LLM agent nodes. - # m-1 fix: log a warning when _node_token_usage is present but - # not a dict so that malformed values surface in production logs - # rather than silently disabling token accounting and cost - # enforcement. - if _tok_info is not None and not isinstance(_tok_info, dict): - self.logger.warning( - "Node %s returned non-dict _node_token_usage (%r); " - "token accounting skipped", - node_name, - type(_tok_info).__name__, - ) - if isinstance(_tok_info, dict): - # m-2 fix: compute _pt and _ct once at the top of the - # isinstance block and reuse them in both _node_usages and - # the cost block, eliminating the double _safe_token_int() - # call that doubled log noise for malformed inputs. - _pt = _safe_token_int(_tok_info.get("prompt_tokens", 0)) - _ct = _safe_token_int(_tok_info.get("completion_tokens", 0)) - self._node_usages.append( - ( - str(_tok_info.get("node_id", node_name)), - str(_tok_info.get("provider", "unknown")), - str(_tok_info.get("model", "unknown")), - _pt, - _ct, - ) - ) + self.state_manager.state.metadata["last_output"] = output_message - # Cost accumulation and enforcement (ADR-2029, issue #15). - # Only when the router supplied a non-empty pricing table; an - # empty dict means cost tracking was not requested and is skipped - # entirely. When the table is non-empty, a missing provider or - # model entry is a hard error — proceeding with an assumed zero - # price is forbidden because it would enable unbounded spend. - if self._pricing: - _provider = str(_tok_info.get("provider", "unknown")) - _model = str(_tok_info.get("model", "unknown")) - # _pt and _ct already computed above (m-2 fix) - - _provider_pricing = self._pricing.get(_provider) - if _provider_pricing is None or not isinstance( - _provider_pricing, dict - ): - raise ExecutionError( - f"Missing pricing entry for provider '{_provider}'", - kind="cost", - reason="missing_pricing_entry", - ) - _model_pricing = _provider_pricing.get(_model) - if _model_pricing is None or not isinstance( - _model_pricing, dict - ): - raise ExecutionError( - f"Missing pricing entry for model '{_model}' " - f"under provider '{_provider}'", - kind="cost", - reason="missing_pricing_entry", + output_str = str(output_message) if output_message else "" + for prefix in ["GOTO_", "ROUTE_"]: + if prefix in output_str: + cmd_part = output_str.split(":", 1)[0] + target = cmd_part.replace(prefix, "").strip().lower() + if target: + self.state_manager.state.metadata["next_node"] = target + self.logger.debug( + f"Parsed routing command from {node_name}: " + f"next_node={target}" ) + break - # Validate that both rate keys exist and are numeric - # (ADR-2029: "proceeding with a missing or zero price is - # forbidden"). A missing key or a non-numeric value is - # treated the same as a missing pricing entry — the - # library must not silently fall back to $0/token. - _prompt_rate_raw = _model_pricing.get("prompt") - _completion_rate_raw = _model_pricing.get("completion") - if _prompt_rate_raw is None or _completion_rate_raw is None: - raise ExecutionError( - f"Incomplete pricing entry for model '{_model}' " - f"under provider '{_provider}': missing 'prompt' " - f"or 'completion' rate key", - kind="cost", - reason="missing_pricing_entry", - ) - try: - # Rates are USD per million tokens (ADR-2029 example: - # gpt-4.1-mini prompt=$0.15/1M, completion=$0.60/1M). - _prompt_rate = float(_prompt_rate_raw) - _completion_rate = float(_completion_rate_raw) - except (TypeError, ValueError) as _rate_err: - raise ExecutionError( - f"Invalid pricing rate for model '{_model}' " - f"under provider '{_provider}': {_rate_err}", - kind="cost", - reason="missing_pricing_entry", - ) from _rate_err - _node_cost = ( - _pt / 1_000_000.0 * _prompt_rate - + _ct / 1_000_000.0 * _completion_rate - ) - self._accumulated_cost += _node_cost - - # C-1 fix (ADR-2029, issue #15): coerce max_cost_usd - # with a narrow try/except BEFORE the broad - # "except Exception" handler so that a malformed value - # (e.g. a non-numeric string) raises ExecutionError - # immediately rather than being silently swallowed by - # the broad handler and causing the graph to return the - # input string as if execution succeeded. - _max_cost = self._limits.get("max_cost_usd") - if _max_cost is not None: - # Bool guard: float(True)==1.0 / float(False)==0.0 - # would silently set a nonsensical budget. Reject - # bools before the float() conversion, matching the - # pattern used for max_depth/max_model_calls/max_tool_calls. - if isinstance(_max_cost, bool): - raise ExecutionError( - f"Invalid max_cost_usd value {_max_cost!r}: " - "bool is not a valid cost limit", - kind="cost", - ) - try: - _max_cost_float = float(_max_cost) - except (TypeError, ValueError) as _cost_err: - # A malformed limit value (e.g. a non-numeric string) - # is a router configuration error, not a missing pricing - # entry. Use reason="" (the documented default for - # non-cost-specific errors per ADR-2029) so router logs - # and alerts correctly identify the root cause. - raise ExecutionError( - f"Invalid max_cost_usd value " - f"{_max_cost!r}: {_cost_err}", - kind="cost", - ) from _cost_err - if self._accumulated_cost > _max_cost_float: - raise ExecutionError( - f"Cost limit exceeded: " - f"{self._accumulated_cost:.6f} USD " - f"> {_max_cost} USD", - kind="cost", - reason="budget_exhausted", - ) - else: - output_message = result - - # Store the output message in state for edge conditions to use - self.state_manager.state.metadata["last_output"] = output_message - - # Parse CleverAgents v2.0 routing commands from agent outputs - # so edge conditions can route to the correct next node. - output_str = str(output_message) if output_message else "" - for prefix in ["GOTO_", "ROUTE_"]: - if prefix in output_str: - # Extract target node from "GOTO_NODE:rest_of_message" - cmd_part = output_str.split(":", 1)[0] - target = cmd_part.replace(prefix, "").strip().lower() - if target: - self.state_manager.state.metadata["next_node"] = target - self.logger.debug( - f"Parsed routing command from {node_name}: next_node={target}" - ) - 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 + except ExecutionError: + raise + except Exception as e: + self.logger.error( + f"Error executing node {node_name}: {e}", exc_info=True + ) + output_message = message # Determine next nodes self.logger.debug( @@ -1577,50 +1614,62 @@ class PureLangGraph: # they are also yielded immediately (fast path) or held back # until the stream completes (slow path). _collected_tokens: List[str] = [] - try: - async for token in node._stream_agent(state): - _collected_tokens.append(token) # always accumulate - if _all_edges_unconditional: - yield token # yield immediately for fast path - full_response = "".join(_collected_tokens) # always correct + # ── Budget pre-flight + ContextVar setup (issue #76) ─────────── + async with self._node_budget_context(node_name): + try: + async for token in node._stream_agent(state): + _collected_tokens.append(token) + if _all_edges_unconditional: + yield token + full_response = "".join(_collected_tokens) - # Update state with the full response only on success. - # Mirrors _execute_from_node which does NOT update state.messages - # on failure — it only sets state.metadata["last_output"] for - # routing purposes. Updating state on failure would persist the - # user's input as an assistant message, corrupting routing - # decisions and memory for subsequent turns (Major #2 fix). - agent_result: dict[str, Any] = { - "messages": [ - { - "role": "assistant", - "content": full_response, - "node": node_name, - "agent": node.config.agent, - } - ], - "current_node": node_name, - "metadata": {"last_agent_node": node_name}, - } - self.state_manager.update_state(agent_result, node_id=node_name) - except ExecutionError: - raise - except Exception as e: # pylint: disable=broad-exception-caught - # M3 fix: re-raise non-ExecutionError exceptions so callers - # can distinguish "stream completed normally" from "stream - # failed mid-way" and map them to HTTP 5xx/4xx. Wrapping - # in ExecutionError mirrors the non-streaming path's - # exception handling and preserves the original cause. - # The exception type is logged at ERROR level; the message - # is deliberately sanitised (no {e} interpolation) to avoid - # leaking sensitive provider details (API key fragments, - # internal URLs) — use `from e` to preserve the cause chain. - self.logger.error( - "Agent node '%s' streaming failed: %s", - node_name, - type(e).__name__, - ) - raise ExecutionError("Agent node streaming failed") from e + agent_result: dict[str, Any] = { + "messages": [ + { + "role": "assistant", + "content": full_response, + "node": node_name, + "agent": node.config.agent, + } + ], + "current_node": node_name, + "metadata": {"last_agent_node": node_name}, + } + self.state_manager.update_state(agent_result, node_id=node_name) + + _tok_info = getattr(node, "_last_stream_usage", None) + if isinstance(_tok_info, dict): + _pt = _safe_token_int(_tok_info.get("prompt_tokens", 0)) + _ct = _safe_token_int(_tok_info.get("completion_tokens", 0)) + self._node_usages.append( + ( + str(_tok_info.get("node_id", node_name)), + str(_tok_info.get("provider", "unknown")), + str(_tok_info.get("model", "unknown")), + _pt, + _ct, + ) + ) + + if self._pricing: + _stream_node_cost = _compute_cost( + self._pricing, + str(_tok_info.get("provider", "unknown")), + str(_tok_info.get("model", "unknown")), + _pt, + _ct, + ) + self._accumulated_cost += _stream_node_cost + self._check_post_node_cost_limit() + except ExecutionError: + raise + except Exception as e: + self.logger.error( + "Agent node '%s' streaming failed: %s", + node_name, + type(e).__name__, + ) + raise ExecutionError("Agent node streaming failed") from e # Update last_output for routing (mirrors non-streaming path). # Note: this is only reached on the success path; the except blocks @@ -1651,105 +1700,6 @@ class PureLangGraph: ) break - # Collect per-node token usage (stored in node._last_stream_usage) - _tok_info = getattr(node, "_last_stream_usage", None) - if isinstance(_tok_info, dict): - _pt = _safe_token_int(_tok_info.get("prompt_tokens", 0)) - _ct = _safe_token_int(_tok_info.get("completion_tokens", 0)) - self._node_usages.append( - ( - str(_tok_info.get("node_id", node_name)), - str(_tok_info.get("provider", "unknown")), - str(_tok_info.get("model", "unknown")), - _pt, - _ct, - ) - ) - - # M3 fix: Cost enforcement in streaming path. - # Mirrors _execute_from_node cost block so that max_cost_usd - # is honoured for streaming requests (the most expensive use case). - if self._pricing: - _stream_provider = str(_tok_info.get("provider", "unknown")) - _stream_model = str(_tok_info.get("model", "unknown")) - _stream_provider_pricing = self._pricing.get(_stream_provider) - if _stream_provider_pricing is None or not isinstance( - _stream_provider_pricing, dict - ): - raise ExecutionError( - f"Missing pricing entry for provider '{_stream_provider}'", - kind="cost", - reason="missing_pricing_entry", - ) - _stream_model_pricing = _stream_provider_pricing.get( - _stream_model - ) - if _stream_model_pricing is None or not isinstance( - _stream_model_pricing, dict - ): - raise ExecutionError( - f"Missing pricing entry for model '{_stream_model}' " - f"under provider '{_stream_provider}'", - kind="cost", - reason="missing_pricing_entry", - ) - _stream_prompt_rate_raw = _stream_model_pricing.get("prompt") - _stream_completion_rate_raw = _stream_model_pricing.get( - "completion" - ) - if ( - _stream_prompt_rate_raw is None - or _stream_completion_rate_raw is None - ): - raise ExecutionError( - f"Incomplete pricing entry for model '{_stream_model}' " - f"under provider '{_stream_provider}': missing 'prompt' " - f"or 'completion' rate key", - kind="cost", - reason="missing_pricing_entry", - ) - try: - _stream_prompt_rate = float(_stream_prompt_rate_raw) - _stream_completion_rate = float(_stream_completion_rate_raw) - except (TypeError, ValueError) as _stream_rate_err: - raise ExecutionError( - f"Invalid pricing rate for model '{_stream_model}' " - f"under provider '{_stream_provider}': " - f"{_stream_rate_err}", - kind="cost", - reason="missing_pricing_entry", - ) from _stream_rate_err - _stream_node_cost = ( - _pt / 1_000_000.0 * _stream_prompt_rate - + _ct / 1_000_000.0 * _stream_completion_rate - ) - self._accumulated_cost += _stream_node_cost - - _stream_max_cost = self._limits.get("max_cost_usd") - if _stream_max_cost is not None: - if isinstance(_stream_max_cost, bool): - raise ExecutionError( - f"Invalid max_cost_usd value {_stream_max_cost!r}: " - "bool is not a valid cost limit", - kind="cost", - ) - try: - _stream_max_cost_float = float(_stream_max_cost) - except (TypeError, ValueError) as _stream_cost_err: - raise ExecutionError( - f"Invalid max_cost_usd value " - f"{_stream_max_cost!r}: {_stream_cost_err}", - kind="cost", - ) from _stream_cost_err - if self._accumulated_cost > _stream_max_cost_float: - raise ExecutionError( - f"Cost limit exceeded: " - f"{self._accumulated_cost:.6f} USD " - f"> {_stream_max_cost} USD", - kind="cost", - reason="budget_exhausted", - ) - # Confirmed terminal: _statically_terminal=True guarantees that # all static successors are END/end sentinels, so # _get_next_nodes() can only return those same targets (or a @@ -1768,149 +1718,61 @@ class PureLangGraph: else: # Intermediate AGENT node: use ainvoke() (node.execute) to avoid # buffering tokens only to discard them (AC2). - try: - result = await node.execute(state) - if isinstance(result, dict): - _tok_info = result.pop("_node_token_usage", None) - self.state_manager.update_state(result, node_id=node_name) - if "messages" in result and result["messages"]: - last_msg = result["messages"][-1] - if isinstance(last_msg, dict) and "content" in last_msg: - full_response = last_msg["content"] + # ── Budget pre-flight + ContextVar setup (issue #76) ─────────── + async with self._node_budget_context(node_name): + try: + result = await node.execute(state) + if isinstance(result, dict): + _tok_info = result.pop("_node_token_usage", None) + self.state_manager.update_state(result, node_id=node_name) + if "messages" in result and result["messages"]: + last_msg = result["messages"][-1] + if isinstance(last_msg, dict) and "content" in last_msg: + full_response = last_msg["content"] + else: + full_response = str(last_msg) else: - full_response = str(last_msg) + full_response = result.get( + "content", result.get("output", message) + ) + if isinstance(_tok_info, dict): + _pt = _safe_token_int(_tok_info.get("prompt_tokens", 0)) + _ct = _safe_token_int( + _tok_info.get("completion_tokens", 0) + ) + self._node_usages.append( + ( + str(_tok_info.get("node_id", node_name)), + str(_tok_info.get("provider", "unknown")), + str(_tok_info.get("model", "unknown")), + _pt, + _ct, + ) + ) + + if self._pricing: + _int_node_cost = _compute_cost( + self._pricing, + str(_tok_info.get("provider", "unknown")), + str(_tok_info.get("model", "unknown")), + _pt, + _ct, + ) + self._accumulated_cost += _int_node_cost + self._check_post_node_cost_limit() else: - full_response = result.get( - "content", result.get("output", message) + full_response = ( + str(result) if result is not None else message ) - if isinstance(_tok_info, dict): - _pt = _safe_token_int(_tok_info.get("prompt_tokens", 0)) - _ct = _safe_token_int(_tok_info.get("completion_tokens", 0)) - self._node_usages.append( - ( - str(_tok_info.get("node_id", node_name)), - str(_tok_info.get("provider", "unknown")), - str(_tok_info.get("model", "unknown")), - _pt, - _ct, - ) - ) - - # M3 fix (intermediate branch): Cost enforcement in - # streaming path for intermediate AGENT nodes. - # Mirrors the terminal AGENT branch cost block so that - # max_cost_usd is honoured for every AGENT node in a - # multi-AGENT graph, not just the terminal one. - if self._pricing: - _int_provider = str( - _tok_info.get("provider", "unknown") - ) - _int_model = str(_tok_info.get("model", "unknown")) - _int_provider_pricing = self._pricing.get(_int_provider) - if _int_provider_pricing is None or not isinstance( - _int_provider_pricing, dict - ): - raise ExecutionError( - f"Missing pricing entry for provider " - f"'{_int_provider}'", - kind="cost", - reason="missing_pricing_entry", - ) - _int_model_pricing = _int_provider_pricing.get( - _int_model - ) - if _int_model_pricing is None or not isinstance( - _int_model_pricing, dict - ): - raise ExecutionError( - f"Missing pricing entry for model " - f"'{_int_model}' under provider " - f"'{_int_provider}'", - kind="cost", - reason="missing_pricing_entry", - ) - _int_prompt_rate_raw = _int_model_pricing.get("prompt") - _int_completion_rate_raw = _int_model_pricing.get( - "completion" - ) - if ( - _int_prompt_rate_raw is None - or _int_completion_rate_raw is None - ): - raise ExecutionError( - f"Incomplete pricing entry for model " - f"'{_int_model}' under provider " - f"'{_int_provider}': missing 'prompt' " - f"or 'completion' rate key", - kind="cost", - reason="missing_pricing_entry", - ) - try: - _int_prompt_rate = float(_int_prompt_rate_raw) - _int_completion_rate = float( - _int_completion_rate_raw - ) - except (TypeError, ValueError) as _int_rate_err: - raise ExecutionError( - f"Invalid pricing rate for model " - f"'{_int_model}' under provider " - f"'{_int_provider}': {_int_rate_err}", - kind="cost", - reason="missing_pricing_entry", - ) from _int_rate_err - _int_node_cost = ( - _pt / 1_000_000.0 * _int_prompt_rate - + _ct / 1_000_000.0 * _int_completion_rate - ) - self._accumulated_cost += _int_node_cost - - _int_max_cost = self._limits.get("max_cost_usd") - if _int_max_cost is not None: - if isinstance(_int_max_cost, bool): - raise ExecutionError( - f"Invalid max_cost_usd value " - f"{_int_max_cost!r}: bool is not a " - "valid cost limit", - kind="cost", - ) - try: - _int_max_cost_float = float(_int_max_cost) - except ( - TypeError, - ValueError, - ) as _int_cost_err: - raise ExecutionError( - f"Invalid max_cost_usd value " - f"{_int_max_cost!r}: {_int_cost_err}", - kind="cost", - ) from _int_cost_err - if self._accumulated_cost > _int_max_cost_float: - raise ExecutionError( - f"Cost limit exceeded: " - f"{self._accumulated_cost:.6f} USD " - f"> {_int_max_cost} USD", - kind="cost", - reason="budget_exhausted", - ) - else: - full_response = str(result) if result is not None else message - except ExecutionError: - raise - except Exception as e: # pylint: disable=broad-exception-caught - # M3 fix: re-raise non-ExecutionError exceptions so callers - # can distinguish "node completed normally" from "node - # failed" and map them to HTTP 5xx/4xx. Wrapping in - # ExecutionError mirrors the non-streaming path's behavior. - # The exception type is logged at ERROR level; the message - # is deliberately sanitised (no {e} interpolation) to avoid - # leaking sensitive provider details — use `from e` to - # preserve the cause chain. - self.logger.error( - "Intermediate agent node '%s' ainvoke failed: %s", - node_name, - type(e).__name__, - ) - raise ExecutionError("Intermediate agent node failed") from e + except ExecutionError: + raise + except Exception as e: + self.logger.error( + "Intermediate agent node '%s' ainvoke failed: %s", + node_name, + type(e).__name__, + ) + raise ExecutionError("Intermediate agent node failed") from e self.state_manager.state.metadata["last_output"] = full_response @@ -2023,45 +1885,47 @@ class PureLangGraph: else: # Non-AGENT node: run with node.execute() (ainvoke path) - output_message: Any = message - try: - result = await node.execute(state) - if isinstance(result, dict): - _tok_info = result.pop("_node_token_usage", None) - self.state_manager.update_state(result, node_id=node_name) - if "messages" in result and result["messages"]: - last_msg = result["messages"][-1] - if isinstance(last_msg, dict) and "content" in last_msg: - output_message = last_msg["content"] + # ── Budget pre-flight + ContextVar setup (issue #76) ─────────── + async with self._node_budget_context(node_name): + output_message: Any = message + try: + result = await node.execute(state) + if isinstance(result, dict): + _tok_info = result.pop("_node_token_usage", None) + self.state_manager.update_state(result, node_id=node_name) + if "messages" in result and result["messages"]: + last_msg = result["messages"][-1] + if isinstance(last_msg, dict) and "content" in last_msg: + output_message = last_msg["content"] + else: + output_message = str(last_msg) else: - output_message = str(last_msg) - else: - output_message = result.get( - "content", result.get("output", message) - ) - if isinstance(_tok_info, dict): - _pt = _safe_token_int(_tok_info.get("prompt_tokens", 0)) - _ct = _safe_token_int(_tok_info.get("completion_tokens", 0)) - self._node_usages.append( - ( - str(_tok_info.get("node_id", node_name)), - str(_tok_info.get("provider", "unknown")), - str(_tok_info.get("model", "unknown")), - _pt, - _ct, + output_message = result.get( + "content", result.get("output", message) ) - ) - else: - output_message = result - except ExecutionError: - raise - except Exception as e: # pylint: disable=broad-exception-caught - self.logger.error( - "Non-agent node '%s' execution failed: %s", - node_name, - type(e).__name__, - ) - output_message = message + if isinstance(_tok_info, dict): + _pt = _safe_token_int(_tok_info.get("prompt_tokens", 0)) + _ct = _safe_token_int(_tok_info.get("completion_tokens", 0)) + self._node_usages.append( + ( + str(_tok_info.get("node_id", node_name)), + str(_tok_info.get("provider", "unknown")), + str(_tok_info.get("model", "unknown")), + _pt, + _ct, + ) + ) + else: + output_message = result + except ExecutionError: + raise + except Exception as e: + self.logger.error( + "Non-agent node '%s' execution failed: %s", + node_name, + type(e).__name__, + ) + output_message = message self.state_manager.state.metadata["last_output"] = output_message diff --git a/src/cleveractors/runtime_dispatch.py b/src/cleveractors/runtime_dispatch.py index ba9d22a..7063a6b 100644 --- a/src/cleveractors/runtime_dispatch.py +++ b/src/cleveractors/runtime_dispatch.py @@ -37,6 +37,7 @@ from cleveractors.core.exceptions import ( AgentCreationError, ConfigurationError, ExecutionError, + _compute_cost, ) from cleveractors.langgraph.nodes import Edge, NodeConfig, NodeType from cleveractors.langgraph.pure_graph import PureGraphConfig, PureLangGraph @@ -948,49 +949,25 @@ async def _execute_llm_stream( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, ) - if executor.pricing: - _llm_provider_pricing = executor.pricing.get(provider) - if _llm_provider_pricing is None or not isinstance( - _llm_provider_pricing, dict - ): - raise ExecutionError( - f"Missing pricing entry for provider '{provider}'", - kind="cost", - reason="missing_pricing_entry", - ) - _llm_model_pricing = _llm_provider_pricing.get(model) - if _llm_model_pricing is None or not isinstance(_llm_model_pricing, dict): - raise ExecutionError( - f"Missing pricing entry for model '{model}' " - f"under provider '{provider}'", - kind="cost", - reason="missing_pricing_entry", - ) - _llm_prompt_rate_raw = _llm_model_pricing.get("prompt") - _llm_completion_rate_raw = _llm_model_pricing.get("completion") - if _llm_prompt_rate_raw is None or _llm_completion_rate_raw is None: - raise ExecutionError( - f"Incomplete pricing entry for model '{model}' " - f"under provider '{provider}': missing 'prompt' " - f"or 'completion' rate key", - kind="cost", - reason="missing_pricing_entry", - ) - try: - _llm_prompt_rate = float(_llm_prompt_rate_raw) - _llm_completion_rate = float(_llm_completion_rate_raw) - except (TypeError, ValueError) as _llm_rate_err: - raise ExecutionError( - f"Invalid pricing rate for model '{model}' " - f"under provider '{provider}': {_llm_rate_err}", - kind="cost", - reason="missing_pricing_entry", - ) from _llm_rate_err - _llm_node_cost = ( - prompt_tokens / 1_000_000.0 * _llm_prompt_rate - + completion_tokens / 1_000_000.0 * _llm_completion_rate + # If max_cost_usd is set but pricing is empty, hard-fail at once + # (consistent with PureLangGraph._get_max_cost_usd_limit). + _llm_max_cost = executor.limits.get("max_cost_usd") + if _llm_max_cost is not None and not executor.pricing: + raise ExecutionError( + "max_cost_usd limit specified but pricing table is empty. " + "Provide pricing data or remove max_cost_usd.", + kind="cost", + reason="missing_pricing_entry", + ) + + if executor.pricing: + _llm_node_cost = _compute_cost( + executor.pricing, + provider, + model, + prompt_tokens, + completion_tokens, ) - _llm_max_cost = executor.limits.get("max_cost_usd") if _llm_max_cost is not None: if isinstance(_llm_max_cost, bool): raise ExecutionError(