116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
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):
|
|
"""Versioned JSON storage envelope for one stopped runtime snapshot."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
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]
|
|
|
|
|
|
_RUN_STATE_ADAPTER = TypeAdapter(RunState)
|
|
|
|
|
|
def dump_run_state(run: RunState) -> dict[str, object]:
|
|
"""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 _inject_v1_budget_defaults(state: dict[str, Any]) -> dict[str, Any]:
|
|
"""Copy a v1 state dict with the one-time step budget defaults applied.
|
|
|
|
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.
|
|
"""
|
|
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
|