feat: add persisted run step budget state

This commit is contained in:
lda
2026-09-05 18:28:54 +07:00 Verified
parent 8a49ecdf0e
commit 3fd1f70f5d
8 changed files with 496 additions and 17 deletions
+53
View File
@@ -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)