feat: enforce step budget during sync dispatch

This commit is contained in:
lda
2026-09-05 18:59:01 +07:00 Verified
parent 3fd1f70f5d
commit ae27006abb
7 changed files with 731 additions and 2 deletions
+30
View File
@@ -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,
)
+5
View File
@@ -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
+13
View File
@@ -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)