feat: enforce step budget during sync dispatch
This commit is contained in:
@@ -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
|
Version-1 envelopes predate step budgets, so they receive the default
|
||||||
limit, a zeroed counter, and an unassigned number per frame exactly once
|
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 = deepcopy(state)
|
||||||
upgraded.setdefault("limits", {"max_steps": RunLimits().max_steps})
|
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():
|
for frame in frames.values():
|
||||||
if isinstance(frame, dict):
|
if isinstance(frame, dict):
|
||||||
frame.setdefault("step_number", None)
|
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
|
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.
|
"""Reject v2 payloads missing budget fields as corrupt state.
|
||||||
|
|
||||||
Unlike v1, a v2 envelope promises budget fields; a missing counter is
|
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:
|
if "limits" not in state or "steps_executed" not in state:
|
||||||
raise ValueError("invalid persisted workflow run state: missing step budget")
|
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: "
|
"invalid persisted workflow run state: "
|
||||||
f"frame {frame_id!r} is missing its step number"
|
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:
|
def _restore_root_alias(run: RunState) -> RunState:
|
||||||
|
|||||||
@@ -134,6 +134,12 @@ class TraceEntry:
|
|||||||
next_node_id: str
|
next_node_id: str
|
||||||
output: dict[str, Any] = field(default_factory=dict)
|
output: dict[str, Any] = field(default_factory=dict)
|
||||||
state_changes: 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)
|
@dataclass(slots=True)
|
||||||
@@ -179,6 +185,12 @@ class InterruptRequest:
|
|||||||
request_schema: dict[str, object] = field(default_factory=_object_schema)
|
request_schema: dict[str, object] = field(default_factory=_object_schema)
|
||||||
resume_schema: dict[str, object] = field(default_factory=_object_schema)
|
resume_schema: dict[str, object] = field(default_factory=_object_schema)
|
||||||
typed: bool = False
|
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:
|
def _default_run_limits() -> RunLimits:
|
||||||
|
|||||||
@@ -22,6 +22,12 @@ from wf_core.runtime.scheduler import (
|
|||||||
)
|
)
|
||||||
from wf_core.tokens import END
|
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(
|
def append_trace(
|
||||||
run: RunState,
|
run: RunState,
|
||||||
@@ -34,7 +40,28 @@ def append_trace(
|
|||||||
next_node_id: str,
|
next_node_id: str,
|
||||||
output: dict[str, Any],
|
output: dict[str, Any],
|
||||||
state_changes: dict[str, Any],
|
state_changes: dict[str, Any],
|
||||||
|
step_number: int | None | Any = _FROM_FRAME,
|
||||||
) -> None:
|
) -> 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(
|
run.trace.append(
|
||||||
TraceEntry(
|
TraceEntry(
|
||||||
frame_id=frame_id,
|
frame_id=frame_id,
|
||||||
@@ -45,6 +72,7 @@ def append_trace(
|
|||||||
next_node_id=next_node_id,
|
next_node_id=next_node_id,
|
||||||
output=output,
|
output=output,
|
||||||
state_changes=state_changes,
|
state_changes=state_changes,
|
||||||
|
step_number=step_number,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -57,6 +85,7 @@ def append_step_result_trace(
|
|||||||
step_type: str,
|
step_type: str,
|
||||||
next_node_id: str,
|
next_node_id: str,
|
||||||
result: StepExecutionResult,
|
result: StepExecutionResult,
|
||||||
|
step_number: int | None | Any = _FROM_FRAME,
|
||||||
) -> None:
|
) -> None:
|
||||||
append_trace(
|
append_trace(
|
||||||
run,
|
run,
|
||||||
@@ -68,6 +97,7 @@ def append_step_result_trace(
|
|||||||
next_node_id=next_node_id,
|
next_node_id=next_node_id,
|
||||||
output=result.output,
|
output=result.output,
|
||||||
state_changes=result.state_changes,
|
state_changes=result.state_changes,
|
||||||
|
step_number=step_number,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -73,6 +73,11 @@ def handle_interrupt_step(
|
|||||||
public_node_id=public_frame.node_id,
|
public_node_id=public_frame.node_id,
|
||||||
route=route,
|
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.interrupt = interrupt_request
|
||||||
run.status = RunStatus.INTERRUPTED
|
run.status = RunStatus.INTERRUPTED
|
||||||
frame.status = FrameStatus.INTERRUPTED
|
frame.status = FrameStatus.INTERRUPTED
|
||||||
|
|||||||
@@ -69,6 +69,15 @@ def resume_interrupt(
|
|||||||
resume_outcome: str,
|
resume_outcome: str,
|
||||||
reducers: Mapping[str, ReducerDefinition] | None = None,
|
reducers: Mapping[str, ReducerDefinition] | None = 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:
|
if run.interrupt is None:
|
||||||
raise WorkflowExecutionError("run is interrupted but has no interrupt request")
|
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.
|
# scope, a concurrent one buffers in the item lineage for barrier merge.
|
||||||
state_changes = commit_foreach_aware_patch(run, frame, patch)
|
state_changes = commit_foreach_aware_patch(run, frame, patch)
|
||||||
next_node_id = index.next_node_id(frame.node_id, resume_outcome)
|
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(
|
append_step_result_trace(
|
||||||
run,
|
run,
|
||||||
frame_id=frame.id,
|
frame_id=frame.id,
|
||||||
@@ -131,6 +143,7 @@ def resume_interrupt(
|
|||||||
output=resume_payload,
|
output=resume_payload,
|
||||||
state_changes=state_changes,
|
state_changes=state_changes,
|
||||||
),
|
),
|
||||||
|
step_number=activation_number,
|
||||||
)
|
)
|
||||||
run.interrupt = None
|
run.interrupt = None
|
||||||
advance_frame(run, frame, outcome=resume_outcome, next_node_id=next_node_id)
|
advance_frame(run, frame, outcome=resume_outcome, next_node_id=next_node_id)
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ from wf_core.models.steps import (
|
|||||||
from wf_core.models.workflow import Workflow
|
from wf_core.models.workflow import Workflow
|
||||||
from wf_core.run_state import ExecutionFrame, FrameStatus, RunState, StepExecutionResult
|
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.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.flow import advance_frame, append_step_result_trace
|
||||||
from wf_core.runtime.ops.foreach import step_foreach
|
from wf_core.runtime.ops.foreach import step_foreach
|
||||||
from wf_core.runtime.ops.handlers import (
|
from wf_core.runtime.ops.handlers import (
|
||||||
@@ -129,6 +130,11 @@ def step_workflow(
|
|||||||
return run
|
return run
|
||||||
index, step = prepared
|
index, step = prepared
|
||||||
frame = run.current_frame()
|
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):
|
if isinstance(step, NodeUse):
|
||||||
node_def = index.node_defs[step.node]
|
node_def = index.node_defs[step.node]
|
||||||
@@ -241,6 +247,9 @@ async def step_workflow_async(
|
|||||||
return run
|
return run
|
||||||
index, step = prepared
|
index, step = prepared
|
||||||
frame = run.current_frame()
|
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):
|
if isinstance(step, NodeUse):
|
||||||
node_def = index.node_defs[step.node]
|
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
|
Only handler awaits run concurrently. Finalization, tracing, and frame
|
||||||
advancement happen afterward in ready-queue order so `RunState` is mutated
|
advancement happen afterward in ready-queue order so `RunState` is mutated
|
||||||
deterministically.
|
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)]
|
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 = []
|
tasks = []
|
||||||
for frame in frames:
|
for frame in frames:
|
||||||
node = _node_use_for_frame(index, frame)
|
node = _node_use_for_frame(index, frame)
|
||||||
|
|||||||
@@ -4,15 +4,27 @@ import pytest
|
|||||||
|
|
||||||
from wf_core import (
|
from wf_core import (
|
||||||
END,
|
END,
|
||||||
|
ConditionNode,
|
||||||
Edge,
|
Edge,
|
||||||
|
EndNode,
|
||||||
|
ForeachNode,
|
||||||
|
InterruptNode,
|
||||||
NodeDef,
|
NodeDef,
|
||||||
NodeUse,
|
NodeUse,
|
||||||
|
PreparedSubgraph,
|
||||||
|
ReducerRef,
|
||||||
|
RunStatus,
|
||||||
SchemaRef,
|
SchemaRef,
|
||||||
|
StateField,
|
||||||
StateSchema,
|
StateSchema,
|
||||||
|
SubgraphNode,
|
||||||
Workflow,
|
Workflow,
|
||||||
WorkflowExecutionError,
|
WorkflowExecutionError,
|
||||||
dump_run_state,
|
dump_run_state,
|
||||||
|
execute_workflow,
|
||||||
load_run_state,
|
load_run_state,
|
||||||
|
resume_workflow,
|
||||||
|
step_workflow,
|
||||||
)
|
)
|
||||||
from wf_core.errors import WorkflowStepLimitExceeded
|
from wf_core.errors import WorkflowStepLimitExceeded
|
||||||
from wf_core.run_codec import load_run_state_with_upgrade
|
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):
|
with pytest.raises(ValueError):
|
||||||
load_run_state_with_upgrade(stored)
|
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)
|
||||||
|
|||||||
Reference in New Issue
Block a user