lda.chat/lda-wf
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs: complete run step budget slice

Browse Source
This commit is contained in:
lda
2026-09-05 23:15:57 +07:00 Verified
parent 715035dce8
commit fd48a67819
19 changed files with 735 additions and 541 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/core/test_run_limits.py
+168
View File
@@ -0,0 +1,168 @@
"""Run limit model and admission tests."""
from __future__ import annotations
from typing import Any, cast
import pytest
from wf_core import (
END,
Edge,
NodeDef,
NodeUse,
RunLimits,
SchemaRef,
StateSchema,
Workflow,
WorkflowExecutionError,
)
from wf_core.errors import WorkflowStepLimitExceeded
from wf_core.runtime.limits import 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=cast(Any, True))
with pytest.raises(TypeError):
RunLimits(max_steps=cast(Any, False))
with pytest.raises(TypeError):
RunLimits(max_steps=cast(Any, "10"))
with pytest.raises(TypeError):
RunLimits(max_steps=cast(Any, 10.0))
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_remaining_delegates_to_run_property() -> None:
"""`remaining_step_attempts` is the computed `steps_remaining` convenience."""
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=4))
admit_step_attempt(run, run.current_frame(), workflow.start)
assert remaining_step_attempts(run) == run.steps_remaining == 3
tests/core/test_run_step_budget.py
+31 -411
View File
@@ -1,6 +1,16 @@
"""Sync step-budget dispatch tests.
Covers run-wide counting for every current step kind, trace numbering,
and the engine-level limits seam. Model/admission unit tests live in
`test_run_limits.py`; codec and migration tests live in
`test_run_step_budget_codec.py`.
"""
from __future__ import annotations
import inspect
from collections.abc import Callable
from typing import Any
import pytest
@@ -15,30 +25,23 @@ from wf_core import (
NodeUse,
PreparedSubgraph,
ReducerRef,
RunLimits,
RunStatus,
SchemaRef,
StateField,
StateSchema,
SubgraphNode,
Workflow,
WorkflowExecutionError,
dump_run_state,
execute_workflow,
execute_workflow_async,
execute_workflow_result_async,
load_run_state,
resume_workflow,
resume_workflow_async,
resume_workflow_result_async,
step_workflow,
)
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.run_state import RunState
from wf_core.runtime.ops.runs import create_run_state
from wf_core.runtime.preparation import prepare_resume
@@ -65,343 +68,15 @@ def _minimal_workflow(name: str = "budget") -> Workflow:
)
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)
def test_v1_smuggled_budget_fields_receive_defaults() -> None:
"""V1 envelopes predate budgets; smuggled fields must not survive."""
from typing import Any, cast
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=7))
stored = _strip_to_v1(dump_run_state(run))
state = cast(dict[str, Any], stored["state"])
state["limits"] = {"max_steps": 999_999}
state["steps_executed"] = 999
restored, upgraded = load_run_state_with_upgrade(stored)
assert upgraded is True
assert restored.limits.max_steps == 10_000
assert restored.steps_executed == 0
def test_v2_bool_max_steps_is_corrupt() -> None:
from typing import Any, cast
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
cast(dict[str, Any], stored["state"])["limits"] = {"max_steps": True}
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_str_max_steps_is_corrupt() -> None:
from typing import Any, cast
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
cast(dict[str, Any], stored["state"])["limits"] = {"max_steps": "10"}
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_negative_steps_executed_is_corrupt() -> None:
from typing import Any, cast
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
cast(dict[str, Any], stored["state"])["steps_executed"] = -100
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_exceeding_steps_executed_is_corrupt() -> None:
from typing import Any, cast
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
cast(dict[str, Any], stored["state"])["steps_executed"] = 5
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_bool_steps_executed_is_corrupt() -> None:
from typing import Any, cast
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
cast(dict[str, Any], stored["state"])["steps_executed"] = True
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_incoherent_frame_step_number_is_corrupt() -> None:
from typing import Any, cast
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
admit_step_attempt(run, run.current_frame(), workflow.start)
stored = dump_run_state(run)
state = cast(dict[str, Any], stored["state"])
frames = cast(dict[str, Any], state["frames"])
cast(dict[str, Any], frames["root"])["step_number"] = 999
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_bool_frame_step_number_is_corrupt() -> None:
from typing import Any, cast
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
admit_step_attempt(run, run.current_frame(), workflow.start)
stored = dump_run_state(run)
state = cast(dict[str, Any], stored["state"])
frames = cast(dict[str, Any], state["frames"])
cast(dict[str, Any], frames["root"])["step_number"] = True
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_pending_frame_none_number_round_trips() -> None:
"""Fresh V2 runs legitimately persist unadmitted (None) frame numbers."""
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
restored, upgraded = load_run_state_with_upgrade(dump_run_state(run))
assert upgraded is False
assert restored.steps_executed == 0
assert restored.frames["root"].step_number is None
# --- Task 2: sync dispatch and trace numbering ---
def _empty_schema() -> SchemaRef:
return SchemaRef(type="object", properties={})
def _ok_handler(payload: dict, _context: object) -> dict:
def _ok_handler(payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "ok", "output": {}}
def _trace_numbers(run) -> list:
def _trace_numbers(run: RunState) -> list[int | None]:
return [entry.step_number for entry in run.trace]
@@ -557,16 +232,10 @@ def _serial_foreach_workflow() -> Workflow:
def test_sync_foreach_controller_and_body_share_counter() -> None:
workflow = _serial_foreach_workflow()
run = execute_workflow(
workflow,
{"items": ["a", "b"]},
{
"record": lambda payload, _ctx: {
"outcome": "ok",
"output": {"seen": payload["value"]},
}
},
)
def record(payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "ok", "output": {"seen": payload["value"]}}
run = execute_workflow(workflow, {"items": ["a", "b"]}, {"record": record})
assert run.status == RunStatus.COMPLETED
assert run.steps_executed == 5
@@ -741,7 +410,7 @@ def test_sync_explicit_end_counts() -> None:
edges=[Edge.model_validate({"from": "finish", "outcome": "done", "to": "end"})],
)
def finish(_payload: dict, _context: object) -> dict:
def finish(_payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "done", "output": {}}
run = execute_workflow(workflow, {}, {"finish": finish})
@@ -766,7 +435,7 @@ def test_sync_legacy_end_creates_no_extra_attempt() -> None:
def test_sync_handler_failure_consumes_attempt() -> None:
workflow = _minimal_workflow()
def explode(_payload: dict, _context: object) -> dict:
def explode(_payload: dict[str, Any], _context: object) -> dict[str, Any]:
raise ValueError("boom")
run = create_run_state(workflow, {})
@@ -801,7 +470,7 @@ def test_sync_handled_error_outcome_counts_once() -> None:
],
)
def fail_soft(_payload: dict, _context: object) -> dict:
def fail_soft(_payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "error", "output": {}}
run = execute_workflow(workflow, {}, {"risky": fail_soft})
@@ -850,8 +519,10 @@ def test_sync_closed_cycle_fails_at_limit() -> None:
workflow = _cyclic_workflow()
calls: list[str] = []
def make(name: str): # type: ignore[no-untyped-def]
def handler(_payload: dict, _context: object) -> dict:
def make(
name: str,
) -> Callable[[dict[str, Any], object], dict[str, Any]]:
def handler(_payload: dict[str, Any], _context: object) -> dict[str, Any]:
calls.append(name)
return {"outcome": "ok", "output": {}}
@@ -925,7 +596,7 @@ def _counting_loop_workflow() -> Workflow:
def test_sync_exiting_loop_completes_within_budget() -> None:
workflow = _counting_loop_workflow()
def bump(payload: dict, _context: object) -> dict:
def bump(payload: dict[str, Any], _context: object) -> dict[str, Any]:
count = payload.get("count", 0)
assert isinstance(count, int)
return {"outcome": "ok", "output": {"count": count + 1}}
@@ -939,9 +610,9 @@ def test_sync_exiting_loop_completes_within_budget() -> None:
def test_sync_denial_never_invokes_handler() -> None:
workflow = _chain_workflow()
b_calls: list[dict] = []
b_calls: list[dict[str, Any]] = []
def b_handler(payload: dict, _context: object) -> dict:
def b_handler(payload: dict[str, Any], _context: object) -> dict[str, Any]:
b_calls.append(payload)
return {"outcome": "ok", "output": {}}
@@ -955,57 +626,6 @@ def test_sync_denial_never_invokes_handler() -> None:
assert b_calls == []
def _strip_budget_fields(stored: dict) -> dict:
state = dict(_strip_to_v1(stored)["state"])
state["trace"] = [
{key: value for key, value in entry.items() if key != "step_number"}
for entry in state.get("trace", [])
]
if state.get("interrupt") is not None:
state["interrupt"] = {
key: value
for key, value in state["interrupt"].items()
if key != "step_number"
}
return {"version": 1, "state": state}
def test_v1_traces_and_interrupt_receive_none_step_numbers() -> None:
workflow = _interrupt_workflow()
run = execute_workflow(workflow, {}, {"work": _ok_handler})
stored = _strip_budget_fields(dump_run_state(run))
restored, upgraded = load_run_state_with_upgrade(stored)
assert upgraded is True
assert restored.trace[0].step_number is None
assert restored.interrupt is not None
assert restored.interrupt.step_number is None
def test_v2_missing_trace_step_number_is_corrupt() -> None:
workflow = _interrupt_workflow()
run = execute_workflow(workflow, {}, {"work": _ok_handler})
stored = dump_run_state(run)
del stored["state"]["trace"][0]["step_number"]
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_missing_interrupt_step_number_is_corrupt() -> None:
workflow = _interrupt_workflow()
run = execute_workflow(workflow, {}, {"work": _ok_handler})
stored = dump_run_state(run)
del stored["state"]["interrupt"]["step_number"]
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
# --- Task 4: engine-level limits seam ---
def test_execute_workflow_accepts_explicit_limits() -> None:
workflow = _chain_workflow()
@@ -1034,7 +654,7 @@ def test_execute_workflow_defaults_to_ten_thousand() -> None:
def test_execute_workflow_enforces_limits() -> None:
workflow = _cyclic_workflow()
def ok_handler(_payload: dict, _context: object) -> dict:
def ok_handler(_payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "ok", "output": {}}
with pytest.raises(WorkflowStepLimitExceeded):
@@ -1049,7 +669,7 @@ def test_execute_workflow_enforces_limits() -> None:
async def test_execute_workflow_async_accepts_explicit_limits() -> None:
workflow = _chain_workflow()
async def ok_async(_payload: dict, _context: object) -> dict:
async def ok_async(_payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "ok", "output": {}}
run = await execute_workflow_async(
@@ -1068,7 +688,7 @@ async def test_execute_workflow_async_accepts_explicit_limits() -> None:
async def test_execute_workflow_result_async_reports_exhaustion() -> None:
workflow = _cyclic_workflow()
async def ok_async(_payload: dict, _context: object) -> dict:
async def ok_async(_payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "ok", "output": {}}
run = await execute_workflow_result_async(
tests/core/test_run_step_budget_async.py
+34 -31
View File
@@ -9,6 +9,7 @@ later sibling commits after the first unhandled result in reservation order.
from __future__ import annotations
import asyncio
from collections.abc import Callable
from typing import Any
import pytest
@@ -21,6 +22,7 @@ from wf_core import (
NodeDef,
NodeUse,
ReducerRef,
RunLimits,
RunStatus,
SchemaRef,
StateField,
@@ -31,11 +33,19 @@ from wf_core import (
step_workflow_async,
)
from wf_core.errors import WorkflowStepLimitExceeded
from wf_core.runtime.limits import RunLimits, remaining_step_attempts
from wf_core.runtime.limits import remaining_step_attempts
from wf_core.runtime.ops.runs import create_run_state
from wf_core.runtime.preparation import prepare_new_run, prepare_resume
async def _noop_record(payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "ok", "output": payload}
def _sync_noop_record(payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "ok", "output": payload}
def _concurrent_workflow(*, max_active: int, name: str = "async_budget") -> Workflow:
return Workflow(
name=name,
@@ -102,17 +112,19 @@ def _concurrent_workflow(*, max_active: int, name: str = "async_budget") -> Work
)
def _prepare_limited_run(
workflow: Workflow, items: list[Any], *, max_steps: int
):
run = create_run_state(workflow, {"items": items}, limits=RunLimits(max_steps=max_steps))
def _prepare_limited_run(workflow: Workflow, items: list[Any], *, max_steps: int):
run = create_run_state(
workflow, {"items": items}, limits=RunLimits(max_steps=max_steps)
)
prepare_new_run(workflow, {"items": items}, run)
index = prepare_resume(workflow, run, resume_payload=None, resume_outcome="submitted")
index = prepare_resume(
workflow, run, resume_payload=None, resume_outcome="submitted"
)
assert index is not None
return run, index
async def _wait_for(predicate, *, timeout: float = 2.0) -> None: # type: ignore[no-untyped-def]
async def _wait_for(predicate: Callable[[], bool], *, timeout: float = 2.0) -> None:
async def _poll() -> None:
while not predicate():
await asyncio.sleep(0.005)
@@ -126,10 +138,7 @@ async def test_async_batch_bounded_to_remaining_budget() -> None:
items = ["a", "b", "c", "d", "e"]
run, index = _prepare_limited_run(workflow, items, max_steps=4)
async def _noop(payload: dict[str, Any], _ctx: object) -> dict[str, Any]:
return {"outcome": "ok", "output": payload}
await step_workflow_async(workflow, run, {"record": _noop}, index=index)
await step_workflow_async(workflow, run, {"record": _noop_record}, index=index)
assert run.steps_executed == 1
assert remaining_step_attempts(run) == 3
assert run.ready_frame_ids == [f"root:each#0:{i}" for i in range(5)]
@@ -189,10 +198,7 @@ async def test_async_batch_denies_first_when_remaining_zero() -> None:
items = ["a", "b", "c"]
run, index = _prepare_limited_run(workflow, items, max_steps=1)
async def _noop(payload: dict[str, Any], _ctx: object) -> dict[str, Any]:
return {"outcome": "ok", "output": payload}
await step_workflow_async(workflow, run, {"record": _noop}, index=index)
await step_workflow_async(workflow, run, {"record": _noop_record}, index=index)
assert run.steps_executed == 1
assert remaining_step_attempts(run) == 0
@@ -224,10 +230,7 @@ async def test_async_batch_numbers_follow_queue_order_not_completion() -> None:
items = ["a", "b", "c"]
run, index = _prepare_limited_run(workflow, items, max_steps=20)
async def _noop(payload: dict[str, Any], _ctx: object) -> dict[str, Any]:
return {"outcome": "ok", "output": payload}
await step_workflow_async(workflow, run, {"record": _noop}, index=index)
await step_workflow_async(workflow, run, {"record": _noop_record}, index=index)
assert run.steps_executed == 1
allow_a = asyncio.Event()
@@ -282,10 +285,7 @@ async def test_async_batch_reservations_kept_after_failure() -> None:
items = ["a", "b", "c"]
run, index = _prepare_limited_run(workflow, items, max_steps=20)
async def _noop(payload: dict[str, Any], _ctx: object) -> dict[str, Any]:
return {"outcome": "ok", "output": payload}
await step_workflow_async(workflow, run, {"record": _noop}, index=index)
await step_workflow_async(workflow, run, {"record": _noop_record}, index=index)
base_steps = run.steps_executed
assert base_steps == 1
@@ -321,10 +321,7 @@ async def test_async_batch_settles_siblings_and_discards_later_commits() -> None
items = ["a", "b", "c"]
run, index = _prepare_limited_run(workflow, items, max_steps=20)
async def _noop(payload: dict[str, Any], _ctx: object) -> dict[str, Any]:
return {"outcome": "ok", "output": payload}
await step_workflow_async(workflow, run, {"record": _noop}, index=index)
await step_workflow_async(workflow, run, {"record": _noop_record}, index=index)
release = asyncio.Event()
started: list[str] = []
@@ -373,7 +370,9 @@ async def test_sync_async_parity_for_serial_execution() -> None:
def _serial_workflow(name: str) -> Workflow:
return Workflow(
name=name,
input_schema=SchemaRef(type="object", properties={"items": {"type": "array"}}),
input_schema=SchemaRef(
type="object", properties={"items": {"type": "array"}}
),
state_schema=StateSchema.from_field_map(
{
"items": StateField(type="array"),
@@ -383,7 +382,9 @@ async def test_sync_async_parity_for_serial_execution() -> None:
),
}
),
output_schema=SchemaRef(type="object", properties={"seen": {"type": "array"}}),
output_schema=SchemaRef(
type="object", properties={"seen": {"type": "array"}}
),
node_defs=[
NodeDef(
name="record",
@@ -426,7 +427,9 @@ async def test_sync_async_parity_for_serial_execution() -> None:
),
],
edges=[
Edge.model_validate({"from": "each", "outcome": "loop", "to": "record"}),
Edge.model_validate(
{"from": "each", "outcome": "loop", "to": "record"}
),
Edge.model_validate({"from": "record", "outcome": "ok", "to": "each"}),
Edge.model_validate({"from": "each", "outcome": "done", "to": END}),
],
@@ -438,7 +441,7 @@ async def test_sync_async_parity_for_serial_execution() -> None:
sync_run = execute_workflow(
sync_workflow,
{"items": ["a", "b"]},
{"record": lambda payload, _ctx: {"outcome": "ok", "output": payload}},
{"record": _sync_noop_record},
)
async def record(payload: dict[str, Any], _ctx: object) -> dict[str, Any]:
tests/core/test_run_step_budget_codec.py
+351
View File
@@ -0,0 +1,351 @@
"""Step budget codec and migration tests."""
from __future__ import annotations
from typing import Any, cast
import pytest
from wf_core import (
END,
Edge,
InterruptNode,
NodeDef,
NodeUse,
RunLimits,
SchemaRef,
StateSchema,
Workflow,
dump_run_state,
execute_workflow,
load_run_state,
)
from wf_core.run_codec import load_run_state_with_upgrade
from wf_core.run_state import RunState
from wf_core.runtime.limits import admit_step_attempt
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 _empty_schema() -> SchemaRef:
return SchemaRef(type="object", properties={})
def _ok_handler(payload: dict[str, Any], _context: object) -> dict[str, Any]:
return {"outcome": "ok", "output": {}}
def _interrupt_workflow() -> Workflow:
return Workflow(
name="interrupt_counts",
input_schema=_empty_schema(),
state_schema=StateSchema.from_field_map({}),
output_schema=_empty_schema(),
outcomes=["ok"],
node_defs=[
NodeDef(
name="work",
input_schema=_empty_schema(),
output_schema=_empty_schema(),
outcomes=["ok"],
)
],
start="ask",
nodes=[
InterruptNode(id="ask", type="interrupt", kind="approval"),
NodeUse(id="work", type="node", node="work"),
],
edges=[
Edge.model_validate({"from": "ask", "outcome": "submitted", "to": "work"}),
Edge.model_validate({"from": "work", "outcome": "ok", "to": END}),
],
)
def _state_dict(stored: dict[str, object]) -> dict[str, Any]:
return cast(dict[str, Any], stored["state"])
def _strip_to_v1(stored: dict[str, object]) -> dict[str, Any]:
state = dict(_state_dict(stored))
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(cast(dict[str, Any], state["frames"])).items()
}
state["frames"] = frames
return {"version": 1, "state": state}
def _strip_budget_fields(stored: dict[str, object]) -> dict[str, Any]:
state = dict(_strip_to_v1(stored)["state"])
state["trace"] = [
{key: value for key, value in entry.items() if key != "step_number"}
for entry in cast(list[dict[str, Any]], state.get("trace", []))
]
if state.get("interrupt") is not None:
state["interrupt"] = {
key: value
for key, value in cast(dict[str, Any], state["interrupt"]).items()
if key != "step_number"
}
return {"version": 1, "state": state}
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 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)
_state_dict(stored).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)
_state_dict(stored).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 cast(dict[str, Any], _state_dict(stored)["frames"])["root"]["step_number"]
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v1_smuggled_budget_fields_receive_defaults() -> None:
"""V1 envelopes predate budgets; smuggled fields must not survive."""
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=7))
stored = _strip_to_v1(dump_run_state(run))
state = cast(dict[str, Any], stored["state"])
state["limits"] = {"max_steps": 999_999}
state["steps_executed"] = 999
restored, upgraded = load_run_state_with_upgrade(stored)
assert upgraded is True
assert restored.limits.max_steps == 10_000
assert restored.steps_executed == 0
def test_v2_bool_max_steps_is_corrupt() -> None:
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
_state_dict(stored)["limits"] = {"max_steps": True}
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_str_max_steps_is_corrupt() -> None:
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
_state_dict(stored)["limits"] = {"max_steps": "10"}
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_unknown_limits_field_is_corrupt() -> None:
"""Strict v2: unknown fields inside limits are corrupt, not preserved."""
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
_state_dict(stored)["limits"] = {"max_steps": 2, "future_quota": 99}
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_negative_steps_executed_is_corrupt() -> None:
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
_state_dict(stored)["steps_executed"] = -100
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_exceeding_steps_executed_is_corrupt() -> None:
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
_state_dict(stored)["steps_executed"] = 5
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_bool_steps_executed_is_corrupt() -> None:
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
stored = dump_run_state(run)
_state_dict(stored)["steps_executed"] = True
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_incoherent_frame_step_number_is_corrupt() -> None:
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
admit_step_attempt(run, run.current_frame(), workflow.start)
stored = dump_run_state(run)
state = _state_dict(stored)
frames = cast(dict[str, Any], state["frames"])
cast(dict[str, Any], frames["root"])["step_number"] = 999
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_bool_frame_step_number_is_corrupt() -> None:
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
admit_step_attempt(run, run.current_frame(), workflow.start)
stored = dump_run_state(run)
state = _state_dict(stored)
frames = cast(dict[str, Any], state["frames"])
cast(dict[str, Any], frames["root"])["step_number"] = True
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_pending_frame_none_number_round_trips() -> None:
"""Fresh V2 runs legitimately persist unadmitted (None) frame numbers."""
workflow = _minimal_workflow()
run = create_run_state(workflow, {}, limits=RunLimits(max_steps=2))
restored, upgraded = load_run_state_with_upgrade(dump_run_state(run))
assert upgraded is False
assert restored.steps_executed == 0
assert restored.frames["root"].step_number is None
def test_v1_traces_and_interrupt_receive_none_step_numbers() -> None:
workflow = _interrupt_workflow()
run = execute_workflow(workflow, {}, {"work": _ok_handler})
stored = _strip_budget_fields(dump_run_state(run))
restored, upgraded = load_run_state_with_upgrade(stored)
assert upgraded is True
assert restored.trace[0].step_number is None
assert restored.interrupt is not None
assert restored.interrupt.step_number is None
def test_v2_missing_trace_step_number_is_corrupt() -> None:
workflow = _interrupt_workflow()
run = execute_workflow(workflow, {}, {"work": _ok_handler})
stored = dump_run_state(run)
del cast(dict[str, Any], cast(list[Any], _state_dict(stored)["trace"])[0])[
"step_number"
]
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_v2_missing_interrupt_step_number_is_corrupt() -> None:
workflow = _interrupt_workflow()
run = execute_workflow(workflow, {}, {"work": _ok_handler})
stored = dump_run_state(run)
del cast(dict[str, Any], _state_dict(stored)["interrupt"])["step_number"]
with pytest.raises(ValueError):
load_run_state_with_upgrade(stored)
def test_restored_run_state_type() -> None:
workflow = _minimal_workflow()
run = create_run_state(workflow, {})
restored = load_run_state(dump_run_state(run))
assert isinstance(restored, RunState)