feat(core): add execution step budgets

burr/core/__init__.py

@@ -30,6 +30,7 @@
     Application,
     ApplicationBuilder,
     ApplicationContext,
+    ApplicationExecutionLimitError,
     ApplicationGraph,
 )
 from burr.core.graph import Graph, GraphBuilder
@@ -42,6 +43,7 @@
     "ApplicationBuilder",
     "ApplicationGraph",
     "ApplicationContext",
+    "ApplicationExecutionLimitError",
     "Condition",
     "default",
     "expr",

burr/core/application.py

@@ -89,6 +89,19 @@
 StateTypeToSet = TypeVar("StateTypeToSet")
 
 
+class ApplicationExecutionLimitError(RuntimeError):
+    """Raised when an application execution call exhausts its step budget."""
+
+    def __init__(self, max_steps: int, last_action: Action, state: State):
+        self.max_steps = max_steps
+        self.last_action = last_action
+        self.state = state
+        super().__init__(
+            f"Application execution reached max_steps={max_steps} after action "
+            f"'{last_action.name}' while another action was still available."
+        )
+
+
 def _validate_result(result: Any, name: str, schema: ActionSchema = DEFAULT_SCHEMA) -> None:
     # TODO -- split out the action schema into input/output schema types
     # Currently they're tied together, but this doesn't make as much sense for single-step actions
