feat: add persisted run step budget state
This commit is contained in:
+13
-1
@@ -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",
|
||||
|
||||
+52
-1
@@ -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"]
|
||||
|
||||
+83
-12
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user