feat: add persisted run step budget state

src/wf_core/__init__.py

@@ -1,3 +1,4 @@
+from .errors import WorkflowStepLimitExceeded
 from .models import (
     ArrayExpression,
     ConditionNode,
@@ -31,7 +32,12 @@ from .models import (
     validate_strict_json_value,
     workflow_ref_from,
 )
-from .run_codec import PersistedRunState, dump_run_state, load_run_state
+from .run_codec import (
+    PersistedRunState,
+    dump_run_state,
+    load_run_state,
+    load_run_state_with_upgrade,
+)
 from .run_state import (
     ExecutionFrame,
     ForeachContext,
@@ -59,6 +65,7 @@ from .runtime import (
     step_workflow,
     step_workflow_async,
 )
+from .runtime.limits import RunLimits, admit_step_attempt, remaining_step_attempts
 from .tokens import END, START
 from .validation import (
     ValidationIssue,
@@ -100,6 +107,7 @@ __all__ = [
     "PreparedSubgraph",
     "ReducerRef",
     "ReducerSpec",
+    "RunLimits",
     "RunState",
     "RunStatus",
     "RuntimeContext",
@@ -117,12 +125,16 @@ __all__ = [
     "Workflow",
     "WorkflowExecutionError",
     "WorkflowRef",
+    "WorkflowStepLimitExceeded",
+    "admit_step_attempt",
     "coerce_node_result",
     "dump_run_state",
     "execute_workflow",
     "execute_workflow_async",
     "execute_workflow_result_async",
     "load_run_state",
+    "load_run_state_with_upgrade",
+    "remaining_step_attempts",
     "resume_workflow",
     "resume_workflow_async",
     "resume_workflow_result_async",

src/wf_core/errors.py

@@ -1,5 +1,56 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from wf_core.run_state import ExecutionFrame, RunState
+
+
 class WorkflowExecutionError(RuntimeError):
     pass
 
 
-__all__ = ["WorkflowExecutionError"]
+class WorkflowStepLimitExceeded(WorkflowExecutionError):
+    """Raised when a run would exceed its persisted step budget.
+
+    Budget exhaustion is a runtime failure, never a routable workflow outcome:
+    no edge may catch it as an ``error`` outcome.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        max_steps: int,
+        steps_executed: int,
+        frame_id: str,
+        scope_id: str,
+        node_id: str,
+    ) -> None:
+        super().__init__(message)
+        self.max_steps = max_steps
+        self.steps_executed = steps_executed
+        self.frame_id = frame_id
+        self.scope_id = scope_id
+        self.node_id = node_id
+
+    @classmethod
+    def from_run(
+        cls, run: RunState, frame: ExecutionFrame, node_id: str
+    ) -> WorkflowStepLimitExceeded:
+        """Build an exhaustion error for the denied dispatch of ``node_id``."""
+        return cls(
+            f"workflow {run.workflow_name!r} exceeded its step budget "
+            f"(max_steps={run.limits.max_steps}, "
+            f"steps_executed={run.steps_executed}, "
+            f"frame_id={frame.id!r}, scope_id={frame.scope_id!r}, "
+            f"next_node_id={node_id!r})",
+            max_steps=run.limits.max_steps,
+            steps_executed=run.steps_executed,
+            frame_id=frame.id,
+            scope_id=frame.scope_id,
+            node_id=node_id,
+        )
+
+
+__all__ = ["WorkflowExecutionError", "WorkflowStepLimitExceeded"]

src/wf_core/run_codec.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+from copy import deepcopy
 from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 
 from .run_state import ROOT_SCOPE_ID, RunState
+from .runtime.limits import RunLimits
 
 
 class PersistedRunState(BaseModel):
@@ -12,7 +14,16 @@ class PersistedRunState(BaseModel):
 
     model_config = ConfigDict(extra="forbid")
 
-    version: Literal[1] = 1
+    version: Literal[2] = 2
+    state: dict[str, Any]
+
+
+class _AnyPersistedRunState(BaseModel):
+    """Loading envelope accepting every currently readable version."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    version: Literal[1, 2]
     state: dict[str, Any]
 
 
@@ -20,25 +31,85 @@ _RUN_STATE_ADAPTER = TypeAdapter(RunState)
 
 
 def dump_run_state(run: RunState) -> dict[str, object]:
-    """Serialize one stopped `RunState` into the durable v1 envelope."""
+    """Serialize one stopped `RunState` into the durable v2 envelope."""
     return PersistedRunState(
         state=_RUN_STATE_ADAPTER.dump_python(run, mode="json")
     ).model_dump(mode="json")
 
 
-def load_run_state(payload: object) -> RunState:
+def _inject_v1_budget_defaults(state: dict[str, Any]) -> dict[str, Any]:
-    """Validate and restore one durable v1 runtime snapshot.
+    """Copy a v1 state dict with the one-time step budget defaults applied.
 
-    The root scope intentionally shares the compatibility ``RunState.state``
+    Version-1 envelopes predate step budgets, so they receive the default
-    dict during runtime. Serialization loses object identity, so restored
+    limit, a zeroed counter, and an unassigned number per frame exactly once
-    snapshots must recreate this alias before resumed writes occur.
+    at load time. Attempts made before the upgrade are outside the new budget.
     """
-    envelope = PersistedRunState.model_validate(payload)
+    upgraded = deepcopy(state)
-    try:
+    upgraded.setdefault("limits", {"max_steps": RunLimits().max_steps})
-        run = _RUN_STATE_ADAPTER.validate_python(envelope.state)
+    upgraded.setdefault("steps_executed", 0)
-    except ValidationError as exc:
+    frames = upgraded.get("frames")
-        raise ValueError("invalid persisted workflow run state") from exc
+    if isinstance(frames, dict):
+        for frame in frames.values():
+            if isinstance(frame, dict):
+                frame.setdefault("step_number", None)
+    return upgraded
+
+
+def _require_v2_budget_fields(state: dict[str, Any]) -> None:
+    """Reject v2 payloads missing budget fields as corrupt state.
+
+    Unlike v1, a v2 envelope promises budget fields; a missing counter is
+    corruption, not another request for defaults.
+    """
+    if "limits" not in state or "steps_executed" not in state:
+        raise ValueError("invalid persisted workflow run state: missing step budget")
+    frames = state.get("frames")
+    if not isinstance(frames, dict):
+        raise ValueError("invalid persisted workflow run state: missing frames")
+    for frame_id, frame in frames.items():
+        if not isinstance(frame, dict) or "step_number" not in frame:
+            raise ValueError(
+                "invalid persisted workflow run state: "
+                f"frame {frame_id!r} is missing its step number"
+            )
+
+
+def _restore_root_alias(run: RunState) -> RunState:
+    """Recreate the root scope compatibility alias lost by serialization."""
     root_scope = run.scopes.get(ROOT_SCOPE_ID)
     if root_scope is not None:
         run.state = root_scope.committed_state
     return run
+
+
+def load_run_state_with_upgrade(payload: object) -> tuple[RunState, bool]:
+    """Validate one durable snapshot, upgrading v1 envelopes exactly once.
+
+    Returns the restored run plus whether a v1-to-v2 upgrade was applied.
+    """
+    envelope = _AnyPersistedRunState.model_validate(payload)
+    if envelope.version == 1:
+        state = _inject_v1_budget_defaults(envelope.state)
+        try:
+            run = _RUN_STATE_ADAPTER.validate_python(state)
+        except ValidationError as exc:
+            raise ValueError("invalid persisted workflow run state") from exc
+        return _restore_root_alias(run), True
+    _require_v2_budget_fields(envelope.state)
+    try:
+        run = _RUN_STATE_ADAPTER.validate_python(envelope.state)
+    except ValidationError as exc:
+        raise ValueError("invalid persisted workflow run state") from exc
+    return _restore_root_alias(run), False
+
+
+def load_run_state(payload: object) -> RunState:
+    """Validate and restore one durable runtime snapshot.
+
+    Version-1 envelopes receive step budget defaults via the one-time upgrade.
+    The root scope intentionally shares the compatibility ``RunState.state``
+    dict during runtime. Serialization loses object identity, so restored
+    snapshots must recreate this alias before resumed writes occur.
+    """
+    run, _ = load_run_state_with_upgrade(payload)
+    return run

src/wf_core/run_state.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 
 from dataclasses import asdict, dataclass, field
 from enum import StrEnum
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from wf_core.models.reducers import ReducerRef
 from wf_core.models.workflow_refs import WorkflowRef
 from wf_core.paths import StatePath
 
+if TYPE_CHECKING:
+    from wf_core.runtime.limits import RunLimits
+
 ROOT_SCOPE_ID = "root"
 ROOT_LINEAGE_ID = "root"
 ROOT_FRAME_ID = "root"
@@ -84,6 +87,7 @@ class ExecutionFrame:
     activated_incoming_edge: str | None = None
     metadata: dict[str, Any] = field(default_factory=dict)
     finished_at_node_id: str | None = None
+    step_number: int | None = None
 
 
 @dataclass(frozen=True, slots=True)
@@ -177,8 +181,25 @@ class InterruptRequest:
     typed: bool = False
 
 
+def _default_run_limits() -> RunLimits:
+    """Build the default budget without a top-level runtime import.
+
+    Importing ``wf_core.runtime.limits`` at module top would execute the
+    ``wf_core.runtime`` package, whose engine imports this module back.
+    """
+    from wf_core.runtime.limits import RunLimits
+
+    return RunLimits()
+
+
 @dataclass(slots=True)
 class RunState:
+    """Mutable execution state for one workflow run.
+
+    ``limits``/``steps_executed`` form the persisted run-wide step budget (see
+    ``wf_core.runtime.limits``); ``steps_remaining`` is computed from them.
+    """
+
     workflow_name: str
     status: RunStatus
     workflow_input: dict[str, Any]
@@ -196,6 +217,13 @@ class RunState:
     activated_incoming_edge: str | None = None
     error: str | None = None
     interrupt: InterruptRequest | None = None
+    limits: RunLimits = field(default_factory=_default_run_limits)
+    steps_executed: int = 0
+
+    @property
+    def steps_remaining(self) -> int:
+        """Unspent step budget, floored at zero; computed, never persisted."""
+        return max(self.limits.max_steps - self.steps_executed, 0)
 
     def current_frame(self) -> ExecutionFrame:
         if self.current_frame_id is None:
@@ -217,3 +245,9 @@ class RunState:
 
     def to_dict(self) -> dict[str, Any]:
         return asdict(self)
+
+
+# Deferred runtime import: binding ``RunLimits`` here (after every class is
+# defined) lets ``wf_core.runtime`` engine modules import this module back
+# without a cycle, and gives pydantic a resolvable annotation for the codec.
+from wf_core.runtime.limits import RunLimits  # noqa: E402

src/wf_core/runtime/limits.py

@@ -0,0 +1,53 @@
+"""Run-wide step budget policy and admission.
+
+One finite, persisted counter (``RunState.steps_executed``) covers every frame
+and subgraph scope in a run. Admission happens immediately before step dispatch:
+an admitted attempt increments the counter and stamps the selected frame with
+its one-based step number; a denied attempt raises without incrementing and
+without invoking any handler. This slice deliberately adds no per-step durable
+checkpoints; the counter is persisted inside the existing stopped-run envelope.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from wf_core.errors import WorkflowStepLimitExceeded
+
+if TYPE_CHECKING:
+    from wf_core.run_state import ExecutionFrame, RunState
+
+
+@dataclass(frozen=True, slots=True)
+class RunLimits:
+    """Immutable step budget captured when a run is created."""
+
+    max_steps: int = 10_000
+
+    def __post_init__(self) -> None:
+        if isinstance(self.max_steps, bool) or not isinstance(self.max_steps, int):
+            raise TypeError("max_steps must be an integer")
+        if self.max_steps < 1:
+            raise ValueError("max_steps must be positive")
+
+
+def admit_step_attempt(run: RunState, frame: ExecutionFrame, node_id: str) -> int:
+    """Admit one step attempt for ``frame`` about to dispatch ``node_id``.
+
+    On success the run-wide counter is incremented, the frame remembers the
+    assigned step number, and that number is returned. When the budget is
+    already exhausted the counter is left untouched, the frame keeps its
+    previous number, and ``WorkflowStepLimitExceeded`` is raised before any
+    handler runs.
+    """
+    if run.steps_executed >= run.limits.max_steps:
+        raise WorkflowStepLimitExceeded.from_run(run, frame, node_id)
+    run.steps_executed += 1
+    frame.step_number = run.steps_executed
+    return frame.step_number
+
+
+def remaining_step_attempts(run: RunState) -> int:
+    """Return the unspent budget, floored at zero (never negative)."""
+    return max(run.limits.max_steps - run.steps_executed, 0)

src/wf_core/runtime/ops/runs.py

@@ -15,6 +15,7 @@ from wf_core.run_state import (
     RunStatus,
     RuntimeScope,
 )
+from wf_core.runtime.limits import RunLimits
 from wf_core.runtime.scheduler import add_frame
 
 
@@ -30,13 +31,19 @@ def initial_state(
     return state
 
 
-def create_run_state(workflow: Workflow, workflow_input: dict[str, object]) -> RunState:
+def create_run_state(
+    workflow: Workflow,
+    workflow_input: dict[str, object],
+    *,
+    limits: RunLimits | None = None,
+) -> RunState:
     state = initial_state(workflow, workflow_input)
     run = RunState(
         workflow_name=workflow.name,
         status=RunStatus.PENDING,
         workflow_input=dict(workflow_input),
         state=state,
+        limits=limits if limits is not None else RunLimits(),
         scopes={
             ROOT_SCOPE_ID: RuntimeScope(
                 id=ROOT_SCOPE_ID,

tests/core/test_run_codec.py

@@ -27,7 +27,7 @@ def test_run_state_codec_round_trips_completed_output() -> None:
     stored = dump_run_state(run)
     restored = load_run_state(stored)
 
-    assert stored["version"] == 1
+    assert stored["version"] == 2
     assert restored.status is RunStatus.COMPLETED
     assert restored.output["echoed"] == "hi"
 

tests/core/test_run_step_budget.py

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+import pytest
+
+from wf_core import (
+    END,
+    Edge,
+    NodeDef,
+    NodeUse,
+    SchemaRef,
+    StateSchema,
+    Workflow,
+    WorkflowExecutionError,
+    dump_run_state,
+    load_run_state,
+)
+from wf_core.errors import WorkflowStepLimitExceeded
+from wf_core.run_codec import load_run_state_with_upgrade
+from wf_core.runtime.limits import (
+    RunLimits,
+    admit_step_attempt,
+    remaining_step_attempts,
+)
+from wf_core.runtime.ops.runs import create_run_state
+
+
+def _minimal_workflow(name: str = "budget") -> Workflow:
+    return Workflow(
+        name=name,
+        input_schema=SchemaRef(type="object", properties={}),
+        state_schema=StateSchema.from_field_map({}),
+        output_schema=SchemaRef(type="object", properties={}),
+        node_defs=[
+            NodeDef(
+                name="finish",
+                input_schema=SchemaRef(type="object", properties={}),
+                output_schema=SchemaRef(type="object", properties={}),
+                outcomes=["ok"],
+            )
+        ],
+        start="finish",
+        nodes=[
+            NodeUse.model_validate({"id": "finish", "type": "node", "node": "finish"})
+        ],
+        edges=[Edge.model_validate({"from": "finish", "outcome": "ok", "to": END})],
+    )
+
+
+def test_run_limits_default() -> None:
+    limits = RunLimits()
+
+    assert limits.max_steps == 10_000
+
+
+def test_run_limits_rejects_non_positive() -> None:
+    with pytest.raises(ValueError):
+        RunLimits(max_steps=0)
+    with pytest.raises(ValueError):
+        RunLimits(max_steps=-3)
+
+
+def test_run_limits_rejects_bool_and_non_int() -> None:
+    with pytest.raises(TypeError):
+        RunLimits(max_steps=True)  # type: ignore[arg-type]
+    with pytest.raises(TypeError):
+        RunLimits(max_steps=False)  # type: ignore[arg-type]
+    with pytest.raises(TypeError):
+        RunLimits(max_steps="10")  # type: ignore[arg-type]
+    with pytest.raises(TypeError):
+        RunLimits(max_steps=10.0)  # type: ignore[arg-type]
+
+
+def test_create_run_state_defaults_to_budget() -> None:
+    workflow = _minimal_workflow()
+
+    run = create_run_state(workflow, {})
+
+    assert run.limits.max_steps == 10_000
+    assert run.steps_executed == 0
+    assert run.steps_remaining == 10_000
+    assert run.current_frame().step_number is None
+    assert remaining_step_attempts(run) == 10_000
+
+
+def test_create_run_state_captures_limits() -> None:
+    workflow = _minimal_workflow()
+    limits = RunLimits(max_steps=5)
+
+    run = create_run_state(workflow, {}, limits=limits)
+
+    assert run.limits.max_steps == 5
+    assert run.steps_remaining == 5
+
+
+def test_budget_of_one() -> None:
+    workflow = _minimal_workflow()
+    limits = RunLimits(max_steps=1)
+    run = create_run_state(workflow, {}, limits=limits)
+    number = admit_step_attempt(run, run.current_frame(), workflow.start)
+
+    assert number == 1
+    assert run.steps_executed == 1
+    assert run.steps_remaining == 0
+    with pytest.raises(WorkflowStepLimitExceeded):
+        admit_step_attempt(run, run.current_frame(), workflow.start)
+
+
+def test_denied_admission_does_not_increment() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {}, limits=RunLimits(max_steps=1))
+    admit_step_attempt(run, run.current_frame(), workflow.start)
+
+    with pytest.raises(WorkflowStepLimitExceeded):
+        admit_step_attempt(run, run.current_frame(), workflow.start)
+
+    assert run.steps_executed == 1
+    assert run.current_frame().step_number == 1
+    assert run.steps_remaining == 0
+    assert remaining_step_attempts(run) == 0
+
+
+def test_admission_assigns_step_numbers() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {}, limits=RunLimits(max_steps=3))
+
+    first = admit_step_attempt(run, run.current_frame(), workflow.start)
+    second = admit_step_attempt(run, run.current_frame(), workflow.start)
+
+    assert first == 1
+    assert second == 2
+    assert run.steps_executed == 2
+    assert run.current_frame().step_number == 2
+    assert run.steps_remaining == 1
+    assert remaining_step_attempts(run) == 1
+
+
+def test_step_limit_error_details() -> None:
+    workflow = _minimal_workflow(name="budget_details")
+    run = create_run_state(workflow, {}, limits=RunLimits(max_steps=1))
+    frame = run.current_frame()
+    admit_step_attempt(run, frame, workflow.start)
+
+    with pytest.raises(WorkflowStepLimitExceeded) as exc_info:
+        admit_step_attempt(run, frame, workflow.start)
+
+    assert isinstance(exc_info.value, WorkflowExecutionError)
+    message = str(exc_info.value)
+    assert "budget_details" in message
+    assert "1" in message
+    assert frame.id in message
+    assert frame.scope_id in message
+    assert workflow.start in message
+
+
+def test_remaining_never_negative() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {}, limits=RunLimits(max_steps=1))
+    run.steps_executed = 5
+
+    assert run.steps_remaining == 0
+    assert remaining_step_attempts(run) == 0
+
+
+def test_dump_writes_v2_envelope() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {}, limits=RunLimits(max_steps=7))
+    admit_step_attempt(run, run.current_frame(), workflow.start)
+
+    stored = dump_run_state(run)
+
+    assert stored["version"] == 2
+
+
+def test_v2_round_trip() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {}, limits=RunLimits(max_steps=7))
+    admit_step_attempt(run, run.current_frame(), workflow.start)
+
+    stored = dump_run_state(run)
+    restored, upgraded = load_run_state_with_upgrade(stored)
+
+    assert upgraded is False
+    assert restored.limits.max_steps == 7
+    assert restored.steps_executed == 1
+    assert restored.steps_remaining == 6
+    assert restored.frames["root"].step_number == 1
+
+    via_legacy = load_run_state(stored)
+    assert via_legacy.limits.max_steps == 7
+    assert via_legacy.steps_executed == 1
+    assert via_legacy.frames["root"].step_number == 1
+
+
+def _strip_to_v1(stored: dict) -> dict:
+    state = dict(stored["state"])
+    state.pop("limits", None)
+    state.pop("steps_executed", None)
+    frames = {
+        frame_id: {key: value for key, value in frame.items() if key != "step_number"}
+        for frame_id, frame in dict(state["frames"]).items()
+    }
+    state["frames"] = frames
+    return {"version": 1, "state": state}
+
+
+def test_v1_payload_receives_defaults_once() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {}, limits=RunLimits(max_steps=7))
+    stored = _strip_to_v1(dump_run_state(run))
+
+    restored, upgraded = load_run_state_with_upgrade(stored)
+
+    assert upgraded is True
+    assert restored.limits.max_steps == 10_000
+    assert restored.steps_executed == 0
+    assert restored.steps_remaining == 10_000
+    assert restored.frames["root"].step_number is None
+
+    via_legacy = load_run_state(stored)
+    assert via_legacy.limits.max_steps == 10_000
+    assert via_legacy.steps_executed == 0
+
+
+def test_v2_missing_limits_is_corrupt() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {})
+    stored = dump_run_state(run)
+    stored["state"].pop("limits")
+
+    with pytest.raises(ValueError):
+        load_run_state_with_upgrade(stored)
+
+
+def test_v2_missing_steps_executed_is_corrupt() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {})
+    stored = dump_run_state(run)
+    stored["state"].pop("steps_executed")
+
+    with pytest.raises(ValueError):
+        load_run_state_with_upgrade(stored)
+
+
+def test_v2_missing_frame_step_number_is_corrupt() -> None:
+    workflow = _minimal_workflow()
+    run = create_run_state(workflow, {})
+    stored = dump_run_state(run)
+    del stored["state"]["frames"]["root"]["step_number"]
+
+    with pytest.raises(ValueError):
+        load_run_state_with_upgrade(stored)