@@ -1262,6 +1275,7 @@ def iterate(
         halt_before: list[str] = None,
         halt_after: list[str] = None,
         inputs: Optional[Dict[str, Any]] = None,
+        max_steps: Optional[int] = None,
     ) -> Generator[
         Tuple[Action, dict, State[ApplicationStateType]],
         None,
@@ -1280,6 +1294,7 @@ def iterate(
         :param halt_after: The list of actions/tags to halt after execution of. It will halt after the execution of the first one it sees.
         :param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world.
             Note that this is only used for the first iteration -- subsequent iterations will not use this.
+        :param max_steps: Maximum actions to execute during this call. ``None`` preserves unbounded execution.
         :return: Each iteration returns the result of running `step`. This generator also returns a tuple of
             [action, result, current state]
         """
@@ -1290,15 +1305,23 @@ def iterate(
             inputs,
         )
         self._validate_halt_conditions(halt_before, halt_after)
+        if max_steps is not None and (
+            not isinstance(max_steps, int) or isinstance(max_steps, bool) or max_steps <= 0
+        ):
+            raise ValueError("max_steps must be a positive integer or None")
 
         result = None
         prior_action: Optional[Action] = None
+        steps_executed = 0
         while self.has_next_action():
             # self.step will only return None if there is no next action, so we can rely on tuple unpacking
             prior_action, result, state = self.step(inputs=inputs)
+            steps_executed += 1
             yield prior_action, result, state
             if self._should_halt_iterate(halt_before, halt_after, prior_action):
                 break
+            if max_steps is not None and steps_executed >= max_steps and self.has_next_action():
+                raise ApplicationExecutionLimitError(max_steps, prior_action, state)
         return self._return_value_iterate(halt_before, halt_after, prior_action, result)
 
     @_call_execute_method_pre_post(ExecuteMethod.aiterate)
@@ -1308,6 +1331,7 @@ async def aiterate(
         halt_before: list[str] = None,
         halt_after: list[str] = None,
         inputs: Optional[Dict[str, Any]] = None,
+        max_steps: Optional[int] = None,
     ) -> AsyncGenerator[Tuple[Action, dict, State[ApplicationStateType]], None]:
         """Returns a generator that calls step() in a row, enabling you to see the state
         of the system as it updates. This is the asynchronous version so it has no capability of t
@@ -1318,19 +1342,28 @@ async def aiterate(
         :param halt_after: The list of actions/tags to halt after execution of. It will halt on the first one.
         :param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world.
             Note that this is only used for the first iteration -- subsequent iterations will not use this.
+        :param max_steps: Maximum actions to execute during this call. ``None`` preserves unbounded execution.
         :return: Each iteration returns the result of running `step`. This returns nothing -- it's an async generator which is not
             allowed to have a return value.
         """
         halt_before, halt_after, inputs = self._process_control_flow_params(
             halt_before, halt_after, inputs
         )
         self._validate_halt_conditions(halt_before, halt_after)
+        if max_steps is not None and (
+            not isinstance(max_steps, int) or isinstance(max_steps, bool) or max_steps <= 0
+        ):
+            raise ValueError("max_steps must be a positive integer or None")
+        steps_executed = 0
         while self.has_next_action():
             # self.step will only return None if there is no next action, so we can rely on tuple unpacking
             prior_action, result, state = await self.astep(inputs=inputs)
+            steps_executed += 1
             yield prior_action, result, state
             if self._should_halt_iterate(halt_before, halt_after, prior_action):
                 break
+            if max_steps is not None and steps_executed >= max_steps and self.has_next_action():
+                raise ApplicationExecutionLimitError(max_steps, prior_action, state)
 
     @_call_execute_method_pre_post(ExecuteMethod.run)
     def run(
@@ -1339,6 +1372,7 @@ def run(
         halt_before: list[str] = None,
         halt_after: list[str] = None,
         inputs: Optional[Dict[str, Any]] = None,
+        max_steps: Optional[int] = None,
     ) -> Tuple[Action, Optional[dict], State[ApplicationStateType]]:
         """Runs your application through until completion. Does
         not give access to the state along the way -- if you want that, use iterate().
@@ -1349,10 +1383,16 @@ def run(
         :param halt_after: The list of actions/tags to halt after execution of. It will halt on the first one.
         :param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world.
             Note that this is only used for the first iteration -- subsequent iterations will not use this.
+        :param max_steps: Maximum actions to execute during this call. ``None`` preserves unbounded execution.
         :return: The final state, and the results of running the actions in the order that they were specified.
         """
         self.validate_correct_async_use()
-        gen = self.iterate(halt_before=halt_before, halt_after=halt_after, inputs=inputs)
+        gen = self.iterate(
+            halt_before=halt_before,
+            halt_after=halt_after,
+            inputs=inputs,
+            max_steps=max_steps,
+        )
         while True:
             try:
                 next(gen)
@@ -1367,6 +1407,7 @@ async def arun(
         halt_before: list[str] = None,
         halt_after: list[str] = None,
         inputs: Optional[Dict[str, Any]] = None,
+        max_steps: Optional[int] = None,
     ) -> Tuple[Action, Optional[dict], State[ApplicationStateType]]:
         """Runs your application through until completion, using async. Does
         not give access to the state along the way -- if you want that, use iterate().
@@ -1376,6 +1417,7 @@ async def arun(
         :param halt_before: The list of actions/tags to halt before execution of. It will halt on the first one.
         :param halt_after: The list of actions/tags to halt after execution of. It will halt on the first one.
         :param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world
+        :param max_steps: Maximum actions to execute during this call. ``None`` preserves unbounded execution.
         :return: The final state, and the results of running the actions in the order that they were specified.
         """
 
@@ -1386,7 +1428,10 @@ async def arun(
         )
         self._validate_halt_conditions(halt_before, halt_after)
         async for prior_action, result, state in self.aiterate(
-            halt_before=halt_before, halt_after=halt_after, inputs=inputs
+            halt_before=halt_before,
+            halt_after=halt_after,
+            inputs=inputs,
+            max_steps=max_steps,
         ):
             pass
         return self._return_value_iterate(halt_before, halt_after, prior_action, result)

docs/concepts/state-machine.rst

@@ -116,6 +116,28 @@ In the async context, this does not return anything (asynchronous generators are
 
 If you want it to (attempt to) run forever, you can pass empty lists to ``halt_after`` and ``halt_before``.
 
+Limiting execution steps
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Agent and tool-use graphs often contain intentional loops. Pass ``max_steps`` to
+``run``/``arun`` or ``iterate``/``aiterate`` to bound the number of actions a
+single call may execute:
+
+.. code-block:: python
+
+    from burr.core import ApplicationExecutionLimitError
+
+    try:
+        action, result, state = application.run(max_steps=20)
+    except ApplicationExecutionLimitError as exc:
+        # The application remains at the state produced by the last allowed step.
+        save_for_review(exc.state)
+
+The exception also exposes ``max_steps`` and ``last_action``. A halt condition
+or the natural end of the graph takes precedence when it is reached on the last
+allowed step. The budget applies only to the current execution call, so the
+application can be inspected and resumed with a new call.
+
 .. note::
     You can add inputs to ``iterate``/``aiterate`` by passing in a dictionary of inputs through the ``inputs`` parameter.
     This will only apply to the first action. Actions that are not the first but require inputs are considered undefined behavior.

docs/reference/application.rst

@@ -41,6 +41,9 @@ and not the ``Application`` class directly.
 .. autoclass:: burr.core.application.ApplicationContext
    :members:
 
+.. autoclass:: burr.core.application.ApplicationExecutionLimitError
+   :members:
+
 ==========
 Graph APIs
 ==========

tests/core/test_application.py

@@ -25,7 +25,7 @@
 
 import pytest
 
-from burr.core import State
+from burr.core import ApplicationExecutionLimitError, State
 from burr.core.action import (
     DEFAULT_SCHEMA,
     Action,
@@ -2171,6 +2171,184 @@ async def test_aiterate():
     assert app.sequence_id == 11
 
 
+def test_run_stops_before_exceeding_max_steps():
+    counter_action = base_counter_action.with_name("counter")
+    app = Application(
+        state=State({}),
+        entrypoint="counter",
+        partition_key="test",
+        uid="test-123",
+        sequence_id=0,
+        graph=Graph(
+            actions=[counter_action],
+            transitions=[Transition(counter_action, counter_action, default)],
+        ),
+    )
+
+    with pytest.raises(ApplicationExecutionLimitError) as exc_info:
+        app.run(max_steps=3)
+
+    assert exc_info.value.max_steps == 3
+    assert exc_info.value.last_action.name == "counter"
+    assert exc_info.value.state["count"] == 3
+    assert app.state["count"] == 3
+
+
+def test_iterate_raises_only_when_requesting_a_step_beyond_the_budget():
+    counter_action = base_counter_action.with_name("counter")
+    app = Application(
+        state=State({}),
+        entrypoint="counter",
+        partition_key="test",
+        uid="test-123",
+        sequence_id=0,
+        graph=Graph(
+            actions=[counter_action],
+            transitions=[Transition(counter_action, counter_action, default)],
+        ),
+    )
+    iterator = app.iterate(max_steps=1)
+
+    action_, result, state = next(iterator)
+
+    assert action_.name == "counter"
+    assert result == {"count": 1}
+    assert state["count"] == 1
+    with pytest.raises(ApplicationExecutionLimitError):
+        next(iterator)
+
+
+def test_run_allows_natural_completion_at_max_steps():
+    counter_action = base_counter_action.with_name("counter")
+    app = Application(
+        state=State({}),
+        entrypoint="counter",
+        partition_key="test",
+        uid="test-123",
+        sequence_id=0,
+        graph=Graph(actions=[counter_action], transitions=[]),
+    )
+
+    action_, result, state = app.run(max_steps=1)
+
+    assert action_.name == "counter"
+    assert result == {"count": 1}
+    assert state["count"] == 1
+
+
+def test_run_can_resume_after_reaching_max_steps():
+    counter_action = base_counter_action.with_name("counter")
+    app = Application(
+        state=State({}),
+        entrypoint="counter",
+        partition_key="test",
+        uid="test-123",
+        sequence_id=0,
+        graph=Graph(
+            actions=[counter_action],
+            transitions=[Transition(counter_action, counter_action, default)],
+        ),
+    )
+    with pytest.raises(ApplicationExecutionLimitError):
+        app.run(max_steps=1)
+
+    action_, result, state = app.run(halt_after=["counter"], max_steps=1)
+
+    assert action_.name == "counter"
+    assert result == {"count": 2}
+    assert state["count"] == 2
+
+
+def test_run_halt_condition_takes_precedence_over_max_steps():
+    counter_action = base_counter_action.with_name("counter")
+    app = Application(
+        state=State({}),
+        entrypoint="counter",
+        partition_key="test",
+        uid="test-123",
+        sequence_id=0,
+        graph=Graph(
+            actions=[counter_action],
+            transitions=[Transition(counter_action, counter_action, default)],
+        ),
+    )
+
+    action_, result, state = app.run(halt_after=["counter"], max_steps=1)
+
+    assert action_.name == "counter"
+    assert result == {"count": 1}
+    assert state["count"] == 1
+
+
+@pytest.mark.parametrize("max_steps", [0, -1, 1.5, "3", True])
+def test_run_rejects_invalid_max_steps(max_steps):
+    counter_action = base_counter_action.with_name("counter")
+    app = Application(
+        state=State({}),
+        entrypoint="counter",
+        partition_key="test",
+        uid="test-123",
+        sequence_id=0,
+        graph=Graph(actions=[counter_action], transitions=[]),
+    )
+
+    with pytest.raises(ValueError, match="max_steps must be a positive integer"):
+        app.run(max_steps=max_steps)
+
+    assert app.sequence_id == 0
+
+
+async def test_arun_stops_before_exceeding_max_steps():
+    counter_action = base_counter_action_async.with_name("counter")
+    app = Application(
+        state=State({}),
+        entrypoint="counter",
+        partition_key="test",
+        uid="test-123",
+        sequence_id=0,
+        graph=Graph(
+            actions=[counter_action],
+            transitions=[Transition(counter_action, counter_action, default)],
+        ),
+    )
+
+    with pytest.raises(ApplicationExecutionLimitError) as exc_info:
+        await app.arun(max_steps=2)
+
+    assert exc_info.value.max_steps == 2
+    assert exc_info.value.last_action.name == "counter"
+    assert exc_info.value.state["count"] == 2
+
+
+async def test_aiterate_raises_after_yield_and_async_halt_can_resume():
+    counter_action = base_counter_action_async.with_name("counter")
+    app = Application(
+        state=State({}),
+        entrypoint="counter",
+        partition_key="test",
+        uid="test-123",
+        sequence_id=0,
+        graph=Graph(
+            actions=[counter_action],
+            transitions=[Transition(counter_action, counter_action, default)],
+        ),
+    )
+    iterator = app.aiterate(max_steps=1)
+
+    action_, result, state = await iterator.__anext__()
+
+    assert action_.name == "counter"
+    assert result == {"count": 1}
+    assert state["count"] == 1
+    with pytest.raises(ApplicationExecutionLimitError):
+        await iterator.__anext__()
+
+    action_, result, state = await app.arun(halt_after=["counter"], max_steps=1)
+    assert action_.name == "counter"
+    assert result == {"count": 2}
+    assert state["count"] == 2
+
+
 async def test_aiterate_halt_before():
     result_action = Result("count").with_name("result")
     counter_action = base_counter_action_async.with_name("counter")