hurui200320
17d99abce7
CI / quality (pull_request) Successful in 36s
CI / lint (pull_request) Successful in 39s
CI / security (pull_request) Successful in 58s
CI / build (pull_request) Successful in 58s
CI / typecheck (pull_request) Successful in 1m2s
CI / integration_tests (pull_request) Successful in 1m14s
CI / unit_tests (pull_request) Successful in 3m35s
CI / coverage (pull_request) Successful in 3m34s
CI / status-check (pull_request) Successful in 3s
CI / lint (push) Successful in 47s
CI / typecheck (push) Successful in 46s
CI / build (push) Successful in 53s
CI / quality (push) Successful in 1m4s
CI / security (push) Successful in 1m7s
CI / integration_tests (push) Successful in 1m13s
CI / unit_tests (push) Successful in 3m43s
CI / coverage (push) Successful in 3m42s
CI / status-check (push) Successful in 3s
AC1: ExecutionError gains kind (categorical: depth, model_calls, tool_calls,
timeout, cost) and reason (sub-code: budget_exhausted or
missing_pricing_entry) fields, both defaulting to empty string. All existing
raise ExecutionError(msg) call sites are backward-compatible.
AC2: PureLangGraph.__init__ accepts limits: dict[str, Any] and
pricing: dict[str, Any] (both default to {}). When limits['max_depth'] is
supplied, a depth breach raises ExecutionError(kind='depth') instead of
silently returning the current message (the old behaviour was a silent
data-loss bug). When no max_depth limit is supplied, the legacy heuristic
(max(2000, len(nodes)*50)) is used with a silent cap for backward
compatibility with existing callers that do not pass limits.
AC3: max_model_calls is checked before each NodeType.AGENT execution;
counter increments after the check so the limit is enforced before the
(limit+1)-th invocation. Breach raises ExecutionError(kind='model_calls').
AC4: max_tool_calls is checked before each NodeType.TOOL execution; same
pattern as model_calls. Breach raises ExecutionError(kind='tool_calls').
AC5: execute() wraps _execute_from_node() in asyncio.wait_for() when
limits['timeout_ms'] is set. asyncio.TimeoutError is caught and re-raised
as ExecutionError(kind='timeout') so the router can map it to HTTP 429.
AC6+AC7: After each LLM node, cost is computed from the pricing table
(rates are USD per million tokens, matching ADR-2029 example values such
as gpt-4.1-mini prompt=/bin/zsh.15/1M). Cost breach raises
ExecutionError(kind='cost', reason='budget_exhausted'). A missing
provider or model entry raises ExecutionError(kind='cost',
reason='missing_pricing_entry') — proceeding with assumed zero cost is
forbidden per ADR-2029. An empty pricing dict ({}) disables all cost
checking (cost tracking was not requested).
AC8: ExecutionError re-exported from cleveractors/__init__.py and
added to __all__.
Wire: runtime_dispatch._execute_graph() now passes executor.limits and
executor.pricing to PureLangGraph so the limits reach the enforcement
layer.
ExecutionError is also re-raised (not swallowed) in _execute_from_node()'s
broad except block so limit errors always propagate to the caller.
New BDD tests (features/execution_limits.feature, 19 scenarios):
- ExecutionError defaults kind and reason to empty strings
- ExecutionError accepts kind and reason keyword arguments
- ExecutionError exported from cleveractors top-level package
- Graph exceeding max_depth raises ExecutionError(kind='depth')
- Two AGENT nodes with max_model_calls=1 raises on 2nd call
- Two TOOL nodes with max_tool_calls=1 raises on 2nd call
- Slow node with timeout_ms=50 raises ExecutionError(kind='timeout')
- LLM node exceeding max_cost_usd raises budget_exhausted
- Unknown provider/model in pricing raises missing_pricing_entry
- Empty pricing table skips cost calculation entirely
- Executor limits/pricing propagation verified
Quality gates: lint pass, typecheck 0 errors,
unit_tests 2317/2317 pass, integration_tests pass,
coverage 97.17% (threshold 96.5%).
ISSUES CLOSED: #15
1340 lines
50 KiB
Python
1340 lines
50 KiB
Python
"""Step definitions for Execution Limits Enforcement BDD tests (elim suffix).
|
|
|
|
Covers ADR-2029 enforcement in PureLangGraph:
|
|
- ExecutionError structured kind/reason fields (issue #15 AC1)
|
|
- max_depth limit (AC2)
|
|
- max_model_calls limit (AC3)
|
|
- max_tool_calls limit (AC4)
|
|
- timeout_ms limit (AC5)
|
|
- max_cost_usd limit — budget_exhausted and missing_pricing_entry (AC6-AC7)
|
|
- ExecutionError export from cleveractors.__all__ (AC8)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from behave import given, then, when
|
|
from features.steps._coverage_utils import _run
|
|
|
|
import cleveractors.langgraph.pure_graph as _pg_module
|
|
from cleveractors.core.exceptions import ExecutionError
|
|
from cleveractors.langgraph.nodes import Edge, NodeConfig, NodeType
|
|
from cleveractors.langgraph.pure_graph import PureGraphConfig, PureLangGraph
|
|
from cleveractors.runtime import create_executor
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Background
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given("the execution limits test context is initialised (elim)")
|
|
def step_elim_init(context: Any) -> None:
|
|
# n-2 fix: removed dead context.elim_no_error variable (never read).
|
|
# n-3 fix: use bare assignments to match project convention (other step
|
|
# files do not use inline type annotations on context attributes).
|
|
context.elim_error = None
|
|
context.elim_result = None
|
|
context.elim_graph = None
|
|
context.elim_limits = {}
|
|
context.elim_pricing = {}
|
|
context.elim_import_ok = False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ExecutionError structured fields
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@when("an ExecutionError is raised with only a message (elim)")
|
|
def step_elim_raise_msg_only(context: Any) -> None:
|
|
try:
|
|
raise ExecutionError("test message")
|
|
except ExecutionError as exc:
|
|
context.elim_error = exc
|
|
|
|
|
|
@when('an ExecutionError is raised with kind "{kind}" and reason "{reason}" (elim)')
|
|
def step_elim_raise_kind_reason(context: Any, kind: str, reason: str) -> None:
|
|
try:
|
|
raise ExecutionError("test message", kind=kind, reason=reason)
|
|
except ExecutionError as exc:
|
|
context.elim_error = exc
|
|
|
|
|
|
@when('an ExecutionError is raised with kind "{kind}" and empty reason (elim)')
|
|
def step_elim_raise_kind_empty_reason(context: Any, kind: str) -> None:
|
|
try:
|
|
raise ExecutionError("test message", kind=kind)
|
|
except ExecutionError as exc:
|
|
context.elim_error = exc
|
|
|
|
|
|
@when("ExecutionError is imported from cleveractors (elim)")
|
|
def step_elim_import(context: Any) -> None:
|
|
try:
|
|
from cleveractors import ExecutionError as ImportedEE # noqa: F401
|
|
|
|
# Verify the imported class is the same and has the expected attributes
|
|
dummy = ImportedEE("test")
|
|
context.elim_error = dummy
|
|
context.elim_import_ok = True
|
|
except Exception as exc:
|
|
context.elim_import_ok = False
|
|
context.elim_error = None
|
|
|
|
|
|
@then("the kind field should be empty string (elim)")
|
|
def step_elim_kind_empty(context: Any) -> None:
|
|
assert context.elim_error is not None
|
|
assert context.elim_error.kind == "", (
|
|
f"Expected kind='', got {context.elim_error.kind!r}"
|
|
)
|
|
|
|
|
|
@then("the reason field should be empty string (elim)")
|
|
def step_elim_reason_empty(context: Any) -> None:
|
|
assert context.elim_error is not None
|
|
assert context.elim_error.reason == "", (
|
|
f"Expected reason='', got {context.elim_error.reason!r}"
|
|
)
|
|
|
|
|
|
@then('the kind field should be "depth" (elim)')
|
|
def step_elim_kind_depth(context: Any) -> None:
|
|
assert context.elim_error is not None
|
|
assert context.elim_error.kind == "depth", (
|
|
f"Expected kind='depth', got {context.elim_error.kind!r}"
|
|
)
|
|
|
|
|
|
@then('the kind field should be "cost" (elim)')
|
|
def step_elim_kind_cost(context: Any) -> None:
|
|
assert context.elim_error is not None
|
|
assert context.elim_error.kind == "cost", (
|
|
f"Expected kind='cost', got {context.elim_error.kind!r}"
|
|
)
|
|
|
|
|
|
@then('the reason field should be "budget_exhausted" (elim)')
|
|
def step_elim_reason_budget_exhausted(context: Any) -> None:
|
|
assert context.elim_error is not None
|
|
assert context.elim_error.reason == "budget_exhausted", (
|
|
f"Expected reason='budget_exhausted', got {context.elim_error.reason!r}"
|
|
)
|
|
|
|
|
|
@then('the reason field should be "missing_pricing_entry" (elim)')
|
|
def step_elim_reason_missing_pricing(context: Any) -> None:
|
|
assert context.elim_error is not None
|
|
assert context.elim_error.reason == "missing_pricing_entry", (
|
|
f"Expected reason='missing_pricing_entry', got {context.elim_error.reason!r}"
|
|
)
|
|
|
|
|
|
@then("the import should succeed and expose kind and reason attributes (elim)")
|
|
def step_elim_import_ok(context: Any) -> None:
|
|
assert context.elim_import_ok, "ExecutionError import from cleveractors failed"
|
|
assert context.elim_error is not None
|
|
assert hasattr(context.elim_error, "kind"), "ExecutionError has no .kind attribute"
|
|
assert hasattr(context.elim_error, "reason"), (
|
|
"ExecutionError has no .reason attribute"
|
|
)
|
|
assert context.elim_error.kind == "", (
|
|
f"Expected kind='', got {context.elim_error.kind!r}"
|
|
)
|
|
assert context.elim_error.reason == "", (
|
|
f"Expected reason='', got {context.elim_error.reason!r}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Depth limit helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_function_node(name: str) -> MagicMock:
|
|
"""Create a mock FUNCTION node that returns a simple dict result."""
|
|
mock = MagicMock()
|
|
mock.config = MagicMock()
|
|
mock.config.type = NodeType.FUNCTION
|
|
mock.config.name = name
|
|
|
|
async def _execute(state: Any) -> dict[str, Any]:
|
|
return {"content": f"output_from_{name}"}
|
|
|
|
mock.execute = AsyncMock(side_effect=_execute)
|
|
return mock
|
|
|
|
|
|
def _make_agent_node(name: str) -> MagicMock:
|
|
"""Create a mock AGENT node that returns a response dict (no token data)."""
|
|
mock = MagicMock()
|
|
mock.config = MagicMock()
|
|
mock.config.type = NodeType.AGENT
|
|
mock.config.name = name
|
|
|
|
async def _execute(state: Any) -> dict[str, Any]:
|
|
return {
|
|
"messages": [{"role": "assistant", "content": f"agent_response_{name}"}]
|
|
}
|
|
|
|
mock.execute = AsyncMock(side_effect=_execute)
|
|
return mock
|
|
|
|
|
|
def _make_tool_node(name: str) -> MagicMock:
|
|
"""Create a mock TOOL node that returns a simple result."""
|
|
mock = MagicMock()
|
|
mock.config = MagicMock()
|
|
mock.config.type = NodeType.TOOL
|
|
mock.config.name = name
|
|
|
|
async def _execute(state: Any) -> dict[str, Any]:
|
|
return {"content": f"tool_output_{name}"}
|
|
|
|
mock.execute = AsyncMock(side_effect=_execute)
|
|
return mock
|
|
|
|
|
|
def _make_llm_node_with_tokens(
|
|
name: str,
|
|
provider: str,
|
|
model: str,
|
|
prompt_tokens: int,
|
|
completion_tokens: int,
|
|
) -> MagicMock:
|
|
"""Create a mock AGENT node whose execute() includes _node_token_usage.
|
|
|
|
The ``_node_token_usage`` dict mirrors what ``Node._execute_agent()``
|
|
would produce for real LLM calls so that ``PureLangGraph._execute_from_node()``
|
|
picks it up for cost accumulation (issue #15).
|
|
"""
|
|
mock = MagicMock()
|
|
mock.config = MagicMock()
|
|
mock.config.type = NodeType.AGENT
|
|
mock.config.name = name
|
|
|
|
async def _execute(state: Any) -> dict[str, Any]:
|
|
return {
|
|
"messages": [{"role": "assistant", "content": "llm_response"}],
|
|
"_node_token_usage": {
|
|
"node_id": name,
|
|
"provider": provider,
|
|
"model": model,
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
},
|
|
}
|
|
|
|
mock.execute = AsyncMock(side_effect=_execute)
|
|
return mock
|
|
|
|
|
|
def _run_graph(
|
|
graph: PureLangGraph, message: str = "hello"
|
|
) -> tuple[Any, ExecutionError | None]:
|
|
"""Execute a PureLangGraph synchronously and return (result, error)."""
|
|
error: ExecutionError | None = None
|
|
result: Any = None
|
|
try:
|
|
result = _run(graph.execute(message))
|
|
except ExecutionError as exc:
|
|
error = exc
|
|
return result, error
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Depth limit scenarios
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given("a linear 3-node graph with max_depth {max_depth:d} (elim)")
|
|
def step_elim_depth_3node(context: Any, max_depth: int) -> None:
|
|
"""Create a linear 3-node graph: start → a → b → c → end.
|
|
|
|
With max_depth=1 and entry_point="start":
|
|
depth 0: start (routes to a)
|
|
depth 1: a (routes to b)
|
|
depth 2: b — exceeds max_depth=1 → raises ExecutionError(kind="depth")
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="depth_test",
|
|
nodes={
|
|
"a": NodeConfig(name="a", type=NodeType.FUNCTION),
|
|
"b": NodeConfig(name="b", type=NodeType.FUNCTION),
|
|
"c": NodeConfig(name="c", type=NodeType.FUNCTION),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="a"),
|
|
Edge(source="a", target="b"),
|
|
Edge(source="b", target="c"),
|
|
Edge(source="c", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_depth": max_depth})
|
|
# Inject mock nodes for controlled execution
|
|
graph.nodes["a"] = _make_function_node("a")
|
|
graph.nodes["b"] = _make_function_node("b")
|
|
graph.nodes["c"] = _make_function_node("c")
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a linear 2-node graph with max_depth {max_depth:d} (elim)")
|
|
def step_elim_depth_2node(context: Any, max_depth: int) -> None:
|
|
"""Create a linear 2-node graph: start → a → end.
|
|
|
|
With max_depth >= 1, this should complete without depth error.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="depth_ok_test",
|
|
nodes={
|
|
"a": NodeConfig(name="a", type=NodeType.FUNCTION),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="a"),
|
|
Edge(source="a", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_depth": max_depth})
|
|
graph.nodes["a"] = _make_function_node("a")
|
|
context.elim_graph = graph
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model-call limit scenarios
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given("a graph with 2 AGENT nodes and max_model_calls {limit:d} (elim)")
|
|
def step_elim_model_calls_2(context: Any, limit: int) -> None:
|
|
"""Graph: start → agent_a → agent_b → end.
|
|
|
|
With max_model_calls=1: agent_a runs (count becomes 1), agent_b is blocked.
|
|
With max_model_calls=2: both agents run successfully.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="model_calls_test",
|
|
nodes={
|
|
"agent_a": NodeConfig(name="agent_a", type=NodeType.AGENT),
|
|
"agent_b": NodeConfig(name="agent_b", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="agent_a"),
|
|
Edge(source="agent_a", target="agent_b"),
|
|
Edge(source="agent_b", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_model_calls": limit})
|
|
graph.nodes["agent_a"] = _make_agent_node("agent_a")
|
|
graph.nodes["agent_b"] = _make_agent_node("agent_b")
|
|
context.elim_graph = graph
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tool-call limit scenarios
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given("a graph with 2 TOOL nodes and max_tool_calls {limit:d} (elim)")
|
|
def step_elim_tool_calls_2(context: Any, limit: int) -> None:
|
|
"""Graph: start → tool_a → tool_b → end.
|
|
|
|
With max_tool_calls=1: tool_a runs (count becomes 1), tool_b is blocked.
|
|
With max_tool_calls=2: both tools run successfully.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="tool_calls_test",
|
|
nodes={
|
|
"tool_a": NodeConfig(name="tool_a", type=NodeType.TOOL),
|
|
"tool_b": NodeConfig(name="tool_b", type=NodeType.TOOL),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="tool_a"),
|
|
Edge(source="tool_a", target="tool_b"),
|
|
Edge(source="tool_b", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_tool_calls": limit})
|
|
graph.nodes["tool_a"] = _make_tool_node("tool_a")
|
|
graph.nodes["tool_b"] = _make_tool_node("tool_b")
|
|
context.elim_graph = graph
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Timeout scenarios
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given("a graph with a slow node and timeout_ms {timeout_ms:d} (elim)")
|
|
def step_elim_timeout_slow(context: Any, timeout_ms: int) -> None:
|
|
"""Graph with a FUNCTION node that sleeps longer than timeout_ms."""
|
|
config = PureGraphConfig(
|
|
name="timeout_test",
|
|
nodes={
|
|
"slow": NodeConfig(name="slow", type=NodeType.FUNCTION),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="slow"),
|
|
Edge(source="slow", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"timeout_ms": timeout_ms})
|
|
|
|
# Sleep for 2 seconds, which exceeds the 50ms timeout in the scenario
|
|
async def _slow_execute(state: Any) -> dict[str, Any]:
|
|
await asyncio.sleep(2.0)
|
|
return {"content": "slow_done"}
|
|
|
|
mock_slow = MagicMock()
|
|
mock_slow.config = MagicMock()
|
|
mock_slow.config.type = NodeType.FUNCTION
|
|
mock_slow.config.name = "slow"
|
|
mock_slow.execute = AsyncMock(side_effect=_slow_execute)
|
|
graph.nodes["slow"] = mock_slow
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a graph with a fast node and timeout_ms {timeout_ms:d} (elim)")
|
|
def step_elim_timeout_fast(context: Any, timeout_ms: int) -> None:
|
|
"""Graph with a FUNCTION node that completes quickly (well within timeout)."""
|
|
config = PureGraphConfig(
|
|
name="fast_timeout_test",
|
|
nodes={
|
|
"fast": NodeConfig(name="fast", type=NodeType.FUNCTION),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="fast"),
|
|
Edge(source="fast", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"timeout_ms": timeout_ms})
|
|
graph.nodes["fast"] = _make_function_node("fast")
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a graph with a fast node and bool True timeout_ms (elim)")
|
|
def step_elim_timeout_bool_true(context: Any) -> None:
|
|
"""Graph with bool True as timeout_ms.
|
|
|
|
float(True)==1.0 would silently set a 1ms timeout, effectively making
|
|
every request time out. The bool guard must reject it with
|
|
ExecutionError(kind="timeout") before asyncio.wait_for() is called.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="bool_timeout_test",
|
|
nodes={
|
|
"fast": NodeConfig(name="fast", type=NodeType.FUNCTION),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="fast"),
|
|
Edge(source="fast", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"timeout_ms": True})
|
|
graph.nodes["fast"] = _make_function_node("fast")
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a graph with a fast node and non-numeric max_depth (elim)")
|
|
def step_elim_max_depth_non_numeric(context: Any) -> None:
|
|
"""Graph with a non-numeric max_depth limit.
|
|
|
|
A non-numeric max_depth (e.g. a string) should raise
|
|
ExecutionError(kind="depth") immediately rather than a bare ValueError.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="non_numeric_depth_test",
|
|
nodes={
|
|
"fast": NodeConfig(name="fast", type=NodeType.FUNCTION),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="fast"),
|
|
Edge(source="fast", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_depth": "not-a-number"})
|
|
graph.nodes["fast"] = _make_function_node("fast")
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a graph with a fast node and bool False max_depth (elim)")
|
|
def step_elim_max_depth_bool_false(context: Any) -> None:
|
|
"""Graph with bool False as max_depth.
|
|
|
|
int(False)==0 would silently set a depth limit of 0, causing every
|
|
execution to fail at the first node. The bool guard must reject it with
|
|
ExecutionError(kind="depth") with a clear diagnostic.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="bool_depth_test",
|
|
nodes={
|
|
"fast": NodeConfig(name="fast", type=NodeType.FUNCTION),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="fast"),
|
|
Edge(source="fast", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_depth": False})
|
|
graph.nodes["fast"] = _make_function_node("fast")
|
|
context.elim_graph = graph
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cost limit scenarios — budget_exhausted
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given(
|
|
"a graph with an LLM node using {prompt_tokens:d} prompt tokens "
|
|
"and max_cost_usd {max_cost:f} (elim)"
|
|
)
|
|
def step_elim_cost_llm_tokens(
|
|
context: Any, prompt_tokens: int, max_cost: float
|
|
) -> None:
|
|
"""Graph with a single LLM node that reports a specific token count.
|
|
|
|
The graph uses default provider/model (openai / gpt-4.1-mini) with the
|
|
token count stored for use when pricing is applied in the When step.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="cost_test",
|
|
nodes={
|
|
"llm": NodeConfig(name="llm", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="llm"),
|
|
Edge(source="llm", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
# Store for When step to wire into graph with correct pricing
|
|
context.elim_cost_prompt_tokens = prompt_tokens
|
|
context.elim_cost_max_cost = max_cost
|
|
context.elim_cost_provider = "openai"
|
|
context.elim_cost_model = "gpt-4.1-mini"
|
|
context.elim_config = config
|
|
|
|
|
|
@when("the graph is executed with an empty pricing table (elim)")
|
|
def step_elim_execute_with_empty_pricing(context: Any) -> None:
|
|
"""Execute with max_cost_usd but empty pricing — cost check should be skipped."""
|
|
config = context.elim_config
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
provider = context.elim_cost_provider
|
|
model = context.elim_cost_model
|
|
max_cost = context.elim_cost_max_cost
|
|
|
|
limits = {"max_cost_usd": max_cost}
|
|
pricing: dict[str, Any] = {} # empty — skips cost calculation
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
graph.nodes["llm"] = _make_llm_node_with_tokens(
|
|
"llm", provider, model, prompt_tokens, 0
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cost limit scenarios — missing_pricing_entry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given('a graph with an LLM node for provider "{provider}" model "{model}" (elim)')
|
|
def step_elim_llm_node_provider_model(context: Any, provider: str, model: str) -> None:
|
|
"""Graph with a single LLM node for a specific provider/model combination."""
|
|
config = PureGraphConfig(
|
|
name="missing_pricing_test",
|
|
nodes={
|
|
"llm": NodeConfig(name="llm", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="llm"),
|
|
Edge(source="llm", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
context.elim_config = config
|
|
context.elim_cost_provider = provider
|
|
context.elim_cost_model = model
|
|
context.elim_cost_prompt_tokens = 1000
|
|
context.elim_cost_max_cost = 10.0 # high limit — shouldn't trigger budget breach
|
|
|
|
|
|
@when('the graph is executed with pricing only for "anthropic" provider (elim)')
|
|
def step_elim_execute_anthropic_only(context: Any) -> None:
|
|
"""Execute with pricing that only covers anthropic (not openai)."""
|
|
config = context.elim_config
|
|
provider = context.elim_cost_provider
|
|
model = context.elim_cost_model
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
|
|
limits = {"max_cost_usd": context.elim_cost_max_cost}
|
|
pricing = {"anthropic": {"claude-3-5-haiku": {"prompt": 0.80, "completion": 4.00}}}
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
graph.nodes["llm"] = _make_llm_node_with_tokens(
|
|
"llm", provider, model, prompt_tokens, 0
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
@when('the graph is executed with pricing for "openai" but not "{model}" (elim)')
|
|
def step_elim_execute_openai_no_model(context: Any, model: str) -> None:
|
|
"""Execute with pricing that has openai but not the specific model."""
|
|
config = context.elim_config
|
|
provider = context.elim_cost_provider
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
|
|
limits = {"max_cost_usd": context.elim_cost_max_cost}
|
|
# Provide openai pricing for a different model, not "unknown-model"
|
|
pricing = {"openai": {"gpt-4.1-mini": {"prompt": 0.15, "completion": 0.60}}}
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
graph.nodes["llm"] = _make_llm_node_with_tokens(
|
|
"llm", provider, model, prompt_tokens, 0
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
@when("the graph is executed with model pricing entry missing the prompt rate (elim)")
|
|
def step_elim_execute_missing_prompt_rate(context: Any) -> None:
|
|
"""Execute with a model pricing entry that lacks the 'prompt' rate key.
|
|
|
|
ADR-2029 forbids proceeding with a missing rate — the library must raise
|
|
``missing_pricing_entry`` rather than silently using a zero-cost fallback.
|
|
"""
|
|
config = context.elim_config
|
|
provider = context.elim_cost_provider
|
|
model = context.elim_cost_model
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
|
|
limits = {"max_cost_usd": context.elim_cost_max_cost}
|
|
# Model entry exists but is missing the 'prompt' rate key
|
|
pricing = {provider: {model: {"completion": 0.60}}}
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
graph.nodes["llm"] = _make_llm_node_with_tokens(
|
|
"llm", provider, model, prompt_tokens, 0
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
@when("the graph is executed with a non-numeric prompt rate (elim)")
|
|
def step_elim_execute_non_numeric_rate(context: Any) -> None:
|
|
"""Execute with a model pricing entry whose 'prompt' rate is a non-numeric value.
|
|
|
|
ADR-2029 forbids proceeding when the rate cannot be interpreted as a number.
|
|
The library must raise ``missing_pricing_entry`` rather than swallowing the
|
|
TypeError and continuing with zero cost.
|
|
"""
|
|
config = context.elim_config
|
|
provider = context.elim_cost_provider
|
|
model = context.elim_cost_model
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
|
|
limits = {"max_cost_usd": context.elim_cost_max_cost}
|
|
# Non-numeric prompt rate — float("not-a-number") would raise ValueError
|
|
pricing = {provider: {model: {"prompt": "not-a-number", "completion": 0.60}}}
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
graph.nodes["llm"] = _make_llm_node_with_tokens(
|
|
"llm", provider, model, prompt_tokens, 0
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Malformed limit value validation scenarios
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given("a graph with an LLM node for non-numeric max_cost_usd validation (elim)")
|
|
def step_elim_non_numeric_max_cost_given(context: Any) -> None:
|
|
"""Set up a graph with an LLM node for non-numeric max_cost_usd validation.
|
|
|
|
The actual (invalid) max_cost_usd value is supplied in the When step so
|
|
the Given step text does not mislead readers about what is being tested.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="non_numeric_cost_test",
|
|
nodes={
|
|
"llm": NodeConfig(name="llm", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="llm"),
|
|
Edge(source="llm", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
context.elim_config = config
|
|
context.elim_cost_provider = "openai"
|
|
context.elim_cost_model = "gpt-4.1-mini"
|
|
context.elim_cost_prompt_tokens = 1000
|
|
|
|
|
|
@given("a graph with an LLM node for bool max_cost_usd validation (elim)")
|
|
def step_elim_bool_max_cost_given(context: Any) -> None:
|
|
"""Set up a graph with an LLM node for bool max_cost_usd validation.
|
|
|
|
The actual (invalid) bool max_cost_usd value is supplied in the When step.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="bool_cost_test",
|
|
nodes={
|
|
"llm": NodeConfig(name="llm", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="llm"),
|
|
Edge(source="llm", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
context.elim_config = config
|
|
context.elim_cost_provider = "openai"
|
|
context.elim_cost_model = "gpt-4.1-mini"
|
|
context.elim_cost_prompt_tokens = 1000
|
|
|
|
|
|
@when("the graph is executed with pricing and non-numeric max_cost_usd (elim)")
|
|
def step_elim_execute_with_non_numeric_max_cost(context: Any) -> None:
|
|
"""Execute with a non-numeric max_cost_usd to trigger C-1 validation."""
|
|
config = context.elim_config
|
|
provider = context.elim_cost_provider
|
|
model = context.elim_cost_model
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
|
|
# Pass a non-numeric max_cost_usd — should raise ExecutionError(kind="cost")
|
|
limits = {"max_cost_usd": "not-a-number"}
|
|
pricing = {
|
|
provider: {
|
|
model: {"prompt": 0.15, "completion": 0.60},
|
|
}
|
|
}
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
graph.nodes["llm"] = _make_llm_node_with_tokens(
|
|
"llm", provider, model, prompt_tokens, 0
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
@when("the graph is executed with pricing and bool False max_cost_usd (elim)")
|
|
def step_elim_execute_with_bool_max_cost(context: Any) -> None:
|
|
"""Execute with a bool False max_cost_usd to trigger bool guard validation."""
|
|
config = context.elim_config
|
|
provider = context.elim_cost_provider
|
|
model = context.elim_cost_model
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
|
|
# Pass bool False as max_cost_usd — float(False)==0.0 would silently set
|
|
# a $0 budget; the bool guard must reject it with ExecutionError(kind="cost")
|
|
limits = {"max_cost_usd": False}
|
|
pricing = {
|
|
provider: {
|
|
model: {"prompt": 0.15, "completion": 0.60},
|
|
}
|
|
}
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
graph.nodes["llm"] = _make_llm_node_with_tokens(
|
|
"llm", provider, model, prompt_tokens, 0
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
@given("a graph with 1 AGENT node and non-numeric max_model_calls (elim)")
|
|
def step_elim_agent_non_numeric_max_model(context: Any) -> None:
|
|
"""Graph with a single AGENT node and a non-numeric max_model_calls limit."""
|
|
config = PureGraphConfig(
|
|
name="invalid_model_calls_test",
|
|
nodes={
|
|
"agent_a": NodeConfig(name="agent_a", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="agent_a"),
|
|
Edge(source="agent_a", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_model_calls": "not-a-number"})
|
|
graph.nodes["agent_a"] = _make_agent_node("agent_a")
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a graph with 1 AGENT node and bool False max_model_calls (elim)")
|
|
def step_elim_agent_bool_max_model(context: Any) -> None:
|
|
"""Graph with a single AGENT node and a bool False max_model_calls limit.
|
|
|
|
M-4 fix: int(False)==0 would silently set the limit to 0, causing every
|
|
model call to fail without a clear error. The bool guard raises
|
|
ExecutionError(kind="model_calls") immediately.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="bool_model_calls_test",
|
|
nodes={
|
|
"agent_a": NodeConfig(name="agent_a", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="agent_a"),
|
|
Edge(source="agent_a", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_model_calls": False})
|
|
graph.nodes["agent_a"] = _make_agent_node("agent_a")
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a graph with 1 TOOL node and non-numeric max_tool_calls (elim)")
|
|
def step_elim_tool_non_numeric_max_tool(context: Any) -> None:
|
|
"""Graph with a single TOOL node and a non-numeric max_tool_calls limit."""
|
|
config = PureGraphConfig(
|
|
name="invalid_tool_calls_test",
|
|
nodes={
|
|
"tool_a": NodeConfig(name="tool_a", type=NodeType.TOOL),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="tool_a"),
|
|
Edge(source="tool_a", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_tool_calls": "not-a-number"})
|
|
graph.nodes["tool_a"] = _make_tool_node("tool_a")
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a graph with 1 TOOL node and bool False max_tool_calls (elim)")
|
|
def step_elim_tool_bool_max_tool(context: Any) -> None:
|
|
"""Graph with a single TOOL node and a bool False max_tool_calls limit."""
|
|
config = PureGraphConfig(
|
|
name="bool_tool_calls_test",
|
|
nodes={
|
|
"tool_a": NodeConfig(name="tool_a", type=NodeType.TOOL),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="tool_a"),
|
|
Edge(source="tool_a", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={"max_tool_calls": False})
|
|
graph.nodes["tool_a"] = _make_tool_node("tool_a")
|
|
context.elim_graph = graph
|
|
|
|
|
|
@given("a parallel graph where one branch raises an ExecutionError (elim)")
|
|
def step_elim_parallel_execution_error(context: Any) -> None:
|
|
"""Graph with parallel execution where one branch raises ExecutionError.
|
|
|
|
M-1 fix: when one parallel branch raises ExecutionError, sibling tasks
|
|
must be cancelled so they do not continue spending budget.
|
|
|
|
The graph uses a "splitter" entry node that fans out to branch_a and
|
|
branch_b in parallel. branch_a yields control (via asyncio.sleep(0))
|
|
and then raises ExecutionError, ensuring branch_b has started executing
|
|
(sleeping) before the exception propagates — this exercises the t.cancel()
|
|
path in the parallel execution block.
|
|
|
|
A cancellation flag is tracked via a try/except asyncio.CancelledError block
|
|
inside _slow_execute so the Then step can assert that branch_b was actually
|
|
cancelled (not just that the exception propagated).
|
|
|
|
Note: the "start" node handler only routes to next_nodes[0], so parallel
|
|
execution must be triggered from a real non-start node. The "splitter"
|
|
node is a FUNCTION node that returns immediately and has two outgoing edges
|
|
to branch_a and branch_b.
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="parallel_cancel_test",
|
|
nodes={
|
|
"splitter": NodeConfig(name="splitter", type=NodeType.FUNCTION),
|
|
"branch_a": NodeConfig(name="branch_a", type=NodeType.FUNCTION),
|
|
"branch_b": NodeConfig(name="branch_b", type=NodeType.FUNCTION),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="splitter"),
|
|
Edge(source="splitter", target="branch_a"),
|
|
Edge(source="splitter", target="branch_b"),
|
|
Edge(source="branch_a", target="end"),
|
|
Edge(source="branch_b", target="end"),
|
|
],
|
|
entry_point="start",
|
|
parallel_execution=True,
|
|
)
|
|
graph = PureLangGraph(config, limits={})
|
|
|
|
# splitter: returns immediately to fan out to branch_a and branch_b
|
|
splitter_mock = MagicMock()
|
|
splitter_mock.config = MagicMock()
|
|
splitter_mock.config.type = NodeType.FUNCTION
|
|
splitter_mock.config.name = "splitter"
|
|
|
|
async def _splitter_execute(state: Any) -> dict[str, Any]:
|
|
return {"content": "split"}
|
|
|
|
splitter_mock.execute = AsyncMock(side_effect=_splitter_execute)
|
|
graph.nodes["splitter"] = splitter_mock
|
|
|
|
# branch_a: yields once (so branch_b can start), then raises ExecutionError
|
|
raise_mock = MagicMock()
|
|
raise_mock.config = MagicMock()
|
|
raise_mock.config.type = NodeType.FUNCTION
|
|
raise_mock.config.name = "branch_a"
|
|
|
|
async def _raising_execute(state: Any) -> dict[str, Any]:
|
|
# Yield to the event loop so branch_b can start its sleep
|
|
await asyncio.sleep(0)
|
|
raise ExecutionError("branch_a limit exceeded", kind="model_calls")
|
|
|
|
raise_mock.execute = AsyncMock(side_effect=_raising_execute)
|
|
graph.nodes["branch_a"] = raise_mock
|
|
|
|
# branch_b: sleeps long enough that it is still running when branch_a raises,
|
|
# ensuring the t.cancel() path is exercised.
|
|
# A cancellation flag is set inside the except CancelledError handler so the
|
|
# Then step can assert that branch_b was actually cancelled.
|
|
branch_b_cancelled = [False]
|
|
slow_mock = MagicMock()
|
|
slow_mock.config = MagicMock()
|
|
slow_mock.config.type = NodeType.FUNCTION
|
|
slow_mock.config.name = "branch_b"
|
|
|
|
async def _slow_execute(state: Any) -> dict[str, Any]:
|
|
try:
|
|
await asyncio.sleep(10.0)
|
|
return {"content": "slow_done"}
|
|
except asyncio.CancelledError:
|
|
branch_b_cancelled[0] = True
|
|
raise
|
|
|
|
slow_mock.execute = AsyncMock(side_effect=_slow_execute)
|
|
graph.nodes["branch_b"] = slow_mock
|
|
context.elim_graph = graph
|
|
context.elim_branch_b_cancelled = branch_b_cancelled
|
|
|
|
|
|
def _make_non_dict_token_usage_node(name: str) -> MagicMock:
|
|
"""Create a mock AGENT node that returns a non-dict _node_token_usage value.
|
|
|
|
Used to verify that m-1 warning is logged and token accounting is skipped
|
|
without raising an error.
|
|
"""
|
|
mock = MagicMock()
|
|
mock.config = MagicMock()
|
|
mock.config.type = NodeType.AGENT
|
|
mock.config.name = name
|
|
|
|
async def _execute(state: Any) -> dict[str, Any]:
|
|
return {
|
|
"messages": [{"role": "assistant", "content": "response"}],
|
|
"_node_token_usage": "not-a-dict", # malformed: should be a dict
|
|
}
|
|
|
|
mock.execute = AsyncMock(side_effect=_execute)
|
|
return mock
|
|
|
|
|
|
@given("a graph with an LLM node that returns non-dict token usage (elim)")
|
|
def step_elim_non_dict_token_usage(context: Any) -> None:
|
|
"""Graph with a node that returns a non-dict _node_token_usage value."""
|
|
config = PureGraphConfig(
|
|
name="non_dict_token_test",
|
|
nodes={
|
|
"llm": NodeConfig(name="llm", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="llm"),
|
|
Edge(source="llm", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={}, pricing={})
|
|
graph.nodes["llm"] = _make_non_dict_token_usage_node("llm")
|
|
context.elim_graph = graph
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cost accumulation across multiple LLM nodes (m-4 fix)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given(
|
|
"a graph with 2 LLM nodes each using {prompt_tokens:d} prompt tokens "
|
|
"and max_cost_usd {max_cost:f} (elim)"
|
|
)
|
|
def step_elim_cost_2llm_nodes(
|
|
context: Any, prompt_tokens: int, max_cost: float
|
|
) -> None:
|
|
"""Graph with two sequential LLM nodes that each report prompt_tokens.
|
|
|
|
With prompt_rate=0.15/1M and 1M tokens per node:
|
|
node_a cost = 0.15 USD (accumulated = 0.15)
|
|
node_b cost = 0.15 USD (accumulated = 0.30 > 0.25) → budget_exhausted
|
|
"""
|
|
config = PureGraphConfig(
|
|
name="cost_accumulation_test",
|
|
nodes={
|
|
"llm_a": NodeConfig(name="llm_a", type=NodeType.AGENT),
|
|
"llm_b": NodeConfig(name="llm_b", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="llm_a"),
|
|
Edge(source="llm_a", target="llm_b"),
|
|
Edge(source="llm_b", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
context.elim_cost_prompt_tokens = prompt_tokens
|
|
context.elim_cost_max_cost = max_cost
|
|
context.elim_cost_provider = "openai"
|
|
context.elim_cost_model = "gpt-4.1-mini"
|
|
context.elim_config = config
|
|
context.elim_cost_num_nodes = 2
|
|
|
|
|
|
@when(
|
|
'the graph is executed with pricing for "openai"/"gpt-4.1-mini" '
|
|
"at {prompt_rate:f} prompt rate (elim)"
|
|
)
|
|
def step_elim_execute_with_pricing(context: Any, prompt_rate: float) -> None:
|
|
"""Wire graph with limits+pricing and execute it."""
|
|
config = context.elim_config
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
provider = context.elim_cost_provider
|
|
model = context.elim_cost_model
|
|
max_cost = context.elim_cost_max_cost
|
|
num_nodes = getattr(context, "elim_cost_num_nodes", 1)
|
|
|
|
limits = {"max_cost_usd": max_cost}
|
|
pricing = {
|
|
provider: {
|
|
model: {"prompt": prompt_rate, "completion": 0.60},
|
|
}
|
|
}
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
for i in range(num_nodes):
|
|
node_name = "llm" if num_nodes == 1 else f"llm_{'ab'[i]}"
|
|
graph.nodes[node_name] = _make_llm_node_with_tokens(
|
|
node_name, provider, model, prompt_tokens, 0
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Completion tokens cost contribution (m-5 fix)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given(
|
|
"a graph with an LLM node using {prompt_tokens:d} prompt tokens "
|
|
"{completion_tokens:d} completion tokens and max_cost_usd {max_cost:f} (elim)"
|
|
)
|
|
def step_elim_cost_llm_completion_tokens(
|
|
context: Any, prompt_tokens: int, completion_tokens: int, max_cost: float
|
|
) -> None:
|
|
"""Graph with a single LLM node that reports both prompt and completion tokens."""
|
|
config = PureGraphConfig(
|
|
name="completion_cost_test",
|
|
nodes={
|
|
"llm": NodeConfig(name="llm", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="llm"),
|
|
Edge(source="llm", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
context.elim_cost_prompt_tokens = prompt_tokens
|
|
context.elim_cost_completion_tokens = completion_tokens
|
|
context.elim_cost_max_cost = max_cost
|
|
context.elim_cost_provider = "openai"
|
|
context.elim_cost_model = "gpt-4.1-mini"
|
|
context.elim_config = config
|
|
|
|
|
|
@when(
|
|
'the graph is executed with pricing for "openai"/"gpt-4.1-mini" '
|
|
"at {prompt_rate:f} prompt {completion_rate:f} completion rate (elim)"
|
|
)
|
|
def step_elim_execute_with_full_pricing(
|
|
context: Any, prompt_rate: float, completion_rate: float
|
|
) -> None:
|
|
"""Wire graph with limits+pricing (both rates) and execute it."""
|
|
config = context.elim_config
|
|
prompt_tokens = context.elim_cost_prompt_tokens
|
|
completion_tokens = getattr(context, "elim_cost_completion_tokens", 0)
|
|
provider = context.elim_cost_provider
|
|
model = context.elim_cost_model
|
|
max_cost = context.elim_cost_max_cost
|
|
|
|
limits = {"max_cost_usd": max_cost}
|
|
pricing = {
|
|
provider: {
|
|
model: {"prompt": prompt_rate, "completion": completion_rate},
|
|
}
|
|
}
|
|
graph = PureLangGraph(config, limits=limits, pricing=pricing)
|
|
graph.nodes["llm"] = _make_llm_node_with_tokens(
|
|
"llm", provider, model, prompt_tokens, completion_tokens
|
|
)
|
|
context.elim_result, context.elim_error = _run_graph(graph)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# limits/pricing propagation from Executor — storage check
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given(
|
|
"an Executor with graph config and limits max_depth 100 pricing for openai (elim)"
|
|
)
|
|
def step_elim_executor_limits(context: Any) -> None:
|
|
"""Store the Executor limits/pricing for propagation verification."""
|
|
graph_config: dict[str, Any] = {
|
|
"type": "graph",
|
|
"name": "propagation_test",
|
|
"route": {
|
|
"nodes": [],
|
|
"edges": [],
|
|
"entry_node": "start",
|
|
},
|
|
}
|
|
limits = {
|
|
"max_depth": 100,
|
|
"max_model_calls": 10,
|
|
"max_tool_calls": 10,
|
|
"timeout_ms": 60000,
|
|
"max_cost_usd": 1.00,
|
|
}
|
|
pricing = {"openai": {"gpt-4.1-mini": {"prompt": 0.15, "completion": 0.60}}}
|
|
context.elim_executor = create_executor(
|
|
graph_config, credentials=None, limits=limits, pricing=pricing
|
|
)
|
|
context.elim_expected_limits = limits
|
|
context.elim_expected_pricing = pricing
|
|
|
|
|
|
@when("the executor is checked for limit/pricing propagation (elim)")
|
|
def step_elim_check_propagation(context: Any) -> None:
|
|
"""Verify limits and pricing are stored on Executor AND propagated to PureLangGraph.
|
|
|
|
m-3 fix: the original scenario only checked Executor attribute storage.
|
|
This step now also patches PureLangGraph.__init__ to capture the kwargs
|
|
actually passed by _execute_graph(), verifying the full propagation path:
|
|
create_executor() → Executor.limits/pricing → _execute_graph() →
|
|
PureLangGraph(limits=..., pricing=...)
|
|
"""
|
|
executor = context.elim_executor
|
|
captured_kwargs: dict[str, Any] = {}
|
|
|
|
original_init = PureLangGraph.__init__
|
|
|
|
def _capturing_init(self: Any, *args: Any, **kwargs: Any) -> None:
|
|
captured_kwargs.update(kwargs)
|
|
original_init(self, *args, **kwargs)
|
|
|
|
# Patch PureLangGraph.__init__ to capture the kwargs passed by _execute_graph
|
|
with patch.object(PureLangGraph, "__init__", _capturing_init):
|
|
try:
|
|
_run(executor.execute("probe"))
|
|
except Exception: # pylint: disable=broad-exception-caught
|
|
pass # graph may fail; we only care about the kwargs captured
|
|
|
|
context.elim_propagation_ok = (
|
|
executor.limits == context.elim_expected_limits
|
|
and executor.pricing == context.elim_expected_pricing
|
|
and captured_kwargs.get("limits") == context.elim_expected_limits
|
|
and captured_kwargs.get("pricing") == context.elim_expected_pricing
|
|
)
|
|
context.elim_captured_kwargs = captured_kwargs
|
|
|
|
|
|
@then("the PureLangGraph should have the correct limits and pricing (elim)")
|
|
def step_elim_propagation_verified(context: Any) -> None:
|
|
assert context.elim_propagation_ok, (
|
|
f"Propagation check failed. "
|
|
f"Executor limits: {getattr(context, 'elim_executor', None) and context.elim_executor.limits!r}, "
|
|
f"PureLangGraph kwargs: {context.elim_captured_kwargs!r}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Executor dispatch integration — full _execute_graph() path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given("an Executor with a 2-node graph and max_depth {max_depth:d} (elim)")
|
|
def step_elim_executor_dispatch_graph(context: Any, max_depth: int) -> None:
|
|
"""Set up an Executor whose graph config has one node plus the implicit
|
|
start/end nodes so that _execute_graph() creates a real PureLangGraph.
|
|
|
|
With max_depth=0 the depth check fires when the entry node is first
|
|
visited (depth 1 > 0) — before the node actually executes — giving us a
|
|
clean ExecutionError(kind="depth") through the full dispatch path:
|
|
Executor.execute() → _execute_graph() → PureLangGraph.execute() →
|
|
_execute_from_node() → depth check → ExecutionError
|
|
"""
|
|
# A minimal graph actor config: one FUNCTION node, explicit edges from
|
|
# start to the node and from the node to end. No agent is needed because
|
|
# the depth-limit check fires BEFORE node.execute() is called.
|
|
graph_config: dict[str, Any] = {
|
|
"type": "graph",
|
|
"name": "dispatch_depth_test",
|
|
"route": {
|
|
"nodes": [{"id": "node_a"}],
|
|
"edges": [
|
|
{"source": "start", "target": "node_a"},
|
|
{"source": "node_a", "target": "end"},
|
|
],
|
|
"entry_node": "start",
|
|
},
|
|
}
|
|
context.elim_executor_for_dispatch = create_executor(
|
|
graph_config,
|
|
credentials=None,
|
|
limits={"max_depth": max_depth},
|
|
pricing={},
|
|
)
|
|
|
|
|
|
@when("the executor execute method is called with a test message (elim)")
|
|
def step_elim_executor_dispatch_execute(context: Any) -> None:
|
|
"""Call executor.execute() and capture any ExecutionError raised."""
|
|
executor = context.elim_executor_for_dispatch
|
|
try:
|
|
_run(executor.execute("test message"))
|
|
context.elim_error = None
|
|
except ExecutionError as exc:
|
|
context.elim_error = exc
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Generic When/Then steps
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@when("the graph is executed with a test message (elim)")
|
|
def step_elim_execute(context: Any) -> None:
|
|
"""Execute whatever graph is stored in context.elim_graph."""
|
|
assert context.elim_graph is not None, "No graph set up in context.elim_graph"
|
|
context.elim_result, context.elim_error = _run_graph(
|
|
context.elim_graph, "test message"
|
|
)
|
|
|
|
|
|
@then('an ExecutionError should be raised with kind "{kind}" (elim)')
|
|
def step_elim_assert_error_kind(context: Any, kind: str) -> None:
|
|
assert context.elim_error is not None, (
|
|
f"Expected an ExecutionError(kind={kind!r}) but no error was raised. "
|
|
f"Result was: {context.elim_result!r}"
|
|
)
|
|
assert isinstance(context.elim_error, ExecutionError), (
|
|
f"Expected ExecutionError, got {type(context.elim_error).__name__}: "
|
|
f"{context.elim_error}"
|
|
)
|
|
assert context.elim_error.kind == kind, (
|
|
f"Expected kind={kind!r}, got kind={context.elim_error.kind!r}. "
|
|
f"Error message: {context.elim_error}"
|
|
)
|
|
|
|
|
|
@then('the reason should be "{reason}" (elim)')
|
|
def step_elim_assert_reason(context: Any, reason: str) -> None:
|
|
assert context.elim_error is not None, "No error was raised"
|
|
assert context.elim_error.reason == reason, (
|
|
f"Expected reason={reason!r}, got reason={context.elim_error.reason!r}"
|
|
)
|
|
|
|
|
|
@then("no ExecutionError should be raised (elim)")
|
|
def step_elim_assert_no_error(context: Any) -> None:
|
|
assert context.elim_error is None, (
|
|
f"Expected no error but got: {type(context.elim_error).__name__}: "
|
|
f"{context.elim_error} (kind={context.elim_error.kind!r})"
|
|
)
|
|
|
|
|
|
@then("the result should be non-None (elim)")
|
|
def step_elim_assert_result_non_none(context: Any) -> None:
|
|
"""m-6 fix: verify the graph returned a non-None result, not just absence of error.
|
|
|
|
A regression that returned None or discarded the output would pass the
|
|
'no ExecutionError' check but fail here.
|
|
"""
|
|
result_value = context.elim_result
|
|
# execute() returns a 3-tuple (output_str, captured_state, node_usages)
|
|
if isinstance(result_value, tuple):
|
|
output_str = result_value[0]
|
|
else:
|
|
output_str = result_value
|
|
assert output_str is not None, (
|
|
"Graph returned None output — expected a non-None result string"
|
|
)
|
|
|
|
|
|
@then("the sibling branch should have been cancelled (elim)")
|
|
def step_elim_assert_sibling_cancelled(context: Any) -> None:
|
|
"""Issue #2 fix: verify that branch_b was actually cancelled when branch_a raised.
|
|
|
|
A regression that removed the explicit t.cancel() loop in the parallel
|
|
execution block would not be detected by checking only that ExecutionError
|
|
was raised — asyncio.gather() propagates the first exception regardless of
|
|
whether siblings are cancelled. This step asserts that the CancelledError
|
|
was delivered to branch_b's coroutine.
|
|
"""
|
|
cancelled_flag = context.elim_branch_b_cancelled
|
|
assert cancelled_flag[0] is True, (
|
|
"branch_b was NOT cancelled after branch_a raised ExecutionError. "
|
|
"The t.cancel() loop in the parallel execution block may have been removed."
|
|
)
|
|
|
|
|
|
@then("a warning should have been logged for non-dict token usage (elim)")
|
|
def step_elim_assert_warning_logged(context: Any) -> None:
|
|
"""Issue #3 fix: verify that a warning was logged when _node_token_usage was non-dict.
|
|
|
|
A regression that silently dropped the non-dict value without logging would
|
|
not be detected by checking only that no ExecutionError was raised. This
|
|
step patches the logger and asserts the warning was called.
|
|
"""
|
|
# Re-run the graph with a patched logger to capture the warning call.
|
|
# We need to re-execute because the graph has already run in the When step;
|
|
# we patch the module-level logger used by PureLangGraph instances.
|
|
config = PureGraphConfig(
|
|
name="non_dict_token_test_verify",
|
|
nodes={
|
|
"llm": NodeConfig(name="llm", type=NodeType.AGENT),
|
|
},
|
|
edges=[
|
|
Edge(source="start", target="llm"),
|
|
Edge(source="llm", target="end"),
|
|
],
|
|
entry_point="start",
|
|
)
|
|
graph = PureLangGraph(config, limits={}, pricing={})
|
|
graph.nodes["llm"] = _make_non_dict_token_usage_node("llm")
|
|
|
|
with patch.object(_pg_module, "_logger") as mock_logger:
|
|
# Also patch the instance logger since PureLangGraph uses self.logger
|
|
with patch.object(graph, "logger") as mock_instance_logger:
|
|
_run_graph(graph)
|
|
# Check that warning was called on the instance logger
|
|
warning_calls = mock_instance_logger.warning.call_args_list
|
|
assert any(
|
|
"non-dict" in str(call).lower()
|
|
or "non_dict" in str(call).lower()
|
|
or "token_usage" in str(call).lower()
|
|
or "token accounting" in str(call).lower()
|
|
for call in warning_calls
|
|
), (
|
|
f"Expected a warning about non-dict _node_token_usage but none was logged. "
|
|
f"Warning calls: {warning_calls}"
|
|
)
|