From ae27006abb4f60ef13d9d2cdd3ddeb0c9709a611 Mon Sep 17 00:00:00 2001 From: lda Date: Sat, 5 Sep 2026 18:59:01 +0700 Subject: [PATCH] feat: enforce step budget during sync dispatch --- src/wf_core/run_codec.py | 32 +- src/wf_core/run_state.py | 12 + src/wf_core/runtime/ops/flow.py | 30 ++ src/wf_core/runtime/ops/handlers.py | 5 + src/wf_core/runtime/ops/interrupts.py | 13 + src/wf_core/runtime/step.py | 16 + tests/core/test_run_step_budget.py | 625 ++++++++++++++++++++++++++ 7 files changed, 731 insertions(+), 2 deletions(-) diff --git a/src/wf_core/run_codec.py b/src/wf_core/run_codec.py index 0fe811d7..c96b1e98 100644 --- a/src/wf_core/run_codec.py +++ b/src/wf_core/run_codec.py @@ -42,7 +42,9 @@ def _inject_v1_budget_defaults(state: dict[str, Any]) -> dict[str, Any]: 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. + at load time. Pre-budget trace entries and any outstanding interrupt keep + an unassigned (``None``) number: attempts made before the upgrade are + outside the new budget. """ upgraded = deepcopy(state) upgraded.setdefault("limits", {"max_steps": RunLimits().max_steps}) @@ -52,6 +54,14 @@ def _inject_v1_budget_defaults(state: dict[str, Any]) -> dict[str, Any]: for frame in frames.values(): if isinstance(frame, dict): frame.setdefault("step_number", None) + trace = upgraded.get("trace") + if isinstance(trace, list): + for entry in trace: + if isinstance(entry, dict): + entry.setdefault("step_number", None) + interrupt = upgraded.get("interrupt") + if isinstance(interrupt, dict): + interrupt.setdefault("step_number", None) return upgraded @@ -59,7 +69,10 @@ 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. + corruption, not another request for defaults. Trace and interrupt entries + always carry the key (``None`` only for upgraded pre-budget history), so a + missing key is likewise corrupt even though the dataclass default would + otherwise mask it. """ if "limits" not in state or "steps_executed" not in state: raise ValueError("invalid persisted workflow run state: missing step budget") @@ -72,6 +85,21 @@ def _require_v2_budget_fields(state: dict[str, Any]) -> None: "invalid persisted workflow run state: " f"frame {frame_id!r} is missing its step number" ) + trace = state.get("trace") + if isinstance(trace, list): + for position, entry in enumerate(trace): + if not isinstance(entry, dict) or "step_number" not in entry: + raise ValueError( + "invalid persisted workflow run state: " + f"trace entry {position!r} is missing its step number" + ) + interrupt = state.get("interrupt") + if interrupt is not None: + if not isinstance(interrupt, dict) or "step_number" not in interrupt: + raise ValueError( + "invalid persisted workflow run state: " + "interrupt is missing its step number" + ) def _restore_root_alias(run: RunState) -> RunState: diff --git a/src/wf_core/run_state.py b/src/wf_core/run_state.py index dc8c9ed1..e5c8cdd0 100644 --- a/src/wf_core/run_state.py +++ b/src/wf_core/run_state.py @@ -134,6 +134,12 @@ class TraceEntry: next_node_id: str output: dict[str, Any] = field(default_factory=dict) state_changes: dict[str, Any] = field(default_factory=dict) + # One-based step number assigned at admission. Every trace emitted for an + # admitted step carries its frame's number; ``None`` only survives on + # pre-budget (v1) entries. Gaps are valid when an attempt fails or + # interrupts before emitting a trace, so ``RunState.steps_executed`` stays + # authoritative for enforcement and resume. + step_number: int | None = None @dataclass(slots=True) @@ -179,6 +185,12 @@ class InterruptRequest: request_schema: dict[str, object] = field(default_factory=_object_schema) resume_schema: dict[str, object] = field(default_factory=_object_schema) typed: bool = False + # Step number of the admitted interrupt activation. The initial interrupt + # trace and the later resume-completion trace both reuse this one number + # because they describe a single activation; resume never admits again. + # ``None`` only survives on pre-budget (v1) checkpoints whose activation + # predates the counter. + step_number: int | None = None def _default_run_limits() -> RunLimits: diff --git a/src/wf_core/runtime/ops/flow.py b/src/wf_core/runtime/ops/flow.py index 920e7cfe..2f7f135e 100644 --- a/src/wf_core/runtime/ops/flow.py +++ b/src/wf_core/runtime/ops/flow.py @@ -22,6 +22,12 @@ from wf_core.runtime.scheduler import ( ) from wf_core.tokens import END +# Sentinel for ``append_trace()``: copy the named frame's admitted step number +# (failing closed when unassigned). Interrupt resume passes its stored +# activation number explicitly instead, so one activation keeps one number +# across its interrupt and resume-completion entries without a second admission. +_FROM_FRAME: Any = object() + def append_trace( run: RunState, @@ -34,7 +40,28 @@ def append_trace( next_node_id: str, output: dict[str, Any], state_changes: dict[str, Any], + step_number: int | None | Any = _FROM_FRAME, ) -> None: + """Append one trace entry carrying its admitted step number. + + By default the number is copied from the named frame, which must have been + assigned by ``admit_step_attempt()`` during this dispatch; otherwise the + trace would silently describe an uncounted step, so fail closed with + ``WorkflowExecutionError``. Pass ``step_number`` explicitly only to reuse a + persisted activation number (interrupt resume-completion). + """ + if step_number is _FROM_FRAME: + frame = run.frames.get(frame_id) + if frame is None: + raise WorkflowExecutionError( + f"cannot trace step for unknown frame {frame_id!r}" + ) + if frame.step_number is None: + raise WorkflowExecutionError( + f"cannot trace unadmitted step for frame {frame_id!r} " + f"at node {node_id!r}; admit the step before dispatch" + ) + step_number = frame.step_number run.trace.append( TraceEntry( frame_id=frame_id, @@ -45,6 +72,7 @@ def append_trace( next_node_id=next_node_id, output=output, state_changes=state_changes, + step_number=step_number, ) ) @@ -57,6 +85,7 @@ def append_step_result_trace( step_type: str, next_node_id: str, result: StepExecutionResult, + step_number: int | None | Any = _FROM_FRAME, ) -> None: append_trace( run, @@ -68,6 +97,7 @@ def append_step_result_trace( next_node_id=next_node_id, output=result.output, state_changes=result.state_changes, + step_number=step_number, ) diff --git a/src/wf_core/runtime/ops/handlers.py b/src/wf_core/runtime/ops/handlers.py index ada236b4..ae0995ac 100644 --- a/src/wf_core/runtime/ops/handlers.py +++ b/src/wf_core/runtime/ops/handlers.py @@ -73,6 +73,11 @@ def handle_interrupt_step( public_node_id=public_frame.node_id, route=route, ) + # The current dispatch was already admitted by step_workflow(), so the + # frame carries this activation's number. Persist it on the request: resume + # completes the same activation without admitting again, and its + # completion trace reuses this stored number. + interrupt_request.step_number = frame.step_number run.interrupt = interrupt_request run.status = RunStatus.INTERRUPTED frame.status = FrameStatus.INTERRUPTED diff --git a/src/wf_core/runtime/ops/interrupts.py b/src/wf_core/runtime/ops/interrupts.py index d56c1e9c..8a9c017a 100644 --- a/src/wf_core/runtime/ops/interrupts.py +++ b/src/wf_core/runtime/ops/interrupts.py @@ -69,6 +69,15 @@ def resume_interrupt( resume_outcome: str, reducers: Mapping[str, ReducerDefinition] | None = None, ) -> None: + """Complete a previously admitted interrupt activation. + + Resume never admits a new step attempt: supplying the external payload + finishes the activation counted at interrupt time. The completion trace + therefore reuses the stored ``InterruptRequest.step_number`` (which may be + ``None`` only for pre-budget legacy activations) instead of the current + frame number, and execution after resume continues from the persisted + cumulative counter. + """ if run.interrupt is None: raise WorkflowExecutionError("run is interrupted but has no interrupt request") @@ -119,6 +128,9 @@ def resume_interrupt( # scope, a concurrent one buffers in the item lineage for barrier merge. state_changes = commit_foreach_aware_patch(run, frame, patch) next_node_id = index.next_node_id(frame.node_id, resume_outcome) + # Reuse the activation's stored number (not the frame's current number, and + # without admitting): both entries describe one admitted activation. + activation_number = run.interrupt.step_number append_step_result_trace( run, frame_id=frame.id, @@ -131,6 +143,7 @@ def resume_interrupt( output=resume_payload, state_changes=state_changes, ), + step_number=activation_number, ) run.interrupt = None advance_frame(run, frame, outcome=resume_outcome, next_node_id=next_node_id) diff --git a/src/wf_core/runtime/step.py b/src/wf_core/runtime/step.py index 3bd505cf..eb7d5c9a 100644 --- a/src/wf_core/runtime/step.py +++ b/src/wf_core/runtime/step.py @@ -16,6 +16,7 @@ from wf_core.models.steps import ( from wf_core.models.workflow import Workflow from wf_core.run_state import ExecutionFrame, FrameStatus, RunState, StepExecutionResult from wf_core.runtime.foreach_state import item_frame_owner, load_foreach_activation +from wf_core.runtime.limits import admit_step_attempt from wf_core.runtime.ops.flow import advance_frame, append_step_result_trace from wf_core.runtime.ops.foreach import step_foreach from wf_core.runtime.ops.handlers import ( @@ -129,6 +130,11 @@ def step_workflow( return run index, step = prepared frame = run.current_frame() + # One admission per selected step, immediately before dispatch: the counter + # increments before any handler runs, so failures and interrupts consume + # their attempt, while a denied dispatch raises before its handler runs. + # prepare_step() returning None (legacy END, interrupted) admits nothing. + admit_step_attempt(run, frame, frame.node_id) if isinstance(step, NodeUse): node_def = index.node_defs[step.node] @@ -241,6 +247,9 @@ async def step_workflow_async( return run index, step = prepared frame = run.current_frame() + # Same single-admission rule as the sync path; only the concurrent foreach + # batch below (Task 3) reserves differently. + admit_step_attempt(run, frame, frame.node_id) if isinstance(step, NodeUse): node_def = index.node_defs[step.node] @@ -312,8 +321,15 @@ async def _step_async_foreach_item_batch( Only handler awaits run concurrently. Finalization, tracing, and frame advancement happen afterward in ready-queue order so `RunState` is mutated deterministically. + + Task 3 will bound the claimed siblings by the remaining budget and pin the + reservation/failure semantics. Until then every frame in the batch is + admitted in ready-queue order before any handler starts, so each trace has + a number; a denied frame raises before any handler in the batch runs. """ frames = [first_frame, *_claim_matching_async_item_frames(run, index, first_frame)] + for frame in frames: + admit_step_attempt(run, frame, frame.node_id) tasks = [] for frame in frames: node = _node_use_for_frame(index, frame) diff --git a/tests/core/test_run_step_budget.py b/tests/core/test_run_step_budget.py index 783e9d9e..f4e79a60 100644 --- a/tests/core/test_run_step_budget.py +++ b/tests/core/test_run_step_budget.py @@ -4,15 +4,27 @@ import pytest from wf_core import ( END, + ConditionNode, Edge, + EndNode, + ForeachNode, + InterruptNode, NodeDef, NodeUse, + PreparedSubgraph, + ReducerRef, + RunStatus, SchemaRef, + StateField, StateSchema, + SubgraphNode, Workflow, WorkflowExecutionError, dump_run_state, + execute_workflow, load_run_state, + resume_workflow, + step_workflow, ) from wf_core.errors import WorkflowStepLimitExceeded from wf_core.run_codec import load_run_state_with_upgrade @@ -249,3 +261,616 @@ def test_v2_missing_frame_step_number_is_corrupt() -> None: with pytest.raises(ValueError): load_run_state_with_upgrade(stored) + + +# --- Task 2: sync dispatch and trace numbering --- + + +def _empty_schema() -> SchemaRef: + return SchemaRef(type="object", properties={}) + + +def _ok_handler(payload: dict, _context: object) -> dict: + return {"outcome": "ok", "output": {}} + + +def _trace_numbers(run) -> list: + return [entry.step_number for entry in run.trace] + + +def _chain_workflow() -> Workflow: + defs = [ + NodeDef( + name="da", + input_schema=_empty_schema(), + output_schema=_empty_schema(), + outcomes=["ok"], + ), + NodeDef( + name="db", + input_schema=_empty_schema(), + output_schema=_empty_schema(), + outcomes=["ok"], + ), + ] + return Workflow( + name="chain", + input_schema=_empty_schema(), + state_schema=StateSchema.from_field_map({}), + output_schema=_empty_schema(), + outcomes=["ok"], + node_defs=defs, + start="a", + nodes=[ + NodeUse(id="a", type="node", node="da"), + NodeUse(id="b", type="node", node="db"), + ], + edges=[ + Edge.model_validate({"from": "a", "outcome": "ok", "to": "b"}), + Edge.model_validate({"from": "b", "outcome": "ok", "to": END}), + ], + ) + + +def test_sync_node_use_counts_and_numbers_trace() -> None: + workflow = _chain_workflow() + + run = execute_workflow(workflow, {}, {"da": _ok_handler, "db": _ok_handler}) + + assert run.status == RunStatus.COMPLETED + assert run.steps_executed == 2 + assert _trace_numbers(run) == [1, 2] + assert [entry.node_id for entry in run.trace] == ["a", "b"] + assert run.steps_remaining == 10_000 - 2 + + +def test_sync_condition_counts() -> None: + workflow = Workflow( + name="condition_counts", + input_schema=SchemaRef( + type="object", properties={"count": {"type": "integer"}} + ), + state_schema=StateSchema.from_field_map({"count": StateField(type="integer")}), + output_schema=_empty_schema(), + outcomes=["ok"], + node_defs=[ + NodeDef( + name="finish", + input_schema=_empty_schema(), + output_schema=_empty_schema(), + outcomes=["done"], + ) + ], + start="pick", + nodes=[ + ConditionNode.model_validate( + { + "id": "pick", + "type": "condition", + "check": { + "op": "lt", + "left": {"path": "state.count"}, + "right": {"value": 1}, + }, + } + ), + NodeUse(id="finish", type="node", node="finish"), + ], + edges=[ + Edge.model_validate({"from": "pick", "outcome": "true", "to": "finish"}), + Edge.model_validate({"from": "pick", "outcome": "false", "to": END}), + Edge.model_validate({"from": "finish", "outcome": "done", "to": END}), + ], + ) + + run = execute_workflow(workflow, {"count": 10}, {"finish": _ok_handler}) + + assert run.status == RunStatus.COMPLETED + assert run.steps_executed == 1 + assert _trace_numbers(run) == [1] + assert run.trace[0].node_id == "pick" + assert run.trace[0].outcome == "false" + + +def _serial_foreach_workflow() -> Workflow: + foreach = ForeachNode.model_validate( + { + "id": "each", + "type": "foreach", + "over": "state.items", + "as": "item", + "mode": "serial", + } + ) + return Workflow( + name="foreach_counts", + input_schema=SchemaRef(type="object", properties={"items": {"type": "array"}}), + state_schema=StateSchema.from_field_map( + { + "items": StateField(type="array"), + "seen": StateField( + type="array", reducer=ReducerRef(name="wf.std.append") + ), + } + ), + output_schema=SchemaRef(type="object", properties={"seen": {"type": "array"}}), + node_defs=[ + NodeDef( + name="record", + input_schema=SchemaRef( + type="object", properties={"value": {}}, required=["value"] + ), + output_schema=SchemaRef( + type="object", properties={"seen": {}}, required=["seen"] + ), + outcomes=["ok"], + ) + ], + start="each", + nodes=[ + foreach, + NodeUse.model_validate( + { + "id": "work", + "type": "node", + "node": "record", + "input": [{"target": "value", "path": "context.item"}], + "output": [{"source": "seen", "target": "state.seen"}], + } + ), + ], + edges=[ + Edge.model_validate({"from": "each", "outcome": "loop", "to": "work"}), + Edge.model_validate({"from": "work", "outcome": "ok", "to": "each"}), + Edge.model_validate({"from": "each", "outcome": "done", "to": END}), + ], + ) + + +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"]}, + } + }, + ) + + assert run.status == RunStatus.COMPLETED + assert run.steps_executed == 5 + assert _trace_numbers(run) == [1, 2, 3, 4, 5] + assert [entry.node_id for entry in run.trace] == [ + "each", + "work", + "each", + "work", + "each", + ] + assert [entry.outcome for entry in run.trace] == [ + "loop", + "ok", + "loop", + "ok", + "done", + ] + + +def _subgraph_parent_workflow() -> Workflow: + node = SubgraphNode.model_validate( + { + "id": "child", + "type": "subgraph", + "workflow": "child.workflow", + "input_schema": _empty_schema(), + "output_schema": _empty_schema(), + "input": [], + "output": [], + } + ) + return Workflow( + name="subgraph_parent", + input_schema=_empty_schema(), + state_schema=StateSchema.from_field_map({}), + output_schema=_empty_schema(), + outcomes=["ok"], + start="child", + nodes=[node], + edges=[Edge.model_validate({"from": "child", "outcome": "ok", "to": END})], + ) + + +def _subgraph_child_workflow() -> Workflow: + return Workflow( + name="child.workflow", + input_schema=_empty_schema(), + state_schema=StateSchema.from_field_map({}), + output_schema=_empty_schema(), + outcomes=["ok"], + node_defs=[ + NodeDef( + name="answer", + input_schema=_empty_schema(), + output_schema=_empty_schema(), + outcomes=["ok"], + ) + ], + start="answer", + nodes=[NodeUse(id="answer", type="node", node="answer")], + edges=[Edge.model_validate({"from": "answer", "outcome": "ok", "to": END})], + ) + + +def test_sync_subgraph_entry_and_return_share_counter() -> None: + parent = _subgraph_parent_workflow() + child = _subgraph_child_workflow() + + run = execute_workflow( + parent, + {}, + {}, + subgraphs={ + "child.workflow": PreparedSubgraph( + workflow=child, registry={"answer": _ok_handler} + ) + }, + ) + + assert run.status == RunStatus.COMPLETED + # Parent entry admits once without emitting a trace (gap), the child body + # admits once, and the parent return admits once more. + assert run.steps_executed == 3 + assert _trace_numbers(run) == [2, 3] + assert run.trace[0].node_id == "answer" + assert run.trace[-1].node_id == "child" + + +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 test_sync_interrupt_activation_counts_once() -> None: + workflow = _interrupt_workflow() + + run = execute_workflow(workflow, {}, {"work": _ok_handler}) + + assert run.status == RunStatus.INTERRUPTED + assert run.steps_executed == 1 + assert _trace_numbers(run) == [1] + assert run.trace[0].outcome == "interrupt" + assert run.interrupt is not None + assert run.interrupt.step_number == 1 + + +def test_sync_interrupt_resume_reuses_activation_number() -> None: + workflow = _interrupt_workflow() + interrupted = execute_workflow(workflow, {}, {"work": _ok_handler}) + + resumed = resume_workflow( + workflow, + interrupted, + {"work": _ok_handler}, + resume_payload={}, + resume_outcome="submitted", + ) + + assert resumed.status == RunStatus.COMPLETED + # Resume completes the admitted activation without a new attempt: both the + # interrupt entry and its completion entry carry number 1. + assert resumed.steps_executed == 2 + assert _trace_numbers(resumed) == [1, 1, 2] + assert resumed.trace[1].node_id == "ask" + assert resumed.trace[1].outcome == "submitted" + assert resumed.interrupt is None + + +def test_sync_explicit_end_counts() -> None: + workflow = Workflow( + name="explicit_end", + input_schema=_empty_schema(), + state_schema=StateSchema.from_field_map({}), + output_schema=_empty_schema(), + outcomes=["done"], + node_defs=[ + NodeDef( + name="finish", + input_schema=_empty_schema(), + output_schema=_empty_schema(), + outcomes=["done"], + ) + ], + start="finish", + nodes=[ + NodeUse(id="finish", type="node", node="finish"), + EndNode(id="end", type="end", outcome="done"), + ], + edges=[Edge.model_validate({"from": "finish", "outcome": "done", "to": "end"})], + ) + + def finish(_payload: dict, _context: object) -> dict: + return {"outcome": "done", "output": {}} + + run = execute_workflow(workflow, {}, {"finish": finish}) + + assert run.status == RunStatus.COMPLETED + assert run.outcome == "done" + assert run.steps_executed == 2 + assert _trace_numbers(run) == [1, 2] + assert run.trace[-1].step_type == "end" + + +def test_sync_legacy_end_creates_no_extra_attempt() -> None: + workflow = _minimal_workflow() + + run = execute_workflow(workflow, {}, {"finish": _ok_handler}) + + assert run.status == RunStatus.COMPLETED + assert run.steps_executed == 1 + assert _trace_numbers(run) == [1] + + +def test_sync_handler_failure_consumes_attempt() -> None: + workflow = _minimal_workflow() + + def explode(_payload: dict, _context: object) -> dict: + raise ValueError("boom") + + run = create_run_state(workflow, {}) + with pytest.raises(ValueError, match="boom"): + step_workflow(workflow, run, {"finish": explode}) + + assert run.steps_executed == 1 + # The attempt failed before any normal trace entry existed (gap, not a recount). + assert run.trace == [] + + +def test_sync_handled_error_outcome_counts_once() -> None: + workflow = Workflow( + name="error_outcome", + input_schema=_empty_schema(), + state_schema=StateSchema.from_field_map({}), + output_schema=_empty_schema(), + outcomes=["ok"], + node_defs=[ + NodeDef( + name="risky", + input_schema=_empty_schema(), + output_schema=_empty_schema(), + outcomes=["ok", "error"], + ) + ], + start="work", + nodes=[NodeUse(id="work", type="node", node="risky")], + edges=[ + Edge.model_validate({"from": "work", "outcome": "ok", "to": END}), + Edge.model_validate({"from": "work", "outcome": "error", "to": END}), + ], + ) + + def fail_soft(_payload: dict, _context: object) -> dict: + return {"outcome": "error", "output": {}} + + run = execute_workflow(workflow, {}, {"risky": fail_soft}) + + assert run.status == RunStatus.COMPLETED + assert run.steps_executed == 1 + assert _trace_numbers(run) == [1] + assert run.trace[0].outcome == "error" + + +def _cyclic_workflow() -> Workflow: + defs = [ + NodeDef( + name="da", + input_schema=_empty_schema(), + output_schema=_empty_schema(), + outcomes=["ok"], + ), + NodeDef( + name="db", + input_schema=_empty_schema(), + output_schema=_empty_schema(), + outcomes=["ok"], + ), + ] + return Workflow( + name="cycle", + input_schema=_empty_schema(), + state_schema=StateSchema.from_field_map({}), + output_schema=_empty_schema(), + outcomes=["ok"], + node_defs=defs, + start="a", + nodes=[ + NodeUse(id="a", type="node", node="da"), + NodeUse(id="b", type="node", node="db"), + ], + edges=[ + Edge.model_validate({"from": "a", "outcome": "ok", "to": "b"}), + Edge.model_validate({"from": "b", "outcome": "ok", "to": "a"}), + ], + ) + + +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: + calls.append(name) + return {"outcome": "ok", "output": {}} + + return handler + + run = create_run_state(workflow, {}, limits=RunLimits(max_steps=3)) + + # Bounded manual stepping: without a budget this cycle would never stop, + # so the test itself caps iterations instead of relying on the engine loop. + with pytest.raises(WorkflowStepLimitExceeded): + for _ in range(10): + step_workflow(workflow, run, {"da": make("a"), "db": make("b")}) + + assert run.steps_executed == 3 + assert _trace_numbers(run) == [1, 2, 3] + assert calls == ["a", "b", "a"] + + +def _counting_loop_workflow() -> Workflow: + return Workflow( + name="counting_loop", + input_schema=SchemaRef( + type="object", properties={"count": {"type": "integer"}} + ), + state_schema=StateSchema.from_field_map({"count": StateField(type="integer")}), + output_schema=_empty_schema(), + outcomes=["ok"], + node_defs=[ + NodeDef( + name="bump", + input_schema=SchemaRef( + type="object", properties={"count": {"type": "integer"}} + ), + output_schema=SchemaRef( + type="object", properties={"count": {"type": "integer"}} + ), + outcomes=["ok"], + ) + ], + start="again", + nodes=[ + ConditionNode.model_validate( + { + "id": "again", + "type": "condition", + "check": { + "op": "lt", + "left": {"path": "state.count"}, + "right": {"value": 2}, + }, + } + ), + NodeUse.model_validate( + { + "id": "bump", + "type": "node", + "node": "bump", + "input": [{"target": "count", "path": "state.count"}], + "output": [{"source": "count", "target": "state.count"}], + } + ), + ], + edges=[ + Edge.model_validate({"from": "again", "outcome": "true", "to": "bump"}), + Edge.model_validate({"from": "bump", "outcome": "ok", "to": "again"}), + Edge.model_validate({"from": "again", "outcome": "false", "to": END}), + ], + ) + + +def test_sync_exiting_loop_completes_within_budget() -> None: + workflow = _counting_loop_workflow() + + def bump(payload: dict, _context: object) -> dict: + count = payload.get("count", 0) + assert isinstance(count, int) + return {"outcome": "ok", "output": {"count": count + 1}} + + run = execute_workflow(workflow, {"count": 0}, {"bump": bump}) + + assert run.status == RunStatus.COMPLETED + assert run.steps_executed == 5 + assert _trace_numbers(run) == [1, 2, 3, 4, 5] + + +def test_sync_denial_never_invokes_handler() -> None: + workflow = _chain_workflow() + b_calls: list[dict] = [] + + def b_handler(payload: dict, _context: object) -> dict: + b_calls.append(payload) + return {"outcome": "ok", "output": {}} + + run = create_run_state(workflow, {}, limits=RunLimits(max_steps=1)) + + with pytest.raises(WorkflowStepLimitExceeded): + resume_workflow(workflow, run, {"da": _ok_handler, "db": b_handler}) + + assert run.steps_executed == 1 + assert _trace_numbers(run) == [1] + 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)