diff --git a/src/wf_core/__init__.py b/src/wf_core/__init__.py index b8a681bc..6f76cc7a 100644 --- a/src/wf_core/__init__.py +++ b/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", diff --git a/src/wf_core/errors.py b/src/wf_core/errors.py index d1f6b10a..63b72eb6 100644 --- a/src/wf_core/errors.py +++ b/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"] diff --git a/src/wf_core/run_codec.py b/src/wf_core/run_codec.py index b01fb29f..0fe811d7 100644 --- a/src/wf_core/run_codec.py +++ b/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: - """Validate and restore one durable v1 runtime snapshot. +def _inject_v1_budget_defaults(state: dict[str, Any]) -> dict[str, Any]: + """Copy a v1 state dict with the one-time step budget defaults applied. - 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. + Version-1 envelopes predate step budgets, so they receive the default + limit, a zeroed counter, and an unassigned number per frame exactly once + at load time. Attempts made before the upgrade are outside the new budget. """ - envelope = PersistedRunState.model_validate(payload) - try: - run = _RUN_STATE_ADAPTER.validate_python(envelope.state) - except ValidationError as exc: - raise ValueError("invalid persisted workflow run state") from exc + upgraded = deepcopy(state) + upgraded.setdefault("limits", {"max_steps": RunLimits().max_steps}) + upgraded.setdefault("steps_executed", 0) + frames = upgraded.get("frames") + 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 diff --git a/src/wf_core/run_state.py b/src/wf_core/run_state.py index e3283ed8..dc8c9ed1 100644 --- a/src/wf_core/run_state.py +++ b/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 diff --git a/src/wf_core/runtime/limits.py b/src/wf_core/runtime/limits.py new file mode 100644 index 00000000..2237ec99 --- /dev/null +++ b/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) diff --git a/src/wf_core/runtime/ops/runs.py b/src/wf_core/runtime/ops/runs.py index 152db0ff..359c93e3 100644 --- a/src/wf_core/runtime/ops/runs.py +++ b/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, diff --git a/tests/core/test_run_codec.py b/tests/core/test_run_codec.py index 01f5c94c..ec22ef14 100644 --- a/tests/core/test_run_codec.py +++ b/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" diff --git a/tests/core/test_run_step_budget.py b/tests/core/test_run_step_budget.py new file mode 100644 index 00000000..783e9d9e --- /dev/null +++ b/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)