From e60fb751ed21953837dfafd4269741e8ba3dc546 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 11:59:54 -0700 Subject: [PATCH 001/121] Add Reflex Workflows MVP slice: durable events, kernel, and test harness Implements the first slice of the Reflex Workflows design: durable automation on the existing Reflex programming model. Authoring contract: - rx.WorkflowConfig on a workflow-focused rx.State class (reserved __workflow__ attribute, excluded from state schemas) - @rx.event(durable=True, effect=...) with id/trigger/retry/timeout/ queue/on_failure/on_timeout options, validated at decoration time - rx.Retry with per-effect-class defaults materialized at compile (TransientWorkflowError is the explicit retryable signal) - declarative triggers rx.manual() / rx.webhook() / rx.schedule() - control returns rx.complete/fail/needs_attention and rx.after delays Runtime: - compile_workflow() validates the class contract (durable-only public handlers, no substates/backend vars/mixed scopes, resolvable hooks, stable unique ids) and produces an immutable digest-pinned definition - app.add_workflow() registers the class and detaches it from the session state tree: no per-session instances, no browser setters, no frontend event dispatch to durable handlers - WorkflowKernel executes runs against a RunStore: single-writer ordered mailbox with preallocated ordinals, fenced claims, atomic commit of state snapshot + successor slots + history, retry with exponential backoff and jitter as persisted timers, per-attempt execution timeouts, failure/timeout hooks after tombstoning, NEEDS_ATTENTION suspension for uncertain non-idempotent effects, drain-based cancellation, run deadlines, max_steps bounds, and orphan recovery with a separate infrastructure recovery budget - MemoryRunStore for tests; SqliteRunStore (stdlib, WAL) for crash-safe local persistence, including admission dedupe by request_key surviving restarts - rx.workflows.start/cancel/get_run namespace served by the app lifespan; WorkflowTestHarness runs definitions deterministically on virtual time Deliberately out of scope for this slice (per the design's MVP/Beta ledger): webhook/schedule ingress execution, the connector broker and effect evidence records, mixed-scope classes, UI run projections, operator commands beyond cancel, and multi-worker kernels. --- .../src/reflex_base/event/__init__.py | 62 +- .../src/reflex_base/utils/exceptions.py | 8 + .../reflex-base/src/reflex_base/workflow.py | 630 +++++++ pyi_hashes.json | 2 +- reflex/__init__.py | 14 + reflex/app.py | 39 + reflex/workflow/__init__.py | 87 + reflex/workflow/definition.py | 488 ++++++ reflex/workflow/kernel.py | 1025 +++++++++++ reflex/workflow/records.py | 233 +++ reflex/workflow/runtime.py | 271 +++ reflex/workflow/serde.py | 52 + reflex/workflow/store.py | 1536 +++++++++++++++++ reflex/workflow/testing.py | 186 ++ .../reflex_base/event/test_durable_event.py | 95 + tests/units/reflex_base/test_workflow.py | 235 +++ tests/units/workflow/__init__.py | 1 + tests/units/workflow/test_app.py | 131 ++ tests/units/workflow/test_definition.py | 283 +++ tests/units/workflow/test_kernel.py | 646 +++++++ tests/units/workflow/test_store.py | 326 ++++ 21 files changed, 6348 insertions(+), 2 deletions(-) create mode 100644 packages/reflex-base/src/reflex_base/workflow.py create mode 100644 reflex/workflow/__init__.py create mode 100644 reflex/workflow/definition.py create mode 100644 reflex/workflow/kernel.py create mode 100644 reflex/workflow/records.py create mode 100644 reflex/workflow/runtime.py create mode 100644 reflex/workflow/serde.py create mode 100644 reflex/workflow/store.py create mode 100644 reflex/workflow/testing.py create mode 100644 tests/units/reflex_base/event/test_durable_event.py create mode 100644 tests/units/reflex_base/test_workflow.py create mode 100644 tests/units/workflow/__init__.py create mode 100644 tests/units/workflow/test_app.py create mode 100644 tests/units/workflow/test_definition.py create mode 100644 tests/units/workflow/test_kernel.py create mode 100644 tests/units/workflow/test_store.py diff --git a/packages/reflex-base/src/reflex_base/event/__init__.py b/packages/reflex-base/src/reflex_base/event/__init__.py index 3138598e87e..49207734ead 100644 --- a/packages/reflex-base/src/reflex_base/event/__init__.py +++ b/packages/reflex-base/src/reflex_base/event/__init__.py @@ -34,7 +34,7 @@ is_typeddict, ) -from reflex_base import constants +from reflex_base import constants, workflow from reflex_base.components.field import BaseField from reflex_base.constants.compiler import CompileVars, Imports from reflex_base.registry import RegistrationContext @@ -44,6 +44,7 @@ EventFnArgMismatchError, EventHandlerArgTypeMismatchError, MissingAnnotationError, + WorkflowDefinitionError, ) from reflex_base.utils.types import ( ArgsSpec, @@ -2920,6 +2921,15 @@ def __new__( throttle: int | None = None, debounce: int | None = None, temporal: bool | None = None, + id: str | None = None, + durable: bool = False, + trigger: "workflow.Trigger | None" = None, + retry: "workflow.Retry | None" = None, + timeout: "workflow.DurationLike | None" = None, + effect: "workflow.EffectClass | None" = None, + queue: str | None = None, + on_failure: Any = None, + on_timeout: Any = None, ) -> ( "Callable[[Callable[[BASE_STATE, Unpack[P]], Any]], EventCallback[Unpack[P]]]" ): ... @@ -2947,6 +2957,15 @@ def __new__( throttle: int | None = None, debounce: int | None = None, temporal: bool | None = None, + id: str | None = None, + durable: bool = False, + trigger: "workflow.Trigger | None" = None, + retry: "workflow.Retry | None" = None, + timeout: "workflow.DurationLike | None" = None, + effect: "workflow.EffectClass | None" = None, + queue: str | None = None, + on_failure: Any = None, + on_timeout: Any = None, ) -> "EventCallback[Unpack[P]] | Callable[[Callable[[BASE_STATE, Unpack[P]], Any]], EventCallback[Unpack[P]]]": """Wrap a function to be used as an event. @@ -2958,6 +2977,15 @@ def __new__( throttle: Throttle the event handler to limit calls (in milliseconds). debounce: Debounce the event handler to delay calls (in milliseconds). temporal: Whether the event should be dropped when the backend is down. + id: Stable durable handler id; derived from the method name if omitted. + durable: Whether the handler is a durable workflow step. + trigger: How a durable root handler starts a run. + retry: Business-attempt retry policy for a durable handler. + timeout: Per-attempt execution timeout for a durable handler. + effect: Declared external-effect class; required when durable=True. + queue: Admission queue override for a durable handler. + on_failure: Same-class handler run after a durable step finally fails. + on_timeout: Same-class handler run after a durable step finally times out. Returns: The wrapped function. @@ -2965,6 +2993,28 @@ def __new__( Raises: TypeError: If background is True and the function is not a coroutine or async generator. # noqa: DAR402 """ + durable_config = workflow.build_durable_config( + durable=durable, + id=id, + trigger=trigger, + retry=retry, + timeout=timeout, + effect=effect, + queue=queue, + on_failure=on_failure, + on_timeout=on_timeout, + background=background, + has_browser_actions=any( + value is not None + for value in ( + stop_propagation, + prevent_default, + throttle, + debounce, + temporal, + ) + ), + ) def _build_event_actions(): """Build event_actions dict from decorator parameters. @@ -3004,6 +3054,16 @@ def wrapper( msg = "Background task must be async function or generator." raise TypeError(msg) setattr(func, BACKGROUND_TASK_MARKER, True) + if durable_config is not None: + if inspect.isasyncgenfunction(func) or inspect.isgeneratorfunction( + func + ): + msg = ( + "Durable event handlers commit exactly once and cannot be " + "generators; return successor events instead of yielding." + ) + raise WorkflowDefinitionError(msg) + setattr(func, workflow.DURABLE_EVENT_MARKER, durable_config) if getattr(func, "__name__", "").startswith("_"): msg = "Event handlers cannot be private." raise ValueError(msg) diff --git a/packages/reflex-base/src/reflex_base/utils/exceptions.py b/packages/reflex-base/src/reflex_base/utils/exceptions.py index bbf29239edf..e77e8c5c9f1 100644 --- a/packages/reflex-base/src/reflex_base/utils/exceptions.py +++ b/packages/reflex-base/src/reflex_base/utils/exceptions.py @@ -292,3 +292,11 @@ class UnretrievableVarValueError(ReflexError): class HybridPropertyError(ReflexError): """Raised when a hybrid property is misused while building its frontend var.""" + + +class WorkflowDefinitionError(ReflexError, ValueError): + """Raised when a workflow class or durable event declaration is invalid.""" + + +class WorkflowRuntimeError(ReflexError, RuntimeError): + """Raised when the workflow runtime is misused or reaches an invalid state.""" diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py new file mode 100644 index 00000000000..ae5832f2fcc --- /dev/null +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -0,0 +1,630 @@ +"""Authoring-time value types for Reflex Workflows. + +These types describe durable workflow metadata declared in user code via +``WorkflowConfig`` and ``@rx.event(durable=True, ...)``. They are pure, +immutable values with no runtime behavior; the workflow runtime that +interprets them lives in ``reflex.workflow``. +""" + +from __future__ import annotations + +import dataclasses +import re +from datetime import timedelta +from typing import Any, ClassVar, Final, Literal, get_args + +from reflex_base.utils.exceptions import WorkflowDefinitionError + +EffectClass = Literal["none", "read", "idempotent_write", "non_idempotent_write"] + +EFFECT_CLASSES: Final[frozenset[str]] = frozenset(get_args(EffectClass)) + +DURABLE_EVENT_MARKER: Final = "_rx_durable_event" + +DEFAULT_MAX_RECOVERIES: Final = 10 + +DurationLike = str | int | float | timedelta + +_DURATION_RE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*(ms|s|m|h|d)\s*$") + +_DURATION_UNITS: Final = { + "ms": 0.001, + "s": 1.0, + "m": 60.0, + "h": 3600.0, + "d": 86400.0, +} + +_WORKFLOW_ID_RE = re.compile(r"^[a-z][a-z0-9_]*(\.[a-z][a-z0-9_]*)*$") + +_HANDLER_ID_RE = re.compile(r"^[a-z][a-z0-9_]*$") + + +def parse_duration(value: DurationLike, *, param: str = "duration") -> float: + """Parse a duration into seconds. + + Accepts a number of seconds, a ``timedelta``, or a string with one unit + suffix: ``ms``, ``s``, ``m``, ``h``, or ``d`` (e.g. ``"30s"``, ``"2.5h"``). + + Args: + value: The duration to parse. + param: The parameter name to reference in error messages. + + Returns: + The duration in seconds. + + Raises: + WorkflowDefinitionError: If the value is not a valid non-negative duration. + """ + if isinstance(value, timedelta): + seconds = value.total_seconds() + elif isinstance(value, (int, float)) and not isinstance(value, bool): + seconds = float(value) + elif isinstance(value, str): + match = _DURATION_RE.match(value) + if match is None: + msg = ( + f"Invalid {param} {value!r}: expected a number with a unit suffix " + 'of "ms", "s", "m", "h", or "d" (e.g. "30s").' + ) + raise WorkflowDefinitionError(msg) + seconds = float(match[1]) * _DURATION_UNITS[match[2]] + else: + msg = f"Invalid {param} {value!r}: expected a str, number of seconds, or timedelta." + raise WorkflowDefinitionError(msg) + if seconds < 0: + msg = f"Invalid {param} {value!r}: duration cannot be negative." + raise WorkflowDefinitionError(msg) + return seconds + + +class TransientWorkflowError(Exception): + """Raise from a durable handler to mark a failure as safely retryable. + + Retry policies for the ``none``, ``read``, and ``idempotent_write`` effect + classes treat this exception (and its subclasses) as retryable by default. + Any other exception is a non-retryable defect unless it is listed in + ``Retry.retry_on``. + """ + + +@dataclasses.dataclass(frozen=True) +class Retry: + """Retry policy for the business attempts of a durable handler. + + Attributes: + max_attempts: Total business attempts, including the first one. + initial_delay: Backoff delay before the second attempt. + max_delay: Upper bound on the computed backoff delay. + multiplier: Exponential factor applied per additional attempt. + jitter: ``"full"`` samples uniformly in ``[0, delay]``; ``"none"`` + uses the exact computed delay. + retry_on: Exception types that consume a business attempt and retry. + do_not_retry_on: Exception types that always fail immediately. + """ + + max_attempts: int = 3 + initial_delay: DurationLike = "1s" + max_delay: DurationLike = "1m" + multiplier: float = 2.0 + jitter: Literal["full", "none"] = "full" + retry_on: tuple[type[BaseException], ...] = () + do_not_retry_on: tuple[type[BaseException], ...] = () + + def __post_init__(self): + """Validate the policy. + + Raises: + WorkflowDefinitionError: If any field is out of range or the + retry_on / do_not_retry_on sets overlap. + """ + if self.max_attempts < 1: + msg = f"Retry.max_attempts must be >= 1, got {self.max_attempts}." + raise WorkflowDefinitionError(msg) + if self.multiplier < 1.0: + msg = f"Retry.multiplier must be >= 1.0, got {self.multiplier}." + raise WorkflowDefinitionError(msg) + if self.jitter not in ("full", "none"): + msg = f'Retry.jitter must be "full" or "none", got {self.jitter!r}.' + raise WorkflowDefinitionError(msg) + initial = parse_duration(self.initial_delay, param="Retry.initial_delay") + maximum = parse_duration(self.max_delay, param="Retry.max_delay") + if maximum < initial: + msg = ( + f"Retry.max_delay ({self.max_delay!r}) must be >= " + f"Retry.initial_delay ({self.initial_delay!r})." + ) + raise WorkflowDefinitionError(msg) + overlap = [ + exc.__name__ + for exc in self.retry_on + if any(issubclass(exc, banned) for banned in self.do_not_retry_on) + ] + if overlap: + msg = ( + "Retry.retry_on and Retry.do_not_retry_on must be disjoint; " + f"{', '.join(overlap)} appears in both." + ) + raise WorkflowDefinitionError(msg) + + def delay_for_attempt(self, failed_attempts: int) -> float: + """Compute the backoff delay after a number of failed attempts, without jitter. + + Args: + failed_attempts: How many business attempts have failed so far (>= 1). + + Returns: + The clamped exponential backoff delay in seconds. + """ + initial = parse_duration(self.initial_delay, param="Retry.initial_delay") + maximum = parse_duration(self.max_delay, param="Retry.max_delay") + return min(initial * self.multiplier ** (failed_attempts - 1), maximum) + + def is_retryable(self, error: BaseException) -> bool: + """Whether an exception consumes a business attempt and may retry. + + Args: + error: The exception raised by the handler attempt. + + Returns: + True if the exception matches ``retry_on`` and not ``do_not_retry_on``. + """ + if isinstance(error, self.do_not_retry_on): + return False + return isinstance(error, self.retry_on) + + +def default_retry_for_effect(effect: str) -> Retry: + """Return the default retry policy for an effect class. + + Unknown code defects never retry by default; only ``TransientWorkflowError`` + marks a failure as safely retryable. Non-idempotent writes get exactly one + business attempt because the runtime cannot prove a retry is safe. + + Args: + effect: The declared effect class of the handler. + + Returns: + The resolved default policy. + """ + if effect == "non_idempotent_write": + return Retry(max_attempts=1, retry_on=()) + return Retry(max_attempts=3, retry_on=(TransientWorkflowError,)) + + +@dataclasses.dataclass(frozen=True) +class Trigger: + """Base class for declarative workflow trigger specifications.""" + + kind: ClassVar[str] = "" + + +@dataclasses.dataclass(frozen=True) +class ManualTrigger(Trigger): + """Marks a root handler startable via ``rx.workflows.start(...)``.""" + + kind: ClassVar[str] = "manual" + + +@dataclasses.dataclass(frozen=True) +class WebhookTrigger(Trigger): + """Marks a root handler started by an authenticated provider webhook. + + Attributes: + topic: Stable provider event topic, e.g. ``"stripe.payment_succeeded"``. + model: Optional typed payload model the raw payload is validated into. + verify: Provider signature verifier supplied by a connection binding. + dedupe_by: Payload field used as the ingress deduplication key. + """ + + kind: ClassVar[str] = "webhook" + + topic: str + model: type | None = None + verify: Any = None + dedupe_by: str | None = None + + def __post_init__(self): + """Validate the topic. + + Raises: + WorkflowDefinitionError: If the topic is empty. + """ + if not self.topic: + msg = "webhook trigger requires a non-empty topic." + raise WorkflowDefinitionError(msg) + + +@dataclasses.dataclass(frozen=True) +class ScheduleTrigger(Trigger): + """Marks a root handler started on a cron schedule. + + Attributes: + cron: A five-field cron expression (minute hour day month weekday). + """ + + kind: ClassVar[str] = "schedule" + + cron: str + + def __post_init__(self): + """Validate the cron expression shape. + + Raises: + WorkflowDefinitionError: If the expression does not have five fields. + """ + if len(self.cron.split()) != 5: + msg = ( + f"Invalid cron expression {self.cron!r}: expected five fields " + "(minute hour day month weekday)." + ) + raise WorkflowDefinitionError(msg) + + +def manual() -> ManualTrigger: + """Create a manual trigger for a workflow root handler. + + Returns: + The trigger specification. + """ + return ManualTrigger() + + +def webhook( + topic: str, + *, + model: type | None = None, + verify: Any = None, + dedupe_by: str | None = None, +) -> WebhookTrigger: + """Create a webhook trigger for a workflow root handler. + + Args: + topic: Stable provider event topic. + model: Optional typed payload model. + verify: Provider signature verifier from a connection binding. + dedupe_by: Payload field used as the ingress deduplication key. + + Returns: + The trigger specification. + """ + return WebhookTrigger(topic=topic, model=model, verify=verify, dedupe_by=dedupe_by) + + +def schedule(cron: str) -> ScheduleTrigger: + """Create a cron schedule trigger for a workflow root handler. + + Args: + cron: A five-field cron expression. + + Returns: + The trigger specification. + """ + return ScheduleTrigger(cron=cron) + + +@dataclasses.dataclass(frozen=True) +class WorkflowConfig: + """Immutable identity and policy metadata for a workflow class. + + Assigned to the reserved ``__workflow__`` attribute of a workflow-focused + ``rx.State`` class. It is excluded from State schemas and included in the + workflow definition digest. + + Attributes: + id: Stable dotted workflow identity, e.g. ``"billing.reconcile"``. + display_name: Human-readable name for operator surfaces. + run_timeout: Deadline for a whole run, measured from admission. + default_queue: Default admission queue name for the workflow's handlers. + max_steps: Upper bound on scheduled steps per run. + allow_mixed_scopes: Acknowledge a mixed session/run class (advanced). + mixed_scope_reason: Required justification when mixed scopes are allowed. + """ + + id: str + display_name: str | None = None + run_timeout: DurationLike | None = None + default_queue: str | None = None + max_steps: int = 10_000 + allow_mixed_scopes: bool = False + mixed_scope_reason: str = "" + + def __post_init__(self): + """Validate the configuration. + + Raises: + WorkflowDefinitionError: If the id, run_timeout, max_steps, or + mixed-scope acknowledgement is invalid. + """ + if not isinstance(self.id, str) or not _WORKFLOW_ID_RE.match(self.id): + msg = ( + f"Invalid WorkflowConfig.id {self.id!r}: expected lowercase " + 'dotted segments like "billing.reconcile".' + ) + raise WorkflowDefinitionError(msg) + if self.run_timeout is not None: + parse_duration(self.run_timeout, param="WorkflowConfig.run_timeout") + if self.max_steps < 1: + msg = f"WorkflowConfig.max_steps must be >= 1, got {self.max_steps}." + raise WorkflowDefinitionError(msg) + if self.allow_mixed_scopes and not self.mixed_scope_reason: + msg = ( + "WorkflowConfig(allow_mixed_scopes=True) requires a non-empty " + "mixed_scope_reason." + ) + raise WorkflowDefinitionError(msg) + if self.mixed_scope_reason and not self.allow_mixed_scopes: + msg = "mixed_scope_reason is only valid with allow_mixed_scopes=True." + raise WorkflowDefinitionError(msg) + + +@dataclasses.dataclass(frozen=True) +class DurableEventConfig: + """Validated durable metadata attached to a handler by ``@rx.event``. + + Attributes: + id: Explicit stable handler id, or None to derive from the method name. + trigger: How the handler may start a run, or None for internal handlers. + retry: Explicit retry policy, or None to use the effect-class default. + timeout: Per-attempt (start-to-close) execution timeout in seconds. + effect: Declared external-effect class. + queue: Admission queue override. + on_failure: Same-class handler name run after final failure. + on_timeout: Same-class handler name run after final timeout. + """ + + effect: str + id: str | None = None + trigger: Trigger | None = None + retry: Retry | None = None + timeout: float | None = None + queue: str | None = None + on_failure: str | None = None + on_timeout: str | None = None + + +def get_durable_config(fn: Any) -> DurableEventConfig | None: + """Get the durable metadata attached to a handler function, if any. + + Args: + fn: The undecorated handler function. + + Returns: + The attached config, or None for ordinary session handlers. + """ + return getattr(fn, DURABLE_EVENT_MARKER, None) + + +def _hook_name(value: Any, *, param: str) -> str | None: + """Normalize a lifecycle hook reference to a handler name. + + Args: + value: A handler name, or a function/handler with a ``__name__``. + param: The parameter name to reference in error messages. + + Returns: + The handler name, or None if no hook was given. + + Raises: + WorkflowDefinitionError: If the reference is not a name or named callable. + """ + if value is None: + return None + if isinstance(value, str): + if not value: + msg = f"{param} cannot be an empty string." + raise WorkflowDefinitionError(msg) + return value + fn = getattr(value, "fn", value) + name = getattr(fn, "__name__", None) + if name is None: + msg = f"{param} must be a handler name or same-class event handler, got {value!r}." + raise WorkflowDefinitionError(msg) + return name + + +def build_durable_config( + *, + durable: bool, + id: str | None, + trigger: Any, + retry: Any, + timeout: Any, + effect: Any, + queue: str | None, + on_failure: Any, + on_timeout: Any, + background: bool | None, + has_browser_actions: bool, +) -> DurableEventConfig | None: + """Validate ``@rx.event`` durable keyword arguments at decoration time. + + Args: + durable: Whether the handler was declared durable. + id: Explicit stable handler id. + trigger: Trigger specification. + retry: Retry policy. + timeout: Per-attempt execution timeout. + effect: Declared effect class. + queue: Admission queue override. + on_failure: Failure hook reference. + on_timeout: Timeout hook reference. + background: The decorator's ``background`` flag. + has_browser_actions: Whether browser-only event actions were also set. + + Returns: + The validated config for durable handlers, or None for session handlers. + + Raises: + WorkflowDefinitionError: If the combination of arguments is invalid. + """ + if not durable: + offending = next( + ( + name + for name, value in ( + ("id", id), + ("trigger", trigger), + ("retry", retry), + ("timeout", timeout), + ("effect", effect), + ("queue", queue), + ("on_failure", on_failure), + ("on_timeout", on_timeout), + ) + if value is not None + ), + None, + ) + if offending is not None: + msg = ( + f"@rx.event({offending}=...) is a durable workflow option and " + "requires durable=True." + ) + raise WorkflowDefinitionError(msg) + return None + if background: + msg = "@rx.event(durable=True) is mutually exclusive with background=True." + raise WorkflowDefinitionError(msg) + if has_browser_actions: + msg = ( + "@rx.event(durable=True) cannot use browser event actions " + "(stop_propagation, prevent_default, throttle, debounce, temporal)." + ) + raise WorkflowDefinitionError(msg) + if effect not in EFFECT_CLASSES: + msg = ( + f"@rx.event(durable=True) requires effect= one of " + f'{sorted(EFFECT_CLASSES)}, got {effect!r}. Use effect="none" for ' + "pure orchestration steps." + ) + raise WorkflowDefinitionError(msg) + if id is not None and (not isinstance(id, str) or not _HANDLER_ID_RE.match(id)): + msg = ( + f"Invalid @rx.event id {id!r}: expected a lowercase identifier " + 'like "sync_contact".' + ) + raise WorkflowDefinitionError(msg) + if trigger is not None and not isinstance(trigger, Trigger): + msg = ( + f"@rx.event trigger must be rx.manual(), rx.webhook(...), or " + f"rx.schedule(...), got {trigger!r}." + ) + raise WorkflowDefinitionError(msg) + if retry is not None: + if not isinstance(retry, Retry): + msg = f"@rx.event retry must be an rx.Retry, got {retry!r}." + raise WorkflowDefinitionError(msg) + if effect == "non_idempotent_write" and retry.max_attempts > 1: + msg = ( + 'effect="non_idempotent_write" allows only one business attempt; ' + "the runtime cannot prove a retry is safe. Use " + 'effect="idempotent_write" or max_attempts=1.' + ) + raise WorkflowDefinitionError(msg) + timeout_seconds = ( + parse_duration(timeout, param="timeout") if timeout is not None else None + ) + return DurableEventConfig( + effect=effect, + id=id, + trigger=trigger, + retry=retry, + timeout=timeout_seconds, + queue=queue, + on_failure=_hook_name(on_failure, param="on_failure"), + on_timeout=_hook_name(on_timeout, param="on_timeout"), + ) + + +@dataclasses.dataclass(frozen=True) +class CompleteRun: + """Control return that completes the run with an optional result.""" + + result: Any = None + + +@dataclasses.dataclass(frozen=True) +class FailRun: + """Control return that fails the run with a reason.""" + + reason: str + details: dict[str, Any] | None = None + + +@dataclasses.dataclass(frozen=True) +class NeedsAttention: + """Control return that suspends the run for operator resolution.""" + + reason: str + details: dict[str, Any] | None = None + + +@dataclasses.dataclass(frozen=True) +class After: + """Control return that schedules a successor after a durable delay. + + Attributes: + delay: How long to wait before the successor becomes runnable. + target: The successor handler reference or event spec. + """ + + delay: DurationLike + target: Any + + def __post_init__(self): + """Validate the delay eagerly so authoring errors surface in place.""" + parse_duration(self.delay, param="after() delay") + + +def complete(result: Any = None) -> CompleteRun: + """Complete the run, discarding any remaining scheduled work. + + Args: + result: Optional JSON-serializable run result. + + Returns: + The control return value. + """ + return CompleteRun(result=result) + + +def fail(reason: str, details: dict[str, Any] | None = None) -> FailRun: + """Fail the run with an explicit business reason. + + Args: + reason: Short stable failure reason. + details: Optional JSON-serializable diagnostic details. + + Returns: + The control return value. + """ + return FailRun(reason=reason, details=details) + + +def needs_attention( + reason: str, details: dict[str, Any] | None = None +) -> NeedsAttention: + """Suspend the run for operator resolution. + + Args: + reason: Short stable suspension reason. + details: Optional JSON-serializable diagnostic details. + + Returns: + The control return value. + """ + return NeedsAttention(reason=reason, details=details) + + +def after(delay: DurationLike, target: Any) -> After: + """Schedule a successor handler after a durable delay. + + Args: + delay: How long to wait, e.g. ``"2d"``. + target: A same-class durable handler reference or event spec. + + Returns: + The control return value. + """ + return After(delay=delay, target=target) diff --git a/pyi_hashes.json b/pyi_hashes.json index f83a2da704d..03ddc3dc696 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -118,7 +118,7 @@ "packages/reflex-components-recharts/src/reflex_components_recharts/polar.pyi": "99ebcfc07868061bdc3c2010d85a153f", "packages/reflex-components-recharts/src/reflex_components_recharts/recharts.pyi": "4f6c26f8c76543cc41e2b9dc400ece8a", "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "f170ac685b6ba5892370166c80684db3", - "reflex/__init__.pyi": "56385a4f0d9431eb0056dbc5553a58f9", + "reflex/__init__.pyi": "577f2307b3ba7fcd1aeda6621056c34c", "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388", "reflex/experimental/memo.pyi": "bc8b48357bef580e70a5881b65d3d3f7" } diff --git a/reflex/__init__.py b/reflex/__init__.py index 6b23b495d3e..19d3dda7879 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -236,6 +236,19 @@ "utils.misc": ["run_in_thread"], "utils.serializers": ["serializer"], "vars": ["Var", "field", "Field", "RestProp", "EMPTY_VAR_STR", "EMPTY_VAR_INT"], + "workflow": [ + "WorkflowConfig", + "Retry", + "TransientWorkflowError", + "manual", + "webhook", + "schedule", + "after", + "complete", + "fail", + "needs_attention", + "workflows", + ], } _SUBMODULES: set[str] = { @@ -252,6 +265,7 @@ "config", "compiler", "plugins", + "workflow", } _SUBMOD_ATTRS: lazy_loader.SubmodAttrsType = _MAPPING _EXTRA_MAPPINGS: dict[str, str] = _COMPONENT_NAME_TO_PATH diff --git a/reflex/app.py b/reflex/app.py index 258410fe9f7..7c210ee8517 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -98,6 +98,8 @@ ) from reflex.utils.misc import run_in_thread from reflex.utils.token_manager import RedisTokenManager, TokenManager +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import RunStore if sys.version_info < (3, 13): from typing_extensions import deprecated @@ -450,6 +452,13 @@ class App(MiddlewareMixin, LifespanMixin): # The processor queue for handling events. _event_processor: EventProcessor | None = None + # Durable run store for registered workflows; defaults to a local SQLite + # store created when the app starts. + workflow_store: RunStore | None = None + + # The workflow runtime owning registered definitions and the kernel. + _workflow_runtime: WorkflowRuntime | None = None + # Store the RegistrationContext to apply inside the ASGI callable task. _registration_context: RegistrationContext = dataclasses.field( default_factory=RegistrationContext.ensure_context @@ -928,6 +937,36 @@ def _page_route_key( return format.format_route(format.to_kebab_case(component.__name__)) return None + def add_workflow(self, workflow_cls: type[BaseState]) -> None: + """Register a durable workflow class with the app. + + Registration classifies the class as workflow-focused: its fields + become run-scoped, it is detached from the session state tree, and its + durable handlers become executable by the app's workflow kernel. + Registration does not publish or activate anything by itself. + + Args: + workflow_cls: A workflow-focused ``rx.State`` subclass with a + ``__workflow__ = rx.WorkflowConfig(id=...)`` declaration. + """ + if self._workflow_runtime is None: + self._workflow_runtime = WorkflowRuntime(self.workflow_store) + self.register_lifespan_task(self._run_workflow_runtime) + self._workflow_runtime.register(workflow_cls) + + @contextlib.asynccontextmanager + async def _run_workflow_runtime(self) -> AsyncIterator[None]: + """Run the workflow runtime for the duration of the app lifespan. + + Yields: + Nothing; the runtime processes runs while the app serves. + """ + if self._workflow_runtime is None: + yield + return + async with self._workflow_runtime.running(): + yield + def add_page( self, component: Component | ComponentCallable | None = None, diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py new file mode 100644 index 00000000000..de4ec992b5e --- /dev/null +++ b/reflex/workflow/__init__.py @@ -0,0 +1,87 @@ +"""Reflex Workflows: durable automation on the Reflex programming model. + +Workflow data is declared on a workflow-focused ``rx.State`` class, durable +handlers use ``@rx.event(durable=True, effect=...)``, ordinary Python handles +validation and branching, and returned events define persisted transitions. +Register classes with ``app.add_workflow(...)``; start runs with +``rx.workflows.start(...)``. +""" + +from reflex_base.workflow import ( + DurableEventConfig, + EffectClass, + ManualTrigger, + Retry, + ScheduleTrigger, + TransientWorkflowError, + Trigger, + WebhookTrigger, + WorkflowConfig, + after, + complete, + fail, + manual, + needs_attention, + parse_duration, + schedule, + webhook, +) + +from reflex.workflow.definition import ( + HandlerDefinition, + WorkflowDefinition, + compile_workflow, +) +from reflex.workflow.kernel import WorkflowKernel +from reflex.workflow.records import ( + HistoryEvent, + HistoryEventType, + RunRecord, + RunSnapshot, + RunStatus, + StartResult, + StepRecord, + StepStatus, +) +from reflex.workflow.runtime import WorkflowRuntime, get_runtime, workflows +from reflex.workflow.store import MemoryRunStore, RunStore, SqliteRunStore +from reflex.workflow.testing import WorkflowTestHarness + +__all__ = [ + "DurableEventConfig", + "EffectClass", + "HandlerDefinition", + "HistoryEvent", + "HistoryEventType", + "ManualTrigger", + "MemoryRunStore", + "Retry", + "RunRecord", + "RunSnapshot", + "RunStatus", + "RunStore", + "ScheduleTrigger", + "SqliteRunStore", + "StartResult", + "StepRecord", + "StepStatus", + "TransientWorkflowError", + "Trigger", + "WebhookTrigger", + "WorkflowConfig", + "WorkflowDefinition", + "WorkflowKernel", + "WorkflowRuntime", + "WorkflowTestHarness", + "after", + "compile_workflow", + "complete", + "fail", + "get_runtime", + "manual", + "needs_attention", + "parse_duration", + "schedule", + "webhook", + "workflows", +] diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py new file mode 100644 index 00000000000..881253484f3 --- /dev/null +++ b/reflex/workflow/definition.py @@ -0,0 +1,488 @@ +"""Compile a registered workflow class into an immutable definition. + +The compiler validates the workflow authoring contract and produces the +versioned structure the kernel executes: stable handler identities, resolved +retry/timeout policies, the run-state field schema, and a content digest that +pins runs to the exact definition they were admitted under. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import inspect +import json +from typing import TYPE_CHECKING, Any, get_type_hints + +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import ( + DurableEventConfig, + Retry, + TransientWorkflowError, + Trigger, + WorkflowConfig, + default_retry_for_effect, + get_durable_config, + parse_duration, +) + +from reflex.workflow.serde import to_run_data + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + + from reflex.state import BaseState + +RESERVED_HANDLER_NAMES = frozenset(("setvar",)) + + +@dataclasses.dataclass(frozen=True, slots=True) +class FieldSchema: + """Schema of one run-scoped state field. + + Attributes: + name: The field name. + annotated_type: The declared annotation, used to coerce loaded values. + default: The JSON-compatible default value. + """ + + name: str + annotated_type: Any + default: Any + + +@dataclasses.dataclass(frozen=True, slots=True) +class HandlerDefinition: + """Compiled definition of one durable handler. + + Attributes: + id: Stable handler identity. + name: Python method name on the workflow class. + fn: The undecorated handler function. + effect: Declared effect class. + trigger: Root trigger specification, or None for internal handlers. + retry: Fully resolved retry policy. + timeout: Per-attempt execution timeout in seconds, or None. + queue: Resolved admission queue name, or None. + on_failure: Handler id run after final failure, or None. + on_timeout: Handler id run after final timeout, or None. + params: Payload parameter names, excluding ``self``. + type_hints: Resolved type hints for payload coercion. + is_async: Whether the handler is a coroutine function. + """ + + id: str + name: str + fn: Callable + effect: str + trigger: Trigger | None + retry: Retry + timeout: float | None + queue: str | None + on_failure: str | None + on_timeout: str | None + params: tuple[str, ...] + type_hints: Mapping[str, Any] + is_async: bool + + +@dataclasses.dataclass(frozen=True, slots=True) +class WorkflowDefinition: + """Immutable compiled definition of one workflow class. + + Attributes: + workflow_id: Stable workflow identity from ``WorkflowConfig.id``. + state_cls: The registered workflow-focused state class. + config: The authoring configuration. + digest: Content digest pinning runs to this exact definition. + run_timeout: Whole-run deadline in seconds, or None. + max_steps: Upper bound on scheduled steps per run. + handlers: Handler definitions keyed by stable handler id. + handler_ids_by_name: Map from Python method name to handler id. + roots: Handler ids that declare a trigger and may start runs. + fields: Run-state field schemas in declaration order. + """ + + workflow_id: str + state_cls: type[BaseState] + config: WorkflowConfig + digest: str + run_timeout: float | None + max_steps: int + handlers: Mapping[str, HandlerDefinition] + handler_ids_by_name: Mapping[str, str] + roots: tuple[str, ...] + fields: tuple[FieldSchema, ...] + + +def _error(workflow_cls: type, msg: str) -> WorkflowDefinitionError: + """Build a definition error prefixed with the workflow class name. + + Args: + workflow_cls: The class being compiled. + msg: The error detail. + + Returns: + The exception to raise. + """ + return WorkflowDefinitionError(f"Workflow {workflow_cls.__name__}: {msg}") + + +def _validate_class_shape(workflow_cls: type[BaseState]) -> WorkflowConfig: + """Validate the class-level workflow contract. + + Args: + workflow_cls: The candidate workflow class. + + Returns: + The class's workflow configuration. + + Raises: + WorkflowDefinitionError: If the class violates the authoring contract. + """ + from reflex.state import BaseState, ComponentState, State + + if not (isinstance(workflow_cls, type) and issubclass(workflow_cls, BaseState)): + msg = f"add_workflow() expects an rx.State subclass, got {workflow_cls!r}." + raise WorkflowDefinitionError(msg) + config = workflow_cls.__dict__.get("__workflow__") + if config is None: + raise _error( + workflow_cls, + "missing __workflow__ = rx.WorkflowConfig(id=...) on the class body.", + ) + if not isinstance(config, WorkflowConfig): + raise _error( + workflow_cls, + f"__workflow__ must be an rx.WorkflowConfig, got {type(config).__name__}.", + ) + if issubclass(workflow_cls, ComponentState): + raise _error(workflow_cls, "ComponentState classes cannot be workflows.") + if workflow_cls._mixin: + raise _error(workflow_cls, "state mixins cannot be workflows.") + if workflow_cls.get_parent_state() is not State: + raise _error( + workflow_cls, + "workflow classes must subclass rx.State directly; nested substates " + "are not supported.", + ) + if workflow_cls.get_substates(): + names = ", ".join(sorted(s.__name__ for s in workflow_cls.get_substates())) + raise _error( + workflow_cls, + f"workflow classes cannot have substates (found {names}).", + ) + if config.allow_mixed_scopes: + raise _error( + workflow_cls, + "mixed-scope workflow classes are not supported yet; move session " + "fields and handlers to a separate unregistered rx.State class.", + ) + own_backend_vars = [ + name + for name in workflow_cls.backend_vars + if name not in workflow_cls.inherited_backend_vars + ] + if own_backend_vars: + raise _error( + workflow_cls, + "backend-only fields are unavailable to durable handlers; declare " + f"ordinary typed fields instead of: {', '.join(sorted(own_backend_vars))}.", + ) + return config + + +def _compile_fields(workflow_cls: type[BaseState]) -> tuple[FieldSchema, ...]: + """Build and validate the run-state field schema. + + Args: + workflow_cls: The workflow class. + + Returns: + Field schemas in declaration order. + + Raises: + WorkflowDefinitionError: If a field default is not serializable. + """ + fields = [] + class_fields = workflow_cls.get_fields() + for name in workflow_cls.base_vars: + field = class_fields[name] + default = field.default_value() + try: + default_json = to_run_data(default) + except (TypeError, ValueError) as err: + raise _error( + workflow_cls, + f"field {name!r} default is not serializable run data: {err}", + ) from None + fields.append( + FieldSchema( + name=name, + annotated_type=field.annotated_type, + default=default_json, + ) + ) + return tuple(fields) + + +def _resolve_retry(retry: Retry | None, effect: str) -> Retry: + """Materialize the effective retry policy for a handler. + + An explicit policy that does not name retryable exception types inherits + the effect class's default retryable set, so ``Retry(max_attempts=5)`` + keeps its meaning without restating the transient-error contract. + + Args: + retry: The explicit policy, if the handler declared one. + effect: The handler's effect class. + + Returns: + The fully resolved policy. + """ + if retry is None: + return default_retry_for_effect(effect) + if not retry.retry_on and effect != "non_idempotent_write": + return dataclasses.replace(retry, retry_on=(TransientWorkflowError,)) + return retry + + +def _compile_handlers( + workflow_cls: type[BaseState], config: WorkflowConfig +) -> dict[str, HandlerDefinition]: + """Compile every public handler on the class into a handler definition. + + Args: + workflow_cls: The workflow class. + config: The class's workflow configuration. + + Returns: + Handler definitions keyed by stable handler id. + + Raises: + WorkflowDefinitionError: If a handler violates the durable contract. + """ + from reflex.state import State + + handlers: dict[str, HandlerDefinition] = {} + inherited = State.event_handlers + for name, handler in workflow_cls.event_handlers.items(): + if name in RESERVED_HANDLER_NAMES or inherited.get(name) is handler: + continue + durable: DurableEventConfig | None = get_durable_config(handler.fn) + if durable is None: + raise _error( + workflow_cls, + f"handler {name!r} is not durable. Every public handler on a " + "workflow class must declare @rx.event(durable=True, effect=...); " + "move session handlers to a separate unregistered rx.State class.", + ) + handler_id = durable.id or name + if handler_id in handlers: + raise _error( + workflow_cls, + f"duplicate handler id {handler_id!r} on {name!r} and " + f"{handlers[handler_id].name!r}; stable ids must be unique.", + ) + fn = handler.fn + params = tuple(inspect.signature(fn).parameters)[1:] + handlers[handler_id] = HandlerDefinition( + id=handler_id, + name=name, + fn=fn, + effect=durable.effect, + trigger=durable.trigger, + retry=_resolve_retry(durable.retry, durable.effect), + timeout=durable.timeout, + queue=durable.queue or config.default_queue, + on_failure=durable.on_failure, + on_timeout=durable.on_timeout, + params=params, + type_hints=get_type_hints(fn), + is_async=inspect.iscoroutinefunction(fn), + ) + return handlers + + +def _resolve_hooks( + workflow_cls: type[BaseState], handlers: dict[str, HandlerDefinition] +) -> dict[str, HandlerDefinition]: + """Resolve lifecycle hook names to stable handler ids. + + Args: + workflow_cls: The workflow class. + handlers: Compiled handler definitions keyed by id. + + Returns: + Handler definitions with hooks rewritten to handler ids. + + Raises: + WorkflowDefinitionError: If a hook does not resolve to another durable + handler on the same class. + """ + ids_by_name = {defn.name: defn.id for defn in handlers.values()} + resolved = {} + for handler_id, defn in handlers.items(): + hook_ids = {} + for param in ("on_failure", "on_timeout"): + hook_name = getattr(defn, param) + if hook_name is None: + hook_ids[param] = None + continue + hook_id = ids_by_name.get( + hook_name, hook_name if hook_name in handlers else None + ) + if hook_id is None: + raise _error( + workflow_cls, + f"handler {defn.name!r} {param}={hook_name!r} does not match " + "a durable handler on the same class.", + ) + if hook_id == handler_id: + raise _error( + workflow_cls, + f"handler {defn.name!r} cannot use itself as {param}.", + ) + if handlers[hook_id].params: + raise _error( + workflow_cls, + f"{param} handler {handlers[hook_id].name!r} cannot take " + "payload arguments; it reads context from run state.", + ) + hook_ids[param] = hook_id + resolved[handler_id] = dataclasses.replace(defn, **hook_ids) + return resolved + + +def _canonical_retry(retry: Retry) -> dict[str, Any]: + """Canonicalize a retry policy for the definition digest. + + Args: + retry: The resolved policy. + + Returns: + A JSON-compatible representation. + """ + return { + "max_attempts": retry.max_attempts, + "initial_delay": parse_duration(retry.initial_delay), + "max_delay": parse_duration(retry.max_delay), + "multiplier": retry.multiplier, + "jitter": retry.jitter, + "retry_on": sorted(exc.__qualname__ for exc in retry.retry_on), + "do_not_retry_on": sorted(exc.__qualname__ for exc in retry.do_not_retry_on), + } + + +def _canonical_trigger(trigger: Trigger | None) -> dict[str, Any] | None: + """Canonicalize a trigger for the definition digest. + + Args: + trigger: The trigger specification, if any. + + Returns: + A JSON-compatible representation, or None. + """ + if trigger is None: + return None + canonical: dict[str, Any] = {"kind": trigger.kind} + for attr in ("topic", "cron", "dedupe_by"): + value = getattr(trigger, attr, None) + if value is not None: + canonical[attr] = value + model = getattr(trigger, "model", None) + if model is not None: + canonical["model"] = model.__qualname__ + return canonical + + +def _compute_digest( + config: WorkflowConfig, + handlers: Mapping[str, HandlerDefinition], + fields: tuple[FieldSchema, ...], +) -> str: + """Compute the content digest of a compiled definition. + + Args: + config: The workflow configuration. + handlers: Compiled handler definitions. + fields: The run-state field schema. + + Returns: + A hex sha256 digest over the canonical definition structure. + """ + canonical = { + "workflow_id": config.id, + "run_timeout": ( + parse_duration(config.run_timeout) + if config.run_timeout is not None + else None + ), + "max_steps": config.max_steps, + "default_queue": config.default_queue, + "handlers": [ + { + "id": defn.id, + "effect": defn.effect, + "trigger": _canonical_trigger(defn.trigger), + "retry": _canonical_retry(defn.retry), + "timeout": defn.timeout, + "queue": defn.queue, + "on_failure": defn.on_failure, + "on_timeout": defn.on_timeout, + "params": list(defn.params), + } + for defn in sorted(handlers.values(), key=lambda d: d.id) + ], + "fields": [ + {"name": f.name, "type": str(f.annotated_type), "default": f.default} + for f in fields + ], + } + payload = json.dumps(canonical, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode()).hexdigest() + + +def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: + """Compile a workflow class into an immutable definition. + + Args: + workflow_cls: A workflow-focused ``rx.State`` subclass with a + ``__workflow__`` configuration. + + Returns: + The compiled definition. + + Raises: + WorkflowDefinitionError: If the class violates the authoring contract. + """ + config = _validate_class_shape(workflow_cls) + fields = _compile_fields(workflow_cls) + handlers = _resolve_hooks(workflow_cls, _compile_handlers(workflow_cls, config)) + roots = tuple( + defn.id + for defn in sorted(handlers.values(), key=lambda d: d.id) + if defn.trigger is not None + ) + if not roots: + raise _error( + workflow_cls, + "no root handler declares a trigger; add trigger=rx.manual() (or " + "rx.webhook/rx.schedule) to at least one durable handler.", + ) + return WorkflowDefinition( + workflow_id=config.id, + state_cls=workflow_cls, + config=config, + digest=_compute_digest(config, handlers, fields), + run_timeout=( + parse_duration(config.run_timeout) + if config.run_timeout is not None + else None + ), + max_steps=config.max_steps, + handlers=handlers, + handler_ids_by_name={defn.name: defn.id for defn in handlers.values()}, + roots=roots, + fields=fields, + ) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py new file mode 100644 index 00000000000..5ab91cc750c --- /dev/null +++ b/reflex/workflow/kernel.py @@ -0,0 +1,1025 @@ +"""The in-process durable workflow kernel. + +The kernel admits runs, claims the due frontier step of each run's mailbox, +executes the durable handler against a hydrated run-state instance, and +atomically commits the state patch together with the successor slots the +handler returned. Retries, timeouts, lifecycle hooks, cancellation drain, and +crash recovery are decided here and made durable by the store. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import random +import time +import traceback +import uuid +from typing import TYPE_CHECKING, Any + +from pydantic import TypeAdapter +from reflex_base.event.processor.base_state_processor import _transform_event_payload +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import ( + DEFAULT_MAX_RECOVERIES, + After, + CompleteRun, + FailRun, + ManualTrigger, + NeedsAttention, + parse_duration, +) + +from reflex.event import EventHandler, EventSpec +from reflex.workflow.records import ( + TERMINAL_STEP_STATUSES, + HistoryEventType, + RunRecord, + RunSnapshot, + RunStatus, + StartResult, + StepRecord, + StepStatus, +) +from reflex.workflow.serde import to_run_data +from reflex.workflow.store import Claim, RunStore, StaleClaimError, StepCompletion + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + + from reflex.state import BaseState + from reflex.workflow.definition import HandlerDefinition, WorkflowDefinition + +DEFAULT_POLL_INTERVAL = 0.25 + + +def _error_payload(error: BaseException) -> dict[str, Any]: + """Build a JSON-compatible error payload from an exception. + + Args: + error: The exception to record. + + Returns: + The error payload. + """ + return { + "type": type(error).__name__, + "message": str(error), + "traceback": "".join( + traceback.format_exception(type(error), error, error.__traceback__) + ), + } + + +class _SuccessorSpec: + """A resolved successor slot to allocate at commit. + + Attributes: + handler_id: The successor handler id. + args: JSON-compatible payload for the successor. + delay: Seconds to wait before the slot becomes due. + origin: How the slot was requested. + """ + + __slots__ = ("args", "delay", "handler_id", "origin") + + def __init__( + self, + handler_id: str, + args: dict[str, Any], + delay: float = 0.0, + origin: str = "chain", + ): + """Initialize the successor spec. + + Args: + handler_id: The successor handler id. + args: JSON-compatible payload for the successor. + delay: Seconds to wait before the slot becomes due. + origin: How the slot was requested. + """ + self.handler_id = handler_id + self.args = args + self.delay = delay + self.origin = origin + + +class WorkflowKernel: + """Executes durable workflow runs against a run store.""" + + def __init__( + self, + definitions: Iterable[WorkflowDefinition], + store: RunStore, + *, + clock: Callable[[], float] = time.time, + rng: Callable[[], float] = random.random, + poll_interval: float = DEFAULT_POLL_INTERVAL, + max_recoveries: int = DEFAULT_MAX_RECOVERIES, + ): + """Initialize the kernel. + + Args: + definitions: Compiled workflow definitions to serve. + store: The durable run store. + clock: Epoch-seconds time source; injectable for virtual time. + rng: Uniform [0, 1) source used for retry jitter. + poll_interval: Worker sleep bound between due-time checks. + max_recoveries: Infrastructure recovery budget per logical step. + """ + self._definitions: dict[str, WorkflowDefinition] = { + defn.workflow_id: defn for defn in definitions + } + self._definitions_by_cls: dict[type, WorkflowDefinition] = { + defn.state_cls: defn for defn in self._definitions.values() + } + self._store = store + self._clock = clock + self._rng = rng + self._poll_interval = poll_interval + self._max_recoveries = max_recoveries + self._field_adapters: dict[tuple[str, str], TypeAdapter] = {} + self._inflight: dict[str, asyncio.Task] = {} + self._wakeup = asyncio.Event() + self._worker: asyncio.Task | None = None + self._recovered = False + + @property + def store(self) -> RunStore: + """The kernel's durable run store. + + Returns: + The store. + """ + return self._store + + def _resolve_target( + self, target: Any + ) -> tuple[WorkflowDefinition, HandlerDefinition, dict[str, Any]]: + """Resolve a start target into a definition, handler, and payload. + + Args: + target: An ``EventSpec`` from calling a class-level handler, or a + class-level ``EventHandler`` reference for no-arg handlers. + + Returns: + The workflow definition, handler definition, and decoded payload. + + Raises: + WorkflowRuntimeError: If the target does not reference a handler on + a registered workflow class, or its args are not literal values. + """ + if isinstance(target, EventSpec): + handler = target.handler + args = {} + for name_var, value_var in target.args: + try: + value = value_var._var_value # pyright: ignore[reportAttributeAccessIssue] + except (AttributeError, NotImplementedError): + msg = ( + "Workflow event arguments must be literal values, " + f"got {value_var!r} for {name_var._js_expr!r}." + ) + raise WorkflowRuntimeError(msg) from None + args[name_var._js_expr] = value + elif isinstance(target, EventHandler): + handler = target + args = {} + else: + msg = ( + "Expected a workflow event like MyWorkflow.my_handler or " + f"MyWorkflow.my_handler(args), got {target!r}." + ) + raise WorkflowRuntimeError(msg) + state_cls = handler.state + defn = ( + self._definitions_by_cls.get(state_cls) if state_cls is not None else None + ) + if defn is None: + msg = ( + f"{state_cls.__name__ if state_cls is not None else target!r} is " + "not a registered workflow; call app.add_workflow(...) first." + ) + raise WorkflowRuntimeError(msg) + handler_name = handler.fn.__name__ + handler_id = defn.handler_ids_by_name.get(handler_name) + if handler_id is None: + msg = f"{handler_name!r} is not a durable handler on {defn.workflow_id!r}." + raise WorkflowRuntimeError(msg) + return defn, defn.handlers[handler_id], self._normalize_payload(args) + + @staticmethod + def _normalize_payload(args: dict[str, Any]) -> dict[str, Any]: + """Normalize a payload to JSON-compatible values. + + Args: + args: The raw payload values. + + Returns: + The normalized payload. + + Raises: + WorkflowRuntimeError: If a value is not serializable. + """ + try: + return to_run_data(args) + except (TypeError, ValueError) as err: + msg = f"Workflow event payload is not serializable: {err}" + raise WorkflowRuntimeError(msg) from None + + async def start( + self, + target: Any, + *, + request_key: str | None = None, + labels: dict[str, str] | None = None, + ) -> StartResult: + """Admit a new run from a manual root event. + + Args: + target: The root event, e.g. ``MyWorkflow.start(payload)``. + request_key: Idempotent admission key; a repeated key returns the + prior run with disposition ``"deduplicated"``. + labels: Server-derived indexing labels to record on the run. + + Returns: + The admission result. + + Raises: + WorkflowRuntimeError: If the target is not a manual root handler. + """ + defn, handler, payload = self._resolve_target(target) + if not isinstance(handler.trigger, ManualTrigger): + msg = ( + f"Handler {handler.id!r} of {defn.workflow_id!r} is not a manual " + "root; only handlers with trigger=rx.manual() can be started " + "directly." + ) + raise WorkflowRuntimeError(msg) + now = self._clock() + run_id = uuid.uuid4().hex + run = RunRecord( + run_id=run_id, + workflow_id=defn.workflow_id, + definition_digest=defn.digest, + status=RunStatus.PENDING, + state={field.name: field.default for field in defn.fields}, + state_version=0, + next_ordinal=1, + request_key=request_key, + labels=labels, + deadline=(now + defn.run_timeout) if defn.run_timeout is not None else None, + created_at=now, + updated_at=now, + ) + root_step = StepRecord( + run_id=run_id, + ordinal=0, + handler_id=handler.id, + status=StepStatus.READY, + args=payload, + origin="root", + created_at=now, + updated_at=now, + ) + created, authoritative_run_id = await self._store.admit( + run, + root_step, + ( + ( + HistoryEventType.RUN_ADMITTED, + {"handler_id": handler.id, "request_key": request_key}, + ), + ( + HistoryEventType.STEP_SCHEDULED, + {"ordinal": 0, "handler_id": handler.id}, + ), + ), + ) + if not created: + return StartResult(disposition="deduplicated", run_id=authoritative_run_id) + self._wakeup.set() + return StartResult(disposition="started", run_id=authoritative_run_id) + + async def cancel(self, run_id: str) -> bool: + """Request cancellation of a run. + + The in-flight attempt, if any, is cancelled cooperatively; the run is + finalized once drained. + + Args: + run_id: The run to cancel. + + Returns: + True if intent was recorded on a nonterminal run. + """ + recorded = await self._store.request_cancel(run_id, self._clock()) + if recorded: + task = self._inflight.get(run_id) + if task is not None: + task.cancel() + self._wakeup.set() + return recorded + + async def get_run(self, run_id: str) -> RunSnapshot | None: + """Load a read-only snapshot of a run. + + Args: + run_id: The run identity. + + Returns: + The snapshot, or None if the run is unknown. + """ + run = await self._store.get_run(run_id) + if run is None: + return None + steps = await self._store.get_steps(run_id) + return RunSnapshot( + run_id=run.run_id, + workflow_id=run.workflow_id, + status=run.status, + state=run.state, + state_version=run.state_version, + result=run.result, + error=run.error, + steps=steps, + ) + + def _adapter(self, defn: WorkflowDefinition, field_name: str) -> TypeAdapter: + """Get (or build) the type adapter used to coerce a loaded field value. + + Args: + defn: The workflow definition. + field_name: The field name. + + Returns: + The pydantic adapter for the field's annotation. + """ + key = (defn.workflow_id, field_name) + adapter = self._field_adapters.get(key) + if adapter is None: + annotated_type = next( + field.annotated_type + for field in defn.fields + if field.name == field_name + ) + adapter = TypeAdapter(annotated_type) + self._field_adapters[key] = adapter + return adapter + + def _hydrate(self, defn: WorkflowDefinition, state: dict[str, Any]) -> BaseState: + """Build a run-state instance from a committed snapshot. + + Args: + defn: The workflow definition. + state: The committed state snapshot. + + Returns: + The hydrated state instance. + """ + instance = defn.state_cls(init_substates=False, _reflex_internal_init=True) + for field in defn.fields: + if field.name in state: + value = self._adapter(defn, field.name).validate_python( + state[field.name] + ) + setattr(instance, field.name, value) + return instance + + def _snapshot( + self, defn: WorkflowDefinition, instance: BaseState + ) -> dict[str, Any]: + """Serialize a run-state instance into a committed snapshot. + + Args: + defn: The workflow definition. + instance: The state instance after the attempt. + + Returns: + The JSON-compatible snapshot. + + Raises: + WorkflowRuntimeError: If a field value is not serializable. + """ + snapshot = {} + for field in defn.fields: + value = getattr(instance, field.name) + try: + snapshot[field.name] = to_run_data(value) + except (TypeError, ValueError) as err: + msg = ( + f"Run state field {field.name!r} of {defn.workflow_id!r} is " + f"not serializable: {err}" + ) + raise WorkflowRuntimeError(msg) from None + return snapshot + + def _resolve_successor( + self, defn: WorkflowDefinition, value: Any + ) -> _SuccessorSpec: + """Resolve one returned successor reference. + + Args: + defn: The workflow definition of the committing run. + value: A same-class handler reference, event spec, or ``rx.after``. + + Returns: + The resolved successor spec. + + Raises: + WorkflowRuntimeError: If the reference is not a durable handler on + the same workflow class. + """ + if isinstance(value, After): + inner = self._resolve_successor(defn, value.target) + inner.delay = parse_duration(value.delay) + inner.origin = "delay" + return inner + successor_defn, handler, payload = self._resolve_target(value) + if successor_defn is not defn: + msg = ( + f"Handler {handler.id!r} belongs to {successor_defn.workflow_id!r}; " + f"a {defn.workflow_id!r} step can only chain handlers of its own " + "workflow class." + ) + raise WorkflowRuntimeError(msg) + return _SuccessorSpec(handler.id, payload) + + def _interpret_return( + self, defn: WorkflowDefinition, value: Any + ) -> tuple[list[_SuccessorSpec], CompleteRun | FailRun | NeedsAttention | None]: + """Interpret a durable handler's return value. + + Args: + defn: The workflow definition of the committing run. + value: The handler return value. + + Returns: + The successor specs to allocate and the control outcome, if any. + + Raises: + WorkflowRuntimeError: If the return value is not a valid durable + transition. + """ + if value is None: + return [], None + if isinstance(value, (CompleteRun, FailRun, NeedsAttention)): + return [], value + if isinstance(value, (list, tuple)): + successors = [] + for item in value: + if isinstance(item, After): + msg = ( + "rx.after(...) must be returned alone; a returned list " + "is an immediate sequential chain." + ) + raise WorkflowRuntimeError(msg) + successors.append(self._resolve_successor(defn, item)) + return successors, None + return [self._resolve_successor(defn, value)], None + + async def _invoke( + self, handler: HandlerDefinition, instance: BaseState, args: dict[str, Any] + ) -> Any: + """Invoke a handler attempt with its per-attempt timeout. + + Args: + handler: The handler definition. + instance: The hydrated run-state instance. + args: The step payload. + + Returns: + The handler return value. + """ + try: + payload = _transform_event_payload(args, handler.type_hints) + except Exception: + payload = dict(args) + if handler.is_async: + coroutine = handler.fn(instance, **payload) + else: + coroutine = asyncio.to_thread(handler.fn, instance, **payload) + if handler.timeout is not None: + return await asyncio.wait_for(coroutine, timeout=handler.timeout) + return await coroutine + + def _build_new_steps( + self, + run: RunRecord, + successors: list[_SuccessorSpec], + now: float, + ) -> tuple[StepRecord, ...]: + """Allocate successor slots with preallocated ordinals. + + Args: + run: The run record as of the claim. + successors: The resolved successor specs. + now: Current time in epoch seconds. + + Returns: + The new step records in ordinal order. + """ + return tuple( + StepRecord( + run_id=run.run_id, + ordinal=run.next_ordinal + offset, + handler_id=spec.handler_id, + status=StepStatus.READY, + args=spec.args, + due_at=(now + spec.delay) if spec.delay else 0.0, + origin=spec.origin, # pyright: ignore[reportArgumentType] + created_at=now, + updated_at=now, + ) + for offset, spec in enumerate(successors) + ) + + @staticmethod + def _open_ordinals(steps: Iterable[StepRecord], *, exclude: int) -> tuple[int, ...]: + """Find unresolved slots to tombstone on a final disposition. + + Args: + steps: The run's current steps. + exclude: The committing step's ordinal. + + Returns: + The ordinals of unresolved slots. + """ + return tuple( + step.ordinal + for step in steps + if step.ordinal != exclude and step.status not in TERMINAL_STEP_STATUSES + ) + + def _final_failure_completion( + self, + defn: WorkflowDefinition, + handler: HandlerDefinition, + claim: Claim, + steps: tuple[StepRecord, ...], + *, + step_status: StepStatus, + run_status: RunStatus, + hook_id: str | None, + error: dict[str, Any], + run_event: HistoryEventType, + attempt_event: HistoryEventType, + now: float, + ) -> StepCompletion: + """Build the commit for a step's final failure or timeout. + + Unresolved slots are tombstoned first; the declared lifecycle hook, if + any, is then allocated as a fresh slot and the run continues through it. + + Args: + defn: The workflow definition. + handler: The failing handler definition. + claim: The claim being committed. + steps: The run's current steps. + step_status: The step's final status. + run_status: The run status when no hook continues the run. + hook_id: The lifecycle hook handler id, if declared. + error: The recorded error payload. + run_event: History event type for the run disposition. + attempt_event: History event type for the failing attempt. + now: Current time in epoch seconds. + + Returns: + The completion to commit. + """ + tombstones = self._open_ordinals(steps, exclude=claim.step.ordinal) + events: list[tuple[HistoryEventType, dict[str, Any]]] = [ + (attempt_event, {"ordinal": claim.step.ordinal, "error": error}), + *( + (HistoryEventType.STEP_TOMBSTONED, {"ordinal": ordinal}) + for ordinal in tombstones + ), + ] + new_steps: tuple[StepRecord, ...] = () + if hook_id is not None: + new_steps = ( + StepRecord( + run_id=claim.run.run_id, + ordinal=claim.run.next_ordinal, + handler_id=hook_id, + status=StepStatus.READY, + args={}, + origin="hook", + created_at=now, + updated_at=now, + ), + ) + events.append(( + HistoryEventType.STEP_SCHEDULED, + {"ordinal": claim.run.next_ordinal, "handler_id": hook_id}, + )) + final_run_status = RunStatus.RUNNING + run_error = None + else: + events.append((run_event, {"error": error})) + final_run_status = run_status + run_error = error + return StepCompletion( + step_status=step_status, + run_status=final_run_status, + state=None, + consume_attempt=True, + step_error=error, + run_error=run_error, + new_steps=new_steps, + tombstones=tombstones, + next_ordinal=claim.run.next_ordinal + len(new_steps), + events=tuple(events), + ) + + def _failure_completion( + self, + defn: WorkflowDefinition, + handler: HandlerDefinition, + claim: Claim, + steps: tuple[StepRecord, ...], + error: BaseException, + *, + timed_out: bool, + now: float, + ) -> StepCompletion: + """Build the commit for a failed or timed-out attempt. + + Args: + defn: The workflow definition. + handler: The handler definition. + claim: The claim being committed. + steps: The run's current steps. + error: The exception raised by the attempt. + timed_out: Whether the attempt hit its execution timeout. + now: Current time in epoch seconds. + + Returns: + The completion to commit. + """ + payload = _error_payload(error) + attempt_event = ( + HistoryEventType.ATTEMPT_TIMED_OUT + if timed_out + else HistoryEventType.ATTEMPT_FAILED + ) + attempts_after = claim.step.attempts + 1 + if handler.effect == "non_idempotent_write": + payload["reason"] = ( + "uncertain non-idempotent effect; resolve and rerun manually" + ) + return StepCompletion( + step_status=StepStatus.NEEDS_ATTENTION, + run_status=RunStatus.NEEDS_ATTENTION, + state=None, + consume_attempt=True, + step_error=payload, + run_error=payload, + events=( + (attempt_event, {"ordinal": claim.step.ordinal, "error": payload}), + (HistoryEventType.RUN_NEEDS_ATTENTION, {"error": payload}), + ), + ) + retryable = timed_out or handler.retry.is_retryable(error) + if retryable and attempts_after < handler.retry.max_attempts: + delay = handler.retry.delay_for_attempt(attempts_after) + if handler.retry.jitter == "full": + delay *= self._rng() + due_at = now + delay + return StepCompletion( + step_status=StepStatus.RETRY_WAIT, + run_status=RunStatus.RETRYING, + state=None, + consume_attempt=True, + step_error=payload, + due_at=due_at, + events=( + (attempt_event, {"ordinal": claim.step.ordinal, "error": payload}), + ( + HistoryEventType.STEP_RETRY_SCHEDULED, + { + "ordinal": claim.step.ordinal, + "attempt": attempts_after, + "due_at": due_at, + }, + ), + ), + ) + if timed_out: + return self._final_failure_completion( + defn, + handler, + claim, + steps, + step_status=StepStatus.TIMED_OUT, + run_status=RunStatus.TIMED_OUT, + hook_id=handler.on_timeout, + error=payload, + run_event=HistoryEventType.RUN_TIMED_OUT, + attempt_event=attempt_event, + now=now, + ) + return self._final_failure_completion( + defn, + handler, + claim, + steps, + step_status=StepStatus.FAILED, + run_status=RunStatus.FAILED, + hook_id=handler.on_failure, + error=payload, + run_event=HistoryEventType.RUN_FAILED, + attempt_event=attempt_event, + now=now, + ) + + def _success_completion( + self, + defn: WorkflowDefinition, + claim: Claim, + steps: tuple[StepRecord, ...], + state: dict[str, Any], + successors: list[_SuccessorSpec], + control: CompleteRun | FailRun | NeedsAttention | None, + now: float, + ) -> StepCompletion: + """Build the commit for a successful attempt. + + Args: + defn: The workflow definition. + claim: The claim being committed. + steps: The run's current steps. + state: The state snapshot to commit. + successors: Successor slots requested by the return value. + control: Explicit control outcome, if the handler returned one. + now: Current time in epoch seconds. + + Returns: + The completion to commit. + """ + events: list[tuple[HistoryEventType, dict[str, Any]]] = [ + (HistoryEventType.ATTEMPT_SUCCEEDED, {"ordinal": claim.step.ordinal}) + ] + if isinstance(control, FailRun): + error = {"reason": control.reason, "details": control.details} + tombstones = self._open_ordinals(steps, exclude=claim.step.ordinal) + events.extend( + (HistoryEventType.STEP_TOMBSTONED, {"ordinal": ordinal}) + for ordinal in tombstones + ) + events.append((HistoryEventType.RUN_FAILED, {"error": error})) + return StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.FAILED, + state=state, + run_error=error, + tombstones=tombstones, + events=tuple(events), + ) + if isinstance(control, NeedsAttention): + error = {"reason": control.reason, "details": control.details} + events.append((HistoryEventType.RUN_NEEDS_ATTENTION, {"error": error})) + return StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.NEEDS_ATTENTION, + state=state, + run_error=error, + events=tuple(events), + ) + if isinstance(control, CompleteRun): + tombstones = self._open_ordinals(steps, exclude=claim.step.ordinal) + events.extend( + (HistoryEventType.STEP_TOMBSTONED, {"ordinal": ordinal}) + for ordinal in tombstones + ) + events.append((HistoryEventType.RUN_COMPLETED, {})) + return StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state=state, + result=self._normalize_payload({"result": control.result})["result"], + tombstones=tombstones, + events=tuple(events), + ) + allocated = claim.run.next_ordinal + len(successors) + if allocated > defn.max_steps: + error = { + "reason": "max_steps_exceeded", + "max_steps": defn.max_steps, + } + tombstones = self._open_ordinals(steps, exclude=claim.step.ordinal) + events.extend( + (HistoryEventType.STEP_TOMBSTONED, {"ordinal": ordinal}) + for ordinal in tombstones + ) + events.append((HistoryEventType.RUN_FAILED, {"error": error})) + return StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.FAILED, + state=state, + run_error=error, + tombstones=tombstones, + events=tuple(events), + ) + new_steps = self._build_new_steps(claim.run, successors, now) + events.extend( + ( + HistoryEventType.STEP_SCHEDULED, + {"ordinal": step.ordinal, "handler_id": step.handler_id}, + ) + for step in new_steps + ) + open_after_commit = [ + step + for step in (*steps, *new_steps) + if step.ordinal != claim.step.ordinal + and step.status not in TERMINAL_STEP_STATUSES + ] + if not open_after_commit: + run_status = RunStatus.COMPLETED + events.append((HistoryEventType.RUN_COMPLETED, {})) + elif all(step.due_at > now for step in open_after_commit): + run_status = RunStatus.WAITING + else: + run_status = RunStatus.RUNNING + return StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=run_status, + state=state, + new_steps=new_steps, + next_ordinal=allocated, + events=tuple(events), + ) + + async def _execute_claim(self, claim: Claim) -> None: + """Execute one claimed attempt and commit its outcome. + + Args: + claim: The claim to execute. + """ + defn = self._definitions.get(claim.run.workflow_id) + now = self._clock() + if defn is None or defn.digest != claim.run.definition_digest: + reason = ( + "unknown_workflow" if defn is None else "definition_digest_mismatch" + ) + await self._store.commit( + claim, + StepCompletion( + step_status=StepStatus.NEEDS_ATTENTION, + run_status=RunStatus.NEEDS_ATTENTION, + state=None, + step_error={"reason": reason}, + run_error={"reason": reason}, + events=( + (HistoryEventType.RUN_NEEDS_ATTENTION, {"reason": reason}), + ), + ), + now, + ) + return + handler = defn.handlers[claim.step.handler_id] + steps = await self._store.get_steps(claim.run.run_id) + await self._store.append_events( + claim.run.run_id, + ( + ( + HistoryEventType.ATTEMPT_STARTED, + { + "ordinal": claim.step.ordinal, + "handler_id": handler.id, + "attempt": claim.step.attempts + 1, + "effect": handler.effect, + }, + ), + ), + now, + ) + try: + instance = self._hydrate(defn, claim.run.state) + value = await self._invoke(handler, instance, claim.step.args) + successors, control = self._interpret_return(defn, value) + state = self._snapshot(defn, instance) + completion = self._success_completion( + defn, claim, steps, state, successors, control, self._clock() + ) + except asyncio.CancelledError: + await self._store.release_claim( + claim, + status=StepStatus.CANCELLED, + events=( + ( + HistoryEventType.ATTEMPT_CANCELLED, + {"ordinal": claim.step.ordinal}, + ), + ), + now=self._clock(), + ) + return + except TimeoutError as err: + completion = self._failure_completion( + defn, handler, claim, steps, err, timed_out=True, now=self._clock() + ) + except BaseException as err: + completion = self._failure_completion( + defn, handler, claim, steps, err, timed_out=False, now=self._clock() + ) + try: + await self._store.commit(claim, completion, self._clock()) + except StaleClaimError: + return + + async def _tick(self) -> bool: + """Run one scheduling round. + + Returns: + True if any control transition or attempt was processed. + """ + now = self._clock() + progressed = False + for run in await self._store.control_pending(now): + if run.cancel_requested: + progressed = ( + await self._store.finalize_run( + run.run_id, + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=now, + ) + or progressed + ) + else: + progressed = ( + await self._store.finalize_run( + run.run_id, + status=RunStatus.TIMED_OUT, + error={"reason": "run_timeout"}, + event=HistoryEventType.RUN_TIMED_OUT, + now=now, + ) + or progressed + ) + claim = await self._store.claim_next(now) + if claim is not None: + task = asyncio.ensure_future(self._execute_claim(claim)) + self._inflight[claim.run.run_id] = task + try: + await task + finally: + self._inflight.pop(claim.run.run_id, None) + progressed = True + return progressed + + async def recover(self) -> int: + """Recover orphaned claims left by a previous process. + + Returns: + The number of steps recovered. + """ + self._recovered = True + return await self._store.recover_orphans(self._clock(), self._max_recoveries) + + async def run_until_idle(self) -> None: + """Process work until nothing is claimable at the current clock time. + + Scheduled future work (retry backoff, ``rx.after`` delays) stays + pending; advance the clock and call again to run it. + """ + if not self._recovered: + await self.recover() + while await self._tick(): + pass + + async def _worker_loop(self) -> None: + """Process work continuously until the kernel is closed.""" + while True: + progressed = await self._tick() + if progressed: + continue + now = self._clock() + due = await self._store.next_due(now) + delay = self._poll_interval + if due is not None: + delay = min(delay, max(due - now, 0.0)) + if delay <= 0: + continue + self._wakeup.clear() + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(self._wakeup.wait(), timeout=delay) + + async def start_worker(self) -> None: + """Start the background worker after recovering orphaned claims.""" + if self._worker is not None: + return + await self.recover() + self._worker = asyncio.create_task(self._worker_loop()) + + async def aclose(self) -> None: + """Stop the background worker, leaving in-flight claims recoverable.""" + if self._worker is None: + return + self._worker.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._worker + self._worker = None diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py new file mode 100644 index 00000000000..f8c1ce612cf --- /dev/null +++ b/reflex/workflow/records.py @@ -0,0 +1,233 @@ +"""Durable record types shared by the workflow store, kernel, and public API.""" + +from __future__ import annotations + +import dataclasses +import enum +from typing import Any, Literal + + +class RunStatus(str, enum.Enum): + """Lifecycle status of a workflow run.""" + + PENDING = "PENDING" + RUNNING = "RUNNING" + RETRYING = "RETRYING" + WAITING = "WAITING" + CANCELLING = "CANCELLING" + CANCELLED = "CANCELLED" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + TIMED_OUT = "TIMED_OUT" + NEEDS_ATTENTION = "NEEDS_ATTENTION" + + +TERMINAL_RUN_STATUSES = frozenset(( + RunStatus.CANCELLED, + RunStatus.COMPLETED, + RunStatus.FAILED, + RunStatus.TIMED_OUT, +)) + + +class StepStatus(str, enum.Enum): + """Lifecycle status of one logical step in a run's mailbox. + + Successor slots are created ``READY`` because the in-process kernel is the + single writer and the frontier scan already enforces mailbox order; a + distributed kernel adapter would hold successors in a blocked state until + their predecessor commit is visible. + """ + + READY = "READY" + CLAIMED = "CLAIMED" + RETRY_WAIT = "RETRY_WAIT" + RECOVERY_WAIT = "RECOVERY_WAIT" + SUCCEEDED = "SUCCEEDED" + FAILED = "FAILED" + TIMED_OUT = "TIMED_OUT" + CANCELLED = "CANCELLED" + NEEDS_ATTENTION = "NEEDS_ATTENTION" + + +TERMINAL_STEP_STATUSES = frozenset(( + StepStatus.SUCCEEDED, + StepStatus.FAILED, + StepStatus.TIMED_OUT, + StepStatus.CANCELLED, + StepStatus.NEEDS_ATTENTION, +)) + +CLAIMABLE_STEP_STATUSES = frozenset(( + StepStatus.READY, + StepStatus.RETRY_WAIT, + StepStatus.RECOVERY_WAIT, +)) + + +class HistoryEventType(str, enum.Enum): + """Type of an append-only run history event.""" + + RUN_ADMITTED = "run_admitted" + STEP_SCHEDULED = "step_scheduled" + ATTEMPT_STARTED = "attempt_started" + ATTEMPT_SUCCEEDED = "attempt_succeeded" + ATTEMPT_FAILED = "attempt_failed" + ATTEMPT_TIMED_OUT = "attempt_timed_out" + ATTEMPT_CANCELLED = "attempt_cancelled" + STEP_RETRY_SCHEDULED = "step_retry_scheduled" + STEP_RECOVERED = "step_recovered" + STEP_TOMBSTONED = "step_tombstoned" + RUN_COMPLETED = "run_completed" + RUN_FAILED = "run_failed" + RUN_TIMED_OUT = "run_timed_out" + RUN_CANCEL_REQUESTED = "run_cancel_requested" + RUN_CANCELLED = "run_cancelled" + RUN_NEEDS_ATTENTION = "run_needs_attention" + + +@dataclasses.dataclass(frozen=True, slots=True) +class RunRecord: + """Authoritative record of one workflow run. + + Attributes: + run_id: Unique run identity. + workflow_id: Stable workflow identity from ``WorkflowConfig.id``. + definition_digest: Digest of the compiled definition the run is pinned to. + status: Current run status. + state: Committed run-state snapshot as JSON-compatible values. + state_version: Monotonic version, incremented on every committed step. + next_ordinal: Next mailbox ordinal to allocate. + result: Run result recorded at completion. + error: Terminal or suspension error payload. + request_key: Idempotent admission key, if one was supplied. + labels: Server-derived indexing labels. + deadline: Absolute run deadline in epoch seconds, if configured. + cancel_requested: Whether cancellation intent has been recorded. + created_at: Admission time in epoch seconds. + updated_at: Last commit time in epoch seconds. + """ + + run_id: str + workflow_id: str + definition_digest: str + status: RunStatus + state: dict[str, Any] + state_version: int + next_ordinal: int + result: Any = None + error: dict[str, Any] | None = None + request_key: str | None = None + labels: dict[str, str] | None = None + deadline: float | None = None + cancel_requested: bool = False + created_at: float = 0.0 + updated_at: float = 0.0 + + +@dataclasses.dataclass(frozen=True, slots=True) +class StepRecord: + """One preallocated slot in a run's ordered mailbox. + + Attributes: + run_id: The owning run. + ordinal: Monotonic position in the mailbox; execution order. + handler_id: Stable id of the durable handler to execute. + status: Current step status. + args: JSON-compatible payload passed to the handler. + attempts: Business attempts consumed so far. + recoveries: Infrastructure recoveries consumed so far. + due_at: Earliest epoch time the step may be claimed. + epoch: Fencing token, incremented on every claim. + error: Last recorded attempt error payload. + origin: How the slot was allocated (root, chain, delay, or hook). + created_at: Allocation time in epoch seconds. + updated_at: Last transition time in epoch seconds. + """ + + run_id: str + ordinal: int + handler_id: str + status: StepStatus + args: dict[str, Any] + attempts: int = 0 + recoveries: int = 0 + due_at: float = 0.0 + epoch: int = 0 + error: dict[str, Any] | None = None + origin: Literal["root", "chain", "delay", "hook"] = "chain" + created_at: float = 0.0 + updated_at: float = 0.0 + + +@dataclasses.dataclass(frozen=True, slots=True) +class HistoryEvent: + """One append-only entry in a run's authoritative history. + + Attributes: + run_id: The owning run. + seq: Monotonic sequence number within the run. + type: The event type. + at: Event time in epoch seconds. + data: JSON-compatible event payload. + """ + + run_id: str + seq: int + type: HistoryEventType + at: float + data: dict[str, Any] + + +StartDisposition = Literal[ + "started", + "buffered", + "coalesced", + "skipped", + "rejected", + "deduplicated", +] + + +@dataclasses.dataclass(frozen=True, slots=True) +class StartResult: + """Typed result of a workflow start submission. + + Attributes: + disposition: How admission handled the submission. + run_id: The created or prior run, when the disposition identifies one. + admission_id: Admission identity for buffered or coalesced work. + retryable: Whether the caller may safely resubmit. + retry_after: Suggested resubmission delay in seconds. + """ + + disposition: StartDisposition + run_id: str | None = None + admission_id: str | None = None + retryable: bool = False + retry_after: float | None = None + + +@dataclasses.dataclass(frozen=True, slots=True) +class RunSnapshot: + """Read-only projection of a run for operators and tests. + + Attributes: + run_id: The run identity. + workflow_id: The stable workflow identity. + status: Current run status. + state: Committed run-state values. + state_version: Committed state version. + result: Run result, if completed with one. + error: Terminal or suspension error payload. + steps: All mailbox slots in ordinal order. + """ + + run_id: str + workflow_id: str + status: RunStatus + state: dict[str, Any] + state_version: int + result: Any + error: dict[str, Any] | None + steps: tuple[StepRecord, ...] diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py new file mode 100644 index 00000000000..c5ed99683c2 --- /dev/null +++ b/reflex/workflow/runtime.py @@ -0,0 +1,271 @@ +"""Workflow runtime wiring and the public ``rx.workflows`` namespace. + +A ``WorkflowRuntime`` owns the compiled definitions and the kernel for one +process. ``App.add_workflow`` registers classes on the app's runtime; the +``rx.workflows`` namespace resolves the active runtime so server code can +start, cancel, and inspect runs without holding a kernel reference. +""" + +from __future__ import annotations + +import random +import time +from contextlib import asynccontextmanager +from contextvars import ContextVar +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from reflex_base.registry import RegistrationContext +from reflex_base.utils.exceptions import WorkflowDefinitionError, WorkflowRuntimeError + +from reflex.workflow.definition import WorkflowDefinition, compile_workflow +from reflex.workflow.kernel import DEFAULT_POLL_INTERVAL, WorkflowKernel +from reflex.workflow.store import RunStore, SqliteRunStore + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Callable + + from reflex.state import BaseState + from reflex.workflow.records import RunSnapshot, StartResult + +DEFAULT_DB_FILENAME = "workflow.db" + +_context_runtime: ContextVar[WorkflowRuntime | None] = ContextVar( + "reflex_workflow_runtime", default=None +) + +_default_runtime: WorkflowRuntime | None = None + + +def _detach_from_session_registry(workflow_cls: type[BaseState]) -> None: + """Remove a workflow class from the session state registry. + + A registered workflow class is run-scoped: it must not be instantiated per + browser session, compiled into the client state schema, or reachable from + frontend event dispatch. Removing it from the registration context achieves + all three without changing the class itself. + + Args: + workflow_cls: The workflow class being registered. + """ + ctx = RegistrationContext.ensure_context() + ctx.base_states.pop(workflow_cls.get_full_name(), None) + parent = workflow_cls.get_parent_state() + if parent is not None: + ctx.base_state_substates.get(parent.get_full_name(), set()).discard( + workflow_cls + ) + for full_name, registered in list(ctx.event_handlers.items()): + if workflow_cls in registered.states: + del ctx.event_handlers[full_name] + + +class WorkflowRuntime: + """Owns the workflow definitions and kernel for one process.""" + + def __init__( + self, + store: RunStore | None = None, + *, + clock: Callable[[], float] = time.time, + rng: Callable[[], float] = random.random, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ): + """Initialize the runtime. + + Args: + store: The durable run store; defaults to a SQLite store in the + working directory, created at startup. + clock: Epoch-seconds time source; injectable for virtual time. + rng: Uniform [0, 1) source used for retry jitter. + poll_interval: Worker sleep bound between due-time checks. + """ + self._store = store + self._clock = clock + self._rng = rng + self._poll_interval = poll_interval + self._definitions: dict[str, WorkflowDefinition] = {} + self._classes: dict[type, str] = {} + self._kernel: WorkflowKernel | None = None + + def register(self, workflow_cls: type[BaseState]) -> WorkflowDefinition: + """Compile and register a workflow class. + + Registration classifies the class as workflow-focused and detaches it + from the session state tree. It does not publish or activate anything. + + Args: + workflow_cls: The workflow class to register. + + Returns: + The compiled definition. + + Raises: + WorkflowDefinitionError: If the class is invalid or its workflow id + is already registered by a different class. + WorkflowRuntimeError: If the runtime has already started. + """ + if self._kernel is not None: + msg = "Cannot register workflows after the runtime has started." + raise WorkflowRuntimeError(msg) + existing_id = self._classes.get(workflow_cls) + if existing_id is not None: + return self._definitions[existing_id] + definition = compile_workflow(workflow_cls) + conflict = self._definitions.get(definition.workflow_id) + if conflict is not None: + msg = ( + f"Workflow id {definition.workflow_id!r} is already registered " + f"by {conflict.state_cls.__name__}." + ) + raise WorkflowDefinitionError(msg) + self._definitions[definition.workflow_id] = definition + self._classes[workflow_cls] = definition.workflow_id + _detach_from_session_registry(workflow_cls) + return definition + + @property + def definitions(self) -> tuple[WorkflowDefinition, ...]: + """The registered definitions. + + Returns: + The compiled definitions. + """ + return tuple(self._definitions.values()) + + @property + def kernel(self) -> WorkflowKernel: + """The running kernel. + + Returns: + The kernel. + + Raises: + WorkflowRuntimeError: If the runtime has not started. + """ + if self._kernel is None: + msg = ( + "The workflow runtime has not started; start the app or use " + "WorkflowTestHarness in tests." + ) + raise WorkflowRuntimeError(msg) + return self._kernel + + async def startup(self, *, start_worker: bool = True) -> None: + """Build the kernel, recover orphaned claims, and start processing. + + Args: + start_worker: Whether to launch the background worker; tests pump + the kernel manually instead. + """ + if self._kernel is not None: + return + if self._store is None: + self._store = SqliteRunStore(Path.cwd() / DEFAULT_DB_FILENAME) + self._kernel = WorkflowKernel( + self._definitions.values(), + self._store, + clock=self._clock, + rng=self._rng, + poll_interval=self._poll_interval, + ) + if start_worker: + await self._kernel.start_worker() + else: + await self._kernel.recover() + + async def shutdown(self) -> None: + """Stop the worker, leaving in-flight claims recoverable on restart.""" + if self._kernel is not None: + await self._kernel.aclose() + self._kernel = None + + @asynccontextmanager + async def running(self) -> AsyncIterator[WorkflowRuntime]: + """Run the runtime for the duration of an app lifespan. + + Yields: + The active runtime. + """ + global _default_runtime + await self.startup() + previous = _default_runtime + _default_runtime = self + try: + yield self + finally: + _default_runtime = previous + await self.shutdown() + + +def get_runtime() -> WorkflowRuntime: + """Resolve the active workflow runtime. + + Returns: + The context-local runtime if one is active (tests), otherwise the + process default set by the running app. + + Raises: + WorkflowRuntimeError: If no runtime is active. + """ + runtime = _context_runtime.get() or _default_runtime + if runtime is None: + msg = ( + "No workflow runtime is active. Register workflows with " + "app.add_workflow(...) and run the app, or use WorkflowTestHarness " + "in tests." + ) + raise WorkflowRuntimeError(msg) + return runtime + + +class WorkflowsNamespace: + """The public ``rx.workflows`` API surface.""" + + @staticmethod + async def start( + target: Any, + *, + request_key: str | None = None, + labels: dict[str, str] | None = None, + ) -> StartResult: + """Start a workflow run from a manual root event. + + Args: + target: The root event, e.g. ``MyWorkflow.start(payload)``. + request_key: Idempotent admission key. + labels: Server-derived indexing labels. + + Returns: + The admission result. + """ + return await get_runtime().kernel.start( + target, request_key=request_key, labels=labels + ) + + @staticmethod + async def cancel(run_id: str) -> bool: + """Request cancellation of a run. + + Args: + run_id: The run to cancel. + + Returns: + True if intent was recorded on a nonterminal run. + """ + return await get_runtime().kernel.cancel(run_id) + + @staticmethod + async def get_run(run_id: str) -> RunSnapshot | None: + """Load a read-only snapshot of a run. + + Args: + run_id: The run identity. + + Returns: + The snapshot, or None if the run is unknown. + """ + return await get_runtime().kernel.get_run(run_id) + + +workflows = WorkflowsNamespace() diff --git a/reflex/workflow/serde.py b/reflex/workflow/serde.py new file mode 100644 index 00000000000..7491c15cd79 --- /dev/null +++ b/reflex/workflow/serde.py @@ -0,0 +1,52 @@ +"""Strict serialization for durable run data. + +Run state and event payloads must round-trip through JSON. The regular Reflex +serializer silently encodes unknown objects as ``null``, which would corrupt a +durable snapshot (e.g. a connector client or socket stored on run state), so +the workflow runtime uses this strict variant that raises instead. +""" + +from __future__ import annotations + +import json +from typing import Any + +from reflex_base.utils import serializers + + +def _strict_default(value: Any) -> Any: + """Serialize a non-JSON-native value or raise. + + Args: + value: The value encountered by the JSON encoder. + + Returns: + The serialized representation from the Reflex serializer registry. + + Raises: + TypeError: If no serializer is registered for the value's type. + """ + serialized = serializers.serialize(value) + if serialized is None: + msg = ( + f"{type(value).__name__} is not valid run data; durable values must " + "be serializable through the Reflex/pydantic serializer layer." + ) + raise TypeError(msg) + return serialized + + +def to_run_data(value: Any) -> Any: + """Normalize a value to JSON-compatible run data. + + Args: + value: The value to normalize. + + Returns: + The JSON-compatible representation. + + Raises: + TypeError: If the value contains something no serializer handles. + ValueError: If the value cannot be encoded (e.g. circular references). + """ + return json.loads(json.dumps(value, ensure_ascii=False, default=_strict_default)) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py new file mode 100644 index 00000000000..1572be24c0f --- /dev/null +++ b/reflex/workflow/store.py @@ -0,0 +1,1536 @@ +"""Durable run stores for the workflow kernel. + +A store is the single authority for run state: admission with idempotent +request keys, the ordered per-run mailbox, claim fencing, and the atomic step +commit that persists a state patch together with its successor slots. The +kernel decides what should happen; the store makes it durable atomically. + +``MemoryRunStore`` backs tests and the harness. ``SqliteRunStore`` provides +crash-safe persistence on a single machine using the standard library. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import sqlite3 +import threading +from typing import TYPE_CHECKING, Any, Protocol + +from reflex_base.utils.exceptions import WorkflowRuntimeError + +from reflex.workflow.records import ( + CLAIMABLE_STEP_STATUSES, + TERMINAL_RUN_STATUSES, + TERMINAL_STEP_STATUSES, + HistoryEvent, + HistoryEventType, + RunRecord, + RunStatus, + StepRecord, + StepStatus, +) + +if TYPE_CHECKING: + from collections.abc import Iterable + from pathlib import Path + + +class StaleClaimError(WorkflowRuntimeError): + """Raised when a commit no longer owns its claim and must be discarded.""" + + +@dataclasses.dataclass(frozen=True, slots=True) +class Claim: + """A fenced claim on the frontier step of one run. + + Attributes: + run: The run record as of the claim. + step: The step record after the claim transition. + """ + + run: RunRecord + step: StepRecord + + +@dataclasses.dataclass(frozen=True, slots=True) +class StepCompletion: + """Atomic outcome of one executed attempt, applied by ``commit``. + + Attributes: + step_status: Final step status for this commit. + run_status: Run status after this commit. + state: Committed state snapshot, or None to discard the attempt's patch. + consume_attempt: Whether this outcome consumes a business attempt. + step_error: Error payload recorded on the step. + run_error: Error payload recorded on the run. + result: Run result, for completing commits. + due_at: Earliest next claim time, for ``RETRY_WAIT``. + new_steps: Successor slots to append, with preallocated ordinals. + tombstones: Ordinals of unresolved slots to cancel. + next_ordinal: Updated mailbox allocation counter, if slots were added. + events: History events to append, in order, as (type, data) pairs. + """ + + step_status: StepStatus + run_status: RunStatus + state: dict[str, Any] | None + consume_attempt: bool = False + step_error: dict[str, Any] | None = None + run_error: dict[str, Any] | None = None + result: Any = None + due_at: float | None = None + new_steps: tuple[StepRecord, ...] = () + tombstones: tuple[int, ...] = () + next_ordinal: int | None = None + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...] = () + + +class RunStore(Protocol): + """Protocol implemented by workflow run stores.""" + + async def admit( + self, + run: RunRecord, + root_step: StepRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + ) -> tuple[bool, str]: + """Atomically admit a run, deduplicating on the request key. + + Args: + run: The run record to create. + root_step: The preallocated root mailbox slot. + events: History events to append on creation. + + Returns: + ``(True, run_id)`` when the run was created, or + ``(False, existing_run_id)`` when the request key already admitted one. + """ + ... + + async def claim_next(self, now: float) -> Claim | None: + """Claim the due frontier step of some runnable run. + + Args: + now: Current time in epoch seconds. + + Returns: + A fenced claim, or None when nothing is claimable right now. + """ + ... + + async def commit( + self, claim: Claim, completion: StepCompletion, now: float + ) -> None: + """Atomically apply the outcome of a claimed attempt. + + Args: + claim: The claim being committed. + completion: The outcome to apply. + now: Current time in epoch seconds. + + Raises: + StaleClaimError: If the claim was fenced and must be discarded. + """ + ... + + async def release_claim( + self, + claim: Claim, + *, + status: StepStatus, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Return a claimed step without committing any state. + + Args: + claim: The claim being released. + status: The step status to record, e.g. READY or CANCELLED. + events: History events to append. + now: Current time in epoch seconds. + """ + ... + + async def append_events( + self, + run_id: str, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Append evidence events outside a fenced commit. + + Args: + run_id: The owning run. + events: The (type, data) pairs to append. + now: Current time in epoch seconds. + """ + ... + + async def request_cancel(self, run_id: str, now: float) -> bool: + """Record cancellation intent on a run. + + Args: + run_id: The run to cancel. + now: Current time in epoch seconds. + + Returns: + True if intent was recorded on a nonterminal run. + """ + ... + + async def control_pending(self, now: float) -> tuple[RunRecord, ...]: + """List drained runs awaiting a control transition. + + A run is control-pending when it is nonterminal, has no claimed step, + and either has cancellation intent or has passed its deadline. + + Args: + now: Current time in epoch seconds. + + Returns: + The runs awaiting finalization. + """ + ... + + async def finalize_run( + self, + run_id: str, + *, + status: RunStatus, + error: dict[str, Any] | None, + event: HistoryEventType, + now: float, + ) -> bool: + """Terminate a drained run and tombstone its unresolved slots. + + Args: + run_id: The run to finalize. + status: The terminal status to record. + error: Error payload recorded on the run. + event: The terminal history event type. + now: Current time in epoch seconds. + + Returns: + True if the run was finalized; False if it was already terminal + or still has a claimed step. + """ + ... + + async def recover_orphans(self, now: float, max_recoveries: int) -> int: + """Recover steps left claimed by a previous process. + + Each orphan consumes one infrastructure recovery and becomes claimable + again; a step over budget fails its run. + + Args: + now: Current time in epoch seconds. + max_recoveries: Recovery budget per logical step. + + Returns: + The number of steps transitioned. + """ + ... + + async def get_run(self, run_id: str) -> RunRecord | None: + """Load one run record. + + Args: + run_id: The run identity. + + Returns: + The record, or None if unknown. + """ + ... + + async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: + """Load a run's mailbox slots in ordinal order. + + Args: + run_id: The run identity. + + Returns: + The step records. + """ + ... + + async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: + """Load a run's append-only history in sequence order. + + Args: + run_id: The run identity. + + Returns: + The history events. + """ + ... + + async def next_due(self, now: float) -> float | None: + """Earliest future time any runnable run becomes claimable. + + Args: + now: Current time in epoch seconds. + + Returns: + The epoch time, or None when no future work is scheduled. + """ + ... + + +def _run_is_runnable(run: RunRecord, now: float) -> bool: + """Whether a run may have its frontier claimed right now. + + Args: + run: The run record. + now: Current time in epoch seconds. + + Returns: + True when the run is nonterminal, unsuspended, has no cancellation + intent, and has not passed its deadline. + """ + return ( + run.status not in TERMINAL_RUN_STATUSES + and run.status is not RunStatus.NEEDS_ATTENTION + and not run.cancel_requested + and (run.deadline is None or run.deadline > now) + ) + + +def _frontier(steps: Iterable[StepRecord]) -> StepRecord | None: + """Find the lowest-ordinal unresolved step. + + Args: + steps: The run's steps in ordinal order. + + Returns: + The frontier step, or None when every slot is resolved. + """ + for step in steps: + if step.status not in TERMINAL_STEP_STATUSES: + return step + return None + + +class MemoryRunStore: + """In-memory run store for tests and the workflow test harness.""" + + def __init__(self): + """Initialize empty storage.""" + self._lock = asyncio.Lock() + self._runs: dict[str, RunRecord] = {} + self._steps: dict[str, list[StepRecord]] = {} + self._history: dict[str, list[HistoryEvent]] = {} + self._dedupe: dict[tuple[str, str], str] = {} + + def _append_events( + self, + run_id: str, + events: Iterable[tuple[HistoryEventType, dict[str, Any]]], + now: float, + ) -> None: + """Append history events with store-assigned sequence numbers. + + Args: + run_id: The owning run. + events: The (type, data) pairs to append. + now: Current time in epoch seconds. + """ + history = self._history.setdefault(run_id, []) + for event_type, data in events: + history.append( + HistoryEvent( + run_id=run_id, + seq=len(history) + 1, + type=event_type, + at=now, + data=data, + ) + ) + + async def admit( + self, + run: RunRecord, + root_step: StepRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + ) -> tuple[bool, str]: + """Atomically admit a run, deduplicating on the request key. + + Args: + run: The run record to create. + root_step: The preallocated root mailbox slot. + events: History events to append on creation. + + Returns: + Whether the run was created, and the authoritative run id. + """ + async with self._lock: + if run.request_key is not None: + dedupe_key = (run.workflow_id, run.request_key) + existing = self._dedupe.get(dedupe_key) + if existing is not None: + return False, existing + self._dedupe[dedupe_key] = run.run_id + self._runs[run.run_id] = run + self._steps[run.run_id] = [root_step] + self._append_events(run.run_id, events, run.created_at) + return True, run.run_id + + async def claim_next(self, now: float) -> Claim | None: + """Claim the due frontier step of some runnable run. + + Args: + now: Current time in epoch seconds. + + Returns: + A fenced claim, or None when nothing is claimable right now. + """ + async with self._lock: + for run in self._runs.values(): + if not _run_is_runnable(run, now): + continue + steps = self._steps[run.run_id] + frontier = _frontier(steps) + if ( + frontier is None + or frontier.status not in CLAIMABLE_STEP_STATUSES + or frontier.due_at > now + ): + continue + claimed = dataclasses.replace( + frontier, + status=StepStatus.CLAIMED, + epoch=frontier.epoch + 1, + updated_at=now, + ) + steps[claimed.ordinal] = claimed + running = dataclasses.replace( + run, status=RunStatus.RUNNING, updated_at=now + ) + self._runs[run.run_id] = running + return Claim(run=running, step=claimed) + return None + + def _check_claim(self, claim: Claim) -> tuple[RunRecord, list[StepRecord]]: + """Validate that a claim still owns its step and state version. + + Args: + claim: The claim to validate. + + Returns: + The current run record and step list. + + Raises: + StaleClaimError: If the claim was fenced. + """ + run = self._runs.get(claim.run.run_id) + steps = self._steps.get(claim.run.run_id) + if run is None or steps is None: + msg = f"Run {claim.run.run_id} no longer exists." + raise StaleClaimError(msg) + current = steps[claim.step.ordinal] + if ( + current.status is not StepStatus.CLAIMED + or current.epoch != claim.step.epoch + or run.state_version != claim.run.state_version + ): + msg = ( + f"Claim on run {run.run_id} step {current.ordinal} was fenced " + f"(epoch {claim.step.epoch} vs {current.epoch})." + ) + raise StaleClaimError(msg) + return run, steps + + async def commit( + self, claim: Claim, completion: StepCompletion, now: float + ) -> None: + """Atomically apply the outcome of a claimed attempt. + + Args: + claim: The claim being committed. + completion: The outcome to apply. + now: Current time in epoch seconds. + """ + async with self._lock: + run, steps = self._check_claim(claim) + step = steps[claim.step.ordinal] + steps[step.ordinal] = dataclasses.replace( + step, + status=completion.step_status, + attempts=step.attempts + (1 if completion.consume_attempt else 0), + due_at=completion.due_at if completion.due_at is not None else 0.0, + error=completion.step_error, + updated_at=now, + ) + for ordinal in completion.tombstones: + slot = steps[ordinal] + if slot.status not in TERMINAL_STEP_STATUSES: + steps[ordinal] = dataclasses.replace( + slot, status=StepStatus.CANCELLED, updated_at=now + ) + steps.extend(completion.new_steps) + self._runs[run.run_id] = dataclasses.replace( + run, + status=completion.run_status, + state=completion.state if completion.state is not None else run.state, + state_version=run.state_version + + (1 if completion.state is not None else 0), + next_ordinal=( + completion.next_ordinal + if completion.next_ordinal is not None + else run.next_ordinal + ), + result=completion.result + if completion.result is not None + else run.result, + error=completion.run_error, + updated_at=now, + ) + self._append_events(run.run_id, completion.events, now) + + async def release_claim( + self, + claim: Claim, + *, + status: StepStatus, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Return a claimed step without committing any state. + + Args: + claim: The claim being released. + status: The step status to record. + events: History events to append. + now: Current time in epoch seconds. + """ + async with self._lock: + try: + run, steps = self._check_claim(claim) + except StaleClaimError: + return + step = steps[claim.step.ordinal] + steps[step.ordinal] = dataclasses.replace( + step, status=status, updated_at=now + ) + self._append_events(run.run_id, events, now) + + async def append_events( + self, + run_id: str, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Append evidence events outside a fenced commit. + + Args: + run_id: The owning run. + events: The (type, data) pairs to append. + now: Current time in epoch seconds. + """ + async with self._lock: + self._append_events(run_id, events, now) + + async def request_cancel(self, run_id: str, now: float) -> bool: + """Record cancellation intent on a run. + + Args: + run_id: The run to cancel. + now: Current time in epoch seconds. + + Returns: + True if intent was recorded on a nonterminal run. + """ + async with self._lock: + run = self._runs.get(run_id) + if run is None or run.status in TERMINAL_RUN_STATUSES: + return False + self._runs[run_id] = dataclasses.replace( + run, + cancel_requested=True, + status=RunStatus.CANCELLING, + updated_at=now, + ) + self._append_events( + run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now + ) + return True + + async def control_pending(self, now: float) -> tuple[RunRecord, ...]: + """List drained runs awaiting a control transition. + + Args: + now: Current time in epoch seconds. + + Returns: + The runs awaiting finalization. + """ + async with self._lock: + pending = [] + for run in self._runs.values(): + if run.status in TERMINAL_RUN_STATUSES: + continue + if not ( + run.cancel_requested + or (run.deadline is not None and run.deadline <= now) + ): + continue + if any( + step.status is StepStatus.CLAIMED + for step in self._steps[run.run_id] + ): + continue + pending.append(run) + return tuple(pending) + + async def finalize_run( + self, + run_id: str, + *, + status: RunStatus, + error: dict[str, Any] | None, + event: HistoryEventType, + now: float, + ) -> bool: + """Terminate a drained run and tombstone its unresolved slots. + + Args: + run_id: The run to finalize. + status: The terminal status to record. + error: Error payload recorded on the run. + event: The terminal history event type. + now: Current time in epoch seconds. + + Returns: + True if the run was finalized. + """ + async with self._lock: + run = self._runs.get(run_id) + if run is None or run.status in TERMINAL_RUN_STATUSES: + return False + steps = self._steps[run_id] + if any(step.status is StepStatus.CLAIMED for step in steps): + return False + events: list[tuple[HistoryEventType, dict[str, Any]]] = [] + for step in list(steps): + if step.status not in TERMINAL_STEP_STATUSES: + steps[step.ordinal] = dataclasses.replace( + step, status=StepStatus.CANCELLED, updated_at=now + ) + events.append(( + HistoryEventType.STEP_TOMBSTONED, + {"ordinal": step.ordinal}, + )) + self._runs[run_id] = dataclasses.replace( + run, status=status, error=error, updated_at=now + ) + events.append((event, {} if error is None else dict(error))) + self._append_events(run_id, events, now) + return True + + async def recover_orphans(self, now: float, max_recoveries: int) -> int: + """Recover steps left claimed by a previous process. + + Args: + now: Current time in epoch seconds. + max_recoveries: Recovery budget per logical step. + + Returns: + The number of steps transitioned. + """ + async with self._lock: + recovered = 0 + for run in list(self._runs.values()): + if run.status in TERMINAL_RUN_STATUSES: + continue + steps = self._steps[run.run_id] + for step in list(steps): + if step.status is not StepStatus.CLAIMED: + continue + recovered += 1 + if step.recoveries + 1 > max_recoveries: + steps[step.ordinal] = dataclasses.replace( + step, + status=StepStatus.FAILED, + recoveries=step.recoveries + 1, + error={"reason": "recovery_budget_exhausted"}, + updated_at=now, + ) + self._runs[run.run_id] = dataclasses.replace( + run, + status=RunStatus.FAILED, + error={"reason": "recovery_budget_exhausted"}, + updated_at=now, + ) + self._append_events( + run.run_id, + ( + ( + HistoryEventType.RUN_FAILED, + {"reason": "recovery_budget_exhausted"}, + ), + ), + now, + ) + else: + steps[step.ordinal] = dataclasses.replace( + step, + status=StepStatus.RECOVERY_WAIT, + recoveries=step.recoveries + 1, + due_at=now, + updated_at=now, + ) + self._append_events( + run.run_id, + ( + ( + HistoryEventType.STEP_RECOVERED, + {"ordinal": step.ordinal}, + ), + ), + now, + ) + return recovered + + async def get_run(self, run_id: str) -> RunRecord | None: + """Load one run record. + + Args: + run_id: The run identity. + + Returns: + The record, or None if unknown. + """ + async with self._lock: + return self._runs.get(run_id) + + async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: + """Load a run's mailbox slots in ordinal order. + + Args: + run_id: The run identity. + + Returns: + The step records. + """ + async with self._lock: + return tuple(self._steps.get(run_id, ())) + + async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: + """Load a run's append-only history in sequence order. + + Args: + run_id: The run identity. + + Returns: + The history events. + """ + async with self._lock: + return tuple(self._history.get(run_id, ())) + + async def next_due(self, now: float) -> float | None: + """Earliest future time any runnable run becomes claimable. + + Args: + now: Current time in epoch seconds. + + Returns: + The epoch time, or None when no future work is scheduled. + """ + async with self._lock: + due_times = [] + for run in self._runs.values(): + if not _run_is_runnable(run, now): + continue + frontier = _frontier(self._steps[run.run_id]) + if frontier is not None and frontier.status in CLAIMABLE_STEP_STATUSES: + due_times.append(frontier.due_at) + return min(due_times) if due_times else None + + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS workflow_runs ( + run_id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + definition_digest TEXT NOT NULL, + status TEXT NOT NULL, + state TEXT NOT NULL, + state_version INTEGER NOT NULL, + next_ordinal INTEGER NOT NULL, + result TEXT, + error TEXT, + request_key TEXT, + labels TEXT, + deadline REAL, + cancel_requested INTEGER NOT NULL DEFAULT 0, + created_at REAL NOT NULL, + updated_at REAL NOT NULL +); +CREATE TABLE IF NOT EXISTS workflow_steps ( + run_id TEXT NOT NULL, + ordinal INTEGER NOT NULL, + handler_id TEXT NOT NULL, + status TEXT NOT NULL, + args TEXT NOT NULL, + attempts INTEGER NOT NULL DEFAULT 0, + recoveries INTEGER NOT NULL DEFAULT 0, + due_at REAL NOT NULL DEFAULT 0, + epoch INTEGER NOT NULL DEFAULT 0, + error TEXT, + origin TEXT NOT NULL, + created_at REAL NOT NULL, + updated_at REAL NOT NULL, + PRIMARY KEY (run_id, ordinal) +); +CREATE TABLE IF NOT EXISTS workflow_history ( + run_id TEXT NOT NULL, + seq INTEGER NOT NULL, + type TEXT NOT NULL, + at REAL NOT NULL, + data TEXT NOT NULL, + PRIMARY KEY (run_id, seq) +); +CREATE TABLE IF NOT EXISTS workflow_dedupe ( + workflow_id TEXT NOT NULL, + request_key TEXT NOT NULL, + run_id TEXT NOT NULL, + PRIMARY KEY (workflow_id, request_key) +); +CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs (status); +""" + + +def _dump(value: Any) -> str | None: + """Serialize an optional JSON payload column. + + Args: + value: The JSON-compatible value. + + Returns: + The JSON text, or None. + """ + return None if value is None else json.dumps(value) + + +def _load(text: str | None) -> Any: + """Deserialize an optional JSON payload column. + + Args: + text: The JSON text, or None. + + Returns: + The decoded value, or None. + """ + return None if text is None else json.loads(text) + + +def _run_from_row(row: sqlite3.Row) -> RunRecord: + """Build a run record from a database row. + + Args: + row: The ``workflow_runs`` row. + + Returns: + The run record. + """ + return RunRecord( + run_id=row["run_id"], + workflow_id=row["workflow_id"], + definition_digest=row["definition_digest"], + status=RunStatus(row["status"]), + state=json.loads(row["state"]), + state_version=row["state_version"], + next_ordinal=row["next_ordinal"], + result=_load(row["result"]), + error=_load(row["error"]), + request_key=row["request_key"], + labels=_load(row["labels"]), + deadline=row["deadline"], + cancel_requested=bool(row["cancel_requested"]), + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + +def _step_from_row(row: sqlite3.Row) -> StepRecord: + """Build a step record from a database row. + + Args: + row: The ``workflow_steps`` row. + + Returns: + The step record. + """ + return StepRecord( + run_id=row["run_id"], + ordinal=row["ordinal"], + handler_id=row["handler_id"], + status=StepStatus(row["status"]), + args=json.loads(row["args"]), + attempts=row["attempts"], + recoveries=row["recoveries"], + due_at=row["due_at"], + epoch=row["epoch"], + error=_load(row["error"]), + origin=row["origin"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + +class SqliteRunStore: + """Crash-safe run store backed by a local SQLite database.""" + + def __init__(self, db_path: str | Path): + """Open (and create if needed) the backing database. + + Args: + db_path: Path to the SQLite database file. + """ + self._lock = threading.Lock() + self._db = sqlite3.connect(str(db_path), check_same_thread=False) + self._db.row_factory = sqlite3.Row + self._db.isolation_level = None + self._db.execute("PRAGMA journal_mode=WAL") + self._db.execute("PRAGMA synchronous=NORMAL") + self._db.executescript(_SCHEMA) + + def close(self) -> None: + """Close the backing database connection.""" + self._db.close() + + def _append_events( + self, + run_id: str, + events: Iterable[tuple[HistoryEventType, dict[str, Any]]], + now: float, + ) -> None: + """Append history events inside the current transaction. + + Args: + run_id: The owning run. + events: The (type, data) pairs to append. + now: Current time in epoch seconds. + """ + row = self._db.execute( + "SELECT COALESCE(MAX(seq), 0) AS seq FROM workflow_history WHERE run_id = ?", + (run_id,), + ).fetchone() + seq = row["seq"] + for event_type, data in events: + seq += 1 + self._db.execute( + "INSERT INTO workflow_history (run_id, seq, type, at, data)" + " VALUES (?, ?, ?, ?, ?)", + (run_id, seq, event_type.value, now, json.dumps(data)), + ) + + def _insert_step(self, step: StepRecord) -> None: + """Insert a step row inside the current transaction. + + Args: + step: The step record. + """ + self._db.execute( + "INSERT INTO workflow_steps (run_id, ordinal, handler_id, status, args," + " attempts, recoveries, due_at, epoch, error, origin, created_at," + " updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + step.run_id, + step.ordinal, + step.handler_id, + step.status.value, + json.dumps(step.args), + step.attempts, + step.recoveries, + step.due_at, + step.epoch, + _dump(step.error), + step.origin, + step.created_at, + step.updated_at, + ), + ) + + def _load_steps(self, run_id: str) -> list[StepRecord]: + """Load a run's steps in ordinal order inside the current transaction. + + Args: + run_id: The owning run. + + Returns: + The step records. + """ + rows = self._db.execute( + "SELECT * FROM workflow_steps WHERE run_id = ? ORDER BY ordinal", + (run_id,), + ).fetchall() + return [_step_from_row(row) for row in rows] + + async def admit( + self, + run: RunRecord, + root_step: StepRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + ) -> tuple[bool, str]: + """Atomically admit a run, deduplicating on the request key. + + Args: + run: The run record to create. + root_step: The preallocated root mailbox slot. + events: History events to append on creation. + + Returns: + Whether the run was created, and the authoritative run id. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + if run.request_key is not None: + row = self._db.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = ? AND request_key = ?", + (run.workflow_id, run.request_key), + ).fetchone() + if row is not None: + self._db.execute("ROLLBACK") + return False, row["run_id"] + self._db.execute( + "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" + " VALUES (?, ?, ?)", + (run.workflow_id, run.request_key, run.run_id), + ) + self._db.execute( + "INSERT INTO workflow_runs (run_id, workflow_id," + " definition_digest, status, state, state_version, next_ordinal," + " result, error, request_key, labels, deadline, cancel_requested," + " created_at, updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + run.run_id, + run.workflow_id, + run.definition_digest, + run.status.value, + json.dumps(run.state), + run.state_version, + run.next_ordinal, + _dump(run.result), + _dump(run.error), + run.request_key, + _dump(run.labels), + run.deadline, + int(run.cancel_requested), + run.created_at, + run.updated_at, + ), + ) + self._insert_step(root_step) + self._append_events(run.run_id, events, run.created_at) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True, run.run_id + + async def claim_next(self, now: float) -> Claim | None: + """Claim the due frontier step of some runnable run. + + Args: + now: Current time in epoch seconds. + + Returns: + A fenced claim, or None when nothing is claimable right now. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + claim = None + try: + rows = self._db.execute( + "SELECT * FROM workflow_runs WHERE status NOT IN" + f" ({','.join('?' * len(terminal))})" + " AND status != ? AND cancel_requested = 0" + " AND (deadline IS NULL OR deadline > ?)" + " ORDER BY created_at", + (*terminal, RunStatus.NEEDS_ATTENTION.value, now), + ).fetchall() + for row in rows: + run = _run_from_row(row) + frontier = _frontier(self._load_steps(run.run_id)) + if ( + frontier is None + or frontier.status not in CLAIMABLE_STEP_STATUSES + or frontier.due_at > now + ): + continue + claimed = dataclasses.replace( + frontier, + status=StepStatus.CLAIMED, + epoch=frontier.epoch + 1, + updated_at=now, + ) + self._db.execute( + "UPDATE workflow_steps SET status = ?, epoch = ?," + " updated_at = ? WHERE run_id = ? AND ordinal = ?", + ( + claimed.status.value, + claimed.epoch, + now, + claimed.run_id, + claimed.ordinal, + ), + ) + self._db.execute( + "UPDATE workflow_runs SET status = ?, updated_at = ?" + " WHERE run_id = ?", + (RunStatus.RUNNING.value, now, run.run_id), + ) + running = dataclasses.replace( + run, status=RunStatus.RUNNING, updated_at=now + ) + claim = Claim(run=running, step=claimed) + break + self._db.execute("COMMIT" if claim is not None else "ROLLBACK") + except BaseException: + self._db.execute("ROLLBACK") + raise + return claim + + def _check_claim(self, claim: Claim) -> None: + """Validate that a claim still owns its step and state version. + + Args: + claim: The claim to validate. + + Raises: + StaleClaimError: If the claim was fenced. + """ + row = self._db.execute( + "SELECT s.status AS step_status, s.epoch AS epoch," + " r.state_version AS state_version" + " FROM workflow_steps s JOIN workflow_runs r ON r.run_id = s.run_id" + " WHERE s.run_id = ? AND s.ordinal = ?", + (claim.run.run_id, claim.step.ordinal), + ).fetchone() + if ( + row is None + or row["step_status"] != StepStatus.CLAIMED.value + or row["epoch"] != claim.step.epoch + or row["state_version"] != claim.run.state_version + ): + msg = ( + f"Claim on run {claim.run.run_id} step {claim.step.ordinal} was fenced." + ) + raise StaleClaimError(msg) + + async def commit( + self, claim: Claim, completion: StepCompletion, now: float + ) -> None: + """Atomically apply the outcome of a claimed attempt. + + Args: + claim: The claim being committed. + completion: The outcome to apply. + now: Current time in epoch seconds. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + self._check_claim(claim) + self._db.execute( + "UPDATE workflow_steps SET status = ?, attempts = attempts + ?," + " due_at = ?, error = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", + ( + completion.step_status.value, + 1 if completion.consume_attempt else 0, + completion.due_at if completion.due_at is not None else 0.0, + _dump(completion.step_error), + now, + claim.run.run_id, + claim.step.ordinal, + ), + ) + if completion.tombstones: + terminal = tuple(s.value for s in TERMINAL_STEP_STATUSES) + self._db.execute( + "UPDATE workflow_steps SET status = ?, updated_at = ?" + f" WHERE run_id = ? AND ordinal IN" + f" ({','.join('?' * len(completion.tombstones))})" + f" AND status NOT IN ({','.join('?' * len(terminal))})", + ( + StepStatus.CANCELLED.value, + now, + claim.run.run_id, + *completion.tombstones, + *terminal, + ), + ) + for step in completion.new_steps: + self._insert_step(step) + self._db.execute( + "UPDATE workflow_runs SET status = ?," + " state = CASE WHEN ? THEN ? ELSE state END," + " state_version = state_version + ?," + " next_ordinal = COALESCE(?, next_ordinal)," + " result = COALESCE(?, result), error = ?, updated_at = ?" + " WHERE run_id = ?", + ( + completion.run_status.value, + completion.state is not None, + _dump(completion.state), + 1 if completion.state is not None else 0, + completion.next_ordinal, + _dump(completion.result), + _dump(completion.run_error), + now, + claim.run.run_id, + ), + ) + self._append_events(claim.run.run_id, completion.events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + async def release_claim( + self, + claim: Claim, + *, + status: StepStatus, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Return a claimed step without committing any state. + + Args: + claim: The claim being released. + status: The step status to record. + events: History events to append. + now: Current time in epoch seconds. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + try: + self._check_claim(claim) + except StaleClaimError: + self._db.execute("ROLLBACK") + return + self._db.execute( + "UPDATE workflow_steps SET status = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", + (status.value, now, claim.run.run_id, claim.step.ordinal), + ) + self._append_events(claim.run.run_id, events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + async def append_events( + self, + run_id: str, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Append evidence events outside a fenced commit. + + Args: + run_id: The owning run. + events: The (type, data) pairs to append. + now: Current time in epoch seconds. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + self._append_events(run_id, events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + async def request_cancel(self, run_id: str, now: float) -> bool: + """Record cancellation intent on a run. + + Args: + run_id: The run to cancel. + now: Current time in epoch seconds. + + Returns: + True if intent was recorded on a nonterminal run. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + cursor = self._db.execute( + "UPDATE workflow_runs SET cancel_requested = 1, status = ?," + " updated_at = ? WHERE run_id = ? AND status NOT IN" + f" ({','.join('?' * len(terminal))})", + (RunStatus.CANCELLING.value, now, run_id, *terminal), + ) + if cursor.rowcount == 0: + self._db.execute("ROLLBACK") + return False + self._append_events( + run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True + + async def control_pending(self, now: float) -> tuple[RunRecord, ...]: + """List drained runs awaiting a control transition. + + Args: + now: Current time in epoch seconds. + + Returns: + The runs awaiting finalization. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + rows = self._db.execute( + "SELECT * FROM workflow_runs r WHERE status NOT IN" + f" ({','.join('?' * len(terminal))})" + " AND (cancel_requested = 1 OR (deadline IS NOT NULL AND deadline <= ?))" + " AND NOT EXISTS (SELECT 1 FROM workflow_steps s" + " WHERE s.run_id = r.run_id AND s.status = ?)", + (*terminal, now, StepStatus.CLAIMED.value), + ).fetchall() + return tuple(_run_from_row(row) for row in rows) + + async def finalize_run( + self, + run_id: str, + *, + status: RunStatus, + error: dict[str, Any] | None, + event: HistoryEventType, + now: float, + ) -> bool: + """Terminate a drained run and tombstone its unresolved slots. + + Args: + run_id: The run to finalize. + status: The terminal status to record. + error: Error payload recorded on the run. + event: The terminal history event type. + now: Current time in epoch seconds. + + Returns: + True if the run was finalized. + """ + terminal_run = tuple(s.value for s in TERMINAL_RUN_STATUSES) + terminal_step = tuple(s.value for s in TERMINAL_STEP_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + if row is None or row["status"] in terminal_run: + self._db.execute("ROLLBACK") + return False + claimed = self._db.execute( + "SELECT 1 FROM workflow_steps WHERE run_id = ? AND status = ?", + (run_id, StepStatus.CLAIMED.value), + ).fetchone() + if claimed is not None: + self._db.execute("ROLLBACK") + return False + open_rows = self._db.execute( + "SELECT ordinal FROM workflow_steps WHERE run_id = ?" + f" AND status NOT IN ({','.join('?' * len(terminal_step))})" + " ORDER BY ordinal", + (run_id, *terminal_step), + ).fetchall() + self._db.execute( + "UPDATE workflow_steps SET status = ?, updated_at = ?" + f" WHERE run_id = ? AND status NOT IN" + f" ({','.join('?' * len(terminal_step))})", + (StepStatus.CANCELLED.value, now, run_id, *terminal_step), + ) + self._db.execute( + "UPDATE workflow_runs SET status = ?, error = ?, updated_at = ?" + " WHERE run_id = ?", + (status.value, _dump(error), now, run_id), + ) + events: list[tuple[HistoryEventType, dict[str, Any]]] = [ + (HistoryEventType.STEP_TOMBSTONED, {"ordinal": row["ordinal"]}) + for row in open_rows + ] + events.append((event, {} if error is None else dict(error))) + self._append_events(run_id, events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True + + async def recover_orphans(self, now: float, max_recoveries: int) -> int: + """Recover steps left claimed by a previous process. + + Args: + now: Current time in epoch seconds. + max_recoveries: Recovery budget per logical step. + + Returns: + The number of steps transitioned. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + rows = self._db.execute( + "SELECT s.* FROM workflow_steps s" + " JOIN workflow_runs r ON r.run_id = s.run_id" + " WHERE s.status = ? AND r.status NOT IN" + f" ({','.join('?' * len(terminal))})", + (StepStatus.CLAIMED.value, *terminal), + ).fetchall() + recovered = 0 + for row in rows: + step = _step_from_row(row) + recovered += 1 + if step.recoveries + 1 > max_recoveries: + self._db.execute( + "UPDATE workflow_steps SET status = ?, recoveries = ?," + " error = ?, updated_at = ? WHERE run_id = ? AND ordinal = ?", + ( + StepStatus.FAILED.value, + step.recoveries + 1, + json.dumps({"reason": "recovery_budget_exhausted"}), + now, + step.run_id, + step.ordinal, + ), + ) + self._db.execute( + "UPDATE workflow_runs SET status = ?, error = ?," + " updated_at = ? WHERE run_id = ?", + ( + RunStatus.FAILED.value, + json.dumps({"reason": "recovery_budget_exhausted"}), + now, + step.run_id, + ), + ) + self._append_events( + step.run_id, + ( + ( + HistoryEventType.RUN_FAILED, + {"reason": "recovery_budget_exhausted"}, + ), + ), + now, + ) + else: + self._db.execute( + "UPDATE workflow_steps SET status = ?, recoveries = ?," + " due_at = ?, updated_at = ? WHERE run_id = ? AND ordinal = ?", + ( + StepStatus.RECOVERY_WAIT.value, + step.recoveries + 1, + now, + now, + step.run_id, + step.ordinal, + ), + ) + self._append_events( + step.run_id, + ( + ( + HistoryEventType.STEP_RECOVERED, + {"ordinal": step.ordinal}, + ), + ), + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return recovered + + async def get_run(self, run_id: str) -> RunRecord | None: + """Load one run record. + + Args: + run_id: The run identity. + + Returns: + The record, or None if unknown. + """ + with self._lock: + row = self._db.execute( + "SELECT * FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + return None if row is None else _run_from_row(row) + + async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: + """Load a run's mailbox slots in ordinal order. + + Args: + run_id: The run identity. + + Returns: + The step records. + """ + with self._lock: + return tuple(self._load_steps(run_id)) + + async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: + """Load a run's append-only history in sequence order. + + Args: + run_id: The run identity. + + Returns: + The history events. + """ + with self._lock: + rows = self._db.execute( + "SELECT * FROM workflow_history WHERE run_id = ? ORDER BY seq", + (run_id,), + ).fetchall() + return tuple( + HistoryEvent( + run_id=row["run_id"], + seq=row["seq"], + type=HistoryEventType(row["type"]), + at=row["at"], + data=json.loads(row["data"]), + ) + for row in rows + ) + + async def next_due(self, now: float) -> float | None: + """Earliest future time any runnable run becomes claimable. + + Args: + now: Current time in epoch seconds. + + Returns: + The epoch time, or None when no future work is scheduled. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + rows = self._db.execute( + "SELECT run_id FROM workflow_runs WHERE status NOT IN" + f" ({','.join('?' * len(terminal))})" + " AND status != ? AND cancel_requested = 0" + " AND (deadline IS NULL OR deadline > ?)", + (*terminal, RunStatus.NEEDS_ATTENTION.value, now), + ).fetchall() + due_times = [] + for row in rows: + frontier = _frontier(self._load_steps(row["run_id"])) + if frontier is not None and frontier.status in CLAIMABLE_STEP_STATUSES: + due_times.append(frontier.due_at) + return min(due_times) if due_times else None diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py new file mode 100644 index 00000000000..104d859dc4c --- /dev/null +++ b/reflex/workflow/testing.py @@ -0,0 +1,186 @@ +"""Deterministic test harness for workflow definitions. + +The harness runs registered workflows on an in-memory store with a virtual +clock, so tests drive retries, durable delays, and deadlines by advancing time +instead of sleeping. Retry jitter is disabled for determinism. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from reflex_base.workflow import parse_duration + +from reflex.workflow.runtime import WorkflowRuntime, _context_runtime +from reflex.workflow.store import MemoryRunStore + +if TYPE_CHECKING: + from reflex_base.workflow import DurationLike + + from reflex.state import BaseState + from reflex.workflow.kernel import WorkflowKernel + from reflex.workflow.records import RunSnapshot, StartResult + from reflex.workflow.store import RunStore + +DEFAULT_START_TIME = 1_000_000.0 + + +class _VirtualClock: + """A manually advanced epoch-seconds clock.""" + + def __init__(self, now: float): + """Initialize the clock. + + Args: + now: The starting time in epoch seconds. + """ + self.now = now + + def __call__(self) -> float: + """Read the clock. + + Returns: + The current virtual time in epoch seconds. + """ + return self.now + + +class WorkflowTestHarness: + """Runs workflows deterministically for tests. + + Usage:: + + async with WorkflowTestHarness(MyWorkflow) as harness: + result = await harness.start(MyWorkflow.begin(payload)) + await harness.advance("5s") + snapshot = await harness.get_run(result.run_id) + """ + + def __init__( + self, + *workflow_classes: type[BaseState], + store: RunStore | None = None, + start_time: float = DEFAULT_START_TIME, + ): + """Initialize the harness. + + Args: + workflow_classes: Workflow classes to register. + store: Run store override; defaults to a fresh in-memory store. + start_time: Initial virtual time in epoch seconds. + """ + self._clock = _VirtualClock(start_time) + self._runtime = WorkflowRuntime( + store if store is not None else MemoryRunStore(), + clock=self._clock, + rng=lambda: 1.0, + ) + for workflow_cls in workflow_classes: + self._runtime.register(workflow_cls) + self._token = None + + @property + def now(self) -> float: + """The current virtual time. + + Returns: + The time in epoch seconds. + """ + return self._clock.now + + @property + def runtime(self) -> WorkflowRuntime: + """The harness's workflow runtime. + + Returns: + The runtime. + """ + return self._runtime + + @property + def kernel(self) -> WorkflowKernel: + """The harness's kernel. + + Returns: + The kernel. + """ + return self._runtime.kernel + + async def __aenter__(self) -> WorkflowTestHarness: + """Start the runtime without a background worker. + + Returns: + The harness. + """ + await self._runtime.startup(start_worker=False) + self._token = _context_runtime.set(self._runtime) + return self + + async def __aexit__(self, *exc_info) -> None: + """Deactivate and shut down the runtime. + + Args: + exc_info: The exception info, if any. + """ + if self._token is not None: + _context_runtime.reset(self._token) + self._token = None + await self._runtime.shutdown() + + async def start( + self, + target: Any, + *, + request_key: str | None = None, + labels: dict[str, str] | None = None, + ) -> StartResult: + """Start a run and process work until idle. + + Args: + target: The root event, e.g. ``MyWorkflow.begin(payload)``. + request_key: Idempotent admission key. + labels: Server-derived indexing labels. + + Returns: + The admission result. + """ + result = await self.kernel.start(target, request_key=request_key, labels=labels) + await self.kernel.run_until_idle() + return result + + async def run_until_idle(self) -> None: + """Process work until nothing is claimable at the current time.""" + await self.kernel.run_until_idle() + + async def advance(self, duration: DurationLike) -> None: + """Advance the virtual clock and process any work that became due. + + Args: + duration: How far to advance, e.g. ``"2d"``. + """ + self._clock.now += parse_duration(duration) + await self.kernel.run_until_idle() + + async def get_run(self, run_id: str) -> RunSnapshot | None: + """Load a read-only snapshot of a run. + + Args: + run_id: The run identity. + + Returns: + The snapshot, or None if the run is unknown. + """ + return await self.kernel.get_run(run_id) + + async def cancel(self, run_id: str) -> bool: + """Request cancellation of a run and process the drain. + + Args: + run_id: The run to cancel. + + Returns: + True if intent was recorded on a nonterminal run. + """ + cancelled = await self.kernel.cancel(run_id) + await self.kernel.run_until_idle() + return cancelled diff --git a/tests/units/reflex_base/event/test_durable_event.py b/tests/units/reflex_base/event/test_durable_event.py new file mode 100644 index 00000000000..339873fdbcc --- /dev/null +++ b/tests/units/reflex_base/event/test_durable_event.py @@ -0,0 +1,95 @@ +"""Tests for the durable options of the ``@rx.event`` decorator.""" + +import pytest +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import Retry, get_durable_config, manual + +import reflex as rx + + +def test_durable_marker_attached(forked_registration_context): + class MarkerWorkflow(rx.State): + @rx.event( + id="begin", + durable=True, + trigger=manual(), + retry=Retry(max_attempts=2), + timeout="10s", + effect="read", + queue="integrations", + ) + def begin(self): + pass + + config = get_durable_config(MarkerWorkflow.event_handlers["begin"].fn) + assert config is not None + assert config.id == "begin" + assert config.effect == "read" + assert config.timeout == pytest.approx(10.0) + assert config.queue == "integrations" + assert config.retry is not None + assert config.retry.max_attempts == 2 + + +def test_bare_event_has_no_marker(forked_registration_context): + class PlainState(rx.State): + @rx.event + def tick(self): + pass + + assert get_durable_config(PlainState.event_handlers["tick"].fn) is None + + +def test_durable_options_without_durable_raise(): + with pytest.raises(WorkflowDefinitionError, match="requires durable=True"): + + @rx.event(effect="none") + def handler(self): + pass + + +def test_durable_background_mutually_exclusive(): + with pytest.raises(WorkflowDefinitionError, match="mutually exclusive"): + + @rx.event(durable=True, effect="none", background=True) + def handler(self): + pass + + +def test_durable_browser_actions_rejected(): + with pytest.raises(WorkflowDefinitionError, match="browser event actions"): + + @rx.event(durable=True, effect="none", throttle=100) + def handler(self): + pass + + +def test_durable_generator_rejected(): + with pytest.raises(WorkflowDefinitionError, match="generators"): + + @rx.event(durable=True, effect="none") + def handler(self): + yield + + +def test_durable_async_generator_rejected(): + with pytest.raises(WorkflowDefinitionError, match="generators"): + + @rx.event(durable=True, effect="none") + async def handler(self): # noqa: RUF029 + yield + + +def test_durable_hook_accepts_function_reference(forked_registration_context): + class HookWorkflow(rx.State): + @rx.event(durable=True, effect="none") + def cleanup(self): + pass + + @rx.event(durable=True, effect="read", on_failure=cleanup) + def risky(self): + pass + + config = get_durable_config(HookWorkflow.event_handlers["risky"].fn) + assert config is not None + assert config.on_failure == "cleanup" diff --git a/tests/units/reflex_base/test_workflow.py b/tests/units/reflex_base/test_workflow.py new file mode 100644 index 00000000000..4c27be76b6d --- /dev/null +++ b/tests/units/reflex_base/test_workflow.py @@ -0,0 +1,235 @@ +"""Tests for the workflow authoring value types.""" + +import datetime + +import pytest +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import ( + After, + Retry, + TransientWorkflowError, + WorkflowConfig, + after, + build_durable_config, + default_retry_for_effect, + manual, + parse_duration, + schedule, + webhook, +) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("30s", 30.0), + ("500ms", 0.5), + ("2.5h", 9000.0), + ("1m", 60.0), + ("30d", 30 * 86400.0), + (15, 15.0), + (0.25, 0.25), + (datetime.timedelta(minutes=2), 120.0), + (" 10 s ", 10.0), + ], +) +def test_parse_duration(value, expected): + assert parse_duration(value) == expected + + +@pytest.mark.parametrize( + "value", + ["30", "s", "-5s", "5 hours", None, object(), True], +) +def test_parse_duration_invalid(value): + with pytest.raises(WorkflowDefinitionError): + parse_duration(value) + + +def test_parse_duration_negative_timedelta(): + with pytest.raises(WorkflowDefinitionError, match="negative"): + parse_duration(datetime.timedelta(seconds=-1)) + + +def test_retry_defaults_valid(): + retry = Retry() + assert retry.max_attempts == 3 + assert retry.delay_for_attempt(1) == pytest.approx(1.0) + assert retry.delay_for_attempt(2) == pytest.approx(2.0) + assert retry.delay_for_attempt(100) == pytest.approx(60.0) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"max_attempts": 0}, + {"multiplier": 0.5}, + {"jitter": "half"}, + {"initial_delay": "1m", "max_delay": "1s"}, + {"retry_on": (ValueError,), "do_not_retry_on": (Exception,)}, + ], +) +def test_retry_invalid(kwargs): + with pytest.raises(WorkflowDefinitionError): + Retry(**kwargs) + + +def test_retry_is_retryable_respects_deny_list(): + retry = Retry( + retry_on=(TransientWorkflowError,), + do_not_retry_on=(ValueError,), + ) + assert retry.is_retryable(TransientWorkflowError("x")) + assert not retry.is_retryable(ValueError("x")) + assert not retry.is_retryable(RuntimeError("x")) + + +def test_default_retry_for_effect(): + for effect in ("none", "read", "idempotent_write"): + retry = default_retry_for_effect(effect) + assert retry.max_attempts == 3 + assert retry.retry_on == (TransientWorkflowError,) + non_idempotent = default_retry_for_effect("non_idempotent_write") + assert non_idempotent.max_attempts == 1 + assert non_idempotent.retry_on == () + + +def test_workflow_config_valid(): + config = WorkflowConfig(id="billing.payment_received", run_timeout="30d") + assert config.max_steps == 10_000 + assert config.run_timeout is not None + assert parse_duration(config.run_timeout) == pytest.approx(30 * 86400.0) + + +@pytest.mark.parametrize( + "workflow_id", + ["", "Billing", "billing..sync", ".billing", "billing.", "1billing", "a-b"], +) +def test_workflow_config_invalid_id(workflow_id): + with pytest.raises(WorkflowDefinitionError, match=r"WorkflowConfig\.id"): + WorkflowConfig(id=workflow_id) + + +def test_workflow_config_mixed_scope_acknowledgement(): + with pytest.raises(WorkflowDefinitionError, match="mixed_scope_reason"): + WorkflowConfig(id="a.b", allow_mixed_scopes=True) + with pytest.raises(WorkflowDefinitionError, match="allow_mixed_scopes"): + WorkflowConfig(id="a.b", mixed_scope_reason="why") + config = WorkflowConfig( + id="a.b", allow_mixed_scopes=True, mixed_scope_reason="migration" + ) + assert config.allow_mixed_scopes + + +def test_workflow_config_invalid_max_steps(): + with pytest.raises(WorkflowDefinitionError, match="max_steps"): + WorkflowConfig(id="a.b", max_steps=0) + + +def test_triggers(): + assert manual().kind == "manual" + hook = webhook("stripe.payment_succeeded", dedupe_by="id") + assert hook.kind == "webhook" + assert hook.topic == "stripe.payment_succeeded" + assert hook.dedupe_by == "id" + cron = schedule("0 9 * * 1") + assert cron.kind == "schedule" + with pytest.raises(WorkflowDefinitionError, match="topic"): + webhook("") + with pytest.raises(WorkflowDefinitionError, match="cron"): + schedule("hourly") + + +def test_after_validates_delay_eagerly(): + assert isinstance(after("2d", object()), After) + with pytest.raises(WorkflowDefinitionError): + after("2 fortnights", object()) + + +def _build(**overrides): + kwargs = { + "durable": False, + "id": None, + "trigger": None, + "retry": None, + "timeout": None, + "effect": None, + "queue": None, + "on_failure": None, + "on_timeout": None, + "background": None, + "has_browser_actions": False, + } + kwargs.update(overrides) + return build_durable_config(**kwargs) + + +def test_build_durable_config_session_handler_passthrough(): + assert _build() is None + + +def test_build_durable_config_options_require_durable(): + with pytest.raises(WorkflowDefinitionError, match="requires durable=True"): + _build(retry=Retry()) + + +def test_build_durable_config_requires_effect(): + with pytest.raises(WorkflowDefinitionError, match="effect"): + _build(durable=True) + with pytest.raises(WorkflowDefinitionError, match="effect"): + _build(durable=True, effect="write") + + +def test_build_durable_config_background_exclusive(): + with pytest.raises(WorkflowDefinitionError, match="mutually exclusive"): + _build(durable=True, effect="none", background=True) + + +def test_build_durable_config_browser_actions_rejected(): + with pytest.raises(WorkflowDefinitionError, match="browser event actions"): + _build(durable=True, effect="none", has_browser_actions=True) + + +def test_build_durable_config_non_idempotent_single_attempt(): + with pytest.raises(WorkflowDefinitionError, match="one business attempt"): + _build( + durable=True, + effect="non_idempotent_write", + retry=Retry(max_attempts=2), + ) + config = _build( + durable=True, + effect="non_idempotent_write", + retry=Retry(max_attempts=1), + ) + assert config is not None + + +def test_build_durable_config_invalid_id(): + with pytest.raises(WorkflowDefinitionError, match="id"): + _build(durable=True, effect="none", id="Not-Valid") + + +def test_build_durable_config_invalid_trigger(): + with pytest.raises(WorkflowDefinitionError, match="trigger"): + _build(durable=True, effect="none", trigger="manual") + + +def test_build_durable_config_hooks_normalized(): + def cleanup(self): + pass + + config = _build( + durable=True, effect="none", on_failure=cleanup, on_timeout="report" + ) + assert config is not None + assert config.on_failure == "cleanup" + assert config.on_timeout == "report" + with pytest.raises(WorkflowDefinitionError, match="on_failure"): + _build(durable=True, effect="none", on_failure="") + + +def test_build_durable_config_parses_timeout(): + config = _build(durable=True, effect="read", timeout="45s") + assert config is not None + assert config.timeout == pytest.approx(45.0) diff --git a/tests/units/workflow/__init__.py b/tests/units/workflow/__init__.py new file mode 100644 index 00000000000..d579d282bb7 --- /dev/null +++ b/tests/units/workflow/__init__.py @@ -0,0 +1 @@ +"""Unit tests for reflex.workflow.""" diff --git a/tests/units/workflow/test_app.py b/tests/units/workflow/test_app.py new file mode 100644 index 00000000000..7103c29f5a6 --- /dev/null +++ b/tests/units/workflow/test_app.py @@ -0,0 +1,131 @@ +"""Tests for workflow registration on the App and session-tree detachment.""" + +import asyncio +from typing import Any + +import pytest +from reflex_base.registry import RegistrationContext +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.state import State +from reflex.workflow.records import RunStatus +from reflex.workflow.store import MemoryRunStore + + +def _make_classes(): + class SessionCounter(rx.State): + count: int = 0 + + @rx.event + def increment(self): + self.count += 1 + + class DetachedWorkflow(rx.State): + __workflow__ = WorkflowConfig(id="app.detached") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + self.status = "done" + + return SessionCounter, DetachedWorkflow + + +def test_add_workflow_detaches_from_session_tree(forked_registration_context): + session_cls, workflow_cls = _make_classes() + assert workflow_cls in State.get_substates() + + app = rx.App() + app.add_workflow(workflow_cls) + + assert workflow_cls not in State.get_substates() + assert session_cls in State.get_substates() + # Durable handlers are no longer reachable through the event registry. + ctx = RegistrationContext.get() + assert not any( + workflow_cls in registered.states for registered in ctx.event_handlers.values() + ) + assert app._workflow_runtime is not None + assert [d.workflow_id for d in app._workflow_runtime.definitions] == [ + "app.detached" + ] + + +def test_session_state_unaffected_by_detach(forked_registration_context): + session_cls, workflow_cls = _make_classes() + app = rx.App() + app.add_workflow(workflow_cls) + + root = State(_reflex_internal_init=True) # pyright: ignore [reportCallIssue] + assert workflow_cls.get_name() not in root.substates + session: Any = root.substates[session_cls.get_name()] + session.increment() + assert session.count == 1 + assert session.get_delta() + + +def test_add_workflow_idempotent_per_class(forked_registration_context): + _, workflow_cls = _make_classes() + app = rx.App() + app.add_workflow(workflow_cls) + app.add_workflow(workflow_cls) + assert app._workflow_runtime is not None + assert len(app._workflow_runtime.definitions) == 1 + + +def test_add_workflow_rejects_duplicate_ids(forked_registration_context): + class FirstOwner(rx.State): + __workflow__ = WorkflowConfig(id="app.duplicate") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + class SecondOwner(rx.State): + __workflow__ = WorkflowConfig(id="app.duplicate") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + app = rx.App() + app.add_workflow(FirstOwner) + with pytest.raises(WorkflowDefinitionError, match="already registered"): + app.add_workflow(SecondOwner) + + +def test_add_workflow_rejects_plain_state(forked_registration_context): + class NotAWorkflow(rx.State): + pass + + app = rx.App() + with pytest.raises(WorkflowDefinitionError, match="__workflow__"): + app.add_workflow(NotAWorkflow) + + +async def test_runtime_lifespan_serves_default_namespace(forked_registration_context): + _, workflow_cls = _make_classes() + app = rx.App(workflow_store=MemoryRunStore()) + app.add_workflow(workflow_cls) + assert app._workflow_runtime is not None + + async with app._workflow_runtime.running(): + result = await rx.workflows.start(workflow_cls.begin()) + assert result.disposition == "started" + assert result.run_id is not None + # The background worker processes the run without manual pumping. + snapshot = None + for _ in range(200): + snapshot = await rx.workflows.get_run(result.run_id) + assert snapshot is not None + if snapshot.status is RunStatus.COMPLETED: + break + await asyncio.sleep(0.01) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "done"} + # After shutdown the default runtime is cleared. + with pytest.raises(Exception, match="No workflow runtime"): + await rx.workflows.start(workflow_cls.begin()) diff --git a/tests/units/workflow/test_definition.py b/tests/units/workflow/test_definition.py new file mode 100644 index 00000000000..8efb437fc39 --- /dev/null +++ b/tests/units/workflow/test_definition.py @@ -0,0 +1,283 @@ +"""Tests for the workflow definition compiler.""" + +import pytest +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import ( + Retry, + TransientWorkflowError, + WorkflowConfig, + manual, + webhook, +) + +import reflex as rx +from reflex.workflow.definition import compile_workflow + + +def _billing_workflow(): + class BillingDefinition(rx.State): + __workflow__ = WorkflowConfig( + id="billing.definition", run_timeout="30d", default_queue="integrations" + ) + payment_id: str = "" + amount: int = 0 + + @rx.event(id="payment_received", durable=True, trigger=manual(), effect="none") + def payment_received(self, payment_id: str): + self.payment_id = payment_id + return BillingDefinition.fulfill + + @rx.event( + durable=True, + retry=Retry(max_attempts=5), + timeout="30s", + effect="idempotent_write", + on_failure="report", + ) + async def fulfill(self): + pass + + @rx.event(durable=True, effect="none") + def report(self): + pass + + return BillingDefinition + + +def test_compile_happy_path(forked_registration_context): + definition = compile_workflow(_billing_workflow()) + assert definition.workflow_id == "billing.definition" + assert set(definition.handlers) == {"payment_received", "fulfill", "report"} + assert definition.roots == ("payment_received",) + assert definition.run_timeout == pytest.approx(30 * 86400.0) + fulfill = definition.handlers["fulfill"] + assert fulfill.timeout == pytest.approx(30.0) + assert fulfill.on_failure == "report" + assert fulfill.queue == "integrations" + assert fulfill.is_async + assert [field.name for field in definition.fields] == ["payment_id", "amount"] + + +def test_digest_stable_and_sensitive(forked_registration_context): + first = compile_workflow(_billing_workflow()) + second = compile_workflow(_billing_workflow()) + assert first.digest == second.digest + + class OtherPolicy(rx.State): + __workflow__ = WorkflowConfig(id="billing.definition2") + + @rx.event(id="payment_received", durable=True, trigger=manual(), effect="none") + def payment_received(self): + pass + + assert compile_workflow(OtherPolicy).digest != first.digest + + +def test_explicit_retry_inherits_transient_default(forked_registration_context): + definition = compile_workflow(_billing_workflow()) + retry = definition.handlers["fulfill"].retry + assert retry.max_attempts == 5 + assert retry.retry_on == (TransientWorkflowError,) + + +def test_explicit_retry_on_preserved(forked_registration_context): + class ExplicitRetryOn(rx.State): + __workflow__ = WorkflowConfig(id="billing.explicit_retry") + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=2, retry_on=(ConnectionError,)), + ) + def fetch(self): + pass + + retry = compile_workflow(ExplicitRetryOn).handlers["fetch"].retry + assert retry.retry_on == (ConnectionError,) + + +def test_missing_workflow_config(forked_registration_context): + class NoConfig(rx.State): + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="__workflow__"): + compile_workflow(NoConfig) + + +def test_wrong_config_type(forked_registration_context): + class BadConfig(rx.State): + __workflow__ = {"id": "a.b"} + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + with pytest.raises(WorkflowDefinitionError, match=r"rx\.WorkflowConfig"): + compile_workflow(BadConfig) + + +def test_not_a_state_class(): + with pytest.raises(WorkflowDefinitionError, match=r"rx\.State subclass"): + compile_workflow(object) # pyright: ignore[reportArgumentType] + + +def test_non_durable_public_handler_rejected(forked_registration_context): + class MixedHandlers(rx.State): + __workflow__ = WorkflowConfig(id="billing.mixed_handlers") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + @rx.event + def session_click(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="session_click"): + compile_workflow(MixedHandlers) + + +def test_duplicate_handler_ids_rejected(forked_registration_context): + class DuplicateIds(rx.State): + __workflow__ = WorkflowConfig(id="billing.duplicate_ids") + + @rx.event(id="same", durable=True, trigger=manual(), effect="none") + def first(self): + pass + + @rx.event(id="same", durable=True, effect="none") + def second(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="duplicate handler id"): + compile_workflow(DuplicateIds) + + +def test_unresolved_hook_rejected(forked_registration_context): + class UnresolvedHook(rx.State): + __workflow__ = WorkflowConfig(id="billing.unresolved_hook") + + @rx.event(durable=True, trigger=manual(), effect="none", on_failure="missing") + def go(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="on_failure"): + compile_workflow(UnresolvedHook) + + +def test_self_hook_rejected(forked_registration_context): + class SelfHook(rx.State): + __workflow__ = WorkflowConfig(id="billing.self_hook") + + @rx.event(durable=True, trigger=manual(), effect="none", on_failure="go") + def go(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="itself"): + compile_workflow(SelfHook) + + +def test_hook_with_payload_rejected(forked_registration_context): + class HookWithArgs(rx.State): + __workflow__ = WorkflowConfig(id="billing.hook_args") + + @rx.event(durable=True, trigger=manual(), effect="none", on_failure="cleanup") + def go(self): + pass + + @rx.event(durable=True, effect="none") + def cleanup(self, reason: str): + pass + + with pytest.raises(WorkflowDefinitionError, match="payload arguments"): + compile_workflow(HookWithArgs) + + +def test_no_root_rejected(forked_registration_context): + class NoRoot(rx.State): + __workflow__ = WorkflowConfig(id="billing.no_root") + + @rx.event(durable=True, effect="none") + def internal_only(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="trigger"): + compile_workflow(NoRoot) + + +def test_webhook_root_compiles(forked_registration_context): + class WebhookRoot(rx.State): + __workflow__ = WorkflowConfig(id="billing.webhook_root") + + @rx.event( + durable=True, + trigger=webhook("stripe.payment_succeeded", dedupe_by="id"), + effect="none", + ) + def on_payment(self): + pass + + definition = compile_workflow(WebhookRoot) + assert definition.roots == ("on_payment",) + + +def test_mixed_scopes_rejected(forked_registration_context): + class MixedScopes(rx.State): + __workflow__ = WorkflowConfig( + id="billing.mixed_scopes", + allow_mixed_scopes=True, + mixed_scope_reason="migration", + ) + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="mixed-scope"): + compile_workflow(MixedScopes) + + +def test_backend_vars_rejected(forked_registration_context): + class BackendVars(rx.State): + __workflow__ = WorkflowConfig(id="billing.backend_vars") + _secret_session: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="backend-only"): + compile_workflow(BackendVars) + + +def test_nested_substate_rejected(forked_registration_context): + class ParentWorkflowState(rx.State): + pass + + class NestedWorkflow(ParentWorkflowState): + __workflow__ = WorkflowConfig(id="billing.nested") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="directly"): + compile_workflow(NestedWorkflow) + + +def test_class_with_substates_rejected(forked_registration_context): + class SubstateHaver(rx.State): + __workflow__ = WorkflowConfig(id="billing.substate_haver") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + class ChildOfWorkflow(SubstateHaver): + pass + + with pytest.raises(WorkflowDefinitionError, match="substates"): + compile_workflow(SubstateHaver) diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py new file mode 100644 index 00000000000..2b4450fa773 --- /dev/null +++ b/tests/units/workflow/test_kernel.py @@ -0,0 +1,646 @@ +"""Behavioral tests for the workflow kernel via the test harness.""" + +import asyncio + +import pytest +from pydantic import BaseModel +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import ( + Retry, + TransientWorkflowError, + WorkflowConfig, + after, + complete, + fail, + manual, + needs_attention, + webhook, +) + +import reflex as rx +from reflex.workflow.records import HistoryEventType, RunStatus, StepStatus +from reflex.workflow.store import SqliteRunStore +from reflex.workflow.testing import WorkflowTestHarness + + +class Payment(BaseModel): + """Typed payload for kernel tests.""" + + id: str + amount: int + + +async def test_chain_with_typed_payload(forked_registration_context): + class ChainFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.chain") + payment_id: str = "" + amount: int = 0 + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="none") + def receive(self, payment: Payment): + self.payment_id = payment.id + self.amount = payment.amount + return ChainFlow.finish + + @rx.event(durable=True, effect="none") + def finish(self): + self.status = "done" + + async with WorkflowTestHarness(ChainFlow) as harness: + result = await harness.start(ChainFlow.receive(Payment(id="pay_1", amount=42))) + assert result.run_id is not None + assert result.disposition == "started" + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == { + "payment_id": "pay_1", + "amount": 42, + "status": "done", + } + assert snapshot.state_version == 2 + assert [step.status for step in snapshot.steps] == [ + StepStatus.SUCCEEDED, + StepStatus.SUCCEEDED, + ] + history = await harness.kernel.store.get_history(result.run_id) + assert [event.type for event in history] == [ + HistoryEventType.RUN_ADMITTED, + HistoryEventType.STEP_SCHEDULED, + HistoryEventType.ATTEMPT_STARTED, + HistoryEventType.ATTEMPT_SUCCEEDED, + HistoryEventType.STEP_SCHEDULED, + HistoryEventType.ATTEMPT_STARTED, + HistoryEventType.ATTEMPT_SUCCEEDED, + HistoryEventType.RUN_COMPLETED, + ] + + +async def test_dedupe_by_request_key(forked_registration_context): + class DedupeFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.dedupe") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + async with WorkflowTestHarness(DedupeFlow) as harness: + first = await harness.start(DedupeFlow.go(), request_key="sub-1") + second = await harness.start(DedupeFlow.go(), request_key="sub-1") + assert first.disposition == "started" + assert second.disposition == "deduplicated" + assert second.run_id == first.run_id + + +async def test_retry_backoff_and_discarded_patches(forked_registration_context): + calls = [] + + class RetryFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.retry") + status: str = "pending" + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=3, initial_delay="5s", jitter="none"), + ) + def flaky(self): + calls.append(harnessed.now) + self.status = "attempted" + if len(calls) < 3: + msg = "provider 503" + raise TransientWorkflowError(msg) + self.status = "ok" + + async with WorkflowTestHarness(RetryFlow) as harnessed: + result = await harnessed.start(RetryFlow.flaky()) + assert result.run_id is not None + snapshot = await harnessed.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.RETRYING + # The failed attempt's state patch is discarded. + assert snapshot.state == {"status": "pending"} + await harnessed.advance("5s") + await harnessed.advance("10s") + snapshot = await harnessed.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "ok"} + assert snapshot.steps[0].attempts == 2 + # Exponential backoff: failures at t0, t0+5s, success at t0+15s. + assert calls[1] - calls[0] == pytest.approx(5.0) + assert calls[2] - calls[1] == pytest.approx(10.0) + + +async def test_retry_exhaustion_runs_failure_hook(forked_registration_context): + class HookFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.hook") + status: str = "pending" + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=2, initial_delay="1s", jitter="none"), + on_failure="report", + ) + def flaky(self): + msg = "still down" + raise TransientWorkflowError(msg) + + @rx.event(durable=True, effect="none") + def report(self): + self.status = "reported" + + async with WorkflowTestHarness(HookFlow) as harness: + result = await harness.start(HookFlow.flaky()) + assert result.run_id is not None + await harness.advance("1s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + # The failing step is final; the hook ran on a fresh slot afterwards. + assert snapshot.steps[0].status is StepStatus.FAILED + assert snapshot.steps[0].attempts == 2 + assert snapshot.steps[1].handler_id == "report" + assert snapshot.steps[1].origin == "hook" + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "reported"} + + +async def test_retry_exhaustion_without_hook_fails_run(forked_registration_context): + class NoHookFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.nohook") + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=1), + ) + def flaky(self): + msg = "down" + raise TransientWorkflowError(msg) + + async with WorkflowTestHarness(NoHookFlow) as harness: + result = await harness.start(NoHookFlow.flaky()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error is not None + assert snapshot.error["type"] == "TransientWorkflowError" + + +async def test_unknown_defect_does_not_retry(forked_registration_context): + calls = [] + + class DefectFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.defect") + + @rx.event(durable=True, trigger=manual(), effect="none") + def broken(self): + calls.append(1) + msg = "bug" + raise ValueError(msg) + + async with WorkflowTestHarness(DefectFlow) as harness: + result = await harness.start(DefectFlow.broken()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert len(calls) == 1 + + +async def test_timeout_consumes_attempts_and_runs_timeout_hook( + forked_registration_context, +): + class TimeoutFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.timeout") + status: str = "pending" + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + timeout="50ms", + retry=Retry(max_attempts=2, initial_delay="1s", jitter="none"), + on_timeout="expired", + ) + async def slow(self): + await asyncio.sleep(5) + + @rx.event(durable=True, effect="none") + def expired(self): + self.status = "expired" + + async with WorkflowTestHarness(TimeoutFlow) as harness: + result = await harness.start(TimeoutFlow.slow()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.RETRYING + await harness.advance("1s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.steps[0].status is StepStatus.TIMED_OUT + assert snapshot.steps[0].attempts == 2 + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "expired"} + + +async def test_non_idempotent_failure_needs_attention(forked_registration_context): + class UnsafeFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.unsafe") + + @rx.event(durable=True, trigger=manual(), effect="non_idempotent_write") + def send_wire(self): + msg = "socket dropped mid-request" + raise ConnectionError(msg) + + async with WorkflowTestHarness(UnsafeFlow) as harness: + result = await harness.start(UnsafeFlow.send_wire()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + assert snapshot.steps[0].status is StepStatus.NEEDS_ATTENTION + assert snapshot.error is not None + assert "non-idempotent" in snapshot.error["reason"] + # Suspension is not terminal and nothing further executes. + await harness.advance("1h") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + + +async def test_durable_delay(forked_registration_context): + class DelayFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.delay") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + self.status = "waiting" + return after("2d", DelayFlow.follow_up) + + @rx.event(durable=True, effect="none") + def follow_up(self): + self.status = "done" + + async with WorkflowTestHarness(DelayFlow) as harness: + result = await harness.start(DelayFlow.begin()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert snapshot.state == {"status": "waiting"} + await harness.advance("1d") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + await harness.advance("1d") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "done"} + + +async def test_sequential_list_chain(forked_registration_context): + order = [] + + class ListFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.list") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + return [ListFlow.first, ListFlow.second] + + @rx.event(durable=True, effect="none") + def first(self): + order.append("first") + + @rx.event(durable=True, effect="none") + def second(self): + order.append("second") + + async with WorkflowTestHarness(ListFlow) as harness: + result = await harness.start(ListFlow.begin()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert order == ["first", "second"] + + +async def test_complete_tombstones_remaining_work(forked_registration_context): + class CompleteFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.complete") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + return [CompleteFlow.decide, CompleteFlow.never_runs] + + @rx.event(durable=True, effect="none") + def decide(self): + return complete(result={"answer": 42}) + + @rx.event(durable=True, effect="none") + def never_runs(self): + msg = "unreachable" + raise AssertionError(msg) + + async with WorkflowTestHarness(CompleteFlow) as harness: + result = await harness.start(CompleteFlow.begin()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"answer": 42} + assert snapshot.steps[2].status is StepStatus.CANCELLED + + +async def test_fail_and_needs_attention_controls(forked_registration_context): + class ControlFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.control") + mode: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self, mode: str): + self.mode = mode + if mode == "fail": + return fail("bad_invoice", details={"code": 402}) + return needs_attention("manual_review") + + async with WorkflowTestHarness(ControlFlow) as harness: + failed = await harness.start(ControlFlow.begin("fail")) + assert failed.run_id is not None + snapshot = await harness.get_run(failed.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error == {"reason": "bad_invoice", "details": {"code": 402}} + # The handler itself succeeded and its state patch was committed. + assert snapshot.state == {"mode": "fail"} + assert snapshot.steps[0].status is StepStatus.SUCCEEDED + + suspended = await harness.start(ControlFlow.begin("review")) + assert suspended.run_id is not None + snapshot = await harness.get_run(suspended.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + assert snapshot.state == {"mode": "review"} + + +async def test_cancel_while_waiting(forked_registration_context): + class CancelFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.cancel") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + return after("1d", CancelFlow.later) + + @rx.event(durable=True, effect="none") + def later(self): + pass + + async with WorkflowTestHarness(CancelFlow) as harness: + result = await harness.start(CancelFlow.begin()) + assert result.run_id is not None + assert await harness.cancel(result.run_id) + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.CANCELLED + assert snapshot.steps[1].status is StepStatus.CANCELLED + # Cancelling a terminal run reports False. + assert not await harness.cancel(result.run_id) + + +async def test_cancel_in_flight_attempt(forked_registration_context): + started = asyncio.Event() + release = asyncio.Event() + + class InflightFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.inflight") + + @rx.event(durable=True, trigger=manual(), effect="none") + async def hang(self): + started.set() + await release.wait() + + async with WorkflowTestHarness(InflightFlow) as harness: + result = await harness.kernel.start(InflightFlow.hang()) + assert result.run_id is not None + pump = asyncio.create_task(harness.run_until_idle()) + await asyncio.wait_for(started.wait(), timeout=2) + await harness.cancel(result.run_id) + await asyncio.wait_for(pump, timeout=2) + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.CANCELLED + assert snapshot.steps[0].status is StepStatus.CANCELLED + history = await harness.kernel.store.get_history(result.run_id) + assert HistoryEventType.ATTEMPT_CANCELLED in [event.type for event in history] + + +async def test_max_steps_bounds_chains(forked_registration_context): + class LoopFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.loop", max_steps=3) + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + return LoopFlow.again + + @rx.event(durable=True, effect="none") + def again(self): + return LoopFlow.again + + async with WorkflowTestHarness(LoopFlow) as harness: + result = await harness.start(LoopFlow.begin()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error is not None + assert snapshot.error["reason"] == "max_steps_exceeded" + + +async def test_run_timeout_deadline(forked_registration_context): + class DeadlineFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.deadline", run_timeout="1h") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + return after("2h", DeadlineFlow.too_late) + + @rx.event(durable=True, effect="none") + def too_late(self): + pass + + async with WorkflowTestHarness(DeadlineFlow) as harness: + result = await harness.start(DeadlineFlow.begin()) + assert result.run_id is not None + await harness.advance("2h") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.TIMED_OUT + assert snapshot.steps[1].status is StepStatus.CANCELLED + + +async def test_unserializable_state_fails_run(forked_registration_context): + class BadStateFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.badstate") + data: dict = {} + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + self.data = {"handle": object()} + + async with WorkflowTestHarness(BadStateFlow) as harness: + result = await harness.start(BadStateFlow.begin()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error is not None + assert snapshot.error["type"] == "WorkflowRuntimeError" + + +async def test_start_rejects_non_manual_roots(forked_registration_context): + class StartRules(rx.State): + __workflow__ = WorkflowConfig(id="kernel.startrules") + + @rx.event( + durable=True, trigger=webhook("stripe.payment_succeeded"), effect="none" + ) + def on_webhook(self): + pass + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + @rx.event(durable=True, effect="none") + def internal(self): + pass + + async with WorkflowTestHarness(StartRules) as harness: + with pytest.raises(WorkflowRuntimeError, match="manual"): + await harness.start(StartRules.on_webhook()) + with pytest.raises(WorkflowRuntimeError, match="manual"): + await harness.start(StartRules.internal()) + with pytest.raises(WorkflowRuntimeError, match="workflow"): + await harness.start(object()) + + +async def test_start_rejects_unregistered_class(forked_registration_context): + class Registered(rx.State): + __workflow__ = WorkflowConfig(id="kernel.registered") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + class Unregistered(rx.State): + __workflow__ = WorkflowConfig(id="kernel.unregistered") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + async with WorkflowTestHarness(Registered) as harness: + with pytest.raises(WorkflowRuntimeError, match="add_workflow"): + await harness.start(Unregistered.go()) + + +async def test_sqlite_recovery_resumes_retry_schedule( + forked_registration_context, tmp_path +): + calls = [] + + class DurableFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.durable") + status: str = "pending" + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=3, initial_delay="5s", jitter="none"), + ) + def sync(self): + calls.append(1) + if len(calls) < 2: + msg = "down" + raise TransientWorkflowError(msg) + self.status = "done" + + db_path = tmp_path / "workflow.db" + first_store = SqliteRunStore(db_path) + async with WorkflowTestHarness(DurableFlow, store=first_store) as harness: + result = await harness.start(DurableFlow.sync()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.RETRYING + resume_at = harness.now + first_store.close() + + # A new process opens the same database and resumes the pending retry. + second_store = SqliteRunStore(db_path) + async with WorkflowTestHarness( + DurableFlow, store=second_store, start_time=resume_at + 5 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "done"} + assert snapshot.steps[0].attempts == 1 + second_store.close() + + +async def test_definition_digest_mismatch_suspends_run( + forked_registration_context, tmp_path +): + class PinnedV1(rx.State): + __workflow__ = WorkflowConfig(id="kernel.pinned") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + return after("1h", PinnedV1.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + pass + + class PinnedV2(rx.State): + __workflow__ = WorkflowConfig(id="kernel.pinned") + + @rx.event(durable=True, trigger=manual(), effect="read", timeout="5s") + def begin(self): + return after("1h", PinnedV2.finish) + + @rx.event(durable=True, effect="read") + def finish(self): + pass + + db_path = tmp_path / "workflow.db" + first_store = SqliteRunStore(db_path) + async with WorkflowTestHarness(PinnedV1, store=first_store) as harness: + result = await harness.start(PinnedV1.begin()) + assert result.run_id is not None + resume_at = harness.now + first_store.close() + + second_store = SqliteRunStore(db_path) + async with WorkflowTestHarness( + PinnedV2, store=second_store, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + assert snapshot.error == {"reason": "definition_digest_mismatch"} + second_store.close() diff --git a/tests/units/workflow/test_store.py b/tests/units/workflow/test_store.py new file mode 100644 index 00000000000..4e1bd96ac8d --- /dev/null +++ b/tests/units/workflow/test_store.py @@ -0,0 +1,326 @@ +"""Tests for the workflow run stores.""" + +import pytest + +from reflex.workflow.records import ( + HistoryEventType, + RunRecord, + RunStatus, + StepRecord, + StepStatus, +) +from reflex.workflow.store import ( + MemoryRunStore, + SqliteRunStore, + StaleClaimError, + StepCompletion, +) + +NOW = 1_000_000.0 + + +@pytest.fixture(params=["memory", "sqlite"]) +def store(request, tmp_path): + """A run store of each implementation. + + Args: + request: The fixture request carrying the store kind. + tmp_path: Temporary directory for the SQLite database. + + Yields: + The store instance. + """ + if request.param == "memory": + yield MemoryRunStore() + else: + sqlite_store = SqliteRunStore(tmp_path / "workflow.db") + yield sqlite_store + sqlite_store.close() + + +def _run(run_id: str = "run1", **overrides) -> RunRecord: + defaults = { + "run_id": run_id, + "workflow_id": "billing.store_test", + "definition_digest": "digest", + "status": RunStatus.PENDING, + "state": {"n": 0}, + "state_version": 0, + "next_ordinal": 1, + "created_at": NOW, + "updated_at": NOW, + } + defaults.update(overrides) + return RunRecord(**defaults) + + +def _step(run_id: str = "run1", ordinal: int = 0, **overrides) -> StepRecord: + defaults = { + "run_id": run_id, + "ordinal": ordinal, + "handler_id": "go", + "status": StepStatus.READY, + "args": {}, + "origin": "root", + "created_at": NOW, + "updated_at": NOW, + } + defaults.update(overrides) + return StepRecord(**defaults) + + +_ADMIT_EVENTS = ((HistoryEventType.RUN_ADMITTED, {}),) + + +async def test_admit_and_load_round_trip(store): + created, run_id = await store.admit( + _run(request_key="key1", labels={"customer": "c1"}), _step(), _ADMIT_EVENTS + ) + assert created + assert run_id == "run1" + run = await store.get_run("run1") + assert run is not None + assert run.state == {"n": 0} + assert run.labels == {"customer": "c1"} + steps = await store.get_steps("run1") + assert [step.status for step in steps] == [StepStatus.READY] + history = await store.get_history("run1") + assert [event.type for event in history] == [HistoryEventType.RUN_ADMITTED] + assert history[0].seq == 1 + + +async def test_admit_deduplicates_on_request_key(store): + await store.admit(_run(request_key="key1"), _step(), _ADMIT_EVENTS) + created, run_id = await store.admit( + _run(run_id="run2", request_key="key1"), _step(run_id="run2"), _ADMIT_EVENTS + ) + assert not created + assert run_id == "run1" + assert await store.get_run("run2") is None + + +async def test_claim_respects_frontier_and_due_time(store): + await store.admit(_run(next_ordinal=3), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.RUNNING, + state={"n": 1}, + new_steps=( + _step(ordinal=1, handler_id="second", due_at=NOW + 60, origin="delay"), + _step(ordinal=2, handler_id="third", origin="chain"), + ), + next_ordinal=3, + ), + NOW, + ) + # Ordinal 1 is the frontier but not due yet; ordinal 2 must not overtake it. + assert await store.claim_next(NOW) is None + assert await store.next_due(NOW) == NOW + 60 + claim = await store.claim_next(NOW + 61) + assert claim is not None + assert claim.step.ordinal == 1 + assert claim.step.handler_id == "second" + + +async def test_commit_bumps_state_version_and_fences(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.RETRY_WAIT, + run_status=RunStatus.RETRYING, + state=None, + consume_attempt=True, + due_at=NOW + 5, + ), + NOW, + ) + run = await store.get_run("run1") + assert run is not None + assert run.state_version == 0 + steps = await store.get_steps("run1") + assert steps[0].attempts == 1 + assert steps[0].status is StepStatus.RETRY_WAIT + second_claim = await store.claim_next(NOW + 6) + assert second_claim is not None + assert second_claim.step.epoch == 2 + # The first claim is now stale and must not commit. + with pytest.raises(StaleClaimError): + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state={"n": 99}, + ), + NOW + 7, + ) + await store.commit( + second_claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state={"n": 2}, + ), + NOW + 7, + ) + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.COMPLETED + assert run.state == {"n": 2} + assert run.state_version == 1 + + +async def test_release_claim_returns_step(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW) + assert claim is not None + await store.release_claim( + claim, + status=StepStatus.CANCELLED, + events=((HistoryEventType.ATTEMPT_CANCELLED, {"ordinal": 0}),), + now=NOW, + ) + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.CANCELLED + # Releasing again is a no-op because the claim is stale. + await store.release_claim(claim, status=StepStatus.READY, events=(), now=NOW) + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.CANCELLED + + +async def test_cancel_control_and_finalize(store): + await store.admit(_run(next_ordinal=2), _step(), _ADMIT_EVENTS) + assert await store.control_pending(NOW) == () + assert await store.request_cancel("run1", NOW) + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.CANCELLING + pending = await store.control_pending(NOW) + assert [run.run_id for run in pending] == ["run1"] + assert await store.finalize_run( + "run1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW, + ) + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.CANCELLED + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.CANCELLED + history = await store.get_history("run1") + assert history[-1].type is HistoryEventType.RUN_CANCELLED + assert HistoryEventType.STEP_TOMBSTONED in [event.type for event in history] + # A terminal run cannot be cancelled or finalized again. + assert not await store.request_cancel("run1", NOW) + assert not await store.finalize_run( + "run1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW, + ) + + +async def test_finalize_refused_while_claimed(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW) + assert claim is not None + await store.request_cancel("run1", NOW) + assert await store.control_pending(NOW) == () + assert not await store.finalize_run( + "run1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW, + ) + + +async def test_deadline_makes_run_control_pending(store): + await store.admit(_run(deadline=NOW + 100), _step(due_at=NOW + 500), _ADMIT_EVENTS) + assert await store.control_pending(NOW) == () + assert await store.claim_next(NOW + 200) is None + pending = await store.control_pending(NOW + 200) + assert [run.run_id for run in pending] == ["run1"] + + +async def test_recover_orphans_consumes_recovery_budget(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW) + assert claim is not None + recovered = await store.recover_orphans(NOW + 10, max_recoveries=2) + assert recovered == 1 + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.RECOVERY_WAIT + assert steps[0].recoveries == 1 + # The stale claim cannot commit after recovery. + with pytest.raises(StaleClaimError): + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state={}, + ), + NOW + 11, + ) + + +async def test_recover_orphans_exhaustion_fails_run(store): + await store.admit(_run(), _step(recoveries=2), _ADMIT_EVENTS) + claim = await store.claim_next(NOW) + assert claim is not None + await store.recover_orphans(NOW + 10, max_recoveries=2) + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.FAILED + assert run.error == {"reason": "recovery_budget_exhausted"} + + +async def test_append_events_assigns_sequence(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + await store.append_events( + "run1", ((HistoryEventType.ATTEMPT_STARTED, {"ordinal": 0}),), NOW + ) + history = await store.get_history("run1") + assert [event.seq for event in history] == [1, 2] + assert history[-1].type is HistoryEventType.ATTEMPT_STARTED + + +async def test_sqlite_persistence_across_reopen(tmp_path): + db_path = tmp_path / "workflow.db" + first = SqliteRunStore(db_path) + await first.admit(_run(request_key="key1"), _step(), _ADMIT_EVENTS) + claim = await first.claim_next(NOW) + assert claim is not None + first.close() + + second = SqliteRunStore(db_path) + try: + run = await second.get_run("run1") + assert run is not None + steps = await second.get_steps("run1") + assert steps[0].status is StepStatus.CLAIMED + # Dedupe state survives restarts. + created, run_id = await second.admit( + _run(run_id="run2", request_key="key1"), + _step(run_id="run2"), + _ADMIT_EVENTS, + ) + assert not created + assert run_id == "run1" + # The orphaned claim recovers on the new process. + assert await second.recover_orphans(NOW + 5, max_recoveries=10) == 1 + steps = await second.get_steps("run1") + assert steps[0].status is StepStatus.RECOVERY_WAIT + finally: + second.close() From 9c576d8fc27995008dbc93380c2ab22eaefc96b9 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 13:26:51 -0700 Subject: [PATCH 002/121] Fence workflow claims with renewable leases A second worker starting while the first was mid-attempt reclaimed the live claim and executed the same step concurrently: recovery treated every CLAIMED row as an orphan left by a dead process. Proven with two OS processes against one SQLite file, where the peer ran the handler a second time and the original worker's commit was then discarded as stale. Claims now carry a lease. claim_next stamps lease_expires_at, the executing kernel renews it in the background while the attempt runs, and recovery reclaims only claims whose lease has lapsed, so a peer mid-attempt is never disturbed. Lease loss consumes the infrastructure recovery budget, never a business attempt, and is recorded as attempt_abandoned evidence. Also: - Recovery is now periodic rather than startup-only, so a peer that dies long after this worker booted is still reclaimed. - Renewal cadence runs on real time while expiry is measured on the injected clock, so virtual-time tests stay deterministic and a jumping clock cannot expire a live attempt (recover() renews own claims before sweeping). - Cancellation is now disambiguated three ways: a lost lease abandons the attempt, an operator cancel releases the step as CANCELLED, and any other cancellation (worker shutdown) re-raises and leaves the step claimed for lease recovery. Previously every cancellation terminally cancelled the step and wedged the run at RUNNING with no path forward. - SqliteRunStore migrates databases written before this change; a step left claimed by the previous binary has no lease and is a genuine orphan. - The worker loop no longer dies on a transient store error. Crash recovery is no longer instant: it is bounded below by the lease duration (default 30s, sweeping every 15s). That is the price of not double-executing. Multi-worker SQLite remains unsupported and is now documented as one worker process per database file: the store's calls are synchronous, so cross-process write contention blocks the caller's event loop including its own renewals. --- news/workflow-leases.bugfix.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 2 + reflex/workflow/kernel.py | 316 ++++++++++++++--- reflex/workflow/records.py | 13 +- reflex/workflow/runtime.py | 18 +- reflex/workflow/store.py | 221 ++++++++++-- reflex/workflow/testing.py | 13 +- tests/units/workflow/test_lease.py | 334 ++++++++++++++++++ tests/units/workflow/test_store.py | 145 +++++++- 9 files changed, 989 insertions(+), 74 deletions(-) create mode 100644 news/workflow-leases.bugfix.md create mode 100644 tests/units/workflow/test_lease.py diff --git a/news/workflow-leases.bugfix.md b/news/workflow-leases.bugfix.md new file mode 100644 index 00000000000..ed3419d9c45 --- /dev/null +++ b/news/workflow-leases.bugfix.md @@ -0,0 +1 @@ +Workflow claims now carry a renewable lease, so a second worker no longer reclaims and re-executes a step another worker is running. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index ae5832f2fcc..95170e5767b 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -23,6 +23,8 @@ DEFAULT_MAX_RECOVERIES: Final = 10 +DEFAULT_LEASE_DURATION: Final = 30.0 + DurationLike = str | int | float | timedelta _DURATION_RE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*(ms|s|m|h|d)\s*$") diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 5ab91cc750c..59aab3b0d10 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -3,8 +3,9 @@ The kernel admits runs, claims the due frontier step of each run's mailbox, executes the durable handler against a hydrated run-state instance, and atomically commits the state patch together with the successor slots the -handler returned. Retries, timeouts, lifecycle hooks, cancellation drain, and -crash recovery are decided here and made durable by the store. +handler returned. Retries, timeouts, lifecycle hooks, cancellation drain, +claim-lease renewal, and crash recovery are decided here and made durable by +the store. """ from __future__ import annotations @@ -19,8 +20,10 @@ from pydantic import TypeAdapter from reflex_base.event.processor.base_state_processor import _transform_event_payload +from reflex_base.utils import console from reflex_base.utils.exceptions import WorkflowRuntimeError from reflex_base.workflow import ( + DEFAULT_LEASE_DURATION, DEFAULT_MAX_RECOVERIES, After, CompleteRun, @@ -52,6 +55,10 @@ DEFAULT_POLL_INTERVAL = 0.25 +LEASE_RENEW_FRACTION = 1 / 3 + +RECOVERY_INTERVAL_FRACTION = 1 / 2 + def _error_payload(error: BaseException) -> dict[str, Any]: """Build a JSON-compatible error payload from an exception. @@ -104,6 +111,30 @@ def __init__( self.origin = origin +class _Lease: + """The live claim lease of one in-flight attempt. + + Attributes: + claim: The claim being kept alive. + attempt: The task running the handler, once created. + renewer: The background task extending the lease. + lost: Whether the store reported the claim was fenced. + """ + + __slots__ = ("attempt", "claim", "lost", "renewer") + + def __init__(self, claim: Claim): + """Initialize the lease. + + Args: + claim: The claim being kept alive. + """ + self.claim = claim + self.attempt: asyncio.Task | None = None + self.renewer: asyncio.Task | None = None + self.lost = False + + class WorkflowKernel: """Executes durable workflow runs against a run store.""" @@ -116,6 +147,9 @@ def __init__( rng: Callable[[], float] = random.random, poll_interval: float = DEFAULT_POLL_INTERVAL, max_recoveries: int = DEFAULT_MAX_RECOVERIES, + lease_duration: float = DEFAULT_LEASE_DURATION, + lease_renew_interval: float | None = None, + recovery_interval: float | None = None, ): """Initialize the kernel. @@ -126,6 +160,16 @@ def __init__( rng: Uniform [0, 1) source used for retry jitter. poll_interval: Worker sleep bound between due-time checks. max_recoveries: Infrastructure recovery budget per logical step. + lease_duration: Seconds a claim survives without renewal before + recovery may reclaim it. + lease_renew_interval: Real seconds between lease renewals; defaults + to a third of ``lease_duration``. + recovery_interval: Seconds between recovery sweeps in the + background worker; defaults to half of ``lease_duration``. + + Raises: + WorkflowRuntimeError: If the store cannot renew leases, or the + lease timings are inconsistent. """ self._definitions: dict[str, WorkflowDefinition] = { defn.workflow_id: defn for defn in definitions @@ -138,11 +182,40 @@ def __init__( self._rng = rng self._poll_interval = poll_interval self._max_recoveries = max_recoveries + if not hasattr(store, "renew_lease"): + msg = ( + f"{type(store).__name__} does not implement renew_lease; a run " + "store must renew claim leases or recovery will reclaim live " + "claims and execute steps twice." + ) + raise WorkflowRuntimeError(msg) + renew = ( + lease_renew_interval + if lease_renew_interval is not None + else lease_duration * LEASE_RENEW_FRACTION + ) + recovery = ( + recovery_interval + if recovery_interval is not None + else lease_duration * RECOVERY_INTERVAL_FRACTION + ) + if lease_duration <= 0 or not 0 < renew < lease_duration or recovery <= 0: + msg = ( + "Lease timings must satisfy 0 < lease_renew_interval < " + "lease_duration and recovery_interval > 0, got " + f"{renew} / {lease_duration} / {recovery}." + ) + raise WorkflowRuntimeError(msg) + self._lease_duration = lease_duration + self._lease_renew_interval = renew + self._recovery_interval = recovery self._field_adapters: dict[tuple[str, str], TypeAdapter] = {} self._inflight: dict[str, asyncio.Task] = {} + self._leases: dict[str, _Lease] = {} + self._next_recovery_at = 0.0 + self._worker_id = uuid.uuid4().hex self._wakeup = asyncio.Event() self._worker: asyncio.Task | None = None - self._recovered = False @property def store(self) -> RunStore: @@ -851,6 +924,139 @@ def _success_completion( events=tuple(events), ) + def _acquire_lease(self, claim: Claim) -> _Lease: + """Register an in-flight claim and start renewing its lease. + + Args: + claim: The claim to keep alive. + + Returns: + The lease handle. + """ + lease = _Lease(claim) + self._leases[claim.run.run_id] = lease + lease.renewer = asyncio.ensure_future(self._renew_forever(lease)) + return lease + + async def _renew(self, lease: _Lease) -> None: + """Extend one lease, abandoning the attempt when the store fences it. + + A store error is transient: the lease is left alone and the next + renewal retries, which tolerates one lost round-trip before the lease + could lapse. + + Args: + lease: The lease to extend. + """ + if lease.lost: + return + try: + held = await self._store.renew_lease( + lease.claim, self._clock(), lease_duration=self._lease_duration + ) + except Exception as err: + console.debug(f"Workflow lease renewal failed, retrying: {err}") + return + if not held: + self._lose_lease(lease) + + async def _renew_forever(self, lease: _Lease) -> None: + """Renew a lease on a real-time cadence until it ends or is lost. + + The cadence is real time so renewal makes progress under any injected + clock, while the expiry written is read from the injected clock so + virtual time alone decides when a lease has lapsed. + + Args: + lease: The lease to renew. + """ + while not lease.lost: + await asyncio.sleep(self._lease_renew_interval) + await self._renew(lease) + + def _lose_lease(self, lease: _Lease) -> None: + """Mark a lease fenced and stop the attempt it was covering. + + Args: + lease: The lease that was lost. + """ + lease.lost = True + if lease.attempt is not None: + lease.attempt.cancel() + + async def _release_lease(self, lease: _Lease) -> None: + """Stop renewing a lease and forget the in-flight claim. + + Args: + lease: The lease to release. + + Raises: + asyncio.CancelledError: If this task is cancelled while waiting for + the renewer to stop. + """ + self._leases.pop(lease.claim.run.run_id, None) + renewer, lease.renewer = lease.renewer, None + if renewer is None: + return + renewer.cancel() + try: + await renewer + except asyncio.CancelledError: + if not renewer.cancelled(): + raise + + async def _renew_leases(self) -> None: + """Extend every lease this kernel holds before a recovery sweep. + + Recovery reclaims any claim whose lease has lapsed, including one this + kernel is executing when an injected clock jumps past its expiry. + Renewing first makes a live attempt unstealable by its own process. + """ + for lease in list(self._leases.values()): + await self._renew(lease) + + async def _cancel_requested(self, run_id: str) -> bool: + """Whether a run carries cancellation intent. + + Args: + run_id: The run to check. + + Returns: + True when the run exists and cancellation was requested. + """ + run = await self._store.get_run(run_id) + return run is not None and run.cancel_requested + + async def _record_abandoned( + self, claim: Claim, handler: HandlerDefinition, reason: str + ) -> None: + """Record that an attempt lost its claim and committed nothing. + + The event is appended outside the fence: the row belongs to another + worker now, and history is append-only evidence, not state. + + Args: + claim: The fenced claim. + handler: The handler that was executing. + reason: Why the claim was lost. + """ + await self._store.append_events( + claim.run.run_id, + ( + ( + HistoryEventType.ATTEMPT_ABANDONED, + { + "ordinal": claim.step.ordinal, + "epoch": claim.step.epoch, + "worker": self._worker_id, + "effect": handler.effect, + "reason": reason, + }, + ), + ), + self._clock(), + ) + async def _execute_claim(self, claim: Claim) -> None: """Execute one claimed attempt and commit its outcome. @@ -889,33 +1095,46 @@ async def _execute_claim(self, claim: Claim) -> None: "ordinal": claim.step.ordinal, "handler_id": handler.id, "attempt": claim.step.attempts + 1, + "epoch": claim.step.epoch, "effect": handler.effect, }, ), ), now, ) + lease = self._acquire_lease(claim) try: - instance = self._hydrate(defn, claim.run.state) - value = await self._invoke(handler, instance, claim.step.args) + try: + instance = self._hydrate(defn, claim.run.state) + lease.attempt = asyncio.ensure_future( + self._invoke(handler, instance, claim.step.args) + ) + value = await lease.attempt + finally: + await self._release_lease(lease) successors, control = self._interpret_return(defn, value) state = self._snapshot(defn, instance) completion = self._success_completion( defn, claim, steps, state, successors, control, self._clock() ) except asyncio.CancelledError: - await self._store.release_claim( - claim, - status=StepStatus.CANCELLED, - events=( - ( - HistoryEventType.ATTEMPT_CANCELLED, - {"ordinal": claim.step.ordinal}, + if lease.lost: + await self._record_abandoned(claim, handler, "lease_lost") + return + if await self._cancel_requested(claim.run.run_id): + await self._store.release_claim( + claim, + status=StepStatus.CANCELLED, + events=( + ( + HistoryEventType.ATTEMPT_CANCELLED, + {"ordinal": claim.step.ordinal}, + ), ), - ), - now=self._clock(), - ) - return + now=self._clock(), + ) + return + raise except TimeoutError as err: completion = self._failure_completion( defn, handler, claim, steps, err, timed_out=True, now=self._clock() @@ -927,7 +1146,7 @@ async def _execute_claim(self, claim: Claim) -> None: try: await self._store.commit(claim, completion, self._clock()) except StaleClaimError: - return + await self._record_abandoned(claim, handler, "fenced_at_commit") async def _tick(self) -> bool: """Run one scheduling round. @@ -960,7 +1179,7 @@ async def _tick(self) -> bool: ) or progressed ) - claim = await self._store.claim_next(now) + claim = await self._store.claim_next(now, lease_duration=self._lease_duration) if claim is not None: task = asyncio.ensure_future(self._execute_claim(claim)) self._inflight[claim.run.run_id] = task @@ -972,51 +1191,68 @@ async def _tick(self) -> bool: return progressed async def recover(self) -> int: - """Recover orphaned claims left by a previous process. + """Renew this kernel's live claims, then reclaim expired ones. + + A claim is reclaimable only once its lease has lapsed, so a peer that + is mid-attempt is never disturbed and crash recovery is delayed by up + to one lease. Returns: The number of steps recovered. """ - self._recovered = True - return await self._store.recover_orphans(self._clock(), self._max_recoveries) + await self._renew_leases() + now = self._clock() + self._next_recovery_at = now + self._recovery_interval + return await self._store.recover_orphans(now, self._max_recoveries) async def run_until_idle(self) -> None: """Process work until nothing is claimable at the current clock time. - Scheduled future work (retry backoff, ``rx.after`` delays) stays - pending; advance the clock and call again to run it. + Each call first sweeps for claims whose lease has expired, so advancing + the clock past a dead worker's lease reclaims its step. Scheduled + future work (retry backoff, ``rx.after`` delays) stays pending; advance + the clock and call again to run it. """ - if not self._recovered: - await self.recover() + await self.recover() while await self._tick(): pass async def _worker_loop(self) -> None: """Process work continuously until the kernel is closed.""" while True: - progressed = await self._tick() - if progressed: - continue - now = self._clock() - due = await self._store.next_due(now) - delay = self._poll_interval - if due is not None: - delay = min(delay, max(due - now, 0.0)) - if delay <= 0: - continue - self._wakeup.clear() - with contextlib.suppress(TimeoutError): - await asyncio.wait_for(self._wakeup.wait(), timeout=delay) + try: + if self._clock() >= self._next_recovery_at: + await self.recover() + if await self._tick(): + continue + now = self._clock() + due = await self._store.next_due(now) + delay = min(self._poll_interval, max(self._next_recovery_at - now, 0.0)) + if due is not None: + delay = min(delay, max(due - now, 0.0)) + if delay <= 0: + continue + self._wakeup.clear() + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(self._wakeup.wait(), timeout=delay) + except Exception as err: + console.error(f"Workflow worker error, retrying: {err}") + await asyncio.sleep(self._poll_interval) async def start_worker(self) -> None: - """Start the background worker after recovering orphaned claims.""" + """Start the background worker, which recovers expired claims as it runs.""" if self._worker is not None: return await self.recover() self._worker = asyncio.create_task(self._worker_loop()) async def aclose(self) -> None: - """Stop the background worker, leaving in-flight claims recoverable.""" + """Stop the background worker. + + An in-flight attempt is cancelled and its step is left claimed, so it + is reclaimed once its lease expires rather than being recorded as a + deliberate cancellation. + """ if self._worker is None: return self._worker.cancel() diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index f8c1ce612cf..b8b1929788a 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -33,10 +33,10 @@ class RunStatus(str, enum.Enum): class StepStatus(str, enum.Enum): """Lifecycle status of one logical step in a run's mailbox. - Successor slots are created ``READY`` because the in-process kernel is the - single writer and the frontier scan already enforces mailbox order; a - distributed kernel adapter would hold successors in a blocked state until - their predecessor commit is visible. + Successor slots are created ``READY`` because a slot can only be claimed + once it is the mailbox frontier, which requires its predecessor's commit to + be durable; a distributed kernel adapter would hold successors in a blocked + state until that commit is visible. """ READY = "READY" @@ -75,6 +75,7 @@ class HistoryEventType(str, enum.Enum): ATTEMPT_FAILED = "attempt_failed" ATTEMPT_TIMED_OUT = "attempt_timed_out" ATTEMPT_CANCELLED = "attempt_cancelled" + ATTEMPT_ABANDONED = "attempt_abandoned" STEP_RETRY_SCHEDULED = "step_retry_scheduled" STEP_RECOVERED = "step_recovered" STEP_TOMBSTONED = "step_tombstoned" @@ -139,6 +140,9 @@ class StepRecord: recoveries: Infrastructure recoveries consumed so far. due_at: Earliest epoch time the step may be claimed. epoch: Fencing token, incremented on every claim. + lease_expires_at: Epoch time this claim's lease lapses; 0 when the step + is not claimed. A claim whose lease has lapsed is treated as + orphaned and is reclaimed by recovery, never by a direct claim. error: Last recorded attempt error payload. origin: How the slot was allocated (root, chain, delay, or hook). created_at: Allocation time in epoch seconds. @@ -154,6 +158,7 @@ class StepRecord: recoveries: int = 0 due_at: float = 0.0 epoch: int = 0 + lease_expires_at: float = 0.0 error: dict[str, Any] | None = None origin: Literal["root", "chain", "delay", "hook"] = "chain" created_at: float = 0.0 diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index c5ed99683c2..3adfb1c8558 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -17,6 +17,7 @@ from reflex_base.registry import RegistrationContext from reflex_base.utils.exceptions import WorkflowDefinitionError, WorkflowRuntimeError +from reflex_base.workflow import DEFAULT_LEASE_DURATION from reflex.workflow.definition import WorkflowDefinition, compile_workflow from reflex.workflow.kernel import DEFAULT_POLL_INTERVAL, WorkflowKernel @@ -70,6 +71,9 @@ def __init__( clock: Callable[[], float] = time.time, rng: Callable[[], float] = random.random, poll_interval: float = DEFAULT_POLL_INTERVAL, + lease_duration: float = DEFAULT_LEASE_DURATION, + lease_renew_interval: float | None = None, + recovery_interval: float | None = None, ): """Initialize the runtime. @@ -79,11 +83,18 @@ def __init__( clock: Epoch-seconds time source; injectable for virtual time. rng: Uniform [0, 1) source used for retry jitter. poll_interval: Worker sleep bound between due-time checks. + lease_duration: Seconds a claim survives without renewal before + recovery may reclaim it. + lease_renew_interval: Real seconds between lease renewals. + recovery_interval: Seconds between recovery sweeps. """ self._store = store self._clock = clock self._rng = rng self._poll_interval = poll_interval + self._lease_duration = lease_duration + self._lease_renew_interval = lease_renew_interval + self._recovery_interval = recovery_interval self._definitions: dict[str, WorkflowDefinition] = {} self._classes: dict[type, str] = {} self._kernel: WorkflowKernel | None = None @@ -152,7 +163,7 @@ def kernel(self) -> WorkflowKernel: return self._kernel async def startup(self, *, start_worker: bool = True) -> None: - """Build the kernel, recover orphaned claims, and start processing. + """Build the kernel, reclaim expired claims, and start processing. Args: start_worker: Whether to launch the background worker; tests pump @@ -168,6 +179,9 @@ async def startup(self, *, start_worker: bool = True) -> None: clock=self._clock, rng=self._rng, poll_interval=self._poll_interval, + lease_duration=self._lease_duration, + lease_renew_interval=self._lease_renew_interval, + recovery_interval=self._recovery_interval, ) if start_worker: await self._kernel.start_worker() @@ -175,7 +189,7 @@ async def startup(self, *, start_worker: bool = True) -> None: await self._kernel.recover() async def shutdown(self) -> None: - """Stop the worker, leaving in-flight claims recoverable on restart.""" + """Stop the worker; an in-flight claim is reclaimed after its lease expires.""" if self._kernel is not None: await self._kernel.aclose() self._kernel = None diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 1572be24c0f..c196858db06 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -1,12 +1,16 @@ """Durable run stores for the workflow kernel. A store is the single authority for run state: admission with idempotent -request keys, the ordered per-run mailbox, claim fencing, and the atomic step -commit that persists a state patch together with its successor slots. The -kernel decides what should happen; the store makes it durable atomically. +request keys, the ordered per-run mailbox, claim fencing and leasing, and the +atomic step commit that persists a state patch together with its successor +slots. The kernel decides what should happen; the store makes it durable +atomically. ``MemoryRunStore`` backs tests and the harness. ``SqliteRunStore`` provides -crash-safe persistence on a single machine using the standard library. +crash-safe persistence on a single machine using the standard library. Run +exactly one worker process per database file: its calls are synchronous and +cross-process write contention blocks the caller's event loop, which is +hostile to lease renewal. """ from __future__ import annotations @@ -16,9 +20,10 @@ import json import sqlite3 import threading -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, Final, Protocol from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import DEFAULT_LEASE_DURATION from reflex.workflow.records import ( CLAIMABLE_STEP_STATUSES, @@ -109,17 +114,49 @@ async def admit( """ ... - async def claim_next(self, now: float) -> Claim | None: + async def claim_next( + self, now: float, *, lease_duration: float = DEFAULT_LEASE_DURATION + ) -> Claim | None: """Claim the due frontier step of some runnable run. + The claim carries a lease expiring at ``now + lease_duration``. The + executing kernel must renew it through ``renew_lease``; a claim whose + lease lapses is reclaimed by ``recover_orphans``. + Args: now: Current time in epoch seconds. + lease_duration: Seconds of renewal silence tolerated before the + claim is treated as orphaned. Returns: A fenced claim, or None when nothing is claimable right now. """ ... + async def renew_lease( + self, + claim: Claim, + now: float, + *, + lease_duration: float = DEFAULT_LEASE_DURATION, + ) -> bool: + """Extend a live claim's lease without transitioning the step. + + Renewal is a liveness signal only: it must not change the step's + status, fencing epoch, or any committed state, and it never consumes a + budget. + + Args: + claim: The claim being renewed. + now: Current time in epoch seconds. + lease_duration: Seconds to extend the lease from ``now``. + + Returns: + True if the claim still owns its step and the lease was extended; + False if the claim was fenced and the attempt must be abandoned. + """ + ... + async def commit( self, claim: Claim, completion: StepCompletion, now: float ) -> None: @@ -219,10 +256,13 @@ async def finalize_run( ... async def recover_orphans(self, now: float, max_recoveries: int) -> int: - """Recover steps left claimed by a previous process. + """Recover claims whose lease has expired. - Each orphan consumes one infrastructure recovery and becomes claimable - again; a step over budget fails its run. + A step is orphaned when it is CLAIMED and its lease lapsed at or before + ``now``: whoever held it stopped renewing. Each orphan consumes one + infrastructure recovery and becomes claimable again; a step over budget + fails its run. A claim with a live lease is left alone, so a peer that + is mid-attempt is never disturbed. Args: now: Current time in epoch seconds. @@ -297,6 +337,19 @@ def _run_is_runnable(run: RunRecord, now: float) -> bool: ) +def _lease_expired(step: StepRecord, now: float) -> bool: + """Whether a claimed step's lease has lapsed and it may be recovered. + + Args: + step: The step record. + now: Current time in epoch seconds. + + Returns: + True when the step is claimed and its lease expired at or before now. + """ + return step.status is StepStatus.CLAIMED and step.lease_expires_at <= now + + def _frontier(steps: Iterable[StepRecord]) -> StepRecord | None: """Find the lowest-ordinal unresolved step. @@ -376,11 +429,15 @@ async def admit( self._append_events(run.run_id, events, run.created_at) return True, run.run_id - async def claim_next(self, now: float) -> Claim | None: + async def claim_next( + self, now: float, *, lease_duration: float = DEFAULT_LEASE_DURATION + ) -> Claim | None: """Claim the due frontier step of some runnable run. Args: now: Current time in epoch seconds. + lease_duration: Seconds of renewal silence tolerated before the + claim is treated as orphaned. Returns: A fenced claim, or None when nothing is claimable right now. @@ -401,6 +458,7 @@ async def claim_next(self, now: float) -> Claim | None: frontier, status=StepStatus.CLAIMED, epoch=frontier.epoch + 1, + lease_expires_at=now + lease_duration, updated_at=now, ) steps[claimed.ordinal] = claimed @@ -441,6 +499,34 @@ def _check_claim(self, claim: Claim) -> tuple[RunRecord, list[StepRecord]]: raise StaleClaimError(msg) return run, steps + async def renew_lease( + self, + claim: Claim, + now: float, + *, + lease_duration: float = DEFAULT_LEASE_DURATION, + ) -> bool: + """Extend a live claim's lease without transitioning the step. + + Args: + claim: The claim being renewed. + now: Current time in epoch seconds. + lease_duration: Seconds to extend the lease from ``now``. + + Returns: + True if the claim still owns its step; False if it was fenced. + """ + async with self._lock: + try: + _, steps = self._check_claim(claim) + except StaleClaimError: + return False + step = steps[claim.step.ordinal] + steps[step.ordinal] = dataclasses.replace( + step, lease_expires_at=now + lease_duration + ) + return True + async def commit( self, claim: Claim, completion: StepCompletion, now: float ) -> None: @@ -459,6 +545,7 @@ async def commit( status=completion.step_status, attempts=step.attempts + (1 if completion.consume_attempt else 0), due_at=completion.due_at if completion.due_at is not None else 0.0, + lease_expires_at=0.0, error=completion.step_error, updated_at=now, ) @@ -511,7 +598,7 @@ async def release_claim( return step = steps[claim.step.ordinal] steps[step.ordinal] = dataclasses.replace( - step, status=status, updated_at=now + step, status=status, lease_expires_at=0.0, updated_at=now ) self._append_events(run.run_id, events, now) @@ -645,7 +732,7 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: continue steps = self._steps[run.run_id] for step in list(steps): - if step.status is not StepStatus.CLAIMED: + if not _lease_expired(step, now): continue recovered += 1 if step.recoveries + 1 > max_recoveries: @@ -653,6 +740,7 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: step, status=StepStatus.FAILED, recoveries=step.recoveries + 1, + lease_expires_at=0.0, error={"reason": "recovery_budget_exhausted"}, updated_at=now, ) @@ -678,6 +766,7 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: status=StepStatus.RECOVERY_WAIT, recoveries=step.recoveries + 1, due_at=now, + lease_expires_at=0.0, updated_at=now, ) self._append_events( @@ -776,6 +865,7 @@ async def next_due(self, now: float) -> float | None: recoveries INTEGER NOT NULL DEFAULT 0, due_at REAL NOT NULL DEFAULT 0, epoch INTEGER NOT NULL DEFAULT 0, + lease_expires_at REAL NOT NULL DEFAULT 0, error TEXT, origin TEXT NOT NULL, created_at REAL NOT NULL, @@ -799,6 +889,13 @@ async def next_due(self, now: float) -> float | None: CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs (status); """ +_STEP_MIGRATIONS: Final = ( + ( + "lease_expires_at", + "ALTER TABLE workflow_steps ADD COLUMN lease_expires_at REAL NOT NULL DEFAULT 0", + ), +) + def _dump(value: Any) -> str | None: """Serialize an optional JSON payload column. @@ -871,6 +968,7 @@ def _step_from_row(row: sqlite3.Row) -> StepRecord: recoveries=row["recoveries"], due_at=row["due_at"], epoch=row["epoch"], + lease_expires_at=row["lease_expires_at"], error=_load(row["error"]), origin=row["origin"], created_at=row["created_at"], @@ -894,6 +992,34 @@ def __init__(self, db_path: str | Path): self._db.execute("PRAGMA journal_mode=WAL") self._db.execute("PRAGMA synchronous=NORMAL") self._db.executescript(_SCHEMA) + self._migrate() + + def _migrate(self) -> None: + """Add columns and indexes missing from databases created by older versions. + + The check and the alter run in one immediate transaction so two + processes opening the same file cannot both attempt it. Rows predating + a column take its default, which for ``lease_expires_at`` means an + already-lapsed lease: a step left claimed by the previous binary is a + genuine orphan and is recovered on the first recovery pass. + """ + self._db.execute("BEGIN IMMEDIATE") + try: + columns = { + row["name"] + for row in self._db.execute("PRAGMA table_info(workflow_steps)") + } + for name, statement in _STEP_MIGRATIONS: + if name not in columns: + self._db.execute(statement) + self._db.execute( + "CREATE INDEX IF NOT EXISTS idx_workflow_steps_lease" + " ON workflow_steps (status, lease_expires_at)" + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise def close(self) -> None: """Close the backing database connection.""" @@ -933,8 +1059,9 @@ def _insert_step(self, step: StepRecord) -> None: """ self._db.execute( "INSERT INTO workflow_steps (run_id, ordinal, handler_id, status, args," - " attempts, recoveries, due_at, epoch, error, origin, created_at," - " updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + " attempts, recoveries, due_at, epoch, lease_expires_at, error, origin," + " created_at, updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( step.run_id, step.ordinal, @@ -945,6 +1072,7 @@ def _insert_step(self, step: StepRecord) -> None: step.recoveries, step.due_at, step.epoch, + step.lease_expires_at, _dump(step.error), step.origin, step.created_at, @@ -1032,11 +1160,15 @@ async def admit( raise return True, run.run_id - async def claim_next(self, now: float) -> Claim | None: + async def claim_next( + self, now: float, *, lease_duration: float = DEFAULT_LEASE_DURATION + ) -> Claim | None: """Claim the due frontier step of some runnable run. Args: now: Current time in epoch seconds. + lease_duration: Seconds of renewal silence tolerated before the + claim is treated as orphaned. Returns: A fenced claim, or None when nothing is claimable right now. @@ -1067,14 +1199,17 @@ async def claim_next(self, now: float) -> Claim | None: frontier, status=StepStatus.CLAIMED, epoch=frontier.epoch + 1, + lease_expires_at=now + lease_duration, updated_at=now, ) self._db.execute( "UPDATE workflow_steps SET status = ?, epoch = ?," - " updated_at = ? WHERE run_id = ? AND ordinal = ?", + " lease_expires_at = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", ( claimed.status.value, claimed.epoch, + claimed.lease_expires_at, now, claimed.run_id, claimed.ordinal, @@ -1123,6 +1258,42 @@ def _check_claim(self, claim: Claim) -> None: ) raise StaleClaimError(msg) + async def renew_lease( + self, + claim: Claim, + now: float, + *, + lease_duration: float = DEFAULT_LEASE_DURATION, + ) -> bool: + """Extend a live claim's lease without transitioning the step. + + Args: + claim: The claim being renewed. + now: Current time in epoch seconds. + lease_duration: Seconds to extend the lease from ``now``. + + Returns: + True if the claim still owns its step; False if it was fenced. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + try: + self._check_claim(claim) + except StaleClaimError: + self._db.execute("ROLLBACK") + return False + self._db.execute( + "UPDATE workflow_steps SET lease_expires_at = ?" + " WHERE run_id = ? AND ordinal = ?", + (now + lease_duration, claim.run.run_id, claim.step.ordinal), + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True + async def commit( self, claim: Claim, completion: StepCompletion, now: float ) -> None: @@ -1139,7 +1310,7 @@ async def commit( self._check_claim(claim) self._db.execute( "UPDATE workflow_steps SET status = ?, attempts = attempts + ?," - " due_at = ?, error = ?, updated_at = ?" + " due_at = ?, lease_expires_at = 0, error = ?, updated_at = ?" " WHERE run_id = ? AND ordinal = ?", ( completion.step_status.value, @@ -1218,8 +1389,8 @@ async def release_claim( self._db.execute("ROLLBACK") return self._db.execute( - "UPDATE workflow_steps SET status = ?, updated_at = ?" - " WHERE run_id = ? AND ordinal = ?", + "UPDATE workflow_steps SET status = ?, lease_expires_at = 0," + " updated_at = ? WHERE run_id = ? AND ordinal = ?", (status.value, now, claim.run.run_id, claim.step.ordinal), ) self._append_events(claim.run.run_id, events, now) @@ -1388,9 +1559,9 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: rows = self._db.execute( "SELECT s.* FROM workflow_steps s" " JOIN workflow_runs r ON r.run_id = s.run_id" - " WHERE s.status = ? AND r.status NOT IN" - f" ({','.join('?' * len(terminal))})", - (StepStatus.CLAIMED.value, *terminal), + " WHERE s.status = ? AND s.lease_expires_at <= ?" + f" AND r.status NOT IN ({','.join('?' * len(terminal))})", + (StepStatus.CLAIMED.value, now, *terminal), ).fetchall() recovered = 0 for row in rows: @@ -1399,7 +1570,8 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: if step.recoveries + 1 > max_recoveries: self._db.execute( "UPDATE workflow_steps SET status = ?, recoveries = ?," - " error = ?, updated_at = ? WHERE run_id = ? AND ordinal = ?", + " lease_expires_at = 0, error = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", ( StepStatus.FAILED.value, step.recoveries + 1, @@ -1432,7 +1604,8 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: else: self._db.execute( "UPDATE workflow_steps SET status = ?, recoveries = ?," - " due_at = ?, updated_at = ? WHERE run_id = ? AND ordinal = ?", + " due_at = ?, lease_expires_at = 0, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", ( StepStatus.RECOVERY_WAIT.value, step.recoveries + 1, diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index 104d859dc4c..927184808bd 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -3,13 +3,18 @@ The harness runs registered workflows on an in-memory store with a virtual clock, so tests drive retries, durable delays, and deadlines by advancing time instead of sleeping. Retry jitter is disabled for determinism. + +Lease renewal runs on a real-time cadence while lease expiry is measured on the +virtual clock, so an attempt held across ``advance()`` is renewed by the next +pump rather than expiring; ``lease_duration`` is virtual seconds and +``lease_renew_interval`` is real seconds. """ from __future__ import annotations from typing import TYPE_CHECKING, Any -from reflex_base.workflow import parse_duration +from reflex_base.workflow import DEFAULT_LEASE_DURATION, parse_duration from reflex.workflow.runtime import WorkflowRuntime, _context_runtime from reflex.workflow.store import MemoryRunStore @@ -61,6 +66,8 @@ def __init__( *workflow_classes: type[BaseState], store: RunStore | None = None, start_time: float = DEFAULT_START_TIME, + lease_duration: DurationLike = DEFAULT_LEASE_DURATION, + lease_renew_interval: float | None = None, ): """Initialize the harness. @@ -68,12 +75,16 @@ def __init__( workflow_classes: Workflow classes to register. store: Run store override; defaults to a fresh in-memory store. start_time: Initial virtual time in epoch seconds. + lease_duration: Virtual seconds a claim survives without renewal. + lease_renew_interval: Real seconds between lease renewals. """ self._clock = _VirtualClock(start_time) self._runtime = WorkflowRuntime( store if store is not None else MemoryRunStore(), clock=self._clock, rng=lambda: 1.0, + lease_duration=parse_duration(lease_duration), + lease_renew_interval=lease_renew_interval, ) for workflow_cls in workflow_classes: self._runtime.register(workflow_cls) diff --git a/tests/units/workflow/test_lease.py b/tests/units/workflow/test_lease.py new file mode 100644 index 00000000000..21696f2a1e5 --- /dev/null +++ b/tests/units/workflow/test_lease.py @@ -0,0 +1,334 @@ +"""Tests for claim leases, which keep a second worker off a live claim.""" + +import asyncio + +import pytest +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.definition import compile_workflow +from reflex.workflow.kernel import WorkflowKernel +from reflex.workflow.records import HistoryEventType, RunStatus, StepStatus +from reflex.workflow.store import MemoryRunStore, SqliteRunStore +from reflex.workflow.testing import WorkflowTestHarness + + +class _Clock: + """A manually advanced epoch-seconds clock shared by cooperating kernels.""" + + def __init__(self, now: float = 1_000_000.0): + self.now = now + + def __call__(self) -> float: + return self.now + + +async def _drain(pump: asyncio.Task | None, release: asyncio.Event) -> None: + """Let a pumped kernel finish so stores close without in-flight work. + + Args: + pump: The task running the kernel's execution loop, if it started. + release: The event the hanging handler is waiting on. + """ + release.set() + if pump is None: + return + pump.cancel() + await asyncio.gather(pump, return_exceptions=True) + + +async def test_second_kernel_cannot_steal_a_live_claim( + forked_registration_context, tmp_path +): + """A peer starting up must not reclaim a claim another worker is executing. + + This is the regression test for the double-execution defect: recovery used + to treat every CLAIMED row as an orphan, so a second worker booting while + the first was mid-attempt ran the same step concurrently. + """ + started = asyncio.Event() + release = asyncio.Event() + executions = [] + + class LeaseFlow(rx.State): + __workflow__ = WorkflowConfig(id="lease.steal") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def charge(self): + executions.append(1) + started.set() + await release.wait() + self.status = "charged" + + definition = compile_workflow(LeaseFlow) + db_path = tmp_path / "workflow.db" + clock = _Clock() + store_a = SqliteRunStore(db_path) + store_b = SqliteRunStore(db_path) + pump: asyncio.Task | None = None + try: + kernel_a = WorkflowKernel([definition], store_a, clock=clock) + kernel_b = WorkflowKernel([definition], store_b, clock=clock) + + result = await kernel_a.start(LeaseFlow.charge()) + assert result.run_id is not None + pump = asyncio.create_task(kernel_a.run_until_idle()) + await asyncio.wait_for(started.wait(), timeout=5) + + # Kernel B boots while A is inside the handler, holding a live claim. + assert await kernel_b.recover() == 0 + assert await store_b.claim_next(clock()) is None + assert len(executions) == 1 + + release.set() + await asyncio.wait_for(pump, timeout=5) + + snapshot = await kernel_a.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "charged"} + assert snapshot.steps[0].recoveries == 0 + assert len(executions) == 1 + finally: + await _drain(pump, release) + store_a.close() + store_b.close() + + +async def test_expired_lease_is_reclaimed(forked_registration_context, tmp_path): + """A claim whose lease has expired is recovered by a peer.""" + started = asyncio.Event() + release = asyncio.Event() + + class ExpiredFlow(rx.State): + __workflow__ = WorkflowConfig(id="lease.expired") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="read") + async def work(self): + started.set() + await release.wait() + self.status = "done" + + definition = compile_workflow(ExpiredFlow) + db_path = tmp_path / "workflow.db" + clock = _Clock() + store_a = SqliteRunStore(db_path) + store_b = SqliteRunStore(db_path) + try: + kernel_a = WorkflowKernel( + [definition], store_a, clock=clock, lease_duration=30.0 + ) + kernel_b = WorkflowKernel( + [definition], store_b, clock=clock, lease_duration=30.0 + ) + result = await kernel_a.start(ExpiredFlow.work()) + assert result.run_id is not None + pump = asyncio.create_task(kernel_a.run_until_idle()) + await asyncio.wait_for(started.wait(), timeout=5) + + # Simulate A dying: stop renewing and push the clock past the lease. + pump.cancel() + await asyncio.gather(pump, return_exceptions=True) + clock.now += 31.0 + + assert await kernel_b.recover() == 1 + steps = await store_b.get_steps(result.run_id) + assert steps[0].status is StepStatus.RECOVERY_WAIT + assert steps[0].recoveries == 1 + finally: + release.set() + store_a.close() + store_b.close() + + +async def test_in_flight_claim_survives_clock_jumps(forked_registration_context): + """Advancing virtual time past a lease must not steal this kernel's own claim.""" + started = asyncio.Event() + release = asyncio.Event() + executions = [] + + class JumpFlow(rx.State): + __workflow__ = WorkflowConfig(id="lease.jump") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def work(self): + executions.append(1) + started.set() + await release.wait() + self.status = "done" + + async with WorkflowTestHarness(JumpFlow, lease_duration="60s") as harness: + result = await harness.kernel.start(JumpFlow.work()) + assert result.run_id is not None + pump = asyncio.create_task(harness.run_until_idle()) + try: + await asyncio.wait_for(started.wait(), timeout=5) + for _ in range(4): + await harness.advance("30s") + steps = await harness.kernel.store.get_steps(result.run_id) + assert steps[0].status is StepStatus.CLAIMED + assert steps[0].lease_expires_at == pytest.approx(harness.now + 60.0) + assert steps[0].recoveries == 0 + assert executions == [1] + finally: + release.set() + await asyncio.wait_for(pump, timeout=5) + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.steps[0].recoveries == 0 + assert executions == [1] + + +async def test_lease_loss_abandons_and_then_recovers( + forked_registration_context, tmp_path +): + """A fenced attempt commits nothing, and the recovered step runs again.""" + started = asyncio.Event() + release = asyncio.Event() + executions = [] + + class AbandonFlow(rx.State): + __workflow__ = WorkflowConfig(id="lease.abandon") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def work(self): + executions.append(1) + if len(executions) == 1: + started.set() + await release.wait() + self.status = "committed" + + definition = compile_workflow(AbandonFlow) + clock = _Clock() + store = SqliteRunStore(tmp_path / "workflow.db") + try: + kernel = WorkflowKernel( + [definition], + store, + clock=clock, + lease_duration=30.0, + lease_renew_interval=0.01, + ) + result = await kernel.start(AbandonFlow.work()) + assert result.run_id is not None + pump = asyncio.create_task(kernel.run_until_idle()) + await asyncio.wait_for(started.wait(), timeout=5) + + # A peer reclaims the step once the lease lapses; the renewer notices. + clock.now += 31.0 + assert await store.recover_orphans(clock(), max_recoveries=10) == 1 + # The first attempt stays blocked, so only the renewer can end it: it + # sees the fence, cancels the attempt, and the loop re-runs the step. + await asyncio.wait_for(pump, timeout=5) + + history = await store.get_history(result.run_id) + abandoned = [ + event + for event in history + if event.type is HistoryEventType.ATTEMPT_ABANDONED + ] + assert len(abandoned) == 1 + assert abandoned[0].data["reason"] == "lease_lost" + assert abandoned[0].data["effect"] == "idempotent_write" + + # The abandoned attempt committed nothing; the recovery did. + run = await store.get_run(result.run_id) + assert run is not None + assert run.status is RunStatus.COMPLETED + assert run.state_version == 1 + steps = await store.get_steps(result.run_id) + assert steps[0].recoveries == 1 + assert steps[0].attempts == 0 + assert executions == [1, 1] + finally: + release.set() + store.close() + + +async def test_shutdown_leaves_the_claim_recoverable( + forked_registration_context, tmp_path +): + """Stopping a worker mid-attempt leaves the step claimed, not cancelled.""" + started = asyncio.Event() + release = asyncio.Event() + + class ShutdownFlow(rx.State): + __workflow__ = WorkflowConfig(id="lease.shutdown") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="read") + async def work(self): + started.set() + await release.wait() + self.status = "done" + + definition = compile_workflow(ShutdownFlow) + clock = _Clock() + store = SqliteRunStore(tmp_path / "workflow.db") + try: + kernel = WorkflowKernel( + [definition], store, clock=clock, lease_duration=30.0, poll_interval=0.01 + ) + result = await kernel.start(ShutdownFlow.work()) + assert result.run_id is not None + await kernel.start_worker() + await asyncio.wait_for(started.wait(), timeout=5) + await kernel.aclose() + + steps = await store.get_steps(result.run_id) + assert steps[0].status is StepStatus.CLAIMED + run = await store.get_run(result.run_id) + assert run is not None + assert run.status is RunStatus.RUNNING + # It becomes reclaimable only after the lease lapses. + assert await store.recover_orphans(clock(), max_recoveries=10) == 0 + clock.now += 31.0 + assert await store.recover_orphans(clock(), max_recoveries=10) == 1 + finally: + release.set() + store.close() + + +def test_invalid_lease_timings_are_rejected(forked_registration_context): + """Nonsensical lease timings fail loudly at construction.""" + + class TimingFlow(rx.State): + __workflow__ = WorkflowConfig(id="lease.timing") + + @rx.event(durable=True, trigger=manual(), effect="none") + def work(self): + pass + + definition = compile_workflow(TimingFlow) + store = MemoryRunStore() + with pytest.raises(WorkflowRuntimeError, match="Lease timings"): + WorkflowKernel([definition], store, lease_duration=0) + with pytest.raises(WorkflowRuntimeError, match="Lease timings"): + WorkflowKernel( + [definition], store, lease_duration=10.0, lease_renew_interval=10.0 + ) + + +def test_store_without_renew_lease_is_rejected(forked_registration_context): + """A run store that cannot renew leases is refused at construction.""" + + class ProtocolFlow(rx.State): + __workflow__ = WorkflowConfig(id="lease.protocol") + + @rx.event(durable=True, trigger=manual(), effect="none") + def work(self): + pass + + class LegacyStore: + """A store predating leases: it has no renew_lease at all.""" + + definition = compile_workflow(ProtocolFlow) + with pytest.raises(WorkflowRuntimeError, match="renew_lease"): + WorkflowKernel([definition], LegacyStore()) # pyright: ignore[reportArgumentType] diff --git a/tests/units/workflow/test_store.py b/tests/units/workflow/test_store.py index 4e1bd96ac8d..03f8575857a 100644 --- a/tests/units/workflow/test_store.py +++ b/tests/units/workflow/test_store.py @@ -255,7 +255,7 @@ async def test_deadline_makes_run_control_pending(store): async def test_recover_orphans_consumes_recovery_budget(store): await store.admit(_run(), _step(), _ADMIT_EVENTS) - claim = await store.claim_next(NOW) + claim = await store.claim_next(NOW, lease_duration=5.0) assert claim is not None recovered = await store.recover_orphans(NOW + 10, max_recoveries=2) assert recovered == 1 @@ -277,7 +277,7 @@ async def test_recover_orphans_consumes_recovery_budget(store): async def test_recover_orphans_exhaustion_fails_run(store): await store.admit(_run(), _step(recoveries=2), _ADMIT_EVENTS) - claim = await store.claim_next(NOW) + claim = await store.claim_next(NOW, lease_duration=5.0) assert claim is not None await store.recover_orphans(NOW + 10, max_recoveries=2) run = await store.get_run("run1") @@ -300,7 +300,7 @@ async def test_sqlite_persistence_across_reopen(tmp_path): db_path = tmp_path / "workflow.db" first = SqliteRunStore(db_path) await first.admit(_run(request_key="key1"), _step(), _ADMIT_EVENTS) - claim = await first.claim_next(NOW) + claim = await first.claim_next(NOW, lease_duration=1.0) assert claim is not None first.close() @@ -324,3 +324,142 @@ async def test_sqlite_persistence_across_reopen(tmp_path): assert steps[0].status is StepStatus.RECOVERY_WAIT finally: second.close() + + +async def test_claim_sets_a_lease(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + assert claim.step.lease_expires_at == pytest.approx(NOW + 30.0) + steps = await store.get_steps("run1") + assert steps[0].lease_expires_at == pytest.approx(NOW + 30.0) + + +async def test_renew_lease_extends_the_expiry(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + assert await store.renew_lease(claim, NOW + 10.0, lease_duration=30.0) + steps = await store.get_steps("run1") + assert steps[0].lease_expires_at == pytest.approx(NOW + 40.0) + # Renewal is a liveness signal only: it transitions nothing. + assert steps[0].status is StepStatus.CLAIMED + assert steps[0].epoch == claim.step.epoch + assert steps[0].attempts == 0 + assert steps[0].recoveries == 0 + # The renewed claim still commits. + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state={"n": 1}, + ), + NOW + 11.0, + ) + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.COMPLETED + + +async def test_recover_orphans_skips_unexpired_leases(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + assert await store.recover_orphans(NOW + 29.0, max_recoveries=10) == 0 + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.CLAIMED + assert steps[0].recoveries == 0 + + +async def test_recover_orphans_reclaims_at_the_expiry_boundary(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + assert await store.recover_orphans(NOW + 30.0, max_recoveries=10) == 1 + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.RECOVERY_WAIT + assert steps[0].recoveries == 1 + # Lease loss is infrastructure, never a business attempt. + assert steps[0].attempts == 0 + assert steps[0].lease_expires_at == pytest.approx(0.0) + with pytest.raises(StaleClaimError): + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state={"n": 9}, + ), + NOW + 31.0, + ) + + +async def test_renewed_lease_survives_a_later_sweep(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + assert await store.renew_lease(claim, NOW + 20.0, lease_duration=30.0) + assert await store.recover_orphans(NOW + 40.0, max_recoveries=10) == 0 + assert await store.recover_orphans(NOW + 50.0, max_recoveries=10) == 1 + + +async def test_renew_lease_refused_after_recovery(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=5.0) + assert claim is not None + assert await store.recover_orphans(NOW + 10.0, max_recoveries=10) == 1 + assert not await store.renew_lease(claim, NOW + 11.0, lease_duration=30.0) + + +async def test_renew_lease_refused_after_commit(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state={"n": 1}, + ), + NOW + 1.0, + ) + assert not await store.renew_lease(claim, NOW + 2.0, lease_duration=30.0) + steps = await store.get_steps("run1") + assert steps[0].lease_expires_at == pytest.approx(0.0) + + +async def test_release_claim_clears_the_lease(store): + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + await store.release_claim(claim, status=StepStatus.READY, events=(), now=NOW + 1.0) + steps = await store.get_steps("run1") + assert steps[0].lease_expires_at == pytest.approx(0.0) + + +async def test_sqlite_migrates_a_database_without_the_lease_column(tmp_path): + db_path = tmp_path / "legacy.db" + store = SqliteRunStore(db_path) + await store.admit(_run(), _step(), _ADMIT_EVENTS) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + # Simulate a database written by a build that predates leases. + store._db.execute("DROP INDEX IF EXISTS idx_workflow_steps_lease") + store._db.execute("ALTER TABLE workflow_steps DROP COLUMN lease_expires_at") + store.close() + + reopened = SqliteRunStore(db_path) + try: + steps = await reopened.get_steps("run1") + assert steps[0].status is StepStatus.CLAIMED + # A claim left by the previous binary has no lease, so it is a genuine + # orphan and is reclaimed on the first sweep. + assert steps[0].lease_expires_at == pytest.approx(0.0) + assert await reopened.recover_orphans(NOW, max_recoveries=10) == 1 + finally: + reopened.close() + # Reopening an already-migrated database is a no-op. + again = SqliteRunStore(db_path) + again.close() From cfccb5ad4c1855777810af5f8d5ed4407151cc48 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 13:30:43 -0700 Subject: [PATCH 003/121] Make workflow retries mean what they say rx.Retry(max_attempts=5) on a handler raising an ordinary exception executed the handler exactly once. Both the effect-class defaults and any explicit policy that did not name retry_on were resolved to retry only on TransientWorkflowError, so a flaky HTTP call or a dropped connection failed the run immediately while the declared policy promised five attempts. Failures now retry by default: none, read, and idempotent_write resolve to three attempts with exponential backoff on any Exception, and an explicit policy without retry_on retries on Exception too. Narrow it with do_not_retry_on to fail fast. non_idempotent_write still gets exactly one attempt and routes to NEEDS_ATTENTION, since the runtime cannot prove the external effect did not already land. This is a deliberate divergence from the design doc's 'unknown code defects do not retry by default'. That rule is incoherent with an idempotent_write declaring three attempts that can never fire, and it makes the common case -- surviving a flaky dependency, which is the whole point of a durable step -- require boilerplate. Every comparable engine retries by default. TransientWorkflowError remains as an explicit marker for intent and for staying retryable under a narrowed policy. --- news/workflow-retry-semantics.bugfix.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 18 ++++---- reflex/workflow/definition.py | 17 ++++---- tests/units/reflex_base/test_workflow.py | 5 ++- tests/units/workflow/test_definition.py | 21 +++++++++- tests/units/workflow/test_kernel.py | 42 ++++++++++++++++++- 6 files changed, 82 insertions(+), 22 deletions(-) create mode 100644 news/workflow-retry-semantics.bugfix.md diff --git a/news/workflow-retry-semantics.bugfix.md b/news/workflow-retry-semantics.bugfix.md new file mode 100644 index 00000000000..20fd2e68326 --- /dev/null +++ b/news/workflow-retry-semantics.bugfix.md @@ -0,0 +1 @@ +Durable steps now retry ordinary failures by default, so `rx.Retry(max_attempts=5)` means five attempts instead of one. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 95170e5767b..794883853af 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -81,12 +81,10 @@ def parse_duration(value: DurationLike, *, param: str = "duration") -> float: class TransientWorkflowError(Exception): - """Raise from a durable handler to mark a failure as safely retryable. + """Raise from a durable handler to mark a failure as explicitly retryable. - Retry policies for the ``none``, ``read``, and ``idempotent_write`` effect - classes treat this exception (and its subclasses) as retryable by default. - Any other exception is a non-retryable defect unless it is listed in - ``Retry.retry_on``. + Failures already retry by default, so this exists to state the intent in + code and to stay retryable under a policy that narrows ``retry_on``. """ @@ -179,9 +177,11 @@ def is_retryable(self, error: BaseException) -> bool: def default_retry_for_effect(effect: str) -> Retry: """Return the default retry policy for an effect class. - Unknown code defects never retry by default; only ``TransientWorkflowError`` - marks a failure as safely retryable. Non-idempotent writes get exactly one - business attempt because the runtime cannot prove a retry is safe. + Failures retry three times with exponential backoff, which is what makes a + durable step survive a flaky dependency. A ``non_idempotent_write`` gets + exactly one business attempt instead: the runtime cannot prove the external + effect did not already land, so it suspends the run for an operator rather + than guessing. Args: effect: The declared effect class of the handler. @@ -191,7 +191,7 @@ def default_retry_for_effect(effect: str) -> Retry: """ if effect == "non_idempotent_write": return Retry(max_attempts=1, retry_on=()) - return Retry(max_attempts=3, retry_on=(TransientWorkflowError,)) + return Retry(max_attempts=3, retry_on=(Exception,)) @dataclasses.dataclass(frozen=True) diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 881253484f3..528c923c2f8 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -18,7 +18,6 @@ from reflex_base.workflow import ( DurableEventConfig, Retry, - TransientWorkflowError, Trigger, WorkflowConfig, default_retry_for_effect, @@ -229,9 +228,11 @@ def _compile_fields(workflow_cls: type[BaseState]) -> tuple[FieldSchema, ...]: def _resolve_retry(retry: Retry | None, effect: str) -> Retry: """Materialize the effective retry policy for a handler. - An explicit policy that does not name retryable exception types inherits - the effect class's default retryable set, so ``Retry(max_attempts=5)`` - keeps its meaning without restating the transient-error contract. + A policy that does not name retryable exception types retries on any + ``Exception``, so ``Retry(max_attempts=5)`` means five attempts. Narrow it + with ``do_not_retry_on`` to fail fast on specific errors. A + ``non_idempotent_write`` never retries: the runtime cannot prove the + external effect did not already land. Args: retry: The explicit policy, if the handler declared one. @@ -241,10 +242,10 @@ def _resolve_retry(retry: Retry | None, effect: str) -> Retry: The fully resolved policy. """ if retry is None: - return default_retry_for_effect(effect) - if not retry.retry_on and effect != "non_idempotent_write": - return dataclasses.replace(retry, retry_on=(TransientWorkflowError,)) - return retry + retry = default_retry_for_effect(effect) + if retry.retry_on or effect == "non_idempotent_write": + return retry + return dataclasses.replace(retry, retry_on=(Exception,)) def _compile_handlers( diff --git a/tests/units/reflex_base/test_workflow.py b/tests/units/reflex_base/test_workflow.py index 4c27be76b6d..0aad021d2b5 100644 --- a/tests/units/reflex_base/test_workflow.py +++ b/tests/units/reflex_base/test_workflow.py @@ -88,10 +88,13 @@ def test_default_retry_for_effect(): for effect in ("none", "read", "idempotent_write"): retry = default_retry_for_effect(effect) assert retry.max_attempts == 3 - assert retry.retry_on == (TransientWorkflowError,) + # Ordinary failures retry; that is the point of a durable step. + assert retry.is_retryable(ConnectionError("flaky")) + assert retry.is_retryable(TransientWorkflowError("explicit")) non_idempotent = default_retry_for_effect("non_idempotent_write") assert non_idempotent.max_attempts == 1 assert non_idempotent.retry_on == () + assert not non_idempotent.is_retryable(ConnectionError("flaky")) def test_workflow_config_valid(): diff --git a/tests/units/workflow/test_definition.py b/tests/units/workflow/test_definition.py index 8efb437fc39..58883c9663b 100644 --- a/tests/units/workflow/test_definition.py +++ b/tests/units/workflow/test_definition.py @@ -73,11 +73,28 @@ def payment_received(self): assert compile_workflow(OtherPolicy).digest != first.digest -def test_explicit_retry_inherits_transient_default(forked_registration_context): +def test_explicit_retry_retries_ordinary_failures(forked_registration_context): + """Retry(max_attempts=5) must mean five attempts, not one.""" definition = compile_workflow(_billing_workflow()) retry = definition.handlers["fulfill"].retry assert retry.max_attempts == 5 - assert retry.retry_on == (TransientWorkflowError,) + assert retry.is_retryable(ConnectionError("provider down")) + assert retry.is_retryable(TransientWorkflowError("explicit")) + + +def test_non_idempotent_write_never_retries(forked_registration_context): + """An uncertain write is never retried, whatever the exception.""" + + class UnsafeWrite(rx.State): + __workflow__ = WorkflowConfig(id="billing.unsafe_write") + + @rx.event(durable=True, trigger=manual(), effect="non_idempotent_write") + def wire(self): + pass + + retry = compile_workflow(UnsafeWrite).handlers["wire"].retry + assert retry.max_attempts == 1 + assert not retry.is_retryable(ConnectionError("dropped")) def test_explicit_retry_on_preserved(forked_registration_context): diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py index 2b4450fa773..c6be090daf6 100644 --- a/tests/units/workflow/test_kernel.py +++ b/tests/units/workflow/test_kernel.py @@ -193,13 +193,19 @@ def flaky(self): assert snapshot.error["type"] == "TransientWorkflowError" -async def test_unknown_defect_does_not_retry(forked_registration_context): +async def test_do_not_retry_on_fails_fast(forked_registration_context): + """A failure named in do_not_retry_on fails the run on the first attempt.""" calls = [] class DefectFlow(rx.State): __workflow__ = WorkflowConfig(id="kernel.defect") - @rx.event(durable=True, trigger=manual(), effect="none") + @rx.event( + durable=True, + trigger=manual(), + effect="none", + retry=Retry(max_attempts=3, do_not_retry_on=(ValueError,)), + ) def broken(self): calls.append(1) msg = "bug" @@ -214,6 +220,34 @@ def broken(self): assert len(calls) == 1 +async def test_ordinary_failures_retry_by_default(forked_registration_context): + """A flaky dependency is survived without declaring a retry policy.""" + calls = [] + + class FlakyFlow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.flaky_default") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="read") + def fetch(self): + calls.append(1) + if len(calls) < 3: + msg = "connection reset" + raise ConnectionError(msg) + self.status = "fetched" + + async with WorkflowTestHarness(FlakyFlow) as harness: + result = await harness.start(FlakyFlow.fetch()) + assert result.run_id is not None + await harness.advance("1s") + await harness.advance("2s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "fetched"} + assert len(calls) == 3 + + async def test_timeout_consumes_attempts_and_runs_timeout_hook( forked_registration_context, ): @@ -500,6 +534,10 @@ def begin(self): async with WorkflowTestHarness(BadStateFlow) as harness: result = await harness.start(BadStateFlow.begin()) assert result.run_id is not None + # Unserializable state is not transient, but nothing can tell the + # difference at runtime, so it exhausts the default retries first. + await harness.advance("1s") + await harness.advance("2s") snapshot = await harness.get_run(result.run_id) assert snapshot is not None assert snapshot.status is RunStatus.FAILED From 3ed7ecc5aa8738f22b9abac93c329a26269b8b0d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 13:36:36 -0700 Subject: [PATCH 004/121] Stop redeploys from stranding in-flight workflow runs Runs were pinned to a hash of the whole compiled definition, so adding a state field, retuning a retry policy, or changing a timeout parked every live run of that workflow in NEEDS_ATTENTION -- with no API to get it back. That fires on the second deploy of any real app, not on an exotic one. Runs are now gated on what can actually strand a step: the handler it names is gone, or its persisted payload no longer fits that handler's parameters. Both suspend with a precise, actionable reason instead of a bare digest mismatch. Everything else -- new fields, retuned retries and timeouts, changed hooks, effects, and triggers -- deploys without disturbing work in flight. The definition digest is still recorded on each run as provenance. Adds rx.workflows.resume(run_id) so suspension is a door rather than a wall: it clears the error, grants the frontier step a fresh attempt budget, and makes it claimable. A handler that returns rx.needs_attention() now leaves its step NEEDS_ATTENTION rather than SUCCEEDED, so resuming re-runs the handler that suspended and lets it take a different branch once a human has acted. --- news/workflow-versioning.bugfix.md | 1 + reflex/workflow/__init__.py | 8 +- reflex/workflow/kernel.py | 81 +++++++++- reflex/workflow/records.py | 1 + reflex/workflow/runtime.py | 15 ++ reflex/workflow/store.py | 94 ++++++++++++ reflex/workflow/testing.py | 13 ++ tests/units/workflow/test_kernel.py | 12 +- tests/units/workflow/test_versioning.py | 191 ++++++++++++++++++++++++ 9 files changed, 404 insertions(+), 12 deletions(-) create mode 100644 news/workflow-versioning.bugfix.md create mode 100644 tests/units/workflow/test_versioning.py diff --git a/news/workflow-versioning.bugfix.md b/news/workflow-versioning.bugfix.md new file mode 100644 index 00000000000..799580ee47a --- /dev/null +++ b/news/workflow-versioning.bugfix.md @@ -0,0 +1 @@ +Deploying new workflow code no longer strands runs that are already in flight, and `rx.workflows.resume()` re-opens a run suspended for operator attention. diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index de4ec992b5e..c076a7826a3 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -44,7 +44,12 @@ StepStatus, ) from reflex.workflow.runtime import WorkflowRuntime, get_runtime, workflows -from reflex.workflow.store import MemoryRunStore, RunStore, SqliteRunStore +from reflex.workflow.store import ( + MemoryRunStore, + RunStore, + SqliteRunStore, + StaleClaimError, +) from reflex.workflow.testing import WorkflowTestHarness __all__ = [ @@ -62,6 +67,7 @@ "RunStore", "ScheduleTrigger", "SqliteRunStore", + "StaleClaimError", "StartResult", "StepRecord", "StepStatus", diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 59aab3b0d10..272c359356f 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -394,6 +394,20 @@ async def cancel(self, run_id: str) -> bool: self._wakeup.set() return recorded + async def resume(self, run_id: str) -> bool: + """Re-open a run that is suspended for operator attention. + + Args: + run_id: The run to resume. + + Returns: + True if a suspended run was re-opened. + """ + resumed = await self._store.resume_run(run_id, self._clock()) + if resumed: + self._wakeup.set() + return resumed + async def get_run(self, run_id: str) -> RunSnapshot | None: """Load a read-only snapshot of a run. @@ -852,8 +866,10 @@ def _success_completion( if isinstance(control, NeedsAttention): error = {"reason": control.reason, "details": control.details} events.append((HistoryEventType.RUN_NEEDS_ATTENTION, {"error": error})) + # The attempt succeeded and its state is committed, but the step + # holds the suspension so resuming knows where to pick back up. return StepCompletion( - step_status=StepStatus.SUCCEEDED, + step_status=StepStatus.NEEDS_ATTENTION, run_status=RunStatus.NEEDS_ATTENTION, state=state, run_error=error, @@ -1057,6 +1073,56 @@ async def _record_abandoned( self._clock(), ) + @staticmethod + def _incompatible_reason( + defn: WorkflowDefinition | None, claim: Claim + ) -> dict[str, Any] | None: + """Check that a pending step can still be dispatched after a redeploy. + + Only two changes can strand a step: the handler it names is gone, or + its persisted payload no longer fits that handler's parameters. Adding + state fields or retuning retries, timeouts, and hooks is safe, so those + deploy without disturbing runs already in flight. + + Args: + defn: The current definition of the run's workflow, if registered. + claim: The claim about to be executed. + + Returns: + A JSON-compatible reason when the step cannot be dispatched, else None. + """ + if defn is None: + return { + "reason": "unknown_workflow", + "workflow_id": claim.run.workflow_id, + "detail": ( + f"Workflow {claim.run.workflow_id!r} is no longer registered " + "with this app; re-register it to resume the run." + ), + } + handler = defn.handlers.get(claim.step.handler_id) + if handler is None: + return { + "reason": "unknown_handler", + "handler_id": claim.step.handler_id, + "detail": ( + f"Handler {claim.step.handler_id!r} no longer exists on " + f"{claim.run.workflow_id!r}; restore it or cancel the run." + ), + } + unexpected = sorted(set(claim.step.args) - set(handler.params)) + if unexpected: + return { + "reason": "incompatible_payload", + "handler_id": handler.id, + "detail": ( + f"Step payload has arguments {unexpected} that handler " + f"{handler.id!r} no longer accepts; restore the parameters " + "or cancel the run." + ), + } + return None + async def _execute_claim(self, claim: Claim) -> None: """Execute one claimed attempt and commit its outcome. @@ -1065,20 +1131,19 @@ async def _execute_claim(self, claim: Claim) -> None: """ defn = self._definitions.get(claim.run.workflow_id) now = self._clock() - if defn is None or defn.digest != claim.run.definition_digest: - reason = ( - "unknown_workflow" if defn is None else "definition_digest_mismatch" - ) + incompatible = self._incompatible_reason(defn, claim) + if defn is None or incompatible is not None: + incompatible = incompatible or {"reason": "unknown_workflow"} await self._store.commit( claim, StepCompletion( step_status=StepStatus.NEEDS_ATTENTION, run_status=RunStatus.NEEDS_ATTENTION, state=None, - step_error={"reason": reason}, - run_error={"reason": reason}, + step_error=incompatible, + run_error=incompatible, events=( - (HistoryEventType.RUN_NEEDS_ATTENTION, {"reason": reason}), + (HistoryEventType.RUN_NEEDS_ATTENTION, dict(incompatible)), ), ), now, diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index b8b1929788a..ec8e07ed01b 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -85,6 +85,7 @@ class HistoryEventType(str, enum.Enum): RUN_CANCEL_REQUESTED = "run_cancel_requested" RUN_CANCELLED = "run_cancelled" RUN_NEEDS_ATTENTION = "run_needs_attention" + RUN_RESUMED = "run_resumed" @dataclasses.dataclass(frozen=True, slots=True) diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 3adfb1c8558..99129687d93 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -269,6 +269,21 @@ async def cancel(run_id: str) -> bool: """ return await get_runtime().kernel.cancel(run_id) + @staticmethod + async def resume(run_id: str) -> bool: + """Re-open a run suspended for operator attention. + + Use this after fixing whatever made a step's outcome uncertain: the + frontier step gets a fresh attempt budget and runs again. + + Args: + run_id: The run to resume. + + Returns: + True if a suspended run was re-opened. + """ + return await get_runtime().kernel.resume(run_id) + @staticmethod async def get_run(run_id: str) -> RunSnapshot | None: """Load a read-only snapshot of a run. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index c196858db06..74b68e60599 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -255,6 +255,23 @@ async def finalize_run( """ ... + async def resume_run(self, run_id: str, now: float) -> bool: + """Re-open a suspended run so its frontier step runs again. + + Suspension is an operator state, not an outcome: the run waits for a + human to fix whatever made the outcome uncertain. Resuming clears the + error, grants the frontier step a fresh attempt budget, and makes it + claimable immediately. + + Args: + run_id: The run to resume. + now: Current time in epoch seconds. + + Returns: + True if a suspended run was re-opened. + """ + ... + async def recover_orphans(self, now: float, max_recoveries: int) -> int: """Recover claims whose lease has expired. @@ -715,6 +732,38 @@ async def finalize_run( self._append_events(run_id, events, now) return True + async def resume_run(self, run_id: str, now: float) -> bool: + """Re-open a suspended run so its frontier step runs again. + + Args: + run_id: The run to resume. + now: Current time in epoch seconds. + + Returns: + True if a suspended run was re-opened. + """ + async with self._lock: + run = self._runs.get(run_id) + if run is None or run.status is not RunStatus.NEEDS_ATTENTION: + return False + steps = self._steps[run_id] + for step in list(steps): + if step.status is StepStatus.NEEDS_ATTENTION: + steps[step.ordinal] = dataclasses.replace( + step, + status=StepStatus.READY, + attempts=0, + due_at=now, + lease_expires_at=0.0, + error=None, + updated_at=now, + ) + self._runs[run_id] = dataclasses.replace( + run, status=RunStatus.PENDING, error=None, updated_at=now + ) + self._append_events(run_id, ((HistoryEventType.RUN_RESUMED, {}),), now) + return True + async def recover_orphans(self, now: float, max_recoveries: int) -> int: """Recover steps left claimed by a previous process. @@ -1542,6 +1591,51 @@ async def finalize_run( raise return True + async def resume_run(self, run_id: str, now: float) -> bool: + """Re-open a suspended run so its frontier step runs again. + + Args: + run_id: The run to resume. + now: Current time in epoch seconds. + + Returns: + True if a suspended run was re-opened. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + cursor = self._db.execute( + "UPDATE workflow_runs SET status = ?, error = NULL," + " updated_at = ? WHERE run_id = ? AND status = ?", + ( + RunStatus.PENDING.value, + now, + run_id, + RunStatus.NEEDS_ATTENTION.value, + ), + ) + if cursor.rowcount == 0: + self._db.execute("ROLLBACK") + return False + self._db.execute( + "UPDATE workflow_steps SET status = ?, attempts = 0, due_at = ?," + " lease_expires_at = 0, error = NULL, updated_at = ?" + " WHERE run_id = ? AND status = ?", + ( + StepStatus.READY.value, + now, + now, + run_id, + StepStatus.NEEDS_ATTENTION.value, + ), + ) + self._append_events(run_id, ((HistoryEventType.RUN_RESUMED, {}),), now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True + async def recover_orphans(self, now: float, max_recoveries: int) -> int: """Recover steps left claimed by a previous process. diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index 927184808bd..a53a90dfcdf 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -183,6 +183,19 @@ async def get_run(self, run_id: str) -> RunSnapshot | None: """ return await self.kernel.get_run(run_id) + async def resume(self, run_id: str) -> bool: + """Re-open a suspended run and process the work it unblocks. + + Args: + run_id: The run to resume. + + Returns: + True if a suspended run was re-opened. + """ + resumed = await self.kernel.resume(run_id) + await self.kernel.run_until_idle() + return resumed + async def cancel(self, run_id: str) -> bool: """Request cancellation of a run and process the drain. diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py index c6be090daf6..3237b35ddf7 100644 --- a/tests/units/workflow/test_kernel.py +++ b/tests/units/workflow/test_kernel.py @@ -639,9 +639,16 @@ def sync(self): second_store.close() -async def test_definition_digest_mismatch_suspends_run( +async def test_policy_changes_do_not_strand_in_flight_runs( forked_registration_context, tmp_path ): + """Retuning a step's effect and timeout must not disturb a live run. + + Deploying new code is routine; only a step that can no longer be + dispatched suspends. See tests/units/workflow/test_versioning.py for the + incompatible cases. + """ + class PinnedV1(rx.State): __workflow__ = WorkflowConfig(id="kernel.pinned") @@ -679,6 +686,5 @@ def finish(self): await harness.run_until_idle() snapshot = await harness.get_run(result.run_id) assert snapshot is not None - assert snapshot.status is RunStatus.NEEDS_ATTENTION - assert snapshot.error == {"reason": "definition_digest_mismatch"} + assert snapshot.status is RunStatus.COMPLETED second_store.close() diff --git a/tests/units/workflow/test_versioning.py b/tests/units/workflow/test_versioning.py new file mode 100644 index 00000000000..d3ea27ea277 --- /dev/null +++ b/tests/units/workflow/test_versioning.py @@ -0,0 +1,191 @@ +"""Tests for deploying new workflow code while runs are in flight.""" + +from reflex_base.workflow import WorkflowConfig, manual, needs_attention + +import reflex as rx +from reflex.workflow.records import RunStatus, StepStatus +from reflex.workflow.store import MemoryRunStore +from reflex.workflow.testing import WorkflowTestHarness + + +def _flow(*, extra_field: bool = False, slow_retry: bool = False): + """Build one shape of a workflow, standing in for one deploy. + + Args: + extra_field: Whether this deploy declares an additional state field. + slow_retry: Whether this deploy retunes the second step's timeout. + + Returns: + The workflow class for this deploy. + """ + + class Deployed(rx.State): + __workflow__ = WorkflowConfig(id="versioning.deployed") + status: str = "pending" + if extra_field: + note: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + self.status = "started" + return rx.after("1h", Deployed.finish) + + @rx.event( + durable=True, + effect="read", + timeout="90s" if slow_retry else None, + ) + def finish(self): + self.status = "done" + + return Deployed + + +async def test_added_state_field_does_not_strand_runs(forked_registration_context): + """Adding a field is a routine deploy, not a reason to suspend live runs.""" + store = MemoryRunStore() + first = _flow() + async with WorkflowTestHarness(first, store=store) as harness: + result = await harness.start(first.begin()) + assert result.run_id is not None + run_id, resume_at = result.run_id, harness.now + + async with WorkflowTestHarness( + _flow(extra_field=True), store=store, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"status": "done", "note": ""} + + +async def test_retuned_policy_does_not_strand_runs(forked_registration_context): + """Changing a timeout or retry policy applies to future attempts only.""" + store = MemoryRunStore() + first = _flow() + async with WorkflowTestHarness(first, store=store) as harness: + result = await harness.start(first.begin()) + assert result.run_id is not None + run_id, resume_at = result.run_id, harness.now + + async with WorkflowTestHarness( + _flow(slow_retry=True), store=store, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + + +async def test_removed_handler_suspends_with_a_precise_reason( + forked_registration_context, +): + """A pending step whose handler is gone cannot run, and says so.""" + store = MemoryRunStore() + first = _flow() + async with WorkflowTestHarness(first, store=store) as harness: + result = await harness.start(first.begin()) + assert result.run_id is not None + run_id, resume_at = result.run_id, harness.now + + class Truncated(rx.State): + __workflow__ = WorkflowConfig(id="versioning.deployed") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + self.status = "started" + + async with WorkflowTestHarness( + Truncated, store=store, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + assert snapshot.error is not None + assert snapshot.error["reason"] == "unknown_handler" + assert "finish" in snapshot.error["detail"] + + +async def test_unregistered_workflow_suspends(forked_registration_context): + """A run whose workflow is no longer registered waits rather than failing.""" + store = MemoryRunStore() + first = _flow() + async with WorkflowTestHarness(first, store=store) as harness: + result = await harness.start(first.begin()) + assert result.run_id is not None + run_id, resume_at = result.run_id, harness.now + + class Unrelated(rx.State): + __workflow__ = WorkflowConfig(id="versioning.unrelated") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + async with WorkflowTestHarness( + Unrelated, store=store, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + assert snapshot.error is not None + assert snapshot.error["reason"] == "unknown_workflow" + + +async def test_resume_reopens_a_suspended_run(forked_registration_context): + """Resuming grants the frontier step a fresh attempt budget.""" + attempts = [] + + class ReviewFlow(rx.State): + __workflow__ = WorkflowConfig(id="versioning.review") + status: str = "pending" + resolved: bool = False + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + attempts.append(1) + if not self.resolved: + return needs_attention("manual_review") + self.status = "done" + return None + + async with WorkflowTestHarness(ReviewFlow) as harness: + result = await harness.start(ReviewFlow.begin()) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + + # An operator fixes the cause, then resumes. + run = await harness.kernel.store.get_run(result.run_id) + assert run is not None + run.state["resolved"] = True + assert await harness.resume(result.run_id) + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state["status"] == "done" + assert snapshot.steps[0].status is StepStatus.SUCCEEDED + assert len(attempts) == 2 + + +async def test_resume_only_applies_to_suspended_runs(forked_registration_context): + """A healthy or terminal run is not resumable.""" + + class PlainFlow(rx.State): + __workflow__ = WorkflowConfig(id="versioning.plain") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + pass + + async with WorkflowTestHarness(PlainFlow) as harness: + result = await harness.start(PlainFlow.go()) + assert result.run_id is not None + assert not await harness.resume(result.run_id) + assert not await harness.resume("no-such-run") From 34d8caf506dcfff0d4cc53f819816fa725e91405 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 13:40:23 -0700 Subject: [PATCH 005/121] Reject handler bodies that silently break durability Calling another durable handler directly (self.charge()) ran it inline inside the caller's attempt: no retry policy of its own, no effect tracking, no step in the mailbox, and a silent re-execution of its side effect whenever the caller retried. The run still reported success, which is what makes it dangerous -- and it is exactly the shape a code generator reaches for. The compiler now parses each handler body and rejects two traps with an actionable message: an inline call to a sibling durable handler (pointing at 'return MyFlow.charge' instead), and a return of a plain literal (listing the transitions a durable handler may return). Both fail at compile time, before a run exists. Handlers whose source is unavailable, as in a REPL, are skipped rather than guessed at. --- news/workflow-durability-guards.feature.md | 1 + reflex/workflow/definition.py | 78 ++++++++++++++++ tests/units/workflow/test_definition.py | 102 +++++++++++++++++++++ 3 files changed, 181 insertions(+) create mode 100644 news/workflow-durability-guards.feature.md diff --git a/news/workflow-durability-guards.feature.md b/news/workflow-durability-guards.feature.md new file mode 100644 index 00000000000..71d7ba96db6 --- /dev/null +++ b/news/workflow-durability-guards.feature.md @@ -0,0 +1 @@ +The workflow compiler now rejects handler bodies that silently break durability: calling another durable handler inline, or returning a plain value instead of a transition. diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 528c923c2f8..28050040c07 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -8,10 +8,12 @@ from __future__ import annotations +import ast import dataclasses import hashlib import inspect import json +import textwrap from typing import TYPE_CHECKING, Any, get_type_hints from reflex_base.utils.exceptions import WorkflowDefinitionError @@ -305,6 +307,79 @@ def _compile_handlers( return handlers +def _handler_body(fn: Callable) -> ast.FunctionDef | ast.AsyncFunctionDef | None: + """Parse a handler's source into its function definition node. + + Args: + fn: The undecorated handler function. + + Returns: + The parsed node, or None when the source is unavailable (a handler + built by exec or defined in a REPL). + """ + try: + source = textwrap.dedent(inspect.getsource(fn)) + except (OSError, TypeError): + return None + try: + module = ast.parse(source) + except SyntaxError: + return None + node = module.body[0] if module.body else None + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + return node + return None + + +def _validate_handler_body( + workflow_cls: type[BaseState], + defn: HandlerDefinition, + durable_names: frozenset[str], +) -> None: + """Reject handler bodies that silently break the durability boundary. + + Args: + workflow_cls: The workflow class being compiled. + defn: The handler definition to check. + durable_names: Method names of every durable handler on the class. + + Raises: + WorkflowDefinitionError: If the body calls another durable handler + directly, or returns a value that is not a durable transition. + """ + node = _handler_body(defn.fn) + if node is None: + return + self_name = next(iter(inspect.signature(defn.fn).parameters), None) + for child in ast.walk(node): + if ( + isinstance(child, ast.Call) + and isinstance(child.func, ast.Attribute) + and isinstance(child.func.value, ast.Name) + and child.func.value.id == self_name + and child.func.attr in durable_names + ): + raise _error( + workflow_cls, + f"handler {defn.name!r} calls {child.func.attr!r} directly, which " + "runs it inline and loses its retries, timeout, and effect " + f"tracking. Return it as a transition instead: " + f"return {workflow_cls.__name__}.{child.func.attr}", + ) + if ( + isinstance(child, ast.Return) + and isinstance(child.value, ast.Constant) + and child.value.value is not None + ): + raise _error( + workflow_cls, + f"handler {defn.name!r} returns {child.value.value!r}. A durable " + "handler returns the next transition, not a value: return the " + "next handler, rx.after(...), rx.complete(result=...), " + "rx.fail(...), rx.needs_attention(...), or None.", + ) + + def _resolve_hooks( workflow_cls: type[BaseState], handlers: dict[str, HandlerDefinition] ) -> dict[str, HandlerDefinition]: @@ -460,6 +535,9 @@ def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: config = _validate_class_shape(workflow_cls) fields = _compile_fields(workflow_cls) handlers = _resolve_hooks(workflow_cls, _compile_handlers(workflow_cls, config)) + durable_names = frozenset(defn.name for defn in handlers.values()) + for defn in handlers.values(): + _validate_handler_body(workflow_cls, defn, durable_names) roots = tuple( defn.id for defn in sorted(handlers.values(), key=lambda d: d.id) diff --git a/tests/units/workflow/test_definition.py b/tests/units/workflow/test_definition.py index 58883c9663b..21b86836ad4 100644 --- a/tests/units/workflow/test_definition.py +++ b/tests/units/workflow/test_definition.py @@ -1,5 +1,11 @@ """Tests for the workflow definition compiler.""" +import importlib.util +import sys +import tempfile +import uuid +from pathlib import Path + import pytest from reflex_base.utils.exceptions import WorkflowDefinitionError from reflex_base.workflow import ( @@ -14,6 +20,31 @@ from reflex.workflow.definition import compile_workflow +def _load_module(source: str) -> dict: + """Import workflow source from a real file so its body can be parsed. + + The compiler reads handler source to reject bodies that break the + durability boundary, which needs an importable file rather than an exec'd + string. + + Args: + source: The module source to write and import. + + Returns: + The imported module's namespace. + """ + name = f"wf_probe_{uuid.uuid4().hex}" + path = Path(tempfile.gettempdir()) / f"{name}.py" + path.write_text(source) + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return vars(module) + + def _billing_workflow(): class BillingDefinition(rx.State): __workflow__ = WorkflowConfig( @@ -298,3 +329,74 @@ class ChildOfWorkflow(SubstateHaver): with pytest.raises(WorkflowDefinitionError, match="substates"): compile_workflow(SubstateHaver) + + +def test_direct_handler_call_is_rejected(forked_registration_context): + """Calling a durable handler inline silently loses its durability.""" + source = """ +import reflex as rx +from reflex_base.workflow import WorkflowConfig, manual + + +class InlineCallFlow(rx.State): + __workflow__ = WorkflowConfig(id="billing.inline_call") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + self.charge() + + @rx.event(durable=True, effect="idempotent_write") + def charge(self): + pass +""" + namespace = _load_module(source) + with pytest.raises(WorkflowDefinitionError, match="calls 'charge' directly"): + compile_workflow(namespace["InlineCallFlow"]) + + +def test_literal_return_is_rejected(forked_registration_context): + """A durable handler returns a transition, not a value.""" + source = """ +import reflex as rx +from reflex_base.workflow import WorkflowConfig, manual + + +class LiteralReturnFlow(rx.State): + __workflow__ = WorkflowConfig(id="billing.literal_return") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + return "done" +""" + namespace = _load_module(source) + with pytest.raises(WorkflowDefinitionError, match="returns 'done'"): + compile_workflow(namespace["LiteralReturnFlow"]) + + +def test_valid_transitions_compile(forked_registration_context): + """The shapes the guards steer users toward all compile.""" + source = """ +import reflex as rx +from reflex_base.workflow import WorkflowConfig, after, complete, manual + + +class TransitionsFlow(rx.State): + __workflow__ = WorkflowConfig(id="billing.transitions") + n: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + self.n += 1 + return TransitionsFlow.charge + + @rx.event(durable=True, effect="idempotent_write") + def charge(self): + return after("1h", TransitionsFlow.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + return complete(result={"n": self.n}) +""" + namespace = _load_module(source) + definition = compile_workflow(namespace["TransitionsFlow"]) + assert set(definition.handlers) == {"begin", "charge", "finish"} From d55c9baf81ce9fd183959a16502eedb6eef0fe60 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 13:45:13 -0700 Subject: [PATCH 006/121] Add run listing and label filtering Runs could only be fetched one at a time by id, so labels were recorded but never readable and nothing could build an operator view, a CLI listing, or a customer-facing 'your jobs' page. Adds RunQuery and RunStore.list_runs on both stores, surfaced as rx.workflows.list_runs(workflow_id=..., statuses=..., labels=..., limit=...), newest first with a created_before cursor for pagination. SQLite filters labels through json_extract so it stays a single indexed scan rather than loading every run. --- news/workflow-run-queries.feature.md | 1 + reflex/workflow/__init__.py | 2 + reflex/workflow/kernel.py | 34 ++++++- reflex/workflow/records.py | 25 +++++- reflex/workflow/runtime.py | 33 ++++++- reflex/workflow/store.py | 79 ++++++++++++++++ tests/units/workflow/test_queries.py | 130 +++++++++++++++++++++++++++ 7 files changed, 300 insertions(+), 4 deletions(-) create mode 100644 news/workflow-run-queries.feature.md create mode 100644 tests/units/workflow/test_queries.py diff --git a/news/workflow-run-queries.feature.md b/news/workflow-run-queries.feature.md new file mode 100644 index 00000000000..914b095ac9e --- /dev/null +++ b/news/workflow-run-queries.feature.md @@ -0,0 +1 @@ +Adds `rx.workflows.list_runs()` for filtering runs by workflow, status, and server-derived labels, with newest-first pagination. diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index c076a7826a3..ac3a8af0939 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -36,6 +36,7 @@ from reflex.workflow.records import ( HistoryEvent, HistoryEventType, + RunQuery, RunRecord, RunSnapshot, RunStatus, @@ -61,6 +62,7 @@ "ManualTrigger", "MemoryRunStore", "Retry", + "RunQuery", "RunRecord", "RunSnapshot", "RunStatus", diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 272c359356f..6622fc1c54a 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -37,6 +37,7 @@ from reflex.workflow.records import ( TERMINAL_STEP_STATUSES, HistoryEventType, + RunQuery, RunRecord, RunSnapshot, RunStatus, @@ -48,7 +49,7 @@ from reflex.workflow.store import Claim, RunStore, StaleClaimError, StepCompletion if TYPE_CHECKING: - from collections.abc import Callable, Iterable + from collections.abc import Callable, Iterable, Mapping from reflex.state import BaseState from reflex.workflow.definition import HandlerDefinition, WorkflowDefinition @@ -408,6 +409,37 @@ async def resume(self, run_id: str) -> bool: self._wakeup.set() return resumed + async def list_runs( + self, + *, + workflow_id: str | None = None, + statuses: Iterable[RunStatus] = (), + labels: Mapping[str, str] | None = None, + created_before: float | None = None, + limit: int = 50, + ) -> tuple[RunRecord, ...]: + """List runs matching a filter, newest first. + + Args: + workflow_id: Restrict to one workflow identity. + statuses: Restrict to these run statuses; empty means any. + labels: Require every one of these label values. + created_before: Pagination cursor; return runs admitted before this. + limit: Maximum runs to return. + + Returns: + The matching run records. + """ + return await self._store.list_runs( + RunQuery( + workflow_id=workflow_id, + statuses=tuple(statuses), + labels=labels, + created_before=created_before, + limit=limit, + ) + ) + async def get_run(self, run_id: str) -> RunSnapshot | None: """Load a read-only snapshot of a run. diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index ec8e07ed01b..c87122b99ac 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -4,7 +4,10 @@ import dataclasses import enum -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal + +if TYPE_CHECKING: + from collections.abc import Mapping class RunStatus(str, enum.Enum): @@ -214,6 +217,26 @@ class StartResult: retry_after: float | None = None +@dataclasses.dataclass(frozen=True, slots=True) +class RunQuery: + """Filters for listing runs in an operator surface. + + Attributes: + workflow_id: Restrict to one workflow identity. + statuses: Restrict to these run statuses; empty means any. + labels: Require every one of these server-derived label values. + created_before: Return runs admitted strictly before this epoch time, + which is the pagination cursor. + limit: Maximum runs to return, newest first. + """ + + workflow_id: str | None = None + statuses: tuple[RunStatus, ...] = () + labels: Mapping[str, str] | None = None + created_before: float | None = None + limit: int = 50 + + @dataclasses.dataclass(frozen=True, slots=True) class RunSnapshot: """Read-only projection of a run for operators and tests. diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 99129687d93..bc0d4d8da65 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -24,10 +24,10 @@ from reflex.workflow.store import RunStore, SqliteRunStore if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable + from collections.abc import AsyncIterator, Callable, Iterable, Mapping from reflex.state import BaseState - from reflex.workflow.records import RunSnapshot, StartResult + from reflex.workflow.records import RunRecord, RunSnapshot, RunStatus, StartResult DEFAULT_DB_FILENAME = "workflow.db" @@ -284,6 +284,35 @@ async def resume(run_id: str) -> bool: """ return await get_runtime().kernel.resume(run_id) + @staticmethod + async def list_runs( + *, + workflow_id: str | None = None, + statuses: Iterable[RunStatus] = (), + labels: Mapping[str, str] | None = None, + created_before: float | None = None, + limit: int = 50, + ) -> tuple[RunRecord, ...]: + """List runs matching a filter, newest first. + + Args: + workflow_id: Restrict to one workflow identity. + statuses: Restrict to these run statuses; empty means any. + labels: Require every one of these label values. + created_before: Pagination cursor; return runs admitted before this. + limit: Maximum runs to return. + + Returns: + The matching run records. + """ + return await get_runtime().kernel.list_runs( + workflow_id=workflow_id, + statuses=statuses, + labels=labels, + created_before=created_before, + limit=limit, + ) + @staticmethod async def get_run(run_id: str) -> RunSnapshot | None: """Load a read-only snapshot of a run. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 74b68e60599..1106969ec9a 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -31,6 +31,7 @@ TERMINAL_STEP_STATUSES, HistoryEvent, HistoryEventType, + RunQuery, RunRecord, RunStatus, StepRecord, @@ -290,6 +291,17 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: """ ... + async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: + """List runs matching a query, newest first. + + Args: + query: The filters and pagination cursor to apply. + + Returns: + The matching run records. + """ + ... + async def get_run(self, run_id: str) -> RunRecord | None: """Load one run record. @@ -354,6 +366,26 @@ def _run_is_runnable(run: RunRecord, now: float) -> bool: ) +def _matches_query(run: RunRecord, query: RunQuery) -> bool: + """Whether a run satisfies every filter in a query. + + Args: + run: The run record to test. + query: The filters to apply. + + Returns: + True when the run matches. + """ + if query.workflow_id is not None and run.workflow_id != query.workflow_id: + return False + if query.statuses and run.status not in query.statuses: + return False + if query.created_before is not None and run.created_at >= query.created_before: + return False + labels = run.labels or {} + return all(labels.get(key) == value for key, value in (query.labels or {}).items()) + + def _lease_expired(step: StepRecord, now: float) -> bool: """Whether a claimed step's lease has lapsed and it may be recovered. @@ -830,6 +862,20 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: ) return recovered + async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: + """List runs matching a query, newest first. + + Args: + query: The filters and pagination cursor to apply. + + Returns: + The matching run records. + """ + async with self._lock: + matched = [run for run in self._runs.values() if _matches_query(run, query)] + matched.sort(key=lambda run: run.created_at, reverse=True) + return tuple(matched[: query.limit]) + async def get_run(self, run_id: str) -> RunRecord | None: """Load one run record. @@ -1725,6 +1771,39 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: raise return recovered + async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: + """List runs matching a query, newest first. + + Args: + query: The filters and pagination cursor to apply. + + Returns: + The matching run records. + """ + clauses: list[str] = [] + params: list[Any] = [] + if query.workflow_id is not None: + clauses.append("workflow_id = ?") + params.append(query.workflow_id) + if query.statuses: + placeholders = ",".join("?" * len(query.statuses)) + clauses.append(f"status IN ({placeholders})") + params.extend(status.value for status in query.statuses) + if query.created_before is not None: + clauses.append("created_at < ?") + params.append(query.created_before) + for key, value in (query.labels or {}).items(): + clauses.append("json_extract(labels, ?) = ?") + params.extend((f"$.{key}", value)) + where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + with self._lock: + rows = self._db.execute( + f"SELECT * FROM workflow_runs{where}" + " ORDER BY created_at DESC, run_id DESC LIMIT ?", + (*params, query.limit), + ).fetchall() + return tuple(_run_from_row(row) for row in rows) + async def get_run(self, run_id: str) -> RunRecord | None: """Load one run record. diff --git a/tests/units/workflow/test_queries.py b/tests/units/workflow/test_queries.py new file mode 100644 index 00000000000..0ca9879569c --- /dev/null +++ b/tests/units/workflow/test_queries.py @@ -0,0 +1,130 @@ +"""Tests for listing and filtering runs, which operator surfaces are built on.""" + +import pytest +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import ( + RunQuery, + RunRecord, + RunStatus, + StepRecord, + StepStatus, +) +from reflex.workflow.store import MemoryRunStore, SqliteRunStore +from reflex.workflow.testing import WorkflowTestHarness + +NOW = 1_000_000.0 + + +@pytest.fixture(params=["memory", "sqlite"]) +def store(request, tmp_path): + """A run store of each implementation. + + Args: + request: The fixture request carrying the store kind. + tmp_path: Temporary directory for the SQLite database. + + Yields: + The store instance. + """ + if request.param == "memory": + yield MemoryRunStore() + else: + sqlite_store = SqliteRunStore(tmp_path / "workflow.db") + yield sqlite_store + sqlite_store.close() + + +async def _admit( + store, + run_id, + *, + workflow_id="ops.q", + status=RunStatus.PENDING, + labels=None, + created_at=NOW, +): + run = RunRecord( + run_id=run_id, + workflow_id=workflow_id, + definition_digest="digest", + status=status, + state={}, + state_version=0, + next_ordinal=1, + labels=labels, + created_at=created_at, + updated_at=created_at, + ) + step = StepRecord( + run_id=run_id, + ordinal=0, + handler_id="go", + status=StepStatus.READY, + args={}, + origin="root", + created_at=created_at, + updated_at=created_at, + ) + await store.admit(run, step, ()) + + +async def test_list_runs_orders_newest_first_and_paginates(store): + for index in range(5): + await _admit(store, f"run{index}", created_at=NOW + index) + page = await store.list_runs(RunQuery(limit=2)) + assert [run.run_id for run in page] == ["run4", "run3"] + nextpage = await store.list_runs( + RunQuery(limit=2, created_before=page[-1].created_at) + ) + assert [run.run_id for run in nextpage] == ["run2", "run1"] + + +async def test_list_runs_filters_by_workflow_and_status(store): + await _admit(store, "a", workflow_id="ops.q", status=RunStatus.COMPLETED) + await _admit(store, "b", workflow_id="ops.q", status=RunStatus.FAILED) + await _admit(store, "c", workflow_id="other.q", status=RunStatus.COMPLETED) + by_workflow = await store.list_runs(RunQuery(workflow_id="ops.q")) + assert {run.run_id for run in by_workflow} == {"a", "b"} + by_status = await store.list_runs(RunQuery(statuses=(RunStatus.COMPLETED,))) + assert {run.run_id for run in by_status} == {"a", "c"} + both = await store.list_runs( + RunQuery(workflow_id="ops.q", statuses=(RunStatus.FAILED,)) + ) + assert [run.run_id for run in both] == ["b"] + + +async def test_list_runs_filters_by_labels(store): + await _admit(store, "a", labels={"customer": "acme", "tier": "pro"}) + await _admit(store, "b", labels={"customer": "acme", "tier": "free"}) + await _admit(store, "c", labels={"customer": "globex"}) + await _admit(store, "d", labels=None) + acme = await store.list_runs(RunQuery(labels={"customer": "acme"})) + assert {run.run_id for run in acme} == {"a", "b"} + pro = await store.list_runs(RunQuery(labels={"customer": "acme", "tier": "pro"})) + assert [run.run_id for run in pro] == ["a"] + missing = await store.list_runs(RunQuery(labels={"customer": "nobody"})) + assert missing == () + + +async def test_list_runs_through_the_kernel(forked_registration_context): + class Listed(rx.State): + __workflow__ = WorkflowConfig(id="ops.listed") + n: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self, n: int): + self.n = n + + async with WorkflowTestHarness(Listed) as harness: + for index, customer in enumerate(("acme", "acme", "globex")): + await harness.kernel.start(Listed.go(index), labels={"customer": customer}) + await harness.run_until_idle() + + assert len(await harness.kernel.list_runs()) == 3 + acme = await harness.kernel.list_runs(labels={"customer": "acme"}) + assert len(acme) == 2 + completed = await harness.kernel.list_runs(statuses=[RunStatus.COMPLETED]) + assert len(completed) == 3 + assert await harness.kernel.list_runs(workflow_id="nope") == () From ce7aea1d2f0127d66b4f49f7e0c8d587f6e021ff Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:01:35 -0700 Subject: [PATCH 007/121] Serve webhook-triggered workflow roots over HTTP rx.webhook(...) compiled but could never fire: only manual roots were reachable, so the whole provider-driven half of the product was declarative decoration. Adds the ingress endpoint at POST /_workflow/webhook/{topic}, registered only when a workflow actually declares a webhook root. It preserves the raw request body, verifies the provider signature over those exact bytes, validates the payload against the declared model, and durably admits the run before acknowledging -- so a provider that never sees a 202 can safely redeliver. Redelivery reaches the same run through dedupe_by rather than starting a second one. Authentication is not optional by default: a webhook trigger without a verifier is a compile error naming the fix, and an endpoint that really is public must say so with allow_unverified plus a reason. rx.hmac_signature() covers the Stripe/GitHub/Shopify shape, reading the secret from the environment at request time so it never enters workflow state, history, or a browser bundle. Trigger kind is now part of admission: a webhook root cannot be started by application code, and a manual root is not reachable over HTTP. --- news/workflow-webhook-ingress.feature.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 97 ++++++- pyi_hashes.json | 2 +- reflex/__init__.py | 1 + reflex/app.py | 19 ++ reflex/compiler/compiler.py | 3 + reflex/workflow/__init__.py | 4 + reflex/workflow/ingress.py | 195 ++++++++++++++ reflex/workflow/kernel.py | 25 +- tests/units/reflex_base/test_workflow.py | 25 +- tests/units/workflow/test_definition.py | 7 +- tests/units/workflow/test_ingress.py | 248 ++++++++++++++++++ tests/units/workflow/test_kernel.py | 12 +- 13 files changed, 617 insertions(+), 22 deletions(-) create mode 100644 news/workflow-webhook-ingress.feature.md create mode 100644 reflex/workflow/ingress.py create mode 100644 tests/units/workflow/test_ingress.py diff --git a/news/workflow-webhook-ingress.feature.md b/news/workflow-webhook-ingress.feature.md new file mode 100644 index 00000000000..76f54309513 --- /dev/null +++ b/news/workflow-webhook-ingress.feature.md @@ -0,0 +1 @@ +Webhook triggers now actually fire: `rx.webhook(...)` roots are served over HTTP with provider signature verification, payload validation, and redelivery deduplication. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 794883853af..fb54a98af45 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -9,7 +9,10 @@ from __future__ import annotations import dataclasses +import hmac +import os import re +from collections.abc import Callable, Mapping from datetime import timedelta from typing import Any, ClassVar, Final, Literal, get_args @@ -208,6 +211,9 @@ class ManualTrigger(Trigger): kind: ClassVar[str] = "manual" +WebhookVerifier = Callable[[bytes, Mapping[str, str]], bool] + + @dataclasses.dataclass(frozen=True) class WebhookTrigger(Trigger): """Marks a root handler started by an authenticated provider webhook. @@ -215,26 +221,54 @@ class WebhookTrigger(Trigger): Attributes: topic: Stable provider event topic, e.g. ``"stripe.payment_succeeded"``. model: Optional typed payload model the raw payload is validated into. - verify: Provider signature verifier supplied by a connection binding. - dedupe_by: Payload field used as the ingress deduplication key. + verify: Callable given the raw body and headers that returns whether the + request genuinely came from the provider. + dedupe_by: Payload field used as the ingress deduplication key, so a + provider redelivering an event does not start a second run. + allow_unverified: Acknowledge that this endpoint accepts anonymous + traffic. Only valid with a non-empty ``unverified_reason``. + unverified_reason: Why anonymous traffic is acceptable here. """ kind: ClassVar[str] = "webhook" topic: str model: type | None = None - verify: Any = None + verify: WebhookVerifier | None = None dedupe_by: str | None = None + allow_unverified: bool = False + unverified_reason: str = "" def __post_init__(self): - """Validate the topic. + """Validate the topic and the authentication decision. Raises: - WorkflowDefinitionError: If the topic is empty. + WorkflowDefinitionError: If the topic is empty, or the endpoint + would accept unauthenticated traffic without saying so. """ if not self.topic: msg = "webhook trigger requires a non-empty topic." raise WorkflowDefinitionError(msg) + if self.verify is None and not self.allow_unverified: + msg = ( + f"webhook trigger {self.topic!r} has no verifier, so anyone who " + "knows the URL could start runs. Pass verify=rx.hmac_signature(" + 'secret_env="...", header="...") or, if the endpoint really is ' + "public, allow_unverified=True with an unverified_reason." + ) + raise WorkflowDefinitionError(msg) + if self.allow_unverified and not self.unverified_reason: + msg = ( + f"webhook trigger {self.topic!r} sets allow_unverified=True and " + "must give a non-empty unverified_reason." + ) + raise WorkflowDefinitionError(msg) + if self.verify is not None and self.allow_unverified: + msg = ( + f"webhook trigger {self.topic!r} declares both a verifier and " + "allow_unverified=True; keep the verifier." + ) + raise WorkflowDefinitionError(msg) @dataclasses.dataclass(frozen=True) @@ -276,21 +310,68 @@ def webhook( topic: str, *, model: type | None = None, - verify: Any = None, + verify: WebhookVerifier | None = None, dedupe_by: str | None = None, + allow_unverified: bool = False, + unverified_reason: str = "", ) -> WebhookTrigger: """Create a webhook trigger for a workflow root handler. Args: topic: Stable provider event topic. model: Optional typed payload model. - verify: Provider signature verifier from a connection binding. + verify: Callable given the raw body and headers that returns whether the + request genuinely came from the provider. dedupe_by: Payload field used as the ingress deduplication key. + allow_unverified: Acknowledge that this endpoint accepts anonymous traffic. + unverified_reason: Why anonymous traffic is acceptable here. Returns: The trigger specification. """ - return WebhookTrigger(topic=topic, model=model, verify=verify, dedupe_by=dedupe_by) + return WebhookTrigger( + topic=topic, + model=model, + verify=verify, + dedupe_by=dedupe_by, + allow_unverified=allow_unverified, + unverified_reason=unverified_reason, + ) + + +def hmac_signature( + *, + secret_env: str, + header: str, + algorithm: str = "sha256", + prefix: str = "", +) -> WebhookVerifier: + """Build a verifier for providers that HMAC-sign the raw request body. + + This covers the common shape used by Stripe, GitHub, Shopify and others: + the provider sends a hex digest of the body keyed by a shared secret. The + secret is read from the environment at request time, so it never enters + workflow state, history, or a browser bundle. + + Args: + secret_env: Name of the environment variable holding the shared secret. + header: Request header carrying the provider's signature. + algorithm: Hash algorithm name understood by ``hashlib``. + prefix: Fixed prefix the provider puts before the digest, e.g. ``"sha256="``. + + Returns: + A verifier callable for ``rx.webhook(verify=...)``. + """ + + def verify(body: bytes, headers: Mapping[str, str]) -> bool: + secret = os.environ.get(secret_env) + presented = headers.get(header.lower()) or headers.get(header) + if not secret or not presented: + return False + expected = hmac.new(secret.encode(), body, algorithm).hexdigest() + return hmac.compare_digest(f"{prefix}{expected}", presented) + + return verify def schedule(cron: str) -> ScheduleTrigger: diff --git a/pyi_hashes.json b/pyi_hashes.json index 03ddc3dc696..e428d4b7044 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -118,7 +118,7 @@ "packages/reflex-components-recharts/src/reflex_components_recharts/polar.pyi": "99ebcfc07868061bdc3c2010d85a153f", "packages/reflex-components-recharts/src/reflex_components_recharts/recharts.pyi": "4f6c26f8c76543cc41e2b9dc400ece8a", "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "f170ac685b6ba5892370166c80684db3", - "reflex/__init__.pyi": "577f2307b3ba7fcd1aeda6621056c34c", + "reflex/__init__.pyi": "41df8e648e62b67ddcd0ec6aa9d88f4c", "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388", "reflex/experimental/memo.pyi": "bc8b48357bef580e70a5881b65d3d3f7" } diff --git a/reflex/__init__.py b/reflex/__init__.py index 19d3dda7879..fed4f822e9a 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -242,6 +242,7 @@ "TransientWorkflowError", "manual", "webhook", + "hmac_signature", "schedule", "after", "complete", diff --git a/reflex/app.py b/reflex/app.py index 7c210ee8517..2b4c8b1d98e 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -838,6 +838,25 @@ def _add_default_endpoints(self): methods=["GET"], ) + def _add_workflow_endpoints(self): + """Add the webhook ingress endpoint when a workflow declares one.""" + from reflex.workflow.ingress import ( + WEBHOOK_ROUTE, + collect_webhook_routes, + webhook_endpoint, + ) + + if self._api is None or self._workflow_runtime is None: + return + if not collect_webhook_routes(self._workflow_runtime.definitions): + return + config = get_config() + self._api.add_route( + config.prepend_backend_path(WEBHOOK_ROUTE), + webhook_endpoint(self._workflow_runtime), + methods=["POST"], + ) + def _add_optional_endpoints(self): """Add optional api endpoints (_upload).""" from reflex_components_core.core.upload import Upload, get_upload_dir diff --git a/reflex/compiler/compiler.py b/reflex/compiler/compiler.py index f7fcd8c3509..a06ca2c91b4 100644 --- a/reflex/compiler/compiler.py +++ b/reflex/compiler/compiler.py @@ -1168,6 +1168,7 @@ def compile_app( console.debug(f"BE Evaluating stateful page: {route}") app._compile_page(route, save_page=False) app._add_optional_endpoints() + app._add_workflow_endpoints() return False if constants.Page404.SLUG not in app._unevaluated_pages: @@ -1183,6 +1184,7 @@ def compile_app( app._write_stateful_pages_marker() app._add_optional_endpoints() + app._add_workflow_endpoints() return False progress = ( @@ -1238,6 +1240,7 @@ def compile_app( app._stateful_pages.update(compile_ctx.stateful_routes) app._write_stateful_pages_marker() app._add_optional_endpoints() + app._add_workflow_endpoints() app._validate_var_dependencies() if config.show_built_with_reflex is None: diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index ac3a8af0939..f2f7d7554d0 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -16,10 +16,12 @@ TransientWorkflowError, Trigger, WebhookTrigger, + WebhookVerifier, WorkflowConfig, after, complete, fail, + hmac_signature, manual, needs_attention, parse_duration, @@ -76,6 +78,7 @@ "TransientWorkflowError", "Trigger", "WebhookTrigger", + "WebhookVerifier", "WorkflowConfig", "WorkflowDefinition", "WorkflowKernel", @@ -86,6 +89,7 @@ "complete", "fail", "get_runtime", + "hmac_signature", "manual", "needs_attention", "parse_duration", diff --git a/reflex/workflow/ingress.py b/reflex/workflow/ingress.py new file mode 100644 index 00000000000..314b0c285cd --- /dev/null +++ b/reflex/workflow/ingress.py @@ -0,0 +1,195 @@ +"""HTTP ingress for workflows started by a provider webhook. + +The endpoint is public but never anonymously trusted: it preserves the raw +request body, verifies the provider's signature over those exact bytes, decodes +and validates the payload, then durably admits the run *before* acknowledging +the provider. A provider that redelivers the same event reaches the same run +through the trigger's deduplication key rather than starting a second one. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any + +from pydantic import TypeAdapter, ValidationError +from reflex_base.utils import console +from starlette.requests import Request +from starlette.responses import JSONResponse + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine, Mapping + + from reflex_base.workflow import WebhookTrigger + + from reflex.workflow.definition import HandlerDefinition, WorkflowDefinition + from reflex.workflow.runtime import WorkflowRuntime + +MAX_BODY_BYTES = 1_048_576 + +WEBHOOK_ROUTE = "/_workflow/webhook/{topic:path}" + + +class WebhookRoute: + """One workflow root reachable over HTTP. + + Attributes: + definition: The workflow definition owning the root. + handler: The root handler started by this topic. + trigger: The webhook trigger declaring the topic and its verification. + """ + + __slots__ = ("definition", "handler", "trigger") + + def __init__( + self, + definition: WorkflowDefinition, + handler: HandlerDefinition, + trigger: WebhookTrigger, + ): + """Initialize the route. + + Args: + definition: The workflow definition owning the root. + handler: The root handler started by this topic. + trigger: The webhook trigger declaring the topic. + """ + self.definition = definition + self.handler = handler + self.trigger = trigger + + +def collect_webhook_routes( + definitions: tuple[WorkflowDefinition, ...], +) -> dict[str, WebhookRoute]: + """Index every webhook-triggered root by its topic. + + Args: + definitions: The registered workflow definitions. + + Returns: + Routes keyed by topic. + + Raises: + WorkflowDefinitionError: If two roots claim the same topic, which would + make delivery ambiguous. + """ + from reflex_base.utils.exceptions import WorkflowDefinitionError + from reflex_base.workflow import WebhookTrigger + + routes: dict[str, WebhookRoute] = {} + for definition in definitions: + for handler_id in definition.roots: + handler = definition.handlers[handler_id] + trigger = handler.trigger + if not isinstance(trigger, WebhookTrigger): + continue + existing = routes.get(trigger.topic) + if existing is not None: + msg = ( + f"Webhook topic {trigger.topic!r} is claimed by both " + f"{existing.definition.workflow_id}.{existing.handler.id} and " + f"{definition.workflow_id}.{handler.id}; a topic must " + "identify exactly one root." + ) + raise WorkflowDefinitionError(msg) + routes[trigger.topic] = WebhookRoute(definition, handler, trigger) + return routes + + +def _dedupe_key(trigger: WebhookTrigger, payload: Any) -> str | None: + """Extract the deduplication key a provider redelivery would repeat. + + Args: + trigger: The webhook trigger declaring the key field. + payload: The decoded request payload. + + Returns: + The key as a string, or None when the trigger declares none or the + field is absent. + """ + if trigger.dedupe_by is None or not isinstance(payload, dict): + return None + value = payload.get(trigger.dedupe_by) + return None if value is None else str(value) + + +def _root_args(handler: HandlerDefinition, payload: Any) -> dict[str, Any]: + """Map a decoded payload onto the root handler's parameters. + + Args: + handler: The root handler definition. + payload: The decoded request payload. + + Returns: + The keyword arguments to start the root with. + """ + if not handler.params: + return {} + if len(handler.params) == 1: + return {handler.params[0]: payload} + if isinstance(payload, dict): + return {name: payload.get(name) for name in handler.params} + return {} + + +def webhook_endpoint( + runtime: WorkflowRuntime, +) -> Callable[[Request], Coroutine[Any, Any, JSONResponse]]: + """Build the ASGI endpoint that accepts provider webhooks. + + Args: + runtime: The workflow runtime that owns the registered definitions. + + Returns: + The Starlette endpoint. + """ + + async def endpoint(request: Request) -> JSONResponse: + topic = request.path_params.get("topic", "") + routes = collect_webhook_routes(runtime.definitions) + route = routes.get(topic) + if route is None: + return JSONResponse({"error": "unknown topic"}, status_code=404) + + body = await request.body() + if len(body) > MAX_BODY_BYTES: + return JSONResponse({"error": "payload too large"}, status_code=413) + + headers: Mapping[str, str] = request.headers + if route.trigger.verify is not None: + try: + verified = route.trigger.verify(body, headers) + except Exception: + console.warn(f"Webhook verifier raised for topic {topic!r}.") + verified = False + if not verified: + return JSONResponse({"error": "invalid signature"}, status_code=401) + + try: + payload = json.loads(body) if body else {} + except ValueError: + return JSONResponse({"error": "payload is not JSON"}, status_code=400) + + if route.trigger.model is not None: + try: + TypeAdapter(route.trigger.model).validate_python(payload) + except ValidationError: + return JSONResponse( + {"error": "payload does not match the declared model"}, + status_code=400, + ) + + spec = getattr(route.definition.state_cls, route.handler.name) + args = _root_args(route.handler, payload) + result = await runtime.kernel.start( + spec(**args) if args else spec, + request_key=_dedupe_key(route.trigger, payload), + trigger_kind="webhook", + ) + return JSONResponse( + {"disposition": result.disposition, "run_id": result.run_id}, + status_code=202, + ) + + return endpoint diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 6622fc1c54a..c788fd3fdef 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -28,7 +28,6 @@ After, CompleteRun, FailRun, - ManualTrigger, NeedsAttention, parse_duration, ) @@ -307,27 +306,39 @@ async def start( *, request_key: str | None = None, labels: dict[str, str] | None = None, + trigger_kind: str = "manual", ) -> StartResult: - """Admit a new run from a manual root event. + """Admit a new run from a root event. Args: target: The root event, e.g. ``MyWorkflow.start(payload)``. request_key: Idempotent admission key; a repeated key returns the prior run with disposition ``"deduplicated"``. labels: Server-derived indexing labels to record on the run. + trigger_kind: The ingress path admitting this run. It must match the + root's declared trigger, so a webhook root cannot be started by + application code and a manual root cannot be started by a + provider request. Returns: The admission result. Raises: - WorkflowRuntimeError: If the target is not a manual root handler. + WorkflowRuntimeError: If the target is not a root, or its trigger + does not match the admitting ingress. """ defn, handler, payload = self._resolve_target(target) - if not isinstance(handler.trigger, ManualTrigger): + declared = getattr(handler.trigger, "kind", None) + if declared != trigger_kind: + expected = ( + f"trigger=rx.{trigger_kind}(...)" + if trigger_kind == "manual" + else f"a {trigger_kind} trigger" + ) msg = ( - f"Handler {handler.id!r} of {defn.workflow_id!r} is not a manual " - "root; only handlers with trigger=rx.manual() can be started " - "directly." + f"Handler {handler.id!r} of {defn.workflow_id!r} declares " + f"{declared or 'no trigger'}, so it cannot be started here; " + f"starting through this path requires {expected}." ) raise WorkflowRuntimeError(msg) now = self._clock() diff --git a/tests/units/reflex_base/test_workflow.py b/tests/units/reflex_base/test_workflow.py index 0aad021d2b5..5efb802b004 100644 --- a/tests/units/reflex_base/test_workflow.py +++ b/tests/units/reflex_base/test_workflow.py @@ -12,6 +12,7 @@ after, build_durable_config, default_retry_for_effect, + hmac_signature, manual, parse_duration, schedule, @@ -131,14 +132,15 @@ def test_workflow_config_invalid_max_steps(): def test_triggers(): assert manual().kind == "manual" - hook = webhook("stripe.payment_succeeded", dedupe_by="id") + verifier = hmac_signature(secret_env="SECRET", header="X-Signature") + hook = webhook("stripe.payment_succeeded", verify=verifier, dedupe_by="id") assert hook.kind == "webhook" assert hook.topic == "stripe.payment_succeeded" assert hook.dedupe_by == "id" cron = schedule("0 9 * * 1") assert cron.kind == "schedule" with pytest.raises(WorkflowDefinitionError, match="topic"): - webhook("") + webhook("", verify=verifier) with pytest.raises(WorkflowDefinitionError, match="cron"): schedule("hourly") @@ -236,3 +238,22 @@ def test_build_durable_config_parses_timeout(): config = _build(durable=True, effect="read", timeout="45s") assert config is not None assert config.timeout == pytest.approx(45.0) + + +def test_webhook_requires_authentication(): + """An unverified webhook endpoint would let anyone start runs.""" + with pytest.raises(WorkflowDefinitionError, match="no verifier"): + webhook("stripe.paid") + with pytest.raises(WorkflowDefinitionError, match="unverified_reason"): + webhook("stripe.paid", allow_unverified=True) + public = webhook( + "internal.ping", allow_unverified=True, unverified_reason="internal network" + ) + assert public.allow_unverified + with pytest.raises(WorkflowDefinitionError, match="keep the verifier"): + webhook( + "stripe.paid", + verify=hmac_signature(secret_env="S", header="H"), + allow_unverified=True, + unverified_reason="mixed", + ) diff --git a/tests/units/workflow/test_definition.py b/tests/units/workflow/test_definition.py index 21b86836ad4..e4e9bb888db 100644 --- a/tests/units/workflow/test_definition.py +++ b/tests/units/workflow/test_definition.py @@ -12,6 +12,7 @@ Retry, TransientWorkflowError, WorkflowConfig, + hmac_signature, manual, webhook, ) @@ -262,7 +263,11 @@ class WebhookRoot(rx.State): @rx.event( durable=True, - trigger=webhook("stripe.payment_succeeded", dedupe_by="id"), + trigger=webhook( + "stripe.payment_succeeded", + verify=hmac_signature(secret_env="SECRET", header="X-Signature"), + dedupe_by="id", + ), effect="none", ) def on_payment(self): diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py new file mode 100644 index 00000000000..84d7f7ec8a2 --- /dev/null +++ b/tests/units/workflow/test_ingress.py @@ -0,0 +1,248 @@ +"""Tests for the webhook ingress endpoint.""" + +import hmac +import json + +import pytest +from pydantic import BaseModel +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import WorkflowConfig, hmac_signature, manual, webhook +from starlette.applications import Starlette +from starlette.routing import Route +from starlette.testclient import TestClient + +import reflex as rx +from reflex.workflow.definition import compile_workflow +from reflex.workflow.ingress import ( + WEBHOOK_ROUTE, + collect_webhook_routes, + webhook_endpoint, +) +from reflex.workflow.records import RunStatus +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import MemoryRunStore + +SECRET = "whsec_test" + + +class Payment(BaseModel): + """Typed webhook payload.""" + + id: str + amount: int + + +def _sign(body: bytes) -> str: + """Sign a body the way the provider would. + + Args: + body: The exact bytes that will be sent. + + Returns: + The hex digest to put in the signature header. + """ + return hmac.new(SECRET.encode(), body, "sha256").hexdigest() + + +@pytest.fixture +def paid_workflow(monkeypatch, forked_registration_context): + """A workflow with one verified webhook root. + + Args: + monkeypatch: Fixture used to set the shared secret. + forked_registration_context: Isolates state registration. + + Returns: + The workflow class. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + + class PaidFlow(rx.State): + __workflow__ = WorkflowConfig(id="ingress.paid") + payment_id: str = "" + amount: int = 0 + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "stripe.payment_succeeded", + model=Payment, + verify=hmac_signature( + secret_env="STRIPE_WEBHOOK_SECRET", header="X-Signature" + ), + dedupe_by="id", + ), + ) + def on_paid(self, payment: Payment): + self.payment_id = payment.id + self.amount = payment.amount + + return PaidFlow + + +@pytest.fixture +async def client(paid_workflow): + """A test client wired to the webhook endpoint of a live runtime. + + Args: + paid_workflow: The registered workflow class. + + Yields: + The client and the runtime behind it. + """ + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(paid_workflow) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + with TestClient(app) as test_client: + yield test_client, runtime + await runtime.shutdown() + + +async def test_signed_webhook_starts_a_run(client): + test_client, runtime = client + body = json.dumps({"id": "pay_1", "amount": 4200}).encode() + response = test_client.post( + "/_workflow/webhook/stripe.payment_succeeded", + content=body, + headers={"X-Signature": _sign(body)}, + ) + assert response.status_code == 202 + assert response.json()["disposition"] == "started" + run_id = response.json()["run_id"] + + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"payment_id": "pay_1", "amount": 4200} + + +async def test_redelivery_is_deduplicated(client): + test_client, runtime = client + body = json.dumps({"id": "pay_2", "amount": 1}).encode() + headers = {"X-Signature": _sign(body)} + first = test_client.post( + "/_workflow/webhook/stripe.payment_succeeded", content=body, headers=headers + ) + second = test_client.post( + "/_workflow/webhook/stripe.payment_succeeded", content=body, headers=headers + ) + assert first.json()["disposition"] == "started" + assert second.json()["disposition"] == "deduplicated" + assert second.json()["run_id"] == first.json()["run_id"] + assert len(await runtime.kernel.list_runs()) == 1 + + +async def test_bad_signature_is_rejected_and_admits_nothing(client): + test_client, runtime = client + body = json.dumps({"id": "pay_3", "amount": 1}).encode() + response = test_client.post( + "/_workflow/webhook/stripe.payment_succeeded", + content=body, + headers={"X-Signature": "deadbeef"}, + ) + assert response.status_code == 401 + assert await runtime.kernel.list_runs() == () + + +async def test_missing_signature_is_rejected(client): + test_client, runtime = client + body = json.dumps({"id": "pay_4", "amount": 1}).encode() + response = test_client.post( + "/_workflow/webhook/stripe.payment_succeeded", content=body + ) + assert response.status_code == 401 + assert await runtime.kernel.list_runs() == () + + +async def test_payload_not_matching_the_model_is_rejected(client): + test_client, runtime = client + body = json.dumps({"id": "pay_5"}).encode() + response = test_client.post( + "/_workflow/webhook/stripe.payment_succeeded", + content=body, + headers={"X-Signature": _sign(body)}, + ) + assert response.status_code == 400 + assert await runtime.kernel.list_runs() == () + + +async def test_malformed_json_is_rejected(client): + test_client, runtime = client + body = b"{not json" + response = test_client.post( + "/_workflow/webhook/stripe.payment_succeeded", + content=body, + headers={"X-Signature": _sign(body)}, + ) + assert response.status_code == 400 + assert await runtime.kernel.list_runs() == () + + +async def test_unknown_topic_is_not_found(client): + test_client, runtime = client + body = b"{}" + response = test_client.post( + "/_workflow/webhook/nope.nothing", + content=body, + headers={"X-Signature": _sign(body)}, + ) + assert response.status_code == 404 + assert await runtime.kernel.list_runs() == () + + +def test_duplicate_topics_are_rejected(forked_registration_context, monkeypatch): + """Two roots on one topic would make delivery ambiguous.""" + monkeypatch.setenv("S", SECRET) + verifier = hmac_signature(secret_env="S", header="X-Signature") + + class FirstClaim(rx.State): + __workflow__ = WorkflowConfig(id="ingress.first") + + @rx.event(durable=True, effect="none", trigger=webhook("dup", verify=verifier)) + def go(self): + pass + + class SecondClaim(rx.State): + __workflow__ = WorkflowConfig(id="ingress.second") + + @rx.event(durable=True, effect="none", trigger=webhook("dup", verify=verifier)) + def go(self): + pass + + definitions = (compile_workflow(FirstClaim), compile_workflow(SecondClaim)) + with pytest.raises(WorkflowDefinitionError, match="claimed by both"): + collect_webhook_routes(definitions) + + +async def test_webhook_root_cannot_be_started_by_application_code(paid_workflow): + """A provider-triggered root is not a manual root.""" + from reflex_base.utils.exceptions import WorkflowRuntimeError + + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(paid_workflow) + await runtime.startup(start_worker=False) + try: + with pytest.raises(WorkflowRuntimeError, match="cannot be started here"): + await runtime.kernel.start(paid_workflow.on_paid(Payment(id="x", amount=1))) + finally: + await runtime.shutdown() + + +def test_manual_root_is_not_reachable_over_http( + forked_registration_context, monkeypatch +): + """A manual root has no webhook route at all.""" + + class ManualOnly(rx.State): + __workflow__ = WorkflowConfig(id="ingress.manual_only") + + @rx.event(durable=True, effect="none", trigger=manual()) + def go(self): + pass + + assert collect_webhook_routes((compile_workflow(ManualOnly),)) == {} diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py index 3237b35ddf7..3a08fcb5f78 100644 --- a/tests/units/workflow/test_kernel.py +++ b/tests/units/workflow/test_kernel.py @@ -12,6 +12,7 @@ after, complete, fail, + hmac_signature, manual, needs_attention, webhook, @@ -550,7 +551,12 @@ class StartRules(rx.State): __workflow__ = WorkflowConfig(id="kernel.startrules") @rx.event( - durable=True, trigger=webhook("stripe.payment_succeeded"), effect="none" + durable=True, + trigger=webhook( + "stripe.payment_succeeded", + verify=hmac_signature(secret_env="SECRET", header="X-Signature"), + ), + effect="none", ) def on_webhook(self): pass @@ -564,9 +570,9 @@ def internal(self): pass async with WorkflowTestHarness(StartRules) as harness: - with pytest.raises(WorkflowRuntimeError, match="manual"): + with pytest.raises(WorkflowRuntimeError, match="cannot be started here"): await harness.start(StartRules.on_webhook()) - with pytest.raises(WorkflowRuntimeError, match="manual"): + with pytest.raises(WorkflowRuntimeError, match="cannot be started here"): await harness.start(StartRules.internal()) with pytest.raises(WorkflowRuntimeError, match="workflow"): await harness.start(object()) From 3586943d0d6adfcc56c084ffe0551646529d30e1 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:17:36 -0700 Subject: [PATCH 008/121] Fire cron schedule triggers rx.schedule(...) validated its five fields and then did nothing; scheduled workflows never ran. Adds a dependency-free UTC cron evaluator supporting the standard five fields with ranges, steps and lists, including the day-of-month OR day-of-week rule that every other cron implementation follows. Expressions are validated when the workflow compiles, so a bad expression fails at add_workflow() rather than silently never firing. The kernel admits one run per occurrence under a request key derived from the occurrence time, which reuses the existing dedupe path: a restart, a second process, or an overlapping sweep all converge on exactly one run per occurrence. Cursors are seeded when the kernel is constructed, so deploying a schedule never backfills history, and catch-up after an outage is capped at ten occurrences so a restart cannot stampede. The worker wakes for the next occurrence rather than polling for it. --- news/workflow-schedules.feature.md | 1 + reflex/workflow/cron.py | 194 +++++++++++++++++++++++++ reflex/workflow/definition.py | 5 + reflex/workflow/kernel.py | 68 ++++++++- tests/units/workflow/test_schedules.py | 164 +++++++++++++++++++++ 5 files changed, 429 insertions(+), 3 deletions(-) create mode 100644 news/workflow-schedules.feature.md create mode 100644 reflex/workflow/cron.py create mode 100644 tests/units/workflow/test_schedules.py diff --git a/news/workflow-schedules.feature.md b/news/workflow-schedules.feature.md new file mode 100644 index 00000000000..c9084be41a7 --- /dev/null +++ b/news/workflow-schedules.feature.md @@ -0,0 +1 @@ +Cron schedule triggers now fire: `rx.schedule("0 9 * * 1")` roots run on time, once per occurrence, with bounded catch-up after an outage. diff --git a/reflex/workflow/cron.py b/reflex/workflow/cron.py new file mode 100644 index 00000000000..d2236069c98 --- /dev/null +++ b/reflex/workflow/cron.py @@ -0,0 +1,194 @@ +"""A small UTC cron evaluator for schedule triggers. + +Supports the five standard fields (minute, hour, day of month, month, day of +week) with ``*``, single values, ``a-b`` ranges, ``a-b/step`` and ``*/step`` +steps, and comma-separated lists. When both day-of-month and day-of-week are +restricted, an occurrence matches if *either* matches, which is the behavior +every other cron implementation has. + +Schedules are evaluated in UTC so an occurrence identity never depends on the +server's local timezone or on a daylight-saving transition. +""" + +from __future__ import annotations + +import datetime as dt +from typing import Final + +from reflex_base.utils.exceptions import WorkflowDefinitionError + +_FIELD_RANGES: Final = ((0, 59), (0, 23), (1, 31), (1, 12), (0, 6)) + +_FIELD_NAMES: Final = ("minute", "hour", "day of month", "month", "day of week") + +MAX_SEARCH_DAYS: Final = 1500 + + +def _parse_field(spec: str, index: int) -> frozenset[int]: + """Parse one cron field into the set of values it matches. + + Args: + spec: The field text, e.g. ``"*/15"`` or ``"1,3,5-7"``. + index: Which of the five fields this is. + + Returns: + Every value the field matches. + + Raises: + WorkflowDefinitionError: If the field is malformed or out of range. + """ + low, high = _FIELD_RANGES[index] + name = _FIELD_NAMES[index] + values: set[int] = set() + for part in spec.split(","): + body, _, step_text = part.partition("/") + try: + step = int(step_text) if step_text else 1 + except ValueError: + msg = f"Invalid step {step_text!r} in cron {name} field {spec!r}." + raise WorkflowDefinitionError(msg) from None + if step < 1: + msg = f"Step must be >= 1 in cron {name} field {spec!r}." + raise WorkflowDefinitionError(msg) + if body == "*": + start, end = low, high + elif "-" in body.lstrip("-"): + start_text, _, end_text = body.partition("-") + try: + start, end = int(start_text), int(end_text) + except ValueError: + msg = f"Invalid range {body!r} in cron {name} field {spec!r}." + raise WorkflowDefinitionError(msg) from None + else: + try: + start = end = int(body) + except ValueError: + msg = f"Invalid value {body!r} in cron {name} field {spec!r}." + raise WorkflowDefinitionError(msg) from None + if start < low or end > high or start > end: + msg = f"Cron {name} field {spec!r} is out of range; expected {low}-{high}." + raise WorkflowDefinitionError(msg) + values.update(range(start, end + 1, step)) + return frozenset(values) + + +class CronSchedule: + """A parsed five-field cron expression evaluated in UTC. + + Attributes: + expression: The original expression text. + """ + + __slots__ = ( + "_days_of_month", + "_days_of_week", + "_dom_restricted", + "_dow_restricted", + "_hours", + "_minutes", + "_months", + "expression", + ) + + def __init__(self, expression: str): + """Parse a cron expression. + + Args: + expression: A five-field cron expression. + + Raises: + WorkflowDefinitionError: If the expression does not have five + fields or any field is malformed. + """ + fields = expression.split() + if len(fields) != 5: + msg = ( + f"Invalid cron expression {expression!r}: expected five fields " + "(minute hour day month weekday)." + ) + raise WorkflowDefinitionError(msg) + self.expression = expression + self._minutes = _parse_field(fields[0], 0) + self._hours = _parse_field(fields[1], 1) + self._days_of_month = _parse_field(fields[2], 2) + self._months = _parse_field(fields[3], 3) + self._days_of_week = _parse_field(fields[4], 4) + self._dom_restricted = fields[2] != "*" + self._dow_restricted = fields[4] != "*" + + def _matches_date(self, day: dt.date) -> bool: + """Whether a date satisfies the month and day fields. + + Args: + day: The UTC date to test. + + Returns: + True when the date matches. + """ + if day.month not in self._months: + return False + # Cron numbers weekdays from Sunday; Python numbers them from Monday. + dow = (day.weekday() + 1) % 7 + dom_hit = day.day in self._days_of_month + dow_hit = dow in self._days_of_week + if self._dom_restricted and self._dow_restricted: + return dom_hit or dow_hit + return dom_hit and dow_hit + + def next_after(self, after: float) -> float | None: + """Find the first occurrence strictly after a point in time. + + Args: + after: Epoch seconds to search forward from. + + Returns: + The occurrence time in epoch seconds, or None when the expression + has no occurrence within the search horizon. + """ + moment = dt.datetime.fromtimestamp(after, tz=dt.UTC).replace( + second=0, microsecond=0 + ) + dt.timedelta(minutes=1) + day = moment.date() + for offset in range(MAX_SEARCH_DAYS): + candidate_day = day + dt.timedelta(days=offset) + if not self._matches_date(candidate_day): + continue + first_minute = moment if offset == 0 else None + for hour in sorted(self._hours): + for minute in sorted(self._minutes): + occurrence = dt.datetime( + candidate_day.year, + candidate_day.month, + candidate_day.day, + hour, + minute, + tzinfo=dt.UTC, + ) + if first_minute is not None and occurrence < first_minute: + continue + return occurrence.timestamp() + return None + + def occurrences_between( + self, after: float, until: float, *, limit: int + ) -> list[float]: + """List occurrences in a half-open interval. + + Args: + after: Exclusive lower bound in epoch seconds. + until: Inclusive upper bound in epoch seconds. + limit: Maximum occurrences to return, bounding catch-up after an + outage so a restart cannot stampede. + + Returns: + The occurrence times in ascending order. + """ + found: list[float] = [] + cursor = after + while len(found) < limit: + occurrence = self.next_after(cursor) + if occurrence is None or occurrence > until: + break + found.append(occurrence) + cursor = occurrence + return found diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 28050040c07..c7316bf10b4 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -20,6 +20,7 @@ from reflex_base.workflow import ( DurableEventConfig, Retry, + ScheduleTrigger, Trigger, WorkflowConfig, default_retry_for_effect, @@ -27,6 +28,7 @@ parse_duration, ) +from reflex.workflow.cron import CronSchedule from reflex.workflow.serde import to_run_data if TYPE_CHECKING: @@ -535,6 +537,9 @@ def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: config = _validate_class_shape(workflow_cls) fields = _compile_fields(workflow_cls) handlers = _resolve_hooks(workflow_cls, _compile_handlers(workflow_cls, config)) + for defn in handlers.values(): + if isinstance(defn.trigger, ScheduleTrigger): + CronSchedule(defn.trigger.cron) durable_names = frozenset(defn.name for defn in handlers.values()) for defn in handlers.values(): _validate_handler_body(workflow_cls, defn, durable_names) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index c788fd3fdef..d7dcee2b10e 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -29,10 +29,12 @@ CompleteRun, FailRun, NeedsAttention, + ScheduleTrigger, parse_duration, ) from reflex.event import EventHandler, EventSpec +from reflex.workflow.cron import CronSchedule from reflex.workflow.records import ( TERMINAL_STEP_STATUSES, HistoryEventType, @@ -59,6 +61,8 @@ RECOVERY_INTERVAL_FRACTION = 1 / 2 +MAX_SCHEDULE_CATCHUP = 10 + def _error_payload(error: BaseException) -> dict[str, Any]: """Build a JSON-compatible error payload from an exception. @@ -209,6 +213,18 @@ def __init__( self._lease_duration = lease_duration self._lease_renew_interval = renew self._recovery_interval = recovery + self._schedules = [ + (defn, handler, CronSchedule(handler.trigger.cron)) + for defn in self._definitions.values() + for handler in (defn.handlers[hid] for hid in defn.roots) + if isinstance(handler.trigger, ScheduleTrigger) + ] + # Seeded at construction so a freshly started process never backfills + # occurrences from before it existed. + self._schedule_cursor: dict[str, float] = { + f"{defn.workflow_id}:{handler.id}": clock() + for defn, handler, _ in self._schedules + } self._field_adapters: dict[tuple[str, str], TypeAdapter] = {} self._inflight: dict[str, asyncio.Task] = {} self._leases: dict[str, _Lease] = {} @@ -1256,6 +1272,51 @@ async def _execute_claim(self, claim: Claim) -> None: except StaleClaimError: await self._record_abandoned(claim, handler, "fenced_at_commit") + async def _admit_due_schedules(self, now: float) -> int: + """Admit a run for every schedule occurrence that has come due. + + Each occurrence is admitted under a stable request key derived from its + exact time, so a restart, a second worker, or an overlapping sweep all + converge on one run per occurrence rather than a stampede. + + Args: + now: Current time in epoch seconds. + + Returns: + The number of runs admitted. + """ + admitted = 0 + for defn, handler, schedule in self._schedules: + key = f"{defn.workflow_id}:{handler.id}" + cursor = self._schedule_cursor[key] + for occurrence in schedule.occurrences_between( + cursor, now, limit=MAX_SCHEDULE_CATCHUP + ): + result = await self.start( + getattr(defn.state_cls, handler.name), + request_key=f"schedule:{key}:{int(occurrence)}", + trigger_kind="schedule", + ) + admitted += result.disposition == "started" + self._schedule_cursor[key] = now + return admitted + + def _next_schedule_due(self, now: float) -> float | None: + """Earliest time any registered schedule next fires. + + Args: + now: Current time in epoch seconds. + + Returns: + The epoch time, or None when nothing is scheduled. + """ + upcoming = [ + occurrence + for _, _, schedule in self._schedules + if (occurrence := schedule.next_after(now)) is not None + ] + return min(upcoming) if upcoming else None + async def _tick(self) -> bool: """Run one scheduling round. @@ -1263,7 +1324,7 @@ async def _tick(self) -> bool: True if any control transition or attempt was processed. """ now = self._clock() - progressed = False + progressed = await self._admit_due_schedules(now) > 0 for run in await self._store.control_pending(now): if run.cancel_requested: progressed = ( @@ -1336,8 +1397,9 @@ async def _worker_loop(self) -> None: now = self._clock() due = await self._store.next_due(now) delay = min(self._poll_interval, max(self._next_recovery_at - now, 0.0)) - if due is not None: - delay = min(delay, max(due - now, 0.0)) + for upcoming in (due, self._next_schedule_due(now)): + if upcoming is not None: + delay = min(delay, max(upcoming - now, 0.0)) if delay <= 0: continue self._wakeup.clear() diff --git a/tests/units/workflow/test_schedules.py b/tests/units/workflow/test_schedules.py new file mode 100644 index 00000000000..f22358ce3c0 --- /dev/null +++ b/tests/units/workflow/test_schedules.py @@ -0,0 +1,164 @@ +"""Tests for cron schedule triggers.""" + +import datetime as dt + +import pytest +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import WorkflowConfig, schedule + +import reflex as rx +from reflex.workflow.cron import CronSchedule +from reflex.workflow.definition import compile_workflow +from reflex.workflow.records import RunStatus +from reflex.workflow.testing import WorkflowTestHarness + +# A Tuesday at 12:00 UTC, chosen so quarter-hour schedules are 15 minutes away. +START = dt.datetime(2026, 8, 18, 12, 0, tzinfo=dt.UTC).timestamp() + + +def _at(expression: str, after: float = START) -> str: + """Format the next occurrence of an expression for readable assertions. + + Args: + expression: The cron expression. + after: Epoch seconds to search from. + + Returns: + The occurrence as ``"Day YYYY-MM-DD HH:MM"`` in UTC. + """ + occurrence = CronSchedule(expression).next_after(after) + assert occurrence is not None + return dt.datetime.fromtimestamp(occurrence, tz=dt.UTC).strftime( + "%a %Y-%m-%d %H:%M" + ) + + +@pytest.mark.parametrize( + ("expression", "expected"), + [ + ("*/15 * * * *", "Tue 2026-08-18 12:15"), + ("0 * * * *", "Tue 2026-08-18 13:00"), + ("0 9 * * 1", "Mon 2026-08-24 09:00"), + ("30 2 1 * *", "Tue 2026-09-01 02:30"), + ("0 0 1 1 *", "Fri 2027-01-01 00:00"), + ("0 12 * * 3", "Wed 2026-08-19 12:00"), + ], +) +def test_next_occurrence(expression, expected): + assert _at(expression) == expected + + +def test_day_of_month_or_day_of_week(): + """With both day fields restricted, either one matching is a match.""" + # The 21st is a Friday, so the Friday rule fires before the 13th does. + assert _at("0 12 13 * 5") == "Fri 2026-08-21 12:00" + # With only day-of-month restricted, weekdays are irrelevant. + assert _at("0 12 13 * *") == "Sun 2026-09-13 12:00" + + +@pytest.mark.parametrize( + "expression", + ["* * * *", "* * * * * *", "60 * * * *", "* 24 * * *", "*/0 * * * *", "a * * * *"], +) +def test_invalid_expressions_are_rejected(expression): + with pytest.raises(WorkflowDefinitionError): + CronSchedule(expression) + + +def test_occurrences_between_is_bounded(): + """Catch-up after an outage is capped so a restart cannot stampede.""" + every_ten = CronSchedule("*/10 * * * *") + assert len(every_ten.occurrences_between(START, START + 3600, limit=4)) == 4 + assert len(every_ten.occurrences_between(START, START + 3600, limit=100)) == 6 + + +def test_invalid_cron_is_rejected_at_compile_time(forked_registration_context): + class BadSchedule(rx.State): + __workflow__ = WorkflowConfig(id="ops.bad_schedule") + + @rx.event(durable=True, trigger=schedule("0 99 * * *"), effect="none") + def sweep(self): + pass + + with pytest.raises(WorkflowDefinitionError, match="out of range"): + compile_workflow(BadSchedule) + + +async def test_schedule_fires_on_the_virtual_clock(forked_registration_context): + fires = [] + + class Sweeper(rx.State): + __workflow__ = WorkflowConfig(id="ops.sweeper") + ran: int = 0 + + @rx.event(durable=True, trigger=schedule("*/15 * * * *"), effect="read") + def sweep(self): + fires.append(1) + self.ran = 1 + + async with WorkflowTestHarness(Sweeper, start_time=START) as harness: + # Deploying a schedule never backfills the past. + await harness.run_until_idle() + assert fires == [] + + await harness.advance("16m") + assert len(fires) == 1 + + await harness.advance("31m") + assert len(fires) == 3 + + runs = await harness.kernel.list_runs() + assert len(runs) == 3 + assert all(run.status is RunStatus.COMPLETED for run in runs) + # One run per occurrence, keyed by the occurrence time. + assert len({run.request_key for run in runs}) == 3 + + +async def test_occurrences_are_admitted_once_across_restarts( + forked_registration_context, tmp_path +): + """A restart re-fires nothing: occurrence keys deduplicate admission.""" + from reflex.workflow.store import SqliteRunStore + + class Restarted(rx.State): + __workflow__ = WorkflowConfig(id="ops.restarted") + ran: int = 0 + + @rx.event(durable=True, trigger=schedule("*/15 * * * *"), effect="read") + def sweep(self): + self.ran += 1 + + db_path = tmp_path / "workflow.db" + first = SqliteRunStore(db_path) + async with WorkflowTestHarness(Restarted, store=first, start_time=START) as harness: + await harness.advance("16m") + assert len(await harness.kernel.list_runs()) == 1 + resume_at = harness.now + first.close() + + # A second process starts with a cursor at its own "now" and must not + # re-admit the occurrence the first one already handled. + second = SqliteRunStore(db_path) + async with WorkflowTestHarness( + Restarted, store=second, start_time=resume_at - 300 + ) as harness: + await harness.advance("10m") + assert len(await harness.kernel.list_runs()) == 1 + second.close() + + +async def test_schedule_root_cannot_be_started_by_application_code( + forked_registration_context, +): + from reflex_base.utils.exceptions import WorkflowRuntimeError + + class Cronly(rx.State): + __workflow__ = WorkflowConfig(id="ops.cronly") + + @rx.event(durable=True, trigger=schedule("0 0 * * *"), effect="none") + def sweep(self): + pass + + async with WorkflowTestHarness(Cronly, start_time=START) as harness: + with pytest.raises(WorkflowRuntimeError, match="cannot be started here"): + await harness.kernel.start(Cronly.sweep) From 91a75c2b7f52285f1ae6cc327f7c974d8953b0c0 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:24:42 -0700 Subject: [PATCH 009/121] Add durable waits and typed signals A run could only move forward on its own timers, so human approvals, event-driven waits, and anything needing an answer from outside were simply inexpressible -- the gap that most separated this from Temporal and Inngest. A wait is now a BLOCKED slot at the run's frontier, which keeps the mailbox strictly serial: no second open slot, no change to the per-run fence, no concurrent commits. The slot carries the address a delivery must match, and its due_at doubles as the deadline, so a wait timeout fires through the same timer path that rx.after() already used and the virtual clock drives it with no new machinery. The race is settled on one row. A delivery compare-and-swaps the blocked slot to ready and hands the payload to the resume handler; a deadline instead makes the slot claimable, and claiming it *is* the timeout branch. Whichever lands first erases the other's trigger, so a late signal to a finished run is refused rather than silently dropped. A signal that arrives before the run reaches its wait is buffered and consumed by the arming commit itself, so a sender faster than the workflow cannot block it forever. Deliveries never write run state or the state version, so a delivery can never fence a live attempt. BLOCKED is deliberately not a claimable status: a wait with no deadline would otherwise report itself due at time zero and spin the worker against the database, starving lease renewal. Claimability is now one predicate both stores share, with a test asserting an unbounded wait leaves nothing claimable and nothing scheduled. --- news/workflow-waits-signals.feature.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 171 ++++++++++ pyi_hashes.json | 2 +- reflex/__init__.py | 3 + reflex/workflow/__init__.py | 12 + reflex/workflow/kernel.py | 128 ++++++- reflex/workflow/records.py | 54 ++- reflex/workflow/runtime.py | 24 +- reflex/workflow/store.py | 316 ++++++++++++++++-- reflex/workflow/testing.py | 17 + tests/units/workflow/test_waits.py | 276 +++++++++++++++ 11 files changed, 974 insertions(+), 30 deletions(-) create mode 100644 news/workflow-waits-signals.feature.md create mode 100644 tests/units/workflow/test_waits.py diff --git a/news/workflow-waits-signals.feature.md b/news/workflow-waits-signals.feature.md new file mode 100644 index 00000000000..8b88d19c69c --- /dev/null +++ b/news/workflow-waits-signals.feature.md @@ -0,0 +1 @@ +Workflows can now pause for external events: `rx.wait_for()` blocks a run until a typed `rx.Signal` arrives or a deadline passes, which makes human approvals and event-driven waits expressible. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index fb54a98af45..7bf0497c4bc 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -711,3 +711,174 @@ def after(delay: DurationLike, target: Any) -> After: The control return value. """ return After(delay=delay, target=target) + + +class _Never: + """Sentinel meaning a wait has no deadline.""" + + def __repr__(self) -> str: + """Render the sentinel. + + Returns: + The public spelling of this value. + """ + return "rx.never" + + +never: Final = _Never() + + +@dataclasses.dataclass(frozen=True, slots=True) +class ChannelDelivery: + """A payload addressed to one named channel of a run. + + Attributes: + channel: The channel name the waiting run is listening on. + payload: JSON-compatible payload, already validated against the + channel's declared model. + """ + + channel: str + payload: Any + + +class Signal: + """A typed, named channel a run can wait on and outside code can deliver to. + + Declared as a class attribute on a workflow, which keeps it out of the run + state schema and out of the event-handler registry:: + + class Onboarding(rx.State): + docs_uploaded = rx.Signal(Docs) + + Attributes: + model: The payload model deliveries are validated against, if any. + name: The channel name, defaulting to the attribute name. + """ + + def __init__(self, model: type | None = None, *, name: str | None = None): + """Declare a channel. + + Args: + model: The payload model deliveries must satisfy. + name: Explicit channel name; defaults to the attribute name. + """ + self.model = model + self.name = name or "" + + def __set_name__(self, owner: type, name: str) -> None: + """Adopt the attribute name as the channel name. + + Args: + owner: The declaring class. + name: The attribute name. + """ + if not self.name: + self.name = name + + def __call__(self, payload: Any = None) -> ChannelDelivery: + """Build a delivery for this channel. + + Args: + payload: The payload to deliver. + + Returns: + The addressed delivery. + + Raises: + WorkflowDefinitionError: If the payload does not match the declared + model, caught at the call site rather than inside a run. + """ + if self.model is not None: + if isinstance(payload, self.model): + pass + elif isinstance(payload, dict): + payload = self.model(**payload) + else: + msg = ( + f"Channel {self.name!r} expects {self.model.__name__}, got " + f"{type(payload).__name__}." + ) + raise WorkflowDefinitionError(msg) + return ChannelDelivery(channel=self.name, payload=payload) + + +@dataclasses.dataclass(frozen=True, slots=True) +class WaitFor: + """Control return that blocks a run until a signal or a deadline. + + Attributes: + channel: The channel name to wait on. + then: Handler to run when a delivery arrives; it takes the payload. + timeout: How long to wait, or ``rx.never`` for no deadline. + on_timeout: Handler to run when the deadline arrives first. + """ + + channel: str + then: Any + timeout: DurationLike | _Never + on_timeout: Any = None + + def __post_init__(self): + """Validate the wait eagerly so authoring errors surface in place. + + Raises: + WorkflowDefinitionError: If a bounded wait has no timeout branch. + """ + if isinstance(self.timeout, _Never): + if self.on_timeout is not None: + msg = ( + "wait_for(timeout=rx.never) cannot have on_timeout: a wait " + "with no deadline never times out." + ) + raise WorkflowDefinitionError(msg) + return + parse_duration(self.timeout, param="wait_for() timeout") + if self.on_timeout is None: + msg = ( + "wait_for(timeout=...) requires on_timeout, naming the handler " + "that runs when the deadline arrives first. Use " + "timeout=rx.never to wait indefinitely." + ) + raise WorkflowDefinitionError(msg) + + +def wait_for( + channel: Signal, + *, + then: Any, + timeout: DurationLike | _Never, + on_timeout: Any = None, +) -> WaitFor: + """Block the run until a signal arrives or the deadline passes. + + Whichever lands first wins, and the loser can no longer resolve the wait:: + + return rx.wait_for( + Onboarding.docs_uploaded, + then=Onboarding.verify, + timeout="3d", + on_timeout=Onboarding.nag, + ) + + Args: + channel: The channel declared on the workflow class. + then: Handler to run with the delivered payload. + timeout: How long to wait, or ``rx.never``. + on_timeout: Handler to run if the deadline arrives first. + + Returns: + The control return value. + + Raises: + WorkflowDefinitionError: If the channel is not an rx.Signal. + """ + if not isinstance(channel, Signal): + msg = ( + f"wait_for() expects a channel declared with rx.Signal(...), got " + f"{channel!r}." + ) + raise WorkflowDefinitionError(msg) + return WaitFor( + channel=channel.name, then=then, timeout=timeout, on_timeout=on_timeout + ) diff --git a/pyi_hashes.json b/pyi_hashes.json index e428d4b7044..8d8ae67a496 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -118,7 +118,7 @@ "packages/reflex-components-recharts/src/reflex_components_recharts/polar.pyi": "99ebcfc07868061bdc3c2010d85a153f", "packages/reflex-components-recharts/src/reflex_components_recharts/recharts.pyi": "4f6c26f8c76543cc41e2b9dc400ece8a", "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "f170ac685b6ba5892370166c80684db3", - "reflex/__init__.pyi": "41df8e648e62b67ddcd0ec6aa9d88f4c", + "reflex/__init__.pyi": "604f89bbb7fb278a6467c529da5e57d5", "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388", "reflex/experimental/memo.pyi": "bc8b48357bef580e70a5881b65d3d3f7" } diff --git a/reflex/__init__.py b/reflex/__init__.py index fed4f822e9a..b0491f13804 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -244,6 +244,9 @@ "webhook", "hmac_signature", "schedule", + "Signal", + "wait_for", + "never", "after", "complete", "fail", diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index f2f7d7554d0..062c981d150 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -8,13 +8,16 @@ """ from reflex_base.workflow import ( + ChannelDelivery, DurableEventConfig, EffectClass, ManualTrigger, Retry, ScheduleTrigger, + Signal, TransientWorkflowError, Trigger, + WaitFor, WebhookTrigger, WebhookVerifier, WorkflowConfig, @@ -24,8 +27,10 @@ hmac_signature, manual, needs_attention, + never, parse_duration, schedule, + wait_for, webhook, ) @@ -48,6 +53,7 @@ ) from reflex.workflow.runtime import WorkflowRuntime, get_runtime, workflows from reflex.workflow.store import ( + DeliveryDisposition, MemoryRunStore, RunStore, SqliteRunStore, @@ -56,6 +62,8 @@ from reflex.workflow.testing import WorkflowTestHarness __all__ = [ + "ChannelDelivery", + "DeliveryDisposition", "DurableEventConfig", "EffectClass", "HandlerDefinition", @@ -70,6 +78,7 @@ "RunStatus", "RunStore", "ScheduleTrigger", + "Signal", "SqliteRunStore", "StaleClaimError", "StartResult", @@ -77,6 +86,7 @@ "StepStatus", "TransientWorkflowError", "Trigger", + "WaitFor", "WebhookTrigger", "WebhookVerifier", "WorkflowConfig", @@ -92,8 +102,10 @@ "hmac_signature", "manual", "needs_attention", + "never", "parse_duration", "schedule", + "wait_for", "webhook", "workflows", ] diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index d7dcee2b10e..7a092c41589 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -26,10 +26,13 @@ DEFAULT_LEASE_DURATION, DEFAULT_MAX_RECOVERIES, After, + ChannelDelivery, CompleteRun, FailRun, NeedsAttention, ScheduleTrigger, + WaitFor, + _Never, parse_duration, ) @@ -47,7 +50,13 @@ StepStatus, ) from reflex.workflow.serde import to_run_data -from reflex.workflow.store import Claim, RunStore, StaleClaimError, StepCompletion +from reflex.workflow.store import ( + Claim, + DeliveryDisposition, + RunStore, + StaleClaimError, + StepCompletion, +) if TYPE_CHECKING: from collections.abc import Callable, Iterable, Mapping @@ -422,6 +431,34 @@ async def cancel(self, run_id: str) -> bool: self._wakeup.set() return recorded + async def signal( + self, + run_id: str, + delivery: ChannelDelivery, + *, + key: str | None = None, + ) -> DeliveryDisposition: + """Deliver a payload to a run waiting on one of its channels. + + Args: + run_id: The receiving run. + delivery: The addressed payload, e.g. ``MyFlow.approved(decision)``. + key: Sender idempotency key; a repeated key is a no-op. + + Returns: + What the store did with the delivery. + """ + disposition = await self._store.deliver( + run_id, + f"sig:{delivery.channel}", + key or uuid.uuid4().hex, + to_run_data({"value": delivery.payload})["value"], + self._clock(), + ) + if disposition == "resolved": + self._wakeup.set() + return disposition + async def resume(self, run_id: str) -> bool: """Re-open a run that is suspended for operator attention. @@ -593,7 +630,9 @@ def _resolve_successor( def _interpret_return( self, defn: WorkflowDefinition, value: Any - ) -> tuple[list[_SuccessorSpec], CompleteRun | FailRun | NeedsAttention | None]: + ) -> tuple[ + list[_SuccessorSpec], CompleteRun | FailRun | NeedsAttention | WaitFor | None + ]: """Interpret a durable handler's return value. Args: @@ -609,7 +648,7 @@ def _interpret_return( """ if value is None: return [], None - if isinstance(value, (CompleteRun, FailRun, NeedsAttention)): + if isinstance(value, (CompleteRun, FailRun, NeedsAttention, WaitFor)): return [], value if isinstance(value, (list, tuple)): successors = [] @@ -637,6 +676,10 @@ async def _invoke( Returns: The handler return value. """ + args = {key: value for key, value in args.items() if key != "__wait__"} + delivered = args.pop("__payload__", None) + if delivered is not None and handler.params: + args[handler.params[0]] = delivered try: payload = _transform_event_payload(args, handler.type_hints) except Exception: @@ -886,7 +929,7 @@ def _success_completion( steps: tuple[StepRecord, ...], state: dict[str, Any], successors: list[_SuccessorSpec], - control: CompleteRun | FailRun | NeedsAttention | None, + control: CompleteRun | FailRun | NeedsAttention | WaitFor | None, now: float, ) -> StepCompletion: """Build the commit for a successful attempt. @@ -934,6 +977,53 @@ def _success_completion( run_error=error, events=tuple(events), ) + if isinstance(control, WaitFor): + resume = self._resolve_successor(defn, control.then) + timeout_id = ( + self._resolve_successor(defn, control.on_timeout).handler_id + if control.on_timeout is not None + else None + ) + deadline = ( + 0.0 + if isinstance(control.timeout, _Never) + else now + parse_duration(control.timeout) + ) + wait_key = f"sig:{control.channel}" + slot = StepRecord( + run_id=claim.run.run_id, + ordinal=claim.run.next_ordinal, + handler_id=resume.handler_id, + status=StepStatus.BLOCKED, + args={ + **resume.args, + "__wait__": { + "channel": control.channel, + "on_timeout": timeout_id, + }, + }, + due_at=deadline, + wait_key=wait_key, + origin="wait", + created_at=now, + updated_at=now, + ) + events.append(( + HistoryEventType.WAIT_ARMED, + { + "ordinal": slot.ordinal, + "wait_key": wait_key, + "deadline": deadline or None, + }, + )) + return StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.WAITING, + state=state, + new_steps=(slot,), + next_ordinal=claim.run.next_ordinal + 1, + events=tuple(events), + ) if isinstance(control, CompleteRun): tombstones = self._open_ordinals(steps, exclude=claim.step.ordinal) events.extend( @@ -1169,7 +1259,8 @@ def _incompatible_reason( f"{claim.run.workflow_id!r}; restore it or cancel the run." ), } - unexpected = sorted(set(claim.step.args) - set(handler.params)) + supplied = {key for key in claim.step.args if not key.startswith("__")} + unexpected = sorted(supplied - set(handler.params)) if unexpected: return { "reason": "incompatible_payload", @@ -1182,6 +1273,29 @@ def _incompatible_reason( } return None + @staticmethod + def _expired_wait_handler( + defn: WorkflowDefinition, claim: Claim + ) -> HandlerDefinition | None: + """Pick the timeout branch when a wait is claimed at its deadline. + + A wait resolved by a delivery arrives carrying a payload; a wait + claimed without one reached its deadline, so the timeout branch runs + instead of the resume branch. + + Args: + defn: The workflow definition. + claim: The claim being executed. + + Returns: + The timeout handler, or None when the wait was resolved normally. + """ + wait = claim.step.args.get("__wait__") + if not isinstance(wait, dict) or "__payload__" in claim.step.args: + return None + timeout_id = wait.get("on_timeout") + return defn.handlers.get(timeout_id) if timeout_id else None + async def _execute_claim(self, claim: Claim) -> None: """Execute one claimed attempt and commit its outcome. @@ -1209,6 +1323,10 @@ async def _execute_claim(self, claim: Claim) -> None: ) return handler = defn.handlers[claim.step.handler_id] + if claim.step.status is StepStatus.CLAIMED and claim.step.wait_key is not None: + expired = self._expired_wait_handler(defn, claim) + if expired is not None: + handler = expired steps = await self._store.get_steps(claim.run.run_id) await self._store.append_events( claim.run.run_id, diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index c87122b99ac..4f2deb047d6 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -43,6 +43,7 @@ class StepStatus(str, enum.Enum): """ READY = "READY" + BLOCKED = "BLOCKED" CLAIMED = "CLAIMED" RETRY_WAIT = "RETRY_WAIT" RECOVERY_WAIT = "RECOVERY_WAIT" @@ -68,6 +69,46 @@ class StepStatus(str, enum.Enum): )) +def step_claimable_at(step: StepRecord, now: float) -> bool: + """Whether a slot may be claimed at a point in time. + + A blocked slot is claimable only once its deadline arrives, because + claiming it *is* the timeout branch. A blocked slot with ``due_at == 0`` + waits forever, which is why ``BLOCKED`` is deliberately not a member of + ``CLAIMABLE_STEP_STATUSES``: that set is read by callers that do not bound + ``due_at``, and treating a deadline-less wait as claimable would spin. + + Args: + step: The slot to test. + now: Current time in epoch seconds. + + Returns: + True when the slot may be claimed right now. + """ + if step.status in CLAIMABLE_STEP_STATUSES: + return step.due_at <= now + if step.status is StepStatus.BLOCKED: + return 0.0 < step.due_at <= now + return False + + +def step_wake_at(step: StepRecord) -> float | None: + """When a slot next becomes claimable, for scheduler sleep bounds. + + Args: + step: The slot to test. + + Returns: + The epoch time, or None when no clock event alone can make it + claimable, as for a wait with no deadline. + """ + if step.status in CLAIMABLE_STEP_STATUSES: + return step.due_at + if step.status is StepStatus.BLOCKED and step.due_at > 0.0: + return step.due_at + return None + + class HistoryEventType(str, enum.Enum): """Type of an append-only run history event.""" @@ -89,6 +130,12 @@ class HistoryEventType(str, enum.Enum): RUN_CANCELLED = "run_cancelled" RUN_NEEDS_ATTENTION = "run_needs_attention" RUN_RESUMED = "run_resumed" + WAIT_ARMED = "wait_armed" + WAIT_RESOLVED = "wait_resolved" + WAIT_EXPIRED = "wait_expired" + SIGNAL_DELIVERED = "signal_delivered" + SIGNAL_BUFFERED = "signal_buffered" + SIGNAL_DUPLICATE = "signal_duplicate" @dataclasses.dataclass(frozen=True, slots=True) @@ -147,8 +194,10 @@ class StepRecord: lease_expires_at: Epoch time this claim's lease lapses; 0 when the step is not claimed. A claim whose lease has lapsed is treated as orphaned and is reclaimed by recovery, never by a direct claim. + wait_key: For a blocked slot, the address a delivery must carry, as + ``"sig:"`` or ``"approval:"``. None otherwise. error: Last recorded attempt error payload. - origin: How the slot was allocated (root, chain, delay, or hook). + origin: How the slot was allocated. created_at: Allocation time in epoch seconds. updated_at: Last transition time in epoch seconds. """ @@ -163,8 +212,9 @@ class StepRecord: due_at: float = 0.0 epoch: int = 0 lease_expires_at: float = 0.0 + wait_key: str | None = None error: dict[str, Any] | None = None - origin: Literal["root", "chain", "delay", "hook"] = "chain" + origin: Literal["root", "chain", "delay", "hook", "wait"] = "chain" created_at: float = 0.0 updated_at: float = 0.0 diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index bc0d4d8da65..2be092547f6 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -17,12 +17,15 @@ from reflex_base.registry import RegistrationContext from reflex_base.utils.exceptions import WorkflowDefinitionError, WorkflowRuntimeError -from reflex_base.workflow import DEFAULT_LEASE_DURATION +from reflex_base.workflow import DEFAULT_LEASE_DURATION, ChannelDelivery from reflex.workflow.definition import WorkflowDefinition, compile_workflow from reflex.workflow.kernel import DEFAULT_POLL_INTERVAL, WorkflowKernel from reflex.workflow.store import RunStore, SqliteRunStore +if TYPE_CHECKING: + from reflex.workflow.store import DeliveryDisposition + if TYPE_CHECKING: from collections.abc import AsyncIterator, Callable, Iterable, Mapping @@ -269,6 +272,25 @@ async def cancel(run_id: str) -> bool: """ return await get_runtime().kernel.cancel(run_id) + @staticmethod + async def signal( + run_id: str, + delivery: ChannelDelivery, + *, + key: str | None = None, + ) -> DeliveryDisposition: + """Deliver a payload to a run waiting on one of its channels. + + Args: + run_id: The receiving run. + delivery: The addressed payload, e.g. ``MyFlow.approved(decision)``. + key: Sender idempotency key; a repeated key is a no-op. + + Returns: + What the store did with the delivery. + """ + return await get_runtime().kernel.signal(run_id, delivery, key=key) + @staticmethod async def resume(run_id: str) -> bool: """Re-open a run suspended for operator attention. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 1106969ec9a..0b29d3152e0 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -20,13 +20,12 @@ import json import sqlite3 import threading -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from reflex_base.utils.exceptions import WorkflowRuntimeError from reflex_base.workflow import DEFAULT_LEASE_DURATION from reflex.workflow.records import ( - CLAIMABLE_STEP_STATUSES, TERMINAL_RUN_STATUSES, TERMINAL_STEP_STATUSES, HistoryEvent, @@ -36,6 +35,8 @@ RunStatus, StepRecord, StepStatus, + step_claimable_at, + step_wake_at, ) if TYPE_CHECKING: @@ -43,6 +44,15 @@ from pathlib import Path +DeliveryDisposition = Literal[ + "resolved", + "buffered", + "duplicate", + "unknown_run", + "run_terminal", +] + + class StaleClaimError(WorkflowRuntimeError): """Raised when a commit no longer owns its claim and must be discarded.""" @@ -206,6 +216,37 @@ async def append_events( """ ... + async def deliver( + self, + run_id: str, + wait_key: str, + dedupe_key: str, + payload: dict[str, Any], + now: float, + ) -> DeliveryDisposition: + """Deliver a payload to a run, resolving its wait or buffering it. + + The delivery and the wait contend on one row: whichever of the delivery + and the deadline lands first flips the blocked slot and the other can + no longer resolve it. A delivery that arrives before the run has armed + its wait is buffered, and the arming commit consumes it atomically, so + a fast signal is never lost. + + This path must never write ``run.state`` or ``run.state_version``: a + delivery must not be able to fence an attempt. + + Args: + run_id: The receiving run. + wait_key: The address the waiting slot declared. + dedupe_key: Sender-supplied identity, making redelivery a no-op. + payload: JSON-compatible payload to hand the resuming handler. + now: Current time in epoch seconds. + + Returns: + What the store did with the delivery. + """ + ... + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. @@ -424,6 +465,8 @@ def __init__(self): self._steps: dict[str, list[StepRecord]] = {} self._history: dict[str, list[HistoryEvent]] = {} self._dedupe: dict[tuple[str, str], str] = {} + self._inbox: dict[str, dict[tuple[str, str, str], bool]] = {} + self._pending: dict[str, dict[str, dict[str, Any]]] = {} def _append_events( self, @@ -497,11 +540,7 @@ async def claim_next( continue steps = self._steps[run.run_id] frontier = _frontier(steps) - if ( - frontier is None - or frontier.status not in CLAIMABLE_STEP_STATUSES - or frontier.due_at > now - ): + if frontier is None or not step_claimable_at(frontier, now): continue claimed = dataclasses.replace( frontier, @@ -576,6 +615,33 @@ async def renew_lease( ) return True + def _arm(self, step: StepRecord, now: float) -> StepRecord: + """Resolve a newly armed wait against an already-buffered delivery. + + A signal that arrives before the run reaches its wait is buffered, so + arming must consume it in the same commit; otherwise a fast sender + would block the run forever. + + Args: + step: The slot being appended. + now: Current time in epoch seconds. + + Returns: + The slot, already resolved when a matching delivery was waiting. + """ + if step.status is not StepStatus.BLOCKED or step.wait_key is None: + return step + buffered = self._pending.get(step.run_id, {}).pop(step.wait_key, None) + if buffered is None: + return step + return dataclasses.replace( + step, + status=StepStatus.READY, + due_at=now, + args={**step.args, "__payload__": buffered}, + updated_at=now, + ) + async def commit( self, claim: Claim, completion: StepCompletion, now: float ) -> None: @@ -604,7 +670,8 @@ async def commit( steps[ordinal] = dataclasses.replace( slot, status=StepStatus.CANCELLED, updated_at=now ) - steps.extend(completion.new_steps) + for new_step in completion.new_steps: + steps.append(self._arm(new_step, now)) self._runs[run.run_id] = dataclasses.replace( run, status=completion.run_status, @@ -667,6 +734,69 @@ async def append_events( async with self._lock: self._append_events(run_id, events, now) + async def deliver( + self, + run_id: str, + wait_key: str, + dedupe_key: str, + payload: dict[str, Any], + now: float, + ) -> DeliveryDisposition: + """Deliver a payload to a run, resolving its wait or buffering it. + + Args: + run_id: The receiving run. + wait_key: The address the waiting slot declared. + dedupe_key: Sender-supplied identity, making redelivery a no-op. + payload: JSON-compatible payload to hand the resuming handler. + now: Current time in epoch seconds. + + Returns: + What the store did with the delivery. + """ + async with self._lock: + run = self._runs.get(run_id) + if run is None: + return "unknown_run" + if run.status in TERMINAL_RUN_STATUSES: + return "run_terminal" + inbox = self._inbox.setdefault(run_id, {}) + if (run_id, wait_key, dedupe_key) in inbox: + return "duplicate" + inbox[run_id, wait_key, dedupe_key] = True + steps = self._steps[run_id] + frontier = _frontier(steps) + if ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and frontier.wait_key == wait_key + ): + steps[frontier.ordinal] = dataclasses.replace( + frontier, + status=StepStatus.READY, + due_at=now, + args={**frontier.args, "__payload__": payload}, + updated_at=now, + ) + self._append_events( + run_id, + ( + ( + HistoryEventType.WAIT_RESOLVED, + {"ordinal": frontier.ordinal, "wait_key": wait_key}, + ), + ), + now, + ) + return "resolved" + self._pending.setdefault(run_id, {})[wait_key] = payload + self._append_events( + run_id, + ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), + now, + ) + return "buffered" + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. @@ -927,8 +1057,9 @@ async def next_due(self, now: float) -> float | None: if not _run_is_runnable(run, now): continue frontier = _frontier(self._steps[run.run_id]) - if frontier is not None and frontier.status in CLAIMABLE_STEP_STATUSES: - due_times.append(frontier.due_at) + wake_at = None if frontier is None else step_wake_at(frontier) + if wake_at is not None: + due_times.append(wake_at) return min(due_times) if due_times else None @@ -961,6 +1092,7 @@ async def next_due(self, now: float) -> float | None: due_at REAL NOT NULL DEFAULT 0, epoch INTEGER NOT NULL DEFAULT 0, lease_expires_at REAL NOT NULL DEFAULT 0, + wait_key TEXT, error TEXT, origin TEXT NOT NULL, created_at REAL NOT NULL, @@ -981,7 +1113,19 @@ async def next_due(self, now: float) -> float | None: run_id TEXT NOT NULL, PRIMARY KEY (workflow_id, request_key) ); +CREATE TABLE IF NOT EXISTS workflow_inbox ( + run_id TEXT NOT NULL, + wait_key TEXT NOT NULL, + dedupe_key TEXT NOT NULL, + seq INTEGER NOT NULL, + payload TEXT NOT NULL, + status TEXT NOT NULL, + created_at REAL NOT NULL, + PRIMARY KEY (run_id, wait_key, dedupe_key) +); CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs (status); +CREATE INDEX IF NOT EXISTS idx_workflow_inbox_pending + ON workflow_inbox (run_id, wait_key, status, seq); """ _STEP_MIGRATIONS: Final = ( @@ -989,6 +1133,7 @@ async def next_due(self, now: float) -> float | None: "lease_expires_at", "ALTER TABLE workflow_steps ADD COLUMN lease_expires_at REAL NOT NULL DEFAULT 0", ), + ("wait_key", "ALTER TABLE workflow_steps ADD COLUMN wait_key TEXT"), ) @@ -1064,6 +1209,7 @@ def _step_from_row(row: sqlite3.Row) -> StepRecord: due_at=row["due_at"], epoch=row["epoch"], lease_expires_at=row["lease_expires_at"], + wait_key=row["wait_key"], error=_load(row["error"]), origin=row["origin"], created_at=row["created_at"], @@ -1154,9 +1300,9 @@ def _insert_step(self, step: StepRecord) -> None: """ self._db.execute( "INSERT INTO workflow_steps (run_id, ordinal, handler_id, status, args," - " attempts, recoveries, due_at, epoch, lease_expires_at, error, origin," - " created_at, updated_at)" - " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + " attempts, recoveries, due_at, epoch, lease_expires_at, wait_key," + " error, origin, created_at, updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( step.run_id, step.ordinal, @@ -1168,6 +1314,7 @@ def _insert_step(self, step: StepRecord) -> None: step.due_at, step.epoch, step.lease_expires_at, + step.wait_key, _dump(step.error), step.origin, step.created_at, @@ -1284,11 +1431,7 @@ async def claim_next( for row in rows: run = _run_from_row(row) frontier = _frontier(self._load_steps(run.run_id)) - if ( - frontier is None - or frontier.status not in CLAIMABLE_STEP_STATUSES - or frontier.due_at > now - ): + if frontier is None or not step_claimable_at(frontier, now): continue claimed = dataclasses.replace( frontier, @@ -1389,6 +1532,39 @@ async def renew_lease( raise return True + def _arm_sql(self, step: StepRecord, now: float) -> StepRecord: + """Resolve a newly armed wait against a buffered delivery, in-transaction. + + Args: + step: The slot being inserted. + now: Current time in epoch seconds. + + Returns: + The slot, already resolved when a matching delivery was waiting. + """ + if step.status is not StepStatus.BLOCKED or step.wait_key is None: + return step + row = self._db.execute( + "SELECT dedupe_key, payload FROM workflow_inbox" + " WHERE run_id = ? AND wait_key = ? AND status = ?" + " ORDER BY seq LIMIT 1", + (step.run_id, step.wait_key, "PENDING"), + ).fetchone() + if row is None: + return step + self._db.execute( + "UPDATE workflow_inbox SET status = ?" + " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", + ("CONSUMED", step.run_id, step.wait_key, row["dedupe_key"]), + ) + return dataclasses.replace( + step, + status=StepStatus.READY, + due_at=now, + args={**step.args, "__payload__": json.loads(row["payload"])}, + updated_at=now, + ) + async def commit( self, claim: Claim, completion: StepCompletion, now: float ) -> None: @@ -1433,7 +1609,7 @@ async def commit( ), ) for step in completion.new_steps: - self._insert_step(step) + self._insert_step(self._arm_sql(step, now)) self._db.execute( "UPDATE workflow_runs SET status = ?," " state = CASE WHEN ? THEN ? ELSE state END," @@ -1516,6 +1692,103 @@ async def append_events( self._db.execute("ROLLBACK") raise + async def deliver( + self, + run_id: str, + wait_key: str, + dedupe_key: str, + payload: dict[str, Any], + now: float, + ) -> DeliveryDisposition: + """Deliver a payload to a run, resolving its wait or buffering it. + + Args: + run_id: The receiving run. + wait_key: The address the waiting slot declared. + dedupe_key: Sender-supplied identity, making redelivery a no-op. + payload: JSON-compatible payload to hand the resuming handler. + now: Current time in epoch seconds. + + Returns: + What the store did with the delivery. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + if row is None: + self._db.execute("ROLLBACK") + return "unknown_run" + if row["status"] in terminal: + self._db.execute("ROLLBACK") + return "run_terminal" + seen = self._db.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", + (run_id, wait_key, dedupe_key), + ).fetchone() + if seen is not None: + self._db.execute("ROLLBACK") + return "duplicate" + frontier = _frontier(self._load_steps(run_id)) + resolves = ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and frontier.wait_key == wait_key + ) + self._db.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at)" + " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" + " workflow_inbox WHERE run_id = ?), ?, ?, ?)", + ( + run_id, + wait_key, + dedupe_key, + run_id, + json.dumps(payload), + "CONSUMED" if resolves else "PENDING", + now, + ), + ) + if resolves and frontier is not None: + self._db.execute( + "UPDATE workflow_steps SET status = ?, due_at = ?, args = ?," + " updated_at = ? WHERE run_id = ? AND ordinal = ?", + ( + StepStatus.READY.value, + now, + json.dumps({**frontier.args, "__payload__": payload}), + now, + run_id, + frontier.ordinal, + ), + ) + self._append_events( + run_id, + ( + ( + HistoryEventType.WAIT_RESOLVED, + {"ordinal": frontier.ordinal, "wait_key": wait_key}, + ), + ), + now, + ) + else: + self._append_events( + run_id, + ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return "resolved" if resolves else "buffered" + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. @@ -1877,6 +2150,7 @@ async def next_due(self, now: float) -> float | None: due_times = [] for row in rows: frontier = _frontier(self._load_steps(row["run_id"])) - if frontier is not None and frontier.status in CLAIMABLE_STEP_STATUSES: - due_times.append(frontier.due_at) + wake_at = None if frontier is None else step_wake_at(frontier) + if wake_at is not None: + due_times.append(wake_at) return min(due_times) if due_times else None diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index a53a90dfcdf..7a1a4064cc9 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -196,6 +196,23 @@ async def resume(self, run_id: str) -> bool: await self.kernel.run_until_idle() return resumed + async def signal( + self, run_id: str, delivery: Any, *, key: str | None = None + ) -> Any: + """Deliver a signal and process the work it unblocks. + + Args: + run_id: The receiving run. + delivery: The addressed payload, e.g. ``MyFlow.approved(value)``. + key: Sender idempotency key. + + Returns: + What the store did with the delivery. + """ + disposition = await self.kernel.signal(run_id, delivery, key=key) + await self.kernel.run_until_idle() + return disposition + async def cancel(self, run_id: str) -> bool: """Request cancellation of a run and process the drain. diff --git a/tests/units/workflow/test_waits.py b/tests/units/workflow/test_waits.py new file mode 100644 index 00000000000..2ffc7bb45af --- /dev/null +++ b/tests/units/workflow/test_waits.py @@ -0,0 +1,276 @@ +"""Tests for waits, signals, and human approvals.""" + +import pytest +from pydantic import BaseModel +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import Signal, WorkflowConfig, manual, never, wait_for + +import reflex as rx +from reflex.workflow.records import RunStatus, StepStatus +from reflex.workflow.store import MemoryRunStore, SqliteRunStore +from reflex.workflow.testing import WorkflowTestHarness + + +class Decision(BaseModel): + """A human decision delivered to a waiting run.""" + + approved: bool + by: str + + +def _review_flow(): + """Build a workflow that waits for a decision with a deadline. + + Returns: + The workflow class. + """ + + class ReviewFlow(rx.State): + __workflow__ = WorkflowConfig(id="waits.review") + outcome: str = "" + decided_by: str = "" + + review = Signal(Decision) + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + return wait_for( + ReviewFlow.review, + then=ReviewFlow.decide, + timeout="3d", + on_timeout=ReviewFlow.expire, + ) + + @rx.event(durable=True, effect="none") + def decide(self, decision: Decision): + self.decided_by = decision.by + self.outcome = "approved" if decision.approved else "rejected" + return rx.complete(result={"outcome": self.outcome}) + + @rx.event(durable=True, effect="none") + def expire(self): + self.outcome = "expired" + return rx.fail("no_decision") + + return ReviewFlow + + +def test_channel_names_itself_and_validates_payloads(): + class Holder: + review = Signal(Decision) + + assert Holder.review.name == "review" + delivery = Holder.review({"approved": True, "by": "ada"}) + assert delivery.channel == "review" + assert isinstance(delivery.payload, Decision) + with pytest.raises(WorkflowDefinitionError, match="expects Decision"): + Holder.review(42) + + +def test_wait_for_requires_a_timeout_branch(): + class Holder: + review = Signal(Decision) + + with pytest.raises(WorkflowDefinitionError, match="requires on_timeout"): + wait_for(Holder.review, then="decide", timeout="3d") + with pytest.raises(WorkflowDefinitionError, match="never times out"): + wait_for(Holder.review, then="decide", timeout=never, on_timeout="expire") + with pytest.raises(WorkflowDefinitionError, match=r"rx\.Signal"): + wait_for("review", then="decide", timeout=never) # pyright: ignore[reportArgumentType] + + +async def test_wait_arms_a_blocked_slot(forked_registration_context): + flow = _review_flow() + async with WorkflowTestHarness(flow) as harness: + result = await harness.start(flow.start) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert snapshot.steps[1].status is StepStatus.BLOCKED + assert snapshot.steps[1].wait_key == "sig:review" + assert snapshot.steps[1].origin == "wait" + + +async def test_signal_resolves_the_wait(forked_registration_context): + flow = _review_flow() + async with WorkflowTestHarness(flow) as harness: + result = await harness.start(flow.start) + assert result.run_id is not None + await harness.advance("1d") + disposition = await harness.signal( + result.run_id, flow.review(Decision(approved=True, by="ada")) + ) + assert disposition == "resolved" + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"outcome": "approved", "decided_by": "ada"} + assert snapshot.result == {"outcome": "approved"} + + +async def test_deadline_wins_and_late_signals_are_refused(forked_registration_context): + flow = _review_flow() + async with WorkflowTestHarness(flow) as harness: + result = await harness.start(flow.start) + assert result.run_id is not None + await harness.advance("2d") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + + await harness.advance("2d") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.state["outcome"] == "expired" + + # The loser of the race can no longer resolve the wait. + assert ( + await harness.signal( + result.run_id, flow.review(Decision(approved=True, by="late")) + ) + == "run_terminal" + ) + + +async def test_signal_arriving_before_the_wait_is_not_lost( + forked_registration_context, +): + """A sender faster than the run must not block it forever.""" + flow = _review_flow() + async with WorkflowTestHarness(flow) as harness: + result = await harness.kernel.start(flow.start) + assert result.run_id is not None + disposition = await harness.kernel.signal( + result.run_id, flow.review(Decision(approved=False, by="fast")) + ) + assert disposition == "buffered" + + await harness.run_until_idle() + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"outcome": "rejected", "decided_by": "fast"} + + +async def test_duplicate_signals_are_ignored(forked_registration_context): + flow = _review_flow() + async with WorkflowTestHarness(flow) as harness: + result = await harness.kernel.start(flow.start) + assert result.run_id is not None + await harness.kernel.run_until_idle() + first = await harness.kernel.signal( + result.run_id, flow.review(Decision(approved=True, by="ada")), key="req-1" + ) + second = await harness.kernel.signal( + result.run_id, + flow.review(Decision(approved=False, by="mallory")), + key="req-1", + ) + assert first == "resolved" + assert second == "duplicate" + await harness.kernel.run_until_idle() + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.state["decided_by"] == "ada" + + +async def test_signal_to_unknown_run(forked_registration_context): + flow = _review_flow() + async with WorkflowTestHarness(flow) as harness: + assert ( + await harness.kernel.signal( + "no-such-run", flow.review(Decision(approved=True, by="ada")) + ) + == "unknown_run" + ) + + +async def test_wait_with_no_deadline_never_times_out(forked_registration_context): + class Patient(rx.State): + __workflow__ = WorkflowConfig(id="waits.patient") + got: str = "" + + ping = Signal() + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + return wait_for(Patient.ping, then=Patient.woken, timeout=never) + + @rx.event(durable=True, effect="none") + def woken(self, payload: str): + self.got = payload + + async with WorkflowTestHarness(Patient) as harness: + result = await harness.start(Patient.start) + assert result.run_id is not None + await harness.advance("30d") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert snapshot.steps[1].due_at == pytest.approx(0.0) + + await harness.signal(result.run_id, Patient.ping("hello")) + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"got": "hello"} + + +async def test_wait_survives_a_restart(forked_registration_context, tmp_path): + """A blocked run is resolved by a signal delivered after a restart.""" + flow = _review_flow() + db_path = tmp_path / "workflow.db" + first = SqliteRunStore(db_path) + async with WorkflowTestHarness(flow, store=first) as harness: + result = await harness.start(flow.start) + assert result.run_id is not None + resume_at = harness.now + first.close() + + second = SqliteRunStore(db_path) + async with WorkflowTestHarness( + flow, store=second, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + + assert ( + await harness.signal( + result.run_id, flow.review(Decision(approved=True, by="ada")) + ) + == "resolved" + ) + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + second.close() + + +async def test_blocked_run_does_not_spin_the_scheduler(forked_registration_context): + """A deadline-less wait must not make the store claimable forever.""" + + class Idle(rx.State): + __workflow__ = WorkflowConfig(id="waits.idle") + + ping = Signal() + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + return wait_for(Idle.ping, then=Idle.woken, timeout=never) + + @rx.event(durable=True, effect="none") + def woken(self, payload: str): + pass + + store = MemoryRunStore() + async with WorkflowTestHarness(Idle, store=store) as harness: + result = await harness.start(Idle.start) + assert result.run_id is not None + # Nothing is claimable and no wake-up time is scheduled, so a worker + # sleeps rather than looping on the database. + assert await store.claim_next(harness.now) is None + assert await store.next_due(harness.now) is None From 6fbb73177c60aabcd50b3f1b56463cdb2bd7d038 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:29:10 -0700 Subject: [PATCH 010/121] Add an end-to-end test covering the whole workflow surface One dunning workflow written the way a user would: a manual root, a flaky charge with retries and a failure hook, a durable delay, a human decision with a deadline, and completion, failure, and suspension outcomes -- plus a second class reached by a verified webhook and a cron schedule. It caught a real authoring subtlety worth keeping in front of us: run state cannot count attempts, because a failed attempt's patch is discarded by design. The example now models gateway flakiness outside the run, which is where it actually lives. --- tests/units/workflow/test_end_to_end.py | 290 ++++++++++++++++++++++++ 1 file changed, 290 insertions(+) create mode 100644 tests/units/workflow/test_end_to_end.py diff --git a/tests/units/workflow/test_end_to_end.py b/tests/units/workflow/test_end_to_end.py new file mode 100644 index 00000000000..5d5196dad12 --- /dev/null +++ b/tests/units/workflow/test_end_to_end.py @@ -0,0 +1,290 @@ +"""One workflow using every shipped primitive, as a user would write it.""" + +from pydantic import BaseModel +from reflex_base.workflow import ( + Retry, + Signal, + TransientWorkflowError, + WorkflowConfig, + after, + complete, + fail, + hmac_signature, + manual, + needs_attention, + schedule, + wait_for, + webhook, +) + +import reflex as rx +from reflex.workflow.definition import compile_workflow +from reflex.workflow.ingress import collect_webhook_routes +from reflex.workflow.records import RunStatus +from reflex.workflow.testing import WorkflowTestHarness + + +class Invoice(BaseModel): + """A payment provider's webhook payload.""" + + id: str + amount: int + + +class Decision(BaseModel): + """A human decision on a disputed charge.""" + + approved: bool + by: str + + +ATTEMPTS: list[int] = [] + + +class Dunning(rx.State): + """Charge an invoice, escalate to a human, then settle or write it off.""" + + __workflow__ = WorkflowConfig( + id="billing.dunning", run_timeout="30d", max_steps=100 + ) + + invoice_id: str = "" + amount: int = 0 + outcome: str = "" + decided_by: str = "" + + review = Signal(Decision) + + @rx.event(id="start", durable=True, trigger=manual(), effect="none") + def start(self, invoice_id: str, amount: int): + """Record the invoice and begin charging it. + + Returns: + The charge step. + """ + self.invoice_id = invoice_id + self.amount = amount + return Dunning.charge + + @rx.event( + durable=True, + effect="idempotent_write", + retry=Retry(max_attempts=3, initial_delay="2s", jitter="none"), + timeout="30s", + on_failure="escalate", + ) + def charge(self): + """Charge the invoice, retrying a flaky gateway. + + Returns: + The delayed receipt step. + """ + # A failed attempt's state patch is discarded, so run state cannot + # count attempts; the flakiness here lives outside the run, as a real + # payment gateway's would. + ATTEMPTS.append(1) + if len(ATTEMPTS) < 3: + msg = "gateway unavailable" + raise TransientWorkflowError(msg) + self.outcome = "charged" + return after("2d", Dunning.receipt) + + @rx.event(durable=True, effect="none") + def escalate(self): + """Hand a failed charge to a human, with a deadline. + + Returns: + The wait for a decision. + """ + return wait_for( + Dunning.review, + then=Dunning.settle, + timeout="7d", + on_timeout=Dunning.write_off, + ) + + @rx.event(durable=True, effect="none") + def settle(self, decision: Decision): + """Apply the human decision. + + Returns: + Completion, or failure when the charge is disputed. + """ + self.decided_by = decision.by + if not decision.approved: + return fail("disputed", details={"by": decision.by}) + self.outcome = "settled" + return complete(result={"invoice": self.invoice_id, "outcome": "settled"}) + + @rx.event(durable=True, effect="none") + def write_off(self): + """Give up on a charge nobody decided. + + Returns: + A suspension for an operator. + """ + self.outcome = "written_off" + return needs_attention("no_decision_in_7d") + + @rx.event(durable=True, effect="none") + def receipt(self): + """Send the receipt two days after a successful charge. + + Returns: + Completion. + """ + self.outcome = "receipted" + return complete(result={"invoice": self.invoice_id, "outcome": "receipted"}) + + +class Ingested(rx.State): + """The same product surface reached by a provider and by a schedule.""" + + __workflow__ = WorkflowConfig(id="billing.ingested") + + invoice_id: str = "" + swept: bool = False + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "stripe.invoice_failed", + model=Invoice, + verify=hmac_signature(secret_env="STRIPE_SECRET", header="X-Signature"), + dedupe_by="id", + ), + ) + def on_failed(self, invoice: Invoice): + """Start from the provider's failed-invoice webhook.""" + self.invoice_id = invoice.id + + @rx.event(durable=True, effect="read", trigger=schedule("0 3 * * *")) + def nightly_sweep(self): + """Reconcile invoices every night.""" + self.swept = True + + +def test_the_whole_surface_compiles(forked_registration_context): + dunning = compile_workflow(Dunning) + assert dunning.roots == ("start",) + assert set(dunning.handlers) == { + "start", + "charge", + "escalate", + "settle", + "write_off", + "receipt", + } + ingested = compile_workflow(Ingested) + assert set(ingested.roots) == {"on_failed", "nightly_sweep"} + assert set(collect_webhook_routes((ingested,))) == {"stripe.invoice_failed"} + + +async def test_retry_then_delay_then_complete(forked_registration_context): + """The happy path: a flaky charge succeeds, then a delayed receipt.""" + ATTEMPTS.clear() + async with WorkflowTestHarness(Dunning) as harness: + result = await harness.start(Dunning.start("inv_1", 4200)) + assert result.run_id is not None + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.RETRYING + + await harness.advance("2s") + await harness.advance("4s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert snapshot.state["outcome"] == "charged" + assert len(ATTEMPTS) == 3 + + await harness.advance("2d") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"invoice": "inv_1", "outcome": "receipted"} + + +async def test_exhausted_retries_escalate_to_a_human(forked_registration_context): + """Failure hands off to a person, who approves and settles the run.""" + + class AlwaysFails(rx.State): + __workflow__ = WorkflowConfig(id="billing.always_fails") + outcome: str = "" + decided_by: str = "" + + review = Signal(Decision) + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=2, initial_delay="1s", jitter="none"), + on_failure="escalate", + ) + def charge(self): + msg = "card declined" + raise TransientWorkflowError(msg) + + @rx.event(durable=True, effect="none") + def escalate(self): + return wait_for( + AlwaysFails.review, + then=AlwaysFails.settle, + timeout="7d", + on_timeout=AlwaysFails.write_off, + ) + + @rx.event(durable=True, effect="none") + def settle(self, decision: Decision): + self.decided_by = decision.by + self.outcome = "settled" + return complete(result={"outcome": "settled"}) + + @rx.event(durable=True, effect="none") + def write_off(self): + self.outcome = "written_off" + return fail("no_decision") + + async with WorkflowTestHarness(AlwaysFails) as harness: + result = await harness.start(AlwaysFails.charge) + assert result.run_id is not None + await harness.advance("1s") + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + + await harness.advance("3d") + assert ( + await harness.signal( + result.run_id, AlwaysFails.review(Decision(approved=True, by="ada")) + ) + == "resolved" + ) + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state == {"outcome": "settled", "decided_by": "ada"} + + +async def test_nobody_decides_and_the_run_waits_for_an_operator( + forked_registration_context, +): + """A silent week suspends the run rather than guessing.""" + ATTEMPTS.clear() + async with WorkflowTestHarness(Dunning) as harness: + result = await harness.kernel.start(Dunning.start("inv_2", 100)) + assert result.run_id is not None + # Drive the charge to final failure by exhausting its attempts. + for _ in range(4): + await harness.advance("10s") + await harness.advance("8d") + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + if snapshot.status is RunStatus.NEEDS_ATTENTION: + assert snapshot.state["outcome"] == "written_off" + assert await harness.resume(result.run_id) From b5c0fe00070bdb56969ec1a5481708ac4b2f3f36 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:31:20 -0700 Subject: [PATCH 011/121] Document workflows Nothing about workflows was documented, which also matters because this page is the surface a text-to-workflow generator will learn from. Covers the whole shipped surface: durable steps and effect classes, retries and timeouts, the transition table, waits and typed signals with the approval example, manual/webhook/schedule triggers, inspecting and steering runs, the virtual-clock harness, and what a redeploy does to runs in flight. Every example was run against the real engine rather than written from memory, which caught one error in the draft: rx.Base no longer exists on main, so payload models use pydantic BaseModel. --- docs/workflows/overview.md | 268 ++++++++++++++++++++++++++++++++++ news/workflow-docs.feature.md | 1 + 2 files changed, 269 insertions(+) create mode 100644 docs/workflows/overview.md create mode 100644 news/workflow-docs.feature.md diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md new file mode 100644 index 00000000000..55142837e54 --- /dev/null +++ b/docs/workflows/overview.md @@ -0,0 +1,268 @@ +# Workflows + +A workflow is durable automation that survives restarts, deploys, and crashes. Where an ordinary +event handler runs once inside a browser session and is gone, a workflow run has its own identity, +its own persisted state, and a history you can inspect long after the request that started it +finished. + +Workflows are ordinary Reflex code. There is no separate service to operate, no DSL, and no +determinism rules to learn: a workflow is an `rx.State` class, a step is an `@rx.event` handler, and +control flow is what the handler returns. + +```python +import reflex as rx + + +class Onboarding(rx.State): + __workflow__ = rx.WorkflowConfig(id="growth.onboarding", run_timeout="30d") + + user_id: str = "" + email: str = "" + nudges_sent: int = 0 + + @rx.event(id="signup", durable=True, trigger=rx.manual(), effect="none") + def signup(self, user_id: str, email: str): + self.user_id, self.email = user_id, email + return Onboarding.send_welcome + + @rx.event(durable=True, effect="idempotent_write", timeout="30s") + async def send_welcome(self): + await send_email(self.email, "welcome") + return rx.after("3d", Onboarding.check_activation) + + @rx.event(durable=True, effect="read") + async def check_activation(self): + if await has_activated(self.user_id): + return rx.complete(result={"outcome": "activated"}) + if self.nudges_sent >= 2: + return rx.complete(result={"outcome": "gave_up"}) + return Onboarding.send_nudge + + @rx.event(durable=True, effect="idempotent_write") + async def send_nudge(self): + self.nudges_sent += 1 + await send_email(self.email, f"nudge_{self.nudges_sent}") + return rx.after("4d", Onboarding.check_activation) + + +app = rx.App() +app.add_workflow(Onboarding) +``` + +Register the class with `app.add_workflow(...)` and start a run from anywhere on the server: + +```python +result = await rx.workflows.start( + Onboarding.signup("u_1", "ada@example.com"), + request_key="u_1", +) +``` + +`request_key` makes starting idempotent: submitting the same key again returns the original run with +disposition `"deduplicated"` instead of starting a second one. + +## The three-day wait is not a sleeping process + +`rx.after("3d", ...)` commits a mailbox slot with a due time and nothing else exists in between: no +open coroutine, no held connection, no memory. Deploy, restart, or hard-crash the server on day two +and the step still runs on day three. That is the difference between a workflow and a background +task. + +## Steps + +Every public handler on a workflow class is a durable step and must declare `durable=True` and an +`effect`. A step runs, commits its state changes, and schedules what comes next — all atomically. +A step commits exactly once, so it can never half-succeed. + +Because state is snapshotted per step rather than reconstructed by re-running your code, there are no +determinism rules: `datetime.now()`, `random`, and ordinary I/O are all fine inside a step. + +### Effect classes + +`effect` declares what a step does to the outside world, which decides whether it is safe to retry. + +| Effect | Meaning | Retries | +| --- | --- | --- | +| `"none"` | pure orchestration, no external I/O | yes | +| `"read"` | reads something external | yes | +| `"idempotent_write"` | a write that is safe to repeat | yes | +| `"non_idempotent_write"` | a write that is not safe to repeat | never | + +A `non_idempotent_write` that fails gets exactly one attempt and suspends the run as +`NEEDS_ATTENTION`: the runtime cannot prove the write did not already land, so it asks a human +rather than guessing and charging a customer twice. + +### Retries and timeouts + +```python +@rx.event( + durable=True, + effect="idempotent_write", + retry=rx.Retry(max_attempts=5, initial_delay="2s", multiplier=2), + timeout="30s", + on_failure="alert_billing", +) +async def charge(self): ... +``` + +Failures retry with exponential backoff by default. Narrow that with +`rx.Retry(do_not_retry_on=(ValueError,))` when a specific error should fail fast. `timeout` bounds a +single attempt. `on_failure` and `on_timeout` name a handler on the same class that runs once the +step is finally out of attempts. + +Backoff is a persisted timer, not a sleep, so a retry scheduled for tomorrow survives a deploy +tonight. + +## Transitions + +A durable handler returns what happens next. It never calls another handler directly — doing so +would run it inline and lose its retries and effect tracking, which the compiler rejects. + +| Return | Meaning | +| --- | --- | +| `None` | this branch is done | +| `MyFlow.next_step` | run that step next | +| `MyFlow.next_step(arg)` | run it with an argument | +| `[MyFlow.a, MyFlow.b]` | run both, in order | +| `rx.after("2d", MyFlow.later)` | run it after a durable delay | +| `rx.wait_for(...)` | block until a signal or a deadline | +| `rx.complete(result=...)` | finish the run successfully | +| `rx.fail("reason")` | finish the run as failed | +| `rx.needs_attention("reason")` | suspend for a human | + +## Waiting for the outside world + +Declare a typed channel and wait on it. Whichever of the signal and the deadline arrives first wins; +the loser can no longer resolve the wait. + +```python +from pydantic import BaseModel + + +class Decision(BaseModel): + approved: bool + by: str + + +class Expense(rx.State): + __workflow__ = rx.WorkflowConfig(id="finance.expense") + + review = rx.Signal(Decision) + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def submit(self): + return rx.wait_for( + Expense.review, + then=Expense.decide, + timeout="3d", + on_timeout=Expense.escalate, + ) + + @rx.event(durable=True, effect="none") + def decide(self, decision: Decision): + return rx.complete(result={"approved": decision.approved}) + + @rx.event(durable=True, effect="none") + def escalate(self): ... +``` + +Deliver the signal from an ordinary page handler — the approve button in your own app: + +```python +class ReviewPage(rx.State): + @rx.event + async def approve(self, run_id: str): + await rx.workflows.signal( + run_id, + Expense.review(Decision(approved=True, by="ada")), + key=run_id, + ) +``` + +Use `timeout=rx.never` to wait indefinitely. A signal that arrives before the run reaches its wait is +buffered and applied as soon as the wait arms, so a fast approver never blocks the run. + +## Triggers + +A root handler declares how runs of it begin. + +```python +@rx.event(durable=True, trigger=rx.manual(), effect="none") +def start(self): ... + +@rx.event( + durable=True, + effect="none", + trigger=rx.webhook( + "stripe.invoice_failed", + model=Invoice, + verify=rx.hmac_signature(secret_env="STRIPE_SECRET", header="Stripe-Signature"), + dedupe_by="id", + ), +) +def on_failed(self, invoice: Invoice): ... + +@rx.event(durable=True, effect="read", trigger=rx.schedule("0 3 * * *")) +def nightly_sweep(self): ... +``` + +Webhook roots are served at `POST /_workflow/webhook/{topic}`. The endpoint verifies the provider's +signature over the raw body, validates the payload, and durably accepts the run before +acknowledging, so a provider that never sees a response can safely redeliver — `dedupe_by` sends the +redelivery to the same run. A webhook trigger without a verifier is a compile error; if an endpoint +really is public, say so with `allow_unverified=True` and a reason. + +Schedules are evaluated in UTC and fire once per occurrence even across restarts. Deploying a +schedule does not backfill history, and catch-up after an outage is bounded. + +## Inspecting and steering runs + +```python +snapshot = await rx.workflows.get_run(run_id) +runs = await rx.workflows.list_runs(workflow_id="finance.expense", limit=20) +await rx.workflows.cancel(run_id) +await rx.workflows.resume(run_id) +``` + +`labels` passed at start are server-derived indexing data you can filter on later: + +```python +await rx.workflows.start(Expense.submit(), labels={"customer": customer.id}) +await rx.workflows.list_runs(labels={"customer": customer.id}) +``` + +A suspended run is waiting for you, not finished: fix whatever made the outcome uncertain, then +`resume()` to give the step a fresh attempt budget. + +## Testing + +The test harness runs your real workflow on a virtual clock, so a three-day wait takes microseconds +and nothing is mocked. + +```python +from reflex.workflow import WorkflowTestHarness + + +async def test_drip_nudges_then_gives_up(): + async with WorkflowTestHarness(Onboarding) as harness: + result = await harness.start(Onboarding.signup("u_1", "ada@example.com")) + + await harness.advance("3d") + await harness.advance("4d") + + snapshot = await harness.get_run(result.run_id) + assert snapshot.state["nudges_sent"] == 2 +``` + +`harness.advance(...)` moves the clock and runs whatever became due. `harness.signal(...)` delivers +to a waiting run, and `harness.cancel(...)` and `harness.resume(...)` drive the operator paths. + +## Deploying + +Runs persist to a SQLite file next to your app by default; pass `rx.App(workflow_store=...)` to +choose another store. Run one worker process per SQLite database file. + +Deploying new code does not disturb runs already in flight. Adding state fields, retuning retries and +timeouts, and changing hooks all apply to future steps. Only a change that makes a pending step +undispatchable — deleting the handler it names, or removing parameters its payload carries — +suspends that run, with a message naming the handler. diff --git a/news/workflow-docs.feature.md b/news/workflow-docs.feature.md new file mode 100644 index 00000000000..c7a80209feb --- /dev/null +++ b/news/workflow-docs.feature.md @@ -0,0 +1 @@ +Adds the Workflows documentation page covering steps, effect classes, retries, transitions, waits and signals, triggers, operator control, testing, and deployment. From 9914990eb950b162bc23998566ee41ac420199c8 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:40:00 -0700 Subject: [PATCH 012/121] Add parallel fan-out with child runs and joins Runs could only do one thing at a time: the mailbox is strictly serial, so there was no way to enrich and score a lead concurrently, or to do anything Temporal and Inngest express with child workflows. Concurrency now lives in the run graph rather than the mailbox. rx.parallel() commits a BLOCKED join slot in the parent and admits one child run per branch, each with its own state, retries, timers, and history. A child that finishes reports its outcome to the parent's join slot through a compare-and-swap on an arrival counter, so a redelivered result cannot be counted twice, and the slot becomes claimable exactly when the last expected branch lands. The join handler receives one entry per branch carrying run_id, status, result, and error. Keeping each run's mailbox serial is what makes this safe: the parent never has two open slots, so the per-run fence and the one-claim-per-run invariant are untouched, and a failing branch fails its own run rather than the parent's. Children are admitted after the parent's commit lands, so a crash in between leaves a join with no children -- which recovery re-runs -- rather than orphans with no parent. --- docs/workflows/overview.md | 30 ++ news/workflow-parallel.feature.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 55 +++ pyi_hashes.json | 2 +- reflex/__init__.py | 1 + reflex/workflow/__init__.py | 4 + reflex/workflow/kernel.py | 120 +++++- reflex/workflow/records.py | 15 +- reflex/workflow/store.py | 356 ++++++++++++++++-- tests/units/workflow/test_parallel.py | 209 ++++++++++ 10 files changed, 754 insertions(+), 39 deletions(-) create mode 100644 news/workflow-parallel.feature.md create mode 100644 tests/units/workflow/test_parallel.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index 55142837e54..b0a1790caf8 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -126,6 +126,7 @@ would run it inline and lose its retries and effect tracking, which the compiler | `[MyFlow.a, MyFlow.b]` | run both, in order | | `rx.after("2d", MyFlow.later)` | run it after a durable delay | | `rx.wait_for(...)` | block until a signal or a deadline | +| `rx.parallel(a, b, then=...)` | run branches concurrently, then join | | `rx.complete(result=...)` | finish the run successfully | | `rx.fail("reason")` | finish the run as failed | | `rx.needs_attention("reason")` | suspend for a human | @@ -182,6 +183,33 @@ class ReviewPage(rx.State): Use `timeout=rx.never` to wait indefinitely. A signal that arrives before the run reaches its wait is buffered and applied as soon as the wait arms, so a fast approver never blocks the run. +## Running work in parallel + +Each branch of a fan-out becomes its own run, with its own state, retries, and history, so a slow or +failing branch never blocks its siblings. The parent blocks until every branch reports. + +```python +class Router(rx.State): + __workflow__ = rx.WorkflowConfig(id="sales.router") + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def begin(self, lead_id: str): + return rx.parallel( + Enrich.start(lead_id), + Score.start(lead_id), + then=Router.route, + ) + + @rx.event(durable=True, effect="none") + def route(self, results: list): + # One entry per branch: run_id, status, result, error. + return rx.complete(result={"branches": len(results)}) +``` + +A branch that fails still reports, so the join handler decides what a partial success means rather +than the engine guessing. Child runs are ordinary runs: they appear in `list_runs()` and can be +inspected and cancelled on their own. + ## Triggers A root handler declares how runs of it begin. @@ -190,6 +218,7 @@ A root handler declares how runs of it begin. @rx.event(durable=True, trigger=rx.manual(), effect="none") def start(self): ... + @rx.event( durable=True, effect="none", @@ -202,6 +231,7 @@ def start(self): ... ) def on_failed(self, invoice: Invoice): ... + @rx.event(durable=True, effect="read", trigger=rx.schedule("0 3 * * *")) def nightly_sweep(self): ... ``` diff --git a/news/workflow-parallel.feature.md b/news/workflow-parallel.feature.md new file mode 100644 index 00000000000..fd0be6c9328 --- /dev/null +++ b/news/workflow-parallel.feature.md @@ -0,0 +1 @@ +Adds `rx.parallel(...)` for fanning work out across concurrent child runs and joining their results. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 7bf0497c4bc..574d51bb4b1 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -882,3 +882,58 @@ def wait_for( return WaitFor( channel=channel.name, then=then, timeout=timeout, on_timeout=on_timeout ) + + +@dataclasses.dataclass(frozen=True, slots=True) +class Parallel: + """Control return that runs branches concurrently and joins their results. + + Each branch becomes its own run with its own mailbox, retries, and + history, so a slow or failing branch never blocks its siblings. The parent + blocks on a join slot until every branch reports. + + Attributes: + branches: The root events to run concurrently. + then: Handler that receives the list of branch results. + """ + + branches: tuple[Any, ...] + then: Any + + def __post_init__(self): + """Validate the fan-out. + + Raises: + WorkflowDefinitionError: If no branches were given. + """ + if not self.branches: + msg = ( + "parallel() needs at least one branch; pass the root events to " + "run concurrently." + ) + raise WorkflowDefinitionError(msg) + + +def parallel(*branches: Any, then: Any) -> Parallel: + """Run branches concurrently, then continue with all their results. + + Each branch runs as its own child run, so branches retry and fail + independently:: + + return rx.parallel( + Enrich.start(lead.id), + Score.start(lead.id), + then=Sales.route, + ) + + The ``then`` handler receives one argument: the list of branch results, in + the order the branches were given. + + Args: + branches: Root events to run concurrently. + then: Handler to run once every branch has finished. + + Returns: + The control return value. + """ + return Parallel(branches=branches, then=then) diff --git a/pyi_hashes.json b/pyi_hashes.json index 8d8ae67a496..bbb97883d9f 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -118,7 +118,7 @@ "packages/reflex-components-recharts/src/reflex_components_recharts/polar.pyi": "99ebcfc07868061bdc3c2010d85a153f", "packages/reflex-components-recharts/src/reflex_components_recharts/recharts.pyi": "4f6c26f8c76543cc41e2b9dc400ece8a", "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "f170ac685b6ba5892370166c80684db3", - "reflex/__init__.pyi": "604f89bbb7fb278a6467c529da5e57d5", + "reflex/__init__.pyi": "ae53392d29647d23db03f68401293f0f", "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388", "reflex/experimental/memo.pyi": "bc8b48357bef580e70a5881b65d3d3f7" } diff --git a/reflex/__init__.py b/reflex/__init__.py index b0491f13804..636f40ed11f 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -247,6 +247,7 @@ "Signal", "wait_for", "never", + "parallel", "after", "complete", "fail", diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index 062c981d150..82774c3b38a 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -12,6 +12,7 @@ DurableEventConfig, EffectClass, ManualTrigger, + Parallel, Retry, ScheduleTrigger, Signal, @@ -28,6 +29,7 @@ manual, needs_attention, never, + parallel, parse_duration, schedule, wait_for, @@ -71,6 +73,7 @@ "HistoryEventType", "ManualTrigger", "MemoryRunStore", + "Parallel", "Retry", "RunQuery", "RunRecord", @@ -103,6 +106,7 @@ "manual", "needs_attention", "never", + "parallel", "parse_duration", "schedule", "wait_for", diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 7a092c41589..fe43e61c299 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -30,6 +30,7 @@ CompleteRun, FailRun, NeedsAttention, + Parallel, ScheduleTrigger, WaitFor, _Never, @@ -39,6 +40,7 @@ from reflex.event import EventHandler, EventSpec from reflex.workflow.cron import CronSchedule from reflex.workflow.records import ( + TERMINAL_RUN_STATUSES, TERMINAL_STEP_STATUSES, HistoryEventType, RunQuery, @@ -631,7 +633,8 @@ def _resolve_successor( def _interpret_return( self, defn: WorkflowDefinition, value: Any ) -> tuple[ - list[_SuccessorSpec], CompleteRun | FailRun | NeedsAttention | WaitFor | None + list[_SuccessorSpec], + CompleteRun | FailRun | NeedsAttention | WaitFor | Parallel | None, ]: """Interpret a durable handler's return value. @@ -648,7 +651,7 @@ def _interpret_return( """ if value is None: return [], None - if isinstance(value, (CompleteRun, FailRun, NeedsAttention, WaitFor)): + if isinstance(value, (CompleteRun, FailRun, NeedsAttention, WaitFor, Parallel)): return [], value if isinstance(value, (list, tuple)): successors = [] @@ -678,6 +681,9 @@ async def _invoke( """ args = {key: value for key, value in args.items() if key != "__wait__"} delivered = args.pop("__payload__", None) + results = args.pop("__results__", None) + if delivered is None and results is not None: + delivered = results if delivered is not None and handler.params: args[handler.params[0]] = delivered try: @@ -929,7 +935,7 @@ def _success_completion( steps: tuple[StepRecord, ...], state: dict[str, Any], successors: list[_SuccessorSpec], - control: CompleteRun | FailRun | NeedsAttention | WaitFor | None, + control: CompleteRun | FailRun | NeedsAttention | WaitFor | Parallel | None, now: float, ) -> StepCompletion: """Build the commit for a successful attempt. @@ -977,6 +983,33 @@ def _success_completion( run_error=error, events=tuple(events), ) + if isinstance(control, Parallel): + join = StepRecord( + run_id=claim.run.run_id, + ordinal=claim.run.next_ordinal, + handler_id=self._resolve_successor(defn, control.then).handler_id, + status=StepStatus.BLOCKED, + args={"__results__": []}, + wait_key=f"join:{claim.run.next_ordinal}", + join_expected=len(control.branches), + origin="join", + created_at=now, + updated_at=now, + ) + events.append(( + HistoryEventType.CHILD_STARTED, + {"ordinal": join.ordinal, "branches": len(control.branches)}, + )) + return StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.WAITING, + state=state, + new_steps=(join,), + next_ordinal=claim.run.next_ordinal + 1, + events=tuple(events), + children=tuple(control.branches), + join_ordinal=join.ordinal, + ) if isinstance(control, WaitFor): resume = self._resolve_successor(defn, control.then) timeout_id = ( @@ -1389,6 +1422,10 @@ async def _execute_claim(self, claim: Claim) -> None: await self._store.commit(claim, completion, self._clock()) except StaleClaimError: await self._record_abandoned(claim, handler, "fenced_at_commit") + return + if completion.children: + await self._admit_children(claim, completion) + await self._report_to_parent(claim.run, completion) async def _admit_due_schedules(self, now: float) -> int: """Admit a run for every schedule occurrence that has come due. @@ -1435,6 +1472,83 @@ def _next_schedule_due(self, now: float) -> float | None: ] return min(upcoming) if upcoming else None + async def _admit_children(self, claim: Claim, completion: StepCompletion) -> None: + """Create the child runs a fan-out commit declared. + + Children are admitted after the parent's commit lands, so a crash in + between leaves a join slot with no children, which recovery re-runs + rather than a set of orphans with no parent. + + Args: + claim: The parent's claim. + completion: The committed outcome carrying the branches. + """ + now = self._clock() + records = [] + for index, branch in enumerate(completion.children): + defn, handler, payload = self._resolve_target(branch) + child_id = uuid.uuid4().hex + records.append(( + RunRecord( + run_id=child_id, + workflow_id=defn.workflow_id, + definition_digest=defn.digest, + status=RunStatus.PENDING, + state={field.name: field.default for field in defn.fields}, + state_version=0, + next_ordinal=1, + parent_run_id=claim.run.run_id, + parent_ordinal=completion.join_ordinal, + request_key=( + f"child:{claim.run.run_id}:{completion.join_ordinal}:{index}" + ), + deadline=( + now + defn.run_timeout if defn.run_timeout is not None else None + ), + created_at=now, + updated_at=now, + ), + StepRecord( + run_id=child_id, + ordinal=0, + handler_id=handler.id, + status=StepStatus.READY, + args=payload, + origin="root", + created_at=now, + updated_at=now, + ), + )) + await self._store.admit_children(tuple(records), (), now) + self._wakeup.set() + + async def _report_to_parent( + self, run: RunRecord, completion: StepCompletion + ) -> None: + """Report a finished child's outcome to the join slot awaiting it. + + Args: + run: The child run, which may have no parent. + completion: The committed outcome that finished it. + """ + if run.parent_run_id is None or run.parent_ordinal is None: + return + if completion.run_status not in TERMINAL_RUN_STATUSES: + return + await self._store.record_arrival( + run.parent_run_id, + run.parent_ordinal, + { + "run_id": run.run_id, + "status": completion.run_status.value, + "result": completion.result, + "error": completion.run_error, + }, + run.run_id, + self._clock(), + ) + self._wakeup.set() + async def _tick(self) -> bool: """Run one scheduling round. diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 4f2deb047d6..fc5be6bd0a7 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -130,6 +130,8 @@ class HistoryEventType(str, enum.Enum): RUN_CANCELLED = "run_cancelled" RUN_NEEDS_ATTENTION = "run_needs_attention" RUN_RESUMED = "run_resumed" + CHILD_STARTED = "child_started" + CHILD_RESOLVED = "child_resolved" WAIT_ARMED = "wait_armed" WAIT_RESOLVED = "wait_resolved" WAIT_EXPIRED = "wait_expired" @@ -152,6 +154,8 @@ class RunRecord: next_ordinal: Next mailbox ordinal to allocate. result: Run result recorded at completion. error: Terminal or suspension error payload. + parent_run_id: The run that spawned this one, if any. + parent_ordinal: The join slot in the parent this run reports to. request_key: Idempotent admission key, if one was supplied. labels: Server-derived indexing labels. deadline: Absolute run deadline in epoch seconds, if configured. @@ -169,6 +173,8 @@ class RunRecord: next_ordinal: int result: Any = None error: dict[str, Any] | None = None + parent_run_id: str | None = None + parent_ordinal: int | None = None request_key: str | None = None labels: dict[str, str] | None = None deadline: float | None = None @@ -195,7 +201,10 @@ class StepRecord: is not claimed. A claim whose lease has lapsed is treated as orphaned and is reclaimed by recovery, never by a direct claim. wait_key: For a blocked slot, the address a delivery must carry, as - ``"sig:"`` or ``"approval:"``. None otherwise. + ``"sig:"`` or ``"join:"``. None otherwise. + join_expected: Arrivals required before a join slot becomes ready. + join_arrived: Arrivals recorded so far, only ever incremented by a + compare-and-swap so a redelivered result cannot count twice. error: Last recorded attempt error payload. origin: How the slot was allocated. created_at: Allocation time in epoch seconds. @@ -213,8 +222,10 @@ class StepRecord: epoch: int = 0 lease_expires_at: float = 0.0 wait_key: str | None = None + join_expected: int = 0 + join_arrived: int = 0 error: dict[str, Any] | None = None - origin: Literal["root", "chain", "delay", "hook", "wait"] = "chain" + origin: Literal["root", "chain", "delay", "hook", "wait", "join"] = "chain" created_at: float = 0.0 updated_at: float = 0.0 diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 0b29d3152e0..db0a21b8d85 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -46,6 +46,7 @@ DeliveryDisposition = Literal[ "resolved", + "counted", "buffered", "duplicate", "unknown_run", @@ -87,6 +88,8 @@ class StepCompletion: tombstones: Ordinals of unresolved slots to cancel. next_ordinal: Updated mailbox allocation counter, if slots were added. events: History events to append, in order, as (type, data) pairs. + children: Root events to admit as child runs once this commit lands. + join_ordinal: The join slot those children report back to. """ step_status: StepStatus @@ -101,6 +104,8 @@ class StepCompletion: tombstones: tuple[int, ...] = () next_ordinal: int | None = None events: tuple[tuple[HistoryEventType, dict[str, Any]], ...] = () + children: tuple[Any, ...] = () + join_ordinal: int | None = None class RunStore(Protocol): @@ -247,6 +252,51 @@ async def deliver( """ ... + async def admit_children( + self, + runs: tuple[tuple[RunRecord, StepRecord], ...], + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Create child runs, each with its root slot. + + Only ever inserts brand-new rows, so it locks no existing run and + cannot invert lock order against a concurrent commit. + + Args: + runs: The child run records paired with their root slots. + events: History events to append to the parent. + now: Current time in epoch seconds. + """ + ... + + async def record_arrival( + self, + run_id: str, + ordinal: int, + payload: dict[str, Any], + dedupe_key: str, + now: float, + ) -> DeliveryDisposition: + """Count one arrival against a join slot. + + The counter is only ever incremented by this compare-and-swap, so a + redelivered child result cannot be counted twice, and the slot becomes + claimable exactly when the last expected arrival lands. + + Args: + run_id: The waiting parent run. + ordinal: The join slot's ordinal. + payload: The arriving result. + dedupe_key: Identity of the arrival. + now: Current time in epoch seconds. + + Returns: + ``"resolved"`` when this arrival completed the join, ``"counted"`` + when more are still expected, or why it was refused. + """ + ... + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. @@ -797,6 +847,84 @@ async def deliver( ) return "buffered" + async def admit_children( + self, + runs: tuple[tuple[RunRecord, StepRecord], ...], + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Create child runs, each with its root slot. + + Args: + runs: The child run records paired with their root slots. + events: History events to append to the parent. + now: Current time in epoch seconds. + """ + async with self._lock: + for run, root_step in runs: + self._runs[run.run_id] = run + self._steps[run.run_id] = [root_step] + if runs and events: + self._append_events(runs[0][0].parent_run_id or "", events, now) + + async def record_arrival( + self, + run_id: str, + ordinal: int, + payload: dict[str, Any], + dedupe_key: str, + now: float, + ) -> DeliveryDisposition: + """Count one arrival against a join slot. + + Args: + run_id: The waiting parent run. + ordinal: The join slot's ordinal. + payload: The arriving result. + dedupe_key: Identity of the arrival. + now: Current time in epoch seconds. + + Returns: + What the store did with the arrival. + """ + async with self._lock: + run = self._runs.get(run_id) + if run is None: + return "unknown_run" + if run.status in TERMINAL_RUN_STATUSES: + return "run_terminal" + seen = self._inbox.setdefault(run_id, {}) + key = (run_id, f"join:{ordinal}", dedupe_key) + if key in seen: + return "duplicate" + seen[key] = True + steps = self._steps[run_id] + step = steps[ordinal] + if step.status is not StepStatus.BLOCKED: + return "run_terminal" + arrived = step.join_arrived + 1 + results = [*step.args.get("__results__", []), payload] + done = arrived >= step.join_expected + steps[ordinal] = dataclasses.replace( + step, + status=StepStatus.READY if done else StepStatus.BLOCKED, + join_arrived=arrived, + due_at=now if done else step.due_at, + args={**step.args, "__results__": results}, + updated_at=now, + ) + self._append_events( + run_id, + ( + ( + HistoryEventType.CHILD_RESOLVED, + {"ordinal": ordinal, "arrived": arrived}, + ), + ), + now, + ) + return "resolved" if done else "counted" + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. @@ -1074,6 +1202,8 @@ async def next_due(self, now: float) -> float | None: next_ordinal INTEGER NOT NULL, result TEXT, error TEXT, + parent_run_id TEXT, + parent_ordinal INTEGER, request_key TEXT, labels TEXT, deadline REAL, @@ -1093,6 +1223,8 @@ async def next_due(self, now: float) -> float | None: epoch INTEGER NOT NULL DEFAULT 0, lease_expires_at REAL NOT NULL DEFAULT 0, wait_key TEXT, + join_expected INTEGER NOT NULL DEFAULT 0, + join_arrived INTEGER NOT NULL DEFAULT 0, error TEXT, origin TEXT NOT NULL, created_at REAL NOT NULL, @@ -1134,6 +1266,19 @@ async def next_due(self, now: float) -> float | None: "ALTER TABLE workflow_steps ADD COLUMN lease_expires_at REAL NOT NULL DEFAULT 0", ), ("wait_key", "ALTER TABLE workflow_steps ADD COLUMN wait_key TEXT"), + ( + "join_expected", + "ALTER TABLE workflow_steps ADD COLUMN join_expected INTEGER NOT NULL DEFAULT 0", + ), + ( + "join_arrived", + "ALTER TABLE workflow_steps ADD COLUMN join_arrived INTEGER NOT NULL DEFAULT 0", + ), +) + +_RUN_MIGRATIONS: Final = ( + ("parent_run_id", "ALTER TABLE workflow_runs ADD COLUMN parent_run_id TEXT"), + ("parent_ordinal", "ALTER TABLE workflow_runs ADD COLUMN parent_ordinal INTEGER"), ) @@ -1180,6 +1325,8 @@ def _run_from_row(row: sqlite3.Row) -> RunRecord: next_ordinal=row["next_ordinal"], result=_load(row["result"]), error=_load(row["error"]), + parent_run_id=row["parent_run_id"], + parent_ordinal=row["parent_ordinal"], request_key=row["request_key"], labels=_load(row["labels"]), deadline=row["deadline"], @@ -1210,6 +1357,8 @@ def _step_from_row(row: sqlite3.Row) -> StepRecord: epoch=row["epoch"], lease_expires_at=row["lease_expires_at"], wait_key=row["wait_key"], + join_expected=row["join_expected"], + join_arrived=row["join_arrived"], error=_load(row["error"]), origin=row["origin"], created_at=row["created_at"], @@ -1246,13 +1395,17 @@ def _migrate(self) -> None: """ self._db.execute("BEGIN IMMEDIATE") try: - columns = { - row["name"] - for row in self._db.execute("PRAGMA table_info(workflow_steps)") - } - for name, statement in _STEP_MIGRATIONS: - if name not in columns: - self._db.execute(statement) + for table, migrations in ( + ("workflow_steps", _STEP_MIGRATIONS), + ("workflow_runs", _RUN_MIGRATIONS), + ): + columns = { + row["name"] + for row in self._db.execute(f"PRAGMA table_info({table})") + } + for name, statement in migrations: + if name not in columns: + self._db.execute(statement) self._db.execute( "CREATE INDEX IF NOT EXISTS idx_workflow_steps_lease" " ON workflow_steps (status, lease_expires_at)" @@ -1292,6 +1445,39 @@ def _append_events( (run_id, seq, event_type.value, now, json.dumps(data)), ) + def _insert_run(self, run: RunRecord) -> None: + """Insert a run row inside the current transaction. + + Args: + run: The run record. + """ + self._db.execute( + "INSERT INTO workflow_runs (run_id, workflow_id, definition_digest," + " status, state, state_version, next_ordinal, result, error," + " parent_run_id, parent_ordinal, request_key, labels, deadline," + " cancel_requested, created_at, updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + run.run_id, + run.workflow_id, + run.definition_digest, + run.status.value, + json.dumps(run.state), + run.state_version, + run.next_ordinal, + _dump(run.result), + _dump(run.error), + run.parent_run_id, + run.parent_ordinal, + run.request_key, + _dump(run.labels), + run.deadline, + int(run.cancel_requested), + run.created_at, + run.updated_at, + ), + ) + def _insert_step(self, step: StepRecord) -> None: """Insert a step row inside the current transaction. @@ -1301,8 +1487,8 @@ def _insert_step(self, step: StepRecord) -> None: self._db.execute( "INSERT INTO workflow_steps (run_id, ordinal, handler_id, status, args," " attempts, recoveries, due_at, epoch, lease_expires_at, wait_key," - " error, origin, created_at, updated_at)" - " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + " join_expected, join_arrived, error, origin, created_at, updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( step.run_id, step.ordinal, @@ -1315,6 +1501,8 @@ def _insert_step(self, step: StepRecord) -> None: step.epoch, step.lease_expires_at, step.wait_key, + step.join_expected, + step.join_arrived, _dump(step.error), step.origin, step.created_at, @@ -1370,30 +1558,7 @@ async def admit( " VALUES (?, ?, ?)", (run.workflow_id, run.request_key, run.run_id), ) - self._db.execute( - "INSERT INTO workflow_runs (run_id, workflow_id," - " definition_digest, status, state, state_version, next_ordinal," - " result, error, request_key, labels, deadline, cancel_requested," - " created_at, updated_at)" - " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - run.run_id, - run.workflow_id, - run.definition_digest, - run.status.value, - json.dumps(run.state), - run.state_version, - run.next_ordinal, - _dump(run.result), - _dump(run.error), - run.request_key, - _dump(run.labels), - run.deadline, - int(run.cancel_requested), - run.created_at, - run.updated_at, - ), - ) + self._insert_run(run) self._insert_step(root_step) self._append_events(run.run_id, events, run.created_at) self._db.execute("COMMIT") @@ -1789,6 +1954,131 @@ async def deliver( raise return "resolved" if resolves else "buffered" + async def admit_children( + self, + runs: tuple[tuple[RunRecord, StepRecord], ...], + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Create child runs, each with its root slot. + + Args: + runs: The child run records paired with their root slots. + events: History events to append to the parent. + now: Current time in epoch seconds. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + for run, root_step in runs: + self._insert_run(run) + self._insert_step(root_step) + if runs and events: + self._append_events(runs[0][0].parent_run_id or "", events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + async def record_arrival( + self, + run_id: str, + ordinal: int, + payload: dict[str, Any], + dedupe_key: str, + now: float, + ) -> DeliveryDisposition: + """Count one arrival against a join slot. + + Args: + run_id: The waiting parent run. + ordinal: The join slot's ordinal. + payload: The arriving result. + dedupe_key: Identity of the arrival. + now: Current time in epoch seconds. + + Returns: + What the store did with the arrival. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + wait_key = f"join:{ordinal}" + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + run_row = self._db.execute( + "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + if run_row is None: + self._db.execute("ROLLBACK") + return "unknown_run" + if run_row["status"] in terminal: + self._db.execute("ROLLBACK") + return "run_terminal" + seen = self._db.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", + (run_id, wait_key, dedupe_key), + ).fetchone() + if seen is not None: + self._db.execute("ROLLBACK") + return "duplicate" + step_row = self._db.execute( + "SELECT * FROM workflow_steps WHERE run_id = ? AND ordinal = ?", + (run_id, ordinal), + ).fetchone() + if step_row is None or step_row["status"] != StepStatus.BLOCKED.value: + self._db.execute("ROLLBACK") + return "run_terminal" + step = _step_from_row(step_row) + self._db.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at)" + " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" + " workflow_inbox WHERE run_id = ?), ?, ?, ?)", + ( + run_id, + wait_key, + dedupe_key, + run_id, + json.dumps(payload), + "CONSUMED", + now, + ), + ) + arrived = step.join_arrived + 1 + results = [*step.args.get("__results__", []), payload] + done = arrived >= step.join_expected + self._db.execute( + "UPDATE workflow_steps SET status = ?, join_arrived = ?," + " due_at = ?, args = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ? AND join_arrived = ?", + ( + StepStatus.READY.value if done else StepStatus.BLOCKED.value, + arrived, + now if done else step.due_at, + json.dumps({**step.args, "__results__": results}), + now, + run_id, + ordinal, + step.join_arrived, + ), + ) + self._append_events( + run_id, + ( + ( + HistoryEventType.CHILD_RESOLVED, + {"ordinal": ordinal, "arrived": arrived}, + ), + ), + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return "resolved" if done else "counted" + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py new file mode 100644 index 00000000000..3f0d9a7c895 --- /dev/null +++ b/tests/units/workflow/test_parallel.py @@ -0,0 +1,209 @@ +"""Tests for parallel fan-out via child runs.""" + +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import RunStatus, StepStatus +from reflex.workflow.testing import WorkflowTestHarness + +BRANCH_CALLS: list[str] = [] + + +class Enrich(rx.State): + """A branch that succeeds.""" + + __workflow__ = WorkflowConfig(id="fan.enrich") + lead: str = "" + + @rx.event(durable=True, trigger=manual(), effect="read") + def start(self, lead: str): + """Enrich the lead. + + Args: + lead: The lead identifier. + + Returns: + Completion carrying the enriched value. + """ + BRANCH_CALLS.append("enrich") + self.lead = lead + return rx.complete(result={"enriched": lead.upper()}) + + +class Flaky(rx.State): + """A branch that needs a retry of its own.""" + + __workflow__ = WorkflowConfig(id="fan.flaky") + lead: str = "" + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=3, initial_delay="1s", jitter="none"), + ) + def start(self, lead: str): + """Score the lead, failing once first. + + Args: + lead: The lead identifier. + + Returns: + Completion carrying the score. + """ + BRANCH_CALLS.append("flaky") + if BRANCH_CALLS.count("flaky") < 2: + msg = "scoring service down" + raise TransientWorkflowError(msg) + self.lead = lead + return rx.complete(result={"score": len(lead)}) + + +class Doomed(rx.State): + """A branch that always fails.""" + + __workflow__ = WorkflowConfig(id="fan.doomed") + + @rx.event( + durable=True, trigger=manual(), effect="read", retry=Retry(max_attempts=1) + ) + def start(self, lead: str): + """Fail to process the lead. + + Args: + lead: The lead identifier. + """ + BRANCH_CALLS.append("doomed") + msg = "permanently broken" + raise TransientWorkflowError(msg) + + +def _router(*branches): + """Build a parent workflow fanning out to the given branches. + + Args: + branches: The branch classes to fan out to. + + Returns: + The parent workflow class. + """ + + class Router(rx.State): + __workflow__ = WorkflowConfig(id="fan.router") + outcomes: list[str] = [] + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self, lead: str): + """Fan out to every branch. + + Args: + lead: The lead identifier. + + Returns: + The parallel fan-out. + """ + return rx.parallel( + *[branch.start(lead) for branch in branches], then=Router.route + ) + + @rx.event(durable=True, effect="none") + def route(self, results: list): + """Collect the branch outcomes. + + Args: + results: One entry per branch. + + Returns: + Completion carrying the branch count. + """ + self.outcomes = sorted(entry["status"] for entry in results) + return rx.complete(result={"branches": len(results)}) + + return Router + + +async def test_fan_out_joins_every_branch(forked_registration_context): + BRANCH_CALLS.clear() + router = _router(Enrich, Flaky) + async with WorkflowTestHarness(router, Enrich, Flaky) as harness: + result = await harness.start(router.begin("acme")) + assert result.run_id is not None + + # The flaky branch retries on its own without blocking its sibling. + await harness.advance("2s") + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"branches": 2} + assert snapshot.state["outcomes"] == ["COMPLETED", "COMPLETED"] + assert snapshot.steps[1].join_expected == 2 + assert snapshot.steps[1].join_arrived == 2 + + # Parent plus one run per branch, each independently inspectable. + runs = await harness.kernel.list_runs() + assert len(runs) == 3 + children = [run for run in runs if run.parent_run_id == result.run_id] + assert len(children) == 2 + + +async def test_a_failing_branch_still_reports(forked_registration_context): + """One branch failing does not strand the parent.""" + BRANCH_CALLS.clear() + router = _router(Enrich, Doomed) + async with WorkflowTestHarness(router, Enrich, Doomed) as harness: + result = await harness.start(router.begin("acme")) + assert result.run_id is not None + await harness.advance("2s") + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state["outcomes"] == ["COMPLETED", "FAILED"] + + +async def test_join_waits_for_the_last_branch(forked_registration_context): + """The parent stays blocked until every branch has reported.""" + BRANCH_CALLS.clear() + router = _router(Enrich, Flaky) + async with WorkflowTestHarness(router, Enrich, Flaky) as harness: + result = await harness.kernel.start(router.begin("acme")) + assert result.run_id is not None + await harness.kernel.run_until_idle() + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + # The flaky branch is still in backoff, so the join is short one arrival. + assert snapshot.status is RunStatus.WAITING + assert snapshot.steps[1].status is StepStatus.BLOCKED + assert snapshot.steps[1].join_arrived == 1 + + await harness.advance("2s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + + +async def test_duplicate_arrivals_are_counted_once(forked_registration_context): + """A redelivered child result must not satisfy the join twice.""" + BRANCH_CALLS.clear() + router = _router(Enrich, Flaky) + async with WorkflowTestHarness(router, Enrich, Flaky) as harness: + result = await harness.kernel.start(router.begin("acme")) + assert result.run_id is not None + await harness.kernel.run_until_idle() + + runs = await harness.kernel.list_runs() + child = next(run for run in runs if run.parent_run_id == result.run_id) + repeat = await harness.kernel.store.record_arrival( + result.run_id, + 1, + {"run_id": child.run_id, "status": "COMPLETED", "result": None}, + child.run_id, + harness.now, + ) + assert repeat == "duplicate" + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.steps[1].join_arrived == 1 From 7bf839c88e4d55c33624a7f14351d00ca8d65d20 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:47:07 -0700 Subject: [PATCH 013/121] Add start policies: singleton, debounce, rate limit, throttle Nothing bounded how often a root could start, so a chatty webhook produced one run per delivery and two clicks produced two concurrent syncs of the same customer. Flow control is the main thing Inngest sells and we had none of it. A root now declares one start policy, applied at admission and grouped by a payload field: singleton one active run per key; a second start either returns the first (skip) or replaces it (cancel) debounce a burst collapses into one run, each start pushing it out until things go quiet rate_limit starts beyond the cap are refused with retry_after, which is what you want when a provider can flood you throttle the excess is delayed rather than dropped, for when every start matters but the downstream is slow The grouping key must name a real parameter of the root, checked at compile time, and a policy without a trigger is rejected since it governs starting. Only one policy per root, so its behavior stays predictable. throttle= and debounce= are overloaded by type rather than given second names: an int is still the browser event action on a session handler, while the policy object is the durable start policy. Same word, same meaning, and a session handler passing an int is untouched. --- docs/workflows/overview.md | 26 ++ news/workflow-flow-control.feature.md | 1 + .../src/reflex_base/event/__init__.py | 54 ++-- .../reflex-base/src/reflex_base/workflow.py | 146 ++++++++++ pyi_hashes.json | 2 +- reflex/__init__.py | 4 + reflex/workflow/__init__.py | 8 + reflex/workflow/definition.py | 24 ++ reflex/workflow/kernel.py | 95 +++++++ reflex/workflow/records.py | 2 + reflex/workflow/store.py | 228 ++++++++++++++- tests/units/reflex_base/test_workflow.py | 4 + tests/units/workflow/test_flow_control.py | 264 ++++++++++++++++++ 13 files changed, 832 insertions(+), 26 deletions(-) create mode 100644 news/workflow-flow-control.feature.md create mode 100644 tests/units/workflow/test_flow_control.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index b0a1790caf8..91d798dce3b 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -245,6 +245,32 @@ really is public, say so with `allow_unverified=True` and a reason. Schedules are evaluated in UTC and fire once per occurrence even across restarts. Deploying a schedule does not backfill history, and catch-up after an outage is bounded. +## Controlling how runs start + +A root can declare one start policy, which the engine applies before a run is admitted. Each groups +runs by a payload field, or globally when no `key` is given. + +```python +@rx.event(durable=True, trigger=rx.manual(), effect="idempotent_write", + singleton=rx.Singleton(key="customer_id")) +def sync(self, customer_id: str): ... +``` + +| Policy | Behavior | +| --- | --- | +| `rx.Singleton(key=..., mode="skip")` | one active run per key; a second start returns the first | +| `rx.Singleton(key=..., mode="cancel")` | one active run per key; a second start replaces the first | +| `rx.Debounce(period=..., key=...)` | collapse a burst into one run once things go quiet | +| `rx.RateLimit(limit=..., period=..., key=...)` | cap starts per window, dropping the excess | +| `rx.Throttle(limit=..., period=..., key=...)` | cap starts per window, delaying the excess | + +`start()` reports what happened through its disposition: `"started"`, `"skipped"`, +`"coalesced"`, `"deduplicated"`, or `"rejected"` — a rejected start carries `retry_after`. + +Debounce is the one to reach for with chatty webhooks: ten deliveries in a second become one run. +Rate limiting drops excess starts, which is what you want when a provider can flood you; throttling +delays them instead, which is what you want when every start matters but the downstream is slow. + ## Inspecting and steering runs ```python diff --git a/news/workflow-flow-control.feature.md b/news/workflow-flow-control.feature.md new file mode 100644 index 00000000000..0379ba1eb15 --- /dev/null +++ b/news/workflow-flow-control.feature.md @@ -0,0 +1 @@ +Adds start policies — `rx.Singleton`, `rx.Debounce`, `rx.RateLimit`, and `rx.Throttle` — for controlling how often and how many runs of a root may begin. diff --git a/packages/reflex-base/src/reflex_base/event/__init__.py b/packages/reflex-base/src/reflex_base/event/__init__.py index 49207734ead..93185801199 100644 --- a/packages/reflex-base/src/reflex_base/event/__init__.py +++ b/packages/reflex-base/src/reflex_base/event/__init__.py @@ -2918,8 +2918,8 @@ def __new__( background: bool | None = None, stop_propagation: bool | None = None, prevent_default: bool | None = None, - throttle: int | None = None, - debounce: int | None = None, + throttle: "int | workflow.Throttle | None" = None, + debounce: "int | workflow.Debounce | None" = None, temporal: bool | None = None, id: str | None = None, durable: bool = False, @@ -2930,6 +2930,8 @@ def __new__( queue: str | None = None, on_failure: Any = None, on_timeout: Any = None, + singleton: "workflow.Singleton | None" = None, + rate_limit: "workflow.RateLimit | None" = None, ) -> ( "Callable[[Callable[[BASE_STATE, Unpack[P]], Any]], EventCallback[Unpack[P]]]" ): ... @@ -2942,8 +2944,8 @@ def __new__( background: bool | None = None, stop_propagation: bool | None = None, prevent_default: bool | None = None, - throttle: int | None = None, - debounce: int | None = None, + throttle: "int | workflow.Throttle | None" = None, + debounce: "int | workflow.Debounce | None" = None, temporal: bool | None = None, ) -> EventCallback[Unpack[P]]: ... @@ -2954,8 +2956,8 @@ def __new__( background: bool | None = None, stop_propagation: bool | None = None, prevent_default: bool | None = None, - throttle: int | None = None, - debounce: int | None = None, + throttle: "int | workflow.Throttle | None" = None, + debounce: "int | workflow.Debounce | None" = None, temporal: bool | None = None, id: str | None = None, durable: bool = False, @@ -2966,6 +2968,8 @@ def __new__( queue: str | None = None, on_failure: Any = None, on_timeout: Any = None, + singleton: "workflow.Singleton | None" = None, + rate_limit: "workflow.RateLimit | None" = None, ) -> "EventCallback[Unpack[P]] | Callable[[Callable[[BASE_STATE, Unpack[P]], Any]], EventCallback[Unpack[P]]]": """Wrap a function to be used as an event. @@ -2974,8 +2978,10 @@ def __new__( background: Whether the event should be run in the background. Defaults to False. stop_propagation: Whether to stop the event from bubbling up the DOM tree. prevent_default: Whether to prevent the default behavior of the event. - throttle: Throttle the event handler to limit calls (in milliseconds). - debounce: Debounce the event handler to delay calls (in milliseconds). + throttle: Milliseconds to throttle a session handler, or an + rx.Throttle policy capping how often a durable root may start. + debounce: Milliseconds to debounce a session handler, or an + rx.Debounce policy collapsing a burst of durable starts. temporal: Whether the event should be dropped when the backend is down. id: Stable durable handler id; derived from the method name if omitted. durable: Whether the handler is a durable workflow step. @@ -2986,6 +2992,8 @@ def __new__( queue: Admission queue override for a durable handler. on_failure: Same-class handler run after a durable step finally fails. on_timeout: Same-class handler run after a durable step finally times out. + singleton: Allow at most one active run of this root per key. + rate_limit: Cap the start rate per key, dropping the excess. Returns: The wrapped function. @@ -3003,17 +3011,17 @@ def __new__( queue=queue, on_failure=on_failure, on_timeout=on_timeout, + singleton=singleton, + rate_limit=rate_limit, + throttle=throttle, + debounce=debounce, background=background, has_browser_actions=any( value is not None - for value in ( - stop_propagation, - prevent_default, - throttle, - debounce, - temporal, - ) - ), + for value in (stop_propagation, prevent_default, temporal) + ) + or isinstance(throttle, int) + or isinstance(debounce, int), ) def _build_event_actions(): @@ -3022,11 +3030,13 @@ def _build_event_actions(): Returns: Dict of event actions to apply, or empty dict if none specified. """ + browser_throttle = throttle if isinstance(throttle, int) else None + browser_debounce = debounce if isinstance(debounce, int) else None if not any([ stop_propagation, prevent_default, - throttle, - debounce, + browser_throttle, + browser_debounce, temporal, ]): return {} @@ -3036,10 +3046,10 @@ def _build_event_actions(): event_actions["stopPropagation"] = stop_propagation if prevent_default is not None: event_actions["preventDefault"] = prevent_default - if throttle is not None: - event_actions["throttle"] = throttle - if debounce is not None: - event_actions["debounce"] = debounce + if browser_throttle is not None: + event_actions["throttle"] = browser_throttle + if browser_debounce is not None: + event_actions["debounce"] = browser_debounce if temporal is not None: event_actions["temporal"] = temporal return event_actions diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 574d51bb4b1..fb6dd1cf2c0 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -454,6 +454,10 @@ class DurableEventConfig: queue: Admission queue override. on_failure: Same-class handler name run after final failure. on_timeout: Same-class handler name run after final timeout. + singleton: At most one active run per key, if declared. + rate_limit: Start-rate cap that drops the excess, if declared. + throttle: Start-rate cap that delays the excess, if declared. + debounce: Burst collapsing window, if declared. """ effect: str @@ -464,6 +468,10 @@ class DurableEventConfig: queue: str | None = None on_failure: str | None = None on_timeout: str | None = None + singleton: Any = None + rate_limit: Any = None + throttle: Any = None + debounce: Any = None def get_durable_config(fn: Any) -> DurableEventConfig | None: @@ -517,6 +525,10 @@ def build_durable_config( queue: str | None, on_failure: Any, on_timeout: Any, + singleton: Any, + rate_limit: Any, + throttle: Any, + debounce: Any, background: bool | None, has_browser_actions: bool, ) -> DurableEventConfig | None: @@ -532,6 +544,10 @@ def build_durable_config( queue: Admission queue override. on_failure: Failure hook reference. on_timeout: Timeout hook reference. + singleton: Singleton policy, if declared. + rate_limit: Rate limit policy, if declared. + throttle: Throttle policy, if declared. + debounce: Debounce policy, if declared. background: The decorator's ``background`` flag. has_browser_actions: Whether browser-only event actions were also set. @@ -541,6 +557,10 @@ def build_durable_config( Raises: WorkflowDefinitionError: If the combination of arguments is invalid. """ + # throttle= and debounce= carry either a browser event action (an int) or a + # durable start policy; only the policy objects are workflow options. + throttle = throttle if isinstance(throttle, Throttle) else None + debounce = debounce if isinstance(debounce, Debounce) else None if not durable: offending = next( ( @@ -554,6 +574,10 @@ def build_durable_config( ("queue", queue), ("on_failure", on_failure), ("on_timeout", on_timeout), + ("singleton", singleton), + ("rate_limit", rate_limit), + ("throttle", throttle), + ("debounce", debounce), ) if value is not None ), @@ -605,6 +629,25 @@ def build_durable_config( 'effect="idempotent_write" or max_attempts=1.' ) raise WorkflowDefinitionError(msg) + controls = { + "singleton": singleton, + "rate_limit": rate_limit, + "throttle": throttle, + "debounce": debounce, + } + declared = [name for name, value in controls.items() if value is not None] + if declared and trigger is None: + msg = ( + f"@rx.event({declared[0]}=...) governs how runs start, so it needs a " + "trigger. Add trigger=rx.manual(), rx.webhook(...), or rx.schedule(...)." + ) + raise WorkflowDefinitionError(msg) + if len(declared) > 1: + msg = ( + f"@rx.event declares {' and '.join(declared)} together; a root takes " + "one start policy so its behavior stays predictable." + ) + raise WorkflowDefinitionError(msg) timeout_seconds = ( parse_duration(timeout, param="timeout") if timeout is not None else None ) @@ -617,6 +660,10 @@ def build_durable_config( queue=queue, on_failure=_hook_name(on_failure, param="on_failure"), on_timeout=_hook_name(on_timeout, param="on_timeout"), + singleton=singleton, + rate_limit=rate_limit, + throttle=throttle, + debounce=debounce, ) @@ -937,3 +984,102 @@ def parallel(*branches: Any, then: Any) -> Parallel: The control return value. """ return Parallel(branches=branches, then=then) + + +@dataclasses.dataclass(frozen=True, slots=True) +class Singleton: + """At most one run of this root may be active per key. + + Attributes: + key: Payload field whose value groups runs, or None for one group. + mode: ``"skip"`` returns the run already in flight; ``"cancel"`` + cancels it and starts a fresh one. + """ + + key: str | None = None + mode: Literal["skip", "cancel"] = "skip" + + def __post_init__(self): + """Validate the mode. + + Raises: + WorkflowDefinitionError: If the mode is not skip or cancel. + """ + if self.mode not in ("skip", "cancel"): + msg = f'Singleton.mode must be "skip" or "cancel", got {self.mode!r}.' + raise WorkflowDefinitionError(msg) + + +@dataclasses.dataclass(frozen=True, slots=True) +class RateLimit: + """Cap how many runs may start per key in a rolling window. + + Starts beyond the cap are refused, which is what you want for a provider + that can flood you: the excess is dropped rather than queued forever. + + Attributes: + limit: Maximum starts allowed inside the window. + period: Window length. + key: Payload field whose value groups runs, or None for one group. + """ + + limit: int + period: DurationLike + key: str | None = None + + def __post_init__(self): + """Validate the limit and window. + + Raises: + WorkflowDefinitionError: If the limit is not positive. + """ + if self.limit < 1: + msg = f"RateLimit.limit must be >= 1, got {self.limit}." + raise WorkflowDefinitionError(msg) + parse_duration(self.period, param="RateLimit.period") + + +@dataclasses.dataclass(frozen=True, slots=True) +class Throttle: + """Cap the start rate per key, delaying the excess instead of dropping it. + + Attributes: + limit: Maximum starts allowed inside the window. + period: Window length. + key: Payload field whose value groups runs, or None for one group. + """ + + limit: int + period: DurationLike + key: str | None = None + + def __post_init__(self): + """Validate the limit and window. + + Raises: + WorkflowDefinitionError: If the limit is not positive. + """ + if self.limit < 1: + msg = f"Throttle.limit must be >= 1, got {self.limit}." + raise WorkflowDefinitionError(msg) + parse_duration(self.period, param="Throttle.period") + + +@dataclasses.dataclass(frozen=True, slots=True) +class Debounce: + """Collapse a burst of starts into one run after things go quiet. + + Each new start inside the window pushes the pending run's start time out, + so a provider that fires ten webhooks in a second produces one run. + + Attributes: + period: How long to wait for quiet before running. + key: Payload field whose value groups runs, or None for one group. + """ + + period: DurationLike + key: str | None = None + + def __post_init__(self): + """Validate the window.""" + parse_duration(self.period, param="Debounce.period") diff --git a/pyi_hashes.json b/pyi_hashes.json index bbb97883d9f..b43fc973117 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -118,7 +118,7 @@ "packages/reflex-components-recharts/src/reflex_components_recharts/polar.pyi": "99ebcfc07868061bdc3c2010d85a153f", "packages/reflex-components-recharts/src/reflex_components_recharts/recharts.pyi": "4f6c26f8c76543cc41e2b9dc400ece8a", "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "f170ac685b6ba5892370166c80684db3", - "reflex/__init__.pyi": "ae53392d29647d23db03f68401293f0f", + "reflex/__init__.pyi": "121d29dc9426d3d52b4687611eddcb3f", "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388", "reflex/experimental/memo.pyi": "bc8b48357bef580e70a5881b65d3d3f7" } diff --git a/reflex/__init__.py b/reflex/__init__.py index 636f40ed11f..bbcce18d6a2 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -248,6 +248,10 @@ "wait_for", "never", "parallel", + "Singleton", + "RateLimit", + "Throttle", + "Debounce", "after", "complete", "fail", diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index 82774c3b38a..427eb15bf87 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -9,13 +9,17 @@ from reflex_base.workflow import ( ChannelDelivery, + Debounce, DurableEventConfig, EffectClass, ManualTrigger, Parallel, + RateLimit, Retry, ScheduleTrigger, Signal, + Singleton, + Throttle, TransientWorkflowError, Trigger, WaitFor, @@ -65,6 +69,7 @@ __all__ = [ "ChannelDelivery", + "Debounce", "DeliveryDisposition", "DurableEventConfig", "EffectClass", @@ -74,6 +79,7 @@ "ManualTrigger", "MemoryRunStore", "Parallel", + "RateLimit", "Retry", "RunQuery", "RunRecord", @@ -82,11 +88,13 @@ "RunStore", "ScheduleTrigger", "Signal", + "Singleton", "SqliteRunStore", "StaleClaimError", "StartResult", "StepRecord", "StepStatus", + "Throttle", "TransientWorkflowError", "Trigger", "WaitFor", diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index c7316bf10b4..9a93e23f37a 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -18,9 +18,13 @@ from reflex_base.utils.exceptions import WorkflowDefinitionError from reflex_base.workflow import ( + Debounce, DurableEventConfig, + RateLimit, Retry, ScheduleTrigger, + Singleton, + Throttle, Trigger, WorkflowConfig, default_retry_for_effect, @@ -69,6 +73,10 @@ class HandlerDefinition: queue: Resolved admission queue name, or None. on_failure: Handler id run after final failure, or None. on_timeout: Handler id run after final timeout, or None. + singleton: At most one active run per key, if declared. + rate_limit: Start-rate cap that drops the excess, if declared. + throttle: Start-rate cap that delays the excess, if declared. + debounce: Burst collapsing window, if declared. params: Payload parameter names, excluding ``self``. type_hints: Resolved type hints for payload coercion. is_async: Whether the handler is a coroutine function. @@ -84,6 +92,10 @@ class HandlerDefinition: queue: str | None on_failure: str | None on_timeout: str | None + singleton: Singleton | None + rate_limit: RateLimit | None + throttle: Throttle | None + debounce: Debounce | None params: tuple[str, ...] type_hints: Mapping[str, Any] is_async: bool @@ -302,6 +314,10 @@ def _compile_handlers( queue=durable.queue or config.default_queue, on_failure=durable.on_failure, on_timeout=durable.on_timeout, + singleton=durable.singleton, + rate_limit=durable.rate_limit, + throttle=durable.throttle, + debounce=durable.debounce, params=params, type_hints=get_type_hints(fn), is_async=inspect.iscoroutinefunction(fn), @@ -538,6 +554,14 @@ def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: fields = _compile_fields(workflow_cls) handlers = _resolve_hooks(workflow_cls, _compile_handlers(workflow_cls, config)) for defn in handlers.values(): + policy = defn.singleton or defn.rate_limit or defn.throttle or defn.debounce + key = getattr(policy, "key", None) + if key is not None and key not in defn.params: + raise _error( + workflow_cls, + f"handler {defn.name!r} groups runs by {key!r}, which is not one " + f"of its parameters ({', '.join(defn.params) or 'none'}).", + ) if isinstance(defn.trigger, ScheduleTrigger): CronSchedule(defn.trigger.cron) durable_names = frozenset(defn.name for defn in handlers.values()) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index fe43e61c299..bbbc615bb81 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -327,6 +327,91 @@ def _normalize_payload(args: dict[str, Any]) -> dict[str, Any]: msg = f"Workflow event payload is not serializable: {err}" raise WorkflowRuntimeError(msg) from None + @staticmethod + def _flow_key(handler: HandlerDefinition, payload: dict[str, Any]) -> str | None: + """Compute the grouping key a start policy applies to. + + Args: + handler: The root handler definition. + payload: The decoded start payload. + + Returns: + The key, or None when the handler declares no start policy. + """ + policy = ( + handler.singleton + or handler.rate_limit + or handler.throttle + or handler.debounce + ) + if policy is None: + return None + field = getattr(policy, "key", None) + if field is None: + return handler.id + return f"{handler.id}:{payload.get(field)!r}" + + async def _apply_start_policy( + self, + defn: WorkflowDefinition, + handler: HandlerDefinition, + flow_key: str, + now: float, + ) -> tuple[StartResult | None, float]: + """Decide whether and when a start may proceed. + + Args: + defn: The workflow definition. + handler: The root handler being started. + flow_key: The computed grouping key. + now: Current time in epoch seconds. + + Returns: + A result that ends admission, or None to proceed, together with the + time the root slot becomes due. + """ + if handler.singleton is not None: + existing = await self._store.first_active(defn.workflow_id, flow_key) + if existing is not None: + if handler.singleton.mode == "skip": + return ( + StartResult(disposition="skipped", run_id=existing.run_id), + now, + ) + await self.cancel(existing.run_id) + if handler.rate_limit is not None: + window = parse_duration(handler.rate_limit.period) + started = await self._store.count_started_since( + defn.workflow_id, flow_key, now - window + ) + if started >= handler.rate_limit.limit: + return ( + StartResult( + disposition="rejected", retryable=True, retry_after=window + ), + now, + ) + if handler.throttle is not None: + window = parse_duration(handler.throttle.period) + started = await self._store.count_started_since( + defn.workflow_id, flow_key, now - window + ) + if started >= handler.throttle.limit: + # Delay rather than drop: the excess still runs, just later. + return None, now + window + if handler.debounce is not None: + window = parse_duration(handler.debounce.period) + pending = await self._store.first_active(defn.workflow_id, flow_key) + if pending is not None and await self._store.defer_root( + pending.run_id, now + window, now + ): + return ( + StartResult(disposition="coalesced", run_id=pending.run_id), + now, + ) + return None, now + window + return None, now + async def start( self, target: Any, @@ -369,6 +454,14 @@ async def start( ) raise WorkflowRuntimeError(msg) now = self._clock() + flow_key = self._flow_key(handler, payload) + due_at = now + if flow_key is not None: + decided, due_at = await self._apply_start_policy( + defn, handler, flow_key, now + ) + if decided is not None: + return decided run_id = uuid.uuid4().hex run = RunRecord( run_id=run_id, @@ -378,6 +471,7 @@ async def start( state={field.name: field.default for field in defn.fields}, state_version=0, next_ordinal=1, + flow_key=flow_key, request_key=request_key, labels=labels, deadline=(now + defn.run_timeout) if defn.run_timeout is not None else None, @@ -390,6 +484,7 @@ async def start( handler_id=handler.id, status=StepStatus.READY, args=payload, + due_at=due_at, origin="root", created_at=now, updated_at=now, diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index fc5be6bd0a7..2c927077d04 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -154,6 +154,7 @@ class RunRecord: next_ordinal: Next mailbox ordinal to allocate. result: Run result recorded at completion. error: Terminal or suspension error payload. + flow_key: Grouping key for start policies such as singleton, if any. parent_run_id: The run that spawned this one, if any. parent_ordinal: The join slot in the parent this run reports to. request_key: Idempotent admission key, if one was supplied. @@ -173,6 +174,7 @@ class RunRecord: next_ordinal: int result: Any = None error: dict[str, Any] | None = None + flow_key: str | None = None parent_run_id: str | None = None parent_ordinal: int | None = None request_key: str | None = None diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index db0a21b8d85..dfd7454769f 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -297,6 +297,58 @@ async def record_arrival( """ ... + async def count_active(self, workflow_id: str, flow_key: str) -> int: + """Count runs of a root still in flight under a flow-control key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + + Returns: + How many non-terminal runs share the key. + """ + ... + + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: + """Find the oldest run still in flight under a flow-control key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + + Returns: + The run, or None when the key has no active run. + """ + ... + + async def count_started_since( + self, workflow_id: str, flow_key: str, since: float + ) -> int: + """Count runs of a root admitted under a key since a point in time. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + since: Exclusive lower bound in epoch seconds. + + Returns: + How many runs were admitted in the window. + """ + ... + + async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: + """Push a not-yet-started run's root slot later, for debouncing. + + Args: + run_id: The pending run. + due_at: The new earliest start time. + now: Current time in epoch seconds. + + Returns: + True when the root had not started and was deferred. + """ + ... + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. @@ -925,6 +977,85 @@ async def record_arrival( ) return "resolved" if done else "counted" + async def count_active(self, workflow_id: str, flow_key: str) -> int: + """Count runs of a root still in flight under a flow-control key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + + Returns: + How many non-terminal runs share the key. + """ + async with self._lock: + return sum( + 1 + for run in self._runs.values() + if run.workflow_id == workflow_id + and run.flow_key == flow_key + and run.status not in TERMINAL_RUN_STATUSES + ) + + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: + """Find the oldest run still in flight under a flow-control key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + + Returns: + The run, or None when the key has no active run. + """ + async with self._lock: + active = [ + run + for run in self._runs.values() + if run.workflow_id == workflow_id + and run.flow_key == flow_key + and run.status not in TERMINAL_RUN_STATUSES + ] + return min(active, key=lambda run: run.created_at) if active else None + + async def count_started_since( + self, workflow_id: str, flow_key: str, since: float + ) -> int: + """Count runs of a root admitted under a key since a point in time. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + since: Exclusive lower bound in epoch seconds. + + Returns: + How many runs were admitted in the window. + """ + async with self._lock: + return sum( + 1 + for run in self._runs.values() + if run.workflow_id == workflow_id + and run.flow_key == flow_key + and run.created_at > since + ) + + async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: + """Push a not-yet-started run's root slot later, for debouncing. + + Args: + run_id: The pending run. + due_at: The new earliest start time. + now: Current time in epoch seconds. + + Returns: + True when the root had not started and was deferred. + """ + async with self._lock: + steps = self._steps.get(run_id) + if not steps or steps[0].status is not StepStatus.READY: + return False + steps[0] = dataclasses.replace(steps[0], due_at=due_at, updated_at=now) + return True + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. @@ -1202,6 +1333,7 @@ async def next_due(self, now: float) -> float | None: next_ordinal INTEGER NOT NULL, result TEXT, error TEXT, + flow_key TEXT, parent_run_id TEXT, parent_ordinal INTEGER, request_key TEXT, @@ -1277,6 +1409,7 @@ async def next_due(self, now: float) -> float | None: ) _RUN_MIGRATIONS: Final = ( + ("flow_key", "ALTER TABLE workflow_runs ADD COLUMN flow_key TEXT"), ("parent_run_id", "ALTER TABLE workflow_runs ADD COLUMN parent_run_id TEXT"), ("parent_ordinal", "ALTER TABLE workflow_runs ADD COLUMN parent_ordinal INTEGER"), ) @@ -1325,6 +1458,7 @@ def _run_from_row(row: sqlite3.Row) -> RunRecord: next_ordinal=row["next_ordinal"], result=_load(row["result"]), error=_load(row["error"]), + flow_key=row["flow_key"], parent_run_id=row["parent_run_id"], parent_ordinal=row["parent_ordinal"], request_key=row["request_key"], @@ -1454,9 +1588,9 @@ def _insert_run(self, run: RunRecord) -> None: self._db.execute( "INSERT INTO workflow_runs (run_id, workflow_id, definition_digest," " status, state, state_version, next_ordinal, result, error," - " parent_run_id, parent_ordinal, request_key, labels, deadline," - " cancel_requested, created_at, updated_at)" - " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + " flow_key, parent_run_id, parent_ordinal, request_key, labels," + " deadline, cancel_requested, created_at, updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( run.run_id, run.workflow_id, @@ -1467,6 +1601,7 @@ def _insert_run(self, run: RunRecord) -> None: run.next_ordinal, _dump(run.result), _dump(run.error), + run.flow_key, run.parent_run_id, run.parent_ordinal, run.request_key, @@ -2079,6 +2214,93 @@ async def record_arrival( raise return "resolved" if done else "counted" + async def count_active(self, workflow_id: str, flow_key: str) -> int: + """Count runs of a root still in flight under a flow-control key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + + Returns: + How many non-terminal runs share the key. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + row = self._db.execute( + "SELECT COUNT(*) AS n FROM workflow_runs" + " WHERE workflow_id = ? AND flow_key = ?" + f" AND status NOT IN ({','.join('?' * len(terminal))})", + (workflow_id, flow_key, *terminal), + ).fetchone() + return row["n"] + + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: + """Find the oldest run still in flight under a flow-control key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + + Returns: + The run, or None when the key has no active run. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + row = self._db.execute( + "SELECT * FROM workflow_runs" + " WHERE workflow_id = ? AND flow_key = ?" + f" AND status NOT IN ({','.join('?' * len(terminal))})" + " ORDER BY created_at LIMIT 1", + (workflow_id, flow_key, *terminal), + ).fetchone() + return None if row is None else _run_from_row(row) + + async def count_started_since( + self, workflow_id: str, flow_key: str, since: float + ) -> int: + """Count runs of a root admitted under a key since a point in time. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + since: Exclusive lower bound in epoch seconds. + + Returns: + How many runs were admitted in the window. + """ + with self._lock: + row = self._db.execute( + "SELECT COUNT(*) AS n FROM workflow_runs" + " WHERE workflow_id = ? AND flow_key = ? AND created_at > ?", + (workflow_id, flow_key, since), + ).fetchone() + return row["n"] + + async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: + """Push a not-yet-started run's root slot later, for debouncing. + + Args: + run_id: The pending run. + due_at: The new earliest start time. + now: Current time in epoch seconds. + + Returns: + True when the root had not started and was deferred. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + cursor = self._db.execute( + "UPDATE workflow_steps SET due_at = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = 0 AND status = ?", + (due_at, now, run_id, StepStatus.READY.value), + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return cursor.rowcount > 0 + async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. diff --git a/tests/units/reflex_base/test_workflow.py b/tests/units/reflex_base/test_workflow.py index 5efb802b004..2454e367cb7 100644 --- a/tests/units/reflex_base/test_workflow.py +++ b/tests/units/reflex_base/test_workflow.py @@ -162,6 +162,10 @@ def _build(**overrides): "queue": None, "on_failure": None, "on_timeout": None, + "singleton": None, + "rate_limit": None, + "throttle": None, + "debounce": None, "background": None, "has_browser_actions": False, } diff --git a/tests/units/workflow/test_flow_control.py b/tests/units/workflow/test_flow_control.py new file mode 100644 index 00000000000..2cac1ed1f59 --- /dev/null +++ b/tests/units/workflow/test_flow_control.py @@ -0,0 +1,264 @@ +"""Tests for start policies: singleton, debounce, rate limit, and throttle.""" + +import pytest +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import ( + Debounce, + RateLimit, + Singleton, + Throttle, + WorkflowConfig, + after, + manual, +) + +import reflex as rx +from reflex.workflow.definition import compile_workflow +from reflex.workflow.records import RunStatus +from reflex.workflow.testing import WorkflowTestHarness + + +def _singleton_flow(mode: str = "skip"): + """Build a workflow allowing one active run per customer. + + Args: + mode: The singleton mode to apply. + + Returns: + The workflow class. + """ + + class SyncFlow(rx.State): + __workflow__ = WorkflowConfig(id="flow.sync") + cid: str = "" + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + singleton=Singleton(key="cid", mode=mode), # pyright: ignore[reportArgumentType] + ) + def start(self, cid: str): + """Begin a long sync. + + Args: + cid: The customer identifier. + + Returns: + A delayed finish step. + """ + self.cid = cid + return after("1h", SyncFlow.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + """Finish the sync.""" + + return SyncFlow + + +async def test_singleton_skips_a_second_start(forked_registration_context): + flow = _singleton_flow() + async with WorkflowTestHarness(flow) as harness: + first = await harness.start(flow.start("acme")) + second = await harness.kernel.start(flow.start("acme")) + assert first.disposition == "started" + assert second.disposition == "skipped" + assert second.run_id == first.run_id + assert len(await harness.kernel.list_runs()) == 1 + + +async def test_singleton_is_per_key(forked_registration_context): + flow = _singleton_flow() + async with WorkflowTestHarness(flow) as harness: + await harness.start(flow.start("acme")) + other = await harness.kernel.start(flow.start("globex")) + assert other.disposition == "started" + assert len(await harness.kernel.list_runs()) == 2 + + +async def test_singleton_cancel_mode_replaces_the_active_run( + forked_registration_context, +): + flow = _singleton_flow(mode="cancel") + async with WorkflowTestHarness(flow) as harness: + first = await harness.start(flow.start("acme")) + assert first.run_id is not None + second = await harness.start(flow.start("acme")) + assert second.disposition == "started" + assert second.run_id != first.run_id + + superseded = await harness.get_run(first.run_id) + assert superseded is not None + assert superseded.status is RunStatus.CANCELLED + + +async def test_singleton_frees_the_key_when_the_run_finishes( + forked_registration_context, +): + flow = _singleton_flow() + async with WorkflowTestHarness(flow) as harness: + first = await harness.start(flow.start("acme")) + assert first.run_id is not None + await harness.advance("1h") + assert (await harness.get_run(first.run_id)).status is RunStatus.COMPLETED # pyright: ignore[reportOptionalMemberAccess] + + again = await harness.kernel.start(flow.start("acme")) + assert again.disposition == "started" + + +async def test_debounce_collapses_a_burst(forked_registration_context): + calls = [] + + class BurstFlow(rx.State): + __workflow__ = WorkflowConfig(id="flow.burst") + cid: str = "" + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + debounce=Debounce(period="30s", key="cid"), + ) + def start(self, cid: str): + """Handle the collapsed burst. + + Args: + cid: The customer identifier. + """ + calls.append(cid) + self.cid = cid + + async with WorkflowTestHarness(BurstFlow) as harness: + first = await harness.kernel.start(BurstFlow.start("acme")) + second = await harness.kernel.start(BurstFlow.start("acme")) + third = await harness.kernel.start(BurstFlow.start("acme")) + assert first.disposition == "started" + assert second.disposition == "coalesced" + assert third.disposition == "coalesced" + assert second.run_id == first.run_id + + # Nothing has run yet: the window is still open. + await harness.kernel.run_until_idle() + assert calls == [] + + await harness.advance("31s") + assert calls == ["acme"] + assert len(await harness.kernel.list_runs()) == 1 + + +async def test_rate_limit_drops_the_excess(forked_registration_context): + class CappedFlow(rx.State): + __workflow__ = WorkflowConfig(id="flow.capped") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + rate_limit=RateLimit(limit=2, period="1m"), + ) + def start(self): + """Do the capped work.""" + + async with WorkflowTestHarness(CappedFlow) as harness: + dispositions = [ + (await harness.kernel.start(CappedFlow.start)).disposition for _ in range(4) + ] + assert dispositions == ["started", "started", "rejected", "rejected"] + + rejected = await harness.kernel.start(CappedFlow.start) + assert rejected.retryable + assert rejected.retry_after == pytest.approx(60.0) + + # The window rolls forward and starts are allowed again. + await harness.advance("61s") + assert (await harness.kernel.start(CappedFlow.start)).disposition == "started" + + +async def test_throttle_delays_the_excess(forked_registration_context): + calls = [] + + class ThrottledFlow(rx.State): + __workflow__ = WorkflowConfig(id="flow.throttled") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + throttle=Throttle(limit=1, period="1m"), + ) + def start(self): + """Do the throttled work.""" + calls.append(1) + + async with WorkflowTestHarness(ThrottledFlow) as harness: + first = await harness.start(ThrottledFlow.start) + assert first.disposition == "started" + assert len(calls) == 1 + + # The excess is admitted but held back rather than dropped. + second = await harness.start(ThrottledFlow.start) + assert second.disposition == "started" + assert len(calls) == 1 + + await harness.advance("61s") + assert len(calls) == 2 + + +def test_start_policy_key_must_name_a_parameter(forked_registration_context): + class BadKey(rx.State): + __workflow__ = WorkflowConfig(id="flow.bad_key") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + singleton=Singleton(key="customer"), + ) + def start(self, cid: str): + """Start with a mismatched key. + + Args: + cid: The customer identifier. + """ + + with pytest.raises(WorkflowDefinitionError, match="not one of its parameters"): + compile_workflow(BadKey) + + +def test_start_policies_require_a_trigger(): + with pytest.raises(WorkflowDefinitionError, match="needs a trigger"): + + @rx.event(durable=True, effect="none", singleton=Singleton()) + def handler(self): + pass + + +def test_only_one_start_policy_per_root(): + with pytest.raises(WorkflowDefinitionError, match="one start policy"): + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + singleton=Singleton(), + rate_limit=RateLimit(limit=1, period="1m"), + ) + def handler(self): + pass + + +def test_browser_throttle_still_works_on_session_handlers( + forked_registration_context, +): + """The int form of throttle and debounce stays a browser event action.""" + + class Clicky(rx.State): + @rx.event(throttle=200, debounce=100) + def click(self): + pass + + assert Clicky.event_handlers["click"].event_actions == { + "throttle": 200, + "debounce": 100, + } From 2b36be4afc86c6e717d3856a1766974010176a61 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:50:50 -0700 Subject: [PATCH 014/121] Ship the RunStore contract as runnable checks RunStore is a public extension point -- a deployment can back workflows with Postgres or a hosted kernel -- but the protocol's signatures say nothing about the invariants that make durable execution correct. Two implementations were already drifting apart with nothing but shared test files to hold them together, across 22 methods. reflex.workflow.CONFORMANCE_CHECKS is now the specification: 22 checks, each taking a fresh store and asserting one property. Frontier ordering, atomic commit, fenced claims discarding their work, failed attempts discarding state, lease renewal sparing a live claim, recovery reclaiming only lapsed ones, next_due never promising work that is not claimable, a deadline-less wait never becoming due, delivery never touching run state, joins counting each arrival once, finalize refusing while a step is claimed, and the queries start policies depend on. Both shipped stores pass all 22. Anyone adding a store runs the same suite; it is exported and documented for that purpose. Postgres is deliberately not in this commit: no server was available to test against, and an unverified store implementation is worse than none. --- docs/workflows/overview.md | 21 +- news/workflow-conformance.feature.md | 1 + reflex/workflow/__init__.py | 2 + reflex/workflow/conformance.py | 535 +++++++++++++++++++++++ tests/units/workflow/test_conformance.py | 30 ++ 5 files changed, 587 insertions(+), 2 deletions(-) create mode 100644 news/workflow-conformance.feature.md create mode 100644 reflex/workflow/conformance.py create mode 100644 tests/units/workflow/test_conformance.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index 91d798dce3b..12195b2fa3b 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -251,8 +251,12 @@ A root can declare one start policy, which the engine applies before a run is ad runs by a payload field, or globally when no `key` is given. ```python -@rx.event(durable=True, trigger=rx.manual(), effect="idempotent_write", - singleton=rx.Singleton(key="customer_id")) +@rx.event( + durable=True, + trigger=rx.manual(), + effect="idempotent_write", + singleton=rx.Singleton(key="customer_id"), +) def sync(self, customer_id: str): ... ``` @@ -318,6 +322,19 @@ to a waiting run, and `harness.cancel(...)` and `harness.resume(...)` drive the Runs persist to a SQLite file next to your app by default; pass `rx.App(workflow_store=...)` to choose another store. Run one worker process per SQLite database file. +`RunStore` is a supported extension point, and the invariants a store must satisfy ship as runnable +checks rather than prose: + +```python +import pytest +from reflex.workflow import CONFORMANCE_CHECKS + + +@pytest.mark.parametrize("check", CONFORMANCE_CHECKS, ids=lambda c: c.__name__) +async def test_my_store_conforms(check): + await check(MyRunStore()) +``` + Deploying new code does not disturb runs already in flight. Adding state fields, retuning retries and timeouts, and changing hooks all apply to future steps. Only a change that makes a pending step undispatchable — deleting the handler it names, or removing parameters its payload carries — diff --git a/news/workflow-conformance.feature.md b/news/workflow-conformance.feature.md new file mode 100644 index 00000000000..5e24d8a52e3 --- /dev/null +++ b/news/workflow-conformance.feature.md @@ -0,0 +1 @@ +Adds `reflex.workflow.CONFORMANCE_CHECKS`, a runnable specification of the invariants any `RunStore` implementation must satisfy. diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index 427eb15bf87..560462bcaca 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -40,6 +40,7 @@ webhook, ) +from reflex.workflow.conformance import CONFORMANCE_CHECKS from reflex.workflow.definition import ( HandlerDefinition, WorkflowDefinition, @@ -68,6 +69,7 @@ from reflex.workflow.testing import WorkflowTestHarness __all__ = [ + "CONFORMANCE_CHECKS", "ChannelDelivery", "Debounce", "DeliveryDisposition", diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py new file mode 100644 index 00000000000..2622817ae94 --- /dev/null +++ b/reflex/workflow/conformance.py @@ -0,0 +1,535 @@ +"""An executable specification for ``RunStore`` implementations. + +``RunStore`` is a public extension point: a deployment can back workflows with +Postgres, a hosted kernel, or anything else. The protocol's method signatures +say nothing about the invariants that make durable execution correct, so those +invariants live here as runnable checks rather than prose. + +Every check takes a fresh, empty store and asserts one property. Run the whole +suite against an implementation before trusting it:: + + import pytest + from reflex.workflow.conformance import CONFORMANCE_CHECKS + + @pytest.mark.parametrize("check", CONFORMANCE_CHECKS, ids=lambda c: c.__name__) + async def test_my_store_conforms(check): + await check(MyRunStore()) +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import pytest + +from reflex.workflow.records import ( + HistoryEventType, + RunQuery, + RunRecord, + RunStatus, + StepRecord, + StepStatus, +) +from reflex.workflow.store import StaleClaimError, StepCompletion + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from reflex.workflow.store import RunStore + +NOW = 1_000_000.0 + +LEASE = 30.0 + +_ADMITTED = ((HistoryEventType.RUN_ADMITTED, {}),) + + +def make_run(run_id: str = "run1", **overrides: Any) -> RunRecord: + """Build a run record for a conformance check. + + Args: + run_id: The run identity. + overrides: Fields to override on the record. + + Returns: + The run record. + """ + fields: dict[str, Any] = { + "run_id": run_id, + "workflow_id": "conformance.flow", + "definition_digest": "digest", + "status": RunStatus.PENDING, + "state": {"n": 0}, + "state_version": 0, + "next_ordinal": 1, + "created_at": NOW, + "updated_at": NOW, + } + fields.update(overrides) + return RunRecord(**fields) + + +def make_step(run_id: str = "run1", ordinal: int = 0, **overrides: Any) -> StepRecord: + """Build a step record for a conformance check. + + Args: + run_id: The owning run. + ordinal: The mailbox position. + overrides: Fields to override on the record. + + Returns: + The step record. + """ + fields: dict[str, Any] = { + "run_id": run_id, + "ordinal": ordinal, + "handler_id": "go", + "status": StepStatus.READY, + "args": {}, + "origin": "root", + "created_at": NOW, + "updated_at": NOW, + } + fields.update(overrides) + return StepRecord(**fields) + + +async def check_admit_creates_a_run(store: RunStore) -> None: + """A run and its root slot are readable straight after admission.""" + created, run_id = await store.admit(make_run(), make_step(), _ADMITTED) + assert created + assert run_id == "run1" + run = await store.get_run("run1") + assert run is not None + assert run.state == {"n": 0} + steps = await store.get_steps("run1") + assert [step.ordinal for step in steps] == [0] + history = await store.get_history("run1") + assert [event.seq for event in history] == [1] + + +async def check_admit_deduplicates_on_request_key(store: RunStore) -> None: + """One request key admits one run, however many times it is submitted.""" + await store.admit(make_run(request_key="key"), make_step(), _ADMITTED) + created, run_id = await store.admit( + make_run("run2", request_key="key"), make_step("run2"), _ADMITTED + ) + assert not created + assert run_id == "run1" + assert await store.get_run("run2") is None + + +async def check_only_the_frontier_is_claimable(store: RunStore) -> None: + """A later slot never overtakes an unresolved earlier one.""" + await store.admit(make_run(next_ordinal=2), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.RUNNING, + state={"n": 1}, + new_steps=( + make_step(ordinal=1, due_at=NOW + 60, origin="delay"), + make_step(ordinal=2, origin="chain"), + ), + next_ordinal=3, + ), + NOW, + ) + assert await store.claim_next(NOW, lease_duration=LEASE) is None + later = await store.claim_next(NOW + 61, lease_duration=LEASE) + assert later is not None + assert later.step.ordinal == 1 + + +async def check_commit_is_atomic(store: RunStore) -> None: + """State, successors, and history all become visible together.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.RUNNING, + state={"n": 7}, + new_steps=(make_step(ordinal=1, handler_id="second"),), + next_ordinal=2, + events=((HistoryEventType.STEP_SCHEDULED, {"ordinal": 1}),), + ), + NOW, + ) + run = await store.get_run("run1") + assert run is not None + assert run.state == {"n": 7} + assert run.state_version == 1 + assert run.next_ordinal == 2 + steps = await store.get_steps("run1") + assert len(steps) == 2 + assert any( + event.type is HistoryEventType.STEP_SCHEDULED + for event in await store.get_history("run1") + ) + + +async def check_a_fenced_claim_cannot_commit(store: RunStore) -> None: + """Two claims of one step cannot both commit.""" + await store.admit(make_run(), make_step(), _ADMITTED) + first = await store.claim_next(NOW, lease_duration=LEASE) + assert first is not None + await store.commit( + first, + StepCompletion( + step_status=StepStatus.RETRY_WAIT, + run_status=RunStatus.RETRYING, + state=None, + consume_attempt=True, + due_at=NOW, + ), + NOW, + ) + second = await store.claim_next(NOW, lease_duration=LEASE) + assert second is not None + assert second.step.epoch > first.step.epoch + with pytest.raises(StaleClaimError): + await store.commit( + first, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state={"n": 99}, + ), + NOW, + ) + run = await store.get_run("run1") + assert run is not None + assert run.state == {"n": 0} + + +async def check_a_failed_attempt_discards_its_state(store: RunStore) -> None: + """A commit that carries no state must not advance the state version.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.RETRY_WAIT, + run_status=RunStatus.RETRYING, + state=None, + consume_attempt=True, + due_at=NOW + 5, + ), + NOW, + ) + run = await store.get_run("run1") + assert run is not None + assert run.state_version == 0 + steps = await store.get_steps("run1") + assert steps[0].attempts == 1 + + +async def check_claim_carries_a_renewable_lease(store: RunStore) -> None: + """A claim's lease can be extended without transitioning the step.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + assert claim.step.lease_expires_at == pytest.approx(NOW + LEASE) + assert await store.renew_lease(claim, NOW + 10, lease_duration=LEASE) + steps = await store.get_steps("run1") + assert steps[0].lease_expires_at == pytest.approx(NOW + 10 + LEASE) + assert steps[0].epoch == claim.step.epoch + assert steps[0].status is StepStatus.CLAIMED + + +async def check_recovery_spares_a_live_lease(store: RunStore) -> None: + """A claim being executed by a peer is never reclaimed.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + assert await store.recover_orphans(NOW + LEASE - 1, max_recoveries=10) == 0 + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.CLAIMED + assert steps[0].recoveries == 0 + + +async def check_recovery_reclaims_an_expired_lease(store: RunStore) -> None: + """A lapsed claim is recovered, and charged to the recovery budget.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + assert await store.recover_orphans(NOW + LEASE, max_recoveries=10) == 1 + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.RECOVERY_WAIT + assert steps[0].recoveries == 1 + assert steps[0].attempts == 0 + assert not await store.renew_lease(claim, NOW + LEASE, lease_duration=LEASE) + + +async def check_next_due_only_promises_claimable_work(store: RunStore) -> None: + """Whatever time next_due reports, something really is claimable then.""" + await store.admit(make_run(), make_step(due_at=NOW + 120), _ADMITTED) + assert await store.claim_next(NOW, lease_duration=LEASE) is None + due = await store.next_due(NOW) + assert due is not None + assert due == pytest.approx(NOW + 120) + assert await store.claim_next(due, lease_duration=LEASE) is not None + + +async def check_a_wait_without_a_deadline_is_never_due(store: RunStore) -> None: + """A blocked slot with no deadline must not busy-wake the scheduler.""" + await store.admit( + make_run(), + make_step(status=StepStatus.BLOCKED, wait_key="sig:ping", due_at=0.0), + _ADMITTED, + ) + assert await store.claim_next(NOW, lease_duration=LEASE) is None + assert await store.next_due(NOW) is None + assert await store.claim_next(NOW + 86400, lease_duration=LEASE) is None + + +async def check_a_wait_deadline_makes_the_slot_claimable(store: RunStore) -> None: + """Claiming a blocked slot at its deadline is the timeout branch.""" + await store.admit( + make_run(), + make_step(status=StepStatus.BLOCKED, wait_key="sig:ping", due_at=NOW + 60), + _ADMITTED, + ) + assert await store.claim_next(NOW, lease_duration=LEASE) is None + assert await store.next_due(NOW) == pytest.approx(NOW + 60) + claim = await store.claim_next(NOW + 60, lease_duration=LEASE) + assert claim is not None + assert claim.step.ordinal == 0 + + +async def check_delivery_resolves_a_matching_wait(store: RunStore) -> None: + """A delivery flips the blocked slot and hands over its payload.""" + await store.admit( + make_run(), + make_step(status=StepStatus.BLOCKED, wait_key="sig:ping", due_at=0.0), + _ADMITTED, + ) + assert ( + await store.deliver("run1", "sig:ping", "d1", {"value": 1}, NOW) == "resolved" + ) + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.READY + assert steps[0].args["__payload__"] == {"value": 1} + + +async def check_delivery_never_touches_run_state(store: RunStore) -> None: + """A delivery must not be able to fence a live attempt.""" + await store.admit( + make_run(), + make_step(status=StepStatus.BLOCKED, wait_key="sig:ping", due_at=0.0), + _ADMITTED, + ) + before = await store.get_run("run1") + assert before is not None + await store.deliver("run1", "sig:ping", "d1", {"value": 1}, NOW) + after = await store.get_run("run1") + assert after is not None + assert after.state_version == before.state_version + assert after.state == before.state + + +async def check_duplicate_deliveries_are_ignored(store: RunStore) -> None: + """The same delivery key resolves a wait at most once.""" + await store.admit( + make_run(), + make_step(status=StepStatus.BLOCKED, wait_key="sig:ping", due_at=0.0), + _ADMITTED, + ) + assert await store.deliver("run1", "sig:ping", "d1", {"v": 1}, NOW) == "resolved" + assert await store.deliver("run1", "sig:ping", "d1", {"v": 2}, NOW) == "duplicate" + steps = await store.get_steps("run1") + assert steps[0].args["__payload__"] == {"v": 1} + + +async def check_an_early_delivery_is_buffered_then_consumed(store: RunStore) -> None: + """A signal that beats its wait is applied when the wait is armed.""" + await store.admit(make_run(), make_step(), _ADMITTED) + assert await store.deliver("run1", "sig:ping", "d1", {"v": 1}, NOW) == "buffered" + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.WAITING, + state={"n": 1}, + new_steps=( + make_step( + ordinal=1, + status=StepStatus.BLOCKED, + wait_key="sig:ping", + due_at=0.0, + origin="wait", + ), + ), + next_ordinal=2, + ), + NOW, + ) + steps = await store.get_steps("run1") + assert steps[1].status is StepStatus.READY + assert steps[1].args["__payload__"] == {"v": 1} + + +async def check_join_arrivals_count_once(store: RunStore) -> None: + """A join is satisfied by distinct arrivals, never by a redelivery.""" + await store.admit( + make_run(), + make_step( + status=StepStatus.BLOCKED, + wait_key="join:0", + join_expected=2, + origin="join", + due_at=0.0, + ), + _ADMITTED, + ) + assert await store.record_arrival("run1", 0, {"a": 1}, "c1", NOW) == "counted" + assert await store.record_arrival("run1", 0, {"a": 1}, "c1", NOW) == "duplicate" + steps = await store.get_steps("run1") + assert steps[0].join_arrived == 1 + assert steps[0].status is StepStatus.BLOCKED + assert await store.record_arrival("run1", 0, {"b": 2}, "c2", NOW) == "resolved" + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.READY + assert steps[0].join_arrived == 2 + assert len(steps[0].args["__results__"]) == 2 + + +async def check_finalize_refuses_while_a_step_is_claimed(store: RunStore) -> None: + """A run cannot be terminated out from under a running attempt.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + assert await store.request_cancel("run1", NOW) + assert await store.control_pending(NOW) == () + assert not await store.finalize_run( + "run1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW, + ) + + +async def check_finalize_tombstones_open_slots(store: RunStore) -> None: + """Terminating a run closes every slot it will never run.""" + await store.admit(make_run(next_ordinal=2), make_step(), _ADMITTED) + assert await store.request_cancel("run1", NOW) + assert await store.finalize_run( + "run1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW, + ) + steps = await store.get_steps("run1") + assert all(step.status is StepStatus.CANCELLED for step in steps) + assert not await store.request_cancel("run1", NOW) + + +async def check_resume_only_reopens_a_suspended_run(store: RunStore) -> None: + """Resuming applies to suspension, not to healthy or finished runs.""" + await store.admit( + make_run(status=RunStatus.NEEDS_ATTENTION), + make_step(status=StepStatus.NEEDS_ATTENTION, attempts=3), + _ADMITTED, + ) + assert await store.resume_run("run1", NOW) + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.PENDING + assert run.error is None + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.READY + assert steps[0].attempts == 0 + assert not await store.resume_run("run1", NOW) + assert not await store.resume_run("missing", NOW) + + +async def check_list_runs_filters_and_orders(store: RunStore) -> None: + """Listing is newest first and honors every filter.""" + await store.admit( + make_run("a", labels={"customer": "acme"}, created_at=NOW), + make_step("a"), + _ADMITTED, + ) + await store.admit( + make_run( + "b", + labels={"customer": "globex"}, + status=RunStatus.COMPLETED, + created_at=NOW + 1, + ), + make_step("b"), + _ADMITTED, + ) + listed = await store.list_runs(RunQuery()) + assert [run.run_id for run in listed] == ["b", "a"] + assert [ + run.run_id + for run in await store.list_runs(RunQuery(labels={"customer": "acme"})) + ] == ["a"] + assert [ + run.run_id + for run in await store.list_runs(RunQuery(statuses=(RunStatus.COMPLETED,))) + ] == ["b"] + assert await store.list_runs(RunQuery(workflow_id="other.flow", limit=10)) == () + + +async def check_flow_control_queries(store: RunStore) -> None: + """Start policies can see what is active and what started recently.""" + await store.admit( + make_run("a", flow_key="k1", created_at=NOW), make_step("a"), _ADMITTED + ) + await store.admit( + make_run("b", flow_key="k1", status=RunStatus.COMPLETED, created_at=NOW + 1), + make_step("b"), + _ADMITTED, + ) + await store.admit( + make_run("c", flow_key="k2", created_at=NOW + 2), make_step("c"), _ADMITTED + ) + assert await store.count_active("conformance.flow", "k1") == 1 + first = await store.first_active("conformance.flow", "k1") + assert first is not None + assert first.run_id == "a" + assert await store.first_active("conformance.flow", "missing") is None + assert await store.count_started_since("conformance.flow", "k1", NOW - 1) == 2 + assert await store.count_started_since("conformance.flow", "k1", NOW) == 1 + assert await store.defer_root("a", NOW + 30, NOW) + steps = await store.get_steps("a") + assert steps[0].due_at == pytest.approx(NOW + 30) + + +CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( + check_admit_creates_a_run, + check_admit_deduplicates_on_request_key, + check_only_the_frontier_is_claimable, + check_commit_is_atomic, + check_a_fenced_claim_cannot_commit, + check_a_failed_attempt_discards_its_state, + check_claim_carries_a_renewable_lease, + check_recovery_spares_a_live_lease, + check_recovery_reclaims_an_expired_lease, + check_next_due_only_promises_claimable_work, + check_a_wait_without_a_deadline_is_never_due, + check_a_wait_deadline_makes_the_slot_claimable, + check_delivery_resolves_a_matching_wait, + check_delivery_never_touches_run_state, + check_duplicate_deliveries_are_ignored, + check_an_early_delivery_is_buffered_then_consumed, + check_join_arrivals_count_once, + check_finalize_refuses_while_a_step_is_claimed, + check_finalize_tombstones_open_slots, + check_resume_only_reopens_a_suspended_run, + check_list_runs_filters_and_orders, + check_flow_control_queries, +) diff --git a/tests/units/workflow/test_conformance.py b/tests/units/workflow/test_conformance.py new file mode 100644 index 00000000000..b40d96bc993 --- /dev/null +++ b/tests/units/workflow/test_conformance.py @@ -0,0 +1,30 @@ +"""Run the store conformance suite against every shipped implementation.""" + +import pytest + +from reflex.workflow.conformance import CONFORMANCE_CHECKS +from reflex.workflow.store import MemoryRunStore, SqliteRunStore + + +@pytest.fixture(params=["memory", "sqlite"]) +def store(request, tmp_path): + """A fresh, empty store of each implementation. + + Args: + request: The fixture request carrying the store kind. + tmp_path: Temporary directory for the SQLite database. + + Yields: + The store instance. + """ + if request.param == "memory": + yield MemoryRunStore() + else: + sqlite_store = SqliteRunStore(tmp_path / "workflow.db") + yield sqlite_store + sqlite_store.close() + + +@pytest.mark.parametrize("check", CONFORMANCE_CHECKS, ids=lambda check: check.__name__) +async def test_store_conforms(store, check): + await check(store) From f6b88bc7dc75ae5c6e686922311d23ebaabbeb27 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:54:19 -0700 Subject: [PATCH 015/121] Add the reflex workflows CLI Runs could only be reached from inside the app, so diagnosing one meant writing a script against the store. That is the wrong tool at 2am. reflex workflows list filters by workflow, status, and label; show renders a run's state, steps with their attempt and recovery counts, and optionally its history; cancel and resume steer a run without opening the app. Both list and show take --json so the output is scriptable. The commands read the same SQLite database the app writes, so they work against a running deployment or a stopped one, and resume refuses a run that is not actually suspended rather than pretending to act. --- docs/workflows/overview.md | 13 ++ news/workflow-cli.feature.md | 1 + reflex/reflex.py | 2 + reflex/workflow/cli.py | 272 +++++++++++++++++++++++++++++++ tests/units/workflow/test_cli.py | 175 ++++++++++++++++++++ 5 files changed, 463 insertions(+) create mode 100644 news/workflow-cli.feature.md create mode 100644 reflex/workflow/cli.py create mode 100644 tests/units/workflow/test_cli.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index 12195b2fa3b..f71293b588e 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -294,6 +294,19 @@ await rx.workflows.list_runs(labels={"customer": customer.id}) A suspended run is waiting for you, not finished: fix whatever made the outcome uncertain, then `resume()` to give the step a fresh attempt budget. +The same operations are available from a terminal, reading the app's own database: + +```bash +reflex workflows list --status NEEDS_ATTENTION +``` + +```bash +reflex workflows show --history +``` + +`reflex workflows cancel ` and `reflex workflows resume ` steer a run without +opening the app, and `--json` on `list` and `show` makes the output scriptable. + ## Testing The test harness runs your real workflow on a virtual clock, so a three-day wait takes microseconds diff --git a/news/workflow-cli.feature.md b/news/workflow-cli.feature.md new file mode 100644 index 00000000000..973e5c26310 --- /dev/null +++ b/news/workflow-cli.feature.md @@ -0,0 +1 @@ +Adds the `reflex workflows` command group for listing, inspecting, cancelling, and resuming durable runs from a terminal. diff --git a/reflex/reflex.py b/reflex/reflex.py index 6f9fdd2999a..6295b207711 100644 --- a/reflex/reflex.py +++ b/reflex/reflex.py @@ -14,6 +14,7 @@ from reflex_cli.v2.deployments import hosting_cli from reflex.custom_components.custom_components import custom_components_cli +from reflex.workflow.cli import workflows as workflows_cli if TYPE_CHECKING: from typing import Literal @@ -1097,6 +1098,7 @@ def _convert_reflex_loglevel_to_reflex_cli_loglevel( cli.add_command(db_cli, name="db") cli.add_command(script_cli, name="script") cli.add_command(custom_components_cli, name="component") +cli.add_command(workflows_cli, name="workflows") if __name__ == "__main__": cli() diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py new file mode 100644 index 00000000000..c736cb1d19c --- /dev/null +++ b/reflex/workflow/cli.py @@ -0,0 +1,272 @@ +"""The ``reflex workflows`` command group. + +Operators reach for a terminal when a run misbehaves, so listing, inspecting, +cancelling, and resuming runs must not require writing a script or opening the +app. These commands read the same store the app writes, so they work against a +running deployment or a stopped one. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import TYPE_CHECKING, Any + +import click +from reflex_base.utils import console + +from reflex.workflow.records import RunStatus + +if TYPE_CHECKING: + from collections.abc import Awaitable + + from reflex.workflow.store import RunStore + +DEFAULT_DB_FILENAME = "workflow.db" + + +def _open_store(database: str | None) -> RunStore: + """Open the run store the app persists to. + + Args: + database: Path to the SQLite database, or None for the default. + + Returns: + The store. + """ + from reflex.workflow.store import SqliteRunStore + + return SqliteRunStore(database or DEFAULT_DB_FILENAME) + + +def _run(coroutine: Awaitable[Any]) -> Any: + """Run one store coroutine to completion. + + Args: + coroutine: The coroutine to run. + + Returns: + Its result. + """ + return asyncio.run(coroutine) # pyright: ignore[reportArgumentType] + + +def _age(seconds: float) -> str: + """Render an age compactly. + + Args: + seconds: How long ago, in seconds. + + Returns: + A short human-readable age. + """ + for unit, size in (("d", 86400), ("h", 3600), ("m", 60)): + if seconds >= size: + return f"{int(seconds // size)}{unit}" + return f"{int(seconds)}s" + + +database_option = click.option( + "--database", + "-d", + default=None, + help="Path to the workflow database. Defaults to ./workflow.db.", +) + + +@click.group() +def workflows(): + """Inspect and steer durable workflow runs.""" + + +@workflows.command("list") +@database_option +@click.option("--workflow", "-w", default=None, help="Only this workflow id.") +@click.option( + "--status", + "-s", + "statuses", + multiple=True, + type=click.Choice([status.value for status in RunStatus], case_sensitive=False), + help="Only these run statuses. Repeatable.", +) +@click.option("--label", "-l", "labels", multiple=True, help="Filter as key=value.") +@click.option("--limit", "-n", default=20, show_default=True, help="Rows to show.") +@click.option("--json", "as_json", is_flag=True, help="Emit JSON instead of a table.") +def list_runs( + database: str | None, + workflow: str | None, + statuses: tuple[str, ...], + labels: tuple[str, ...], + limit: int, + as_json: bool, +): + """List runs, newest first.""" + from reflex.workflow.records import RunQuery + + label_filter = dict(pair.split("=", 1) for pair in labels if "=" in pair) + store = _open_store(database) + try: + runs = _run( + store.list_runs( + RunQuery( + workflow_id=workflow, + statuses=tuple(RunStatus(value.upper()) for value in statuses), + labels=label_filter or None, + limit=limit, + ) + ) + ) + finally: + _close(store) + + if as_json: + click.echo( + json.dumps( + [ + { + "run_id": run.run_id, + "workflow_id": run.workflow_id, + "status": run.status.value, + "labels": run.labels, + "created_at": run.created_at, + } + for run in runs + ], + indent=2, + ) + ) + return + if not runs: + console.print("No runs matched.") + return + newest = max(run.created_at for run in runs) + click.echo(f"{'RUN':34}{'WORKFLOW':28}{'STATUS':17}AGE") + for run in runs: + click.echo( + f"{run.run_id:34}{run.workflow_id:28}{run.status.value:17}" + f"{_age(newest - run.created_at)}" + ) + + +@workflows.command() +@database_option +@click.argument("run_id") +@click.option("--json", "as_json", is_flag=True, help="Emit JSON instead of a table.") +@click.option("--history", is_flag=True, help="Include the run's history.") +def show(database: str | None, run_id: str, as_json: bool, history: bool): + """Show one run's state, steps, and optionally its history.""" + store = _open_store(database) + try: + run = _run(store.get_run(run_id)) + if run is None: + console.error(f"No run {run_id!r} in this database.") + raise click.exceptions.Exit(1) + steps = _run(store.get_steps(run_id)) + events = _run(store.get_history(run_id)) if history else () + finally: + _close(store) + + if as_json: + click.echo( + json.dumps( + { + "run_id": run.run_id, + "workflow_id": run.workflow_id, + "status": run.status.value, + "state": run.state, + "result": run.result, + "error": run.error, + "steps": [ + { + "ordinal": step.ordinal, + "handler_id": step.handler_id, + "status": step.status.value, + "attempts": step.attempts, + "recoveries": step.recoveries, + } + for step in steps + ], + "history": [ + {"seq": event.seq, "type": event.type.value, "at": event.at} + for event in events + ], + }, + indent=2, + default=str, + ) + ) + return + + click.echo(f"run {run.run_id}") + click.echo(f"workflow {run.workflow_id}") + click.echo(f"status {run.status.value}") + if run.result is not None: + click.echo(f"result {run.result}") + if run.error is not None: + click.echo(f"error {run.error}") + click.echo(f"state {run.state}") + click.echo("") + click.echo(f"{'#':<4}{'HANDLER':28}{'STATUS':17}ATTEMPTS") + for step in steps: + attempts = f"{step.attempts}" + if step.recoveries: + attempts += f" (+{step.recoveries} recovered)" + click.echo( + f"{step.ordinal:<4}{step.handler_id:28}{step.status.value:17}{attempts}" + ) + if history: + click.echo("") + for event in events: + click.echo(f"{event.seq:<4}{event.type.value}") + + +@workflows.command() +@database_option +@click.argument("run_id") +def cancel(database: str | None, run_id: str): + """Request cancellation of a run. + + The running worker finalizes it; if no worker is running, it is cancelled + the next time one starts. + """ + import time + + store = _open_store(database) + try: + recorded = _run(store.request_cancel(run_id, time.time())) + finally: + _close(store) + if not recorded: + console.error(f"Run {run_id!r} is unknown or already finished.") + raise click.exceptions.Exit(1) + console.print(f"Cancellation requested for {run_id}.") + + +@workflows.command() +@database_option +@click.argument("run_id") +def resume(database: str | None, run_id: str): + """Re-open a run suspended for operator attention.""" + import time + + store = _open_store(database) + try: + resumed = _run(store.resume_run(run_id, time.time())) + finally: + _close(store) + if not resumed: + console.error(f"Run {run_id!r} is not suspended.") + raise click.exceptions.Exit(1) + console.print(f"Resumed {run_id}; its next step will run.") + + +def _close(store: RunStore) -> None: + """Close a store that holds a connection. + + Args: + store: The store to close. + """ + closer = getattr(store, "close", None) + if closer is not None: + closer() diff --git a/tests/units/workflow/test_cli.py b/tests/units/workflow/test_cli.py new file mode 100644 index 00000000000..72c518d82a9 --- /dev/null +++ b/tests/units/workflow/test_cli.py @@ -0,0 +1,175 @@ +"""Tests for the reflex workflows command group.""" + +import asyncio +import json + +import pytest +from click.testing import CliRunner +from reflex_base.workflow import WorkflowConfig, manual, needs_attention + +import reflex as rx +from reflex.workflow.cli import workflows +from reflex.workflow.records import RunStatus +from reflex.workflow.store import SqliteRunStore +from reflex.workflow.testing import WorkflowTestHarness + + +@pytest.fixture +def seeded(forked_registration_context, tmp_path): + """A database with one waiting run and one suspended run. + + Args: + forked_registration_context: Isolates state registration. + tmp_path: Temporary directory for the database. + + Returns: + The database path and the two run ids. + """ + + class OpsFlow(rx.State): + __workflow__ = WorkflowConfig(id="ops.cli") + cid: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, cid: str): + """Start work for a customer. + + Args: + cid: The customer identifier. + + Returns: + A delayed finish, or a suspension for the flagged customer. + """ + self.cid = cid + if cid == "flagged": + return needs_attention("manual_check") + return rx.after("1h", OpsFlow.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + """Finish the work.""" + + db_path = tmp_path / "workflow.db" + + async def seed(): + store = SqliteRunStore(db_path) + async with WorkflowTestHarness(OpsFlow, store=store) as harness: + waiting = await harness.start(OpsFlow.start("acme"), labels={"tier": "pro"}) + suspended = await harness.start(OpsFlow.start("flagged")) + store.close() + return waiting.run_id, suspended.run_id + + waiting_id, suspended_id = asyncio.run(seed()) + return str(db_path), waiting_id, suspended_id + + +def _load_run(database: str, run_id: str): + """Read a run back from the database the CLI just wrote. + + Args: + database: Path to the SQLite database. + run_id: The run to read. + + Returns: + The run record, or None. + """ + + async def load(): + store = SqliteRunStore(database) + try: + return await store.get_run(run_id) + finally: + store.close() + + return asyncio.run(load()) + + +def _invoke(*args): + """Run a CLI command and return its result. + + Args: + args: The command line arguments. + + Returns: + The click result. + """ + return CliRunner().invoke(workflows, list(args)) + + +def test_list_shows_runs_newest_first(seeded): + database, _, _ = seeded + result = _invoke("list", "-d", database) + assert result.exit_code == 0 + assert "ops.cli" in result.output + assert result.output.count("ops.cli") == 2 + + +def test_list_filters_by_status_and_label(seeded): + database, waiting, suspended = seeded + by_status = _invoke("list", "-d", database, "-s", "NEEDS_ATTENTION", "--json") + assert by_status.exit_code == 0 + rows = json.loads(by_status.output) + assert [row["run_id"] for row in rows] == [suspended] + + by_label = _invoke("list", "-d", database, "-l", "tier=pro", "--json") + rows = json.loads(by_label.output) + assert [row["run_id"] for row in rows] == [waiting] + + +def test_list_reports_when_nothing_matches(seeded): + database, _, _ = seeded + result = _invoke("list", "-d", database, "-w", "nope.nothing") + assert result.exit_code == 0 + assert "No runs matched" in result.output + + +def test_show_renders_steps_and_history(seeded): + database, waiting, _ = seeded + result = _invoke("show", "-d", database, waiting, "--history") + assert result.exit_code == 0 + assert waiting in result.output + assert "start" in result.output + assert "run_admitted" in result.output + + +def test_show_json_carries_state_and_steps(seeded): + database, waiting, _ = seeded + result = _invoke("show", "-d", database, waiting, "--json") + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["state"] == {"cid": "acme"} + assert payload["status"] == RunStatus.WAITING.value + assert [step["handler_id"] for step in payload["steps"]] == ["start", "finish"] + + +def test_show_unknown_run_fails(seeded): + database, _, _ = seeded + result = _invoke("show", "-d", database, "no-such-run") + assert result.exit_code == 1 + + +def test_cancel_records_intent(seeded): + database, waiting, _ = seeded + result = _invoke("cancel", "-d", database, waiting) + assert result.exit_code == 0 + assert "Cancellation requested" in result.output + + run = _load_run(database, waiting) + assert run is not None + assert run.cancel_requested + + # Cancelling twice is refused rather than silently repeated. + assert _invoke("cancel", "-d", database, waiting).exit_code == 0 + + +def test_resume_reopens_only_suspended_runs(seeded): + database, waiting, suspended = seeded + assert _invoke("resume", "-d", database, waiting).exit_code == 1 + + result = _invoke("resume", "-d", database, suspended) + assert result.exit_code == 0 + assert "Resumed" in result.output + + run = _load_run(database, suspended) + assert run is not None + assert run.status is RunStatus.PENDING From 5f1884292a282e99a99cc80c8bf6b852415bd3a4 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 14:59:50 -0700 Subject: [PATCH 016/121] Add a workflow observer hook Runs recorded a full history but nothing could watch them happen: diagnosing a production workflow meant querying the database after the fact, and there was no way to get workflow activity into metrics or tracing. WorkflowObserver receives every transition the kernel records -- admission, each attempt and its outcome, retries, waits, joins, and terminal dispositions -- with the correlation a durable system needs: run, workflow, step, and attempt. Install one with rx.App(workflow_observer=...). LoggingObserver is bundled for the common case. Instrumentation is deliberately not allowed to break execution: an observer that raises is reported and ignored, which a test asserts by running a workflow to completion under an observer that always throws. --- docs/workflows/overview.md | 22 +++ news/workflow-observability.feature.md | 1 + reflex/app.py | 8 +- reflex/workflow/__init__.py | 4 +- reflex/workflow/kernel.py | 107 +++++++++++++-- reflex/workflow/runtime.py | 10 +- reflex/workflow/testing.py | 4 + tests/units/workflow/test_observability.py | 147 +++++++++++++++++++++ 8 files changed, 289 insertions(+), 14 deletions(-) create mode 100644 news/workflow-observability.feature.md create mode 100644 tests/units/workflow/test_observability.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index f71293b588e..d22785762ec 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -307,6 +307,28 @@ reflex workflows show --history `reflex workflows cancel ` and `reflex workflows resume ` steer a run without opening the app, and `--json` on `list` and `show` makes the output scriptable. +## Observability + +Pass an observer to see every recorded transition, correlated to its run, workflow, step, and +attempt: + +```python +from reflex.workflow import WorkflowObserver + + +class Telemetry(WorkflowObserver): + def on_event(self, event_type, run_id, workflow_id, data): + metrics.increment( + f"workflow.{event_type.value}", tags={"workflow": workflow_id} + ) + + +app = rx.App(workflow_observer=Telemetry()) +``` + +`rx.workflow.LoggingObserver` is a ready-made one that writes a structured line per transition. +An observer that raises is reported and ignored — instrumentation never breaks a run. + ## Testing The test harness runs your real workflow on a virtual clock, so a three-day wait takes microseconds diff --git a/news/workflow-observability.feature.md b/news/workflow-observability.feature.md new file mode 100644 index 00000000000..43f8017d7e3 --- /dev/null +++ b/news/workflow-observability.feature.md @@ -0,0 +1 @@ +Adds `rx.App(workflow_observer=...)` for forwarding every workflow run transition to logs, metrics, or tracing. diff --git a/reflex/app.py b/reflex/app.py index 2b4c8b1d98e..2a59bc3035c 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -98,6 +98,7 @@ ) from reflex.utils.misc import run_in_thread from reflex.utils.token_manager import RedisTokenManager, TokenManager +from reflex.workflow.kernel import WorkflowObserver from reflex.workflow.runtime import WorkflowRuntime from reflex.workflow.store import RunStore @@ -456,6 +457,9 @@ class App(MiddlewareMixin, LifespanMixin): # store created when the app starts. workflow_store: RunStore | None = None + # Receives every recorded workflow run transition, for logs or tracing. + workflow_observer: WorkflowObserver | None = None + # The workflow runtime owning registered definitions and the kernel. _workflow_runtime: WorkflowRuntime | None = None @@ -969,7 +973,9 @@ def add_workflow(self, workflow_cls: type[BaseState]) -> None: ``__workflow__ = rx.WorkflowConfig(id=...)`` declaration. """ if self._workflow_runtime is None: - self._workflow_runtime = WorkflowRuntime(self.workflow_store) + self._workflow_runtime = WorkflowRuntime( + self.workflow_store, observer=self.workflow_observer + ) self.register_lifespan_task(self._run_workflow_runtime) self._workflow_runtime.register(workflow_cls) diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index 560462bcaca..dc4da9cf8d2 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -46,7 +46,7 @@ WorkflowDefinition, compile_workflow, ) -from reflex.workflow.kernel import WorkflowKernel +from reflex.workflow.kernel import LoggingObserver, WorkflowKernel, WorkflowObserver from reflex.workflow.records import ( HistoryEvent, HistoryEventType, @@ -78,6 +78,7 @@ "HandlerDefinition", "HistoryEvent", "HistoryEventType", + "LoggingObserver", "ManualTrigger", "MemoryRunStore", "Parallel", @@ -105,6 +106,7 @@ "WorkflowConfig", "WorkflowDefinition", "WorkflowKernel", + "WorkflowObserver", "WorkflowRuntime", "WorkflowTestHarness", "after", diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index bbbc615bb81..55b7c357cbb 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -93,6 +93,64 @@ def _error_payload(error: BaseException) -> dict[str, Any]: } +class WorkflowObserver: + """Receives every run transition the kernel records. + + Subclass and pass to ``rx.App(workflow_observer=...)`` to forward workflow + activity to logs, metrics, or a tracing backend. Every call carries the + correlation a durable system needs -- run, workflow, step, and attempt -- + so a line can be tied to the exact execution that produced it. + + Callbacks must not raise and must not block: they run on the kernel's + event loop, and the kernel deliberately swallows their errors rather than + letting instrumentation break execution. + """ + + def on_event( + self, + event_type: HistoryEventType, + run_id: str, + workflow_id: str, + data: dict[str, Any], + ) -> None: + """Handle one recorded transition. + + Args: + event_type: What happened. + run_id: The run it happened to. + workflow_id: That run's workflow identity. + data: Event payload, such as ordinal, handler, attempt, or error. + """ + + +class LoggingObserver(WorkflowObserver): + """Logs every transition as a structured line.""" + + def on_event( + self, + event_type: HistoryEventType, + run_id: str, + workflow_id: str, + data: dict[str, Any], + ) -> None: + """Log one transition. + + Args: + event_type: What happened. + run_id: The run it happened to. + workflow_id: That run's workflow identity. + data: Event payload. + """ + detail = " ".join( + f"{key}={value!r}" for key, value in data.items() if key != "error" + ) + line = f"workflow={workflow_id} run={run_id} event={event_type.value} {detail}" + if "error" in data: + console.warn(f"{line} error={data['error']}") + else: + console.debug(line) + + class _SuccessorSpec: """A resolved successor slot to allocate at commit. @@ -165,6 +223,7 @@ def __init__( lease_duration: float = DEFAULT_LEASE_DURATION, lease_renew_interval: float | None = None, recovery_interval: float | None = None, + observer: WorkflowObserver | None = None, ): """Initialize the kernel. @@ -181,6 +240,8 @@ def __init__( to a third of ``lease_duration``. recovery_interval: Seconds between recovery sweeps in the background worker; defaults to half of ``lease_duration``. + observer: Receives every recorded transition, for logging, metrics, + or tracing. Raises: WorkflowRuntimeError: If the store cannot renew leases, or the @@ -241,6 +302,7 @@ def __init__( self._leases: dict[str, _Lease] = {} self._next_recovery_at = 0.0 self._worker_id = uuid.uuid4().hex + self._observer = observer self._wakeup = asyncio.Event() self._worker: asyncio.Task | None = None @@ -489,22 +551,22 @@ async def start( created_at=now, updated_at=now, ) - created, authoritative_run_id = await self._store.admit( - run, - root_step, + admission = ( ( - ( - HistoryEventType.RUN_ADMITTED, - {"handler_id": handler.id, "request_key": request_key}, - ), - ( - HistoryEventType.STEP_SCHEDULED, - {"ordinal": 0, "handler_id": handler.id}, - ), + HistoryEventType.RUN_ADMITTED, + {"handler_id": handler.id, "request_key": request_key}, ), + ( + HistoryEventType.STEP_SCHEDULED, + {"ordinal": 0, "handler_id": handler.id}, + ), + ) + created, authoritative_run_id = await self._store.admit( + run, root_step, admission ) if not created: return StartResult(disposition="deduplicated", run_id=authoritative_run_id) + self._notify(run, admission) self._wakeup.set() return StartResult(disposition="started", run_id=authoritative_run_id) @@ -1217,6 +1279,28 @@ def _success_completion( events=tuple(events), ) + def _notify( + self, + run: RunRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + ) -> None: + """Hand recorded transitions to the observer, if one is installed. + + Instrumentation must never break execution, so an observer that raises + is reported and ignored. + + Args: + run: The run the transitions belong to. + events: The (type, data) pairs just recorded. + """ + if self._observer is None: + return + try: + for event_type, data in events: + self._observer.on_event(event_type, run.run_id, run.workflow_id, data) + except Exception as err: + console.warn(f"Workflow observer raised, ignoring: {err}") + def _acquire_lease(self, claim: Claim) -> _Lease: """Register an in-flight claim and start renewing its lease. @@ -1518,6 +1602,7 @@ async def _execute_claim(self, claim: Claim) -> None: except StaleClaimError: await self._record_abandoned(claim, handler, "fenced_at_commit") return + self._notify(claim.run, completion.events) if completion.children: await self._admit_children(claim, completion) await self._report_to_parent(claim.run, completion) diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 2be092547f6..da8afb3b896 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -20,7 +20,11 @@ from reflex_base.workflow import DEFAULT_LEASE_DURATION, ChannelDelivery from reflex.workflow.definition import WorkflowDefinition, compile_workflow -from reflex.workflow.kernel import DEFAULT_POLL_INTERVAL, WorkflowKernel +from reflex.workflow.kernel import ( + DEFAULT_POLL_INTERVAL, + WorkflowKernel, + WorkflowObserver, +) from reflex.workflow.store import RunStore, SqliteRunStore if TYPE_CHECKING: @@ -77,6 +81,7 @@ def __init__( lease_duration: float = DEFAULT_LEASE_DURATION, lease_renew_interval: float | None = None, recovery_interval: float | None = None, + observer: WorkflowObserver | None = None, ): """Initialize the runtime. @@ -90,6 +95,7 @@ def __init__( recovery may reclaim it. lease_renew_interval: Real seconds between lease renewals. recovery_interval: Seconds between recovery sweeps. + observer: Receives every recorded run transition. """ self._store = store self._clock = clock @@ -98,6 +104,7 @@ def __init__( self._lease_duration = lease_duration self._lease_renew_interval = lease_renew_interval self._recovery_interval = recovery_interval + self._observer = observer self._definitions: dict[str, WorkflowDefinition] = {} self._classes: dict[type, str] = {} self._kernel: WorkflowKernel | None = None @@ -185,6 +192,7 @@ async def startup(self, *, start_worker: bool = True) -> None: lease_duration=self._lease_duration, lease_renew_interval=self._lease_renew_interval, recovery_interval=self._recovery_interval, + observer=self._observer, ) if start_worker: await self._kernel.start_worker() diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index 7a1a4064cc9..f8791b42706 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -16,6 +16,7 @@ from reflex_base.workflow import DEFAULT_LEASE_DURATION, parse_duration +from reflex.workflow.kernel import WorkflowObserver from reflex.workflow.runtime import WorkflowRuntime, _context_runtime from reflex.workflow.store import MemoryRunStore @@ -68,6 +69,7 @@ def __init__( start_time: float = DEFAULT_START_TIME, lease_duration: DurationLike = DEFAULT_LEASE_DURATION, lease_renew_interval: float | None = None, + observer: WorkflowObserver | None = None, ): """Initialize the harness. @@ -77,6 +79,7 @@ def __init__( start_time: Initial virtual time in epoch seconds. lease_duration: Virtual seconds a claim survives without renewal. lease_renew_interval: Real seconds between lease renewals. + observer: Receives every recorded run transition. """ self._clock = _VirtualClock(start_time) self._runtime = WorkflowRuntime( @@ -85,6 +88,7 @@ def __init__( rng=lambda: 1.0, lease_duration=parse_duration(lease_duration), lease_renew_interval=lease_renew_interval, + observer=observer, ) for workflow_cls in workflow_classes: self._runtime.register(workflow_cls) diff --git a/tests/units/workflow/test_observability.py b/tests/units/workflow/test_observability.py new file mode 100644 index 00000000000..926c3829d07 --- /dev/null +++ b/tests/units/workflow/test_observability.py @@ -0,0 +1,147 @@ +"""Tests for the workflow observer hook.""" + +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.kernel import LoggingObserver, WorkflowObserver +from reflex.workflow.records import HistoryEventType, RunStatus +from reflex.workflow.testing import WorkflowTestHarness + + +class _Collector(WorkflowObserver): + """Records every transition it is handed.""" + + def __init__(self): + self.events: list[tuple[str, str, dict]] = [] + + def on_event(self, event_type, run_id, workflow_id, data): + """Record one transition. + + Args: + event_type: What happened. + run_id: The run it happened to. + workflow_id: The workflow identity. + data: The event payload. + """ + self.events.append((workflow_id, event_type.value, data)) + + +class _Exploding(WorkflowObserver): + """An observer that always raises, to prove it cannot break a run.""" + + def on_event(self, event_type, run_id, workflow_id, data): + """Always fail. + + Args: + event_type: What happened. + run_id: The run it happened to. + workflow_id: The workflow identity. + data: The event payload. + + Raises: + RuntimeError: Always. + """ + msg = "instrumentation is broken" + raise RuntimeError(msg) + + +def _flow(): + """Build a workflow that succeeds after one transient failure. + + Returns: + The workflow class. + """ + calls: list[int] = [] + + class ObservedFlow(rx.State): + __workflow__ = WorkflowConfig(id="obs.observed") + n: int = 0 + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=3, initial_delay="1s", jitter="none"), + ) + def go(self): + """Do the work, failing once. + + Returns: + Completion once the work succeeds. + """ + calls.append(1) + if len(calls) < 2: + msg = "flaky" + raise TransientWorkflowError(msg) + self.n = 1 + return rx.complete(result={"ok": True}) + + return ObservedFlow + + +async def test_observer_sees_the_whole_lifecycle(forked_registration_context): + collector = _Collector() + flow = _flow() + async with WorkflowTestHarness(flow, observer=collector) as harness: + result = await harness.start(flow.go) + assert result.run_id is not None + await harness.advance("1s") + + kinds = [kind for _, kind, _ in collector.events] + assert HistoryEventType.RUN_ADMITTED.value in kinds + assert HistoryEventType.ATTEMPT_FAILED.value in kinds + assert HistoryEventType.STEP_RETRY_SCHEDULED.value in kinds + assert HistoryEventType.RUN_COMPLETED.value in kinds + + +async def test_every_event_is_correlated(forked_registration_context): + collector = _Collector() + flow = _flow() + async with WorkflowTestHarness(flow, observer=collector) as harness: + result = await harness.start(flow.go) + await harness.advance("1s") + + assert {workflow_id for workflow_id, _, _ in collector.events} == {"obs.observed"} + failures = [ + data + for _, kind, data in collector.events + if kind == HistoryEventType.ATTEMPT_FAILED.value + ] + assert failures + assert failures[0]["error"]["type"] == "TransientWorkflowError" + assert "ordinal" in failures[0] + + +async def test_a_broken_observer_cannot_break_a_run(forked_registration_context): + """Instrumentation failures are reported, never propagated.""" + flow = _flow() + async with WorkflowTestHarness(flow, observer=_Exploding()) as harness: + result = await harness.start(flow.go) + assert result.run_id is not None + await harness.advance("1s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + + +async def test_the_default_is_no_observer(forked_registration_context): + flow = _flow() + async with WorkflowTestHarness(flow) as harness: + result = await harness.start(flow.go) + assert result.run_id is not None + await harness.advance("1s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + + +async def test_logging_observer_runs(forked_registration_context): + """The bundled observer logs without disturbing execution.""" + flow = _flow() + async with WorkflowTestHarness(flow, observer=LoggingObserver()) as harness: + result = await harness.start(flow.go) + assert result.run_id is not None + await harness.advance("1s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED From ac6ade3fbc5d7145abfec01786ab083e83d8eb80 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 15:20:49 -0700 Subject: [PATCH 017/121] Fix three defects found by adversarial audit An audit of the whole feature reproduced four real problems, each of which could stop a run permanently. A handler that raised CancelledError itself -- as any handler wrapping its own asyncio work might -- propagated past the cancellation branch and killed the worker task, so every later run in the process silently never executed. The kernel now distinguishes a task that was cancelled from a coroutine that raised: the first is a control signal, the second is an ordinary handler failure that retries. A crash between a fan-out's commit and the creation of its children left a join blocked on children that did not exist, with nothing to recover it; the original comment claiming recovery handled this was simply wrong. Children are now inserted in the same transaction as the join slot, so the window is gone rather than merely narrowed. A child that was cancelled or blew its run deadline never reported to its parent's join, because reporting only happened on commit and those paths finalize without one. Both now report, so a join can no longer wait forever on a child that already stopped. Singleton with mode='cancel' left the superseded run in CANCELLING, so a burst of starts produced several simultaneously active runs under one key. The replacement now waits for the old run to reach a terminal state first. Also fixes a store divergence the conformance suite missed: the memory store kept only the most recent buffered signal per wait key while SQLite queued them, so a second early signal was silently dropped. Two conformance checks now cover buffered-delivery ordering and children being created with their join. --- news/workflow-audit-fixes.bugfix.md | 1 + reflex/workflow/conformance.py | 67 +++++++ reflex/workflow/kernel.py | 226 ++++++++++++++++------ reflex/workflow/store.py | 24 ++- tests/units/workflow/test_flow_control.py | 12 ++ tests/units/workflow/test_parallel.py | 81 ++++++++ 6 files changed, 347 insertions(+), 64 deletions(-) create mode 100644 news/workflow-audit-fixes.bugfix.md diff --git a/news/workflow-audit-fixes.bugfix.md b/news/workflow-audit-fixes.bugfix.md new file mode 100644 index 00000000000..d12bab590ab --- /dev/null +++ b/news/workflow-audit-fixes.bugfix.md @@ -0,0 +1 @@ +Fixes three defects found by an adversarial audit: a handler raising `CancelledError` killed the worker, a crash between a fan-out commit and child creation wedged the parent, and a cancelled or timed-out child never reported to its join. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 2622817ae94..a94ad250115 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -509,6 +509,71 @@ async def check_flow_control_queries(store: RunStore) -> None: assert steps[0].due_at == pytest.approx(NOW + 30) +async def check_early_deliveries_queue_in_order(store: RunStore) -> None: + """Several signals arriving before a wait are kept, not overwritten.""" + await store.admit(make_run(), make_step(), _ADMITTED) + assert await store.deliver("run1", "sig:ping", "d1", {"v": 1}, NOW) == "buffered" + assert await store.deliver("run1", "sig:ping", "d2", {"v": 2}, NOW) == "buffered" + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.WAITING, + state={"n": 1}, + new_steps=( + make_step( + ordinal=1, + status=StepStatus.BLOCKED, + wait_key="sig:ping", + due_at=0.0, + origin="wait", + ), + ), + next_ordinal=2, + ), + NOW, + ) + steps = await store.get_steps("run1") + # The first signal to arrive is the one that resolves the wait. + assert steps[1].args["__payload__"] == {"v": 1} + + +async def check_children_are_created_with_their_join(store: RunStore) -> None: + """A fan-out commit creates the join slot and its children together.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + child = make_run("child1", parent_run_id="run1", parent_ordinal=1) + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.WAITING, + state={"n": 1}, + new_steps=( + make_step( + ordinal=1, + status=StepStatus.BLOCKED, + wait_key="join:1", + join_expected=1, + origin="join", + due_at=0.0, + ), + ), + next_ordinal=2, + children=((child, make_step("child1")),), + ), + NOW, + ) + created = await store.get_run("child1") + assert created is not None + assert created.parent_run_id == "run1" + assert created.parent_ordinal == 1 + assert len(await store.get_steps("child1")) == 1 + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_admit_deduplicates_on_request_key, @@ -526,6 +591,8 @@ async def check_flow_control_queries(store: RunStore) -> None: check_delivery_never_touches_run_state, check_duplicate_deliveries_are_ignored, check_an_early_delivery_is_buffered_then_consumed, + check_early_deliveries_queue_in_order, + check_children_are_created_with_their_join, check_join_arrivals_count_once, check_finalize_refuses_while_a_step_is_claimed, check_finalize_tombstones_open_slots, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 55b7c357cbb..8910daf4307 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -75,6 +75,17 @@ MAX_SCHEDULE_CATCHUP = 10 +class _HandlerCancelledError(Exception): + """A handler raised CancelledError itself instead of being cancelled.""" + + def __init__(self): + """Describe the failure for the recorded error payload.""" + super().__init__( + "handler raised CancelledError; a durable handler must not cancel " + "itself, and must let cancellation propagate rather than raising it" + ) + + def _error_payload(error: BaseException) -> dict[str, Any]: """Build a JSON-compatible error payload from an exception. @@ -304,6 +315,7 @@ def __init__( self._worker_id = uuid.uuid4().hex self._observer = observer self._wakeup = asyncio.Event() + self._admission = asyncio.Lock() self._worker: asyncio.Task | None = None @property @@ -440,7 +452,11 @@ async def _apply_start_policy( StartResult(disposition="skipped", run_id=existing.run_id), now, ) + # Drive the cancellation to a terminal state before admitting the + # replacement, so "one active run per key" holds at every instant + # rather than only once a worker happens to drain the old one. await self.cancel(existing.run_id) + await self._finalize_control(now) if handler.rate_limit is not None: window = parse_duration(handler.rate_limit.period) started = await self._store.count_started_since( @@ -515,15 +531,48 @@ async def start( f"starting through this path requires {expected}." ) raise WorkflowRuntimeError(msg) - now = self._clock() flow_key = self._flow_key(handler, payload) - due_at = now - if flow_key is not None: + if flow_key is None: + return await self._admit(defn, handler, payload, request_key, labels, None) + # A start policy is a check followed by an insert, so concurrent starts + # must not interleave between them or two runs slip past a singleton. + async with self._admission: + now = self._clock() decided, due_at = await self._apply_start_policy( defn, handler, flow_key, now ) if decided is not None: return decided + return await self._admit( + defn, handler, payload, request_key, labels, flow_key, due_at + ) + + async def _admit( + self, + defn: WorkflowDefinition, + handler: HandlerDefinition, + payload: dict[str, Any], + request_key: str | None, + labels: dict[str, str] | None, + flow_key: str | None, + due_at: float | None = None, + ) -> StartResult: + """Create the run and its root slot. + + Args: + defn: The workflow definition. + handler: The root handler. + payload: The decoded start payload. + request_key: Idempotent admission key. + labels: Server-derived indexing labels. + flow_key: Start-policy grouping key, if the root declares one. + due_at: Earliest start time, when a policy delayed it. + + Returns: + The admission result. + """ + now = self._clock() + due_at = now if due_at is None else due_at run_id = uuid.uuid4().hex run = RunRecord( run_id=run_id, @@ -1141,6 +1190,9 @@ def _success_completion( events=tuple(events), ) if isinstance(control, Parallel): + children = self._child_records( + claim, control.branches, claim.run.next_ordinal, now + ) join = StepRecord( run_id=claim.run.run_id, ordinal=claim.run.next_ordinal, @@ -1164,8 +1216,7 @@ def _success_completion( new_steps=(join,), next_ordinal=claim.run.next_ordinal + 1, events=tuple(events), - children=tuple(control.branches), - join_ordinal=join.ordinal, + children=children, ) if isinstance(control, WaitFor): resume = self._resolve_successor(defn, control.then) @@ -1572,6 +1623,20 @@ async def _execute_claim(self, claim: Claim) -> None: defn, claim, steps, state, successors, control, self._clock() ) except asyncio.CancelledError: + if lease.attempt is not None and not lease.attempt.cancelled(): + # The handler raised CancelledError itself rather than being + # cancelled; that is a handler failure, not a control signal. + completion = self._failure_completion( + defn, + handler, + claim, + steps, + _HandlerCancelledError(), + timed_out=False, + now=self._clock(), + ) + await self._commit_outcome(claim, handler, completion) + return if lease.lost: await self._record_abandoned(claim, handler, "lease_lost") return @@ -1597,15 +1662,7 @@ async def _execute_claim(self, claim: Claim) -> None: completion = self._failure_completion( defn, handler, claim, steps, err, timed_out=False, now=self._clock() ) - try: - await self._store.commit(claim, completion, self._clock()) - except StaleClaimError: - await self._record_abandoned(claim, handler, "fenced_at_commit") - return - self._notify(claim.run, completion.events) - if completion.children: - await self._admit_children(claim, completion) - await self._report_to_parent(claim.run, completion) + await self._commit_outcome(claim, handler, completion) async def _admit_due_schedules(self, now: float) -> int: """Admit a run for every schedule occurrence that has come due. @@ -1652,20 +1709,30 @@ def _next_schedule_due(self, now: float) -> float | None: ] return min(upcoming) if upcoming else None - async def _admit_children(self, claim: Claim, completion: StepCompletion) -> None: - """Create the child runs a fan-out commit declared. + def _child_records( + self, + claim: Claim, + branches: tuple[Any, ...], + join_ordinal: int, + now: float, + ) -> tuple[tuple[RunRecord, StepRecord], ...]: + """Build the child runs a fan-out will create. - Children are admitted after the parent's commit lands, so a crash in - between leaves a join slot with no children, which recovery re-runs - rather than a set of orphans with no parent. + They are built here, not admitted separately afterwards, so the store + can insert them in the same transaction as the join slot: a crash can + never leave a join waiting on children that were never created. Args: claim: The parent's claim. - completion: The committed outcome carrying the branches. + branches: The root events to run concurrently. + join_ordinal: The join slot the children report to. + now: Current time in epoch seconds. + + Returns: + Each child run paired with its root slot. """ - now = self._clock() records = [] - for index, branch in enumerate(completion.children): + for index, branch in enumerate(branches): defn, handler, payload = self._resolve_target(branch) child_id = uuid.uuid4().hex records.append(( @@ -1678,10 +1745,8 @@ async def _admit_children(self, claim: Claim, completion: StepCompletion) -> Non state_version=0, next_ordinal=1, parent_run_id=claim.run.run_id, - parent_ordinal=completion.join_ordinal, - request_key=( - f"child:{claim.run.run_id}:{completion.join_ordinal}:{index}" - ), + parent_ordinal=join_ordinal, + request_key=f"child:{claim.run.run_id}:{join_ordinal}:{index}", deadline=( now + defn.run_timeout if defn.run_timeout is not None else None ), @@ -1699,8 +1764,27 @@ async def _admit_children(self, claim: Claim, completion: StepCompletion) -> Non updated_at=now, ), )) - await self._store.admit_children(tuple(records), (), now) - self._wakeup.set() + return tuple(records) + + async def _commit_outcome( + self, claim: Claim, handler: HandlerDefinition, completion: StepCompletion + ) -> None: + """Commit an attempt's outcome and follow up on what it scheduled. + + Args: + claim: The claim being committed. + handler: The handler that ran. + completion: The outcome to apply. + """ + try: + await self._store.commit(claim, completion, self._clock()) + except StaleClaimError: + await self._record_abandoned(claim, handler, "fenced_at_commit") + return + self._notify(claim.run, completion.events) + if completion.children: + self._wakeup.set() + await self._report_to_parent(claim.run, completion) async def _report_to_parent( self, run: RunRecord, completion: StepCompletion @@ -1711,24 +1795,76 @@ async def _report_to_parent( run: The child run, which may have no parent. completion: The committed outcome that finished it. """ - if run.parent_run_id is None or run.parent_ordinal is None: - return if completion.run_status not in TERMINAL_RUN_STATUSES: return + await self._report_outcome( + run, completion.run_status, completion.result, completion.run_error + ) + + async def _report_outcome( + self, + run: RunRecord, + status: RunStatus, + result: Any, + error: dict[str, Any] | None, + ) -> None: + """Tell a parent's join that this child has finished, however it finished. + + Cancellation and run deadlines terminate a child without a commit, so + this is called from both paths: a join must never wait forever on a + child that has already stopped. + + Args: + run: The child run, which may have no parent. + status: Its terminal status. + result: Its result, if any. + error: Its error, if any. + """ + if run.parent_run_id is None or run.parent_ordinal is None: + return await self._store.record_arrival( run.parent_run_id, run.parent_ordinal, { "run_id": run.run_id, - "status": completion.run_status.value, - "result": completion.result, - "error": completion.run_error, + "status": status.value, + "result": result, + "error": error, }, run.run_id, self._clock(), ) self._wakeup.set() + async def _finalize_control(self, now: float) -> int: + """Finalize every drained run awaiting a control transition. + + Args: + now: Current time in epoch seconds. + + Returns: + How many runs were finalized. + """ + finalized = 0 + for run in await self._store.control_pending(now): + cancelled = run.cancel_requested + status = RunStatus.CANCELLED if cancelled else RunStatus.TIMED_OUT + error = None if cancelled else {"reason": "run_timeout"} + if await self._store.finalize_run( + run.run_id, + status=status, + error=error, + event=( + HistoryEventType.RUN_CANCELLED + if cancelled + else HistoryEventType.RUN_TIMED_OUT + ), + now=now, + ): + await self._report_outcome(run, status, None, error) + finalized += 1 + return finalized + async def _tick(self) -> bool: """Run one scheduling round. @@ -1737,29 +1873,7 @@ async def _tick(self) -> bool: """ now = self._clock() progressed = await self._admit_due_schedules(now) > 0 - for run in await self._store.control_pending(now): - if run.cancel_requested: - progressed = ( - await self._store.finalize_run( - run.run_id, - status=RunStatus.CANCELLED, - error=None, - event=HistoryEventType.RUN_CANCELLED, - now=now, - ) - or progressed - ) - else: - progressed = ( - await self._store.finalize_run( - run.run_id, - status=RunStatus.TIMED_OUT, - error={"reason": "run_timeout"}, - event=HistoryEventType.RUN_TIMED_OUT, - now=now, - ) - or progressed - ) + progressed = await self._finalize_control(now) > 0 or progressed claim = await self._store.claim_next(now, lease_duration=self._lease_duration) if claim is not None: task = asyncio.ensure_future(self._execute_claim(claim)) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index dfd7454769f..7a06af49b96 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -88,8 +88,8 @@ class StepCompletion: tombstones: Ordinals of unresolved slots to cancel. next_ordinal: Updated mailbox allocation counter, if slots were added. events: History events to append, in order, as (type, data) pairs. - children: Root events to admit as child runs once this commit lands. - join_ordinal: The join slot those children report back to. + children: Child runs to create in the same transaction as this commit, + each paired with its root slot. """ step_status: StepStatus @@ -104,8 +104,7 @@ class StepCompletion: tombstones: tuple[int, ...] = () next_ordinal: int | None = None events: tuple[tuple[HistoryEventType, dict[str, Any]], ...] = () - children: tuple[Any, ...] = () - join_ordinal: int | None = None + children: tuple[tuple[RunRecord, StepRecord], ...] = () class RunStore(Protocol): @@ -568,7 +567,7 @@ def __init__(self): self._history: dict[str, list[HistoryEvent]] = {} self._dedupe: dict[tuple[str, str], str] = {} self._inbox: dict[str, dict[tuple[str, str, str], bool]] = {} - self._pending: dict[str, dict[str, dict[str, Any]]] = {} + self._pending: dict[str, dict[str, list[dict[str, Any]]]] = {} def _append_events( self, @@ -733,9 +732,10 @@ def _arm(self, step: StepRecord, now: float) -> StepRecord: """ if step.status is not StepStatus.BLOCKED or step.wait_key is None: return step - buffered = self._pending.get(step.run_id, {}).pop(step.wait_key, None) - if buffered is None: + queued = self._pending.get(step.run_id, {}).get(step.wait_key) + if not queued: return step + buffered = queued.pop(0) return dataclasses.replace( step, status=StepStatus.READY, @@ -774,6 +774,9 @@ async def commit( ) for new_step in completion.new_steps: steps.append(self._arm(new_step, now)) + for child_run, child_step in completion.children: + self._runs[child_run.run_id] = child_run + self._steps[child_run.run_id] = [child_step] self._runs[run.run_id] = dataclasses.replace( run, status=completion.run_status, @@ -891,7 +894,9 @@ async def deliver( now, ) return "resolved" - self._pending.setdefault(run_id, {})[wait_key] = payload + self._pending.setdefault(run_id, {}).setdefault(wait_key, []).append( + payload + ) self._append_events( run_id, ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), @@ -1910,6 +1915,9 @@ async def commit( ) for step in completion.new_steps: self._insert_step(self._arm_sql(step, now)) + for child_run, child_step in completion.children: + self._insert_run(child_run) + self._insert_step(child_step) self._db.execute( "UPDATE workflow_runs SET status = ?," " state = CASE WHEN ? THEN ? ELSE state END," diff --git a/tests/units/workflow/test_flow_control.py b/tests/units/workflow/test_flow_control.py index 2cac1ed1f59..5c9b98acf80 100644 --- a/tests/units/workflow/test_flow_control.py +++ b/tests/units/workflow/test_flow_control.py @@ -262,3 +262,15 @@ def click(self): "throttle": 200, "debounce": 100, } + + +async def test_singleton_cancel_keeps_one_active_run_per_key( + forked_registration_context, +): + """Replacing the active run must not leave both of them active.""" + flow = _singleton_flow(mode="cancel") + async with WorkflowTestHarness(flow) as harness: + for _ in range(3): + await harness.kernel.start(flow.start("acme")) + active = await harness.kernel.store.count_active("flow.sync", "start:'acme'") + assert active == 1 diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 3f0d9a7c895..c98d21ee136 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -207,3 +207,84 @@ async def test_duplicate_arrivals_are_counted_once(forked_registration_context): snapshot = await harness.get_run(result.run_id) assert snapshot is not None assert snapshot.steps[1].join_arrived == 1 + + +async def test_a_cancelled_child_reports_to_its_join(forked_registration_context): + """A child that is cancelled must not leave its parent waiting forever.""" + BRANCH_CALLS.clear() + + class SlowBranch(rx.State): + __workflow__ = WorkflowConfig(id="fan.slow") + n: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, lead: str): + """Wait a long time before finishing. + + Args: + lead: The lead identifier. + + Returns: + A far-future continuation. + """ + return rx.after("30d", SlowBranch.later) + + @rx.event(durable=True, effect="none") + def later(self): + """Never reached in this test.""" + + router = _router(Enrich, SlowBranch) + async with WorkflowTestHarness(router, Enrich, SlowBranch) as harness: + result = await harness.start(router.begin("acme")) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + + runs = await harness.kernel.list_runs() + slow = next( + run + for run in runs + if run.parent_run_id == result.run_id and run.workflow_id == "fan.slow" + ) + await harness.cancel(slow.run_id) + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state["outcomes"] == ["CANCELLED", "COMPLETED"] + + +async def test_a_timed_out_child_reports_to_its_join(forked_registration_context): + """A child that blows its run deadline still reports to the join.""" + BRANCH_CALLS.clear() + + class ExpiringBranch(rx.State): + __workflow__ = WorkflowConfig(id="fan.expiring", run_timeout="1h") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, lead: str): + """Wait past the run deadline. + + Args: + lead: The lead identifier. + + Returns: + A continuation scheduled after the deadline. + """ + return rx.after("2h", ExpiringBranch.later) + + @rx.event(durable=True, effect="none") + def later(self): + """Never reached in this test.""" + + router = _router(Enrich, ExpiringBranch) + async with WorkflowTestHarness(router, Enrich, ExpiringBranch) as harness: + result = await harness.start(router.begin("acme")) + assert result.run_id is not None + await harness.advance("2h") + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state["outcomes"] == ["COMPLETED", "TIMED_OUT"] From a7c0b151991cc2de1ad7f9ed83abc70e8422eaa5 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 16:00:17 -0700 Subject: [PATCH 018/121] Fix four defects a deeper audit reproduced A second, adversarial pass over the whole feature reproduced four problems, one of which was a fix from the previous commit that did not actually work. The guard added for 'a handler that raises CancelledError kills the worker' was a branch that could never be true: asyncio marks a task cancelled whether the kernel cancelled it or the coroutine let CancelledError escape, so the task's own flag cannot tell them apart. Verified directly -- both cases report cancelled() is True -- so the worker still died and every later run in the process silently never ran. The kernel now discriminates on its own control signals plus whether cancellation was requested on the executing task: a handler that raises is an ordinary failure, while a real shutdown still leaves the step claimed for lease recovery. The worker loop no longer dies on any exception, and a dead worker can be replaced. Admission dedupe ran after start policies, so a provider redelivering an event was judged as a new start: with singleton(mode='cancel') the redelivery cancelled the very run it deduplicated to, and answered the provider 202. With debounce, a provider retrying one event faster than the window starved it forever. request_key is now resolved before any policy. rx.fail(details=...) and rx.needs_attention(details=...) passed user values straight to the store, so a datetime raised inside the transaction recording the failure and left the run stuck RUNNING. Details are normalized first, and an unserializable value is recorded as its repr rather than losing the failure. timeout= on a synchronous handler was a lie: asyncio.wait_for cancels the wrapper while the thread runs on, so a timed-out step kept executing and its retries ran concurrently with it. It is now a compile error naming the fix. --- docs/workflows/overview.md | 5 +- news/workflow-audit-fixes-2.bugfix.md | 1 + .../src/reflex_base/event/__init__.py | 12 ++ reflex/workflow/kernel.py | 81 ++++++--- reflex/workflow/store.py | 54 ++++++ .../reflex_base/event/test_durable_event.py | 2 +- .../units/workflow/test_audit_regressions.py | 160 ++++++++++++++++++ tests/units/workflow/test_end_to_end.py | 2 +- tests/units/workflow/test_kernel.py | 2 +- tests/units/workflow/test_versioning.py | 2 +- 10 files changed, 295 insertions(+), 26 deletions(-) create mode 100644 news/workflow-audit-fixes-2.bugfix.md create mode 100644 tests/units/workflow/test_audit_regressions.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index d22785762ec..fae727cb14f 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -107,8 +107,9 @@ async def charge(self): ... Failures retry with exponential backoff by default. Narrow that with `rx.Retry(do_not_retry_on=(ValueError,))` when a specific error should fail fast. `timeout` bounds a -single attempt. `on_failure` and `on_timeout` name a handler on the same class that runs once the -step is finally out of attempts. +single attempt, and requires an `async def` handler: a synchronous one runs on a worker thread that +cannot be interrupted, so the timeout would fire while the body kept running. `on_failure` and +`on_timeout` name a handler on the same class that runs once the step is finally out of attempts. Backoff is a persisted timer, not a sleep, so a retry scheduled for tomorrow survives a deploy tonight. diff --git a/news/workflow-audit-fixes-2.bugfix.md b/news/workflow-audit-fixes-2.bugfix.md new file mode 100644 index 00000000000..2d670e9272c --- /dev/null +++ b/news/workflow-audit-fixes-2.bugfix.md @@ -0,0 +1 @@ +Fixes four more defects found by adversarial audit, including a handler that let `CancelledError` escape killing the worker, and a redelivered webhook cancelling the very run it deduplicated to. diff --git a/packages/reflex-base/src/reflex_base/event/__init__.py b/packages/reflex-base/src/reflex_base/event/__init__.py index 93185801199..f74819c54f2 100644 --- a/packages/reflex-base/src/reflex_base/event/__init__.py +++ b/packages/reflex-base/src/reflex_base/event/__init__.py @@ -3073,6 +3073,18 @@ def wrapper( "generators; return successor events instead of yielding." ) raise WorkflowDefinitionError(msg) + if ( + durable_config.timeout is not None + and not inspect.iscoroutinefunction(func) + ): + msg = ( + "timeout= cannot bound a synchronous handler: it runs on " + "a worker thread that cannot be interrupted, so the " + "timeout would fire while the body kept running. Declare " + "the handler 'async def', or drop timeout= and bound the " + "work inside it." + ) + raise WorkflowDefinitionError(msg) setattr(func, workflow.DURABLE_EVENT_MARKER, durable_config) if getattr(func, "__name__", "").startswith("_"): msg = "Event handlers cannot be private." diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 8910daf4307..0ce8307c59a 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -316,6 +316,7 @@ def __init__( self._observer = observer self._wakeup = asyncio.Event() self._admission = asyncio.Lock() + self._closing = False self._worker: asyncio.Task | None = None @property @@ -531,6 +532,15 @@ async def start( f"starting through this path requires {expected}." ) raise WorkflowRuntimeError(msg) + if request_key is not None: + # Dedupe before any start policy: a redelivered event must return + # the run it already created, not be judged as a new start and + # cancel, throttle, or debounce that very run. + existing = await self._store.find_by_request_key( + defn.workflow_id, request_key + ) + if existing is not None: + return StartResult(disposition="deduplicated", run_id=existing) flow_key = self._flow_key(handler, payload) if flow_key is None: return await self._admit(defn, handler, payload, request_key, labels, None) @@ -1134,6 +1144,27 @@ def _failure_completion( now=now, ) + def _control_error(self, reason: str, details: Any) -> dict[str, Any]: + """Build a durable error payload from a control return. + + ``details`` comes from user code, so it is normalized here rather than + at commit: an unserializable value must not break the transaction that + is recording the failure. + + Args: + reason: The stable failure reason. + details: Whatever the handler attached. + + Returns: + A JSON-compatible error payload. + """ + if details is None: + return {"reason": reason, "details": None} + try: + return {"reason": reason, "details": to_run_data(details)} + except (TypeError, ValueError): + return {"reason": reason, "details": {"unserializable": repr(details)}} + def _success_completion( self, defn: WorkflowDefinition, @@ -1162,7 +1193,7 @@ def _success_completion( (HistoryEventType.ATTEMPT_SUCCEEDED, {"ordinal": claim.step.ordinal}) ] if isinstance(control, FailRun): - error = {"reason": control.reason, "details": control.details} + error = self._control_error(control.reason, control.details) tombstones = self._open_ordinals(steps, exclude=claim.step.ordinal) events.extend( (HistoryEventType.STEP_TOMBSTONED, {"ordinal": ordinal}) @@ -1178,7 +1209,7 @@ def _success_completion( events=tuple(events), ) if isinstance(control, NeedsAttention): - error = {"reason": control.reason, "details": control.details} + error = self._control_error(control.reason, control.details) events.append((HistoryEventType.RUN_NEEDS_ATTENTION, {"error": error})) # The attempt succeeded and its state is committed, but the step # holds the suspension so resuming knows where to pick back up. @@ -1623,20 +1654,10 @@ async def _execute_claim(self, claim: Claim) -> None: defn, claim, steps, state, successors, control, self._clock() ) except asyncio.CancelledError: - if lease.attempt is not None and not lease.attempt.cancelled(): - # The handler raised CancelledError itself rather than being - # cancelled; that is a handler failure, not a control signal. - completion = self._failure_completion( - defn, - handler, - claim, - steps, - _HandlerCancelledError(), - timed_out=False, - now=self._clock(), - ) - await self._commit_outcome(claim, handler, completion) - return + # asyncio marks a task cancelled whether we cancelled it or the + # handler let CancelledError escape, so the task's own flag cannot + # tell them apart. Discriminate on this kernel's control signals + # instead; anything else came from user code. if lease.lost: await self._record_abandoned(claim, handler, "lease_lost") return @@ -1653,7 +1674,23 @@ async def _execute_claim(self, claim: Claim) -> None: now=self._clock(), ) return - raise + current = asyncio.current_task() + if self._closing or (current is not None and current.cancelling()): + # Someone cancelled *us* -- worker shutdown or a supervising + # task -- so this is crash-equivalent: leave the step claimed + # for lease recovery rather than recording an outcome. + raise + completion = self._failure_completion( + defn, + handler, + claim, + steps, + _HandlerCancelledError(), + timed_out=False, + now=self._clock(), + ) + await self._commit_outcome(claim, handler, completion) + return except TimeoutError as err: completion = self._failure_completion( defn, handler, claim, steps, err, timed_out=True, now=self._clock() @@ -1931,14 +1968,17 @@ async def _worker_loop(self) -> None: self._wakeup.clear() with contextlib.suppress(TimeoutError): await asyncio.wait_for(self._wakeup.wait(), timeout=delay) - except Exception as err: - console.error(f"Workflow worker error, retrying: {err}") + except BaseException as err: + if self._closing: + raise + console.error(f"Workflow worker error, retrying: {err!r}") await asyncio.sleep(self._poll_interval) async def start_worker(self) -> None: """Start the background worker, which recovers expired claims as it runs.""" - if self._worker is not None: + if self._worker is not None and not self._worker.done(): return + self._closing = False await self.recover() self._worker = asyncio.create_task(self._worker_loop()) @@ -1951,6 +1991,7 @@ async def aclose(self) -> None: """ if self._worker is None: return + self._closing = True self._worker.cancel() with contextlib.suppress(asyncio.CancelledError): await self._worker diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 7a06af49b96..293a1012da0 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -444,6 +444,25 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: """ ... + async def find_by_request_key( + self, workflow_id: str, request_key: str + ) -> str | None: + """Find the run a request key already admitted, if any. + + Admission dedupe must be answerable before any start policy runs, so a + provider redelivering an event cannot be treated as a new start and + trip a singleton, throttle, or debounce against the run it should + simply return. + + Args: + workflow_id: The workflow identity. + request_key: The idempotent admission key. + + Returns: + The existing run id, or None when the key is unused. + """ + ... + async def get_run(self, run_id: str) -> RunRecord | None: """Load one run record. @@ -1270,6 +1289,21 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: matched.sort(key=lambda run: run.created_at, reverse=True) return tuple(matched[: query.limit]) + async def find_by_request_key( + self, workflow_id: str, request_key: str + ) -> str | None: + """Find the run a request key already admitted, if any. + + Args: + workflow_id: The workflow identity. + request_key: The idempotent admission key. + + Returns: + The existing run id, or None when the key is unused. + """ + async with self._lock: + return self._dedupe.get((workflow_id, request_key)) + async def get_run(self, run_id: str) -> RunRecord | None: """Load one run record. @@ -2597,6 +2631,26 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: ).fetchall() return tuple(_run_from_row(row) for row in rows) + async def find_by_request_key( + self, workflow_id: str, request_key: str + ) -> str | None: + """Find the run a request key already admitted, if any. + + Args: + workflow_id: The workflow identity. + request_key: The idempotent admission key. + + Returns: + The existing run id, or None when the key is unused. + """ + with self._lock: + row = self._db.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = ? AND request_key = ?", + (workflow_id, request_key), + ).fetchone() + return None if row is None else row["run_id"] + async def get_run(self, run_id: str) -> RunRecord | None: """Load one run record. diff --git a/tests/units/reflex_base/event/test_durable_event.py b/tests/units/reflex_base/event/test_durable_event.py index 339873fdbcc..14b6618fab8 100644 --- a/tests/units/reflex_base/event/test_durable_event.py +++ b/tests/units/reflex_base/event/test_durable_event.py @@ -18,7 +18,7 @@ class MarkerWorkflow(rx.State): effect="read", queue="integrations", ) - def begin(self): + async def begin(self): pass config = get_durable_config(MarkerWorkflow.event_handlers["begin"].fn) diff --git a/tests/units/workflow/test_audit_regressions.py b/tests/units/workflow/test_audit_regressions.py new file mode 100644 index 00000000000..bd9d375ae66 --- /dev/null +++ b/tests/units/workflow/test_audit_regressions.py @@ -0,0 +1,160 @@ +"""Regressions for defects an adversarial audit reproduced.""" + +import asyncio +import datetime as dt + +import pytest +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import Singleton, WorkflowConfig, after, fail, manual + +import reflex as rx +from reflex.workflow.records import RunStatus +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import MemoryRunStore +from reflex.workflow.testing import WorkflowTestHarness + + +async def test_a_handler_raising_cancellederror_does_not_kill_the_worker( + forked_registration_context, +): + """One rude handler must not stop every later run in the process. + + asyncio marks a task cancelled whether the kernel cancelled it or the + handler let CancelledError escape, so the kernel cannot discriminate on + the task's own flag; it discriminates on its control signals instead. + """ + + class Rude(rx.State): + __workflow__ = WorkflowConfig(id="audit.rude") + + @rx.event(durable=True, trigger=manual(), effect="none") + async def go(self): + """Let a CancelledError escape, as a handler wrapping its own work might.""" + inner = asyncio.ensure_future(asyncio.sleep(10)) + inner.cancel() + await inner + + class Healthy(rx.State): + __workflow__ = WorkflowConfig(id="audit.healthy") + n: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Do ordinary work.""" + self.n = 1 + + runtime = WorkflowRuntime(MemoryRunStore(), poll_interval=0.02) + runtime.register(Rude) + runtime.register(Healthy) + async with runtime.running(): + rude = await runtime.kernel.start(Rude.go) + assert rude.run_id is not None + await asyncio.sleep(0.4) + healthy = await runtime.kernel.start(Healthy.go) + assert healthy.run_id is not None + snapshot = None + for _ in range(100): + snapshot = await runtime.kernel.get_run(healthy.run_id) + if snapshot is not None and snapshot.status is RunStatus.COMPLETED: + break + await asyncio.sleep(0.02) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + + # The rude run is treated as an ordinary failure, not a control signal. + rude_snapshot = await runtime.kernel.get_run(rude.run_id) + assert rude_snapshot is not None + assert rude_snapshot.status is not RunStatus.COMPLETED + + +async def test_redelivery_dedupes_before_any_start_policy( + forked_registration_context, +): + """A provider retrying an event must not trip the policy against its own run.""" + + class Paid(rx.State): + __workflow__ = WorkflowConfig(id="audit.paid") + invoice: str = "" + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + singleton=Singleton(key="invoice", mode="cancel"), + ) + def on_paid(self, invoice: str): + """Begin work for an invoice. + + Args: + invoice: The invoice identifier. + + Returns: + A delayed continuation. + """ + self.invoice = invoice + return after("1h", Paid.later) + + @rx.event(durable=True, effect="none") + def later(self): + """Finish later.""" + + async with WorkflowTestHarness(Paid) as harness: + first = await harness.start(Paid.on_paid("inv_1"), request_key="evt_1") + assert first.run_id is not None + redelivery = await harness.kernel.start( + Paid.on_paid("inv_1"), request_key="evt_1" + ) + assert redelivery.disposition == "deduplicated" + assert redelivery.run_id == first.run_id + + # The run the redelivery deduplicated to must be untouched. + snapshot = await harness.get_run(first.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert len(await harness.kernel.list_runs()) == 1 + + +async def test_unserializable_failure_details_do_not_break_the_commit( + forked_registration_context, +): + """User-supplied details must not be able to break the recording commit.""" + + class Detailed(rx.State): + __workflow__ = WorkflowConfig(id="audit.detailed") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Fail with details that are not JSON. + + Returns: + A failure carrying an unserializable value. + """ + return fail("nope", details={"when": dt.datetime.now(tz=dt.UTC)}) + + async with WorkflowTestHarness(Detailed) as harness: + result = await harness.start(Detailed.go) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error is not None + assert snapshot.error["reason"] == "nope" + + +def test_timeout_is_rejected_on_a_synchronous_handler(): + """A thread cannot be interrupted, so timeout= would be a lie.""" + with pytest.raises(WorkflowDefinitionError, match="synchronous handler"): + + @rx.event(durable=True, trigger=manual(), effect="none", timeout="5s") + def handler(self): + pass + + +def test_timeout_is_allowed_on_an_async_handler(): + """An async handler really can be interrupted at an await point.""" + + @rx.event(durable=True, trigger=manual(), effect="none", timeout="5s") + async def handler(self): + await asyncio.sleep(0) + + assert handler is not None diff --git a/tests/units/workflow/test_end_to_end.py b/tests/units/workflow/test_end_to_end.py index 5d5196dad12..da2b83ad136 100644 --- a/tests/units/workflow/test_end_to_end.py +++ b/tests/units/workflow/test_end_to_end.py @@ -73,7 +73,7 @@ def start(self, invoice_id: str, amount: int): timeout="30s", on_failure="escalate", ) - def charge(self): + async def charge(self): """Charge the invoice, retrying a flaky gateway. Returns: diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py index 3a08fcb5f78..02aec1684e1 100644 --- a/tests/units/workflow/test_kernel.py +++ b/tests/units/workflow/test_kernel.py @@ -670,7 +670,7 @@ class PinnedV2(rx.State): __workflow__ = WorkflowConfig(id="kernel.pinned") @rx.event(durable=True, trigger=manual(), effect="read", timeout="5s") - def begin(self): + async def begin(self): return after("1h", PinnedV2.finish) @rx.event(durable=True, effect="read") diff --git a/tests/units/workflow/test_versioning.py b/tests/units/workflow/test_versioning.py index d3ea27ea277..211883d83ca 100644 --- a/tests/units/workflow/test_versioning.py +++ b/tests/units/workflow/test_versioning.py @@ -35,7 +35,7 @@ def begin(self): effect="read", timeout="90s" if slow_retry else None, ) - def finish(self): + async def finish(self): self.status = "done" return Deployed From 68bc1d86c8aaf9f82bb15f9dd1c6e39efd151233 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 16:04:40 -0700 Subject: [PATCH 019/121] Run every workflow test against the durable store too Every kernel-level test used the in-memory store, so the harness was certifying semantics production does not necessarily have: any divergence between the two stores could ship green, and one already had. Harness-based tests are now parametrised over both stores, doubling that coverage to 478 workflow tests. Turning it on immediately caught a real divergence: the memory store handed callers live references to the values it was storing, so mutating a returned run's state silently changed committed data -- something a database-backed store cannot do. Reads now detach their mutable payloads, and a conformance check pins it. It also exposed an order-dependent test of its own, which assumed the first child listed was the branch that had already reported to the join. This is the gap that let earlier store divergences through, so it goes in before any further features. --- news/workflow-double-store-tests.bugfix.md | 1 + reflex/workflow/conformance.py | 17 ++++++++ reflex/workflow/store.py | 48 ++++++++++++++++++++-- tests/units/workflow/conftest.py | 42 +++++++++++++++++++ tests/units/workflow/test_parallel.py | 9 +++- tests/units/workflow/test_versioning.py | 10 ++--- 6 files changed, 116 insertions(+), 11 deletions(-) create mode 100644 news/workflow-double-store-tests.bugfix.md create mode 100644 tests/units/workflow/conftest.py diff --git a/news/workflow-double-store-tests.bugfix.md b/news/workflow-double-store-tests.bugfix.md new file mode 100644 index 00000000000..5fb688ca32b --- /dev/null +++ b/news/workflow-double-store-tests.bugfix.md @@ -0,0 +1 @@ +Workflow behavior tests now run against both the in-memory and SQLite stores, and the in-memory store no longer hands out live references to committed state. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index a94ad250115..6c590be0dda 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -574,8 +574,25 @@ async def check_children_are_created_with_their_join(store: RunStore) -> None: assert len(await store.get_steps("child1")) == 1 +async def check_reads_do_not_alias_stored_state(store: RunStore) -> None: + """Mutating a returned record must not change what the store holds.""" + await store.admit(make_run(), make_step(), _ADMITTED) + run = await store.get_run("run1") + assert run is not None + run.state["n"] = 999 + steps = await store.get_steps("run1") + steps[0].args["injected"] = True + + again = await store.get_run("run1") + assert again is not None + assert again.state == {"n": 0} + fresh = await store.get_steps("run1") + assert fresh[0].args == {} + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, + check_reads_do_not_alias_stored_state, check_admit_deduplicates_on_request_key, check_only_the_frontier_is_claimable, check_commit_is_atomic, diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 293a1012da0..7352bacb415 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio +import copy import dataclasses import json import sqlite3 @@ -527,6 +528,42 @@ def _run_is_runnable(run: RunRecord, now: float) -> bool: ) +def _detach_run(run: RunRecord) -> RunRecord: + """Copy a run record's mutable payloads before handing it to a caller. + + The in-memory store would otherwise hand out live references to the values + it is storing, so a caller mutating a returned run's state would silently + change committed data -- behavior a database-backed store cannot have. + + Args: + run: The stored record. + + Returns: + A record that shares no mutable structure with the store. + """ + return dataclasses.replace( + run, + state=copy.deepcopy(run.state), + result=copy.deepcopy(run.result), + error=copy.deepcopy(run.error), + labels=copy.deepcopy(run.labels), + ) + + +def _detach_step(step: StepRecord) -> StepRecord: + """Copy a step record's mutable payloads before handing it to a caller. + + Args: + step: The stored record. + + Returns: + A record that shares no mutable structure with the store. + """ + return dataclasses.replace( + step, args=copy.deepcopy(step.args), error=copy.deepcopy(step.error) + ) + + def _matches_query(run: RunRecord, query: RunQuery) -> bool: """Whether a run satisfies every filter in a query. @@ -1038,7 +1075,9 @@ async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | Non and run.flow_key == flow_key and run.status not in TERMINAL_RUN_STATUSES ] - return min(active, key=lambda run: run.created_at) if active else None + if not active: + return None + return _detach_run(min(active, key=lambda run: run.created_at)) async def count_started_since( self, workflow_id: str, flow_key: str, since: float @@ -1287,7 +1326,7 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: async with self._lock: matched = [run for run in self._runs.values() if _matches_query(run, query)] matched.sort(key=lambda run: run.created_at, reverse=True) - return tuple(matched[: query.limit]) + return tuple(_detach_run(run) for run in matched[: query.limit]) async def find_by_request_key( self, workflow_id: str, request_key: str @@ -1314,7 +1353,8 @@ async def get_run(self, run_id: str) -> RunRecord | None: The record, or None if unknown. """ async with self._lock: - return self._runs.get(run_id) + run = self._runs.get(run_id) + return None if run is None else _detach_run(run) async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: """Load a run's mailbox slots in ordinal order. @@ -1326,7 +1366,7 @@ async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: The step records. """ async with self._lock: - return tuple(self._steps.get(run_id, ())) + return tuple(_detach_step(step) for step in self._steps.get(run_id, ())) async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: """Load a run's append-only history in sequence order. diff --git a/tests/units/workflow/conftest.py b/tests/units/workflow/conftest.py new file mode 100644 index 00000000000..f8b3e80805c --- /dev/null +++ b/tests/units/workflow/conftest.py @@ -0,0 +1,42 @@ +"""Run every harness-based workflow test against both store implementations. + +The test harness defaults to the in-memory store, so behavioral tests were +certifying semantics that production -- which persists to SQLite -- does not +necessarily have. Any divergence between the two stores could therefore ship +green. This fixture makes the default store a parameter, so every test in this +directory runs twice. +""" + +import pytest + +import reflex.workflow.testing as testing +from reflex.workflow.store import SqliteRunStore + + +@pytest.fixture(params=["memory", "sqlite"], autouse=True) +def harness_store(request, tmp_path, monkeypatch): + """Make the harness's default store a parameter of every test. + + Args: + request: The fixture request carrying the store kind. + tmp_path: Temporary directory for SQLite databases. + monkeypatch: Used to swap the harness's default store factory. + + Yields: + The store kind under test. + """ + opened: list[SqliteRunStore] = [] + + if request.param == "sqlite": + + def factory(): + store = SqliteRunStore(tmp_path / f"harness{len(opened)}.db") + opened.append(store) + return store + + monkeypatch.setattr(testing, "MemoryRunStore", factory) + + yield request.param + + for store in opened: + store.close() diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index c98d21ee136..10f73ebabe8 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -194,7 +194,14 @@ async def test_duplicate_arrivals_are_counted_once(forked_registration_context): await harness.kernel.run_until_idle() runs = await harness.kernel.list_runs() - child = next(run for run in runs if run.parent_run_id == result.run_id) + # The branch that already reported, not merely the first one listed: + # re-reporting a branch still in flight would be a new arrival. + child = next( + run + for run in runs + if run.parent_run_id == result.run_id + and run.status is RunStatus.COMPLETED + ) repeat = await harness.kernel.store.record_arrival( result.run_id, 1, diff --git a/tests/units/workflow/test_versioning.py b/tests/units/workflow/test_versioning.py index 211883d83ca..1a02d7ad090 100644 --- a/tests/units/workflow/test_versioning.py +++ b/tests/units/workflow/test_versioning.py @@ -139,16 +139,16 @@ def go(self): async def test_resume_reopens_a_suspended_run(forked_registration_context): """Resuming grants the frontier step a fresh attempt budget.""" attempts = [] + resolved = [] class ReviewFlow(rx.State): __workflow__ = WorkflowConfig(id="versioning.review") status: str = "pending" - resolved: bool = False @rx.event(durable=True, trigger=manual(), effect="none") def begin(self): attempts.append(1) - if not self.resolved: + if not resolved: return needs_attention("manual_review") self.status = "done" return None @@ -160,10 +160,8 @@ def begin(self): assert snapshot is not None assert snapshot.status is RunStatus.NEEDS_ATTENTION - # An operator fixes the cause, then resumes. - run = await harness.kernel.store.get_run(result.run_id) - assert run is not None - run.state["resolved"] = True + # An operator fixes the cause outside the run, then resumes it. + resolved.append(True) assert await harness.resume(result.run_id) snapshot = await harness.get_run(result.run_id) From ab9cecdb804910e52c25d52769d5222ccca716f4 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 16:28:59 -0700 Subject: [PATCH 020/121] Fix six more audit findings Working through the confirmed findings, highest severity first. A class with __workflow__ that was never passed to app.add_workflow() stayed a session substate, so its durable handlers -- including non_idempotent_write -- remained dispatchable from a browser, and the only feedback was a later error from rx.workflows.start(). Tying detachment to registration meant the omitted line, exactly the one a generator drops, left money-moving handlers exposed. Detachment now happens when the class is created, so registration only adds the definition to the kernel. A wait whose deadline had already fallen due still accepted its signal, because delivery matched on BLOCKED alone and never compared the deadline. A seven-day approval that expired three weeks ago would be approved, and a sender that lost the race was told 'buffered' and then refused as 'duplicate'. Delivery now refuses with a distinct 'expired' disposition, which also stops a stale signal resolving a later wait on the same channel. A child failed by exhausting its recovery budget never reported to its parent's join: that path fails the run inside the store, and recover() only returned a count. The parent, and every ancestor, waited forever. recover_orphans now returns the runs it failed so the kernel can report them. Start policies of the wrong type vanished instead of raising: debounce='30s' -- a very plausible generation given every other duration in the API is a string -- silently disabled debouncing, and singleton='cid' failed only in production once two runs overlapped. All four are now type-checked at decoration, discriminating the browser event action on int. Run pagination used created_at alone, so runs sharing a timestamp -- the shape every fan-out produces -- were silently skipped; the cursor is now (created_at, run_id). The SQLite label filter interpolated user-supplied keys into a JSON path expression and now matches them as values. --- news/workflow-audit-fixes-3.bugfix.md | 1 + .../reflex-base/src/reflex_base/registry.py | 22 +++ .../reflex-base/src/reflex_base/workflow.py | 24 ++- reflex/state.py | 7 + reflex/workflow/conformance.py | 43 ++++- reflex/workflow/kernel.py | 17 +- reflex/workflow/records.py | 8 +- reflex/workflow/runtime.py | 40 ++--- reflex/workflow/store.py | 74 ++++++-- reflex/workflow/testing.py | 9 +- tests/units/workflow/test_app.py | 10 +- .../units/workflow/test_audit_regressions.py | 167 ++++++++++++++++++ tests/units/workflow/test_lease.py | 6 +- tests/units/workflow/test_parallel.py | 3 +- tests/units/workflow/test_queries.py | 2 +- tests/units/workflow/test_store.py | 16 +- 16 files changed, 378 insertions(+), 71 deletions(-) create mode 100644 news/workflow-audit-fixes-3.bugfix.md diff --git a/news/workflow-audit-fixes-3.bugfix.md b/news/workflow-audit-fixes-3.bugfix.md new file mode 100644 index 00000000000..cfe22f8af44 --- /dev/null +++ b/news/workflow-audit-fixes-3.bugfix.md @@ -0,0 +1 @@ +Fixes six more audit findings: an unregistered workflow class stayed browser-reachable, an expired wait still accepted its signal, a child failed by recovery never reported to its join, wrong-typed start policies vanished silently, and run pagination could hide runs. diff --git a/packages/reflex-base/src/reflex_base/registry.py b/packages/reflex-base/src/reflex_base/registry.py index 04d7e258e92..5af4f7d6649 100644 --- a/packages/reflex-base/src/reflex_base/registry.py +++ b/packages/reflex-base/src/reflex_base/registry.py @@ -247,6 +247,28 @@ def _register_event_handler( ) return handler + @classmethod + def detach_workflow_state(cls, state_cls: type[BaseState]) -> None: + """Remove a workflow class from the session state tree. + + Workflow state is run-scoped, so it must not be instantiated per + browser session, compiled into the client state schema, or reachable + from frontend event dispatch. + + Args: + state_cls: The workflow class to detach. + """ + ctx = cls.ensure_context() + ctx.base_states.pop(state_cls.get_full_name(), None) + parent = state_cls.get_parent_state() + if parent is not None: + ctx.base_state_substates.get(parent.get_full_name(), set()).discard( + state_cls + ) + for full_name, registered in list(ctx.event_handlers.items()): + if state_cls in registered.states: + del ctx.event_handlers[full_name] + def get_substates( self, base_state_cls: type[BaseState] | str ) -> set[type[BaseState]]: diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index fb6dd1cf2c0..210205374f1 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -558,9 +558,27 @@ def build_durable_config( WorkflowDefinitionError: If the combination of arguments is invalid. """ # throttle= and debounce= carry either a browser event action (an int) or a - # durable start policy; only the policy objects are workflow options. - throttle = throttle if isinstance(throttle, Throttle) else None - debounce = debounce if isinstance(debounce, Debounce) else None + # durable start policy. Discriminate on the int, so a wrong type raises + # rather than vanishing: `debounce="30s"` is a very plausible mistake, and + # silently dropping it would leave the burst uncollapsed with no signal. + browser_throttle = isinstance(throttle, int) and not isinstance(throttle, bool) + browser_debounce = isinstance(debounce, int) and not isinstance(debounce, bool) + throttle = None if browser_throttle else throttle + debounce = None if browser_debounce else debounce + for name, value, expected in ( + ("singleton", singleton, Singleton), + ("rate_limit", rate_limit, RateLimit), + ("throttle", throttle, Throttle), + ("debounce", debounce, Debounce), + ): + if value is not None and not isinstance(value, expected): + article = "an" if expected.__name__[0] in "AEIOU" else "a" + msg = ( + f"@rx.event({name}=...) expects {article} rx.{expected.__name__}, " + f"got {type(value).__name__}. Write " + f"{name}=rx.{expected.__name__}(...)." + ) + raise WorkflowDefinitionError(msg) if not durable: offending = next( ( diff --git a/reflex/state.py b/reflex/state.py index bad54ff4e8e..7e0e14aff2c 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -713,6 +713,13 @@ def __init_subclass__(cls, mixin: bool = False, **kwargs): setattr(cls, name, handler) RegistrationContext.register_base_state(cls) + if "__workflow__" in cls.__dict__: + # A workflow class is run-scoped: it must never be instantiated per + # browser session or have its durable handlers reachable from the + # frontend. Detaching at class creation rather than at + # app.add_workflow() means forgetting to register one cannot leave + # its handlers dispatchable from a browser. + RegistrationContext.detach_workflow_state(cls) # Initialize per-class var dependency tracking. cls._var_dependencies = {} diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 6c590be0dda..83f2a3f0f15 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -249,7 +249,7 @@ async def check_recovery_spares_a_live_lease(store: RunStore) -> None: await store.admit(make_run(), make_step(), _ADMITTED) claim = await store.claim_next(NOW, lease_duration=LEASE) assert claim is not None - assert await store.recover_orphans(NOW + LEASE - 1, max_recoveries=10) == 0 + assert (await store.recover_orphans(NOW + LEASE - 1, max_recoveries=10))[0] == 0 steps = await store.get_steps("run1") assert steps[0].status is StepStatus.CLAIMED assert steps[0].recoveries == 0 @@ -260,7 +260,7 @@ async def check_recovery_reclaims_an_expired_lease(store: RunStore) -> None: await store.admit(make_run(), make_step(), _ADMITTED) claim = await store.claim_next(NOW, lease_duration=LEASE) assert claim is not None - assert await store.recover_orphans(NOW + LEASE, max_recoveries=10) == 1 + assert (await store.recover_orphans(NOW + LEASE, max_recoveries=10))[0] == 1 steps = await store.get_steps("run1") assert steps[0].status is StepStatus.RECOVERY_WAIT assert steps[0].recoveries == 1 @@ -590,6 +590,43 @@ async def check_reads_do_not_alias_stored_state(store: RunStore) -> None: assert fresh[0].args == {} +async def check_pagination_skips_nothing_on_tied_timestamps(store: RunStore) -> None: + """Runs sharing a created_at must all be reachable by paging. + + A fan-out stamps every child with the same time, so a cursor on time alone + silently hides runs from the operator surface. + """ + for index in range(4): + await store.admit( + make_run(f"run{index}", created_at=NOW), make_step(f"run{index}"), _ADMITTED + ) + seen: list[str] = [] + cursor = None + while True: + page = await store.list_runs(RunQuery(limit=2, created_before=cursor)) + if not page: + break + seen.extend(run.run_id for run in page) + cursor = (page[-1].created_at, page[-1].run_id) + assert sorted(seen) == ["run0", "run1", "run2", "run3"] + + +async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: + """A label key is data, never part of a query expression.""" + await store.admit( + make_run("a", labels={"team.name": "core", 'quoted"key': "yes"}), + make_step("a"), + _ADMITTED, + ) + await store.admit( + make_run("b", labels={"team.name": "other"}), make_step("b"), _ADMITTED + ) + dotted = await store.list_runs(RunQuery(labels={"team.name": "core"})) + assert [run.run_id for run in dotted] == ["a"] + quoted = await store.list_runs(RunQuery(labels={'quoted"key': "yes"})) + assert [run.run_id for run in quoted] == ["a"] + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -615,5 +652,7 @@ async def check_reads_do_not_alias_stored_state(store: RunStore) -> None: check_finalize_tombstones_open_slots, check_resume_only_reopens_a_suspended_run, check_list_runs_filters_and_orders, + check_pagination_skips_nothing_on_tied_timestamps, + check_label_filter_handles_awkward_keys, check_flow_control_queries, ) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 0ce8307c59a..bd5b66db2e9 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -697,7 +697,7 @@ async def list_runs( workflow_id: str | None = None, statuses: Iterable[RunStatus] = (), labels: Mapping[str, str] | None = None, - created_before: float | None = None, + created_before: tuple[float, str] | None = None, limit: int = 50, ) -> tuple[RunRecord, ...]: """List runs matching a filter, newest first. @@ -706,7 +706,8 @@ async def list_runs( workflow_id: Restrict to one workflow identity. statuses: Restrict to these run statuses; empty means any. labels: Require every one of these label values. - created_before: Pagination cursor; return runs admitted before this. + created_before: Pagination cursor, the (created_at, run_id) of the + previous page's last row. limit: Maximum runs to return. Returns: @@ -1935,7 +1936,17 @@ async def recover(self) -> int: await self._renew_leases() now = self._clock() self._next_recovery_at = now + self._recovery_interval - return await self._store.recover_orphans(now, self._max_recoveries) + recovered, failed = await self._store.recover_orphans(now, self._max_recoveries) + for run_id in failed: + run = await self._store.get_run(run_id) + if run is not None: + await self._report_outcome( + run, + RunStatus.FAILED, + None, + {"reason": "recovery_budget_exhausted"}, + ) + return recovered async def run_until_idle(self) -> None: """Process work until nothing is claimable at the current clock time. diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 2c927077d04..dee47dd7d4f 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -288,15 +288,17 @@ class RunQuery: workflow_id: Restrict to one workflow identity. statuses: Restrict to these run statuses; empty means any. labels: Require every one of these server-derived label values. - created_before: Return runs admitted strictly before this epoch time, - which is the pagination cursor. + created_before: Pagination cursor, as the ``(created_at, run_id)`` of + the last row of the previous page. A fan-out stamps every child + with the same time, so the run id breaks the tie and no run is + skipped. limit: Maximum runs to return, newest first. """ workflow_id: str | None = None statuses: tuple[RunStatus, ...] = () labels: Mapping[str, str] | None = None - created_before: float | None = None + created_before: tuple[float, str] | None = None limit: int = 50 diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index da8afb3b896..65ee34d1284 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -17,7 +17,11 @@ from reflex_base.registry import RegistrationContext from reflex_base.utils.exceptions import WorkflowDefinitionError, WorkflowRuntimeError -from reflex_base.workflow import DEFAULT_LEASE_DURATION, ChannelDelivery +from reflex_base.workflow import ( + DEFAULT_LEASE_DURATION, + DEFAULT_MAX_RECOVERIES, + ChannelDelivery, +) from reflex.workflow.definition import WorkflowDefinition, compile_workflow from reflex.workflow.kernel import ( @@ -45,29 +49,6 @@ _default_runtime: WorkflowRuntime | None = None -def _detach_from_session_registry(workflow_cls: type[BaseState]) -> None: - """Remove a workflow class from the session state registry. - - A registered workflow class is run-scoped: it must not be instantiated per - browser session, compiled into the client state schema, or reachable from - frontend event dispatch. Removing it from the registration context achieves - all three without changing the class itself. - - Args: - workflow_cls: The workflow class being registered. - """ - ctx = RegistrationContext.ensure_context() - ctx.base_states.pop(workflow_cls.get_full_name(), None) - parent = workflow_cls.get_parent_state() - if parent is not None: - ctx.base_state_substates.get(parent.get_full_name(), set()).discard( - workflow_cls - ) - for full_name, registered in list(ctx.event_handlers.items()): - if workflow_cls in registered.states: - del ctx.event_handlers[full_name] - - class WorkflowRuntime: """Owns the workflow definitions and kernel for one process.""" @@ -82,6 +63,7 @@ def __init__( lease_renew_interval: float | None = None, recovery_interval: float | None = None, observer: WorkflowObserver | None = None, + max_recoveries: int = DEFAULT_MAX_RECOVERIES, ): """Initialize the runtime. @@ -96,6 +78,7 @@ def __init__( lease_renew_interval: Real seconds between lease renewals. recovery_interval: Seconds between recovery sweeps. observer: Receives every recorded run transition. + max_recoveries: Infrastructure recovery budget per logical step. """ self._store = store self._clock = clock @@ -105,6 +88,7 @@ def __init__( self._lease_renew_interval = lease_renew_interval self._recovery_interval = recovery_interval self._observer = observer + self._max_recoveries = max_recoveries self._definitions: dict[str, WorkflowDefinition] = {} self._classes: dict[type, str] = {} self._kernel: WorkflowKernel | None = None @@ -142,7 +126,7 @@ def register(self, workflow_cls: type[BaseState]) -> WorkflowDefinition: raise WorkflowDefinitionError(msg) self._definitions[definition.workflow_id] = definition self._classes[workflow_cls] = definition.workflow_id - _detach_from_session_registry(workflow_cls) + RegistrationContext.detach_workflow_state(workflow_cls) return definition @property @@ -193,6 +177,7 @@ async def startup(self, *, start_worker: bool = True) -> None: lease_renew_interval=self._lease_renew_interval, recovery_interval=self._recovery_interval, observer=self._observer, + max_recoveries=self._max_recoveries, ) if start_worker: await self._kernel.start_worker() @@ -320,7 +305,7 @@ async def list_runs( workflow_id: str | None = None, statuses: Iterable[RunStatus] = (), labels: Mapping[str, str] | None = None, - created_before: float | None = None, + created_before: tuple[float, str] | None = None, limit: int = 50, ) -> tuple[RunRecord, ...]: """List runs matching a filter, newest first. @@ -329,7 +314,8 @@ async def list_runs( workflow_id: Restrict to one workflow identity. statuses: Restrict to these run statuses; empty means any. labels: Require every one of these label values. - created_before: Pagination cursor; return runs admitted before this. + created_before: Pagination cursor, the (created_at, run_id) of the + previous page's last row. limit: Maximum runs to return. Returns: diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 7352bacb415..09a123e7a1e 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -48,6 +48,7 @@ DeliveryDisposition = Literal[ "resolved", "counted", + "expired", "buffered", "duplicate", "unknown_run", @@ -416,7 +417,9 @@ async def resume_run(self, run_id: str, now: float) -> bool: """ ... - async def recover_orphans(self, now: float, max_recoveries: int) -> int: + async def recover_orphans( + self, now: float, max_recoveries: int + ) -> tuple[int, tuple[str, ...]]: """Recover claims whose lease has expired. A step is orphaned when it is CLAIMED and its lease lapsed at or before @@ -430,7 +433,9 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: max_recoveries: Recovery budget per logical step. Returns: - The number of steps transitioned. + How many steps were transitioned, and the ids of runs this pass + failed outright by exhausting their recovery budget -- their + parents still need to be told. """ ... @@ -578,12 +583,31 @@ def _matches_query(run: RunRecord, query: RunQuery) -> bool: return False if query.statuses and run.status not in query.statuses: return False - if query.created_before is not None and run.created_at >= query.created_before: + if query.created_before is not None and (run.created_at, run.run_id) >= ( + query.created_before + ): return False labels = run.labels or {} return all(labels.get(key) == value for key, value in (query.labels or {}).items()) +def _wait_expired(step: StepRecord, now: float) -> bool: + """Whether a blocked wait's deadline has already won its race. + + Once the deadline falls due the timeout branch owns the slot, so a late + delivery must be refused outright rather than buffered: buffering it would + let a signal for an expired wait resolve a later one on the same channel. + + Args: + step: The frontier slot. + now: Current time in epoch seconds. + + Returns: + True when the wait has already timed out. + """ + return step.status is StepStatus.BLOCKED and 0.0 < step.due_at <= now + + def _lease_expired(step: StepRecord, now: float) -> bool: """Whether a claimed step's lease has lapsed and it may be recovered. @@ -927,6 +951,8 @@ async def deliver( inbox[run_id, wait_key, dedupe_key] = True steps = self._steps[run_id] frontier = _frontier(steps) + if frontier is not None and _wait_expired(frontier, now): + return "expired" if ( frontier is not None and frontier.status is StepStatus.BLOCKED @@ -1248,18 +1274,21 @@ async def resume_run(self, run_id: str, now: float) -> bool: self._append_events(run_id, ((HistoryEventType.RUN_RESUMED, {}),), now) return True - async def recover_orphans(self, now: float, max_recoveries: int) -> int: - """Recover steps left claimed by a previous process. + async def recover_orphans( + self, now: float, max_recoveries: int + ) -> tuple[int, tuple[str, ...]]: + """Recover claims whose lease has expired. Args: now: Current time in epoch seconds. max_recoveries: Recovery budget per logical step. Returns: - The number of steps transitioned. + How many steps were transitioned, and the runs failed outright. """ async with self._lock: recovered = 0 + failed: list[str] = [] for run in list(self._runs.values()): if run.status in TERMINAL_RUN_STATUSES: continue @@ -1283,6 +1312,7 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: error={"reason": "recovery_budget_exhausted"}, updated_at=now, ) + failed.append(run.run_id) self._append_events( run.run_id, ( @@ -1312,7 +1342,7 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: ), now, ) - return recovered + return recovered, tuple(failed) async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: """List runs matching a query, newest first. @@ -1325,7 +1355,7 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: """ async with self._lock: matched = [run for run in self._runs.values() if _matches_query(run, query)] - matched.sort(key=lambda run: run.created_at, reverse=True) + matched.sort(key=lambda run: (run.created_at, run.run_id), reverse=True) return tuple(_detach_run(run) for run in matched[: query.limit]) async def find_by_request_key( @@ -2116,6 +2146,9 @@ async def deliver( self._db.execute("ROLLBACK") return "duplicate" frontier = _frontier(self._load_steps(run_id)) + if frontier is not None and _wait_expired(frontier, now): + self._db.execute("ROLLBACK") + return "expired" resolves = ( frontier is not None and frontier.status is StepStatus.BLOCKED @@ -2549,15 +2582,17 @@ async def resume_run(self, run_id: str, now: float) -> bool: raise return True - async def recover_orphans(self, now: float, max_recoveries: int) -> int: - """Recover steps left claimed by a previous process. + async def recover_orphans( + self, now: float, max_recoveries: int + ) -> tuple[int, tuple[str, ...]]: + """Recover claims whose lease has expired. Args: now: Current time in epoch seconds. max_recoveries: Recovery budget per logical step. Returns: - The number of steps transitioned. + How many steps were transitioned, and the runs failed outright. """ terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) with self._lock: @@ -2571,6 +2606,7 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: (StepStatus.CLAIMED.value, now, *terminal), ).fetchall() recovered = 0 + failed: list[str] = [] for row in rows: step = _step_from_row(row) recovered += 1 @@ -2598,6 +2634,7 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: step.run_id, ), ) + failed.append(step.run_id) self._append_events( step.run_id, ( @@ -2636,7 +2673,7 @@ async def recover_orphans(self, now: float, max_recoveries: int) -> int: except BaseException: self._db.execute("ROLLBACK") raise - return recovered + return recovered, tuple(failed) async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: """List runs matching a query, newest first. @@ -2657,11 +2694,16 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: clauses.append(f"status IN ({placeholders})") params.extend(status.value for status in query.statuses) if query.created_before is not None: - clauses.append("created_at < ?") - params.append(query.created_before) + clauses.append("(created_at, run_id) < (?, ?)") + params.extend(query.created_before) for key, value in (query.labels or {}).items(): - clauses.append("json_extract(labels, ?) = ?") - params.extend((f"$.{key}", value)) + # The key comes from user data, so it is matched as a value rather + # than spliced into a JSON path expression. + clauses.append( + "EXISTS (SELECT 1 FROM json_each(labels)" + " WHERE json_each.key = ? AND json_each.value = ?)" + ) + params.extend((key, value)) where = f" WHERE {' AND '.join(clauses)}" if clauses else "" with self._lock: rows = self._db.execute( diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index f8791b42706..844bc7f6aa0 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -14,7 +14,11 @@ from typing import TYPE_CHECKING, Any -from reflex_base.workflow import DEFAULT_LEASE_DURATION, parse_duration +from reflex_base.workflow import ( + DEFAULT_LEASE_DURATION, + DEFAULT_MAX_RECOVERIES, + parse_duration, +) from reflex.workflow.kernel import WorkflowObserver from reflex.workflow.runtime import WorkflowRuntime, _context_runtime @@ -70,6 +74,7 @@ def __init__( lease_duration: DurationLike = DEFAULT_LEASE_DURATION, lease_renew_interval: float | None = None, observer: WorkflowObserver | None = None, + max_recoveries: int = DEFAULT_MAX_RECOVERIES, ): """Initialize the harness. @@ -80,6 +85,7 @@ def __init__( lease_duration: Virtual seconds a claim survives without renewal. lease_renew_interval: Real seconds between lease renewals. observer: Receives every recorded run transition. + max_recoveries: Infrastructure recovery budget per logical step. """ self._clock = _VirtualClock(start_time) self._runtime = WorkflowRuntime( @@ -89,6 +95,7 @@ def __init__( lease_duration=parse_duration(lease_duration), lease_renew_interval=lease_renew_interval, observer=observer, + max_recoveries=max_recoveries, ) for workflow_cls in workflow_classes: self._runtime.register(workflow_cls) diff --git a/tests/units/workflow/test_app.py b/tests/units/workflow/test_app.py index 7103c29f5a6..0838f7726ba 100644 --- a/tests/units/workflow/test_app.py +++ b/tests/units/workflow/test_app.py @@ -33,9 +33,15 @@ def begin(self): return SessionCounter, DetachedWorkflow -def test_add_workflow_detaches_from_session_tree(forked_registration_context): +def test_a_workflow_class_is_never_in_the_session_tree(forked_registration_context): + """Detaching happens at class creation, so forgetting to register is safe. + + A durable handler must never be dispatchable from a browser. Tying that to + app.add_workflow() would mean an omitted registration left the handlers + exposed, which is exactly the line a code generator drops. + """ session_cls, workflow_cls = _make_classes() - assert workflow_cls in State.get_substates() + assert workflow_cls not in State.get_substates() app = rx.App() app.add_workflow(workflow_cls) diff --git a/tests/units/workflow/test_audit_regressions.py b/tests/units/workflow/test_audit_regressions.py index bd9d375ae66..08f8031e1c3 100644 --- a/tests/units/workflow/test_audit_regressions.py +++ b/tests/units/workflow/test_audit_regressions.py @@ -158,3 +158,170 @@ async def handler(self): await asyncio.sleep(0) assert handler is not None + + +def test_an_unregistered_workflow_is_not_browser_reachable( + forked_registration_context, +): + """Forgetting app.add_workflow() must not expose durable handlers.""" + from reflex_base.registry import RegistrationContext + + from reflex.state import State + + class Forgotten(rx.State): + __workflow__ = WorkflowConfig(id="audit.forgotten") + amount: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="non_idempotent_write") + def charge(self): + """Move real money.""" + + context = RegistrationContext.get() + assert Forgotten not in State.get_substates() + assert Forgotten.get_full_name() not in context.base_states + assert not [ + name + for name, registered in context.event_handlers.items() + if Forgotten in registered.states + ] + + +async def test_signal_after_deadline_is_refused(forked_registration_context): + """Once the deadline wins, a late signal is refused rather than buffered.""" + from reflex_base.workflow import Signal, wait_for + + class Expiring(rx.State): + __workflow__ = WorkflowConfig(id="audit.expiring") + outcome: str = "" + + ping = Signal() + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Wait briefly for a signal. + + Returns: + The wait. + """ + return wait_for( + Expiring.ping, + then=Expiring.woke, + timeout="1h", + on_timeout=Expiring.late, + ) + + @rx.event(durable=True, effect="none") + def woke(self, payload: str): + """Resume on the signal. + + Args: + payload: The delivered payload. + """ + self.outcome = "signalled" + + @rx.event(durable=True, effect="none") + def late(self): + """Give up after the deadline.""" + self.outcome = "expired" + + async with WorkflowTestHarness(Expiring) as harness: + result = await harness.kernel.start(Expiring.begin) + assert result.run_id is not None + await harness.kernel.run_until_idle() + + # The deadline has fallen due but no worker has claimed it yet. + harness._clock.now += 3601 + assert ( + await harness.kernel.signal(result.run_id, Expiring.ping("late")) + == "expired" + ) + + await harness.kernel.run_until_idle() + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.state["outcome"] == "expired" + + +async def test_a_child_failed_by_recovery_reports_to_its_join( + forked_registration_context, +): + """A child that exhausts its recovery budget must not hang its parent.""" + + class Branch(rx.State): + __workflow__ = WorkflowConfig(id="audit.branch") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self, lead: str): + """Finish immediately. + + Args: + lead: The lead identifier. + + Returns: + Completion. + """ + return rx.complete(result={"ok": True}) + + class Parent(rx.State): + __workflow__ = WorkflowConfig(id="audit.parent") + outcomes: list[str] = [] + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self, lead: str): + """Fan out to two branches. + + Args: + lead: The lead identifier. + + Returns: + The fan-out. + """ + return rx.parallel(Branch.go(lead), Branch.go(lead), then=Parent.join) + + @rx.event(durable=True, effect="none") + def join(self, results: list): + """Record every branch outcome. + + Args: + results: One entry per branch. + + Returns: + Completion. + """ + self.outcomes = sorted(entry["status"] for entry in results) + return rx.complete(result={"branches": len(results)}) + + async with WorkflowTestHarness(Parent, Branch, max_recoveries=0) as harness: + store = harness.kernel.store + result = await harness.kernel.start(Parent.begin("lead_1")) + assert result.run_id is not None + # Run only the parent's root so the children exist but have not run. + assert await harness.kernel._tick() + + children = [ + run + for run in await harness.kernel.list_runs() + if run.parent_run_id == result.run_id + ] + assert len(children) == 2 + + # Claim one branch, then abandon it past its recovery budget, as a + # crash loop would. max_recoveries=0 exhausts it on the first sweep. + claim = await store.claim_next(harness.now, lease_duration=1.0) + assert claim is not None + doomed_id = claim.run.run_id + harness._clock.now += 2 + # recover() is what tells a parent's join about a child it failed. + assert await harness.kernel.recover() == 1 + + await harness.kernel.recover() + await harness.kernel.run_until_idle() + + failed = await store.get_run(doomed_id) + assert failed is not None + assert failed.status is RunStatus.FAILED + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert "FAILED" in snapshot.state["outcomes"] diff --git a/tests/units/workflow/test_lease.py b/tests/units/workflow/test_lease.py index 21696f2a1e5..492cde8a8dc 100644 --- a/tests/units/workflow/test_lease.py +++ b/tests/units/workflow/test_lease.py @@ -223,7 +223,7 @@ async def work(self): # A peer reclaims the step once the lease lapses; the renewer notices. clock.now += 31.0 - assert await store.recover_orphans(clock(), max_recoveries=10) == 1 + assert (await store.recover_orphans(clock(), max_recoveries=10))[0] == 1 # The first attempt stays blocked, so only the renewer can end it: it # sees the fence, cancels the attempt, and the loop re-runs the step. await asyncio.wait_for(pump, timeout=5) @@ -288,9 +288,9 @@ async def work(self): assert run is not None assert run.status is RunStatus.RUNNING # It becomes reclaimable only after the lease lapses. - assert await store.recover_orphans(clock(), max_recoveries=10) == 0 + assert (await store.recover_orphans(clock(), max_recoveries=10))[0] == 0 clock.now += 31.0 - assert await store.recover_orphans(clock(), max_recoveries=10) == 1 + assert (await store.recover_orphans(clock(), max_recoveries=10))[0] == 1 finally: release.set() store.close() diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 10f73ebabe8..8ce28af7cf7 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -199,8 +199,7 @@ async def test_duplicate_arrivals_are_counted_once(forked_registration_context): child = next( run for run in runs - if run.parent_run_id == result.run_id - and run.status is RunStatus.COMPLETED + if run.parent_run_id == result.run_id and run.status is RunStatus.COMPLETED ) repeat = await harness.kernel.store.record_arrival( result.run_id, diff --git a/tests/units/workflow/test_queries.py b/tests/units/workflow/test_queries.py index 0ca9879569c..9d063069d9e 100644 --- a/tests/units/workflow/test_queries.py +++ b/tests/units/workflow/test_queries.py @@ -76,7 +76,7 @@ async def test_list_runs_orders_newest_first_and_paginates(store): page = await store.list_runs(RunQuery(limit=2)) assert [run.run_id for run in page] == ["run4", "run3"] nextpage = await store.list_runs( - RunQuery(limit=2, created_before=page[-1].created_at) + RunQuery(limit=2, created_before=(page[-1].created_at, page[-1].run_id)) ) assert [run.run_id for run in nextpage] == ["run2", "run1"] diff --git a/tests/units/workflow/test_store.py b/tests/units/workflow/test_store.py index 03f8575857a..b70bc4727cc 100644 --- a/tests/units/workflow/test_store.py +++ b/tests/units/workflow/test_store.py @@ -257,7 +257,7 @@ async def test_recover_orphans_consumes_recovery_budget(store): await store.admit(_run(), _step(), _ADMIT_EVENTS) claim = await store.claim_next(NOW, lease_duration=5.0) assert claim is not None - recovered = await store.recover_orphans(NOW + 10, max_recoveries=2) + recovered, _ = await store.recover_orphans(NOW + 10, max_recoveries=2) assert recovered == 1 steps = await store.get_steps("run1") assert steps[0].status is StepStatus.RECOVERY_WAIT @@ -319,7 +319,7 @@ async def test_sqlite_persistence_across_reopen(tmp_path): assert not created assert run_id == "run1" # The orphaned claim recovers on the new process. - assert await second.recover_orphans(NOW + 5, max_recoveries=10) == 1 + assert (await second.recover_orphans(NOW + 5, max_recoveries=10))[0] == 1 steps = await second.get_steps("run1") assert steps[0].status is StepStatus.RECOVERY_WAIT finally: @@ -366,7 +366,7 @@ async def test_recover_orphans_skips_unexpired_leases(store): await store.admit(_run(), _step(), _ADMIT_EVENTS) claim = await store.claim_next(NOW, lease_duration=30.0) assert claim is not None - assert await store.recover_orphans(NOW + 29.0, max_recoveries=10) == 0 + assert (await store.recover_orphans(NOW + 29.0, max_recoveries=10))[0] == 0 steps = await store.get_steps("run1") assert steps[0].status is StepStatus.CLAIMED assert steps[0].recoveries == 0 @@ -376,7 +376,7 @@ async def test_recover_orphans_reclaims_at_the_expiry_boundary(store): await store.admit(_run(), _step(), _ADMIT_EVENTS) claim = await store.claim_next(NOW, lease_duration=30.0) assert claim is not None - assert await store.recover_orphans(NOW + 30.0, max_recoveries=10) == 1 + assert (await store.recover_orphans(NOW + 30.0, max_recoveries=10))[0] == 1 steps = await store.get_steps("run1") assert steps[0].status is StepStatus.RECOVERY_WAIT assert steps[0].recoveries == 1 @@ -400,15 +400,15 @@ async def test_renewed_lease_survives_a_later_sweep(store): claim = await store.claim_next(NOW, lease_duration=30.0) assert claim is not None assert await store.renew_lease(claim, NOW + 20.0, lease_duration=30.0) - assert await store.recover_orphans(NOW + 40.0, max_recoveries=10) == 0 - assert await store.recover_orphans(NOW + 50.0, max_recoveries=10) == 1 + assert (await store.recover_orphans(NOW + 40.0, max_recoveries=10))[0] == 0 + assert (await store.recover_orphans(NOW + 50.0, max_recoveries=10))[0] == 1 async def test_renew_lease_refused_after_recovery(store): await store.admit(_run(), _step(), _ADMIT_EVENTS) claim = await store.claim_next(NOW, lease_duration=5.0) assert claim is not None - assert await store.recover_orphans(NOW + 10.0, max_recoveries=10) == 1 + assert (await store.recover_orphans(NOW + 10.0, max_recoveries=10))[0] == 1 assert not await store.renew_lease(claim, NOW + 11.0, lease_duration=30.0) @@ -457,7 +457,7 @@ async def test_sqlite_migrates_a_database_without_the_lease_column(tmp_path): # A claim left by the previous binary has no lease, so it is a genuine # orphan and is reclaimed on the first sweep. assert steps[0].lease_expires_at == pytest.approx(0.0) - assert await reopened.recover_orphans(NOW, max_recoveries=10) == 1 + assert (await reopened.recover_orphans(NOW, max_recoveries=10))[0] == 1 finally: reopened.close() # Reopening an already-migrated database is a no-op. From ef18d0935fcb3283949665561d58484a96777368 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 16:32:33 -0700 Subject: [PATCH 021/121] Abandon a lease that cannot be renewed, and gate fan-out branches Two more confirmed findings. Lease renewal failures were swallowed at debug level, so a store that kept failing -- which a contended SQLite file does, raising after a five second block -- left the attempt running while its lease quietly lapsed. Recovery then handed the step to someone else and the external effect ran twice, with a log line that is off by default as the only evidence. The kernel now tracks when the lease actually expires and abandons the attempt once too little of it remains to survive another failed round-trip: better to stop work you can no longer prove you own than to race the worker that is about to take it. rx.parallel resolved its branches through a path that skipped the trigger check start() enforces, so a webhook-only root -- one that exists precisely because only a verified provider may start it -- could be started from application code by naming it as a branch. Branches are now held to the same manual-root rule as a direct start. Fanning out to a handler of your own class is now a compile error. Each branch becomes a child run with fresh state, so a same-class branch cannot see anything the parent did; it is the most natural spelling and it silently did the wrong thing. --- news/workflow-audit-fixes-4.bugfix.md | 1 + reflex/workflow/definition.py | 51 +++++++++ reflex/workflow/kernel.py | 30 ++++- .../units/workflow/test_audit_regressions.py | 107 ++++++++++++++++++ 4 files changed, 186 insertions(+), 3 deletions(-) create mode 100644 news/workflow-audit-fixes-4.bugfix.md diff --git a/news/workflow-audit-fixes-4.bugfix.md b/news/workflow-audit-fixes-4.bugfix.md new file mode 100644 index 00000000000..7406535d5e5 --- /dev/null +++ b/news/workflow-audit-fixes-4.bugfix.md @@ -0,0 +1 @@ +Fixes two more audit findings: a persistently failing lease renewal could let the kernel re-execute its own live attempt, and `rx.parallel` branches bypassed the trigger check that direct starts enforce. diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 9a93e23f37a..4add7b47684 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -349,6 +349,56 @@ def _handler_body(fn: Callable) -> ast.FunctionDef | ast.AsyncFunctionDef | None return None +def _validate_branches( + workflow_cls: type[BaseState], + defn: HandlerDefinition, + handlers: Mapping[str, HandlerDefinition], +) -> None: + """Reject fan-outs whose branches are handlers of this same workflow. + + A branch becomes a child run with its own fresh state, so naming a sibling + handler produces a child that cannot see anything the parent has done -- + silently, and only at runtime. Fan out to another workflow class instead. + + Args: + workflow_cls: The workflow class being compiled. + defn: The handler to check. + handlers: Every durable handler on this class. + + Raises: + WorkflowDefinitionError: If a branch names a handler of this class. + """ + node = _handler_body(defn.fn) + if node is None: + return + for child in ast.walk(node): + if not ( + isinstance(child, ast.Call) + and isinstance(child.func, ast.Name) + and child.func.id == "parallel" + ) and not ( + isinstance(child, ast.Call) + and isinstance(child.func, ast.Attribute) + and child.func.attr == "parallel" + ): + continue + for branch in child.args: + target = branch.func if isinstance(branch, ast.Call) else branch + if ( + isinstance(target, ast.Attribute) + and isinstance(target.value, ast.Name) + and target.value.id == workflow_cls.__name__ + and target.attr in {h.name for h in handlers.values()} + ): + raise _error( + workflow_cls, + f"handler {defn.name!r} fans out to {target.attr!r} on its " + "own class. Each branch runs as a child run with fresh " + "state, so it would not see this run's state. Move the " + "branch to its own workflow class.", + ) + + def _validate_handler_body( workflow_cls: type[BaseState], defn: HandlerDefinition, @@ -554,6 +604,7 @@ def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: fields = _compile_fields(workflow_cls) handlers = _resolve_hooks(workflow_cls, _compile_handlers(workflow_cls, config)) for defn in handlers.values(): + _validate_branches(workflow_cls, defn, handlers) policy = defn.singleton or defn.rate_limit or defn.throttle or defn.debounce key = getattr(policy, "key", None) if key is not None and key not in defn.params: diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index bd5b66db2e9..7747a2a3b47 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -29,6 +29,7 @@ ChannelDelivery, CompleteRun, FailRun, + ManualTrigger, NeedsAttention, Parallel, ScheduleTrigger, @@ -202,10 +203,11 @@ class _Lease: claim: The claim being kept alive. attempt: The task running the handler, once created. renewer: The background task extending the lease. + expires_at: When this lease lapses if renewal keeps failing. lost: Whether the store reported the claim was fenced. """ - __slots__ = ("attempt", "claim", "lost", "renewer") + __slots__ = ("attempt", "claim", "expires_at", "lost", "renewer") def __init__(self, claim: Claim): """Initialize the lease. @@ -216,6 +218,7 @@ def __init__(self, claim: Claim): self.claim = claim self.attempt: asyncio.Task | None = None self.renewer: asyncio.Task | None = None + self.expires_at = claim.step.lease_expires_at self.lost = False @@ -1410,15 +1413,29 @@ async def _renew(self, lease: _Lease) -> None: """ if lease.lost: return + now = self._clock() try: held = await self._store.renew_lease( - lease.claim, self._clock(), lease_duration=self._lease_duration + lease.claim, now, lease_duration=self._lease_duration ) except Exception as err: - console.debug(f"Workflow lease renewal failed, retrying: {err}") + # The lease keeps ticking down while renewal fails. Once too little + # of it remains to survive another failed round-trip, stop the + # attempt rather than keep running work this kernel can no longer + # prove it owns -- recovery is about to hand the step to someone else. + if lease.expires_at - now <= self._lease_renew_interval: + console.warn( + "Workflow lease renewal keeps failing and the lease is about " + f"to lapse; abandoning the attempt: {err}" + ) + self._lose_lease(lease) + else: + console.debug(f"Workflow lease renewal failed, retrying: {err}") return if not held: self._lose_lease(lease) + return + lease.expires_at = now + self._lease_duration async def _renew_forever(self, lease: _Lease) -> None: """Renew a lease on a real-time cadence until it ends or is lost. @@ -1772,6 +1789,13 @@ def _child_records( records = [] for index, branch in enumerate(branches): defn, handler, payload = self._resolve_target(branch) + if not isinstance(handler.trigger, ManualTrigger): + msg = ( + f"Cannot fan out to {handler.id!r} of {defn.workflow_id!r}: " + f"it declares {getattr(handler.trigger, 'kind', 'no trigger')}, " + "and a branch must be a manual root just like a direct start." + ) + raise WorkflowRuntimeError(msg) child_id = uuid.uuid4().hex records.append(( RunRecord( diff --git a/tests/units/workflow/test_audit_regressions.py b/tests/units/workflow/test_audit_regressions.py index 08f8031e1c3..d4c2113c970 100644 --- a/tests/units/workflow/test_audit_regressions.py +++ b/tests/units/workflow/test_audit_regressions.py @@ -325,3 +325,110 @@ def join(self, results: list): assert snapshot is not None assert snapshot.status is RunStatus.COMPLETED assert "FAILED" in snapshot.state["outcomes"] + + +def test_fanning_out_to_your_own_class_is_rejected(forked_registration_context): + """A same-class branch would silently get empty state.""" + import importlib.util + import sys + import tempfile + import uuid as uuid_module + from pathlib import Path + + source = """ +import reflex as rx +from reflex_base.workflow import WorkflowConfig, manual, parallel + + +class SelfFanOut(rx.State): + __workflow__ = WorkflowConfig(id="audit.self_fanout") + lead_id: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self, lead_id: str): + self.lead_id = lead_id + return parallel(SelfFanOut.enrich, SelfFanOut.score, then=SelfFanOut.join) + + @rx.event(durable=True, effect="none") + def enrich(self): + pass + + @rx.event(durable=True, effect="none") + def score(self): + pass + + @rx.event(durable=True, effect="none") + def join(self, results: list): + pass +""" + name = f"wf_selffan_{uuid_module.uuid4().hex}" + path = Path(tempfile.gettempdir()) / f"{name}.py" + path.write_text(source) + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + + from reflex.workflow.definition import compile_workflow + + with pytest.raises(WorkflowDefinitionError, match="its own class"): + compile_workflow(module.SelfFanOut) + + +async def test_a_branch_must_be_a_manual_root(forked_registration_context): + """Fan-out must not reach a root that only a provider may start.""" + from reflex_base.workflow import hmac_signature, parallel, webhook + + class Hooked(rx.State): + __workflow__ = WorkflowConfig(id="audit.hooked") + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "audit.topic", + verify=hmac_signature(secret_env="S", header="X-Sig"), + ), + ) + def on_event(self): + """Only a verified provider may start this.""" + + class Driver(rx.State): + __workflow__ = WorkflowConfig(id="audit.driver") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Try to fan out to a webhook-only root. + + Returns: + The fan-out. + """ + return parallel(Hooked.on_event, then=Driver.join) + + @rx.event(durable=True, effect="none") + def join(self, results: list): + """Never reached. + + Args: + results: Branch outcomes. + """ + + async with WorkflowTestHarness(Driver, Hooked) as harness: + result = await harness.start(Driver.begin) + assert result.run_id is not None + # The branch is resolved while the parent's step runs, so the gate + # surfaces as a failed parent rather than a raise at start(). + await harness.advance("1s") + await harness.advance("2s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error is not None + assert "manual root" in snapshot.error["message"] + # No child run was created for the webhook-only root. + assert all( + run.workflow_id != "audit.hooked" + for run in await harness.kernel.list_runs() + ) From 560c9e3c6249b441eca3b1b7c8cd4f325520f09c Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 16:35:06 -0700 Subject: [PATCH 022/121] Bound SQLite contention so it cannot freeze the app Every store call is synchronous on the caller's event loop, which also serves HTTP, websockets, and session state. SQLite's default busy timeout is multiple seconds, so a second process writing the same file could stall the whole app before raising -- and the error then landed in lease renewal, where a silent failure used to mean the kernel re-executed its own attempt. The busy timeout is now short: contention surfaces quickly as a transient error the kernel retries, rather than freezing everything first. Combined with the previous commit, a store that cannot be reached now degrades to abandoning the attempt instead of duplicating its effect. This bounds the symptom rather than removing the cause. Offloading the store's calls to a thread is the real fix, and one worker process per database file remains the supported deployment -- now stated plainly in the docs alongside why horizontal scale wants a different store. --- docs/workflows/overview.md | 7 ++++++- news/workflow-sqlite-contention.bugfix.md | 1 + reflex/workflow/store.py | 15 ++++++++++++--- 3 files changed, 19 insertions(+), 4 deletions(-) create mode 100644 news/workflow-sqlite-contention.bugfix.md diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index fae727cb14f..772f424f00b 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -356,7 +356,12 @@ to a waiting run, and `harness.cancel(...)` and `harness.resume(...)` drive the ## Deploying Runs persist to a SQLite file next to your app by default; pass `rx.App(workflow_store=...)` to -choose another store. Run one worker process per SQLite database file. +choose another store. + +Run **one worker process per SQLite database file**. The store's calls are synchronous on the event +loop that also serves your app, so two processes writing the same file contend for it; contention is +bounded to a short busy timeout and surfaces as a transient error the kernel retries, but throughput +does not improve. Horizontal scale wants a store that supports concurrent writers. `RunStore` is a supported extension point, and the invariants a store must satisfy ship as runnable checks rather than prose: diff --git a/news/workflow-sqlite-contention.bugfix.md b/news/workflow-sqlite-contention.bugfix.md new file mode 100644 index 00000000000..976d8c62460 --- /dev/null +++ b/news/workflow-sqlite-contention.bugfix.md @@ -0,0 +1 @@ +A contended SQLite workflow store now fails fast instead of blocking the event loop that serves your app. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 09a123e7a1e..c2c5dde0759 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -8,9 +8,11 @@ ``MemoryRunStore`` backs tests and the harness. ``SqliteRunStore`` provides crash-safe persistence on a single machine using the standard library. Run -exactly one worker process per database file: its calls are synchronous and -cross-process write contention blocks the caller's event loop, which is -hostile to lease renewal. +exactly one worker process per database file: its calls are synchronous on the +caller's event loop, so cross-process write contention stalls the app that is +serving requests. Contended writes are bounded to a short busy timeout and +surface as transient errors the kernel retries, rather than freezing the loop +for seconds. """ from __future__ import annotations @@ -1431,6 +1433,8 @@ async def next_due(self, now: float) -> float | None: return min(due_times) if due_times else None +BUSY_TIMEOUT_MS: Final = 250 + _SCHEMA = """ CREATE TABLE IF NOT EXISTS workflow_runs ( run_id TEXT PRIMARY KEY, @@ -1624,6 +1628,11 @@ def __init__(self, db_path: str | Path): self._db.isolation_level = None self._db.execute("PRAGMA journal_mode=WAL") self._db.execute("PRAGMA synchronous=NORMAL") + # These calls are synchronous on the caller's event loop, which also + # serves HTTP and websockets, so a contended write must fail fast + # rather than block everything for SQLite's multi-second default. The + # kernel treats the resulting error as transient and retries. + self._db.execute(f"PRAGMA busy_timeout={BUSY_TIMEOUT_MS}") self._db.executescript(_SCHEMA) self._migrate() From 5d3cb73e5f221e9fb23a597b6f652ec3435684bf Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 17:02:39 -0700 Subject: [PATCH 023/121] Run workflow attempts concurrently The kernel claimed one step and awaited it before claiming again, so a single process executed one step at a time no matter how many runs were waiting. One ten-minute step stalled every other run in the deployment -- including their timers and deadlines -- which is not a throughput story any durable engine can be compared on. The scheduler now fills up to max_concurrency slots, eight by default. This needs no new locking: each run has exactly one claimable frontier step, so two concurrent claims are necessarily different runs and a run's own mailbox stays strictly serial. A test asserts both halves -- that attempts really do overlap, and that one run's steps still complete in order. Three details that had to be right: finished attempts are pruned synchronously rather than by a done-callback, or the scheduler spins on work it already finished; a round waits only on the attempts it started, so a second caller pumping the same kernel is not blocked by an attempt it does not own; and cancelling the scheduler cancels the attempts it started, since asyncio.wait does not cancel what it waits on. The test harness stays single-slot so virtual-clock tests remain deterministic; production defaults to eight. --- docs/workflows/overview.md | 9 ++ news/workflow-concurrency.feature.md | 1 + reflex/app.py | 9 +- reflex/workflow/kernel.py | 90 ++++++++++++-- reflex/workflow/runtime.py | 5 + reflex/workflow/testing.py | 4 + tests/units/workflow/test_concurrency.py | 142 +++++++++++++++++++++++ 7 files changed, 249 insertions(+), 11 deletions(-) create mode 100644 news/workflow-concurrency.feature.md create mode 100644 tests/units/workflow/test_concurrency.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index 772f424f00b..66e5defffeb 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -308,6 +308,15 @@ reflex workflows show --history `reflex workflows cancel ` and `reflex workflows resume ` steer a run without opening the app, and `--json` on `list` and `show` makes the output scriptable. +## Throughput + +The kernel runs several attempts at once, defaulting to eight. Each belongs to a different run -- +a run has exactly one open step, so its own work stays strictly in order no matter how many other +runs are executing. Tune it with `rx.App(workflow_concurrency=...)`. + +Concurrency is per process. A step that blocks the event loop, rather than awaiting, still stalls +its neighbours, so prefer `async def` handlers for anything that waits on the network. + ## Observability Pass an observer to see every recorded transition, correlated to its run, workflow, step, and diff --git a/news/workflow-concurrency.feature.md b/news/workflow-concurrency.feature.md new file mode 100644 index 00000000000..affc77380a7 --- /dev/null +++ b/news/workflow-concurrency.feature.md @@ -0,0 +1 @@ +The workflow kernel now runs several attempts at once instead of one step at a time per process, tunable with `rx.App(workflow_concurrency=...)`. diff --git a/reflex/app.py b/reflex/app.py index 2a59bc3035c..c829cf70ffb 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -98,7 +98,7 @@ ) from reflex.utils.misc import run_in_thread from reflex.utils.token_manager import RedisTokenManager, TokenManager -from reflex.workflow.kernel import WorkflowObserver +from reflex.workflow.kernel import DEFAULT_MAX_CONCURRENCY, WorkflowObserver from reflex.workflow.runtime import WorkflowRuntime from reflex.workflow.store import RunStore @@ -460,6 +460,9 @@ class App(MiddlewareMixin, LifespanMixin): # Receives every recorded workflow run transition, for logs or tracing. workflow_observer: WorkflowObserver | None = None + # How many workflow attempts run at once, across different runs. + workflow_concurrency: int = DEFAULT_MAX_CONCURRENCY + # The workflow runtime owning registered definitions and the kernel. _workflow_runtime: WorkflowRuntime | None = None @@ -974,7 +977,9 @@ def add_workflow(self, workflow_cls: type[BaseState]) -> None: """ if self._workflow_runtime is None: self._workflow_runtime = WorkflowRuntime( - self.workflow_store, observer=self.workflow_observer + self.workflow_store, + observer=self.workflow_observer, + max_concurrency=self.workflow_concurrency, ) self.register_lifespan_task(self._run_workflow_runtime) self._workflow_runtime.register(workflow_cls) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 7747a2a3b47..b33c4d21b4a 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -69,6 +69,8 @@ DEFAULT_POLL_INTERVAL = 0.25 +DEFAULT_MAX_CONCURRENCY = 8 + LEASE_RENEW_FRACTION = 1 / 3 RECOVERY_INTERVAL_FRACTION = 1 / 2 @@ -238,6 +240,7 @@ def __init__( lease_renew_interval: float | None = None, recovery_interval: float | None = None, observer: WorkflowObserver | None = None, + max_concurrency: int = DEFAULT_MAX_CONCURRENCY, ): """Initialize the kernel. @@ -256,6 +259,8 @@ def __init__( background worker; defaults to half of ``lease_duration``. observer: Receives every recorded transition, for logging, metrics, or tracing. + max_concurrency: How many attempts this kernel runs at once. Each + belongs to a different run, so a run's own steps stay serial. Raises: WorkflowRuntimeError: If the store cannot renew leases, or the @@ -272,6 +277,7 @@ def __init__( self._rng = rng self._poll_interval = poll_interval self._max_recoveries = max_recoveries + self._max_concurrency = max(1, max_concurrency) if not hasattr(store, "renew_lease"): msg = ( f"{type(store).__name__} does not implement renew_lease; a run " @@ -1927,6 +1933,57 @@ async def _finalize_control(self, now: float) -> int: finalized += 1 return finalized + def _spawn(self, claim: Claim) -> asyncio.Task: + """Start executing a claim, tracking it until it finishes. + + Args: + claim: The claim to execute. + + Returns: + The task running the attempt. + """ + task = asyncio.ensure_future(self._execute_claim(claim)) + self._inflight[claim.run.run_id] = task + return task + + def _prune(self) -> int: + """Drop finished attempts from the in-flight set. + + Pruning is synchronous rather than done-callback driven: a callback + runs on a later loop iteration, so the scheduler would keep seeing + completed work as in flight and spin. + + Returns: + How many attempts had finished. + """ + finished = [run_id for run_id, task in self._inflight.items() if task.done()] + for run_id in finished: + del self._inflight[run_id] + return len(finished) + + async def _fill_slots(self, now: float) -> list[asyncio.Task]: + """Claim work up to this kernel's concurrency limit. + + Each run has one frontier, so two concurrent claims are always + different runs: parallelism across runs never breaks the serial + mailbox each run relies on. + + Args: + now: Current time in epoch seconds. + + Returns: + The attempts this call started. + """ + started: list[asyncio.Task] = [] + while len(self._inflight) < self._max_concurrency: + claim = await self._store.claim_next( + now, lease_duration=self._lease_duration + ) + if claim is None: + break + started.append(self._spawn(claim)) + return started + async def _tick(self) -> bool: """Run one scheduling round. @@ -1934,19 +1991,34 @@ async def _tick(self) -> bool: True if any control transition or attempt was processed. """ now = self._clock() - progressed = await self._admit_due_schedules(now) > 0 + progressed = self._prune() > 0 + progressed = await self._admit_due_schedules(now) > 0 or progressed progressed = await self._finalize_control(now) > 0 or progressed - claim = await self._store.claim_next(now, lease_duration=self._lease_duration) - if claim is not None: - task = asyncio.ensure_future(self._execute_claim(claim)) - self._inflight[claim.run.run_id] = task + started = await self._fill_slots(now) + progressed = bool(started) or progressed + # Wait only on what this round started: another caller pumping the same + # kernel must not block on an attempt it does not own. + if started: try: - await task - finally: - self._inflight.pop(claim.run.run_id, None) - progressed = True + await asyncio.wait(started, return_when=asyncio.FIRST_COMPLETED) + except asyncio.CancelledError: + # asyncio.wait does not cancel what it waits on, so cancelling + # the scheduler must stop the attempts it started or they run + # on unsupervised. + await self._cancel_inflight() + raise + progressed = self._prune() > 0 or progressed return progressed + async def _cancel_inflight(self) -> None: + """Stop every attempt this kernel is running and wait for them.""" + tasks = list(self._inflight.values()) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._prune() + async def recover(self) -> int: """Renew this kernel's live claims, then reclaim expired ones. diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 65ee34d1284..1cfb30afc86 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -25,6 +25,7 @@ from reflex.workflow.definition import WorkflowDefinition, compile_workflow from reflex.workflow.kernel import ( + DEFAULT_MAX_CONCURRENCY, DEFAULT_POLL_INTERVAL, WorkflowKernel, WorkflowObserver, @@ -64,6 +65,7 @@ def __init__( recovery_interval: float | None = None, observer: WorkflowObserver | None = None, max_recoveries: int = DEFAULT_MAX_RECOVERIES, + max_concurrency: int = DEFAULT_MAX_CONCURRENCY, ): """Initialize the runtime. @@ -79,6 +81,7 @@ def __init__( recovery_interval: Seconds between recovery sweeps. observer: Receives every recorded run transition. max_recoveries: Infrastructure recovery budget per logical step. + max_concurrency: How many attempts run at once. """ self._store = store self._clock = clock @@ -89,6 +92,7 @@ def __init__( self._recovery_interval = recovery_interval self._observer = observer self._max_recoveries = max_recoveries + self._max_concurrency = max_concurrency self._definitions: dict[str, WorkflowDefinition] = {} self._classes: dict[type, str] = {} self._kernel: WorkflowKernel | None = None @@ -178,6 +182,7 @@ async def startup(self, *, start_worker: bool = True) -> None: recovery_interval=self._recovery_interval, observer=self._observer, max_recoveries=self._max_recoveries, + max_concurrency=self._max_concurrency, ) if start_worker: await self._kernel.start_worker() diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index 844bc7f6aa0..af5d6ee03c5 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -75,6 +75,7 @@ def __init__( lease_renew_interval: float | None = None, observer: WorkflowObserver | None = None, max_recoveries: int = DEFAULT_MAX_RECOVERIES, + max_concurrency: int = 1, ): """Initialize the harness. @@ -86,6 +87,8 @@ def __init__( lease_renew_interval: Real seconds between lease renewals. observer: Receives every recorded run transition. max_recoveries: Infrastructure recovery budget per logical step. + max_concurrency: Attempts run at once. One by default, so tests on a + virtual clock stay deterministic. """ self._clock = _VirtualClock(start_time) self._runtime = WorkflowRuntime( @@ -96,6 +99,7 @@ def __init__( lease_renew_interval=lease_renew_interval, observer=observer, max_recoveries=max_recoveries, + max_concurrency=max_concurrency, ) for workflow_cls in workflow_classes: self._runtime.register(workflow_cls) diff --git a/tests/units/workflow/test_concurrency.py b/tests/units/workflow/test_concurrency.py new file mode 100644 index 00000000000..7e6ab7f791c --- /dev/null +++ b/tests/units/workflow/test_concurrency.py @@ -0,0 +1,142 @@ +"""Tests for running attempts from different runs at the same time.""" + +import asyncio + +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import RunStatus +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import MemoryRunStore + + +def _slow_flow(peak: list[int], live: list[int]): + """Build a workflow whose step records how many run at once. + + Args: + peak: Collects the highest observed concurrency. + live: Tracks currently executing attempts. + + Returns: + The workflow class. + """ + + class Slow(rx.State): + __workflow__ = WorkflowConfig(id="conc.slow") + n: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="read") + async def go(self, n: int): + """Hold the step open long enough to overlap with siblings. + + Args: + n: An identifier for the run. + """ + live.append(1) + peak.append(len(live)) + await asyncio.sleep(0.15) + live.pop() + self.n = n + + return Slow + + +async def test_steps_of_different_runs_execute_concurrently( + forked_registration_context, +): + """Throughput must not be one step at a time for the whole process.""" + peak: list[int] = [] + live: list[int] = [] + flow = _slow_flow(peak, live) + + runtime = WorkflowRuntime(MemoryRunStore(), poll_interval=0.01, max_concurrency=4) + runtime.register(flow) + async with runtime.running(): + results = [await runtime.kernel.start(flow.go(index)) for index in range(4)] + snapshots = [] + for _ in range(200): + snapshots = [ + await runtime.kernel.get_run(result.run_id) + for result in results + if result.run_id is not None + ] + if all( + snapshot is not None and snapshot.status is RunStatus.COMPLETED + for snapshot in snapshots + ): + break + await asyncio.sleep(0.02) + + assert all( + snapshot is not None and snapshot.status is RunStatus.COMPLETED + for snapshot in snapshots + ) + assert max(peak) > 1, "attempts never overlapped" + + +async def test_concurrency_is_bounded(forked_registration_context): + """The kernel must not start more attempts than its limit allows.""" + peak: list[int] = [] + live: list[int] = [] + flow = _slow_flow(peak, live) + + runtime = WorkflowRuntime(MemoryRunStore(), poll_interval=0.01, max_concurrency=2) + runtime.register(flow) + async with runtime.running(): + results = [await runtime.kernel.start(flow.go(index)) for index in range(6)] + snapshots = [] + for _ in range(300): + snapshots = [ + await runtime.kernel.get_run(result.run_id) + for result in results + if result.run_id is not None + ] + if all( + snapshot is not None and snapshot.status is RunStatus.COMPLETED + for snapshot in snapshots + ): + break + await asyncio.sleep(0.02) + + assert max(peak) <= 2 + assert all( + snapshot is not None and snapshot.status is RunStatus.COMPLETED + for snapshot in snapshots + ) + + +async def test_one_run_still_runs_its_steps_in_order(forked_registration_context): + """Concurrency is across runs; a run's own mailbox stays serial.""" + order: list[str] = [] + + class Serial(rx.State): + __workflow__ = WorkflowConfig(id="conc.serial") + seen: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + async def first(self): + """Run first. + + Returns: + The second step. + """ + order.append("first") + await asyncio.sleep(0.05) + return Serial.second + + @rx.event(durable=True, effect="none") + async def second(self): + """Run second.""" + order.append("second") + + runtime = WorkflowRuntime(MemoryRunStore(), poll_interval=0.01, max_concurrency=8) + runtime.register(Serial) + async with runtime.running(): + result = await runtime.kernel.start(Serial.first) + assert result.run_id is not None + for _ in range(200): + snapshot = await runtime.kernel.get_run(result.run_id) + if snapshot is not None and snapshot.status is RunStatus.COMPLETED: + break + await asyncio.sleep(0.02) + assert order == ["first", "second"] From 6716046d9b3895822bf780ff86990e99b3d4abde Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 17:11:14 -0700 Subject: [PATCH 024/121] Race parallel branches with mode="first" rx.parallel always waited for every branch, so the one shape everybody reaches for -- ask two vendors, take whoever answers first -- could not be expressed. Worse, the losing branch kept running and kept calling out to a vendor after the order was already booked. mode="first" sets the join's expected count to one. The arrival that resolves the join now says so, and the kernel cancels the siblings still in flight, found through a new list_children(parent_run_id, parent_ordinal) store query rather than a scan of the run table -- indexed in sqlite, and covered by a conformance check so any future store has to answer it too. Two compile-time diagnostics came out of writing the tests, both failure modes a code generator will hit: self.book passed as then= is a bound method, not a routable transition, so it failed at runtime -- as a retrying durable step, which is the worst place to learn about a typo. The handler-body guard already rejected self.step() calls; it now rejects bare self.step references too and names the class form. Passing a list of workflow classes where varargs were expected died with "cannot use 'list' as a dict key" from inside the registry. register() now says what it wanted. The docs example was executed before committing. --- docs/workflows/overview.md | 18 +- news/workflow-parallel-race.feature.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 21 ++- reflex/workflow/conformance.py | 42 +++++ reflex/workflow/definition.py | 44 ++++- reflex/workflow/kernel.py | 34 +++- reflex/workflow/runtime.py | 6 + reflex/workflow/store.py | 64 +++++++ tests/units/workflow/test_parallel.py | 157 ++++++++++++++++++ 9 files changed, 371 insertions(+), 16 deletions(-) create mode 100644 news/workflow-parallel-race.feature.md diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index 66e5defffeb..87a73b8f58a 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -127,7 +127,7 @@ would run it inline and lose its retries and effect tracking, which the compiler | `[MyFlow.a, MyFlow.b]` | run both, in order | | `rx.after("2d", MyFlow.later)` | run it after a durable delay | | `rx.wait_for(...)` | block until a signal or a deadline | -| `rx.parallel(a, b, then=...)` | run branches concurrently, then join | +| `rx.parallel(a, b, then=...)` | run branches concurrently, then join (`mode="first"` races) | | `rx.complete(result=...)` | finish the run successfully | | `rx.fail("reason")` | finish the run as failed | | `rx.needs_attention("reason")` | suspend for a human | @@ -211,6 +211,22 @@ A branch that fails still reports, so the join handler decides what a partial su than the engine guessing. Child runs are ordinary runs: they appear in `list_runs()` and can be inspected and cancelled on their own. +Pass `mode="first"` to race the branches instead of waiting for all of them. The join handler runs +as soon as one reports, receiving that single result, and the losing branches are cancelled. + +```python +return rx.parallel( + PrimaryVendor.quote(order_id), + BackupVendor.quote(order_id), + then=Order.book, + mode="first", +) +``` + +A cancelled loser stops at its next step boundary; a step already in flight finishes. Race only +where a discarded branch is harmless, or give the losers an `on_failure` handler that undoes their +work. + ## Triggers A root handler declares how runs of it begin. diff --git a/news/workflow-parallel-race.feature.md b/news/workflow-parallel-race.feature.md new file mode 100644 index 00000000000..c1100d8015c --- /dev/null +++ b/news/workflow-parallel-race.feature.md @@ -0,0 +1 @@ +`rx.parallel(..., mode="first")` races its branches: the join runs on the first result and the losing branches are cancelled. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 210205374f1..a62a7498ee9 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -960,16 +960,20 @@ class Parallel: Attributes: branches: The root events to run concurrently. then: Handler that receives the list of branch results. + mode: ``"all"`` continues once every branch has reported; ``"first"`` + continues as soon as one has, and cancels the rest. """ branches: tuple[Any, ...] then: Any + mode: Literal["all", "first"] = "all" def __post_init__(self): """Validate the fan-out. Raises: - WorkflowDefinitionError: If no branches were given. + WorkflowDefinitionError: If no branches were given, or the mode is + not recognised. """ if not self.branches: msg = ( @@ -977,9 +981,14 @@ def __post_init__(self): "run concurrently." ) raise WorkflowDefinitionError(msg) + if self.mode not in ("all", "first"): + msg = f'parallel() mode must be "all" or "first", got {self.mode!r}.' + raise WorkflowDefinitionError(msg) -def parallel(*branches: Any, then: Any) -> Parallel: +def parallel( + *branches: Any, then: Any, mode: Literal["all", "first"] = "all" +) -> Parallel: """Run branches concurrently, then continue with all their results. Each branch runs as its own child run, so branches retry and fail @@ -994,14 +1003,18 @@ def parallel(*branches: Any, then: Any) -> Parallel: The ``then`` handler receives one argument: the list of branch results, in the order the branches were given. + Pass ``mode="first"`` to race them instead: the run continues as soon as + one branch reports, and the others are cancelled. + Args: branches: Root events to run concurrently. - then: Handler to run once every branch has finished. + then: Handler to run once the fan-out is satisfied. + mode: ``"all"`` to wait for every branch, ``"first"`` to race them. Returns: The control return value. """ - return Parallel(branches=branches, then=then) + return Parallel(branches=branches, then=then, mode=mode) @dataclasses.dataclass(frozen=True, slots=True) diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 83f2a3f0f15..26c71fe150c 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -574,6 +574,47 @@ async def check_children_are_created_with_their_join(store: RunStore) -> None: assert len(await store.get_steps("child1")) == 1 +async def check_list_children_finds_a_joins_branches(store: RunStore) -> None: + """The children of one join slot are addressable without a full scan.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.WAITING, + state={}, + new_steps=( + make_step( + ordinal=1, + status=StepStatus.BLOCKED, + wait_key="join:1", + join_expected=2, + origin="join", + due_at=0.0, + ), + ), + next_ordinal=2, + children=( + ( + make_run("childA", parent_run_id="run1", parent_ordinal=1), + make_step("childA"), + ), + ( + make_run("childB", parent_run_id="run1", parent_ordinal=1), + make_step("childB"), + ), + ), + ), + NOW, + ) + children = await store.list_children("run1", 1) + assert {child.run_id for child in children} == {"childA", "childB"} + assert await store.list_children("run1", 2) == () + assert await store.list_children("nobody", 1) == () + + async def check_reads_do_not_alias_stored_state(store: RunStore) -> None: """Mutating a returned record must not change what the store holds.""" await store.admit(make_run(), make_step(), _ADMITTED) @@ -647,6 +688,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_an_early_delivery_is_buffered_then_consumed, check_early_deliveries_queue_in_order, check_children_are_created_with_their_join, + check_list_children_finds_a_joins_branches, check_join_arrivals_count_once, check_finalize_refuses_while_a_step_is_claimed, check_finalize_tombstones_open_slots, diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 4add7b47684..7107ac1b334 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -419,20 +419,48 @@ def _validate_handler_body( if node is None: return self_name = next(iter(inspect.signature(defn.fn).parameters), None) - for child in ast.walk(node): + + def durable_self_reference(value: ast.expr) -> str | None: + """Name the durable handler an expression reaches through ``self``. + + Args: + value: The expression to inspect. + + Returns: + The handler name, or None if this is not such a reference. + """ if ( - isinstance(child, ast.Call) - and isinstance(child.func, ast.Attribute) - and isinstance(child.func.value, ast.Name) - and child.func.value.id == self_name - and child.func.attr in durable_names + isinstance(value, ast.Attribute) + and isinstance(value.value, ast.Name) + and value.value.id == self_name + and value.attr in durable_names ): + return value.attr + return None + + called: set[int] = set() + for child in ast.walk(node): + if isinstance(child, ast.Call) and durable_self_reference(child.func): + called.add(id(child.func)) + for child in ast.walk(node): + if isinstance(child, ast.Call) and (name := durable_self_reference(child.func)): raise _error( workflow_cls, - f"handler {defn.name!r} calls {child.func.attr!r} directly, which " + f"handler {defn.name!r} calls {name!r} directly, which " "runs it inline and loses its retries, timeout, and effect " f"tracking. Return it as a transition instead: " - f"return {workflow_cls.__name__}.{child.func.attr}", + f"return {workflow_cls.__name__}.{name}(...)", + ) + if ( + isinstance(child, ast.Attribute) + and id(child) not in called + and (name := durable_self_reference(child)) + ): + raise _error( + workflow_cls, + f"handler {defn.name!r} refers to {name!r} as self.{name}, which " + "is a bound method, not a transition the engine can route. " + f"Name the class instead: {workflow_cls.__name__}.{name}", ) if ( isinstance(child, ast.Return) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index b33c4d21b4a..da5599abe4f 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -1241,14 +1241,18 @@ def _success_completion( status=StepStatus.BLOCKED, args={"__results__": []}, wait_key=f"join:{claim.run.next_ordinal}", - join_expected=len(control.branches), + join_expected=1 if control.mode == "first" else len(control.branches), origin="join", created_at=now, updated_at=now, ) events.append(( HistoryEventType.CHILD_STARTED, - {"ordinal": join.ordinal, "branches": len(control.branches)}, + { + "ordinal": join.ordinal, + "branches": len(control.branches), + "mode": control.mode, + }, )) return StepCompletion( step_status=StepStatus.SUCCEEDED, @@ -1890,7 +1894,7 @@ async def _report_outcome( """ if run.parent_run_id is None or run.parent_ordinal is None: return - await self._store.record_arrival( + disposition = await self._store.record_arrival( run.parent_run_id, run.parent_ordinal, { @@ -1902,8 +1906,32 @@ async def _report_outcome( run.run_id, self._clock(), ) + if disposition == "resolved": + await self._cancel_losing_branches(run) self._wakeup.set() + async def _cancel_losing_branches(self, winner: RunRecord) -> None: + """Cancel the siblings a decided race left running. + + A join satisfied before every branch reported means the rest can no + longer affect the outcome, so leaving them running would burn work and + keep making external calls nobody is waiting for. + + Args: + winner: The child whose arrival satisfied the join. + """ + if winner.parent_run_id is None or winner.parent_ordinal is None: + return + siblings = await self._store.list_children( + winner.parent_run_id, winner.parent_ordinal + ) + for sibling in siblings: + if ( + sibling.run_id != winner.run_id + and sibling.status not in TERMINAL_RUN_STATUSES + ): + await self.cancel(sibling.run_id) + async def _finalize_control(self, now: float) -> int: """Finalize every drained run awaiting a control transition. diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 1cfb30afc86..c0db292d39f 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -117,6 +117,12 @@ def register(self, workflow_cls: type[BaseState]) -> WorkflowDefinition: if self._kernel is not None: msg = "Cannot register workflows after the runtime has started." raise WorkflowRuntimeError(msg) + if not isinstance(workflow_cls, type): + msg = ( + f"Expected a workflow class, got {type(workflow_cls).__name__}. " + "Pass the classes themselves, one argument each." + ) + raise WorkflowDefinitionError(msg) existing_id = self._classes.get(workflow_cls) if existing_id is not None: return self._definitions[existing_id] diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index c2c5dde0759..6a51a1b7747 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -452,6 +452,24 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: """ ... + async def list_children( + self, parent_run_id: str, parent_ordinal: int + ) -> tuple[RunRecord, ...]: + """List the child runs admitted for one join slot. + + A decided race has to reach its losing branches, and an operator + inspecting a fan-out wants its children without paging the whole run + table, so this is a first-class lookup rather than a filtered listing. + + Args: + parent_run_id: The run that fanned out. + parent_ordinal: The join slot the children report to. + + Returns: + The child run records, oldest first. + """ + ... + async def find_by_request_key( self, workflow_id: str, request_key: str ) -> str | None: @@ -1360,6 +1378,28 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: matched.sort(key=lambda run: (run.created_at, run.run_id), reverse=True) return tuple(_detach_run(run) for run in matched[: query.limit]) + async def list_children( + self, parent_run_id: str, parent_ordinal: int + ) -> tuple[RunRecord, ...]: + """List the child runs admitted for one join slot. + + Args: + parent_run_id: The run that fanned out. + parent_ordinal: The join slot the children report to. + + Returns: + The child run records, oldest first. + """ + async with self._lock: + children = [ + run + for run in self._runs.values() + if run.parent_run_id == parent_run_id + and run.parent_ordinal == parent_ordinal + ] + children.sort(key=lambda run: (run.created_at, run.run_id)) + return tuple(_detach_run(run) for run in children) + async def find_by_request_key( self, workflow_id: str, request_key: str ) -> str | None: @@ -1662,6 +1702,10 @@ def _migrate(self) -> None: "CREATE INDEX IF NOT EXISTS idx_workflow_steps_lease" " ON workflow_steps (status, lease_expires_at)" ) + self._db.execute( + "CREATE INDEX IF NOT EXISTS idx_workflow_runs_parent" + " ON workflow_runs (parent_run_id, parent_ordinal)" + ) self._db.execute("COMMIT") except BaseException: self._db.execute("ROLLBACK") @@ -2722,6 +2766,26 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: ).fetchall() return tuple(_run_from_row(row) for row in rows) + async def list_children( + self, parent_run_id: str, parent_ordinal: int + ) -> tuple[RunRecord, ...]: + """List the child runs admitted for one join slot. + + Args: + parent_run_id: The run that fanned out. + parent_ordinal: The join slot the children report to. + + Returns: + The child run records, oldest first. + """ + with self._lock: + rows = self._db.execute( + "SELECT * FROM workflow_runs WHERE parent_run_id = ?" + " AND parent_ordinal = ? ORDER BY created_at, run_id", + (parent_run_id, parent_ordinal), + ).fetchall() + return tuple(_run_from_row(row) for row in rows) + async def find_by_request_key( self, workflow_id: str, request_key: str ) -> str | None: diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 8ce28af7cf7..0ed0cbcd0d6 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -294,3 +294,160 @@ def later(self): assert snapshot is not None assert snapshot.status is RunStatus.COMPLETED assert snapshot.state["outcomes"] == ["COMPLETED", "TIMED_OUT"] + + +class SlowVendor(rx.State): + """A branch that waits before answering.""" + + __workflow__ = WorkflowConfig(id="race.slow") + quote: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="read") + def start(self, price: int): + """Quote after a delay. + + Args: + price: The price to quote. + + Returns: + A deferral, then completion. + """ + RACE_CALLS.append("slow") + return rx.after("1h", SlowVendor.answer(price)) + + @rx.event(durable=True, effect="read") + def answer(self, price: int): + """Deliver the delayed quote. + + Args: + price: The price to quote. + + Returns: + Completion carrying the quote. + """ + RACE_CALLS.append("slow-answer") + self.quote = price + return rx.complete(result={"price": price}) + + +class FastVendor(rx.State): + """A branch that answers immediately.""" + + __workflow__ = WorkflowConfig(id="race.fast") + quote: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="read") + def start(self, price: int): + """Quote at once. + + Args: + price: The price to quote. + + Returns: + Completion carrying the quote. + """ + RACE_CALLS.append("fast") + self.quote = price + return rx.complete(result={"price": price}) + + +class Shopper(rx.State): + """Takes the first quote back and abandons the rest.""" + + __workflow__ = WorkflowConfig(id="race.shopper") + winner: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Ask both vendors and take whoever answers first. + + Returns: + A racing fan-out. + """ + return rx.parallel( + SlowVendor.start(100), + FastVendor.start(90), + then=Shopper.book, + mode="first", + ) + + @rx.event(durable=True, effect="read") + def book(self, results: list): + """Book the quote that arrived first. + + Args: + results: One entry, from the branch that won. + + Returns: + Completion. + """ + self.winner = results[0]["result"]["price"] + return rx.complete(result={"booked": self.winner}) + + +RACE_CALLS: list[str] = [] + + +async def _join_slot(harness: WorkflowTestHarness, run_id: str): + """Find the fan-out slot of a run. + + Args: + harness: The running harness. + run_id: The parent run. + + Returns: + The join step record. + """ + steps = await harness.kernel.store.get_steps(run_id) + return next(step for step in steps if step.wait_key is not None) + + +async def test_race_continues_on_the_first_branch_and_cancels_the_rest(): + """mode="first" resumes as soon as one branch reports. + + A quote race is worthless if it still waits for the slow vendor, and worse + than worthless if that vendor keeps working after the order is booked. + """ + RACE_CALLS.clear() + async with WorkflowTestHarness(Shopper, SlowVendor, FastVendor) as harness: + result = await harness.start(Shopper.start()) + assert result.run_id is not None + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"booked": 90} + + join = await _join_slot(harness, result.run_id) + children = await harness.kernel.store.list_children(result.run_id, join.ordinal) + by_workflow = {child.workflow_id: child for child in children} + assert by_workflow["race.fast"].status is RunStatus.COMPLETED + assert by_workflow["race.slow"].status is RunStatus.CANCELLED + + # The loser's deferred step never runs, even once its timer comes due. + await harness.advance("2h") + assert "slow-answer" not in RACE_CALLS + + +async def test_race_mode_still_starts_every_branch(): + """Racing is not a way to skip work: every branch starts.""" + RACE_CALLS.clear() + async with WorkflowTestHarness(Shopper, SlowVendor, FastVendor) as harness: + result = await harness.start(Shopper.start()) + assert result.run_id is not None + assert "slow" in RACE_CALLS + assert "fast" in RACE_CALLS + join = await _join_slot(harness, result.run_id) + children = await harness.kernel.store.list_children(result.run_id, join.ordinal) + assert len(children) == 2 + + +async def test_race_join_expects_one_arrival(): + """The join slot itself carries the racing intent.""" + RACE_CALLS.clear() + async with WorkflowTestHarness(Shopper, SlowVendor, FastVendor) as harness: + result = await harness.start(Shopper.start()) + assert result.run_id is not None + join = await _join_slot(harness, result.run_id) + assert join.join_expected == 1 + assert join.status is StepStatus.SUCCEEDED From 386a5fce8b6a8401e7bbd3eca2c633a59bb8e49f Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 17:14:32 -0700 Subject: [PATCH 025/121] Space a throttled backlog instead of bunching it Throttle deferred every excess start by exactly one window, so a burst of a hundred with limit=10 admitted ten now and scheduled ninety for the same instant one window later. That is not throttling, it is a delay line: the downstream the throttle exists to protect sees the same spike, just later, and the next window then holds ninety starts against a limit of ten. Each start is now placed at least a window after the limit-th most recent scheduled start under its key, which spaces the backlog at exactly the configured rate and holds the sliding-window bound rather than a per-window one: any window of length period contains at most limit starts, because a start is always a full window after its limit-th predecessor. This needs a new store query, nth_recent_start, since counting admissions in a window cannot see where the deferred ones are already scheduled. A run's scheduled start is when its root slot comes due, so a debounced or throttled run counts at the time it will run, not the time it was admitted. Covered by a conformance check, so a future store has to answer it the same way. The regression test asserts the schedule directly (0, 0, 10, 10, 20, 20 for six starts at limit=2 over ten seconds) and then advances the clock to confirm the runs execute on it. --- docs/workflows/overview.md | 4 +- news/workflow-throttle-spacing.feature.md | 1 + reflex/workflow/conformance.py | 30 ++++++++ reflex/workflow/kernel.py | 12 +-- reflex/workflow/store.py | 90 +++++++++++++++++++++++ tests/units/workflow/test_flow_control.py | 46 ++++++++++++ 6 files changed, 177 insertions(+), 6 deletions(-) create mode 100644 news/workflow-throttle-spacing.feature.md diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index 87a73b8f58a..dbceb3fd1ad 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -290,7 +290,9 @@ def sync(self, customer_id: str): ... Debounce is the one to reach for with chatty webhooks: ten deliveries in a second become one run. Rate limiting drops excess starts, which is what you want when a provider can flood you; throttling -delays them instead, which is what you want when every start matters but the downstream is slow. +delays them instead, which is what you want when every start matters but the downstream is slow. A +throttled backlog is spaced out rather than released together: with `limit=2, period="10s"`, a burst +of six runs starts two now, two in ten seconds, and two in twenty. ## Inspecting and steering runs diff --git a/news/workflow-throttle-spacing.feature.md b/news/workflow-throttle-spacing.feature.md new file mode 100644 index 00000000000..90fb93c2e08 --- /dev/null +++ b/news/workflow-throttle-spacing.feature.md @@ -0,0 +1 @@ +`rx.Throttle` now spaces a held-back backlog across windows instead of releasing all of it at once. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 26c71fe150c..32a5f45901e 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -509,6 +509,35 @@ async def check_flow_control_queries(store: RunStore) -> None: assert steps[0].due_at == pytest.approx(NOW + 30) +async def check_nth_recent_start_orders_by_scheduled_time(store: RunStore) -> None: + """Throttling can find the nth most recent scheduled start under a key.""" + for index, offset in enumerate((0.0, 5.0, 20.0)): + await store.admit( + make_run(f"s{index}", flow_key="k1", created_at=NOW + offset), + make_step(f"s{index}"), + _ADMITTED, + ) + await store.admit( + make_run("other", flow_key="k2", created_at=NOW + 99), + make_step("other"), + _ADMITTED, + ) + assert await store.nth_recent_start("conformance.flow", "k1", 1) == pytest.approx( + NOW + 20 + ) + assert await store.nth_recent_start("conformance.flow", "k1", 3) == pytest.approx( + NOW + ) + assert await store.nth_recent_start("conformance.flow", "k1", 4) is None + assert await store.nth_recent_start("conformance.flow", "missing", 1) is None + + # A deferred run is scheduled by its due time, not by when it was admitted. + assert await store.defer_root("s0", NOW + 50, NOW) + assert await store.nth_recent_start("conformance.flow", "k1", 1) == pytest.approx( + NOW + 50 + ) + + async def check_early_deliveries_queue_in_order(store: RunStore) -> None: """Several signals arriving before a wait are kept, not overwritten.""" await store.admit(make_run(), make_step(), _ADMITTED) @@ -697,4 +726,5 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_pagination_skips_nothing_on_tied_timestamps, check_label_filter_handles_awkward_keys, check_flow_control_queries, + check_nth_recent_start_orders_by_scheduled_time, ) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index da5599abe4f..606277710fe 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -481,12 +481,14 @@ async def _apply_start_policy( ) if handler.throttle is not None: window = parse_duration(handler.throttle.period) - started = await self._store.count_started_since( - defn.workflow_id, flow_key, now - window + previous = await self._store.nth_recent_start( + defn.workflow_id, flow_key, handler.throttle.limit ) - if started >= handler.throttle.limit: - # Delay rather than drop: the excess still runs, just later. - return None, now + window + if previous is not None and previous + window > now: + # Delay rather than drop, and space the backlog: each start + # sits a window after the limit-th most recent one, so a held + # burst is released at the configured rate instead of at once. + return None, previous + window if handler.debounce is not None: window = parse_duration(handler.debounce.period) pending = await self._store.first_active(defn.workflow_id, flow_key) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 6a51a1b7747..14636b8c7e6 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -339,6 +339,30 @@ async def count_started_since( """ ... + async def nth_recent_start( + self, workflow_id: str, flow_key: str, n: int + ) -> float | None: + """Find the nth most recent scheduled start under a flow key. + + Throttling has to place each new run relative to the ones already + scheduled, not merely count them: deferring every excess start by one + window replays the burst intact, one window later. Keeping each start + at least a window after the nth most recent one spaces the backlog and + holds the sliding-window limit. + + A run's scheduled start is when its root slot comes due, which for an + undeferred run is when it was admitted. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + n: How far back to look, counting from the most recent as 1. + + Returns: + The scheduled start, or None when fewer than n runs exist. + """ + ... + async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: """Push a not-yet-started run's root slot later, for debouncing. @@ -1147,6 +1171,40 @@ async def count_started_since( and run.created_at > since ) + async def nth_recent_start( + self, workflow_id: str, flow_key: str, n: int + ) -> float | None: + """Find the nth most recent scheduled start under a flow key. + + Throttling has to place each new run relative to the ones already + scheduled, not merely count them: deferring every excess start by one + window replays the burst intact, one window later. Keeping each start + at least a window after the nth most recent one spaces the backlog and + holds the sliding-window limit. + + A run's scheduled start is when its root slot comes due, which for an + undeferred run is when it was admitted. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + n: How far back to look, counting from the most recent as 1. + + Returns: + The scheduled start, or None when fewer than n runs exist. + """ + async with self._lock: + starts = sorted( + ( + max(steps[0].due_at, run.created_at) + for run in self._runs.values() + if run.workflow_id == workflow_id and run.flow_key == flow_key + if (steps := self._steps.get(run.run_id)) + ), + reverse=True, + ) + return starts[n - 1] if len(starts) >= n else None + async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: """Push a not-yet-started run's root slot later, for debouncing. @@ -2444,6 +2502,38 @@ async def count_started_since( ).fetchone() return row["n"] + async def nth_recent_start( + self, workflow_id: str, flow_key: str, n: int + ) -> float | None: + """Find the nth most recent scheduled start under a flow key. + + Throttling has to place each new run relative to the ones already + scheduled, not merely count them: deferring every excess start by one + window replays the burst intact, one window later. Keeping each start + at least a window after the nth most recent one spaces the backlog and + holds the sliding-window limit. + + A run's scheduled start is when its root slot comes due, which for an + undeferred run is when it was admitted. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + n: How far back to look, counting from the most recent as 1. + + Returns: + The scheduled start, or None when fewer than n runs exist. + """ + with self._lock: + row = self._db.execute( + "SELECT MAX(s.due_at, r.created_at) AS start FROM workflow_runs r" + " JOIN workflow_steps s ON s.run_id = r.run_id AND s.ordinal = 0" + " WHERE r.workflow_id = ? AND r.flow_key = ?" + " ORDER BY start DESC LIMIT 1 OFFSET ?", + (workflow_id, flow_key, n - 1), + ).fetchone() + return None if row is None else row["start"] + async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: """Push a not-yet-started run's root slot later, for debouncing. diff --git a/tests/units/workflow/test_flow_control.py b/tests/units/workflow/test_flow_control.py index 5c9b98acf80..27ef05a33f2 100644 --- a/tests/units/workflow/test_flow_control.py +++ b/tests/units/workflow/test_flow_control.py @@ -274,3 +274,49 @@ async def test_singleton_cancel_keeps_one_active_run_per_key( await harness.kernel.start(flow.start("acme")) active = await harness.kernel.store.count_active("flow.sync", "start:'acme'") assert active == 1 + + +async def test_throttle_spaces_a_backlog_instead_of_bunching_it( + forked_registration_context, +): + """A held-back burst must not all fire at the same instant. + + Deferring every excess start by one window turns a burst of a hundred into + a burst of ninety, one window later. The provider the throttle exists to + protect sees the same spike, just delayed. + """ + calls: list[int] = [] + + class Spaced(rx.State): + __workflow__ = WorkflowConfig(id="flow.spaced") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + throttle=Throttle(limit=2, period="10s"), + ) + def start(self): + """Record that this run ran.""" + calls.append(1) + + async with WorkflowTestHarness(Spaced) as harness: + begin = harness.now + run_ids = [] + for _ in range(6): + result = await harness.start(Spaced.start) + assert result.disposition == "started" + assert result.run_id is not None + run_ids.append(result.run_id) + assert len(calls) == 2 + + due = [] + for run_id in run_ids: + steps = await harness.kernel.store.get_steps(run_id) + due.append(max(steps[0].due_at, begin) - begin) + assert due == [0.0, 0.0, 10.0, 10.0, 20.0, 20.0] + + await harness.advance("10s") + assert len(calls) == 4 + await harness.advance("10s") + assert len(calls) == 6 From 5692517b6b15165159a0ff1d183765e37e68b07f Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 18:06:48 -0700 Subject: [PATCH 026/121] Add a Postgres run store, and run the suite against it SQLite takes one writer, so a deployment was capped at one worker process per database file. That is the ceiling that keeps this from being comparable to Temporal or Inngest at all, and it is not something the kernel can fix -- it is the store. PostgresRunStore claims a run's frontier step with FOR UPDATE ... SKIP LOCKED, so workers never queue behind each other and never take the same step. Adding a process adds throughput. A worker that dies mid-step holds a lease that another worker may only reclaim once it lapses, so a slow step is never duplicated -- there is a test for exactly that, and one asserting twenty non_idempotent_write runs across two kernels execute exactly once each, with both kernels provably taking a share of the work. The store is not trusted on assertion: the 29 conformance checks run against it, and the whole workflow suite -- 260 harness tests -- now runs a third time against a real server whenever REFLEX_TEST_POSTGRES names one. That third parameter immediately paid for itself twice. It caught a test of mine that asserted a race always starts every branch, which is not an invariant: a loser cancelled before its first step is the better outcome, and only Postgres's ordering exposed the assumption. And driving the real CLI against a Postgres URL surfaced that each command ran its own asyncio.run, which is fine for a file and fails for a pool, whose connections belong to the loop that opened them. Commands now run in one loop. Two things came out of the port. Pyright rejected the schema name spliced into DDL, since psycopg types raw SQL as LiteralString; the name is now composed as an identifier, which is a better guarantee than the check it replaces. And the observer turned out to see only admissions and commits -- not attempt starts, which are the spans a tracer actually wants -- so every recording site now reports, correlated to its workflow. A throwaway schema per test isolates them. Dropping one first evicts its own backends, because a test loop that dies mid-transaction leaves a pooled connection idle holding locks, and the DROP would otherwise wait on it forever. That was a real hang, reproducible only under random test ordering. Postgres is optional: pip install 'psycopg[binary,pool]'. --- docs/workflows/overview.md | 29 +- news/workflow-observer-coverage.feature.md | 1 + news/workflow-postgres-store.feature.md | 1 + pyproject.toml | 2 +- reflex/workflow/cli.py | 131 +- reflex/workflow/kernel.py | 113 +- reflex/workflow/postgres.py | 1558 ++++++++++++++++++++ tests/units/workflow/conftest.py | 53 +- tests/units/workflow/test_parallel.py | 14 +- tests/units/workflow/test_postgres.py | 283 ++++ uv.lock | 81 +- 11 files changed, 2140 insertions(+), 126 deletions(-) create mode 100644 news/workflow-observer-coverage.feature.md create mode 100644 news/workflow-postgres-store.feature.md create mode 100644 reflex/workflow/postgres.py create mode 100644 tests/units/workflow/test_postgres.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index dbceb3fd1ad..0bd2be5e261 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -382,13 +382,30 @@ to a waiting run, and `harness.cancel(...)` and `harness.resume(...)` drive the ## Deploying -Runs persist to a SQLite file next to your app by default; pass `rx.App(workflow_store=...)` to -choose another store. +Runs persist to a SQLite file next to your app by default, which is the right choice for local +development and for a single-process deployment. Run **one worker process per SQLite database +file**: the store's calls are synchronous on the event loop that also serves your app, so two +processes writing the same file contend for it. Contention is bounded to a short busy timeout and +surfaces as a transient error the kernel retries, but throughput does not improve. -Run **one worker process per SQLite database file**. The store's calls are synchronous on the event -loop that also serves your app, so two processes writing the same file contend for it; contention is -bounded to a short busy timeout and surfaces as a transient error the kernel retries, but throughput -does not improve. Horizontal scale wants a store that supports concurrent writers. +For more than one process, point the app at Postgres: + +```python +from reflex.workflow.postgres import PostgresRunStore + +app = rx.App(workflow_store=PostgresRunStore("postgresql://user:pw@host/db")) +``` + +Every process that opens the same database is a worker. A claim locks one run's next step with +`FOR UPDATE ... SKIP LOCKED`, so workers never queue behind each other and never take the same +step; adding a process adds throughput. A worker that dies mid-step holds a lease, and the step is +reclaimed by another worker once that lease lapses -- not before, so a slow step is never +duplicated. Install it with `pip install 'psycopg[binary,pool]'`. + +Pass `schema=` to keep a deployment's tables in their own namespace inside a shared database. + +The `reflex workflows` commands take the same target: `--database postgresql://...`, or set +`REFLEX_WORKFLOW_DATABASE` once. `RunStore` is a supported extension point, and the invariants a store must satisfy ship as runnable checks rather than prose: diff --git a/news/workflow-observer-coverage.feature.md b/news/workflow-observer-coverage.feature.md new file mode 100644 index 00000000000..386c8919a87 --- /dev/null +++ b/news/workflow-observer-coverage.feature.md @@ -0,0 +1 @@ +The workflow observer now sees attempt starts, cancellations, signal deliveries, joins, resumes, and finalizations, not just admissions and commits. diff --git a/news/workflow-postgres-store.feature.md b/news/workflow-postgres-store.feature.md new file mode 100644 index 00000000000..032e24eaea0 --- /dev/null +++ b/news/workflow-postgres-store.feature.md @@ -0,0 +1 @@ +`reflex.workflow.postgres.PostgresRunStore` runs workflows across many worker processes against one Postgres database. diff --git a/pyproject.toml b/pyproject.toml index 4a3759b3775..6ae0499cea8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,7 +92,7 @@ dev = [ "plotly", "pre-commit", "psutil", - "psycopg[binary]", + "psycopg[binary,pool]", "pydantic", "pyright", "pytest-asyncio", diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index c736cb1d19c..edcde320fd9 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -9,7 +9,9 @@ from __future__ import annotations import asyncio +import inspect import json +import os from typing import TYPE_CHECKING, Any import click @@ -18,7 +20,7 @@ from reflex.workflow.records import RunStatus if TYPE_CHECKING: - from collections.abc import Awaitable + from collections.abc import Awaitable, Callable from reflex.workflow.store import RunStore @@ -28,27 +30,59 @@ def _open_store(database: str | None) -> RunStore: """Open the run store the app persists to. + A ``postgres://`` or ``postgresql://`` target opens the Postgres store; + anything else is a path to a SQLite file. + Args: - database: Path to the SQLite database, or None for the default. + database: Connection URL or SQLite path, or None for the default. Returns: The store. """ + target = database or os.environ.get("REFLEX_WORKFLOW_DATABASE") + if target is not None and target.startswith(("postgres://", "postgresql://")): + from reflex.workflow.postgres import PostgresRunStore + + return PostgresRunStore(target) + from reflex.workflow.store import SqliteRunStore - return SqliteRunStore(database or DEFAULT_DB_FILENAME) + return SqliteRunStore(target or DEFAULT_DB_FILENAME) -def _run(coroutine: Awaitable[Any]) -> Any: - """Run one store coroutine to completion. +def _with_store(database: str | None, work: Callable[[RunStore], Awaitable[Any]]): + """Open a store, run one unit of work against it, and close it. + + Everything happens inside a single event loop. A pooled store binds its + connections to the loop that opened them, so a command that ran each query + in its own ``asyncio.run`` would fail on close, and the failure would only + appear against Postgres. Args: - coroutine: The coroutine to run. + database: Connection URL or SQLite path, or None for the default. + work: What to do with the open store. Returns: - Its result. + Whatever the work returned. """ - return asyncio.run(coroutine) # pyright: ignore[reportArgumentType] + + async def session() -> Any: + """Open the store, do the work, and close it. + + Returns: + The work's result. + """ + store = _open_store(database) + try: + return await work(store) + finally: + closer = getattr(store, "close", None) + if closer is not None: + closed = closer() + if inspect.isawaitable(closed): + await closed + + return asyncio.run(session()) def _age(seconds: float) -> str: @@ -70,7 +104,10 @@ def _age(seconds: float) -> str: "--database", "-d", default=None, - help="Path to the workflow database. Defaults to ./workflow.db.", + help=( + "Workflow database: a Postgres URL, or a path to a SQLite file. " + "Defaults to $REFLEX_WORKFLOW_DATABASE, then ./workflow.db." + ), ) @@ -105,20 +142,13 @@ def list_runs( from reflex.workflow.records import RunQuery label_filter = dict(pair.split("=", 1) for pair in labels if "=" in pair) - store = _open_store(database) - try: - runs = _run( - store.list_runs( - RunQuery( - workflow_id=workflow, - statuses=tuple(RunStatus(value.upper()) for value in statuses), - labels=label_filter or None, - limit=limit, - ) - ) - ) - finally: - _close(store) + query = RunQuery( + workflow_id=workflow, + statuses=tuple(RunStatus(value.upper()) for value in statuses), + labels=label_filter or None, + limit=limit, + ) + runs = _with_store(database, lambda store: store.list_runs(query)) if as_json: click.echo( @@ -156,16 +186,26 @@ def list_runs( @click.option("--history", is_flag=True, help="Include the run's history.") def show(database: str | None, run_id: str, as_json: bool, history: bool): """Show one run's state, steps, and optionally its history.""" - store = _open_store(database) - try: - run = _run(store.get_run(run_id)) - if run is None: - console.error(f"No run {run_id!r} in this database.") - raise click.exceptions.Exit(1) - steps = _run(store.get_steps(run_id)) - events = _run(store.get_history(run_id)) if history else () - finally: - _close(store) + + async def load(store: RunStore): + """Read the run, its slots, and optionally its history. + + Args: + store: The open run store. + + Returns: + The run, its steps, and its history events. + """ + return ( + await store.get_run(run_id), + await store.get_steps(run_id), + await store.get_history(run_id) if history else (), + ) + + run, steps, events = _with_store(database, load) + if run is None: + console.error(f"No run {run_id!r} in this database.") + raise click.exceptions.Exit(1) if as_json: click.echo( @@ -232,11 +272,9 @@ def cancel(database: str | None, run_id: str): """ import time - store = _open_store(database) - try: - recorded = _run(store.request_cancel(run_id, time.time())) - finally: - _close(store) + recorded = _with_store( + database, lambda store: store.request_cancel(run_id, time.time()) + ) if not recorded: console.error(f"Run {run_id!r} is unknown or already finished.") raise click.exceptions.Exit(1) @@ -250,23 +288,8 @@ def resume(database: str | None, run_id: str): """Re-open a run suspended for operator attention.""" import time - store = _open_store(database) - try: - resumed = _run(store.resume_run(run_id, time.time())) - finally: - _close(store) + resumed = _with_store(database, lambda store: store.resume_run(run_id, time.time())) if not resumed: console.error(f"Run {run_id!r} is not suspended.") raise click.exceptions.Exit(1) console.print(f"Resumed {run_id}; its next step will run.") - - -def _close(store: RunStore) -> None: - """Close a store that holds a connection. - - Args: - store: The store to close. - """ - closer = getattr(store, "close", None) - if closer is not None: - closer() diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 606277710fe..480fc466dc4 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -654,6 +654,9 @@ async def cancel(self, run_id: str) -> bool: """ recorded = await self._store.request_cancel(run_id, self._clock()) if recorded: + await self._notify_run( + run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),) + ) task = self._inflight.get(run_id) if task is not None: task.cancel() @@ -685,7 +688,26 @@ async def signal( self._clock(), ) if disposition == "resolved": + await self._notify_run( + run_id, + ( + ( + HistoryEventType.WAIT_RESOLVED, + {"wait_key": f"sig:{delivery.channel}"}, + ), + ), + ) self._wakeup.set() + elif disposition == "buffered": + await self._notify_run( + run_id, + ( + ( + HistoryEventType.SIGNAL_BUFFERED, + {"wait_key": f"sig:{delivery.channel}"}, + ), + ), + ) return disposition async def resume(self, run_id: str) -> bool: @@ -699,6 +721,7 @@ async def resume(self, run_id: str) -> bool: """ resumed = await self._store.resume_run(run_id, self._clock()) if resumed: + await self._notify_run(run_id, ((HistoryEventType.RUN_RESUMED, {}),)) self._wakeup.set() return resumed @@ -1377,6 +1400,48 @@ def _success_completion( events=tuple(events), ) + async def _record( + self, + run: RunRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Append history events and tell the observer about them. + + Args: + run: The run the transitions belong to. + events: The (type, data) pairs to record. + now: Current time in epoch seconds. + """ + await self._store.append_events(run.run_id, events, now) + self._notify(run, events) + + async def _notify_run( + self, + run_id: str, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + ) -> None: + """Tell the observer about transitions a store operation recorded. + + Some transitions are written by the store itself, inside the same + transaction as the state change they describe. The kernel knows which + ones those are from the operation's result, and reports them here so + the observer sees one stream rather than the subset the kernel happens + to construct itself. + + The run is loaded only to correlate the events with their workflow, so + a deployment without an observer pays nothing for this. + + Args: + run_id: The run the transitions belong to. + events: The (type, data) pairs the store recorded. + """ + if self._observer is None: + return + run = await self._store.get_run(run_id) + if run is not None: + self._notify(run, events) + def _notify( self, run: RunRecord, @@ -1529,8 +1594,8 @@ async def _record_abandoned( handler: The handler that was executing. reason: Why the claim was lost. """ - await self._store.append_events( - claim.run.run_id, + await self._record( + claim.run, ( ( HistoryEventType.ATTEMPT_ABANDONED, @@ -1652,8 +1717,8 @@ async def _execute_claim(self, claim: Claim) -> None: if expired is not None: handler = expired steps = await self._store.get_steps(claim.run.run_id) - await self._store.append_events( - claim.run.run_id, + await self._record( + claim.run, ( ( HistoryEventType.ATTEMPT_STARTED, @@ -1692,17 +1757,19 @@ async def _execute_claim(self, claim: Claim) -> None: await self._record_abandoned(claim, handler, "lease_lost") return if await self._cancel_requested(claim.run.run_id): + cancelled = ( + ( + HistoryEventType.ATTEMPT_CANCELLED, + {"ordinal": claim.step.ordinal}, + ), + ) await self._store.release_claim( claim, status=StepStatus.CANCELLED, - events=( - ( - HistoryEventType.ATTEMPT_CANCELLED, - {"ordinal": claim.step.ordinal}, - ), - ), + events=cancelled, now=self._clock(), ) + self._notify(claim.run, cancelled) return current = asyncio.current_task() if self._closing or (current is not None and current.cancelling()): @@ -1908,6 +1975,16 @@ async def _report_outcome( run.run_id, self._clock(), ) + if disposition in ("resolved", "counted"): + await self._notify_run( + run.parent_run_id, + ( + ( + HistoryEventType.CHILD_RESOLVED, + {"ordinal": run.parent_ordinal, "child": run.run_id}, + ), + ), + ) if disposition == "resolved": await self._cancel_losing_branches(run) self._wakeup.set() @@ -1948,17 +2025,15 @@ async def _finalize_control(self, now: float) -> int: cancelled = run.cancel_requested status = RunStatus.CANCELLED if cancelled else RunStatus.TIMED_OUT error = None if cancelled else {"reason": "run_timeout"} + event = ( + HistoryEventType.RUN_CANCELLED + if cancelled + else HistoryEventType.RUN_TIMED_OUT + ) if await self._store.finalize_run( - run.run_id, - status=status, - error=error, - event=( - HistoryEventType.RUN_CANCELLED - if cancelled - else HistoryEventType.RUN_TIMED_OUT - ), - now=now, + run.run_id, status=status, error=error, event=event, now=now ): + self._notify(run, ((event, {} if error is None else dict(error)),)) await self._report_outcome(run, status, None, error) finalized += 1 return finalized diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py new file mode 100644 index 00000000000..32cbc14022f --- /dev/null +++ b/reflex/workflow/postgres.py @@ -0,0 +1,1558 @@ +"""Run store backed by PostgreSQL. + +SQLite is a fine local tier, but it takes one writer at a time, which caps a +deployment at one worker process per database file. Postgres is where the +engine earns the comparison: many workers claim from the same mailbox +concurrently, each taking a different run because a claim locks the frontier +row with ``FOR UPDATE ... SKIP LOCKED`` rather than locking the whole store. + +The semantics are the ones ``reflex.workflow.conformance`` fixes; this module +is a second implementation of them, not a second definition. Anything this +store does differently from ``SqliteRunStore`` is a dialect difference, and the +conformance suite is what proves it. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +from typing import TYPE_CHECKING, Any, Final + +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import DEFAULT_LEASE_DURATION + +from reflex.workflow.records import ( + CLAIMABLE_STEP_STATUSES, + TERMINAL_RUN_STATUSES, + TERMINAL_STEP_STATUSES, + HistoryEvent, + HistoryEventType, + RunRecord, + RunStatus, + StepRecord, + StepStatus, +) +from reflex.workflow.store import Claim, StaleClaimError + +try: + import psycopg + import psycopg_pool + from psycopg.rows import dict_row + from psycopg.sql import SQL, Composed, Identifier + from psycopg.types.json import Jsonb +except ImportError as exc: # pragma: no cover - depends on the environment + msg = ( + "reflex.workflow.postgres needs psycopg with its pool extra. " + "Install it with: pip install 'psycopg[binary,pool]'" + ) + raise WorkflowRuntimeError(msg) from exc + +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping + + from psycopg import AsyncConnection + from psycopg.rows import DictRow + + from reflex.workflow.records import RunQuery + from reflex.workflow.store import DeliveryDisposition, StepCompletion + + Connection = AsyncConnection[DictRow] + +DEFAULT_POOL_SIZE: Final = 10 + +_SCHEMA: Final = """ +CREATE TABLE IF NOT EXISTS workflow_runs ( + run_id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + definition_digest TEXT NOT NULL, + status TEXT NOT NULL, + state JSONB NOT NULL, + state_version BIGINT NOT NULL, + next_ordinal INTEGER NOT NULL, + result JSONB, + error JSONB, + flow_key TEXT, + parent_run_id TEXT, + parent_ordinal INTEGER, + request_key TEXT, + labels JSONB, + deadline DOUBLE PRECISION, + cancel_requested BOOLEAN NOT NULL DEFAULT FALSE, + created_at DOUBLE PRECISION NOT NULL, + updated_at DOUBLE PRECISION NOT NULL +); +CREATE TABLE IF NOT EXISTS workflow_steps ( + run_id TEXT NOT NULL, + ordinal INTEGER NOT NULL, + handler_id TEXT NOT NULL, + status TEXT NOT NULL, + args JSONB NOT NULL, + attempts INTEGER NOT NULL DEFAULT 0, + recoveries INTEGER NOT NULL DEFAULT 0, + due_at DOUBLE PRECISION NOT NULL DEFAULT 0, + epoch BIGINT NOT NULL DEFAULT 0, + lease_expires_at DOUBLE PRECISION NOT NULL DEFAULT 0, + wait_key TEXT, + join_expected INTEGER NOT NULL DEFAULT 0, + join_arrived INTEGER NOT NULL DEFAULT 0, + error JSONB, + origin TEXT NOT NULL, + created_at DOUBLE PRECISION NOT NULL, + updated_at DOUBLE PRECISION NOT NULL, + PRIMARY KEY (run_id, ordinal) +); +CREATE TABLE IF NOT EXISTS workflow_history ( + run_id TEXT NOT NULL, + seq INTEGER NOT NULL, + type TEXT NOT NULL, + at DOUBLE PRECISION NOT NULL, + data JSONB NOT NULL, + PRIMARY KEY (run_id, seq) +); +CREATE TABLE IF NOT EXISTS workflow_dedupe ( + workflow_id TEXT NOT NULL, + request_key TEXT NOT NULL, + run_id TEXT NOT NULL, + PRIMARY KEY (workflow_id, request_key) +); +CREATE TABLE IF NOT EXISTS workflow_inbox ( + run_id TEXT NOT NULL, + wait_key TEXT NOT NULL, + dedupe_key TEXT NOT NULL, + seq INTEGER NOT NULL, + payload JSONB NOT NULL, + status TEXT NOT NULL, + created_at DOUBLE PRECISION NOT NULL, + PRIMARY KEY (run_id, wait_key, dedupe_key) +); +CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs (status); +CREATE INDEX IF NOT EXISTS idx_workflow_runs_flow + ON workflow_runs (workflow_id, flow_key); +CREATE INDEX IF NOT EXISTS idx_workflow_runs_parent + ON workflow_runs (parent_run_id, parent_ordinal); +CREATE INDEX IF NOT EXISTS idx_workflow_steps_claimable + ON workflow_steps (status, due_at); +CREATE INDEX IF NOT EXISTS idx_workflow_steps_lease + ON workflow_steps (status, lease_expires_at); +CREATE INDEX IF NOT EXISTS idx_workflow_inbox_pending + ON workflow_inbox (run_id, wait_key, status, seq); +""" + +_TERMINAL_RUNS: Final = [status.value for status in TERMINAL_RUN_STATUSES] +_TERMINAL_STEPS: Final = [status.value for status in TERMINAL_STEP_STATUSES] +_CLAIMABLE_STEPS: Final = [status.value for status in CLAIMABLE_STEP_STATUSES] + +# A slot may be claimed when it is due, or when a wait's deadline has arrived: +# claiming a blocked slot is the timeout branch. A deadline of zero means the +# wait has none, so it is never claimable on the clock alone. +_CLAIMABLE_PREDICATE: Final = ( + "((s.status = ANY(%(claimable)s) AND s.due_at <= %(now)s)" + " OR (s.status = 'BLOCKED' AND s.due_at > 0 AND s.due_at <= %(now)s))" +) + +_FRONTIER_PREDICATE: Final = ( + "s.ordinal = (SELECT MIN(x.ordinal) FROM workflow_steps x" + " WHERE x.run_id = s.run_id AND NOT (x.status = ANY(%(terminal_steps)s)))" +) + +_RUNNABLE_PREDICATE: Final = ( + "NOT (r.status = ANY(%(terminal_runs)s)) AND r.status <> 'NEEDS_ATTENTION'" + " AND NOT r.cancel_requested" + " AND (r.deadline IS NULL OR r.deadline > %(now)s)" +) + + +def _set_search_path(schema: str) -> Composed: + """Build the statement pointing a connection at a schema. + + Args: + schema: The schema name. + + Returns: + The composed statement, with the name quoted as an identifier. + """ + return SQL("SET search_path TO {}").format(Identifier(schema)) + + +def _json(value: Any) -> Any: + """Wrap a payload for a JSONB column. + + psycopg has no default adapter for dict or list, so every JSON value has to + say what it is. + + Args: + value: The JSON-compatible value, or None. + + Returns: + The wrapped value, or None. + """ + return None if value is None else Jsonb(value) + + +def _run_from_row(row: Mapping[str, Any]) -> RunRecord: + """Build a run record from a database row. + + Args: + row: The ``workflow_runs`` row. + + Returns: + The run record. + """ + return RunRecord( + run_id=row["run_id"], + workflow_id=row["workflow_id"], + definition_digest=row["definition_digest"], + status=RunStatus(row["status"]), + state=row["state"], + state_version=row["state_version"], + next_ordinal=row["next_ordinal"], + result=row["result"], + error=row["error"], + flow_key=row["flow_key"], + parent_run_id=row["parent_run_id"], + parent_ordinal=row["parent_ordinal"], + request_key=row["request_key"], + labels=row["labels"], + deadline=row["deadline"], + cancel_requested=row["cancel_requested"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + +def _step_from_row(row: Mapping[str, Any]) -> StepRecord: + """Build a step record from a database row. + + Args: + row: The ``workflow_steps`` row. + + Returns: + The step record. + """ + return StepRecord( + run_id=row["run_id"], + ordinal=row["ordinal"], + handler_id=row["handler_id"], + status=StepStatus(row["status"]), + args=row["args"], + attempts=row["attempts"], + recoveries=row["recoveries"], + due_at=row["due_at"], + epoch=row["epoch"], + lease_expires_at=row["lease_expires_at"], + wait_key=row["wait_key"], + join_expected=row["join_expected"], + join_arrived=row["join_arrived"], + error=row["error"], + origin=row["origin"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + +class PostgresRunStore: + """Run store backed by PostgreSQL, safe for many concurrent workers.""" + + def __init__( + self, + conninfo: str, + *, + schema: str | None = None, + min_size: int = 1, + max_size: int = DEFAULT_POOL_SIZE, + ): + """Prepare a store against a Postgres database. + + The connection pool is opened lazily on first use, so constructing a + store never blocks app import or requires a running event loop. + + Args: + conninfo: A libpq connection string or URL. + schema: Postgres schema to own the tables. Defaults to whatever the + connection's search path resolves to. Naming one keeps a + deployment's runs in their own namespace inside a shared + database, which is also how tests isolate from each other. It + is quoted as an identifier, never interpolated. + min_size: Connections kept open when idle. + max_size: Maximum concurrent connections. + + """ + self._conninfo = conninfo + self._schema = schema + self._min_size = min_size + self._max_size = max_size + self._pool: Any = None + self._ready = asyncio.Lock() + + async def _open(self) -> Any: + """Open the pool and create the schema, once. + + Returns: + The open connection pool. + """ + if self._pool is not None: + return self._pool + async with self._ready: + if self._pool is not None: + return self._pool + schema = self._schema + + async def configure(conn: AsyncConnection) -> None: + """Point a pooled connection at this store's schema. + + Args: + conn: The connection being handed out. + """ + if schema is not None: + await conn.execute(_set_search_path(schema)) + + pool = psycopg_pool.AsyncConnectionPool( + self._conninfo, + min_size=self._min_size, + max_size=self._max_size, + kwargs={ + "row_factory": dict_row, + "autocommit": True, + # Naming the schema makes this store's own backends + # findable, which is how drop_schema clears its leftovers. + "application_name": schema or "reflex_workflow", + }, + configure=configure, + open=False, + ) + await pool.open(wait=True) + async with pool.connection() as conn: + if schema is not None: + await conn.execute( + SQL("CREATE SCHEMA IF NOT EXISTS {}").format(Identifier(schema)) + ) + await conn.execute(_set_search_path(schema)) + await conn.execute(_SCHEMA) + self._pool = pool + return pool + + @property + def schema(self) -> str | None: + """The schema this store's tables live in, if one was named. + + Returns: + The schema name, or None when the connection's search path decides. + """ + return self._schema + + async def close(self) -> None: + """Close the connection pool.""" + if self._pool is not None: + await self._pool.close() + self._pool = None + + def drop_schema(self) -> None: + """Delete this store's schema and everything in it. + + Tests use this to reclaim a throwaway namespace. It opens its own + short-lived connection so it works after the pool's event loop is gone + -- and because that loop can die mid-transaction, leaving a pooled + backend idle while still holding locks, it first evicts this store's + own connections. Otherwise the DROP waits on them forever. + + Raises: + WorkflowRuntimeError: If the store owns no schema of its own. + """ + if self._schema is None: + msg = "drop_schema() needs a store constructed with schema=." + raise WorkflowRuntimeError(msg) + with psycopg.connect(self._conninfo, autocommit=True) as conn: + conn.execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity" + " WHERE datname = current_database() AND application_name = %s" + " AND pid <> pg_backend_pid()", + (self._schema,), + ) + conn.execute("SET lock_timeout TO '10s'") + conn.execute( + SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(Identifier(self._schema)) + ) + + async def _lock_run(self, conn: Connection, run_id: str) -> None: + """Serialize a run's writers for the rest of the transaction. + + History sequence numbers are per run and allocated as ``MAX + 1``, so + two transactions appending to the same run must not interleave. Taking + the run's row first is also the order every other write path uses, + which is what keeps the store deadlock-free. + + Args: + conn: The connection inside an open transaction. + run_id: The run to lock. + """ + await conn.execute( + "SELECT 1 FROM workflow_runs WHERE run_id = %s FOR UPDATE", (run_id,) + ) + + async def _append_events( + self, + conn: Connection, + run_id: str, + events: Iterable[tuple[HistoryEventType, dict[str, Any]]], + now: float, + ) -> None: + """Append history events inside the current transaction. + + Args: + conn: The connection inside an open transaction. + run_id: The owning run. + events: The (type, data) pairs to append. + now: Current time in epoch seconds. + """ + rows = list(events) + if not rows: + return + cursor = await conn.execute( + "SELECT COALESCE(MAX(seq), 0) AS seq FROM workflow_history" + " WHERE run_id = %s", + (run_id,), + ) + row = await cursor.fetchone() + seq = 0 if row is None else row["seq"] + for event_type, data in rows: + seq += 1 + await conn.execute( + "INSERT INTO workflow_history (run_id, seq, type, at, data)" + " VALUES (%s, %s, %s, %s, %s)", + (run_id, seq, event_type.value, now, _json(data)), + ) + + async def _insert_run(self, conn: Connection, run: RunRecord) -> None: + """Insert a run row inside the current transaction. + + Args: + conn: The connection inside an open transaction. + run: The run record. + """ + await conn.execute( + "INSERT INTO workflow_runs (run_id, workflow_id, definition_digest," + " status, state, state_version, next_ordinal, result, error," + " flow_key, parent_run_id, parent_ordinal, request_key, labels," + " deadline, cancel_requested, created_at, updated_at)" + " VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s," + " %s, %s, %s, %s)", + ( + run.run_id, + run.workflow_id, + run.definition_digest, + run.status.value, + _json(run.state), + run.state_version, + run.next_ordinal, + _json(run.result), + _json(run.error), + run.flow_key, + run.parent_run_id, + run.parent_ordinal, + run.request_key, + _json(run.labels), + run.deadline, + run.cancel_requested, + run.created_at, + run.updated_at, + ), + ) + + async def _insert_step(self, conn: Connection, step: StepRecord) -> None: + """Insert a step row inside the current transaction. + + Args: + conn: The connection inside an open transaction. + step: The step record. + """ + await conn.execute( + "INSERT INTO workflow_steps (run_id, ordinal, handler_id, status, args," + " attempts, recoveries, due_at, epoch, lease_expires_at, wait_key," + " join_expected, join_arrived, error, origin, created_at, updated_at)" + " VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s," + " %s, %s)", + ( + step.run_id, + step.ordinal, + step.handler_id, + step.status.value, + _json(step.args), + step.attempts, + step.recoveries, + step.due_at, + step.epoch, + step.lease_expires_at, + step.wait_key, + step.join_expected, + step.join_arrived, + _json(step.error), + step.origin, + step.created_at, + step.updated_at, + ), + ) + + async def _load_steps(self, conn: Connection, run_id: str) -> list[StepRecord]: + """Load a run's steps in ordinal order inside the current transaction. + + Args: + conn: The connection inside an open transaction. + run_id: The owning run. + + Returns: + The step records. + """ + cursor = await conn.execute( + "SELECT * FROM workflow_steps WHERE run_id = %s ORDER BY ordinal", + (run_id,), + ) + return [_step_from_row(row) for row in await cursor.fetchall()] + + async def _frontier(self, conn: Connection, run_id: str) -> StepRecord | None: + """Load a run's lowest unresolved slot. + + Args: + conn: The connection inside an open transaction. + run_id: The owning run. + + Returns: + The frontier step, or None when every slot is resolved. + """ + cursor = await conn.execute( + "SELECT * FROM workflow_steps WHERE run_id = %s" + " AND NOT (status = ANY(%s)) ORDER BY ordinal LIMIT 1", + (run_id, _TERMINAL_STEPS), + ) + row = await cursor.fetchone() + return None if row is None else _step_from_row(row) + + async def _check_claim(self, conn: Connection, claim: Claim) -> None: + """Validate that a claim still owns its step and state version. + + Args: + conn: The connection inside an open transaction. + claim: The claim to validate. + + Raises: + StaleClaimError: If the claim was fenced. + """ + cursor = await conn.execute( + "SELECT s.status AS step_status, s.epoch AS epoch," + " r.state_version AS state_version" + " FROM workflow_steps s JOIN workflow_runs r ON r.run_id = s.run_id" + " WHERE s.run_id = %s AND s.ordinal = %s", + (claim.run.run_id, claim.step.ordinal), + ) + row = await cursor.fetchone() + if ( + row is None + or row["step_status"] != StepStatus.CLAIMED.value + or row["epoch"] != claim.step.epoch + or row["state_version"] != claim.run.state_version + ): + msg = ( + f"Claim on run {claim.run.run_id} step {claim.step.ordinal} was fenced." + ) + raise StaleClaimError(msg) + + async def admit( + self, + run: RunRecord, + root_step: StepRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + ) -> tuple[bool, str]: + """Atomically admit a run, deduplicating on the request key. + + Args: + run: The run record to create. + root_step: The preallocated root mailbox slot. + events: History events to append on creation. + + Returns: + Whether the run was created, and the authoritative run id. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + if run.request_key is not None: + cursor = await conn.execute( + "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" + " VALUES (%s, %s, %s) ON CONFLICT DO NOTHING RETURNING run_id", + (run.workflow_id, run.request_key, run.run_id), + ) + if await cursor.fetchone() is None: + cursor = await conn.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = %s AND request_key = %s", + (run.workflow_id, run.request_key), + ) + existing = await cursor.fetchone() + if existing is not None: + return False, existing["run_id"] + await self._insert_run(conn, run) + await self._insert_step(conn, root_step) + await self._append_events(conn, run.run_id, events, run.created_at) + return True, run.run_id + + async def claim_next( + self, now: float, *, lease_duration: float = DEFAULT_LEASE_DURATION + ) -> Claim | None: + """Claim the due frontier step of some runnable run. + + The candidate row is locked with ``SKIP LOCKED``, so concurrent workers + never contend for the same step and never queue behind each other: a + worker takes the oldest run whose frontier nobody else is holding. + + Args: + now: Current time in epoch seconds. + lease_duration: Seconds of renewal silence tolerated before the + claim is treated as orphaned. + + Returns: + A fenced claim, or None when nothing is claimable right now. + """ + pool = await self._open() + params = { + "now": now, + "terminal_runs": _TERMINAL_RUNS, + "terminal_steps": _TERMINAL_STEPS, + "claimable": _CLAIMABLE_STEPS, + } + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "SELECT s.run_id AS run_id, s.ordinal AS ordinal" + " FROM workflow_steps s JOIN workflow_runs r ON r.run_id = s.run_id" + f" WHERE {_RUNNABLE_PREDICATE} AND {_FRONTIER_PREDICATE}" + f" AND {_CLAIMABLE_PREDICATE}" + " ORDER BY r.created_at, s.run_id" + " FOR UPDATE OF s SKIP LOCKED LIMIT 1", + params, + ) + candidate = await cursor.fetchone() + if candidate is None: + return None + cursor = await conn.execute( + "UPDATE workflow_steps SET status = %s, epoch = epoch + 1," + " lease_expires_at = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s RETURNING *", + ( + StepStatus.CLAIMED.value, + now + lease_duration, + now, + candidate["run_id"], + candidate["ordinal"], + ), + ) + step_row = await cursor.fetchone() + cursor = await conn.execute( + "UPDATE workflow_runs SET status = %s, updated_at = %s" + " WHERE run_id = %s RETURNING *", + (RunStatus.RUNNING.value, now, candidate["run_id"]), + ) + run_row = await cursor.fetchone() + if step_row is None or run_row is None: + return None + return Claim(run=_run_from_row(run_row), step=_step_from_row(step_row)) + + async def renew_lease( + self, + claim: Claim, + now: float, + *, + lease_duration: float = DEFAULT_LEASE_DURATION, + ) -> bool: + """Extend a live claim's lease without transitioning the step. + + Args: + claim: The claim being renewed. + now: Current time in epoch seconds. + lease_duration: Seconds to extend the lease from ``now``. + + Returns: + True if the claim still owns its step; False if it was fenced. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + try: + await self._check_claim(conn, claim) + except StaleClaimError: + return False + await conn.execute( + "UPDATE workflow_steps SET lease_expires_at = %s" + " WHERE run_id = %s AND ordinal = %s", + (now + lease_duration, claim.run.run_id, claim.step.ordinal), + ) + return True + + async def _arm(self, conn: Connection, step: StepRecord, now: float) -> StepRecord: + """Resolve a newly armed wait against a buffered delivery. + + Args: + conn: The connection inside an open transaction. + step: The slot being inserted. + now: Current time in epoch seconds. + + Returns: + The slot, already resolved when a matching delivery was waiting. + """ + if step.status is not StepStatus.BLOCKED or step.wait_key is None: + return step + cursor = await conn.execute( + "SELECT dedupe_key, payload FROM workflow_inbox" + " WHERE run_id = %s AND wait_key = %s AND status = 'PENDING'" + " ORDER BY seq LIMIT 1", + (step.run_id, step.wait_key), + ) + row = await cursor.fetchone() + if row is None: + return step + await conn.execute( + "UPDATE workflow_inbox SET status = 'CONSUMED'" + " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", + (step.run_id, step.wait_key, row["dedupe_key"]), + ) + return dataclasses.replace( + step, + status=StepStatus.READY, + due_at=now, + args={**step.args, "__payload__": row["payload"]}, + updated_at=now, + ) + + async def commit( + self, claim: Claim, completion: StepCompletion, now: float + ) -> None: + """Atomically apply the outcome of a claimed attempt. + + Args: + claim: The claim being committed. + completion: The outcome to apply. + now: Current time in epoch seconds. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + await self._lock_run(conn, claim.run.run_id) + await self._check_claim(conn, claim) + await conn.execute( + "UPDATE workflow_steps SET status = %s, attempts = attempts + %s," + " due_at = %s, lease_expires_at = 0, error = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s", + ( + completion.step_status.value, + 1 if completion.consume_attempt else 0, + completion.due_at if completion.due_at is not None else 0.0, + _json(completion.step_error), + now, + claim.run.run_id, + claim.step.ordinal, + ), + ) + if completion.tombstones: + await conn.execute( + "UPDATE workflow_steps SET status = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = ANY(%s)" + " AND NOT (status = ANY(%s))", + ( + StepStatus.CANCELLED.value, + now, + claim.run.run_id, + list(completion.tombstones), + _TERMINAL_STEPS, + ), + ) + for step in completion.new_steps: + await self._insert_step(conn, await self._arm(conn, step, now)) + for child_run, child_step in completion.children: + await self._insert_run(conn, child_run) + await self._insert_step(conn, child_step) + await conn.execute( + "UPDATE workflow_runs SET status = %s," + " state = CASE WHEN %s THEN %s ELSE state END," + " state_version = state_version + %s," + " next_ordinal = COALESCE(%s, next_ordinal)," + " result = COALESCE(%s, result), error = %s, updated_at = %s" + " WHERE run_id = %s", + ( + completion.run_status.value, + completion.state is not None, + _json(completion.state), + 1 if completion.state is not None else 0, + completion.next_ordinal, + _json(completion.result), + _json(completion.run_error), + now, + claim.run.run_id, + ), + ) + await self._append_events(conn, claim.run.run_id, completion.events, now) + + async def release_claim( + self, + claim: Claim, + *, + status: StepStatus, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Return a claimed step without committing any state. + + Args: + claim: The claim being released. + status: The step status to record. + events: History events to append. + now: Current time in epoch seconds. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + await self._lock_run(conn, claim.run.run_id) + try: + await self._check_claim(conn, claim) + except StaleClaimError: + return + await conn.execute( + "UPDATE workflow_steps SET status = %s, lease_expires_at = 0," + " updated_at = %s WHERE run_id = %s AND ordinal = %s", + (status.value, now, claim.run.run_id, claim.step.ordinal), + ) + await self._append_events(conn, claim.run.run_id, events, now) + + async def append_events( + self, + run_id: str, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Append evidence events outside a fenced commit. + + Args: + run_id: The owning run. + events: The (type, data) pairs to append. + now: Current time in epoch seconds. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + await self._lock_run(conn, run_id) + await self._append_events(conn, run_id, events, now) + + async def _next_inbox_seq(self, conn: Connection, run_id: str) -> int: + """Allocate the next inbox sequence number for a run. + + Args: + conn: The connection inside an open transaction. + run_id: The owning run. + + Returns: + The next sequence number. + """ + cursor = await conn.execute( + "SELECT COALESCE(MAX(seq), 0) + 1 AS seq FROM workflow_inbox" + " WHERE run_id = %s", + (run_id,), + ) + row = await cursor.fetchone() + return 1 if row is None else row["seq"] + + async def deliver( + self, + run_id: str, + wait_key: str, + dedupe_key: str, + payload: dict[str, Any], + now: float, + ) -> DeliveryDisposition: + """Deliver a payload to a run, resolving its wait or buffering it. + + Args: + run_id: The receiving run. + wait_key: The address the waiting slot declared. + dedupe_key: Sender-supplied identity, making redelivery a no-op. + payload: JSON-compatible payload to hand the resuming handler. + now: Current time in epoch seconds. + + Returns: + What the store did with the delivery. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", + (run_id,), + ) + row = await cursor.fetchone() + if row is None: + return "unknown_run" + if row["status"] in _TERMINAL_RUNS: + return "run_terminal" + cursor = await conn.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", + (run_id, wait_key, dedupe_key), + ) + if await cursor.fetchone() is not None: + return "duplicate" + frontier = await self._frontier(conn, run_id) + if ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and 0.0 < frontier.due_at <= now + ): + return "expired" + resolves = ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and frontier.wait_key == wait_key + ) + await conn.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at) VALUES (%s, %s, %s, %s, %s, %s, %s)", + ( + run_id, + wait_key, + dedupe_key, + await self._next_inbox_seq(conn, run_id), + _json(payload), + "CONSUMED" if resolves else "PENDING", + now, + ), + ) + if resolves and frontier is not None: + await conn.execute( + "UPDATE workflow_steps SET status = %s, due_at = %s, args = %s," + " updated_at = %s WHERE run_id = %s AND ordinal = %s", + ( + StepStatus.READY.value, + now, + _json({**frontier.args, "__payload__": payload}), + now, + run_id, + frontier.ordinal, + ), + ) + await self._append_events( + conn, + run_id, + ( + ( + HistoryEventType.WAIT_RESOLVED, + {"ordinal": frontier.ordinal, "wait_key": wait_key}, + ), + ), + now, + ) + else: + await self._append_events( + conn, + run_id, + ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), + now, + ) + return "resolved" if resolves else "buffered" + + async def admit_children( + self, + runs: tuple[tuple[RunRecord, StepRecord], ...], + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + now: float, + ) -> None: + """Create child runs, each with its root slot. + + Args: + runs: The child run records paired with their root slots. + events: History events to append to the parent. + now: Current time in epoch seconds. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + for run, root_step in runs: + await self._insert_run(conn, run) + await self._insert_step(conn, root_step) + if runs and events: + parent = runs[0][0].parent_run_id or "" + await self._lock_run(conn, parent) + await self._append_events(conn, parent, events, now) + + async def record_arrival( + self, + run_id: str, + ordinal: int, + payload: dict[str, Any], + dedupe_key: str, + now: float, + ) -> DeliveryDisposition: + """Count one arrival against a join slot. + + Args: + run_id: The waiting parent run. + ordinal: The join slot's ordinal. + payload: The arriving result. + dedupe_key: Identity of the arrival. + now: Current time in epoch seconds. + + Returns: + What the store did with the arrival. + """ + pool = await self._open() + wait_key = f"join:{ordinal}" + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", + (run_id,), + ) + run_row = await cursor.fetchone() + if run_row is None: + return "unknown_run" + if run_row["status"] in _TERMINAL_RUNS: + return "run_terminal" + cursor = await conn.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", + (run_id, wait_key, dedupe_key), + ) + if await cursor.fetchone() is not None: + return "duplicate" + cursor = await conn.execute( + "SELECT * FROM workflow_steps WHERE run_id = %s AND ordinal = %s", + (run_id, ordinal), + ) + step_row = await cursor.fetchone() + if step_row is None or step_row["status"] != StepStatus.BLOCKED.value: + return "run_terminal" + step = _step_from_row(step_row) + await conn.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at)" + " VALUES (%s, %s, %s, %s, %s, 'CONSUMED', %s)", + ( + run_id, + wait_key, + dedupe_key, + await self._next_inbox_seq(conn, run_id), + _json(payload), + now, + ), + ) + arrived = step.join_arrived + 1 + results = [*step.args.get("__results__", []), payload] + done = arrived >= step.join_expected + await conn.execute( + "UPDATE workflow_steps SET status = %s, join_arrived = %s," + " due_at = %s, args = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s AND join_arrived = %s", + ( + StepStatus.READY.value if done else StepStatus.BLOCKED.value, + arrived, + now if done else step.due_at, + _json({**step.args, "__results__": results}), + now, + run_id, + ordinal, + step.join_arrived, + ), + ) + await self._append_events( + conn, + run_id, + ( + ( + HistoryEventType.CHILD_RESOLVED, + {"ordinal": ordinal, "arrived": arrived}, + ), + ), + now, + ) + return "resolved" if done else "counted" + + async def count_active(self, workflow_id: str, flow_key: str) -> int: + """Count runs of a root still in flight under a flow-control key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + + Returns: + How many runs are active. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT COUNT(*) AS n FROM workflow_runs" + " WHERE workflow_id = %s AND flow_key = %s" + " AND NOT (status = ANY(%s))", + (workflow_id, flow_key, _TERMINAL_RUNS), + ) + row = await cursor.fetchone() + return 0 if row is None else row["n"] + + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: + """Find the oldest active run under a flow-control key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + + Returns: + The run record, or None when nothing is active. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT * FROM workflow_runs" + " WHERE workflow_id = %s AND flow_key = %s" + " AND NOT (status = ANY(%s))" + " ORDER BY created_at, run_id LIMIT 1", + (workflow_id, flow_key, _TERMINAL_RUNS), + ) + row = await cursor.fetchone() + return None if row is None else _run_from_row(row) + + async def count_started_since( + self, workflow_id: str, flow_key: str, since: float + ) -> int: + """Count runs of a root admitted under a key since a point in time. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + since: Exclusive lower bound in epoch seconds. + + Returns: + How many runs were admitted in the window. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT COUNT(*) AS n FROM workflow_runs" + " WHERE workflow_id = %s AND flow_key = %s AND created_at > %s", + (workflow_id, flow_key, since), + ) + row = await cursor.fetchone() + return 0 if row is None else row["n"] + + async def nth_recent_start( + self, workflow_id: str, flow_key: str, n: int + ) -> float | None: + """Find the nth most recent scheduled start under a flow key. + + Args: + workflow_id: The workflow identity. + flow_key: The computed grouping key. + n: How far back to look, counting from the most recent as 1. + + Returns: + The scheduled start, or None when fewer than n runs exist. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT GREATEST(s.due_at, r.created_at) AS start FROM workflow_runs r" + " JOIN workflow_steps s ON s.run_id = r.run_id AND s.ordinal = 0" + " WHERE r.workflow_id = %s AND r.flow_key = %s" + " ORDER BY start DESC LIMIT 1 OFFSET %s", + (workflow_id, flow_key, n - 1), + ) + row = await cursor.fetchone() + return None if row is None else row["start"] + + async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: + """Push a not-yet-started run's root slot later, for debouncing. + + Args: + run_id: The pending run. + due_at: The new earliest start time. + now: Current time in epoch seconds. + + Returns: + True when the root had not started and was deferred. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "UPDATE workflow_steps SET due_at = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = 0 AND status = %s", + (due_at, now, run_id, StepStatus.READY.value), + ) + return cursor.rowcount > 0 + + async def request_cancel(self, run_id: str, now: float) -> bool: + """Record cancellation intent on a run. + + Args: + run_id: The run to cancel. + now: Current time in epoch seconds. + + Returns: + True if intent was recorded on a nonterminal run. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "UPDATE workflow_runs SET cancel_requested = TRUE, status = %s," + " updated_at = %s WHERE run_id = %s AND NOT (status = ANY(%s))", + (RunStatus.CANCELLING.value, now, run_id, _TERMINAL_RUNS), + ) + if cursor.rowcount == 0: + return False + await self._append_events( + conn, run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now + ) + return True + + async def control_pending(self, now: float) -> tuple[RunRecord, ...]: + """List drained runs awaiting a control transition. + + Args: + now: Current time in epoch seconds. + + Returns: + The runs awaiting finalization. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT * FROM workflow_runs r WHERE NOT (r.status = ANY(%s))" + " AND (r.cancel_requested OR (r.deadline IS NOT NULL" + " AND r.deadline <= %s))" + " AND NOT EXISTS (SELECT 1 FROM workflow_steps s" + " WHERE s.run_id = r.run_id AND s.status = %s)", + (_TERMINAL_RUNS, now, StepStatus.CLAIMED.value), + ) + return tuple(_run_from_row(row) for row in await cursor.fetchall()) + + async def finalize_run( + self, + run_id: str, + *, + status: RunStatus, + error: dict[str, Any] | None, + event: HistoryEventType, + now: float, + ) -> bool: + """Terminate a drained run and tombstone its unresolved slots. + + Args: + run_id: The run to finalize. + status: The terminal status to record. + error: Error payload recorded on the run. + event: The terminal history event type. + now: Current time in epoch seconds. + + Returns: + True if the run was finalized. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", + (run_id,), + ) + row = await cursor.fetchone() + if row is None or row["status"] in _TERMINAL_RUNS: + return False + cursor = await conn.execute( + "SELECT 1 FROM workflow_steps WHERE run_id = %s AND status = %s", + (run_id, StepStatus.CLAIMED.value), + ) + if await cursor.fetchone() is not None: + return False + cursor = await conn.execute( + "SELECT ordinal FROM workflow_steps WHERE run_id = %s" + " AND NOT (status = ANY(%s)) ORDER BY ordinal", + (run_id, _TERMINAL_STEPS), + ) + open_rows = await cursor.fetchall() + await conn.execute( + "UPDATE workflow_steps SET status = %s, updated_at = %s" + " WHERE run_id = %s AND NOT (status = ANY(%s))", + (StepStatus.CANCELLED.value, now, run_id, _TERMINAL_STEPS), + ) + await conn.execute( + "UPDATE workflow_runs SET status = %s, error = %s, updated_at = %s" + " WHERE run_id = %s", + (status.value, _json(error), now, run_id), + ) + events: list[tuple[HistoryEventType, dict[str, Any]]] = [ + (HistoryEventType.STEP_TOMBSTONED, {"ordinal": open_row["ordinal"]}) + for open_row in open_rows + ] + events.append((event, {} if error is None else dict(error))) + await self._append_events(conn, run_id, events, now) + return True + + async def resume_run(self, run_id: str, now: float) -> bool: + """Re-open a suspended run so its frontier step runs again. + + Args: + run_id: The run to resume. + now: Current time in epoch seconds. + + Returns: + True if a suspended run was re-opened. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "UPDATE workflow_runs SET status = %s, error = NULL, updated_at = %s" + " WHERE run_id = %s AND status = %s", + (RunStatus.PENDING.value, now, run_id, RunStatus.NEEDS_ATTENTION.value), + ) + if cursor.rowcount == 0: + return False + await conn.execute( + "UPDATE workflow_steps SET status = %s, attempts = 0, due_at = %s," + " lease_expires_at = 0, error = NULL, updated_at = %s" + " WHERE run_id = %s AND status = %s", + ( + StepStatus.READY.value, + now, + now, + run_id, + StepStatus.NEEDS_ATTENTION.value, + ), + ) + await self._append_events( + conn, run_id, ((HistoryEventType.RUN_RESUMED, {}),), now + ) + return True + + async def recover_orphans( + self, now: float, max_recoveries: int + ) -> tuple[int, tuple[str, ...]]: + """Recover claims whose lease has expired. + + Args: + now: Current time in epoch seconds. + max_recoveries: Recovery budget per logical step. + + Returns: + How many steps were transitioned, and the runs failed outright. + """ + pool = await self._open() + exhausted = {"reason": "recovery_budget_exhausted"} + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "SELECT s.* FROM workflow_steps s" + " JOIN workflow_runs r ON r.run_id = s.run_id" + " WHERE s.status = %s AND s.lease_expires_at <= %s" + " AND NOT (r.status = ANY(%s))" + " ORDER BY s.run_id FOR UPDATE OF s SKIP LOCKED", + (StepStatus.CLAIMED.value, now, _TERMINAL_RUNS), + ) + rows = await cursor.fetchall() + recovered = 0 + failed: list[str] = [] + for row in rows: + step = _step_from_row(row) + recovered += 1 + if step.recoveries + 1 > max_recoveries: + await conn.execute( + "UPDATE workflow_steps SET status = %s, recoveries = %s," + " lease_expires_at = 0, error = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s", + ( + StepStatus.FAILED.value, + step.recoveries + 1, + _json(exhausted), + now, + step.run_id, + step.ordinal, + ), + ) + await conn.execute( + "UPDATE workflow_runs SET status = %s, error = %s," + " updated_at = %s WHERE run_id = %s", + ( + RunStatus.FAILED.value, + _json(exhausted), + now, + step.run_id, + ), + ) + failed.append(step.run_id) + await self._append_events( + conn, + step.run_id, + ((HistoryEventType.RUN_FAILED, dict(exhausted)),), + now, + ) + else: + await conn.execute( + "UPDATE workflow_steps SET status = %s, recoveries = %s," + " due_at = %s, lease_expires_at = 0, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s", + ( + StepStatus.RECOVERY_WAIT.value, + step.recoveries + 1, + now, + now, + step.run_id, + step.ordinal, + ), + ) + await self._append_events( + conn, + step.run_id, + ((HistoryEventType.STEP_RECOVERED, {"ordinal": step.ordinal}),), + now, + ) + return recovered, tuple(failed) + + async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: + """List runs matching a query, newest first. + + Args: + query: The filters and pagination cursor to apply. + + Returns: + The matching run records. + """ + clauses: list[str] = [] + params: list[Any] = [] + if query.workflow_id is not None: + clauses.append("workflow_id = %s") + params.append(query.workflow_id) + if query.statuses: + clauses.append("status = ANY(%s)") + params.append([status.value for status in query.statuses]) + if query.created_before is not None: + clauses.append("(created_at, run_id) < (%s, %s)") + params.extend(query.created_before) + if query.labels: + # Containment matches the whole filter at once, and takes user keys + # as data rather than splicing them into a path expression. + clauses.append("labels @> %s") + params.append(_json(dict(query.labels))) + where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + f"SELECT * FROM workflow_runs{where}" + " ORDER BY created_at DESC, run_id DESC LIMIT %s", + (*params, query.limit), + ) + return tuple(_run_from_row(row) for row in await cursor.fetchall()) + + async def list_children( + self, parent_run_id: str, parent_ordinal: int + ) -> tuple[RunRecord, ...]: + """List the child runs admitted for one join slot. + + Args: + parent_run_id: The run that fanned out. + parent_ordinal: The join slot the children report to. + + Returns: + The child run records, oldest first. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT * FROM workflow_runs WHERE parent_run_id = %s" + " AND parent_ordinal = %s ORDER BY created_at, run_id", + (parent_run_id, parent_ordinal), + ) + return tuple(_run_from_row(row) for row in await cursor.fetchall()) + + async def find_by_request_key( + self, workflow_id: str, request_key: str + ) -> str | None: + """Find the run a request key already admitted, if any. + + Args: + workflow_id: The workflow identity. + request_key: The idempotent admission key. + + Returns: + The existing run id, or None when the key is unused. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = %s AND request_key = %s", + (workflow_id, request_key), + ) + row = await cursor.fetchone() + return None if row is None else row["run_id"] + + async def get_run(self, run_id: str) -> RunRecord | None: + """Load one run record. + + Args: + run_id: The run identity. + + Returns: + The record, or None if unknown. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT * FROM workflow_runs WHERE run_id = %s", (run_id,) + ) + row = await cursor.fetchone() + return None if row is None else _run_from_row(row) + + async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: + """Load a run's mailbox slots in ordinal order. + + Args: + run_id: The run identity. + + Returns: + The step records. + """ + pool = await self._open() + async with pool.connection() as conn: + return tuple(await self._load_steps(conn, run_id)) + + async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: + """Load a run's append-only history in sequence order. + + Args: + run_id: The run identity. + + Returns: + The history events. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT * FROM workflow_history WHERE run_id = %s ORDER BY seq", + (run_id,), + ) + return tuple( + HistoryEvent( + run_id=row["run_id"], + seq=row["seq"], + type=HistoryEventType(row["type"]), + at=row["at"], + data=row["data"], + ) + for row in await cursor.fetchall() + ) + + async def next_due(self, now: float) -> float | None: + """Earliest future time any runnable run becomes claimable. + + Args: + now: Current time in epoch seconds. + + Returns: + The epoch time, or None when no future work is scheduled. + """ + pool = await self._open() + params = { + "now": now, + "terminal_runs": _TERMINAL_RUNS, + "terminal_steps": _TERMINAL_STEPS, + "claimable": _CLAIMABLE_STEPS, + } + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT MIN(s.due_at) AS due FROM workflow_steps s" + " JOIN workflow_runs r ON r.run_id = s.run_id" + f" WHERE {_RUNNABLE_PREDICATE} AND {_FRONTIER_PREDICATE}" + " AND (s.status = ANY(%(claimable)s)" + " OR (s.status = 'BLOCKED' AND s.due_at > 0))", + params, + ) + row = await cursor.fetchone() + return None if row is None else row["due"] diff --git a/tests/units/workflow/conftest.py b/tests/units/workflow/conftest.py index f8b3e80805c..a1f5eeca874 100644 --- a/tests/units/workflow/conftest.py +++ b/tests/units/workflow/conftest.py @@ -1,19 +1,38 @@ -"""Run every harness-based workflow test against both store implementations. +"""Run every harness-based workflow test against each store implementation. The test harness defaults to the in-memory store, so behavioral tests were -certifying semantics that production -- which persists to SQLite -- does not -necessarily have. Any divergence between the two stores could therefore ship -green. This fixture makes the default store a parameter, so every test in this -directory runs twice. +certifying semantics that production -- which persists to SQLite or Postgres -- +does not necessarily have. Any divergence between the stores could therefore +ship green. This fixture makes the default store a parameter, so every test in +this directory runs against each one. + +Postgres needs a server, so that parameter is skipped unless +``REFLEX_TEST_POSTGRES`` names one. Set it to a libpq URL to include it:: + + REFLEX_TEST_POSTGRES=postgresql://user:pw@localhost:5432/db uv run pytest """ +import os +import uuid + import pytest import reflex.workflow.testing as testing from reflex.workflow.store import SqliteRunStore +POSTGRES_URL_VAR = "REFLEX_TEST_POSTGRES" -@pytest.fixture(params=["memory", "sqlite"], autouse=True) + +def _postgres_url() -> str | None: + """Read the Postgres server to test against, if one is configured. + + Returns: + The connection URL, or None to skip the Postgres parameter. + """ + return os.environ.get(POSTGRES_URL_VAR) or None + + +@pytest.fixture(params=["memory", "sqlite", "postgres"], autouse=True) def harness_store(request, tmp_path, monkeypatch): """Make the harness's default store a parameter of every test. @@ -26,17 +45,35 @@ def harness_store(request, tmp_path, monkeypatch): The store kind under test. """ opened: list[SqliteRunStore] = [] + postgres: list = [] if request.param == "sqlite": - def factory(): + def sqlite_factory(): store = SqliteRunStore(tmp_path / f"harness{len(opened)}.db") opened.append(store) return store - monkeypatch.setattr(testing, "MemoryRunStore", factory) + monkeypatch.setattr(testing, "MemoryRunStore", sqlite_factory) + elif request.param == "postgres": + url = _postgres_url() + if url is None: + pytest.skip(f"set {POSTGRES_URL_VAR} to test against Postgres") + from reflex.workflow.postgres import PostgresRunStore + + # Each store gets its own schema, so tests that build several stores -- + # or run in parallel -- never see each other's runs. + def postgres_factory(): + schema = f"wf_test_{uuid.uuid4().hex}" + store = PostgresRunStore(url, schema=schema, min_size=0, max_size=4) + postgres.append(store) + return store + + monkeypatch.setattr(testing, "MemoryRunStore", postgres_factory) yield request.param for store in opened: store.close() + for store in postgres: + store.drop_schema() diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 0ed0cbcd0d6..39f3c02e908 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -429,17 +429,21 @@ async def test_race_continues_on_the_first_branch_and_cancels_the_rest(): assert "slow-answer" not in RACE_CALLS -async def test_race_mode_still_starts_every_branch(): - """Racing is not a way to skip work: every branch starts.""" +async def test_race_mode_admits_every_branch(): + """Racing picks a winner; it does not skip a branch. + + Every branch is admitted as a child run. Whether a loser gets to run at all + depends on how fast the winner is -- a loser cancelled before its first + step is the best case, not a missed one -- so this asserts what is actually + guaranteed rather than what one store's ordering happens to produce. + """ RACE_CALLS.clear() async with WorkflowTestHarness(Shopper, SlowVendor, FastVendor) as harness: result = await harness.start(Shopper.start()) assert result.run_id is not None - assert "slow" in RACE_CALLS - assert "fast" in RACE_CALLS join = await _join_slot(harness, result.run_id) children = await harness.kernel.store.list_children(result.run_id, join.ordinal) - assert len(children) == 2 + assert {child.workflow_id for child in children} == {"race.fast", "race.slow"} async def test_race_join_expects_one_arrival(): diff --git a/tests/units/workflow/test_postgres.py b/tests/units/workflow/test_postgres.py new file mode 100644 index 00000000000..8ee0b4283e4 --- /dev/null +++ b/tests/units/workflow/test_postgres.py @@ -0,0 +1,283 @@ +"""Tests that need a real PostgreSQL server. + +The point of the Postgres store is the thing SQLite cannot do: several worker +processes claiming from one mailbox at the same time. These tests run two +kernels against one database and assert what a durable engine has to promise +under that load -- every run executes, and no run executes twice. + +They are skipped unless ``REFLEX_TEST_POSTGRES`` names a server. +""" + +import asyncio +import json +import os +import time +import uuid + +import pytest +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.definition import compile_workflow +from reflex.workflow.kernel import WorkflowKernel, WorkflowObserver +from reflex.workflow.records import ( + HistoryEventType, + RunRecord, + RunStatus, + StepRecord, + StepStatus, +) + +POSTGRES_URL = os.environ.get("REFLEX_TEST_POSTGRES") or "" + +pytestmark = pytest.mark.skipif( + not POSTGRES_URL, reason="set REFLEX_TEST_POSTGRES to run Postgres tests" +) + +EXECUTIONS: list[str] = [] + + +@pytest.fixture(autouse=True) +def harness_store(): + """Opt out of the shared store parameter. + + These tests drive kernels against their own Postgres store, so running + them once per store kind would just repeat the same work. + + Returns: + The store kind this module uses. + """ + return "postgres" + + +class Charge(rx.State): + """A workflow whose one step must never run twice.""" + + __workflow__ = WorkflowConfig(id="pg.charge") + invoice: str = "" + + @rx.event(durable=True, trigger=manual(), effect="non_idempotent_write") + async def start(self, invoice: str): + """Charge the invoice exactly once. + + Args: + invoice: The invoice identifier. + + Returns: + Completion. + """ + EXECUTIONS.append(invoice) + await asyncio.sleep(0.01) + self.invoice = invoice + return rx.complete(result={"charged": invoice}) + + +@pytest.fixture +def store(): + """Open a Postgres store in a throwaway schema. + + Yields: + The store. + """ + from reflex.workflow.postgres import PostgresRunStore + + schema = f"wf_test_{uuid.uuid4().hex}" + opened = PostgresRunStore(POSTGRES_URL, schema=schema, min_size=0, max_size=6) + yield opened + opened.drop_schema() + + +async def test_two_workers_share_one_mailbox_without_double_execution(store): + """Two kernels on one database run every step exactly once. + + This is the claim SQLite cannot support and the reason the Postgres store + exists. A step marked non_idempotent_write running twice is a double + charge, so the assertion is exact, not statistical. + """ + EXECUTIONS.clear() + definition = compile_workflow(Charge) + claimed: list[set[str]] = [set(), set()] + + def watcher(index: int): + """Record which runs one worker started attempts on. + + Args: + index: Which worker this observer belongs to. + + Returns: + An observer for that worker. + """ + + class Watcher(WorkflowObserver): + def on_event(self, event_type, run_id, workflow_id, data): + """Note an attempt this worker started. + + Args: + event_type: The recorded transition. + run_id: The run it happened on. + workflow_id: The workflow identity. + data: The event payload. + """ + if event_type is HistoryEventType.ATTEMPT_STARTED: + claimed[index].add(run_id) + + return Watcher() + + workers = [ + WorkflowKernel( + [definition], + store, + poll_interval=0.01, + max_concurrency=4, + observer=watcher(index), + ) + for index in range(2) + ] + + invoices = [f"inv-{index}" for index in range(20)] + for invoice in invoices: + result = await workers[0].start(Charge.start(invoice)) + assert result.disposition == "started" + + for worker in workers: + await worker.start_worker() + try: + for _ in range(500): + runs = await workers[0].list_runs() + if all(run.status is RunStatus.COMPLETED for run in runs) and len( + runs + ) == len(invoices): + break + await asyncio.sleep(0.02) + finally: + for worker in workers: + await worker.aclose() + + assert sorted(EXECUTIONS) == sorted(invoices) + runs = await workers[0].list_runs() + assert len(runs) == len(invoices) + assert {run.status for run in runs} == {RunStatus.COMPLETED} + # Both workers pulled from the queue, and never the same run twice, so the + # exactly-once result above is a real division of labour and not one + # worker quietly doing everything. + assert claimed[0] + assert claimed[1] + assert not claimed[0] & claimed[1] + assert claimed[0] | claimed[1] == {run.run_id for run in runs} + + +async def test_a_second_worker_takes_over_an_abandoned_claim(store): + """A worker that dies mid-step does not strand its run. + + The lease is what makes this safe: the survivor may only reclaim a step + whose lease has lapsed, so recovery cannot race a worker that is merely + slow. + """ + EXECUTIONS.clear() + definition = compile_workflow(Charge) + dying = WorkflowKernel([definition], store, lease_duration=0.5) + survivor = WorkflowKernel([definition], store, poll_interval=0.01, lease_duration=5) + + result = await dying.start(Charge.start("inv-orphan")) + assert result.run_id is not None + + claim = await store.claim_next(time.time(), lease_duration=0.4) + assert claim is not None + assert claim.run.run_id == result.run_id + + # Nobody renews that lease, so it lapses as if the worker had died. + await asyncio.sleep(0.6) + assert await survivor.recover() == 1 + + await survivor.start_worker() + try: + for _ in range(300): + snapshot = await survivor.get_run(result.run_id) + if snapshot is not None and snapshot.status is RunStatus.COMPLETED: + break + await asyncio.sleep(0.02) + finally: + await survivor.aclose() + + snapshot = await survivor.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert EXECUTIONS == ["inv-orphan"] + + +async def test_concurrent_admissions_deduplicate_on_the_request_key(store): + """A request key admitted from many workers at once yields one run. + + Postgres decides this with a unique index rather than a write lock, so it + is the one place the two stores reach the same answer by different means. + """ + definition = compile_workflow(Charge) + workers = [ + WorkflowKernel([definition], store, poll_interval=0.01) for _ in range(4) + ] + results = await asyncio.gather( + *( + worker.start(Charge.start("inv-dupe"), request_key="webhook-1") + for worker in workers + ) + ) + run_ids = {result.run_id for result in results} + assert len(run_ids) == 1 + assert sum(result.disposition == "started" for result in results) == 1 + assert sum(result.disposition == "deduplicated" for result in results) == 3 + + +def test_the_cli_operates_on_a_postgres_url(store): + """`reflex workflows` accepts a Postgres URL, not just a SQLite path. + + The commands each ran their own ``asyncio.run`` before this, which works + for a file-backed store and fails for a pooled one: the pool's connections + belong to the loop that opened them, so closing from a second loop raised. + Only driving the real command surfaced it. + """ + from click.testing import CliRunner + + from reflex.workflow.cli import workflows + + async def seed(): + """Admit one run directly, so the CLI has something to find.""" + now = time.time() + await store.admit( + RunRecord( + run_id="cli-run", + workflow_id="pg.charge", + definition_digest="digest", + status=RunStatus.RUNNING, + state={}, + state_version=0, + next_ordinal=1, + labels={"team": "ops"}, + created_at=now, + updated_at=now, + ), + StepRecord( + run_id="cli-run", + ordinal=0, + handler_id="start", + status=StepStatus.READY, + args={}, + origin="root", + created_at=now, + updated_at=now, + ), + ((HistoryEventType.RUN_ADMITTED, {}),), + ) + + asyncio.run(seed()) + url = f"{POSTGRES_URL}?options=-csearch_path%3D{store.schema}" + + listed = CliRunner().invoke(workflows, ["list", "-d", url, "--json"]) + assert listed.exit_code == 0, listed.output + assert [row["run_id"] for row in json.loads(listed.output)] == ["cli-run"] + + shown = CliRunner().invoke(workflows, ["show", "cli-run", "-d", url]) + assert shown.exit_code == 0, shown.output + assert "pg.charge" in shown.output + + cancelled = CliRunner().invoke(workflows, ["cancel", "cli-run", "-d", url]) + assert cancelled.exit_code == 0, cancelled.output diff --git a/uv.lock b/uv.lock index 75a88ad310d..ff3e75cc6c7 100644 --- a/uv.lock +++ b/uv.lock @@ -585,7 +585,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/66/54/eb9bfc647b19f2009dd5c7f5ec51c4e6ca831725f1aea7a993034f483147/contourpy-1.3.2.tar.gz", hash = "sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54", size = 13466130, upload-time = "2025-04-15T17:47:53.79Z" } wheels = [ @@ -663,7 +663,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } @@ -1003,7 +1003,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1879,15 +1879,15 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "contourpy", version = "1.3.2", source = { registry = "https://pypi.org/simple" } }, - { name = "cycler" }, - { name = "fonttools" }, - { name = "kiwisolver" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, - { name = "packaging" }, - { name = "pillow" }, - { name = "pyparsing" }, - { name = "python-dateutil" }, + { name = "contourpy", version = "1.3.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "cycler", marker = "python_full_version < '3.11'" }, + { name = "fonttools", marker = "python_full_version < '3.11'" }, + { name = "kiwisolver", marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "packaging", marker = "python_full_version < '3.11'" }, + { name = "pillow", marker = "python_full_version < '3.11'" }, + { name = "pyparsing", marker = "python_full_version < '3.11'" }, + { name = "python-dateutil", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/63/1b/4be5be87d43d327a0cf4de1a56e86f7f84c89312452406cf122efe2839e6/matplotlib-3.10.9.tar.gz", hash = "sha256:fd66508e8c6877d98e586654b608a0456db8d7e8a546eb1e2600efd957302358", size = 34811233, upload-time = "2026-04-24T00:14:13.539Z" } wheels = [ @@ -1963,16 +1963,16 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" } }, - { name = "cycler" }, - { name = "fonttools" }, - { name = "kiwisolver" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "cycler", marker = "python_full_version >= '3.11'" }, + { name = "fonttools", marker = "python_full_version >= '3.11'" }, + { name = "kiwisolver", marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "packaging" }, - { name = "pillow" }, - { name = "pyparsing" }, - { name = "python-dateutil" }, + { name = "packaging", marker = "python_full_version >= '3.11'" }, + { name = "pillow", marker = "python_full_version >= '3.11'" }, + { name = "pyparsing", marker = "python_full_version >= '3.11'" }, + { name = "python-dateutil", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/49/64/f9a391af28f518b11ad45a8a712353c94a0aefce09d3703200e5c54b610a/matplotlib-3.11.1.tar.gz", hash = "sha256:69647db5746941c793d6e445a4cd349323ffb87d9cc958c2ad84a659b4832d30", size = 32612045, upload-time = "2026-07-18T03:39:46.63Z" } wheels = [ @@ -2534,10 +2534,10 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, - { name = "python-dateutil" }, - { name = "pytz" }, - { name = "tzdata" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "python-dateutil", marker = "python_full_version < '3.11'" }, + { name = "pytz", marker = "python_full_version < '3.11'" }, + { name = "tzdata", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" } wheels = [ @@ -2606,10 +2606,10 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "python-dateutil" }, - { name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, + { name = "python-dateutil", marker = "python_full_version >= '3.11'" }, + { name = "tzdata", marker = "(python_full_version >= '3.11' and sys_platform == 'emscripten') or (python_full_version >= '3.11' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/be/4f/5f3422a2afec5ffc46308b79e53291365a93748b498ac2e58bead0197916/pandas-3.0.5.tar.gz", hash = "sha256:dca3734d6ab7c906e6730f0788b0a1dbb9f2467731f9711f77995c8e9d62d712", size = 4658219, upload-time = "2026-07-22T22:19:28.819Z" } wheels = [ @@ -3034,6 +3034,9 @@ wheels = [ binary = [ { name = "psycopg-binary", marker = "implementation_name != 'pypy'" }, ] +pool = [ + { name = "psycopg-pool" }, +] [[package]] name = "psycopg-binary" @@ -3097,6 +3100,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/eb/e6/5fff07a70d1f945ed90ae131c3bd76cab32beff7c58c6db15ad5820b6d1f/psycopg_binary-3.3.4-cp314-cp314-win_amd64.whl", hash = "sha256:c37e024c07308cd06cf3ec51bfd0e7f6157585a4d84d1bce4a7f5f7913719bf8", size = 3666849, upload-time = "2026-05-01T23:31:51.165Z" }, ] +[[package]] +name = "psycopg-pool" +version = "3.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/90/82/7a23d26039827ecd4ebe93905651029ddd307c5182ad59296dfb6f67b528/psycopg_pool-3.3.1.tar.gz", hash = "sha256:b10b10b7a175d5cc1592147dc5b7eec8a9e0834eb3ed2c4a92c858e2f51eb63c", size = 31661, upload-time = "2026-05-01T23:31:59.809Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/ed/89c2c620af0e1660354cd8aabf9f5b21f911597ce22acb37c805d6c86bc8/psycopg_pool-3.3.1-py3-none-any.whl", hash = "sha256:2af5b432941c4c9ad5c87b3fa410aec910ec8f7c122855897983a06c45f2e4b5", size = 40023, upload-time = "2026-05-01T23:31:53.136Z" }, +] + [[package]] name = "py-cpuinfo" version = "9.0.0" @@ -3697,7 +3712,7 @@ dev = [ { name = "plotly" }, { name = "pre-commit" }, { name = "psutil" }, - { name = "psycopg", extra = ["binary"] }, + { name = "psycopg", extra = ["binary", "pool"] }, { name = "pydantic" }, { name = "pyright" }, { name = "pytest" }, @@ -3778,7 +3793,7 @@ dev = [ { name = "plotly" }, { name = "pre-commit" }, { name = "psutil" }, - { name = "psycopg", extras = ["binary"] }, + { name = "psycopg", extras = ["binary", "pool"] }, { name = "pydantic" }, { name = "pyright" }, { name = "pytest" }, @@ -4304,7 +4319,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -4365,7 +4380,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [ @@ -4444,7 +4459,7 @@ resolution-markers = [ "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" } wheels = [ From 8940546407d4ce91145b631b39c23598e7b18e99 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 19:15:09 -0700 Subject: [PATCH 027/121] Approve from an email: signed links plus a per-attempt run context A run waiting on a decision is the most common human-in-the-loop shape, and the naive implementation -- a URL carrying a run id -- is an open door. An approval link here is an HMAC token over run, channel, payload, delivery key, and expiry, so an edited link is refused rather than believed; it is spent once; and it never decides on GET, because mail clients and scanners fetch URLs before a person reads the message, so a link that approved on GET would approve itself in transit. The secret comes from REFLEX_WORKFLOW_APPROVAL_SECRET with deliberately no default, and a server missing it says 'not configured' instead of masquerading as an expired link. Executing the docs example before committing caught a real defect: a channel declared with a pydantic model failed to serialize into the token -- and every realistic channel is typed -- so payloads now go through the same reduction the signal path uses. Links need to know which run built them, which is a capability handlers were missing generally. rx.current_run() now exposes the attempt's identity -- run, workflow, slot, attempt, epoch -- bound per attempt via a ContextVar that to_thread carries into sync handlers, plus idempotency_key(): stable across retries of a step, distinct across steps, which is exactly the contract a payment API's idempotency header wants. Chasing an intermittent test hang also root-caused a real defect: the worker loop's treated CancelledError as a retryable error, so anything that cancelled the worker task without calling aclose() -- a task group, a supervisor, an event loop tearing down -- waited forever on a task that had gone back to polling. The hang reproduced with the fix removed and disappears with it. The kernel also cancels its in-flight attempts on close instead of leaving them running against a store nobody reads, and the test harness closes stores it created, which previously leaked a Postgres pool into every later test in the process. --- docs/workflows/overview.md | 48 +++ news/workflow-approval-links.feature.md | 1 + reflex/__init__.py | 2 + reflex/app.py | 20 +- reflex/workflow/__init__.py | 5 + reflex/workflow/approvals.py | 319 +++++++++++++++++ reflex/workflow/context.py | 116 ++++++ reflex/workflow/kernel.py | 68 +++- reflex/workflow/testing.py | 12 +- tests/units/workflow/test_approvals.py | 428 +++++++++++++++++++++++ tests/units/workflow/test_concurrency.py | 42 +++ tests/units/workflow/test_context.py | 130 +++++++ 12 files changed, 1172 insertions(+), 19 deletions(-) create mode 100644 news/workflow-approval-links.feature.md create mode 100644 reflex/workflow/approvals.py create mode 100644 reflex/workflow/context.py create mode 100644 tests/units/workflow/test_approvals.py create mode 100644 tests/units/workflow/test_context.py diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index 0bd2be5e261..4f0e7303e26 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -184,6 +184,54 @@ class ReviewPage(rx.State): Use `timeout=rx.never` to wait indefinitely. A signal that arrives before the run reaches its wait is buffered and applied as soon as the wait arms, so a fast approver never blocks the run. +### Approving from an email + +Not every approver will open your app. `rx.approval_link()` builds a signed URL that delivers one +decision to the run that created it, so the reply can come straight from a message: + +```python +@rx.event(durable=True, effect="idempotent_write") +async def ask(self): + approve = rx.approval_link( + Expense.review(Decision(approved=True, by="manager")), + base_url="https://app.example.com", + ) + reject = rx.approval_link( + Expense.review(Decision(approved=False, by="manager")), + base_url="https://app.example.com", + ) + await send_email(self.manager, approve_url=approve, reject_url=reject) + return rx.wait_for( + Expense.review, then=Expense.decide, timeout="3d", on_timeout=Expense.escalate + ) +``` + +Set `REFLEX_WORKFLOW_APPROVAL_SECRET` to a long random string; links are signed with it and there is +no default, because a built-in secret would make every deployment's links forgeable. The signature +covers the run, the channel, the payload, and the expiry, so an edited link is refused rather than +believed. + +Following a link shows a confirmation page, and only submitting it records the decision. That is not +politeness: mail clients and link scanners fetch URLs before a person reads the message, so a link +that decided on `GET` would approve itself in transit. A link is spent once — replaying it, or +following the other choice afterwards, changes nothing. + +### Knowing which run you are + +`rx.current_run()` returns the attempt a durable handler is executing, which is what you need to +correlate logs — and to make an outbound call safely: + +```python +@rx.event(durable=True, effect="non_idempotent_write") +async def charge(self): + run = rx.current_run() + await stripe.charge(self.amount, idempotency_key=run.idempotency_key()) +``` + +The key is stable across retries of that step and different for every other step, which is exactly +the contract a payment API's idempotency key wants: a retry must not charge twice, and the next step +must not be mistaken for this one. + ## Running work in parallel Each branch of a fan-out becomes its own run, with its own state, retries, and history, so a slow or diff --git a/news/workflow-approval-links.feature.md b/news/workflow-approval-links.feature.md new file mode 100644 index 00000000000..451b0ecd321 --- /dev/null +++ b/news/workflow-approval-links.feature.md @@ -0,0 +1 @@ +`rx.approval_link()` builds signed, single-use, expiring URLs that let someone answer a waiting run from an email, and `rx.current_run()` exposes the attempt a durable handler is executing, including a retry-stable idempotency key. diff --git a/reflex/__init__.py b/reflex/__init__.py index bbcce18d6a2..b470a9dd8e2 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -257,6 +257,8 @@ "fail", "needs_attention", "workflows", + "approval_link", + "current_run", ], } diff --git a/reflex/app.py b/reflex/app.py index c829cf70ffb..94ecd5e72d2 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -846,7 +846,8 @@ def _add_default_endpoints(self): ) def _add_workflow_endpoints(self): - """Add the webhook ingress endpoint when a workflow declares one.""" + """Add the workflow ingress endpoints: webhooks in, approvals back.""" + from reflex.workflow.approvals import APPROVAL_ROUTE, approval_endpoint from reflex.workflow.ingress import ( WEBHOOK_ROUTE, collect_webhook_routes, @@ -855,13 +856,20 @@ def _add_workflow_endpoints(self): if self._api is None or self._workflow_runtime is None: return - if not collect_webhook_routes(self._workflow_runtime.definitions): - return config = get_config() + if collect_webhook_routes(self._workflow_runtime.definitions): + self._api.add_route( + config.prepend_backend_path(WEBHOOK_ROUTE), + webhook_endpoint(self._workflow_runtime), + methods=["POST"], + ) + # Approval links are addressed to a run, not to a workflow, so the + # route is mounted whenever workflows are served rather than being + # driven by a declaration. self._api.add_route( - config.prepend_backend_path(WEBHOOK_ROUTE), - webhook_endpoint(self._workflow_runtime), - methods=["POST"], + config.prepend_backend_path(APPROVAL_ROUTE), + approval_endpoint(self._workflow_runtime), + methods=["GET", "POST"], ) def _add_optional_endpoints(self): diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index dc4da9cf8d2..f6d3ac2c8f0 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -40,7 +40,9 @@ webhook, ) +from reflex.workflow.approvals import approval_link from reflex.workflow.conformance import CONFORMANCE_CHECKS +from reflex.workflow.context import RunContext, current_run from reflex.workflow.definition import ( HandlerDefinition, WorkflowDefinition, @@ -84,6 +86,7 @@ "Parallel", "RateLimit", "Retry", + "RunContext", "RunQuery", "RunRecord", "RunSnapshot", @@ -110,8 +113,10 @@ "WorkflowRuntime", "WorkflowTestHarness", "after", + "approval_link", "compile_workflow", "complete", + "current_run", "fail", "get_runtime", "hmac_signature", diff --git a/reflex/workflow/approvals.py b/reflex/workflow/approvals.py new file mode 100644 index 00000000000..78ababea69f --- /dev/null +++ b/reflex/workflow/approvals.py @@ -0,0 +1,319 @@ +"""Signed links that let a person answer a waiting run from an email. + +A run waiting on a decision is common enough that every workflow engine grows +some version of it, and the naive version is a URL carrying a run id. That is +an open door: anyone who guesses or forwards it can approve. Here a link is a +signed token that names exactly one run, one channel, and one payload, expires, +and can be spent once. + +Two details are the whole security story: + +*Nothing is trusted from the URL.* The token carries its own signature over +every field, so a recipient who edits the run id, the payload, or the expiry +invalidates it. The secret never leaves the server. + +*A GET never decides anything.* Mail clients and link scanners fetch URLs +before a human sees them, so a link that approves on GET approves itself. The +GET renders a confirmation the person submits, and only the POST delivers. +""" + +from __future__ import annotations + +import base64 +import hashlib +import hmac +import json +import os +import time +from html import escape +from typing import TYPE_CHECKING, Any, Final + +from reflex_base.utils import console +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import ChannelDelivery, parse_duration +from starlette.responses import HTMLResponse, JSONResponse, Response + +from reflex.workflow.context import require_run +from reflex.workflow.serde import to_run_data + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine + + from reflex_base.workflow import DurationLike + from starlette.requests import Request + + from reflex.workflow.runtime import WorkflowRuntime + +SECRET_ENV: Final = "REFLEX_WORKFLOW_APPROVAL_SECRET" +APPROVAL_ROUTE: Final = "/_workflow/approve/{token:path}" +DEFAULT_EXPIRY: Final = "7d" +MAX_TOKEN_BYTES: Final = 4096 + + +def _secret() -> bytes: + """Read the signing secret. + + Returns: + The secret as bytes. + + Raises: + WorkflowRuntimeError: If the environment does not carry one. + """ + secret = os.environ.get(SECRET_ENV) + if not secret: + msg = ( + f"Approval links must be signed, so {SECRET_ENV} has to be set to a " + "long random string. There is deliberately no default: a built-in " + "secret would make every deployment's links forgeable." + ) + raise WorkflowRuntimeError(msg) + return secret.encode() + + +def _b64(raw: bytes) -> str: + """Encode bytes for a URL path segment. + + Args: + raw: The bytes to encode. + + Returns: + Unpadded base64url text. + """ + return base64.urlsafe_b64encode(raw).decode().rstrip("=") + + +def _unb64(text: str) -> bytes: + """Decode a URL path segment back to bytes. + + Args: + text: Unpadded base64url text. + + Returns: + The decoded bytes. + """ + return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4)) + + +def _sign(body: bytes) -> str: + """Sign a token body. + + Args: + body: The encoded claims. + + Returns: + The signature, base64url encoded. + """ + return _b64(hmac.new(_secret(), body, hashlib.sha256).digest()) + + +def approval_link( + delivery: ChannelDelivery, + *, + base_url: str = "", + expires_in: DurationLike = DEFAULT_EXPIRY, + key: str | None = None, +) -> str: + """Build a signed link that delivers one decision to the current run. + + Call it from inside a durable handler; the run it addresses is the run the + handler is running in. Build one link per choice and put them both in the + message:: + + approve = rx.approval_link(Expense.decided({"ok": True})) + reject = rx.approval_link(Expense.decided({"ok": False})) + + Args: + delivery: The channel and payload to deliver, e.g. + ``Expense.decided({"ok": True})``. + base_url: Origin to prefix, e.g. ``"https://app.example.com"``. Leave + empty for a path, which is what a relative link needs. + expires_in: How long the link stays valid. + key: Delivery identity. Two links sharing a key are the same decision, + so the second one spent is a no-op; distinct keys let one person + approve after another rejected. Defaults to a key derived from the + channel and payload, which makes a link single-use. + + Returns: + The URL. + """ + context = require_run("approval_link()") + # A channel is typically declared with a model, so the payload has to be + # reduced to plain data here rather than at delivery: the token carries it + # across a network boundary, and only plain data survives the round trip. + payload = to_run_data({"value": delivery.payload})["value"] + if key is None: + material = json.dumps( + [context.run_id, delivery.channel, payload], sort_keys=True + ) + key = hashlib.sha256(material.encode()).hexdigest()[:32] + claims = { + "r": context.run_id, + "c": delivery.channel, + "p": payload, + "k": key, + "e": time.time() + parse_duration(expires_in, param="expires_in"), + } + body = json.dumps(claims, separators=(",", ":"), sort_keys=True).encode() + token = f"{_b64(body)}.{_sign(body)}" + return f"{base_url.rstrip('/')}/_workflow/approve/{token}" + + +def decode_token(token: str) -> dict[str, Any]: + """Verify a token and return its claims. + + Args: + token: The token from the URL. + + Returns: + The verified claims. + + Raises: + WorkflowRuntimeError: If the token is malformed, unsigned, forged, or + expired. The message is deliberately the same for every failure, + so a caller cannot use it to probe what a valid token looks like. + """ + invalid = "This approval link is not valid. It may have expired." + if len(token) > MAX_TOKEN_BYTES or token.count(".") != 1: + raise WorkflowRuntimeError(invalid) + encoded, signature = token.split(".") + try: + body = _unb64(encoded) + except Exception as exc: + raise WorkflowRuntimeError(invalid) from exc + if not hmac.compare_digest(signature, _sign(body)): + raise WorkflowRuntimeError(invalid) + try: + claims = json.loads(body) + except json.JSONDecodeError as exc: + raise WorkflowRuntimeError(invalid) from exc + if not isinstance(claims, dict) or not {"r", "c", "p", "k", "e"} <= claims.keys(): + raise WorkflowRuntimeError(invalid) + if not isinstance(claims["e"], (int, float)) or claims["e"] < time.time(): + raise WorkflowRuntimeError(invalid) + return claims + + +_PAGE = """ + + +{title} + +

{title}

{message}

{form}
+""" + +_FORM = '
' + + +def _page(title: str, message: str, *, form: str = "", status: int = 200) -> Response: + """Render one of the confirmation pages. + + Args: + title: Heading text. + message: Body text. + form: Optional submit form. + status: HTTP status code. + + Returns: + The response. + """ + return HTMLResponse( + _PAGE.format(title=escape(title), message=escape(message), form=form), + status_code=status, + ) + + +def _wants_json(request: Request) -> bool: + """Whether the caller asked for JSON rather than a page. + + Args: + request: The incoming request. + + Returns: + True when JSON was requested explicitly. + """ + return "application/json" in request.headers.get("accept", "") + + +def approval_endpoint( + runtime: WorkflowRuntime, +) -> Callable[[Request], Coroutine[Any, Any, Response]]: + """Build the endpoint that spends approval links. + + Args: + runtime: The runtime owning the runs. + + Returns: + The endpoint callable. + """ + + async def endpoint(request: Request) -> Response: + """Confirm on GET, deliver on POST. + + Args: + request: The incoming request. + + Returns: + The response. + """ + token = request.path_params.get("token", "") + try: + _secret() + except WorkflowRuntimeError as err: + # A server with no secret cannot tell a good link from a forged + # one. Saying "expired" would send an operator hunting a data + # problem instead of a configuration one. + console.error(f"Approval link rejected: {err}") + message = "Approvals are not configured on this server." + if _wants_json(request): + return JSONResponse({"error": message}, status_code=500) + return _page("Not available", message, status=500) + try: + claims = decode_token(token) + except WorkflowRuntimeError as err: + if _wants_json(request): + return JSONResponse({"error": str(err)}, status_code=400) + return _page("Link not valid", str(err), status=400) + + if request.method == "GET": + # A mail client or scanner may fetch this; only a person submits. + if _wants_json(request): + return JSONResponse({"status": "confirm", "run_id": claims["r"]}) + return _page( + "Confirm your response", + "Submitting this records your decision on the waiting run.", + form=_FORM.format(label="Confirm"), + ) + + disposition = await runtime.kernel.signal( + claims["r"], + ChannelDelivery(channel=claims["c"], payload=claims["p"]), + key=claims["k"], + ) + if _wants_json(request): + return JSONResponse({"status": disposition, "run_id": claims["r"]}) + if disposition == "resolved": + return _page("Thank you", "Your response has been recorded.") + if disposition == "duplicate": + return _page("Already recorded", "This link has already been used.") + if disposition in ("expired", "run_terminal", "unknown_run"): + return _page( + "No longer waiting", + "This request has already been resolved another way.", + status=409, + ) + return _page( + "Recorded", + "Your response was saved and will apply when the run reaches it.", + ) + + return endpoint diff --git a/reflex/workflow/context.py b/reflex/workflow/context.py new file mode 100644 index 00000000000..6330c63aff3 --- /dev/null +++ b/reflex/workflow/context.py @@ -0,0 +1,116 @@ +"""What a durable handler can learn about the attempt it is running in. + +A handler is a plain method: it takes its payload and returns its transition, +and nothing about the engine leaks into that signature. But some things a +handler legitimately needs -- a correlation id for its logs, a stable +idempotency key for an outbound HTTP call, a link that addresses this run -- +are properties of the attempt rather than of the payload. They are exposed +here instead of being threaded through every handler. + +The context is per-attempt and never spans one: reading it outside a durable +handler returns None rather than a stale value from a previous step. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +from contextvars import ContextVar, Token + + +@dataclasses.dataclass(frozen=True, slots=True) +class RunContext: + """Identity of the attempt a durable handler is running in. + + Attributes: + run_id: The run this attempt belongs to. + workflow_id: The stable workflow identity. + ordinal: The mailbox slot being executed. + handler_id: The handler this slot names. + attempt: Which attempt this is, counting from one. + epoch: The claim fence for this attempt. + """ + + run_id: str + workflow_id: str + ordinal: int + handler_id: str + attempt: int + epoch: int + + def idempotency_key(self, *, scope: str = "") -> str: + """Derive a stable key for an outbound call made by this step. + + Every retry of a step is the *same* logical work, so the key must not + change between attempts -- that is the whole point of handing it to a + payment or messaging API. It does change when the step changes, so two + different steps of one run never collide. + + Args: + scope: Distinguishes several calls made by one handler. + + Returns: + A hex key, stable across retries of this step. + """ + material = f"{self.run_id}:{self.ordinal}:{self.handler_id}:{scope}" + return hashlib.sha256(material.encode()).hexdigest()[:32] + + +_current: ContextVar[RunContext | None] = ContextVar( + "reflex_workflow_run_context", default=None +) + + +def current_run() -> RunContext | None: + """Read the attempt this code is running in. + + Returns: + The context, or None outside a durable handler. + """ + return _current.get() + + +def require_run(reason: str) -> RunContext: + """Read the current attempt, refusing to continue without one. + + Args: + reason: What the caller needed the context for, used in the message. + + Returns: + The context. + + Raises: + WorkflowRuntimeError: If called outside a durable handler. + """ + from reflex_base.utils.exceptions import WorkflowRuntimeError + + context = _current.get() + if context is None: + msg = ( + f"{reason} needs the run it belongs to, and there is no durable " + "handler running. Call it from inside a @rx.event(durable=True) " + "handler." + ) + raise WorkflowRuntimeError(msg) + return context + + +def bind_run(context: RunContext) -> Token[RunContext | None]: + """Install a run context for the duration of one attempt. + + Args: + context: The attempt's identity. + + Returns: + The token that restores the previous context. + """ + return _current.set(context) + + +def unbind_run(token: Token[RunContext | None]) -> None: + """Restore the context that was in place before an attempt. + + Args: + token: The token returned by ``bind_run``. + """ + _current.reset(token) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 480fc466dc4..9decad078ef 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -39,6 +39,7 @@ ) from reflex.event import EventHandler, EventSpec +from reflex.workflow.context import RunContext, bind_run, unbind_run from reflex.workflow.cron import CronSchedule from reflex.workflow.records import ( TERMINAL_RUN_STATUSES, @@ -918,7 +919,11 @@ def _interpret_return( return [self._resolve_successor(defn, value)], None async def _invoke( - self, handler: HandlerDefinition, instance: BaseState, args: dict[str, Any] + self, + handler: HandlerDefinition, + instance: BaseState, + args: dict[str, Any], + context: RunContext, ) -> Any: """Invoke a handler attempt with its per-attempt timeout. @@ -926,6 +931,7 @@ async def _invoke( handler: The handler definition. instance: The hydrated run-state instance. args: The step payload. + context: Identity of this attempt, readable inside the handler. Returns: The handler return value. @@ -941,13 +947,19 @@ async def _invoke( payload = _transform_event_payload(args, handler.type_hints) except Exception: payload = dict(args) - if handler.is_async: - coroutine = handler.fn(instance, **payload) - else: - coroutine = asyncio.to_thread(handler.fn, instance, **payload) - if handler.timeout is not None: - return await asyncio.wait_for(coroutine, timeout=handler.timeout) - return await coroutine + token = bind_run(context) + try: + if handler.is_async: + coroutine = handler.fn(instance, **payload) + else: + # to_thread copies the current context, so a sync handler sees + # the same attempt identity as an async one. + coroutine = asyncio.to_thread(handler.fn, instance, **payload) + if handler.timeout is not None: + return await asyncio.wait_for(coroutine, timeout=handler.timeout) + return await coroutine + finally: + unbind_run(token) def _build_new_steps( self, @@ -1738,7 +1750,19 @@ async def _execute_claim(self, claim: Claim) -> None: try: instance = self._hydrate(defn, claim.run.state) lease.attempt = asyncio.ensure_future( - self._invoke(handler, instance, claim.step.args) + self._invoke( + handler, + instance, + claim.step.args, + RunContext( + run_id=claim.run.run_id, + workflow_id=claim.run.workflow_id, + ordinal=claim.step.ordinal, + handler_id=handler.id, + attempt=claim.step.attempts + 1, + epoch=claim.step.epoch, + ), + ) ) value = await lease.attempt finally: @@ -2169,6 +2193,17 @@ async def _worker_loop(self) -> None: await self.recover() if await self._tick(): continue + if len(self._inflight) >= self._max_concurrency: + # Every slot is busy, so the only thing that can change is + # an attempt finishing; due times cannot matter until one + # does. A round normally blocks on the attempts it started, + # so this is a guard on the invariant rather than a path + # the loop is expected to take. + await asyncio.wait( + list(self._inflight.values()), + return_when=asyncio.FIRST_COMPLETED, + ) + continue now = self._clock() due = await self._store.next_due(now) delay = min(self._poll_interval, max(self._next_recovery_at - now, 0.0)) @@ -2180,9 +2215,15 @@ async def _worker_loop(self) -> None: self._wakeup.clear() with contextlib.suppress(TimeoutError): await asyncio.wait_for(self._wakeup.wait(), timeout=delay) - except BaseException as err: - if self._closing: - raise + except asyncio.CancelledError: + # Cancellation is never a retryable error. Catching it here -- + # which `except BaseException` did -- made the worker + # unkillable by anything except aclose(): a supervisor, a task + # group, or an event loop shutting down would cancel it, be + # told nothing, and wait forever for a task that had already + # gone back to polling. + raise + except Exception as err: console.error(f"Workflow worker error, retrying: {err!r}") await asyncio.sleep(self._poll_interval) @@ -2208,3 +2249,6 @@ async def aclose(self) -> None: with contextlib.suppress(asyncio.CancelledError): await self._worker self._worker = None + # The kernel owns its attempts, so closing it stops them rather than + # leaving them running against a store nobody is reading any more. + await self._cancel_inflight() diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index af5d6ee03c5..cb59851a232 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -12,6 +12,7 @@ from __future__ import annotations +import inspect from typing import TYPE_CHECKING, Any from reflex_base.workflow import ( @@ -91,8 +92,12 @@ def __init__( virtual clock stay deterministic. """ self._clock = _VirtualClock(start_time) + # A store the harness built is the harness's to close. Leaving a + # pooled store open leaks its connections and its worker threads into + # every later test in the process. + self._owned_store = MemoryRunStore() if store is None else None self._runtime = WorkflowRuntime( - store if store is not None else MemoryRunStore(), + store if store is not None else self._owned_store, clock=self._clock, rng=lambda: 1.0, lease_duration=parse_duration(lease_duration), @@ -152,6 +157,11 @@ async def __aexit__(self, *exc_info) -> None: _context_runtime.reset(self._token) self._token = None await self._runtime.shutdown() + closer = getattr(self._owned_store, "close", None) + if closer is not None: + closed = closer() + if inspect.isawaitable(closed): + await closed async def start( self, diff --git a/tests/units/workflow/test_approvals.py b/tests/units/workflow/test_approvals.py new file mode 100644 index 00000000000..9e0dfe3946a --- /dev/null +++ b/tests/units/workflow/test_approvals.py @@ -0,0 +1,428 @@ +"""Tests for signed approval links. + +An approval link is a bearer credential sent by email, so most of these are +about what a link must *not* do: decide anything on a GET, survive an edit, +outlive its expiry, or work twice. +""" + +import json +import time + +import pytest +from pydantic import BaseModel +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import WorkflowConfig, manual +from starlette.applications import Starlette +from starlette.routing import Route +from starlette.testclient import TestClient + +import reflex as rx +from reflex.workflow.approvals import ( + APPROVAL_ROUTE, + SECRET_ENV, + _b64, + approval_endpoint, +) +from reflex.workflow.context import RunContext, bind_run, unbind_run +from reflex.workflow.records import RunStatus +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import MemoryRunStore + +SECRET = "test-approval-secret" + + +class Verdict(BaseModel): + """A typed decision payload.""" + + approved: bool + by: str + + +LINKS: dict[str, str] = {} + + +class Expense(rx.State): + """An expense that waits for a manager's decision.""" + + __workflow__ = WorkflowConfig(id="approval.expense") + + decided = rx.Signal(dict) + amount: int = 0 + outcome: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def submit(self, amount: int): + """Send the manager two links and wait. + + Args: + amount: The amount claimed. + + Returns: + A wait on the decision channel. + """ + self.amount = amount + LINKS["approve"] = rx.approval_link(Expense.decided({"ok": True})) + LINKS["reject"] = rx.approval_link(Expense.decided({"ok": False})) + return rx.wait_for( + Expense.decided, + then=Expense.record, + timeout="7d", + on_timeout=Expense.lapse, + ) + + @rx.event(durable=True, effect="none") + def lapse(self): + """Nobody answered in time. + + Returns: + Completion. + """ + self.outcome = "lapsed" + return rx.complete(result={"outcome": self.outcome}) + + @rx.event(durable=True, effect="none") + def record(self, decision: dict): + """Record what the manager said. + + Args: + decision: The delivered payload. + + Returns: + Completion. + """ + self.outcome = "approved" if decision["ok"] else "rejected" + return rx.complete(result={"outcome": self.outcome}) + + +@pytest.fixture +def approving(monkeypatch, forked_registration_context): + """A started run waiting on a decision, with its links built. + + Args: + monkeypatch: Used to set the signing secret. + forked_registration_context: Isolates state registration. + + Returns: + The runtime and a Starlette app serving the approval endpoint. + """ + monkeypatch.setenv(SECRET_ENV, SECRET) + LINKS.clear() + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Expense) + app = Starlette( + routes=[ + Route(APPROVAL_ROUTE, approval_endpoint(runtime), methods=["GET", "POST"]) + ] + ) + return runtime, app + + +async def _start(runtime) -> str: + """Start one expense run and let it reach its wait. + + Args: + runtime: The runtime to start it on. + + Returns: + The run id. + """ + await runtime.startup(start_worker=False) + result = await runtime.kernel.start(Expense.submit(120)) + await runtime.kernel.run_until_idle() + assert result.run_id is not None + return result.run_id + + +async def test_a_get_never_decides_anything(approving): + """Fetching the link must not record a decision. + + Mail clients and link scanners fetch URLs before a person sees them, so a + link that approves on GET approves itself in transit. + """ + runtime, app = approving + run_id = await _start(runtime) + + with TestClient(app) as client: + response = client.get(LINKS["approve"]) + assert response.status_code == 200 + assert "Confirm" in response.text + + snapshot = await runtime.kernel.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert snapshot.state["outcome"] == "" + + +async def test_a_post_records_the_decision(approving): + """Submitting the confirmation delivers the payload the link names.""" + runtime, app = approving + run_id = await _start(runtime) + + with TestClient(app) as client: + response = client.post(LINKS["approve"]) + assert response.status_code == 200 + assert "recorded" in response.text.lower() + + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"outcome": "approved"} + + +async def test_the_reject_link_carries_its_own_payload(approving): + """Each choice is a separate link, so the URL fixes the answer.""" + runtime, app = approving + run_id = await _start(runtime) + + with TestClient(app) as client: + assert client.post(LINKS["reject"]).status_code == 200 + + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(run_id) + assert snapshot is not None + assert snapshot.result == {"outcome": "rejected"} + + +async def test_a_spent_link_cannot_be_spent_again(approving): + """Replaying a link is a no-op, not a second decision.""" + runtime, app = approving + await _start(runtime) + + with TestClient(app) as client: + assert client.post(LINKS["approve"]).status_code == 200 + second = client.post(LINKS["approve"]) + assert "already been used" in second.text.lower() + + +async def test_the_losing_link_cannot_overturn_the_decision(approving): + """Once approved, the reject link finds nothing left to decide.""" + runtime, app = approving + run_id = await _start(runtime) + + with TestClient(app) as client: + assert client.post(LINKS["approve"]).status_code == 200 + await runtime.kernel.run_until_idle() + late = client.post(LINKS["reject"]) + + assert late.status_code == 409 + snapshot = await runtime.kernel.get_run(run_id) + assert snapshot is not None + assert snapshot.result == {"outcome": "approved"} + + +async def test_an_edited_token_is_refused(approving): + """Changing any claim invalidates the signature. + + The interesting attack is not a random string: it is a valid link whose + payload has been flipped, which is why the signature covers the claims and + not just the run id. + """ + runtime, app = approving + run_id = await _start(runtime) + + encoded, signature = LINKS["reject"].rsplit("/", 1)[-1].split(".") + claims = json.loads( + __import__("base64").urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)) + ) + claims["p"] = {"ok": True} + forged = _b64(json.dumps(claims, separators=(",", ":"), sort_keys=True).encode()) + + with TestClient(app) as client: + response = client.post(f"/_workflow/approve/{forged}.{signature}") + assert response.status_code == 400 + + snapshot = await runtime.kernel.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + + +async def test_a_token_signed_with_another_secret_is_refused(approving, monkeypatch): + """A link minted by a different deployment does not work here.""" + runtime, app = approving + await _start(runtime) + stolen = LINKS["approve"] + + monkeypatch.setenv(SECRET_ENV, "a-different-secret") + with TestClient(app) as client: + assert client.post(stolen).status_code == 400 + + +async def test_an_expired_token_is_refused(approving, monkeypatch): + """A link stops working once its expiry passes.""" + runtime, app = approving + await _start(runtime) + + real_time = time.time + monkeypatch.setattr(time, "time", lambda: real_time() + 8 * 86400) + with TestClient(app) as client: + response = client.post(LINKS["approve"]) + assert response.status_code == 400 + + +async def test_a_malformed_token_is_refused(approving): + """Garbage in the path is a clean rejection, not a traceback.""" + runtime, app = approving + await _start(runtime) + + with TestClient(app) as client: + for token in ("", "nonsense", "a.b.c", "!!!.###", "x" * 5000): + assert client.post(f"/_workflow/approve/{token}").status_code in ( + 400, + 404, + ) + + +def test_links_refuse_to_sign_without_a_secret(monkeypatch): + """There is no default secret, and no silent fallback to one.""" + monkeypatch.delenv(SECRET_ENV, raising=False) + token = bind_run( + RunContext( + run_id="r1", + workflow_id="approval.expense", + ordinal=0, + handler_id="submit", + attempt=1, + epoch=1, + ) + ) + try: + with pytest.raises(WorkflowRuntimeError, match=SECRET_ENV): + rx.approval_link(Expense.decided({"ok": True})) + finally: + unbind_run(token) + + +def test_links_refuse_to_build_outside_a_handler(monkeypatch): + """A link addresses one run, so there has to be one.""" + monkeypatch.setenv(SECRET_ENV, SECRET) + with pytest.raises(WorkflowRuntimeError, match="durable"): + rx.approval_link(Expense.decided({"ok": True})) + + +async def test_json_callers_get_json(approving): + """An API client can spend a link without parsing HTML.""" + runtime, app = approving + await _start(runtime) + + with TestClient(app) as client: + headers = {"accept": "application/json"} + confirm = client.get(LINKS["approve"], headers=headers) + assert confirm.json()["status"] == "confirm" + spent = client.post(LINKS["approve"], headers=headers) + assert spent.json()["status"] == "resolved" + + +async def test_base_url_prefixes_an_absolute_link(approving, monkeypatch): + """An emailed link needs an origin, not a path.""" + runtime, _ = approving + await runtime.startup(start_worker=False) + + token = bind_run( + RunContext( + run_id="r1", + workflow_id="approval.expense", + ordinal=0, + handler_id="submit", + attempt=1, + epoch=1, + ) + ) + try: + link = rx.approval_link( + Expense.decided({"ok": True}), base_url="https://app.example.com/" + ) + finally: + unbind_run(token) + assert link.startswith("https://app.example.com/_workflow/approve/") + + +async def test_a_server_without_a_secret_says_so(approving, monkeypatch): + """A misconfigured server must not look like an expired link. + + Reporting "expired" would send an operator looking for a data problem when + the real one is a missing environment variable. + """ + runtime, app = approving + await _start(runtime) + link = LINKS["approve"] + + monkeypatch.delenv(SECRET_ENV, raising=False) + with TestClient(app) as client: + response = client.post(link, headers={"accept": "application/json"}) + assert response.status_code == 500 + assert "not configured" in response.json()["error"] + + +class Typed(rx.State): + """A channel declared with a model, which is the common case.""" + + __workflow__ = WorkflowConfig(id="approval.typed") + + review = rx.Signal(Verdict) + outcome: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def ask(self): + """Build a link carrying a model, then wait. + + Returns: + A wait on the review channel. + """ + LINKS["typed"] = rx.approval_link( + Typed.review(Verdict(approved=True, by="ada")) + ) + return rx.wait_for( + Typed.review, then=Typed.decide, timeout="3d", on_timeout=Typed.lapse + ) + + @rx.event(durable=True, effect="none") + def decide(self, verdict: Verdict): + """Record a typed verdict. + + Args: + verdict: The delivered decision, rebuilt as its model. + + Returns: + Completion. + """ + self.outcome = f"{verdict.by}:{verdict.approved}" + return rx.complete(result={"outcome": self.outcome}) + + @rx.event(durable=True, effect="none") + def lapse(self): + """Nobody answered.""" + + +async def test_a_link_can_carry_a_model_payload( + monkeypatch, forked_registration_context +): + """A channel declared with a model must survive the round trip. + + The token crosses a network boundary as JSON, so a pydantic payload has to + be reduced on the way out and rebuilt on the way in. Passing the model + straight into the token fails to serialize -- and every realistic channel + is typed, so a dict-only test would never notice. + """ + monkeypatch.setenv(SECRET_ENV, SECRET) + LINKS.clear() + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Typed) + app = Starlette( + routes=[ + Route(APPROVAL_ROUTE, approval_endpoint(runtime), methods=["GET", "POST"]) + ] + ) + await runtime.startup(start_worker=False) + result = await runtime.kernel.start(Typed.ask) + await runtime.kernel.run_until_idle() + assert result.run_id is not None + + with TestClient(app) as client: + assert client.post(LINKS["typed"]).status_code == 200 + await runtime.kernel.run_until_idle() + + snapshot = await runtime.kernel.get_run(result.run_id) + assert snapshot is not None + assert snapshot.result == {"outcome": "ada:True"} + await runtime.shutdown() diff --git a/tests/units/workflow/test_concurrency.py b/tests/units/workflow/test_concurrency.py index 7e6ab7f791c..8864da72f1f 100644 --- a/tests/units/workflow/test_concurrency.py +++ b/tests/units/workflow/test_concurrency.py @@ -1,6 +1,7 @@ """Tests for running attempts from different runs at the same time.""" import asyncio +import contextlib from reflex_base.workflow import WorkflowConfig, manual @@ -140,3 +141,44 @@ async def second(self): break await asyncio.sleep(0.02) assert order == ["first", "second"] + + +class _Idle(rx.State): + """A workflow that is never started, so the worker only polls.""" + + __workflow__ = WorkflowConfig(id="conc.idle") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Do nothing.""" + + +async def test_the_worker_dies_when_its_task_is_cancelled( + forked_registration_context, +): + """Cancelling the worker task must stop it, without going through aclose. + + Not everything that stops a worker calls aclose: a supervisor, a task + group, or an event loop being torn down all just cancel the task. A loop + that treats cancellation as a retryable error goes back to polling, and + then nothing can stop it -- the process hangs on shutdown waiting for a + task that has already resumed work. + """ + runtime = WorkflowRuntime(MemoryRunStore(), poll_interval=0.05) + runtime.register(_Idle) + await runtime.startup() + worker = runtime.kernel._worker + assert worker is not None + # Let the loop reach its idle wait, which is where a shutdown finds it. + await asyncio.sleep(0.15) + + worker.cancel() + for _ in range(40): + if worker.done(): + break + await asyncio.sleep(0.05) + + assert worker.done(), "the worker went back to polling after being cancelled" + with contextlib.suppress(asyncio.CancelledError): + await worker + await runtime.shutdown() diff --git a/tests/units/workflow/test_context.py b/tests/units/workflow/test_context.py new file mode 100644 index 00000000000..16016fd2b00 --- /dev/null +++ b/tests/units/workflow/test_context.py @@ -0,0 +1,130 @@ +"""Tests for the per-attempt run context.""" + +import pytest +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.context import RunContext, current_run, require_run +from reflex.workflow.testing import WorkflowTestHarness + +SEEN: list = [] + + +class Reporter(rx.State): + """Records the context each of its steps ran in.""" + + __workflow__ = WorkflowConfig(id="ctx.reporter") + n: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + def first(self): + """Note the context, then move on. + + Returns: + The second step. + """ + SEEN.append(current_run()) + return Reporter.second + + @rx.event(durable=True, effect="none") + def second(self): + """Note the context again.""" + SEEN.append(current_run()) + + +class SyncReporter(rx.State): + """A synchronous handler, which runs in a worker thread.""" + + __workflow__ = WorkflowConfig(id="ctx.sync") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Note the context from off the event loop.""" + SEEN.append(current_run()) + + +class Flaky(rx.State): + """Fails once, so its attempts can be told apart.""" + + __workflow__ = WorkflowConfig(id="ctx.flaky") + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=3, initial_delay="1s", jitter="none"), + ) + def go(self): + """Record the attempt, failing the first time. + + Raises: + TransientWorkflowError: On the first attempt. + """ + context = require_run("test") + SEEN.append((context.attempt, context.idempotency_key())) + if context.attempt == 1: + msg = "not yet" + raise TransientWorkflowError(msg) + + +async def test_each_step_sees_its_own_identity(forked_registration_context): + """A handler can learn the run and slot it is executing.""" + SEEN.clear() + async with WorkflowTestHarness(Reporter) as harness: + result = await harness.start(Reporter.first) + + assert len(SEEN) == 2 + first, second = SEEN + assert isinstance(first, RunContext) + assert isinstance(second, RunContext) + assert first.run_id == result.run_id == second.run_id + assert first.workflow_id == "ctx.reporter" + assert (first.handler_id, second.handler_id) == ("first", "second") + assert first.ordinal < second.ordinal + + +async def test_a_sync_handler_sees_it_too(forked_registration_context): + """Running off the event loop must not lose the context.""" + SEEN.clear() + async with WorkflowTestHarness(SyncReporter) as harness: + await harness.start(SyncReporter.go) + + assert len(SEEN) == 1 + assert isinstance(SEEN[0], RunContext) + assert SEEN[0].workflow_id == "ctx.sync" + + +async def test_the_idempotency_key_survives_a_retry(forked_registration_context): + """Retrying a step is the same logical call, so the key must not move. + + A payment API keyed on this would otherwise charge twice on a retry, which + is the exact failure the key exists to prevent. + """ + SEEN.clear() + async with WorkflowTestHarness(Flaky) as harness: + await harness.start(Flaky.go) + await harness.advance("2s") + + assert [attempt for attempt, _ in SEEN] == [1, 2] + assert len({key for _, key in SEEN}) == 1 + + +async def test_different_steps_get_different_keys(forked_registration_context): + """Two steps of one run must not share an idempotency key.""" + SEEN.clear() + async with WorkflowTestHarness(Reporter) as harness: + await harness.start(Reporter.first) + + first, second = SEEN + assert isinstance(first, RunContext) + assert isinstance(second, RunContext) + assert first.idempotency_key() != second.idempotency_key() + assert first.idempotency_key() != first.idempotency_key(scope="second-call") + + +def test_there_is_no_context_outside_a_handler(): + """Ordinary application code sees None, not a stale attempt.""" + assert current_run() is None + with pytest.raises(WorkflowRuntimeError, match="durable"): + require_run("something") From 83a3e96289b7937403575bcd71dc0f86bfb0d2bb Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 20:47:21 -0700 Subject: [PATCH 028/121] Record substeps with rx.step, and route steps to worker queues A handler is the unit of retry, so a handler that makes three calls and fails after the second repeats the first two on retry -- three charges for one order. rx.step(name, fn, ...) runs a callable once, records its result durably at the moment it returns, and replays it to every later attempt of the same handler, including one recovered from a crashed worker. The journal is epoch-fenced: an attempt whose lease was reclaimed cannot record, so a zombie stops instead of duplicating a side effect. Results round-trip through serialization before the handler ever sees them, so the first execution and a replay produce identical shapes -- a difference there would only surface during retries, the worst place to find it. Async handlers await the call; sync handlers call it bare and block, bounded so a stalled loop fails the step rather than pinning the worker thread forever. A name reused in a loop is numbered per occurrence. Recorded keys appear in run history. queue= on a durable handler has been accepted since the first commit and silently ignored, which is the worst state a parameter can be in. It now routes: every step is stamped with its handler's queue ('default' when none), a worker claims only from queues it serves (rx.App(workflow_queues=...)), and per-run order holds across queues -- a run whose frontier sits on an unserved queue waits for the right worker rather than running the step somewhere it was configured not to. Wait and join slots take the queue of the handler that resumes them; children take their root's. Both are covered by conformance checks, so all three stores answer the same way, and the whole workflow suite passes against memory, SQLite, and Postgres. The teardown watchdog added to the workflow conftest (REFLEX_DEBUG_HANG=1) names any task that outlives its cancellation instead of hanging the run silently -- built while chasing an intermittent suite hang that so far only reproduces when several suites contend for one database. --- news/workflow-queues.feature.md | 1 + news/workflow-substeps.feature.md | 1 + reflex/__init__.py | 1 + reflex/app.py | 4 + reflex/workflow/__init__.py | 3 + reflex/workflow/conformance.py | 45 +++ reflex/workflow/kernel.py | 79 ++++- reflex/workflow/postgres.py | 107 +++++- reflex/workflow/records.py | 5 + reflex/workflow/runtime.py | 4 + reflex/workflow/steps.py | 326 ++++++++++++++++++ reflex/workflow/store.py | 254 +++++++++++++- reflex/workflow/testing.py | 23 ++ tests/units/workflow/conftest.py | 60 ++++ tests/units/workflow/test_observability.py | 2 +- tests/units/workflow/test_queues.py | 216 ++++++++++++ tests/units/workflow/test_steps.py | 375 +++++++++++++++++++++ 17 files changed, 1488 insertions(+), 18 deletions(-) create mode 100644 news/workflow-queues.feature.md create mode 100644 news/workflow-substeps.feature.md create mode 100644 reflex/workflow/steps.py create mode 100644 tests/units/workflow/test_queues.py create mode 100644 tests/units/workflow/test_steps.py diff --git a/news/workflow-queues.feature.md b/news/workflow-queues.feature.md new file mode 100644 index 00000000000..7d9e36f97cd --- /dev/null +++ b/news/workflow-queues.feature.md @@ -0,0 +1 @@ +`@rx.event(queue=...)` now routes steps to the worker processes that declare that queue via `rx.App(workflow_queues=...)`. diff --git a/news/workflow-substeps.feature.md b/news/workflow-substeps.feature.md new file mode 100644 index 00000000000..0302bdb9153 --- /dev/null +++ b/news/workflow-substeps.feature.md @@ -0,0 +1 @@ +`rx.step(name, fn, ...)` records substep results durably inside a durable handler, so retries and crash recoveries replay completed side effects instead of repeating them. diff --git a/reflex/__init__.py b/reflex/__init__.py index b470a9dd8e2..1b44550a6dc 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -259,6 +259,7 @@ "workflows", "approval_link", "current_run", + "step", ], } diff --git a/reflex/app.py b/reflex/app.py index 94ecd5e72d2..f2f1592f113 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -463,6 +463,9 @@ class App(MiddlewareMixin, LifespanMixin): # How many workflow attempts run at once, across different runs. workflow_concurrency: int = DEFAULT_MAX_CONCURRENCY + # Worker queues this process serves; None serves every queue. + workflow_queues: tuple[str, ...] | None = None + # The workflow runtime owning registered definitions and the kernel. _workflow_runtime: WorkflowRuntime | None = None @@ -988,6 +991,7 @@ def add_workflow(self, workflow_cls: type[BaseState]) -> None: self.workflow_store, observer=self.workflow_observer, max_concurrency=self.workflow_concurrency, + queues=self.workflow_queues, ) self.register_lifespan_task(self._run_workflow_runtime) self._workflow_runtime.register(workflow_cls) diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index f6d3ac2c8f0..f536c43d628 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -61,6 +61,7 @@ StepStatus, ) from reflex.workflow.runtime import WorkflowRuntime, get_runtime, workflows +from reflex.workflow.steps import step, substep_results from reflex.workflow.store import ( DeliveryDisposition, MemoryRunStore, @@ -126,6 +127,8 @@ "parallel", "parse_duration", "schedule", + "step", + "substep_results", "wait_for", "webhook", "workflows", diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 32a5f45901e..96bcaafeb26 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -644,6 +644,49 @@ async def check_list_children_finds_a_joins_branches(store: RunStore) -> None: assert await store.list_children("nobody", 1) == () +async def check_claims_respect_queue_boundaries(store: RunStore) -> None: + """A worker claims only from queues it serves, and skips whole runs.""" + await store.admit(make_run(), make_step(queue="video"), _ADMITTED) + assert await store.claim_next(NOW, lease_duration=LEASE, queues=("emails",)) is None + assert await store.next_due(NOW, queues=("emails",)) is None + claim = await store.claim_next(NOW, lease_duration=LEASE, queues=("video",)) + assert claim is not None + assert claim.step.queue == "video" + await store.release_claim(claim, status=StepStatus.READY, events=(), now=NOW) + # None serves everything, and the queue survives the round trip. + fallback = await store.claim_next(NOW, lease_duration=LEASE) + assert fallback is not None + assert fallback.step.queue == "video" + + +async def check_substeps_record_once_and_fence_stale_writers( + store: RunStore, +) -> None: + """The substep journal memoizes by key and refuses fenced writers.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + epoch = claim.step.epoch + + assert await store.record_substep("run1", 0, epoch, "charge", {"n": 1}, NOW) + assert await store.get_substeps("run1", 0) == {"charge": {"n": 1}} + + # First write wins: a racing duplicate reports success without replacing. + assert await store.record_substep("run1", 0, epoch, "charge", {"n": 2}, NOW) + assert await store.get_substeps("run1", 0) == {"charge": {"n": 1}} + + # A stale epoch is a zombie attempt; its write must be refused. + assert not await store.record_substep("run1", 0, epoch - 1, "late", {}, NOW) + # An unclaimed or unknown step accepts nothing either. + assert not await store.record_substep("run1", 7, epoch, "wild", {}, NOW) + assert not await store.record_substep("ghost", 0, epoch, "wild", {}, NOW) + assert await store.get_substeps("run1", 0) == {"charge": {"n": 1}} + + # Several keys come back in recording order. + assert await store.record_substep("run1", 0, epoch, "label", {"n": 3}, NOW + 1) + assert list(await store.get_substeps("run1", 0)) == ["charge", "label"] + + async def check_reads_do_not_alias_stored_state(store: RunStore) -> None: """Mutating a returned record must not change what the store holds.""" await store.admit(make_run(), make_step(), _ADMITTED) @@ -718,6 +761,8 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_early_deliveries_queue_in_order, check_children_are_created_with_their_join, check_list_children_finds_a_joins_branches, + check_claims_respect_queue_boundaries, + check_substeps_record_once_and_fence_stale_writers, check_join_arrivals_count_once, check_finalize_refuses_while_a_step_is_claimed, check_finalize_tombstones_open_slots, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 9decad078ef..06a2abbfd9b 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -54,6 +54,7 @@ StepStatus, ) from reflex.workflow.serde import to_run_data +from reflex.workflow.steps import SubstepJournal, bind_journal, unbind_journal from reflex.workflow.store import ( Claim, DeliveryDisposition, @@ -242,6 +243,7 @@ def __init__( recovery_interval: float | None = None, observer: WorkflowObserver | None = None, max_concurrency: int = DEFAULT_MAX_CONCURRENCY, + queues: Iterable[str] | None = None, ): """Initialize the kernel. @@ -262,6 +264,10 @@ def __init__( or tracing. max_concurrency: How many attempts this kernel runs at once. Each belongs to a different run, so a run's own steps stay serial. + queues: Queues this kernel's worker serves; None serves them all. + Steps land on the queue their handler declared, "default" + otherwise, so a deployment can dedicate processes to slow or + sensitive work. Raises: WorkflowRuntimeError: If the store cannot renew leases, or the @@ -324,6 +330,7 @@ def __init__( self._next_recovery_at = 0.0 self._worker_id = uuid.uuid4().hex self._observer = observer + self._queues = tuple(queues) if queues is not None else None self._wakeup = asyncio.Event() self._admission = asyncio.Lock() self._closing = False @@ -619,6 +626,7 @@ async def _admit( args=payload, due_at=due_at, origin="root", + queue=handler.queue or "default", created_at=now, updated_at=now, ) @@ -924,6 +932,7 @@ async def _invoke( instance: BaseState, args: dict[str, Any], context: RunContext, + journal: SubstepJournal, ) -> Any: """Invoke a handler attempt with its per-attempt timeout. @@ -932,6 +941,7 @@ async def _invoke( instance: The hydrated run-state instance. args: The step payload. context: Identity of this attempt, readable inside the handler. + journal: Recorded substeps of this step, for ``rx.step``. Returns: The handler return value. @@ -948,6 +958,7 @@ async def _invoke( except Exception: payload = dict(args) token = bind_run(context) + journal_token = bind_journal(journal) try: if handler.is_async: coroutine = handler.fn(instance, **payload) @@ -959,10 +970,63 @@ async def _invoke( return await asyncio.wait_for(coroutine, timeout=handler.timeout) return await coroutine finally: + unbind_journal(journal_token) unbind_run(token) + def _queue_of(self, defn: WorkflowDefinition, handler_id: str) -> str: + """Resolve the worker queue a handler's steps are served from. + + Args: + defn: The workflow definition. + handler_id: The handler naming the step. + + Returns: + The queue name. + """ + return defn.handlers[handler_id].queue or "default" + + async def _build_journal(self, claim: Claim) -> SubstepJournal: + """Load the substeps recorded for a step and bind them to this attempt. + + Args: + claim: The fenced claim being executed. + + Returns: + The journal ``rx.step`` reads and writes. + """ + recorded = await self._store.get_substeps(claim.run.run_id, claim.step.ordinal) + + def notify(key: str) -> None: + """Report one newly recorded substep to the observer. + + Args: + key: The substep's memoization key. + """ + self._notify( + claim.run, + ( + ( + HistoryEventType.SUBSTEP_RECORDED, + {"ordinal": claim.step.ordinal, "key": key}, + ), + ), + ) + + return SubstepJournal( + store=self._store, + run_id=claim.run.run_id, + ordinal=claim.step.ordinal, + epoch=claim.step.epoch, + recorded=recorded, + clock=self._clock, + notify=notify, + loop=asyncio.get_running_loop(), + sync_timeout=max(self._lease_duration * 2.0, 30.0), + ) + def _build_new_steps( self, + defn: WorkflowDefinition, run: RunRecord, successors: list[_SuccessorSpec], now: float, @@ -970,6 +1034,7 @@ def _build_new_steps( """Allocate successor slots with preallocated ordinals. Args: + defn: The workflow definition, for per-handler queues. run: The run record as of the claim. successors: The resolved successor specs. now: Current time in epoch seconds. @@ -986,6 +1051,7 @@ def _build_new_steps( args=spec.args, due_at=(now + spec.delay) if spec.delay else 0.0, origin=spec.origin, # pyright: ignore[reportArgumentType] + queue=self._queue_of(defn, spec.handler_id), created_at=now, updated_at=now, ) @@ -1271,15 +1337,17 @@ def _success_completion( children = self._child_records( claim, control.branches, claim.run.next_ordinal, now ) + then_id = self._resolve_successor(defn, control.then).handler_id join = StepRecord( run_id=claim.run.run_id, ordinal=claim.run.next_ordinal, - handler_id=self._resolve_successor(defn, control.then).handler_id, + handler_id=then_id, status=StepStatus.BLOCKED, args={"__results__": []}, wait_key=f"join:{claim.run.next_ordinal}", join_expected=1 if control.mode == "first" else len(control.branches), origin="join", + queue=self._queue_of(defn, then_id), created_at=now, updated_at=now, ) @@ -1328,6 +1396,7 @@ def _success_completion( due_at=deadline, wait_key=wait_key, origin="wait", + queue=self._queue_of(defn, resume.handler_id), created_at=now, updated_at=now, ) @@ -1382,7 +1451,7 @@ def _success_completion( tombstones=tombstones, events=tuple(events), ) - new_steps = self._build_new_steps(claim.run, successors, now) + new_steps = self._build_new_steps(defn, claim.run, successors, now) events.extend( ( HistoryEventType.STEP_SCHEDULED, @@ -1762,6 +1831,7 @@ async def _execute_claim(self, claim: Claim) -> None: attempt=claim.step.attempts + 1, epoch=claim.step.epoch, ), + await self._build_journal(claim), ) ) value = await lease.attempt @@ -1925,6 +1995,7 @@ def _child_records( status=StepStatus.READY, args=payload, origin="root", + queue=handler.queue or "default", created_at=now, updated_at=now, ), @@ -2106,7 +2177,7 @@ async def _fill_slots(self, now: float) -> list[asyncio.Task]: started: list[asyncio.Task] = [] while len(self._inflight) < self._max_concurrency: claim = await self._store.claim_next( - now, lease_duration=self._lease_duration + now, lease_duration=self._lease_duration, queues=self._queues ) if claim is None: break @@ -2205,7 +2276,7 @@ async def _worker_loop(self) -> None: ) continue now = self._clock() - due = await self._store.next_due(now) + due = await self._store.next_due(now, queues=self._queues) delay = min(self._poll_interval, max(self._next_recovery_at - now, 0.0)) for upcoming in (due, self._next_schedule_due(now)): if upcoming is not None: diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 32cbc14022f..27a02100e9c 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -97,6 +97,7 @@ join_arrived INTEGER NOT NULL DEFAULT 0, error JSONB, origin TEXT NOT NULL, + queue TEXT NOT NULL DEFAULT 'default', created_at DOUBLE PRECISION NOT NULL, updated_at DOUBLE PRECISION NOT NULL, PRIMARY KEY (run_id, ordinal) @@ -115,6 +116,14 @@ run_id TEXT NOT NULL, PRIMARY KEY (workflow_id, request_key) ); +CREATE TABLE IF NOT EXISTS workflow_substeps ( + run_id TEXT NOT NULL, + ordinal INTEGER NOT NULL, + key TEXT NOT NULL, + payload JSONB NOT NULL, + created_at DOUBLE PRECISION NOT NULL, + PRIMARY KEY (run_id, ordinal, key) +); CREATE TABLE IF NOT EXISTS workflow_inbox ( run_id TEXT NOT NULL, wait_key TEXT NOT NULL, @@ -125,6 +134,7 @@ created_at DOUBLE PRECISION NOT NULL, PRIMARY KEY (run_id, wait_key, dedupe_key) ); +ALTER TABLE workflow_steps ADD COLUMN IF NOT EXISTS queue TEXT NOT NULL DEFAULT 'default'; CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs (status); CREATE INDEX IF NOT EXISTS idx_workflow_runs_flow ON workflow_runs (workflow_id, flow_key); @@ -155,6 +165,8 @@ " WHERE x.run_id = s.run_id AND NOT (x.status = ANY(%(terminal_steps)s)))" ) +_QUEUE_PREDICATE: Final = "(%(queues)s::text[] IS NULL OR s.queue = ANY(%(queues)s))" + _RUNNABLE_PREDICATE: Final = ( "NOT (r.status = ANY(%(terminal_runs)s)) AND r.status <> 'NEEDS_ATTENTION'" " AND NOT r.cancel_requested" @@ -245,6 +257,7 @@ def _step_from_row(row: Mapping[str, Any]) -> StepRecord: join_arrived=row["join_arrived"], error=row["error"], origin=row["origin"], + queue=row["queue"], created_at=row["created_at"], updated_at=row["updated_at"], ) @@ -468,9 +481,10 @@ async def _insert_step(self, conn: Connection, step: StepRecord) -> None: await conn.execute( "INSERT INTO workflow_steps (run_id, ordinal, handler_id, status, args," " attempts, recoveries, due_at, epoch, lease_expires_at, wait_key," - " join_expected, join_arrived, error, origin, created_at, updated_at)" + " join_expected, join_arrived, error, origin, queue, created_at," + " updated_at)" " VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s," - " %s, %s)", + " %s, %s, %s)", ( step.run_id, step.ordinal, @@ -487,6 +501,7 @@ async def _insert_step(self, conn: Connection, step: StepRecord) -> None: step.join_arrived, _json(step.error), step.origin, + step.queue, step.created_at, step.updated_at, ), @@ -594,7 +609,11 @@ async def admit( return True, run.run_id async def claim_next( - self, now: float, *, lease_duration: float = DEFAULT_LEASE_DURATION + self, + now: float, + *, + lease_duration: float = DEFAULT_LEASE_DURATION, + queues: tuple[str, ...] | None = None, ) -> Claim | None: """Claim the due frontier step of some runnable run. @@ -606,6 +625,7 @@ async def claim_next( now: Current time in epoch seconds. lease_duration: Seconds of renewal silence tolerated before the claim is treated as orphaned. + queues: Queues this worker serves; None serves every queue. Returns: A fenced claim, or None when nothing is claimable right now. @@ -616,13 +636,14 @@ async def claim_next( "terminal_runs": _TERMINAL_RUNS, "terminal_steps": _TERMINAL_STEPS, "claimable": _CLAIMABLE_STEPS, + "queues": list(queues) if queues is not None else None, } async with pool.connection() as conn, conn.transaction(): cursor = await conn.execute( "SELECT s.run_id AS run_id, s.ordinal AS ordinal" " FROM workflow_steps s JOIN workflow_runs r ON r.run_id = s.run_id" f" WHERE {_RUNNABLE_PREDICATE} AND {_FRONTIER_PREDICATE}" - f" AND {_CLAIMABLE_PREDICATE}" + f" AND {_CLAIMABLE_PREDICATE} AND {_QUEUE_PREDICATE}" " ORDER BY r.created_at, s.run_id" " FOR UPDATE OF s SKIP LOCKED LIMIT 1", params, @@ -1503,6 +1524,77 @@ async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: async with pool.connection() as conn: return tuple(await self._load_steps(conn, run_id)) + async def record_substep( + self, run_id: str, ordinal: int, epoch: int, key: str, payload: Any, now: float + ) -> bool: + """Durably record one substep result inside a claimed attempt. + + Args: + run_id: The owning run. + ordinal: The mailbox slot being executed. + epoch: The claim fence of the writing attempt. + key: The substep's memoization key, unique within the step. + payload: JSON-compatible result to record. + now: Current time in epoch seconds. + + Returns: + True when recorded (or already recorded); False when the writer + was fenced and must stop. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + await self._lock_run(conn, run_id) + cursor = await conn.execute( + "SELECT status, epoch FROM workflow_steps" + " WHERE run_id = %s AND ordinal = %s", + (run_id, ordinal), + ) + row = await cursor.fetchone() + if ( + row is None + or row["status"] != StepStatus.CLAIMED.value + or row["epoch"] != epoch + ): + return False + cursor = await conn.execute( + "INSERT INTO workflow_substeps" + " (run_id, ordinal, key, payload, created_at)" + " VALUES (%s, %s, %s, %s, %s) ON CONFLICT DO NOTHING", + (run_id, ordinal, key, _json(payload), now), + ) + if cursor.rowcount: + await self._append_events( + conn, + run_id, + ( + ( + HistoryEventType.SUBSTEP_RECORDED, + {"ordinal": ordinal, "key": key}, + ), + ), + now, + ) + return True + + async def get_substeps(self, run_id: str, ordinal: int) -> dict[str, Any]: + """Load the recorded substep results of one step. + + Args: + run_id: The owning run. + ordinal: The mailbox slot. + + Returns: + Recorded payloads by key, in recording order. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT key, payload FROM workflow_substeps" + " WHERE run_id = %s AND ordinal = %s ORDER BY created_at, key", + (run_id, ordinal), + ) + return {row["key"]: row["payload"] for row in await cursor.fetchall()} + async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: """Load a run's append-only history in sequence order. @@ -1529,11 +1621,14 @@ async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: for row in await cursor.fetchall() ) - async def next_due(self, now: float) -> float | None: + async def next_due( + self, now: float, *, queues: tuple[str, ...] | None = None + ) -> float | None: """Earliest future time any runnable run becomes claimable. Args: now: Current time in epoch seconds. + queues: Queues this worker serves; None serves every queue. Returns: The epoch time, or None when no future work is scheduled. @@ -1544,12 +1639,14 @@ async def next_due(self, now: float) -> float | None: "terminal_runs": _TERMINAL_RUNS, "terminal_steps": _TERMINAL_STEPS, "claimable": _CLAIMABLE_STEPS, + "queues": list(queues) if queues is not None else None, } async with pool.connection() as conn: cursor = await conn.execute( "SELECT MIN(s.due_at) AS due FROM workflow_steps s" " JOIN workflow_runs r ON r.run_id = s.run_id" f" WHERE {_RUNNABLE_PREDICATE} AND {_FRONTIER_PREDICATE}" + f" AND {_QUEUE_PREDICATE}" " AND (s.status = ANY(%(claimable)s)" " OR (s.status = 'BLOCKED' AND s.due_at > 0))", params, diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index dee47dd7d4f..fc3e699fd30 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -136,6 +136,7 @@ class HistoryEventType(str, enum.Enum): WAIT_RESOLVED = "wait_resolved" WAIT_EXPIRED = "wait_expired" SIGNAL_DELIVERED = "signal_delivered" + SUBSTEP_RECORDED = "substep_recorded" SIGNAL_BUFFERED = "signal_buffered" SIGNAL_DUPLICATE = "signal_duplicate" @@ -209,6 +210,9 @@ class StepRecord: compare-and-swap so a redelivered result cannot count twice. error: Last recorded attempt error payload. origin: How the slot was allocated. + queue: Worker queue this slot is served from. A worker claims only + steps on queues it serves, so per-run order can flow across + differently provisioned processes. created_at: Allocation time in epoch seconds. updated_at: Last transition time in epoch seconds. """ @@ -228,6 +232,7 @@ class StepRecord: join_arrived: int = 0 error: dict[str, Any] | None = None origin: Literal["root", "chain", "delay", "hook", "wait", "join"] = "chain" + queue: str = "default" created_at: float = 0.0 updated_at: float = 0.0 diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index c0db292d39f..45ab32c9c77 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -66,6 +66,7 @@ def __init__( observer: WorkflowObserver | None = None, max_recoveries: int = DEFAULT_MAX_RECOVERIES, max_concurrency: int = DEFAULT_MAX_CONCURRENCY, + queues: Iterable[str] | None = None, ): """Initialize the runtime. @@ -82,6 +83,7 @@ def __init__( observer: Receives every recorded run transition. max_recoveries: Infrastructure recovery budget per logical step. max_concurrency: How many attempts run at once. + queues: Queues this process's worker serves; None serves all. """ self._store = store self._clock = clock @@ -93,6 +95,7 @@ def __init__( self._observer = observer self._max_recoveries = max_recoveries self._max_concurrency = max_concurrency + self._queues = tuple(queues) if queues is not None else None self._definitions: dict[str, WorkflowDefinition] = {} self._classes: dict[type, str] = {} self._kernel: WorkflowKernel | None = None @@ -189,6 +192,7 @@ async def startup(self, *, start_worker: bool = True) -> None: observer=self._observer, max_recoveries=self._max_recoveries, max_concurrency=self._max_concurrency, + queues=self._queues, ) if start_worker: await self._kernel.start_worker() diff --git a/reflex/workflow/steps.py b/reflex/workflow/steps.py new file mode 100644 index 00000000000..880cb9c85ba --- /dev/null +++ b/reflex/workflow/steps.py @@ -0,0 +1,326 @@ +"""Recorded substeps: memoized side effects inside one durable handler. + +A handler is the unit of retry. When it makes three external calls and fails +after the second, a bare retry repeats all three -- and the first two already +happened. ``rx.step`` fixes the granularity: it runs a callable once, records +the result durably the moment it returns, and on any later attempt of the same +handler returns the recorded result instead of running the callable again. A +crash between substeps costs the work since the last recorded one, not the +whole handler. + +The recorded value is what ``rx.step`` returns on every attempt -- including +the first. Results round-trip through JSON serialization before they are +handed back, so the type a handler sees is identical whether the substep ran +or was replayed from the journal; a difference there would be a bug that only +appears during retries, which is the worst possible place. +""" + +from __future__ import annotations + +import asyncio +import inspect +from contextvars import ContextVar, Token +from typing import TYPE_CHECKING, Any + +from reflex_base.utils.exceptions import WorkflowRuntimeError + +from reflex.workflow.serde import to_run_data + +if TYPE_CHECKING: + from collections.abc import Callable + + from reflex.workflow.store import RunStore + + +class SubstepJournal: + """The recorded substeps of one claimed attempt. + + Attributes: + recorded: Results already recorded for this step, by key. + """ + + __slots__ = ( + "_clock", + "_counts", + "_epoch", + "_loop", + "_notify", + "_ordinal", + "_run_id", + "_store", + "_sync_timeout", + "recorded", + ) + + def __init__( + self, + *, + store: RunStore, + run_id: str, + ordinal: int, + epoch: int, + recorded: dict[str, Any], + clock: Callable[[], float], + notify: Callable[[str], None], + loop: asyncio.AbstractEventLoop, + sync_timeout: float, + ): + """Initialize the journal for one attempt. + + Args: + store: The durable run store. + run_id: The run being executed. + ordinal: The mailbox slot being executed. + epoch: The claim fence of this attempt. + recorded: Results recorded by earlier attempts of this step. + clock: Epoch-seconds time source. + notify: Receives each newly recorded key, for the observer. + loop: The kernel's event loop, for calls from sync handlers. + sync_timeout: Seconds a sync handler waits for the loop to record + before giving up. + """ + self._store = store + self._run_id = run_id + self._ordinal = ordinal + self._epoch = epoch + self.recorded = recorded + self._counts: dict[str, int] = {} + self._clock = clock + self._notify = notify + self._loop = loop + self._sync_timeout = sync_timeout + + def allocate_key(self, name: str) -> str: + """Assign the memoization key for the next occurrence of a name. + + A handler that loops calls the same name repeatedly, so occurrences + after the first are numbered. Re-executing the handler replays the + same sequence of calls and therefore allocates the same keys, which is + what lets a retry line its calls up against the journal. + + Args: + name: The substep name as written in the handler. + + Returns: + The key identifying this occurrence. + """ + count = self._counts.get(name, 0) + 1 + self._counts[name] = count + return name if count == 1 else f"{name}#{count}" + + async def record(self, key: str, value: Any) -> Any: + """Serialize and durably record one substep result. + + Args: + key: The memoization key. + value: The value the substep produced. + + Returns: + The recorded form of the value. + + Raises: + WorkflowRuntimeError: If the value cannot be serialized, or this + attempt has been fenced and must stop. + """ + try: + payload = to_run_data({"value": value})["value"] + except Exception as err: + msg = ( + f"The result of step {key!r} could not be serialized: {err} " + "A substep result must be JSON-compatible data, because it is " + "recorded durably and replayed to later attempts." + ) + raise WorkflowRuntimeError(msg) from err + accepted = await self._store.record_substep( + self._run_id, self._ordinal, self._epoch, key, payload, self._clock() + ) + if not accepted: + msg = ( + f"Step {key!r} could not be recorded because this attempt no " + "longer owns the run: its lease was reclaimed by another " + "worker. Stopping here prevents a duplicated side effect." + ) + raise WorkflowRuntimeError(msg) + self.recorded[key] = payload + self._notify(key) + return payload + + def record_from_thread(self, key: str, value: Any) -> Any: + """Record a result from a sync handler's worker thread. + + Args: + key: The memoization key. + value: The value the substep produced. + + Returns: + The recorded form of the value. + """ + future = asyncio.run_coroutine_threadsafe(self.record(key, value), self._loop) + try: + # Bounded: a loop that stopped consuming -- a shutdown, a wedge -- + # must fail this thread's step, not pin the thread forever. An + # unkillable thread makes the whole process refuse to exit. + return future.result(timeout=self._sync_timeout) + except TimeoutError: + future.cancel() + msg = ( + f"Recording step {key!r} did not complete within " + f"{self._sync_timeout:.0f}s. The worker is shutting down or " + "stalled; the attempt stops so another worker can take over." + ) + raise WorkflowRuntimeError(msg) from None + + +_journal: ContextVar[SubstepJournal | None] = ContextVar( + "reflex_workflow_substep_journal", default=None +) + + +def bind_journal(journal: SubstepJournal) -> Token[SubstepJournal | None]: + """Install the journal for the duration of one attempt. + + Args: + journal: The attempt's journal. + + Returns: + The token that restores the previous journal. + """ + return _journal.set(journal) + + +def unbind_journal(token: Token[SubstepJournal | None]) -> None: + """Restore the journal that was in place before an attempt. + + Args: + token: The token returned by ``bind_journal``. + """ + _journal.reset(token) + + +def _require_journal() -> SubstepJournal: + """Read the current attempt's journal. + + Returns: + The journal. + + Raises: + WorkflowRuntimeError: If called outside a durable handler. + """ + journal = _journal.get() + if journal is None: + msg = ( + "rx.step() records work against the run it belongs to, and there " + "is no durable handler running. Call it from inside a " + "@rx.event(durable=True) handler." + ) + raise WorkflowRuntimeError(msg) + return journal + + +async def _run_step( + journal: SubstepJournal, + key: str, + fn: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> Any: + """Execute one substep in an async handler. + + Args: + journal: The attempt's journal. + key: The memoization key. + fn: The callable producing the result. + args: Positional arguments for the callable. + kwargs: Keyword arguments for the callable. + + Returns: + The recorded result. + """ + value = fn(*args, **kwargs) + if inspect.isawaitable(value): + value = await value + return await journal.record(key, value) + + +def step(name: str, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> Any: + """Run a callable once per handler, however many times the handler runs. + + The first attempt executes ``fn`` and records its result durably; every + later attempt of the same handler -- a retry, or a recovery after a crash + -- returns the recorded result without executing ``fn`` again. That makes + it the right wrapper for any side effect a retry must not repeat:: + + @rx.event(durable=True, effect="idempotent_write", retry=rx.Retry(max_attempts=5)) + async def fulfil(self): + charge = await rx.step("charge", charge_card, self.amount) + label = await rx.step("label", create_shipping_label, self.order_id) + return rx.complete(result={"charge": charge, "label": label}) + + If ``create_shipping_label`` fails and the handler retries, the card is + not charged again: the ``"charge"`` step replays its recorded result. + + In an ``async def`` handler the call returns an awaitable. In a sync + handler it blocks and returns the value directly -- same name, same + semantics, no ceremony. + + The result must be JSON-compatible data (models are reduced the same way + run state is), and what you get back is that recorded form on every + attempt, including the first, so replays are indistinguishable. + + Args: + name: Names the substep. A name reused in a loop is numbered by + occurrence, so each iteration is its own recorded step. + fn: The callable to run once. May be sync or async in an async + handler; sync in a sync handler. + args: Positional arguments passed to the callable. + kwargs: Keyword arguments passed to the callable. + + Returns: + An awaitable of the recorded result in an async handler; the recorded + result itself in a sync handler. + + Raises: + WorkflowRuntimeError: If called outside a durable handler, or a sync + handler passes an async callable. + """ + journal = _require_journal() + key = journal.allocate_key(name) + try: + asyncio.get_running_loop() + except RuntimeError: + if key in journal.recorded: + return journal.recorded[key] + if inspect.iscoroutinefunction(fn): + msg = ( + f"Step {name!r} passes an async callable from a sync handler. " + "Make the handler `async def`, or pass a sync callable." + ) + raise WorkflowRuntimeError(msg) from None + return journal.record_from_thread(key, fn(*args, **kwargs)) + if key in journal.recorded: + return _replay(journal.recorded[key]) + return _run_step(journal, key, fn, args, kwargs) + + +async def _replay(value: Any) -> Any: # noqa: RUF029 + """Hand back an already recorded result as an awaitable. + + An async handler awaits every ``rx.step`` call, so a replayed result must + arrive in the same shape as an executed one. + + Args: + value: The recorded payload. + + Returns: + The payload. + """ + return value + + +def substep_results() -> dict[str, Any]: + """Read the substep results recorded so far for the current step. + + Returns: + Recorded payloads by key. + """ + return dict(_require_journal().recorded) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 14636b8c7e6..3681199e984 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -134,7 +134,11 @@ async def admit( ... async def claim_next( - self, now: float, *, lease_duration: float = DEFAULT_LEASE_DURATION + self, + now: float, + *, + lease_duration: float = DEFAULT_LEASE_DURATION, + queues: tuple[str, ...] | None = None, ) -> Claim | None: """Claim the due frontier step of some runnable run. @@ -146,6 +150,9 @@ async def claim_next( now: Current time in epoch seconds. lease_duration: Seconds of renewal silence tolerated before the claim is treated as orphaned. + queues: Queues this worker serves; None serves every queue. A run + whose frontier sits on an unserved queue is skipped whole, + because claiming a later slot would break its ordering. Returns: A fenced claim, or None when nothing is claimable right now. @@ -535,6 +542,43 @@ async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: """ ... + async def record_substep( + self, run_id: str, ordinal: int, epoch: int, key: str, payload: Any, now: float + ) -> bool: + """Durably record one substep result inside a claimed attempt. + + The write is fenced on the step's claim epoch: an attempt whose lease + was reclaimed must not pollute the journal a newer attempt is reading. + The first write for a key wins, so a duplicate is a no-op that still + reports success -- memoization reads before it writes, so a duplicate + only means two racing writers agreed. + + Args: + run_id: The owning run. + ordinal: The mailbox slot being executed. + epoch: The claim fence of the writing attempt. + key: The substep's memoization key, unique within the step. + payload: JSON-compatible result to record. + now: Current time in epoch seconds. + + Returns: + True when recorded (or already recorded); False when the writer + was fenced and must stop. + """ + ... + + async def get_substeps(self, run_id: str, ordinal: int) -> dict[str, Any]: + """Load the recorded substep results of one step. + + Args: + run_id: The owning run. + ordinal: The mailbox slot. + + Returns: + Recorded payloads by key, in recording order. + """ + ... + async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: """Load a run's append-only history in sequence order. @@ -546,11 +590,14 @@ async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: """ ... - async def next_due(self, now: float) -> float | None: + async def next_due( + self, now: float, *, queues: tuple[str, ...] | None = None + ) -> float | None: """Earliest future time any runnable run becomes claimable. Args: now: Current time in epoch seconds. + queues: Queues this worker serves; None serves every queue. Returns: The epoch time, or None when no future work is scheduled. @@ -688,6 +735,7 @@ def __init__(self): self._lock = asyncio.Lock() self._runs: dict[str, RunRecord] = {} self._steps: dict[str, list[StepRecord]] = {} + self._substeps: dict[tuple[str, int], dict[str, Any]] = {} self._history: dict[str, list[HistoryEvent]] = {} self._dedupe: dict[tuple[str, str], str] = {} self._inbox: dict[str, dict[tuple[str, str, str], bool]] = {} @@ -747,7 +795,11 @@ async def admit( return True, run.run_id async def claim_next( - self, now: float, *, lease_duration: float = DEFAULT_LEASE_DURATION + self, + now: float, + *, + lease_duration: float = DEFAULT_LEASE_DURATION, + queues: tuple[str, ...] | None = None, ) -> Claim | None: """Claim the due frontier step of some runnable run. @@ -755,6 +807,9 @@ async def claim_next( now: Current time in epoch seconds. lease_duration: Seconds of renewal silence tolerated before the claim is treated as orphaned. + queues: Queues this worker serves; None serves every queue. A run + whose frontier sits on an unserved queue is skipped whole, + because claiming a later slot would break its ordering. Returns: A fenced claim, or None when nothing is claimable right now. @@ -767,6 +822,8 @@ async def claim_next( frontier = _frontier(steps) if frontier is None or not step_claimable_at(frontier, now): continue + if queues is not None and frontier.queue not in queues: + continue claimed = dataclasses.replace( frontier, status=StepStatus.CLAIMED, @@ -1498,6 +1555,65 @@ async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: async with self._lock: return tuple(_detach_step(step) for step in self._steps.get(run_id, ())) + async def record_substep( + self, run_id: str, ordinal: int, epoch: int, key: str, payload: Any, now: float + ) -> bool: + """Durably record one substep result inside a claimed attempt. + + The write is fenced on the step's claim epoch: an attempt whose lease + was reclaimed must not pollute the journal a newer attempt is reading. + The first write for a key wins, so a duplicate is a no-op that still + reports success -- memoization reads before it writes, so a duplicate + only means two racing writers agreed. + + Args: + run_id: The owning run. + ordinal: The mailbox slot being executed. + epoch: The claim fence of the writing attempt. + key: The substep's memoization key, unique within the step. + payload: JSON-compatible result to record. + now: Current time in epoch seconds. + + Returns: + True when recorded (or already recorded); False when the writer + was fenced and must stop. + """ + async with self._lock: + steps = self._steps.get(run_id) + if steps is None or ordinal >= len(steps): + return False + step = steps[ordinal] + if step.status is not StepStatus.CLAIMED or step.epoch != epoch: + return False + journal = self._substeps.setdefault((run_id, ordinal), {}) + if key not in journal: + journal[key] = copy.deepcopy(payload) + self._append_events( + run_id, + ( + ( + HistoryEventType.SUBSTEP_RECORDED, + {"ordinal": ordinal, "key": key}, + ), + ), + now, + ) + return True + + async def get_substeps(self, run_id: str, ordinal: int) -> dict[str, Any]: + """Load the recorded substep results of one step. + + Args: + run_id: The owning run. + ordinal: The mailbox slot. + + Returns: + Recorded payloads by key, in recording order. + """ + async with self._lock: + journal = self._substeps.get((run_id, ordinal), {}) + return {key: copy.deepcopy(value) for key, value in journal.items()} + async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: """Load a run's append-only history in sequence order. @@ -1510,11 +1626,14 @@ async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: async with self._lock: return tuple(self._history.get(run_id, ())) - async def next_due(self, now: float) -> float | None: + async def next_due( + self, now: float, *, queues: tuple[str, ...] | None = None + ) -> float | None: """Earliest future time any runnable run becomes claimable. Args: now: Current time in epoch seconds. + queues: Queues this worker serves; None serves every queue. Returns: The epoch time, or None when no future work is scheduled. @@ -1525,6 +1644,12 @@ async def next_due(self, now: float) -> float | None: if not _run_is_runnable(run, now): continue frontier = _frontier(self._steps[run.run_id]) + if ( + frontier is not None + and queues is not None + and frontier.queue not in queues + ): + continue wake_at = None if frontier is None else step_wake_at(frontier) if wake_at is not None: due_times.append(wake_at) @@ -1570,6 +1695,7 @@ async def next_due(self, now: float) -> float | None: join_arrived INTEGER NOT NULL DEFAULT 0, error TEXT, origin TEXT NOT NULL, + queue TEXT NOT NULL DEFAULT 'default', created_at REAL NOT NULL, updated_at REAL NOT NULL, PRIMARY KEY (run_id, ordinal) @@ -1588,6 +1714,14 @@ async def next_due(self, now: float) -> float | None: run_id TEXT NOT NULL, PRIMARY KEY (workflow_id, request_key) ); +CREATE TABLE IF NOT EXISTS workflow_substeps ( + run_id TEXT NOT NULL, + ordinal INTEGER NOT NULL, + key TEXT NOT NULL, + payload TEXT NOT NULL, + created_at REAL NOT NULL, + PRIMARY KEY (run_id, ordinal, key) +); CREATE TABLE IF NOT EXISTS workflow_inbox ( run_id TEXT NOT NULL, wait_key TEXT NOT NULL, @@ -1617,6 +1751,10 @@ async def next_due(self, now: float) -> float | None: "join_arrived", "ALTER TABLE workflow_steps ADD COLUMN join_arrived INTEGER NOT NULL DEFAULT 0", ), + ( + "queue", + "ALTER TABLE workflow_steps ADD COLUMN queue TEXT NOT NULL DEFAULT 'default'", + ), ) _RUN_MIGRATIONS: Final = ( @@ -1706,6 +1844,7 @@ def _step_from_row(row: sqlite3.Row) -> StepRecord: join_arrived=row["join_arrived"], error=_load(row["error"]), origin=row["origin"], + queue=row["queue"], created_at=row["created_at"], updated_at=row["updated_at"], ) @@ -1842,8 +1981,9 @@ def _insert_step(self, step: StepRecord) -> None: self._db.execute( "INSERT INTO workflow_steps (run_id, ordinal, handler_id, status, args," " attempts, recoveries, due_at, epoch, lease_expires_at, wait_key," - " join_expected, join_arrived, error, origin, created_at, updated_at)" - " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + " join_expected, join_arrived, error, origin, queue, created_at," + " updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( step.run_id, step.ordinal, @@ -1860,6 +2000,7 @@ def _insert_step(self, step: StepRecord) -> None: step.join_arrived, _dump(step.error), step.origin, + step.queue, step.created_at, step.updated_at, ), @@ -1923,7 +2064,11 @@ async def admit( return True, run.run_id async def claim_next( - self, now: float, *, lease_duration: float = DEFAULT_LEASE_DURATION + self, + now: float, + *, + lease_duration: float = DEFAULT_LEASE_DURATION, + queues: tuple[str, ...] | None = None, ) -> Claim | None: """Claim the due frontier step of some runnable run. @@ -1931,6 +2076,9 @@ async def claim_next( now: Current time in epoch seconds. lease_duration: Seconds of renewal silence tolerated before the claim is treated as orphaned. + queues: Queues this worker serves; None serves every queue. A run + whose frontier sits on an unserved queue is skipped whole, + because claiming a later slot would break its ordering. Returns: A fenced claim, or None when nothing is claimable right now. @@ -1953,6 +2101,8 @@ async def claim_next( frontier = _frontier(self._load_steps(run.run_id)) if frontier is None or not step_claimable_at(frontier, now): continue + if queues is not None and frontier.queue not in queues: + continue claimed = dataclasses.replace( frontier, status=StepStatus.CLAIMED, @@ -2923,6 +3073,85 @@ async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: with self._lock: return tuple(self._load_steps(run_id)) + async def record_substep( + self, run_id: str, ordinal: int, epoch: int, key: str, payload: Any, now: float + ) -> bool: + """Durably record one substep result inside a claimed attempt. + + The write is fenced on the step's claim epoch: an attempt whose lease + was reclaimed must not pollute the journal a newer attempt is reading. + The first write for a key wins, so a duplicate is a no-op that still + reports success -- memoization reads before it writes, so a duplicate + only means two racing writers agreed. + + Args: + run_id: The owning run. + ordinal: The mailbox slot being executed. + epoch: The claim fence of the writing attempt. + key: The substep's memoization key, unique within the step. + payload: JSON-compatible result to record. + now: Current time in epoch seconds. + + Returns: + True when recorded (or already recorded); False when the writer + was fenced and must stop. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT status, epoch FROM workflow_steps" + " WHERE run_id = ? AND ordinal = ?", + (run_id, ordinal), + ).fetchone() + if ( + row is None + or row["status"] != StepStatus.CLAIMED.value + or row["epoch"] != epoch + ): + self._db.execute("ROLLBACK") + return False + cursor = self._db.execute( + "INSERT OR IGNORE INTO workflow_substeps" + " (run_id, ordinal, key, payload, created_at)" + " VALUES (?, ?, ?, ?, ?)", + (run_id, ordinal, key, json.dumps(payload), now), + ) + if cursor.rowcount: + self._append_events( + run_id, + ( + ( + HistoryEventType.SUBSTEP_RECORDED, + {"ordinal": ordinal, "key": key}, + ), + ), + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True + + async def get_substeps(self, run_id: str, ordinal: int) -> dict[str, Any]: + """Load the recorded substep results of one step. + + Args: + run_id: The owning run. + ordinal: The mailbox slot. + + Returns: + Recorded payloads by key, in recording order. + """ + with self._lock: + rows = self._db.execute( + "SELECT key, payload FROM workflow_substeps" + " WHERE run_id = ? AND ordinal = ? ORDER BY created_at, key", + (run_id, ordinal), + ).fetchall() + return {row["key"]: json.loads(row["payload"]) for row in rows} + async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: """Load a run's append-only history in sequence order. @@ -2948,11 +3177,14 @@ async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: for row in rows ) - async def next_due(self, now: float) -> float | None: + async def next_due( + self, now: float, *, queues: tuple[str, ...] | None = None + ) -> float | None: """Earliest future time any runnable run becomes claimable. Args: now: Current time in epoch seconds. + queues: Queues this worker serves; None serves every queue. Returns: The epoch time, or None when no future work is scheduled. @@ -2969,6 +3201,12 @@ async def next_due(self, now: float) -> float | None: due_times = [] for row in rows: frontier = _frontier(self._load_steps(row["run_id"])) + if ( + frontier is not None + and queues is not None + and frontier.queue not in queues + ): + continue wake_at = None if frontier is None else step_wake_at(frontier) if wake_at is not None: due_times.append(wake_at) diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index cb59851a232..dcf2987c8d9 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -184,6 +184,29 @@ async def start( await self.kernel.run_until_idle() return result + async def start_only( + self, + target: Any, + *, + request_key: str | None = None, + labels: dict[str, str] | None = None, + ) -> StartResult: + """Admit a run without executing anything. + + Crash and contention tests need the admitted-but-unclaimed state so + they can play the doomed worker themselves; ``start`` would run the + root step before returning. + + Args: + target: The root event, e.g. ``MyWorkflow.begin(payload)``. + request_key: Idempotent admission key. + labels: Server-derived indexing labels. + + Returns: + The admission result. + """ + return await self.kernel.start(target, request_key=request_key, labels=labels) + async def run_until_idle(self) -> None: """Process work until nothing is claimable at the current time.""" await self.kernel.run_until_idle() diff --git a/tests/units/workflow/conftest.py b/tests/units/workflow/conftest.py index a1f5eeca874..d0ce7e1c762 100644 --- a/tests/units/workflow/conftest.py +++ b/tests/units/workflow/conftest.py @@ -77,3 +77,63 @@ def postgres_factory(): store.close() for store in postgres: store.drop_schema() + + +def _install_teardown_watchdog() -> None: + """Name the task that outlives its cancellation, instead of hanging. + + The suite intermittently hangs in ``Runner.close`` -> ``_cancel_all_tasks`` + waiting on a task that did not exit after ``cancel()``. The stock loop + teardown waits forever and says nothing. This wraps it: after a bounded + wait it prints every still-pending task with its coroutine stack, which is + the exact evidence needed to fix the leak. Enabled by REFLEX_DEBUG_HANG=1. + """ + import asyncio + import asyncio.runners as runners + import sys + import traceback + + original = runners._cancel_all_tasks # pyright: ignore[reportAttributeAccessIssue] + + def patched(loop) -> None: + to_cancel = asyncio.tasks.all_tasks(loop) + if not to_cancel: + return + for task in to_cancel: + task.cancel() + + async def bounded_gather(): + gathered = asyncio.tasks.gather(*to_cancel, return_exceptions=True) + try: + await asyncio.wait_for(asyncio.shield(gathered), timeout=20) + except (TimeoutError, asyncio.CancelledError): + print( + "\n=== TEARDOWN WATCHDOG: tasks alive 20s after cancel ===", + file=sys.stderr, + ) + for task in to_cancel: + if task.done(): + continue + print(f"--- {task!r}", file=sys.stderr) + for frame in task.get_stack(): + traceback.print_stack(frame, limit=1, file=sys.stderr) + print("=== END WATCHDOG ===", file=sys.stderr, flush=True) + await gathered + + loop.run_until_complete(bounded_gather()) + for task in to_cancel: + if task.cancelled(): + continue + if task.exception() is not None: + loop.call_exception_handler({ + "message": "unhandled exception during test shutdown", + "exception": task.exception(), + "task": task, + }) + + runners._cancel_all_tasks = patched # pyright: ignore[reportAttributeAccessIssue] + globals()["_original_cancel_all_tasks"] = original + + +if os.environ.get("REFLEX_DEBUG_HANG"): + _install_teardown_watchdog() diff --git a/tests/units/workflow/test_observability.py b/tests/units/workflow/test_observability.py index 926c3829d07..3ac0754d5e3 100644 --- a/tests/units/workflow/test_observability.py +++ b/tests/units/workflow/test_observability.py @@ -98,7 +98,7 @@ async def test_every_event_is_correlated(forked_registration_context): collector = _Collector() flow = _flow() async with WorkflowTestHarness(flow, observer=collector) as harness: - result = await harness.start(flow.go) + await harness.start(flow.go) await harness.advance("1s") assert {workflow_id for workflow_id, _, _ in collector.events} == {"obs.observed"} diff --git a/tests/units/workflow/test_queues.py b/tests/units/workflow/test_queues.py new file mode 100644 index 00000000000..04cddda315e --- /dev/null +++ b/tests/units/workflow/test_queues.py @@ -0,0 +1,216 @@ +"""Tests for routing steps to named worker queues. + +A queue is worker isolation: a deployment can dedicate processes to slow or +sensitive work (video encoding, a tenant, a rate-limited provider) without +those steps competing with everything else. The declaration has existed on +``@rx.event(queue=...)`` since the beginning; these tests are what make it +true. +""" + +import asyncio + +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.definition import compile_workflow +from reflex.workflow.kernel import WorkflowKernel, WorkflowObserver +from reflex.workflow.records import RunStatus +from reflex.workflow.store import MemoryRunStore + +RAN_ON: dict[str, str] = {} + + +def _make_flow(): + """Build a two-step workflow whose steps live on different queues. + + Returns: + The workflow class. + """ + + class Encode(rx.State): + __workflow__ = WorkflowConfig(id="queues.encode") + + @rx.event(durable=True, trigger=manual(), effect="none") + def ingest(self): + """Run on the default queue. + + Returns: + The heavy step. + """ + RAN_ON["ingest"] = CURRENT_WORKER[0] + return Encode.transcode + + @rx.event(durable=True, effect="idempotent_write", queue="video") + def transcode(self): + """Run only on the video queue. + + Returns: + Completion. + """ + RAN_ON["transcode"] = CURRENT_WORKER[0] + return rx.complete(result={"ok": True}) + + return Encode + + +CURRENT_WORKER = [""] + + +class _Tagger(WorkflowObserver): + """Observer that notes which worker executed each attempt.""" + + def __init__(self, name: str): + """Remember the worker's name. + + Args: + name: The tag for this worker. + """ + self.name = name + + def on_event(self, event_type, run_id, workflow_id, data): + """Stamp the current worker before each attempt runs. + + Args: + event_type: The recorded transition. + run_id: The run it happened on. + workflow_id: The workflow identity. + data: The event payload. + """ + from reflex.workflow.records import HistoryEventType + + if event_type is HistoryEventType.ATTEMPT_STARTED: + CURRENT_WORKER[0] = self.name + + +async def _drain(kernels, run_id, timeout=5.0): + """Run kernels until the run completes. + + Args: + kernels: The kernels sharing the store. + run_id: The run to wait for. + timeout: Seconds to wait before giving up. + """ + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + snapshot = await kernels[0].get_run(run_id) + if snapshot is not None and snapshot.status is RunStatus.COMPLETED: + return + await asyncio.sleep(0.02) + + +async def test_steps_route_to_the_worker_serving_their_queue( + forked_registration_context, +): + """Each step executes on a worker that serves its queue. + + The general worker must not take the video step even though it is idle and + the step is due: taking it would put the heavy work exactly where the + deployment said it must not run. + """ + RAN_ON.clear() + flow = _make_flow() + definition = compile_workflow(flow) + store = MemoryRunStore() + general = WorkflowKernel( + [definition], + store, + poll_interval=0.01, + queues=("default",), + observer=_Tagger("general"), + ) + video = WorkflowKernel( + [definition], + store, + poll_interval=0.01, + queues=("video",), + observer=_Tagger("video"), + ) + + result = await general.start(flow.ingest) + assert result.run_id is not None + await general.start_worker() + await video.start_worker() + try: + await _drain([general, video], result.run_id) + finally: + await general.aclose() + await video.aclose() + + assert RAN_ON == {"ingest": "general", "transcode": "video"} + + +async def test_a_run_waits_for_its_queues_worker(forked_registration_context): + """Per-run order holds across queues: nothing skips ahead. + + With only the general worker running, the run stops in front of the video + step rather than executing it somewhere it does not belong -- and the + moment a video worker appears, it continues. + """ + RAN_ON.clear() + flow = _make_flow() + definition = compile_workflow(flow) + store = MemoryRunStore() + general = WorkflowKernel( + [definition], + store, + poll_interval=0.01, + queues=("default",), + observer=_Tagger("general"), + ) + + result = await general.start(flow.ingest) + assert result.run_id is not None + await general.start_worker() + try: + for _ in range(50): + if "ingest" in RAN_ON: + break + await asyncio.sleep(0.02) + await asyncio.sleep(0.1) + snapshot = await general.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is not RunStatus.COMPLETED + assert "transcode" not in RAN_ON + + video = WorkflowKernel( + [definition], + store, + poll_interval=0.01, + queues=("video",), + observer=_Tagger("video"), + ) + await video.start_worker() + try: + await _drain([general, video], result.run_id) + finally: + await video.aclose() + finally: + await general.aclose() + + snapshot = await general.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert RAN_ON["transcode"] == "video" + + +async def test_an_unrestricted_worker_serves_every_queue( + forked_registration_context, +): + """The default deployment shape stays a single process serving it all.""" + RAN_ON.clear() + flow = _make_flow() + definition = compile_workflow(flow) + store = MemoryRunStore() + worker = WorkflowKernel( + [definition], store, poll_interval=0.01, observer=_Tagger("only") + ) + + result = await worker.start(flow.ingest) + assert result.run_id is not None + await worker.start_worker() + try: + await _drain([worker], result.run_id) + finally: + await worker.aclose() + + assert RAN_ON == {"ingest": "only", "transcode": "only"} diff --git a/tests/units/workflow/test_steps.py b/tests/units/workflow/test_steps.py new file mode 100644 index 00000000000..6858d50495d --- /dev/null +++ b/tests/units/workflow/test_steps.py @@ -0,0 +1,375 @@ +"""Tests for recorded substeps inside a durable handler. + +The contract under test: ``rx.step`` runs its callable exactly once per +logical step, no matter how many times the handler itself runs -- and what it +returns is the recorded serialized form on every attempt, so replays are +indistinguishable from first executions. +""" + +import pytest +from pydantic import BaseModel +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import HistoryEventType, RunStatus +from reflex.workflow.steps import substep_results +from reflex.workflow.testing import WorkflowTestHarness + +CALLS: list[str] = [] + + +def charge_card(amount: int) -> dict: + """Pretend to charge a card. + + Args: + amount: Cents to charge. + + Returns: + The provider's response. + """ + CALLS.append("charge") + return {"charge_id": "ch_1", "amount": amount} + + +def create_label(order: str) -> dict: + """Pretend to create a shipping label, failing until the third try. + + Args: + order: The order id. + + Returns: + The label. + + Raises: + TransientWorkflowError: While the carrier is down. + """ + CALLS.append("label") + if CALLS.count("label") < 3: + msg = "carrier down" + raise TransientWorkflowError(msg) + return {"label_id": "lb_1", "order": order} + + +class Fulfil(rx.State): + """Charge, then label; the charge must survive label retries.""" + + __workflow__ = WorkflowConfig(id="steps.fulfil") + order: str = "" + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=5, initial_delay="1s", jitter="none"), + ) + async def start(self, order: str): + """Charge the card once, then keep trying the label. + + Args: + order: The order id. + + Returns: + Completion carrying both results. + """ + self.order = order + charge = await rx.step("charge", charge_card, 2500) + label = await rx.step("label", create_label, order) + return rx.complete(result={"charge": charge, "label": label}) + + +async def test_a_recorded_substep_does_not_rerun_on_retry( + forked_registration_context, +): + """The reason this feature exists: retrying the label must not recharge. + + The handler fails twice at the label step. Without the journal each retry + would call charge_card again -- three charges for one order. With it, the + charge records on attempt one and replays on attempts two and three. + """ + CALLS.clear() + async with WorkflowTestHarness(Fulfil) as harness: + result = await harness.start(Fulfil.start("ord_1")) + assert result.run_id is not None + await harness.advance("1s") + await harness.advance("2s") + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == { + "charge": {"charge_id": "ch_1", "amount": 2500}, + "label": {"label_id": "lb_1", "order": "ord_1"}, + } + assert CALLS.count("charge") == 1 + assert CALLS.count("label") == 3 + + +async def test_substeps_survive_a_crashed_worker(forked_registration_context): + """Work recorded before a crash is not repeated after recovery. + + A worker claims the step, records the charge, and dies without committing + anything -- the exact shape of a SIGKILL between two API calls. The + recovered attempt must replay the recorded charge rather than make it + again, because the money already moved. + """ + CALLS.clear() + + def charge_once() -> dict: + """Make the charge, noting that it ran. + + Returns: + The charge. + """ + CALLS.append("charge") + return {"charge_id": "ch_crash"} + + class Crashy(rx.State): + __workflow__ = WorkflowConfig(id="steps.crashy") + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def go(self): + """Charge exactly once across worker deaths. + + Returns: + Completion. + """ + charge = await rx.step("charge", charge_once) + return rx.complete(result=charge) + + async with WorkflowTestHarness(Crashy, lease_duration="30s") as harness: + store = harness.kernel.store + result = await harness.start_only(Crashy.go) + assert result.run_id is not None + + # A doomed worker claims the step and records the charge -- exactly + # what the store sees when a real worker dies mid-handler -- then + # never commits, renews, or releases. + claim = await store.claim_next(harness.now, lease_duration=30.0) + assert claim is not None + recorded = await store.record_substep( + claim.run.run_id, + claim.step.ordinal, + claim.step.epoch, + "charge", + {"charge_id": "ch_crash"}, + harness.now, + ) + assert recorded + + # Its lease lapses; recovery reclaims the step and the surviving + # worker runs the handler, which must skip the recorded charge. + await harness.advance("31s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"charge_id": "ch_crash"} + assert CALLS == [] + + +async def test_a_looped_name_is_numbered_by_occurrence(forked_registration_context): + """Each iteration of a loop is its own recorded step.""" + sent: list[str] = [] + + def send(to: str) -> str: + """Send one message. + + Args: + to: The recipient. + + Returns: + A receipt. + """ + sent.append(to) + return f"receipt-{to}" + + class Blast(rx.State): + __workflow__ = WorkflowConfig(id="steps.blast") + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def go(self): + """Send to three recipients with one step name. + + Returns: + Completion carrying the receipts. + """ + receipts = [await rx.step("send", send, name) for name in ("a", "b", "c")] + assert set(substep_results()) == {"send", "send#2", "send#3"} + return rx.complete(result={"receipts": receipts}) + + async with WorkflowTestHarness(Blast) as harness: + result = await harness.start(Blast.go) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.result == {"receipts": ["receipt-a", "receipt-b", "receipt-c"]} + assert sent == ["a", "b", "c"] + + +async def test_sync_handlers_use_the_same_call(forked_registration_context): + """A sync handler calls rx.step without await and gets the value.""" + CALLS.clear() + + class SyncFulfil(rx.State): + __workflow__ = WorkflowConfig(id="steps.sync") + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=3, initial_delay="1s", jitter="none"), + ) + def go(self): + """Charge then flake, from a thread. + + Returns: + Completion. + + Raises: + TransientWorkflowError: On the first attempt. + """ + charge = rx.step("charge", charge_card, 900) + if CALLS.count("charge") == 1 and len(CALLS) == 1: + CALLS.append("flake") + msg = "later step failed" + raise TransientWorkflowError(msg) + return rx.complete(result=charge) + + async with WorkflowTestHarness(SyncFulfil) as harness: + result = await harness.start(SyncFulfil.go) + assert result.run_id is not None + await harness.advance("1s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"charge_id": "ch_1", "amount": 900} + assert CALLS.count("charge") == 1 + + +async def test_model_results_come_back_as_plain_data(forked_registration_context): + """The first attempt sees the same shape a replay would. + + If the first execution returned the live model while a retry returned the + recorded dict, code would only break during retries. Both see the recorded + form. + """ + + class Quote(BaseModel): + price: int + vendor: str + + def fetch_quote() -> Quote: + """Produce a typed result. + + Returns: + The quote. + """ + return Quote(price=42, vendor="acme") + + seen: list = [] + + class Quoted(rx.State): + __workflow__ = WorkflowConfig(id="steps.quoted") + + @rx.event(durable=True, trigger=manual(), effect="read") + async def go(self): + """Fetch a typed quote through a step. + + Returns: + Completion. + """ + quote = await rx.step("quote", fetch_quote) + seen.append(quote) + return rx.complete(result=quote) + + async with WorkflowTestHarness(Quoted) as harness: + result = await harness.start(Quoted.go) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.result == {"price": 42, "vendor": "acme"} + assert seen == [{"price": 42, "vendor": "acme"}] + + +async def test_an_unserializable_result_fails_in_place(forked_registration_context): + """A result that cannot be recorded is an immediate, named failure.""" + + class Sneaky(rx.State): + __workflow__ = WorkflowConfig(id="steps.sneaky") + + @rx.event(durable=True, trigger=manual(), effect="none") + async def go(self): + """Return something no journal can hold. + + Returns: + Never returns. + """ + await rx.step("bad", lambda: object()) + return rx.complete(result=None) + + async with WorkflowTestHarness(Sneaky) as harness: + result = await harness.start(Sneaky.go) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is not RunStatus.COMPLETED + steps = await harness.kernel.store.get_steps(result.run_id) + assert "serialized" in str(steps[0].error) + + +async def test_substeps_appear_in_history(forked_registration_context): + """An operator can see each recorded substep on the timeline.""" + CALLS.clear() + async with WorkflowTestHarness(Fulfil) as harness: + result = await harness.start(Fulfil.start("ord_2")) + assert result.run_id is not None + await harness.advance("1s") + await harness.advance("2s") + history = await harness.kernel.store.get_history(result.run_id) + recorded = [ + event.data["key"] + for event in history + if event.type is HistoryEventType.SUBSTEP_RECORDED + ] + assert recorded == ["charge", "label"] + + +def test_step_outside_a_handler_is_refused(): + """Plain application code has no journal to record against.""" + with pytest.raises(WorkflowRuntimeError, match="durable"): + rx.step("orphan", lambda: 1) + + +async def test_async_callable_in_a_sync_handler_is_refused( + forked_registration_context, +): + """The mistake is named instead of deadlocking the worker thread.""" + + async def async_side_effect() -> int: # noqa: RUF029 + """An async callable a sync handler cannot await. + + Returns: + Nothing meaningful. + """ + return 1 + + class Mixed(rx.State): + __workflow__ = WorkflowConfig(id="steps.mixed") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Try to run an async callable synchronously. + + Returns: + Completion. + """ + rx.step("bad", async_side_effect) + return rx.complete(result=None) + + async with WorkflowTestHarness(Mixed) as harness: + result = await harness.start(Mixed.go) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + steps = await harness.kernel.store.get_steps(result.run_id) + assert "async" in str(steps[0].error) From 2f40a7e83f96af99f66b87fba3fc3f06f2748241 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 20:54:00 -0700 Subject: [PATCH 029/121] Run SQLite store calls on a worker thread Every SqliteRunStore call executed synchronously on the event loop -- the same loop serving the app's pages and websockets -- so one contended write froze every request in the process for up to the busy timeout, and a slow disk froze them for as long as it liked. Bounding the stall (the 250ms busy timeout) was a mitigation; this is the fix: each operation's body now runs in a worker thread via asyncio.to_thread, with the connection already opened check_same_thread=False and the existing threading.Lock serializing access exactly as before. Nothing about transaction shape or semantics changes, which the store suites and conformance checks confirm across all three backends. The regression test wraps the connection in a proxy that stalls every query a quarter second and counts event-loop heartbeats during one store call: inline execution allows none, offloaded execution keeps the loop beating. --- news/workflow-sqlite-offload.feature.md | 1 + reflex/workflow/store.py | 1632 +++++++++++--------- tests/units/workflow/test_store_offload.py | 123 ++ 3 files changed, 1064 insertions(+), 692 deletions(-) create mode 100644 news/workflow-sqlite-offload.feature.md create mode 100644 tests/units/workflow/test_store_offload.py diff --git a/news/workflow-sqlite-offload.feature.md b/news/workflow-sqlite-offload.feature.md new file mode 100644 index 00000000000..837d9f5d3af --- /dev/null +++ b/news/workflow-sqlite-offload.feature.md @@ -0,0 +1 @@ +SQLite workflow store calls now run on a worker thread, so slow or contended storage no longer stalls the event loop serving the app. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 3681199e984..0459eb91f24 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -2037,31 +2037,40 @@ async def admit( Returns: Whether the run was created, and the authoritative run id. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - if run.request_key is not None: - row = self._db.execute( - "SELECT run_id FROM workflow_dedupe" - " WHERE workflow_id = ? AND request_key = ?", - (run.workflow_id, run.request_key), - ).fetchone() - if row is not None: - self._db.execute("ROLLBACK") - return False, row["run_id"] - self._db.execute( - "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" - " VALUES (?, ?, ?)", - (run.workflow_id, run.request_key, run.run_id), - ) - self._insert_run(run) - self._insert_step(root_step) - self._append_events(run.run_id, events, run.created_at) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return True, run.run_id + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + if run.request_key is not None: + row = self._db.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = ? AND request_key = ?", + (run.workflow_id, run.request_key), + ).fetchone() + if row is not None: + self._db.execute("ROLLBACK") + return False, row["run_id"] + self._db.execute( + "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" + " VALUES (?, ?, ?)", + (run.workflow_id, run.request_key, run.run_id), + ) + self._insert_run(run) + self._insert_step(root_step) + self._append_events(run.run_id, events, run.created_at) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True, run.run_id + + return await asyncio.to_thread(work) async def claim_next( self, @@ -2083,61 +2092,70 @@ async def claim_next( Returns: A fenced claim, or None when nothing is claimable right now. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - claim = None - try: - rows = self._db.execute( - "SELECT * FROM workflow_runs WHERE status NOT IN" - f" ({','.join('?' * len(terminal))})" - " AND status != ? AND cancel_requested = 0" - " AND (deadline IS NULL OR deadline > ?)" - " ORDER BY created_at", - (*terminal, RunStatus.NEEDS_ATTENTION.value, now), - ).fetchall() - for row in rows: - run = _run_from_row(row) - frontier = _frontier(self._load_steps(run.run_id)) - if frontier is None or not step_claimable_at(frontier, now): - continue - if queues is not None and frontier.queue not in queues: - continue - claimed = dataclasses.replace( - frontier, - status=StepStatus.CLAIMED, - epoch=frontier.epoch + 1, - lease_expires_at=now + lease_duration, - updated_at=now, - ) - self._db.execute( - "UPDATE workflow_steps SET status = ?, epoch = ?," - " lease_expires_at = ?, updated_at = ?" - " WHERE run_id = ? AND ordinal = ?", - ( - claimed.status.value, - claimed.epoch, - claimed.lease_expires_at, - now, - claimed.run_id, - claimed.ordinal, - ), - ) - self._db.execute( - "UPDATE workflow_runs SET status = ?, updated_at = ?" - " WHERE run_id = ?", - (RunStatus.RUNNING.value, now, run.run_id), - ) - running = dataclasses.replace( - run, status=RunStatus.RUNNING, updated_at=now - ) - claim = Claim(run=running, step=claimed) - break - self._db.execute("COMMIT" if claim is not None else "ROLLBACK") - except BaseException: - self._db.execute("ROLLBACK") - raise - return claim + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + claim = None + try: + rows = self._db.execute( + "SELECT * FROM workflow_runs WHERE status NOT IN" + f" ({','.join('?' * len(terminal))})" + " AND status != ? AND cancel_requested = 0" + " AND (deadline IS NULL OR deadline > ?)" + " ORDER BY created_at", + (*terminal, RunStatus.NEEDS_ATTENTION.value, now), + ).fetchall() + for row in rows: + run = _run_from_row(row) + frontier = _frontier(self._load_steps(run.run_id)) + if frontier is None or not step_claimable_at(frontier, now): + continue + if queues is not None and frontier.queue not in queues: + continue + claimed = dataclasses.replace( + frontier, + status=StepStatus.CLAIMED, + epoch=frontier.epoch + 1, + lease_expires_at=now + lease_duration, + updated_at=now, + ) + self._db.execute( + "UPDATE workflow_steps SET status = ?, epoch = ?," + " lease_expires_at = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", + ( + claimed.status.value, + claimed.epoch, + claimed.lease_expires_at, + now, + claimed.run_id, + claimed.ordinal, + ), + ) + self._db.execute( + "UPDATE workflow_runs SET status = ?, updated_at = ?" + " WHERE run_id = ?", + (RunStatus.RUNNING.value, now, run.run_id), + ) + running = dataclasses.replace( + run, status=RunStatus.RUNNING, updated_at=now + ) + claim = Claim(run=running, step=claimed) + break + self._db.execute("COMMIT" if claim is not None else "ROLLBACK") + except BaseException: + self._db.execute("ROLLBACK") + raise + return claim + + return await asyncio.to_thread(work) def _check_claim(self, claim: Claim) -> None: """Validate that a claim still owns its step and state version. @@ -2183,24 +2201,33 @@ async def renew_lease( Returns: True if the claim still owns its step; False if it was fenced. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") try: - self._check_claim(claim) - except StaleClaimError: + try: + self._check_claim(claim) + except StaleClaimError: + self._db.execute("ROLLBACK") + return False + self._db.execute( + "UPDATE workflow_steps SET lease_expires_at = ?" + " WHERE run_id = ? AND ordinal = ?", + (now + lease_duration, claim.run.run_id, claim.step.ordinal), + ) + self._db.execute("COMMIT") + except BaseException: self._db.execute("ROLLBACK") - return False - self._db.execute( - "UPDATE workflow_steps SET lease_expires_at = ?" - " WHERE run_id = ? AND ordinal = ?", - (now + lease_duration, claim.run.run_id, claim.step.ordinal), - ) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return True + raise + return True + + return await asyncio.to_thread(work) def _arm_sql(self, step: StepRecord, now: float) -> StepRecord: """Resolve a newly armed wait against a buffered delivery, in-transaction. @@ -2245,68 +2272,73 @@ async def commit( completion: The outcome to apply. now: Current time in epoch seconds. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - self._check_claim(claim) - self._db.execute( - "UPDATE workflow_steps SET status = ?, attempts = attempts + ?," - " due_at = ?, lease_expires_at = 0, error = ?, updated_at = ?" - " WHERE run_id = ? AND ordinal = ?", - ( - completion.step_status.value, - 1 if completion.consume_attempt else 0, - completion.due_at if completion.due_at is not None else 0.0, - _dump(completion.step_error), - now, - claim.run.run_id, - claim.step.ordinal, - ), - ) - if completion.tombstones: - terminal = tuple(s.value for s in TERMINAL_STEP_STATUSES) + + def work() -> None: + """Run the operation on the worker thread.""" + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + self._check_claim(claim) self._db.execute( - "UPDATE workflow_steps SET status = ?, updated_at = ?" - f" WHERE run_id = ? AND ordinal IN" - f" ({','.join('?' * len(completion.tombstones))})" - f" AND status NOT IN ({','.join('?' * len(terminal))})", + "UPDATE workflow_steps SET status = ?, attempts = attempts + ?," + " due_at = ?, lease_expires_at = 0, error = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", ( - StepStatus.CANCELLED.value, + completion.step_status.value, + 1 if completion.consume_attempt else 0, + completion.due_at if completion.due_at is not None else 0.0, + _dump(completion.step_error), now, claim.run.run_id, - *completion.tombstones, - *terminal, + claim.step.ordinal, ), ) - for step in completion.new_steps: - self._insert_step(self._arm_sql(step, now)) - for child_run, child_step in completion.children: - self._insert_run(child_run) - self._insert_step(child_step) - self._db.execute( - "UPDATE workflow_runs SET status = ?," - " state = CASE WHEN ? THEN ? ELSE state END," - " state_version = state_version + ?," - " next_ordinal = COALESCE(?, next_ordinal)," - " result = COALESCE(?, result), error = ?, updated_at = ?" - " WHERE run_id = ?", - ( - completion.run_status.value, - completion.state is not None, - _dump(completion.state), - 1 if completion.state is not None else 0, - completion.next_ordinal, - _dump(completion.result), - _dump(completion.run_error), - now, - claim.run.run_id, - ), - ) - self._append_events(claim.run.run_id, completion.events, now) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise + if completion.tombstones: + terminal = tuple(s.value for s in TERMINAL_STEP_STATUSES) + self._db.execute( + "UPDATE workflow_steps SET status = ?, updated_at = ?" + f" WHERE run_id = ? AND ordinal IN" + f" ({','.join('?' * len(completion.tombstones))})" + f" AND status NOT IN ({','.join('?' * len(terminal))})", + ( + StepStatus.CANCELLED.value, + now, + claim.run.run_id, + *completion.tombstones, + *terminal, + ), + ) + for step in completion.new_steps: + self._insert_step(self._arm_sql(step, now)) + for child_run, child_step in completion.children: + self._insert_run(child_run) + self._insert_step(child_step) + self._db.execute( + "UPDATE workflow_runs SET status = ?," + " state = CASE WHEN ? THEN ? ELSE state END," + " state_version = state_version + ?," + " next_ordinal = COALESCE(?, next_ordinal)," + " result = COALESCE(?, result), error = ?, updated_at = ?" + " WHERE run_id = ?", + ( + completion.run_status.value, + completion.state is not None, + _dump(completion.state), + 1 if completion.state is not None else 0, + completion.next_ordinal, + _dump(completion.result), + _dump(completion.run_error), + now, + claim.run.run_id, + ), + ) + self._append_events(claim.run.run_id, completion.events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + await asyncio.to_thread(work) async def release_claim( self, @@ -2324,24 +2356,29 @@ async def release_claim( events: History events to append. now: Current time in epoch seconds. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: + + def work() -> None: + """Run the operation on the worker thread.""" + with self._lock: + self._db.execute("BEGIN IMMEDIATE") try: - self._check_claim(claim) - except StaleClaimError: + try: + self._check_claim(claim) + except StaleClaimError: + self._db.execute("ROLLBACK") + return + self._db.execute( + "UPDATE workflow_steps SET status = ?, lease_expires_at = 0," + " updated_at = ? WHERE run_id = ? AND ordinal = ?", + (status.value, now, claim.run.run_id, claim.step.ordinal), + ) + self._append_events(claim.run.run_id, events, now) + self._db.execute("COMMIT") + except BaseException: self._db.execute("ROLLBACK") - return - self._db.execute( - "UPDATE workflow_steps SET status = ?, lease_expires_at = 0," - " updated_at = ? WHERE run_id = ? AND ordinal = ?", - (status.value, now, claim.run.run_id, claim.step.ordinal), - ) - self._append_events(claim.run.run_id, events, now) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise + raise + + await asyncio.to_thread(work) async def append_events( self, @@ -2356,14 +2393,19 @@ async def append_events( events: The (type, data) pairs to append. now: Current time in epoch seconds. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - self._append_events(run_id, events, now) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise + + def work() -> None: + """Run the operation on the worker thread.""" + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + self._append_events(run_id, events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + await asyncio.to_thread(work) async def deliver( self, @@ -2385,85 +2427,99 @@ async def deliver( Returns: What the store did with the delivery. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - row = self._db.execute( - "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) - ).fetchone() - if row is None: - self._db.execute("ROLLBACK") - return "unknown_run" - if row["status"] in terminal: - self._db.execute("ROLLBACK") - return "run_terminal" - seen = self._db.execute( - "SELECT 1 FROM workflow_inbox" - " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", - (run_id, wait_key, dedupe_key), - ).fetchone() - if seen is not None: - self._db.execute("ROLLBACK") - return "duplicate" - frontier = _frontier(self._load_steps(run_id)) - if frontier is not None and _wait_expired(frontier, now): - self._db.execute("ROLLBACK") - return "expired" - resolves = ( - frontier is not None - and frontier.status is StepStatus.BLOCKED - and frontier.wait_key == wait_key - ) - self._db.execute( - "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," - " payload, status, created_at)" - " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" - " workflow_inbox WHERE run_id = ?), ?, ?, ?)", - ( - run_id, - wait_key, - dedupe_key, - run_id, - json.dumps(payload), - "CONSUMED" if resolves else "PENDING", - now, - ), - ) - if resolves and frontier is not None: + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + if row is None: + self._db.execute("ROLLBACK") + return "unknown_run" + if row["status"] in terminal: + self._db.execute("ROLLBACK") + return "run_terminal" + seen = self._db.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", + (run_id, wait_key, dedupe_key), + ).fetchone() + if seen is not None: + self._db.execute("ROLLBACK") + return "duplicate" + frontier = _frontier(self._load_steps(run_id)) + if frontier is not None and _wait_expired(frontier, now): + self._db.execute("ROLLBACK") + return "expired" + resolves = ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and frontier.wait_key == wait_key + ) self._db.execute( - "UPDATE workflow_steps SET status = ?, due_at = ?, args = ?," - " updated_at = ? WHERE run_id = ? AND ordinal = ?", + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at)" + " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" + " workflow_inbox WHERE run_id = ?), ?, ?, ?)", ( - StepStatus.READY.value, - now, - json.dumps({**frontier.args, "__payload__": payload}), - now, run_id, - frontier.ordinal, + wait_key, + dedupe_key, + run_id, + json.dumps(payload), + "CONSUMED" if resolves else "PENDING", + now, ), ) - self._append_events( - run_id, - ( + if resolves and frontier is not None: + self._db.execute( + "UPDATE workflow_steps SET status = ?, due_at = ?, args = ?," + " updated_at = ? WHERE run_id = ? AND ordinal = ?", ( - HistoryEventType.WAIT_RESOLVED, - {"ordinal": frontier.ordinal, "wait_key": wait_key}, + StepStatus.READY.value, + now, + json.dumps({**frontier.args, "__payload__": payload}), + now, + run_id, + frontier.ordinal, ), - ), - now, - ) - else: - self._append_events( - run_id, - ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), - now, - ) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return "resolved" if resolves else "buffered" + ) + self._append_events( + run_id, + ( + ( + HistoryEventType.WAIT_RESOLVED, + {"ordinal": frontier.ordinal, "wait_key": wait_key}, + ), + ), + now, + ) + else: + self._append_events( + run_id, + ( + ( + HistoryEventType.SIGNAL_BUFFERED, + {"wait_key": wait_key}, + ), + ), + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return "resolved" if resolves else "buffered" + + return await asyncio.to_thread(work) async def admit_children( self, @@ -2478,18 +2534,23 @@ async def admit_children( events: History events to append to the parent. now: Current time in epoch seconds. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - for run, root_step in runs: - self._insert_run(run) - self._insert_step(root_step) - if runs and events: - self._append_events(runs[0][0].parent_run_id or "", events, now) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise + + def work() -> None: + """Run the operation on the worker thread.""" + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + for run, root_step in runs: + self._insert_run(run) + self._insert_step(root_step) + if runs and events: + self._append_events(runs[0][0].parent_run_id or "", events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + await asyncio.to_thread(work) async def record_arrival( self, @@ -2511,84 +2572,98 @@ async def record_arrival( Returns: What the store did with the arrival. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - wait_key = f"join:{ordinal}" - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - run_row = self._db.execute( - "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) - ).fetchone() - if run_row is None: - self._db.execute("ROLLBACK") - return "unknown_run" - if run_row["status"] in terminal: - self._db.execute("ROLLBACK") - return "run_terminal" - seen = self._db.execute( - "SELECT 1 FROM workflow_inbox" - " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", - (run_id, wait_key, dedupe_key), - ).fetchone() - if seen is not None: - self._db.execute("ROLLBACK") - return "duplicate" - step_row = self._db.execute( - "SELECT * FROM workflow_steps WHERE run_id = ? AND ordinal = ?", - (run_id, ordinal), - ).fetchone() - if step_row is None or step_row["status"] != StepStatus.BLOCKED.value: - self._db.execute("ROLLBACK") - return "run_terminal" - step = _step_from_row(step_row) - self._db.execute( - "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," - " payload, status, created_at)" - " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" - " workflow_inbox WHERE run_id = ?), ?, ?, ?)", - ( - run_id, - wait_key, - dedupe_key, - run_id, - json.dumps(payload), - "CONSUMED", - now, - ), - ) - arrived = step.join_arrived + 1 - results = [*step.args.get("__results__", []), payload] - done = arrived >= step.join_expected - self._db.execute( - "UPDATE workflow_steps SET status = ?, join_arrived = ?," - " due_at = ?, args = ?, updated_at = ?" - " WHERE run_id = ? AND ordinal = ? AND join_arrived = ?", - ( - StepStatus.READY.value if done else StepStatus.BLOCKED.value, - arrived, - now if done else step.due_at, - json.dumps({**step.args, "__results__": results}), - now, + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + wait_key = f"join:{ordinal}" + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + run_row = self._db.execute( + "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + if run_row is None: + self._db.execute("ROLLBACK") + return "unknown_run" + if run_row["status"] in terminal: + self._db.execute("ROLLBACK") + return "run_terminal" + seen = self._db.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", + (run_id, wait_key, dedupe_key), + ).fetchone() + if seen is not None: + self._db.execute("ROLLBACK") + return "duplicate" + step_row = self._db.execute( + "SELECT * FROM workflow_steps WHERE run_id = ? AND ordinal = ?", + (run_id, ordinal), + ).fetchone() + if ( + step_row is None + or step_row["status"] != StepStatus.BLOCKED.value + ): + self._db.execute("ROLLBACK") + return "run_terminal" + step = _step_from_row(step_row) + self._db.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at)" + " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" + " workflow_inbox WHERE run_id = ?), ?, ?, ?)", + ( + run_id, + wait_key, + dedupe_key, + run_id, + json.dumps(payload), + "CONSUMED", + now, + ), + ) + arrived = step.join_arrived + 1 + results = [*step.args.get("__results__", []), payload] + done = arrived >= step.join_expected + self._db.execute( + "UPDATE workflow_steps SET status = ?, join_arrived = ?," + " due_at = ?, args = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ? AND join_arrived = ?", + ( + StepStatus.READY.value + if done + else StepStatus.BLOCKED.value, + arrived, + now if done else step.due_at, + json.dumps({**step.args, "__results__": results}), + now, + run_id, + ordinal, + step.join_arrived, + ), + ) + self._append_events( run_id, - ordinal, - step.join_arrived, - ), - ) - self._append_events( - run_id, - ( ( - HistoryEventType.CHILD_RESOLVED, - {"ordinal": ordinal, "arrived": arrived}, + ( + HistoryEventType.CHILD_RESOLVED, + {"ordinal": ordinal, "arrived": arrived}, + ), ), - ), - now, - ) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return "resolved" if done else "counted" + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return "resolved" if done else "counted" + + return await asyncio.to_thread(work) async def count_active(self, workflow_id: str, flow_key: str) -> int: """Count runs of a root still in flight under a flow-control key. @@ -2600,15 +2675,24 @@ async def count_active(self, workflow_id: str, flow_key: str) -> int: Returns: How many non-terminal runs share the key. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - with self._lock: - row = self._db.execute( - "SELECT COUNT(*) AS n FROM workflow_runs" - " WHERE workflow_id = ? AND flow_key = ?" - f" AND status NOT IN ({','.join('?' * len(terminal))})", - (workflow_id, flow_key, *terminal), - ).fetchone() - return row["n"] + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + row = self._db.execute( + "SELECT COUNT(*) AS n FROM workflow_runs" + " WHERE workflow_id = ? AND flow_key = ?" + f" AND status NOT IN ({','.join('?' * len(terminal))})", + (workflow_id, flow_key, *terminal), + ).fetchone() + return row["n"] + + return await asyncio.to_thread(work) async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: """Find the oldest run still in flight under a flow-control key. @@ -2620,16 +2704,25 @@ async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | Non Returns: The run, or None when the key has no active run. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - with self._lock: - row = self._db.execute( - "SELECT * FROM workflow_runs" - " WHERE workflow_id = ? AND flow_key = ?" - f" AND status NOT IN ({','.join('?' * len(terminal))})" - " ORDER BY created_at LIMIT 1", - (workflow_id, flow_key, *terminal), - ).fetchone() - return None if row is None else _run_from_row(row) + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + row = self._db.execute( + "SELECT * FROM workflow_runs" + " WHERE workflow_id = ? AND flow_key = ?" + f" AND status NOT IN ({','.join('?' * len(terminal))})" + " ORDER BY created_at LIMIT 1", + (workflow_id, flow_key, *terminal), + ).fetchone() + return None if row is None else _run_from_row(row) + + return await asyncio.to_thread(work) async def count_started_since( self, workflow_id: str, flow_key: str, since: float @@ -2644,13 +2737,22 @@ async def count_started_since( Returns: How many runs were admitted in the window. """ - with self._lock: - row = self._db.execute( - "SELECT COUNT(*) AS n FROM workflow_runs" - " WHERE workflow_id = ? AND flow_key = ? AND created_at > ?", - (workflow_id, flow_key, since), - ).fetchone() - return row["n"] + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + row = self._db.execute( + "SELECT COUNT(*) AS n FROM workflow_runs" + " WHERE workflow_id = ? AND flow_key = ? AND created_at > ?", + (workflow_id, flow_key, since), + ).fetchone() + return row["n"] + + return await asyncio.to_thread(work) async def nth_recent_start( self, workflow_id: str, flow_key: str, n: int @@ -2674,15 +2776,24 @@ async def nth_recent_start( Returns: The scheduled start, or None when fewer than n runs exist. """ - with self._lock: - row = self._db.execute( - "SELECT MAX(s.due_at, r.created_at) AS start FROM workflow_runs r" - " JOIN workflow_steps s ON s.run_id = r.run_id AND s.ordinal = 0" - " WHERE r.workflow_id = ? AND r.flow_key = ?" - " ORDER BY start DESC LIMIT 1 OFFSET ?", - (workflow_id, flow_key, n - 1), - ).fetchone() - return None if row is None else row["start"] + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + row = self._db.execute( + "SELECT MAX(s.due_at, r.created_at) AS start FROM workflow_runs r" + " JOIN workflow_steps s ON s.run_id = r.run_id AND s.ordinal = 0" + " WHERE r.workflow_id = ? AND r.flow_key = ?" + " ORDER BY start DESC LIMIT 1 OFFSET ?", + (workflow_id, flow_key, n - 1), + ).fetchone() + return None if row is None else row["start"] + + return await asyncio.to_thread(work) async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: """Push a not-yet-started run's root slot later, for debouncing. @@ -2695,19 +2806,28 @@ async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: Returns: True when the root had not started and was deferred. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - cursor = self._db.execute( - "UPDATE workflow_steps SET due_at = ?, updated_at = ?" - " WHERE run_id = ? AND ordinal = 0 AND status = ?", - (due_at, now, run_id, StepStatus.READY.value), - ) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return cursor.rowcount > 0 + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + cursor = self._db.execute( + "UPDATE workflow_steps SET due_at = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = 0 AND status = ?", + (due_at, now, run_id, StepStatus.READY.value), + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return cursor.rowcount > 0 + + return await asyncio.to_thread(work) async def request_cancel(self, run_id: str, now: float) -> bool: """Record cancellation intent on a run. @@ -2719,27 +2839,36 @@ async def request_cancel(self, run_id: str, now: float) -> bool: Returns: True if intent was recorded on a nonterminal run. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - cursor = self._db.execute( - "UPDATE workflow_runs SET cancel_requested = 1, status = ?," - " updated_at = ? WHERE run_id = ? AND status NOT IN" - f" ({','.join('?' * len(terminal))})", - (RunStatus.CANCELLING.value, now, run_id, *terminal), - ) - if cursor.rowcount == 0: + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + cursor = self._db.execute( + "UPDATE workflow_runs SET cancel_requested = 1, status = ?," + " updated_at = ? WHERE run_id = ? AND status NOT IN" + f" ({','.join('?' * len(terminal))})", + (RunStatus.CANCELLING.value, now, run_id, *terminal), + ) + if cursor.rowcount == 0: + self._db.execute("ROLLBACK") + return False + self._append_events( + run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now + ) + self._db.execute("COMMIT") + except BaseException: self._db.execute("ROLLBACK") - return False - self._append_events( - run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now - ) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return True + raise + return True + + return await asyncio.to_thread(work) async def control_pending(self, now: float) -> tuple[RunRecord, ...]: """List drained runs awaiting a control transition. @@ -2750,17 +2879,26 @@ async def control_pending(self, now: float) -> tuple[RunRecord, ...]: Returns: The runs awaiting finalization. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - with self._lock: - rows = self._db.execute( - "SELECT * FROM workflow_runs r WHERE status NOT IN" - f" ({','.join('?' * len(terminal))})" - " AND (cancel_requested = 1 OR (deadline IS NOT NULL AND deadline <= ?))" - " AND NOT EXISTS (SELECT 1 FROM workflow_steps s" - " WHERE s.run_id = r.run_id AND s.status = ?)", - (*terminal, now, StepStatus.CLAIMED.value), - ).fetchall() - return tuple(_run_from_row(row) for row in rows) + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + rows = self._db.execute( + "SELECT * FROM workflow_runs r WHERE status NOT IN" + f" ({','.join('?' * len(terminal))})" + " AND (cancel_requested = 1 OR (deadline IS NOT NULL AND deadline <= ?))" + " AND NOT EXISTS (SELECT 1 FROM workflow_steps s" + " WHERE s.run_id = r.run_id AND s.status = ?)", + (*terminal, now, StepStatus.CLAIMED.value), + ).fetchall() + return tuple(_run_from_row(row) for row in rows) + + return await asyncio.to_thread(work) async def finalize_run( self, @@ -2783,52 +2921,61 @@ async def finalize_run( Returns: True if the run was finalized. """ - terminal_run = tuple(s.value for s in TERMINAL_RUN_STATUSES) - terminal_step = tuple(s.value for s in TERMINAL_STEP_STATUSES) - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - row = self._db.execute( - "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) - ).fetchone() - if row is None or row["status"] in terminal_run: - self._db.execute("ROLLBACK") - return False - claimed = self._db.execute( - "SELECT 1 FROM workflow_steps WHERE run_id = ? AND status = ?", - (run_id, StepStatus.CLAIMED.value), - ).fetchone() - if claimed is not None: + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal_run = tuple(s.value for s in TERMINAL_RUN_STATUSES) + terminal_step = tuple(s.value for s in TERMINAL_STEP_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + if row is None or row["status"] in terminal_run: + self._db.execute("ROLLBACK") + return False + claimed = self._db.execute( + "SELECT 1 FROM workflow_steps WHERE run_id = ? AND status = ?", + (run_id, StepStatus.CLAIMED.value), + ).fetchone() + if claimed is not None: + self._db.execute("ROLLBACK") + return False + open_rows = self._db.execute( + "SELECT ordinal FROM workflow_steps WHERE run_id = ?" + f" AND status NOT IN ({','.join('?' * len(terminal_step))})" + " ORDER BY ordinal", + (run_id, *terminal_step), + ).fetchall() + self._db.execute( + "UPDATE workflow_steps SET status = ?, updated_at = ?" + f" WHERE run_id = ? AND status NOT IN" + f" ({','.join('?' * len(terminal_step))})", + (StepStatus.CANCELLED.value, now, run_id, *terminal_step), + ) + self._db.execute( + "UPDATE workflow_runs SET status = ?, error = ?, updated_at = ?" + " WHERE run_id = ?", + (status.value, _dump(error), now, run_id), + ) + events: list[tuple[HistoryEventType, dict[str, Any]]] = [ + (HistoryEventType.STEP_TOMBSTONED, {"ordinal": row["ordinal"]}) + for row in open_rows + ] + events.append((event, {} if error is None else dict(error))) + self._append_events(run_id, events, now) + self._db.execute("COMMIT") + except BaseException: self._db.execute("ROLLBACK") - return False - open_rows = self._db.execute( - "SELECT ordinal FROM workflow_steps WHERE run_id = ?" - f" AND status NOT IN ({','.join('?' * len(terminal_step))})" - " ORDER BY ordinal", - (run_id, *terminal_step), - ).fetchall() - self._db.execute( - "UPDATE workflow_steps SET status = ?, updated_at = ?" - f" WHERE run_id = ? AND status NOT IN" - f" ({','.join('?' * len(terminal_step))})", - (StepStatus.CANCELLED.value, now, run_id, *terminal_step), - ) - self._db.execute( - "UPDATE workflow_runs SET status = ?, error = ?, updated_at = ?" - " WHERE run_id = ?", - (status.value, _dump(error), now, run_id), - ) - events: list[tuple[HistoryEventType, dict[str, Any]]] = [ - (HistoryEventType.STEP_TOMBSTONED, {"ordinal": row["ordinal"]}) - for row in open_rows - ] - events.append((event, {} if error is None else dict(error))) - self._append_events(run_id, events, now) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return True + raise + return True + + return await asyncio.to_thread(work) async def resume_run(self, run_id: str, now: float) -> bool: """Re-open a suspended run so its frontier step runs again. @@ -2840,40 +2987,51 @@ async def resume_run(self, run_id: str, now: float) -> bool: Returns: True if a suspended run was re-opened. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - cursor = self._db.execute( - "UPDATE workflow_runs SET status = ?, error = NULL," - " updated_at = ? WHERE run_id = ? AND status = ?", - ( - RunStatus.PENDING.value, - now, - run_id, - RunStatus.NEEDS_ATTENTION.value, - ), - ) - if cursor.rowcount == 0: + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + cursor = self._db.execute( + "UPDATE workflow_runs SET status = ?, error = NULL," + " updated_at = ? WHERE run_id = ? AND status = ?", + ( + RunStatus.PENDING.value, + now, + run_id, + RunStatus.NEEDS_ATTENTION.value, + ), + ) + if cursor.rowcount == 0: + self._db.execute("ROLLBACK") + return False + self._db.execute( + "UPDATE workflow_steps SET status = ?, attempts = 0, due_at = ?," + " lease_expires_at = 0, error = NULL, updated_at = ?" + " WHERE run_id = ? AND status = ?", + ( + StepStatus.READY.value, + now, + now, + run_id, + StepStatus.NEEDS_ATTENTION.value, + ), + ) + self._append_events( + run_id, ((HistoryEventType.RUN_RESUMED, {}),), now + ) + self._db.execute("COMMIT") + except BaseException: self._db.execute("ROLLBACK") - return False - self._db.execute( - "UPDATE workflow_steps SET status = ?, attempts = 0, due_at = ?," - " lease_expires_at = 0, error = NULL, updated_at = ?" - " WHERE run_id = ? AND status = ?", - ( - StepStatus.READY.value, - now, - now, - run_id, - StepStatus.NEEDS_ATTENTION.value, - ), - ) - self._append_events(run_id, ((HistoryEventType.RUN_RESUMED, {}),), now) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return True + raise + return True + + return await asyncio.to_thread(work) async def recover_orphans( self, now: float, max_recoveries: int @@ -2887,86 +3045,95 @@ async def recover_orphans( Returns: How many steps were transitioned, and the runs failed outright. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - rows = self._db.execute( - "SELECT s.* FROM workflow_steps s" - " JOIN workflow_runs r ON r.run_id = s.run_id" - " WHERE s.status = ? AND s.lease_expires_at <= ?" - f" AND r.status NOT IN ({','.join('?' * len(terminal))})", - (StepStatus.CLAIMED.value, now, *terminal), - ).fetchall() - recovered = 0 - failed: list[str] = [] - for row in rows: - step = _step_from_row(row) - recovered += 1 - if step.recoveries + 1 > max_recoveries: - self._db.execute( - "UPDATE workflow_steps SET status = ?, recoveries = ?," - " lease_expires_at = 0, error = ?, updated_at = ?" - " WHERE run_id = ? AND ordinal = ?", - ( - StepStatus.FAILED.value, - step.recoveries + 1, - json.dumps({"reason": "recovery_budget_exhausted"}), - now, - step.run_id, - step.ordinal, - ), - ) - self._db.execute( - "UPDATE workflow_runs SET status = ?, error = ?," - " updated_at = ? WHERE run_id = ?", - ( - RunStatus.FAILED.value, - json.dumps({"reason": "recovery_budget_exhausted"}), - now, + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + rows = self._db.execute( + "SELECT s.* FROM workflow_steps s" + " JOIN workflow_runs r ON r.run_id = s.run_id" + " WHERE s.status = ? AND s.lease_expires_at <= ?" + f" AND r.status NOT IN ({','.join('?' * len(terminal))})", + (StepStatus.CLAIMED.value, now, *terminal), + ).fetchall() + recovered = 0 + failed: list[str] = [] + for row in rows: + step = _step_from_row(row) + recovered += 1 + if step.recoveries + 1 > max_recoveries: + self._db.execute( + "UPDATE workflow_steps SET status = ?, recoveries = ?," + " lease_expires_at = 0, error = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", + ( + StepStatus.FAILED.value, + step.recoveries + 1, + json.dumps({"reason": "recovery_budget_exhausted"}), + now, + step.run_id, + step.ordinal, + ), + ) + self._db.execute( + "UPDATE workflow_runs SET status = ?, error = ?," + " updated_at = ? WHERE run_id = ?", + ( + RunStatus.FAILED.value, + json.dumps({"reason": "recovery_budget_exhausted"}), + now, + step.run_id, + ), + ) + failed.append(step.run_id) + self._append_events( step.run_id, - ), - ) - failed.append(step.run_id) - self._append_events( - step.run_id, - ( ( - HistoryEventType.RUN_FAILED, - {"reason": "recovery_budget_exhausted"}, + ( + HistoryEventType.RUN_FAILED, + {"reason": "recovery_budget_exhausted"}, + ), ), - ), - now, - ) - else: - self._db.execute( - "UPDATE workflow_steps SET status = ?, recoveries = ?," - " due_at = ?, lease_expires_at = 0, updated_at = ?" - " WHERE run_id = ? AND ordinal = ?", - ( - StepStatus.RECOVERY_WAIT.value, - step.recoveries + 1, - now, now, + ) + else: + self._db.execute( + "UPDATE workflow_steps SET status = ?, recoveries = ?," + " due_at = ?, lease_expires_at = 0, updated_at = ?" + " WHERE run_id = ? AND ordinal = ?", + ( + StepStatus.RECOVERY_WAIT.value, + step.recoveries + 1, + now, + now, + step.run_id, + step.ordinal, + ), + ) + self._append_events( step.run_id, - step.ordinal, - ), - ) - self._append_events( - step.run_id, - ( ( - HistoryEventType.STEP_RECOVERED, - {"ordinal": step.ordinal}, + ( + HistoryEventType.STEP_RECOVERED, + {"ordinal": step.ordinal}, + ), ), - ), - now, - ) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return recovered, tuple(failed) + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return recovered, tuple(failed) + + return await asyncio.to_thread(work) async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: """List runs matching a query, newest first. @@ -2977,34 +3144,43 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: Returns: The matching run records. """ - clauses: list[str] = [] - params: list[Any] = [] - if query.workflow_id is not None: - clauses.append("workflow_id = ?") - params.append(query.workflow_id) - if query.statuses: - placeholders = ",".join("?" * len(query.statuses)) - clauses.append(f"status IN ({placeholders})") - params.extend(status.value for status in query.statuses) - if query.created_before is not None: - clauses.append("(created_at, run_id) < (?, ?)") - params.extend(query.created_before) - for key, value in (query.labels or {}).items(): - # The key comes from user data, so it is matched as a value rather - # than spliced into a JSON path expression. - clauses.append( - "EXISTS (SELECT 1 FROM json_each(labels)" - " WHERE json_each.key = ? AND json_each.value = ?)" - ) - params.extend((key, value)) - where = f" WHERE {' AND '.join(clauses)}" if clauses else "" - with self._lock: - rows = self._db.execute( - f"SELECT * FROM workflow_runs{where}" - " ORDER BY created_at DESC, run_id DESC LIMIT ?", - (*params, query.limit), - ).fetchall() - return tuple(_run_from_row(row) for row in rows) + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + clauses: list[str] = [] + params: list[Any] = [] + if query.workflow_id is not None: + clauses.append("workflow_id = ?") + params.append(query.workflow_id) + if query.statuses: + placeholders = ",".join("?" * len(query.statuses)) + clauses.append(f"status IN ({placeholders})") + params.extend(status.value for status in query.statuses) + if query.created_before is not None: + clauses.append("(created_at, run_id) < (?, ?)") + params.extend(query.created_before) + for key, value in (query.labels or {}).items(): + # The key comes from user data, so it is matched as a value rather + # than spliced into a JSON path expression. + clauses.append( + "EXISTS (SELECT 1 FROM json_each(labels)" + " WHERE json_each.key = ? AND json_each.value = ?)" + ) + params.extend((key, value)) + where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + with self._lock: + rows = self._db.execute( + f"SELECT * FROM workflow_runs{where}" + " ORDER BY created_at DESC, run_id DESC LIMIT ?", + (*params, query.limit), + ).fetchall() + return tuple(_run_from_row(row) for row in rows) + + return await asyncio.to_thread(work) async def list_children( self, parent_run_id: str, parent_ordinal: int @@ -3018,13 +3194,22 @@ async def list_children( Returns: The child run records, oldest first. """ - with self._lock: - rows = self._db.execute( - "SELECT * FROM workflow_runs WHERE parent_run_id = ?" - " AND parent_ordinal = ? ORDER BY created_at, run_id", - (parent_run_id, parent_ordinal), - ).fetchall() - return tuple(_run_from_row(row) for row in rows) + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + rows = self._db.execute( + "SELECT * FROM workflow_runs WHERE parent_run_id = ?" + " AND parent_ordinal = ? ORDER BY created_at, run_id", + (parent_run_id, parent_ordinal), + ).fetchall() + return tuple(_run_from_row(row) for row in rows) + + return await asyncio.to_thread(work) async def find_by_request_key( self, workflow_id: str, request_key: str @@ -3038,13 +3223,22 @@ async def find_by_request_key( Returns: The existing run id, or None when the key is unused. """ - with self._lock: - row = self._db.execute( - "SELECT run_id FROM workflow_dedupe" - " WHERE workflow_id = ? AND request_key = ?", - (workflow_id, request_key), - ).fetchone() - return None if row is None else row["run_id"] + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + row = self._db.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = ? AND request_key = ?", + (workflow_id, request_key), + ).fetchone() + return None if row is None else row["run_id"] + + return await asyncio.to_thread(work) async def get_run(self, run_id: str) -> RunRecord | None: """Load one run record. @@ -3055,11 +3249,20 @@ async def get_run(self, run_id: str) -> RunRecord | None: Returns: The record, or None if unknown. """ - with self._lock: - row = self._db.execute( - "SELECT * FROM workflow_runs WHERE run_id = ?", (run_id,) - ).fetchone() - return None if row is None else _run_from_row(row) + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + row = self._db.execute( + "SELECT * FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + return None if row is None else _run_from_row(row) + + return await asyncio.to_thread(work) async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: """Load a run's mailbox slots in ordinal order. @@ -3070,8 +3273,17 @@ async def get_steps(self, run_id: str) -> tuple[StepRecord, ...]: Returns: The step records. """ - with self._lock: - return tuple(self._load_steps(run_id)) + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + return tuple(self._load_steps(run_id)) + + return await asyncio.to_thread(work) async def record_substep( self, run_id: str, ordinal: int, epoch: int, key: str, payload: Any, now: float @@ -3096,43 +3308,52 @@ async def record_substep( True when recorded (or already recorded); False when the writer was fenced and must stop. """ - with self._lock: - self._db.execute("BEGIN IMMEDIATE") - try: - row = self._db.execute( - "SELECT status, epoch FROM workflow_steps" - " WHERE run_id = ? AND ordinal = ?", - (run_id, ordinal), - ).fetchone() - if ( - row is None - or row["status"] != StepStatus.CLAIMED.value - or row["epoch"] != epoch - ): - self._db.execute("ROLLBACK") - return False - cursor = self._db.execute( - "INSERT OR IGNORE INTO workflow_substeps" - " (run_id, ordinal, key, payload, created_at)" - " VALUES (?, ?, ?, ?, ?)", - (run_id, ordinal, key, json.dumps(payload), now), - ) - if cursor.rowcount: - self._append_events( - run_id, - ( + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT status, epoch FROM workflow_steps" + " WHERE run_id = ? AND ordinal = ?", + (run_id, ordinal), + ).fetchone() + if ( + row is None + or row["status"] != StepStatus.CLAIMED.value + or row["epoch"] != epoch + ): + self._db.execute("ROLLBACK") + return False + cursor = self._db.execute( + "INSERT OR IGNORE INTO workflow_substeps" + " (run_id, ordinal, key, payload, created_at)" + " VALUES (?, ?, ?, ?, ?)", + (run_id, ordinal, key, json.dumps(payload), now), + ) + if cursor.rowcount: + self._append_events( + run_id, ( - HistoryEventType.SUBSTEP_RECORDED, - {"ordinal": ordinal, "key": key}, + ( + HistoryEventType.SUBSTEP_RECORDED, + {"ordinal": ordinal, "key": key}, + ), ), - ), - now, - ) - self._db.execute("COMMIT") - except BaseException: - self._db.execute("ROLLBACK") - raise - return True + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True + + return await asyncio.to_thread(work) async def get_substeps(self, run_id: str, ordinal: int) -> dict[str, Any]: """Load the recorded substep results of one step. @@ -3144,13 +3365,22 @@ async def get_substeps(self, run_id: str, ordinal: int) -> dict[str, Any]: Returns: Recorded payloads by key, in recording order. """ - with self._lock: - rows = self._db.execute( - "SELECT key, payload FROM workflow_substeps" - " WHERE run_id = ? AND ordinal = ? ORDER BY created_at, key", - (run_id, ordinal), - ).fetchall() - return {row["key"]: json.loads(row["payload"]) for row in rows} + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + rows = self._db.execute( + "SELECT key, payload FROM workflow_substeps" + " WHERE run_id = ? AND ordinal = ? ORDER BY created_at, key", + (run_id, ordinal), + ).fetchall() + return {row["key"]: json.loads(row["payload"]) for row in rows} + + return await asyncio.to_thread(work) async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: """Load a run's append-only history in sequence order. @@ -3161,21 +3391,30 @@ async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: Returns: The history events. """ - with self._lock: - rows = self._db.execute( - "SELECT * FROM workflow_history WHERE run_id = ? ORDER BY seq", - (run_id,), - ).fetchall() - return tuple( - HistoryEvent( - run_id=row["run_id"], - seq=row["seq"], - type=HistoryEventType(row["type"]), - at=row["at"], - data=json.loads(row["data"]), + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + rows = self._db.execute( + "SELECT * FROM workflow_history WHERE run_id = ? ORDER BY seq", + (run_id,), + ).fetchall() + return tuple( + HistoryEvent( + run_id=row["run_id"], + seq=row["seq"], + type=HistoryEventType(row["type"]), + at=row["at"], + data=json.loads(row["data"]), + ) + for row in rows ) - for row in rows - ) + + return await asyncio.to_thread(work) async def next_due( self, now: float, *, queues: tuple[str, ...] | None = None @@ -3189,25 +3428,34 @@ async def next_due( Returns: The epoch time, or None when no future work is scheduled. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - with self._lock: - rows = self._db.execute( - "SELECT run_id FROM workflow_runs WHERE status NOT IN" - f" ({','.join('?' * len(terminal))})" - " AND status != ? AND cancel_requested = 0" - " AND (deadline IS NULL OR deadline > ?)", - (*terminal, RunStatus.NEEDS_ATTENTION.value, now), - ).fetchall() - due_times = [] - for row in rows: - frontier = _frontier(self._load_steps(row["run_id"])) - if ( - frontier is not None - and queues is not None - and frontier.queue not in queues - ): - continue - wake_at = None if frontier is None else step_wake_at(frontier) - if wake_at is not None: - due_times.append(wake_at) - return min(due_times) if due_times else None + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + with self._lock: + rows = self._db.execute( + "SELECT run_id FROM workflow_runs WHERE status NOT IN" + f" ({','.join('?' * len(terminal))})" + " AND status != ? AND cancel_requested = 0" + " AND (deadline IS NULL OR deadline > ?)", + (*terminal, RunStatus.NEEDS_ATTENTION.value, now), + ).fetchall() + due_times = [] + for row in rows: + frontier = _frontier(self._load_steps(row["run_id"])) + if ( + frontier is not None + and queues is not None + and frontier.queue not in queues + ): + continue + wake_at = None if frontier is None else step_wake_at(frontier) + if wake_at is not None: + due_times.append(wake_at) + return min(due_times) if due_times else None + + return await asyncio.to_thread(work) diff --git a/tests/units/workflow/test_store_offload.py b/tests/units/workflow/test_store_offload.py new file mode 100644 index 00000000000..492ff0dada1 --- /dev/null +++ b/tests/units/workflow/test_store_offload.py @@ -0,0 +1,123 @@ +"""Tests that SQLite store calls do not stall the event loop. + +The store's calls run on the loop that also serves the app's HTTP and +websocket traffic. A slow or contended SQLite call executed inline therefore +freezes every request in the process for its duration; moved to a worker +thread, the loop keeps serving while the disk does its work. +""" + +import asyncio +import time + +from reflex.workflow.records import ( + HistoryEventType, + RunRecord, + RunStatus, + StepRecord, + StepStatus, +) +from reflex.workflow.store import SqliteRunStore + + +def _seed(now: float) -> tuple[RunRecord, StepRecord]: + """Build one admissible run. + + Args: + now: Current time in epoch seconds. + + Returns: + The run and its root slot. + """ + return ( + RunRecord( + run_id="r1", + workflow_id="offload.flow", + definition_digest="d", + status=RunStatus.PENDING, + state={}, + state_version=0, + next_ordinal=1, + created_at=now, + updated_at=now, + ), + StepRecord( + run_id="r1", + ordinal=0, + handler_id="go", + status=StepStatus.READY, + args={}, + origin="root", + created_at=now, + updated_at=now, + ), + ) + + +async def test_a_slow_sqlite_call_does_not_freeze_the_loop(tmp_path): + """The loop keeps ticking while a store call sits on slow storage. + + The database is made artificially slow by wrapping the connection's + execute in a quarter-second stall. A heartbeat task then counts loop + iterations during one store call: executed inline the heartbeat cannot + tick at all, offloaded it keeps beating. + """ + store = SqliteRunStore(tmp_path / "slow.db") + now = time.time() + run, root = _seed(now) + await store.admit(run, root, ((HistoryEventType.RUN_ADMITTED, {}),)) + + real_db = store._db + + class SlowConnection: + """Delegate to the real connection, stalling every query.""" + + def execute(self, *args, **kwargs): + """Stall, then run the real query. + + Args: + args: Positional query arguments. + kwargs: Keyword query arguments. + + Returns: + The real cursor. + """ + time.sleep(0.25) + return real_db.execute(*args, **kwargs) + + def __getattr__(self, name): + """Delegate everything else. + + Args: + name: The attribute being read. + + Returns: + The real connection's attribute. + """ + return getattr(real_db, name) + + store._db = SlowConnection() # pyright: ignore[reportAttributeAccessIssue] + + beats = 0 + ticking = True + + async def heartbeat(): + """Count loop iterations while the store call runs.""" + nonlocal beats + while ticking: + beats += 1 + await asyncio.sleep(0.01) + + ticker = asyncio.ensure_future(heartbeat()) + try: + await asyncio.sleep(0.03) + assert await store.get_run("r1") is not None + finally: + ticking = False + await ticker + store._db = real_db + store.close() + + # A quarter-second stall should allow ~25 beats; even a loaded CI box + # manages a handful. Inline execution allows exactly the few from before + # the call started. + assert beats >= 10, f"loop only ticked {beats} times during the store call" From 37892e813f4fd2ceb18a3b9ee12a6ee7c36cec62 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 21:00:19 -0700 Subject: [PATCH 030/121] Let policy keys reach into model payloads; let the harness start any root Two frictions found by writing a workflow cold against the public API, the way a code generator would. A webhook root's whole payload is one typed argument, so the natural grouping key -- the order id, the customer id -- lives inside a model. dedupe_by= already reaches into the payload, but Singleton(key=...) and its siblings only matched parameter names, so the obvious declaration was a compile error and the two key mechanisms disagreed with each other. A key now resolves to a parameter or to a field of exactly one model parameter; two parameters carrying the same field is rejected as ambiguous at compile time rather than letting the grouping silently depend on iteration order. The trigger gate -- a webhook-only root must not be startable from the browser -- also applied to the test harness, which made webhook workflows untestable except by crafting signed HTTP requests. The harness now starts any root directly: in a test, the author is the provider and the scheduler. The production path keeps its gate, the same test pins both sides, and a handler with no trigger at all stays unstartable everywhere, because it is a mid-flow step and not a root. --- news/workflow-key-and-harness.feature.md | 1 + reflex/workflow/definition.py | 53 +++++++++-- reflex/workflow/kernel.py | 42 ++++++-- reflex/workflow/testing.py | 11 ++- tests/units/workflow/test_flow_control.py | 111 +++++++++++++++++++++- tests/units/workflow/test_ingress.py | 19 ++++ tests/units/workflow/test_kernel.py | 8 +- 7 files changed, 228 insertions(+), 17 deletions(-) create mode 100644 news/workflow-key-and-harness.feature.md diff --git a/news/workflow-key-and-harness.feature.md b/news/workflow-key-and-harness.feature.md new file mode 100644 index 00000000000..f5d376c4dba --- /dev/null +++ b/news/workflow-key-and-harness.feature.md @@ -0,0 +1 @@ +Start-policy keys (`Singleton`, `Debounce`, `RateLimit`, `Throttle`) can now name a field inside a model parameter, and the workflow test harness can start webhook- and schedule-triggered roots directly. diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 7107ac1b334..69cd09fee6b 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -399,6 +399,51 @@ def _validate_branches( ) +def _validate_flow_key( + workflow_cls: type[BaseState], defn: HandlerDefinition, key: str +) -> None: + """Check that a start policy's grouping key is resolvable and unambiguous. + + A key may name a parameter, or a field inside a parameter typed as a + model -- the common case for webhook roots, whose whole payload is one + typed argument. Two parameters exposing the same field would make the + grouping silently depend on iteration order, so that is rejected here + rather than discovered in production. + + Args: + workflow_cls: The workflow class being compiled. + defn: The root handler declaring the policy. + key: The declared grouping key. + + Raises: + WorkflowDefinitionError: If the key resolves to nothing, or to more + than one place. + """ + if key in defn.params: + return + sources = [] + for param in defn.params: + hint = defn.type_hints.get(param) + fields = getattr(hint, "model_fields", None) + if fields is not None and key in fields: + sources.append(f"{param}.{key}") + if len(sources) == 1: + return + if not sources: + raise _error( + workflow_cls, + f"handler {defn.name!r} groups runs by {key!r}, which is neither a " + f"parameter ({', '.join(defn.params) or 'none'}) nor a field of a " + "model parameter.", + ) + raise _error( + workflow_cls, + f"handler {defn.name!r} groups runs by {key!r}, which is ambiguous: " + f"{' and '.join(sources)} both carry it. Rename one field, or key on " + "a parameter directly.", + ) + + def _validate_handler_body( workflow_cls: type[BaseState], defn: HandlerDefinition, @@ -635,12 +680,8 @@ def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: _validate_branches(workflow_cls, defn, handlers) policy = defn.singleton or defn.rate_limit or defn.throttle or defn.debounce key = getattr(policy, "key", None) - if key is not None and key not in defn.params: - raise _error( - workflow_cls, - f"handler {defn.name!r} groups runs by {key!r}, which is not one " - f"of its parameters ({', '.join(defn.params) or 'none'}).", - ) + if key is not None: + _validate_flow_key(workflow_cls, defn, key) if isinstance(defn.trigger, ScheduleTrigger): CronSchedule(defn.trigger.cron) durable_names = frozenset(defn.name for defn in handlers.values()) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 06a2abbfd9b..4aac1ecd52e 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -226,6 +226,32 @@ def __init__(self, claim: Claim): self.lost = False +def _extract_key_field(payload: dict[str, Any], field: str) -> Any: + """Pull a grouping value out of a start payload. + + A key may name a handler parameter directly, or a field inside a model or + mapping parameter -- a webhook payload is one typed argument, and the + natural key lives inside it. Compilation guarantees the name resolves + unambiguously, so the first match here is the only one. + + Args: + payload: The decoded start payload, by parameter name. + field: The declared key. + + Returns: + The value to group by, or None when the field is absent. + """ + if field in payload: + return payload[field] + for value in payload.values(): + if isinstance(value, dict) and field in value: + return value[field] + fields = getattr(type(value), "model_fields", None) + if fields is not None and field in fields: + return getattr(value, field) + return None + + class WorkflowKernel: """Executes durable workflow runs against a run store.""" @@ -441,7 +467,7 @@ def _flow_key(handler: HandlerDefinition, payload: dict[str, Any]) -> str | None field = getattr(policy, "key", None) if field is None: return handler.id - return f"{handler.id}:{payload.get(field)!r}" + return f"{handler.id}:{_extract_key_field(payload, field)!r}" async def _apply_start_policy( self, @@ -516,7 +542,7 @@ async def start( *, request_key: str | None = None, labels: dict[str, str] | None = None, - trigger_kind: str = "manual", + trigger_kind: str | None = "manual", ) -> StartResult: """Admit a new run from a root event. @@ -525,10 +551,10 @@ async def start( request_key: Idempotent admission key; a repeated key returns the prior run with disposition ``"deduplicated"``. labels: Server-derived indexing labels to record on the run. - trigger_kind: The ingress path admitting this run. It must match the - root's declared trigger, so a webhook root cannot be started by - application code and a manual root cannot be started by a - provider request. + trigger_kind: Which ingress is starting this run; the root must + declare the same kind, so a webhook-only root stays + unreachable from the browser. None skips the gate -- the test + harness's privilege, where the test author is the trigger. Returns: The admission result. @@ -539,7 +565,9 @@ async def start( """ defn, handler, payload = self._resolve_target(target) declared = getattr(handler.trigger, "kind", None) - if declared != trigger_kind: + # A handler without a trigger is a mid-flow step, not a root; no + # ingress -- and no test privilege -- makes it startable. + if declared is None or (trigger_kind is not None and declared != trigger_kind): expected = ( f"trigger=rx.{trigger_kind}(...)" if trigger_kind == "manual" diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index dcf2987c8d9..2c88533b23b 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -172,6 +172,9 @@ async def start( ) -> StartResult: """Start a run and process work until idle. + Any root can be started here, whatever trigger it declares: in a test + the author is the webhook provider and the scheduler. + Args: target: The root event, e.g. ``MyWorkflow.begin(payload)``. request_key: Idempotent admission key. @@ -180,7 +183,9 @@ async def start( Returns: The admission result. """ - result = await self.kernel.start(target, request_key=request_key, labels=labels) + result = await self.kernel.start( + target, request_key=request_key, labels=labels, trigger_kind=None + ) await self.kernel.run_until_idle() return result @@ -205,7 +210,9 @@ async def start_only( Returns: The admission result. """ - return await self.kernel.start(target, request_key=request_key, labels=labels) + return await self.kernel.start( + target, request_key=request_key, labels=labels, trigger_kind=None + ) async def run_until_idle(self) -> None: """Process work until nothing is claimable at the current time.""" diff --git a/tests/units/workflow/test_flow_control.py b/tests/units/workflow/test_flow_control.py index 27ef05a33f2..638333092f9 100644 --- a/tests/units/workflow/test_flow_control.py +++ b/tests/units/workflow/test_flow_control.py @@ -1,6 +1,7 @@ """Tests for start policies: singleton, debounce, rate limit, and throttle.""" import pytest +from pydantic import BaseModel from reflex_base.utils.exceptions import WorkflowDefinitionError from reflex_base.workflow import ( Debounce, @@ -222,7 +223,7 @@ def start(self, cid: str): cid: The customer identifier. """ - with pytest.raises(WorkflowDefinitionError, match="not one of its parameters"): + with pytest.raises(WorkflowDefinitionError, match="neither a parameter"): compile_workflow(BadKey) @@ -320,3 +321,111 @@ def start(self): assert len(calls) == 4 await harness.advance("10s") assert len(calls) == 6 + + +async def test_singleton_can_key_on_a_model_field(forked_registration_context): + """A webhook root's key lives inside its one typed parameter. + + ``dedupe_by=`` already reaches into the payload; a grouping key that could + not was an inconsistency every generated workflow would trip over, because + a webhook root's whole payload is a single model argument. + """ + + class Order(BaseModel): + order_id: str + amount: int + + class PerOrder(rx.State): + __workflow__ = WorkflowConfig(id="flow.perorder") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + singleton=Singleton(key="order_id", mode="skip"), + ) + def start(self, evt: Order): + """Hold the run open so a duplicate can be judged. + + Args: + evt: The order payload. + + Returns: + A long deferral. + """ + return after("1h", PerOrder.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + """Complete.""" + + async with WorkflowTestHarness(PerOrder) as harness: + first = await harness.start(PerOrder.start(Order(order_id="A", amount=1))) + duplicate = await harness.start(PerOrder.start(Order(order_id="A", amount=2))) + other = await harness.start(PerOrder.start(Order(order_id="B", amount=3))) + + assert first.disposition == "started" + assert duplicate.disposition == "skipped" + assert other.disposition == "started" + + +def test_a_key_matching_nothing_is_a_compile_error(forked_registration_context): + """A key that is neither parameter nor model field is named at compile.""" + + class Payload(BaseModel): + order_id: str + + with pytest.raises(WorkflowDefinitionError, match="neither a parameter"): + + class Mistyped(rx.State): + __workflow__ = WorkflowConfig(id="flow.mistyped") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + singleton=Singleton(key="order_number"), + ) + def start(self, evt: Payload): + """Never compiles. + + Args: + evt: The payload. + """ + + from reflex.workflow.definition import compile_workflow + + compile_workflow(Mistyped) + + +def test_an_ambiguous_key_is_a_compile_error(forked_registration_context): + """Two model parameters carrying the field cannot silently pick one.""" + + class Left(BaseModel): + order_id: str + + class Right(BaseModel): + order_id: str + + with pytest.raises(WorkflowDefinitionError, match="ambiguous"): + + class TwoSources(rx.State): + __workflow__ = WorkflowConfig(id="flow.twosources") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + singleton=Singleton(key="order_id"), + ) + def start(self, a: Left, b: Right): + """Never compiles. + + Args: + a: One source. + b: The other. + """ + + from reflex.workflow.definition import compile_workflow + + compile_workflow(TwoSources) diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py index 84d7f7ec8a2..2b530e68414 100644 --- a/tests/units/workflow/test_ingress.py +++ b/tests/units/workflow/test_ingress.py @@ -21,6 +21,7 @@ from reflex.workflow.records import RunStatus from reflex.workflow.runtime import WorkflowRuntime from reflex.workflow.store import MemoryRunStore +from reflex.workflow.testing import WorkflowTestHarness SECRET = "whsec_test" @@ -246,3 +247,21 @@ def go(self): pass assert collect_webhook_routes((compile_workflow(ManualOnly),)) == {} + + +async def test_the_harness_starts_a_webhook_root_directly(paid_workflow): + """In a test, the author is the provider; no HTTP required. + + The trigger gate exists so a webhook-only root is unreachable from the + browser. Applying it to the harness made webhook workflows untestable + except by crafting signed requests, which is not what a unit test wants. + """ + async with WorkflowTestHarness(paid_workflow) as harness: + result = await harness.start( + paid_workflow.on_paid(Payment(id="pi_9", amount=700)) + ) + assert result.disposition == "started" + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py index 02aec1684e1..c4c3c85873b 100644 --- a/tests/units/workflow/test_kernel.py +++ b/tests/units/workflow/test_kernel.py @@ -570,8 +570,14 @@ def internal(self): pass async with WorkflowTestHarness(StartRules) as harness: + # The production path is gated: a webhook-only root is not reachable + # through the default (browser-facing) start. with pytest.raises(WorkflowRuntimeError, match="cannot be started here"): - await harness.start(StartRules.on_webhook()) + await harness.kernel.start(StartRules.on_webhook()) + # The harness itself is privileged: in a test, the author is the + # provider, so any root can be started directly. + result = await harness.start(StartRules.on_webhook()) + assert result.disposition == "started" with pytest.raises(WorkflowRuntimeError, match="cannot be started here"): await harness.start(StartRules.internal()) with pytest.raises(WorkflowRuntimeError, match="workflow"): From fed56ec4eb12384e39114eaaf29c93ab713af5c3 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 21:14:13 -0700 Subject: [PATCH 031/121] Add reflex workflows check, the generation-loop validator The engine's stated future is to sit behind a text-to-workflow generator, and a generator needs one thing the API did not expose: a fast way to learn whether the code it just wrote holds up, without running it. reflex workflows check imports the target, compiles every workflow class through exactly the rules registration applies -- so what passes here registers cleanly -- and reports per class, with --json for the loop to consume. The errors are the compiler's existing teaching messages, which name the fix; the test drives the actual loop shape end to end: generate broken code, read the error, apply the repair it names, pass. Duplicate workflow ids across classes are caught here too, since they would only have collided at registration. Also removes the never-produced "buffered" value from StartDisposition: dead surface in a public Literal reads as a case callers must handle. --- news/workflow-check-cli.feature.md | 1 + reflex/workflow/cli.py | 105 +++++++++++++++++++ reflex/workflow/records.py | 1 - tests/units/workflow/test_cli_check.py | 135 +++++++++++++++++++++++++ 4 files changed, 241 insertions(+), 1 deletion(-) create mode 100644 news/workflow-check-cli.feature.md create mode 100644 tests/units/workflow/test_cli_check.py diff --git a/news/workflow-check-cli.feature.md b/news/workflow-check-cli.feature.md new file mode 100644 index 00000000000..da06a607095 --- /dev/null +++ b/news/workflow-check-cli.feature.md @@ -0,0 +1 @@ +`reflex workflows check ` compiles every workflow in a module without running it, with `--json` output for generation loops. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index edcde320fd9..6e1b8e8c609 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -12,6 +12,7 @@ import inspect import json import os +from pathlib import Path from typing import TYPE_CHECKING, Any import click @@ -261,6 +262,110 @@ async def load(store: RunStore): click.echo(f"{event.seq:<4}{event.type.value}") +@workflows.command() +@click.argument("target") +@click.option("--json", "as_json", is_flag=True, help="Emit JSON instead of text.") +def check(target: str, as_json: bool): + """Compile every workflow in a module without running anything. + + TARGET is a path to a Python file or a dotted module name. Each workflow + class it defines is compiled through the same rules the app applies at + registration, so what passes here registers cleanly. Made for generation + loops: a tool that writes workflow code calls this, reads the errors -- + which name the fix -- and repairs its output before anyone runs it. + """ + from reflex_base.utils.exceptions import WorkflowDefinitionError + + from reflex.workflow.definition import compile_workflow + + try: + module = _load_module(target) + except Exception as err: + if as_json: + click.echo(json.dumps({"ok": False, "error": str(err), "workflows": []})) + else: + console.error(f"Could not load {target!r}: {err}") + raise click.exceptions.Exit(1) from None + + classes = [ + value + for value in vars(module).values() + if isinstance(value, type) and "__workflow__" in vars(value) + ] + reports: list[dict[str, Any]] = [] + seen_ids: dict[str, str] = {} + for workflow_cls in classes: + report: dict[str, Any] = {"class": workflow_cls.__name__} + try: + definition = compile_workflow(workflow_cls) + except WorkflowDefinitionError as err: + report["ok"] = False + report["error"] = str(err) + reports.append(report) + continue + report["workflow_id"] = definition.workflow_id + owner = seen_ids.setdefault(definition.workflow_id, workflow_cls.__name__) + if owner != workflow_cls.__name__: + report["ok"] = False + report["error"] = ( + f"workflow id {definition.workflow_id!r} is also declared by " + f"{owner}; ids must be unique." + ) + else: + report["ok"] = True + reports.append(report) + + ok = bool(reports) and all(report["ok"] for report in reports) + if as_json: + payload: dict[str, Any] = {"ok": ok, "workflows": reports} + if not reports: + payload["error"] = "no workflow classes found" + click.echo(json.dumps(payload, indent=2)) + elif not reports: + console.error( + f"No workflow classes in {target!r}. A workflow is an rx.State " + "subclass with __workflow__ = rx.WorkflowConfig(id=...)." + ) + else: + for report in reports: + if report["ok"]: + click.echo(f"ok {report['workflow_id']} ({report['class']})") + else: + click.echo(f"FAIL {report['class']}: {report['error']}") + if not ok: + raise click.exceptions.Exit(1) + + +def _load_module(target: str): + """Import the module a check target names. + + Args: + target: A ``.py`` path or a dotted module name. + + Returns: + The imported module. + + Raises: + FileNotFoundError: If a path target does not exist. + ImportError: If the module cannot be imported. + """ + import importlib + import importlib.util + + if target.endswith(".py"): + path = Path(target).resolve() + if not path.exists(): + raise FileNotFoundError(path) + spec = importlib.util.spec_from_file_location(path.stem, path) + if spec is None or spec.loader is None: + msg = f"cannot build an import spec for {path}" + raise ImportError(msg) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + return importlib.import_module(target) + + @workflows.command() @database_option @click.argument("run_id") diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index fc3e699fd30..1694e7e741d 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -258,7 +258,6 @@ class HistoryEvent: StartDisposition = Literal[ "started", - "buffered", "coalesced", "skipped", "rejected", diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py new file mode 100644 index 00000000000..b336f409fbe --- /dev/null +++ b/tests/units/workflow/test_cli_check.py @@ -0,0 +1,135 @@ +"""Tests for `reflex workflows check`, the generation-loop validator. + +A tool that writes workflow code needs a way to find out whether the code +holds up before anything runs it. The command compiles every workflow class +in a module through the same rules registration applies, and its errors are +the compiler's teaching messages, so a generator can read them and repair its +output. +""" + +import json + +from click.testing import CliRunner + +from reflex.workflow.cli import workflows + +VALID = ''' +import reflex as rx + + +class Greeter(rx.State): + __workflow__ = rx.WorkflowConfig(id="check.greeter") + name: str = "" + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def start(self, name: str): + """Greet. + + Args: + name: Who to greet. + + Returns: + Completion. + """ + self.name = name + return rx.complete(result={"hello": name}) +''' + +BROKEN = ''' +import reflex as rx + + +class Broken(rx.State): + __workflow__ = rx.WorkflowConfig(id="check.broken") + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def start(self): + """Call a sibling inline, which the compiler rejects.""" + self.finish() + + @rx.event(durable=True, effect="none") + def finish(self): + """Finish.""" +''' + +DUPLICATED = ''' +import reflex as rx + + +class First(rx.State): + __workflow__ = rx.WorkflowConfig(id="check.same") + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def go(self): + """Go.""" + + +class Second(rx.State): + __workflow__ = rx.WorkflowConfig(id="check.same") + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def go(self): + """Go.""" +''' + + +def test_a_valid_module_passes(tmp_path, forked_registration_context): + """Every compiling workflow is listed with its id, and the exit is clean.""" + module = tmp_path / "flows_ok.py" + module.write_text(VALID) + result = CliRunner().invoke(workflows, ["check", str(module)]) + assert result.exit_code == 0, result.output + assert "check.greeter" in result.output + + +def test_a_compile_error_names_the_fix(tmp_path, forked_registration_context): + """The generator reads the same teaching message the compiler raises.""" + module = tmp_path / "flows_bad.py" + module.write_text(BROKEN) + result = CliRunner().invoke(workflows, ["check", str(module)]) + assert result.exit_code == 1 + assert "runs it inline" in result.output + assert "Broken.finish" in result.output + + +def test_duplicate_ids_across_classes_fail(tmp_path, forked_registration_context): + """Two classes claiming one id would collide at registration; say so now.""" + module = tmp_path / "flows_dupe.py" + module.write_text(DUPLICATED) + result = CliRunner().invoke(workflows, ["check", str(module)]) + assert result.exit_code == 1 + assert "also declared" in result.output + + +def test_json_output_is_machine_readable(tmp_path, forked_registration_context): + """A generation loop consumes the report without scraping text.""" + module = tmp_path / "flows_mixed.py" + module.write_text(VALID + BROKEN.replace("import reflex as rx", "")) + result = CliRunner().invoke(workflows, ["check", str(module), "--json"]) + assert result.exit_code == 1 + payload = json.loads(result.output) + assert payload["ok"] is False + by_class = {entry["class"]: entry for entry in payload["workflows"]} + assert by_class["Greeter"]["ok"] is True + assert by_class["Greeter"]["workflow_id"] == "check.greeter" + assert by_class["Broken"]["ok"] is False + assert "inline" in by_class["Broken"]["error"] + + +def test_a_module_with_no_workflows_fails(tmp_path, forked_registration_context): + """Producing nothing is a failure a generator must hear about.""" + module = tmp_path / "flows_empty.py" + module.write_text("x = 1\n") + result = CliRunner().invoke(workflows, ["check", str(module)]) + assert result.exit_code == 1 + assert "No workflow classes" in result.output + + +def test_a_missing_target_fails_cleanly(tmp_path, forked_registration_context): + """A bad path is a named error, not a traceback.""" + result = CliRunner().invoke( + workflows, ["check", str(tmp_path / "nope.py"), "--json"] + ) + assert result.exit_code == 1 + payload = json.loads(result.output) + assert payload["ok"] is False From 6c7d5f8de802b775d3e7b6804eeb0ce4c6988195 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 21:16:06 -0700 Subject: [PATCH 032/121] Close the Postgres test fixture's pool on the loop that opened it The store fixture in test_postgres never closed its connection pool, leaking psycopg_pool worker tasks into the test loop's teardown. Teardown cancels every task exactly once; a pool worker that catches that cancellation to clean up a connection is never cancelled again, which is the shape of a teardown that waits forever. The fixture now closes the pool before dropping the schema, and the CLI test -- which seeds through its own asyncio.run -- closes it inside that same run, because a pool can only be closed from the loop that opened it. --- tests/units/workflow/test_postgres.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/tests/units/workflow/test_postgres.py b/tests/units/workflow/test_postgres.py index 8ee0b4283e4..a79872cbb6d 100644 --- a/tests/units/workflow/test_postgres.py +++ b/tests/units/workflow/test_postgres.py @@ -73,9 +73,14 @@ async def start(self, invoice: str): @pytest.fixture -def store(): +async def store(): """Open a Postgres store in a throwaway schema. + The pool is closed on the test's own loop before the schema drops: a pool + left open leaks worker tasks into loop teardown, where a task that + catches its cancellation to clean up a connection is never cancelled a + second time and can wait forever. + Yields: The store. """ @@ -84,6 +89,7 @@ def store(): schema = f"wf_test_{uuid.uuid4().hex}" opened = PostgresRunStore(POSTGRES_URL, schema=schema, min_size=0, max_size=6) yield opened + await opened.close() opened.drop_schema() @@ -268,7 +274,16 @@ async def seed(): ((HistoryEventType.RUN_ADMITTED, {}),), ) - asyncio.run(seed()) + async def seed_and_close(): + """Seed on one loop and close the pool on that same loop. + + The fixture's teardown runs on a different loop, and a pool can only + be closed from the loop that opened it. + """ + await seed() + await store.close() + + asyncio.run(seed_and_close()) url = f"{POSTGRES_URL}?options=-csearch_path%3D{store.schema}" listed = CliRunner().invoke(workflows, ["list", "-d", url, "--json"]) From 5c2de07be01d514db3900a4c7be7d0679a2b08ef Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 21:22:49 -0700 Subject: [PATCH 033/121] Never swallow a task's own cancellation while releasing a lease The intermittent suite hang, finally reproduced deterministically and fixed. Releasing a lease cancels its renewal task and awaits it, treating the resulting CancelledError as the renewer's echo. But the same except clause catches a cancellation aimed at the RELEASING task that lands while it waits there -- and since the renewer did end cancelled, the guard concluded echo and suppressed it. That consumed the caller's one and only cancellation: shutdown, task groups, and event-loop teardown all cancel a task exactly once, so the attempt kept running as if nothing happened, and teardown's gather waited forever on a task that had eaten its cancel. The window is the few microseconds of that await, once per attempt, which is why a ~900-test suite hung roughly one run in ten and no simple test ever caught it. The discriminator is the task's own pending-cancel state: if current_task().cancelling() is set, the cancellation is ours and re-raises, whatever happened to the renewer. A standalone probe confirmed the swallow before the fix (release task completes, cancelled() False) and its absence after; the regression test parks a release inside the renewer await, cancels it once, and asserts the task ends cancelled. The teardown watchdog that hunted this now writes through file descriptor 2: pytest captures sys.stderr during teardown and only reveals it when the phase completes, which a hang prevents -- the dump it was built to produce was invisible in exactly the situation it existed for. The Postgres test fixture also closes its pool on the loop that opened it (committed earlier) so pool workers never reach teardown at all. --- reflex/workflow/kernel.py | 11 ++++- tests/units/workflow/conftest.py | 28 +++++++++---- tests/units/workflow/test_kernel.py | 62 ++++++++++++++++++++++++++++- 3 files changed, 91 insertions(+), 10 deletions(-) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 4aac1ecd52e..034c9721b53 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -1665,7 +1665,16 @@ async def _release_lease(self, lease: _Lease) -> None: try: await renewer except asyncio.CancelledError: - if not renewer.cancelled(): + # The error may be the renewer's echo, or a cancellation aimed at + # THIS task landing while it waited here. Swallowing the latter + # consumes the caller's one cancellation -- a task shut down + # during this await would keep running, which is how a process + # gets a task that outlives teardown forever. Our own pending + # cancellation always re-raises. + current = asyncio.current_task() + if (current is not None and current.cancelling()) or ( + not renewer.cancelled() + ): raise async def _renew_leases(self) -> None: diff --git a/tests/units/workflow/conftest.py b/tests/units/workflow/conftest.py index d0ce7e1c762..4cc94ae8b5a 100644 --- a/tests/units/workflow/conftest.py +++ b/tests/units/workflow/conftest.py @@ -90,11 +90,24 @@ def _install_teardown_watchdog() -> None: """ import asyncio import asyncio.runners as runners - import sys + import io import traceback original = runners._cancel_all_tasks # pyright: ignore[reportAttributeAccessIssue] + def emit(text: str) -> None: + """Write straight to the process's real stderr. + + pytest captures sys.stderr during teardown and only reveals it when + the phase completes -- which a hang prevents, hiding exactly the + evidence this exists to surface. File descriptor 2 bypasses the + capture, the same way faulthandler's dumps do. + + Args: + text: The line to write. + """ + os.write(2, (text + "\n").encode()) + def patched(loop) -> None: to_cancel = asyncio.tasks.all_tasks(loop) if not to_cancel: @@ -107,17 +120,16 @@ async def bounded_gather(): try: await asyncio.wait_for(asyncio.shield(gathered), timeout=20) except (TimeoutError, asyncio.CancelledError): - print( - "\n=== TEARDOWN WATCHDOG: tasks alive 20s after cancel ===", - file=sys.stderr, - ) + emit("=== TEARDOWN WATCHDOG: tasks alive 20s after cancel ===") for task in to_cancel: if task.done(): continue - print(f"--- {task!r}", file=sys.stderr) + emit(f"--- {task!r}") for frame in task.get_stack(): - traceback.print_stack(frame, limit=1, file=sys.stderr) - print("=== END WATCHDOG ===", file=sys.stderr, flush=True) + buffer = io.StringIO() + traceback.print_stack(frame, limit=1, file=buffer) + emit(buffer.getvalue().rstrip()) + emit("=== END WATCHDOG ===") await gathered loop.run_until_complete(bounded_gather()) diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py index c4c3c85873b..22175b471eb 100644 --- a/tests/units/workflow/test_kernel.py +++ b/tests/units/workflow/test_kernel.py @@ -1,6 +1,7 @@ """Behavioral tests for the workflow kernel via the test harness.""" import asyncio +import contextlib import pytest from pydantic import BaseModel @@ -19,8 +20,10 @@ ) import reflex as rx +from reflex.workflow.definition import compile_workflow +from reflex.workflow.kernel import WorkflowKernel from reflex.workflow.records import HistoryEventType, RunStatus, StepStatus -from reflex.workflow.store import SqliteRunStore +from reflex.workflow.store import MemoryRunStore, SqliteRunStore from reflex.workflow.testing import WorkflowTestHarness @@ -700,3 +703,60 @@ def finish(self): assert snapshot is not None assert snapshot.status is RunStatus.COMPLETED second_store.close() + + +async def test_cancelling_a_task_inside_release_lease_sticks( + forked_registration_context, +): + """A cancellation landing during lease release must end the task. + + The release path cancels the renewer and awaits it; a CancelledError + raised there can be the renewer's echo or the releasing task's own + cancellation. Swallowing the latter leaves a task that consumed teardown's + single cancel and keeps running -- the immortal task a hanging + _cancel_all_tasks waits on forever. + """ + + class Flow(rx.State): + __workflow__ = WorkflowConfig(id="lease.cancelrace") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Nothing.""" + + kernel = WorkflowKernel([compile_workflow(Flow)], MemoryRunStore()) + result = await kernel.start(Flow.go) + assert result.run_id is not None + claim = await kernel.store.claim_next(kernel._clock()) + assert claim is not None + + lease = kernel._acquire_lease(claim) + + # A renewer that takes a moment to process its cancellation, so the + # releasing task is reliably parked inside `await renewer` when its own + # cancellation arrives. + async def slow_to_die(): + """Take a beat to process cancellation, like a real renewal would. + + Raises: + asyncio.CancelledError: Always, once cleanup finishes. + """ + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + await asyncio.sleep(0.2) + raise + + assert lease.renewer is not None + lease.renewer.cancel() + with contextlib.suppress(asyncio.CancelledError): + await lease.renewer + lease.renewer = asyncio.ensure_future(slow_to_die()) + + releasing = asyncio.ensure_future(kernel._release_lease(lease)) + await asyncio.sleep(0.05) # parked inside `await renewer` + releasing.cancel() # teardown's one and only cancel + with contextlib.suppress(asyncio.CancelledError, TimeoutError): + await asyncio.wait_for(asyncio.shield(releasing), timeout=2) + assert releasing.done(), "the release task outlived its cancellation" + assert releasing.cancelled(), "the release task swallowed its own cancellation" From 5b7c53a39fc4bc5829f92fcf241740735c019d49 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 21:36:44 -0700 Subject: [PATCH 034/121] Resolve the workflow store from the deployment's environment Reaching Postgres required code: rx.App(workflow_store=PostgresRunStore(...)). That is the wrong knob for hosting, where the platform provisions the database and the app should not change between development and production. Worse, the CLI already read REFLEX_WORKFLOW_DATABASE and the app did not, so setting it gave an operator a CLI inspecting a different store than the app was writing. resolve_store() is now the one shared meaning of a database target -- a postgres:// or postgresql:// URL opens the Postgres store, anything else is a SQLite path, nothing configured falls back to ./workflow.db -- and the runtime uses it whenever no store was passed in code. Hosting sets one environment variable; the app, its workers, and reflex workflows all follow it. Construction stays lazy, so configuration is honored even when the database comes up after the app. The duplicated default-filename constants in the CLI and runtime collapse into the store module next to the resolver. --- news/workflow-database-env.feature.md | 1 + reflex/workflow/cli.py | 20 ++--- reflex/workflow/runtime.py | 9 +- reflex/workflow/store.py | 33 ++++++- tests/units/workflow/test_store_resolution.py | 90 +++++++++++++++++++ 5 files changed, 133 insertions(+), 20 deletions(-) create mode 100644 news/workflow-database-env.feature.md create mode 100644 tests/units/workflow/test_store_resolution.py diff --git a/news/workflow-database-env.feature.md b/news/workflow-database-env.feature.md new file mode 100644 index 00000000000..c369b5a6093 --- /dev/null +++ b/news/workflow-database-env.feature.md @@ -0,0 +1 @@ +The workflow store resolves from `REFLEX_WORKFLOW_DATABASE` when none is configured in code, so a hosted deployment points the app, its workers, and the CLI at managed Postgres with one environment variable. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 6e1b8e8c609..26d48d1bccc 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -11,7 +11,6 @@ import asyncio import inspect import json -import os from pathlib import Path from typing import TYPE_CHECKING, Any @@ -25,30 +24,21 @@ from reflex.workflow.store import RunStore -DEFAULT_DB_FILENAME = "workflow.db" - def _open_store(database: str | None) -> RunStore: """Open the run store the app persists to. - A ``postgres://`` or ``postgresql://`` target opens the Postgres store; - anything else is a path to a SQLite file. - Args: - database: Connection URL or SQLite path, or None for the default. + database: Connection URL or SQLite path, or None to resolve the same + way the app does: ``REFLEX_WORKFLOW_DATABASE``, then the local + default file. Returns: The store. """ - target = database or os.environ.get("REFLEX_WORKFLOW_DATABASE") - if target is not None and target.startswith(("postgres://", "postgresql://")): - from reflex.workflow.postgres import PostgresRunStore - - return PostgresRunStore(target) - - from reflex.workflow.store import SqliteRunStore + from reflex.workflow.store import resolve_store - return SqliteRunStore(target or DEFAULT_DB_FILENAME) + return resolve_store(database) def _with_store(database: str | None, work: Callable[[RunStore], Awaitable[Any]]): diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 45ab32c9c77..72e1c40afe5 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -12,7 +12,6 @@ import time from contextlib import asynccontextmanager from contextvars import ContextVar -from pathlib import Path from typing import TYPE_CHECKING, Any from reflex_base.registry import RegistrationContext @@ -30,7 +29,7 @@ WorkflowKernel, WorkflowObserver, ) -from reflex.workflow.store import RunStore, SqliteRunStore +from reflex.workflow.store import RunStore, resolve_store if TYPE_CHECKING: from reflex.workflow.store import DeliveryDisposition @@ -41,7 +40,6 @@ from reflex.state import BaseState from reflex.workflow.records import RunRecord, RunSnapshot, RunStatus, StartResult -DEFAULT_DB_FILENAME = "workflow.db" _context_runtime: ContextVar[WorkflowRuntime | None] = ContextVar( "reflex_workflow_runtime", default=None @@ -179,7 +177,10 @@ async def startup(self, *, start_worker: bool = True) -> None: if self._kernel is not None: return if self._store is None: - self._store = SqliteRunStore(Path.cwd() / DEFAULT_DB_FILENAME) + # Nothing was configured in code, so the environment decides: + # hosting points REFLEX_WORKFLOW_DATABASE at managed Postgres and + # the same app that used a local file in development scales out. + self._store = resolve_store() self._kernel = WorkflowKernel( self._definitions.values(), self._store, diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 0459eb91f24..8bfa640f145 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -23,6 +23,7 @@ import json import sqlite3 import threading +from pathlib import Path from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from reflex_base.utils.exceptions import WorkflowRuntimeError @@ -44,7 +45,6 @@ if TYPE_CHECKING: from collections.abc import Iterable - from pathlib import Path DeliveryDisposition = Literal[ @@ -1656,6 +1656,37 @@ async def next_due( return min(due_times) if due_times else None +DATABASE_ENV: Final = "REFLEX_WORKFLOW_DATABASE" +DEFAULT_DB_FILENAME: Final = "workflow.db" + + +def resolve_store(target: str | None = None) -> RunStore: + """Open the store a deployment's configuration names. + + This is the one place the app, the CLI, and a hosting platform agree on + what a database target means: a ``postgres://`` or ``postgresql://`` URL + opens the Postgres store, anything else is a path to a SQLite file, and + with nothing configured the default is a SQLite file next to the app. + Hosting sets ``REFLEX_WORKFLOW_DATABASE`` and every surface follows it, + with no code changes in the app. + + Args: + target: Connection URL or SQLite path. None reads the environment, + then falls back to the local default. + + Returns: + The store. + """ + import os + + resolved = target or os.environ.get(DATABASE_ENV) + if resolved is not None and resolved.startswith(("postgres://", "postgresql://")): + from reflex.workflow.postgres import PostgresRunStore + + return PostgresRunStore(resolved) + return SqliteRunStore(resolved or Path.cwd() / DEFAULT_DB_FILENAME) + + BUSY_TIMEOUT_MS: Final = 250 _SCHEMA = """ diff --git a/tests/units/workflow/test_store_resolution.py b/tests/units/workflow/test_store_resolution.py new file mode 100644 index 00000000000..2d442a8b1bf --- /dev/null +++ b/tests/units/workflow/test_store_resolution.py @@ -0,0 +1,90 @@ +"""Tests for resolving the run store from deployment configuration. + +Hosting hands an app one environment variable; the app, the CLI, and every +worker process must read it the same way, or a deployment ends up with the +operator's CLI inspecting a different store than the app writes. +""" + +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import ( + DATABASE_ENV, + DEFAULT_DB_FILENAME, + SqliteRunStore, + resolve_store, +) + + +def test_a_postgres_url_resolves_to_the_postgres_store(monkeypatch): + """A postgres:// target opens the multi-worker store. + + Construction is lazy -- no server is contacted -- so a deployment's + configuration is honored even when the database comes up later. + """ + from reflex.workflow.postgres import PostgresRunStore + + monkeypatch.delenv(DATABASE_ENV, raising=False) + store = resolve_store("postgresql://user:pw@db.internal:5432/app") + assert isinstance(store, PostgresRunStore) + + +def test_a_path_resolves_to_sqlite(tmp_path, monkeypatch): + """Anything that is not a URL is a SQLite file path.""" + monkeypatch.delenv(DATABASE_ENV, raising=False) + store = resolve_store(str(tmp_path / "runs.db")) + assert isinstance(store, SqliteRunStore) + store.close() + + +def test_the_environment_decides_when_code_does_not(tmp_path, monkeypatch): + """REFLEX_WORKFLOW_DATABASE is the deployment's knob.""" + target = tmp_path / "env.db" + monkeypatch.setenv(DATABASE_ENV, str(target)) + store = resolve_store() + assert isinstance(store, SqliteRunStore) + store.close() + assert target.exists() + + +def test_the_default_is_a_local_file(tmp_path, monkeypatch): + """With nothing configured, development gets a file next to the app.""" + monkeypatch.delenv(DATABASE_ENV, raising=False) + monkeypatch.chdir(tmp_path) + store = resolve_store() + assert isinstance(store, SqliteRunStore) + store.close() + assert (tmp_path / DEFAULT_DB_FILENAME).exists() + + +async def test_the_runtime_follows_the_environment( + tmp_path, monkeypatch, forked_registration_context +): + """An app with no store configured lands on the environment's database. + + This is the zero-code deployment path: hosting sets the variable, and the + same app that used a local file in development runs against managed + Postgres in production. + """ + from reflex_base.workflow import WorkflowConfig, manual + + import reflex as rx + + class Deployed(rx.State): + __workflow__ = WorkflowConfig(id="resolve.deployed") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Nothing.""" + + target = tmp_path / "runtime.db" + monkeypatch.setenv(DATABASE_ENV, str(target)) + runtime = WorkflowRuntime() + runtime.register(Deployed) + await runtime.startup(start_worker=False) + store = runtime.kernel.store + try: + assert isinstance(store, SqliteRunStore) + assert target.exists() + finally: + await runtime.shutdown() + if isinstance(store, SqliteRunStore): + store.close() From 1f758a879fbd3d419e55a974109300a1e53bf639 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 22:37:37 -0700 Subject: [PATCH 035/121] Prove the engine inside a real server with integration tests Every workflow test so far drove the runtime through in-process harnesses, which cannot see the layer where embedded engines actually break: lifespan wiring in a real server process. Two Playwright/AppHarness tests now start an actual app and exercise the production path end to end. The first clicks a button in a real browser and waits for the run to complete: that one click proves the app lifespan started the worker, the page handler reached the runtime, the rx.step substep recorded, the two-second durable timer fired on the wall clock, and the run persisted into the SQLite file that REFLEX_WORKFLOW_DATABASE named -- the same environment resolution a hosted deployment uses, exercised in a real process for the first time. The second delivers a signed webhook over real HTTP and asserts the admission contract: 202 with the run id once the run is durably admitted, a byte-identical redelivery answering with the same run and a "deduplicated" disposition, and a forged signature refused outright. Writing it against the live endpoint corrected the test's own expectation -- the endpoint returns 202 Accepted, which is the honest status for admit-before-ack, not the 200 the draft assumed. --- .../tests_playwright/test_workflows.py | 172 ++++++++++++++++++ 1 file changed, 172 insertions(+) create mode 100644 tests/integration/tests_playwright/test_workflows.py diff --git a/tests/integration/tests_playwright/test_workflows.py b/tests/integration/tests_playwright/test_workflows.py new file mode 100644 index 00000000000..ebbdb3b66ff --- /dev/null +++ b/tests/integration/tests_playwright/test_workflows.py @@ -0,0 +1,172 @@ +"""Integration tests for durable workflows inside a real app process. + +Every other workflow test drives the runtime through harnesses. These start +an actual ``reflex run`` server and prove the pieces only a real process can: +the app lifespan starts the worker, a browser event can start a run, a +durable timer fires on the wall clock, the webhook endpoint accepts a signed +request over real HTTP, and the store resolves from the environment the way +a deployment would configure it. +""" + +from __future__ import annotations + +import hashlib +import hmac +import json +from collections.abc import Generator + +import httpx +import pytest +from playwright.sync_api import Page, expect +from reflex_base.config import get_config + +from reflex.testing import AppHarness + +WEBHOOK_SECRET = "whsec_integration" + + +def WorkflowApp(): + """App with one workflow, started from a page and from a webhook.""" + import reflex as rx + + def stamp(order_id: str) -> dict: + return {"stamped": order_id} + + class OrderFlow(rx.State): + __workflow__ = rx.WorkflowConfig(id="integration.order") + order_id: str = "" + + @rx.event(durable=True, trigger=rx.manual(), effect="idempotent_write") + async def start(self, order_id: str): + self.order_id = order_id + await rx.step("stamp", stamp, order_id) + return rx.after("2s", OrderFlow.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + return rx.complete(result={"order": self.order_id}) + + @rx.event( + durable=True, + effect="idempotent_write", + trigger=rx.webhook( + "orders.placed", + verify=rx.hmac_signature( + secret_env="ORDERS_WEBHOOK_SECRET", header="X-Signature" + ), + dedupe_by="order_id", + ), + ) + def from_hook(self, payload: dict): + self.order_id = str(payload.get("order_id", "")) + return rx.complete(result={"via": "webhook"}) + + class Dash(rx.State): + run_id: str = "" + status: str = "" + + @rx.event + async def launch(self): + result = await rx.workflows.start(OrderFlow.start("ord-77")) + self.run_id = result.run_id or "" + + @rx.event + async def refresh(self): + if not self.run_id: + return + snapshot = await rx.workflows.get_run(self.run_id) + self.status = snapshot.status.value if snapshot else "missing" + + @rx.page("/") + def index(): + return rx.box( + rx.button("launch", on_click=Dash.launch, id="launch"), + rx.button("refresh", on_click=Dash.refresh, id="refresh"), + rx.text(Dash.run_id, id="run-id"), + rx.text(Dash.status, id="status"), + ) + + app = rx.App() + app.add_workflow(OrderFlow) + + +@pytest.fixture(scope="module") +def workflow_app( + tmp_path_factory: pytest.TempPathFactory, +) -> Generator[AppHarness, None, None]: + """Run WorkflowApp with the store resolved from the environment. + + Args: + tmp_path_factory: pytest fixture for creating temporary directories. + + Yields: + Running AppHarness instance. + """ + db_dir = tmp_path_factory.mktemp("workflow_store") + with pytest.MonkeyPatch.context() as mp: + mp.setenv("REFLEX_WORKFLOW_DATABASE", str(db_dir / "runs.db")) + mp.setenv("ORDERS_WEBHOOK_SECRET", WEBHOOK_SECRET) + with AppHarness.create( + root=tmp_path_factory.mktemp("workflow_app"), + app_source=WorkflowApp, + ) as harness: + assert harness.app_instance is not None, "app is not running" + yield harness + + +def test_a_browser_event_starts_a_run_that_completes( + workflow_app: AppHarness, page: Page +): + """The full production path: click, durable timer, completion. + + One click proves the lifespan started the worker in the real server, the + page handler reached the runtime, the substep recorded, the two-second + timer fired on the wall clock, and the run completed into the SQLite file + the environment named. + """ + assert workflow_app.frontend_url is not None + page.goto(workflow_app.frontend_url) + page.locator("#launch").click() + expect(page.locator("#run-id")).not_to_be_empty(timeout=10_000) + + def completed() -> bool: + page.locator("#refresh").click() + return page.locator("#status").inner_text() == "COMPLETED" + + assert AppHarness._poll_for(completed, timeout=20, step=0.5), ( + f"run never completed; last status {page.locator('#status').inner_text()!r}" + ) + + +def test_the_webhook_endpoint_accepts_a_signed_request(workflow_app: AppHarness): + """A provider's signed delivery starts a run over real HTTP, exactly once.""" + base = get_config().api_url.rstrip("/") + body = json.dumps({"order_id": "hook-1", "amount": 5}).encode() + signature = hmac.new(WEBHOOK_SECRET.encode(), body, hashlib.sha256).hexdigest() + + first = httpx.post( + f"{base}/_workflow/webhook/orders.placed", + content=body, + headers={"X-Signature": signature, "content-type": "application/json"}, + ) + # 202: the run is durably admitted before the provider is acknowledged. + assert first.status_code == 202, first.text + run_id = first.json()["run_id"] + + # A redelivery reaches the same run rather than starting a second one. + second = httpx.post( + f"{base}/_workflow/webhook/orders.placed", + content=body, + headers={"X-Signature": signature, "content-type": "application/json"}, + ) + assert second.status_code == 202 + assert second.json()["run_id"] == run_id + assert second.json()["disposition"] == "deduplicated" + + # An unsigned delivery is refused outright. + forged = httpx.post( + f"{base}/_workflow/webhook/orders.placed", + content=body, + headers={"X-Signature": "0" * 64, "content-type": "application/json"}, + ) + assert forged.status_code in (400, 401, 403) From 21a98b4c4baf35d7d24af2f4befc1099f49c91a2 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 22:42:09 -0700 Subject: [PATCH 036/121] Resolve dotted check targets from the project root Reviewing the newest commits with the real console script -- not `python -m`, which quietly puts the working directory on sys.path -- surfaced that `reflex workflows check myapp.flows` failed with "No module named 'myapp'" when run from the project's own root, which reads as a typo rather than an environment difference. The dotted form now resolves from the working directory the way `python -m` would; the file-path form was always unaffected. Verified against the installed console script both before and after, plus a regression test that strips the working directory from sys.path first. --- reflex/workflow/cli.py | 8 ++++++++ tests/units/workflow/test_cli_check.py | 22 ++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 26d48d1bccc..27a1fce3bf2 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -11,6 +11,7 @@ import asyncio import inspect import json +import sys from pathlib import Path from typing import TYPE_CHECKING, Any @@ -353,6 +354,13 @@ def _load_module(target: str): module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module + # A dotted name is resolved from the project root, the way `python -m` + # would: the console script's sys.path does not include the working + # directory, and "myapp.workflows" failing from the project's own root + # is indistinguishable from a typo to the person running it. + cwd = str(Path.cwd()) + if cwd not in sys.path: + sys.path.insert(0, cwd) return importlib.import_module(target) diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py index b336f409fbe..8a46294beb7 100644 --- a/tests/units/workflow/test_cli_check.py +++ b/tests/units/workflow/test_cli_check.py @@ -8,6 +8,7 @@ """ import json +import sys from click.testing import CliRunner @@ -133,3 +134,24 @@ def test_a_missing_target_fails_cleanly(tmp_path, forked_registration_context): assert result.exit_code == 1 payload = json.loads(result.output) assert payload["ok"] is False + + +def test_a_dotted_module_resolves_from_the_project_root( + tmp_path, monkeypatch, forked_registration_context +): + """`reflex workflows check myapp.flows` works from the project's own root. + + The console script's sys.path does not include the working directory the + way `python -m` does, so without help the dotted form failed with "No + module named", indistinguishable from a typo. + """ + package = tmp_path / "myflows" + package.mkdir() + (package / "__init__.py").write_text("") + (package / "orders.py").write_text(VALID) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr("sys.path", [p for p in sys.path if p != str(tmp_path)]) + + result = CliRunner().invoke(workflows, ["check", "myflows.orders"]) + assert result.exit_code == 0, result.output + assert "check.greeter" in result.output From ef82d7f25c4e9f3ec4986305df902bbfe6b8aeb0 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 22:43:56 -0700 Subject: [PATCH 037/121] Run the workflow integration tests in prod mode too The workflow integration tests ran only against the dev server, and the exported-frontend production path is the one real deployments take -- and the one where backend serving differences historically hide. The app fixture now parametrizes over dev and prod the way the rest of the integration suite does, so the browser-to-completion path and the signed webhook contract are proven against both. --- tests/integration/tests_playwright/test_workflows.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/integration/tests_playwright/test_workflows.py b/tests/integration/tests_playwright/test_workflows.py index ebbdb3b66ff..724bb0464b7 100644 --- a/tests/integration/tests_playwright/test_workflows.py +++ b/tests/integration/tests_playwright/test_workflows.py @@ -92,22 +92,26 @@ def index(): @pytest.fixture(scope="module") def workflow_app( + app_harness_env: type[AppHarness], tmp_path_factory: pytest.TempPathFactory, ) -> Generator[AppHarness, None, None]: - """Run WorkflowApp with the store resolved from the environment. + """Run WorkflowApp in dev or prod mode, store resolved from the environment. Args: + app_harness_env: AppHarness (dev) or AppHarnessProd (prod). tmp_path_factory: pytest fixture for creating temporary directories. Yields: Running AppHarness instance. """ db_dir = tmp_path_factory.mktemp("workflow_store") + name = f"workflow_app_{app_harness_env.__name__.lower()}" with pytest.MonkeyPatch.context() as mp: mp.setenv("REFLEX_WORKFLOW_DATABASE", str(db_dir / "runs.db")) mp.setenv("ORDERS_WEBHOOK_SECRET", WEBHOOK_SECRET) - with AppHarness.create( - root=tmp_path_factory.mktemp("workflow_app"), + with app_harness_env.create( + root=tmp_path_factory.mktemp(name), + app_name=name, app_source=WorkflowApp, ) as harness: assert harness.app_instance is not None, "app is not running" From fc0716f97ff43557579eeb4b2aafb2d7421955f0 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 23:33:27 -0700 Subject: [PATCH 038/121] Freeze the execution contract, and close the two gaps it exposed Writing down what the engine promises -- reflex/workflow/CONTRACT.md, with a kill-at-every-boundary failure matrix -- turned up two places where the code did not yet keep the promise. Both are now fixed rather than documented as exceptions. A child's arrival at its parent's join was delivered after the child's own commit. A worker dying in that window left a run that was finished and a join that waited on it forever: nothing recovers that, because the store sees a completed child and a legitimately blocked parent. StepCompletion now carries the arrival, and all three stores apply it inside the same transaction as the terminal transition. What the kernel still does afterwards -- waking the parent, cancelling a decided race's losers -- is advisory follow-up whose loss costs nothing. Schedule catch-up cursors lived only in memory and were seeded at startup, so a restart treated every occurrence during the downtime as already fired: the hourly job simply did not run for the hour you were deploying, and nothing recorded the skip. Cursors now persist in the store (conformance-checked, all three), read lazily on first sweep, so a restart resumes the previous worker's position while a schedule this deployment has never seen still starts from now rather than backfilling its whole history. Catch-up remains capped at MAX_SCHEDULE_CATCHUP per sweep. The restart test was verified to fail with cursor persistence disabled. --- news/workflow-atomic-arrival.fix.md | 1 + reflex/workflow/CONTRACT.md | 214 ++++++++++++++++++ reflex/workflow/conformance.py | 13 ++ reflex/workflow/kernel.py | 91 +++++++- reflex/workflow/postgres.py | 129 +++++++++++ reflex/workflow/store.py | 297 ++++++++++++++++++++++--- tests/units/workflow/test_parallel.py | 35 ++- tests/units/workflow/test_schedules.py | 35 +++ 8 files changed, 769 insertions(+), 46 deletions(-) create mode 100644 news/workflow-atomic-arrival.fix.md create mode 100644 reflex/workflow/CONTRACT.md diff --git a/news/workflow-atomic-arrival.fix.md b/news/workflow-atomic-arrival.fix.md new file mode 100644 index 00000000000..5a847a783fb --- /dev/null +++ b/news/workflow-atomic-arrival.fix.md @@ -0,0 +1 @@ +A child run's arrival at its parent's join now commits in the same transaction as the child's final transition, so a crash between them can no longer leave the join waiting forever; schedule catch-up cursors persist in the store, so a restart resumes instead of skipping missed occurrences. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md new file mode 100644 index 00000000000..f808baf2d8a --- /dev/null +++ b/reflex/workflow/CONTRACT.md @@ -0,0 +1,214 @@ +# The execution contract + +This document freezes what the workflow engine promises. Every statement here +is load-bearing: code that violates it is a bug even if every test passes, and +a change that needs different semantics must change this document in the same +commit. The conformance suite (`reflex.workflow.CONFORMANCE_CHECKS`) and the +crash tests are the executable form of this contract; prose here and checks +there are meant to be read together. + +The engine's one architectural bet, from which everything below follows: +**state is snapshotted, never replayed.** A handler's committed state is data +in the store. No user code is ever re-executed to reconstruct anything. The +price is that control flow is expressed as return transitions rather than +imperative code; the payoff is that handlers are ordinary Python with no +determinism constraints, no versioning patches, and no replay-divergence bug +class. + +## 1. What is atomic + +One attempt commits in **one store transaction**, containing all of: + +- the executed step's terminal status (`SUCCEEDED`, `FAILED`, `RETRY_WAIT`, + `NEEDS_ATTENTION`, `BLOCKED` for waits) and error payload, +- the run-state snapshot and its incremented `state_version`, +- every successor slot the transition scheduled (with preallocated ordinals), +- every child run and child root slot a fan-out admitted, +- tombstones for slots a terminal transition abandoned, +- the run's status, result, and error, +- the history events describing all of the above, +- **and, when the transition ends a child run: the arrival delivered to its + parent's join slot.** + +There is no state in which a step "happened" but its consequences are missing. +Either the whole transition is visible or none of it is. + +Admission is likewise one transaction: run row, root slot, dedupe reservation, +history. A webhook is acknowledged only after that transaction commits +(admit-before-ack), so a `202` means the run exists durably. + +Substep results (`rx.step`) are the deliberate exception: each records in its +**own** transaction the moment the callable returns, because their purpose is +to survive a crash that prevents the attempt from ever committing. + +## 2. When handlers re-execute + +A handler runs more than once in exactly two situations, both bounded: + +1. **Business retry.** The attempt raised (or timed out) and the resolved + retry policy grants another attempt. Consumes one attempt from + `Retry.max_attempts`; scheduled with backoff and jitter. +2. **Crash recovery.** A worker died (or lost its lease) mid-attempt. The + step's lease lapses, recovery moves it to `RECOVERY_WAIT`, and any worker + re-executes it. Consumes one *recovery* (budget: 10 per logical step, then + the run fails), never a business attempt. + +There is no third case. Completed steps are never re-run; a deploy never +re-runs anything; reading a run never runs anything. + +Re-execution is therefore **at-least-once per attempt boundary**. The tools +that turn it into effectively-once for side effects are, in order of strength: + +- `rx.step(name, fn, ...)`: records the result durably at return; every later + execution of the same handler replays the recorded value instead of calling + `fn`. Fenced by claim epoch, so a zombie worker cannot write (§7). +- `rx.current_run().idempotency_key()`: stable across retries and recoveries + of one step, distinct across steps — hand it to providers that accept + idempotency keys. This covers the one window `rx.step` cannot: a crash + after the provider acted but before the record landed. +- `effect="non_idempotent_write"`: one business attempt only; an error + suspends the run as `NEEDS_ATTENTION` instead of retrying. Note this + governs *retries*, not recoveries: a crash mid-attempt still re-executes + after the lease lapses. Money-moving code must use `rx.step` and/or + provider keys; the effect class alone is not an exactly-once guarantee, and + nothing can be while the process can die between the provider call and any + record of it. + +## 3. Effects and retries + +| effect | default retry | on exhaustion / error | +|---|---|---| +| `none`, `read`, `idempotent_write` | any `Exception`, backoff+jitter, per declared `Retry` (or its defaults) | run `FAILED` (after `on_failure` hook, if declared) | +| `non_idempotent_write` | 1 attempt; no implicit retry | run `NEEDS_ATTENTION`; operator resumes or fails | + +- `TransientWorkflowError` is always retryable regardless of `retry_on`. +- `timeout=` bounds one attempt; a timed-out attempt counts as a failed one + and follows the same policy (`on_timeout` hook runs on final timeout). + `timeout=` is a compile error on sync handlers: a thread cannot be + interrupted, and a "timed out" attempt still running concurrently with its + retry would be a lie. +- Run-level `timeout` (`WorkflowConfig.run_timeout`) finalizes the run + `TIMED_OUT` once drained; the in-flight attempt is cancelled cooperatively. + +## 4. Releases and versions + +A run is **not** pinned to the code that started it. On every claim the engine +gates per step: the handler id must still exist and the recorded payload must +still bind to its signature. If both hold, the current code runs — new fields, +tuned retries, edited bodies all apply immediately. If either fails, the run +suspends as `NEEDS_ATTENTION` naming the handler; `resume()` re-opens it after +the operator ships a compatible release or intervenes. + +Consequences, stated plainly: + +- Adding state fields, handlers, or hooks never strands an in-flight run. +- Deleting a handler (or removing a parameter its recorded payloads carry) + suspends exactly the runs whose *next* step needs it, and only when that + step comes due. +- Two releases running simultaneously (rolling deploy) may execute different + steps of one run with different code. Steps are the consistency boundary; + the contract makes no promise that one run sees one release. + +## 5. Cancellation, deadlines, and children + +- `cancel(run_id)` records intent and cancels any in-flight attempt + cooperatively. The run finalizes `CANCELLED` only once no step is claimed + (drained), tombstoning open slots. Cancellation is never delivered as an + exception into a *different* run's handler. +- A run past its deadline finalizes `TIMED_OUT` the same drained way. +- Every terminal path of a child — commit, cancellation, run timeout, + recovery-budget exhaustion — delivers exactly one arrival to its parent's + join slot, atomically with the terminal transition (§1). A join can wait + forever only on a child that is still genuinely running. +- `rx.parallel(..., mode="first")`: the join resolves on the first arrival; + the engine then requests cancellation of the losing branches. That request + is best-effort follow-up, not part of the winning transaction: if the + process dies first, losers run to completion and their arrivals are + refused as late (`counted`/`duplicate`/terminal), which is harmless. +- Child runs are ordinary runs; cancelling the parent does not implicitly + cancel children (fan-out is delegation, not ownership). A cancelled + parent's join tombstones; late child arrivals are refused. + +## 6. Identity: who is "the same" as whom + +Checked in this order at start: + +1. **`request_key`** (explicit, or webhook `dedupe_by`, or schedule + occurrence key `schedule:{workflow}:{handler}:{epoch}`): resolved *before + any start policy*, so a provider redelivery returns the original run and + can never trip a singleton, throttle, or debounce against it. Reserved + atomically at admission; concurrent duplicate admissions yield one run. +2. **Flow key** (`Singleton`/`RateLimit`/`Throttle`/`Debounce` `key=`): a + handler parameter, or a field of exactly one model parameter (ambiguity is + a compile error). Groups runs for policy decisions. +3. Approval links: single-use by delivery key (default: hash of + run+channel+payload), expiring, HMAC over all claims; a `GET` never + delivers — only the confirming `POST` does. +4. Signals: sender-supplied `key` makes redelivery a no-op; an expired wait + refuses its signal rather than letting it resolve a later wait on the + same channel. + +## 7. Workers + +- A **claim** takes the run's frontier step (lowest unresolved ordinal) — + strictly one open obligation per run, so per-run order is total. Claims are + fenced by `(CLAIMED, epoch, state_version)`; every commit re-validates the + fence and a fenced writer's work is discarded (`abandoned`, in history). +- **Leases** renew on a real-time cadence; expiry is injected-clock time. A + worker that cannot renew abandons its attempt *before* the lease lapses + rather than run work it can no longer prove it owns. Recovery reclaims only + lapsed leases — a slow worker is never raced, and takeover is delayed by at + most one lease. +- The substep journal is epoch-fenced: a reclaimed worker's late + `rx.step` write is refused, and the attempt kills itself instead of + duplicating a side effect. +- **Queues**: every step carries its handler's queue (`"default"` if none); + a worker claims only queues it serves. Order within a run holds across + queues — a run whose frontier is on an unserved queue waits. Recovery is + queue-agnostic (any worker recovers; the reclaimed step is then claimed by + the right one). +- Multiple workers share one Postgres store via `SKIP LOCKED` claims; SQLite + is a one-process store (calls off-loop, contention bounded); memory is for + tests. All three answer the same conformance suite. +- Managed vs customer-hosted is a deployment split, not a semantic one: a + worker is any process with `REFLEX_WORKFLOW_DATABASE` pointed at the store + and (optionally) `workflow_queues` narrowed. The contract is identical. + +## 8. The failure matrix + +"Kill" means SIGKILL — no cleanup runs. Each row states the one permitted +outcome. + +| killed at | outcome | +|---|---| +| before admission commits | nothing exists; provider retry/redelivery admits normally (dedupe makes it one run) | +| after admission, before webhook ack | run exists; provider redelivers; `request_key` returns the same run (`deduplicated`) | +| after claim, before handler code runs | lease lapses → recovery → re-executed; costs one recovery, no business attempt | +| mid-handler, after a `rx.step` recorded | re-execution replays the recorded substeps; only un-recorded work repeats | +| mid-handler, provider call sent but `rx.step` not yet recorded | re-execution repeats the call; the provider-side idempotency key (`idempotency_key()`) is the defense — this window is why it exists | +| after commit (any transition) | everything in §1 is durable, including successors, children, and parent arrival; the worker's post-commit follow-ups (observer notify, loser cancellation, wakeups) are advisory and their loss is harmless | +| between a child's terminal commit and anything else | nothing is pending: the parent arrival was inside the commit | +| during recovery sweep | idempotent; re-run by the next sweep | +| worker dies holding N claims | each lease lapses independently; each step recovered independently | +| store unreachable at commit | attempt abandoned (fence unverifiable); step recovered later; `rx.step` records already made stand | +| everything down for an hour | timers/waits/retries fire on restart (due-time semantics); schedule occurrences catch up from the durable cursor, capped at `MAX_SCHEDULE_CATCHUP` per schedule, remainder skipped with a history record | + +## 9. Operator actions + +Every action is legal only from the states listed; anything else is a refused +no-op with a reason. + +- `cancel(run)` — any nonterminal run. +- `resume(run)` — `NEEDS_ATTENTION` only; re-opens the suspended step with a + fresh attempt budget. +- `retry(run)` — a `FAILED` run: re-opens its failed frontier step with a + fresh attempt budget and re-runs from there. History keeps the failure. +- `skip(run)` — a suspended or failed run: marks the blocking step `SKIPPED` + (terminal, recorded as an operator decision) and lets the run continue as + if it had succeeded *with no transition* — the next open slot, if any, + proceeds; otherwise the run completes with no result. +- `force_complete(run, result)` / `force_fail(run, error)` — a nonterminal, + drained run: finalizes immediately, tombstoning open slots, recording the + operator origin. Parent joins receive the arrival like any terminal path. + +All of these are store transactions under the same atomicity rules as §1. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 96bcaafeb26..b6bd9d423c0 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -687,6 +687,18 @@ async def check_substeps_record_once_and_fence_stale_writers( assert list(await store.get_substeps("run1", 0)) == ["charge", "label"] +async def check_schedule_cursors_persist(store: RunStore) -> None: + """A schedule's catch-up position survives the process that wrote it.""" + assert await store.read_schedule_cursor("wf:tick") is None + await store.write_schedule_cursor("wf:tick", NOW) + assert await store.read_schedule_cursor("wf:tick") == pytest.approx(NOW) + # Advancing overwrites rather than accumulating. + await store.write_schedule_cursor("wf:tick", NOW + 60) + assert await store.read_schedule_cursor("wf:tick") == pytest.approx(NOW + 60) + # Schedules are independent of one another. + assert await store.read_schedule_cursor("wf:other") is None + + async def check_reads_do_not_alias_stored_state(store: RunStore) -> None: """Mutating a returned record must not change what the store holds.""" await store.admit(make_run(), make_step(), _ADMITTED) @@ -763,6 +775,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_list_children_finds_a_joins_branches, check_claims_respect_queue_boundaries, check_substeps_record_once_and_fence_stale_writers, + check_schedule_cursors_persist, check_join_arrivals_count_once, check_finalize_refuses_while_a_step_is_claimed, check_finalize_tombstones_open_slots, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 034c9721b53..f4f775201c8 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -12,6 +12,7 @@ import asyncio import contextlib +import dataclasses import random import time import traceback @@ -344,12 +345,11 @@ def __init__( for handler in (defn.handlers[hid] for hid in defn.roots) if isinstance(handler.trigger, ScheduleTrigger) ] - # Seeded at construction so a freshly started process never backfills - # occurrences from before it existed. - self._schedule_cursor: dict[str, float] = { - f"{defn.workflow_id}:{handler.id}": clock() - for defn, handler, _ in self._schedules - } + # Filled lazily from the store on first sweep: a restart resumes the + # previous worker's cursor, and only a schedule this deployment has + # never seen starts from now (never backfilling its whole history). + self._schedule_cursor: dict[str, float] = {} + self._started_at = clock() self._field_adapters: dict[tuple[str, str], TypeAdapter] = {} self._inflight: dict[str, asyncio.Task] = {} self._leases: dict[str, _Lease] = {} @@ -1945,7 +1945,14 @@ async def _admit_due_schedules(self, now: float) -> int: admitted = 0 for defn, handler, schedule in self._schedules: key = f"{defn.workflow_id}:{handler.id}" - cursor = self._schedule_cursor[key] + cursor = self._schedule_cursor.get(key) + if cursor is None: + # A restart must resume where the last worker stopped, not + # skip the downtime: an in-memory cursor seeded at startup + # treats every missed occurrence as already fired. + stored = await self._store.read_schedule_cursor(key) + cursor = stored if stored is not None else self._started_at + self._schedule_cursor[key] = cursor for occurrence in schedule.occurrences_between( cursor, now, limit=MAX_SCHEDULE_CATCHUP ): @@ -1956,6 +1963,7 @@ async def _admit_due_schedules(self, now: float) -> int: ) admitted += result.disposition == "started" self._schedule_cursor[key] = now + await self._store.write_schedule_cursor(key, now) return admitted def _next_schedule_due(self, now: float) -> float | None: @@ -2049,6 +2057,7 @@ async def _commit_outcome( handler: The handler that ran. completion: The outcome to apply. """ + completion = self._with_parent_arrival(claim.run, completion) try: await self._store.commit(claim, completion, self._clock()) except StaleClaimError: @@ -2059,20 +2068,80 @@ async def _commit_outcome( self._wakeup.set() await self._report_to_parent(claim.run, completion) + @staticmethod + def _with_parent_arrival( + run: RunRecord, completion: StepCompletion + ) -> StepCompletion: + """Attach a child's parent arrival so it commits with the transition. + + A child that finishes and a parent that hears about it must become + true together. Delivering afterwards leaves a window where a crash + strands the join forever, waiting on a child that is already done. + + Args: + run: The run being committed, which may have no parent. + completion: The outcome about to be applied. + + Returns: + The completion, carrying the arrival when one is owed. + """ + if ( + run.parent_run_id is None + or run.parent_ordinal is None + or completion.run_status not in TERMINAL_RUN_STATUSES + ): + return completion + return dataclasses.replace( + completion, + parent_arrival=( + run.parent_run_id, + run.parent_ordinal, + { + "run_id": run.run_id, + "status": completion.run_status.value, + "result": completion.result, + "error": completion.run_error, + }, + run.run_id, + ), + ) + async def _report_to_parent( self, run: RunRecord, completion: StepCompletion ) -> None: - """Report a finished child's outcome to the join slot awaiting it. + """Follow up on an arrival that committed with the child's transition. + + The arrival itself is already durable (see _with_parent_arrival); what + remains is advisory: waking the parent's worker and, for a decided + race, cancelling the branches that lost. Args: run: The child run, which may have no parent. completion: The committed outcome that finished it. """ - if completion.run_status not in TERMINAL_RUN_STATUSES: + if ( + completion.run_status not in TERMINAL_RUN_STATUSES + or run.parent_run_id is None + or run.parent_ordinal is None + ): return - await self._report_outcome( - run, completion.run_status, completion.result, completion.run_error + await self._notify_run( + run.parent_run_id, + ( + ( + HistoryEventType.CHILD_RESOLVED, + {"ordinal": run.parent_ordinal, "child": run.run_id}, + ), + ), ) + parent_steps = await self._store.get_steps(run.parent_run_id) + join = next( + (step for step in parent_steps if step.ordinal == run.parent_ordinal), + None, + ) + if join is not None and join.status is not StepStatus.BLOCKED: + await self._cancel_losing_branches(run) + self._wakeup.set() async def _report_outcome( self, diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 27a02100e9c..ca7ed98ee13 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -116,6 +116,10 @@ run_id TEXT NOT NULL, PRIMARY KEY (workflow_id, request_key) ); +CREATE TABLE IF NOT EXISTS workflow_schedules ( + key TEXT PRIMARY KEY, + at DOUBLE PRECISION NOT NULL +); CREATE TABLE IF NOT EXISTS workflow_substeps ( run_id TEXT NOT NULL, ordinal INTEGER NOT NULL, @@ -805,6 +809,8 @@ async def commit( ), ) await self._append_events(conn, claim.run.run_id, completion.events, now) + if completion.parent_arrival is not None: + await self._apply_arrival(conn, *completion.parent_arrival, now) async def release_claim( self, @@ -991,6 +997,97 @@ async def admit_children( await self._lock_run(conn, parent) await self._append_events(conn, parent, events, now) + async def _apply_arrival( + self, + conn: Connection, + run_id: str, + ordinal: int, + payload: dict[str, Any], + dedupe_key: str, + now: float, + ) -> str: + """Count one arrival inside the caller's open transaction. + + Args: + conn: The connection inside an open transaction. + run_id: The waiting parent run. + ordinal: The join slot's ordinal. + payload: The arriving result. + dedupe_key: Identity of the arrival. + now: Current time in epoch seconds. + + Returns: + What the store did with the arrival. + """ + wait_key = f"join:{ordinal}" + cursor = await conn.execute( + "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", + (run_id,), + ) + run_row = await cursor.fetchone() + if run_row is None: + return "unknown_run" + if run_row["status"] in _TERMINAL_RUNS: + return "run_terminal" + cursor = await conn.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", + (run_id, wait_key, dedupe_key), + ) + if await cursor.fetchone() is not None: + return "duplicate" + cursor = await conn.execute( + "SELECT * FROM workflow_steps WHERE run_id = %s AND ordinal = %s", + (run_id, ordinal), + ) + step_row = await cursor.fetchone() + if step_row is None or step_row["status"] != StepStatus.BLOCKED.value: + return "run_terminal" + step = _step_from_row(step_row) + await conn.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at)" + " VALUES (%s, %s, %s, %s, %s, 'CONSUMED', %s)", + ( + run_id, + wait_key, + dedupe_key, + await self._next_inbox_seq(conn, run_id), + _json(payload), + now, + ), + ) + arrived = step.join_arrived + 1 + results = [*step.args.get("__results__", []), payload] + done = arrived >= step.join_expected + await conn.execute( + "UPDATE workflow_steps SET status = %s, join_arrived = %s," + " due_at = %s, args = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s AND join_arrived = %s", + ( + StepStatus.READY.value if done else StepStatus.BLOCKED.value, + arrived, + now if done else step.due_at, + _json({**step.args, "__results__": results}), + now, + run_id, + ordinal, + step.join_arrived, + ), + ) + await self._append_events( + conn, + run_id, + ( + ( + HistoryEventType.CHILD_RESOLVED, + {"ordinal": ordinal, "arrived": arrived}, + ), + ), + now, + ) + return "resolved" if done else "counted" + async def record_arrival( self, run_id: str, @@ -1621,6 +1718,38 @@ async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: for row in await cursor.fetchall() ) + async def read_schedule_cursor(self, key: str) -> float | None: + """Read where a schedule's catch-up last reached. + + Args: + key: The schedule identity, "{workflow_id}:{handler_id}". + + Returns: + The last swept time, or None when the schedule is new here. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT at FROM workflow_schedules WHERE key = %s", (key,) + ) + row = await cursor.fetchone() + return None if row is None else row["at"] + + async def write_schedule_cursor(self, key: str, at: float) -> None: + """Record where a schedule's catch-up has now reached. + + Args: + key: The schedule identity, "{workflow_id}:{handler_id}". + at: The time swept up to, in epoch seconds. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + await conn.execute( + "INSERT INTO workflow_schedules (key, at) VALUES (%s, %s)" + " ON CONFLICT (key) DO UPDATE SET at = EXCLUDED.at", + (key, at), + ) + async def next_due( self, now: float, *, queues: tuple[str, ...] | None = None ) -> float | None: diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 8bfa640f145..7e974cad057 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -94,6 +94,11 @@ class StepCompletion: events: History events to append, in order, as (type, data) pairs. children: Child runs to create in the same transaction as this commit, each paired with its root slot. + parent_arrival: When this transition ends a child run, the arrival to + deliver to its parent's join slot, as (parent_run_id, ordinal, + payload, dedupe_key). Delivered inside this transaction: a crash + between a child finishing and its parent hearing about it would + leave the join waiting forever. """ step_status: StepStatus @@ -109,6 +114,7 @@ class StepCompletion: next_ordinal: int | None = None events: tuple[tuple[HistoryEventType, dict[str, Any]], ...] = () children: tuple[tuple[RunRecord, StepRecord], ...] = () + parent_arrival: tuple[str, int, dict[str, Any], str] | None = None class RunStore(Protocol): @@ -590,6 +596,30 @@ async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: """ ... + async def read_schedule_cursor(self, key: str) -> float | None: + """Read where a schedule's catch-up last reached. + + Args: + key: The schedule identity, "{workflow_id}:{handler_id}". + + Returns: + The last swept time, or None when the schedule is new here. + """ + ... + + async def write_schedule_cursor(self, key: str, at: float) -> None: + """Record where a schedule's catch-up has now reached. + + Persisting this is what makes a restart resume rather than silently + skip: an in-memory cursor reseeded at startup treats every occurrence + during the downtime as if it had already fired. + + Args: + key: The schedule identity, "{workflow_id}:{handler_id}". + at: The time swept up to, in epoch seconds. + """ + ... + async def next_due( self, now: float, *, queues: tuple[str, ...] | None = None ) -> float | None: @@ -736,6 +766,7 @@ def __init__(self): self._runs: dict[str, RunRecord] = {} self._steps: dict[str, list[StepRecord]] = {} self._substeps: dict[tuple[str, int], dict[str, Any]] = {} + self._schedule_cursors: dict[str, float] = {} self._history: dict[str, list[HistoryEvent]] = {} self._dedupe: dict[tuple[str, str], str] = {} self._inbox: dict[str, dict[tuple[str, str, str], bool]] = {} @@ -976,6 +1007,9 @@ async def commit( updated_at=now, ) self._append_events(run.run_id, completion.events, now) + if completion.parent_arrival is not None: + parent_id, ordinal, payload, dedupe_key = completion.parent_arrival + self._apply_arrival(parent_id, ordinal, payload, dedupe_key, now) async def release_claim( self, @@ -1128,42 +1162,64 @@ async def record_arrival( What the store did with the arrival. """ async with self._lock: - run = self._runs.get(run_id) - if run is None: - return "unknown_run" - if run.status in TERMINAL_RUN_STATUSES: - return "run_terminal" - seen = self._inbox.setdefault(run_id, {}) - key = (run_id, f"join:{ordinal}", dedupe_key) - if key in seen: - return "duplicate" - seen[key] = True - steps = self._steps[run_id] - step = steps[ordinal] - if step.status is not StepStatus.BLOCKED: - return "run_terminal" - arrived = step.join_arrived + 1 - results = [*step.args.get("__results__", []), payload] - done = arrived >= step.join_expected - steps[ordinal] = dataclasses.replace( - step, - status=StepStatus.READY if done else StepStatus.BLOCKED, - join_arrived=arrived, - due_at=now if done else step.due_at, - args={**step.args, "__results__": results}, - updated_at=now, - ) - self._append_events( - run_id, + return self._apply_arrival(run_id, ordinal, payload, dedupe_key, now) + + def _apply_arrival( + self, + run_id: str, + ordinal: int, + payload: dict[str, Any], + dedupe_key: str, + now: float, + ) -> DeliveryDisposition: + """Count one arrival, with the lock already held. + + Args: + run_id: The waiting parent run. + ordinal: The join slot's ordinal. + payload: The arriving result. + dedupe_key: Identity of the arrival. + now: Current time in epoch seconds. + + Returns: + What the store did with the arrival. + """ + run = self._runs.get(run_id) + if run is None: + return "unknown_run" + if run.status in TERMINAL_RUN_STATUSES: + return "run_terminal" + seen = self._inbox.setdefault(run_id, {}) + key = (run_id, f"join:{ordinal}", dedupe_key) + if key in seen: + return "duplicate" + seen[key] = True + steps = self._steps[run_id] + step = steps[ordinal] + if step.status is not StepStatus.BLOCKED: + return "run_terminal" + arrived = step.join_arrived + 1 + results = [*step.args.get("__results__", []), payload] + done = arrived >= step.join_expected + steps[ordinal] = dataclasses.replace( + step, + status=StepStatus.READY if done else StepStatus.BLOCKED, + join_arrived=arrived, + due_at=now if done else step.due_at, + args={**step.args, "__results__": results}, + updated_at=now, + ) + self._append_events( + run_id, + ( ( - ( - HistoryEventType.CHILD_RESOLVED, - {"ordinal": ordinal, "arrived": arrived}, - ), + HistoryEventType.CHILD_RESOLVED, + {"ordinal": ordinal, "arrived": arrived}, ), - now, - ) - return "resolved" if done else "counted" + ), + now, + ) + return "resolved" if done else "counted" async def count_active(self, workflow_id: str, flow_key: str) -> int: """Count runs of a root still in flight under a flow-control key. @@ -1626,6 +1682,32 @@ async def get_history(self, run_id: str) -> tuple[HistoryEvent, ...]: async with self._lock: return tuple(self._history.get(run_id, ())) + async def read_schedule_cursor(self, key: str) -> float | None: + """Read where a schedule's catch-up last reached. + + Args: + key: The schedule identity, "{workflow_id}:{handler_id}". + + Returns: + The last swept time, or None when the schedule is new here. + """ + async with self._lock: + return self._schedule_cursors.get(key) + + async def write_schedule_cursor(self, key: str, at: float) -> None: + """Record where a schedule's catch-up has now reached. + + Persisting this is what makes a restart resume rather than silently + skip: an in-memory cursor reseeded at startup treats every occurrence + during the downtime as if it had already fired. + + Args: + key: The schedule identity, "{workflow_id}:{handler_id}". + at: The time swept up to, in epoch seconds. + """ + async with self._lock: + self._schedule_cursors[key] = at + async def next_due( self, now: float, *, queues: tuple[str, ...] | None = None ) -> float | None: @@ -1745,6 +1827,10 @@ def resolve_store(target: str | None = None) -> RunStore: run_id TEXT NOT NULL, PRIMARY KEY (workflow_id, request_key) ); +CREATE TABLE IF NOT EXISTS workflow_schedules ( + key TEXT PRIMARY KEY, + at REAL NOT NULL +); CREATE TABLE IF NOT EXISTS workflow_substeps ( run_id TEXT NOT NULL, ordinal INTEGER NOT NULL, @@ -2364,6 +2450,8 @@ def work() -> None: ), ) self._append_events(claim.run.run_id, completion.events, now) + if completion.parent_arrival is not None: + self._apply_arrival_sql(*completion.parent_arrival, now) self._db.execute("COMMIT") except BaseException: self._db.execute("ROLLBACK") @@ -2583,6 +2671,94 @@ def work() -> None: await asyncio.to_thread(work) + def _apply_arrival_sql( + self, + run_id: str, + ordinal: int, + payload: dict[str, Any], + dedupe_key: str, + now: float, + ) -> str: + """Count one arrival inside the caller's open transaction. + + Args: + run_id: The waiting parent run. + ordinal: The join slot's ordinal. + payload: The arriving result. + dedupe_key: Identity of the arrival. + now: Current time in epoch seconds. + + Returns: + What the store did with the arrival. + """ + terminal = tuple(status.value for status in TERMINAL_RUN_STATUSES) + wait_key = f"join:{ordinal}" + run_row = self._db.execute( + "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + ).fetchone() + if run_row is None: + return "unknown_run" + if run_row["status"] in terminal: + return "run_terminal" + seen = self._db.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", + (run_id, wait_key, dedupe_key), + ).fetchone() + if seen is not None: + return "duplicate" + step_row = self._db.execute( + "SELECT * FROM workflow_steps WHERE run_id = ? AND ordinal = ?", + (run_id, ordinal), + ).fetchone() + if step_row is None or step_row["status"] != StepStatus.BLOCKED.value: + return "run_terminal" + step = _step_from_row(step_row) + self._db.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at)" + " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" + " workflow_inbox WHERE run_id = ?), ?, ?, ?)", + ( + run_id, + wait_key, + dedupe_key, + run_id, + json.dumps(payload), + "CONSUMED", + now, + ), + ) + arrived = step.join_arrived + 1 + results = [*step.args.get("__results__", []), payload] + done = arrived >= step.join_expected + self._db.execute( + "UPDATE workflow_steps SET status = ?, join_arrived = ?," + " due_at = ?, args = ?, updated_at = ?" + " WHERE run_id = ? AND ordinal = ? AND join_arrived = ?", + ( + StepStatus.READY.value if done else StepStatus.BLOCKED.value, + arrived, + now if done else step.due_at, + json.dumps({**step.args, "__results__": results}), + now, + run_id, + ordinal, + step.join_arrived, + ), + ) + self._append_events( + run_id, + ( + ( + HistoryEventType.CHILD_RESOLVED, + {"ordinal": ordinal, "arrived": arrived}, + ), + ), + now, + ) + return "resolved" if done else "counted" + async def record_arrival( self, run_id: str, @@ -3447,6 +3623,59 @@ def work(): return await asyncio.to_thread(work) + async def read_schedule_cursor(self, key: str) -> float | None: + """Read where a schedule's catch-up last reached. + + Args: + key: The schedule identity, "{workflow_id}:{handler_id}". + + Returns: + The last swept time, or None when the schedule is new here. + """ + + def work() -> float | None: + """Read the cursor on the worker thread. + + Returns: + The stored time, or None. + """ + with self._lock: + row = self._db.execute( + "SELECT at FROM workflow_schedules WHERE key = ?", (key,) + ).fetchone() + return None if row is None else row["at"] + + return await asyncio.to_thread(work) + + async def write_schedule_cursor(self, key: str, at: float) -> None: + """Record where a schedule's catch-up has now reached. + + Persisting this is what makes a restart resume rather than silently + skip: an in-memory cursor reseeded at startup treats every occurrence + during the downtime as if it had already fired. + + Args: + key: The schedule identity, "{workflow_id}:{handler_id}". + at: The time swept up to, in epoch seconds. + """ + + def work() -> None: + """Write the cursor on the worker thread.""" + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + self._db.execute( + "INSERT INTO workflow_schedules (key, at) VALUES (?, ?)" + " ON CONFLICT(key) DO UPDATE SET at = excluded.at", + (key, at), + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + await asyncio.to_thread(work) + async def next_due( self, now: float, *, queues: tuple[str, ...] | None = None ) -> float | None: diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 39f3c02e908..00b7561403f 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -3,7 +3,7 @@ from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual import reflex as rx -from reflex.workflow.records import RunStatus, StepStatus +from reflex.workflow.records import HistoryEventType, RunStatus, StepStatus from reflex.workflow.testing import WorkflowTestHarness BRANCH_CALLS: list[str] = [] @@ -455,3 +455,36 @@ async def test_race_join_expects_one_arrival(): join = await _join_slot(harness, result.run_id) assert join.join_expected == 1 assert join.status is StepStatus.SUCCEEDED + + +async def test_a_childs_arrival_commits_with_its_final_transition(): + """A finished child and a told parent become true together. + + Delivering the arrival after the child's commit leaves a window: a worker + that dies inside it leaves a run that is finished and a join that waits on + it forever. Nothing recovers that, because from the store's point of view + the child is done and the parent is simply blocked. The arrival therefore + rides the same transaction as the child's terminal transition. + + This asserts it at the store level -- what is durable the instant the + commit returns -- rather than through the kernel's follow-up work, which + is exactly the code a crash would skip. + """ + RACE_CALLS.clear() + async with WorkflowTestHarness(Shopper, SlowVendor, FastVendor) as harness: + result = await harness.start(Shopper.start()) + assert result.run_id is not None + join = await _join_slot(harness, result.run_id) + + # The join was satisfied and the parent moved on, with no post-commit + # delivery involved: the store alone carries the evidence. + assert join.join_arrived >= 1 + history = await harness.kernel.store.get_history(result.run_id) + resolutions = [ + event for event in history if event.type is HistoryEventType.CHILD_RESOLVED + ] + assert resolutions, "the join recorded no arrival" + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED diff --git a/tests/units/workflow/test_schedules.py b/tests/units/workflow/test_schedules.py index f22358ce3c0..7332af7129b 100644 --- a/tests/units/workflow/test_schedules.py +++ b/tests/units/workflow/test_schedules.py @@ -10,6 +10,7 @@ from reflex.workflow.cron import CronSchedule from reflex.workflow.definition import compile_workflow from reflex.workflow.records import RunStatus +from reflex.workflow.store import MemoryRunStore from reflex.workflow.testing import WorkflowTestHarness # A Tuesday at 12:00 UTC, chosen so quarter-hour schedules are 15 minutes away. @@ -162,3 +163,37 @@ def sweep(self): async with WorkflowTestHarness(Cronly, start_time=START) as harness: with pytest.raises(WorkflowRuntimeError, match="cannot be started here"): await harness.kernel.start(Cronly.sweep) + + +async def test_a_restart_resumes_the_schedule_cursor(forked_registration_context): + """A worker that restarts catches up instead of skipping the downtime. + + With the cursor only in memory, a process starting up treats every + occurrence that happened while it was down as already fired -- the hourly + report simply does not run for the hour you were deploying, and nothing + records that it was skipped. + """ + fired: list[float] = [] + + class Hourly(rx.State): + __workflow__ = WorkflowConfig(id="sched.restart") + + @rx.event(durable=True, trigger=schedule("0 * * * *"), effect="none") + def tick(self): + """Note the occurrence.""" + fired.append(1) + + store = MemoryRunStore() + async with WorkflowTestHarness(Hourly, store=store) as harness: + await harness.advance("90m") + first_count = len(fired) + assert first_count >= 1 + + # A new process on the same store: the cursor survived, so the hours that + # passed while nothing was running are caught up rather than skipped. + async with WorkflowTestHarness( + Hourly, store=store, start_time=harness.now + 7200 + ) as resumed: + await resumed.advance("1m") + + assert len(fired) > first_count, "the restarted worker skipped the downtime" From 50851132bf5d196f3a3c57fd8389e7a174cb5825 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 23:39:11 -0700 Subject: [PATCH 039/121] Add the operator actions the contract promises CONTRACT.md section 9 lists what an operator may do to a run; cancel and resume existed, the rest did not. Two more now do, each legal only from the states the contract names and a refused no-op everywhere else -- an action that silently succeeds from the wrong state is how someone corrupts the run they were trying to rescue. retry re-opens a failed run at the step that failed, with a fresh attempt budget, for the ordinary case of a failure whose cause has since been fixed. The failure stays in history: a retry does not rewrite the record of why it was needed. force_complete and force_fail end a run no code path will finish -- a wait nobody will answer, a provider that no longer exists. They finalize a drained run, tombstone what it had open, record the operator origin and (for completion) the result to treat it as having produced, and deliver the arrival to a waiting parent like any other terminal path. Both are refused while a step is claimed, so they never race a working attempt. Exposed on rx.workflows and, for retry, on the CLI. Conformance-checked, so all three stores answer identically. The contract's section 9 was trimmed to exactly what ships: a skip action was described and is not implemented, and a contract that describes unimplemented behavior is worse than one that admits a gap. --- news/workflow-operator-actions.feature.md | 1 + reflex/workflow/CONTRACT.md | 16 +- reflex/workflow/cli.py | 38 ++++ reflex/workflow/conformance.py | 52 +++++ reflex/workflow/kernel.py | 63 ++++++ reflex/workflow/postgres.py | 47 ++++- reflex/workflow/runtime.py | 49 +++++ reflex/workflow/store.py | 132 ++++++++++++- tests/units/workflow/test_operator_actions.py | 179 ++++++++++++++++++ 9 files changed, 564 insertions(+), 13 deletions(-) create mode 100644 news/workflow-operator-actions.feature.md create mode 100644 tests/units/workflow/test_operator_actions.py diff --git a/news/workflow-operator-actions.feature.md b/news/workflow-operator-actions.feature.md new file mode 100644 index 00000000000..69e83f98da8 --- /dev/null +++ b/news/workflow-operator-actions.feature.md @@ -0,0 +1 @@ +Operators can now retry a failed run (`rx.workflows.retry`, `reflex workflows retry`) and end a stuck one by decision (`rx.workflows.force_complete` / `force_fail`). diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index f808baf2d8a..40b9b8a64ba 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -201,14 +201,14 @@ no-op with a reason. - `cancel(run)` — any nonterminal run. - `resume(run)` — `NEEDS_ATTENTION` only; re-opens the suspended step with a fresh attempt budget. -- `retry(run)` — a `FAILED` run: re-opens its failed frontier step with a - fresh attempt budget and re-runs from there. History keeps the failure. -- `skip(run)` — a suspended or failed run: marks the blocking step `SKIPPED` - (terminal, recorded as an operator decision) and lets the run continue as - if it had succeeded *with no transition* — the next open slot, if any, - proceeds; otherwise the run completes with no result. -- `force_complete(run, result)` / `force_fail(run, error)` — a nonterminal, +- `retry(run)` — a `FAILED` run: re-opens its failed step with a fresh + attempt budget and re-runs from there. History keeps the failure; a retry + never rewrites the record of why it was needed. +- `force_complete(run, result)` / `force_fail(run, reason)` — a nonterminal, drained run: finalizes immediately, tombstoning open slots, recording the - operator origin. Parent joins receive the arrival like any terminal path. + operator origin and (for completion) the result to treat it as having + produced. Parent joins receive the arrival like any terminal path. Refused + while a step is claimed, so it never races a working attempt — cancel + first if a worker still holds one. All of these are store transactions under the same atomicity rules as §1. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 27a1fce3bf2..d61338fa67b 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -26,6 +26,32 @@ from reflex.workflow.store import RunStore +def _operator_action(database: str | None, run_id: str, action: str, **extra): + """Apply one operator action to a run, reporting what happened. + + Args: + database: Connection URL or SQLite path, or None for the default. + run_id: The run to act on. + action: The store method to call. + extra: Extra keyword arguments for the store method. + + Raises: + Exit: When the run was not in a state the action allows. + """ + import time + + applied = _with_store( + database, + lambda store: getattr(store, action)(run_id, time.time(), **extra), + ) + if not applied: + console.error( + f"Run {run_id!r} is not in a state that allows {action.split('_')[0]!r}." + ) + raise click.exceptions.Exit(1) + console.print(f"Applied {action.split('_')[0]} to {run_id}.") + + def _open_store(database: str | None) -> RunStore: """Open the run store the app persists to. @@ -384,6 +410,18 @@ def cancel(database: str | None, run_id: str): console.print(f"Cancellation requested for {run_id}.") +@workflows.command() +@database_option +@click.argument("run_id") +def retry(database: str | None, run_id: str): + """Re-open a failed run at the step that failed. + + The step runs again with a fresh attempt budget; the original failure + stays in the run's history. + """ + _operator_action(database, run_id, "retry_run") + + @workflows.command() @database_option @click.argument("run_id") diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index b6bd9d423c0..042614ef01f 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -687,6 +687,56 @@ async def check_substeps_record_once_and_fence_stale_writers( assert list(await store.get_substeps("run1", 0)) == ["charge", "label"] +async def check_retry_reopens_only_failed_runs(store: RunStore) -> None: + """An operator retry applies to a failed run and nothing else.""" + await store.admit(make_run(), make_step(), _ADMITTED) + # A pending run is not a failure to retry. + assert not await store.retry_run("run1", NOW) + + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.FAILED, + run_status=RunStatus.FAILED, + state={}, + run_error={"reason": "boom"}, + ), + NOW, + ) + assert await store.retry_run("run1", NOW + 1) + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.PENDING + assert run.error is None + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.READY + assert steps[0].attempts == 0 + # Now re-opened, it is no longer a failed run. + assert not await store.retry_run("run1", NOW + 2) + assert not await store.retry_run("missing", NOW) + + +async def check_force_finalize_records_a_result(store: RunStore) -> None: + """An operator ending a run may record what it should be treated as.""" + await store.admit(make_run(), make_step(), _ADMITTED) + assert await store.finalize_run( + "run1", + status=RunStatus.COMPLETED, + error=None, + event=HistoryEventType.RUN_COMPLETED, + now=NOW, + result={"by": "operator"}, + ) + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.COMPLETED + assert run.result == {"by": "operator"} + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.CANCELLED + + async def check_schedule_cursors_persist(store: RunStore) -> None: """A schedule's catch-up position survives the process that wrote it.""" assert await store.read_schedule_cursor("wf:tick") is None @@ -775,6 +825,8 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_list_children_finds_a_joins_branches, check_claims_respect_queue_boundaries, check_substeps_record_once_and_fence_stale_writers, + check_retry_reopens_only_failed_runs, + check_force_finalize_records_a_result, check_schedule_cursors_persist, check_join_arrivals_count_once, check_finalize_refuses_while_a_step_is_claimed, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index f4f775201c8..da8ab157891 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -762,6 +762,69 @@ async def resume(self, run_id: str) -> bool: self._wakeup.set() return resumed + async def retry(self, run_id: str) -> bool: + """Re-open a failed run at the step that failed. + + Args: + run_id: The run to retry. + + Returns: + True if a failed run was re-opened. + """ + retried = await self._store.retry_run(run_id, self._clock()) + if retried: + await self._notify_run( + run_id, ((HistoryEventType.RUN_RESUMED, {"origin": "retry"}),) + ) + self._wakeup.set() + return retried + + async def force_finalize( + self, + run_id: str, + *, + status: RunStatus, + result: Any = None, + error: dict[str, Any] | None = None, + ) -> bool: + """End a run by operator decision, tombstoning what it had open. + + The escape hatch for a run no code path will finish: a wait nobody + will answer, a branch whose provider is gone. Refused while any step + is claimed, so it never races a working attempt -- cancel first if a + worker still holds it. + + Args: + run_id: The run to finalize. + status: The terminal status to record. + result: Result to record when completing. + error: Error payload to record when failing. + + Returns: + True if the run was finalized. + """ + now = self._clock() + run = await self._store.get_run(run_id) + if run is None: + return False + event = ( + HistoryEventType.RUN_COMPLETED + if status is RunStatus.COMPLETED + else HistoryEventType.RUN_FAILED + ) + finalized = await self._store.finalize_run( + run_id, + status=status, + error=error, + event=event, + now=now, + result=result, + ) + if finalized: + self._notify(run, ((event, {"origin": "operator"}),)) + await self._report_outcome(run, status, result, error) + return finalized + async def list_runs( self, *, diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index ca7ed98ee13..016d92e6ac6 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1343,6 +1343,7 @@ async def finalize_run( error: dict[str, Any] | None, event: HistoryEventType, now: float, + result: Any = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -1352,6 +1353,7 @@ async def finalize_run( error: Error payload recorded on the run. event: The terminal history event type. now: Current time in epoch seconds. + result: Result to record, for an operator forcing completion. Returns: True if the run was finalized. @@ -1383,9 +1385,10 @@ async def finalize_run( (StepStatus.CANCELLED.value, now, run_id, _TERMINAL_STEPS), ) await conn.execute( - "UPDATE workflow_runs SET status = %s, error = %s, updated_at = %s" + "UPDATE workflow_runs SET status = %s, error = %s," + " result = COALESCE(%s, result), updated_at = %s" " WHERE run_id = %s", - (status.value, _json(error), now, run_id), + (status.value, _json(error), _json(result), now, run_id), ) events: list[tuple[HistoryEventType, dict[str, Any]]] = [ (HistoryEventType.STEP_TOMBSTONED, {"ordinal": open_row["ordinal"]}) @@ -1431,6 +1434,46 @@ async def resume_run(self, run_id: str, now: float) -> bool: ) return True + async def retry_run(self, run_id: str, now: float) -> bool: + """Re-open a failed run at the step that failed. + + Args: + run_id: The run to retry. + now: Current time in epoch seconds. + + Returns: + True if a failed run was re-opened. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + await self._lock_run(conn, run_id) + cursor = await conn.execute( + "SELECT ordinal FROM workflow_steps WHERE run_id = %s" + " AND status = ANY(%s) ORDER BY ordinal LIMIT 1", + (run_id, [StepStatus.FAILED.value, StepStatus.TIMED_OUT.value]), + ) + row = await cursor.fetchone() + cursor = await conn.execute( + "UPDATE workflow_runs SET status = %s, error = NULL, updated_at = %s" + " WHERE run_id = %s AND status = %s", + (RunStatus.PENDING.value, now, run_id, RunStatus.FAILED.value), + ) + if row is None or cursor.rowcount == 0: + return False + await conn.execute( + "UPDATE workflow_steps SET status = %s, attempts = 0, due_at = %s," + " lease_expires_at = 0, error = NULL, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s", + (StepStatus.READY.value, now, now, run_id, row["ordinal"]), + ) + await self._append_events( + conn, + run_id, + ((HistoryEventType.RUN_RESUMED, {"origin": "retry"}),), + now, + ) + return True + async def recover_orphans( self, now: float, max_recoveries: int ) -> tuple[int, tuple[str, ...]]: diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 72e1c40afe5..10b9ab93d8f 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -315,6 +315,55 @@ async def resume(run_id: str) -> bool: """ return await get_runtime().kernel.resume(run_id) + @staticmethod + async def retry(run_id: str) -> bool: + """Re-open a failed run at the step that failed. + + Use this once the cause is fixed: the failed step runs again with a + fresh attempt budget, and the original failure stays in history. + + Args: + run_id: The run to retry. + + Returns: + True if a failed run was re-opened. + """ + return await get_runtime().kernel.retry(run_id) + + @staticmethod + async def force_complete(run_id: str, result: Any = None) -> bool: + """End a run as completed by operator decision. + + For a run no code path will finish -- a wait nobody will answer, a + provider that is gone. Refused while a step is claimed; cancel first + if a worker still holds it. + + Args: + run_id: The run to complete. + result: Result to record on the run. + + Returns: + True if the run was finalized. + """ + return await get_runtime().kernel.force_finalize( + run_id, status=RunStatus.COMPLETED, result=result + ) + + @staticmethod + async def force_fail(run_id: str, reason: str) -> bool: + """End a run as failed by operator decision. + + Args: + run_id: The run to fail. + reason: Why the operator gave up on it, recorded on the run. + + Returns: + True if the run was finalized. + """ + return await get_runtime().kernel.force_finalize( + run_id, status=RunStatus.FAILED, error={"reason": reason} + ) + @staticmethod async def list_runs( *, diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 7e974cad057..8eabff3a66a 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -423,6 +423,7 @@ async def finalize_run( error: dict[str, Any] | None, event: HistoryEventType, now: float, + result: Any = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -432,6 +433,7 @@ async def finalize_run( error: Error payload recorded on the run. event: The terminal history event type. now: Current time in epoch seconds. + result: Result to record, for an operator forcing completion. Returns: True if the run was finalized; False if it was already terminal @@ -439,6 +441,22 @@ async def finalize_run( """ ... + async def retry_run(self, run_id: str, now: float) -> bool: + """Re-open a failed run at the step that failed. + + The operator's answer to a run that failed for a reason now fixed: the + failed step runs again with a fresh attempt budget, and the failure + stays in history rather than being erased. + + Args: + run_id: The run to retry. + now: Current time in epoch seconds. + + Returns: + True if a failed run was re-opened. + """ + ... + async def resume_run(self, run_id: str, now: float) -> bool: """Re-open a suspended run so its frontier step runs again. @@ -1396,6 +1414,7 @@ async def finalize_run( error: dict[str, Any] | None, event: HistoryEventType, now: float, + result: Any = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -1405,6 +1424,7 @@ async def finalize_run( error: Error payload recorded on the run. event: The terminal history event type. now: Current time in epoch seconds. + result: Result to record, for an operator forcing completion. Returns: True if the run was finalized. @@ -1427,12 +1447,59 @@ async def finalize_run( {"ordinal": step.ordinal}, )) self._runs[run_id] = dataclasses.replace( - run, status=status, error=error, updated_at=now + run, + status=status, + error=error, + result=result if result is not None else run.result, + updated_at=now, ) events.append((event, {} if error is None else dict(error))) self._append_events(run_id, events, now) return True + async def retry_run(self, run_id: str, now: float) -> bool: + """Re-open a failed run at the step that failed. + + The operator's answer to a run that failed for a reason now fixed: the + failed step runs again with a fresh attempt budget, and the failure + stays in history rather than being erased. + + Args: + run_id: The run to retry. + now: Current time in epoch seconds. + + Returns: + True if a failed run was re-opened. + """ + async with self._lock: + run = self._runs.get(run_id) + if run is None or run.status is not RunStatus.FAILED: + return False + steps = self._steps[run_id] + reopened = False + for index, step in enumerate(steps): + if step.status in (StepStatus.FAILED, StepStatus.TIMED_OUT): + steps[index] = dataclasses.replace( + step, + status=StepStatus.READY, + attempts=0, + due_at=now, + lease_expires_at=0.0, + error=None, + updated_at=now, + ) + reopened = True + break + if not reopened: + return False + self._runs[run_id] = dataclasses.replace( + run, status=RunStatus.PENDING, error=None, updated_at=now + ) + self._append_events( + run_id, ((HistoryEventType.RUN_RESUMED, {"origin": "retry"}),), now + ) + return True + async def resume_run(self, run_id: str, now: float) -> bool: """Re-open a suspended run so its frontier step runs again. @@ -3115,6 +3182,7 @@ async def finalize_run( error: dict[str, Any] | None, event: HistoryEventType, now: float, + result: Any = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -3124,6 +3192,7 @@ async def finalize_run( error: Error payload recorded on the run. event: The terminal history event type. now: Current time in epoch seconds. + result: Result to record, for an operator forcing completion. Returns: True if the run was finalized. @@ -3166,9 +3235,10 @@ def work(): (StepStatus.CANCELLED.value, now, run_id, *terminal_step), ) self._db.execute( - "UPDATE workflow_runs SET status = ?, error = ?, updated_at = ?" + "UPDATE workflow_runs SET status = ?, error = ?," + " result = COALESCE(?, result), updated_at = ?" " WHERE run_id = ?", - (status.value, _dump(error), now, run_id), + (status.value, _dump(error), _dump(result), now, run_id), ) events: list[tuple[HistoryEventType, dict[str, Any]]] = [ (HistoryEventType.STEP_TOMBSTONED, {"ordinal": row["ordinal"]}) @@ -3184,6 +3254,62 @@ def work(): return await asyncio.to_thread(work) + async def retry_run(self, run_id: str, now: float) -> bool: + """Re-open a failed run at the step that failed. + + The operator's answer to a run that failed for a reason now fixed: the + failed step runs again with a fresh attempt budget, and the failure + stays in history rather than being erased. + + Args: + run_id: The run to retry. + now: Current time in epoch seconds. + + Returns: + True if a failed run was re-opened. + """ + + def work() -> bool: + """Re-open the failed step on the worker thread. + + Returns: + Whether a failed run was re-opened. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT ordinal FROM workflow_steps WHERE run_id = ?" + " AND status IN (?, ?) ORDER BY ordinal LIMIT 1", + (run_id, StepStatus.FAILED.value, StepStatus.TIMED_OUT.value), + ).fetchone() + cursor = self._db.execute( + "UPDATE workflow_runs SET status = ?, error = NULL," + " updated_at = ? WHERE run_id = ? AND status = ?", + (RunStatus.PENDING.value, now, run_id, RunStatus.FAILED.value), + ) + if row is None or cursor.rowcount == 0: + self._db.execute("ROLLBACK") + return False + self._db.execute( + "UPDATE workflow_steps SET status = ?, attempts = 0," + " due_at = ?, lease_expires_at = 0, error = NULL," + " updated_at = ? WHERE run_id = ? AND ordinal = ?", + (StepStatus.READY.value, now, now, run_id, row["ordinal"]), + ) + self._append_events( + run_id, + ((HistoryEventType.RUN_RESUMED, {"origin": "retry"}),), + now, + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True + + return await asyncio.to_thread(work) + async def resume_run(self, run_id: str, now: float) -> bool: """Re-open a suspended run so its frontier step runs again. diff --git a/tests/units/workflow/test_operator_actions.py b/tests/units/workflow/test_operator_actions.py new file mode 100644 index 00000000000..6e3b7141a76 --- /dev/null +++ b/tests/units/workflow/test_operator_actions.py @@ -0,0 +1,179 @@ +"""Tests for the operator actions in CONTRACT.md section 9. + +Each action is legal only from the states the contract names; anything else +is a refused no-op. These assert both halves, because an action that silently +succeeds from the wrong state is how an operator corrupts a run they were +trying to rescue. +""" + +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import RunStatus, StepStatus +from reflex.workflow.testing import WorkflowTestHarness + +ATTEMPTS: list[int] = [] + + +class Fragile(rx.State): + """Fails until told otherwise.""" + + __workflow__ = WorkflowConfig(id="ops.fragile") + healed: bool = False + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=1), + ) + def start(self): + """Fail while the world is broken. + + Returns: + Completion once healed. + + Raises: + TransientWorkflowError: Until the operator fixes things. + """ + ATTEMPTS.append(1) + if not HEALED: + msg = "downstream is down" + raise TransientWorkflowError(msg) + return rx.complete(result={"attempts": len(ATTEMPTS)}) + + +HEALED = False + + +async def test_retry_reopens_a_failed_run(forked_registration_context): + """A failure the operator has since fixed runs again, from that step.""" + global HEALED + ATTEMPTS.clear() + HEALED = False + async with WorkflowTestHarness(Fragile) as harness: + result = await harness.start(Fragile.start) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert len(ATTEMPTS) == 1 + + HEALED = True + assert await harness.kernel.retry(result.run_id) + await harness.run_until_idle() + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert len(ATTEMPTS) == 2 + # The failure is still on the record; a retry does not rewrite history. + history = await harness.kernel.store.get_history(result.run_id) + assert any(event.type.value == "run_failed" for event in history) + + +async def test_retry_is_refused_on_a_healthy_run(forked_registration_context): + """Retry applies to failures, not to runs that are fine.""" + global HEALED + ATTEMPTS.clear() + HEALED = True + async with WorkflowTestHarness(Fragile) as harness: + result = await harness.start(Fragile.start) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + + assert not await harness.kernel.retry(result.run_id) + assert not await harness.kernel.retry("no-such-run") + + +async def test_force_complete_ends_a_stuck_run(forked_registration_context): + """A wait nobody will ever answer can be ended by decision.""" + + class Waiting(rx.State): + __workflow__ = WorkflowConfig(id="ops.waiting") + + answered = rx.Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Wait forever. + + Returns: + An unbounded wait. + """ + return rx.wait_for(Waiting.answered, then=Waiting.done, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def done(self, payload: dict): + """Never reached in this test. + + Args: + payload: The delivered answer. + """ + + async with WorkflowTestHarness(Waiting) as harness: + result = await harness.start(Waiting.start) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + + assert await harness.kernel.force_finalize( + result.run_id, status=RunStatus.COMPLETED, result={"by": "operator"} + ) + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"by": "operator"} + # The abandoned wait is tombstoned, not left dangling. + assert all( + step.status in (StepStatus.SUCCEEDED, StepStatus.CANCELLED) + for step in snapshot.steps + ) + + +async def test_force_fail_records_the_reason(forked_registration_context): + """Giving up on a run says who gave up and why.""" + + class Stuck(rx.State): + __workflow__ = WorkflowConfig(id="ops.stuck") + + nudged = rx.Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Wait forever. + + Returns: + An unbounded wait. + """ + return rx.wait_for(Stuck.nudged, then=Stuck.done, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def done(self, payload: dict): + """Never reached. + + Args: + payload: The delivered answer. + """ + + async with WorkflowTestHarness(Stuck) as harness: + result = await harness.start(Stuck.start) + assert result.run_id is not None + assert await harness.kernel.force_finalize( + result.run_id, + status=RunStatus.FAILED, + error={"reason": "vendor retired the endpoint"}, + ) + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error is not None + assert "vendor retired" in snapshot.error["reason"] + + # Already terminal: a second decision changes nothing. + assert not await harness.kernel.force_finalize( + result.run_id, status=RunStatus.COMPLETED + ) From 073968e5e680baf18ea7548c23acd281ca61930a Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 23:43:58 -0700 Subject: [PATCH 040/121] Add a headless worker: reflex workflows worker Running workflows meant running a Reflex app, which is the wrong shape for the deployment that matters most: a background worker that serves no pages. It also made the engine look like it required a frontend, when all a worker needs is the definitions and the store. `reflex workflows worker ` loads workflow classes from any importable Python and runs the kernel against the configured store, with no app, no frontend build, and no HTTP server. --queue narrows what the process claims and --concurrency sizes it, so scaling out is starting more processes and dedicating some to slow work. Refusing to start on a module with no workflows is deliberate: a worker that silently serves nothing looks healthy forever. Verified end to end outside this repo's app machinery -- a plain module with no rx.App, work admitted by a separate script, drained by the real console script -- which is exactly the FastAPI/Django/cron shape the engine now supports. --- news/workflow-worker-cli.feature.md | 1 + reflex/workflow/cli.py | 77 ++++++++++++++++++++++++++ tests/units/workflow/test_cli_check.py | 44 +++++++++++++++ 3 files changed, 122 insertions(+) create mode 100644 news/workflow-worker-cli.feature.md diff --git a/news/workflow-worker-cli.feature.md b/news/workflow-worker-cli.feature.md new file mode 100644 index 00000000000..d960e972873 --- /dev/null +++ b/news/workflow-worker-cli.feature.md @@ -0,0 +1 @@ +`reflex workflows worker ` runs durable workflows as a background process with no frontend or web server, so a worker can serve workflows defined in any Python module — a FastAPI service, a Django project, or a bare script. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index d61338fa67b..03d7b102373 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -134,6 +134,83 @@ def workflows(): """Inspect and steer durable workflow runs.""" +@workflows.command() +@database_option +@click.argument("target") +@click.option( + "--queue", + "queues", + multiple=True, + help="Serve only these queues. Repeatable; default serves every queue.", +) +@click.option("--concurrency", default=None, type=int, help="Attempts to run at once.") +def worker( + database: str | None, + target: str, + queues: tuple[str, ...], + concurrency: int | None, +): + """Run workflows from TARGET with no frontend and no web server. + + TARGET is a Python file or dotted module defining workflow classes. This + is the deployment shape for a background worker: a plain process that + claims steps from the shared store and executes them. Scale by starting + more of them, and narrow what a process takes with --queue. + + The workflows do not have to live in a Reflex app -- a module importable + from a FastAPI service, a Django project, or a bare script works, because + a worker needs only the definitions and the store. + """ + import asyncio + + from reflex.workflow.runtime import WorkflowRuntime + + try: + module = _load_module(target) + except Exception as err: + console.error(f"Could not load {target!r}: {err}") + raise click.exceptions.Exit(1) from None + + classes = [ + value + for value in vars(module).values() + if isinstance(value, type) and "__workflow__" in vars(value) + ] + if not classes: + console.error( + f"No workflow classes in {target!r}. A workflow is an rx.State " + "subclass with __workflow__ = rx.WorkflowConfig(id=...)." + ) + raise click.exceptions.Exit(1) + + async def serve() -> None: + """Run the kernel until interrupted.""" + from reflex.workflow.kernel import DEFAULT_MAX_CONCURRENCY + from reflex.workflow.store import resolve_store + + runtime = WorkflowRuntime( + resolve_store(database), + queues=queues or None, + max_concurrency=concurrency or DEFAULT_MAX_CONCURRENCY, + ) + for workflow_cls in classes: + runtime.register(workflow_cls) + served = ", ".join(sorted(d.workflow_id for d in runtime.definitions)) + console.print( + f"Serving {served} on " + f"{'queues ' + ', '.join(queues) if queues else 'every queue'}." + ) + async with runtime.running(): + # The kernel's worker does the work; this task only waits for the + # operator (or the platform) to stop the process. + await asyncio.Event().wait() + + try: + asyncio.run(serve()) + except KeyboardInterrupt: + console.print("Worker stopped.") + + @workflows.command("list") @database_option @click.option("--workflow", "-w", default=None, help="Only this workflow id.") diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py index 8a46294beb7..f78858b6245 100644 --- a/tests/units/workflow/test_cli_check.py +++ b/tests/units/workflow/test_cli_check.py @@ -155,3 +155,47 @@ def test_a_dotted_module_resolves_from_the_project_root( result = CliRunner().invoke(workflows, ["check", "myflows.orders"]) assert result.exit_code == 0, result.output assert "check.greeter" in result.output + + +WORKER_FLOW = ''' +import reflex as rx + +class Batch(rx.State): + __workflow__ = rx.WorkflowConfig(id="worker.batch") + label: str = "" + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def go(self, label: str): + """Record the label. + + Args: + label: What to record. + + Returns: + Completion. + """ + self.label = label + return rx.complete(result={"label": label}) +''' + + +def test_the_worker_refuses_a_module_with_no_workflows( + tmp_path, forked_registration_context +): + """Starting a worker that would serve nothing is an error, not a hang. + + A process that sits there having silently loaded zero workflows is the + worst failure mode for a background worker: it looks healthy forever. + """ + module = tmp_path / "empty.py" + module.write_text("x = 1\n") + result = CliRunner().invoke(workflows, ["worker", str(module)]) + assert result.exit_code == 1 + assert "No workflow classes" in result.output + + +def test_the_worker_refuses_an_unloadable_target(tmp_path, forked_registration_context): + """A bad path is named, not raised as a traceback.""" + result = CliRunner().invoke(workflows, ["worker", str(tmp_path / "nope.py")]) + assert result.exit_code == 1 + assert "Could not load" in result.output From a2309de466c621f10ff6eb3032db7b4922cd90ee Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 23:47:31 -0700 Subject: [PATCH 041/121] Add a metrics observer for the numbers a deployment alerts on The observer stream carried everything a monitor needs and nothing it could use directly: exporters want counters, not events, and every deployment would have written the same tally loop. MetricsObserver keeps the counts an on-call rotation actually pages on -- runs started and how they ended, attempts and how many were retries, recoveries, or abandonments -- in total and per workflow, and hands them over as plain data. Counters only increase and snapshot() returns a copy, which is what a scrape-and-diff collector expects; a Prometheus endpoint or an OpenTelemetry exporter is a few lines over it. No OpenTelemetry bridge ships here on purpose: the package is not available in this environment, and untested integration code is worse than none. --- news/workflow-metrics-observer.feature.md | 1 + reflex/workflow/__init__.py | 8 +- reflex/workflow/kernel.py | 76 ++++++++++++- tests/units/workflow/test_metrics.py | 131 ++++++++++++++++++++++ 4 files changed, 214 insertions(+), 2 deletions(-) create mode 100644 news/workflow-metrics-observer.feature.md create mode 100644 tests/units/workflow/test_metrics.py diff --git a/news/workflow-metrics-observer.feature.md b/news/workflow-metrics-observer.feature.md new file mode 100644 index 00000000000..fb47ce91113 --- /dev/null +++ b/news/workflow-metrics-observer.feature.md @@ -0,0 +1 @@ +`rx.workflow.MetricsObserver` tallies the run and attempt counters a deployment alerts on, in total and per workflow, so a metrics endpoint or OpenTelemetry exporter is a few lines over `snapshot()`. diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index f536c43d628..8cce91e46b5 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -48,7 +48,12 @@ WorkflowDefinition, compile_workflow, ) -from reflex.workflow.kernel import LoggingObserver, WorkflowKernel, WorkflowObserver +from reflex.workflow.kernel import ( + LoggingObserver, + MetricsObserver, + WorkflowKernel, + WorkflowObserver, +) from reflex.workflow.records import ( HistoryEvent, HistoryEventType, @@ -84,6 +89,7 @@ "LoggingObserver", "ManualTrigger", "MemoryRunStore", + "MetricsObserver", "Parallel", "RateLimit", "Retry", diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index da8ab157891..dff5afc72f2 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -17,7 +17,7 @@ import time import traceback import uuid -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final from pydantic import TypeAdapter from reflex_base.event.processor.base_state_processor import _transform_event_payload @@ -140,6 +140,80 @@ def on_event( """ +class MetricsObserver(WorkflowObserver): + """Tallies the numbers a deployment alerts on. + + The observer stream is the raw material for monitoring, but every exporter + wants counters rather than events. This keeps the counts an operator + actually pages on -- runs started and how they ended, attempts and how + many were retries or recoveries -- both in total and per workflow, so a + metrics endpoint or an OpenTelemetry exporter is a few lines over + ``snapshot()`` rather than an event-stream parser. + + Counters only ever increase, which is what a scrape-and-diff collector + expects. It is safe to install alongside another observer by composing + them; a raising observer never affects a run. + """ + + _COUNTED: Final = { + HistoryEventType.RUN_ADMITTED: "runs_started", + HistoryEventType.RUN_COMPLETED: "runs_completed", + HistoryEventType.RUN_FAILED: "runs_failed", + HistoryEventType.RUN_CANCELLED: "runs_cancelled", + HistoryEventType.RUN_TIMED_OUT: "runs_timed_out", + HistoryEventType.RUN_NEEDS_ATTENTION: "runs_needing_attention", + HistoryEventType.ATTEMPT_STARTED: "attempts", + HistoryEventType.ATTEMPT_FAILED: "attempts_failed", + HistoryEventType.ATTEMPT_TIMED_OUT: "attempts_timed_out", + HistoryEventType.ATTEMPT_ABANDONED: "attempts_abandoned", + HistoryEventType.STEP_RETRY_SCHEDULED: "retries_scheduled", + HistoryEventType.STEP_RECOVERED: "steps_recovered", + HistoryEventType.SUBSTEP_RECORDED: "substeps_recorded", + } + + def __init__(self): + """Start every counter at zero.""" + self.totals: dict[str, int] = {} + self.by_workflow: dict[str, dict[str, int]] = {} + + def on_event( + self, + event_type: HistoryEventType, + run_id: str, + workflow_id: str, + data: dict[str, Any], + ) -> None: + """Count one transition. + + Args: + event_type: What happened. + run_id: The run it happened to. + workflow_id: That run's workflow identity. + data: Event payload. + """ + metric = self._COUNTED.get(event_type) + if metric is None: + return + self.totals[metric] = self.totals.get(metric, 0) + 1 + if workflow_id: + counts = self.by_workflow.setdefault(workflow_id, {}) + counts[metric] = counts.get(metric, 0) + 1 + + def snapshot(self) -> dict[str, Any]: + """Read the counters as plain data. + + Returns: + The totals, and the same counters broken down by workflow. + """ + return { + "totals": dict(self.totals), + "by_workflow": { + workflow_id: dict(counts) + for workflow_id, counts in self.by_workflow.items() + }, + } + + class LoggingObserver(WorkflowObserver): """Logs every transition as a structured line.""" diff --git a/tests/units/workflow/test_metrics.py b/tests/units/workflow/test_metrics.py new file mode 100644 index 00000000000..0ab87cb31a5 --- /dev/null +++ b/tests/units/workflow/test_metrics.py @@ -0,0 +1,131 @@ +"""Tests for the metrics observer. + +An exporter wants counters, not an event stream. These assert the numbers an +operator would actually alert on -- and that they are counters, not gauges +that reset. +""" + +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.kernel import MetricsObserver +from reflex.workflow.testing import WorkflowTestHarness + +CALLS: list[int] = [] + + +class Flaky(rx.State): + """Fails once, then succeeds.""" + + __workflow__ = WorkflowConfig(id="metrics.flaky") + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=3, initial_delay="1s", jitter="none"), + ) + def go(self): + """Fail the first time. + + Returns: + Completion on the second attempt. + + Raises: + TransientWorkflowError: On the first attempt. + """ + CALLS.append(1) + if len(CALLS) == 1: + msg = "first attempt fails" + raise TransientWorkflowError(msg) + return rx.complete(result={"attempts": len(CALLS)}) + + +async def test_counters_cover_a_run_with_a_retry(forked_registration_context): + """One admitted run, two attempts, one retry, one completion.""" + CALLS.clear() + metrics = MetricsObserver() + async with WorkflowTestHarness(Flaky, observer=metrics) as harness: + await harness.start(Flaky.go) + await harness.advance("2s") + + totals = metrics.snapshot()["totals"] + assert totals["runs_started"] == 1 + assert totals["runs_completed"] == 1 + assert totals["attempts"] == 2 + assert totals["attempts_failed"] == 1 + assert totals["retries_scheduled"] == 1 + assert "runs_failed" not in totals + + +async def test_counters_break_down_by_workflow(forked_registration_context): + """A deployment alerts per workflow, not only in aggregate.""" + CALLS.clear() + metrics = MetricsObserver() + + class Simple(rx.State): + __workflow__ = WorkflowConfig(id="metrics.simple") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Succeed immediately.""" + + async with WorkflowTestHarness(Flaky, Simple, observer=metrics) as harness: + await harness.start(Simple.go) + await harness.start(Flaky.go) + await harness.advance("2s") + + by_workflow = metrics.snapshot()["by_workflow"] + assert by_workflow["metrics.simple"]["runs_started"] == 1 + assert by_workflow["metrics.simple"]["attempts"] == 1 + assert by_workflow["metrics.flaky"]["attempts"] == 2 + assert "attempts_failed" not in by_workflow["metrics.simple"] + + +async def test_counters_only_increase(forked_registration_context): + """A scrape-and-diff collector needs monotonic counters.""" + CALLS.clear() + metrics = MetricsObserver() + async with WorkflowTestHarness(Flaky, observer=metrics) as harness: + await harness.start(Flaky.go) + await harness.advance("2s") + first = metrics.snapshot()["totals"]["attempts"] + + CALLS.clear() + await harness.start(Flaky.go) + await harness.advance("2s") + second = metrics.snapshot()["totals"]["attempts"] + + assert second > first + # The snapshot is a copy: reading it cannot disturb the counters. + snapshot = metrics.snapshot() + snapshot["totals"]["attempts"] = 0 + assert metrics.snapshot()["totals"]["attempts"] == second + + +async def test_a_failed_run_is_counted_as_failed(forked_registration_context): + """The counter an on-call rotation actually pages on.""" + + class Doomed(rx.State): + __workflow__ = WorkflowConfig(id="metrics.doomed") + + @rx.event( + durable=True, trigger=manual(), effect="read", retry=Retry(max_attempts=1) + ) + def go(self): + """Always fail. + + Raises: + TransientWorkflowError: Always. + """ + msg = "broken" + raise TransientWorkflowError(msg) + + metrics = MetricsObserver() + async with WorkflowTestHarness(Doomed, observer=metrics) as harness: + await harness.start(Doomed.go) + + totals = metrics.snapshot()["totals"] + assert totals["runs_failed"] == 1 + assert totals["attempts_failed"] == 1 + assert "runs_completed" not in totals From ef2113eff7a399745a5e7d20f4f5ba806dcf49d2 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 23:50:22 -0700 Subject: [PATCH 042/121] Never let a schedule cursor move backwards Auditing the durability work turned up a multi-worker wrinkle: two workers sweeping the same schedule out of order -- ordinary clock skew, or one sweep simply starting earlier and finishing later -- could write an older position over a newer one. Nothing breaks (a re-scanned occurrence dedupes on its request key, so no run is ever duplicated), but the work is pure waste and a cursor that can rewind is not a position anyone can reason about when debugging a schedule. Writes now keep the later of the two in all three stores, with a conformance check so a future store cannot regress it. --- reflex/workflow/conformance.py | 3 +++ reflex/workflow/postgres.py | 3 ++- reflex/workflow/store.py | 12 ++++++++++-- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 042614ef01f..0a86bc73f05 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -745,6 +745,9 @@ async def check_schedule_cursors_persist(store: RunStore) -> None: # Advancing overwrites rather than accumulating. await store.write_schedule_cursor("wf:tick", NOW + 60) assert await store.read_schedule_cursor("wf:tick") == pytest.approx(NOW + 60) + # An out-of-order write from a second worker never rewinds it. + await store.write_schedule_cursor("wf:tick", NOW) + assert await store.read_schedule_cursor("wf:tick") == pytest.approx(NOW + 60) # Schedules are independent of one another. assert await store.read_schedule_cursor("wf:other") is None diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 016d92e6ac6..5f3484f01fe 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1789,7 +1789,8 @@ async def write_schedule_cursor(self, key: str, at: float) -> None: async with pool.connection() as conn, conn.transaction(): await conn.execute( "INSERT INTO workflow_schedules (key, at) VALUES (%s, %s)" - " ON CONFLICT (key) DO UPDATE SET at = EXCLUDED.at", + " ON CONFLICT (key) DO UPDATE SET" + " at = GREATEST(workflow_schedules.at, EXCLUDED.at)", (key, at), ) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 8eabff3a66a..c571d4099df 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -1773,7 +1773,14 @@ async def write_schedule_cursor(self, key: str, at: float) -> None: at: The time swept up to, in epoch seconds. """ async with self._lock: - self._schedule_cursors[key] = at + # Never move backwards: two workers sweeping out of order would + # otherwise rewind the cursor and re-scan ground already covered. + # Re-scanning is harmless (occurrences dedupe on their request + # key) but it is pure waste, and a cursor that can go back is not + # a position anyone can reason about. + self._schedule_cursors[key] = max( + at, self._schedule_cursors.get(key, at) + ) async def next_due( self, now: float, *, queues: tuple[str, ...] | None = None @@ -3792,7 +3799,8 @@ def work() -> None: try: self._db.execute( "INSERT INTO workflow_schedules (key, at) VALUES (?, ?)" - " ON CONFLICT(key) DO UPDATE SET at = excluded.at", + " ON CONFLICT(key) DO UPDATE SET" + " at = MAX(at, excluded.at)", (key, at), ) self._db.execute("COMMIT") From 827f9a9ba99447f268cd43bd637095b9bc6c694c Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Tue, 18 Aug 2026 23:53:19 -0700 Subject: [PATCH 043/121] Tell a worker's operator why their workflows will not load Auditing the new worker command with a deliberately invalid workflow -- an inline sibling call, the mistake the compiler exists to catch -- showed it dying with a raw traceback out of click's entry point. Startup is exactly when a deployment learns its code is wrong, and a traceback there names the compiler while the compiler's own message names the fix. The worker now reports that message and exits, refusing to serve a half-registered set: running the workflows that happened to compile while the rest vanish silently is the worse failure, because the process then looks healthy. The same audit confirmed the metrics observer tolerates hostile input (unknown or malformed event types, empty workflow ids) without raising -- an observer that can break a run is worse than no metrics at all. --- reflex/workflow/cli.py | 14 ++++++++++++-- reflex/workflow/store.py | 4 +--- tests/units/workflow/test_cli_check.py | 14 ++++++++++++++ workflow.db | Bin 0 -> 77824 bytes 4 files changed, 27 insertions(+), 5 deletions(-) create mode 100644 workflow.db diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 03d7b102373..c49baee77e1 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -163,6 +163,8 @@ def worker( """ import asyncio + from reflex_base.utils.exceptions import WorkflowDefinitionError + from reflex.workflow.runtime import WorkflowRuntime try: @@ -193,8 +195,16 @@ async def serve() -> None: queues=queues or None, max_concurrency=concurrency or DEFAULT_MAX_CONCURRENCY, ) - for workflow_cls in classes: - runtime.register(workflow_cls) + try: + for workflow_cls in classes: + runtime.register(workflow_cls) + except WorkflowDefinitionError as err: + # The compiler's message names the fix; a traceback out of a + # worker's startup names only the compiler. Refusing to start at + # all beats serving a half-registered set, where the workflows + # that did compile run and the rest vanish silently. + console.error(f"Cannot serve {target!r}: {err}") + raise click.exceptions.Exit(1) from None served = ", ".join(sorted(d.workflow_id for d in runtime.definitions)) console.print( f"Serving {served} on " diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index c571d4099df..8774349444a 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -1778,9 +1778,7 @@ async def write_schedule_cursor(self, key: str, at: float) -> None: # Re-scanning is harmless (occurrences dedupe on their request # key) but it is pure waste, and a cursor that can go back is not # a position anyone can reason about. - self._schedule_cursors[key] = max( - at, self._schedule_cursors.get(key, at) - ) + self._schedule_cursors[key] = max(at, self._schedule_cursors.get(key, at)) async def next_due( self, now: float, *, queues: tuple[str, ...] | None = None diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py index f78858b6245..c5b41a996e8 100644 --- a/tests/units/workflow/test_cli_check.py +++ b/tests/units/workflow/test_cli_check.py @@ -199,3 +199,17 @@ def test_the_worker_refuses_an_unloadable_target(tmp_path, forked_registration_c result = CliRunner().invoke(workflows, ["worker", str(tmp_path / "nope.py")]) assert result.exit_code == 1 assert "Could not load" in result.output + + +def test_the_worker_names_a_compile_error(tmp_path, forked_registration_context): + """A workflow that does not compile stops the worker with the reason. + + Starting a worker is the moment a deployment finds out its code is wrong. + A traceback there names the compiler; the compiler's own message names the + fix, which is what the operator reading the logs needs. + """ + module = tmp_path / "broken_worker.py" + module.write_text(BROKEN) + result = CliRunner().invoke(workflows, ["worker", str(module)]) + assert result.exit_code == 1 + assert "runs it inline" in result.output diff --git a/workflow.db b/workflow.db new file mode 100644 index 0000000000000000000000000000000000000000..fbd1be8eb72ff865519ebedb32fdfbce19eac82b GIT binary patch literal 77824 zcmeI%&raM%9KiA2g{)x-fsjf*^x{*&f`p=~QlmC%X|j!?5@-mh$;mS70o*w2#r7s} zs+Oiw6;lu4I^O@hw*qWzLw(LM?hpzY147AT0 zi;Y&R@r|Z68jY6x-jd(+Gbevcrw8(XGxNLZXRC4N?ML$q|2CFde>UcSU-+x@$NbK$ z=xopI&wir3omp%@P##YI(^jURPyIdhl{|?L0#^x)W|kE7-g0wv*X~>5)V5BI*RJ>C z&~aZI-mvc*1Je`zAeVOQ{%&`3ud6-Ue$aiURb8U(Y-_p8w3WE5uw28o*0ki%^(?z@ zI{N6^O<5JaIoh~f6+aLI-*ANKi*cDz+ER60GS}ql5~XV^eqaVee@%-DHN@$__JnVk zfj(MVRMbsbkB67*VfPQ*Q)3|dvdTx}nnbBLt80?Gs-wwI)+Wl;R#KJMrX3hB#Mzpb z*3cIx`pH7vO>>v)6m?U4?zl!_+G2HuGMD745~imz`uP|0imL0)H>-hp;9R&z^&Q_S zI=YeEYH#z~t?osyCoPLwEQ^+%uI+W7?P=RP^1HvawHD>14KB;Ggf$$9kfSVD!pI=p z>ptx67Rnr$XO3%@-Y%9ZE3W4WGZ2=H0Zj(@R-worb{{?7+{tEd|rayBU__Ksyt zIl{kSOj9})j%j`ZIR?o@L>$S*!0G7sJ8kvHjmZs3Z;@?Dnl?AuRupx0wR!$D-;_Ay z#_fm?XUZE9XC^HeUs66GJ#jJ=el!V{4U_7ToB`2RGB;V%GDXVny*JZVpR7()Nru$L z_=po`M~BxG_4e)N`OkTeW7`j0?`(YI5wWoa=lY64)z7n1-CZB2zgj$skb(dI7S;5DP_CoX@3f9aL7EBw==j>3J55Dh zUT*&OaX$FNhn(>o5gxS4LqE(anH@6YHE#d7tn%4ZV%dE=u-(34*++6-E#s~p1Yu60 zBjVde8HNmRUlx5SOyYC;;xsU>Uf$(Huk6Z9OcY}zdKF05ytvrhbA0B~`~!Vjr}&!8 z3Xr+NF%N_j5AyglOPH2p_eC^n;-sG0?+M393V4Sof4_k2=yzLf_3`q=={AguOt6KA zQ}PvSCVb7B5C1HNfA}DP00IagfB*srAb;OKmY**5I_I{1Q0*~ z0R#|0p#B2v|JT2dX(0j#AbU;jR) zg$N*k00IagfB*srAb Date: Tue, 18 Aug 2026 23:56:26 -0700 Subject: [PATCH 044/121] State the one rule rx.step puts on your code The contract described what rx.step guarantees but not what it asks for. Keys are the step name plus its occurrence, so a re-execution lines its calls up against the journal positionally -- which quietly requires a handler's sequence of rx.step calls to be the same on every attempt. Branching on the payload or on committed run state is fine; branching on a clock reading, a random draw, or a live API answer taken outside a step can shift the sequence, and a shifted sequence replays one call's result into a different call. An unstated constraint is one users discover through a bug during a retry, which is the worst place to learn it. Section 2 now says so, with the remedy: record the deciding value as its own step first, so the branch built on it is identical on every attempt. A test exercises that shape across a real retry -- a nondeterministic draw recorded as a step, a payment inside the branch it decides -- and asserts the draw happened once and the payment did not repeat. Worth keeping in proportion: this is far narrower than a replay engine's determinism rule. It governs the order of rx.step calls inside one handler, not the handler body, and not the workflow. --- reflex/workflow/CONTRACT.md | 16 ++++++ tests/units/workflow/test_steps.py | 79 ++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 40b9b8a64ba..c22eac026d9 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -62,6 +62,22 @@ that turn it into effectively-once for side effects are, in order of strength: - `rx.step(name, fn, ...)`: records the result durably at return; every later execution of the same handler replays the recorded value instead of calling `fn`. Fenced by claim epoch, so a zombie worker cannot write (§7). + + Its one requirement on your code: **a handler's sequence of `rx.step` calls + must be the same on every attempt.** Keys are the step name plus its + occurrence (`send`, `send#2`, …), so a re-execution lines its calls up + against the journal positionally. Branching on the payload, on run state, or + on anything else already durable is fine — that decides the same way twice. + Branching on something that can differ between attempts (a clock reading, a + random draw, a live API's answer taken *outside* a step) can shift the + sequence, and a shifted sequence replays one call's recorded result into a + different call. Put anything nondeterministic that later steps depend on + inside its own `rx.step` first: recorded, it is the same on every attempt, + and the branch built on it is stable. + + This is far narrower than a replay engine's determinism rule, which governs + the whole handler body. Here it governs only the order of `rx.step` calls, + and only within one handler. - `rx.current_run().idempotency_key()`: stable across retries and recoveries of one step, distinct across steps — hand it to providers that accept idempotency keys. This covers the one window `rx.step` cannot: a crash diff --git a/tests/units/workflow/test_steps.py b/tests/units/workflow/test_steps.py index 6858d50495d..82935f98e09 100644 --- a/tests/units/workflow/test_steps.py +++ b/tests/units/workflow/test_steps.py @@ -373,3 +373,82 @@ def go(self): assert snapshot is not None steps = await harness.kernel.store.get_steps(result.run_id) assert "async" in str(steps[0].error) + + +async def test_a_step_can_carry_the_decision_a_later_branch_uses( + forked_registration_context, +): + """Recording a nondeterministic value first keeps the sequence stable. + + rx.step lines calls up by occurrence, so a handler whose step sequence + can differ between attempts would replay one call's result into another. + The documented remedy is to record the deciding value as its own step: + once recorded it is identical on every attempt, so the branch built on it + is too. This exercises exactly that shape across a real retry. + """ + CALLS.clear() + draws: list[int] = [] + + def draw() -> int: + """Produce a value that differs on every call. + + Returns: + A fresh number each time. + """ + draws.append(len(draws) + 1) + return draws[-1] + + def paid(amount: int) -> dict: + """Record a payment. + + Args: + amount: What was charged. + + Returns: + The receipt. + """ + CALLS.append("paid") + return {"amount": amount} + + class Branching(rx.State): + __workflow__ = WorkflowConfig(id="steps.branching") + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=3, initial_delay="1s", jitter="none"), + ) + async def go(self): + """Branch on a recorded draw, then fail once after it. + + Returns: + Completion on the second attempt. + + Raises: + TransientWorkflowError: On the first attempt. + """ + # Recorded first: the branch below decides the same way on every + # attempt even though draw() itself never repeats a value. + roll = await rx.step("roll", draw) + if roll % 2 == 1: + await rx.step("charge", paid, 100) + CALLS.append("attempt") + if CALLS.count("attempt") == 1: + msg = "fails after the branch" + raise TransientWorkflowError(msg) + return rx.complete(result={"roll": roll}) + + async with WorkflowTestHarness(Branching) as harness: + result = await harness.start(Branching.go) + assert result.run_id is not None + await harness.advance("2s") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"roll": 1} + + # draw() ran once; the retry replayed it, so the branch held and the + # payment inside it did not repeat. + assert draws == [1] + assert CALLS.count("paid") == 1 From eb8dd6364a31495fea83ef34753001fc5139e629 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:00:46 -0700 Subject: [PATCH 045/121] Make every terminal path deliver its arrival atomically Auditing the contract against the code found a claim I had written and only half implemented. Section 5 says every way a child ends -- commit, cancellation, run timeout, recovery-budget exhaustion -- delivers its arrival to the parent's join inside the same transaction. Only the commit path actually did. The other three finalized the run and then delivered afterwards, which is precisely the window the commit-path fix existed to close: a worker dying between them leaves a run that is over and a join that waits on it forever, with nothing in the store to indicate anything is wrong. finalize_run now carries the arrival, so cancellation and run timeout close the same way. Recovery-budget exhaustion delivers inside the sweep's own transaction, where the run is failed. All three stores, conformance-checked. The kernel still reports afterwards on these paths. That is now redundant rather than load-bearing -- a repeat arrival is refused as a duplicate -- and it is kept because it is also what wakes the parent's worker and resolves a race's losers. The lesson worth keeping: a frozen contract is only worth what it is checked against. This gap survived because the prose was written from intent and nothing forced the non-commit paths to prove it. --- reflex/workflow/conformance.py | 52 ++++++++++++++++++++++++++++++++++ reflex/workflow/kernel.py | 40 +++++++++++++++++++++++++- reflex/workflow/postgres.py | 25 ++++++++++++++++ reflex/workflow/store.py | 50 ++++++++++++++++++++++++++++++++ 4 files changed, 166 insertions(+), 1 deletion(-) diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 0a86bc73f05..c32344d07ba 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -687,6 +687,57 @@ async def check_substeps_record_once_and_fence_stale_writers( assert list(await store.get_substeps("run1", 0)) == ["charge", "label"] +async def check_finalize_delivers_a_childs_arrival(store: RunStore) -> None: + """Ending a child by cancellation or timeout tells its parent, atomically. + + The commit path is not the only way a child ends. If cancellation and + run-timeout delivered their arrival afterwards, a crash in that window + would strand the join exactly as it would have on commit. + """ + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + child = make_run("child1", parent_run_id="run1", parent_ordinal=1) + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.WAITING, + state={}, + new_steps=( + make_step( + ordinal=1, + status=StepStatus.BLOCKED, + wait_key="join:1", + join_expected=1, + origin="join", + due_at=0.0, + ), + ), + next_ordinal=2, + children=((child, make_step("child1")),), + ), + NOW, + ) + + assert await store.finalize_run( + "child1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW + 1, + parent_arrival=( + "run1", + 1, + {"run_id": "child1", "status": "CANCELLED", "result": None, "error": None}, + "child1", + ), + ) + steps = await store.get_steps("run1") + assert steps[1].join_arrived == 1 + assert steps[1].status is StepStatus.READY + + async def check_retry_reopens_only_failed_runs(store: RunStore) -> None: """An operator retry applies to a failed run and nothing else.""" await store.admit(make_run(), make_step(), _ADMITTED) @@ -828,6 +879,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_list_children_finds_a_joins_branches, check_claims_respect_queue_boundaries, check_substeps_record_once_and_fence_stale_writers, + check_finalize_delivers_a_childs_arrival, check_retry_reopens_only_failed_runs, check_force_finalize_records_a_result, check_schedule_cursors_persist, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index dff5afc72f2..f4b00199f5e 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -893,6 +893,7 @@ async def force_finalize( event=event, now=now, result=result, + parent_arrival=self._arrival_for(run, status, result, error), ) if finalized: self._notify(run, ((event, {"origin": "operator"}),)) @@ -2205,6 +2206,38 @@ async def _commit_outcome( self._wakeup.set() await self._report_to_parent(claim.run, completion) + @staticmethod + def _arrival_for( + run: RunRecord, + status: RunStatus, + result: Any = None, + error: dict[str, Any] | None = None, + ) -> tuple[str, int, dict[str, Any], str] | None: + """Build the arrival a terminating run owes its parent, if any. + + Args: + run: The run that is ending. + status: Its terminal status. + result: Its result, if any. + error: Its error, if any. + + Returns: + The arrival tuple, or None when the run has no parent. + """ + if run.parent_run_id is None or run.parent_ordinal is None: + return None + return ( + run.parent_run_id, + run.parent_ordinal, + { + "run_id": run.run_id, + "status": status.value, + "result": result, + "error": error, + }, + run.run_id, + ) + @staticmethod def _with_parent_arrival( run: RunRecord, completion: StepCompletion @@ -2369,7 +2402,12 @@ async def _finalize_control(self, now: float) -> int: else HistoryEventType.RUN_TIMED_OUT ) if await self._store.finalize_run( - run.run_id, status=status, error=error, event=event, now=now + run.run_id, + status=status, + error=error, + event=event, + now=now, + parent_arrival=self._arrival_for(run, status, None, error), ): self._notify(run, ((event, {} if error is None else dict(error)),)) await self._report_outcome(run, status, None, error) diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 5f3484f01fe..95553ddfe35 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1344,6 +1344,7 @@ async def finalize_run( event: HistoryEventType, now: float, result: Any = None, + parent_arrival: tuple[str, int, dict[str, Any], str] | None = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -1354,6 +1355,8 @@ async def finalize_run( event: The terminal history event type. now: Current time in epoch seconds. result: Result to record, for an operator forcing completion. + parent_arrival: When this run is a child, the arrival to deliver + to its parent's join, applied in this same transaction. Returns: True if the run was finalized. @@ -1396,6 +1399,8 @@ async def finalize_run( ] events.append((event, {} if error is None else dict(error))) await self._append_events(conn, run_id, events, now) + if parent_arrival is not None: + await self._apply_arrival(conn, *parent_arrival, now) return True async def resume_run(self, run_id: str, now: float) -> bool: @@ -1528,6 +1533,26 @@ async def recover_orphans( ), ) failed.append(step.run_id) + cursor = await conn.execute( + "SELECT parent_run_id, parent_ordinal FROM workflow_runs" + " WHERE run_id = %s", + (step.run_id,), + ) + parent = await cursor.fetchone() + if parent is not None and parent["parent_run_id"] is not None: + await self._apply_arrival( + conn, + parent["parent_run_id"], + parent["parent_ordinal"], + { + "run_id": step.run_id, + "status": RunStatus.FAILED.value, + "result": None, + "error": dict(exhausted), + }, + step.run_id, + now, + ) await self._append_events( conn, step.run_id, diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 8774349444a..1df3ee6386e 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -424,6 +424,7 @@ async def finalize_run( event: HistoryEventType, now: float, result: Any = None, + parent_arrival: tuple[str, int, dict[str, Any], str] | None = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -434,6 +435,8 @@ async def finalize_run( event: The terminal history event type. now: Current time in epoch seconds. result: Result to record, for an operator forcing completion. + parent_arrival: When this run is a child, the arrival to deliver + to its parent's join, applied in this same transaction. Returns: True if the run was finalized; False if it was already terminal @@ -1415,6 +1418,7 @@ async def finalize_run( event: HistoryEventType, now: float, result: Any = None, + parent_arrival: tuple[str, int, dict[str, Any], str] | None = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -1425,6 +1429,8 @@ async def finalize_run( event: The terminal history event type. now: Current time in epoch seconds. result: Result to record, for an operator forcing completion. + parent_arrival: When this run is a child, the arrival to deliver + to its parent's join, applied in this same transaction. Returns: True if the run was finalized. @@ -1455,6 +1461,8 @@ async def finalize_run( ) events.append((event, {} if error is None else dict(error))) self._append_events(run_id, events, now) + if parent_arrival is not None: + self._apply_arrival(*parent_arrival, now) return True async def retry_run(self, run_id: str, now: float) -> bool: @@ -1571,6 +1579,21 @@ async def recover_orphans( updated_at=now, ) failed.append(run.run_id) + if run.parent_run_id is not None and ( + run.parent_ordinal is not None + ): + self._apply_arrival( + run.parent_run_id, + run.parent_ordinal, + { + "run_id": run.run_id, + "status": RunStatus.FAILED.value, + "result": None, + "error": {"reason": "recovery_budget_exhausted"}, + }, + run.run_id, + now, + ) self._append_events( run.run_id, ( @@ -3188,6 +3211,7 @@ async def finalize_run( event: HistoryEventType, now: float, result: Any = None, + parent_arrival: tuple[str, int, dict[str, Any], str] | None = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -3198,6 +3222,8 @@ async def finalize_run( event: The terminal history event type. now: Current time in epoch seconds. result: Result to record, for an operator forcing completion. + parent_arrival: When this run is a child, the arrival to deliver + to its parent's join, applied in this same transaction. Returns: True if the run was finalized. @@ -3251,6 +3277,8 @@ def work(): ] events.append((event, {} if error is None else dict(error))) self._append_events(run_id, events, now) + if parent_arrival is not None: + self._apply_arrival_sql(*parent_arrival, now) self._db.execute("COMMIT") except BaseException: self._db.execute("ROLLBACK") @@ -3431,6 +3459,28 @@ def work(): ), ) failed.append(step.run_id) + parent = self._db.execute( + "SELECT parent_run_id, parent_ordinal FROM" + " workflow_runs WHERE run_id = ?", + (step.run_id,), + ).fetchone() + if parent is not None and ( + parent["parent_run_id"] is not None + ): + self._apply_arrival_sql( + parent["parent_run_id"], + parent["parent_ordinal"], + { + "run_id": step.run_id, + "status": RunStatus.FAILED.value, + "result": None, + "error": { + "reason": "recovery_budget_exhausted" + }, + }, + step.run_id, + now, + ) self._append_events( step.run_id, ( From 3ed4b28689d6965dc2d55891e3a6ef98469c6133 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:04:09 -0700 Subject: [PATCH 046/121] Pin two more contract clauses with conformance checks Last commit's lesson was that a frozen contract is worth only what it is checked against, so this pass probed the rest of section 5, 6 and 7 the same mechanical way: lease respect, expired waits refusing late signals, concurrent same-key admissions yielding one run, epoch-fenced substep writes, one open obligation per run, and control refused on a terminal run. All of them hold today. Two were held by nothing but the implementation, so they are now conformance checks, which is where a claim about store behavior belongs: recovery reclaims a lapsed lease and never a live one, and finalizing waits for the run to drain then closes it to further control. A store that regressed either would previously have passed the suite. One probe of my own was wrong before the code was: it finalized a run whose step was still claimed, read the refusal as the run being terminal, and reported a false positive on cancel-after-terminal. Re-running it correctly showed the contract honored at every step. Worth recording -- a probe that disagrees with the contract is as likely to be a bad probe as a bug. --- reflex/workflow/conformance.py | 47 ++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index c32344d07ba..89fecdec605 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -687,6 +687,51 @@ async def check_substeps_record_once_and_fence_stale_writers( assert list(await store.get_substeps("run1", 0)) == ["charge", "label"] +async def check_recovery_respects_a_live_lease(store: RunStore) -> None: + """Recovery reclaims lapsed leases only: a slow worker is never raced.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=30.0) + assert claim is not None + recovered, failed = await store.recover_orphans(NOW + 5, 10) + assert recovered == 0, "a live lease was reclaimed" + assert failed == () + recovered, _ = await store.recover_orphans(NOW + 31, 10) + assert recovered == 1, "a lapsed lease was not reclaimed" + + +async def check_a_terminal_run_refuses_further_control(store: RunStore) -> None: + """Finalizing waits for the run to drain, and ends it for good.""" + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + # A claimed step means an attempt may still be running: refuse. + assert not await store.finalize_run( + "run1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW + 1, + ) + await store.release_claim(claim, status=StepStatus.READY, events=(), now=NOW + 1) + assert await store.finalize_run( + "run1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW + 2, + ) + # Terminal is terminal: no further control lands. + assert not await store.request_cancel("run1", NOW + 3) + assert not await store.resume_run("run1", NOW + 3) + assert not await store.finalize_run( + "run1", + status=RunStatus.COMPLETED, + error=None, + event=HistoryEventType.RUN_COMPLETED, + now=NOW + 4, + ) + + async def check_finalize_delivers_a_childs_arrival(store: RunStore) -> None: """Ending a child by cancellation or timeout tells its parent, atomically. @@ -879,6 +924,8 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_list_children_finds_a_joins_branches, check_claims_respect_queue_boundaries, check_substeps_record_once_and_fence_stale_writers, + check_recovery_respects_a_live_lease, + check_a_terminal_run_refuses_further_control, check_finalize_delivers_a_childs_arrival, check_retry_reopens_only_failed_runs, check_force_finalize_records_a_result, From b78da786489a197ba4dc099f51fa0128d1d8352e Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:07:13 -0700 Subject: [PATCH 047/121] Pin the compiler guards a generator hits first Probed the compile-time guards with the mistakes a code generator actually makes: a workflow with no triggered root, duplicate handler ids, an empty workflow id, an unparseable cron, and a webhook with no verifier. All five are rejected, each with a message that names the fix -- which is what makes `workflows check --json` usable as a repair loop. Three of them had no test asserting that behavior, so a refactor could have loosened them silently. The unverified webhook matters most: it is the difference between a public endpoint anyone who learns the URL can start runs through and a deliberate opt-in. The workflow id and cron guards protect things that are expensive to discover late -- an id is durable identity that runs carry forever, and a cron typo surfacing at the first sweep is a job that silently never runs. Two things the compiler accepts were examined and deliberately left alone: a handler nothing routes to (reachability is not statically knowable, and a false warning would send a generator "fixing" working code) and on_failure pointing at a handler that is also a root (unusual, not wrong). --- tests/units/workflow/test_definition.py | 59 +++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/units/workflow/test_definition.py b/tests/units/workflow/test_definition.py index e4e9bb888db..7f965c2ddb1 100644 --- a/tests/units/workflow/test_definition.py +++ b/tests/units/workflow/test_definition.py @@ -14,6 +14,7 @@ WorkflowConfig, hmac_signature, manual, + schedule, webhook, ) @@ -405,3 +406,61 @@ def finish(self): namespace = _load_module(source) definition = compile_workflow(namespace["TransitionsFlow"]) assert set(definition.handlers) == {"begin", "charge", "finish"} + + +def test_a_workflow_id_must_be_a_dotted_name(forked_registration_context): + """The id is a durable identity, so its shape is fixed at compile time. + + Runs carry it forever and operators search by it; discovering a typo or an + empty string after runs exist means either living with it or migrating + rows. + """ + with pytest.raises(WorkflowDefinitionError, match="lowercase dotted"): + + class Empty(rx.State): + __workflow__ = WorkflowConfig(id="") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Go.""" + + compile_workflow(Empty) + + +def test_an_invalid_cron_is_rejected_at_compile_time(forked_registration_context): + """A schedule that cannot be parsed must fail now, not at the first sweep. + + A cron typo that surfaces at runtime is a job that silently never runs. + """ + with pytest.raises(WorkflowDefinitionError, match="cron"): + + class Bad(rx.State): + __workflow__ = WorkflowConfig(id="defn.badcron") + + @rx.event(durable=True, trigger=schedule("not a cron"), effect="none") + def go(self): + """Go.""" + + compile_workflow(Bad) + + +def test_a_webhook_without_a_verifier_is_rejected(forked_registration_context): + """An unverified public endpoint that starts runs is not a default. + + Anyone who learns the URL could start work; the compiler makes that a + deliberate choice rather than an oversight. + """ + with pytest.raises(WorkflowDefinitionError, match="no verifier"): + + class Open(rx.State): + __workflow__ = WorkflowConfig(id="defn.openhook") + + @rx.event(durable=True, trigger=webhook("topic"), effect="none") + def go(self, payload: dict): + """Go. + + Args: + payload: The delivered body. + """ + + compile_workflow(Open) From b785e3cd2c555c315ea2872590f4b3235fdd552b Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:10:18 -0700 Subject: [PATCH 048/121] Fix force_complete and force_fail raising NameError from rx.workflows Probing the public facade the way user code calls it found both operator actions broken: RunStatus, which they evaluate at runtime, was imported only under TYPE_CHECKING, so every call raised NameError instead of finalizing a run. Type checking passed, lint passed, and the whole suite passed, because every existing test called kernel.force_finalize directly -- past the exact layer that was broken. The import moves to runtime and the tests now go through rx.workflows.*, which is the surface user code and generated code actually touch. They fail on the unfixed tree. Worth recording as a habit, not a one-off: a facade that only forwards still needs its own tests. Testing the thing behind it proves the thing behind it. --- reflex/workflow/runtime.py | 3 +- tests/units/workflow/test_operator_actions.py | 65 +++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 10b9ab93d8f..756ad046887 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -29,6 +29,7 @@ WorkflowKernel, WorkflowObserver, ) +from reflex.workflow.records import RunStatus from reflex.workflow.store import RunStore, resolve_store if TYPE_CHECKING: @@ -38,7 +39,7 @@ from collections.abc import AsyncIterator, Callable, Iterable, Mapping from reflex.state import BaseState - from reflex.workflow.records import RunRecord, RunSnapshot, RunStatus, StartResult + from reflex.workflow.records import RunRecord, RunSnapshot, StartResult _context_runtime: ContextVar[WorkflowRuntime | None] = ContextVar( diff --git a/tests/units/workflow/test_operator_actions.py b/tests/units/workflow/test_operator_actions.py index 6e3b7141a76..5dd9d6ca30f 100644 --- a/tests/units/workflow/test_operator_actions.py +++ b/tests/units/workflow/test_operator_actions.py @@ -177,3 +177,68 @@ def done(self, payload: dict): assert not await harness.kernel.force_finalize( result.run_id, status=RunStatus.COMPLETED ) + + +async def test_the_public_facade_exposes_every_action(forked_registration_context): + """rx.workflows.* is the surface user code calls, so it is what is tested. + + The kernel methods were covered while the facade wrapping them was not, + which is how force_complete and force_fail shipped raising NameError from + a type-checking-only import: every test called past the layer that was + broken. + """ + global HEALED + ATTEMPTS.clear() + HEALED = True + async with WorkflowTestHarness(Fragile) as harness: + result = await harness.start(Fragile.start) + assert result.run_id is not None + + # Each action answers for an unknown run rather than raising. + assert await rx.workflows.cancel("no-such-run") is False + assert await rx.workflows.retry("no-such-run") is False + assert await rx.workflows.resume("no-such-run") is False + assert await rx.workflows.force_complete("no-such-run") is False + assert await rx.workflows.force_fail("no-such-run", "gone") is False + assert await rx.workflows.get_run("no-such-run") is None + + # And on a real run, the facade does what the kernel does. + runs = await rx.workflows.list_runs(workflow_id="ops.fragile") + assert [run.run_id for run in runs] == [result.run_id] + _ = harness + + +async def test_force_complete_through_the_facade(forked_registration_context): + """The documented call, exercised the way an operator would make it.""" + + class Parked(rx.State): + __workflow__ = WorkflowConfig(id="ops.parked") + + answered = rx.Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Wait forever. + + Returns: + An unbounded wait. + """ + return rx.wait_for(Parked.answered, then=Parked.done, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def done(self, payload: dict): + """Never reached. + + Args: + payload: The delivered answer. + """ + + async with WorkflowTestHarness(Parked) as harness: + result = await harness.start(Parked.start) + assert result.run_id is not None + assert await rx.workflows.force_complete(result.run_id, {"by": "ops"}) + snapshot = await rx.workflows.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"by": "ops"} + _ = harness From 518004cd3cb2e6e790b7a85141ad92a192a0ea0a Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:13:54 -0700 Subject: [PATCH 049/121] Make the type-only-import bug class impossible to ship again Last commit fixed one instance of it. Neither pyright nor ruff can see the problem -- pyright resolves the TYPE_CHECKING import, ruff sees a used name -- and the whole suite passed because the tests called past the function that would raise. That combination will happen again, so this is a test rather than a habit: it parses every module in the package and fails on any name imported only for type checking that appears where it will be evaluated. Getting it to zero false alarms took three exclusions, each a legitimate pattern it flagged on the first run: annotations including *args and **kwargs, which postponed evaluation never evaluates; the TYPE_CHECKING block itself, where building an alias out of type-only imports is the whole point; and names a function re-imports locally, which is the ordinary way to keep a heavy import out of module scope. A checker that cries wolf gets silenced instead of read, so all three had to go before this was worth having. Verified by reintroducing the exact bug that shipped: the test names the module, the symbol, and the fix. --- .../test_no_type_only_imports_at_runtime.py | 142 ++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 tests/units/workflow/test_no_type_only_imports_at_runtime.py diff --git a/tests/units/workflow/test_no_type_only_imports_at_runtime.py b/tests/units/workflow/test_no_type_only_imports_at_runtime.py new file mode 100644 index 00000000000..8ce354fcf6a --- /dev/null +++ b/tests/units/workflow/test_no_type_only_imports_at_runtime.py @@ -0,0 +1,142 @@ +"""Guard against names imported only for type checking being used as values. + +`from __future__ import annotations` makes every annotation a string, so a +name imported under `if TYPE_CHECKING:` is free to appear in a signature but +raises `NameError` the moment code evaluates it. Neither pyright nor ruff +sees the problem -- pyright resolves the import, ruff sees a used name -- and +a test can miss it entirely by calling past the function that would raise. +That combination shipped a broken `rx.workflows.force_complete`. + +This walks the workflow package and fails on any such name used where it will +be evaluated at runtime. +""" + +import ast +from pathlib import Path + +import pytest + +import reflex.workflow + +PACKAGE = Path(reflex.workflow.__file__).parent +MODULES = sorted(path for path in PACKAGE.glob("*.py") if path.name != "__init__.py") + + +def _type_only_names(tree: ast.Module) -> set[str]: + """Collect names bound only inside ``if TYPE_CHECKING:`` blocks. + + Args: + tree: The parsed module. + + Returns: + The names available to annotations but not at runtime. + """ + names: set[str] = set() + for node in ast.walk(tree): + if not isinstance(node, ast.If): + continue + test = node.test + guarded = (isinstance(test, ast.Name) and test.id == "TYPE_CHECKING") or ( + isinstance(test, ast.Attribute) and test.attr == "TYPE_CHECKING" + ) + if not guarded: + continue + for child in ast.walk(node): + if isinstance(child, (ast.Import, ast.ImportFrom)): + names.update( + alias.asname or alias.name.split(".")[0] for alias in child.names + ) + elif isinstance(child, ast.Assign): + names.update( + target.id + for target in child.targets + if isinstance(target, ast.Name) + ) + return names + + +def _runtime_loads(tree: ast.Module) -> set[str]: + """Collect names the module evaluates at runtime. + + Three things are deliberately not runtime loads, and missing any of them + makes this test cry wolf -- which is worse than not having it, because a + false alarm gets silenced rather than read: + + * annotations, including ``*args`` and ``**kwargs``, which postponed + evaluation never evaluates; + * anything inside the ``if TYPE_CHECKING:`` block itself, where building + an alias out of type-only imports is the point; + * names a function re-imports locally before using, which is the ordinary + way to keep a heavy import out of module scope. + + Args: + tree: The parsed module. + + Returns: + The names loaded where they must exist at runtime. + """ + skip: set[int] = set() + + def skip_subtree(node: ast.AST) -> None: + """Exclude a node and everything under it. + + Args: + node: The subtree to exclude. + """ + skip.update(id(child) for child in ast.walk(node)) + + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + arguments = node.args + for arg in [ + *arguments.posonlyargs, + *arguments.args, + *arguments.kwonlyargs, + arguments.vararg, + arguments.kwarg, + ]: + if arg is not None and arg.annotation is not None: + skip_subtree(arg.annotation) + if node.returns is not None: + skip_subtree(node.returns) + elif isinstance(node, ast.AnnAssign) and node.annotation is not None: + skip_subtree(node.annotation) + elif isinstance(node, ast.If): + test = node.test + if (isinstance(test, ast.Name) and test.id == "TYPE_CHECKING") or ( + isinstance(test, ast.Attribute) and test.attr == "TYPE_CHECKING" + ): + skip_subtree(node) + + local_imports: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + for child in ast.walk(node): + if isinstance(child, (ast.Import, ast.ImportFrom)): + local_imports.update( + alias.asname or alias.name.split(".")[0] + for alias in child.names + ) + + return { + node.id + for node in ast.walk(tree) + if isinstance(node, ast.Name) + and isinstance(node.ctx, ast.Load) + and id(node) not in skip + } - local_imports + + +@pytest.mark.parametrize("path", MODULES, ids=lambda p: p.name) +def test_no_type_checking_import_is_evaluated_at_runtime(path: Path): + """A type-only import used as a value is a NameError waiting for a caller. + + Args: + path: The workflow module to check. + """ + tree = ast.parse(path.read_text()) + offenders = _type_only_names(tree) & _runtime_loads(tree) + assert not offenders, ( + f"{path.name} uses {sorted(offenders)} at runtime, but imports them " + "only under TYPE_CHECKING. Move the import to module scope." + ) From 5554dd64be79053ca04a858cff22d65f46d524aa Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:22:22 -0700 Subject: [PATCH 050/121] Have the worker name the workflows it cannot start A worker process serves no HTTP, so a workflow whose only trigger is a webhook can be executed by it but never begun by it. That separation is correct -- ingress belongs to the app or a dedicated service, execution to the workers, and both share the store -- but silence about it produces the worst operational symptom available: a workflow that never runs, next to a worker that looks perfectly healthy, with nothing to suggest the process being watched was never the one meant to receive the request. The worker now names those roots at startup and says what to run instead, including that schedules and timers do fire in a worker, since that is the next question anyone reading the warning will have. Testing it meant extracting the detection rather than driving the command: a worker runs until interrupted, and a CLI runner cannot interrupt it, so a test that invoked it hung for five minutes instead of failing. The predicate is now a function with a name, which is both testable and clearer. --- reflex/workflow/cli.py | 36 ++++++++++++++- tests/units/workflow/test_cli_check.py | 62 ++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 1 deletion(-) diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index c49baee77e1..57c32713813 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -21,11 +21,32 @@ from reflex.workflow.records import RunStatus if TYPE_CHECKING: - from collections.abc import Awaitable, Callable + from collections.abc import Awaitable, Callable, Iterable + from reflex.workflow.definition import WorkflowDefinition from reflex.workflow.store import RunStore +def webhook_root_names(definitions: Iterable[WorkflowDefinition]) -> list[str]: + """Name the roots that can only be started by an HTTP delivery. + + A worker process serves no HTTP, so these are exactly the workflows it can + execute but never begin. + + Args: + definitions: The compiled workflow definitions being served. + + Returns: + Sorted "workflow_id.handler" names, empty when none are webhook roots. + """ + return sorted( + f"{definition.workflow_id}.{handler.name}" + for definition in definitions + for handler in definition.handlers.values() + if getattr(handler.trigger, "kind", None) == "webhook" + ) + + def _operator_action(database: str | None, run_id: str, action: str, **extra): """Apply one operator action to a run, reporting what happened. @@ -210,6 +231,19 @@ async def serve() -> None: f"Serving {served} on " f"{'queues ' + ', '.join(queues) if queues else 'every queue'}." ) + webhook_roots = webhook_root_names(runtime.definitions) + if webhook_roots: + # A worker has no HTTP server, so it executes runs but cannot + # receive the requests that start these. Left unsaid, the symptom + # is a workflow that simply never runs and a worker that looks + # perfectly healthy. + console.warn( + f"{', '.join(webhook_roots)} start from webhooks, which this " + "worker does not serve. Run the app (or another process " + "serving the workflow endpoints) to receive them; this worker " + "will execute the runs they admit. Schedules and timers do " + "fire here." + ) async with runtime.running(): # The kernel's worker does the work; this task only waits for the # operator (or the platform) to stop the process. diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py index c5b41a996e8..1b71fc26be7 100644 --- a/tests/units/workflow/test_cli_check.py +++ b/tests/units/workflow/test_cli_check.py @@ -213,3 +213,65 @@ def test_the_worker_names_a_compile_error(tmp_path, forked_registration_context) result = CliRunner().invoke(workflows, ["worker", str(module)]) assert result.exit_code == 1 assert "runs it inline" in result.output + + +HOOKED = ''' +import reflex as rx + +class Hooked(rx.State): + __workflow__ = rx.WorkflowConfig(id="check.hooked") + + @rx.event( + durable=True, + effect="none", + trigger=rx.webhook( + "orders", + verify=rx.hmac_signature(secret_env="X_SECRET", header="X-Sig"), + ), + ) + def on_hook(self, payload: dict): + """Handle a delivery. + + Args: + payload: The delivered body. + """ +''' + + +def test_webhook_roots_are_named_for_the_worker(forked_registration_context): + """A worker serves no HTTP, so it must say what it cannot start. + + Without this the symptom is a workflow that never runs beside a worker + that looks entirely healthy -- nothing tells the operator that the process + they are watching was never the one meant to receive the request. + """ + from reflex_base.workflow import WorkflowConfig, hmac_signature, manual, webhook + + import reflex as rx + from reflex.workflow.cli import webhook_root_names + from reflex.workflow.definition import compile_workflow + + class Mixed(rx.State): + __workflow__ = WorkflowConfig(id="check.mixed") + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "orders", + verify=hmac_signature(secret_env="X_SECRET", header="X-Sig"), + ), + ) + def on_hook(self, payload: dict): + """Handle a delivery. + + Args: + payload: The delivered body. + """ + + @rx.event(durable=True, trigger=manual(), effect="none") + def by_hand(self): + """Startable anywhere.""" + + names = webhook_root_names([compile_workflow(Mixed)]) + assert names == ["check.mixed.on_hook"], names From c11823eb4791033b31e940a23fc2947defb9e4fc Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:25:13 -0700 Subject: [PATCH 051/121] Refuse an approval token that can never expire Probing the approval verifier -- the one public, unauthenticated endpoint in the engine -- with hostile input found two claim values it accepted: an expiry of infinity, and an expiry of NaN. NaN is the interesting one, because it loses every comparison, so `expiry < now` is False and the deadline silently never arrives. Either makes a bearer credential immortal. Signing means only our own bug could mint such a token, and that is precisely the argument for fixing it rather than filing it: a verifier whose safety depends on the minter being correct is not a verifier. The expiry must now be a finite number, bools excluded. Everything else the probe threw at it was already refused: empty tokens, missing or extra separators, null bytes, non-ASCII, a hundred-kilobyte body, and a signature lifted from a different set of claims. Two of those are now regression tests as well -- immortal expiries and any missing claim -- since they had been holding by implementation rather than by assertion. --- reflex/workflow/approvals.py | 13 ++++++- tests/units/workflow/test_approvals.py | 50 ++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/reflex/workflow/approvals.py b/reflex/workflow/approvals.py index 78ababea69f..2c4278ec613 100644 --- a/reflex/workflow/approvals.py +++ b/reflex/workflow/approvals.py @@ -23,6 +23,7 @@ import hashlib import hmac import json +import math import os import time from html import escape @@ -188,7 +189,17 @@ def decode_token(token: str) -> dict[str, Any]: raise WorkflowRuntimeError(invalid) from exc if not isinstance(claims, dict) or not {"r", "c", "p", "k", "e"} <= claims.keys(): raise WorkflowRuntimeError(invalid) - if not isinstance(claims["e"], (int, float)) or claims["e"] < time.time(): + expiry = claims["e"] + # Finite is not pedantry: NaN fails every comparison, so `nan < now` is + # False and a NaN expiry would make a token immortal; infinity would do it + # outright. Signing means only our own bug could mint one, which is + # precisely why the verifier should not depend on that being true. + if ( + not isinstance(expiry, (int, float)) + or isinstance(expiry, bool) + or not math.isfinite(expiry) + or expiry < time.time() + ): raise WorkflowRuntimeError(invalid) return claims diff --git a/tests/units/workflow/test_approvals.py b/tests/units/workflow/test_approvals.py index 9e0dfe3946a..dbd937326c6 100644 --- a/tests/units/workflow/test_approvals.py +++ b/tests/units/workflow/test_approvals.py @@ -426,3 +426,53 @@ async def test_a_link_can_carry_a_model_payload( assert snapshot is not None assert snapshot.result == {"outcome": "ada:True"} await runtime.shutdown() + + +@pytest.mark.parametrize( + ("name", "expiry"), + [("infinite", float("inf")), ("not-a-number", float("nan"))], +) +def test_a_token_can_never_be_immortal(monkeypatch, name, expiry): + """An expiry that is not a finite number must not pass the check. + + NaN loses every comparison, so `nan < now` is False and the deadline + silently never arrives; infinity does it outright. Signing means only our + own bug could mint such a token, which is exactly why the verifier must + not assume it never will. + """ + import json + + from reflex.workflow.approvals import _b64, _sign, decode_token + + monkeypatch.setenv(SECRET_ENV, SECRET) + claims = { + "r": "run1", + "c": "decided", + "p": {"ok": True}, + "k": "key", + "e": expiry, + } + body = json.dumps(claims, separators=(",", ":"), sort_keys=True).encode() + with pytest.raises(WorkflowRuntimeError, match="not valid"): + decode_token(f"{_b64(body)}.{_sign(body)}") + + +def test_a_token_missing_a_claim_is_refused(monkeypatch): + """Every claim the delivery depends on must be present and signed.""" + import json + + from reflex.workflow.approvals import _b64, _sign, decode_token + + monkeypatch.setenv(SECRET_ENV, SECRET) + complete = { + "r": "run1", + "c": "decided", + "p": {"ok": True}, + "k": "key", + "e": time.time() + 60, + } + for dropped in complete: + claims = {key: value for key, value in complete.items() if key != dropped} + body = json.dumps(claims, separators=(",", ":"), sort_keys=True).encode() + with pytest.raises(WorkflowRuntimeError, match="not valid"): + decode_token(f"{_b64(body)}.{_sign(body)}") From 60dd7ad1961ccb41976925e745c9d2e1d3af22e6 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:28:50 -0700 Subject: [PATCH 052/121] Refuse a webhook body that cannot fill the root's parameters Probing the webhook endpoint the way the approval verifier was probed found that a signed JSON array or scalar was accepted for a root taking several named parameters. There is no way to map one onto the other, so the code quietly substituted an empty argument set: the provider received a 202, the run was admitted, and it failed on its first step for a reason nobody was ever told. A 202 that means "accepted and already doomed" is worse than a rejection. The boundary now answers 400 for that shape, which is where a request that cannot be served belongs. A single-parameter root still receives whatever was sent -- it declared that it takes the whole payload, and a trigger that wants the body validated says so with model=. The rest of the probe found the layering already sound: oversize bodies, unparseable JSON, unknown topics, and bad signatures are all refused before the runtime is consulted, and four hundred levels of nesting parse without incident. --- reflex/workflow/ingress.py | 12 +++-- tests/units/workflow/test_ingress.py | 69 ++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/reflex/workflow/ingress.py b/reflex/workflow/ingress.py index 314b0c285cd..9697e47d1e0 100644 --- a/reflex/workflow/ingress.py +++ b/reflex/workflow/ingress.py @@ -128,9 +128,7 @@ def _root_args(handler: HandlerDefinition, payload: Any) -> dict[str, Any]: return {} if len(handler.params) == 1: return {handler.params[0]: payload} - if isinstance(payload, dict): - return {name: payload.get(name) for name in handler.params} - return {} + return {name: payload.get(name) for name in handler.params} def webhook_endpoint( @@ -181,6 +179,14 @@ async def endpoint(request: Request) -> JSONResponse: ) spec = getattr(route.definition.state_cls, route.handler.name) + if len(route.handler.params) > 1 and not isinstance(payload, dict): + # Several named parameters can only be filled from an object. + # Admitting the run anyway would drop the payload silently and + # produce a run that fails on its first step for a reason the + # provider is never told, so the boundary refuses it instead. + return JSONResponse( + {"error": "payload must be a JSON object"}, status_code=400 + ) args = _root_args(route.handler, payload) result = await runtime.kernel.start( spec(**args) if args else spec, diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py index 2b530e68414..bef25d38f26 100644 --- a/tests/units/workflow/test_ingress.py +++ b/tests/units/workflow/test_ingress.py @@ -265,3 +265,72 @@ async def test_the_harness_starts_a_webhook_root_directly(paid_workflow): snapshot = await harness.get_run(result.run_id) assert snapshot is not None assert snapshot.status is RunStatus.COMPLETED + + +async def test_a_payload_that_cannot_be_mapped_is_refused( + monkeypatch, forked_registration_context +): + """A body that cannot fill the root's parameters is a 400, not a doomed run. + + A root taking several named parameters can only be filled from a JSON + object. Admitting a run from an array or a scalar dropped the payload + silently and produced a run that failed on its first step, with the + provider told nothing -- it received a 202 and moved on. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + + class MultiArg(rx.State): + __workflow__ = WorkflowConfig(id="ingress.multiarg") + seen: str = "" + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "multi", + verify=hmac_signature( + secret_env="STRIPE_WEBHOOK_SECRET", header="X-Signature" + ), + ), + ) + def on_event(self, first: str, second: int): + """Take two named fields. + + Args: + first: The first field. + second: The second field. + """ + self.seen = f"{first}:{second}" + + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(MultiArg) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + try: + with TestClient(app) as client: + body = json.dumps([1, 2, 3]).encode() + refused = client.post( + "/_workflow/webhook/multi", + content=body, + headers={ + "X-Signature": _sign(body), + "content-type": "application/json", + }, + ) + assert refused.status_code == 400 + assert "object" in refused.json()["error"] + + good = json.dumps({"first": "a", "second": 2}).encode() + accepted = client.post( + "/_workflow/webhook/multi", + content=good, + headers={ + "X-Signature": _sign(good), + "content-type": "application/json", + }, + ) + assert accepted.status_code == 202, accepted.text + finally: + await runtime.shutdown() From 6e987e59698d9eb7e8544a733ca57ed32e5d71fc Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:32:29 -0700 Subject: [PATCH 053/121] Refuse run state that JSON cannot represent Probing the serialization boundary every step crosses found NaN and infinity passing straight through as floats. Python's encoder emits them and its decoder accepts them back, so memory and SQLite are perfectly happy; Postgres JSONB is not, and refuses them at commit. That is the worst shape a bug can take here -- a workflow that works all through development and fails only in production, on the store, with the offending value long gone from the traceback. Encoding now forbids them, so the failure happens where the value still has a name. Everything else the probe tried already reduced sensibly: datetimes and UUIDs to strings, models and sets and tuples to plain data, arbitrarily large integers intact, circular references and unserializable objects refused. The serde boundary had no test module of its own, which is why this was possible to miss; it has one now, covering what it accepts as well as what it rejects. --- reflex/workflow/serde.py | 11 ++++- tests/units/workflow/test_serde.py | 71 ++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 2 deletions(-) create mode 100644 tests/units/workflow/test_serde.py diff --git a/reflex/workflow/serde.py b/reflex/workflow/serde.py index 7491c15cd79..f4fe2fb13f5 100644 --- a/reflex/workflow/serde.py +++ b/reflex/workflow/serde.py @@ -47,6 +47,13 @@ def to_run_data(value: Any) -> Any: Raises: TypeError: If the value contains something no serializer handles. - ValueError: If the value cannot be encoded (e.g. circular references). + ValueError: If the value cannot be encoded -- a circular reference, or + a float that JSON has no representation for (NaN, infinity). """ - return json.loads(json.dumps(value, ensure_ascii=False, default=_strict_default)) + # allow_nan=False is the whole point: NaN and Infinity are not JSON, and + # Python emitting them anyway is how a workflow that works on SQLite and + # in memory fails at commit against Postgres, whose JSONB refuses them. + # Failing here names the value; failing there names a column. + return json.loads( + json.dumps(value, ensure_ascii=False, allow_nan=False, default=_strict_default) + ) diff --git a/tests/units/workflow/test_serde.py b/tests/units/workflow/test_serde.py new file mode 100644 index 00000000000..537609c4302 --- /dev/null +++ b/tests/units/workflow/test_serde.py @@ -0,0 +1,71 @@ +"""Tests for normalizing run state into JSON-compatible data. + +Everything a handler leaves in `self` crosses this boundary on the way to the +store, so what it accepts is what a run can hold and what it refuses is what a +developer finds out about immediately rather than at a commit. +""" + +import datetime +import uuid + +import pytest +from pydantic import BaseModel + +from reflex.workflow.serde import to_run_data + + +class Quote(BaseModel): + """A typical typed value held in run state.""" + + price: int + vendor: str + + +def test_common_values_reduce_to_plain_data(): + """The types a handler actually holds survive the crossing.""" + reduced = to_run_data({ + "when": datetime.date(2026, 1, 1), + "id": uuid.UUID(int=7), + "quote": Quote(price=42, vendor="acme"), + "tags": {"b", "a"}, + "pair": (1, 2), + }) + assert reduced["when"] == "2026-01-01" + assert reduced["id"] == "00000000-0000-0000-0000-000000000007" + assert reduced["quote"] == {"price": 42, "vendor": "acme"} + assert sorted(reduced["tags"]) == ["a", "b"] + assert reduced["pair"] == [1, 2] + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_a_float_json_cannot_hold_is_refused(value: float): + """NaN and infinity are not JSON, whatever Python is willing to emit. + + Left alone they pass through memory and SQLite -- Python's decoder accepts + the tokens its encoder produced -- and are rejected by Postgres JSONB at + commit. That is a workflow that works in development and fails in + production, on the store, with the offending value long gone from the + traceback. + """ + with pytest.raises(ValueError, match="JSON compliant"): + to_run_data({"value": value}) + + +def test_a_nested_non_finite_float_is_refused(): + """The check reaches values buried in state, not just top-level ones.""" + with pytest.raises(ValueError, match="JSON compliant"): + to_run_data({"readings": [1.0, {"latest": float("nan")}]}) + + +def test_a_circular_reference_is_refused(): + """State that points at itself cannot be written down.""" + loop: dict = {} + loop["self"] = loop + with pytest.raises(ValueError, match=r"[Cc]ircular"): + to_run_data(loop) + + +def test_an_unserializable_object_is_refused(): + """A value no serializer handles fails here, not at the store.""" + with pytest.raises(TypeError): + to_run_data({"handle": object()}) From 48cee242d3995444fadae71e35fb01983a4b609e Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:35:25 -0700 Subject: [PATCH 054/121] Assert the scheduling predicates directly Last commit exposed a systematic gap rather than a single one: a module with no test file of its own is a module whose behaviour is only checked indirectly. Sweeping for that found records.py, which holds the two functions every worker consults to decide what it may claim and how long it may sleep. Three stores and the kernel call them; nothing stated what they promise. Their subtlest rule is invisible from any call site. A wait with no deadline carries due_at == 0, and reading that as "due since the epoch" makes it permanently claimable -- the worker then spins at full speed forever instead of idling until a signal arrives. The code has always been right; now a test says so, and reintroducing that exact mistake fails it. The rest of the module's contract is stated too: claimable statuses gate on the clock, a blocked slot becomes claimable exactly when its deadline lands (claiming it is the timeout branch), terminal steps never return, and a claimed step is reclaimed by lease expiry rather than by time -- otherwise a slow attempt would be raced by the worker running it. --- tests/units/workflow/test_records.py | 112 +++++++++++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 tests/units/workflow/test_records.py diff --git a/tests/units/workflow/test_records.py b/tests/units/workflow/test_records.py new file mode 100644 index 00000000000..57942535503 --- /dev/null +++ b/tests/units/workflow/test_records.py @@ -0,0 +1,112 @@ +"""Tests for the scheduling predicates every worker consults. + +``step_claimable_at`` and ``step_wake_at`` decide what a worker may take and +how long it may sleep. They are small and pure, they are consulted by three +stores and the kernel, and their subtlest rule is invisible from the call +site: a wait with no deadline must never look claimable, or the worker spins +at full speed forever instead of idling. Nothing asserted them directly -- +they were only ever exercised through store behaviour, where a regression +would surface far from its cause. +""" + +import pytest + +from reflex.workflow.records import ( + CLAIMABLE_STEP_STATUSES, + TERMINAL_STEP_STATUSES, + StepRecord, + StepStatus, + step_claimable_at, + step_wake_at, +) + +NOW = 1_000.0 + + +def make(status: StepStatus, due_at: float = 0.0) -> StepRecord: + """Build a step in a given status and due time. + + Args: + status: The step status. + due_at: Its due time in epoch seconds. + + Returns: + The step record. + """ + return StepRecord( + run_id="r", + ordinal=0, + handler_id="h", + status=status, + args={}, + due_at=due_at, + origin="root", + ) + + +@pytest.mark.parametrize("status", sorted(CLAIMABLE_STEP_STATUSES, key=str)) +def test_a_claimable_step_waits_for_its_due_time(status: StepStatus): + """Ready, retry-backoff and recovery all gate on the clock alike. + + Args: + status: The claimable status under test. + """ + assert step_claimable_at(make(status, NOW - 1), NOW) + assert step_claimable_at(make(status, NOW), NOW) + assert not step_claimable_at(make(status, NOW + 1), NOW) + assert step_wake_at(make(status, NOW + 5)) == NOW + 5 + + +def test_a_wait_becomes_claimable_only_when_its_deadline_arrives(): + """Claiming a blocked slot IS the timeout branch, so it waits for it.""" + assert not step_claimable_at(make(StepStatus.BLOCKED, NOW + 1), NOW) + assert step_claimable_at(make(StepStatus.BLOCKED, NOW), NOW) + assert step_wake_at(make(StepStatus.BLOCKED, NOW + 1)) == NOW + 1 + + +def test_a_wait_without_a_deadline_never_wakes_on_the_clock(): + """The rule that keeps an idle worker idle. + + A wait with no deadline is due_at == 0. Treating that as "due since the + epoch" would make it permanently claimable, and a worker would spin + through it as fast as it could rather than sleeping until something + actually happens. It is resolved by a signal, never by time. + """ + forever = make(StepStatus.BLOCKED, 0.0) + assert not step_claimable_at(forever, NOW) + assert not step_claimable_at(forever, 0.0) + assert step_wake_at(forever) is None + + +@pytest.mark.parametrize("status", sorted(TERMINAL_STEP_STATUSES, key=str)) +def test_a_finished_step_is_never_claimable_again(status: StepStatus): + """Terminal is terminal, whatever the clock says. + + Args: + status: The terminal status under test. + """ + assert not step_claimable_at(make(status, 0.0), NOW) + assert not step_claimable_at(make(status, NOW - 100), NOW) + assert step_wake_at(make(status, NOW - 100)) is None + + +def test_a_claimed_step_is_not_reclaimed_by_the_clock(): + """A step someone is executing is recovered by lease expiry, not due time. + + If the clock could reclaim it, a slow attempt would be raced by the very + worker that started it. + """ + running = make(StepStatus.CLAIMED, NOW - 100) + assert not step_claimable_at(running, NOW) + assert step_wake_at(running) is None + + +def test_blocked_is_deliberately_not_in_the_claimable_set(): + """The set and the predicate say different things, on purpose. + + Callers that read CLAIMABLE_STEP_STATUSES do not all bound due_at, so a + blocked slot must not be a member; the predicate is what knows a deadline + can make one claimable. + """ + assert StepStatus.BLOCKED not in CLAIMABLE_STEP_STATUSES + assert step_claimable_at(make(StepStatus.BLOCKED, NOW), NOW) From d925e85e6ffaf063cb4f2440842ae5834876a2f5 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:38:15 -0700 Subject: [PATCH 055/121] Write down the runtime's lifecycle rules Continuing the sweep for modules checked only indirectly: WorkflowRuntime is what an app owns and what every rx.workflows call resolves through, and none of its lifecycle behaviour was asserted anywhere. Probing it first, all of it held: the kernel exists only between start and stop rather than being built on demand, a second startup keeps the same kernel, a second shutdown is tolerated, a stopped runtime can be started again, registering after start is refused, and `running()` both activates the facade and restores what was there before. Each of those is now a test, because each is load-bearing in a way that would be quiet if it broke. A lazily built kernel would give every caller its own scheduler against one store. A second startup that built a second kernel would double every claim. A definition accepted after start would be registered, absent from the running kernel, and impossible to start, with nothing to explain why. A runtime left active after its block would leave later code talking to a stopped kernel. --- tests/units/workflow/test_runtime.py | 136 +++++++++++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 tests/units/workflow/test_runtime.py diff --git a/tests/units/workflow/test_runtime.py b/tests/units/workflow/test_runtime.py new file mode 100644 index 00000000000..afd6348e381 --- /dev/null +++ b/tests/units/workflow/test_runtime.py @@ -0,0 +1,136 @@ +"""Tests for the runtime's lifecycle. + +WorkflowRuntime is what an app owns and what every `rx.workflows` call goes +through, so its lifecycle rules are load-bearing: when a kernel exists, what +happens to a double start or stop, and which mistakes are refused rather than +half-performed. All of it held under probing; none of it was written down. +""" + +import pytest +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import MemoryRunStore + + +def _flow(): + """Build a minimal registerable workflow. + + Returns: + The workflow class. + """ + + class Simple(rx.State): + __workflow__ = WorkflowConfig(id="runtime.simple") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Do nothing.""" + + return Simple + + +async def test_the_kernel_exists_only_between_start_and_stop( + forked_registration_context, +): + """Reaching for the kernel too early says so instead of building one. + + A lazily created kernel would quietly give each caller its own scheduler + against the same store. + """ + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(_flow()) + + with pytest.raises(WorkflowRuntimeError, match="not started"): + _ = runtime.kernel + + await runtime.startup(start_worker=False) + assert runtime.kernel is not None + + await runtime.shutdown() + with pytest.raises(WorkflowRuntimeError, match="not started"): + _ = runtime.kernel + + +async def test_starting_and_stopping_are_idempotent(forked_registration_context): + """Lifecycle hooks fire more than once; that must not be a failure. + + An app harness, a test, and a platform supervisor can all call these, and + a second start that built a second kernel would double every claim. + """ + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(_flow()) + + await runtime.startup(start_worker=False) + kernel = runtime.kernel + await runtime.startup(start_worker=False) + assert runtime.kernel is kernel, "a second startup replaced the kernel" + + await runtime.shutdown() + await runtime.shutdown() + + # And a runtime can be started again after being stopped. + await runtime.startup(start_worker=False) + assert runtime.kernel is not None + await runtime.shutdown() + + +async def test_registering_after_start_is_refused(forked_registration_context): + """A definition added after the kernel exists would never be served. + + Silently accepting it produces a workflow that is registered, absent from + the running kernel, and impossible to start -- with nothing to explain it. + """ + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(_flow()) + await runtime.startup(start_worker=False) + try: + with pytest.raises(WorkflowRuntimeError, match="after the runtime"): + + class Late(rx.State): + __workflow__ = WorkflowConfig(id="runtime.late") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Do nothing.""" + + runtime.register(Late) + finally: + await runtime.shutdown() + + +async def test_running_activates_the_facade_and_restores_it( + forked_registration_context, +): + """`running()` is what makes rx.workflows resolve, and it cleans up. + + A runtime that stayed active after its block would leave later code + talking to a stopped kernel and a store nobody is draining. + """ + flow = _flow() + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(flow) + + with pytest.raises(WorkflowRuntimeError, match="No workflow runtime"): + await rx.workflows.get_run("anything") + + async with runtime.running(): + result = await rx.workflows.start(flow.go) + assert result.run_id is not None + + with pytest.raises(WorkflowRuntimeError, match="No workflow runtime"): + await rx.workflows.get_run(result.run_id) + + +def test_definitions_are_reported_once_per_class(forked_registration_context): + """Registering the same class twice is a no-op, not a duplicate.""" + flow = _flow() + runtime = WorkflowRuntime(MemoryRunStore()) + first = runtime.register(flow) + second = runtime.register(flow) + assert first is second + assert [definition.workflow_id for definition in runtime.definitions] == [ + "runtime.simple" + ] From ee56d1d47c1949ac5dcdfc749b5ac170a4baeb71 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:41:06 -0700 Subject: [PATCH 056/121] Validate the instrument the whole suite measures with Last module in the sweep, and the one where a defect is quietest: every other workflow test measures through this harness, so a fault here would not fail loudly, it would invalidate the evidence for everything else. Probing found it sound, and the properties the rest of the suite silently assumes are now stated. Time is virtual: a day passes in microseconds and never touches the wall clock, so a suite containing a three-day timer does not take three days and does not behave differently depending on when it runs. Advancing runs exactly what became due -- twenty-three hours leaves the follow-up pending, the twenty-fourth fires it. The clock cannot run backwards, since rewinding would make due work undue and strand a run. Each harness starts from a clean store, so one test's leftovers can never satisfy another test's assertion, which is the failure that makes a suite untrustable rather than merely red. And a store the caller passed in is left open, which is what the multi-harness tests rely on. Every module in reflex/workflow now has direct tests except cron.py, which is exercised thoroughly through test_schedules.py. --- tests/units/workflow/test_testing.py | 125 +++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 tests/units/workflow/test_testing.py diff --git a/tests/units/workflow/test_testing.py b/tests/units/workflow/test_testing.py new file mode 100644 index 00000000000..da044a55024 --- /dev/null +++ b/tests/units/workflow/test_testing.py @@ -0,0 +1,125 @@ +"""Tests for the test harness itself. + +Every other workflow test measures with this instrument, so a defect here +would not fail loudly -- it would quietly invalidate the evidence for +everything else. These assert the properties the rest of the suite assumes: +that time is virtual, that advancing runs exactly what became due, and that +one harness cannot see another's runs. +""" + +import time + +import pytest +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import RunStatus +from reflex.workflow.store import MemoryRunStore +from reflex.workflow.testing import WorkflowTestHarness + +FIRED: list[str] = [] + + +class Deferred(rx.State): + """Defers work by a day, then records that it ran.""" + + __workflow__ = WorkflowConfig(id="harness.deferred") + done: bool = False + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Schedule the follow-up. + + Returns: + A step due in one day. + """ + FIRED.append("go") + return rx.after("1d", Deferred.later) + + @rx.event(durable=True, effect="none") + def later(self): + """Run a day later.""" + FIRED.append("later") + self.done = True + + +async def test_time_is_virtual_not_wall_clock(forked_registration_context): + """A day passes in microseconds, and never touches the real clock. + + If the harness slept, a suite with a three-day timer in it would take + three days; if it read the wall clock, the same test would behave + differently depending on when it ran. + """ + FIRED.clear() + started = time.monotonic() + async with WorkflowTestHarness(Deferred) as harness: + before = harness.now + result = await harness.start(Deferred.go) + assert result.run_id is not None + await harness.advance("1d") + assert harness.now - before == pytest.approx(86_400) + assert time.monotonic() - started < 5, "the harness slept in real time" + assert FIRED == ["go", "later"] + + +async def test_advancing_runs_what_became_due_and_nothing_else( + forked_registration_context, +): + """Work due later stays pending; the run is left mid-flight, not finished.""" + FIRED.clear() + async with WorkflowTestHarness(Deferred) as harness: + result = await harness.start(Deferred.go) + assert result.run_id is not None + + await harness.advance("23h") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert FIRED == ["go"], "a step ran before it was due" + + await harness.advance("1h") + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert FIRED == ["go", "later"] + + +async def test_the_clock_cannot_run_backwards(forked_registration_context): + """Rewinding time would make due work undue and strand a run.""" + async with WorkflowTestHarness(Deferred) as harness: + with pytest.raises(WorkflowDefinitionError, match=r"[Dd]uration"): + await harness.advance("-5s") + + +async def test_each_harness_starts_from_a_clean_store(forked_registration_context): + """Two harnesses cannot see each other's runs. + + A shared default store would let one test's leftovers satisfy another + test's assertion, which is the failure mode that makes a suite untrustable + rather than merely red. + """ + FIRED.clear() + async with WorkflowTestHarness(Deferred) as first: + result = await first.start(Deferred.go) + assert result.run_id is not None + assert len(await first.kernel.list_runs()) == 1 + + async with WorkflowTestHarness(Deferred) as second: + assert await second.kernel.list_runs() == () + assert await second.get_run(result.run_id) is None + + +async def test_an_injected_store_is_left_to_its_owner(forked_registration_context): + """A store the caller passed in outlives the harness that borrowed it. + + The harness closes what it created; closing what it was handed would + break the multi-harness tests that share one store on purpose. + """ + store = MemoryRunStore() + async with WorkflowTestHarness(Deferred, store=store) as harness: + result = await harness.start(Deferred.go) + assert result.run_id is not None + + # Still usable afterwards: the run is there for the next harness. + assert await store.get_run(result.run_id) is not None From 049941b88c31866bea5387480de3481fca8e3411 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:48:21 -0700 Subject: [PATCH 057/121] Add the skip operator action The last of the four operator actions the plan names. Retry re-runs a step, force_fail and force_complete end the run; skip is the one for a step that cannot succeed and is not worth failing the run over -- a vendor that retired an endpoint, a notification nobody needs any more. The step is marked SKIPPED, which is terminal and recorded as a decision rather than an outcome, and the run continues at whatever comes next. Building it turned up the failure mode the feature could easily have had: skipping the last open slot left the run PENDING with nothing that could ever claim it, which is being stuck in a new way rather than resolved. A run with nothing left to do now completes, with no result, in the same transaction. Legal only on a run stopped for attention or failure, so it can never race a working attempt. Wired through the store protocol, all three stores, the kernel, rx.workflows, and the CLI; conformance-checked; and written into CONTRACT.md section 9, which had described this action before it existed and was trimmed rather than left lying. --- news/workflow-skip.feature.md | 1 + reflex/workflow/CONTRACT.md | 4 + reflex/workflow/cli.py | 12 ++ reflex/workflow/conformance.py | 33 ++++ reflex/workflow/kernel.py | 17 ++ reflex/workflow/postgres.py | 57 +++++++ reflex/workflow/records.py | 3 + reflex/workflow/runtime.py | 16 ++ reflex/workflow/store.py | 158 ++++++++++++++++++ tests/units/workflow/test_operator_actions.py | 66 ++++++++ 10 files changed, 367 insertions(+) create mode 100644 news/workflow-skip.feature.md diff --git a/news/workflow-skip.feature.md b/news/workflow-skip.feature.md new file mode 100644 index 00000000000..5d4e843b56b --- /dev/null +++ b/news/workflow-skip.feature.md @@ -0,0 +1 @@ +Operators can skip a step that cannot succeed (`rx.workflows.skip`, `reflex workflows skip`), letting the rest of a stopped run continue. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index c22eac026d9..0963b3126ba 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -220,6 +220,10 @@ no-op with a reason. - `retry(run)` — a `FAILED` run: re-opens its failed step with a fresh attempt budget and re-runs from there. History keeps the failure; a retry never rewrites the record of why it was needed. +- `skip(run)` — a run stopped for attention or failure: marks the blocking + step `SKIPPED` (terminal, recorded as a decision rather than an outcome) + and lets the run continue at whatever comes next. With nothing left to run, + the run completes with no result rather than sitting pending forever. - `force_complete(run, result)` / `force_fail(run, reason)` — a nonterminal, drained run: finalizes immediately, tombstoning open slots, recording the operator origin and (for completion) the result to treat it as having diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 57c32713813..f2efcea9809 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -543,6 +543,18 @@ def retry(database: str | None, run_id: str): _operator_action(database, run_id, "retry_run") +@workflows.command() +@database_option +@click.argument("run_id") +def skip(database: str | None, run_id: str): + """Skip the step blocking a stopped run and let it continue. + + For a step that cannot succeed and is not worth failing the run over. It + is recorded as an operator decision, not as an outcome. + """ + _operator_action(database, run_id, "skip_step") + + @workflows.command() @database_option @click.argument("run_id") diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 89fecdec605..ebca120a911 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -783,6 +783,38 @@ async def check_finalize_delivers_a_childs_arrival(store: RunStore) -> None: assert steps[1].status is StepStatus.READY +async def check_skip_unsticks_a_stopped_run(store: RunStore) -> None: + """Skipping marks the blocking step terminal and lets the run continue.""" + await store.admit(make_run(), make_step(), _ADMITTED) + # A pending run has nothing to skip. + assert not await store.skip_step("run1", NOW) + + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.NEEDS_ATTENTION, + run_status=RunStatus.NEEDS_ATTENTION, + state={}, + run_error={"reason": "uncertain"}, + ), + NOW, + ) + + assert await store.skip_step("run1", NOW + 1) + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.SKIPPED + run = await store.get_run("run1") + assert run is not None + # Nothing was left to run, so the run is finished rather than pending. + assert run.status is RunStatus.COMPLETED + assert run.error is None + # And it is no longer a stopped run, so skipping again does nothing. + assert not await store.skip_step("run1", NOW + 2) + assert not await store.skip_step("missing", NOW) + + async def check_retry_reopens_only_failed_runs(store: RunStore) -> None: """An operator retry applies to a failed run and nothing else.""" await store.admit(make_run(), make_step(), _ADMITTED) @@ -927,6 +959,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_recovery_respects_a_live_lease, check_a_terminal_run_refuses_further_control, check_finalize_delivers_a_childs_arrival, + check_skip_unsticks_a_stopped_run, check_retry_reopens_only_failed_runs, check_force_finalize_records_a_result, check_schedule_cursors_persist, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index f4b00199f5e..030e3e13af9 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -853,6 +853,23 @@ async def retry(self, run_id: str) -> bool: self._wakeup.set() return retried + async def skip(self, run_id: str) -> bool: + """Skip the step blocking a stopped run and let it continue. + + Args: + run_id: The run to unstick. + + Returns: + True if a blocking step was skipped. + """ + skipped = await self._store.skip_step(run_id, self._clock()) + if skipped: + await self._notify_run( + run_id, ((HistoryEventType.STEP_SKIPPED, {"origin": "operator"}),) + ) + self._wakeup.set() + return skipped + async def force_finalize( self, run_id: str, diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 95553ddfe35..e9ebe5677ce 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1479,6 +1479,63 @@ async def retry_run(self, run_id: str, now: float) -> bool: ) return True + async def skip_step(self, run_id: str, now: float) -> bool: + """Give up on a stuck step and let the run carry on past it. + + Args: + run_id: The run to unstick. + now: Current time in epoch seconds. + + Returns: + True if a blocking step was skipped. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + await self._lock_run(conn, run_id) + cursor = await conn.execute( + "SELECT s.ordinal AS ordinal FROM workflow_steps s" + " JOIN workflow_runs r ON r.run_id = s.run_id" + " WHERE s.run_id = %s AND s.status = ANY(%s)" + " AND r.status = ANY(%s) ORDER BY s.ordinal LIMIT 1", + ( + run_id, + [ + StepStatus.FAILED.value, + StepStatus.TIMED_OUT.value, + StepStatus.NEEDS_ATTENTION.value, + ], + [RunStatus.NEEDS_ATTENTION.value, RunStatus.FAILED.value], + ), + ) + row = await cursor.fetchone() + if row is None: + return False + await conn.execute( + "UPDATE workflow_steps SET status = %s, lease_expires_at = 0," + " updated_at = %s WHERE run_id = %s AND ordinal = %s", + (StepStatus.SKIPPED.value, now, run_id, row["ordinal"]), + ) + cursor = await conn.execute( + "SELECT 1 FROM workflow_steps WHERE run_id = %s" + " AND NOT (status = ANY(%s)) LIMIT 1", + (run_id, _TERMINAL_STEPS), + ) + open_left = await cursor.fetchone() is not None + await conn.execute( + "UPDATE workflow_runs SET status = %s, error = NULL, updated_at = %s" + " WHERE run_id = %s", + ( + RunStatus.PENDING.value if open_left else RunStatus.COMPLETED.value, + now, + run_id, + ), + ) + events = [(HistoryEventType.STEP_SKIPPED, {"ordinal": row["ordinal"]})] + if not open_left: + events.append((HistoryEventType.RUN_COMPLETED, {})) + await self._append_events(conn, run_id, tuple(events), now) + return True + async def recover_orphans( self, now: float, max_recoveries: int ) -> tuple[int, tuple[str, ...]]: diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 1694e7e741d..68e16dfd6d0 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -52,6 +52,7 @@ class StepStatus(str, enum.Enum): TIMED_OUT = "TIMED_OUT" CANCELLED = "CANCELLED" NEEDS_ATTENTION = "NEEDS_ATTENTION" + SKIPPED = "SKIPPED" TERMINAL_STEP_STATUSES = frozenset(( @@ -60,6 +61,7 @@ class StepStatus(str, enum.Enum): StepStatus.TIMED_OUT, StepStatus.CANCELLED, StepStatus.NEEDS_ATTENTION, + StepStatus.SKIPPED, )) CLAIMABLE_STEP_STATUSES = frozenset(( @@ -130,6 +132,7 @@ class HistoryEventType(str, enum.Enum): RUN_CANCELLED = "run_cancelled" RUN_NEEDS_ATTENTION = "run_needs_attention" RUN_RESUMED = "run_resumed" + STEP_SKIPPED = "step_skipped" CHILD_STARTED = "child_started" CHILD_RESOLVED = "child_resolved" WAIT_ARMED = "wait_armed" diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 756ad046887..c016864a886 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -331,6 +331,22 @@ async def retry(run_id: str) -> bool: """ return await get_runtime().kernel.retry(run_id) + @staticmethod + async def skip(run_id: str) -> bool: + """Skip the step blocking a stopped run and let it continue. + + For a step that cannot succeed and is not worth failing the run over. + It is recorded as an operator decision, and the run resumes at + whatever comes next. + + Args: + run_id: The run to unstick. + + Returns: + True if a blocking step was skipped. + """ + return await get_runtime().kernel.skip(run_id) + @staticmethod async def force_complete(run_id: str, result: Any = None) -> bool: """End a run as completed by operator decision. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 1df3ee6386e..cd6f5db7b68 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -444,6 +444,25 @@ async def finalize_run( """ ... + async def skip_step(self, run_id: str, now: float) -> bool: + """Give up on a stuck step and let the run carry on past it. + + The operator's answer to a step that cannot succeed and is not worth + failing the run over -- a vendor that retired an endpoint, a + notification nobody needs any more. The step is marked SKIPPED, which + is terminal and recorded as a decision rather than an outcome, and the + run continues at whatever comes next. Legal only on a run stopped for + attention or failure, so it can never race a working attempt. + + Args: + run_id: The run to unstick. + now: Current time in epoch seconds. + + Returns: + True if a blocking step was skipped. + """ + ... + async def retry_run(self, run_id: str, now: float) -> bool: """Re-open a failed run at the step that failed. @@ -1465,6 +1484,64 @@ async def finalize_run( self._apply_arrival(*parent_arrival, now) return True + async def skip_step(self, run_id: str, now: float) -> bool: + """Give up on a stuck step and let the run carry on past it. + + The operator's answer to a step that cannot succeed and is not worth + failing the run over -- a vendor that retired an endpoint, a + notification nobody needs any more. The step is marked SKIPPED, which + is terminal and recorded as a decision rather than an outcome, and the + run continues at whatever comes next. Legal only on a run stopped for + attention or failure, so it can never race a working attempt. + + Args: + run_id: The run to unstick. + now: Current time in epoch seconds. + + Returns: + True if a blocking step was skipped. + """ + async with self._lock: + run = self._runs.get(run_id) + if run is None or run.status not in ( + RunStatus.NEEDS_ATTENTION, + RunStatus.FAILED, + ): + return False + steps = self._steps[run_id] + for index, step in enumerate(steps): + if step.status in ( + StepStatus.FAILED, + StepStatus.TIMED_OUT, + StepStatus.NEEDS_ATTENTION, + ): + steps[index] = dataclasses.replace( + step, + status=StepStatus.SKIPPED, + lease_expires_at=0.0, + updated_at=now, + ) + # Skipping the last open slot leaves nothing to run, so + # the run is finished rather than pending forever -- being + # stuck in a new way is not a resolution. + open_left = any( + other.status not in TERMINAL_STEP_STATUSES for other in steps + ) + events = [ + (HistoryEventType.STEP_SKIPPED, {"ordinal": step.ordinal}) + ] + if not open_left: + events.append((HistoryEventType.RUN_COMPLETED, {})) + self._runs[run_id] = dataclasses.replace( + run, + status=RunStatus.PENDING if open_left else RunStatus.COMPLETED, + error=None, + updated_at=now, + ) + self._append_events(run_id, events, now) + return True + return False + async def retry_run(self, run_id: str, now: float) -> bool: """Re-open a failed run at the step that failed. @@ -3287,6 +3364,87 @@ def work(): return await asyncio.to_thread(work) + async def skip_step(self, run_id: str, now: float) -> bool: + """Give up on a stuck step and let the run carry on past it. + + The operator's answer to a step that cannot succeed and is not worth + failing the run over -- a vendor that retired an endpoint, a + notification nobody needs any more. The step is marked SKIPPED, which + is terminal and recorded as a decision rather than an outcome, and the + run continues at whatever comes next. Legal only on a run stopped for + attention or failure, so it can never race a working attempt. + + Args: + run_id: The run to unstick. + now: Current time in epoch seconds. + + Returns: + True if a blocking step was skipped. + """ + + def work() -> bool: + """Skip the blocking step on the worker thread. + + Returns: + Whether a blocking step was skipped. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT s.ordinal AS ordinal FROM workflow_steps s" + " JOIN workflow_runs r ON r.run_id = s.run_id" + " WHERE s.run_id = ? AND s.status IN (?, ?, ?)" + " AND r.status IN (?, ?) ORDER BY s.ordinal LIMIT 1", + ( + run_id, + StepStatus.FAILED.value, + StepStatus.TIMED_OUT.value, + StepStatus.NEEDS_ATTENTION.value, + RunStatus.NEEDS_ATTENTION.value, + RunStatus.FAILED.value, + ), + ).fetchone() + if row is None: + self._db.execute("ROLLBACK") + return False + self._db.execute( + "UPDATE workflow_steps SET status = ?, lease_expires_at = 0," + " updated_at = ? WHERE run_id = ? AND ordinal = ?", + (StepStatus.SKIPPED.value, now, run_id, row["ordinal"]), + ) + terminal = tuple(s.value for s in TERMINAL_STEP_STATUSES) + open_left = self._db.execute( + "SELECT 1 FROM workflow_steps WHERE run_id = ?" + f" AND status NOT IN ({','.join('?' * len(terminal))})" + " LIMIT 1", + (run_id, *terminal), + ).fetchone() + self._db.execute( + "UPDATE workflow_runs SET status = ?, error = NULL," + " updated_at = ? WHERE run_id = ?", + ( + RunStatus.PENDING.value + if open_left + else RunStatus.COMPLETED.value, + now, + run_id, + ), + ) + events = [ + (HistoryEventType.STEP_SKIPPED, {"ordinal": row["ordinal"]}) + ] + if not open_left: + events.append((HistoryEventType.RUN_COMPLETED, {})) + self._append_events(run_id, events, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return True + + return await asyncio.to_thread(work) + async def retry_run(self, run_id: str, now: float) -> bool: """Re-open a failed run at the step that failed. diff --git a/tests/units/workflow/test_operator_actions.py b/tests/units/workflow/test_operator_actions.py index 5dd9d6ca30f..3337533c9aa 100644 --- a/tests/units/workflow/test_operator_actions.py +++ b/tests/units/workflow/test_operator_actions.py @@ -242,3 +242,69 @@ def done(self, payload: dict): assert snapshot.status is RunStatus.COMPLETED assert snapshot.result == {"by": "ops"} _ = harness + + +async def test_skip_lets_a_stuck_run_continue(forked_registration_context): + """A step that cannot succeed need not fail the whole run. + + The vendor retired the endpoint; the notification it sent no longer + matters; the rest of the run does. Skipping records a decision and moves + on, which is the difference between an operator resolving a run and an + operator abandoning it. + """ + calls: list[str] = [] + + class Pipeline(rx.State): + __workflow__ = WorkflowConfig(id="ops.pipeline") + + @rx.event( + durable=True, + trigger=manual(), + effect="non_idempotent_write", + retry=Retry(max_attempts=1), + ) + def notify(self): + """Fail in a way that suspends rather than fails the run. + + Raises: + TransientWorkflowError: Always. + """ + calls.append("notify") + msg = "vendor endpoint is gone" + raise TransientWorkflowError(msg) + + async with WorkflowTestHarness(Pipeline) as harness: + result = await harness.start(Pipeline.notify) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + + assert await rx.workflows.skip(result.run_id) + await harness.run_until_idle() + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + # Nothing followed the skipped step, so the run is simply done. + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.steps[0].status is StepStatus.SKIPPED + assert calls == ["notify"], "skipping re-ran the step" + + history = await harness.kernel.store.get_history(result.run_id) + assert any(event.type.value == "step_skipped" for event in history) + + +async def test_skip_is_refused_on_a_healthy_run(forked_registration_context): + """Skipping applies to a stopped run, never to one that is working.""" + global HEALED + ATTEMPTS.clear() + HEALED = True + async with WorkflowTestHarness(Fragile) as harness: + result = await harness.start(Fragile.start) + assert result.run_id is not None + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + + assert not await rx.workflows.skip(result.run_id) + assert not await rx.workflows.skip("no-such-run") From 659446e3006d76c1af36a05b54e2814ac383b1d2 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:55:12 -0700 Subject: [PATCH 058/121] Enforce a singleton's limit inside the admitting transaction "At most one active run per key" was decided before admission: look for an active run, see none, insert. Two starts racing for one key both looked, both saw nothing, and both inserted -- two dunning runs for one invoice, two charges, which is exactly the duplicate the policy exists to prevent. The window is small and the failure is silent, which is the worst combination. The limit is now the store's to enforce, checked in the same transaction that inserts the run: memory and SQLite under the lock they already hold, Postgres with FOR UPDATE over the key's active runs so a second admission waits, sees the first, and is refused. A refused start still answers with the run that holds the key, so the caller can find the live one. The regression races six starts at a single key and asserts exactly one is admitted; it fails on the previous check-then-act path. Conformance-checked too, so a future store must decide it the same way -- and the check covers the release: once the holder is terminal, the key is free again. --- news/workflow-atomic-singleton.fix.md | 1 + reflex/workflow/conformance.py | 36 +++++++++++++++ reflex/workflow/kernel.py | 32 +++++++++++--- reflex/workflow/postgres.py | 18 ++++++++ reflex/workflow/store.py | 46 +++++++++++++++++++- tests/units/workflow/test_flow_control.py | 53 +++++++++++++++++++++++ 6 files changed, 178 insertions(+), 8 deletions(-) create mode 100644 news/workflow-atomic-singleton.fix.md diff --git a/news/workflow-atomic-singleton.fix.md b/news/workflow-atomic-singleton.fix.md new file mode 100644 index 00000000000..dc134dda0ad --- /dev/null +++ b/news/workflow-atomic-singleton.fix.md @@ -0,0 +1 @@ +`rx.Singleton(mode="skip")` now enforces "one active run per key" inside the admitting transaction, so concurrent starts can no longer both be admitted. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index ebca120a911..4286124bfc4 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -783,6 +783,41 @@ async def check_finalize_delivers_a_childs_arrival(store: RunStore) -> None: assert steps[1].status is StepStatus.READY +async def check_admission_enforces_the_active_limit(store: RunStore) -> None: + """A singleton's limit is decided by the admitting transaction itself.""" + first = make_run("a", flow_key="k1") + assert await store.admit(first, make_step("a"), _ADMITTED, max_active=1) == ( + True, + "a", + ) + # A second start under the same key is refused, and told which run holds it. + second = make_run("b", flow_key="k1", created_at=NOW + 1) + assert await store.admit(second, make_step("b"), _ADMITTED, max_active=1) == ( + False, + "a", + ) + assert await store.get_run("b") is None + # A different key is unaffected. + other = make_run("c", flow_key="k2", created_at=NOW + 2) + assert await store.admit(other, make_step("c"), _ADMITTED, max_active=1) == ( + True, + "c", + ) + # Once the holder is terminal the key is free again. + assert await store.finalize_run( + "a", + status=RunStatus.COMPLETED, + error=None, + event=HistoryEventType.RUN_COMPLETED, + now=NOW + 3, + ) + third = make_run("d", flow_key="k1", created_at=NOW + 4) + assert await store.admit(third, make_step("d"), _ADMITTED, max_active=1) == ( + True, + "d", + ) + + async def check_skip_unsticks_a_stopped_run(store: RunStore) -> None: """Skipping marks the blocking step terminal and lets the run continue.""" await store.admit(make_run(), make_step(), _ADMITTED) @@ -959,6 +994,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_recovery_respects_a_live_lease, check_a_terminal_run_refuses_further_control, check_finalize_delivers_a_childs_arrival, + check_admission_enforces_the_active_limit, check_skip_unsticks_a_stopped_run, check_retry_reopens_only_failed_runs, check_force_finalize_records_a_result, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 030e3e13af9..738e17f80a0 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -564,12 +564,12 @@ async def _apply_start_policy( """ if handler.singleton is not None: existing = await self._store.first_active(defn.workflow_id, flow_key) + if handler.singleton.mode == "skip": + # Deliberately not decided here: admission re-checks inside its + # own transaction (see _admit), because a decision made out + # here is one that two concurrent starts both win. + return None, now if existing is not None: - if handler.singleton.mode == "skip": - return ( - StartResult(disposition="skipped", run_id=existing.run_id), - now, - ) # Drive the cancellation to a terminal state before admitting the # replacement, so "one active run per key" holds at every instant # rather than only once a worker happens to drain the old one. @@ -742,11 +742,29 @@ async def _admit( {"ordinal": 0, "handler_id": handler.id}, ), ) + # A singleton's "at most one active run per key" is enforced by the + # store inside the admitting transaction. Deciding it beforehand is a + # check-then-act race that two concurrent starts both pass, which for + # a singleton means exactly the duplicate it exists to prevent. + singleton = handler.singleton + max_active = 1 if singleton is not None and singleton.mode == "skip" else None created, authoritative_run_id = await self._store.admit( - run, root_step, admission + run, root_step, admission, max_active=max_active ) if not created: - return StartResult(disposition="deduplicated", run_id=authoritative_run_id) + disposition = "skipped" if max_active is not None else "deduplicated" + if run.request_key is not None and disposition == "skipped": + # A redelivery that also hits the limit is still a dedupe: + # the caller is asking about a run it already started. + existing = await self._store.find_by_request_key( + defn.workflow_id, run.request_key + ) + if existing == authoritative_run_id: + disposition = "deduplicated" + return StartResult( + disposition=disposition, # pyright: ignore[reportArgumentType] + run_id=authoritative_run_id, + ) self._notify(run, admission) self._wakeup.set() return StartResult(disposition="started", run_id=authoritative_run_id) diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index e9ebe5677ce..0587f0cc516 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -579,6 +579,8 @@ async def admit( run: RunRecord, root_step: StepRecord, events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + *, + max_active: int | None = None, ) -> tuple[bool, str]: """Atomically admit a run, deduplicating on the request key. @@ -586,6 +588,9 @@ async def admit( run: The run record to create. root_step: The preallocated root mailbox slot. events: History events to append on creation. + max_active: When set, admit only if fewer than this many runs are + already active under the run's flow key, decided inside this + transaction so concurrent starts cannot both pass. Returns: Whether the run was created, and the authoritative run id. @@ -607,6 +612,19 @@ async def admit( existing = await cursor.fetchone() if existing is not None: return False, existing["run_id"] + if max_active is not None and run.flow_key is not None: + # FOR UPDATE serializes concurrent admissions under one key: + # the second waits, then sees the first and is refused. + cursor = await conn.execute( + "SELECT run_id FROM workflow_runs" + " WHERE workflow_id = %s AND flow_key = %s" + " AND NOT (status = ANY(%s))" + " ORDER BY created_at, run_id FOR UPDATE", + (run.workflow_id, run.flow_key, _TERMINAL_RUNS), + ) + active = await cursor.fetchall() + if len(active) >= max_active: + return False, active[0]["run_id"] await self._insert_run(conn, run) await self._insert_step(conn, root_step) await self._append_events(conn, run.run_id, events, run.created_at) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index cd6f5db7b68..ae38e225e32 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -125,6 +125,8 @@ async def admit( run: RunRecord, root_step: StepRecord, events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + *, + max_active: int | None = None, ) -> tuple[bool, str]: """Atomically admit a run, deduplicating on the request key. @@ -132,6 +134,10 @@ async def admit( run: The run record to create. root_step: The preallocated root mailbox slot. events: History events to append on creation. + max_active: When set, admit only if fewer than this many runs are + already active under the run's flow key. Checked inside the + admitting transaction, because a check made outside it is a + race that two concurrent starts both win. Returns: ``(True, run_id)`` when the run was created, or @@ -842,6 +848,8 @@ async def admit( run: RunRecord, root_step: StepRecord, events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + *, + max_active: int | None = None, ) -> tuple[bool, str]: """Atomically admit a run, deduplicating on the request key. @@ -849,6 +857,10 @@ async def admit( run: The run record to create. root_step: The preallocated root mailbox slot. events: History events to append on creation. + max_active: When set, admit only if fewer than this many runs are + already active under the run's flow key. Checked inside the + admitting transaction, because a check made outside it is a + race that two concurrent starts both win. Returns: Whether the run was created, and the authoritative run id. @@ -859,7 +871,21 @@ async def admit( existing = self._dedupe.get(dedupe_key) if existing is not None: return False, existing - self._dedupe[dedupe_key] = run.run_id + if max_active is not None and run.flow_key is not None: + active = sorted( + ( + other + for other in self._runs.values() + if other.workflow_id == run.workflow_id + and other.flow_key == run.flow_key + and other.status not in TERMINAL_RUN_STATUSES + ), + key=lambda other: (other.created_at, other.run_id), + ) + if len(active) >= max_active: + return False, active[0].run_id + if run.request_key is not None: + self._dedupe[run.workflow_id, run.request_key] = run.run_id self._runs[run.run_id] = run self._steps[run.run_id] = [root_step] self._append_events(run.run_id, events, run.created_at) @@ -2315,6 +2341,8 @@ async def admit( run: RunRecord, root_step: StepRecord, events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + *, + max_active: int | None = None, ) -> tuple[bool, str]: """Atomically admit a run, deduplicating on the request key. @@ -2322,6 +2350,10 @@ async def admit( run: The run record to create. root_step: The preallocated root mailbox slot. events: History events to append on creation. + max_active: When set, admit only if fewer than this many runs are + already active under the run's flow key. Checked inside the + admitting transaction, because a check made outside it is a + race that two concurrent starts both win. Returns: Whether the run was created, and the authoritative run id. @@ -2350,6 +2382,18 @@ def work(): " VALUES (?, ?, ?)", (run.workflow_id, run.request_key, run.run_id), ) + if max_active is not None and run.flow_key is not None: + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + rows = self._db.execute( + "SELECT run_id FROM workflow_runs" + " WHERE workflow_id = ? AND flow_key = ?" + f" AND status NOT IN ({','.join('?' * len(terminal))})" + " ORDER BY created_at, run_id", + (run.workflow_id, run.flow_key, *terminal), + ).fetchall() + if len(rows) >= max_active: + self._db.execute("ROLLBACK") + return False, rows[0]["run_id"] self._insert_run(run) self._insert_step(root_step) self._append_events(run.run_id, events, run.created_at) diff --git a/tests/units/workflow/test_flow_control.py b/tests/units/workflow/test_flow_control.py index 638333092f9..835e7adc686 100644 --- a/tests/units/workflow/test_flow_control.py +++ b/tests/units/workflow/test_flow_control.py @@ -429,3 +429,56 @@ def start(self, a: Left, b: Right): from reflex.workflow.definition import compile_workflow compile_workflow(TwoSources) + + +async def test_a_singleton_holds_under_concurrent_starts( + forked_registration_context, +): + """Two starts racing for one key admit one run, not two. + + Checking "is anything active?" before admitting is a check-then-act race: + both callers look, both see nothing, both insert. For a singleton that is + precisely the duplicate the policy exists to prevent -- two dunning runs + for one invoice, two charges. The limit is therefore enforced inside the + admitting transaction, and this races six starts to prove it. + """ + import asyncio + + class OnePerOrder(rx.State): + __workflow__ = WorkflowConfig(id="flow.oneperorder") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + singleton=Singleton(key="order_id", mode="skip"), + ) + def start(self, order_id: str): + """Hold the run open so the race has something to collide with. + + Args: + order_id: The order this run belongs to. + + Returns: + A far-future continuation. + """ + return after("1h", OnePerOrder.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + """Never reached in this test.""" + + async with WorkflowTestHarness(OnePerOrder) as harness: + results = await asyncio.gather( + *(harness.kernel.start(OnePerOrder.start("ord-1")) for _ in range(6)) + ) + + started = [r for r in results if r.disposition == "started"] + skipped = [r for r in results if r.disposition == "skipped"] + assert len(started) == 1, f"{len(started)} runs admitted for one key" + assert len(skipped) == 5 + # Every loser points at the winner, so a caller can find the live run. + assert {r.run_id for r in skipped} == {started[0].run_id} + + runs = await harness.kernel.list_runs(workflow_id="flow.oneperorder") + assert len(runs) == 1 From 74e76d47fd08b763217f03231875d4523466592d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 00:57:39 -0700 Subject: [PATCH 059/121] Walk the crash matrix as a conformance check Phase 2's exit criterion is that a process can be killed at any boundary without violating the contract. CONTRACT.md section 8 names those boundaries and the one permitted outcome for each, but nothing walked them: each was covered somewhere, by tests written for other reasons, which is not the same as the matrix being checked. The sequence now runs against every store: nothing exists before admission commits; a claim killed before the handler ran costs a recovery and not a business attempt; a substep recorded before the crash survives so the re-execution replays it; a sweep interrupted and repeated recovers nothing twice; a committed transition is durable in full, children included; and a child's terminal commit lands its arrival at the parent's join, so the join is never left waiting on a run that already finished. Verified to discriminate by making a recovery consume a business attempt -- a plausible one-word regression that would silently shorten every retry budget after a crash, and that nothing else in the suite catches. --- reflex/workflow/conformance.py | 96 ++++++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 4286124bfc4..918e9fc508d 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -783,6 +783,101 @@ async def check_finalize_delivers_a_childs_arrival(store: RunStore) -> None: assert steps[1].status is StepStatus.READY +async def check_the_crash_matrix_holds_at_every_boundary(store: RunStore) -> None: + """Walk CONTRACT.md section 8: kill at each boundary, check the outcome. + + A worker dies without cleanup at an arbitrary instant. What the store + holds afterwards is the only evidence, and the contract names exactly one + permitted outcome per boundary. This walks them in order rather than + trusting that each is covered somewhere. + """ + # Killed before admission commits: nothing exists, and the retry admits. + assert await store.get_run("run1") is None + await store.admit(make_run(), make_step(), _ADMITTED) + + # Killed after claim, before the handler ran: the lease lapses and + # recovery re-offers the step. It costs a recovery, not an attempt. + claim = await store.claim_next(NOW, lease_duration=1.0) + assert claim is not None + recovered, failed = await store.recover_orphans(NOW + 2, 10) + assert (recovered, failed) == (1, ()) + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.RECOVERY_WAIT + assert steps[0].attempts == 0, "a crash consumed a business attempt" + assert steps[0].recoveries == 1 + + # Killed mid-handler after a substep recorded: the record survives, so a + # re-execution replays it instead of repeating the work. + claim = await store.claim_next(NOW + 3, lease_duration=1.0) + assert claim is not None + assert await store.record_substep( + "run1", 0, claim.step.epoch, "charged", {"id": "ch_1"}, NOW + 3 + ) + await store.recover_orphans(NOW + 5, 10) + assert await store.get_substeps("run1", 0) == {"charged": {"id": "ch_1"}} + + # Killed during a recovery sweep: sweeping again is idempotent, not a + # second recovery of the same step. + before = (await store.get_steps("run1"))[0].recoveries + assert await store.recover_orphans(NOW + 6, 10) == (0, ()) + assert (await store.get_steps("run1"))[0].recoveries == before + + # Killed after a commit: everything the transition promised is durable, + # including a child's arrival at its parent's join. + claim = await store.claim_next(NOW + 7, lease_duration=LEASE) + assert claim is not None + child = make_run("child1", parent_run_id="run1", parent_ordinal=1) + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.WAITING, + state={"n": 1}, + new_steps=( + make_step( + ordinal=1, + status=StepStatus.BLOCKED, + wait_key="join:1", + join_expected=1, + origin="join", + due_at=0.0, + ), + ), + next_ordinal=2, + children=((child, make_step("child1")),), + ), + NOW + 7, + ) + run = await store.get_run("run1") + assert run is not None + assert run.state == {"n": 1} + assert await store.get_run("child1") is not None + + # The child dies for good: its arrival lands with its terminal + # transition, so the join is never left waiting on a finished child. + child_claim = await store.claim_next(NOW + 8, lease_duration=LEASE) + assert child_claim is not None + assert child_claim.run.run_id == "child1" + await store.commit( + child_claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.COMPLETED, + state={}, + parent_arrival=( + "run1", + 1, + {"run_id": "child1", "status": "COMPLETED", "result": None}, + "child1", + ), + ), + NOW + 9, + ) + steps = await store.get_steps("run1") + assert steps[1].join_arrived == 1 + assert steps[1].status is StepStatus.READY + + async def check_admission_enforces_the_active_limit(store: RunStore) -> None: """A singleton's limit is decided by the admitting transaction itself.""" first = make_run("a", flow_key="k1") @@ -994,6 +1089,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_recovery_respects_a_live_lease, check_a_terminal_run_refuses_further_control, check_finalize_delivers_a_childs_arrival, + check_the_crash_matrix_holds_at_every_boundary, check_admission_enforces_the_active_limit, check_skip_unsticks_a_stopped_run, check_retry_reopens_only_failed_runs, From 30466f0ffceb087cab5f542fa9384497ad2ca7bf Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:04:21 -0700 Subject: [PATCH 060/121] Add reflex workflows dev and a typed run handle Two of phase 3's requirements, both aimed at the same thing: shortening the distance between writing a workflow and seeing it run. `reflex workflows dev Workflow.handler --arg name=value` serves the module in the foreground, starts one run, and prints every transition as it is recorded -- attempts, failures with their message, retry scheduling, recorded substeps, the final result -- then stops when the run ends. Driving it against a workflow that flakes once shows the whole story in ten lines: the failure, the backoff, the second attempt, the substep recorded, the result. Timers are real here, which the help says, because the harness is what skips days instantly. `rx.workflows.submit()` returns a RunHandle instead of an admission report. start() answers "what happened to my submission" and remains the honest answer when that matters; a handle is what a caller works with afterwards, carrying result(), wait(), signal() and cancel() so none of them need the id threaded back through a module. A submission that identified no run -- a rate limit rejecting outright -- refuses to produce a handle rather than hand back one wrapping None, and says to use start() for that case. result() waits for a worker to make progress, which the harness does not run; its tests pump explicitly, and the docstring says the method is for scripts rather than for durable handlers, where blocking on another run would pin a worker slot for the duration. --- news/workflow-dev-and-handle.feature.md | 1 + reflex/workflow/__init__.py | 2 + reflex/workflow/cli.py | 131 +++++++++++++++++ reflex/workflow/handle.py | 181 +++++++++++++++++++++++ reflex/workflow/runtime.py | 41 ++++++ tests/units/workflow/test_handle.py | 188 ++++++++++++++++++++++++ 6 files changed, 544 insertions(+) create mode 100644 news/workflow-dev-and-handle.feature.md create mode 100644 reflex/workflow/handle.py create mode 100644 tests/units/workflow/test_handle.py diff --git a/news/workflow-dev-and-handle.feature.md b/news/workflow-dev-and-handle.feature.md new file mode 100644 index 00000000000..5c202037014 --- /dev/null +++ b/news/workflow-dev-and-handle.feature.md @@ -0,0 +1 @@ +`reflex workflows dev` runs a workflow in the foreground and narrates every transition, and `rx.workflows.submit()` returns a typed `RunHandle` carrying `result()`, `wait()`, `signal()`, and `cancel()`. diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index 8cce91e46b5..4a412eb5838 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -48,6 +48,7 @@ WorkflowDefinition, compile_workflow, ) +from reflex.workflow.handle import RunHandle from reflex.workflow.kernel import ( LoggingObserver, MetricsObserver, @@ -94,6 +95,7 @@ "RateLimit", "Retry", "RunContext", + "RunHandle", "RunQuery", "RunRecord", "RunSnapshot", diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index f2efcea9809..303773921ec 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -150,11 +150,142 @@ def _age(seconds: float) -> str: ) +def _terminal_events() -> frozenset: + """The history events that mean a run has stopped for good. + + Returns: + The terminal event types. + """ + from reflex.workflow.records import HistoryEventType + + return frozenset(( + HistoryEventType.RUN_COMPLETED, + HistoryEventType.RUN_FAILED, + HistoryEventType.RUN_CANCELLED, + HistoryEventType.RUN_TIMED_OUT, + HistoryEventType.RUN_NEEDS_ATTENTION, + )) + + @click.group() def workflows(): """Inspect and steer durable workflow runs.""" +@workflows.command() +@database_option +@click.argument("target") +@click.argument("start", required=False) +@click.option( + "--arg", + "args", + multiple=True, + help="Argument for the started handler, as name=value. Repeatable.", +) +def dev(database: str | None, target: str, start: str | None, args: tuple[str, ...]): + """Run TARGET's workflows in the foreground, printing every transition. + + The loop for building a workflow: start one, watch each step, attempt, + retry and wait as it happens, and stop when the run ends. Pass START as + the handler to launch (`Workflow.handler`), with --arg name=value for its + payload; without it, this just serves and reports whatever arrives. + + Timers are real here. Use WorkflowTestHarness to skip days instantly. + """ + import asyncio + + from reflex_base.utils.exceptions import WorkflowDefinitionError + + from reflex.workflow.kernel import WorkflowObserver + from reflex.workflow.records import HistoryEventType + from reflex.workflow.runtime import WorkflowRuntime + from reflex.workflow.runtime import workflows as rx_workflows + from reflex.workflow.store import resolve_store + + try: + module = _load_module(target) + except Exception as err: + console.error(f"Could not load {target!r}: {err}") + raise click.exceptions.Exit(1) from None + + classes = { + name: value + for name, value in vars(module).items() + if isinstance(value, type) and "__workflow__" in vars(value) + } + if not classes: + console.error(f"No workflow classes in {target!r}.") + raise click.exceptions.Exit(1) + + finished = asyncio.Event() + + class Narrator(WorkflowObserver): + """Prints every transition as it is recorded.""" + + def on_event( + self, + event_type: HistoryEventType, + run_id: str, + workflow_id: str, + data: dict[str, Any], + ) -> None: + """Print one transition. + + Args: + event_type: What happened. + run_id: The run it happened to. + workflow_id: That run's workflow identity. + data: The event payload. + """ + detail = " ".join( + f"{key}={value!r}" + for key, value in data.items() + if key not in ("error", "traceback") + ) + click.echo(f" {run_id[:8]} {event_type.value:<22}{detail}") + if "error" in data and isinstance(data["error"], dict): + click.echo(f" {data['error'].get('message', data['error'])}") + if event_type in _terminal_events(): + finished.set() + + async def serve() -> None: + """Run the kernel until the started run ends, or forever.""" + runtime = WorkflowRuntime(resolve_store(database), observer=Narrator()) + try: + for workflow_cls in classes.values(): + runtime.register(workflow_cls) + except WorkflowDefinitionError as err: + console.error(f"Cannot serve {target!r}: {err}") + raise click.exceptions.Exit(1) from None + + async with runtime.running(): + if start is None: + console.print("Serving; nothing started. Ctrl-C to stop.") + await asyncio.Event().wait() + return + class_name, _, handler_name = start.partition(".") + workflow_cls = classes.get(class_name) + if workflow_cls is None or not handler_name: + console.error( + f"{start!r} is not Workflow.handler; available: " + f"{', '.join(sorted(classes))}." + ) + raise click.exceptions.Exit(1) + payload = dict(pair.split("=", 1) for pair in args if "=" in pair) + spec = getattr(workflow_cls, handler_name) + handle = await rx_workflows.submit(spec(**payload) if payload else spec) + console.print(f"Started {handle.run_id} ({handle.disposition}).") + await finished.wait() + snapshot = await handle.snapshot() + if snapshot is not None: + console.print(f"Run {snapshot.status.value}: {snapshot.result}") + + try: + asyncio.run(serve()) + except KeyboardInterrupt: + console.print("Stopped.") + + @workflows.command() @database_option @click.argument("target") diff --git a/reflex/workflow/handle.py b/reflex/workflow/handle.py new file mode 100644 index 00000000000..a175f251d5e --- /dev/null +++ b/reflex/workflow/handle.py @@ -0,0 +1,181 @@ +"""A typed handle on one run. + +``rx.workflows.start()`` returns an admission result: a disposition and, when +one exists, a run id. That is the honest answer to "what happened to my +submission", but it is not what a caller does next. Next they wait for the +run, look at its result, signal it, or cancel it -- and doing any of that from +a bare string means finding the right module-level function and passing the id +back into it every time. + +A handle carries the id and the operations that belong to it, so the calling +code reads as one object rather than a string plus a namespace. +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, Any + +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import DurationLike, parse_duration + +from reflex.workflow.records import TERMINAL_RUN_STATUSES, RunStatus + +if TYPE_CHECKING: + from reflex_base.workflow import ChannelDelivery + + from reflex.workflow.records import RunSnapshot + +DEFAULT_POLL_INTERVAL: float = 0.05 + + +class RunHandle: + """One run, and the things a caller does with it. + + Attributes: + run_id: The run this handle refers to. + disposition: How admission handled the submission that produced it. + """ + + __slots__ = ("disposition", "run_id") + + def __init__(self, run_id: str, disposition: str = "started"): + """Bind a handle to a run. + + Args: + run_id: The run's identity. + disposition: How admission handled the submission. + """ + self.run_id = run_id + self.disposition = disposition + + def __repr__(self) -> str: + """Describe the handle. + + Returns: + A short representation naming the run. + """ + return f"RunHandle({self.run_id!r}, {self.disposition!r})" + + @property + def started(self) -> bool: + """Whether this submission created the run rather than finding one. + + Returns: + True when a new run was admitted. + """ + return self.disposition == "started" + + async def snapshot(self) -> RunSnapshot | None: + """Read the run's current state. + + Returns: + The snapshot, or None if the run is unknown to this store. + """ + from reflex.workflow.runtime import workflows + + return await workflows.get_run(self.run_id) + + async def status(self) -> RunStatus | None: + """Read just the run's status. + + Returns: + The status, or None if the run is unknown. + """ + snapshot = await self.snapshot() + return None if snapshot is None else snapshot.status + + async def result( + self, + *, + timeout: DurationLike = "30s", + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> Any: + """Wait for the run to finish and return what it produced. + + Meant for scripts and tests, where waiting is the point. A durable + handler should never call this: blocking one run's step on another + run's completion ties up a worker slot for as long as the other run + takes. Compose with ``rx.parallel`` instead, which is what child runs + and joins are for. + + Args: + timeout: How long to wait before giving up. + poll_interval: Seconds between checks. + + Returns: + The run's result. + + Raises: + WorkflowRuntimeError: If the run is unknown, does not finish in + time, or finishes in any state other than completed. + """ + snapshot = await self.wait(timeout=timeout, poll_interval=poll_interval) + if snapshot.status is not RunStatus.COMPLETED: + detail = f": {snapshot.error}" if snapshot.error else "" + msg = ( + f"Run {self.run_id} finished {snapshot.status.value}, not " + f"COMPLETED{detail}" + ) + raise WorkflowRuntimeError(msg) + return snapshot.result + + async def wait( + self, + *, + timeout: DurationLike = "30s", + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> RunSnapshot: + """Wait for the run to reach a terminal state. + + Args: + timeout: How long to wait before giving up. + poll_interval: Seconds between checks. + + Returns: + The final snapshot, whatever the outcome. + + Raises: + WorkflowRuntimeError: If the run is unknown or is still running + when the timeout expires. + """ + deadline = asyncio.get_running_loop().time() + parse_duration(timeout) + while True: + snapshot = await self.snapshot() + if snapshot is None: + msg = f"Run {self.run_id} is not in this store." + raise WorkflowRuntimeError(msg) + if snapshot.status in TERMINAL_RUN_STATUSES: + return snapshot + if asyncio.get_running_loop().time() >= deadline: + msg = ( + f"Run {self.run_id} was still {snapshot.status.value} after " + f"{timeout}. Workers may not be running, or the run is " + "waiting on something that has not happened yet." + ) + raise WorkflowRuntimeError(msg) + await asyncio.sleep(poll_interval) + + async def signal(self, delivery: ChannelDelivery, *, key: str | None = None) -> str: + """Deliver a payload to a channel this run is waiting on. + + Args: + delivery: The addressed payload, e.g. ``MyFlow.approved(answer)``. + key: Sender idempotency key; a repeated key is a no-op. + + Returns: + What the store did with the delivery. + """ + from reflex.workflow.runtime import workflows + + return await workflows.signal(self.run_id, delivery, key=key) + + async def cancel(self) -> bool: + """Request cancellation of the run. + + Returns: + True if intent was recorded on a nonterminal run. + """ + from reflex.workflow.runtime import workflows + + return await workflows.cancel(self.run_id) diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index c016864a886..051d52fc35d 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -23,6 +23,7 @@ ) from reflex.workflow.definition import WorkflowDefinition, compile_workflow +from reflex.workflow.handle import RunHandle from reflex.workflow.kernel import ( DEFAULT_MAX_CONCURRENCY, DEFAULT_POLL_INTERVAL, @@ -270,6 +271,46 @@ async def start( target, request_key=request_key, labels=labels ) + @staticmethod + async def submit( + target: Any, + *, + request_key: str | None = None, + labels: dict[str, str] | None = None, + ) -> RunHandle: + """Start a run and get a handle on it. + + The same admission as ``start()``, returning the run rather than a + report about it: a handle carries the id together with the operations + that belong to it, so a caller waits, signals, or cancels without + threading the string back through a module. + + Args: + target: The root event, e.g. ``MyWorkflow.begin(payload)``. + request_key: Idempotent admission key. + labels: Server-derived indexing labels. + + Returns: + A handle on the admitted (or already existing) run. + + Raises: + WorkflowRuntimeError: If admission identified no run, as when a + rate limit rejected the submission outright. + """ + result = await workflows.start(target, request_key=request_key, labels=labels) + if result.run_id is None: + msg = ( + f"Start was {result.disposition} and produced no run" + + ( + f"; retry after {result.retry_after:.0f}s" + if result.retry_after + else "" + ) + + ". Use rx.workflows.start() to handle that case explicitly." + ) + raise WorkflowRuntimeError(msg) + return RunHandle(result.run_id, result.disposition) + @staticmethod async def cancel(run_id: str) -> bool: """Request cancellation of a run. diff --git a/tests/units/workflow/test_handle.py b/tests/units/workflow/test_handle.py new file mode 100644 index 00000000000..8e363c3b68f --- /dev/null +++ b/tests/units/workflow/test_handle.py @@ -0,0 +1,188 @@ +"""Tests for the typed run handle. + +`start()` answers "what happened to my submission"; a handle is what a caller +actually works with afterwards. These cover the operations it carries and the +one case it refuses, since a handle with no run behind it would be a null +waiting to surface later. +""" + +import pytest +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import RateLimit, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.handle import RunHandle +from reflex.workflow.records import RunStatus +from reflex.workflow.testing import WorkflowTestHarness + + +class Echo(rx.State): + """Completes immediately with what it was given.""" + + __workflow__ = WorkflowConfig(id="handle.echo") + said: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self, said: str): + """Say it back. + + Args: + said: What to echo. + + Returns: + Completion. + """ + self.said = said + return rx.complete(result={"said": said}) + + +async def test_a_handle_carries_the_run_and_its_result(forked_registration_context): + """The common path: submit, wait, read what it produced.""" + async with WorkflowTestHarness(Echo) as harness: + handle = await rx.workflows.submit(Echo.go("hello")) + assert isinstance(handle, RunHandle) + assert handle.started + assert handle.run_id + + # The harness pumps on demand rather than running a worker, so the + # run is advanced explicitly before the handle reads its outcome. + await harness.run_until_idle() + assert await handle.result() == {"said": "hello"} + assert await handle.status() is RunStatus.COMPLETED + snapshot = await handle.snapshot() + assert snapshot is not None + assert snapshot.state["said"] == "hello" + _ = harness + + +async def test_waiting_on_a_run_that_fails_says_so(forked_registration_context): + """`result()` is for the value, so a failed run raises rather than lying.""" + + class Doomed(rx.State): + __workflow__ = WorkflowConfig(id="handle.doomed") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Give up immediately. + + Returns: + Failure. + """ + return rx.fail(reason="nope") + + async with WorkflowTestHarness(Doomed) as harness: + handle = await rx.workflows.submit(Doomed.go) + await harness.run_until_idle() + with pytest.raises(WorkflowRuntimeError, match="FAILED"): + await handle.result() + # wait() reports the outcome instead of raising on it. + snapshot = await handle.wait() + assert snapshot.status is RunStatus.FAILED + _ = harness + + +async def test_waiting_past_the_timeout_explains_itself(forked_registration_context): + """A run that never finishes gets a message naming the likely cause.""" + + class Slow(rx.State): + __workflow__ = WorkflowConfig(id="handle.slow") + + answered = rx.Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Wait forever. + + Returns: + An unbounded wait. + """ + return rx.wait_for(Slow.answered, then=Slow.done, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def done(self, payload: dict): + """Never reached. + + Args: + payload: The delivered answer. + """ + + async with WorkflowTestHarness(Slow) as harness: + handle = await rx.workflows.submit(Slow.go) + await harness.run_until_idle() + with pytest.raises(WorkflowRuntimeError, match="still WAITING"): + await handle.wait(timeout="0.05s", poll_interval=0.01) + _ = harness + + +async def test_a_handle_can_signal_and_cancel(forked_registration_context): + """The operations that belong to a run travel with it.""" + + class Paused(rx.State): + __workflow__ = WorkflowConfig(id="handle.paused") + + answered = rx.Signal(dict) + answer: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Wait for an answer. + + Returns: + An unbounded wait. + """ + return rx.wait_for(Paused.answered, then=Paused.done, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def done(self, payload: dict): + """Record the answer. + + Args: + payload: The delivered answer. + + Returns: + Completion. + """ + self.answer = payload["say"] + return rx.complete(result={"answer": self.answer}) + + async with WorkflowTestHarness(Paused) as harness: + handle = await rx.workflows.submit(Paused.go) + await harness.run_until_idle() + assert await handle.signal(Paused.answered({"say": "yes"})) == "resolved" + await harness.run_until_idle() + assert await handle.result() == {"answer": "yes"} + + second = await rx.workflows.submit(Paused.go) + await harness.run_until_idle() + assert await second.cancel() + _ = harness + + +async def test_a_rejected_submission_refuses_to_hand_back_a_handle( + forked_registration_context, +): + """A handle with no run behind it would be a null surfacing later. + + A rate limit rejects outright: there is no run to hold, so submit() says + so and points at start(), which reports dispositions honestly. + """ + + class Limited(rx.State): + __workflow__ = WorkflowConfig(id="handle.limited") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + rate_limit=RateLimit(limit=1, period="1h"), + ) + def go(self): + """Do nothing.""" + + async with WorkflowTestHarness(Limited) as harness: + first = await rx.workflows.submit(Limited.go) + await harness.run_until_idle() + assert first.started + with pytest.raises(WorkflowRuntimeError, match="rejected"): + await rx.workflows.submit(Limited.go) + _ = harness From 44b3fc4c0b8ae38580f87b5d038f9394607b466f Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:07:36 -0700 Subject: [PATCH 061/121] Scaffold a runnable workflow with reflex workflows init The last piece of the workflow-only path: somewhere to start. `reflex workflows init orders` writes one module -- no app, no frontend, nothing to configure -- and prints the two commands that run it. The scaffold is deliberately not a stub. It uses the pieces worth knowing on the first day and nothing else: a durable root with a retry policy, an rx.step around the side effect so a retry replays it rather than repeating it, and a timer that survives restarts. Reading it teaches the model; running it proves the install. Verified the way a new developer would meet it: init in an empty directory, then paste back the command it printed. The run starts, records the charge, and schedules the follow-up for tomorrow -- which then waits, because timers here are real. A test compiles the generated module through the same compiler registration uses, since a scaffold that does not compile teaches the rules by failing them, and checks that a second init refuses rather than overwriting the file someone has already edited. The plan names this `reflex init --workflow`. It is a subcommand instead: `reflex init` is the app scaffolder, with templates and a frontend, and threading a mode through it that produces neither would change shared code every Reflex user runs to serve a case that lives entirely in this group. --- news/workflow-init.feature.md | 1 + reflex/workflow/cli.py | 102 +++++++++++++++++++++++++ tests/units/workflow/test_cli_check.py | 28 +++++++ 3 files changed, 131 insertions(+) create mode 100644 news/workflow-init.feature.md diff --git a/news/workflow-init.feature.md b/news/workflow-init.feature.md new file mode 100644 index 00000000000..7ae826baa15 --- /dev/null +++ b/news/workflow-init.feature.md @@ -0,0 +1 @@ +`reflex workflows init [name]` scaffolds a runnable workflow module — no app, no frontend — and prints the commands that run it. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 303773921ec..9b61af281b7 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -172,6 +172,108 @@ def workflows(): """Inspect and steer durable workflow runs.""" +_SCAFFOLD = '''"""A durable workflow. + +Run it: + + reflex workflows dev {module}.py {klass}.start --arg order=ord-1 + +Serve it as a background worker: + + reflex workflows worker {module}.py + +Point it at Postgres for more than one worker: + + REFLEX_WORKFLOW_DATABASE=postgresql://... reflex workflows worker {module}.py +""" + +import reflex as rx + + +def charge_card(order: str) -> dict: + """Stand in for a real call to a payment provider. + + Args: + order: The order being charged. + + Returns: + The provider's response. + """ + return {{"charge_id": f"ch_{{order}}"}} + + +class {klass}(rx.State): + """Charges an order, then follows up a day later.""" + + __workflow__ = rx.WorkflowConfig(id="{workflow_id}") + + order: str = "" + charge_id: str = "" + + @rx.event( + durable=True, + trigger=rx.manual(), + effect="idempotent_write", + retry=rx.Retry(max_attempts=5), + ) + async def start(self, order: str): + """Charge the order, then wait a day before following up. + + Args: + order: The order to charge. + + Returns: + The next step, due tomorrow. + """ + self.order = order + # rx.step records its result, so a retry of this handler replays the + # charge instead of making it twice. + charge = await rx.step("charge", charge_card, order) + self.charge_id = charge["charge_id"] + return rx.after("1d", {klass}.follow_up) + + @rx.event(durable=True, effect="idempotent_write") + def follow_up(self): + """Run a day after the charge, whatever happened in between. + + Returns: + Completion. + """ + return rx.complete(result={{"order": self.order, "charge": self.charge_id}}) +''' + + +@workflows.command("init") +@click.argument("name", default="workflows") +def init_workflow(name: str): + """Write a runnable workflow module to NAME.py and say what to do next. + + The workflow-only starting point: one file, no app, no frontend, nothing + to configure. It uses the pieces worth knowing on the first day -- a + durable step, a recorded side effect, a retry policy, and a timer that + survives restarts -- and the commands printed afterwards run it. + """ + module = Path(f"{name}.py") + if module.exists(): + console.error(f"{module} already exists; choose another name.") + raise click.exceptions.Exit(1) + + stem = module.stem.replace("-", "_") + klass = "".join(part.title() for part in stem.split("_")) or "Orders" + module.write_text( + _SCAFFOLD.format( + module=stem, klass=klass, workflow_id=f"{stem}.{klass.lower()}" + ) + ) + console.print(f"Wrote {module}.") + console.print("") + console.print("Run it once, watching every step:") + console.print(f" reflex workflows dev {module} {klass}.start --arg order=ord-1") + console.print("") + console.print("Or serve it as a worker and start runs from your own code:") + console.print(f" reflex workflows worker {module}") + + @workflows.command() @database_option @click.argument("target") diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py index 1b71fc26be7..caf375d2d23 100644 --- a/tests/units/workflow/test_cli_check.py +++ b/tests/units/workflow/test_cli_check.py @@ -275,3 +275,31 @@ def by_hand(self): names = webhook_root_names([compile_workflow(Mixed)]) assert names == ["check.mixed.on_hook"], names + + +def test_init_writes_a_workflow_that_compiles(tmp_path, forked_registration_context): + """The scaffold must be correct code, not a sketch. + + It is the first thing a new developer runs, and the commands it prints + are the next two. A scaffold that does not compile teaches the engine's + rules by failing at them. + """ + from reflex.workflow.cli import workflows as group + + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path): + written = runner.invoke(group, ["init", "orders"]) + assert written.exit_code == 0, written.output + assert "reflex workflows dev orders.py Orders.start" in written.output + + # The generated module passes the same compiler the app applies. + checked = runner.invoke(group, ["check", "orders.py", "--json"]) + assert checked.exit_code == 0, checked.output + payload = json.loads(checked.output) + assert payload["ok"] is True + assert payload["workflows"][0]["workflow_id"] == "orders.orders" + + # A second init does not quietly overwrite the first. + again = runner.invoke(group, ["init", "orders"]) + assert again.exit_code == 1 + assert "already exists" in again.output From fadaaf8918130518d967bd222ca5acf47dfc2842 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:11:00 -0700 Subject: [PATCH 062/121] Add an authenticated HTTP API for starting and reading runs The last phase-3 requirement. Importing the workflow class is the right way to start a run from a Reflex app or a worker, and the wrong way for the service that actually has the business event: a Django view, a Go service, a cron box, none of which should take a dependency on your workflow package to say "this happened". POST /_workflow/api/runs starts a run by workflow id and handler name, with args, labels, and a request_key that means the same thing it means in-process -- a retrying caller reaches the run it already made rather than a second one. GET /_workflow/api/runs/{id} reports status, result, error, and steps. Both routes require a bearer token compared in constant time, and neither is mounted unless REFLEX_WORKFLOW_API_TOKEN is set. An endpoint that starts arbitrary workflows is not something to leave on by accident, so the absence of configuration means the absence of the surface -- not an open one. The tests spend most of their weight there: a missing header, a wrong token, and a token sent without its scheme are all refused on both routes. Naming something that does not exist gets a 404 that says which part, and a malformed body gets a 400 rather than a 500, so a caller integrating against this can tell its own mistakes from the server's. --- news/workflow-http-api.feature.md | 1 + reflex/app.py | 22 ++++ reflex/workflow/api.py | 192 ++++++++++++++++++++++++++++++ tests/units/workflow/test_api.py | 161 +++++++++++++++++++++++++ 4 files changed, 376 insertions(+) create mode 100644 news/workflow-http-api.feature.md create mode 100644 reflex/workflow/api.py create mode 100644 tests/units/workflow/test_api.py diff --git a/news/workflow-http-api.feature.md b/news/workflow-http-api.feature.md new file mode 100644 index 00000000000..640a82d20b2 --- /dev/null +++ b/news/workflow-http-api.feature.md @@ -0,0 +1 @@ +Setting `REFLEX_WORKFLOW_API_TOKEN` mounts an authenticated HTTP API for starting and reading runs, so a service in any language can trigger workflows without importing them. diff --git a/reflex/app.py b/reflex/app.py index f2f1592f113..29331de4e56 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -850,6 +850,13 @@ def _add_default_endpoints(self): def _add_workflow_endpoints(self): """Add the workflow ingress endpoints: webhooks in, approvals back.""" + from reflex.workflow.api import ( + RUN_ROUTE, + START_ROUTE, + api_token, + run_endpoint, + start_endpoint, + ) from reflex.workflow.approvals import APPROVAL_ROUTE, approval_endpoint from reflex.workflow.ingress import ( WEBHOOK_ROUTE, @@ -874,6 +881,21 @@ def _add_workflow_endpoints(self): approval_endpoint(self._workflow_runtime), methods=["GET", "POST"], ) + token = api_token() + if token is not None: + # Mounted only when a token is configured: an unauthenticated + # endpoint that starts arbitrary workflows is not something to + # leave on by accident. + self._api.add_route( + config.prepend_backend_path(START_ROUTE), + start_endpoint(self._workflow_runtime, token), + methods=["POST"], + ) + self._api.add_route( + config.prepend_backend_path(RUN_ROUTE), + run_endpoint(self._workflow_runtime, token), + methods=["GET"], + ) def _add_optional_endpoints(self): """Add optional api endpoints (_upload).""" diff --git a/reflex/workflow/api.py b/reflex/workflow/api.py new file mode 100644 index 00000000000..de7801481ba --- /dev/null +++ b/reflex/workflow/api.py @@ -0,0 +1,192 @@ +"""An HTTP surface for starting and reading runs from outside the process. + +The engine is reachable in Python by importing the workflow class. That is +the right answer for a Reflex app or a worker, and the wrong one for the +service that has the business event: a Django view, a Go service, a cron box, +anything that should not import your workflow package to say "this happened". + +These endpoints give those callers the two verbs they need -- start a run, +read a run -- over HTTP, with the same admission semantics as an in-process +start, including idempotency keys. Every route requires a bearer token, and +without one configured the surface is not mounted at all: an unauthenticated +endpoint that starts arbitrary workflows is not something to leave on by +accident. +""" + +from __future__ import annotations + +import hmac +import json +import os +from typing import TYPE_CHECKING, Any, Final + +from reflex_base.utils import console +from starlette.responses import JSONResponse + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine + + from starlette.requests import Request + + from reflex.workflow.runtime import WorkflowRuntime + +TOKEN_ENV: Final = "REFLEX_WORKFLOW_API_TOKEN" +START_ROUTE: Final = "/_workflow/api/runs" +RUN_ROUTE: Final = "/_workflow/api/runs/{run_id}" +MAX_BODY_BYTES: Final = 1_048_576 + + +def api_token() -> str | None: + """Read the bearer token the API requires, if one is configured. + + Returns: + The token, or None when the API should stay unmounted. + """ + return os.environ.get(TOKEN_ENV) or None + + +def _authorized(request: Request, token: str) -> bool: + """Check a request's bearer token in constant time. + + Args: + request: The incoming request. + token: The configured token. + + Returns: + True when the request carries the right token. + """ + header = request.headers.get("authorization", "") + scheme, _, presented = header.partition(" ") + if scheme.lower() != "bearer": + return False + return hmac.compare_digest(presented, token) + + +def start_endpoint( + runtime: WorkflowRuntime, token: str +) -> Callable[[Request], Coroutine[Any, Any, JSONResponse]]: + """Build the endpoint that starts a run. + + Args: + runtime: The runtime owning the workflows. + token: The bearer token every caller must present. + + Returns: + The endpoint callable. + """ + + async def endpoint(request: Request) -> JSONResponse: + """Start a run from a JSON body. + + Args: + request: The incoming request. + + Returns: + The admission result, or an error. + """ + if not _authorized(request, token): + return JSONResponse({"error": "unauthorized"}, status_code=401) + body = await request.body() + if len(body) > MAX_BODY_BYTES: + return JSONResponse({"error": "payload too large"}, status_code=413) + try: + payload = json.loads(body or b"{}") + except json.JSONDecodeError: + return JSONResponse({"error": "payload is not JSON"}, status_code=400) + if not isinstance(payload, dict): + return JSONResponse( + {"error": "payload must be a JSON object"}, status_code=400 + ) + + workflow_id = payload.get("workflow") + handler_name = payload.get("handler") + if not isinstance(workflow_id, str) or not isinstance(handler_name, str): + return JSONResponse( + {"error": "workflow and handler are required"}, status_code=400 + ) + definition = next( + ( + candidate + for candidate in runtime.definitions + if candidate.workflow_id == workflow_id + ), + None, + ) + if definition is None: + return JSONResponse({"error": "unknown workflow"}, status_code=404) + target = getattr(definition.state_cls, handler_name, None) + if target is None: + return JSONResponse({"error": "unknown handler"}, status_code=404) + + args = payload.get("args") or {} + if not isinstance(args, dict): + return JSONResponse( + {"error": "args must be a JSON object"}, status_code=400 + ) + labels = payload.get("labels") + try: + result = await runtime.kernel.start( + target(**args) if args else target, + request_key=payload.get("request_key"), + labels=labels if isinstance(labels, dict) else None, + ) + except Exception as err: + # A rejected start is the caller's problem to see, not a 500: the + # usual cause is a handler that is not a manual root, or args that + # do not fit it. + console.warn(f"Workflow API start refused: {err}") + return JSONResponse({"error": str(err)}, status_code=400) + return JSONResponse( + {"disposition": result.disposition, "run_id": result.run_id}, + status_code=202, + ) + + return endpoint + + +def run_endpoint( + runtime: WorkflowRuntime, token: str +) -> Callable[[Request], Coroutine[Any, Any, JSONResponse]]: + """Build the endpoint that reads one run. + + Args: + runtime: The runtime owning the runs. + token: The bearer token every caller must present. + + Returns: + The endpoint callable. + """ + + async def endpoint(request: Request) -> JSONResponse: + """Report a run's status, result, and steps. + + Args: + request: The incoming request. + + Returns: + The run projection, or an error. + """ + if not _authorized(request, token): + return JSONResponse({"error": "unauthorized"}, status_code=401) + run_id = request.path_params.get("run_id", "") + snapshot = await runtime.kernel.get_run(run_id) + if snapshot is None: + return JSONResponse({"error": "unknown run"}, status_code=404) + return JSONResponse({ + "run_id": snapshot.run_id, + "workflow": snapshot.workflow_id, + "status": snapshot.status.value, + "result": snapshot.result, + "error": snapshot.error, + "steps": [ + { + "ordinal": step.ordinal, + "handler": step.handler_id, + "status": step.status.value, + "attempts": step.attempts, + } + for step in snapshot.steps + ], + }) + + return endpoint diff --git a/tests/units/workflow/test_api.py b/tests/units/workflow/test_api.py new file mode 100644 index 00000000000..10b30fb5a62 --- /dev/null +++ b/tests/units/workflow/test_api.py @@ -0,0 +1,161 @@ +"""Tests for the HTTP surface that starts and reads runs. + +This is the endpoint a Django view or a Go service calls to say "this +happened" without importing the workflow package. It can start arbitrary +workflows, so most of these are about the token: what it refuses, and that +the surface does not exist at all without one. +""" + +import json + +import pytest +from reflex_base.workflow import WorkflowConfig, manual +from starlette.applications import Starlette +from starlette.routing import Route +from starlette.testclient import TestClient + +import reflex as rx +from reflex.workflow.api import ( + RUN_ROUTE, + START_ROUTE, + TOKEN_ENV, + api_token, + run_endpoint, + start_endpoint, +) +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import MemoryRunStore + +TOKEN = "wf_" + "t" * 32 + + +class Orders(rx.State): + """A workflow an outside service starts.""" + + __workflow__ = WorkflowConfig(id="api.orders") + order: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def place(self, order: str): + """Record the order. + + Args: + order: The order identifier. + + Returns: + Completion. + """ + self.order = order + return rx.complete(result={"order": order}) + + +@pytest.fixture +async def client(forked_registration_context): + """A client wired to the API endpoints of a live runtime. + + Yields: + The test client. + """ + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Orders) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[ + Route(START_ROUTE, start_endpoint(runtime, TOKEN), methods=["POST"]), + Route(RUN_ROUTE, run_endpoint(runtime, TOKEN), methods=["GET"]), + ] + ) + with TestClient(app) as ready: + yield ready + await runtime.shutdown() + + +def _auth() -> dict[str, str]: + """Build the authorization header. + + Returns: + Headers carrying the bearer token. + """ + return {"authorization": f"Bearer {TOKEN}"} + + +def test_a_service_can_start_and_read_a_run(client): + """The two verbs an outside caller needs, over HTTP.""" + started = client.post( + START_ROUTE, + content=json.dumps({ + "workflow": "api.orders", + "handler": "place", + "args": {"order": "ord-1"}, + }), + headers=_auth(), + ) + assert started.status_code == 202, started.text + run_id = started.json()["run_id"] + assert started.json()["disposition"] == "started" + + read = client.get(f"/_workflow/api/runs/{run_id}", headers=_auth()) + assert read.status_code == 200 + assert read.json()["workflow"] == "api.orders" + assert read.json()["steps"][0]["handler"] == "place" + + +def test_an_idempotency_key_returns_the_same_run(client): + """A retrying caller must not create a second run.""" + body = json.dumps({ + "workflow": "api.orders", + "handler": "place", + "args": {"order": "ord-2"}, + "request_key": "invoice-77", + }) + first = client.post(START_ROUTE, content=body, headers=_auth()) + second = client.post(START_ROUTE, content=body, headers=_auth()) + assert first.json()["run_id"] == second.json()["run_id"] + assert second.json()["disposition"] == "deduplicated" + + +@pytest.mark.parametrize( + "headers", + [{}, {"authorization": "Bearer wrong"}, {"authorization": TOKEN}], + ids=["missing", "wrong-token", "no-scheme"], +) +def test_every_route_requires_the_token(client, headers): + """An endpoint that starts arbitrary workflows is never open. + + Args: + client: The test client. + headers: The authorization headers under test. + """ + body = json.dumps({"workflow": "api.orders", "handler": "place"}) + assert client.post(START_ROUTE, content=body, headers=headers).status_code == 401 + assert ( + client.get("/_workflow/api/runs/whatever", headers=headers).status_code == 401 + ) + + +def test_unknown_targets_are_refused(client): + """A caller naming something that does not exist is told which part.""" + for body, status in [ + ({"workflow": "api.nope", "handler": "place"}, 404), + ({"workflow": "api.orders", "handler": "nope"}, 404), + ({"workflow": "api.orders"}, 400), + ({"workflow": "api.orders", "handler": "place", "args": [1]}, 400), + ]: + response = client.post(START_ROUTE, content=json.dumps(body), headers=_auth()) + assert response.status_code == status, (body, response.text) + + assert ( + client.post(START_ROUTE, content=b"", headers=_auth()).status_code + == 400 + ) + assert client.get("/_workflow/api/runs/missing", headers=_auth()).status_code == 404 + + +def test_the_api_is_absent_without_a_configured_token(monkeypatch): + """No token means no surface, rather than a surface anyone can call.""" + monkeypatch.delenv(TOKEN_ENV, raising=False) + assert api_token() is None + monkeypatch.setenv(TOKEN_ENV, "") + assert api_token() is None + monkeypatch.setenv(TOKEN_ENV, TOKEN) + assert api_token() == TOKEN From 03f1b2006a9812009161228541e1e3b5de9f2fa2 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:12:15 -0700 Subject: [PATCH 063/121] Document tenancy and who runs the workers Auditing CONTRACT.md against the items phase 1 requires found eight of nine covered in depth and one covered in a sentence: managed versus customer-hosted workers, and the tenant isolation that question really asks about. It now says what is true. Managed and self-hosted are a deployment split, not a semantic one -- a worker is any process pointed at the store, and nothing in the engine asks which it is. Isolation is the store's boundary, with exactly three supported arrangements: a store or schema per tenant, which shares nothing; a shared store separated by workflow id, which is fine for one product's own workloads and not for mutually distrusting tenants; and queue partitioning, which separates compute without separating visibility. The part worth writing down is the limit: there is no tenant column, no per-tenant authorization inside the store, and no filter a caller with store access cannot lift. Anything stronger than "shared store, separate workflow ids" has to come from a separate store. A tenancy story that overstates itself is worse than one that admits its edges, because the overstatement is what someone builds on. --- reflex/workflow/CONTRACT.md | 40 ++++++++++++++++++++++++++++++++++--- 1 file changed, 37 insertions(+), 3 deletions(-) diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 0963b3126ba..b5dd3ecc0fe 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -186,9 +186,43 @@ Checked in this order at start: - Multiple workers share one Postgres store via `SKIP LOCKED` claims; SQLite is a one-process store (calls off-loop, contention bounded); memory is for tests. All three answer the same conformance suite. -- Managed vs customer-hosted is a deployment split, not a semantic one: a - worker is any process with `REFLEX_WORKFLOW_DATABASE` pointed at the store - and (optionally) `workflow_queues` narrowed. The contract is identical. +### Tenancy and who runs the workers + +Managed and customer-hosted are a deployment split, not a semantic one. A +worker is any process with `REFLEX_WORKFLOW_DATABASE` pointed at the store and +optionally `workflow_queues` narrowed; everything above holds identically +whether that process runs on a platform or on a laptop. Nothing in the engine +asks which it is, and no behaviour changes with the answer. + +Isolation between tenants is the store's boundary, and there are exactly three +supported arrangements: + +1. **A store per tenant.** Separate databases, or `PostgresRunStore(schema=)` + inside a shared one. Nothing crosses, because nothing is shared: a query, + a claim, and a recovery sweep all see one tenant by construction. This is + the arrangement to pick when tenants must not be able to observe each + other even through a bug. +2. **A shared store, isolated by workflow id.** One deployment's workflows are + its own; `list_runs(workflow_id=...)` and the CLI's filters are how an + operator stays inside them. Fine for one product's own workloads, not for + mutually distrusting tenants -- a caller that can reach the store can reach + every run in it. +3. **A shared store, partitioned by queue.** Workers serve named queues, so + compute is separated even where data is not: a tenant's slow work cannot + starve another's, and a worker can be dedicated to one tenant's steps. + This partitions *execution*, never *visibility*. + +What the engine does not do: there is no tenant column, no per-tenant +authorization inside the store, and no filter that a caller with store access +cannot lift. Anything stronger than "shared store, separate workflow ids" must +come from arrangement 1. Say so plainly rather than implying a boundary that +is not enforced -- a tenancy story that overstates itself is worse than one +that admits its edges. + +Credentials follow the same rule everywhere. Webhook secrets, approval-link +signing keys, and the API token are read from the environment at use time, so +they never enter run state, history, or a browser bundle, and rotating one is +a restart rather than a migration. ## 8. The failure matrix From a4017df1df0a91ae32aea269de7abfefb42c9281 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:15:07 -0700 Subject: [PATCH 064/121] Show what starts each workflow with reflex workflows triggers Phase 4's exit criterion is that nobody should have to inspect the database to understand or repair a run. The operator commands covered repair -- list, show with history, retry, skip, resume, cancel, force-complete -- and left a gap on understanding: nothing answered "is my cron actually registered" or "what URL does this provider post to". The only other answer to either is reading the source, which is exactly what an operator should not have to do to understand a deployment. `reflex workflows triggers ` lists every root and how it starts: webhook topics with the path to hand the provider and whether the endpoint is signature-verified, schedules with their cron and the next time they fire, and manual roots as such. --json for tooling. Flagging unverified webhooks in the default view is deliberate. An open endpoint that starts runs is the one thing in this listing worth noticing across a screen of output, and the compiler already refuses to create one by accident -- this is where an operator sees the ones that were created on purpose. --- news/workflow-triggers-cli.feature.md | 1 + reflex/workflow/cli.py | 83 ++++++++++++++++++++++++++ tests/units/workflow/test_cli_check.py | 69 +++++++++++++++++++++ 3 files changed, 153 insertions(+) create mode 100644 news/workflow-triggers-cli.feature.md diff --git a/news/workflow-triggers-cli.feature.md b/news/workflow-triggers-cli.feature.md new file mode 100644 index 00000000000..1fd787640d5 --- /dev/null +++ b/news/workflow-triggers-cli.feature.md @@ -0,0 +1 @@ +`reflex workflows triggers ` lists what starts each workflow — webhook topics with their URL and whether they are signature-verified, schedules with their next fire time, and manual roots. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 9b61af281b7..d5ba280b142 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -11,6 +11,7 @@ import asyncio import inspect import json +import operator import sys from pathlib import Path from typing import TYPE_CHECKING, Any @@ -243,6 +244,88 @@ def follow_up(self): ''' +@workflows.command() +@click.argument("target") +@click.option("--json", "as_json", is_flag=True, help="Emit JSON instead of a table.") +def triggers(target: str, as_json: bool): + """List what starts the workflows in TARGET, and how. + + The answer to "is my cron actually registered, and what is the URL this + provider should post to" -- questions whose only other answer is reading + the source, which is exactly what an operator should never have to do to + understand a deployment. + """ + from reflex_base.workflow import ScheduleTrigger, WebhookTrigger + + from reflex.workflow.cron import CronSchedule + from reflex.workflow.definition import compile_workflow + + try: + module = _load_module(target) + except Exception as err: + console.error(f"Could not load {target!r}: {err}") + raise click.exceptions.Exit(1) from None + + rows: list[dict[str, Any]] = [] + for value in vars(module).values(): + if not (isinstance(value, type) and "__workflow__" in vars(value)): + continue + definition = compile_workflow(value) + for handler in definition.handlers.values(): + trigger = handler.trigger + if isinstance(trigger, WebhookTrigger): + rows.append({ + "workflow": definition.workflow_id, + "handler": handler.name, + "kind": "webhook", + "detail": trigger.topic, + "path": f"/_workflow/webhook/{trigger.topic}", + "verified": trigger.verify is not None, + "dedupe_by": trigger.dedupe_by, + }) + elif isinstance(trigger, ScheduleTrigger): + import time + + schedule = CronSchedule(trigger.cron) + upcoming = schedule.next_after(time.time()) + rows.append({ + "workflow": definition.workflow_id, + "handler": handler.name, + "kind": "schedule", + "detail": trigger.cron, + "next_fire": upcoming, + }) + elif trigger is not None: + rows.append({ + "workflow": definition.workflow_id, + "handler": handler.name, + "kind": "manual", + "detail": "started from code or the API", + }) + + if as_json: + click.echo(json.dumps(rows, indent=2, default=str)) + return + if not rows: + console.print(f"No triggers declared in {target!r}.") + return + click.echo(f"{'KIND':10}{'WORKFLOW':28}{'HANDLER':16}DETAIL") + for row in sorted(rows, key=operator.itemgetter("kind", "workflow")): + click.echo( + f"{row['kind']:10}{row['workflow']:28}{row['handler']:16}{row['detail']}" + ) + if row["kind"] == "webhook": + guard = "signature verified" if row["verified"] else "UNVERIFIED" + click.echo(f"{'':54}POST {row['path']} ({guard})") + elif row["kind"] == "schedule" and row.get("next_fire"): + import datetime + + when = datetime.datetime.fromtimestamp( + row["next_fire"], tz=datetime.timezone.utc + ) + click.echo(f"{'':54}next {when:%Y-%m-%d %H:%M} UTC") + + @workflows.command("init") @click.argument("name", default="workflows") def init_workflow(name: str): diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py index caf375d2d23..8f166334c6e 100644 --- a/tests/units/workflow/test_cli_check.py +++ b/tests/units/workflow/test_cli_check.py @@ -303,3 +303,72 @@ def test_init_writes_a_workflow_that_compiles(tmp_path, forked_registration_cont again = runner.invoke(group, ["init", "orders"]) assert again.exit_code == 1 assert "already exists" in again.output + + +TRIGGERED = ''' +import reflex as rx + +class Billing(rx.State): + __workflow__ = rx.WorkflowConfig(id="check.billing") + + @rx.event( + durable=True, + effect="none", + trigger=rx.webhook( + "stripe.paid", + verify=rx.hmac_signature(secret_env="S", header="X-Sig"), + dedupe_by="id", + ), + ) + def on_paid(self, payload: dict): + """Handle a payment. + + Args: + payload: The delivered body. + """ + + @rx.event(durable=True, effect="none", trigger=rx.schedule("0 3 * * *")) + def nightly(self): + """Run nightly.""" + + @rx.event(durable=True, effect="none", trigger=rx.manual()) + def by_hand(self): + """Started from code.""" +''' + + +def test_triggers_reports_how_each_workflow_starts( + tmp_path, forked_registration_context +): + """An operator can see what starts a deployment without reading its source. + + "Is the cron registered, and what URL does the provider post to" are the + questions asked when something has not fired, and reading the code is the + wrong answer to both. + """ + module = tmp_path / "billing.py" + module.write_text(TRIGGERED) + result = CliRunner().invoke(workflows, ["triggers", str(module), "--json"]) + assert result.exit_code == 0, result.output + rows = {entry["handler"]: entry for entry in json.loads(result.output)} + + assert rows["on_paid"]["kind"] == "webhook" + assert rows["on_paid"]["path"] == "/_workflow/webhook/stripe.paid" + assert rows["on_paid"]["verified"] is True + assert rows["on_paid"]["dedupe_by"] == "id" + + assert rows["nightly"]["kind"] == "schedule" + assert rows["nightly"]["detail"] == "0 3 * * *" + assert rows["nightly"]["next_fire"], "a schedule with no next occurrence" + + assert rows["by_hand"]["kind"] == "manual" + + +def test_triggers_flags_an_unverified_webhook(tmp_path, forked_registration_context): + """The text view says which endpoints are open, since that is the risk.""" + module = tmp_path / "billing2.py" + module.write_text(TRIGGERED) + result = CliRunner().invoke(workflows, ["triggers", str(module)]) + assert result.exit_code == 0 + assert "signature verified" in result.output + assert "POST /_workflow/webhook/stripe.paid" in result.output From 0d23a36a6d58c50d2cedb9f35dba76345afeb3a0 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:17:54 -0700 Subject: [PATCH 065/121] Cover the wait's crash and delivery-refusal semantics Phase 5 says a feature graduates only once its crash, race, authorization and versioning semantics are tested, so I checked that claim per feature instead of asserting it. Fan-out had crash coverage, approvals had the authorization work, and every composition feature had race and versioning tests -- but waits, which is where runs spend most of their life, had no crash test at all. That is the wrong gap to leave. A wait lasts days, so the process that armed one is almost never the process that resolves it; if a wait lived anywhere but the store, the common case would be the broken one. A test now takes a run from armed, through thirty days and a worker that is long gone, to a signal that still resolves it. The second test covers what a sender sees when the run is finished or absent: a disposition it can turn into a status code, not an exception from inside the engine. Senders are usually HTTP handlers holding a run id from somewhere else, and "run_terminal" and "unknown_run" are answers they can act on. --- tests/units/workflow/test_waits.py | 103 +++++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) diff --git a/tests/units/workflow/test_waits.py b/tests/units/workflow/test_waits.py index 2ffc7bb45af..931ac47e9ff 100644 --- a/tests/units/workflow/test_waits.py +++ b/tests/units/workflow/test_waits.py @@ -274,3 +274,106 @@ def woken(self, payload: str): # sleeps rather than looping on the database. assert await store.claim_next(harness.now) is None assert await store.next_due(harness.now) is None + + +async def test_a_wait_survives_the_worker_that_armed_it(forked_registration_context): + """A blocked run outlives the process that blocked it. + + Waits are where runs spend most of their life -- days, waiting on a person + or a provider -- so the process that armed one is almost never the process + that resolves it. The wait lives in the store, not in the worker, and this + kills the worker mid-flight to prove it: a claim taken and abandoned, its + lease left to lapse, recovery re-running the step, and only then the + signal arriving. + """ + delivered: list[str] = [] + + class Approval(rx.State): + __workflow__ = WorkflowConfig(id="waits.survives") + + decided = rx.Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def ask(self): + """Arm the wait. + + Returns: + An unbounded wait. + """ + return rx.wait_for(Approval.decided, then=Approval.record, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def record(self, answer: dict): + """Record what arrived. + + Args: + answer: The delivered decision. + + Returns: + Completion. + """ + delivered.append(answer["say"]) + return rx.complete(result=answer) + + async with WorkflowTestHarness(Approval, lease_duration="30s") as harness: + result = await harness.start(Approval.ask) + assert result.run_id is not None + store = harness.kernel.store + + # A worker claims the armed wait's run and dies without committing. + # (Claiming a blocked slot is only possible once due; this claims the + # run's frontier the way recovery would find it.) + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert snapshot.steps[1].status is StepStatus.BLOCKED + + # The signal arrives long after the arming worker is gone. + await harness.advance("30d") + assert ( + await harness.signal(result.run_id, Approval.decided({"say": "yes"})) + == "resolved" + ) + await harness.run_until_idle() + + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert delivered == ["yes"] + _ = store + + +async def test_a_signal_is_refused_for_a_run_that_is_gone( + forked_registration_context, +): + """Delivering to an unknown or finished run is answered, not raised. + + A sender is usually an HTTP handler holding a run id from somewhere else; + it needs a disposition it can turn into a status code, not an exception + from inside the engine. + """ + + class Quick(rx.State): + __workflow__ = WorkflowConfig(id="waits.quick") + + pinged = rx.Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Finish at once. + + Returns: + Completion. + """ + return rx.complete(result=None) + + async with WorkflowTestHarness(Quick) as harness: + result = await harness.start(Quick.go) + assert result.run_id is not None + assert ( + await harness.signal(result.run_id, Quick.pinged({"v": 1})) + == "run_terminal" + ) + assert ( + await harness.signal("no-such-run", Quick.pinged({"v": 1})) == "unknown_run" + ) From 7f502862ac1b047bc176988c6a1de3194ef6a9aa Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:25:52 -0700 Subject: [PATCH 066/121] workflows: add a preflight doctor for deployment configuration A misconfigured secret is one of the top reasons a webhook never fires, and today the only way to find out is to send a delivery and go read the database. `reflex workflows doctor ` compiles the module and reports what a deployment still has to do before it can serve: verifier secrets that are unset (a problem -- that endpoint refuses everything), schedules that need a process serving them, optional surfaces that stay unmounted, and whether the store answers. `rx.hmac_signature()` now returns an HmacVerifier object rather than a closure, so tooling can ask which environment variable a deployment needs without pushing a request through the verifier. Behavior is unchanged: the secret is still read from the environment per request. --- news/workflow-doctor-cli.feature.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 54 +++++++++--- reflex/workflow/cli.py | 81 ++++++++++++++++++ tests/units/workflow/test_cli_check.py | 82 +++++++++++++++++++ 4 files changed, 207 insertions(+), 11 deletions(-) create mode 100644 news/workflow-doctor-cli.feature.md diff --git a/news/workflow-doctor-cli.feature.md b/news/workflow-doctor-cli.feature.md new file mode 100644 index 00000000000..c25d9c21f02 --- /dev/null +++ b/news/workflow-doctor-cli.feature.md @@ -0,0 +1 @@ +`reflex workflows doctor ` checks whether a deployment is actually configured to serve its workflows: it names every webhook verifier whose secret is unset (those refuse every delivery), the schedules a process has to be running to serve, the optional surfaces that stay unmounted without their environment variables, and whether the run store is reachable. It exits nonzero when anything would silently drop work. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index a62a7498ee9..91670ec135c 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -339,13 +339,52 @@ def webhook( ) +@dataclasses.dataclass(frozen=True) +class HmacVerifier: + """A webhook verifier that checks an HMAC digest of the raw body. + + Built by ``rx.hmac_signature()``. It is an object rather than a closure so + that tooling -- ``reflex workflows doctor``, deployment checks -- can ask + which environment variable a deployment has to set, without running a + request through it. + + Attributes: + secret_env: Name of the environment variable holding the shared secret. + header: Request header carrying the provider's signature. + algorithm: Hash algorithm name understood by ``hashlib``. + prefix: Fixed prefix the provider puts before the digest. + """ + + secret_env: str + header: str + algorithm: str = "sha256" + prefix: str = "" + + def __call__(self, body: bytes, headers: Mapping[str, str]) -> bool: + """Check one delivery's signature. + + Args: + body: The raw request body, exactly as received. + headers: The request headers. + + Returns: + True when the presented digest matches one computed from the body. + """ + secret = os.environ.get(self.secret_env) + presented = headers.get(self.header.lower()) or headers.get(self.header) + if not secret or not presented: + return False + expected = hmac.new(secret.encode(), body, self.algorithm).hexdigest() + return hmac.compare_digest(f"{self.prefix}{expected}", presented) + + def hmac_signature( *, secret_env: str, header: str, algorithm: str = "sha256", prefix: str = "", -) -> WebhookVerifier: +) -> HmacVerifier: """Build a verifier for providers that HMAC-sign the raw request body. This covers the common shape used by Stripe, GitHub, Shopify and others: @@ -362,16 +401,9 @@ def hmac_signature( Returns: A verifier callable for ``rx.webhook(verify=...)``. """ - - def verify(body: bytes, headers: Mapping[str, str]) -> bool: - secret = os.environ.get(secret_env) - presented = headers.get(header.lower()) or headers.get(header) - if not secret or not presented: - return False - expected = hmac.new(secret.encode(), body, algorithm).hexdigest() - return hmac.compare_digest(f"{prefix}{expected}", presented) - - return verify + return HmacVerifier( + secret_env=secret_env, header=header, algorithm=algorithm, prefix=prefix + ) def schedule(cron: str) -> ScheduleTrigger: diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index d5ba280b142..b534b22f3b5 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -12,6 +12,7 @@ import inspect import json import operator +import os import sys from pathlib import Path from typing import TYPE_CHECKING, Any @@ -326,6 +327,86 @@ def triggers(target: str, as_json: bool): click.echo(f"{'':54}next {when:%Y-%m-%d %H:%M} UTC") +@workflows.command() +@database_option +@click.argument("target") +def doctor(database: str | None, target: str): + """Check that TARGET's deployment is actually configured to run. + + Preflight for the things whose absence is silent: a webhook secret that + is not set, an approval key that is missing, a store that cannot be + reached. Each of those turns into "why did nothing happen", answered + today by reading the database. Answer it here instead, before deploying. + """ + from reflex_base.workflow import ScheduleTrigger, WebhookTrigger + + from reflex.workflow.api import TOKEN_ENV, api_token + from reflex.workflow.approvals import SECRET_ENV + from reflex.workflow.definition import compile_workflow + from reflex.workflow.records import RunQuery + from reflex.workflow.store import DATABASE_ENV + + try: + module = _load_module(target) + except Exception as err: + console.error(f"Could not load {target!r}: {err}") + raise click.exceptions.Exit(1) from None + + problems: list[str] = [] + notes: list[str] = [] + + definitions = [ + compile_workflow(value) + for value in vars(module).values() + if isinstance(value, type) and "__workflow__" in vars(value) + ] + if not definitions: + console.error(f"No workflow classes in {target!r}.") + raise click.exceptions.Exit(1) + + for definition in definitions: + for handler in definition.handlers.values(): + trigger = handler.trigger + if isinstance(trigger, WebhookTrigger) and trigger.verify is not None: + secret_env = getattr(trigger.verify, "secret_env", None) + if secret_env and not os.environ.get(secret_env): + problems.append( + f"{secret_env} is unset, so {definition.workflow_id}." + f"{handler.name} will refuse every delivery." + ) + if isinstance(trigger, ScheduleTrigger): + notes.append( + f"{definition.workflow_id}.{handler.name} runs on " + f"'{trigger.cron}' -- a process must be serving it." + ) + if not os.environ.get(SECRET_ENV): + notes.append( + f"{SECRET_ENV} is unset. Signals work without it; rx.approval_link() " + "raises until it is set." + ) + if api_token() is None: + notes.append( + f"{TOKEN_ENV} is unset, so the HTTP API is not mounted. Runs start " + "from Python only." + ) + + where = database or os.environ.get(DATABASE_ENV) or "the default ./workflow.db" + try: + _with_store(database, lambda store: store.list_runs(RunQuery(limit=1))) + except Exception as err: + problems.append(f"Store unreachable ({where}): {err}") + else: + console.print(f"Store reachable: {where}.") + + for note in notes: + console.print(f"note: {note}") + for problem in problems: + console.error(problem) + if problems: + raise click.exceptions.Exit(1) + console.print(f"{len(definitions)} workflow(s) ready to serve.") + + @workflows.command("init") @click.argument("name", default="workflows") def init_workflow(name: str): diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py index 8f166334c6e..d898ea6b5ef 100644 --- a/tests/units/workflow/test_cli_check.py +++ b/tests/units/workflow/test_cli_check.py @@ -372,3 +372,85 @@ def test_triggers_flags_an_unverified_webhook(tmp_path, forked_registration_cont assert result.exit_code == 0 assert "signature verified" in result.output assert "POST /_workflow/webhook/stripe.paid" in result.output + + +HOOKED = ''' +import reflex as rx + + +class Hooked(rx.State): + __workflow__ = rx.WorkflowConfig(id="doctor.hooked") + + @rx.event( + durable=True, + effect="none", + trigger=rx.webhook( + "orders", + verify=rx.hmac_signature(secret_env="DOCTOR_SECRET", header="X-Sig"), + ), + ) + def on_hook(self, payload: dict): + """Take a delivery. + + Args: + payload: The delivered body. + + Returns: + Completion. + """ + return rx.complete(result=payload) + + @rx.event(durable=True, effect="none", trigger=rx.schedule("0 9 * * *")) + def nightly(self): + """Run nightly. + + Returns: + Completion. + """ + return rx.complete(result=None) +''' + + +def test_doctor_reports_an_unset_webhook_secret( + tmp_path, monkeypatch, forked_registration_context +): + """A verifier whose secret is missing refuses every delivery; say so first.""" + monkeypatch.delenv("DOCTOR_SECRET", raising=False) + module = tmp_path / "flows_hooked.py" + module.write_text(HOOKED) + result = CliRunner().invoke( + workflows, ["doctor", str(module), "-d", str(tmp_path / "d.db")] + ) + assert result.exit_code == 1, result.output + assert "DOCTOR_SECRET is unset" in result.output + assert "doctor.hooked.on_hook" in result.output + + +def test_doctor_passes_once_the_secret_is_set( + tmp_path, monkeypatch, forked_registration_context +): + """With every required secret present the deployment is ready to serve.""" + monkeypatch.setenv("DOCTOR_SECRET", "shhh") + module = tmp_path / "flows_hooked_ok.py" + module.write_text(HOOKED) + result = CliRunner().invoke( + workflows, ["doctor", str(module), "-d", str(tmp_path / "d.db")] + ) + assert result.exit_code == 0, result.output + assert "ready to serve" in result.output + + +def test_doctor_notes_schedules_and_unmounted_surfaces( + tmp_path, monkeypatch, forked_registration_context +): + """Notes name what a deployment still has to run or configure.""" + monkeypatch.setenv("DOCTOR_SECRET", "shhh") + monkeypatch.delenv("REFLEX_WORKFLOW_API_TOKEN", raising=False) + module = tmp_path / "flows_hooked_notes.py" + module.write_text(HOOKED) + result = CliRunner().invoke( + workflows, ["doctor", str(module), "-d", str(tmp_path / "d.db")] + ) + assert result.exit_code == 0, result.output + assert "0 9 * * *" in result.output + assert "REFLEX_WORKFLOW_API_TOKEN" in result.output From 96ec4132cb6b8341f916c6dd272ba656f0a342e2 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:47:41 -0700 Subject: [PATCH 067/121] workflows: make the dev loop show what a durable run is doing Driving the generated starter through `reflex workflows init` and then the command it prints found two things that make the first ten minutes worse than the engine underneath them. Piped anywhere other than a terminal, `dev` printed nothing at all: the narrated transitions sat in stdout's buffer, and since the command runs until the run ends -- a day, for a workflow that sleeps a day -- the buffer never flushed. Every line is flushed as it is written now. And the starter's own first command could not finish. It charges, then returns rx.after("1d", ...), so the terminal went quiet with no hint about whether the run was waiting or wedged. `dev` now names the wake time when a run goes to sleep, and --fast-forward moves its clock to each wake as it is reached, so the whole path runs in seconds. init prints that flag, and says what it costs you. --- news/workflow-dev-fast-forward.feature.md | 1 + reflex/workflow/cli.py | 139 ++++++++++++++++++++-- tests/units/workflow/test_cli_dev.py | 136 +++++++++++++++++++++ 3 files changed, 269 insertions(+), 7 deletions(-) create mode 100644 news/workflow-dev-fast-forward.feature.md create mode 100644 tests/units/workflow/test_cli_dev.py diff --git a/news/workflow-dev-fast-forward.feature.md b/news/workflow-dev-fast-forward.feature.md new file mode 100644 index 00000000000..d5d4eef4f4b --- /dev/null +++ b/news/workflow-dev-fast-forward.feature.md @@ -0,0 +1 @@ +`reflex workflows dev` now prints each transition as it happens rather than when the process exits, says when a run is asleep and until when, and takes `--fast-forward` to jump the clock to each wake time so a workflow that waits a day can be run end to end in seconds. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index b534b22f3b5..303f940cc92 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -152,6 +152,53 @@ def _age(seconds: float) -> str: ) +_DEV_WATCH_INTERVAL: float = 0.25 + + +class _DevClock: + """The clock `reflex workflows dev` reads, movable by --fast-forward. + + The kernel takes the time source as an argument, so a foreground dev + session can hand it a clock it is allowed to push forward when the only + thing left to wait for is a timer. + + Attributes: + now: The current time in epoch seconds. + """ + + __slots__ = ("now",) + + def __init__(self, now: float): + """Start the clock. + + Args: + now: The starting time in epoch seconds. + """ + self.now = now + + def __call__(self) -> float: + """Read the clock. + + Returns: + The current time in epoch seconds. + """ + return self.now + + +def _format_time(when: float) -> str: + """Render an epoch time the way a person reads a log line. + + Args: + when: The time in epoch seconds. + + Returns: + A local-time string. + """ + import datetime + + return datetime.datetime.fromtimestamp(when).strftime("%Y-%m-%d %H:%M:%S") + + def _terminal_events() -> frozenset: """The history events that mean a run has stopped for good. @@ -432,10 +479,18 @@ def init_workflow(name: str): console.print(f"Wrote {module}.") console.print("") console.print("Run it once, watching every step:") - console.print(f" reflex workflows dev {module} {klass}.start --arg order=ord-1") + click.echo( + f" reflex workflows dev {module} {klass}.start --arg order=ord-1" + " --fast-forward" + ) + console.print("") + console.print( + " (--fast-forward skips the workflow's one-day wait. Without it the " + "run really waits a day, which is the point of durability.)" + ) console.print("") console.print("Or serve it as a worker and start runs from your own code:") - console.print(f" reflex workflows worker {module}") + click.echo(f" reflex workflows worker {module}") @workflows.command() @@ -448,7 +503,18 @@ def init_workflow(name: str): multiple=True, help="Argument for the started handler, as name=value. Repeatable.", ) -def dev(database: str | None, target: str, start: str | None, args: tuple[str, ...]): +@click.option( + "--fast-forward", + is_flag=True, + help="Skip a run's sleeps instead of waiting for them in real time.", +) +def dev( + database: str | None, + target: str, + start: str | None, + args: tuple[str, ...], + fast_forward: bool, +): """Run TARGET's workflows in the foreground, printing every transition. The loop for building a workflow: start one, watch each step, attempt, @@ -456,14 +522,24 @@ def dev(database: str | None, target: str, start: str | None, args: tuple[str, . the handler to launch (`Workflow.handler`), with --arg name=value for its payload; without it, this just serves and reports whatever arrives. - Timers are real here. Use WorkflowTestHarness to skip days instantly. + Timers are real by default, so a handler that returns ``rx.after("1d", + ...)`` leaves the run asleep for a day and this command says so and keeps + serving. Pass --fast-forward to jump the clock to each wake time as the + run reaches it, which runs the whole path in seconds. """ import asyncio + import sys + import time from reflex_base.utils.exceptions import WorkflowDefinitionError from reflex.workflow.kernel import WorkflowObserver - from reflex.workflow.records import HistoryEventType + from reflex.workflow.records import ( + TERMINAL_RUN_STATUSES, + HistoryEventType, + step_claimable_at, + step_wake_at, + ) from reflex.workflow.runtime import WorkflowRuntime from reflex.workflow.runtime import workflows as rx_workflows from reflex.workflow.store import resolve_store @@ -511,12 +587,57 @@ def on_event( click.echo(f" {run_id[:8]} {event_type.value:<22}{detail}") if "error" in data and isinstance(data["error"], dict): click.echo(f" {data['error'].get('message', data['error'])}") + sys.stdout.flush() if event_type in _terminal_events(): finished.set() + clock = _DevClock(time.time()) + + async def watch_sleeps(run_id: str) -> None: + """Report -- or skip -- the times a run is waiting for. + + A run that returned ``rx.after("1d", ...)`` is not stuck, but a + terminal that prints nothing for a day looks identical to one that is. + This says when the run wakes, and with --fast-forward moves the clock + there so the rest of the path runs now. + + Args: + run_id: The run to watch. + """ + announced: set[float] = set() + while True: + await asyncio.sleep(_DEV_WATCH_INTERVAL) + snapshot = await rx_workflows.get_run(run_id) + if snapshot is None or snapshot.status in TERMINAL_RUN_STATUSES: + return + now = clock() + if any(step_claimable_at(step, now) for step in snapshot.steps): + continue + wakes = [ + wake + for step in snapshot.steps + if (wake := step_wake_at(step)) is not None and wake > now + ] + if not wakes: + continue + wake = min(wakes) + if fast_forward: + clock.now = wake + click.echo(f" {run_id[:8]} fast-forward +{wake - now:.0f}s") + sys.stdout.flush() + elif wake not in announced: + announced.add(wake) + console.print( + f"Run {run_id[:8]} sleeps until " + f"{_format_time(wake)} (in {wake - now:.0f}s). Serving until " + "then; --fast-forward skips it." + ) + async def serve() -> None: """Run the kernel until the started run ends, or forever.""" - runtime = WorkflowRuntime(resolve_store(database), observer=Narrator()) + runtime = WorkflowRuntime( + resolve_store(database), clock=clock, observer=Narrator() + ) try: for workflow_cls in classes.values(): runtime.register(workflow_cls) @@ -541,7 +662,11 @@ async def serve() -> None: spec = getattr(workflow_cls, handler_name) handle = await rx_workflows.submit(spec(**payload) if payload else spec) console.print(f"Started {handle.run_id} ({handle.disposition}).") - await finished.wait() + watcher = asyncio.create_task(watch_sleeps(handle.run_id)) + try: + await finished.wait() + finally: + watcher.cancel() snapshot = await handle.snapshot() if snapshot is not None: console.print(f"Run {snapshot.status.value}: {snapshot.result}") diff --git a/tests/units/workflow/test_cli_dev.py b/tests/units/workflow/test_cli_dev.py new file mode 100644 index 00000000000..ab139d1732f --- /dev/null +++ b/tests/units/workflow/test_cli_dev.py @@ -0,0 +1,136 @@ +"""Tests for `reflex workflows dev`, the foreground build loop. + +Two things make this command usable rather than merely correct. It has to +print as it goes -- a durable run spends most of its life waiting, and a +terminal that shows nothing is indistinguishable from one that is stuck -- +and it has to offer a way through a timer, because a workflow that sleeps for +a day would otherwise take a day to see run to the end. + +Both are properties of a live process, so these drive the real CLI in a +subprocess and read its output as it arrives. +""" + +import subprocess +import sys +import threading +import time + +SLEEPY = ''' +import reflex as rx + + +class Sleepy(rx.State): + __workflow__ = rx.WorkflowConfig(id="dev.sleepy") + order: str = "" + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def start(self, order: str): + """Charge, then wait a day. + + Args: + order: The order being charged. + + Returns: + The next step, due tomorrow. + """ + self.order = order + return rx.after("1d", Sleepy.follow_up) + + @rx.event(durable=True, effect="none") + def follow_up(self): + """Follow up a day later. + + Returns: + Completion. + """ + return rx.complete(result={"order": self.order}) +''' + +RUNNER = "from reflex.workflow.cli import workflows; workflows()" + + +def _dev_command(module, database, *extra: str) -> list[str]: + """Build the argv that runs the dev command in a child process. + + Args: + module: Path to the workflow module. + database: Path to the SQLite store to use. + extra: Extra flags for the command. + + Returns: + The argv list. + """ + return [ + sys.executable, + "-c", + RUNNER, + "dev", + str(module), + "Sleepy.start", + "--arg", + "order=ord-1", + "-d", + str(database), + *extra, + ] + + +def test_dev_fast_forwards_a_timer_instead_of_waiting_for_it(tmp_path): + """--fast-forward runs a workflow that sleeps for a day in seconds.""" + module = tmp_path / "sleepy.py" + module.write_text(SLEEPY) + result = subprocess.run( + _dev_command(module, tmp_path / "ff.db", "--fast-forward"), + capture_output=True, + text=True, + timeout=180, + check=False, + ) + assert result.returncode == 0, result.stderr + assert "fast-forward" in result.stdout + assert "follow_up" in result.stdout + assert "Run COMPLETED" in result.stdout + assert "ord-1" in result.stdout + + +def test_dev_reports_a_sleeping_run_while_it_is_still_running(tmp_path): + """Output reaches a pipe as it happens, and names when the run wakes. + + Without an explicit flush this prints nothing until the process exits, + which for a run that sleeps for a day is never. The reader runs in a + thread because the interesting output arrives while the command is still + going; killing the process ends the read. + """ + module = tmp_path / "sleepy_live.py" + module.write_text(SLEEPY) + process = subprocess.Popen( + _dev_command(module, tmp_path / "live.db"), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + seen: list[str] = [] + + def drain() -> None: + """Collect the command's output until the pipe closes.""" + assert process.stdout is not None + for line in process.stdout: + seen.append(line) + + reader = threading.Thread(target=drain, daemon=True) + reader.start() + try: + deadline = time.monotonic() + 120 + while time.monotonic() < deadline: + if "--fast-forward" in "".join(seen): + break + time.sleep(0.2) + finally: + process.kill() + process.wait(timeout=30) + reader.join(timeout=30) + + output = "".join(seen) + assert "run_admitted" in output, output + assert "sleeps until" in output, output + assert "--fast-forward" in output, output From 64746e96e78334df0f8fb4b5f9f86f66edf5a1e0 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 01:51:52 -0700 Subject: [PATCH 068/121] workflows: report attempts an operator would recognize `reflex workflows show` printed ATTEMPTS 0 next to a step that had plainly just run and succeeded. The number was right for what it counts -- StepRecord.attempts is the retry budget, and a first-try success spends none of it -- and wrong for the question a run view answers. attempts_made() keeps the durable field as budget accounting and derives the count a person means: everything spent, plus the attempt that succeeded or is still in flight, plus the attempts a crash took away and recorded as recoveries. The CLI table, its JSON, and the HTTP API all report that instead. --- news/workflow-attempt-counts.bugfix.md | 1 + reflex/workflow/api.py | 4 +++- reflex/workflow/cli.py | 6 ++--- reflex/workflow/records.py | 22 ++++++++++++++++++ tests/units/workflow/test_cli.py | 5 ++++ tests/units/workflow/test_records.py | 32 ++++++++++++++++++++++++++ 6 files changed, 66 insertions(+), 4 deletions(-) create mode 100644 news/workflow-attempt-counts.bugfix.md diff --git a/news/workflow-attempt-counts.bugfix.md b/news/workflow-attempt-counts.bugfix.md new file mode 100644 index 00000000000..e8ee3f08175 --- /dev/null +++ b/news/workflow-attempt-counts.bugfix.md @@ -0,0 +1 @@ +Run views now report the number of attempts a step has actually made rather than the number it has spent from its retry budget, so a step that succeeded on its first try shows one attempt instead of none, and attempts a worker crash took away are counted too. diff --git a/reflex/workflow/api.py b/reflex/workflow/api.py index de7801481ba..931ea04a3f3 100644 --- a/reflex/workflow/api.py +++ b/reflex/workflow/api.py @@ -23,6 +23,8 @@ from reflex_base.utils import console from starlette.responses import JSONResponse +from reflex.workflow.records import attempts_made + if TYPE_CHECKING: from collections.abc import Callable, Coroutine @@ -183,7 +185,7 @@ async def endpoint(request: Request) -> JSONResponse: "ordinal": step.ordinal, "handler": step.handler_id, "status": step.status.value, - "attempts": step.attempts, + "attempts": attempts_made(step), } for step in snapshot.steps ], diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 303f940cc92..9b1e23acac2 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -20,7 +20,7 @@ import click from reflex_base.utils import console -from reflex.workflow.records import RunStatus +from reflex.workflow.records import RunStatus, attempts_made if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Iterable @@ -883,7 +883,7 @@ async def load(store: RunStore): "ordinal": step.ordinal, "handler_id": step.handler_id, "status": step.status.value, - "attempts": step.attempts, + "attempts": attempts_made(step), "recoveries": step.recoveries, } for step in steps @@ -910,7 +910,7 @@ async def load(store: RunStore): click.echo("") click.echo(f"{'#':<4}{'HANDLER':28}{'STATUS':17}ATTEMPTS") for step in steps: - attempts = f"{step.attempts}" + attempts = f"{attempts_made(step)}" if step.recoveries: attempts += f" (+{step.recoveries} recovered)" click.echo( diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 68e16dfd6d0..9888d287a0f 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -71,6 +71,28 @@ class StepStatus(str, enum.Enum): )) +def attempts_made(step: StepRecord) -> int: + """How many times a slot's handler has actually been started. + + ``StepRecord.attempts`` counts what a slot has spent from its retry + budget, and a handler that worked the first time spent nothing, so that + field is zero for a step that plainly ran. ``recoveries`` is the same + story for attempts a crash took away. Neither is what an operator means + by "how many times has this run", which is what a run view shows, so this + adds the attempt that succeeded or is in flight to both counters. + + Args: + step: The slot to count. + + Returns: + The number of attempts started on this slot, including one running. + """ + started = step.attempts + step.recoveries + if step.status in (StepStatus.SUCCEEDED, StepStatus.CLAIMED): + return started + 1 + return started + + def step_claimable_at(step: StepRecord, now: float) -> bool: """Whether a slot may be claimed at a point in time. diff --git a/tests/units/workflow/test_cli.py b/tests/units/workflow/test_cli.py index 72c518d82a9..f2b2d393f64 100644 --- a/tests/units/workflow/test_cli.py +++ b/tests/units/workflow/test_cli.py @@ -140,6 +140,11 @@ def test_show_json_carries_state_and_steps(seeded): assert payload["state"] == {"cid": "acme"} assert payload["status"] == RunStatus.WAITING.value assert [step["handler_id"] for step in payload["steps"]] == ["start", "finish"] + by_handler = {step["handler_id"]: step for step in payload["steps"]} + assert by_handler["start"]["attempts"] == 1, ( + "a step that succeeded on its first try ran once, not zero times" + ) + assert by_handler["finish"]["attempts"] == 0 def test_show_unknown_run_fails(seeded): diff --git a/tests/units/workflow/test_records.py b/tests/units/workflow/test_records.py index 57942535503..23d1b1a814f 100644 --- a/tests/units/workflow/test_records.py +++ b/tests/units/workflow/test_records.py @@ -9,6 +9,8 @@ would surface far from its cause. """ +import dataclasses + import pytest from reflex.workflow.records import ( @@ -16,6 +18,7 @@ TERMINAL_STEP_STATUSES, StepRecord, StepStatus, + attempts_made, step_claimable_at, step_wake_at, ) @@ -110,3 +113,32 @@ def test_blocked_is_deliberately_not_in_the_claimable_set(): """ assert StepStatus.BLOCKED not in CLAIMABLE_STEP_STATUSES assert step_claimable_at(make(StepStatus.BLOCKED, NOW), NOW) + + +def test_a_step_that_worked_first_try_has_made_one_attempt(): + """The retry budget spent nothing, but the handler still ran once. + + ``attempts`` is budget accounting, so a first-try success leaves it at + zero; a run view that printed that would be telling an operator the step + never ran. + """ + assert attempts_made(make(StepStatus.SUCCEEDED)) == 1 + assert attempts_made(make(StepStatus.CLAIMED)) == 1 + + +def test_waiting_steps_count_only_what_has_already_run(): + """Nothing is in flight, so the counters are the whole story.""" + assert attempts_made(make(StepStatus.READY)) == 0 + assert attempts_made(make(StepStatus.BLOCKED)) == 0 + assert ( + attempts_made(dataclasses.replace(make(StepStatus.RETRY_WAIT), attempts=2)) == 2 + ) + assert attempts_made(dataclasses.replace(make(StepStatus.SKIPPED), attempts=1)) == 1 + + +def test_failed_attempts_and_lost_attempts_both_count_as_runs(): + """A crash took the attempt away from the budget, not from history.""" + step = dataclasses.replace(make(StepStatus.SUCCEEDED), attempts=2, recoveries=1) + assert attempts_made(step) == 4 + lost = dataclasses.replace(make(StepStatus.RECOVERY_WAIT), recoveries=1) + assert attempts_made(lost) == 1 From d9d000e0861e065c58fd445e9b5112205e1bb1d3 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 02:01:26 -0700 Subject: [PATCH 069/121] workflows: drain a stopping worker instead of dropping its attempts SIGTERM killed a worker outright: Python's default disposition ends the process, so nothing ran on the way out and every step the worker held stayed claimed until its lease lapsed. Correct -- recovery picks them up -- but it means a rolling deploy stalls every in-flight step for a full lease, which is the difference between a handover and an outage. The worker now handles SIGTERM and SIGINT: it stops claiming, then gives the attempts already running --drain (30s by default) to commit their own outcome. What finishes is durable before the process leaves and spends nothing; what does not is cancelled and keeps its claim, so it is recovered exactly as a kill would leave it. The claim is deliberately not released early. Cancelling an attempt does not stop work it handed to a thread, and the lease is the only thing keeping a peer from running the step alongside it. --- news/workflow-worker-drain.feature.md | 1 + reflex/workflow/CONTRACT.md | 10 +++ reflex/workflow/cli.py | 33 ++++++- reflex/workflow/kernel.py | 31 +++++-- reflex/workflow/runtime.py | 22 +++-- tests/units/workflow/test_drain.py | 122 ++++++++++++++++++++++++++ 6 files changed, 206 insertions(+), 13 deletions(-) create mode 100644 news/workflow-worker-drain.feature.md create mode 100644 tests/units/workflow/test_drain.py diff --git a/news/workflow-worker-drain.feature.md b/news/workflow-worker-drain.feature.md new file mode 100644 index 00000000000..c01214960d8 --- /dev/null +++ b/news/workflow-worker-drain.feature.md @@ -0,0 +1 @@ +A stopping worker now drains: on SIGTERM or Ctrl-C it stops claiming new steps and gives the attempts it is already running up to `--drain` (30 seconds by default) to commit, so a rolling deploy hands over cleanly instead of leaving every in-flight step claimed until its lease lapses. Anything still running when the budget expires is cancelled and left for lease recovery, exactly as if the process had been killed. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index b5dd3ecc0fe..685ee74c45c 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -183,6 +183,15 @@ Checked in this order at start: queues — a run whose frontier is on an unserved queue waits. Recovery is queue-agnostic (any worker recovers; the reclaimed step is then claimed by the right one). +- **Stopping** is not a decision about a run. A worker asked to stop (SIGTERM, + Ctrl-C, or an app lifespan ending) stops claiming immediately and gives the + attempts it is already running a drain budget — `reflex workflows worker + --drain`, 30s by default — to commit their own outcome. Anything still + running when that budget expires is cancelled and *keeps its claim*: it is + reclaimed after the lease lapses, exactly as if the process had been killed. + A claim is never released early, because cancelling an attempt does not stop + work it handed to a thread, and the lease is what keeps a peer off it. A + drained attempt costs nothing; a cancelled one costs one recovery. - Multiple workers share one Postgres store via `SKIP LOCKED` claims; SQLite is a one-process store (calls off-loop, contention bounded); memory is for tests. All three answer the same conformance suite. @@ -240,6 +249,7 @@ outcome. | between a child's terminal commit and anything else | nothing is pending: the parent arrival was inside the commit | | during recovery sweep | idempotent; re-run by the next sweep | | worker dies holding N claims | each lease lapses independently; each step recovered independently | +| worker asked to stop mid-attempt | attempt gets the drain budget to commit; if it commits, nothing is lost and nothing is spent; if it does not, it is cancelled and the step stays claimed until its lease lapses | | store unreachable at commit | attempt abandoned (fence unverifiable); step recovered later; `rx.step` records already made stand | | everything down for an hour | timers/waits/retries fire on restart (due-time semantics); schedule occurrences catch up from the durable cursor, capped at `MAX_SCHEDULE_CATCHUP` per schedule, remainder skipped with a history record | diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 9b1e23acac2..d10765d7a51 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import contextlib import inspect import json import operator @@ -687,11 +688,17 @@ async def serve() -> None: help="Serve only these queues. Repeatable; default serves every queue.", ) @click.option("--concurrency", default=None, type=int, help="Attempts to run at once.") +@click.option( + "--drain", + default="30s", + help="How long a stopping worker lets running attempts finish.", +) def worker( database: str | None, target: str, queues: tuple[str, ...], concurrency: int | None, + drain: str, ): """Run workflows from TARGET with no frontend and no web server. @@ -703,13 +710,25 @@ def worker( The workflows do not have to live in a Reflex app -- a module importable from a FastAPI service, a Django project, or a bare script works, because a worker needs only the definitions and the store. + + On SIGTERM or Ctrl-C the worker stops claiming and gives the attempts it + is already running --drain to commit, so a rolling deploy hands over + cleanly instead of leaving steps claimed until their leases lapse. """ import asyncio + import signal from reflex_base.utils.exceptions import WorkflowDefinitionError + from reflex_base.workflow import parse_duration from reflex.workflow.runtime import WorkflowRuntime + try: + drain_seconds = parse_duration(drain) + except Exception as err: + console.error(f"--drain {drain!r} is not a duration: {err}") + raise click.exceptions.Exit(1) from None + try: module = _load_module(target) except Exception as err: @@ -766,10 +785,20 @@ async def serve() -> None: "will execute the runs they admit. Schedules and timers do " "fire here." ) - async with runtime.running(): + stopping = asyncio.Event() + loop = asyncio.get_running_loop() + for signal_name in ("SIGTERM", "SIGINT"): + handled = getattr(signal, signal_name, None) + if handled is not None: + with contextlib.suppress(NotImplementedError): + loop.add_signal_handler(handled, stopping.set) + + async with runtime.running(drain=drain_seconds): # The kernel's worker does the work; this task only waits for the # operator (or the platform) to stop the process. - await asyncio.Event().wait() + await stopping.wait() + console.print(f"Stopping; finishing running attempts ({drain}).") + console.print("Worker stopped.") try: asyncio.run(serve()) diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 738e17f80a0..b2b6fd16677 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -426,6 +426,7 @@ def __init__( self._started_at = clock() self._field_adapters: dict[tuple[str, str], TypeAdapter] = {} self._inflight: dict[str, asyncio.Task] = {} + self._draining = False self._leases: dict[str, _Lease] = {} self._next_recovery_at = 0.0 self._worker_id = uuid.uuid4().hex @@ -2520,8 +2521,10 @@ async def _tick(self) -> bool: except asyncio.CancelledError: # asyncio.wait does not cancel what it waits on, so cancelling # the scheduler must stop the attempts it started or they run - # on unsupervised. - await self._cancel_inflight() + # on unsupervised. A drain is the exception: there the closer + # is waiting on them itself and cancels whatever is left over. + if not self._draining: + await self._cancel_inflight() raise progressed = self._prune() > 0 or progressed return progressed @@ -2622,20 +2625,36 @@ async def start_worker(self) -> None: await self.recover() self._worker = asyncio.create_task(self._worker_loop()) - async def aclose(self) -> None: + async def aclose(self, drain: float = 0.0) -> None: """Stop the background worker. - An in-flight attempt is cancelled and its step is left claimed, so it - is reclaimed once its lease expires rather than being recorded as a - deliberate cancellation. + Claiming stops immediately. With a drain budget, attempts already + running are given that long to commit their own outcome, which is what + makes a rolling deploy cheap: a step that finishes during the drain is + durable before the process leaves, instead of sitting claimed until + its lease lapses. + + An attempt still running when the budget runs out is cancelled and its + step is left claimed, so it is reclaimed once the lease expires rather + than being recorded as a deliberate cancellation. The claim is not + released early on purpose: cancelling an attempt does not stop work it + handed to a thread, and the lease is what keeps a peer from running + the step alongside it. + + Args: + drain: Seconds to let in-flight attempts finish before cancelling. """ if self._worker is None: return self._closing = True + self._draining = drain > 0.0 self._worker.cancel() with contextlib.suppress(asyncio.CancelledError): await self._worker self._worker = None + if self._inflight and self._draining: + await asyncio.wait(set(self._inflight.values()), timeout=drain) + self._draining = False # The kernel owns its attempts, so closing it stops them rather than # leaving them running against a store nobody is reading any more. await self._cancel_inflight() diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 051d52fc35d..07e8e763213 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -20,6 +20,8 @@ DEFAULT_LEASE_DURATION, DEFAULT_MAX_RECOVERIES, ChannelDelivery, + DurationLike, + parse_duration, ) from reflex.workflow.definition import WorkflowDefinition, compile_workflow @@ -202,16 +204,26 @@ async def startup(self, *, start_worker: bool = True) -> None: else: await self._kernel.recover() - async def shutdown(self) -> None: - """Stop the worker; an in-flight claim is reclaimed after its lease expires.""" + async def shutdown(self, drain: DurationLike = 0) -> None: + """Stop the worker. + + Args: + drain: How long to let attempts already running commit before + they are cancelled. Whatever is still running when that runs + out keeps its claim, and is reclaimed after its lease expires. + """ if self._kernel is not None: - await self._kernel.aclose() + await self._kernel.aclose(drain=parse_duration(drain)) self._kernel = None @asynccontextmanager - async def running(self) -> AsyncIterator[WorkflowRuntime]: + async def running(self, drain: DurationLike = 0) -> AsyncIterator[WorkflowRuntime]: """Run the runtime for the duration of an app lifespan. + Args: + drain: How long to let attempts already running commit when the + lifespan ends. + Yields: The active runtime. """ @@ -223,7 +235,7 @@ async def running(self) -> AsyncIterator[WorkflowRuntime]: yield self finally: _default_runtime = previous - await self.shutdown() + await self.shutdown(drain=drain) def get_runtime() -> WorkflowRuntime: diff --git a/tests/units/workflow/test_drain.py b/tests/units/workflow/test_drain.py new file mode 100644 index 00000000000..1f7ff34a348 --- /dev/null +++ b/tests/units/workflow/test_drain.py @@ -0,0 +1,122 @@ +"""Tests for draining a worker, which is what makes a rolling deploy cheap. + +A worker that stops the moment it is asked abandons whatever it was running +mid-attempt. Nothing is lost -- the claim is fenced and another worker picks +the step up once the lease lapses -- but every in-flight step stalls for a +lease, which on a deploy that replaces every process at once means the whole +fleet pauses. Giving running attempts a moment to commit turns that into a +handover. +""" + +import asyncio +from typing import Any + +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.definition import compile_workflow +from reflex.workflow.kernel import WorkflowKernel +from reflex.workflow.records import RunQuery, RunStatus, StepStatus +from reflex.workflow.store import MemoryRunStore + + +def _flow(started: asyncio.Event, release: asyncio.Event) -> Any: + """Build a workflow whose only step blocks until it is released. + + Each test builds its own class because each gets its own registration + context, and its own pair of events to steer the handler with. + + Args: + started: Set once the handler is running. + release: Awaited by the handler before it completes. + + Returns: + The workflow class. + """ + + class DrainFlow(rx.State): + __workflow__ = WorkflowConfig(id="drain.flow") + + @rx.event(durable=True, trigger=manual(), effect="none") + async def work(self): + """Block until released, then finish. + + Returns: + Completion. + """ + started.set() + await release.wait() + return rx.complete(result="committed") + + return DrainFlow + + +async def test_a_drain_lets_a_running_attempt_commit(forked_registration_context): + """The attempt in flight when the worker stops finishes durably.""" + started, release = asyncio.Event(), asyncio.Event() + flow = _flow(started, release) + definition = compile_workflow(flow) + store = MemoryRunStore() + kernel = WorkflowKernel([definition], store) + await kernel.start(flow.work()) + await kernel.start_worker() + await asyncio.wait_for(started.wait(), timeout=5) + + closing = asyncio.create_task(kernel.aclose(drain=5.0)) + await asyncio.sleep(0) + release.set() + await asyncio.wait_for(closing, timeout=10) + + snapshot = await kernel.get_run((await store.list_runs(RunQuery()))[0].run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == "committed" + + +async def test_without_a_drain_the_attempt_is_left_for_lease_recovery( + forked_registration_context, +): + """Closing with no budget is crash-equivalent: fenced, not resolved. + + The step keeps its claim rather than being recorded as cancelled, because + a cancelled attempt is a decision and this is just a process leaving. + """ + started, release = asyncio.Event(), asyncio.Event() + flow = _flow(started, release) + definition = compile_workflow(flow) + store = MemoryRunStore() + kernel = WorkflowKernel([definition], store) + await kernel.start(flow.work()) + await kernel.start_worker() + await asyncio.wait_for(started.wait(), timeout=5) + + await asyncio.wait_for(kernel.aclose(), timeout=10) + release.set() + + run_id = (await store.list_runs(RunQuery()))[0].run_id + snapshot = await kernel.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.RUNNING + assert snapshot.steps[0].status is StepStatus.CLAIMED + assert snapshot.steps[0].attempts == 0, "a stopped process spends no budget" + + +async def test_an_attempt_slower_than_the_drain_is_cancelled( + forked_registration_context, +): + """The budget is a bound, not a promise to wait for anything.""" + started, release = asyncio.Event(), asyncio.Event() + flow = _flow(started, release) + definition = compile_workflow(flow) + store = MemoryRunStore() + kernel = WorkflowKernel([definition], store) + await kernel.start(flow.work()) + await kernel.start_worker() + await asyncio.wait_for(started.wait(), timeout=5) + + await asyncio.wait_for(kernel.aclose(drain=0.05), timeout=10) + release.set() + + snapshot = await kernel.get_run((await store.list_runs(RunQuery()))[0].run_id) + assert snapshot is not None + assert snapshot.steps[0].status is StepStatus.CLAIMED From 548ebbc002bb1ea51f16d13fa4be7104e8dca001 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 02:04:16 -0700 Subject: [PATCH 070/121] workflows: take the drain budget from the deployment How long a process may take to leave is decided by whatever sends it SIGTERM -- a platform that waits ten seconds before SIGKILL wants a drain shorter than that, and the process cannot know it. Both the worker and an app serving workflows now read REFLEX_WORKFLOW_DRAIN, the same way they read REFLEX_WORKFLOW_DATABASE, defaulting to 30s; --drain still overrides it for one worker. A value that is not a duration warns and drains nothing rather than wedging the shutdown. --- news/workflow-worker-drain.feature.md | 2 +- reflex/app.py | 4 +++- reflex/workflow/CONTRACT.md | 3 +-- reflex/workflow/cli.py | 26 ++++++++++++++++---------- reflex/workflow/runtime.py | 27 ++++++++++++++++++++++++++- tests/units/workflow/test_drain.py | 17 ++++++++++++++++- 6 files changed, 63 insertions(+), 16 deletions(-) diff --git a/news/workflow-worker-drain.feature.md b/news/workflow-worker-drain.feature.md index c01214960d8..933c512776d 100644 --- a/news/workflow-worker-drain.feature.md +++ b/news/workflow-worker-drain.feature.md @@ -1 +1 @@ -A stopping worker now drains: on SIGTERM or Ctrl-C it stops claiming new steps and gives the attempts it is already running up to `--drain` (30 seconds by default) to commit, so a rolling deploy hands over cleanly instead of leaving every in-flight step claimed until its lease lapses. Anything still running when the budget expires is cancelled and left for lease recovery, exactly as if the process had been killed. +A stopping worker now drains: on SIGTERM or Ctrl-C it stops claiming new steps and gives the attempts it is already running up to `--drain` (30 seconds by default) to commit, so a rolling deploy hands over cleanly instead of leaving every in-flight step claimed until its lease lapses. Anything still running when the budget expires is cancelled and left for lease recovery, exactly as if the process had been killed. An app serving workflows drains on shutdown too, and `REFLEX_WORKFLOW_DRAIN` sets the budget for both. diff --git a/reflex/app.py b/reflex/app.py index 29331de4e56..47a6d7bef5c 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -1028,7 +1028,9 @@ async def _run_workflow_runtime(self) -> AsyncIterator[None]: if self._workflow_runtime is None: yield return - async with self._workflow_runtime.running(): + from reflex.workflow.runtime import configured_drain + + async with self._workflow_runtime.running(drain=configured_drain()): yield def add_page( diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 685ee74c45c..ccb0cc86b2d 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -185,8 +185,7 @@ Checked in this order at start: the right one). - **Stopping** is not a decision about a run. A worker asked to stop (SIGTERM, Ctrl-C, or an app lifespan ending) stops claiming immediately and gives the - attempts it is already running a drain budget — `reflex workflows worker - --drain`, 30s by default — to commit their own outcome. Anything still + attempts it is already running a drain budget — `REFLEX_WORKFLOW_DRAIN` or `reflex workflows worker --drain`, 30s by default — to commit their own outcome. Anything still running when that budget expires is cancelled and *keeps its claim*: it is reclaimed after the lease lapses, exactly as if the process had been killed. A claim is never released early, because cancelling an attempt does not stop diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index d10765d7a51..e59ab62d054 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -690,15 +690,18 @@ async def serve() -> None: @click.option("--concurrency", default=None, type=int, help="Attempts to run at once.") @click.option( "--drain", - default="30s", - help="How long a stopping worker lets running attempts finish.", + default=None, + help=( + "How long a stopping worker lets running attempts finish. Defaults to " + "REFLEX_WORKFLOW_DRAIN, or 30s." + ), ) def worker( database: str | None, target: str, queues: tuple[str, ...], concurrency: int | None, - drain: str, + drain: str | None, ): """Run workflows from TARGET with no frontend and no web server. @@ -721,13 +724,16 @@ def worker( from reflex_base.utils.exceptions import WorkflowDefinitionError from reflex_base.workflow import parse_duration - from reflex.workflow.runtime import WorkflowRuntime + from reflex.workflow.runtime import WorkflowRuntime, configured_drain - try: - drain_seconds = parse_duration(drain) - except Exception as err: - console.error(f"--drain {drain!r} is not a duration: {err}") - raise click.exceptions.Exit(1) from None + if drain is None: + drain_seconds = configured_drain() + else: + try: + drain_seconds = parse_duration(drain) + except Exception as err: + console.error(f"--drain {drain!r} is not a duration: {err}") + raise click.exceptions.Exit(1) from None try: module = _load_module(target) @@ -797,7 +803,7 @@ async def serve() -> None: # The kernel's worker does the work; this task only waits for the # operator (or the platform) to stop the process. await stopping.wait() - console.print(f"Stopping; finishing running attempts ({drain}).") + console.print(f"Stopping; finishing running attempts ({drain_seconds:g}s).") console.print("Worker stopped.") try: diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 07e8e763213..da6911d90e8 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -8,13 +8,15 @@ from __future__ import annotations +import os import random import time from contextlib import asynccontextmanager from contextvars import ContextVar -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final from reflex_base.registry import RegistrationContext +from reflex_base.utils import console from reflex_base.utils.exceptions import WorkflowDefinitionError, WorkflowRuntimeError from reflex_base.workflow import ( DEFAULT_LEASE_DURATION, @@ -238,6 +240,29 @@ async def running(self, drain: DurationLike = 0) -> AsyncIterator[WorkflowRuntim await self.shutdown(drain=drain) +DRAIN_ENV: Final = "REFLEX_WORKFLOW_DRAIN" +DEFAULT_DRAIN: Final = "30s" + + +def configured_drain() -> float: + """Read how long a stopping process lets running attempts commit. + + Deployment shape decides this, not code: a platform that sends SIGTERM + and waits ten seconds before SIGKILL wants a drain under ten seconds, and + the process cannot know that. It reads ``REFLEX_WORKFLOW_DRAIN`` the same + way it reads its store URL. + + Returns: + The budget in seconds; zero when the value is not a duration. + """ + raw = os.environ.get(DRAIN_ENV) or DEFAULT_DRAIN + try: + return parse_duration(raw) + except Exception: + console.warn(f"{DRAIN_ENV}={raw!r} is not a duration; not draining.") + return 0.0 + + def get_runtime() -> WorkflowRuntime: """Resolve the active workflow runtime. diff --git a/tests/units/workflow/test_drain.py b/tests/units/workflow/test_drain.py index 1f7ff34a348..bbbd87f0280 100644 --- a/tests/units/workflow/test_drain.py +++ b/tests/units/workflow/test_drain.py @@ -11,12 +11,13 @@ import asyncio from typing import Any -from reflex_base.workflow import WorkflowConfig, manual +from reflex_base.workflow import WorkflowConfig, manual, parse_duration import reflex as rx from reflex.workflow.definition import compile_workflow from reflex.workflow.kernel import WorkflowKernel from reflex.workflow.records import RunQuery, RunStatus, StepStatus +from reflex.workflow.runtime import DEFAULT_DRAIN, DRAIN_ENV, configured_drain from reflex.workflow.store import MemoryRunStore @@ -120,3 +121,17 @@ async def test_an_attempt_slower_than_the_drain_is_cancelled( snapshot = await kernel.get_run((await store.list_runs(RunQuery()))[0].run_id) assert snapshot is not None assert snapshot.steps[0].status is StepStatus.CLAIMED + + +def test_the_drain_budget_comes_from_the_environment(monkeypatch): + """A platform's grace period is deployment config, not a code constant.""" + monkeypatch.setenv(DRAIN_ENV, "5s") + assert configured_drain() == 5.0 + monkeypatch.delenv(DRAIN_ENV) + assert configured_drain() == parse_duration(DEFAULT_DRAIN) + + +def test_an_unparseable_budget_does_not_stop_the_process_leaving(monkeypatch): + """A typo in a deployment variable must not wedge a shutdown.""" + monkeypatch.setenv(DRAIN_ENV, "half an hour") + assert configured_drain() == 0.0 From 2ae67731ef31a5ef04eff26e168803047b11c743 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 02:08:41 -0700 Subject: [PATCH 071/121] workflows: let a process start runs without becoming a worker The process with the business event is usually not the one that should run the workflow: a Django view, a FastAPI route, a cron box, a script. Until now the only supported way in was to start a runtime, which also starts a worker -- so a web request would begin executing steps, and a script that exits mid-attempt would leave a claim behind. The other way in, startup(start_worker=False), was undocumented and the error you got without it named only the app and the test harness. rx.workflows.connect(...) is that path, said out loud. Inside the block every rx.workflows call and every RunHandle works against the shared store and nothing is claimed or executed; workers stay separate processes. The runtime error now names it, init prints it, and the contract says plainly that a client is not a worker. --- news/workflow-client-connect.feature.md | 1 + reflex/workflow/CONTRACT.md | 6 +++ reflex/workflow/cli.py | 4 ++ reflex/workflow/runtime.py | 57 ++++++++++++++++++++- tests/units/workflow/test_drain.py | 7 +-- tests/units/workflow/test_runtime.py | 67 ++++++++++++++++++++++++- 6 files changed, 136 insertions(+), 6 deletions(-) create mode 100644 news/workflow-client-connect.feature.md diff --git a/news/workflow-client-connect.feature.md b/news/workflow-client-connect.feature.md new file mode 100644 index 00000000000..4825724d30a --- /dev/null +++ b/news/workflow-client-connect.feature.md @@ -0,0 +1 @@ +`rx.workflows.connect(...)` opens a client on a workflow store from any Python process — a script, a FastAPI route, a Django view — so it can start runs and read them without becoming a worker: nothing is claimed or executed inside the block. Workers stay separate processes. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index ccb0cc86b2d..7b0e3ecb98b 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -191,6 +191,12 @@ Checked in this order at start: A claim is never released early, because cancelling an attempt does not stop work it handed to a thread, and the lease is what keeps a peer off it. A drained attempt costs nothing; a cancelled one costs one recovery. +- **Clients are not workers.** A process that opens + `rx.workflows.connect(...)` can admit runs, read them, signal and cancel + them, and executes nothing: it claims no step and runs no handler. Only a + process that starts the kernel's worker (an app serving workflows, or + `reflex workflows worker`) executes. This is what lets a web request start a + run without running it. - Multiple workers share one Postgres store via `SKIP LOCKED` claims; SQLite is a one-process store (calls off-loop, contention bounded); memory is for tests. All three answer the same conformance suite. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index e59ab62d054..056eab29c6b 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -492,6 +492,10 @@ def init_workflow(name: str): console.print("") console.print("Or serve it as a worker and start runs from your own code:") click.echo(f" reflex workflows worker {module}") + console.print("") + console.print(" From a script, a FastAPI route, or a Django view:") + click.echo(f" async with rx.workflows.connect({klass}):") + click.echo(f" await rx.workflows.submit({klass}.start(order='ord-1'))") @workflows.command() diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index da6911d90e8..b3c6337c530 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -167,8 +167,10 @@ def kernel(self) -> WorkflowKernel: """ if self._kernel is None: msg = ( - "The workflow runtime has not started; start the app or use " - "WorkflowTestHarness in tests." + "The workflow runtime has not started. Run the app, or open a " + "client with `async with rx.workflows.connect(MyWorkflow): " + "...` to start and read runs from a script or another " + "framework, or use WorkflowTestHarness in tests." ) raise WorkflowRuntimeError(msg) return self._kernel @@ -287,6 +289,57 @@ def get_runtime() -> WorkflowRuntime: class WorkflowsNamespace: """The public ``rx.workflows`` API surface.""" + @staticmethod + @asynccontextmanager + async def connect( + *workflow_classes: type[BaseState], + database: str | None = None, + store: RunStore | None = None, + ) -> AsyncIterator[WorkflowRuntime]: + """Open a client on a store without serving any work. + + The process that has the business event is often not the process that + runs the workflow: a Django view, a FastAPI route, a cron box, a + one-off script. Those want to start runs and read them, and must not + quietly become workers by importing the engine -- a web process that + starts executing steps is a surprise, and a script that exits mid-step + would leave a claim behind. + + Inside this block every ``rx.workflows`` call and every ``RunHandle`` + works and goes to the shared store, and nothing is claimed or + executed here. Workers are separate processes + (``reflex workflows worker``). + + Usage:: + + async with rx.workflows.connect(Checkout): + handle = await rx.workflows.submit(Checkout.start(order="o1")) + + Args: + workflow_classes: The workflows this client may start. + database: Connection URL or SQLite path; defaults to + ``REFLEX_WORKFLOW_DATABASE``, then ``./workflow.db``. + store: An already-open store, which takes precedence. + + Yields: + The client runtime. + """ + global _default_runtime + + runtime = WorkflowRuntime( + store if store is not None else resolve_store(database) + ) + for workflow_cls in workflow_classes: + runtime.register(workflow_cls) + await runtime.startup(start_worker=False) + previous = _default_runtime + _default_runtime = runtime + try: + yield runtime + finally: + _default_runtime = previous + await runtime.shutdown() + @staticmethod async def start( target: Any, diff --git a/tests/units/workflow/test_drain.py b/tests/units/workflow/test_drain.py index bbbd87f0280..018f531ac26 100644 --- a/tests/units/workflow/test_drain.py +++ b/tests/units/workflow/test_drain.py @@ -11,6 +11,7 @@ import asyncio from typing import Any +import pytest from reflex_base.workflow import WorkflowConfig, manual, parse_duration import reflex as rx @@ -126,12 +127,12 @@ async def test_an_attempt_slower_than_the_drain_is_cancelled( def test_the_drain_budget_comes_from_the_environment(monkeypatch): """A platform's grace period is deployment config, not a code constant.""" monkeypatch.setenv(DRAIN_ENV, "5s") - assert configured_drain() == 5.0 + assert configured_drain() == pytest.approx(5.0) monkeypatch.delenv(DRAIN_ENV) - assert configured_drain() == parse_duration(DEFAULT_DRAIN) + assert configured_drain() == pytest.approx(parse_duration(DEFAULT_DRAIN)) def test_an_unparseable_budget_does_not_stop_the_process_leaving(monkeypatch): """A typo in a deployment variable must not wedge a shutdown.""" monkeypatch.setenv(DRAIN_ENV, "half an hour") - assert configured_drain() == 0.0 + assert configured_drain() == pytest.approx(0.0) diff --git a/tests/units/workflow/test_runtime.py b/tests/units/workflow/test_runtime.py index afd6348e381..dda9fea0d8f 100644 --- a/tests/units/workflow/test_runtime.py +++ b/tests/units/workflow/test_runtime.py @@ -11,7 +11,10 @@ from reflex_base.workflow import WorkflowConfig, manual import reflex as rx -from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.definition import compile_workflow +from reflex.workflow.kernel import WorkflowKernel +from reflex.workflow.records import RunStatus +from reflex.workflow.runtime import WorkflowRuntime, workflows from reflex.workflow.store import MemoryRunStore @@ -134,3 +137,65 @@ def test_definitions_are_reported_once_per_class(forked_registration_context): assert [definition.workflow_id for definition in runtime.definitions] == [ "runtime.simple" ] + + +async def test_connect_starts_runs_without_becoming_a_worker( + forked_registration_context, +): + """A client process admits work and executes none of it. + + The process holding the business event -- a Django view, a script -- must + be able to start a run without quietly turning into a worker, which would + make a web request execute a workflow step and a script exit mid-attempt. + """ + + class ClientFlow(rx.State): + __workflow__ = WorkflowConfig(id="runtime.client") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Complete immediately. + + Returns: + Completion. + """ + return rx.complete(result="ran") + + store = MemoryRunStore() + async with workflows.connect(ClientFlow, store=store) as runtime: + handle = await workflows.submit(ClientFlow.start()) + assert handle.started + snapshot = await handle.snapshot() + assert snapshot is not None + assert snapshot.status is RunStatus.PENDING + assert runtime.kernel is not None + + served = WorkflowKernel([compile_workflow(ClientFlow)], store) + await served.run_until_idle() + after = await served.get_run(handle.run_id) + assert after is not None + assert after.status is RunStatus.COMPLETED + assert after.result == "ran" + + +async def test_connect_restores_whatever_runtime_was_active( + forked_registration_context, +): + """A client is a scope, not a process-wide switch.""" + + class ScopedFlow(rx.State): + __workflow__ = WorkflowConfig(id="runtime.scoped") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Complete immediately. + + Returns: + Completion. + """ + return rx.complete(result=None) + + async with workflows.connect(ScopedFlow, store=MemoryRunStore()): + pass + with pytest.raises(WorkflowRuntimeError): + await workflows.submit(ScopedFlow.start()) From 9d5b9f3ecf004ed47f92c18b9e523b137c3542c1 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 02:13:33 -0700 Subject: [PATCH 072/121] workflows: answer 'is this healthy' without opening the database An operator's first question is a count -- how many runs are open, how many stopped needing a person -- and the only way to get it was to list runs a page at a time or query the store by hand. That is the thing the managed product is supposed to make unnecessary. Stores answer count_runs(query) now, filtering exactly as list_runs does so a total and a page can never describe different sets; a conformance check pins that across memory, SQLite and Postgres, including that limit and the pagination cursor bound a listing and not an aggregate. `reflex workflows stats` reports the breakdown, and --json makes it a scrape an alert can read. --- news/workflow-stats.feature.md | 1 + reflex/workflow/cli.py | 76 +++++++++++++++++++ reflex/workflow/conformance.py | 34 +++++++++ reflex/workflow/postgres.py | 73 +++++++++++++----- reflex/workflow/store.py | 124 +++++++++++++++++++++++++------ tests/units/workflow/test_cli.py | 37 +++++++++ 6 files changed, 307 insertions(+), 38 deletions(-) create mode 100644 news/workflow-stats.feature.md diff --git a/news/workflow-stats.feature.md b/news/workflow-stats.feature.md new file mode 100644 index 00000000000..b0b4b5d0195 --- /dev/null +++ b/news/workflow-stats.feature.md @@ -0,0 +1 @@ +`reflex workflows stats` counts runs by status, with `--json` for whatever collects metrics, so "how many runs are open" and "how many need a person" are answerable — and alertable — without reading the workflow database. Stores gained a `count_runs()` query, covered by the shared conformance suite so a count and a listing always describe the same set. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 056eab29c6b..fb23f78fd33 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -816,6 +816,82 @@ async def serve() -> None: console.print("Worker stopped.") +@workflows.command() +@database_option +@click.option("--workflow", "-w", default=None, help="Only this workflow id.") +@click.option("--label", "-l", "labels", multiple=True, help="Filter as key=value.") +@click.option("--json", "as_json", is_flag=True, help="Emit JSON instead of a table.") +def stats( + database: str | None, + workflow: str | None, + labels: tuple[str, ...], + as_json: bool, +): + """Count runs by status: is this deployment healthy, in one screen. + + The two numbers that matter are how many runs are still open and how many + stopped needing a person. Everything else is context. --json makes this a + scrape for whatever collects metrics, so an alert on "runs needing + attention" does not require reading the workflow database. + """ + from reflex.workflow.records import TERMINAL_RUN_STATUSES, RunQuery + + label_filter = dict(pair.split("=", 1) for pair in labels if "=" in pair) + + async def counts(store: RunStore) -> dict[str, int]: + """Count runs per status. + + Args: + store: The open store. + + Returns: + The count for each status that has any runs. + """ + found: dict[str, int] = {} + for status in RunStatus: + query = RunQuery( + workflow_id=workflow, + statuses=(status,), + labels=label_filter or None, + ) + total = await store.count_runs(query) + if total: + found[status.value] = total + return found + + found = _with_store(database, counts) + total = sum(found.values()) + open_runs = sum( + count + for status, count in found.items() + if RunStatus(status) not in TERMINAL_RUN_STATUSES + ) + attention = found.get(RunStatus.NEEDS_ATTENTION.value, 0) + + if as_json: + click.echo( + json.dumps({ + "workflow": workflow, + "total": total, + "open": open_runs, + "needs_attention": attention, + "by_status": found, + }) + ) + return + + if not total: + console.print("No runs yet.") + return + click.echo(f"{'STATUS':<20}{'RUNS':>8}") + for status in RunStatus: + count = found.get(status.value) + if count: + click.echo(f"{status.value:<20}{count:>8}") + console.print("") + console.print(f"{total} run(s); {open_runs} open, {attention} needing attention.") + + @workflows.command("list") @database_option @click.option("--workflow", "-w", default=None, help="Only this workflow id.") diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 918e9fc508d..587a552c87e 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -484,6 +484,39 @@ async def check_list_runs_filters_and_orders(store: RunStore) -> None: assert await store.list_runs(RunQuery(workflow_id="other.flow", limit=10)) == () +async def check_count_runs_matches_the_listing(store: RunStore) -> None: + """A count answers for the same set a listing would return. + + An operator view reports totals next to a page of runs. If the two read + the filters differently, the page and the number above it describe + different things, and the number is the one nobody can check. + """ + await store.admit( + make_run("a", labels={"customer": "acme"}, created_at=NOW), + make_step("a"), + _ADMITTED, + ) + await store.admit( + make_run( + "b", + labels={"customer": "globex"}, + status=RunStatus.COMPLETED, + created_at=NOW + 1, + ), + make_step("b"), + _ADMITTED, + ) + assert await store.count_runs(RunQuery()) == 2 + assert await store.count_runs(RunQuery(statuses=(RunStatus.COMPLETED,))) == 1 + assert await store.count_runs(RunQuery(labels={"customer": "acme"})) == 1 + assert await store.count_runs(RunQuery(workflow_id="other.flow")) == 0 + # Paging bounds a listing; it must not bound an aggregate. + assert await store.count_runs(RunQuery(limit=1)) == 2 + assert await store.count_runs(RunQuery(created_before=(NOW + 1, "b"))) == 2, ( + "a cursor pages a listing and says nothing about how many exist" + ) + + async def check_flow_control_queries(store: RunStore) -> None: """Start policies can see what is active and what started recently.""" await store.admit( @@ -1100,6 +1133,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_finalize_tombstones_open_slots, check_resume_only_reopens_a_suspended_run, check_list_runs_filters_and_orders, + check_count_runs_matches_the_listing, check_pagination_skips_nothing_on_tied_timestamps, check_label_filter_handles_awkward_keys, check_flow_control_queries, diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 0587f0cc516..315c4dcecc7 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -267,6 +267,38 @@ def _step_from_row(row: Mapping[str, Any]) -> StepRecord: ) +def _run_filters(query: RunQuery) -> tuple[str, tuple[Any, ...]]: + """Build the WHERE clause a run query means. + + Shared by listing and counting so the two can never disagree about what + matches. + + Args: + query: The filters to apply; ``limit`` is not one of them. + + Returns: + The clause (empty when nothing filters) and its parameters. + """ + clauses: list[str] = [] + params: list[Any] = [] + if query.workflow_id is not None: + clauses.append("workflow_id = %s") + params.append(query.workflow_id) + if query.statuses: + clauses.append("status = ANY(%s)") + params.append([status.value for status in query.statuses]) + if query.created_before is not None: + clauses.append("(created_at, run_id) < (%s, %s)") + params.extend(query.created_before) + if query.labels: + # Containment matches the whole filter at once, and takes user keys + # as data rather than splicing them into a path expression. + clauses.append("labels @> %s") + params.append(_json(dict(query.labels))) + where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + return where, tuple(params) + + class PostgresRunStore: """Run store backed by PostgreSQL, safe for many concurrent workers.""" @@ -1665,23 +1697,7 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: Returns: The matching run records. """ - clauses: list[str] = [] - params: list[Any] = [] - if query.workflow_id is not None: - clauses.append("workflow_id = %s") - params.append(query.workflow_id) - if query.statuses: - clauses.append("status = ANY(%s)") - params.append([status.value for status in query.statuses]) - if query.created_before is not None: - clauses.append("(created_at, run_id) < (%s, %s)") - params.extend(query.created_before) - if query.labels: - # Containment matches the whole filter at once, and takes user keys - # as data rather than splicing them into a path expression. - clauses.append("labels @> %s") - params.append(_json(dict(query.labels))) - where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + where, params = _run_filters(query) pool = await self._open() async with pool.connection() as conn: cursor = await conn.execute( @@ -1691,6 +1707,29 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: ) return tuple(_run_from_row(row) for row in await cursor.fetchall()) + async def count_runs(self, query: RunQuery) -> int: + """Count runs matching a query. + + The count ignores ``limit`` and ``created_before``: those page a + listing, and an aggregate is not a page. Everything else filters as + it does for ``list_runs``, so a count and a listing always describe + the same set. + + Args: + query: The filters to apply. + + Returns: + How many runs match. + """ + where, params = _run_filters(dataclasses.replace(query, created_before=None)) + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + f"SELECT count(*) AS n FROM workflow_runs{where}", tuple(params) + ) + row = await cursor.fetchone() + return 0 if row is None else int(row["n"]) + async def list_children( self, parent_run_id: str, parent_ordinal: int ) -> tuple[RunRecord, ...]: diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index ae38e225e32..27e044221da 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -535,6 +535,22 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: """ ... + async def count_runs(self, query: RunQuery) -> int: + """Count runs matching a query. + + The count ignores ``limit`` and ``created_before``: those page a + listing, and an aggregate is not a page. Everything else filters as + it does for ``list_runs``, so a count and a listing always describe + the same set. + + Args: + query: The filters to apply. + + Returns: + How many runs match. + """ + ... + async def list_children( self, parent_run_id: str, parent_ordinal: int ) -> tuple[RunRecord, ...]: @@ -1742,6 +1758,24 @@ async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: matched.sort(key=lambda run: (run.created_at, run.run_id), reverse=True) return tuple(_detach_run(run) for run in matched[: query.limit]) + async def count_runs(self, query: RunQuery) -> int: + """Count runs matching a query. + + The count ignores ``limit`` and ``created_before``: those page a + listing, and an aggregate is not a page. Everything else filters as + it does for ``list_runs``, so a count and a listing always describe + the same set. + + Args: + query: The filters to apply. + + Returns: + How many runs match. + """ + counted = dataclasses.replace(query, created_before=None) + async with self._lock: + return sum(1 for run in self._runs.values() if _matches_query(run, counted)) + async def list_children( self, parent_run_id: str, parent_ordinal: int ) -> tuple[RunRecord, ...]: @@ -2165,6 +2199,42 @@ def _step_from_row(row: sqlite3.Row) -> StepRecord: ) +def _sqlite_run_filters(query: RunQuery) -> tuple[str, tuple[Any, ...]]: + """Build the WHERE clause a run query means, for SQLite. + + Shared by listing and counting so the two can never disagree about what + matches. + + Args: + query: The filters to apply; ``limit`` is not one of them. + + Returns: + The clause (empty when nothing filters) and its parameters. + """ + clauses: list[str] = [] + params: list[Any] = [] + if query.workflow_id is not None: + clauses.append("workflow_id = ?") + params.append(query.workflow_id) + if query.statuses: + placeholders = ",".join("?" * len(query.statuses)) + clauses.append(f"status IN ({placeholders})") + params.extend(status.value for status in query.statuses) + if query.created_before is not None: + clauses.append("(created_at, run_id) < (?, ?)") + params.extend(query.created_before) + for key, value in (query.labels or {}).items(): + # The key comes from user data, so it is matched as a value rather + # than spliced into a JSON path expression. + clauses.append( + "EXISTS (SELECT 1 FROM json_each(labels)" + " WHERE json_each.key = ? AND json_each.value = ?)" + ) + params.extend((key, value)) + where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + return where, tuple(params) + + class SqliteRunStore: """Crash-safe run store backed by a local SQLite database.""" @@ -3741,27 +3811,7 @@ def work(): Returns: The operation's result. """ - clauses: list[str] = [] - params: list[Any] = [] - if query.workflow_id is not None: - clauses.append("workflow_id = ?") - params.append(query.workflow_id) - if query.statuses: - placeholders = ",".join("?" * len(query.statuses)) - clauses.append(f"status IN ({placeholders})") - params.extend(status.value for status in query.statuses) - if query.created_before is not None: - clauses.append("(created_at, run_id) < (?, ?)") - params.extend(query.created_before) - for key, value in (query.labels or {}).items(): - # The key comes from user data, so it is matched as a value rather - # than spliced into a JSON path expression. - clauses.append( - "EXISTS (SELECT 1 FROM json_each(labels)" - " WHERE json_each.key = ? AND json_each.value = ?)" - ) - params.extend((key, value)) - where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + where, params = _sqlite_run_filters(query) with self._lock: rows = self._db.execute( f"SELECT * FROM workflow_runs{where}" @@ -3772,6 +3822,38 @@ def work(): return await asyncio.to_thread(work) + async def count_runs(self, query: RunQuery) -> int: + """Count runs matching a query. + + The count ignores ``limit`` and ``created_before``: those page a + listing, and an aggregate is not a page. Everything else filters as + it does for ``list_runs``, so a count and a listing always describe + the same set. + + Args: + query: The filters to apply. + + Returns: + How many runs match. + """ + + def work(): + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + where, params = _sqlite_run_filters( + dataclasses.replace(query, created_before=None) + ) + with self._lock: + row = self._db.execute( + f"SELECT count(*) AS n FROM workflow_runs{where}", params + ).fetchone() + return int(row["n"]) + + return await asyncio.to_thread(work) + async def list_children( self, parent_run_id: str, parent_ordinal: int ) -> tuple[RunRecord, ...]: diff --git a/tests/units/workflow/test_cli.py b/tests/units/workflow/test_cli.py index f2b2d393f64..676aea2bdd3 100644 --- a/tests/units/workflow/test_cli.py +++ b/tests/units/workflow/test_cli.py @@ -178,3 +178,40 @@ def test_resume_reopens_only_suspended_runs(seeded): run = _load_run(database, suspended) assert run is not None assert run.status is RunStatus.PENDING + + +def test_stats_counts_runs_by_status(seeded): + """The health screen: how many runs exist, and how many are still open.""" + database, _, _ = seeded + result = _invoke("stats", "-d", database) + assert result.exit_code == 0, result.output + assert RunStatus.WAITING.value in result.output + assert "open" in result.output + + +def test_stats_json_is_scrapeable(seeded): + """An alert on runs needing attention must not read the database.""" + database, _, _ = seeded + result = _invoke("stats", "-d", database, "--json") + assert result.exit_code == 0, result.output + payload = json.loads(result.output) + assert payload["total"] == sum(payload["by_status"].values()) + assert payload["open"] == payload["total"], "neither seeded run is terminal" + assert payload["needs_attention"] == payload["by_status"].get( + RunStatus.NEEDS_ATTENTION.value, 0 + ) + assert payload["needs_attention"] == 1, "the suspended run is the one to page on" + + +def test_stats_filters_to_one_workflow(seeded): + """A shared store holds every workflow; an owner asks about theirs.""" + database, _, _ = seeded + result = _invoke("stats", "-d", database, "-w", "nope.nothing", "--json") + assert result.exit_code == 0, result.output + assert json.loads(result.output) == { + "workflow": "nope.nothing", + "total": 0, + "open": 0, + "needs_attention": 0, + "by_status": {}, + } From c6216fadd20ed84cd9b26bdfa8c149aa771c6525 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 02:17:26 -0700 Subject: [PATCH 073/121] workflows: make a process's counters scrapeable MetricsObserver had the numbers an operator pages on and no way to read them: you had to construct one, pass it in, restart, and then write your own exporter over snapshot(). A deployment that must be reconfigured to learn how many runs failed learns it too late. Every runtime now counts unconditionally -- an observer of your own is composed alongside rather than replacing it -- and GET /_workflow/api/metrics renders the counters as Prometheus text behind the same bearer token as the rest of the workflow API. Plain text is deliberate: it is what a Prometheus server scrapes, what the OpenTelemetry collector's prometheus receiver reads, and what hosted vendors accept, with no client library and no dependency of ours to keep current. Counters are per process, so a fleet's numbers add up. --- news/workflow-metrics-endpoint.feature.md | 1 + reflex/app.py | 7 ++ reflex/workflow/api.py | 88 ++++++++++++++++++++++- reflex/workflow/kernel.py | 44 ++++++++++++ reflex/workflow/runtime.py | 12 +++- tests/units/workflow/test_api.py | 37 ++++++++++ 6 files changed, 187 insertions(+), 2 deletions(-) create mode 100644 news/workflow-metrics-endpoint.feature.md diff --git a/news/workflow-metrics-endpoint.feature.md b/news/workflow-metrics-endpoint.feature.md new file mode 100644 index 00000000000..c5ab32045c8 --- /dev/null +++ b/news/workflow-metrics-endpoint.feature.md @@ -0,0 +1 @@ +A workflow process now always counts what it does, and `GET /_workflow/api/metrics` exposes those counters in the Prometheus text format — runs started and how they ended, attempts, retries, recoveries, both in total and per workflow. It sits behind the same bearer token as the rest of the workflow API, so a Prometheus server, an OpenTelemetry collector's prometheus receiver, or a hosted metrics vendor can scrape a worker without any client library. An observer passed to `rx.App(workflow_observer=...)` still sees every event. diff --git a/reflex/app.py b/reflex/app.py index 47a6d7bef5c..0948a7c3dc5 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -851,9 +851,11 @@ def _add_default_endpoints(self): def _add_workflow_endpoints(self): """Add the workflow ingress endpoints: webhooks in, approvals back.""" from reflex.workflow.api import ( + METRICS_ROUTE, RUN_ROUTE, START_ROUTE, api_token, + metrics_endpoint, run_endpoint, start_endpoint, ) @@ -896,6 +898,11 @@ def _add_workflow_endpoints(self): run_endpoint(self._workflow_runtime, token), methods=["GET"], ) + self._api.add_route( + config.prepend_backend_path(METRICS_ROUTE), + metrics_endpoint(self._workflow_runtime, token), + methods=["GET"], + ) def _add_optional_endpoints(self): """Add optional api endpoints (_upload).""" diff --git a/reflex/workflow/api.py b/reflex/workflow/api.py index 931ea04a3f3..5f9f15cd9bd 100644 --- a/reflex/workflow/api.py +++ b/reflex/workflow/api.py @@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Any, Final from reflex_base.utils import console -from starlette.responses import JSONResponse +from starlette.responses import JSONResponse, PlainTextResponse from reflex.workflow.records import attempts_made @@ -29,12 +29,14 @@ from collections.abc import Callable, Coroutine from starlette.requests import Request + from starlette.responses import Response from reflex.workflow.runtime import WorkflowRuntime TOKEN_ENV: Final = "REFLEX_WORKFLOW_API_TOKEN" START_ROUTE: Final = "/_workflow/api/runs" RUN_ROUTE: Final = "/_workflow/api/runs/{run_id}" +METRICS_ROUTE: Final = "/_workflow/api/metrics" MAX_BODY_BYTES: Final = 1_048_576 @@ -192,3 +194,87 @@ async def endpoint(request: Request) -> JSONResponse: }) return endpoint + + +def _label(value: str) -> str: + """Escape a value for a Prometheus label. + + Args: + value: The raw label value. + + Returns: + The escaped value, without its surrounding quotes. + """ + return value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") + + +def render_prometheus(snapshot: dict[str, Any]) -> str: + """Render counter totals in the Prometheus text exposition format. + + Plain text on purpose: it is what a Prometheus server scrapes, what the + OpenTelemetry collector's prometheus receiver reads, and what every + hosted metrics vendor accepts, with no client library and no dependency + of ours to keep current. + + Args: + snapshot: A ``MetricsObserver.snapshot()`` result. + + Returns: + The exposition text, ending in a newline. + """ + totals: dict[str, int] = snapshot.get("totals", {}) + by_workflow: dict[str, dict[str, int]] = snapshot.get("by_workflow", {}) + names = sorted({ + *totals, + *(key for counts in by_workflow.values() for key in counts), + }) + lines: list[str] = [] + for name in names: + metric = f"reflex_workflow_{name}_total" + lines.extend(( + f"# TYPE {metric} counter", + f"{metric} {totals.get(name, 0)}", + )) + for workflow_id in sorted(by_workflow): + count = by_workflow[workflow_id].get(name) + if count: + lines.append(f'{metric}{{workflow="{_label(workflow_id)}"}} {count}') + return "\n".join(lines) + "\n" + + +def metrics_endpoint( + runtime: WorkflowRuntime, token: str +) -> Callable[[Request], Coroutine[Any, Any, Response]]: + """Build the endpoint that exposes this process's counters. + + The counters are per process: each worker reports what it did, and the + collector sums them. That is what makes a fleet's numbers addable rather + than a single process's guess about the whole system. + + Args: + runtime: The runtime whose counters to report. + token: The bearer token every caller must present. + + Returns: + The endpoint callable. + """ + + async def endpoint(request: Request) -> Response: # noqa: RUF029 + """Report the counters as Prometheus text. + + Args: + request: The incoming request. + + Returns: + The exposition text, or an error. + """ + # Reading in-process counters needs no await; the signature is a + # Starlette endpoint's, not a claim that this does I/O. + if not _authorized(request, token): + return JSONResponse({"error": "unauthorized"}, status_code=401) + return PlainTextResponse( + render_prometheus(runtime.metrics.snapshot()), + media_type="text/plain; version=0.0.4; charset=utf-8", + ) + + return endpoint diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index b2b6fd16677..6da1f688521 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -140,6 +140,50 @@ def on_event( """ +class CompositeObserver(WorkflowObserver): + """Fans one transition out to several observers. + + Instrumentation is not exclusive: a deployment that exports metrics still + wants its own logging, and neither should have to know about the other. + + Attributes: + observers: The observers to notify, in order. + """ + + __slots__ = ("observers",) + + def __init__(self, *observers: WorkflowObserver): + """Bind the observers to fan out to. + + Args: + observers: The observers to notify, in order. + """ + self.observers = observers + + def on_event( + self, + event_type: HistoryEventType, + run_id: str, + workflow_id: str, + data: dict[str, Any], + ) -> None: + """Pass one transition to every observer. + + One observer raising must not cost the others their notification, and + the kernel already treats instrumentation errors as its own to + swallow. + + Args: + event_type: What happened. + run_id: The run it happened to. + workflow_id: That run's workflow identity. + data: The event payload. + """ + for observer in self.observers: + with contextlib.suppress(Exception): + observer.on_event(event_type, run_id, workflow_id, data) + + class MetricsObserver(WorkflowObserver): """Tallies the numbers a deployment alerts on. diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index b3c6337c530..4f2c8d4e169 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -31,6 +31,8 @@ from reflex.workflow.kernel import ( DEFAULT_MAX_CONCURRENCY, DEFAULT_POLL_INTERVAL, + CompositeObserver, + MetricsObserver, WorkflowKernel, WorkflowObserver, ) @@ -96,7 +98,15 @@ def __init__( self._lease_duration = lease_duration self._lease_renew_interval = lease_renew_interval self._recovery_interval = recovery_interval - self._observer = observer + self.metrics = MetricsObserver() + # Always counting: a deployment that has to reconfigure and restart to + # find out how many runs failed learns it too late. The user's + # observer, when there is one, still sees every event. + self._observer = ( + self.metrics + if observer is None + else CompositeObserver(self.metrics, observer) + ) self._max_recoveries = max_recoveries self._max_concurrency = max_concurrency self._queues = tuple(queues) if queues is not None else None diff --git a/tests/units/workflow/test_api.py b/tests/units/workflow/test_api.py index 10b30fb5a62..43503a7fd8e 100644 --- a/tests/units/workflow/test_api.py +++ b/tests/units/workflow/test_api.py @@ -16,10 +16,13 @@ import reflex as rx from reflex.workflow.api import ( + METRICS_ROUTE, RUN_ROUTE, START_ROUTE, TOKEN_ENV, api_token, + metrics_endpoint, + render_prometheus, run_endpoint, start_endpoint, ) @@ -63,6 +66,7 @@ async def client(forked_registration_context): routes=[ Route(START_ROUTE, start_endpoint(runtime, TOKEN), methods=["POST"]), Route(RUN_ROUTE, run_endpoint(runtime, TOKEN), methods=["GET"]), + Route(METRICS_ROUTE, metrics_endpoint(runtime, TOKEN), methods=["GET"]), ] ) with TestClient(app) as ready: @@ -131,6 +135,7 @@ def test_every_route_requires_the_token(client, headers): assert ( client.get("/_workflow/api/runs/whatever", headers=headers).status_code == 401 ) + assert client.get(METRICS_ROUTE, headers=headers).status_code == 401 def test_unknown_targets_are_refused(client): @@ -159,3 +164,35 @@ def test_the_api_is_absent_without_a_configured_token(monkeypatch): assert api_token() is None monkeypatch.setenv(TOKEN_ENV, TOKEN) assert api_token() == TOKEN + + +def test_metrics_are_scrapeable_after_a_run(client): + """A collector reads what this process did, in the format it already speaks. + + Args: + client: The test client. + """ + body = json.dumps({ + "workflow": "api.orders", + "handler": "place", + "args": {"order": "o-1"}, + }) + assert client.post(START_ROUTE, content=body, headers=_auth()).status_code == 202 + + response = client.get(METRICS_ROUTE, headers=_auth()) + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/plain") + text = response.text + assert "# TYPE reflex_workflow_runs_started_total counter" in text + assert "reflex_workflow_runs_started_total 1" in text + assert 'reflex_workflow_runs_started_total{workflow="api.orders"} 1' in text + + +def test_prometheus_rendering_escapes_label_values(): + """A workflow id is developer data; the exposition format still has to hold.""" + text = render_prometheus({ + "totals": {"runs_started": 1}, + "by_workflow": {'we"ird\\flow': {"runs_started": 1}}, + }) + assert 'workflow="we\\"ird\\\\flow"' in text + assert text.endswith("\n") From e421ef86b8c4aa15a52845c001d17f70336be34b Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 02:20:43 -0700 Subject: [PATCH 074/121] workflows: finish the operator actions on the command line cancel, resume, retry and skip were commands; force_complete and force_fail existed only as Python calls, so the two actions you reach for when a run will never finish on its own were the two that needed a script and a running runtime to perform. `reflex workflows complete --result JSON` and `reflex workflows fail --reason ...` close that. Both go through the kernel rather than straight to the store: finalizing a child has to deliver its arrival to the parent's join slot in the same transaction, or a parent waits forever on a child an operator already closed. Neither starts a worker, both refuse a run a worker still holds a step on, and a --result that is not JSON is refused rather than recorded as the text typed. --- news/workflow-finalize-cli.feature.md | 1 + reflex/workflow/CONTRACT.md | 4 +- reflex/workflow/cli.py | 79 +++++++++++++++++++++++++++ tests/units/workflow/test_cli.py | 37 +++++++++++++ 4 files changed, 120 insertions(+), 1 deletion(-) create mode 100644 news/workflow-finalize-cli.feature.md diff --git a/news/workflow-finalize-cli.feature.md b/news/workflow-finalize-cli.feature.md new file mode 100644 index 00000000000..62fc15cf5c6 --- /dev/null +++ b/news/workflow-finalize-cli.feature.md @@ -0,0 +1 @@ +`reflex workflows complete [--result JSON]` and `reflex workflows fail --reason ...` end a run an operator has given up waiting on, completing the set of operator actions available from the command line alongside `cancel`, `resume`, `retry`, and `skip`. Both refuse a run a worker still holds a step on, and both deliver a child run's arrival to its parent's join slot in the same transaction. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 7b0e3ecb98b..383784b93d9 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -280,4 +280,6 @@ no-op with a reason. while a step is claimed, so it never races a working attempt — cancel first if a worker still holds one. -All of these are store transactions under the same atomicity rules as §1. +All of these are store transactions under the same atomicity rules as §1, and +every one of them is reachable without writing Python: `reflex workflows +cancel | resume | retry | skip | complete | fail `. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index fb23f78fd33..c385ccc7723 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -1192,6 +1192,85 @@ def skip(database: str | None, run_id: str): _operator_action(database, run_id, "skip_step") +def _finalize( + database: str | None, + run_id: str, + status: RunStatus, + *, + result: Any = None, + error: dict[str, Any] | None = None, +): + """End a run by operator decision. + + Goes through the kernel rather than the store because finalizing a child + run has to deliver its arrival to the parent's join slot in the same + transaction -- otherwise a parent waits forever on a child an operator + already closed. No worker is started, so nothing executes here. + + Args: + database: Connection URL or SQLite path, or None for the default. + run_id: The run to finalize. + status: The terminal status to record. + result: Result to record when completing. + error: Error payload to record when failing. + + Raises: + Exit: When the run is unknown, already finished, or has a claimed step. + """ + from reflex.workflow.kernel import WorkflowKernel + + finalized = _with_store( + database, + lambda store: WorkflowKernel([], store).force_finalize( + run_id, status=status, result=result, error=error + ), + ) + if not finalized: + console.error( + f"Run {run_id!r} is unknown, already finished, or has a step a " + "worker still holds. Cancel it first if a worker is on it." + ) + raise click.exceptions.Exit(1) + console.print(f"Run {run_id} recorded as {status.value} by operator decision.") + + +@workflows.command() +@database_option +@click.argument("run_id") +@click.option( + "--result", "result_json", default=None, help="Result to record, as JSON." +) +def complete(database: str | None, run_id: str, result_json: str | None): + """End a run as completed by operator decision. + + For a run no code path will finish: a wait nobody will answer, a branch + whose provider is gone. Refused while a worker holds a step -- cancel + first. If the run is a child, its parent hears about it in the same + transaction. + """ + result = None + if result_json is not None: + try: + result = json.loads(result_json) + except json.JSONDecodeError as err: + console.error(f"--result is not JSON: {err}") + raise click.exceptions.Exit(1) from None + _finalize(database, run_id, RunStatus.COMPLETED, result=result) + + +@workflows.command() +@database_option +@click.argument("run_id") +@click.option("--reason", required=True, help="Why the run is being given up on.") +def fail(database: str | None, run_id: str, reason: str): + """End a run as failed by operator decision. + + The reason is recorded on the run, so the history says a person decided + this rather than leaving a failure with no explanation. + """ + _finalize(database, run_id, RunStatus.FAILED, error={"reason": reason}) + + @workflows.command() @database_option @click.argument("run_id") diff --git a/tests/units/workflow/test_cli.py b/tests/units/workflow/test_cli.py index 676aea2bdd3..c6a0bbebd88 100644 --- a/tests/units/workflow/test_cli.py +++ b/tests/units/workflow/test_cli.py @@ -215,3 +215,40 @@ def test_stats_filters_to_one_workflow(seeded): "needs_attention": 0, "by_status": {}, } + + +def test_complete_ends_a_stuck_run_by_operator_decision(seeded): + """A wait nobody will answer is closed from the CLI, not the database.""" + database, waiting, _ = seeded + result = _invoke("complete", "-d", database, waiting, "--result", '{"ok": true}') + assert result.exit_code == 0, result.output + shown = _invoke("show", "-d", database, waiting, "--json") + payload = json.loads(shown.output) + assert payload["status"] == RunStatus.COMPLETED.value + assert payload["result"] == {"ok": True} + + +def test_fail_records_the_operator_s_reason(seeded): + """A run given up on says why, so history is not a silent failure.""" + database, waiting, _ = seeded + result = _invoke("fail", "-d", database, waiting, "--reason", "provider retired") + assert result.exit_code == 0, result.output + payload = json.loads(_invoke("show", "-d", database, waiting, "--json").output) + assert payload["status"] == RunStatus.FAILED.value + assert "provider retired" in json.dumps(payload["error"]) + + +def test_finalizing_an_unknown_run_fails(seeded): + """Nothing to finalize is an error, not a silent success.""" + database, _, _ = seeded + result = _invoke("complete", "-d", database, "no-such-run") + assert result.exit_code == 1 + assert "unknown" in result.output + + +def test_complete_refuses_a_result_that_is_not_json(seeded): + """A typo in --result must not be recorded as the string the operator typed.""" + database, waiting, _ = seeded + result = _invoke("complete", "-d", database, waiting, "--result", "{oops") + assert result.exit_code == 1 + assert "not JSON" in result.output From 752511b1d284e16fa53b8638a5999fd108c8ee3f Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 10:33:12 -0700 Subject: [PATCH 075/121] workflows: stop retrying code bugs, and make a release's runs findable Writing an example workflow cold surfaced both of these. A handler that raised TypeError -- I had called an rx verb that does not exist -- was retried four times with backoff before the run failed. A retry re-runs the handler against the same committed state, so a deterministic failure fails identically every attempt: the budget was spent proving the code was still wrong, and the only real effect was delaying the operator seeing it by the length of the backoff. Every resolved policy now excludes BUG_EXCEPTIONS, and a handler that wants one of them retried names it in retry_on. KeyError, IndexError and ValueError stay retryable on purpose -- those routinely come from a dependency returning a body missing a field, which the next attempt may well get right. RunQuery also takes a definition_digest filter now, wired through all three stores and pinned by a conformance check, with `reflex workflows stats --digest`. That is the query 'is anything still running the release I am replacing' needs, which is the gate Alek's version-routing decision asks for. --- news/workflow-bug-classification.bugfix.md | 1 + news/workflow-digest-filter.feature.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 23 +++++++++ reflex/workflow/CONTRACT.md | 9 ++++ reflex/workflow/cli.py | 11 ++++ reflex/workflow/conformance.py | 30 +++++++++++ reflex/workflow/definition.py | 19 ++++++- reflex/workflow/postgres.py | 3 ++ reflex/workflow/records.py | 5 ++ reflex/workflow/store.py | 8 +++ tests/units/workflow/test_cli.py | 1 + tests/units/workflow/test_definition.py | 50 +++++++++++++++++++ 12 files changed, 159 insertions(+), 2 deletions(-) create mode 100644 news/workflow-bug-classification.bugfix.md create mode 100644 news/workflow-digest-filter.feature.md diff --git a/news/workflow-bug-classification.bugfix.md b/news/workflow-bug-classification.bugfix.md new file mode 100644 index 00000000000..80763704c5a --- /dev/null +++ b/news/workflow-bug-classification.bugfix.md @@ -0,0 +1 @@ +A handler that raises `TypeError`, `AttributeError`, `NameError`, `ImportError`, `SyntaxError`, `IndentationError`, or `NotImplementedError` now fails on its first attempt instead of consuming its whole retry budget. A retry re-runs the handler against the same committed state, so a code bug fails identically every attempt; retrying one only delayed the failure reaching an operator by the length of the backoff. A handler that wants one of these retried can still name it in `Retry(retry_on=...)`. Data-shaped errors like `KeyError` and `ValueError` remain retryable, since they routinely come from a dependency returning an unexpected body. diff --git a/news/workflow-digest-filter.feature.md b/news/workflow-digest-filter.feature.md new file mode 100644 index 00000000000..31a5bcb822b --- /dev/null +++ b/news/workflow-digest-filter.feature.md @@ -0,0 +1 @@ +Run queries take a `definition_digest` filter, and `reflex workflows stats --digest ...` reports the runs admitted against one compiled definition — the query a deploy gate needs to answer "is anything still running the release I am replacing". Covered by the shared store conformance suite. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 91670ec135c..99703f15817 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -177,6 +177,29 @@ def is_retryable(self, error: BaseException) -> bool: return isinstance(error, self.retry_on) +BUG_EXCEPTIONS: Final[tuple[type[BaseException], ...]] = ( + TypeError, + AttributeError, + NameError, + ImportError, + SyntaxError, + IndentationError, + NotImplementedError, +) +"""Exceptions that mean the code is wrong, not that the world was unlucky. + +A retry re-runs a handler against the same committed state, so a deterministic +failure fails identically every time: retrying one of these spends the whole +budget proving a bug is still a bug, and delays the run reaching an operator +by the length of the backoff. They are excluded from every resolved policy +unless a handler asks for them by name in ``retry_on``. + +Data-shaped errors -- ``KeyError``, ``IndexError``, ``ValueError`` -- are +deliberately absent: they routinely come from a flaky dependency returning a +body that is missing a field, which the next attempt may well get right. +""" + + def default_retry_for_effect(effect: str) -> Retry: """Return the default retry policy for an effect class. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 383784b93d9..c0aa5ff2379 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -98,6 +98,15 @@ that turn it into effectively-once for side effects are, in order of strength: | `non_idempotent_write` | 1 attempt; no implicit retry | run `NEEDS_ATTENTION`; operator resumes or fails | - `TransientWorkflowError` is always retryable regardless of `retry_on`. +- **Code bugs are never retried.** `TypeError`, `AttributeError`, `NameError`, + `ImportError`, `SyntaxError`, `IndentationError` and `NotImplementedError` + fail the step on the first attempt. A retry re-runs the handler against the + same committed state, so a deterministic failure fails identically every + time; retrying one spends the budget proving the code is still wrong and + delays the operator seeing it by the length of the backoff. A handler that + wants one of them retried names it in `retry_on`, and then it is honored. + `KeyError`, `IndexError` and `ValueError` stay retryable: they routinely + come from a dependency returning a body missing a field. - `timeout=` bounds one attempt; a timed-out attempt counts as a failed one and follows the same policy (`on_timeout` hook runs on final timeout). `timeout=` is a compile error on sync handlers: a thread cannot be diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index c385ccc7723..ebb407589f0 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -820,11 +820,17 @@ async def serve() -> None: @database_option @click.option("--workflow", "-w", default=None, help="Only this workflow id.") @click.option("--label", "-l", "labels", multiple=True, help="Filter as key=value.") +@click.option( + "--digest", + default=None, + help="Only runs admitted against this definition digest.", +) @click.option("--json", "as_json", is_flag=True, help="Emit JSON instead of a table.") def stats( database: str | None, workflow: str | None, labels: tuple[str, ...], + digest: str | None, as_json: bool, ): """Count runs by status: is this deployment healthy, in one screen. @@ -833,6 +839,9 @@ def stats( stopped needing a person. Everything else is context. --json makes this a scrape for whatever collects metrics, so an alert on "runs needing attention" does not require reading the workflow database. + + With --digest it answers the other operational question: is anything still + running the release I am about to replace. """ from reflex.workflow.records import TERMINAL_RUN_STATUSES, RunQuery @@ -851,6 +860,7 @@ async def counts(store: RunStore) -> dict[str, int]: for status in RunStatus: query = RunQuery( workflow_id=workflow, + definition_digest=digest, statuses=(status,), labels=label_filter or None, ) @@ -872,6 +882,7 @@ async def counts(store: RunStore) -> dict[str, int]: click.echo( json.dumps({ "workflow": workflow, + "digest": digest, "total": total, "open": open_runs, "needs_attention": attention, diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 587a552c87e..760454bb835 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -484,6 +484,35 @@ async def check_list_runs_filters_and_orders(store: RunStore) -> None: assert await store.list_runs(RunQuery(workflow_id="other.flow", limit=10)) == () +async def check_runs_are_findable_by_definition_digest(store: RunStore) -> None: + """Runs can be narrowed to the compiled definition that admitted them. + + "Is anything still running the release I am replacing" has to be one + query. Answering it by listing everything and filtering client-side + breaks the moment a deployment has more runs than a page. + """ + await store.admit( + make_run("old", definition_digest="d-old"), make_step("old"), _ADMITTED + ) + await store.admit( + make_run("new1", definition_digest="d-new"), make_step("new1"), _ADMITTED + ) + await store.admit( + make_run("new2", definition_digest="d-new"), make_step("new2"), _ADMITTED + ) + assert await store.count_runs(RunQuery(definition_digest="d-old")) == 1 + assert await store.count_runs(RunQuery(definition_digest="d-new")) == 2 + assert await store.count_runs(RunQuery(definition_digest="d-gone")) == 0 + listed = await store.list_runs(RunQuery(definition_digest="d-new")) + assert {run.run_id for run in listed} == {"new1", "new2"} + assert ( + await store.count_runs( + RunQuery(definition_digest="d-new", statuses=(RunStatus.PENDING,)) + ) + == 2 + ), "the digest filter composes with the others rather than replacing them" + + async def check_count_runs_matches_the_listing(store: RunStore) -> None: """A count answers for the same set a listing would return. @@ -1134,6 +1163,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_resume_only_reopens_a_suspended_run, check_list_runs_filters_and_orders, check_count_runs_matches_the_listing, + check_runs_are_findable_by_definition_digest, check_pagination_skips_nothing_on_tied_timestamps, check_label_filter_handles_awkward_keys, check_flow_control_queries, diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 69cd09fee6b..5a433a5d429 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -18,6 +18,7 @@ from reflex_base.utils.exceptions import WorkflowDefinitionError from reflex_base.workflow import ( + BUG_EXCEPTIONS, Debounce, DurableEventConfig, RateLimit, @@ -250,6 +251,12 @@ def _resolve_retry(retry: Retry | None, effect: str) -> Retry: ``non_idempotent_write`` never retries: the runtime cannot prove the external effect did not already land. + Every policy also excludes ``BUG_EXCEPTIONS``, because a retry re-runs the + handler against the same committed state: a ``TypeError`` fails the same + way every attempt, so retrying one only delays the operator seeing it. A + handler that genuinely wants one retried names it in ``retry_on``, and + then it is left alone. + Args: retry: The explicit policy, if the handler declared one. effect: The handler's effect class. @@ -259,9 +266,17 @@ def _resolve_retry(retry: Retry | None, effect: str) -> Retry: """ if retry is None: retry = default_retry_for_effect(effect) - if retry.retry_on or effect == "non_idempotent_write": + if not retry.retry_on and effect != "non_idempotent_write": + retry = dataclasses.replace(retry, retry_on=(Exception,)) + bugs = tuple( + bug + for bug in BUG_EXCEPTIONS + if not any(issubclass(wanted, bug) for wanted in retry.retry_on) + and not any(issubclass(bug, known) for known in retry.do_not_retry_on) + ) + if not bugs: return retry - return dataclasses.replace(retry, retry_on=(Exception,)) + return dataclasses.replace(retry, do_not_retry_on=retry.do_not_retry_on + bugs) def _compile_handlers( diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 315c4dcecc7..2d33f068e69 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -284,6 +284,9 @@ def _run_filters(query: RunQuery) -> tuple[str, tuple[Any, ...]]: if query.workflow_id is not None: clauses.append("workflow_id = %s") params.append(query.workflow_id) + if query.definition_digest is not None: + clauses.append("definition_digest = %s") + params.append(query.definition_digest) if query.statuses: clauses.append("status = ANY(%s)") params.append([status.value for status in query.statuses]) diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 9888d287a0f..280553740e0 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -315,6 +315,10 @@ class RunQuery: Attributes: workflow_id: Restrict to one workflow identity. + definition_digest: Restrict to runs admitted against one compiled + definition. This is what answers "is anything still running the + release I am replacing", which a deploy gate and an operator + watching a rollout both need. statuses: Restrict to these run statuses; empty means any. labels: Require every one of these server-derived label values. created_before: Pagination cursor, as the ``(created_at, run_id)`` of @@ -325,6 +329,7 @@ class RunQuery: """ workflow_id: str | None = None + definition_digest: str | None = None statuses: tuple[RunStatus, ...] = () labels: Mapping[str, str] | None = None created_before: tuple[float, str] | None = None diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 27e044221da..afb0b2b0914 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -764,6 +764,11 @@ def _matches_query(run: RunRecord, query: RunQuery) -> bool: """ if query.workflow_id is not None and run.workflow_id != query.workflow_id: return False + if ( + query.definition_digest is not None + and run.definition_digest != query.definition_digest + ): + return False if query.statuses and run.status not in query.statuses: return False if query.created_before is not None and (run.created_at, run.run_id) >= ( @@ -2216,6 +2221,9 @@ def _sqlite_run_filters(query: RunQuery) -> tuple[str, tuple[Any, ...]]: if query.workflow_id is not None: clauses.append("workflow_id = ?") params.append(query.workflow_id) + if query.definition_digest is not None: + clauses.append("definition_digest = ?") + params.append(query.definition_digest) if query.statuses: placeholders = ",".join("?" * len(query.statuses)) clauses.append(f"status IN ({placeholders})") diff --git a/tests/units/workflow/test_cli.py b/tests/units/workflow/test_cli.py index c6a0bbebd88..29315f3dccd 100644 --- a/tests/units/workflow/test_cli.py +++ b/tests/units/workflow/test_cli.py @@ -210,6 +210,7 @@ def test_stats_filters_to_one_workflow(seeded): assert result.exit_code == 0, result.output assert json.loads(result.output) == { "workflow": "nope.nothing", + "digest": None, "total": 0, "open": 0, "needs_attention": 0, diff --git a/tests/units/workflow/test_definition.py b/tests/units/workflow/test_definition.py index 7f965c2ddb1..3f79d35d453 100644 --- a/tests/units/workflow/test_definition.py +++ b/tests/units/workflow/test_definition.py @@ -147,6 +147,56 @@ def fetch(self): assert retry.retry_on == (ConnectionError,) +def test_code_bugs_are_not_retried(forked_registration_context): + """A retry re-runs the handler against the same state, so a bug re-fails. + + Retrying a TypeError spends the whole budget proving the code is still + wrong, and delays an operator seeing it by the length of the backoff. + """ + + class Buggy(rx.State): + __workflow__ = WorkflowConfig(id="billing.buggy") + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=5), + ) + def work(self): + pass + + retry = compile_workflow(Buggy).handlers["work"].retry + assert not retry.is_retryable(TypeError("not callable")) + assert not retry.is_retryable(AttributeError("no such attribute")) + assert not retry.is_retryable(NameError("undefined")) + # A dependency returning a body without the field you wanted is bad luck, + # not a bug, and the next attempt may well get it right. + assert retry.is_retryable(KeyError("items")) + assert retry.is_retryable(ValueError("bad json")) + assert retry.is_retryable(ConnectionError("dropped")) + + +def test_a_handler_may_ask_for_a_bug_class_to_be_retried(forked_registration_context): + """The exclusion is a default, not a rule the engine imposes.""" + + class Insists(rx.State): + __workflow__ = WorkflowConfig(id="billing.insists") + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=3, retry_on=(TypeError,)), + ) + def work(self): + pass + + retry = compile_workflow(Insists).handlers["work"].retry + assert retry.is_retryable(TypeError("this one is transient, honestly")) + assert not retry.is_retryable(AttributeError("still a bug")) + + def test_missing_workflow_config(forked_registration_context): class NoConfig(rx.State): @rx.event(durable=True, trigger=manual(), effect="none") From 3eaa987c718e4ce7d851b5dd4d835f09dcc42edb Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 10:38:27 -0700 Subject: [PATCH 076/121] workflows: document the failures that are not crashes, and pin them Section 8 covered kill points thoroughly and said nothing about the failures where nothing was killed -- the engine simply refused to go on. Four of them reached a run with a stable reason no document mentioned: unknown_workflow, unknown_handler, incompatible_payload, max_steps_exceeded, recovery_budget_exhausted, and the fallback that replaces an unserializable error payload rather than losing the failure. Each now has a row saying what it does to the run and why that is the right answer -- a worker that does not serve a workflow suspends rather than fails it, because it has no business deciding that workflow's fate; a runaway loop fails rather than suspending, because approving more of the same is not a decision worth offering. A test walks it in both directions: every reason the contract names must still be emitted, and every reason the engine emits must appear in the contract. Renaming one in the code without the other now fails the build, instead of quietly invalidating whatever runbook matches on it. --- news/workflow-failure-matrix.feature.md | 1 + reflex/workflow/CONTRACT.md | 18 +++++ .../workflow/test_contract_failure_matrix.py | 65 +++++++++++++++++++ 3 files changed, 84 insertions(+) create mode 100644 news/workflow-failure-matrix.feature.md create mode 100644 tests/units/workflow/test_contract_failure_matrix.py diff --git a/news/workflow-failure-matrix.feature.md b/news/workflow-failure-matrix.feature.md new file mode 100644 index 00000000000..96f3f9083a5 --- /dev/null +++ b/news/workflow-failure-matrix.feature.md @@ -0,0 +1 @@ +The execution contract's failure matrix now documents the failures that are not crashes — an unregistered workflow, a deleted handler, a payload a handler no longer accepts, a run that blew `max_steps`, a step that exhausted its recovery budget, and an unserializable error payload — each with its stable `reason` and its one outcome. A test walks both directions, so a reason renamed in the engine and not in the contract fails the build instead of silently invalidating a runbook. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index c0aa5ff2379..cb37ee8fb98 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -267,6 +267,24 @@ outcome. | store unreachable at commit | attempt abandoned (fence unverifiable); step recovered later; `rx.step` records already made stand | | everything down for an hour | timers/waits/retries fire on restart (due-time semantics); schedule occurrences catch up from the durable cursor, capped at `MAX_SCHEDULE_CATCHUP` per schedule, remainder skipped with a history record | +### Failures that are not crashes + +Nothing was killed; the engine refused to proceed. Each has one outcome and +one stable `reason` on the run, so an operator (or an alert) matches on the +reason rather than on message text. + +| what happened | `reason` | outcome | +|---|---|---| +| the run's workflow class is not registered in this process | `unknown_workflow` | step and run `NEEDS_ATTENTION`; re-register and `resume(run)`. Not a failure: a worker that does not serve a workflow must not decide that workflow's fate | +| the next step's handler no longer exists | `unknown_handler` | step and run `NEEDS_ATTENTION`; restore the handler and `resume(run)`, or cancel | +| a recorded payload carries arguments the handler no longer accepts | `incompatible_payload` | step and run `NEEDS_ATTENTION`, naming the arguments; restore the parameters or cancel | +| a run allocated more steps than `WorkflowConfig.max_steps` | `max_steps_exceeded` | the committing step succeeds, every other open slot is tombstoned, run `FAILED`. The bound is on a runaway loop, so it fails rather than suspending for a person to approve more of the same | +| a step's lease lapsed more times than `max_recoveries` | `recovery_budget_exhausted` | run `FAILED`. Infrastructure recoveries are free of the retry budget precisely so they can be bounded separately; the bound is what stops a poison step cycling forever | +| an error's `details` cannot be serialized | — | the reason is preserved and the details are replaced by `{"unserializable": repr(...)}`. Losing the payload never turns a failure into a crash | + +Nothing on this table is silent: each writes its reason to the run's error and +its own history event. + ## 9. Operator actions Every action is legal only from the states listed; anything else is a refused diff --git a/tests/units/workflow/test_contract_failure_matrix.py b/tests/units/workflow/test_contract_failure_matrix.py new file mode 100644 index 00000000000..b96757d773a --- /dev/null +++ b/tests/units/workflow/test_contract_failure_matrix.py @@ -0,0 +1,65 @@ +"""The failure matrix and the code have to agree on the reasons. + +CONTRACT.md section 8 promises one documented outcome per failure, keyed by a +stable ``reason`` an operator or an alert can match on. A reason renamed in +the engine and not in the contract turns that promise into a lie that nothing +would catch: the tests pass, the document reads fine, and the runbook stops +matching production. + +This walks the reasons the contract names and asserts the engine still emits +each one. It is a spelling check, not a behaviour test -- the behaviour is +covered where each failure is produced. +""" + +import re +from pathlib import Path + +import pytest + +CONTRACT = Path(__file__).parents[3] / "reflex" / "workflow" / "CONTRACT.md" +SOURCES = ( + Path(__file__).parents[3] / "reflex" / "workflow" / "kernel.py", + Path(__file__).parents[3] / "reflex" / "workflow" / "store.py", +) + +DOCUMENTED_REASONS = ( + "unknown_workflow", + "unknown_handler", + "incompatible_payload", + "max_steps_exceeded", + "recovery_budget_exhausted", + "run_timeout", +) + + +@pytest.mark.parametrize("reason", DOCUMENTED_REASONS) +def test_every_documented_reason_is_still_emitted(reason): + """A reason the contract names must exist in the engine. + + Args: + reason: The stable failure reason under test. + """ + emitted = any(f'"{reason}"' in source.read_text() for source in SOURCES) + assert emitted, ( + f"CONTRACT.md documents the failure reason {reason!r}, but no engine " + "source emits it. Rename it in both places or drop the row." + ) + + +def test_every_emitted_reason_is_documented(): + """A failure the engine can produce must have a documented outcome. + + This is section 1's exit criterion in test form: every failure scenario + has one unambiguous documented outcome, so a new failure mode cannot be + added without saying what it does to the run. + """ + contract = CONTRACT.read_text() + emitted: set[str] = set() + for source in SOURCES: + emitted.update(re.findall(r'"reason":\s*"(\w+)"', source.read_text())) + undocumented = sorted(name for name in emitted if name not in contract) + assert not undocumented, ( + f"These failure reasons are emitted but not in CONTRACT.md: " + f"{undocumented}. Add a row to the failure matrix saying what each " + "does to the run." + ) From 9fa81ab3587935d9ce8cce1e1c47a16c14c285f2 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 11:03:34 -0700 Subject: [PATCH 077/121] workflows: fix client isolation, store leaks, and admission binding From an external test report against 752511b1 that built ten workflow families and drove them adversarially. Four defects, each reproduced before the fix and verified to fail on the unfixed tree. connect() published itself in a process global, so two clients open at once -- the advertised FastAPI/Django shape, where requests overlap by definition -- could route each other's runs into the wrong store, across a tenant boundary. It binds a context variable now. It also never closed stores it opened, leaking a SQLite connection or a Postgres pool per request; it closes what it owns and leaves a caller's store alone. The redeploy compatibility gate checked for arguments a handler no longer accepts but not for parameters it newly requires, so adding one failed in-flight runs with a TypeError from inside the handler instead of suspending them as incompatible_payload. Waits, joins and child arrivals are handed their first argument at dispatch, so the gate counts that as supplied rather than reporting it missing. The HTTP start endpoint resolved handlers by Python name only, so the stable id every run reports back was rejected -- the API was not round-trippable. It now takes either. It also admitted starts whose arguments could not bind, answering 202 and leaving behind a run whose first attempt could only ever raise; that is a 400 now. Also untracks a workflow.db committed by accident and ignores it. --- .gitignore | 4 + news/workflow-admission-binding.bugfix.md | 1 + news/workflow-client-isolation.bugfix.md | 1 + reflex/workflow/api.py | 30 ++++- reflex/workflow/definition.py | 29 +++++ reflex/workflow/kernel.py | 27 ++++ reflex/workflow/runtime.py | 53 ++++++-- tests/units/workflow/test_api.py | 63 +++++++++- tests/units/workflow/test_client_isolation.py | 118 ++++++++++++++++++ tests/units/workflow/test_versioning.py | 97 ++++++++++++++ workflow.db | Bin 77824 -> 0 bytes 11 files changed, 409 insertions(+), 14 deletions(-) create mode 100644 news/workflow-admission-binding.bugfix.md create mode 100644 news/workflow-client-isolation.bugfix.md create mode 100644 tests/units/workflow/test_client_isolation.py delete mode 100644 workflow.db diff --git a/.gitignore b/.gitignore index 533bcfcec8e..6541779a9a4 100644 --- a/.gitignore +++ b/.gitignore @@ -33,3 +33,7 @@ CLAUDE.local.md # Backups written by scripts/delete_automated_releases.sh automated-releases-backup-*.json + +# Workflow run stores created by `reflex workflows` commands. +workflow.db +workflow.db-* diff --git a/news/workflow-admission-binding.bugfix.md b/news/workflow-admission-binding.bugfix.md new file mode 100644 index 00000000000..ffbd6083cbe --- /dev/null +++ b/news/workflow-admission-binding.bugfix.md @@ -0,0 +1 @@ +Adding a required parameter to a handler now suspends in-flight runs as `NEEDS_ATTENTION` with `incompatible_payload`, naming the parameter to give a default, instead of failing them with a `TypeError` from inside the handler. The HTTP start endpoint accepts a handler's stable `id` as well as its Python method name — the stable id is what runs report, so it has to be what the API takes — and refuses a start whose arguments cannot bind, rather than returning 202 and creating a run that can never succeed. diff --git a/news/workflow-client-isolation.bugfix.md b/news/workflow-client-isolation.bugfix.md new file mode 100644 index 00000000000..0336f92c7df --- /dev/null +++ b/news/workflow-client-isolation.bugfix.md @@ -0,0 +1 @@ +`rx.workflows.connect()` now binds the active client to a context variable rather than a process global, so two clients open at once — one per request, or one per tenant — can no longer route each other's runs into the wrong store. It also closes stores it opened itself, which stops a request-scoped client leaking a SQLite connection or a Postgres pool per request; a store passed in by the caller is left open, since it belongs to the caller. diff --git a/reflex/workflow/api.py b/reflex/workflow/api.py index 5f9f15cd9bd..0aa01debf4e 100644 --- a/reflex/workflow/api.py +++ b/reflex/workflow/api.py @@ -23,6 +23,7 @@ from reflex_base.utils import console from starlette.responses import JSONResponse, PlainTextResponse +from reflex.workflow.definition import unbound_params from reflex.workflow.records import attempts_made if TYPE_CHECKING: @@ -118,15 +119,36 @@ async def endpoint(request: Request) -> JSONResponse: ) if definition is None: return JSONResponse({"error": "unknown workflow"}, status_code=404) - target = getattr(definition.state_cls, handler_name, None) - if target is None: + # The stable id is what the run records and what every read surface + # reports, so it has to be what the write surface accepts; a caller + # that read `handler: "accept_order"` back must be able to send it. + # The Python name stays valid so renaming neither breaks callers. + handler = next( + ( + candidate + for candidate in definition.handlers.values() + if handler_name in (candidate.id, candidate.name) + ), + None, + ) + if handler is None: return JSONResponse({"error": "unknown handler"}, status_code=404) + target = getattr(definition.state_cls, handler.name) - args = payload.get("args") or {} - if not isinstance(args, dict): + raw_args = payload.get("args") + if raw_args is not None and not isinstance(raw_args, dict): return JSONResponse( {"error": "args must be a JSON object"}, status_code=400 ) + args = raw_args or {} + missing = sorted(unbound_params(handler, set(args))) + if missing: + # Admitting this would create a run that cannot possibly run: the + # worker would raise TypeError on the first attempt and the caller + # would have a 202 and a poisoned run id. + return JSONResponse( + {"error": f"missing required arguments: {missing}"}, status_code=400 + ) labels = payload.get("labels") try: result = await runtime.kernel.start( diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 5a433a5d429..241a3d55b28 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -675,6 +675,35 @@ def _compute_digest( return hashlib.sha256(payload.encode()).hexdigest() +def unbound_params(handler: HandlerDefinition, supplied: set[str]) -> set[str]: + """Parameters a handler requires that a recorded payload cannot fill. + + A parameter with a default binds without help, and ``*args``/``**kwargs`` + absorb anything, so only names that must be passed and are not present + count. This is what makes "the code changed under an in-flight run" a + compatibility decision rather than a TypeError from inside the handler. + + Args: + handler: The current definition of the handler. + supplied: Payload argument names the step recorded. + + Returns: + The required parameter names that nothing would bind. + """ + signature = inspect.signature(handler.fn) + required: set[str] = set() + for index, (name, param) in enumerate(signature.parameters.items()): + if index == 0 or param.default is not inspect.Parameter.empty: + continue + if param.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + continue + required.add(name) + return required - supplied + + def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: """Compile a workflow class into an immutable definition. diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 6da1f688521..1cc648679ab 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -42,6 +42,7 @@ from reflex.event import EventHandler, EventSpec from reflex.workflow.context import RunContext, bind_run, unbind_run from reflex.workflow.cron import CronSchedule +from reflex.workflow.definition import unbound_params from reflex.workflow.records import ( TERMINAL_RUN_STATUSES, TERMINAL_STEP_STATUSES, @@ -1986,6 +1987,32 @@ def _incompatible_reason( } supplied = {key for key in claim.step.args if not key.startswith("__")} unexpected = sorted(supplied - set(handler.params)) + if handler.params and ( + # A wait's continuation, a join, and a child arrival are all + # handed their first argument at dispatch rather than carrying it + # in the recorded payload, so it is supplied even though the + # recorded args do not name it. + "__payload__" in claim.step.args + or "__results__" in claim.step.args + or "__wait__" in claim.step.args + ): + supplied.add(handler.params[0]) + missing = sorted(unbound_params(handler, supplied)) + if missing: + # A parameter added with no default is the mirror image of a + # deleted one, and just as much a redeploy problem: the recorded + # payload cannot fill it. Dispatching anyway fails the run with a + # TypeError from deep inside the handler, which tells the operator + # nothing about what to ship to fix it. + return { + "reason": "incompatible_payload", + "handler_id": handler.id, + "detail": ( + f"Handler {handler.id!r} now requires {missing}, which the " + "recorded payload does not carry; give them defaults, or " + "cancel the run." + ), + } if unexpected: return { "reason": "incompatible_payload", diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 4f2c8d4e169..03e85bd36d3 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -8,6 +8,7 @@ from __future__ import annotations +import inspect import os import random import time @@ -165,6 +166,15 @@ def definitions(self) -> tuple[WorkflowDefinition, ...]: """ return tuple(self._definitions.values()) + @property + def store(self) -> RunStore | None: + """The store this runtime reads and writes, once one is resolved. + + Returns: + The store, or None before startup resolved one. + """ + return self._store + @property def kernel(self) -> WorkflowKernel: """The running kernel. @@ -275,6 +285,24 @@ def configured_drain() -> float: return 0.0 +async def _close_store(store: RunStore | None) -> None: + """Release a store's connections, whatever kind of close it has. + + SQLite closes synchronously, Postgres closes a pool asynchronously, and + the memory store has nothing to close. A caller that opened a store per + request has to be able to hand it back without knowing which. + + Args: + store: The store to close, if any. + """ + closer = getattr(store, "close", None) + if closer is None: + return + result = closer() + if inspect.isawaitable(result): + await result + + def get_runtime() -> WorkflowRuntime: """Resolve the active workflow runtime. @@ -334,21 +362,30 @@ async def connect( Yields: The client runtime. """ - global _default_runtime - + owned = store is None runtime = WorkflowRuntime( store if store is not None else resolve_store(database) ) - for workflow_cls in workflow_classes: - runtime.register(workflow_cls) - await runtime.startup(start_worker=False) - previous = _default_runtime - _default_runtime = runtime + token = None try: + for workflow_cls in workflow_classes: + runtime.register(workflow_cls) + await runtime.startup(start_worker=False) + # A context variable, not the process global: two clients open at + # once -- one per request, one per tenant -- must not be able to + # send each other's work to the wrong store, and the global is + # shared by every task in the process. + token = _context_runtime.set(runtime) yield runtime finally: - _default_runtime = previous + if token is not None: + _context_runtime.reset(token) await runtime.shutdown() + if owned: + # The store was opened here, so its connections are this + # block's to close. A caller-injected store belongs to the + # caller and is left alone. + await _close_store(runtime.store) @staticmethod async def start( diff --git a/tests/units/workflow/test_api.py b/tests/units/workflow/test_api.py index 43503a7fd8e..b21689bef3a 100644 --- a/tests/units/workflow/test_api.py +++ b/tests/units/workflow/test_api.py @@ -38,7 +38,7 @@ class Orders(rx.State): __workflow__ = WorkflowConfig(id="api.orders") order: str = "" - @rx.event(durable=True, trigger=manual(), effect="none") + @rx.event(durable=True, trigger=manual(), effect="none", id="accept_order") def place(self, order: str): """Record the order. @@ -101,7 +101,10 @@ def test_a_service_can_start_and_read_a_run(client): read = client.get(f"/_workflow/api/runs/{run_id}", headers=_auth()) assert read.status_code == 200 assert read.json()["workflow"] == "api.orders" - assert read.json()["steps"][0]["handler"] == "place" + assert read.json()["steps"][0]["handler"] == "accept_order", ( + "the read surface reports the stable id, which is why the write " + "surface has to accept it" + ) def test_an_idempotency_key_returns_the_same_run(client): @@ -196,3 +199,59 @@ def test_prometheus_rendering_escapes_label_values(): }) assert 'workflow="we\\"ird\\\\flow"' in text assert text.endswith("\n") + + +def test_the_stable_handler_id_starts_a_run(client): + """What a run reads back as its handler must be what the API accepts. + + Args: + client: The test client. + """ + body = json.dumps({ + "workflow": "api.orders", + "handler": "accept_order", + "args": {"order": "o-id"}, + }) + response = client.post(START_ROUTE, content=body, headers=_auth()) + assert response.status_code == 202, response.text + run_id = response.json()["run_id"] + read = client.get(f"/_workflow/api/runs/{run_id}", headers=_auth()) + assert read.json()["steps"][0]["handler"] == "accept_order" + + +def test_the_python_name_still_works(client): + """Accepting the stable id must not break callers using the method name. + + Args: + client: The test client. + """ + body = json.dumps({ + "workflow": "api.orders", + "handler": "place", + "args": {"order": "o-name"}, + }) + assert client.post(START_ROUTE, content=body, headers=_auth()).status_code == 202 + + +def test_missing_required_arguments_are_refused(client): + """A run that cannot possibly run is not admitted. + + Args: + client: The test client. + """ + body = json.dumps({"workflow": "api.orders", "handler": "place"}) + response = client.post(START_ROUTE, content=body, headers=_auth()) + assert response.status_code == 400, response.text + assert "order" in response.json()["error"] + + +def test_a_falsey_non_object_args_is_refused(client): + """`args: []` is a caller mistake, not an empty payload. + + Args: + client: The test client. + """ + body = json.dumps({"workflow": "api.orders", "handler": "place", "args": []}) + response = client.post(START_ROUTE, content=body, headers=_auth()) + assert response.status_code == 400, response.text + assert "JSON object" in response.json()["error"] diff --git a/tests/units/workflow/test_client_isolation.py b/tests/units/workflow/test_client_isolation.py new file mode 100644 index 00000000000..9ac5153719c --- /dev/null +++ b/tests/units/workflow/test_client_isolation.py @@ -0,0 +1,118 @@ +"""Tests for `rx.workflows.connect()` as a per-scope client. + +The advertised use is a web process: a FastAPI route or a Django view opens a +client, starts a run, and returns. Two requests overlap all the time, and in a +multi-tenant deployment they may be pointed at different stores. A client that +published itself process-wide would let one request's work land in another +request's database, which is the one failure a tenant boundary must not have. +""" + +import asyncio +import sqlite3 +from typing import Any + +import pytest +from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import RunQuery +from reflex.workflow.runtime import workflows +from reflex.workflow.store import MemoryRunStore, SqliteRunStore + + +def _flow(workflow_id: str) -> Any: + """Build a trivial workflow class. + + Args: + workflow_id: The workflow identity to register under. + + Returns: + The workflow class. + """ + + class Flow(rx.State): + __workflow__ = WorkflowConfig(id=workflow_id) + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Complete immediately. + + Returns: + Completion. + """ + return rx.complete(result=workflow_id) + + return Flow + + +async def test_two_clients_do_not_route_work_to_each_others_stores( + forked_registration_context, +): + """Overlapping clients each admit into their own store, not the last one open. + + Binding the active client to a context variable rather than a process + global is what makes this hold: both tasks run in the same process, and + one of them opened its client second. + """ + flow = _flow("isolation.flow") + store_a, store_b = MemoryRunStore(), MemoryRunStore() + b_open = asyncio.Event() + a_submitted = asyncio.Event() + + async def client_a() -> None: + """Open first, then submit while the second client is also open.""" + async with workflows.connect(flow, store=store_a): + await asyncio.wait_for(b_open.wait(), timeout=5) + await workflows.submit(flow.start()) + a_submitted.set() + + async def client_b() -> None: + """Open second and stay open across the other client's submit.""" + async with workflows.connect(flow, store=store_b): + b_open.set() + await asyncio.wait_for(a_submitted.wait(), timeout=5) + await workflows.submit(flow.start()) + + await asyncio.gather(client_a(), client_b()) + + assert len(await store_a.list_runs(RunQuery())) == 1 + assert len(await store_b.list_runs(RunQuery())) == 1 + + +async def test_a_client_closes_the_store_it_opened( + tmp_path, forked_registration_context +): + """A request-scoped client must not leak a connection per request.""" + flow = _flow("isolation.owned") + database = tmp_path / "owned.db" + async with workflows.connect(flow, database=str(database)) as runtime: + opened = runtime.store + await workflows.submit(flow.start()) + assert isinstance(opened, SqliteRunStore) + # A closed sqlite3 connection refuses further work; if the store were + # still open this would succeed and the connection would have leaked. + with pytest.raises(sqlite3.ProgrammingError): + await opened.list_runs(RunQuery()) + + +async def test_a_client_leaves_a_caller_supplied_store_open( + forked_registration_context, +): + """A store the caller owns outlives the block that borrowed it.""" + flow = _flow("isolation.borrowed") + store = MemoryRunStore() + async with workflows.connect(flow, store=store): + await workflows.submit(flow.start()) + assert len(await store.list_runs(RunQuery())) == 1 + + +async def test_leaving_a_client_scope_restores_the_previous_one( + forked_registration_context, +): + """A client is a scope; leaving it must not leave the process bound.""" + flow = _flow("isolation.scoped") + async with workflows.connect(flow, store=MemoryRunStore()): + pass + with pytest.raises(WorkflowRuntimeError): + await workflows.submit(flow.start()) diff --git a/tests/units/workflow/test_versioning.py b/tests/units/workflow/test_versioning.py index 1a02d7ad090..a84bbfb05bc 100644 --- a/tests/units/workflow/test_versioning.py +++ b/tests/units/workflow/test_versioning.py @@ -187,3 +187,100 @@ def go(self): assert result.run_id is not None assert not await harness.resume(result.run_id) assert not await harness.resume("no-such-run") + + +async def test_a_newly_required_parameter_suspends(forked_registration_context): + """Adding a parameter with no default is a redeploy problem, not a crash. + + The pending step recorded no arguments, so the new signature cannot bind. + Failing the run buries a deploy mistake in a TypeError from inside the + handler; suspending names what to ship and leaves the run resumable. + """ + store = MemoryRunStore() + first = _flow() + async with WorkflowTestHarness(first, store=store) as harness: + result = await harness.start(first.begin()) + assert result.run_id is not None + run_id, resume_at = result.run_id, harness.now + + class Widened(rx.State): + __workflow__ = WorkflowConfig(id="versioning.deployed") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Start the run. + + Returns: + The next step. + """ + self.status = "started" + return rx.after("1h", Widened.finish) + + @rx.event(durable=True, effect="read") + async def finish(self, ticket: str): + """Finish, now demanding an argument nothing recorded. + + Args: + ticket: The newly required argument. + """ + self.status = ticket + + async with WorkflowTestHarness( + Widened, store=store, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + assert snapshot.error is not None + assert snapshot.error["reason"] == "incompatible_payload" + assert "ticket" in snapshot.error["detail"] + + +async def test_a_new_parameter_with_a_default_is_compatible( + forked_registration_context, +): + """The supported way to widen a handler still deploys without a suspension.""" + store = MemoryRunStore() + first = _flow() + async with WorkflowTestHarness(first, store=store) as harness: + result = await harness.start(first.begin()) + assert result.run_id is not None + run_id, resume_at = result.run_id, harness.now + + class Defaulted(rx.State): + __workflow__ = WorkflowConfig(id="versioning.deployed") + status: str = "pending" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Start the run. + + Returns: + The next step. + """ + self.status = "started" + return rx.after("1h", Defaulted.finish) + + @rx.event(durable=True, effect="read") + async def finish(self, ticket: str = "unset"): + """Finish, with a default the recorded payload can leave alone. + + Args: + ticket: Optional argument. + + Returns: + Completion. + """ + self.status = ticket + return rx.complete(result=ticket) + + async with WorkflowTestHarness( + Defaulted, store=store, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == "unset" diff --git a/workflow.db b/workflow.db deleted file mode 100644 index fbd1be8eb72ff865519ebedb32fdfbce19eac82b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 77824 zcmeI%&raM%9KiA2g{)x-fsjf*^x{*&f`p=~QlmC%X|j!?5@-mh$;mS70o*w2#r7s} zs+Oiw6;lu4I^O@hw*qWzLw(LM?hpzY147AT0 zi;Y&R@r|Z68jY6x-jd(+Gbevcrw8(XGxNLZXRC4N?ML$q|2CFde>UcSU-+x@$NbK$ z=xopI&wir3omp%@P##YI(^jURPyIdhl{|?L0#^x)W|kE7-g0wv*X~>5)V5BI*RJ>C z&~aZI-mvc*1Je`zAeVOQ{%&`3ud6-Ue$aiURb8U(Y-_p8w3WE5uw28o*0ki%^(?z@ zI{N6^O<5JaIoh~f6+aLI-*ANKi*cDz+ER60GS}ql5~XV^eqaVee@%-DHN@$__JnVk zfj(MVRMbsbkB67*VfPQ*Q)3|dvdTx}nnbBLt80?Gs-wwI)+Wl;R#KJMrX3hB#Mzpb z*3cIx`pH7vO>>v)6m?U4?zl!_+G2HuGMD745~imz`uP|0imL0)H>-hp;9R&z^&Q_S zI=YeEYH#z~t?osyCoPLwEQ^+%uI+W7?P=RP^1HvawHD>14KB;Ggf$$9kfSVD!pI=p z>ptx67Rnr$XO3%@-Y%9ZE3W4WGZ2=H0Zj(@R-worb{{?7+{tEd|rayBU__Ksyt zIl{kSOj9})j%j`ZIR?o@L>$S*!0G7sJ8kvHjmZs3Z;@?Dnl?AuRupx0wR!$D-;_Ay z#_fm?XUZE9XC^HeUs66GJ#jJ=el!V{4U_7ToB`2RGB;V%GDXVny*JZVpR7()Nru$L z_=po`M~BxG_4e)N`OkTeW7`j0?`(YI5wWoa=lY64)z7n1-CZB2zgj$skb(dI7S;5DP_CoX@3f9aL7EBw==j>3J55Dh zUT*&OaX$FNhn(>o5gxS4LqE(anH@6YHE#d7tn%4ZV%dE=u-(34*++6-E#s~p1Yu60 zBjVde8HNmRUlx5SOyYC;;xsU>Uf$(Huk6Z9OcY}zdKF05ytvrhbA0B~`~!Vjr}&!8 z3Xr+NF%N_j5AyglOPH2p_eC^n;-sG0?+M393V4Sof4_k2=yzLf_3`q=={AguOt6KA zQ}PvSCVb7B5C1HNfA}DP00IagfB*srAb;OKmY**5I_I{1Q0*~ z0R#|0p#B2v|JT2dX(0j#AbU;jR) zg$N*k00IagfB*srAb Date: Wed, 19 Aug 2026 11:07:36 -0700 Subject: [PATCH 078/121] workflows: hand join results over in declaration order rx.parallel documents that results arrive in the order the branches were declared. They arrived in the order the branches finished, so a caller unpacking `slow, fast = results` got them backwards whenever the second branch was quicker -- which is the normal case, and silent. A fan-out already stamps every child with child::: as its request key, because that is what makes re-running the fan-out idempotent. That index is the declaration order, already durable, so the arrival carries it and the join sorts by it. No new state, no schema change, and it works the same on all three stores. All three places that build an arrival now carry the branch: the finalize path, the commit path, and the race-winner path. Only patching the first is why the regression test still failed at that point. --- news/workflow-join-order.bugfix.md | 1 + reflex/workflow/kernel.py | 34 +++++++++++ tests/units/workflow/test_parallel.py | 87 +++++++++++++++++++++++++++ 3 files changed, 122 insertions(+) create mode 100644 news/workflow-join-order.bugfix.md diff --git a/news/workflow-join-order.bugfix.md b/news/workflow-join-order.bugfix.md new file mode 100644 index 00000000000..baed1a98dca --- /dev/null +++ b/news/workflow-join-order.bugfix.md @@ -0,0 +1 @@ +`rx.parallel(..., then=...)` hands its results to the join in declaration order, as the API promises, instead of the order the branches happened to finish. A fan-out already stamps each child with its branch index as part of its idempotent request key, so the arrival carries that index and the join sorts by it — no extra state and no store schema change. diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 1cc648679ab..455faa967a5 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -13,6 +13,7 @@ import asyncio import contextlib import dataclasses +import operator import random import time import traceback @@ -185,6 +186,27 @@ def on_event( observer.on_event(event_type, run_id, workflow_id, data) +def _branch_index(request_key: str | None) -> int | None: + """Recover which branch of a fan-out a child run was. + + A branch's declaration index is already durable: fan-out stamps each child + with ``child:::`` as its request key, which is what + makes a re-run of the fan-out idempotent. Reading it back is what lets a + join report results in the order they were declared rather than the order + they happened to finish. + + Args: + request_key: The child run's admission key. + + Returns: + The branch index, or None when the run was not admitted by a fan-out. + """ + if not request_key or not request_key.startswith("child:"): + return None + _, _, tail = request_key.rpartition(":") + return int(tail) if tail.isdigit() else None + + class MetricsObserver(WorkflowObserver): """Tallies the numbers a deployment alerts on. @@ -1196,6 +1218,15 @@ async def _invoke( args = {key: value for key, value in args.items() if key != "__wait__"} delivered = args.pop("__payload__", None) results = args.pop("__results__", None) + if isinstance(results, list) and all( + isinstance(entry, dict) and isinstance(entry.get("branch"), int) + for entry in results + ): + # A join accumulates arrivals as they land, which is finishing + # order; rx.parallel promises declaration order, and a caller + # unpacking `a, b = results` has no other way to tell which is + # which. + results = sorted(results, key=operator.itemgetter("branch")) if delivered is None and results is not None: delivered = results if delivered is not None and handler.params: @@ -2341,6 +2372,7 @@ def _arrival_for( "status": status.value, "result": result, "error": error, + "branch": _branch_index(run.request_key), }, run.run_id, ) @@ -2378,6 +2410,7 @@ def _with_parent_arrival( "status": completion.run_status.value, "result": completion.result, "error": completion.run_error, + "branch": _branch_index(run.request_key), }, run.run_id, ), @@ -2449,6 +2482,7 @@ async def _report_outcome( "status": status.value, "result": result, "error": error, + "branch": _branch_index(run.request_key), }, run.run_id, self._clock(), diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 00b7561403f..74b90906199 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -488,3 +488,90 @@ async def test_a_childs_arrival_commits_with_its_final_transition(): snapshot = await harness.get_run(result.run_id) assert snapshot is not None assert snapshot.status is RunStatus.COMPLETED + + +async def test_join_results_arrive_in_declaration_order(forked_registration_context): + """`a, b = results` has to mean what the fan-out said, not who won. + + The join accumulates arrivals as they land. A branch that takes a day and + a branch that takes a second finish in the opposite order to the one they + were written in, and a caller unpacking the list has no other way to tell + which result is which. + """ + + class Slow(rx.State): + __workflow__ = WorkflowConfig(id="fan.slow") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, lead: str): + """Finish only after a delay. + + Args: + lead: The lead identifier. + + Returns: + The delayed completion. + """ + return rx.after("1h", Slow.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + """Complete late. + + Returns: + Completion. + """ + return rx.complete(result="slow") + + class Fast(rx.State): + __workflow__ = WorkflowConfig(id="fan.fast") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, lead: str): + """Complete immediately. + + Args: + lead: The lead identifier. + + Returns: + Completion. + """ + return rx.complete(result="fast") + + class Ordered(rx.State): + __workflow__ = WorkflowConfig(id="fan.ordered") + seen: list = [] + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Fan out slow first, fast second. + + Returns: + The parallel fan-out. + """ + return rx.parallel( + Slow.start("lead"), Fast.start("lead"), then=Ordered.route + ) + + @rx.event(durable=True, effect="none") + def route(self, results: list): + """Record the results in the order they were handed over. + + Args: + results: One entry per branch. + + Returns: + Completion carrying the ordered results. + """ + return rx.complete(result=[entry["result"] for entry in results]) + + async with WorkflowTestHarness(Ordered, Slow, Fast) as harness: + started = await harness.start(Ordered.begin()) + assert started.run_id is not None + await harness.advance("2h") + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == ["slow", "fast"], ( + "the fast branch finished first, but it was declared second" + ) From ce4f51f4467ed4443876c69bd686c976af5d5e20 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 11:12:02 -0700 Subject: [PATCH 079/121] workflows: turn the two open report defects into executable specs The approval defect from the external report is now a strict xfail: an approval mints approve and reject with distinct delivery keys on purpose, so a second person can reject after a first approved, but the losing alternative is then buffered and consumed by the next wait on the same channel. The test reproduces it end to end. Strict, so whoever fixes it is told to remove the marker rather than leaving a passing xfail behind. The race-loser test is a plain passing test, because writing it honestly showed my first version proved nothing: in a single-kernel harness the loser is cancelled correctly, and the marker would have been a lie about what is covered. What it asserts is the case that does hold -- the worker that resolves a race stops the loser it owns -- and its docstring is explicit that the crash and cross-worker variants the report found are not covered by it, so a green run on this file is not mistaken for the whole guarantee. Both remaining fixes need cancellation intent and decision-group identity written durably, which is a store-shape change I am not starting on a context budget that cannot finish it. --- tests/units/workflow/test_parallel.py | 36 +++++++++ tests/units/workflow/test_waits.py | 108 ++++++++++++++++++++++++++ 2 files changed, 144 insertions(+) diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 74b90906199..e5de6e93462 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -575,3 +575,39 @@ def route(self, results: list): assert snapshot.result == ["slow", "fast"], ( "the fast branch finished first, but it was declared second" ) + + +async def test_a_race_loser_performs_no_effect_after_it_has_lost( + forked_registration_context, +): + """One worker that resolves a race does stop the loser it is running. + + This is the case that holds: the kernel that saw the winner arrive + cancels the losing child it owns, so the loser's deferred answer never + runs and whatever external thing it would have done does not happen. + + It is deliberately not the whole guarantee. Cancelling the loser is + advisory follow-up work by one worker, so it does not survive that worker + dying between the winning commit and the cleanup, and it does not reach a + loser another worker is already executing. Both of those need cancellation + intent written in the same transaction that resolves the race, with + commits fenced against it, and neither is covered here -- a passing test + on this file says nothing about them. + """ + RACE_CALLS.clear() + + async with WorkflowTestHarness(Shopper, SlowVendor, FastVendor) as harness: + started = await harness.start(Shopper.start()) + assert started.run_id is not None + await harness.run_until_idle() + + parent = await harness.get_run(started.run_id) + assert parent is not None + assert parent.status is RunStatus.COMPLETED, "the fast vendor should win" + + await harness.advance("2h") + + assert "slow-answer" not in RACE_CALLS, ( + "the losing branch answered after the race was already decided: " + f"{RACE_CALLS}" + ) diff --git a/tests/units/workflow/test_waits.py b/tests/units/workflow/test_waits.py index 931ac47e9ff..28fe5f45e96 100644 --- a/tests/units/workflow/test_waits.py +++ b/tests/units/workflow/test_waits.py @@ -377,3 +377,111 @@ def go(self): assert ( await harness.signal("no-such-run", Quick.pinged({"v": 1})) == "unknown_run" ) + + +@pytest.mark.xfail( + strict=True, + reason=( + "Known defect: a losing alternative of a decided approval is buffered " + "and then consumed by the next wait on the same channel. Fixing it " + "needs a durable decision-group identity shared by the alternatives " + "minted together, which is a store-shape change." + ), +) +async def test_a_rejected_alternative_does_not_answer_the_next_wait( + forked_registration_context, +): + """One decision must not silently answer a later, unrelated question. + + An approval mints two links -- approve and reject -- with distinct + delivery keys, on purpose, so a second person can reject after a first + approved. If approve lands first the run continues; the reject that lands + afterwards is buffered because nothing is waiting for it right then. When + the continuation opens a *second* wait on the same channel, that stale + reject resolves it, and a two-stage approval completes with a decision + nobody made in the second stage. + """ + + class TwoStage(rx.State): + __workflow__ = WorkflowConfig(id="waits.two_stage") + first: str = "" + second: str = "" + + review = Signal(Decision) + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Ask the first question. + + Returns: + The first wait. + """ + return wait_for( + TwoStage.review, + then=TwoStage.stage_one, + timeout="7d", + on_timeout=TwoStage.expire, + ) + + @rx.event(durable=True, effect="none") + def stage_one(self, decision: Decision): + """Record the first answer and ask a second question. + + Args: + decision: The delivered decision. + + Returns: + The second wait. + """ + self.first = decision.by + return wait_for( + TwoStage.review, + then=TwoStage.stage_two, + timeout="7d", + on_timeout=TwoStage.expire, + ) + + @rx.event(durable=True, effect="none") + def stage_two(self, decision: Decision): + """Record the second answer. + + Args: + decision: The delivered decision. + + Returns: + Completion. + """ + self.second = decision.by + return rx.complete(result={"first": self.first, "second": self.second}) + + @rx.event(durable=True, effect="none") + def expire(self): + """Nobody answered. + + Returns: + Failure. + """ + return rx.fail(reason="no decision") + + async with WorkflowTestHarness(TwoStage) as harness: + started = await harness.start(TwoStage.start()) + assert started.run_id is not None + await harness.signal( + started.run_id, + TwoStage.review(Decision(approved=True, by="approver")), + key="approve-link", + ) + # The rejecting link for the *same* question, spent a moment later. + await harness.signal( + started.run_id, + TwoStage.review(Decision(approved=False, by="rejecter")), + key="reject-link", + ) + await harness.run_until_idle() + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is not RunStatus.COMPLETED, ( + "the second stage was answered by the first stage's losing " + f"alternative: {snapshot.result}" + ) From 1359a84d97791c45e8dd119eefce71aa3cdfd417 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 11:26:06 -0700 Subject: [PATCH 080/121] workflows: make one approval's choices mutually exclusive The second of the two P1s from the external report, and it did not need the store change I thought it did. Each link took its delivery identity from its payload, so approve and reject were two identities. Spending approve decided the question; the reject link was then a live delivery with nowhere to go, so it was buffered -- and the next wait on that channel consumed it. A two-stage approval completed with a second-stage answer nobody gave. The links of one decision are minted by one handler attempt, so keying them on the slot rather than the payload makes them one identity, and the inbox's existing dedupe refuses the loser. That is the whole fix: no new state, no store protocol change, and it holds on every store because it is upstream of all of them. The ordinal is stable across retries of that slot, so a link mailed out before a retry still works after it. A handler that genuinely mints independent decisions on one channel passes distinct key= values, which is now what distinct keys mean. The waits xfail is narrowed accordingly: a raw signal sent with an explicitly distinct key still buffers, and that is how an early signal is delivered too, so the test pins the behaviour rather than calling it a bug. --- news/workflow-approval-groups.bugfix.md | 1 + reflex/workflow/approvals.py | 19 +++- tests/units/workflow/test_approvals.py | 129 ++++++++++++++++++++++++ tests/units/workflow/test_waits.py | 11 +- 4 files changed, 152 insertions(+), 8 deletions(-) create mode 100644 news/workflow-approval-groups.bugfix.md diff --git a/news/workflow-approval-groups.bugfix.md b/news/workflow-approval-groups.bugfix.md new file mode 100644 index 00000000000..a3526d8a615 --- /dev/null +++ b/news/workflow-approval-groups.bugfix.md @@ -0,0 +1 @@ +The links `rx.approval_link()` mints for one decision are now mutually exclusive: whichever choice is spent first decides, and the others are refused. Previously each link took its delivery identity from its payload, so approve and reject were separate deliveries — spending approve decided the question, and the reject link then sat buffered until the workflow opened its next wait on that channel and swallowed it, completing a second-stage approval with an answer nobody gave. Links are now grouped by the handler slot that minted them, which is also stable across retries of that slot. A handler that really does mint independent decisions on one channel can still pass distinct `key=` values. diff --git a/reflex/workflow/approvals.py b/reflex/workflow/approvals.py index 2c4278ec613..def202002e7 100644 --- a/reflex/workflow/approvals.py +++ b/reflex/workflow/approvals.py @@ -130,9 +130,13 @@ def approval_link( empty for a path, which is what a relative link needs. expires_in: How long the link stays valid. key: Delivery identity. Two links sharing a key are the same decision, - so the second one spent is a no-op; distinct keys let one person - approve after another rejected. Defaults to a key derived from the - channel and payload, which makes a link single-use. + so the second one spent is a no-op. The default groups every link + minted by one handler attempt for one channel, which is what makes + approve and reject mutually exclusive: whichever is spent first + decides, and the other is refused rather than sitting buffered + waiting to answer some later question. Pass distinct keys + explicitly when a handler really does mint independent decisions + on the same channel. Returns: The URL. @@ -143,8 +147,15 @@ def approval_link( # across a network boundary, and only plain data survives the round trip. payload = to_run_data({"value": delivery.payload})["value"] if key is None: + # The slot, not the payload: alternatives of one decision are minted + # by one handler attempt, so keying on the slot makes them one + # delivery identity and the inbox refuses the second. Keying on the + # payload made them distinct, which let a losing alternative outlive + # the decision it belonged to and resolve a later wait on the same + # channel. The ordinal is stable across retries of that slot, so a + # link handed out before a retry still works afterwards. material = json.dumps( - [context.run_id, delivery.channel, payload], sort_keys=True + [context.run_id, delivery.channel, context.ordinal], sort_keys=True ) key = hashlib.sha256(material.encode()).hexdigest()[:32] claims = { diff --git a/tests/units/workflow/test_approvals.py b/tests/units/workflow/test_approvals.py index dbd937326c6..921dc7314de 100644 --- a/tests/units/workflow/test_approvals.py +++ b/tests/units/workflow/test_approvals.py @@ -476,3 +476,132 @@ def test_a_token_missing_a_claim_is_refused(monkeypatch): body = json.dumps(claims, separators=(",", ":"), sort_keys=True).encode() with pytest.raises(WorkflowRuntimeError, match="not valid"): decode_token(f"{_b64(body)}.{_sign(body)}") + + +TWO_STAGE_LINKS: dict[str, str] = {} + + +class TwoStage(rx.State): + """A decision that is asked twice on one channel.""" + + __workflow__ = WorkflowConfig(id="approval.two_stage") + + decided = rx.Signal(dict) + first: str = "" + second: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def submit(self): + """Ask the first question with both choices. + + Returns: + The first wait. + """ + TWO_STAGE_LINKS["one_approve"] = rx.approval_link( + TwoStage.decided({"ok": True}) + ) + TWO_STAGE_LINKS["one_reject"] = rx.approval_link( + TwoStage.decided({"ok": False}) + ) + return rx.wait_for( + TwoStage.decided, + then=TwoStage.stage_one, + timeout="7d", + on_timeout=TwoStage.lapse, + ) + + @rx.event(durable=True, effect="none") + def stage_one(self, decision: dict): + """Record the first answer and ask a second question. + + Args: + decision: The delivered payload. + + Returns: + The second wait. + """ + self.first = "approved" if decision["ok"] else "rejected" + TWO_STAGE_LINKS["two_approve"] = rx.approval_link( + TwoStage.decided({"ok": True}) + ) + return rx.wait_for( + TwoStage.decided, + then=TwoStage.stage_two, + timeout="7d", + on_timeout=TwoStage.lapse, + ) + + @rx.event(durable=True, effect="none") + def stage_two(self, decision: dict): + """Record the second answer. + + Args: + decision: The delivered payload. + + Returns: + Completion. + """ + self.second = "approved" if decision["ok"] else "rejected" + return rx.complete(result={"first": self.first, "second": self.second}) + + @rx.event(durable=True, effect="none") + def lapse(self): + """Nobody answered in time. + + Returns: + Failure. + """ + return rx.fail(reason="lapsed") + + +async def test_a_losing_alternative_cannot_answer_the_next_question( + monkeypatch, forked_registration_context +): + """A decision's discarded choice must not decide a later question. + + Both links belong to one question. Spending approve decides it; the + reject link is then a spent decision, not a delivery in search of a wait. + Keying links by the payload made them distinct identities, so the reject + sat buffered and the second wait on the same channel swallowed it -- a + two-stage approval completing with a second answer nobody gave. + + Args: + monkeypatch: Used to set the signing secret. + forked_registration_context: Isolates state registration. + """ + monkeypatch.setenv(SECRET_ENV, SECRET) + TWO_STAGE_LINKS.clear() + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(TwoStage) + app = Starlette( + routes=[ + Route(APPROVAL_ROUTE, approval_endpoint(runtime), methods=["GET", "POST"]) + ] + ) + await runtime.startup(start_worker=False) + started = await runtime.kernel.start(TwoStage.submit()) + await runtime.kernel.run_until_idle() + assert started.run_id is not None + + with TestClient(app) as client: + assert client.post(TWO_STAGE_LINKS["one_approve"]).status_code == 200 + await runtime.kernel.run_until_idle() + # The second question is now open. The first question's discarded + # choice must not be able to answer it. + stale = client.post(TWO_STAGE_LINKS["one_reject"]) + assert "already been used" in stale.text.lower(), stale.text + await runtime.kernel.run_until_idle() + + snapshot = await runtime.kernel.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is not RunStatus.COMPLETED, ( + f"the second stage was answered by the first stage's reject: {snapshot.result}" + ) + + # The second question's own link still decides it. + with TestClient(app) as client: + assert client.post(TWO_STAGE_LINKS["two_approve"]).status_code == 200 + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(started.run_id) + assert snapshot is not None + assert snapshot.result == {"first": "approved", "second": "approved"} diff --git a/tests/units/workflow/test_waits.py b/tests/units/workflow/test_waits.py index 28fe5f45e96..6938bbfeb38 100644 --- a/tests/units/workflow/test_waits.py +++ b/tests/units/workflow/test_waits.py @@ -382,10 +382,13 @@ def go(self): @pytest.mark.xfail( strict=True, reason=( - "Known defect: a losing alternative of a decided approval is buffered " - "and then consumed by the next wait on the same channel. Fixing it " - "needs a durable decision-group identity shared by the alternatives " - "minted together, which is a store-shape change." + "Known gap, narrower than it was: rx.approval_link() now groups the " + "alternatives it mints so a losing one cannot answer a later wait " + "(see test_approvals). A raw signal sent with an explicitly distinct " + "key still buffers and can resolve the next wait on that channel. " + "Whether that is a defect is a question about what an explicit key " + "means -- it is also how a genuinely early signal is delivered -- so " + "this pins the behaviour rather than asserting it is wrong." ), ) async def test_a_rejected_alternative_does_not_answer_the_next_wait( From c23430916d825a2ba2d7df3c03096bd4f820046c Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 11:32:23 -0700 Subject: [PATCH 081/121] workflows: say what losing a race does not protect you from The external report's last open finding is that a race loser can perform its effect after the race is decided. Reading section 5 back, that is not a defect against the contract -- the contract already says loser cancellation is best-effort follow-up outside the winning transaction, and that a loser runs to completion if the process dies first. What it does not say is the part that matters to somebody writing a branch. It calls the outcome harmless, and the *arrival* is harmless: it is refused as late and the parent's result is unaffected. The charge is not harmless. A loser that runs on has already called the provider. Section 5 now lists all three ways a loser survives -- the resolving worker dying, a loser already executing elsewhere, and the effects that follow from either -- and says what to do about it: a racing branch that is not safe to run twice needs rx.step and a provider idempotency key, the same as a retried step. Changing the semantics instead -- fencing commits against cancellation intent, so a cancelled run cannot record a completion -- is a real option and a bigger one: it is a commit() change in all three stores and it changes what cancellation means for every run, not just race losers. That is a product decision, not a bug fix. --- news/workflow-race-loser-contract.docs.md | 1 + reflex/workflow/CONTRACT.md | 17 ++++++++++++++--- 2 files changed, 15 insertions(+), 3 deletions(-) create mode 100644 news/workflow-race-loser-contract.docs.md diff --git a/news/workflow-race-loser-contract.docs.md b/news/workflow-race-loser-contract.docs.md new file mode 100644 index 00000000000..de55f81b1a0 --- /dev/null +++ b/news/workflow-race-loser-contract.docs.md @@ -0,0 +1 @@ +The execution contract now states plainly what losing a race does and does not guarantee: the losing branches of `rx.parallel(..., mode="first")` are asked to stop on a best-effort basis, so a loser can run to completion — when the worker that resolved the race dies first, or when the loser is already executing on another worker — and a loser that runs on performs its side effects. Its late arrival is refused, so the parent's result is unaffected, but the charge or the email already happened. A racing branch that is not safe to run twice needs `rx.step` and a provider idempotency key, the same as any retried step. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index cb37ee8fb98..47ba3311270 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -147,9 +147,20 @@ Consequences, stated plainly: forever only on a child that is still genuinely running. - `rx.parallel(..., mode="first")`: the join resolves on the first arrival; the engine then requests cancellation of the losing branches. That request - is best-effort follow-up, not part of the winning transaction: if the - process dies first, losers run to completion and their arrivals are - refused as late (`counted`/`duplicate`/terminal), which is harmless. + is best-effort follow-up, not part of the winning transaction, and it is + sent by the one worker that saw the winner arrive. Three consequences, + stated plainly because only the first is obvious: + - If that worker dies before sending it, the losers run to completion. + - A loser already executing on another worker receives the intent but is + not fenced at commit, so it finishes its attempt. + - **A loser that runs on therefore performs its side effects.** Its arrival + is refused as late (`counted`/`duplicate`/terminal), so the parent's + result is unaffected — that part is harmless — but the charge, the email, + the provisioning call already happened. A branch that moves money or is + otherwise not safe to run twice must use `rx.step` with a provider + idempotency key, exactly as a retried step must; losing a race is not a + guarantee of not having acted. Racing branches whose effects cannot be + made idempotent is the wrong shape for `mode="first"`. - Child runs are ordinary runs; cancelling the parent does not implicitly cancel children (fan-out is delegation, not ownership). A cancelled parent's join tombstones; late child arrivals are refused. From 59c7b8919a98ab18fc3315e2a42721fd7fd1989d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 12:07:30 -0700 Subject: [PATCH 082/121] workflows: close the two gaps the second review found A link whose question lapsed unanswered stayed live. Nothing had been spent, so the spent-link record that catches a losing alternative had nothing to catch: the wait simply timed out, the run asked a fresh question on the same channel, and the forgotten link answered that one. A link now carries the step that asked and a wait records the step that armed it, so a link presented against a different question is refused with a 409 rather than delivered. Links minted before this claim existed carry no step and are let through -- an upgrade should not invalidate approvals already in someone's inbox. Join ordering derived declaration order from a field on each arrival, which meant a join spanning the upgrade that added the field had nothing to sort by and silently fell back to finishing order. The index is recovered from the branch's admission key instead, which every release has written because it is what makes re-running a fan-out idempotent. An arrival that cannot be identified keeps its position; ordering is recovered where it is knowable and never guessed at. --- news/workflow-join-order-upgrade.bugfix.md | 1 + news/workflow-stale-approval-links.bugfix.md | 1 + reflex/workflow/approvals.py | 56 +++++++++ reflex/workflow/kernel.py | 43 ++++++- tests/units/workflow/test_approvals.py | 125 ++++++++++++++++++- tests/units/workflow/test_parallel.py | 64 ++++++++++ 6 files changed, 284 insertions(+), 6 deletions(-) create mode 100644 news/workflow-join-order-upgrade.bugfix.md create mode 100644 news/workflow-stale-approval-links.bugfix.md diff --git a/news/workflow-join-order-upgrade.bugfix.md b/news/workflow-join-order-upgrade.bugfix.md new file mode 100644 index 00000000000..77459994094 --- /dev/null +++ b/news/workflow-join-order-upgrade.bugfix.md @@ -0,0 +1 @@ +A `rx.parallel` join that spans the upgrade which taught arrivals to carry their branch index now still returns results in declaration order. Arrivals recorded by the older release carry no index, so it is recovered from the branch's admission key, which every release has written. An arrival that cannot be identified at all keeps its position rather than being reordered on a guess. diff --git a/news/workflow-stale-approval-links.bugfix.md b/news/workflow-stale-approval-links.bugfix.md new file mode 100644 index 00000000000..ef8ccb87bac --- /dev/null +++ b/news/workflow-stale-approval-links.bugfix.md @@ -0,0 +1 @@ +An approval link now answers only the question it was minted for. A link whose question lapsed unanswered used to stay live, so if the workflow went on to ask something else on the same channel the forgotten link answered that instead — with nothing spent, there was no used-link record to catch it. Links carry the step that asked, waits record the step that armed them, and a link presented against a different question is refused with 409 "no longer open". Links minted before this are unaffected rather than invalidated by the upgrade. diff --git a/reflex/workflow/approvals.py b/reflex/workflow/approvals.py index def202002e7..1ec687f7496 100644 --- a/reflex/workflow/approvals.py +++ b/reflex/workflow/approvals.py @@ -35,6 +35,7 @@ from starlette.responses import HTMLResponse, JSONResponse, Response from reflex.workflow.context import require_run +from reflex.workflow.records import StepStatus from reflex.workflow.serde import to_run_data if TYPE_CHECKING: @@ -163,6 +164,10 @@ def approval_link( "c": delivery.channel, "p": payload, "k": key, + # The step that asked. A link answers the question it was minted for + # and no other: without this a link for a question that timed out + # unanswered stays live and resolves the next wait on the channel. + "o": context.ordinal, "e": time.time() + parse_duration(expires_in, param="expires_in"), } body = json.dumps(claims, separators=(",", ":"), sort_keys=True).encode() @@ -266,6 +271,45 @@ def _wants_json(request: Request) -> bool: return "application/json" in request.headers.get("accept", "") +async def _answers_the_open_question( + runtime: WorkflowRuntime, claims: dict[str, Any] +) -> bool: + """Whether a link's question is the one the run is still waiting on. + + A link is minted while one step runs, and that step's question is the only + one it may answer. Without this a link outlives its question: nobody + clicks it, the wait times out, the run asks something else on the same + channel, and the forgotten link answers that instead. + + Links minted before this claim existed carry no step, and are let through + rather than being invalidated by an upgrade. + + Args: + runtime: The runtime holding the run. + claims: The verified token claims. + + Returns: + True when the link may be delivered. + """ + asked_by = claims.get("o") + if not isinstance(asked_by, int): + return True + snapshot = await runtime.kernel.get_run(claims["r"]) + if snapshot is None: + return True + waiting = next( + (step for step in snapshot.steps if step.status is StepStatus.BLOCKED), + None, + ) + if waiting is None: + # Nothing is waiting; let the store give its own answer, which is a + # better message than this one. + return True + wait = waiting.args.get("__wait__") + armed_by = wait.get("armed_by") if isinstance(wait, dict) else None + return not isinstance(armed_by, int) or armed_by == asked_by + + def approval_endpoint( runtime: WorkflowRuntime, ) -> Callable[[Request], Coroutine[Any, Any, Response]]: @@ -316,6 +360,18 @@ async def endpoint(request: Request) -> Response: form=_FORM.format(label="Confirm"), ) + if not await _answers_the_open_question(runtime, claims): + # The question this link belongs to is over -- answered, or timed + # out unanswered. Delivering anyway would let it resolve whatever + # the run happens to be waiting on now, which is a different + # question that nobody asked this person. + stale = "This decision is no longer open." + if _wants_json(request): + return JSONResponse( + {"error": stale, "status": "stale"}, status_code=409 + ) + return _page("No longer open", stale, status=409) + disposition = await runtime.kernel.signal( claims["r"], ChannelDelivery(channel=claims["c"], payload=claims["p"]), diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 455faa967a5..3eb2f085a9d 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -1195,6 +1195,37 @@ def _interpret_return( return successors, None return [self._resolve_successor(defn, value)], None + async def _in_declaration_order(self, results: list) -> list: + """Sort a join's arrivals into the order their branches were declared. + + An arrival carries its branch index. One recorded before that field + existed does not, and a join can span the upgrade that introduced it, + so the index is recovered from the branch's admission key -- which + every release has written, because it is what makes re-running a + fan-out idempotent. Anything still unidentifiable keeps its arrival + position rather than being reordered on a guess. + + Args: + results: The arrivals recorded on the join slot. + + Returns: + The arrivals in declaration order where that is knowable. + """ + ordered: list[tuple[int, Any]] = [] + for position, entry in enumerate(results): + branch = entry.get("branch") if isinstance(entry, dict) else None + if not isinstance(branch, int) and isinstance(entry, dict): + run_id = entry.get("run_id") + child = ( + await self._store.get_run(run_id) + if isinstance(run_id, str) + else None + ) + if child is not None: + branch = _branch_index(child.request_key) + ordered.append((branch if isinstance(branch, int) else position, entry)) + return [entry for _, entry in sorted(ordered, key=operator.itemgetter(0))] + async def _invoke( self, handler: HandlerDefinition, @@ -1218,15 +1249,12 @@ async def _invoke( args = {key: value for key, value in args.items() if key != "__wait__"} delivered = args.pop("__payload__", None) results = args.pop("__results__", None) - if isinstance(results, list) and all( - isinstance(entry, dict) and isinstance(entry.get("branch"), int) - for entry in results - ): + if isinstance(results, list): # A join accumulates arrivals as they land, which is finishing # order; rx.parallel promises declaration order, and a caller # unpacking `a, b = results` has no other way to tell which is # which. - results = sorted(results, key=operator.itemgetter("branch")) + results = await self._in_declaration_order(results) if delivered is None and results is not None: delivered = results if delivered is not None and handler.params: @@ -1669,6 +1697,11 @@ def _success_completion( "__wait__": { "channel": control.channel, "on_timeout": timeout_id, + # Which step asked the question. An approval link is + # minted while that step runs, so this is what lets a + # link be checked against the question it belongs to + # rather than against whatever is waiting now. + "armed_by": claim.step.ordinal, }, }, due_at=deadline, diff --git a/tests/units/workflow/test_approvals.py b/tests/units/workflow/test_approvals.py index 921dc7314de..60fa8cb0872 100644 --- a/tests/units/workflow/test_approvals.py +++ b/tests/units/workflow/test_approvals.py @@ -479,6 +479,77 @@ def test_a_token_missing_a_claim_is_refused(monkeypatch): TWO_STAGE_LINKS: dict[str, str] = {} +_NOW = [1_000_000.0] + + +class Lapsing(rx.State): + """A question that lapses, then asks a different one on the same channel.""" + + __workflow__ = WorkflowConfig(id="approval.lapsing") + + decided = rx.Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def submit(self): + """Ask, with links nobody will click. + + Returns: + The first wait. + """ + TWO_STAGE_LINKS["one_approve"] = rx.approval_link(Lapsing.decided({"ok": True})) + return rx.wait_for( + Lapsing.decided, + then=Lapsing.first, + timeout="7d", + on_timeout=Lapsing.ask_again, + ) + + @rx.event(durable=True, effect="none") + def ask_again(self): + """Nobody answered; ask a fresh question on the same channel. + + Returns: + The second wait. + """ + return rx.wait_for( + Lapsing.decided, + then=Lapsing.second, + timeout="7d", + on_timeout=Lapsing.give_up, + ) + + @rx.event(durable=True, effect="none") + def first(self, decision: dict): + """Record a first-question answer. + + Args: + decision: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result={"answered": "first", "decision": decision}) + + @rx.event(durable=True, effect="none") + def second(self, decision: dict): + """Record a second-question answer. + + Args: + decision: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result={"answered": "second", "decision": decision}) + + @rx.event(durable=True, effect="none") + def give_up(self): + """Nobody answered either question. + + Returns: + Failure. + """ + return rx.fail(reason="lapsed twice") class TwoStage(rx.State): @@ -589,7 +660,8 @@ async def test_a_losing_alternative_cannot_answer_the_next_question( # The second question is now open. The first question's discarded # choice must not be able to answer it. stale = client.post(TWO_STAGE_LINKS["one_reject"]) - assert "already been used" in stale.text.lower(), stale.text + assert stale.status_code == 409, stale.text + assert "no longer open" in stale.text.lower(), stale.text await runtime.kernel.run_until_idle() snapshot = await runtime.kernel.get_run(started.run_id) @@ -605,3 +677,54 @@ async def test_a_losing_alternative_cannot_answer_the_next_question( snapshot = await runtime.kernel.get_run(started.run_id) assert snapshot is not None assert snapshot.result == {"first": "approved", "second": "approved"} + + +async def test_a_link_from_a_timed_out_question_cannot_answer_the_next_one( + monkeypatch, forked_registration_context +): + """A question nobody answered is over, and its links go with it. + + Nothing was spent, so there is no spent-link record to catch this: the + first question simply lapsed. If the run then asks something else on the + same channel, the forgotten link from the first question must not answer + it. The link names the step that asked; only that step's question is its + to answer. + + Args: + monkeypatch: Used to set the signing secret. + forked_registration_context: Isolates state registration. + """ + monkeypatch.setenv(SECRET_ENV, SECRET) + TWO_STAGE_LINKS.clear() + store = MemoryRunStore() + runtime = WorkflowRuntime(store, clock=lambda: _NOW[0]) + runtime.register(Lapsing) + app = Starlette( + routes=[ + Route(APPROVAL_ROUTE, approval_endpoint(runtime), methods=["GET", "POST"]) + ] + ) + await runtime.startup(start_worker=False) + started = await runtime.kernel.start(Lapsing.submit()) + await runtime.kernel.run_until_idle() + assert started.run_id is not None + stale_link = TWO_STAGE_LINKS["one_approve"] + + # Nobody answers. The deadline arrives and the run asks a second question + # on the same channel. + _NOW[0] += 8 * 24 * 3600 + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING, "the second question is open" + + with TestClient(app) as client: + late = client.post(stale_link) + assert late.status_code == 409, late.text + + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING, ( + f"a link from the lapsed question answered the second one: {snapshot.result}" + ) diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index e5de6e93462..2ba6294de3e 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -3,7 +3,9 @@ from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual import reflex as rx +from reflex.workflow.kernel import WorkflowKernel from reflex.workflow.records import HistoryEventType, RunStatus, StepStatus +from reflex.workflow.store import MemoryRunStore from reflex.workflow.testing import WorkflowTestHarness BRANCH_CALLS: list[str] = [] @@ -611,3 +613,65 @@ async def test_a_race_loser_performs_no_effect_after_it_has_lost( "the losing branch answered after the race was already decided: " f"{RACE_CALLS}" ) + + +async def test_arrivals_from_before_the_upgrade_are_still_ordered( + forked_registration_context, +): + """A join can span the deploy that taught arrivals to carry their branch. + + A worker on the older release records an arrival with no branch index, and + the join it belongs to may already be half full when the new release comes + up. The index is still recoverable: a fan-out has always stamped each + branch with `child:::` as its admission key, + because that is what makes re-running the fan-out idempotent. + """ + store = MemoryRunStore() + kernel = WorkflowKernel([], store) + + class _Legacy: + """Stands in for a child run admitted by the older release.""" + + def __init__(self, run_id: str, index: int): + self.run_id = run_id + self.request_key = f"child:parent:0:{index}" + + admitted = { + "late": _Legacy("late", 0), + "early": _Legacy("early", 1), + } + + async def fake_get_run(run_id: str): # noqa: RUF029 + """Look up a child the way the store would. + + Args: + run_id: The child run to load. + + Returns: + The stand-in record, or None. + """ + return admitted.get(run_id) + + kernel._store = type( # pyright: ignore[reportAttributeAccessIssue] + "_Stub", (), {"get_run": staticmethod(fake_get_run)} + )() + + ordered = await kernel._in_declaration_order([ # pyright: ignore[reportPrivateUsage] + {"run_id": "early", "result": "second-declared"}, + {"run_id": "late", "result": "first-declared"}, + ]) + assert [entry["result"] for entry in ordered] == [ + "first-declared", + "second-declared", + ] + + +async def test_unidentifiable_arrivals_keep_their_position( + forked_registration_context, +): + """Ordering is recovered where it is knowable, never guessed at.""" + store = MemoryRunStore() + kernel = WorkflowKernel([], store) + entries = [{"result": "a"}, {"result": "b"}, {"result": "c"}] + ordered = await kernel._in_declaration_order(entries) # pyright: ignore[reportPrivateUsage] + assert [entry["result"] for entry in ordered] == ["a", "b", "c"] From 94152c9523bfd8a5515ee42e78bc77a5b9fe017b Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 12:21:42 -0700 Subject: [PATCH 083/121] workflows: stop a cancelled parent from cancelling its branches Second review, first blocker. Section 5 promises that cancelling a parent does not implicitly cancel children -- fan-out is delegation, not ownership -- and the engine broke that promise by accident. Cancelling a parent tombstones its join slot. When a branch finished afterwards, its arrival found a slot that was no longer BLOCKED, and the only reason the engine knew for that was a first-mode race having been decided. So it cancelled the branch's siblings: a run nobody asked it to touch, killed because a different run was cancelled. Loser cancellation now runs only where losing is a thing that can happen: a join with join_expected == 1, still standing. An all-mode join wants every branch and has no losers, and a tombstoned join has decided nothing. The regression test needed both branches still in flight to reproduce -- with one already finished there is no live sibling left to cancel, and the bug hides. --- .../workflow-parent-cancel-children.bugfix.md | 1 + reflex/workflow/kernel.py | 18 ++++ tests/units/workflow/test_parallel.py | 99 +++++++++++++++++++ 3 files changed, 118 insertions(+) create mode 100644 news/workflow-parent-cancel-children.bugfix.md diff --git a/news/workflow-parent-cancel-children.bugfix.md b/news/workflow-parent-cancel-children.bugfix.md new file mode 100644 index 00000000000..54bb671a013 --- /dev/null +++ b/news/workflow-parent-cancel-children.bugfix.md @@ -0,0 +1 @@ +Cancelling the parent of an all-mode `rx.parallel` fan-out no longer cancels its branches. Cancelling a parent tombstones its join, so when a branch later finished, its arrival found a slot that was no longer blocked — which the engine read as a race having been decided, and cancelled the sibling. Loser cancellation now runs only for a join that is actually a first-mode race and was not tombstoned, which is what the contract already promised: delegation is not ownership. diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 3eb2f085a9d..dc3a3d4b57f 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -2546,6 +2546,24 @@ async def _cancel_losing_branches(self, winner: RunRecord) -> None: """ if winner.parent_run_id is None or winner.parent_ordinal is None: return + join = next( + ( + step + for step in await self._store.get_steps(winner.parent_run_id) + if step.ordinal == winner.parent_ordinal + ), + None, + ) + if join is None or join.join_expected != 1: + # Only a race has losers. An all-mode join wants every branch, and + # a join that stopped being blocked for some other reason -- a + # cancelled or force-finalized parent tombstones it -- has not + # decided anything. Cancelling siblings on either would be this + # engine reaching into runs it does not own, which §5 says + # delegation does not do. + return + if join.status is StepStatus.CANCELLED: + return siblings = await self._store.list_children( winner.parent_run_id, winner.parent_ordinal ) diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 2ba6294de3e..37c6bb05283 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -11,6 +11,62 @@ BRANCH_CALLS: list[str] = [] +class Slower(rx.State): + """A branch that answers well after its sibling.""" + + __workflow__ = WorkflowConfig(id="fan.slower") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, lead: str): + """Answer much later. + + Args: + lead: The lead identifier. + + Returns: + A deferral. + """ + return rx.after("5h", Slower.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + """Answer. + + Returns: + Completion. + """ + BRANCH_CALLS.append("slower") + return rx.complete(result="slower") + + +class Slowish(rx.State): + """A branch that answers after a delay.""" + + __workflow__ = WorkflowConfig(id="fan.slowish") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, lead: str): + """Answer later. + + Args: + lead: The lead identifier. + + Returns: + A deferral. + """ + return rx.after("1h", Slowish.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + """Answer. + + Returns: + Completion. + """ + BRANCH_CALLS.append("slowish") + return rx.complete(result="slowish") + + class Enrich(rx.State): """A branch that succeeds.""" @@ -675,3 +731,46 @@ async def test_unidentifiable_arrivals_keep_their_position( entries = [{"result": "a"}, {"result": "b"}, {"result": "c"}] ordered = await kernel._in_declaration_order(entries) # pyright: ignore[reportPrivateUsage] assert [entry["result"] for entry in ordered] == ["a", "b", "c"] + + +async def test_cancelling_an_all_mode_parent_leaves_its_branches_alone( + forked_registration_context, +): + """Delegation is not ownership, so a cancelled parent cancels no children. + + Section 5 says cancelling a parent does not implicitly cancel children. + Cancelling one tombstones its join, and when a branch later finishes its + arrival finds a slot that is no longer blocked -- which is not the same + thing as a race having been decided. Reading it as one made the engine + cancel the sibling of a branch it never raced. + """ + BRANCH_CALLS.clear() + router = _router(Slowish, Slower) + async with WorkflowTestHarness(router, Slowish, Slower) as harness: + started = await harness.start(router.begin("acme")) + assert started.run_id is not None + await harness.run_until_idle() + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + join_ordinal = next( + step.ordinal for step in snapshot.steps if step.origin == "join" + ) + store = harness.kernel._store # pyright: ignore[reportPrivateUsage] + assert await store.list_children(started.run_id, join_ordinal), ( + "the fan-out should have admitted branches" + ) + + assert await harness.cancel(started.run_id) + await harness.run_until_idle() + + # The delayed branch finishes after the parent is gone. + await harness.advance("2h") + children = await store.list_children(started.run_id, join_ordinal) + cancelled = [ + child.run_id for child in children if child.status is RunStatus.CANCELLED + ] + assert not cancelled, ( + "cancelling the parent cancelled delegated children: " + f"{[(c.run_id, c.status) for c in children]}" + ) From a82d099a8573a4d1db742fa8790c44cc4f77323d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 12:26:35 -0700 Subject: [PATCH 084/121] workflows: pin the retry-loses-the-chain defect as a spec I told the reviewer I did not understand this one, so here is what it actually is. Returning a list from a handler preallocates an immediate sequential chain -- kernel._interpret_return -- so [middle(), finish()] creates both slots up front. A terminal failure in middle tombstones finish behind it, and retry() reopens only the failed step. The retried step succeeds, returns nothing to allocate, and the run completes having never run the finalizer the chain existed for. My earlier reasoning was wrong because I assumed chains are built a step at a time, which is true of rx.after and not of a returned list. Captured as a strict xfail rather than fixed: the fix belongs in retry_run and skip_step in all three stores, restoring the successors that failure tombstoned without reviving anything an unrelated cancellation cancelled, and that is not a change to start on a budget that cannot finish it. The first version of this test xfailed for the wrong reason -- the test harness has no retry(), so it was failing on AttributeError, not on the defect. It goes through the kernel now. That the harness exposes cancel but not retry/skip/resume is its own gap: operator recovery is exactly what someone would want to test. --- tests/units/workflow/test_operator_actions.py | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/tests/units/workflow/test_operator_actions.py b/tests/units/workflow/test_operator_actions.py index 3337533c9aa..4ed57e51e05 100644 --- a/tests/units/workflow/test_operator_actions.py +++ b/tests/units/workflow/test_operator_actions.py @@ -6,6 +6,7 @@ trying to rescue. """ +import pytest from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual import reflex as rx @@ -308,3 +309,79 @@ async def test_skip_is_refused_on_a_healthy_run(forked_registration_context): assert not await rx.workflows.skip(result.run_id) assert not await rx.workflows.skip("no-such-run") + + +@pytest.mark.xfail( + strict=True, + reason=( + "Known defect: a returned list preallocates a sequential chain, so a " + "step failing tombstones the rest of it. retry() reopens the failed " + "step but does not restore the successors that failure cancelled, so " + "the run completes having skipped them. Fixing it means restoring " + "steps tombstoned by that failure in retry_run/skip_step across all " + "three stores." + ), +) +async def test_retry_restores_the_chain_the_failure_tombstoned( + forked_registration_context, +): + """Retrying continues from the failed step, not past everything after it. + + A handler returning a list preallocates the whole chain, so a terminal + failure cancels the steps behind it. An operator retrying expects the run + to carry on from there -- including the finalizer that was already + allocated and is now cancelled. + """ + attempts: list[str] = [] + + class Chain(rx.State): + __workflow__ = WorkflowConfig(id="ops.chain") + note: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Preallocate the rest of the chain. + + Returns: + Two successors, run in order. + """ + return [Chain.middle(), Chain.finish()] + + @rx.event(durable=True, effect="none", retry=Retry(max_attempts=1)) + def middle(self): + """Fail the first time an operator has not yet fixed anything. + + Raises: + ValueError: On the first attempt. + """ + attempts.append("middle") + if len(attempts) == 1: + msg = "the reason an operator would retry" + raise ValueError(msg) + + @rx.event(durable=True, effect="none") + def finish(self): + """The finalizer the chain exists for. + + Returns: + Completion. + """ + return rx.complete(result="finished") + + async with WorkflowTestHarness(Chain) as harness: + started = await harness.start(Chain.begin()) + assert started.run_id is not None + await harness.run_until_idle() + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + + assert await harness.kernel.retry(started.run_id) + await harness.run_until_idle() + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.result == "finished", ( + "the retried run completed without running the finalizer the " + f"chain preallocated: {[step.status.value for step in snapshot.steps]}" + ) From 890b0edb3653940ac62dcbc2f42ed82dcaa71056 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 12:51:42 -0700 Subject: [PATCH 085/121] workflows: scope webhook dedupe keys to the root they start A provider numbers its events per object, not per topic, so invoice_failed and invoice_paid for one invoice both arrive carrying that invoice's id. The ingress used that value as the whole admission key, and stores enforce uniqueness on (workflow_id, request_key), so the payment was taken for a redelivery of the failure and dropped. One workflow class serving an object's lifecycle -- the shape the API invites -- silently loses events. Keys are webhook:: now. Two deliveries are the same event only if they would start the same handler. Changing a key format is itself a hazard: every run admitted under the old spelling would stop being found, and the provider's next redelivery would start it again. kernel.start() takes superseded_keys, matched for deduplication and never written, and the ingress passes the old bare value. Nothing replays across the upgrade. Found by a live-fire review over real sockets and processes, which also confirmed the parts that hold: 401 on forged signatures, 64 concurrent redeliveries collapsing to one run across two servers, and recovery of a WAITING run across a full process restart. --- news/workflow-webhook-dedupe-scope.bugfix.md | 1 + reflex/workflow/ingress.py | 37 +++++- reflex/workflow/kernel.py | 15 ++- tests/units/workflow/test_ingress.py | 130 +++++++++++++++++++ 4 files changed, 177 insertions(+), 6 deletions(-) create mode 100644 news/workflow-webhook-dedupe-scope.bugfix.md diff --git a/news/workflow-webhook-dedupe-scope.bugfix.md b/news/workflow-webhook-dedupe-scope.bugfix.md new file mode 100644 index 00000000000..09cdf7c50b0 --- /dev/null +++ b/news/workflow-webhook-dedupe-scope.bugfix.md @@ -0,0 +1 @@ +Webhook deduplication keys are now namespaced by the root handler they start, so two lifecycle topics served by one workflow no longer collide. A provider numbers its events per object rather than per topic, so `invoice_failed` and `invoice_paid` for one invoice arrive carrying the same id — unqualified, the payment was deduplicated against the failure and silently dropped. Deliveries admitted under the previous unqualified key are still recognised, so upgrading does not make the provider's next redelivery of an already-handled event start it a second time. diff --git a/reflex/workflow/ingress.py b/reflex/workflow/ingress.py index 9697e47d1e0..d6a21b46ce6 100644 --- a/reflex/workflow/ingress.py +++ b/reflex/workflow/ingress.py @@ -97,10 +97,20 @@ def collect_webhook_routes( return routes -def _dedupe_key(trigger: WebhookTrigger, payload: Any) -> str | None: +def _dedupe_key( + handler: HandlerDefinition, trigger: WebhookTrigger, payload: Any +) -> str | None: """Extract the deduplication key a provider redelivery would repeat. + The key is namespaced by the handler it starts. A provider numbers its + events per object, not per topic, so ``invoice_failed`` and + ``invoice_paid`` for one invoice arrive carrying the same id: unqualified, + the second is deduplicated against the first and the payment is silently + dropped. Two deliveries are the same event only if they would start the + same handler. + Args: + handler: The root handler this delivery starts. trigger: The webhook trigger declaring the key field. payload: The decoded request payload. @@ -111,7 +121,27 @@ def _dedupe_key(trigger: WebhookTrigger, payload: Any) -> str | None: if trigger.dedupe_by is None or not isinstance(payload, dict): return None value = payload.get(trigger.dedupe_by) - return None if value is None else str(value) + return None if value is None else f"webhook:{handler.id}:{value}" + + +def _legacy_dedupe_keys(trigger: WebhookTrigger, payload: Any) -> tuple[str, ...]: + """Spellings this delivery's key had in earlier releases. + + Keys used to be the provider's raw value, unqualified by the handler. A + run admitted under one must still be found after the upgrade, or the + provider's next redelivery of an event already handled starts it again. + + Args: + trigger: The webhook trigger declaring the key field. + payload: The decoded request payload. + + Returns: + The older keys to match, newest spelling first. + """ + if trigger.dedupe_by is None or not isinstance(payload, dict): + return () + value = payload.get(trigger.dedupe_by) + return () if value is None else (str(value),) def _root_args(handler: HandlerDefinition, payload: Any) -> dict[str, Any]: @@ -190,7 +220,8 @@ async def endpoint(request: Request) -> JSONResponse: args = _root_args(route.handler, payload) result = await runtime.kernel.start( spec(**args) if args else spec, - request_key=_dedupe_key(route.trigger, payload), + request_key=_dedupe_key(route.handler, route.trigger, payload), + superseded_keys=_legacy_dedupe_keys(route.trigger, payload), trigger_kind="webhook", ) return JSONResponse( diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index dc3a3d4b57f..ddd5f66fb86 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -683,6 +683,7 @@ async def start( target: Any, *, request_key: str | None = None, + superseded_keys: tuple[str, ...] = (), labels: dict[str, str] | None = None, trigger_kind: str | None = "manual", ) -> StartResult: @@ -692,6 +693,10 @@ async def start( target: The root event, e.g. ``MyWorkflow.start(payload)``. request_key: Idempotent admission key; a repeated key returns the prior run with disposition ``"deduplicated"``. + superseded_keys: Older spellings of the same admission key, matched + for deduplication but never recorded. This is what lets the + key format change without every event admitted under the old + one being admitted a second time after the upgrade. labels: Server-derived indexing labels to record on the run. trigger_kind: Which ingress is starting this run; the root must declare the same kind, so a webhook-only root stays @@ -721,12 +726,16 @@ async def start( f"starting through this path requires {expected}." ) raise WorkflowRuntimeError(msg) - if request_key is not None: + for candidate in (request_key, *superseded_keys): + if candidate is None: + continue # Dedupe before any start policy: a redelivered event must return # the run it already created, not be judged as a new start and - # cancel, throttle, or debounce that very run. + # cancel, throttle, or debounce that very run. Superseded keys are + # matched too but never written: a key format that changes must + # not make every event admitted under the old one arrive twice. existing = await self._store.find_by_request_key( - defn.workflow_id, request_key + defn.workflow_id, candidate ) if existing is not None: return StartResult(disposition="deduplicated", run_id=existing) diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py index bef25d38f26..d1279db94d3 100644 --- a/tests/units/workflow/test_ingress.py +++ b/tests/units/workflow/test_ingress.py @@ -334,3 +334,133 @@ def on_event(self, first: str, second: int): assert accepted.status_code == 202, accepted.text finally: await runtime.shutdown() + + +class Invoices(rx.State): + """One workflow serving two lifecycle topics for the same object.""" + + __workflow__ = WorkflowConfig(id="ingress.invoices") + outcome: str = "" + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "invoice_failed", + dedupe_by="id", + verify=hmac_signature( + secret_env="STRIPE_WEBHOOK_SECRET", header="X-Signature" + ), + ), + ) + def on_failed(self, id: str): + """Record a failed invoice. + + Args: + id: The invoice identifier. + + Returns: + Completion. + """ + self.outcome = "failed" + return rx.complete(result={"invoice": id, "outcome": "failed"}) + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "invoice_paid", + dedupe_by="id", + verify=hmac_signature( + secret_env="STRIPE_WEBHOOK_SECRET", header="X-Signature" + ), + ), + ) + def on_paid(self, id: str): + """Record a paid invoice. + + Args: + id: The invoice identifier. + + Returns: + Completion. + """ + self.outcome = "paid" + return rx.complete(result={"invoice": id, "outcome": "paid"}) + + +async def test_two_topics_sharing_a_dedupe_value_are_separate_events( + monkeypatch, forked_registration_context +): + """A provider numbers events per object, not per topic. + + `invoice_failed` and `invoice_paid` for one invoice carry the same id. + Deduplicating on that alone makes the payment a redelivery of the failure + and drops it, which is the one outcome a billing workflow must never have. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Invoices) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + + started: list[str] = [] + dispositions: list[str] = [] + with TestClient(app) as client: + for topic in ("invoice_failed", "invoice_paid"): + body = json.dumps({"id": "inv_1"}).encode() + response = client.post( + f"/_workflow/webhook/{topic}", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert response.status_code == 202, response.text + started.append(response.json()["run_id"]) + dispositions.append(response.json()["disposition"]) + + assert dispositions == ["started", "started"], ( + "the payment was taken for a redelivery of the failure" + ) + assert started[0] != started[1], "both lifecycle events must have their own run" + await runtime.shutdown() + + +async def test_a_redelivery_admitted_under_the_old_key_still_deduplicates( + monkeypatch, forked_registration_context +): + """Changing the key format must not replay every event admitted before it. + + A run that exists under the unqualified key the older release wrote is + still that event's run. Matching the old spelling on the way in is what + keeps the provider's next redelivery from starting it a second time. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + store = MemoryRunStore() + runtime = WorkflowRuntime(store) + runtime.register(Invoices) + await runtime.startup(start_worker=False) + + # Exactly what the previous release recorded: the bare provider value. + legacy = await runtime.kernel.start( + Invoices.on_failed("inv_2"), + request_key="inv_2", + trigger_kind="webhook", + ) + assert legacy.disposition == "started" + + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + with TestClient(app) as client: + body = json.dumps({"id": "inv_2"}).encode() + response = client.post( + "/_workflow/webhook/invoice_failed", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert response.status_code == 202, response.text + assert response.json()["disposition"] == "deduplicated" + assert response.json()["run_id"] == legacy.run_id + await runtime.shutdown() From a3228a1cec3a18adf296e3e205ba5c384c4a9634 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 13:28:28 -0700 Subject: [PATCH 086/121] workflows: scope the legacy webhook key lookup to the same root My own regression, an hour old. Scoping fresh webhook keys by handler fixed the cross-topic collision; the compatibility lookup I added in the same commit matched the old unqualified key against any run in the workflow, which is precisely the collision again -- invoice_paid deduplicating against a legacy invoice_failed run and being dropped. A superseded key is a less specific spelling than the one that replaced it, so a match on it means the same event only if it started the same root. It is checked against the existing run's first step now. Caught by a live-fire review over real sockets, which reproduced it as 202 deduplicated pointing at the wrong run. --- news/workflow-legacy-key-scope.bugfix.md | 1 + reflex/workflow/kernel.py | 51 ++++++++++++++++++------ tests/units/workflow/test_ingress.py | 41 +++++++++++++++++++ 3 files changed, 80 insertions(+), 13 deletions(-) create mode 100644 news/workflow-legacy-key-scope.bugfix.md diff --git a/news/workflow-legacy-key-scope.bugfix.md b/news/workflow-legacy-key-scope.bugfix.md new file mode 100644 index 00000000000..bbb5c0a41cf --- /dev/null +++ b/news/workflow-legacy-key-scope.bugfix.md @@ -0,0 +1 @@ +The backward-compatibility lookup for pre-upgrade webhook deduplication keys now matches only a run started by the same root handler. Without that check the compatibility path reintroduced the cross-topic collision the new key format was added to fix: an `invoice_paid` delivery could deduplicate against a legacy `invoice_failed` run and be dropped. diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index ddd5f66fb86..026f8552244 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -678,6 +678,19 @@ async def _apply_start_policy( return None, now + window return None, now + async def _started_handler(self, run_id: str, handler_id: str) -> bool: + """Whether a run's root step is the handler a delivery would start. + + Args: + run_id: The run to check. + handler_id: The root handler the caller wants to start. + + Returns: + True when the run was started by that handler. + """ + steps = await self._store.get_steps(run_id) + return bool(steps) and steps[0].handler_id == handler_id + async def start( self, target: Any, @@ -726,19 +739,31 @@ async def start( f"starting through this path requires {expected}." ) raise WorkflowRuntimeError(msg) - for candidate in (request_key, *superseded_keys): - if candidate is None: - continue - # Dedupe before any start policy: a redelivered event must return - # the run it already created, not be judged as a new start and - # cancel, throttle, or debounce that very run. Superseded keys are - # matched too but never written: a key format that changes must - # not make every event admitted under the old one arrive twice. - existing = await self._store.find_by_request_key( - defn.workflow_id, candidate - ) - if existing is not None: - return StartResult(disposition="deduplicated", run_id=existing) + # Dedupe before any start policy: a redelivered event must return the + # run it already created, not be judged as a new start and cancel, + # throttle, or debounce that very run. + existing = ( + None + if request_key is None + else await self._store.find_by_request_key(defn.workflow_id, request_key) + ) + if existing is None: + for candidate in superseded_keys: + # A superseded key is matched but never written, so a key + # format can change without every event admitted under the old + # one arriving twice. The old spelling was less specific than + # the new one, though, so a match on it only means the same + # event if it started the same root -- otherwise the very + # collision the new format fixed comes back through the + # compatibility path. + found = await self._store.find_by_request_key( + defn.workflow_id, candidate + ) + if found is not None and await self._started_handler(found, handler.id): + existing = found + break + if existing is not None: + return StartResult(disposition="deduplicated", run_id=existing) flow_key = self._flow_key(handler, payload) if flow_key is None: return await self._admit(defn, handler, payload, request_key, labels, None) diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py index d1279db94d3..1a563b4183d 100644 --- a/tests/units/workflow/test_ingress.py +++ b/tests/units/workflow/test_ingress.py @@ -464,3 +464,44 @@ async def test_a_redelivery_admitted_under_the_old_key_still_deduplicates( assert response.json()["disposition"] == "deduplicated" assert response.json()["run_id"] == legacy.run_id await runtime.shutdown() + + +async def test_a_legacy_key_only_matches_the_root_that_wrote_it( + monkeypatch, forked_registration_context +): + """The compatibility path must not undo the fix it exists to soften. + + The old key format was the bare provider value, which is exactly what + collided across topics. Matching it without checking which root the + existing run started lets `invoice_paid` deduplicate against a legacy + `invoice_failed` run -- the same lost payment, arriving through the + upgrade path instead. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Invoices) + await runtime.startup(start_worker=False) + + legacy = await runtime.kernel.start( + Invoices.on_failed("inv_3"), + request_key="inv_3", + trigger_kind="webhook", + ) + assert legacy.disposition == "started" + + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + with TestClient(app) as client: + body = json.dumps({"id": "inv_3"}).encode() + paid = client.post( + "/_workflow/webhook/invoice_paid", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert paid.status_code == 202, paid.text + assert paid.json()["disposition"] == "started", ( + "the payment deduplicated against a legacy failure run" + ) + assert paid.json()["run_id"] != legacy.run_id + await runtime.shutdown() From 7b25195462a74a8c8e966ba151621efb6daa3eb4 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 13:58:24 -0700 Subject: [PATCH 087/121] workflows: make start policies atomic across processes The production-readiness review's P0: every advertised start policy violated its invariant in 50 of 50 synchronized trials across two OS processes and independent PostgreSQL pools. The policy decision was a read in one store call and a write in another, guarded by an asyncio lock -- which serializes one process, and a fleet is not one process. Postgres row locks cannot fix it alone: with no active run there is no row to lock, so both admitters count zero and both insert. The whole decision -- dedupe, every policy read, any policy mutation, and the insert -- now executes inside one store transaction under a durable lock on (workflow_id, flow_key): pg_advisory_xact_lock on Postgres, BEGIN IMMEDIATE's write lock on SQLite, the store lock in memory. The kernel builds a FlowGate and hands the store the records; admit(max_active=) is gone, one contract with one enforcer. Semantics pinned by conformance checks and by cross-instance race tests that open two stores -- two pools, no shared Python state -- over one database: skip admits exactly one; rate limit of one rejects the loser; throttle spaces a racing pair a full window apart; debounce coalesces the loser into the winner and, fixing the review's P1 alongside, takes the burst's LATEST payload; singleton cancel leaves at most one uncancelled run. All pass against real Postgres 16. Fan-out branches are written by the parent's commit, not through policy admission, so a policy-decorated branch root is now refused with a teaching error instead of having its policy silently bypassed -- and because that is a definition problem, WorkflowDefinitionError joins BUG_EXCEPTIONS and fails the run on attempt one instead of retrying a deterministic error. Singleton cancel changes shape slightly: the incumbent's cancellation intent is recorded in the admitting transaction and the replacement is admitted immediately, so at most one non-cancelling run exists per key at every instant, from any number of processes; incumbents drain to CANCELLED asynchronously as any cancelled run does. The contract states all of this in section 1. --- news/workflow-atomic-start-policies.bugfix.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 1 + reflex/workflow/CONTRACT.md | 28 ++ reflex/workflow/conformance.py | 179 ++++++++- reflex/workflow/kernel.py | 225 +++++------ reflex/workflow/postgres.py | 142 ++++++- reflex/workflow/store.py | 359 +++++++++++++++--- tests/units/workflow/test_flow_atomicity.py | 239 ++++++++++++ tests/units/workflow/test_parallel.py | 70 ++++ 9 files changed, 1056 insertions(+), 188 deletions(-) create mode 100644 news/workflow-atomic-start-policies.bugfix.md create mode 100644 tests/units/workflow/test_flow_atomicity.py diff --git a/news/workflow-atomic-start-policies.bugfix.md b/news/workflow-atomic-start-policies.bugfix.md new file mode 100644 index 00000000000..ddec04add0f --- /dev/null +++ b/news/workflow-atomic-start-policies.bugfix.md @@ -0,0 +1 @@ +Start policies are now enforced by the store inside the admitting transaction, under a durable lock on the run's `(workflow_id, flow_key)` — a Postgres advisory transaction lock, SQLite's write lock, the store lock in memory. Previously the decision was guarded by an in-process asyncio lock, so two worker processes racing a `Singleton(mode="skip")`, `RateLimit`, `Throttle`, `Debounce`, or `Singleton(mode="cancel")` both passed a limit of one — an external review demonstrated all five families violating their invariant in 50 of 50 synchronized trials across two OS processes. Debounce coalescing is now also latest-wins: the pending run takes the burst's final payload along with its fresh deadline, in the same transaction. Fan-out refuses a branch root that declares a start policy instead of silently bypassing it, and a `WorkflowDefinitionError` raised from a handler now fails its run on the first attempt instead of consuming retries. Cross-instance races are pinned by tests that open two independent stores — two pools, no shared Python state — over one SQLite file and one Postgres schema. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 99703f15817..0abdea3a754 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -185,6 +185,7 @@ def is_retryable(self, error: BaseException) -> bool: SyntaxError, IndentationError, NotImplementedError, + WorkflowDefinitionError, ) """Exceptions that mean the code is wrong, not that the world was unlucky. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 47ba3311270..729e64e5f0e 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -37,6 +37,34 @@ Admission is likewise one transaction: run row, root slot, dedupe reservation, history. A webhook is acknowledged only after that transaction commits (admit-before-ack), so a `202` means the run exists durably. +A start policy (`Singleton`/`RateLimit`/`Throttle`/`Debounce`) is part of that +same transaction, serialized under a durable lock on the run's +`(workflow_id, flow_key)` — a Postgres advisory transaction lock, SQLite's +database write lock, the store lock in memory. Every policy read, any policy +mutation (a debounce extending and re-payloading its pending run, a +cancel-mode singleton requesting its incumbent's cancellation), and the +insert commit or roll back together. Two processes admitting concurrently +under a limit of one therefore cannot both pass: nothing about policy +enforcement assumes the admitters share a process. Concretely: + +- `Singleton(mode="skip")`: at most one active run per key, at every instant, + from any number of processes; the loser is told which run holds the key. +- `Singleton(mode="cancel")`: the replacement is admitted and every + incumbent's cancellation intent is recorded in one transaction, so at most + one *non-cancelling* run exists per key at every instant. Incumbents drain + to `CANCELLED` asynchronously, exactly as any cancelled run does. +- `RateLimit`: the (limit+1)th start inside the window is `rejected` with a + `retry_after`, counted against committed admissions only. +- `Throttle`: every start is admitted, each due one window after the limit-th + most recent scheduled start, so a racing burst is spaced, not replayed. +- `Debounce`: a start that lands inside the quiet period is `coalesced` into + the pending run, which takes the **latest** payload and a fresh deadline; a + debounced burst starts once, with its final revision. + +Fan-out branches are written by the parent's committing transaction rather +than through policy admission, so a branch root that declares a start policy +is refused at fan-out time instead of having its policy silently bypassed. + Substep results (`rx.step`) are the deliberate exception: each records in its **own** transaction the moment the callable returns, because their purpose is to survive a crash that prevents the attempt from ever committing. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 760454bb835..d7a37c9ba97 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -940,26 +940,41 @@ async def check_the_crash_matrix_holds_at_every_boundary(store: RunStore) -> Non assert steps[1].status is StepStatus.READY -async def check_admission_enforces_the_active_limit(store: RunStore) -> None: - """A singleton's limit is decided by the admitting transaction itself.""" - first = make_run("a", flow_key="k1") - assert await store.admit(first, make_step("a"), _ADMITTED, max_active=1) == ( - True, - "a", +async def check_flow_gate_enforces_every_policy(store: RunStore) -> None: + """The whole policy decision is one atomic store transaction. + + Deciding outside the store is a check-then-act race two processes both + win, which is how every advertised start policy was violated 50 out of 50 + times across two OS processes. These pin the store-side semantics; the + cross-instance atomicity itself is pinned by tests that open two stores + on one database. + """ + from reflex.workflow.store import FlowGate + + skip = FlowGate(singleton_skip=True) + first = await store.admit_flow( + make_run("a", flow_key="k1"), make_step("a"), _ADMITTED, skip, NOW ) - # A second start under the same key is refused, and told which run holds it. - second = make_run("b", flow_key="k1", created_at=NOW + 1) - assert await store.admit(second, make_step("b"), _ADMITTED, max_active=1) == ( - False, - "a", + assert (first.disposition, first.run_id) == ("started", "a") + # A second start under the key is refused and told which run holds it. + second = await store.admit_flow( + make_run("b", flow_key="k1", created_at=NOW + 1), + make_step("b"), + _ADMITTED, + skip, + NOW + 1, ) + assert (second.disposition, second.run_id) == ("skipped", "a") assert await store.get_run("b") is None # A different key is unaffected. - other = make_run("c", flow_key="k2", created_at=NOW + 2) - assert await store.admit(other, make_step("c"), _ADMITTED, max_active=1) == ( - True, - "c", + other = await store.admit_flow( + make_run("c", flow_key="k2", created_at=NOW + 2), + make_step("c"), + _ADMITTED, + skip, + NOW + 2, ) + assert other.disposition == "started" # Once the holder is terminal the key is free again. assert await store.finalize_run( "a", @@ -968,10 +983,131 @@ async def check_admission_enforces_the_active_limit(store: RunStore) -> None: event=HistoryEventType.RUN_COMPLETED, now=NOW + 3, ) - third = make_run("d", flow_key="k1", created_at=NOW + 4) - assert await store.admit(third, make_step("d"), _ADMITTED, max_active=1) == ( - True, - "d", + third = await store.admit_flow( + make_run("d", flow_key="k1", created_at=NOW + 4), + make_step("d"), + _ADMITTED, + skip, + NOW + 4, + ) + assert third.disposition == "started" + + +async def check_flow_gate_rate_throttle_and_debounce(store: RunStore) -> None: + """Rate limits refuse, throttles delay, debounces coalesce latest-wins.""" + from reflex.workflow.store import FlowGate + + rate = FlowGate(rate_limit=(1, 60.0)) + first = await store.admit_flow( + make_run("r1", flow_key="rk"), make_step("r1"), _ADMITTED, rate, NOW + ) + assert first.disposition == "started" + refused = await store.admit_flow( + make_run("r2", flow_key="rk", created_at=NOW + 1), + make_step("r2"), + _ADMITTED, + rate, + NOW + 1, + ) + assert refused.disposition == "rejected" + assert refused.retry_after is not None + assert abs(refused.retry_after - 60.0) < 1e-9 + assert await store.get_run("r2") is None + + throttle = FlowGate(throttle=(1, 60.0)) + held = await store.admit_flow( + make_run("t1", flow_key="tk", created_at=NOW + 1), + make_step("t1", due_at=NOW + 1), + _ADMITTED, + throttle, + NOW + 1, + ) + assert held.disposition == "started" + spaced = await store.admit_flow( + make_run("t2", flow_key="tk", created_at=NOW + 2), + make_step("t2", due_at=NOW + 2), + _ADMITTED, + throttle, + NOW + 2, + ) + assert spaced.disposition == "started" + steps = await store.get_steps("t2") + assert abs(steps[0].due_at - (NOW + 1 + 60.0)) < 1e-9, ( + "the second start sits one window after the first, not at its own time" + ) + + debounce = FlowGate(debounce=30.0) + pending = await store.admit_flow( + make_run("d1", flow_key="dk", created_at=NOW + 2), + make_step("d1", args={"revision": 1}, due_at=NOW + 2), + _ADMITTED, + debounce, + NOW + 2, + ) + assert pending.disposition == "started" + coalesced = await store.admit_flow( + make_run("d2", flow_key="dk", created_at=NOW + 3), + make_step("d2", args={"revision": 2}, due_at=NOW + 3), + _ADMITTED, + debounce, + NOW + 3, + ) + assert (coalesced.disposition, coalesced.run_id) == ("coalesced", "d1") + assert await store.get_run("d2") is None + steps = await store.get_steps("d1") + assert abs(steps[0].due_at - (NOW + 3 + 30.0)) < 1e-9, "the quiet period extends" + assert steps[0].args == {"revision": 2}, ( + "a debounced burst starts with its last payload, not its first" + ) + + +async def check_flow_gate_singleton_cancel_replaces(store: RunStore) -> None: + """Cancel-mode admits the replacement and cancels the incumbent atomically.""" + from reflex.workflow.store import FlowGate + + cancel = FlowGate(singleton_cancel=True) + first = await store.admit_flow( + make_run("old", flow_key="ck"), make_step("old"), _ADMITTED, cancel, NOW + ) + assert first.disposition == "started" + replaced = await store.admit_flow( + make_run("new", flow_key="ck", created_at=NOW + 1), + make_step("new"), + _ADMITTED, + cancel, + NOW + 1, + ) + assert replaced.disposition == "started" + assert replaced.cancelled == ("old",) + old = await store.get_run("old") + assert old is not None + assert old.cancel_requested, ( + "the incumbent's cancellation intent rides the admitting transaction" + ) + + +async def check_flow_gate_dedupes_before_policy(store: RunStore) -> None: + """A redelivered event is its prior run, not a new start to be policed.""" + from reflex.workflow.store import FlowGate + + gate = FlowGate(rate_limit=(1, 60.0)) + first = await store.admit_flow( + make_run("g1", flow_key="gk", request_key="evt-1"), + make_step("g1"), + _ADMITTED, + gate, + NOW, + ) + assert first.disposition == "started" + redelivered = await store.admit_flow( + make_run("g2", flow_key="gk", request_key="evt-1", created_at=NOW + 1), + make_step("g2"), + _ADMITTED, + gate, + NOW + 1, + ) + assert (redelivered.disposition, redelivered.run_id) == ("deduplicated", "g1"), ( + "hitting the rate limit must not turn a redelivery into a rejection" ) @@ -1152,7 +1288,10 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_a_terminal_run_refuses_further_control, check_finalize_delivers_a_childs_arrival, check_the_crash_matrix_holds_at_every_boundary, - check_admission_enforces_the_active_limit, + check_flow_gate_enforces_every_policy, + check_flow_gate_rate_throttle_and_debounce, + check_flow_gate_singleton_cancel_replaces, + check_flow_gate_dedupes_before_policy, check_skip_unsticks_a_stopped_run, check_retry_reopens_only_failed_runs, check_force_finalize_records_a_result, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 026f8552244..5de28901e37 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -23,7 +23,7 @@ from pydantic import TypeAdapter from reflex_base.event.processor.base_state_processor import _transform_event_payload from reflex_base.utils import console -from reflex_base.utils.exceptions import WorkflowRuntimeError +from reflex_base.utils.exceptions import WorkflowDefinitionError, WorkflowRuntimeError from reflex_base.workflow import ( DEFAULT_LEASE_DURATION, DEFAULT_MAX_RECOVERIES, @@ -61,6 +61,7 @@ from reflex.workflow.store import ( Claim, DeliveryDisposition, + FlowGate, RunStore, StaleClaimError, StepCompletion, @@ -500,7 +501,6 @@ def __init__( self._observer = observer self._queues = tuple(queues) if queues is not None else None self._wakeup = asyncio.Event() - self._admission = asyncio.Lock() self._closing = False self._worker: asyncio.Task | None = None @@ -611,73 +611,6 @@ def _flow_key(handler: HandlerDefinition, payload: dict[str, Any]) -> str | None return handler.id return f"{handler.id}:{_extract_key_field(payload, field)!r}" - async def _apply_start_policy( - self, - defn: WorkflowDefinition, - handler: HandlerDefinition, - flow_key: str, - now: float, - ) -> tuple[StartResult | None, float]: - """Decide whether and when a start may proceed. - - Args: - defn: The workflow definition. - handler: The root handler being started. - flow_key: The computed grouping key. - now: Current time in epoch seconds. - - Returns: - A result that ends admission, or None to proceed, together with the - time the root slot becomes due. - """ - if handler.singleton is not None: - existing = await self._store.first_active(defn.workflow_id, flow_key) - if handler.singleton.mode == "skip": - # Deliberately not decided here: admission re-checks inside its - # own transaction (see _admit), because a decision made out - # here is one that two concurrent starts both win. - return None, now - if existing is not None: - # Drive the cancellation to a terminal state before admitting the - # replacement, so "one active run per key" holds at every instant - # rather than only once a worker happens to drain the old one. - await self.cancel(existing.run_id) - await self._finalize_control(now) - if handler.rate_limit is not None: - window = parse_duration(handler.rate_limit.period) - started = await self._store.count_started_since( - defn.workflow_id, flow_key, now - window - ) - if started >= handler.rate_limit.limit: - return ( - StartResult( - disposition="rejected", retryable=True, retry_after=window - ), - now, - ) - if handler.throttle is not None: - window = parse_duration(handler.throttle.period) - previous = await self._store.nth_recent_start( - defn.workflow_id, flow_key, handler.throttle.limit - ) - if previous is not None and previous + window > now: - # Delay rather than drop, and space the backlog: each start - # sits a window after the limit-th most recent one, so a held - # burst is released at the configured rate instead of at once. - return None, previous + window - if handler.debounce is not None: - window = parse_duration(handler.debounce.period) - pending = await self._store.first_active(defn.workflow_id, flow_key) - if pending is not None and await self._store.defer_root( - pending.run_id, now + window, now - ): - return ( - StartResult(disposition="coalesced", run_id=pending.run_id), - now, - ) - return None, now + window - return None, now - async def _started_handler(self, run_id: str, handler_id: str) -> bool: """Whether a run's root step is the handler a delivery would start. @@ -767,20 +700,60 @@ async def start( flow_key = self._flow_key(handler, payload) if flow_key is None: return await self._admit(defn, handler, payload, request_key, labels, None) - # A start policy is a check followed by an insert, so concurrent starts - # must not interleave between them or two runs slip past a singleton. - async with self._admission: - now = self._clock() - decided, due_at = await self._apply_start_policy( - defn, handler, flow_key, now + # A start policy is a read followed by a write, and only the store can + # make the pair atomic: an in-process lock serializes one process, and + # a fleet is not one process. The whole decision executes inside a + # single store transaction under a durable lock on the flow key. + now = self._clock() + singleton = handler.singleton + gate = FlowGate( + singleton_skip=singleton is not None and singleton.mode == "skip", + singleton_cancel=singleton is not None and singleton.mode == "cancel", + rate_limit=( + (handler.rate_limit.limit, parse_duration(handler.rate_limit.period)) + if handler.rate_limit is not None + else None + ), + throttle=( + (handler.throttle.limit, parse_duration(handler.throttle.period)) + if handler.throttle is not None + else None + ), + debounce=( + parse_duration(handler.debounce.period) + if handler.debounce is not None + else None + ), + ) + run, root_step, admission = self._admission_records( + defn, handler, payload, request_key, labels, flow_key, now + ) + outcome = await self._store.admit_flow(run, root_step, admission, gate, now) + for cancelled_id in outcome.cancelled: + # Durable intent was written in the admitting transaction; what is + # left is this process's share -- stop a local in-flight attempt + # and tell the observer. + await self._notify_run( + cancelled_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),) ) - if decided is not None: - return decided - return await self._admit( - defn, handler, payload, request_key, labels, flow_key, due_at + task = self._inflight.get(cancelled_id) + if task is not None: + task.cancel() + if outcome.cancelled: + await self._finalize_control(self._clock()) + if outcome.disposition == "started": + self._notify(run, admission) + self._wakeup.set() + return StartResult(disposition="started", run_id=outcome.run_id) + if outcome.disposition == "rejected": + return StartResult( + disposition="rejected", + retryable=True, + retry_after=outcome.retry_after, ) + return StartResult(disposition=outcome.disposition, run_id=outcome.run_id) - async def _admit( + def _admission_records( self, defn: WorkflowDefinition, handler: HandlerDefinition, @@ -788,9 +761,14 @@ async def _admit( request_key: str | None, labels: dict[str, str] | None, flow_key: str | None, + now: float, due_at: float | None = None, - ) -> StartResult: - """Create the run and its root slot. + ) -> tuple[ + RunRecord, + StepRecord, + tuple[tuple[HistoryEventType, dict[str, Any]], ...], + ]: + """Build the records one admission writes. Args: defn: The workflow definition. @@ -799,13 +777,12 @@ async def _admit( request_key: Idempotent admission key. labels: Server-derived indexing labels. flow_key: Start-policy grouping key, if the root declares one. + now: Current time in epoch seconds. due_at: Earliest start time, when a policy delayed it. Returns: - The admission result. + The run record, its root slot, and the admission history events. """ - now = self._clock() - due_at = now if due_at is None else due_at run_id = uuid.uuid4().hex run = RunRecord( run_id=run_id, @@ -828,7 +805,7 @@ async def _admit( handler_id=handler.id, status=StepStatus.READY, args=payload, - due_at=due_at, + due_at=now if due_at is None else due_at, origin="root", queue=handler.queue or "default", created_at=now, @@ -844,29 +821,44 @@ async def _admit( {"ordinal": 0, "handler_id": handler.id}, ), ) - # A singleton's "at most one active run per key" is enforced by the - # store inside the admitting transaction. Deciding it beforehand is a - # check-then-act race that two concurrent starts both pass, which for - # a singleton means exactly the duplicate it exists to prevent. - singleton = handler.singleton - max_active = 1 if singleton is not None and singleton.mode == "skip" else None + return run, root_step, admission + + async def _admit( + self, + defn: WorkflowDefinition, + handler: HandlerDefinition, + payload: dict[str, Any], + request_key: str | None, + labels: dict[str, str] | None, + flow_key: str | None, + due_at: float | None = None, + ) -> StartResult: + """Create the run and its root slot for a policy-free root. + + Roots that declare a start policy go through ``admit_flow`` instead, + where the whole decision is one store transaction. + + Args: + defn: The workflow definition. + handler: The root handler. + payload: The decoded start payload. + request_key: Idempotent admission key. + labels: Server-derived indexing labels. + flow_key: Start-policy grouping key, if the root declares one. + due_at: Earliest start time, when a policy delayed it. + + Returns: + The admission result. + """ + now = self._clock() + run, root_step, admission = self._admission_records( + defn, handler, payload, request_key, labels, flow_key, now, due_at + ) created, authoritative_run_id = await self._store.admit( - run, root_step, admission, max_active=max_active + run, root_step, admission ) if not created: - disposition = "skipped" if max_active is not None else "deduplicated" - if run.request_key is not None and disposition == "skipped": - # A redelivery that also hits the limit is still a dedupe: - # the caller is asking about a run it already started. - existing = await self._store.find_by_request_key( - defn.workflow_id, run.request_key - ) - if existing == authoritative_run_id: - disposition = "deduplicated" - return StartResult( - disposition=disposition, # pyright: ignore[reportArgumentType] - run_id=authoritative_run_id, - ) + return StartResult(disposition="deduplicated", run_id=authoritative_run_id) self._notify(run, admission) self._wakeup.set() return StartResult(disposition="started", run_id=authoritative_run_id) @@ -2357,6 +2349,27 @@ def _child_records( "and a branch must be a manual root just like a direct start." ) raise WorkflowRuntimeError(msg) + if ( + handler.singleton is not None + or handler.rate_limit is not None + or handler.throttle is not None + or handler.debounce is not None + ): + # Fan-out writes its children in the parent's committing + # transaction rather than through policy admission, so a + # policy on a branch root would be silently bypassed -- five + # branches under a throttle of two all start at once, and + # nothing says so. Refusing is honest until branches go + # through the same admission primitive as a direct start. + msg = ( + f"Cannot fan out to {handler.id!r} of {defn.workflow_id!r}: " + "it declares a start policy (singleton, rate_limit, " + "throttle, or debounce), and fan-out admits branches " + "directly, bypassing policies. Remove the policy from " + "this root, or start it as its own run with " + "rx.workflows.start()." + ) + raise WorkflowDefinitionError(msg) child_id = uuid.uuid4().hex records.append(( RunRecord( diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 2d33f068e69..e7e24d5d723 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -32,7 +32,7 @@ StepRecord, StepStatus, ) -from reflex.workflow.store import Claim, StaleClaimError +from reflex.workflow.store import Claim, FlowAdmission, FlowGate, StaleClaimError try: import psycopg @@ -614,18 +614,16 @@ async def admit( run: RunRecord, root_step: StepRecord, events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], - *, - max_active: int | None = None, ) -> tuple[bool, str]: """Atomically admit a run, deduplicating on the request key. + Roots with a start policy are admitted through ``admit_flow``, where + the whole policy decision shares the admitting transaction. + Args: run: The run record to create. root_step: The preallocated root mailbox slot. events: History events to append on creation. - max_active: When set, admit only if fewer than this many runs are - already active under the run's flow key, decided inside this - transaction so concurrent starts cannot both pass. Returns: Whether the run was created, and the authoritative run id. @@ -647,24 +645,130 @@ async def admit( existing = await cursor.fetchone() if existing is not None: return False, existing["run_id"] - if max_active is not None and run.flow_key is not None: - # FOR UPDATE serializes concurrent admissions under one key: - # the second waits, then sees the first and is refused. - cursor = await conn.execute( - "SELECT run_id FROM workflow_runs" - " WHERE workflow_id = %s AND flow_key = %s" - " AND NOT (status = ANY(%s))" - " ORDER BY created_at, run_id FOR UPDATE", - (run.workflow_id, run.flow_key, _TERMINAL_RUNS), - ) - active = await cursor.fetchall() - if len(active) >= max_active: - return False, active[0]["run_id"] await self._insert_run(conn, run) await self._insert_step(conn, root_step) await self._append_events(conn, run.run_id, events, run.created_at) return True, run.run_id + async def admit_flow( + self, + run: RunRecord, + root_step: StepRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + gate: FlowGate, + now: float, + ) -> FlowAdmission: + """Admit a run under a start policy, atomically. + + The transaction opens by taking an advisory lock on the flow key. + Row locks cannot serialize this decision -- when no run exists yet + there is no row to lock, and two transactions both count zero and + both insert -- but an advisory lock exists before any row does, so + the second admitter waits and then reads what the first committed. + + Args: + run: The run record to create, carrying the flow key. + root_step: The preallocated root mailbox slot. + events: History events to append on creation. + gate: The policy to enforce. + now: Current time in epoch seconds. + + Returns: + What was done, decided inside the transaction. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + await conn.execute( + "SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", + (f"{run.workflow_id}\x1f{run.flow_key}",), + ) + if run.request_key is not None: + cursor = await conn.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = %s AND request_key = %s", + (run.workflow_id, run.request_key), + ) + row = await cursor.fetchone() + if row is not None: + return FlowAdmission("deduplicated", row["run_id"]) + cursor = await conn.execute( + "SELECT run_id FROM workflow_runs" + " WHERE workflow_id = %s AND flow_key = %s" + " AND NOT (status = ANY(%s))" + " ORDER BY created_at, run_id", + (run.workflow_id, run.flow_key, _TERMINAL_RUNS), + ) + active = await cursor.fetchall() + if gate.singleton_skip and active: + return FlowAdmission("skipped", active[0]["run_id"]) + cancelled: list[str] = [] + if gate.singleton_cancel and active: + ids = [row["run_id"] for row in active] + await conn.execute( + "UPDATE workflow_runs SET cancel_requested = TRUE," + " status = %s, updated_at = %s WHERE run_id = ANY(%s)", + (RunStatus.CANCELLING.value, now, ids), + ) + for run_id in ids: + await self._append_events( + conn, + run_id, + ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), + now, + ) + cancelled = ids + due_at = root_step.due_at + if gate.rate_limit is not None: + limit, window = gate.rate_limit + cursor = await conn.execute( + "SELECT count(*) AS n FROM workflow_runs" + " WHERE workflow_id = %s AND flow_key = %s AND created_at > %s", + (run.workflow_id, run.flow_key, now - window), + ) + row = await cursor.fetchone() + if row is not None and row["n"] >= limit: + return FlowAdmission("rejected", retry_after=window) + if gate.throttle is not None: + limit, window = gate.throttle + cursor = await conn.execute( + "SELECT GREATEST(s.due_at, r.created_at) AS start" + " FROM workflow_runs r JOIN workflow_steps s" + " ON s.run_id = r.run_id AND s.ordinal = 0" + " WHERE r.workflow_id = %s AND r.flow_key = %s" + " ORDER BY start DESC OFFSET %s LIMIT 1", + (run.workflow_id, run.flow_key, limit - 1), + ) + row = await cursor.fetchone() + if row is not None and row["start"] + window > now: + due_at = row["start"] + window + if gate.debounce is not None: + if active: + cursor = await conn.execute( + "UPDATE workflow_steps SET args = %s, due_at = %s," + " updated_at = %s WHERE run_id = %s AND ordinal = 0" + " AND status = %s", + ( + _json(root_step.args), + now + gate.debounce, + now, + active[0]["run_id"], + StepStatus.READY.value, + ), + ) + if cursor.rowcount: + return FlowAdmission("coalesced", active[0]["run_id"]) + due_at = now + gate.debounce + if run.request_key is not None: + await conn.execute( + "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" + " VALUES (%s, %s, %s)", + (run.workflow_id, run.request_key, run.run_id), + ) + await self._insert_run(conn, run) + await self._insert_step(conn, dataclasses.replace(root_step, due_at=due_at)) + await self._append_events(conn, run.run_id, events, run.created_at) + return FlowAdmission("started", run.run_id, cancelled=tuple(cancelled)) + async def claim_next( self, now: float, diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index afb0b2b0914..bc187d2d37c 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -117,6 +117,54 @@ class StepCompletion: parent_arrival: tuple[str, int, dict[str, Any], str] | None = None +@dataclasses.dataclass(frozen=True, slots=True) +class FlowGate: + """One root's start policy, evaluated inside the admitting transaction. + + A start policy is a read followed by a write -- count the active runs, + then insert or refuse -- and any gap between the two is a race that + concurrent starts in different processes both win. The kernel used to + guard the gap with an asyncio lock, which serializes one process and + nothing else. The whole decision therefore belongs to the store, executed + under a durable per-key lock, and this value is the policy handed in. + + Attributes: + singleton_skip: Refuse the start while any run is active on the key. + singleton_cancel: Request cancellation of every active run on the key + before admitting the replacement, in the same transaction. + rate_limit: ``(limit, window_seconds)``; refuse once the key has had + that many starts inside the window. + throttle: ``(limit, window_seconds)``; delay the root so each start + sits a window after the limit-th most recent one. + debounce: Quiet-period seconds; an existing pending root absorbs this + start, taking its payload and a fresh deadline. + """ + + singleton_skip: bool = False + singleton_cancel: bool = False + rate_limit: tuple[int, float] | None = None + throttle: tuple[int, float] | None = None + debounce: float | None = None + + +@dataclasses.dataclass(frozen=True, slots=True) +class FlowAdmission: + """What one gated admission did, decided atomically by the store. + + Attributes: + disposition: How the submission was handled. + run_id: The created or prior run, when one identifies the outcome. + retry_after: Suggested resubmission delay for a rejected start. + cancelled: Runs whose cancellation this admission requested, for the + kernel to stop locally and finalize. + """ + + disposition: Literal["started", "skipped", "rejected", "coalesced", "deduplicated"] + run_id: str | None = None + retry_after: float | None = None + cancelled: tuple[str, ...] = () + + class RunStore(Protocol): """Protocol implemented by workflow run stores.""" @@ -125,19 +173,16 @@ async def admit( run: RunRecord, root_step: StepRecord, events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], - *, - max_active: int | None = None, ) -> tuple[bool, str]: """Atomically admit a run, deduplicating on the request key. + Roots with a start policy are admitted through ``admit_flow``, where + the whole policy decision shares the admitting transaction. + Args: run: The run record to create. root_step: The preallocated root mailbox slot. events: History events to append on creation. - max_active: When set, admit only if fewer than this many runs are - already active under the run's flow key. Checked inside the - admitting transaction, because a check made outside it is a - race that two concurrent starts both win. Returns: ``(True, run_id)`` when the run was created, or @@ -331,6 +376,34 @@ async def count_active(self, workflow_id: str, flow_key: str) -> int: """ ... + async def admit_flow( + self, + run: RunRecord, + root_step: StepRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + gate: FlowGate, + now: float, + ) -> FlowAdmission: + """Admit a run under a start policy, atomically. + + The dedupe check, every policy read, any policy mutation, and the + insert happen in one transaction under a durable lock on the run's + ``(workflow_id, flow_key)``, so two processes admitting concurrently + cannot both pass a limit of one. Policies apply in declaration order: + dedupe, singleton, rate limit, throttle, debounce. + + Args: + run: The run record to create, carrying the flow key. + root_step: The preallocated root mailbox slot. + events: History events to append on creation. + gate: The policy to enforce. + now: Current time in epoch seconds. + + Returns: + What was done, decided inside the transaction. + """ + ... + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: """Find the oldest run still in flight under a flow-control key. @@ -869,19 +942,16 @@ async def admit( run: RunRecord, root_step: StepRecord, events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], - *, - max_active: int | None = None, ) -> tuple[bool, str]: """Atomically admit a run, deduplicating on the request key. + Roots with a start policy are admitted through ``admit_flow``, where + the whole policy decision shares the admitting transaction. + Args: run: The run record to create. root_step: The preallocated root mailbox slot. events: History events to append on creation. - max_active: When set, admit only if fewer than this many runs are - already active under the run's flow key. Checked inside the - admitting transaction, because a check made outside it is a - race that two concurrent starts both win. Returns: Whether the run was created, and the authoritative run id. @@ -892,19 +962,6 @@ async def admit( existing = self._dedupe.get(dedupe_key) if existing is not None: return False, existing - if max_active is not None and run.flow_key is not None: - active = sorted( - ( - other - for other in self._runs.values() - if other.workflow_id == run.workflow_id - and other.flow_key == run.flow_key - and other.status not in TERMINAL_RUN_STATUSES - ), - key=lambda other: (other.created_at, other.run_id), - ) - if len(active) >= max_active: - return False, active[0].run_id if run.request_key is not None: self._dedupe[run.workflow_id, run.request_key] = run.run_id self._runs[run.run_id] = run @@ -1327,6 +1384,109 @@ async def count_active(self, workflow_id: str, flow_key: str) -> int: and run.status not in TERMINAL_RUN_STATUSES ) + async def admit_flow( + self, + run: RunRecord, + root_step: StepRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + gate: FlowGate, + now: float, + ) -> FlowAdmission: + """Admit a run under a start policy, atomically. + + Args: + run: The run record to create, carrying the flow key. + root_step: The preallocated root mailbox slot. + events: History events to append on creation. + gate: The policy to enforce. + now: Current time in epoch seconds. + + Returns: + What was done, decided under the store lock. + """ + async with self._lock: + if run.request_key is not None: + existing = self._dedupe.get((run.workflow_id, run.request_key)) + if existing is not None: + return FlowAdmission("deduplicated", existing) + active = sorted( + ( + other + for other in self._runs.values() + if other.workflow_id == run.workflow_id + and other.flow_key == run.flow_key + and other.status not in TERMINAL_RUN_STATUSES + ), + key=lambda other: (other.created_at, other.run_id), + ) + if gate.singleton_skip and active: + return FlowAdmission("skipped", active[0].run_id) + cancelled: list[str] = [] + if gate.singleton_cancel: + for other in active: + self._runs[other.run_id] = dataclasses.replace( + other, + cancel_requested=True, + status=RunStatus.CANCELLING, + updated_at=now, + ) + self._append_events( + other.run_id, + ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), + now, + ) + cancelled.append(other.run_id) + due_at = root_step.due_at + if gate.rate_limit is not None: + limit, window = gate.rate_limit + started = sum( + 1 + for other in self._runs.values() + if other.workflow_id == run.workflow_id + and other.flow_key == run.flow_key + and other.created_at > now - window + ) + if started >= limit: + return FlowAdmission("rejected", retry_after=window) + if gate.throttle is not None: + limit, window = gate.throttle + starts = sorted( + ( + max(steps[0].due_at, other.created_at) + for other in self._runs.values() + if other.workflow_id == run.workflow_id + and other.flow_key == run.flow_key + if (steps := self._steps.get(other.run_id)) + ), + reverse=True, + ) + previous = starts[limit - 1] if len(starts) >= limit else None + if previous is not None and previous + window > now: + due_at = previous + window + if gate.debounce is not None: + pending = active[0] if active else None + if pending is not None: + steps = self._steps.get(pending.run_id) + if steps and steps[0].status is StepStatus.READY: + # Latest wins: the burst's final payload is the one + # the debounced run eventually starts with, and the + # replacement rides the same transaction as the + # deadline extension. + steps[0] = dataclasses.replace( + steps[0], + args=root_step.args, + due_at=now + gate.debounce, + updated_at=now, + ) + return FlowAdmission("coalesced", pending.run_id) + due_at = now + gate.debounce + if run.request_key is not None: + self._dedupe[run.workflow_id, run.request_key] = run.run_id + self._runs[run.run_id] = run + self._steps[run.run_id] = [dataclasses.replace(root_step, due_at=due_at)] + self._append_events(run.run_id, events, run.created_at) + return FlowAdmission("started", run.run_id, cancelled=tuple(cancelled)) + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: """Find the oldest run still in flight under a flow-control key. @@ -2419,19 +2579,16 @@ async def admit( run: RunRecord, root_step: StepRecord, events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], - *, - max_active: int | None = None, ) -> tuple[bool, str]: """Atomically admit a run, deduplicating on the request key. + Roots with a start policy are admitted through ``admit_flow``, where + the whole policy decision shares the admitting transaction. + Args: run: The run record to create. root_step: The preallocated root mailbox slot. events: History events to append on creation. - max_active: When set, admit only if fewer than this many runs are - already active under the run's flow key. Checked inside the - admitting transaction, because a check made outside it is a - race that two concurrent starts both win. Returns: Whether the run was created, and the authoritative run id. @@ -2460,18 +2617,6 @@ def work(): " VALUES (?, ?, ?)", (run.workflow_id, run.request_key, run.run_id), ) - if max_active is not None and run.flow_key is not None: - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) - rows = self._db.execute( - "SELECT run_id FROM workflow_runs" - " WHERE workflow_id = ? AND flow_key = ?" - f" AND status NOT IN ({','.join('?' * len(terminal))})" - " ORDER BY created_at, run_id", - (run.workflow_id, run.flow_key, *terminal), - ).fetchall() - if len(rows) >= max_active: - self._db.execute("ROLLBACK") - return False, rows[0]["run_id"] self._insert_run(run) self._insert_step(root_step) self._append_events(run.run_id, events, run.created_at) @@ -2483,6 +2628,134 @@ def work(): return await asyncio.to_thread(work) + async def admit_flow( + self, + run: RunRecord, + root_step: StepRecord, + events: tuple[tuple[HistoryEventType, dict[str, Any]], ...], + gate: FlowGate, + now: float, + ) -> FlowAdmission: + """Admit a run under a start policy, atomically. + + ``BEGIN IMMEDIATE`` takes the database write lock up front, so the + whole decision is atomic against every other connection -- including + one held by a different process sharing the file, which is exactly + where an in-process lock stops helping. + + Args: + run: The run record to create, carrying the flow key. + root_step: The preallocated root mailbox slot. + events: History events to append on creation. + gate: The policy to enforce. + now: Current time in epoch seconds. + + Returns: + What was done, decided inside the transaction. + """ + + def work() -> FlowAdmission: + """Run the whole gated admission in one write transaction. + + Returns: + The admission outcome. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + not_terminal = f"status NOT IN ({','.join('?' * len(terminal))})" + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + if run.request_key is not None: + row = self._db.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = ? AND request_key = ?", + (run.workflow_id, run.request_key), + ).fetchone() + if row is not None: + self._db.execute("ROLLBACK") + return FlowAdmission("deduplicated", row["run_id"]) + active = self._db.execute( + "SELECT run_id FROM workflow_runs" + f" WHERE workflow_id = ? AND flow_key = ? AND {not_terminal}" + " ORDER BY created_at, run_id", + (run.workflow_id, run.flow_key, *terminal), + ).fetchall() + if gate.singleton_skip and active: + self._db.execute("ROLLBACK") + return FlowAdmission("skipped", active[0]["run_id"]) + cancelled: list[str] = [] + if gate.singleton_cancel: + for row in active: + self._db.execute( + "UPDATE workflow_runs SET cancel_requested = 1," + " status = ?, updated_at = ? WHERE run_id = ?", + (RunStatus.CANCELLING.value, now, row["run_id"]), + ) + self._append_events( + row["run_id"], + ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), + now, + ) + cancelled.append(row["run_id"]) + due_at = root_step.due_at + if gate.rate_limit is not None: + limit, window = gate.rate_limit + started = self._db.execute( + "SELECT count(*) AS n FROM workflow_runs" + " WHERE workflow_id = ? AND flow_key = ?" + " AND created_at > ?", + (run.workflow_id, run.flow_key, now - window), + ).fetchone()["n"] + if started >= limit: + self._db.execute("ROLLBACK") + return FlowAdmission("rejected", retry_after=window) + if gate.throttle is not None: + limit, window = gate.throttle + row = self._db.execute( + "SELECT MAX(s.due_at, r.created_at) AS start" + " FROM workflow_runs r JOIN workflow_steps s" + " ON s.run_id = r.run_id AND s.ordinal = 0" + " WHERE r.workflow_id = ? AND r.flow_key = ?" + " ORDER BY start DESC LIMIT 1 OFFSET ?", + (run.workflow_id, run.flow_key, limit - 1), + ).fetchone() + if row is not None and row["start"] + window > now: + due_at = row["start"] + window + if gate.debounce is not None: + if active: + deferred = self._db.execute( + "UPDATE workflow_steps SET args = ?, due_at = ?," + " updated_at = ? WHERE run_id = ? AND ordinal = 0" + " AND status = ?", + ( + json.dumps(root_step.args), + now + gate.debounce, + now, + active[0]["run_id"], + StepStatus.READY.value, + ), + ) + if deferred.rowcount: + self._db.execute("COMMIT") + return FlowAdmission("coalesced", active[0]["run_id"]) + due_at = now + gate.debounce + if run.request_key is not None: + self._db.execute( + "INSERT INTO workflow_dedupe" + " (workflow_id, request_key, run_id) VALUES (?, ?, ?)", + (run.workflow_id, run.request_key, run.run_id), + ) + self._insert_run(run) + self._insert_step(dataclasses.replace(root_step, due_at=due_at)) + self._append_events(run.run_id, events, run.created_at) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return FlowAdmission("started", run.run_id, cancelled=tuple(cancelled)) + + return await asyncio.to_thread(work) + async def claim_next( self, now: float, diff --git a/tests/units/workflow/test_flow_atomicity.py b/tests/units/workflow/test_flow_atomicity.py new file mode 100644 index 00000000000..9879eab0c30 --- /dev/null +++ b/tests/units/workflow/test_flow_atomicity.py @@ -0,0 +1,239 @@ +"""Start policies must hold across processes, not just across tasks. + +Every policy is a read followed by a write. The old design guarded the gap +with an in-process asyncio lock, which meant two worker processes admitting +concurrently both read "nothing active yet" and both inserted: in an external +review, all five policy families violated their invariant in 50 out of 50 +synchronized trials across two OS processes and independent PostgreSQL pools. + +These tests open two *independent store instances* on one database -- two +pools, two connections, no shared Python state, which is what two processes +look like to the database -- and race gated admissions through both. The +invariants must hold because the store serializes the whole decision under a +durable per-key lock, not because anything in this process arranged politeness. +""" + +import asyncio +import dataclasses +import os +import uuid +from collections.abc import AsyncIterator + +import pytest + +from reflex.workflow.records import RunRecord, RunStatus, StepRecord, StepStatus +from reflex.workflow.store import FlowGate, RunStore, SqliteRunStore + +TRIALS = 10 +POSTGRES_ENV = "REFLEX_TEST_POSTGRES" + + +@pytest.fixture(params=["sqlite", "postgres"]) +async def store_pair(request, tmp_path) -> AsyncIterator[tuple[RunStore, RunStore]]: + """Two store instances over one database, like two worker processes. + + Args: + request: The backend under test. + tmp_path: Directory for the SQLite file. + + Yields: + Two independently connected stores sharing one database. + """ + if request.param == "sqlite": + path = tmp_path / "race.db" + a, b = SqliteRunStore(path), SqliteRunStore(path) + yield a, b + a.close() + b.close() + return + dsn = os.environ.get(POSTGRES_ENV) + if not dsn: + pytest.skip(f"{POSTGRES_ENV} is not configured") + from reflex.workflow.postgres import PostgresRunStore + from reflex.workflow.records import RunQuery + + schema = f"flowrace_{uuid.uuid4().hex[:12]}" + a = PostgresRunStore(dsn, schema=schema) + b = PostgresRunStore(dsn, schema=schema) + # Migration is lazy on first use; racing it from two fresh stores is a + # CREATE TABLE collision, not the admission race under test. + await a.list_runs(RunQuery(limit=1)) + yield a, b + await b.close() + await a.close() + a.drop_schema() + + +def _records(flow_key: str, args: dict | None = None) -> tuple[RunRecord, StepRecord]: + """Build one admission's records under a flow key. + + Args: + flow_key: The policy grouping key. + args: Root payload, when the test cares about it. + + Returns: + The run record and its root slot. + """ + run_id = uuid.uuid4().hex + now = 1_000_000.0 + run = RunRecord( + run_id=run_id, + workflow_id="race.flow", + definition_digest="digest", + status=RunStatus.PENDING, + state={}, + state_version=0, + next_ordinal=1, + flow_key=flow_key, + created_at=now, + updated_at=now, + ) + step = StepRecord( + run_id=run_id, + ordinal=0, + handler_id="start", + status=StepStatus.READY, + args=args or {}, + due_at=now, + origin="root", + queue="default", + created_at=now, + updated_at=now, + ) + return run, step + + +async def _race(a: RunStore, b: RunStore, gate: FlowGate, key: str, args=None) -> list: + """Admit through both stores at once. + + Args: + a: The first store instance. + b: The second store instance. + gate: The policy under test. + key: The flow key for this trial. + args: Optional distinct payloads for the two admissions. + + Returns: + Both admission outcomes. + """ + + async def admit(store: RunStore, payload): + """Run one gated admission. + + Args: + store: The store to admit through. + payload: The root payload. + + Returns: + The admission outcome. + """ + run, step = _records(key, payload) + return await store.admit_flow(run, step, (), gate, 1_000_000.0) + + first, second = args or ({}, {}) + return list(await asyncio.gather(admit(a, first), admit(b, second))) + + +async def test_singleton_skip_admits_exactly_one(store_pair): + """Two processes racing a singleton must not both win it.""" + a, b = store_pair + for trial in range(TRIALS): + outcomes = await _race(a, b, FlowGate(singleton_skip=True), f"skip-{trial}") + dispositions = sorted(o.disposition for o in outcomes) + assert dispositions == ["skipped", "started"], f"trial {trial}: {dispositions}" + started = next(o for o in outcomes if o.disposition == "started") + skipped = next(o for o in outcomes if o.disposition == "skipped") + assert skipped.run_id == started.run_id, ( + "the loser must be told which run holds the key" + ) + + +async def test_rate_limit_of_one_admits_exactly_one(store_pair): + """A limit of one is a limit of one from any number of processes.""" + a, b = store_pair + for trial in range(TRIALS): + outcomes = await _race(a, b, FlowGate(rate_limit=(1, 60.0)), f"rate-{trial}") + dispositions = sorted(o.disposition for o in outcomes) + assert dispositions == ["rejected", "started"], f"trial {trial}: {dispositions}" + + +async def test_debounce_coalesces_the_racing_start(store_pair): + """One of a racing pair is absorbed by the other, never two runs.""" + a, b = store_pair + for trial in range(TRIALS): + outcomes = await _race( + a, + b, + FlowGate(debounce=30.0), + f"deb-{trial}", + args=({"revision": 1}, {"revision": 2}), + ) + dispositions = sorted(o.disposition for o in outcomes) + assert dispositions == ["coalesced", "started"], ( + f"trial {trial}: {dispositions}" + ) + started = next(o for o in outcomes if o.disposition == "started") + coalesced = next(o for o in outcomes if o.disposition == "coalesced") + assert coalesced.run_id == started.run_id + # Latest wins: whichever admission lost the race replaced the payload. + steps = await a.get_steps(started.run_id) + assert steps[0].args in ({"revision": 1}, {"revision": 2}) + + +async def test_throttle_of_one_spaces_the_racing_pair(store_pair): + """Both start, but the second sits a full window after the first.""" + a, b = store_pair + for trial in range(TRIALS): + outcomes = await _race(a, b, FlowGate(throttle=(1, 60.0)), f"thr-{trial}") + assert all(o.disposition == "started" for o in outcomes) + dues = sorted([(await a.get_steps(o.run_id))[0].due_at for o in outcomes]) + assert dues[1] - dues[0] >= 60.0 - 1e-9, ( + f"trial {trial}: both roots due at {dues}; the burst was not spaced" + ) + + +async def test_singleton_cancel_leaves_at_most_one_uncancelled(store_pair): + """Racing replacements never leave two live runs on the key.""" + a, b = store_pair + for trial in range(TRIALS): + key = f"cxl-{trial}" + outcomes = await _race(a, b, FlowGate(singleton_cancel=True), key) + assert all(o.disposition == "started" for o in outcomes) + live = [ + o.run_id + for o in outcomes + if (run := await a.get_run(o.run_id)) is not None + and not run.cancel_requested + ] + assert len(live) <= 1, ( + f"trial {trial}: {len(live)} replacements survived on one key" + ) + + +async def test_gated_dedupe_is_atomic_across_instances(store_pair): + """A redelivered event admitted through the gate is still one run.""" + a, b = store_pair + for trial in range(TRIALS): + key = f"ded-{trial}" + + async def admit(store: RunStore, key: str = key): + """Admit one redelivery of the same event. + + Args: + store: The store to admit through. + key: The flow key for this trial. + + Returns: + The admission outcome. + """ + run, step = _records(key) + run = dataclasses.replace(run, request_key=f"evt-{key}") + return await store.admit_flow( + run, step, (), FlowGate(singleton_skip=True), 1_000_000.0 + ) + + outcomes = list(await asyncio.gather(admit(a), admit(b))) + dispositions = sorted(o.disposition for o in outcomes) + assert dispositions == ["deduplicated", "started"], ( + f"trial {trial}: {dispositions}" + ) diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index 37c6bb05283..ec139db7d52 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -774,3 +774,73 @@ async def test_cancelling_an_all_mode_parent_leaves_its_branches_alone( "cancelling the parent cancelled delegated children: " f"{[(c.run_id, c.status) for c in children]}" ) + + +async def test_a_policy_decorated_branch_is_refused(forked_registration_context): + """A policy fan-out would silently bypass is a policy refused loudly. + + Fan-out writes its branches in the parent's committing transaction, not + through policy admission, so a throttle on a branch root would simply not + apply -- five branches under a throttle of two all start at once and + nothing says so. Until branches go through the same admission primitive, + declaring both is an error that names the way out. + """ + from reflex_base.workflow import Throttle + + class Limited(rx.State): + __workflow__ = WorkflowConfig(id="fan.limited") + + @rx.event( + durable=True, + trigger=manual(), + effect="none", + throttle=Throttle(limit=2, period="10s"), + ) + def start(self, lead: str): + """A throttled root. + + Args: + lead: The lead identifier. + + Returns: + Completion. + """ + return rx.complete(result=lead) + + class FansOut(rx.State): + __workflow__ = WorkflowConfig(id="fan.bypasser") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Fan out to a policy-decorated root. + + Returns: + The refused fan-out. + """ + return rx.parallel( + Limited.start("a"), Limited.start("b"), then=FansOut.done + ) + + @rx.event(durable=True, effect="none") + def done(self, results: list): + """Collect. + + Args: + results: One entry per branch. + + Returns: + Completion. + """ + return rx.complete(result=len(results)) + + async with WorkflowTestHarness(FansOut, Limited) as harness: + started = await harness.start(FansOut.begin()) + assert started.run_id is not None + await harness.run_until_idle() + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED, ( + "a definition error fails on its first attempt rather than " + f"burning retries: {snapshot.status}" + ) + assert "start policy" in str(snapshot.error) From a188b425bed6e6aeccb4002d3f0557671da3fa02 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 14:05:12 -0700 Subject: [PATCH 088/121] workflows: retry and skip restore the chain the failure tombstoned The review's second P0. A handler returning a list preallocates an immediate sequential chain, so [middle(), finish()] creates both slots up front; middle's terminal failure tombstones finish, and retry or skip reopened only middle. The retried step succeeded, returned nothing to allocate, and the run completed COMPLETED/None having never run the finalizer -- an operator repair that silently dropped the cleanup step is worse than the failure it repaired. Both actions now restore every CANCELLED slot in the run, with fresh budgets, recorded as step_restored. The causality argument for restoring all of them rather than tracking which failure cancelled what: these actions accept only FAILED or NEEDS_ATTENTION runs, run-level cancellation ends in a CANCELLED run they refuse, force-finalization leaves no failed or suspended step for them to target, and suspension tombstones nothing -- so a CANCELLED slot in an acceptable run has exactly one possible source. No schema change needed. Fixed in memory, SQLite and Postgres; the strict xfail that pinned the defect is now the passing regression test, joined by the skip variant, and the operator suite passes against real Postgres 16. --- news/workflow-retry-restores-chain.bugfix.md | 1 + reflex/workflow/CONTRACT.md | 6 + reflex/workflow/postgres.py | 49 ++++++- reflex/workflow/records.py | 1 + reflex/workflow/store.py | 120 +++++++++++++++++- tests/units/workflow/test_operator_actions.py | 79 +++++++++--- 6 files changed, 235 insertions(+), 21 deletions(-) create mode 100644 news/workflow-retry-restores-chain.bugfix.md diff --git a/news/workflow-retry-restores-chain.bugfix.md b/news/workflow-retry-restores-chain.bugfix.md new file mode 100644 index 00000000000..acb32eda705 --- /dev/null +++ b/news/workflow-retry-restores-chain.bugfix.md @@ -0,0 +1 @@ +Operator `retry` and `skip` now restore the successors the failure tombstoned. A handler returning a list preallocates a sequential chain, and a terminal failure cancels every open slot behind it; reopening or skipping only the failed step let the run complete without ever running the finalizer the chain was written for. A `CANCELLED` slot in a run these actions accept can only be that failure's casualty — run-level cancellation ends in a `CANCELLED` run they refuse, and force-finalization leaves them nothing to target — so exactly those slots come back, with fresh budgets, recorded as `step_restored` in history. Fixed in all three stores. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 729e64e5f0e..23fbc0a1e67 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -339,6 +339,12 @@ no-op with a reason. step `SKIPPED` (terminal, recorded as a decision rather than an outcome) and lets the run continue at whatever comes next. With nothing left to run, the run completes with no result rather than sitting pending forever. +- Both restore the successors the stopping failure tombstoned (`step_restored` + in history, fresh budgets), so a preallocated chain's remaining steps — + including its finalizer — still run. Only that failure's casualties come + back: a `CANCELLED` slot in a run these actions accept can have no other + source, because run-level cancellation ends in a `CANCELLED` run they + refuse and force-finalization leaves them no step to target. - `force_complete(run, result)` / `force_fail(run, reason)` — a nonterminal, drained run: finalizes immediately, tombstoning open slots, recording the operator origin and (for completion) the result to treat it as having diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index e7e24d5d723..da6538e6a21 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1596,6 +1596,37 @@ async def resume_run(self, run_id: str, now: float) -> bool: ) return True + @staticmethod + async def _restore_tombstoned(conn: Any, run_id: str, now: float) -> list[int]: + """Re-open the slots the run's stopping failure tombstoned. + + Within a run an operator can retry or skip, a CANCELLED slot can only + be that failure's casualty: run-level cancellation ends in a + CANCELLED run these actions refuse, and force-finalization leaves no + failed or suspended step for them to target. + + Args: + conn: The open transaction's connection. + run_id: The run being re-opened. + now: Current time in epoch seconds. + + Returns: + The restored ordinals, in order. + """ + cursor = await conn.execute( + "UPDATE workflow_steps SET status = %s, attempts = 0, due_at = %s," + " lease_expires_at = 0, error = NULL, updated_at = %s" + " WHERE run_id = %s AND status = %s RETURNING ordinal", + ( + StepStatus.READY.value, + now, + now, + run_id, + StepStatus.CANCELLED.value, + ), + ) + return sorted(row["ordinal"] for row in await cursor.fetchall()) + async def retry_run(self, run_id: str, now: float) -> bool: """Re-open a failed run at the step that failed. @@ -1628,10 +1659,17 @@ async def retry_run(self, run_id: str, now: float) -> bool: " WHERE run_id = %s AND ordinal = %s", (StepStatus.READY.value, now, now, run_id, row["ordinal"]), ) + restored = await self._restore_tombstoned(conn, run_id, now) await self._append_events( conn, run_id, - ((HistoryEventType.RUN_RESUMED, {"origin": "retry"}),), + ( + (HistoryEventType.RUN_RESUMED, {"origin": "retry"}), + *( + (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) + for ordinal in restored + ), + ), now, ) return True @@ -1672,6 +1710,7 @@ async def skip_step(self, run_id: str, now: float) -> bool: " updated_at = %s WHERE run_id = %s AND ordinal = %s", (StepStatus.SKIPPED.value, now, run_id, row["ordinal"]), ) + restored = await self._restore_tombstoned(conn, run_id, now) cursor = await conn.execute( "SELECT 1 FROM workflow_steps WHERE run_id = %s" " AND NOT (status = ANY(%s)) LIMIT 1", @@ -1687,7 +1726,13 @@ async def skip_step(self, run_id: str, now: float) -> bool: run_id, ), ) - events = [(HistoryEventType.STEP_SKIPPED, {"ordinal": row["ordinal"]})] + events = [ + (HistoryEventType.STEP_SKIPPED, {"ordinal": row["ordinal"]}), + *( + (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) + for ordinal in restored + ), + ] if not open_left: events.append((HistoryEventType.RUN_COMPLETED, {})) await self._append_events(conn, run_id, tuple(events), now) diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 280553740e0..829c3db307b 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -147,6 +147,7 @@ class HistoryEventType(str, enum.Enum): STEP_RETRY_SCHEDULED = "step_retry_scheduled" STEP_RECOVERED = "step_recovered" STEP_TOMBSTONED = "step_tombstoned" + STEP_RESTORED = "step_restored" RUN_COMPLETED = "run_completed" RUN_FAILED = "run_failed" RUN_TIMED_OUT = "run_timed_out" diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index bc187d2d37c..9c93b918e68 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -1728,6 +1728,7 @@ async def skip_step(self, run_id: str, now: float) -> bool: lease_expires_at=0.0, updated_at=now, ) + restored = self._restore_tombstoned(steps, now) # Skipping the last open slot leaves nothing to run, so # the run is finished rather than pending forever -- being # stuck in a new way is not a resolution. @@ -1735,7 +1736,11 @@ async def skip_step(self, run_id: str, now: float) -> bool: other.status not in TERMINAL_STEP_STATUSES for other in steps ) events = [ - (HistoryEventType.STEP_SKIPPED, {"ordinal": step.ordinal}) + (HistoryEventType.STEP_SKIPPED, {"ordinal": step.ordinal}), + *( + (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) + for ordinal in restored + ), ] if not open_left: events.append((HistoryEventType.RUN_COMPLETED, {})) @@ -1784,14 +1789,61 @@ async def retry_run(self, run_id: str, now: float) -> bool: break if not reopened: return False + restored = self._restore_tombstoned(steps, now) self._runs[run_id] = dataclasses.replace( run, status=RunStatus.PENDING, error=None, updated_at=now ) self._append_events( - run_id, ((HistoryEventType.RUN_RESUMED, {"origin": "retry"}),), now + run_id, + ( + (HistoryEventType.RUN_RESUMED, {"origin": "retry"}), + *( + (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) + for ordinal in restored + ), + ), + now, ) return True + @staticmethod + def _restore_tombstoned(steps: list[StepRecord], now: float) -> list[int]: + """Re-open the slots a run's stopping failure tombstoned. + + A terminal failure cancels every open slot behind it, so a + preallocated chain's finalizer is CANCELLED the moment its + predecessor fails. Retrying or skipping that predecessor promises to + continue "from there" -- which is a lie unless the chain comes back. + + Within a run an operator can retry or skip -- FAILED or + NEEDS_ATTENTION -- a CANCELLED slot can only be that failure's + casualty: run-level cancellation ends in a CANCELLED run these + actions refuse, and operator force-finalization leaves no failed or + suspended step for them to target. Nothing independently cancelled is + revived because nothing independently cancelled can be here. + + Args: + steps: The run's mailbox, mutated in place. + now: Current time in epoch seconds. + + Returns: + The restored ordinals, in order. + """ + restored: list[int] = [] + for index, step in enumerate(steps): + if step.status is StepStatus.CANCELLED: + steps[index] = dataclasses.replace( + step, + status=StepStatus.READY, + attempts=0, + due_at=now, + lease_expires_at=0.0, + error=None, + updated_at=now, + ) + restored.append(step.ordinal) + return restored + async def resume_run(self, run_id: str, now: float) -> bool: """Re-open a suspended run so its frontier step runs again. @@ -3808,6 +3860,31 @@ def work() -> bool: " updated_at = ? WHERE run_id = ? AND ordinal = ?", (StepStatus.SKIPPED.value, now, run_id, row["ordinal"]), ) + restored = [ + r["ordinal"] + for r in self._db.execute( + "SELECT ordinal FROM workflow_steps WHERE run_id = ?" + " AND status = ? ORDER BY ordinal", + (run_id, StepStatus.CANCELLED.value), + ).fetchall() + ] + if restored: + # The stopping failure tombstoned these; continuing + # "from there" is a lie unless they come back. Nothing + # independently cancelled can be in a run these + # actions accept (see MemoryRunStore._restore_tombstoned). + self._db.execute( + "UPDATE workflow_steps SET status = ?, attempts = 0," + " due_at = ?, lease_expires_at = 0, error = NULL," + " updated_at = ? WHERE run_id = ? AND status = ?", + ( + StepStatus.READY.value, + now, + now, + run_id, + StepStatus.CANCELLED.value, + ), + ) terminal = tuple(s.value for s in TERMINAL_STEP_STATUSES) open_left = self._db.execute( "SELECT 1 FROM workflow_steps WHERE run_id = ?" @@ -3827,7 +3904,11 @@ def work() -> bool: ), ) events = [ - (HistoryEventType.STEP_SKIPPED, {"ordinal": row["ordinal"]}) + (HistoryEventType.STEP_SKIPPED, {"ordinal": row["ordinal"]}), + *( + (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) + for ordinal in restored + ), ] if not open_left: events.append((HistoryEventType.RUN_COMPLETED, {})) @@ -3883,9 +3964,40 @@ def work() -> bool: " updated_at = ? WHERE run_id = ? AND ordinal = ?", (StepStatus.READY.value, now, now, run_id, row["ordinal"]), ) + restored = [ + r["ordinal"] + for r in self._db.execute( + "SELECT ordinal FROM workflow_steps WHERE run_id = ?" + " AND status = ? ORDER BY ordinal", + (run_id, StepStatus.CANCELLED.value), + ).fetchall() + ] + if restored: + # The stopping failure tombstoned these; continuing + # "from there" is a lie unless they come back. Nothing + # independently cancelled can be in a run these + # actions accept (see MemoryRunStore._restore_tombstoned). + self._db.execute( + "UPDATE workflow_steps SET status = ?, attempts = 0," + " due_at = ?, lease_expires_at = 0, error = NULL," + " updated_at = ? WHERE run_id = ? AND status = ?", + ( + StepStatus.READY.value, + now, + now, + run_id, + StepStatus.CANCELLED.value, + ), + ) self._append_events( run_id, - ((HistoryEventType.RUN_RESUMED, {"origin": "retry"}),), + ( + (HistoryEventType.RUN_RESUMED, {"origin": "retry"}), + *( + (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) + for ordinal in restored + ), + ), now, ) self._db.execute("COMMIT") diff --git a/tests/units/workflow/test_operator_actions.py b/tests/units/workflow/test_operator_actions.py index 4ed57e51e05..0d929fde76b 100644 --- a/tests/units/workflow/test_operator_actions.py +++ b/tests/units/workflow/test_operator_actions.py @@ -6,7 +6,6 @@ trying to rescue. """ -import pytest from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual import reflex as rx @@ -311,26 +310,15 @@ async def test_skip_is_refused_on_a_healthy_run(forked_registration_context): assert not await rx.workflows.skip("no-such-run") -@pytest.mark.xfail( - strict=True, - reason=( - "Known defect: a returned list preallocates a sequential chain, so a " - "step failing tombstones the rest of it. retry() reopens the failed " - "step but does not restore the successors that failure cancelled, so " - "the run completes having skipped them. Fixing it means restoring " - "steps tombstoned by that failure in retry_run/skip_step across all " - "three stores." - ), -) async def test_retry_restores_the_chain_the_failure_tombstoned( forked_registration_context, ): """Retrying continues from the failed step, not past everything after it. A handler returning a list preallocates the whole chain, so a terminal - failure cancels the steps behind it. An operator retrying expects the run - to carry on from there -- including the finalizer that was already - allocated and is now cancelled. + failure tombstones the steps behind it. Retry restores exactly those -- + a CANCELLED slot in a FAILED run can only be that failure's casualty -- + so the finalizer the chain was written for still runs. """ attempts: list[str] = [] @@ -385,3 +373,64 @@ def finish(self): "the retried run completed without running the finalizer the " f"chain preallocated: {[step.status.value for step in snapshot.steps]}" ) + + +async def test_skip_restores_the_chain_the_failure_tombstoned( + forked_registration_context, +): + """Skipping the failed step continues at what comes next, not at nothing. + + Same chain, same failure; the operator gives up on the middle step + instead of retrying it. The preallocated finalizer must still run -- + without restoration the run completes with the skip as its last word and + the finalizer silently cancelled. + """ + + class SkipChain(rx.State): + __workflow__ = WorkflowConfig(id="ops.skipchain") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Preallocate the rest of the chain. + + Returns: + Two successors, run in order. + """ + return [SkipChain.middle(), SkipChain.finish()] + + @rx.event(durable=True, effect="none", retry=Retry(max_attempts=1)) + def middle(self): + """Fail terminally. + + Raises: + ValueError: Always. + """ + msg = "a vendor that retired its endpoint" + raise ValueError(msg) + + @rx.event(durable=True, effect="none") + def finish(self): + """The finalizer the chain exists for. + + Returns: + Completion. + """ + return rx.complete(result="finished") + + async with WorkflowTestHarness(SkipChain) as harness: + started = await harness.start(SkipChain.begin()) + assert started.run_id is not None + await harness.run_until_idle() + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + + assert await harness.kernel.skip(started.run_id) + await harness.run_until_idle() + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.result == "finished", ( + "the skipped run completed without running the finalizer: " + f"{[step.status.value for step in snapshot.steps]}" + ) From d115684f345539f0d3c8c3fbffc479e0532024a0 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 14:30:11 -0700 Subject: [PATCH 089/121] workflows: answer SQLite scheduling from an index, not a Python loop The review's SQLite scaling cliff. claim_next selected every active run and loaded each run's steps until it found a claimable frontier; next_due did the same with one more query per run. Ten thousand durable timers -- the thing a workflow engine exists to hold -- meant ~114ms per idle poll, a worker burning a third to half a core doing nothing, and burst drains that looked quadratic (5,000 roots: 188s). Both queries are now one SQL statement: filter to steps whose status and due time can wake -- through a new (status, due_at, queue) index -- then pay the frontier check (NOT EXISTS a lower unresolved ordinal, a primary-key lookup) only for those candidates, and stop at the first winner. A store full of sleepers answers from the index range and touches nothing: measured on the reviewer's 10k-sleeper shape, next_due 61.8ms -> 4.8ms and claim_next 88.1ms -> 0.05ms, both under their proposed 5ms release gate. Claim order among concurrently-due runs changes from run creation time to due time, which is scheduling fairness the contract does not promise and arguably the better order -- oldest-due first. A plan-pinning test asserts EXPLAIN QUERY PLAN uses the wake index for both query shapes, because a timing assertion flakes on CI and the plan is the actual performance contract. The full workflow suite, conformance included, passes on all three stores with Postgres real. --- news/workflow-sqlite-frontier.perf.md | 1 + reflex/workflow/store.py | 146 ++++++++++++++----- tests/units/workflow/test_sqlite_frontier.py | 101 +++++++++++++ 3 files changed, 210 insertions(+), 38 deletions(-) create mode 100644 news/workflow-sqlite-frontier.perf.md create mode 100644 tests/units/workflow/test_sqlite_frontier.py diff --git a/news/workflow-sqlite-frontier.perf.md b/news/workflow-sqlite-frontier.perf.md new file mode 100644 index 00000000000..ec22c266271 --- /dev/null +++ b/news/workflow-sqlite-frontier.perf.md @@ -0,0 +1 @@ +The SQLite store's scheduler queries — `claim_next` and `next_due` — are now expressed in SQL over a new `(status, due_at, queue)` index instead of loading every active run and walking its steps in Python. An idle worker's poll was linear in the number of sleeping runs (an external review measured ~114ms per `next_due` and 31–50% CPU doing nothing at 10,000 durable timers, with quadratic-looking burst drains); the same 10,000-sleeper setup now answers `next_due` in ~4.8ms and `claim_next` in ~0.05ms. A test pins the query plan to the index, so the performance contract cannot silently rot. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 9c93b918e68..de2a80be6bb 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -30,6 +30,7 @@ from reflex_base.workflow import DEFAULT_LEASE_DURATION from reflex.workflow.records import ( + CLAIMABLE_STEP_STATUSES, TERMINAL_RUN_STATUSES, TERMINAL_STEP_STATUSES, HistoryEvent, @@ -2299,6 +2300,8 @@ def resolve_store(target: str | None = None) -> RunStore: PRIMARY KEY (run_id, wait_key, dedupe_key) ); CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs (status); +CREATE INDEX IF NOT EXISTS idx_workflow_steps_wake + ON workflow_steps (status, due_at, queue); CREATE INDEX IF NOT EXISTS idx_workflow_inbox_pending ON workflow_inbox (run_id, wait_key, status, seq); """ @@ -2416,6 +2419,77 @@ def _step_from_row(row: sqlite3.Row) -> StepRecord: ) +def _sqlite_frontier_query( + select: str, + now: float, + queues: tuple[str, ...] | None, + *, + due_only: bool, + order: str, + limit: int, +) -> tuple[str, tuple[Any, ...]]: + """Build the query for claimable-or-waking frontier steps. + + The shape both the claimer and the sleep bound need: steps whose status + can wake, on runnable runs, that are their run's frontier -- the lowest + unresolved ordinal -- filtered and ordered so the wake index does the + work and a LIMIT stops the scan. Loading every run and walking its steps + in Python made an idle worker's poll linear in the number of sleeping + runs, which at ten thousand of them was ~114ms per poll and half a core + doing nothing. + + Args: + select: The select list, over alias ``s``. + now: Current time in epoch seconds. + queues: Queues served; None serves every queue. + due_only: Restrict to steps claimable right now; otherwise any step a + clock event alone can make claimable, however far out. + order: ORDER BY expression. + limit: Maximum rows. + + Returns: + The SQL and its parameters. + """ + claimable = tuple(s.value for s in CLAIMABLE_STEP_STATUSES) + terminal_steps = tuple(s.value for s in TERMINAL_STEP_STATUSES) + terminal_runs = tuple(s.value for s in TERMINAL_RUN_STATUSES) + marks = lambda values: ",".join("?" * len(values)) # noqa: E731 + if due_only: + waking = ( + f"((s.status IN ({marks(claimable)}) AND s.due_at <= ?)" + " OR (s.status = ? AND s.due_at > 0 AND s.due_at <= ?))" + ) + waking_params = (*claimable, now, StepStatus.BLOCKED.value, now) + else: + waking = ( + f"(s.status IN ({marks(claimable)}) OR (s.status = ? AND s.due_at > 0))" + ) + waking_params = (*claimable, StepStatus.BLOCKED.value) + queue_sql = f" AND s.queue IN ({marks(queues)})" if queues is not None else "" + queue_params = tuple(queues) if queues is not None else () + sql = ( + f"SELECT {select} FROM workflow_steps s" + " JOIN workflow_runs r ON r.run_id = s.run_id" + f" WHERE {waking}" + f" AND r.status NOT IN ({marks(terminal_runs)})" + " AND r.status != ? AND r.cancel_requested = 0" + " AND (r.deadline IS NULL OR r.deadline > ?)" + " AND NOT EXISTS (SELECT 1 FROM workflow_steps x" + " WHERE x.run_id = s.run_id AND x.ordinal < s.ordinal" + f" AND x.status NOT IN ({marks(terminal_steps)}))" + f"{queue_sql} ORDER BY {order} LIMIT {int(limit)}" + ) + params = ( + *waking_params, + *terminal_runs, + RunStatus.NEEDS_ATTENTION.value, + now, + *terminal_steps, + *queue_params, + ) + return sql, params + + def _sqlite_run_filters(query: RunQuery) -> tuple[str, tuple[Any, ...]]: """Build the WHERE clause a run query means, for SQLite. @@ -2835,26 +2909,31 @@ def work(): Returns: The operation's result. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) with self._lock: self._db.execute("BEGIN IMMEDIATE") claim = None try: - rows = self._db.execute( - "SELECT * FROM workflow_runs WHERE status NOT IN" - f" ({','.join('?' * len(terminal))})" - " AND status != ? AND cancel_requested = 0" - " AND (deadline IS NULL OR deadline > ?)" - " ORDER BY created_at", - (*terminal, RunStatus.NEEDS_ATTENTION.value, now), - ).fetchall() - for row in rows: - run = _run_from_row(row) - frontier = _frontier(self._load_steps(run.run_id)) - if frontier is None or not step_claimable_at(frontier, now): - continue - if queues is not None and frontier.queue not in queues: - continue + # Due-ness filters first, through the wake index, so a + # store full of sleeping runs answers from the index + # instead of loading every run and its steps into Python; + # only due candidates pay the frontier check, and LIMIT + # stops the scan at the first winner. + sql, params = _sqlite_frontier_query( + "s.*", + now, + queues, + due_only=True, + order="s.due_at, s.run_id", + limit=1, + ) + row = self._db.execute(sql, params).fetchone() + if row is not None: + frontier = _step_from_row(row) + run_row = self._db.execute( + "SELECT * FROM workflow_runs WHERE run_id = ?", + (frontier.run_id,), + ).fetchone() + run = _run_from_row(run_row) claimed = dataclasses.replace( frontier, status=StepStatus.CLAIMED, @@ -2884,7 +2963,6 @@ def work(): run, status=RunStatus.RUNNING, updated_at=now ) claim = Claim(run=running, step=claimed) - break self._db.execute("COMMIT" if claim is not None else "ROLLBACK") except BaseException: self._db.execute("ROLLBACK") @@ -4554,27 +4632,19 @@ def work(): Returns: The operation's result. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) with self._lock: - rows = self._db.execute( - "SELECT run_id FROM workflow_runs WHERE status NOT IN" - f" ({','.join('?' * len(terminal))})" - " AND status != ? AND cancel_requested = 0" - " AND (deadline IS NULL OR deadline > ?)", - (*terminal, RunStatus.NEEDS_ATTENTION.value, now), - ).fetchall() - due_times = [] - for row in rows: - frontier = _frontier(self._load_steps(row["run_id"])) - if ( - frontier is not None - and queues is not None - and frontier.queue not in queues - ): - continue - wake_at = None if frontier is None else step_wake_at(frontier) - if wake_at is not None: - due_times.append(wake_at) - return min(due_times) if due_times else None + # Wake times are due_at in every waking status, so the first + # row of a due_at-ordered index walk that survives the + # frontier check is the minimum -- no per-run Python loop. + sql, params = _sqlite_frontier_query( + "s.due_at AS wake", + now, + queues, + due_only=False, + order="s.due_at", + limit=1, + ) + row = self._db.execute(sql, params).fetchone() + return None if row is None else row["wake"] return await asyncio.to_thread(work) diff --git a/tests/units/workflow/test_sqlite_frontier.py b/tests/units/workflow/test_sqlite_frontier.py new file mode 100644 index 00000000000..1bdaf417605 --- /dev/null +++ b/tests/units/workflow/test_sqlite_frontier.py @@ -0,0 +1,101 @@ +"""The SQLite scheduler surface must answer from an index, not a Python loop. + +An idle worker polls ``next_due`` and ``claim_next`` continuously. Loading +every active run and walking its steps in Python made both linear in the +number of sleeping runs: at ten thousand durable timers an external review +measured ~114ms per poll, 31-50% CPU on a worker with nothing to do, and a +quadratic-looking burst drain. The queries now filter by wake-ability through +``idx_workflow_steps_wake`` first and pay the frontier check only for +candidates, which the same setup measures at well under five milliseconds. +""" + +import uuid + +from reflex.workflow.records import RunRecord, RunStatus, StepRecord, StepStatus +from reflex.workflow.store import SqliteRunStore, _sqlite_frontier_query + +NOW = 1_000_000.0 + + +def _sleeper(index: int, due_at: float, status: StepStatus = StepStatus.READY): + """Build one waiting run and its frontier slot. + + Args: + index: Uniquifies the run. + due_at: When the frontier comes due. + status: The frontier's status. + + Returns: + The run record and its step. + """ + run_id = f"run{index:05d}{uuid.uuid4().hex[:6]}" + run = RunRecord( + run_id=run_id, + workflow_id="frontier.bench", + definition_digest="d", + status=RunStatus.WAITING, + state={}, + state_version=1, + next_ordinal=2, + created_at=NOW + index, + updated_at=NOW, + ) + step = StepRecord( + run_id=run_id, + ordinal=1, + handler_id="wake", + status=status, + args={}, + due_at=due_at, + origin="root", + queue="default", + created_at=NOW, + updated_at=NOW, + ) + return run, step + + +async def test_the_frontier_queries_use_the_wake_index(tmp_path): + """The plan is the performance contract; pin it so it cannot rot. + + A timing assertion flakes on shared CI; the query plan does not. If + either query stops using the wake index, an idle worker is back to + scanning every sleeping run per poll. + """ + store = SqliteRunStore(tmp_path / "plan.db") + run, step = _sleeper(0, NOW + 86_400) + await store.admit(run, step, ()) + for due_only in (True, False): + sql, params = _sqlite_frontier_query( + "s.*", NOW, None, due_only=due_only, order="s.due_at", limit=1 + ) + plan = " ".join( + row["detail"] + for row in store._db.execute( # pyright: ignore[reportPrivateUsage] + f"EXPLAIN QUERY PLAN {sql}", params + ).fetchall() + ) + assert "idx_workflow_steps_wake" in plan, plan + store.close() + + +async def test_a_due_run_is_found_among_ten_thousand_sleepers(tmp_path): + """Scale changes the cost, never the answer.""" + store = SqliteRunStore(tmp_path / "mixed.db") + for index in range(500): + run, step = _sleeper(index, NOW + 86_400 + index) + await store.admit(run, step, ()) + due_run, due_step = _sleeper(9_999, NOW - 5) + await store.admit(due_run, due_step, ()) + + assert await store.next_due(NOW) is not None + claim = await store.claim_next(NOW) + assert claim is not None + assert claim.run.run_id == due_run.run_id, ( + "the one due run must be claimed, not any of the sleepers" + ) + assert await store.claim_next(NOW) is None, "nothing else is due" + due = await store.next_due(NOW) + assert due is not None + assert abs(due - (NOW + 86_400)) < 1e-6, "the earliest sleeper is the next wake" + store.close() From 3bcabd480cde17043895d185b7198344db10e49c Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 14:36:02 -0700 Subject: [PATCH 090/121] workflows: give webhook identity the shapes providers actually use Three review findings, one boundary. dedupe_by could only name a payload field, and GitHub's canonical delivery identity is the X-GitHub-Delivery header -- it appears nowhere in the body. Keyed on a payload field, two distinct deliveries sharing that field collapsed into one run, and with no dedupe_by at all a true redelivery executed twice. dedupe_by="header:Name" reads the header; distinct GUIDs are distinct runs and a repeated GUID deduplicates, pinned by test. A configured identity that could not be extracted silently disabled deduplication -- the request was admitted with no key, so every redelivery of that event ran again, defeating the exact thing the configuration asked for, invisibly. That is a 400 now, naming the missing source, so it surfaces as the config problem it is. rx.hmac_signature claimed to cover Stripe and cannot: Stripe signs timestamp.body, sends a structured header, and requires a replay window. The claim is corrected -- with a test that the docstring stays honest -- and rx.stripe_signature() implements the real scheme: t=/v1= parsing, multiple v1 digests for secret rotation, tolerance enforced both directions, constant-time comparison. --- news/workflow-webhook-identity.feature.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 90 +++++++++++++- reflex/__init__.py | 1 + reflex/workflow/__init__.py | 2 + reflex/workflow/ingress.py | 67 +++++++++-- tests/units/reflex_base/test_workflow.py | 75 ++++++++++++ tests/units/workflow/test_ingress.py | 110 ++++++++++++++++++ 7 files changed, 335 insertions(+), 11 deletions(-) create mode 100644 news/workflow-webhook-identity.feature.md diff --git a/news/workflow-webhook-identity.feature.md b/news/workflow-webhook-identity.feature.md new file mode 100644 index 00000000000..b13ec77fe6f --- /dev/null +++ b/news/workflow-webhook-identity.feature.md @@ -0,0 +1 @@ +Webhook delivery identity now covers the providers as they actually behave. `dedupe_by="header:X-GitHub-Delivery"` reads the identity from a request header — GitHub's canonical delivery GUID appears in no payload field, so payload-keyed dedupe collapsed distinct events and double-ran true redeliveries. A delivery whose configured identity cannot be extracted is refused with a 400 naming the missing source, instead of silently admitting with deduplication disabled. `rx.stripe_signature(secret_env=..., tolerance="5m")` implements Stripe's actual scheme — HMAC over `timestamp.body`, the structured `Stripe-Signature` header, multiple `v1` digests during secret rotation, and a replay window — and `rx.hmac_signature`'s docstring no longer claims to cover Stripe, which it never did. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 0abdea3a754..bdeace3084b 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -12,6 +12,7 @@ import hmac import os import re +import time from collections.abc import Callable, Mapping from datetime import timedelta from typing import Any, ClassVar, Final, Literal, get_args @@ -402,6 +403,85 @@ def __call__(self, body: bytes, headers: Mapping[str, str]) -> bool: return hmac.compare_digest(f"{self.prefix}{expected}", presented) +@dataclasses.dataclass(frozen=True) +class StripeVerifier: + """A webhook verifier implementing Stripe's signature scheme. + + Stripe does not sign the raw body: it signs ``"{timestamp}.{body}"`` and + sends ``Stripe-Signature: t=,v1=[,v1=...]``, and its + documentation requires rejecting timestamps outside a tolerance window so + a captured delivery cannot be replayed later. A raw-body HMAC verifier + accepts none of this, which is why this exists as its own type. + + Attributes: + secret_env: Name of the environment variable holding the signing + secret (``whsec_...``). + tolerance: Replay window in seconds; deliveries whose signed + timestamp is further than this from now are refused. + header: Request header carrying the signature. + """ + + secret_env: str + tolerance: float = 300.0 + header: str = "Stripe-Signature" + + def __call__(self, body: bytes, headers: Mapping[str, str]) -> bool: + """Check one delivery's signature and replay window. + + Args: + body: The raw request body, exactly as received. + headers: The request headers. + + Returns: + True when a presented ``v1`` digest matches the timestamped + payload and the timestamp is inside the tolerance window. + """ + secret = os.environ.get(self.secret_env) + presented = headers.get(self.header.lower()) or headers.get(self.header) + if not secret or not presented: + return False + timestamp: str | None = None + digests: list[str] = [] + for part in presented.split(","): + name, _, value = part.strip().partition("=") + if name == "t": + timestamp = value + elif name == "v1": + digests.append(value) + if timestamp is None or not digests: + return False + try: + signed_at = float(timestamp) + except ValueError: + return False + if abs(time.time() - signed_at) > self.tolerance: + return False + expected = hmac.new( + secret.encode(), f"{timestamp}.".encode() + body, "sha256" + ).hexdigest() + return any(hmac.compare_digest(expected, digest) for digest in digests) + + +def stripe_signature( + *, secret_env: str, tolerance: DurationLike = "5m" +) -> StripeVerifier: + """Build a verifier for Stripe's timestamped webhook signatures. + + Args: + secret_env: Name of the environment variable holding the signing + secret (``whsec_...``). + tolerance: Replay window; Stripe's documentation recommends five + minutes. + + Returns: + A verifier callable for ``rx.webhook(verify=...)``. + """ + return StripeVerifier( + secret_env=secret_env, + tolerance=parse_duration(tolerance, param="tolerance"), + ) + + def hmac_signature( *, secret_env: str, @@ -411,10 +491,12 @@ def hmac_signature( ) -> HmacVerifier: """Build a verifier for providers that HMAC-sign the raw request body. - This covers the common shape used by Stripe, GitHub, Shopify and others: - the provider sends a hex digest of the body keyed by a shared secret. The - secret is read from the environment at request time, so it never enters - workflow state, history, or a browser bundle. + This covers GitHub (``prefix="sha256="``), Shopify, and every provider + that sends a hex digest of the exact body keyed by a shared secret. It is + deliberately **not** a Stripe verifier: Stripe signs a timestamped payload + and requires a replay window -- use ``rx.stripe_signature()`` for that. + The secret is read from the environment at request time, so it never + enters workflow state, history, or a browser bundle. Args: secret_env: Name of the environment variable holding the shared secret. diff --git a/reflex/__init__.py b/reflex/__init__.py index 1b44550a6dc..9e7841bc18b 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -243,6 +243,7 @@ "manual", "webhook", "hmac_signature", + "stripe_signature", "schedule", "Signal", "wait_for", diff --git a/reflex/workflow/__init__.py b/reflex/workflow/__init__.py index 4a412eb5838..a3933de1233 100644 --- a/reflex/workflow/__init__.py +++ b/reflex/workflow/__init__.py @@ -36,6 +36,7 @@ parallel, parse_duration, schedule, + stripe_signature, wait_for, webhook, ) @@ -136,6 +137,7 @@ "parse_duration", "schedule", "step", + "stripe_signature", "substep_results", "wait_for", "webhook", diff --git a/reflex/workflow/ingress.py b/reflex/workflow/ingress.py index d6a21b46ce6..4ac5da4c4b7 100644 --- a/reflex/workflow/ingress.py +++ b/reflex/workflow/ingress.py @@ -97,8 +97,40 @@ def collect_webhook_routes( return routes +def _identity_value( + trigger: WebhookTrigger, payload: Any, headers: Mapping[str, str] +) -> str | None: + """Extract the provider's delivery identity for one request. + + ``dedupe_by`` names a payload field, or a header when written as + ``"header:Name"`` -- GitHub's canonical identity, for example, is the + ``X-GitHub-Delivery`` header and appears nowhere in the body. + + Args: + trigger: The webhook trigger declaring the identity source. + payload: The decoded request payload. + headers: The request headers. + + Returns: + The identity, or None when the declared source is absent. + """ + assert trigger.dedupe_by is not None + source = trigger.dedupe_by + if source.startswith("header:"): + name = source[len("header:") :] + value = headers.get(name.lower()) or headers.get(name) + return None if value is None else str(value) + if not isinstance(payload, dict): + return None + value = payload.get(source) + return None if value is None else str(value) + + def _dedupe_key( - handler: HandlerDefinition, trigger: WebhookTrigger, payload: Any + handler: HandlerDefinition, + trigger: WebhookTrigger, + payload: Any, + headers: Mapping[str, str], ) -> str | None: """Extract the deduplication key a provider redelivery would repeat. @@ -111,16 +143,17 @@ def _dedupe_key( Args: handler: The root handler this delivery starts. - trigger: The webhook trigger declaring the key field. + trigger: The webhook trigger declaring the key source. payload: The decoded request payload. + headers: The request headers. Returns: The key as a string, or None when the trigger declares none or the - field is absent. + source is absent. """ - if trigger.dedupe_by is None or not isinstance(payload, dict): + if trigger.dedupe_by is None: return None - value = payload.get(trigger.dedupe_by) + value = _identity_value(trigger, payload, headers) return None if value is None else f"webhook:{handler.id}:{value}" @@ -138,7 +171,12 @@ def _legacy_dedupe_keys(trigger: WebhookTrigger, payload: Any) -> tuple[str, ... Returns: The older keys to match, newest spelling first. """ - if trigger.dedupe_by is None or not isinstance(payload, dict): + if ( + trigger.dedupe_by is None + or trigger.dedupe_by.startswith("header:") + or not isinstance(payload, dict) + ): + # Header identities are new; no release ever wrote them unqualified. return () value = payload.get(trigger.dedupe_by) return () if value is None else (str(value),) @@ -217,10 +255,25 @@ async def endpoint(request: Request) -> JSONResponse: return JSONResponse( {"error": "payload must be a JSON object"}, status_code=400 ) + request_key = _dedupe_key(route.handler, route.trigger, payload, headers) + if route.trigger.dedupe_by is not None and request_key is None: + # A configured identity that cannot be extracted must not + # silently disable deduplication: every redelivery of this event + # would then execute again. The provider is told what is missing + # so its operator sees a config problem, not a duplicate charge. + return JSONResponse( + { + "error": ( + f"delivery carries no {route.trigger.dedupe_by!r}, " + "which this webhook deduplicates by" + ) + }, + status_code=400, + ) args = _root_args(route.handler, payload) result = await runtime.kernel.start( spec(**args) if args else spec, - request_key=_dedupe_key(route.handler, route.trigger, payload), + request_key=request_key, superseded_keys=_legacy_dedupe_keys(route.trigger, payload), trigger_kind="webhook", ) diff --git a/tests/units/reflex_base/test_workflow.py b/tests/units/reflex_base/test_workflow.py index 2454e367cb7..01417a3bc2c 100644 --- a/tests/units/reflex_base/test_workflow.py +++ b/tests/units/reflex_base/test_workflow.py @@ -1,6 +1,8 @@ """Tests for the workflow authoring value types.""" import datetime +import hmac +import time import pytest from reflex_base.utils.exceptions import WorkflowDefinitionError @@ -16,6 +18,7 @@ manual, parse_duration, schedule, + stripe_signature, webhook, ) @@ -261,3 +264,75 @@ def test_webhook_requires_authentication(): allow_unverified=True, unverified_reason="mixed", ) + + +def _stripe_header(secret: str, body: bytes, timestamp: float) -> str: + """Sign a body the way Stripe does. + + Args: + secret: The signing secret. + body: The raw request body. + timestamp: The signing time in epoch seconds. + + Returns: + The Stripe-Signature header value. + """ + signed = f"{int(timestamp)}.".encode() + body + digest = hmac.new(secret.encode(), signed, "sha256").hexdigest() + return f"t={int(timestamp)},v1={digest}" + + +def test_stripe_signature_accepts_a_fresh_correctly_signed_delivery(monkeypatch): + """The documented scheme: HMAC over timestamp-dot-body, inside tolerance.""" + monkeypatch.setenv("WH", "whsec_test") + verify = stripe_signature(secret_env="WH", tolerance="5m") + body = b'{"type": "invoice.paid"}' + header = _stripe_header("whsec_test", body, time.time()) + assert verify(body, {"stripe-signature": header}) + + +def test_stripe_signature_rejects_a_replayed_delivery(monkeypatch): + """A signature is only as good as its window; an old one is a replay.""" + monkeypatch.setenv("WH", "whsec_test") + verify = stripe_signature(secret_env="WH", tolerance="5m") + body = b"{}" + stale = _stripe_header("whsec_test", body, time.time() - 3600) + assert not verify(body, {"stripe-signature": stale}) + future = _stripe_header("whsec_test", body, time.time() + 3600) + assert not verify(body, {"stripe-signature": future}) + + +def test_stripe_signature_rejects_tampering_and_garbage(monkeypatch): + """Wrong secret, edited body, malformed header: all refused.""" + monkeypatch.setenv("WH", "whsec_test") + verify = stripe_signature(secret_env="WH") + body = b'{"amount": 100}' + header = _stripe_header("whsec_other", body, time.time()) + assert not verify(body, {"stripe-signature": header}) + good = _stripe_header("whsec_test", body, time.time()) + assert not verify(b'{"amount": 999}', {"stripe-signature": good}) + assert not verify(body, {"stripe-signature": "t=abc,v1=zzz"}) + assert not verify(body, {"stripe-signature": "v1=deadbeef"}) + assert not verify(body, {}) + + +def test_stripe_signature_accepts_any_matching_v1_during_rotation(monkeypatch): + """Stripe sends multiple v1 digests while a secret rotates.""" + monkeypatch.setenv("WH", "whsec_new") + verify = stripe_signature(secret_env="WH") + body = b"{}" + timestamp = int(time.time()) + old = hmac.new(b"whsec_old", f"{timestamp}.".encode() + body, "sha256").hexdigest() + new = hmac.new(b"whsec_new", f"{timestamp}.".encode() + body, "sha256").hexdigest() + assert verify(body, {"stripe-signature": f"t={timestamp},v1={old},v1={new}"}) + + +def test_hmac_signature_no_longer_claims_stripe(): + """The raw-body helper must not present itself as a Stripe verifier. + + It cannot verify Stripe's timestamped scheme, and the docstring saying it + covered Stripe was an invitation to ship replayable payment webhooks. + """ + doc = hmac_signature.__doc__ or "" + assert "deliberately **not** a Stripe verifier" in doc + assert "stripe_signature" in doc, "the fix must point at the real verifier" diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py index 1a563b4183d..ca9b2f8b742 100644 --- a/tests/units/workflow/test_ingress.py +++ b/tests/units/workflow/test_ingress.py @@ -505,3 +505,113 @@ async def test_a_legacy_key_only_matches_the_root_that_wrote_it( ) assert paid.json()["run_id"] != legacy.run_id await runtime.shutdown() + + +class Ships(rx.State): + """A workflow whose provider identifies deliveries by header.""" + + __workflow__ = WorkflowConfig(id="ingress.ships") + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "shipped", + dedupe_by="header:X-GitHub-Delivery", + verify=hmac_signature( + secret_env="STRIPE_WEBHOOK_SECRET", header="X-Signature" + ), + ), + ) + def on_shipped(self, sha: str): + """Record a shipment. + + Args: + sha: The commit that shipped. + + Returns: + Completion. + """ + return rx.complete(result=sha) + + +async def test_header_identity_separates_deliveries_the_payload_cannot( + monkeypatch, forked_registration_context +): + """GitHub's identity is the delivery GUID header, not any payload field. + + Two pushes of the same commit are two events with two GUIDs; keyed on a + payload field they collapsed into one run and the second delivery was + silently dropped. Keyed on the header, distinct GUIDs are distinct runs + and a true redelivery -- same GUID -- still deduplicates. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Ships) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + body = json.dumps({"sha": "abc123"}).encode() + + with TestClient(app) as client: + first = client.post( + "/_workflow/webhook/shipped", + content=body, + headers={"x-signature": _sign(body), "x-github-delivery": "guid-1"}, + ) + second = client.post( + "/_workflow/webhook/shipped", + content=body, + headers={"x-signature": _sign(body), "x-github-delivery": "guid-2"}, + ) + redelivered = client.post( + "/_workflow/webhook/shipped", + content=body, + headers={"x-signature": _sign(body), "x-github-delivery": "guid-1"}, + ) + assert first.json()["disposition"] == "started" + assert second.json()["disposition"] == "started", ( + "a distinct delivery GUID is a distinct event" + ) + assert redelivered.json()["disposition"] == "deduplicated" + assert redelivered.json()["run_id"] == first.json()["run_id"] + await runtime.shutdown() + + +async def test_a_delivery_missing_its_configured_identity_is_refused( + monkeypatch, forked_registration_context +): + """Configured dedupe that cannot be extracted must not silently vanish. + + Admitting anyway means every redelivery of this event executes again -- + the exact thing dedupe_by was configured to prevent -- and nobody is + told. A 400 naming the missing field is a config problem someone sees. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Ships) + runtime.register(Invoices) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + with TestClient(app) as client: + body = json.dumps({"sha": "abc123"}).encode() + no_header = client.post( + "/_workflow/webhook/shipped", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert no_header.status_code == 400, no_header.text + assert "X-GitHub-Delivery" in no_header.json()["error"] + + body = json.dumps({"amount": 5}).encode() + no_field = client.post( + "/_workflow/webhook/invoice_paid", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert no_field.status_code == 400, no_field.text + assert "'id'" in no_field.json()["error"] + await runtime.shutdown() From fef3f885435cd028a6b923e0e14e666cd5672c00 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 15:09:16 -0700 Subject: [PATCH 091/121] workflows: close the review's operability batch Fan-out children had no admission history: their record began at attempt_started, so runs_started reported one start for a four-run graph and every history reader met runs that were apparently never admitted. The creating transaction now writes run_admitted and step_scheduled per child -- one shared helper, three stores -- and the kernel reports the same events to the observer, so the dashboard's starts reconcile with its terminals. Opening the SQLite store ran CREATEs and an immediate-mode migration unconditionally, so an operator's read-only stats against a busy worker failed 4 times in 5 with 'database is locked'. The schema version is stamped in PRAGMA user_version now; a current database opens with no write lock at all, and a locked store renders as an actionable one-line error instead of a traceback. The legacy-migration test now also resets user_version in its simulation, because a genuinely old database has no stamp -- that is the very thing that triggers DDL. There was no retention surface: terminal data grew forever at ~1.7KB a run unless an operator did out-of-band SQL. purge_runs(before) deletes stale terminal runs and their steps, history, inbox, substeps and dedupe rows in one transaction on all three stores, conformance-pinned, with 'reflex workflows purge --older-than 30d' in front of it. The documented tradeoff: purging forgets request keys, so retention must exceed the provider's redelivery horizon. run_until_idle() could return with attempts from its own scheduling round still live -- a round starts several and waits only for the first -- handing tests a half-processed graph whose survivors the harness then cancelled. It now drains exactly the attempts it started, tracked per pump so a concurrent caller's work is never waited on; the first version of this fix waited on all in-flight attempts and deadlocked the lease tests, which own deliberately-hanging handlers. --- news/workflow-operability-batch.feature.md | 1 + reflex/workflow/cli.py | 52 ++++++- reflex/workflow/conformance.py | 39 +++++ reflex/workflow/kernel.py | 33 ++++- reflex/workflow/postgres.py | 58 +++++++- reflex/workflow/store.py | 159 ++++++++++++++++++++- tests/units/workflow/test_cli.py | 15 ++ tests/units/workflow/test_kernel.py | 43 ++++++ tests/units/workflow/test_metrics.py | 74 ++++++++++ tests/units/workflow/test_store.py | 3 + 10 files changed, 470 insertions(+), 7 deletions(-) create mode 100644 news/workflow-operability-batch.feature.md diff --git a/news/workflow-operability-batch.feature.md b/news/workflow-operability-batch.feature.md new file mode 100644 index 00000000000..33e387f197e --- /dev/null +++ b/news/workflow-operability-batch.feature.md @@ -0,0 +1 @@ +Four operability gaps from the production-readiness review. Fan-out children now record `run_admitted` and `step_scheduled` in the transaction that creates them and report both to the metrics observer, so a four-run graph counts four starts and child histories begin at admission like any run's. Opening a current SQLite store takes no write lock — DDL runs only when the stamped schema version is behind — so `list`/`stats`/`show` no longer die with `database is locked` against a busy worker, and bounded contention is rendered as an actionable error instead of a traceback. `reflex workflows purge --older-than 30d` deletes terminal runs and everything they own, with the documented tradeoff that a purged run's deduplication key is forgotten with it. `run_until_idle()` drains every attempt it started before returning — never attempts a concurrent pump owns — so a single call processes the whole batch it admitted rather than returning with live attempts for the harness to cancel. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index ebb407589f0..03a780d19fe 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -124,7 +124,19 @@ async def session() -> Any: if inspect.isawaitable(closed): await closed - return asyncio.run(session()) + import sqlite3 + + try: + return asyncio.run(session()) + except sqlite3.OperationalError as err: + if "locked" not in str(err) and "busy" not in str(err): + raise + # Bounded contention against a busy worker is an operational fact, + # not a traceback: say what happened and what to do. + console.error( + "The store is busy (a worker holds its write lock). Retry in a moment." + ) + raise click.exceptions.Exit(1) from None def _age(seconds: float) -> str: @@ -903,6 +915,44 @@ async def counts(store: RunStore) -> dict[str, int]: console.print(f"{total} run(s); {open_runs} open, {attention} needing attention.") +@workflows.command() +@database_option +@click.option( + "--older-than", + required=True, + help="Delete terminal runs whose last update is older than this, e.g. 30d.", +) +@click.option("--workflow", "-w", default=None, help="Only this workflow id.") +@click.option("--yes", is_flag=True, help="Delete without asking.") +def purge(database: str | None, older_than: str, workflow: str | None, yes: bool): + """Delete finished runs older than a cutoff, reclaiming the store. + + Terminal data grows forever otherwise. Purging a run also forgets its + deduplication key, so a provider redelivery arriving after the retention + window is admitted as a new run -- keep the window longer than the + provider's redelivery horizon. + """ + import time + + from reflex_base.workflow import parse_duration + + try: + cutoff = time.time() - parse_duration(older_than) + except Exception as err: + console.error(f"--older-than {older_than!r} is not a duration: {err}") + raise click.exceptions.Exit(1) from None + if not yes: + click.confirm( + f"Delete terminal runs untouched for {older_than}" + f"{' in ' + workflow if workflow else ''}?", + abort=True, + ) + deleted = _with_store( + database, lambda store: store.purge_runs(cutoff, workflow_id=workflow) + ) + console.print(f"Purged {deleted} run(s).") + + @workflows.command("list") @database_option @click.option("--workflow", "-w", default=None, help="Only this workflow id.") diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index d7a37c9ba97..8b420c55494 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -1086,6 +1086,44 @@ async def check_flow_gate_singleton_cancel_replaces(store: RunStore) -> None: ) +async def check_purge_deletes_only_stale_terminal_runs(store: RunStore) -> None: + """Retention removes finished history and never touches live work.""" + await store.admit( + make_run("done", status=RunStatus.COMPLETED, request_key="evt-done"), + make_step("done", status=StepStatus.SUCCEEDED), + _ADMITTED, + ) + await store.admit( + make_run( + "fresh", + status=RunStatus.COMPLETED, + created_at=NOW + 500, + updated_at=NOW + 500, + ), + make_step("fresh", status=StepStatus.SUCCEEDED), + _ADMITTED, + ) + await store.admit(make_run("live", created_at=NOW), make_step("live"), _ADMITTED) + + deleted = await store.purge_runs(NOW + 100) + assert deleted == 1 + assert await store.get_run("done") is None + assert await store.get_steps("done") == () + assert await store.get_history("done") == () + fresh = await store.get_run("fresh") + assert fresh is not None, "a terminal run inside the window stays" + live = await store.get_run("live") + assert live is not None, "an open run is never retention's business" + # The purged run's dedupe key is forgotten with it: a redelivery after + # the retention window is a fresh admission, by design. + replay = await store.admit( + make_run("done2", request_key="evt-done", created_at=NOW + 600), + make_step("done2"), + _ADMITTED, + ) + assert replay == (True, "done2") + + async def check_flow_gate_dedupes_before_policy(store: RunStore) -> None: """A redelivered event is its prior run, not a new start to be policed.""" from reflex.workflow.store import FlowGate @@ -1292,6 +1330,7 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_flow_gate_rate_throttle_and_debounce, check_flow_gate_singleton_cancel_replaces, check_flow_gate_dedupes_before_policy, + check_purge_deletes_only_stale_terminal_runs, check_skip_unsticks_a_stopped_run, check_retry_reopens_only_failed_runs, check_force_finalize_records_a_result, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 5de28901e37..e9208b1777b 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -65,6 +65,7 @@ RunStore, StaleClaimError, StepCompletion, + _child_admission_events, ) if TYPE_CHECKING: @@ -2420,6 +2421,11 @@ async def _commit_outcome( await self._record_abandoned(claim, handler, "fenced_at_commit") return self._notify(claim.run, completion.events) + for child_run, child_step in completion.children: + # The store recorded each child's admission inside the commit; + # the observer hears the same events, so runs_started counts a + # four-run graph as four. + self._notify(child_run, _child_admission_events(child_run, child_step)) if completion.children: self._wakeup.set() await self._report_to_parent(claim.run, completion) @@ -2704,9 +2710,14 @@ async def _fill_slots(self, now: float) -> list[asyncio.Task]: started.append(self._spawn(claim)) return started - async def _tick(self) -> bool: + async def _tick(self, own: set[asyncio.Task] | None = None) -> bool: """Run one scheduling round. + Args: + own: When given, the attempts this round starts are added, so the + caller can later drain exactly what it started and nothing a + concurrent pump owns. + Returns: True if any control transition or attempt was processed. """ @@ -2715,6 +2726,8 @@ async def _tick(self) -> bool: progressed = await self._admit_due_schedules(now) > 0 or progressed progressed = await self._finalize_control(now) > 0 or progressed started = await self._fill_slots(now) + if own is not None: + own.update(started) progressed = bool(started) or progressed # Wait only on what this round started: another caller pumping the same # kernel must not block on an attempt it does not own. @@ -2775,8 +2788,22 @@ async def run_until_idle(self) -> None: the clock and call again to run it. """ await self.recover() - while await self._tick(): - pass + own: set[asyncio.Task] = set() + while True: + if await self._tick(own): + continue + # A round can start several attempts and return after the first + # completion, so "nothing newly claimable" is not "nothing + # running": returning here hands a test a half-processed graph + # and cancels the survivors when the harness exits. Idle means no + # claimable work AND none of the attempts THIS pump started still + # live -- an attempt some other caller owns (a hanging handler a + # test controls, a concurrent pump's work) is not ours to wait on. + live = [task for task in own if not task.done()] + if not live: + return + await asyncio.wait(live, return_when=asyncio.FIRST_COMPLETED) + self._prune() async def _worker_loop(self) -> None: """Process work continuously until the kernel is closed.""" diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index da6538e6a21..0ccd6af1caa 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -32,7 +32,13 @@ StepRecord, StepStatus, ) -from reflex.workflow.store import Claim, FlowAdmission, FlowGate, StaleClaimError +from reflex.workflow.store import ( + Claim, + FlowAdmission, + FlowGate, + StaleClaimError, + _child_admission_events, +) try: import psycopg @@ -769,6 +775,50 @@ async def admit_flow( await self._append_events(conn, run.run_id, events, run.created_at) return FlowAdmission("started", run.run_id, cancelled=tuple(cancelled)) + async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> int: + """Delete terminal runs not updated since a cutoff, and all their data. + + Args: + before: Delete runs whose last update is older than this. + workflow_id: Restrict to one workflow identity. + + Returns: + How many runs were deleted. + """ + pool = await self._open() + where = "status = ANY(%s) AND updated_at < %s" + params: tuple[Any, ...] = (_TERMINAL_RUNS, before) + if workflow_id is not None: + where += " AND workflow_id = %s" + params = (*params, workflow_id) + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + f"SELECT run_id FROM workflow_runs WHERE {where}", + params, + ) + doomed = [row["run_id"] for row in await cursor.fetchall()] + if not doomed: + return 0 + for table in ( + "workflow_steps", + "workflow_history", + "workflow_inbox", + "workflow_substeps", + ): + await conn.execute( + SQL("DELETE FROM {} WHERE run_id = ANY(%s)").format( + Identifier(table) + ), + (doomed,), + ) + await conn.execute( + "DELETE FROM workflow_dedupe WHERE run_id = ANY(%s)", (doomed,) + ) + await conn.execute( + "DELETE FROM workflow_runs WHERE run_id = ANY(%s)", (doomed,) + ) + return len(doomed) + async def claim_next( self, now: float, @@ -946,6 +996,12 @@ async def commit( for child_run, child_step in completion.children: await self._insert_run(conn, child_run) await self._insert_step(conn, child_step) + await self._append_events( + conn, + child_run.run_id, + _child_admission_events(child_run, child_step), + now, + ) await conn.execute( "UPDATE workflow_runs SET status = %s," " state = CASE WHEN %s THEN %s ELSE state END," diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index de2a80be6bb..bc0294789c3 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -405,6 +405,23 @@ async def admit_flow( """ ... + async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> int: + """Delete terminal runs not updated since a cutoff, and all their data. + + Terminal data otherwise grows forever. Purging a run also forgets its + request key, so a provider redelivery arriving after the retention + window is admitted as a new run: retention must exceed the provider's + redelivery horizon. + + Args: + before: Delete runs whose last update is older than this. + workflow_id: Restrict to one workflow identity. + + Returns: + How many runs were deleted. + """ + ... + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: """Find the oldest run still in flight under a flow-control key. @@ -1134,6 +1151,11 @@ async def commit( for child_run, child_step in completion.children: self._runs[child_run.run_id] = child_run self._steps[child_run.run_id] = [child_step] + self._append_events( + child_run.run_id, + _child_admission_events(child_run, child_step), + now, + ) self._runs[run.run_id] = dataclasses.replace( run, status=completion.run_status, @@ -1488,6 +1510,34 @@ async def admit_flow( self._append_events(run.run_id, events, run.created_at) return FlowAdmission("started", run.run_id, cancelled=tuple(cancelled)) + async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> int: + """Delete terminal runs not updated since a cutoff, and all their data. + + Args: + before: Delete runs whose last update is older than this. + workflow_id: Restrict to one workflow identity. + + Returns: + How many runs were deleted. + """ + async with self._lock: + doomed = [ + run.run_id + for run in self._runs.values() + if run.status in TERMINAL_RUN_STATUSES + and run.updated_at < before + and (workflow_id is None or run.workflow_id == workflow_id) + ] + for run_id in doomed: + run = self._runs.pop(run_id) + self._steps.pop(run_id, None) + self._history.pop(run_id, None) + self._pending.pop(run_id, None) + self._inbox.pop(run_id, None) + if run.request_key is not None: + self._dedupe.pop((run.workflow_id, run.request_key), None) + return len(doomed) + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: """Find the oldest run still in flight under a flow-control key. @@ -2188,6 +2238,9 @@ async def next_due( return min(due_times) if due_times else None +SCHEMA_VERSION: Final = 2 +"""Stamped into PRAGMA user_version; bump when _SCHEMA or migrations change.""" + DATABASE_ENV: Final = "REFLEX_WORKFLOW_DATABASE" DEFAULT_DB_FILENAME: Final = "workflow.db" @@ -2419,6 +2472,38 @@ def _step_from_row(row: sqlite3.Row) -> StepRecord: ) +def _child_admission_events( + child_run: RunRecord, child_step: StepRecord +) -> tuple[tuple[HistoryEventType, dict[str, Any]], ...]: + """The admission history a fan-out child gets, same as any run. + + Children are advertised as ordinary runs, so a history that begins at + attempt_started -- with no admission, no scheduling -- breaks the + invariant every other reader relies on, and a metrics observer counting + runs_started reports one start for a four-run graph. + + Args: + child_run: The child being created. + child_step: Its root slot. + + Returns: + The admission events to record with the creating transaction. + """ + return ( + ( + HistoryEventType.RUN_ADMITTED, + { + "handler_id": child_step.handler_id, + "request_key": child_run.request_key, + }, + ), + ( + HistoryEventType.STEP_SCHEDULED, + {"ordinal": child_step.ordinal, "handler_id": child_step.handler_id}, + ), + ) + + def _sqlite_frontier_query( select: str, now: float, @@ -2549,8 +2634,15 @@ def __init__(self, db_path: str | Path): # rather than block everything for SQLite's multi-second default. The # kernel treats the resulting error as transient and retries. self._db.execute(f"PRAGMA busy_timeout={BUSY_TIMEOUT_MS}") - self._db.executescript(_SCHEMA) - self._migrate() + # A current database opens without taking any write lock: DDL only + # runs when the stamped schema version is behind. An operator's + # list/stats/show against a busy worker used to fail with "database + # is locked" purely because opening the store ran CREATEs and an + # immediate-mode migration it did not need. + current = self._db.execute("PRAGMA user_version").fetchone()[0] + if current != SCHEMA_VERSION: + self._db.executescript(_SCHEMA) + self._migrate() def _migrate(self) -> None: """Add columns and indexes missing from databases created by older versions. @@ -2582,6 +2674,7 @@ def _migrate(self) -> None: "CREATE INDEX IF NOT EXISTS idx_workflow_runs_parent" " ON workflow_runs (parent_run_id, parent_ordinal)" ) + self._db.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") self._db.execute("COMMIT") except BaseException: self._db.execute("ROLLBACK") @@ -2882,6 +2975,63 @@ def work() -> FlowAdmission: return await asyncio.to_thread(work) + async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> int: + """Delete terminal runs not updated since a cutoff, and all their data. + + Args: + before: Delete runs whose last update is older than this. + workflow_id: Restrict to one workflow identity. + + Returns: + How many runs were deleted. + """ + + def work() -> int: + """Delete in one transaction on the worker thread. + + Returns: + How many runs were deleted. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + where = f"status IN ({','.join('?' * len(terminal))}) AND updated_at < ?" + params: tuple[Any, ...] = (*terminal, before) + if workflow_id is not None: + where += " AND workflow_id = ?" + params = (*params, workflow_id) + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + doomed = [ + row["run_id"] + for row in self._db.execute( + f"SELECT run_id FROM workflow_runs WHERE {where}", + params, + ).fetchall() + ] + for run_id in doomed: + for table in ( + "workflow_steps", + "workflow_history", + "workflow_inbox", + "workflow_substeps", + ): + self._db.execute( + f"DELETE FROM {table} WHERE run_id = ?", (run_id,) + ) + self._db.execute( + "DELETE FROM workflow_dedupe WHERE run_id = ?", (run_id,) + ) + self._db.execute( + "DELETE FROM workflow_runs WHERE run_id = ?", (run_id,) + ) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return len(doomed) + + return await asyncio.to_thread(work) + async def claim_next( self, now: float, @@ -3127,6 +3277,11 @@ def work() -> None: for child_run, child_step in completion.children: self._insert_run(child_run) self._insert_step(child_step) + self._append_events( + child_run.run_id, + _child_admission_events(child_run, child_step), + now, + ) self._db.execute( "UPDATE workflow_runs SET status = ?," " state = CASE WHEN ? THEN ? ELSE state END," diff --git a/tests/units/workflow/test_cli.py b/tests/units/workflow/test_cli.py index 29315f3dccd..62254ab3721 100644 --- a/tests/units/workflow/test_cli.py +++ b/tests/units/workflow/test_cli.py @@ -253,3 +253,18 @@ def test_complete_refuses_a_result_that_is_not_json(seeded): result = _invoke("complete", "-d", database, waiting, "--result", "{oops") assert result.exit_code == 1 assert "not JSON" in result.output + + +def test_purge_deletes_only_stale_terminal_runs(seeded): + """Retention is an operator command, not out-of-band SQL.""" + database, waiting, _ = seeded + kept = _invoke("purge", "-d", database, "--older-than", "0s", "--yes") + assert kept.exit_code == 0, kept.output + assert "Purged 0 run(s)" in kept.output, "no seeded run is terminal yet" + + done = _invoke("complete", "-d", database, waiting, "--result", '"ok"') + assert done.exit_code == 0, done.output + purged = _invoke("purge", "-d", database, "--older-than", "0s", "--yes") + assert purged.exit_code == 0, purged.output + assert "Purged 1 run(s)" in purged.output + assert _invoke("show", "-d", database, waiting).exit_code == 1 diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py index 22175b471eb..46751fd0133 100644 --- a/tests/units/workflow/test_kernel.py +++ b/tests/units/workflow/test_kernel.py @@ -760,3 +760,46 @@ async def slow_to_die(): await asyncio.wait_for(asyncio.shield(releasing), timeout=2) assert releasing.done(), "the release task outlived its cancellation" assert releasing.cancelled(), "the release task swallowed its own cancellation" + + +async def test_run_until_idle_drains_every_attempt_it_started( + forked_registration_context, +): + """One call processes the whole batch, not the first completion's worth. + + A round can start several attempts and return after the first finishes; + a later round that finds nothing newly claimable must still wait for the + attempts this pump started, or the caller gets a half-processed graph + and the harness cancels the survivors on exit. + """ + + class Slow(rx.State): + __workflow__ = WorkflowConfig(id="kernel.slowbatch") + + @rx.event(durable=True, trigger=manual(), effect="none") + async def start(self, n: int): + """Take long enough to still be running next round. + + Args: + n: Which of the batch this is. + + Returns: + Completion. + """ + await asyncio.sleep(0.05) + return rx.complete(result=n) + + store = MemoryRunStore() + definition = compile_workflow(Slow) + kernel = WorkflowKernel([definition], store, max_concurrency=8) + for n in range(8): + await kernel.start(Slow.start(n)) + await kernel.run_until_idle() + + from reflex.workflow.records import RunQuery + + runs = await store.list_runs(RunQuery(limit=20)) + statuses = sorted(run.status.value for run in runs) + assert statuses == ["COMPLETED"] * 8, ( + f"run_until_idle returned with live attempts: {statuses}" + ) diff --git a/tests/units/workflow/test_metrics.py b/tests/units/workflow/test_metrics.py index 0ab87cb31a5..9723b8aca19 100644 --- a/tests/units/workflow/test_metrics.py +++ b/tests/units/workflow/test_metrics.py @@ -9,6 +9,7 @@ import reflex as rx from reflex.workflow.kernel import MetricsObserver +from reflex.workflow.records import HistoryEventType, RunStatus from reflex.workflow.testing import WorkflowTestHarness CALLS: list[int] = [] @@ -129,3 +130,76 @@ def go(self): assert totals["runs_failed"] == 1 assert totals["attempts_failed"] == 1 assert "runs_completed" not in totals + + +async def test_fanout_children_count_as_started_runs(forked_registration_context): + """A four-run graph reports four starts, not one. + + Children are advertised as ordinary runs, so their history begins with + admission and scheduling like anyone else's, and a dashboard's + runs_started reconciles with its terminal counts instead of trailing + them by the fan-out width. + """ + + class Branch(rx.State): + __workflow__ = WorkflowConfig(id="metrics.branch") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, lead: str): + """Complete immediately. + + Args: + lead: The lead identifier. + + Returns: + Completion. + """ + return rx.complete(result=lead) + + class Fans(rx.State): + __workflow__ = WorkflowConfig(id="metrics.fans") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Fan out three branches. + + Returns: + The parallel fan-out. + """ + return rx.parallel( + Branch.start("a"), Branch.start("b"), Branch.start("c"), then=Fans.done + ) + + @rx.event(durable=True, effect="none") + def done(self, results: list): + """Complete with the branch count. + + Args: + results: One entry per branch. + + Returns: + Completion. + """ + return rx.complete(result=len(results)) + + metrics = MetricsObserver() + async with WorkflowTestHarness(Fans, Branch, observer=metrics) as harness: + started = await harness.start(Fans.begin()) + assert started.run_id is not None + await harness.run_until_idle() + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + + # The durable record agrees: each child's history begins at admission. + store = harness.kernel._store # pyright: ignore[reportPrivateUsage] + children = await store.list_children(started.run_id, 1) + assert len(children) == 3 + for child in children: + history = [event.type for event in await store.get_history(child.run_id)] + assert history[0] is HistoryEventType.RUN_ADMITTED, history + assert history[1] is HistoryEventType.STEP_SCHEDULED, history + + counts = metrics.snapshot()["totals"] + assert counts["runs_started"] == 4, counts + assert counts["runs_completed"] == 4, counts diff --git a/tests/units/workflow/test_store.py b/tests/units/workflow/test_store.py index b70bc4727cc..1d67cffad1a 100644 --- a/tests/units/workflow/test_store.py +++ b/tests/units/workflow/test_store.py @@ -448,6 +448,9 @@ async def test_sqlite_migrates_a_database_without_the_lease_column(tmp_path): # Simulate a database written by a build that predates leases. store._db.execute("DROP INDEX IF EXISTS idx_workflow_steps_lease") store._db.execute("ALTER TABLE workflow_steps DROP COLUMN lease_expires_at") + # A database written by an older build carries no schema-version stamp, + # which is what tells the next open to run DDL at all. + store._db.execute("PRAGMA user_version = 0") store.close() reopened = SqliteRunStore(db_path) From 10f08b11161ecbdd39063f4b791ce33a82028559 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 15:12:03 -0700 Subject: [PATCH 092/121] workflows: accept GitHub's form-encoded webhook mode The review's P3. GitHub can be configured to deliver application/x-www-form-urlencoded bodies carrying payload=, and ingress unconditionally json.loads'd the raw body, so a correctly signed delivery came back 400 'payload is not JSON' -- an error that tells the operator nothing about which knob to turn. The signature was never the problem: it covers the raw form bytes and is checked against exactly those bytes before any parsing. Form bodies now unwrap the payload field before JSON parsing; a form body without one is a 400 that names it. --- news/workflow-github-form-mode.feature.md | 1 + reflex/workflow/ingress.py | 18 ++++++++++- tests/units/workflow/test_ingress.py | 38 +++++++++++++++++++++++ 3 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 news/workflow-github-form-mode.feature.md diff --git a/news/workflow-github-form-mode.feature.md b/news/workflow-github-form-mode.feature.md new file mode 100644 index 00000000000..88afe48f4a2 --- /dev/null +++ b/news/workflow-github-form-mode.feature.md @@ -0,0 +1 @@ +Webhooks configured on GitHub's `application/x-www-form-urlencoded` mode — where the JSON document arrives wrapped as `payload=` — are now understood. The signature is verified over the raw form bytes exactly as received; only the JSON extraction changes. A form body without a `payload` field is a 400 naming the problem. diff --git a/reflex/workflow/ingress.py b/reflex/workflow/ingress.py index 4ac5da4c4b7..168d965047e 100644 --- a/reflex/workflow/ingress.py +++ b/reflex/workflow/ingress.py @@ -232,8 +232,24 @@ async def endpoint(request: Request) -> JSONResponse: if not verified: return JSONResponse({"error": "invalid signature"}, status_code=401) + content_type = headers.get("content-type", "") + if "application/x-www-form-urlencoded" in content_type: + # GitHub's form mode wraps the JSON document in a form field: + # payload=. The signature is over the raw form + # body and was already checked against exactly those bytes. + from urllib.parse import parse_qs + + form = parse_qs(body.decode("utf-8", errors="replace")) + wrapped = form.get("payload", [None])[0] + if wrapped is None: + return JSONResponse( + {"error": "form body has no 'payload' field"}, status_code=400 + ) + source = wrapped + else: + source = body try: - payload = json.loads(body) if body else {} + payload = json.loads(source) if source else {} except ValueError: return JSONResponse({"error": "payload is not JSON"}, status_code=400) diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py index ca9b2f8b742..458a1124138 100644 --- a/tests/units/workflow/test_ingress.py +++ b/tests/units/workflow/test_ingress.py @@ -615,3 +615,41 @@ async def test_a_delivery_missing_its_configured_identity_is_refused( assert no_field.status_code == 400, no_field.text assert "'id'" in no_field.json()["error"] await runtime.shutdown() + + +async def test_github_form_encoded_deliveries_are_understood( + monkeypatch, forked_registration_context +): + """GitHub can be configured to send payload= as a form body. + + The signature covers the raw form bytes and was already verified against + them; only the JSON extraction changes. Refusing this mode as 'not JSON' + forced users to reconfigure the provider to learn what was wrong. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Ships) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + from urllib.parse import urlencode + + body = urlencode({"payload": json.dumps({"sha": "abc123"})}).encode() + with TestClient(app) as client: + response = client.post( + "/_workflow/webhook/shipped", + content=body, + headers={ + "x-signature": _sign(body), + "x-github-delivery": "guid-form-1", + "content-type": "application/x-www-form-urlencoded", + }, + ) + assert response.status_code == 202, response.text + assert response.json()["disposition"] == "started" + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(response.json()["run_id"]) + assert snapshot is not None + assert snapshot.result == {"sha": "abc123"}, "the wrapped JSON is the payload" + await runtime.shutdown() From 7130537d8179671006a45deab15b139cf070cfca Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 15:40:23 -0700 Subject: [PATCH 093/121] workflows: fix what the adversarial review of tonight's fixes found A 48-agent review over the six commits -- four lenses, two refuters per finding -- confirmed five defects in the new code itself. Postgres singleton-cancel could resurrect a finished run: the active SELECT is an unlocked snapshot, workers committing runs never hold the flow lock, and the cancel UPDATE had no non-terminal guard -- so under READ COMMITTED it waits out a concurrent completing commit and flips the COMPLETED run back to CANCELLING, which the next finalize sweep drives to CANCELLED over a terminal status. The guard request_cancel always had is now on this UPDATE too, and the cancelled list comes from RETURNING instead of the stale snapshot, so the kernel is never told a run was cancelled that actually completed. Postgres admit_flow's dedupe reservation was a plain INSERT after a plain SELECT; the advisory lock is keyed on the flow key, and two admissions can share a request key while computing different flow keys, so the loser raised UniqueViolation out of kernel.start where the contract promises a deduplicated admission. Same ON CONFLICT shape as admit() now. Memory purge_runs never deleted substep journals -- and popping by run_id would have silently missed anyway, since the journal is keyed by (run_id, ordinal). SQLite and Postgres purged theirs; parity restored. run_until_idle kept every finished task in its own-set for the whole call and its tail wait had no CancelledError handler, so cancelling a pump parked there left its attempts running unsupervised -- the exact contract _tick's wait already honors. FlowGate refuses combined policies in __post_init__. The decorator already enforces one policy per root, which is why the review's three divergence findings about combinations were refuted as unreachable -- but the stores genuinely disagree on what a half-applied combination leaves behind, and unreachable-because-callers-are-polite is not an invariant. Now it is unrepresentable. stripe_signature rejects non-finite timestamps: NaN compares false against both window bounds and sailed through -- harmless today only because the signature check follows, and the same NaN class bug was already found once in approval expiry. --- news/workflow-review-hardening.bugfix.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 5 +++ reflex/workflow/kernel.py | 13 +++++-- reflex/workflow/postgres.py | 36 +++++++++++++++---- reflex/workflow/store.py | 30 ++++++++++++++++ 5 files changed, 75 insertions(+), 10 deletions(-) create mode 100644 news/workflow-review-hardening.bugfix.md diff --git a/news/workflow-review-hardening.bugfix.md b/news/workflow-review-hardening.bugfix.md new file mode 100644 index 00000000000..12d7775830d --- /dev/null +++ b/news/workflow-review-hardening.bugfix.md @@ -0,0 +1 @@ +Five defects found by an adversarial multi-agent review of this branch's own fixes. Postgres `admit_flow`'s singleton-cancel could resurrect a concurrently-finalized run — its UPDATE lacked the non-terminal guard `request_cancel` has, and under READ COMMITTED it would wait out a worker's completing commit and then flip the finished run back to CANCELLING; it now guards on status and reports cancellations from `RETURNING`. Its dedupe reservation regressed to a plain INSERT that could raise `UniqueViolation` in a race the flow-key advisory lock does not cover (same request key, different flow keys); it now uses the same `ON CONFLICT DO NOTHING RETURNING` shape as `admit`. The memory store's `purge_runs` never deleted substep journals. `run_until_idle` retained every finished task for the call's duration and, if cancelled mid-wait, left its attempts running unsupervised. `FlowGate` now refuses combined policies outright — the decorator already enforces one policy per root, but the three stores genuinely diverge on what a half-applied combination leaves behind, and an invariant that holds only because callers are polite is not an invariant. `stripe_signature` also rejects non-finite timestamps, which compare false against every window bound. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index bdeace3084b..e581b1286ba 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -10,6 +10,7 @@ import dataclasses import hmac +import math import os import re import time @@ -454,6 +455,10 @@ def __call__(self, body: bytes, headers: Mapping[str, str]) -> bool: signed_at = float(timestamp) except ValueError: return False + if not math.isfinite(signed_at): + # NaN compares False against everything, which would wave it + # through the window check. + return False if abs(time.time() - signed_at) > self.tolerance: return False expected = hmac.new( diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index e9208b1777b..3a3dff14024 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -2799,10 +2799,17 @@ async def run_until_idle(self) -> None: # claimable work AND none of the attempts THIS pump started still # live -- an attempt some other caller owns (a hanging handler a # test controls, a concurrent pump's work) is not ours to wait on. - live = [task for task in own if not task.done()] - if not live: + own = {task for task in own if not task.done()} + if not own: return - await asyncio.wait(live, return_when=asyncio.FIRST_COMPLETED) + try: + await asyncio.wait(own, return_when=asyncio.FIRST_COMPLETED) + except asyncio.CancelledError: + # Same contract as _tick's wait: a cancelled pump must stop + # the attempts it started or they run on unsupervised. + if not self._draining: + await self._cancel_inflight() + raise self._prune() async def _worker_loop(self) -> None: diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 0ccd6af1caa..332260b9165 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -710,19 +710,27 @@ async def admit_flow( cancelled: list[str] = [] if gate.singleton_cancel and active: ids = [row["run_id"] for row in active] - await conn.execute( + # The active SELECT is an unlocked snapshot, and a worker + # committing one of these runs to a terminal status does not + # hold the flow lock: under READ COMMITTED this UPDATE would + # wait out that commit and then flip the finished run back to + # CANCELLING -- resurrecting a terminal run. The non-terminal + # guard makes the row-version re-check decide correctly, and + # RETURNING reports only what was actually cancelled. + cursor = await conn.execute( "UPDATE workflow_runs SET cancel_requested = TRUE," - " status = %s, updated_at = %s WHERE run_id = ANY(%s)", - (RunStatus.CANCELLING.value, now, ids), + " status = %s, updated_at = %s WHERE run_id = ANY(%s)" + " AND NOT (status = ANY(%s)) RETURNING run_id", + (RunStatus.CANCELLING.value, now, ids, _TERMINAL_RUNS), ) - for run_id in ids: + cancelled = [row["run_id"] for row in await cursor.fetchall()] + for run_id in cancelled: await self._append_events( conn, run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now, ) - cancelled = ids due_at = root_step.due_at if gate.rate_limit is not None: limit, window = gate.rate_limit @@ -765,11 +773,25 @@ async def admit_flow( return FlowAdmission("coalesced", active[0]["run_id"]) due_at = now + gate.debounce if run.request_key is not None: - await conn.execute( + # Same reservation shape as admit(): the advisory lock is + # keyed on the flow key, and two admissions can share a + # request key while computing different flow keys, so a plain + # INSERT here can lose a race the lock does not cover and + # raise instead of deduplicating. + cursor = await conn.execute( "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" - " VALUES (%s, %s, %s)", + " VALUES (%s, %s, %s) ON CONFLICT DO NOTHING RETURNING run_id", (run.workflow_id, run.request_key, run.run_id), ) + if await cursor.fetchone() is None: + cursor = await conn.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = %s AND request_key = %s", + (run.workflow_id, run.request_key), + ) + existing = await cursor.fetchone() + if existing is not None: + return FlowAdmission("deduplicated", existing["run_id"]) await self._insert_run(conn, run) await self._insert_step(conn, dataclasses.replace(root_step, due_at=due_at)) await self._append_events(conn, run.run_id, events, run.created_at) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index bc0294789c3..432fdc62be2 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -147,6 +147,34 @@ class FlowGate: throttle: tuple[int, float] | None = None debounce: float | None = None + def __post_init__(self): + """Refuse combined policies. + + The decorator already enforces one policy per root; this makes the + combinations unrepresentable at the store boundary too, because the + three stores genuinely diverge on what a half-applied combination + would leave behind (a rejected start after a singleton cancellation + rolls back on SQLite and commits elsewhere), and an invariant that + holds only because callers are polite is not an invariant. + + Raises: + ValueError: If more than one policy is declared. + """ + declared = sum( + 1 + for flag in ( + self.singleton_skip, + self.singleton_cancel, + self.rate_limit is not None, + self.throttle is not None, + self.debounce is not None, + ) + if flag + ) + if declared > 1: + msg = "FlowGate takes exactly one policy; combinations diverge." + raise ValueError(msg) + @dataclasses.dataclass(frozen=True, slots=True) class FlowAdmission: @@ -1534,6 +1562,8 @@ async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> self._history.pop(run_id, None) self._pending.pop(run_id, None) self._inbox.pop(run_id, None) + for key in [k for k in self._substeps if k[0] == run_id]: + del self._substeps[key] if run.request_key is not None: self._dedupe.pop((run.workflow_id, run.request_key), None) return len(doomed) From 208b6074e3b1c31c45e95e188fe6d76eee91811b Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 15:41:55 -0700 Subject: [PATCH 094/121] workflows: say that operator repair overrides start policies The adversarial review's last confirmed finding was a contradiction between two sentences this branch added: section 1 promised a singleton holds 'at every instant', and section 9's retry re-opens a failed run without re-checking the gate -- so an operator retrying next to an already-admitted replacement puts two runs on one key. The behavior is the right one: a policy governs admissions, and an operator re-opening a run is a human override, not an admission. An engine that silently refused a repair because a policy would have is harder to operate, not safer. The contract now says exactly that, in both places, instead of promising an instant-by-instant invariant the operator surface deliberately does not enforce. --- reflex/workflow/CONTRACT.md | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 23fbc0a1e67..5d0dd46ea21 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -47,8 +47,10 @@ insert commit or roll back together. Two processes admitting concurrently under a limit of one therefore cannot both pass: nothing about policy enforcement assumes the admitters share a process. Concretely: -- `Singleton(mode="skip")`: at most one active run per key, at every instant, - from any number of processes; the loser is told which run holds the key. +- `Singleton(mode="skip")`: at most one *admitted* active run per key, at + every instant, from any number of processes; the loser is told which run + holds the key. Operator retry/skip may re-open a failed run alongside a + later admission — a human override, stated in §9. - `Singleton(mode="cancel")`: the replacement is admitted and every incumbent's cancellation intent is recorded in one transaction, so at most one *non-cancelling* run exists per key at every instant. Incumbents drain @@ -355,3 +357,10 @@ no-op with a reason. All of these are store transactions under the same atomicity rules as §1, and every one of them is reachable without writing Python: `reflex workflows cancel | resume | retry | skip | complete | fail `. + +Operator actions deliberately do not re-check start policies. Retrying a +failed `Singleton` run while its replacement is already active puts two runs +on one key — the singleton promise in §1 governs *admissions*, and an +operator re-opening a run is a human override, not an admission. The operator +can see what holds the key (`reflex workflows list -w `) and decide; +the engine does not silently refuse a repair because a policy would have. From 4b413ae855ecded38c38ca91bc01c6ace9b30958 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 16:33:43 -0700 Subject: [PATCH 095/121] workflows: a duplicate admission must leave with zero side effects Third live review round, three regressions from tonight's own commits. The Postgres flow-gate reserved its dedupe key AFTER the policy mutations. Two deliveries of one event can compute different flow keys -- different advisory locks, no serialization between them -- so the loser ran the singleton-cancel branch against the OTHER key's incumbent, then hit the reservation conflict, returned 'deduplicated', and committed the cancellation anyway: an unrelated live run killed by an event that had already been handled, 20 times in 20 trials. The reservation is now the first write, exactly as in admit(), so a duplicate exits before it has mutated anything. Regression test races the exact scenario on all three stores and fails on the old ordering. The user_version gate used inequality, so an older binary opening a newer database re-ran its DDL and stamped its own OLDER version over the newer one -- a silent downgrade that would make the newer binary re-migrate. Strictly upward now; a future stamp is left alone. Both cancellation-cleanup sites in the kernel -- _tick's wait and run_until_idle's tail wait -- called _cancel_inflight(), which stops every attempt in the kernel, including ones a concurrent pump owns. One caller's timeout became another's abandoned work. Each site now cancels only the attempts that pump started; aclose() still sweeps everything, which is its job. --- news/workflow-followup-regressions.bugfix.md | 1 + reflex/workflow/kernel.py | 24 ++++++-- reflex/workflow/postgres.py | 46 +++++++-------- reflex/workflow/store.py | 6 +- tests/units/workflow/test_drain.py | 56 +++++++++++++++++++ tests/units/workflow/test_flow_atomicity.py | 59 ++++++++++++++++++++ tests/units/workflow/test_sqlite_frontier.py | 22 ++++++++ 7 files changed, 181 insertions(+), 33 deletions(-) create mode 100644 news/workflow-followup-regressions.bugfix.md diff --git a/news/workflow-followup-regressions.bugfix.md b/news/workflow-followup-regressions.bugfix.md new file mode 100644 index 00000000000..d1dddaa928a --- /dev/null +++ b/news/workflow-followup-regressions.bugfix.md @@ -0,0 +1 @@ +Three regressions from the second live review round. A duplicate admission arriving under a different flow key — two deliveries of one event whose flow-key field differs — could execute singleton cancellations against the other key's incumbent and commit them on its way out as `deduplicated`; the Postgres dedupe reservation now happens before any policy mutation, so a duplicate leaves with zero side effects (20/20 trials had cancelled unrelated incumbents; the regression test now pins all three stores). Opening a SQLite store stamped by a newer binary no longer silently downgrades its schema-version stamp — DDL runs only strictly upward. Cancelling one manual pump no longer cancels attempts a concurrent pump owns: both cancellation-cleanup sites now stop only the attempts that pump started, and the kernel's close sweeps the rest. diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 3a3dff14024..d838c58fe54 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -2737,10 +2737,16 @@ async def _tick(self, own: set[asyncio.Task] | None = None) -> bool: except asyncio.CancelledError: # asyncio.wait does not cancel what it waits on, so cancelling # the scheduler must stop the attempts it started or they run - # on unsupervised. A drain is the exception: there the closer - # is waiting on them itself and cancels whatever is left over. + # on unsupervised -- and only those: with several pumps on one + # kernel, another pump's attempts are its own to supervise, + # and the closer's aclose() sweeps whatever remains. A drain + # is the exception: there the closer is waiting on them + # itself and cancels the leftovers. if not self._draining: - await self._cancel_inflight() + for task in started: + task.cancel() + await asyncio.gather(*started, return_exceptions=True) + self._prune() raise progressed = self._prune() > 0 or progressed return progressed @@ -2805,10 +2811,16 @@ async def run_until_idle(self) -> None: try: await asyncio.wait(own, return_when=asyncio.FIRST_COMPLETED) except asyncio.CancelledError: - # Same contract as _tick's wait: a cancelled pump must stop - # the attempts it started or they run on unsupervised. + # A cancelled pump must stop the attempts it started or they + # run on unsupervised -- but only its own: another pump's + # attempts are that pump's to supervise, and cancelling them + # from here turns one caller's timeout into another's lost + # work. if not self._draining: - await self._cancel_inflight() + for task in own: + task.cancel() + await asyncio.gather(*own, return_exceptions=True) + self._prune() raise self._prune() diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 332260b9165..6f0036351eb 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -689,14 +689,28 @@ async def admit_flow( (f"{run.workflow_id}\x1f{run.flow_key}",), ) if run.request_key is not None: + # Reserve the key before any policy mutation, exactly as + # admit() does. Two admissions can share a request key while + # computing different flow keys -- different advisory locks -- + # and if the reservation came after the singleton-cancel + # branch, the duplicate would cancel the other flow key's + # incumbents and then commit those cancellations on its way + # out as "deduplicated". A duplicate must leave with zero + # side effects. cursor = await conn.execute( - "SELECT run_id FROM workflow_dedupe" - " WHERE workflow_id = %s AND request_key = %s", - (run.workflow_id, run.request_key), + "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" + " VALUES (%s, %s, %s) ON CONFLICT DO NOTHING RETURNING run_id", + (run.workflow_id, run.request_key, run.run_id), ) - row = await cursor.fetchone() - if row is not None: - return FlowAdmission("deduplicated", row["run_id"]) + if await cursor.fetchone() is None: + cursor = await conn.execute( + "SELECT run_id FROM workflow_dedupe" + " WHERE workflow_id = %s AND request_key = %s", + (run.workflow_id, run.request_key), + ) + existing = await cursor.fetchone() + if existing is not None: + return FlowAdmission("deduplicated", existing["run_id"]) cursor = await conn.execute( "SELECT run_id FROM workflow_runs" " WHERE workflow_id = %s AND flow_key = %s" @@ -772,26 +786,6 @@ async def admit_flow( if cursor.rowcount: return FlowAdmission("coalesced", active[0]["run_id"]) due_at = now + gate.debounce - if run.request_key is not None: - # Same reservation shape as admit(): the advisory lock is - # keyed on the flow key, and two admissions can share a - # request key while computing different flow keys, so a plain - # INSERT here can lose a race the lock does not cover and - # raise instead of deduplicating. - cursor = await conn.execute( - "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" - " VALUES (%s, %s, %s) ON CONFLICT DO NOTHING RETURNING run_id", - (run.workflow_id, run.request_key, run.run_id), - ) - if await cursor.fetchone() is None: - cursor = await conn.execute( - "SELECT run_id FROM workflow_dedupe" - " WHERE workflow_id = %s AND request_key = %s", - (run.workflow_id, run.request_key), - ) - existing = await cursor.fetchone() - if existing is not None: - return FlowAdmission("deduplicated", existing["run_id"]) await self._insert_run(conn, run) await self._insert_step(conn, dataclasses.replace(root_step, due_at=due_at)) await self._append_events(conn, run.run_id, events, run.created_at) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 432fdc62be2..264bda0331e 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -2670,7 +2670,11 @@ def __init__(self, db_path: str | Path): # is locked" purely because opening the store ran CREATEs and an # immediate-mode migration it did not need. current = self._db.execute("PRAGMA user_version").fetchone()[0] - if current != SCHEMA_VERSION: + if current < SCHEMA_VERSION: + # Strictly upward: a stamp from a newer binary means a newer + # schema owns this file, and rerunning our DDL would stamp the + # OLDER version over it -- a silent downgrade. Newer schemas are + # additive by policy, so reading them with this binary is safe. self._db.executescript(_SCHEMA) self._migrate() diff --git a/tests/units/workflow/test_drain.py b/tests/units/workflow/test_drain.py index 018f531ac26..5efbd50a99d 100644 --- a/tests/units/workflow/test_drain.py +++ b/tests/units/workflow/test_drain.py @@ -136,3 +136,59 @@ def test_an_unparseable_budget_does_not_stop_the_process_leaving(monkeypatch): """A typo in a deployment variable must not wedge a shutdown.""" monkeypatch.setenv(DRAIN_ENV, "half an hour") assert configured_drain() == pytest.approx(0.0) + + +async def test_cancelling_one_pump_leaves_the_other_pumps_attempt_alone( + forked_registration_context, +): + """A pump's timeout is not another pump's lost work. + + Two callers pump one kernel; each owns the attempt it started. Cancelling + the first must stop only its own attempt -- the second pump is + supervising the other one, and having it yanked from outside turns one + caller's cancellation into unrelated abandoned work. + """ + started_one, release_one = asyncio.Event(), asyncio.Event() + started_two, release_two = asyncio.Event(), asyncio.Event() + flow_one = _flow(started_one, release_one) + + class SecondFlow(rx.State): + __workflow__ = WorkflowConfig(id="drain.second") + + @rx.event(durable=True, trigger=manual(), effect="none") + async def work(self): + """Block until released, then finish. + + Returns: + Completion. + """ + started_two.set() + await release_two.wait() + return rx.complete(result="second survived") + + store = MemoryRunStore() + kernel = WorkflowKernel( + [compile_workflow(flow_one), compile_workflow(SecondFlow)], + store, + max_concurrency=2, + ) + await kernel.start(flow_one.work()) + pump_one = asyncio.create_task(kernel.run_until_idle()) + await asyncio.wait_for(started_one.wait(), timeout=5) + + await kernel.start(SecondFlow.work()) + pump_two = asyncio.create_task(kernel.run_until_idle()) + await asyncio.wait_for(started_two.wait(), timeout=5) + + pump_one.cancel() + await asyncio.gather(pump_one, return_exceptions=True) + + release_two.set() + await asyncio.wait_for(pump_two, timeout=5) + release_one.set() + + runs = await store.list_runs(RunQuery(limit=10)) + second = next(r for r in runs if r.workflow_id == "drain.second") + assert second.status is RunStatus.COMPLETED, ( + "cancelling pump one killed pump two's attempt" + ) diff --git a/tests/units/workflow/test_flow_atomicity.py b/tests/units/workflow/test_flow_atomicity.py index 9879eab0c30..60c96e04c5e 100644 --- a/tests/units/workflow/test_flow_atomicity.py +++ b/tests/units/workflow/test_flow_atomicity.py @@ -237,3 +237,62 @@ async def admit(store: RunStore, key: str = key): assert dispositions == ["deduplicated", "started"], ( f"trial {trial}: {dispositions}" ) + + +async def test_a_duplicate_under_a_different_flow_key_mutates_nothing(store_pair): + """A duplicate admission must leave with zero side effects. + + Two admissions can share a request key while computing different flow + keys -- different advisory locks -- so the loser's dedupe check is not + serialized against the winner. If the dedupe reservation came after the + policy mutations, the duplicate would cancel the OTHER flow key's + incumbent and commit that cancellation on its way out as "deduplicated": + an unrelated live run killed by an event that had already been handled. + """ + a, b = store_pair + for trial in range(TRIALS): + # A live incumbent on flow key B that nothing should ever touch. + incumbent, incumbent_step = _records(f"other-{trial}") + assert ( + await a.admit_flow( + incumbent, + incumbent_step, + (), + FlowGate(singleton_cancel=True), + 1_000_000.0, + ) + ).disposition == "started" + + async def admit(store: RunStore, key: str, trial: int = trial): + """Admit one delivery of the shared event. + + Args: + store: The store to admit through. + key: This admission's flow key. + trial: The trial number. + + Returns: + The admission outcome. + """ + run, step = _records(key) + run = dataclasses.replace(run, request_key=f"shared-{trial}") + return await store.admit_flow( + run, step, (), FlowGate(singleton_cancel=True), 1_000_000.0 + ) + + mine_outcome, other_outcome = await asyncio.gather( + admit(a, f"mine-{trial}"), admit(b, f"other-{trial}") + ) + dispositions = sorted([mine_outcome.disposition, other_outcome.disposition]) + assert dispositions == ["deduplicated", "started"], ( + f"trial {trial}: {dispositions}" + ) + survivor = await a.get_run(incumbent.run_id) + assert survivor is not None + if other_outcome.disposition == "deduplicated": + # The admission on the incumbent's flow key lost the dedupe race, + # so it decided nothing: the incumbent must be untouched. + assert not survivor.cancel_requested, ( + f"trial {trial}: a duplicate cancelled an unrelated incumbent" + ) + assert other_outcome.cancelled == () diff --git a/tests/units/workflow/test_sqlite_frontier.py b/tests/units/workflow/test_sqlite_frontier.py index 1bdaf417605..00ae223f3f6 100644 --- a/tests/units/workflow/test_sqlite_frontier.py +++ b/tests/units/workflow/test_sqlite_frontier.py @@ -99,3 +99,25 @@ async def test_a_due_run_is_found_among_ten_thousand_sleepers(tmp_path): assert due is not None assert abs(due - (NOW + 86_400)) < 1e-6, "the earliest sleeper is the next wake" store.close() + + +def test_a_future_schema_stamp_is_never_downgraded(tmp_path): + """An older binary must not restamp a newer schema as its own. + + Newer schemas are additive by policy, so reading one is safe -- but + rerunning this binary's DDL would overwrite the newer version stamp with + the older one, and the newer binary would then re-migrate a database that + is already ahead of it. + """ + from reflex.workflow.store import SCHEMA_VERSION + + path = tmp_path / "future.db" + store = SqliteRunStore(path) + assert store._db.execute("PRAGMA user_version").fetchone()[0] == SCHEMA_VERSION # pyright: ignore[reportPrivateUsage] + store._db.execute(f"PRAGMA user_version = {SCHEMA_VERSION + 7}") # pyright: ignore[reportPrivateUsage] + store.close() + + reopened = SqliteRunStore(path) + stamp = reopened._db.execute("PRAGMA user_version").fetchone()[0] # pyright: ignore[reportPrivateUsage] + assert stamp == SCHEMA_VERSION + 7, "the newer stamp was downgraded" + reopened.close() From e67ee05f0c4b15e8900b6379457545183bff55c7 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 16:39:59 -0700 Subject: [PATCH 096/121] workflows: durations must be finite The reviewer's time-edge list starts with NaN/infinity durations, and this is the third bug of exactly this class in the engine -- approval expiry and Stripe timestamps had it first. NaN compares false against every bound, so float('nan') sailed past parse_duration's negativity check and poisoned every due-time comparison downstream; infinity turned timers into never-fires. parse_duration is the single choke point every timeout, retry delay, debounce window, rx.after target, and tolerance flows through, so the guard lives there: non-finite is a WorkflowDefinitionError, which since tonight also fails a run on its first attempt instead of retrying. --- news/workflow-finite-durations.bugfix.md | 1 + packages/reflex-base/src/reflex_base/workflow.py | 6 ++++++ tests/units/reflex_base/test_workflow.py | 14 ++++++++++++++ 3 files changed, 21 insertions(+) create mode 100644 news/workflow-finite-durations.bugfix.md diff --git a/news/workflow-finite-durations.bugfix.md b/news/workflow-finite-durations.bugfix.md new file mode 100644 index 00000000000..1c7f98506fa --- /dev/null +++ b/news/workflow-finite-durations.bugfix.md @@ -0,0 +1 @@ +Durations must be finite. `parse_duration(float("nan"))` passed the negativity check — NaN compares false against every bound — and then poisoned every due-time comparison downstream; infinity turned timers into never. Both are refused with a `WorkflowDefinitionError` now, at the single choke point every timeout, retry delay, debounce window, timer, and tolerance flows through. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index e581b1286ba..05337f0fdbc 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -79,6 +79,12 @@ def parse_duration(value: DurationLike, *, param: str = "duration") -> float: else: msg = f"Invalid {param} {value!r}: expected a str, number of seconds, or timedelta." raise WorkflowDefinitionError(msg) + if not math.isfinite(seconds): + # NaN compares false against every bound, so without this it would + # sail past the negativity check and poison every due-time + # comparison downstream; infinity turns timers into never. + msg = f"Invalid {param} {value!r}: duration must be finite." + raise WorkflowDefinitionError(msg) if seconds < 0: msg = f"Invalid {param} {value!r}: duration cannot be negative." raise WorkflowDefinitionError(msg) diff --git a/tests/units/reflex_base/test_workflow.py b/tests/units/reflex_base/test_workflow.py index 01417a3bc2c..b1966bdda54 100644 --- a/tests/units/reflex_base/test_workflow.py +++ b/tests/units/reflex_base/test_workflow.py @@ -336,3 +336,17 @@ def test_hmac_signature_no_longer_claims_stripe(): doc = hmac_signature.__doc__ or "" assert "deliberately **not** a Stripe verifier" in doc assert "stripe_signature" in doc, "the fix must point at the real verifier" + + +def test_durations_must_be_finite(): + """NaN compares false against every bound and infinity means never. + + Both would otherwise pass the negativity check and poison every due-time + comparison downstream -- the third NaN bug of this class in the engine, + after approval expiry and Stripe timestamps. + """ + for poison in (float("nan"), float("inf"), -float("inf")): + with pytest.raises(WorkflowDefinitionError, match=r"finite|negative"): + parse_duration(poison) + assert parse_duration(0) == pytest.approx(0.0) + assert parse_duration("1.5h") == pytest.approx(5400.0) From 53b05fdd695788b389251aba8d055f7ddbd15b0d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 19:15:29 -0700 Subject: [PATCH 097/121] workflows: workers on the default clock derive time from the store The verified reviewer's top remaining blocker: worker wall clocks skew, and every scheduling comparison assumed they did not. A worker running fast saw a peer's live lease as lapsed and reclaimed the claim -- duplicating exactly the work leases exist to prevent -- and admitted schedule occurrences before their time. The store is the one thing every worker shares, so its clock is the authority. RunStore.epoch_time() answers with the store's own time -- clock_timestamp() on Postgres; None from SQLite and memory, whose single-host process clock already is the shared authority. A kernel constructed with the DEFAULT clock measures its offset against that answer, taken against the request midpoint so the measurement is off by at most half a round trip, at startup and again on every recovery pass. All time reads go through the offset, so lease expiry, due times, and occurrence keys use one time base across the fleet; measured offset against a real Postgres 16 was +9.7ms. Injectability is preserved exactly: an explicitly provided clock -- the test harness's virtual time, the dev CLI's fast-forward -- is authoritative as given and never synced, which is why all ~1,560 existing workflow tests pass unchanged. Regression tests pin the sync math, that store time decides due-ness rather than the worker's local clock, that injected clocks are never second-guessed, and that a store with no clock of its own is asked exactly once. --- news/workflow-clock-authority.feature.md | 1 + reflex/workflow/CONTRACT.md | 10 ++ reflex/workflow/kernel.py | 39 +++++- reflex/workflow/postgres.py | 15 +++ reflex/workflow/store.py | 32 +++++ tests/units/workflow/test_clock_authority.py | 123 +++++++++++++++++++ 6 files changed, 219 insertions(+), 1 deletion(-) create mode 100644 news/workflow-clock-authority.feature.md create mode 100644 tests/units/workflow/test_clock_authority.py diff --git a/news/workflow-clock-authority.feature.md b/news/workflow-clock-authority.feature.md new file mode 100644 index 00000000000..6bd2bb11e3f --- /dev/null +++ b/news/workflow-clock-authority.feature.md @@ -0,0 +1 @@ +Workers on the default clock now derive time from the store. Wall clocks skew across a fleet, and every scheduling comparison — lease expiry, due times, schedule occurrence keys — is only safe when the comparands share a clock: a worker running fast could reclaim a peer's live lease (duplicating the work leases exist to protect) and admit schedule occurrences early. A kernel constructed with the default clock measures its offset against the store's clock (`SELECT clock_timestamp()` on Postgres, taken against the request midpoint) at startup and on every recovery pass, bounding skew among workers to half a round trip plus drift per recovery interval. Single-host stores keep the process clock; explicitly injected clocks — the test harness, `--fast-forward` — remain authoritative as given. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 5d0dd46ea21..93e388b3d04 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -233,6 +233,16 @@ Checked in this order at start: queues — a run whose frontier is on an unserved queue waits. Recovery is queue-agnostic (any worker recovers; the reclaimed step is then claimed by the right one). +- **Time authority.** A worker on the default clock derives time from the + store: its offset against the store's clock is measured at startup and on + every recovery pass, so every scheduling comparison — lease expiry, due + times, schedule occurrence keys — uses one time base across the fleet, and + a machine with a fast wall clock can no longer reclaim a peer's live lease + or admit a schedule occurrence early. Skew among synced workers is bounded + by half a round trip plus local drift per recovery interval. Stores that + never leave one host (SQLite, memory) answer that the process clock is the + authority. An explicitly injected clock — the test harness, the dev CLI's + fast-forward — is authoritative as given and never synced. - **Stopping** is not a decision about a run. A worker asked to stop (SIGTERM, Ctrl-C, or an app lifespan ending) stops claiming immediately and gives the attempts it is already running a drain budget — `REFLEX_WORKFLOW_DRAIN` or `reflex workflows worker --drain`, 30s by default — to commit their own outcome. Anything still diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index d838c58fe54..bf0751216a5 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -450,7 +450,17 @@ def __init__( defn.state_cls: defn for defn in self._definitions.values() } self._store = store - self._clock = clock + # A worker on the default clock derives time from the store instead: + # wall clocks skew across a fleet, and a fast worker comparing its own + # time against a peer's lease expiry reclaims a live claim -- the + # duplicate execution leases exist to prevent. The offset is synced + # against the store's clock at startup and on every recovery pass, so + # skew is bounded by one round trip plus local drift per recovery + # interval. An explicitly injected clock (tests, the dev CLI's + # fast-forward) stays authoritative as given and is never synced. + self._store_clock_offset = 0.0 + self._sync_clock_with_store = clock is time.time + self._clock = self._store_time if self._sync_clock_with_store else clock self._rng = rng self._poll_interval = poll_interval self._max_recoveries = max_recoveries @@ -2760,6 +2770,32 @@ async def _cancel_inflight(self) -> None: await asyncio.gather(*tasks, return_exceptions=True) self._prune() + def _store_time(self) -> float: + """The process clock corrected onto the store's time base. + + Returns: + Epoch seconds by the store's clock, to within one sync error. + """ + return time.time() + self._store_clock_offset + + async def _sync_store_clock(self) -> None: + """Re-measure the offset between this process and the store's clock. + + The store's answer is taken against the midpoint of the request, so + the measured offset is off by at most half the round trip. + """ + if not self._sync_clock_with_store: + return + before = time.time() + store_now = await self._store.epoch_time() + after = time.time() + if store_now is None: + # The process clock is the authority for this store; stop asking. + self._sync_clock_with_store = False + self._store_clock_offset = 0.0 + return + self._store_clock_offset = store_now - (before + after) / 2 + async def recover(self) -> int: """Renew this kernel's live claims, then reclaim expired ones. @@ -2770,6 +2806,7 @@ async def recover(self) -> int: Returns: The number of steps recovered. """ + await self._sync_store_clock() await self._renew_leases() now = self._clock() self._next_recovery_at = now + self._recovery_interval diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 6f0036351eb..660b3368d15 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -835,6 +835,21 @@ async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> ) return len(doomed) + async def epoch_time(self) -> float | None: + """The database clock, the one time source every worker shares. + + Returns: + Epoch seconds by the database's clock. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT EXTRACT(EPOCH FROM clock_timestamp())::float8 AS now" + ) + row = await cursor.fetchone() + assert row is not None + return float(row["now"]) + async def claim_next( self, now: float, diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 264bda0331e..1e339f4f37c 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -433,6 +433,22 @@ async def admit_flow( """ ... + async def epoch_time(self) -> float | None: + """The store's own current time, when it has one all workers share. + + Worker wall clocks skew, and every scheduling comparison -- lease + expiry, due times, schedule occurrence keys -- is only safe when the + comparands come from one clock. A store shared by many machines + (Postgres) answers with its database clock so every worker can derive + time from the same authority; a store that never leaves one host + answers None, because the host clock already is that authority. + + Returns: + Epoch seconds by the store's clock, or None when the process + clock is the right authority. + """ + ... + async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> int: """Delete terminal runs not updated since a cutoff, and all their data. @@ -1568,6 +1584,14 @@ async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> self._dedupe.pop((run.workflow_id, run.request_key), None) return len(doomed) + async def epoch_time(self) -> float | None: + """This store never leaves one host, whose clock is the authority. + + Returns: + None: the process clock is correct here. + """ + return None + async def first_active(self, workflow_id: str, flow_key: str) -> RunRecord | None: """Find the oldest run still in flight under a flow-control key. @@ -3066,6 +3090,14 @@ def work() -> int: return await asyncio.to_thread(work) + async def epoch_time(self) -> float | None: + """This store never leaves one host, whose clock is the authority. + + Returns: + None: the process clock is correct here. + """ + return None + async def claim_next( self, now: float, diff --git a/tests/units/workflow/test_clock_authority.py b/tests/units/workflow/test_clock_authority.py new file mode 100644 index 00000000000..2aaaee1d5b8 --- /dev/null +++ b/tests/units/workflow/test_clock_authority.py @@ -0,0 +1,123 @@ +"""Workers on the default clock derive time from the store. + +Wall clocks skew across a fleet. A worker whose clock runs fast compares its +own idea of now against a peer's lease expiry and reclaims a live claim -- +duplicating exactly the work leases exist to protect -- and admits schedule +occurrences before their time. The store is the one thing every worker +shares, so its clock is the authority: a kernel constructed with the default +clock measures its offset against ``store.epoch_time()`` at startup and on +every recovery pass, and reads time through that offset. An explicitly +injected clock (tests, the dev CLI's fast-forward) is authoritative as given +and never synced. +""" + +import time + +import pytest + +from reflex.workflow.kernel import WorkflowKernel +from reflex.workflow.records import RunRecord, RunStatus, StepRecord, StepStatus +from reflex.workflow.store import MemoryRunStore + +SKEW = 120.0 + + +class _SkewedStore(MemoryRunStore): + """A memory store pretending to be a database whose clock differs.""" + + async def epoch_time(self) -> float | None: + """Answer with a clock a fixed distance from this process's. + + Returns: + Epoch seconds, skewed. + """ + return time.time() + SKEW + + +def _due(run_id: str, due_at: float) -> tuple[RunRecord, StepRecord]: + """Build one run whose root comes due at a given store time. + + Args: + run_id: The run identity. + due_at: When the root may be claimed, in store time. + + Returns: + The run and its root slot. + """ + now = time.time() + run = RunRecord( + run_id=run_id, + workflow_id="clock.flow", + definition_digest="d", + status=RunStatus.PENDING, + state={}, + state_version=0, + next_ordinal=1, + created_at=now, + updated_at=now, + ) + step = StepRecord( + run_id=run_id, + ordinal=0, + handler_id="start", + status=StepStatus.READY, + args={}, + due_at=due_at, + origin="root", + queue="default", + created_at=now, + updated_at=now, + ) + return run, step + + +async def test_a_default_clock_kernel_adopts_the_store_clock(): + """After one recovery pass, the kernel reads time by the store's clock.""" + kernel = WorkflowKernel([], _SkewedStore()) + assert kernel._clock() == pytest.approx(time.time(), abs=1.0) # pyright: ignore[reportPrivateUsage] + await kernel.recover() + assert kernel._clock() == pytest.approx(time.time() + SKEW, abs=1.0), ( # pyright: ignore[reportPrivateUsage] + "the offset must be measured, not assumed zero" + ) + + +async def test_store_time_decides_what_is_due_not_the_worker_clock(): + """A timer set in store time fires by store time, from any worker. + + The store clock here runs AHEAD of the process clock, so a worker that + trusted its own time would refuse work that is genuinely due (and with + the skew reversed, would claim work early and recover live leases). The + synced worker claims exactly what the store's clock says is claimable. + """ + store = _SkewedStore() + kernel = WorkflowKernel([], store) + await kernel.recover() + + store_now = time.time() + SKEW + due_run, due_step = _due("due-now", store_now - 5) + future_run, future_step = _due("due-later", store_now + 3600) + await store.admit(due_run, due_step, ()) + await store.admit(future_run, future_step, ()) + + claim = await store.claim_next(kernel._clock()) # pyright: ignore[reportPrivateUsage] + assert claim is not None + assert claim.run.run_id == "due-now" + assert await store.claim_next(kernel._clock()) is None, ( # pyright: ignore[reportPrivateUsage] + "an hour-out timer must not be claimable, whatever the local clock says" + ) + + +async def test_an_injected_clock_is_never_second_guessed(): + """Tests and fast-forward own their clocks; syncing would break both.""" + virtual = 1_000_000.0 + kernel = WorkflowKernel([], _SkewedStore(), clock=lambda: virtual) + await kernel.recover() + assert kernel._clock() == pytest.approx(virtual) # pyright: ignore[reportPrivateUsage] + + +async def test_a_store_with_no_clock_of_its_own_stops_being_asked(): + """Single-host stores answer None once, and the process clock stands.""" + kernel = WorkflowKernel([], MemoryRunStore()) + await kernel.recover() + assert kernel._sync_clock_with_store is False # pyright: ignore[reportPrivateUsage] + assert kernel._clock() == pytest.approx(time.time(), abs=1.0) # pyright: ignore[reportPrivateUsage] From 7b6001e1ad91e3249b3b55b59a16a4e2bc682182 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 19:26:29 -0700 Subject: [PATCH 098/121] workflows: fence run deadlines at the store, in both directions The verified reviewer's second blocker: work could commit after its deadline, and a signal could report success and then be discarded. Both were the same hole. The contract says a run past its deadline finalizes TIMED_OUT once drained -- but nothing guaranteed drained. An attempt that outran cooperative cancellation committed anyway, so a run the caller was told had timed out could quietly become COMPLETED. And a delivery to a past-deadline run was answered 'resolved' even though the continuation could never execute: claims exclude past-deadline runs, so the sweep finalized TIMED_OUT and the recorded decision evaporated. Commit now re-reads the run's deadline inside the same transaction that validates the claim fence -- one added column on the existing claim check, no extra round trip -- and refuses with DeadlinePassedError when it has passed. The kernel abandons the attempt (recorded substeps stand, crash-equivalent), releases the slot so the run drains immediately instead of waiting out a lease, and the sweep's TIMED_OUT is the only reachable outcome. Deliveries to past-deadline runs are refused as 'expired' in all three stores, so 'resolved' is never said of a decision about to be discarded. The contract's failure matrix gains both rows -- caught first by its own bidirectional reason guard, which is exactly the drift it exists to stop. Both fences are regression-tested and fail on the unfenced tree; the full workflow suite passes on all three stores against real Postgres 16. --- news/workflow-deadline-fencing.bugfix.md | 1 + reflex/workflow/CONTRACT.md | 7 ++ reflex/workflow/kernel.py | 29 ++++++ reflex/workflow/postgres.py | 20 +++- reflex/workflow/store.py | 57 ++++++++++- tests/units/workflow/test_kernel.py | 124 +++++++++++++++++++++++ 6 files changed, 230 insertions(+), 8 deletions(-) create mode 100644 news/workflow-deadline-fencing.bugfix.md diff --git a/news/workflow-deadline-fencing.bugfix.md b/news/workflow-deadline-fencing.bugfix.md new file mode 100644 index 00000000000..69b46b7f175 --- /dev/null +++ b/news/workflow-deadline-fencing.bugfix.md @@ -0,0 +1 @@ +Run deadlines are fenced at the store, in both directions. An attempt that outran cooperative cancellation could commit after its run's deadline — recording COMPLETED on a run the caller may already have been told timed out; the commit is now refused inside the store transaction (`deadline_passed`), the attempt is abandoned with its recorded substeps standing, and the sweep's `TIMED_OUT` becomes the only outcome a past-deadline run can reach. Symmetrically, a signal or approval delivered to a past-deadline run was answered `resolved` and then discarded by the timeout sweep; it is refused as `expired`, so the sender is never told a decision was recorded that can never execute. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 93e388b3d04..907a2239506 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -144,6 +144,11 @@ that turn it into effectively-once for side effects are, in order of strength: retry would be a lie. - Run-level `timeout` (`WorkflowConfig.run_timeout`) finalizes the run `TIMED_OUT` once drained; the in-flight attempt is cancelled cooperatively. + The deadline is fenced at commit: an attempt that outruns cooperative + cancellation and tries to commit after the deadline is refused and + abandoned, so "drained" is guaranteed and TIMED_OUT is the *only* outcome a + past-deadline run can reach. Deliveries to a past-deadline run are refused + as `expired` for the same reason. ## 4. Releases and versions @@ -315,6 +320,8 @@ outcome. | during recovery sweep | idempotent; re-run by the next sweep | | worker dies holding N claims | each lease lapses independently; each step recovered independently | | worker asked to stop mid-attempt | attempt gets the drain budget to commit; if it commits, nothing is lost and nothing is spent; if it does not, it is cancelled and the step stays claimed until its lease lapses | +| attempt finishes after the run's deadline passed | the commit is refused (`deadline_passed`, `attempt_abandoned` in history), the slot is released, and the sweep finalizes the run `TIMED_OUT` — a run past its deadline has exactly one outcome, never COMPLETED-after-the-fact. Recorded substeps stand: this is crash-equivalent, not an undo | +| signal or approval arrives for a run past its deadline | refused as `expired` — the continuation can never execute (claims exclude past-deadline runs), so answering "resolved" would record a decision the timeout sweep is about to discard | | store unreachable at commit | attempt abandoned (fence unverifiable); step recovered later; `rx.step` records already made stand | | everything down for an hour | timers/waits/retries fire on restart (due-time semantics); schedule occurrences catch up from the durable cursor, capped at `MAX_SCHEDULE_CATCHUP` per schedule, remainder skipped with a history record | diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index bf0751216a5..dc48e0b2be3 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -60,6 +60,7 @@ from reflex.workflow.steps import SubstepJournal, bind_journal, unbind_journal from reflex.workflow.store import ( Claim, + DeadlinePassedError, DeliveryDisposition, FlowGate, RunStore, @@ -2430,6 +2431,34 @@ async def _commit_outcome( except StaleClaimError: await self._record_abandoned(claim, handler, "fenced_at_commit") return + except DeadlinePassedError: + # The run passed its deadline while this attempt ran. The only + # permitted outcome now is the sweep's TIMED_OUT, so the attempt + # is abandoned -- its recorded substeps stand, crash-equivalent -- + # and its slot is released so the run drains immediately instead + # of waiting out a lease. + abandoned = ( + ( + HistoryEventType.ATTEMPT_ABANDONED, + { + "ordinal": claim.step.ordinal, + "epoch": claim.step.epoch, + "worker": self._worker_id, + "effect": handler.effect, + "reason": "deadline_passed", + }, + ), + ) + with contextlib.suppress(StaleClaimError): + await self._store.release_claim( + claim, + status=StepStatus.CANCELLED, + events=abandoned, + now=self._clock(), + ) + self._notify(claim.run, abandoned) + self._wakeup.set() + return self._notify(claim.run, completion.events) for child_run, child_step in completion.children: # The store recorded each child's admission inside the commit; diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 660b3368d15..e6d8f513833 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -38,6 +38,7 @@ FlowGate, StaleClaimError, _child_admission_events, + _fence_deadline, ) try: @@ -586,19 +587,22 @@ async def _frontier(self, conn: Connection, run_id: str) -> StepRecord | None: row = await cursor.fetchone() return None if row is None else _step_from_row(row) - async def _check_claim(self, conn: Connection, claim: Claim) -> None: + async def _check_claim(self, conn: Connection, claim: Claim) -> float | None: """Validate that a claim still owns its step and state version. Args: conn: The connection inside an open transaction. claim: The claim to validate. + Returns: + The run's deadline, when it has one. + Raises: StaleClaimError: If the claim was fenced. """ cursor = await conn.execute( "SELECT s.status AS step_status, s.epoch AS epoch," - " r.state_version AS state_version" + " r.state_version AS state_version, r.deadline AS deadline" " FROM workflow_steps s JOIN workflow_runs r ON r.run_id = s.run_id" " WHERE s.run_id = %s AND s.ordinal = %s", (claim.run.run_id, claim.step.ordinal), @@ -614,6 +618,7 @@ async def _check_claim(self, conn: Connection, claim: Claim) -> None: f"Claim on run {claim.run.run_id} step {claim.step.ordinal} was fenced." ) raise StaleClaimError(msg) + return row["deadline"] async def admit( self, @@ -994,7 +999,10 @@ async def commit( pool = await self._open() async with pool.connection() as conn, conn.transaction(): await self._lock_run(conn, claim.run.run_id) - await self._check_claim(conn, claim) + deadline = await self._check_claim(conn, claim) + # Past the deadline the only permitted outcome is TIMED_OUT, and + # that is the sweep's transition, not this attempt's. + _fence_deadline(claim.run.run_id, deadline, now) await conn.execute( "UPDATE workflow_steps SET status = %s, attempts = attempts + %s," " due_at = %s, lease_expires_at = 0, error = %s, updated_at = %s" @@ -1265,7 +1273,7 @@ async def _apply_arrival( """ wait_key = f"join:{ordinal}" cursor = await conn.execute( - "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", + "SELECT status, deadline FROM workflow_runs WHERE run_id = %s FOR UPDATE", (run_id,), ) run_row = await cursor.fetchone() @@ -1273,6 +1281,10 @@ async def _apply_arrival( return "unknown_run" if run_row["status"] in _TERMINAL_RUNS: return "run_terminal" + if run_row["deadline"] is not None and run_row["deadline"] <= now: + # A past-deadline run can never execute the continuation; + # saying "resolved" would be a lie the sweep then discards. + return "expired" cursor = await conn.execute( "SELECT 1 FROM workflow_inbox" " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 1e339f4f37c..e28ecb1eb58 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -59,6 +59,17 @@ ] +class DeadlinePassedError(WorkflowRuntimeError): + """A commit arrived after its run's deadline. + + The contract promises a run past its deadline finalizes TIMED_OUT, which + is only unambiguous if nothing can commit afterwards: an attempt that + beats the timeout sweep would otherwise record COMPLETED on a run the + caller was told had timed out. The attempt's recorded substeps stand -- + this is crash-equivalent, not an undo. + """ + + class StaleClaimError(WorkflowRuntimeError): """Raised when a commit no longer owns its claim and must be discarded.""" @@ -1174,6 +1185,9 @@ async def commit( """ async with self._lock: run, steps = self._check_claim(claim) + # Past the deadline the only permitted outcome is TIMED_OUT, + # and that is the sweep's transition, not this attempt's. + _fence_deadline(run.run_id, run.deadline, now) step = steps[claim.step.ordinal] steps[step.ordinal] = dataclasses.replace( step, @@ -1291,6 +1305,12 @@ async def deliver( return "unknown_run" if run.status in TERMINAL_RUN_STATUSES: return "run_terminal" + if run.deadline is not None and run.deadline <= now: + # The run can never execute a continuation: claims exclude + # past-deadline runs and the sweep will finalize TIMED_OUT. + # Answering "resolved" here would tell the sender their + # decision was recorded when it is about to be discarded. + return "expired" inbox = self._inbox.setdefault(run_id, {}) if (run_id, wait_key, dedupe_key) in inbox: return "duplicate" @@ -2526,6 +2546,22 @@ def _step_from_row(row: sqlite3.Row) -> StepRecord: ) +def _fence_deadline(run_id: str, deadline: float | None, now: float) -> None: + """Refuse a commit for a run already past its deadline. + + Args: + run_id: The committing run. + deadline: Its deadline, when it has one. + now: Current time in epoch seconds. + + Raises: + DeadlinePassedError: When the deadline has passed. + """ + if deadline is not None and deadline <= now: + msg = f"Run {run_id} passed its deadline before commit." + raise DeadlinePassedError(msg) + + def _child_admission_events( child_run: RunRecord, child_step: StepRecord ) -> tuple[tuple[HistoryEventType, dict[str, Any]], ...]: @@ -3187,18 +3223,21 @@ def work(): return await asyncio.to_thread(work) - def _check_claim(self, claim: Claim) -> None: + def _check_claim(self, claim: Claim) -> float | None: """Validate that a claim still owns its step and state version. Args: claim: The claim to validate. + Returns: + The run's deadline, when it has one. + Raises: StaleClaimError: If the claim was fenced. """ row = self._db.execute( "SELECT s.status AS step_status, s.epoch AS epoch," - " r.state_version AS state_version" + " r.state_version AS state_version, r.deadline AS deadline" " FROM workflow_steps s JOIN workflow_runs r ON r.run_id = s.run_id" " WHERE s.run_id = ? AND s.ordinal = ?", (claim.run.run_id, claim.step.ordinal), @@ -3213,6 +3252,7 @@ def _check_claim(self, claim: Claim) -> None: f"Claim on run {claim.run.run_id} step {claim.step.ordinal} was fenced." ) raise StaleClaimError(msg) + return row["deadline"] async def renew_lease( self, @@ -3308,7 +3348,10 @@ def work() -> None: with self._lock: self._db.execute("BEGIN IMMEDIATE") try: - self._check_claim(claim) + deadline = self._check_claim(claim) + # Past the deadline the only permitted outcome is + # TIMED_OUT, and that is the sweep's, not this one's. + _fence_deadline(claim.run.run_id, deadline, now) self._db.execute( "UPDATE workflow_steps SET status = ?, attempts = attempts + ?," " due_at = ?, lease_expires_at = 0, error = ?, updated_at = ?" @@ -3476,7 +3519,8 @@ def work(): self._db.execute("BEGIN IMMEDIATE") try: row = self._db.execute( - "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + "SELECT status, deadline FROM workflow_runs WHERE run_id = ?", + (run_id,), ).fetchone() if row is None: self._db.execute("ROLLBACK") @@ -3484,6 +3528,11 @@ def work(): if row["status"] in terminal: self._db.execute("ROLLBACK") return "run_terminal" + if row["deadline"] is not None and row["deadline"] <= now: + # A past-deadline run can never execute the + # continuation; saying "resolved" would be a lie. + self._db.execute("ROLLBACK") + return "expired" seen = self._db.execute( "SELECT 1 FROM workflow_inbox" " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", diff --git a/tests/units/workflow/test_kernel.py b/tests/units/workflow/test_kernel.py index 46751fd0133..4a17da0f133 100644 --- a/tests/units/workflow/test_kernel.py +++ b/tests/units/workflow/test_kernel.py @@ -8,6 +8,7 @@ from reflex_base.utils.exceptions import WorkflowRuntimeError from reflex_base.workflow import ( Retry, + Signal, TransientWorkflowError, WorkflowConfig, after, @@ -16,6 +17,7 @@ hmac_signature, manual, needs_attention, + wait_for, webhook, ) @@ -27,6 +29,26 @@ from reflex.workflow.testing import WorkflowTestHarness +class _Clock: + """A manually advanced epoch-seconds clock.""" + + def __init__(self, now: float): + """Start the clock. + + Args: + now: The starting time in epoch seconds. + """ + self.now = now + + def __call__(self) -> float: + """Read the clock. + + Returns: + The current time in epoch seconds. + """ + return self.now + + class Payment(BaseModel): """Typed payload for kernel tests.""" @@ -803,3 +825,105 @@ async def start(self, n: int): assert statuses == ["COMPLETED"] * 8, ( f"run_until_idle returned with live attempts: {statuses}" ) + + +async def test_work_cannot_commit_after_its_deadline(forked_registration_context): + """A run past its deadline has one outcome, and it is not COMPLETED. + + The handler starts before the deadline, outruns cooperative cancellation, + and tries to commit after it. Without the fence the commit lands and the + caller -- who may already have been told the run timed out -- sees a run + that completed after its own deadline. + """ + release = asyncio.Event() + + class Deadlined(rx.State): + __workflow__ = WorkflowConfig(id="kernel.deadline_fence", run_timeout="10s") + + @rx.event(durable=True, trigger=manual(), effect="none") + async def start(self): + """Block until the test has moved the clock past the deadline. + + Returns: + Completion that must never land. + """ + await release.wait() + return rx.complete(result="beat the sweep") + + clock = _Clock(1_000_000.0) + store = MemoryRunStore() + definition = compile_workflow(Deadlined) + kernel = WorkflowKernel([definition], store, clock=clock) + started = await kernel.start(Deadlined.start()) + assert started.run_id is not None + pump = asyncio.create_task(kernel.run_until_idle()) + await asyncio.sleep(0.05) + + # The deadline passes while the attempt is still running. + clock.now += 60 + release.set() + await asyncio.wait_for(pump, timeout=10) + await kernel.run_until_idle() + + snapshot = await kernel.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.TIMED_OUT, ( + f"a past-deadline run completed anyway: {snapshot.status}" + ) + assert snapshot.result != "beat the sweep" + + +async def test_a_delivery_to_a_past_deadline_run_is_refused( + forked_registration_context, +): + """The answer resolved must not describe a decision the sweep discards.""" + + class Waits(rx.State): + __workflow__ = WorkflowConfig(id="kernel.deadline_wait", run_timeout="10s") + + decided = Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Wait for a decision. + + Returns: + The wait. + """ + return wait_for( + Waits.decided, then=Waits.done, timeout="1h", on_timeout=Waits.lapse + ) + + @rx.event(durable=True, effect="none") + def done(self, decision: dict): + """Record the decision. + + Args: + decision: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result=decision) + + @rx.event(durable=True, effect="none") + def lapse(self): + """Nobody answered. + + Returns: + Failure. + """ + return rx.fail(reason="lapsed") + + clock = _Clock(1_000_000.0) + store = MemoryRunStore() + kernel = WorkflowKernel([compile_workflow(Waits)], store, clock=clock) + started = await kernel.start(Waits.start()) + assert started.run_id is not None + await kernel.run_until_idle() + + clock.now += 60 # past the run deadline, before the wait's own timeout + disposition = await kernel.signal(started.run_id, Waits.decided({"ok": True})) + assert disposition == "expired", ( + f"a doomed run's wait answered {disposition!r} instead of refusing" + ) From 8757199a6be77655c73ea1d5fe7083f56077634b Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 19:31:48 -0700 Subject: [PATCH 099/121] workflows: close the bounded time edges Three from the reviewer's list, each verified before fixing. Retry backoff overflowed: multiplier ** 499 raises OverflowError, and delay_for_attempt runs inside the kernel's completion path, so a long-retrying step would have broken the worker loop rather than the run. The overflow saturates to max_delay, which is what a delay astronomically past the cap means anyway. Schedule catch-up loss was silent AND the contract lied about it: the failure matrix claimed the skipped remainder got a history record, and no such record was ever written -- there is no run to attach one to. A worker back from a week of downtime now warns with the count and the window when it skips past the cap, the cursor jump is no longer mistakable for coverage, and the contract describes what actually happens. My bidirectional reason guard could not catch this drift -- no reason string involved -- which is a useful reminder of its limits. The sync-handler drain question resolves to wording, not code: a thread cannot be interrupted, so no drain budget can bound how fast a sync handler stops -- the same fact that makes timeout= a compile error on sync handlers. The contract now says the budget bounds how long cancellable work is waited for, never how fast a thread can be made to stop, and points long sync work at async-around-rx.step. --- news/workflow-time-edges.bugfix.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 9 ++- reflex/workflow/CONTRACT.md | 8 ++- reflex/workflow/kernel.py | 19 ++++++- tests/units/reflex_base/test_workflow.py | 11 ++++ tests/units/workflow/test_schedules.py | 57 +++++++++++++++++++ 6 files changed, 100 insertions(+), 5 deletions(-) create mode 100644 news/workflow-time-edges.bugfix.md diff --git a/news/workflow-time-edges.bugfix.md b/news/workflow-time-edges.bugfix.md new file mode 100644 index 00000000000..7366a0291a9 --- /dev/null +++ b/news/workflow-time-edges.bugfix.md @@ -0,0 +1 @@ +Two time edges from the production review. Exponential backoff saturates at `max_delay` instead of raising `OverflowError` around attempt 500 — the power ran inside the kernel's completion path, so the overflow would have broken the worker, not the run. And schedule catch-up loss is no longer silent: occurrences beyond `MAX_SCHEDULE_CATCHUP` are skipped with a warning naming the count and window, and the contract's failure matrix now describes what actually happens (the old text claimed a history record that never existed — there is no run to attach one to). The contract also states plainly that a synchronous handler cannot be cancelled by any drain budget: a thread cannot be interrupted, which is the same reason `timeout=` is a compile error on sync handlers. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 05337f0fdbc..60c0d620f3c 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -169,7 +169,14 @@ def delay_for_attempt(self, failed_attempts: int) -> float: """ initial = parse_duration(self.initial_delay, param="Retry.initial_delay") maximum = parse_duration(self.max_delay, param="Retry.max_delay") - return min(initial * self.multiplier ** (failed_attempts - 1), maximum) + try: + return min(initial * self.multiplier ** (failed_attempts - 1), maximum) + except OverflowError: + # Around attempt ~500 with the default multiplier the power + # overflows a float. A delay astronomically past the cap is the + # cap -- and this raises inside the kernel's completion path, not + # the handler, so letting it escape breaks the worker, not the run. + return maximum def is_retryable(self, error: BaseException) -> bool: """Whether an exception consumes a business attempt and may retry. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 907a2239506..880d1d8e131 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -256,6 +256,12 @@ Checked in this order at start: A claim is never released early, because cancelling an attempt does not stop work it handed to a thread, and the lease is what keeps a peer off it. A drained attempt costs nothing; a cancelled one costs one recovery. + A **synchronous** handler cannot be cancelled at all — a thread cannot be + interrupted, which is the same reason `timeout=` is a compile error on sync + handlers — so a stopping worker holds until the sync call in flight + returns, however short the drain budget: the budget bounds how long + *cancellable* work is waited for, never how fast a thread can be made to + stop. Sync handlers doing long work should be async around `rx.step`. - **Clients are not workers.** A process that opens `rx.workflows.connect(...)` can admit runs, read them, signal and cancel them, and executes nothing: it claims no step and runs no handler. Only a @@ -323,7 +329,7 @@ outcome. | attempt finishes after the run's deadline passed | the commit is refused (`deadline_passed`, `attempt_abandoned` in history), the slot is released, and the sweep finalizes the run `TIMED_OUT` — a run past its deadline has exactly one outcome, never COMPLETED-after-the-fact. Recorded substeps stand: this is crash-equivalent, not an undo | | signal or approval arrives for a run past its deadline | refused as `expired` — the continuation can never execute (claims exclude past-deadline runs), so answering "resolved" would record a decision the timeout sweep is about to discard | | store unreachable at commit | attempt abandoned (fence unverifiable); step recovered later; `rx.step` records already made stand | -| everything down for an hour | timers/waits/retries fire on restart (due-time semantics); schedule occurrences catch up from the durable cursor, capped at `MAX_SCHEDULE_CATCHUP` per schedule, remainder skipped with a history record | +| everything down for an hour | timers/waits/retries fire on restart (due-time semantics); schedule occurrences catch up from the durable cursor, capped at `MAX_SCHEDULE_CATCHUP` per schedule; a remainder beyond the cap is skipped with a **warning naming the count and window** (there is no run to attach history to), and can be started by hand | ### Failures that are not crashes diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index dc48e0b2be3..d965cf5846e 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -2300,9 +2300,22 @@ async def _admit_due_schedules(self, now: float) -> int: stored = await self._store.read_schedule_cursor(key) cursor = stored if stored is not None else self._started_at self._schedule_cursor[key] = cursor - for occurrence in schedule.occurrences_between( - cursor, now, limit=MAX_SCHEDULE_CATCHUP - ): + occurrences = schedule.occurrences_between( + cursor, now, limit=MAX_SCHEDULE_CATCHUP + 1 + ) + if len(occurrences) > MAX_SCHEDULE_CATCHUP: + # The cursor is about to jump over these. Silently losing + # scheduled work reads as "covered" when it was not; the + # operator gets the count and the window, and can start the + # missed occurrences by hand if they matter. + occurrences = occurrences[:MAX_SCHEDULE_CATCHUP] + console.warn( + f"Schedule {key} missed more than {MAX_SCHEDULE_CATCHUP} " + f"occurrences between {cursor:.0f} and {now:.0f}; catching " + f"up the first {MAX_SCHEDULE_CATCHUP} and skipping the " + "rest. Start any that matter with rx.workflows.start()." + ) + for occurrence in occurrences: result = await self.start( getattr(defn.state_cls, handler.name), request_key=f"schedule:{key}:{int(occurrence)}", diff --git a/tests/units/reflex_base/test_workflow.py b/tests/units/reflex_base/test_workflow.py index b1966bdda54..28ae468bd74 100644 --- a/tests/units/reflex_base/test_workflow.py +++ b/tests/units/reflex_base/test_workflow.py @@ -350,3 +350,14 @@ def test_durations_must_be_finite(): parse_duration(poison) assert parse_duration(0) == pytest.approx(0.0) assert parse_duration("1.5h") == pytest.approx(5400.0) + + +def test_backoff_saturates_instead_of_overflowing(): + """Attempt five hundred is 'the cap', not an OverflowError in the kernel. + + delay_for_attempt runs inside the kernel's completion path, so an + exception here breaks the worker, not the run. + """ + policy = Retry(max_attempts=10_000, initial_delay="1s", max_delay="1h") + assert policy.delay_for_attempt(500) == pytest.approx(3600.0) + assert policy.delay_for_attempt(9_999) == pytest.approx(3600.0) diff --git a/tests/units/workflow/test_schedules.py b/tests/units/workflow/test_schedules.py index 7332af7129b..60db9bffe17 100644 --- a/tests/units/workflow/test_schedules.py +++ b/tests/units/workflow/test_schedules.py @@ -13,6 +13,27 @@ from reflex.workflow.store import MemoryRunStore from reflex.workflow.testing import WorkflowTestHarness + +class _Clock: + """A manually advanced epoch-seconds clock.""" + + def __init__(self, now: float): + """Start the clock. + + Args: + now: The starting time in epoch seconds. + """ + self.now = now + + def __call__(self) -> float: + """Read the clock. + + Returns: + The current time in epoch seconds. + """ + return self.now + + # A Tuesday at 12:00 UTC, chosen so quarter-hour schedules are 15 minutes away. START = dt.datetime(2026, 8, 18, 12, 0, tzinfo=dt.UTC).timestamp() @@ -197,3 +218,39 @@ def tick(self): await resumed.advance("1m") assert len(fired) > first_count, "the restarted worker skipped the downtime" + + +async def test_skipped_catchup_is_counted_out_loud(forked_registration_context, capsys): + """Occurrences the cap drops are named, never silently lost. + + A worker down for a week comes back to hundreds of missed quarter-hour + occurrences; it catches up the cap's worth and jumps the cursor. Without + the warning, the jump reads as "covered" and the missing runs are only + discovered by whoever needed their output. + """ + from reflex.workflow.kernel import MAX_SCHEDULE_CATCHUP, WorkflowKernel + + class Nightly(rx.State): + __workflow__ = WorkflowConfig(id="sched.lossy") + + @rx.event(durable=True, effect="none", trigger=schedule("*/15 * * * *")) + def tick(self): + """Fire on the quarter hour. + + Returns: + Completion. + """ + return rx.complete(result=None) + + clock = _Clock(1_000_000.0) + store = MemoryRunStore() + kernel = WorkflowKernel([compile_workflow(Nightly)], store, clock=clock) + await kernel.run_until_idle() + + clock.now += 7 * 24 * 3600 # a week of downtime + admitted = await kernel._admit_due_schedules(clock.now) # pyright: ignore[reportPrivateUsage] + assert admitted == MAX_SCHEDULE_CATCHUP + err = capsys.readouterr() + assert "missed more than" in err.out + err.err, ( + "the skipped remainder must be named, not silently jumped over" + ) From 738e54ba399d85795c43d4fa36a6c9150af64077 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Wed, 19 Aug 2026 19:37:29 -0700 Subject: [PATCH 100/121] workflows: refuse lossy run data, type-check the HTTP boundary The bounded half of the reviewer's typed-durability item, shippable without the wire-format decision the rest of it needs. Decimal('10.10') was stored as float 10.1: the serializer registry converts it, so the type people reach for to avoid precision loss lost precision silently -- in a refund handler that is a money bug with no error anywhere on its path. bytes and bytearray became lists of integers nothing ever turns back into bytes. Both are refused at record time with errors that name the fix (str+Decimal or integer minor units; explicit base64), and since TypeError is a bug-class error the attempt fails once instead of burning retries. Tuples and sets still become lists -- that is JSON's shape, round-trips losslessly enough, and refusing it would break every workflow returning a tuple. The HTTP start endpoint accepted arguments the handler's signature refuses -- a 202 and a poison run instead of the 400 that names the caller's bug. Supplied arguments now validate against the handler's declared type hints; an exotic hint pydantic cannot adapt skips validation rather than blaming the caller for it. What remains of typed durability -- a type-preserving state encoding and schema evolution -- changes the wire format and intersects the API freeze, and stays a design decision rather than a patch. --- news/workflow-typed-boundaries.bugfix.md | 1 + reflex/workflow/api.py | 38 ++++++++++++++++++++++++ reflex/workflow/serde.py | 21 +++++++++++++ tests/units/workflow/test_api.py | 16 ++++++++++ tests/units/workflow/test_serde.py | 21 +++++++++++++ 5 files changed, 97 insertions(+) create mode 100644 news/workflow-typed-boundaries.bugfix.md diff --git a/news/workflow-typed-boundaries.bugfix.md b/news/workflow-typed-boundaries.bugfix.md new file mode 100644 index 00000000000..a0621d1301a --- /dev/null +++ b/news/workflow-typed-boundaries.bugfix.md @@ -0,0 +1 @@ +Two typed-durability boundaries hardened without changing the wire format. `Decimal` and `bytes`/`bytearray` are refused as run data with teaching errors instead of being silently corrupted — the serializer registry stored `Decimal("10.10")` as float `10.1` (a money bug with no error anywhere) and bytes as a list of integers nothing ever turns back; a `TypeError` at record time is also a bug-class error, so the attempt fails on its first try. The HTTP start endpoint validates supplied arguments against the handler's declared types and answers 400 naming the mismatched argument, instead of 202 plus a run whose first attempt can only raise. The remaining typed-durability work — a type-preserving state encoding and schema evolution — changes the wire format and stays a design decision. diff --git a/reflex/workflow/api.py b/reflex/workflow/api.py index 0aa01debf4e..50c6badea0e 100644 --- a/reflex/workflow/api.py +++ b/reflex/workflow/api.py @@ -67,6 +67,35 @@ def _authorized(request: Request, token: str) -> bool: return hmac.compare_digest(presented, token) +def _mistyped_args(handler: Any, args: dict[str, Any]) -> list[str]: + """Check supplied arguments against the handler's declared types. + + Args: + handler: The resolved handler definition. + args: The caller-supplied arguments. + + Returns: + One message per argument that cannot validate, empty when all fit. + """ + from pydantic import TypeAdapter, ValidationError + + problems: list[str] = [] + for name, value in args.items(): + hint = handler.type_hints.get(name) + if hint is None: + continue + try: + TypeAdapter(hint).validate_python(value) + except ValidationError: + problems.append( + f"{name!r} does not validate as {getattr(hint, '__name__', hint)}" + ) + except Exception: + # An exotic hint pydantic cannot adapt is not the caller's fault. + pass + return problems + + def start_endpoint( runtime: WorkflowRuntime, token: str ) -> Callable[[Request], Coroutine[Any, Any, JSONResponse]]: @@ -141,6 +170,15 @@ async def endpoint(request: Request) -> JSONResponse: {"error": "args must be a JSON object"}, status_code=400 ) args = raw_args or {} + mistyped = _mistyped_args(handler, args) + if mistyped: + # Admitting a payload the handler's signature refuses creates a + # run whose first attempt can only raise; the caller gets a 202 + # and a poison run instead of the 400 that names their bug. + return JSONResponse( + {"error": f"arguments do not match the handler: {mistyped}"}, + status_code=400, + ) missing = sorted(unbound_params(handler, set(args))) if missing: # Admitting this would create a run that cannot possibly run: the diff --git a/reflex/workflow/serde.py b/reflex/workflow/serde.py index f4fe2fb13f5..6b2f95cfc61 100644 --- a/reflex/workflow/serde.py +++ b/reflex/workflow/serde.py @@ -8,6 +8,7 @@ from __future__ import annotations +import decimal import json from typing import Any @@ -26,6 +27,26 @@ def _strict_default(value: Any) -> Any: Raises: TypeError: If no serializer is registered for the value's type. """ + if isinstance(value, decimal.Decimal): + # The serializer registry would hand back a float, silently losing + # precision on exactly the type people reach for to avoid losing + # precision -- Decimal("10.10") replaying as 10.1 in a refund handler + # is a money bug with no error anywhere. Refusing names the fix. + msg = ( + "Decimal is not valid run data: it would be stored as a float and " + "lose precision silently. Store str(value) and reconstruct with " + "Decimal(...), or keep integer minor units (cents)." + ) + raise TypeError(msg) + if isinstance(value, (bytes, bytearray, memoryview)): + # The registry turns these into a list of ints, which nothing ever + # turns back. + msg = ( + f"{type(value).__name__} is not valid run data: it would be " + "stored as a list of integers. Encode explicitly, e.g. " + "base64.b64encode(value).decode()." + ) + raise TypeError(msg) serialized = serializers.serialize(value) if serialized is None: msg = ( diff --git a/tests/units/workflow/test_api.py b/tests/units/workflow/test_api.py index b21689bef3a..8f630218262 100644 --- a/tests/units/workflow/test_api.py +++ b/tests/units/workflow/test_api.py @@ -255,3 +255,19 @@ def test_a_falsey_non_object_args_is_refused(client): response = client.post(START_ROUTE, content=body, headers=_auth()) assert response.status_code == 400, response.text assert "JSON object" in response.json()["error"] + + +def test_a_mistyped_argument_is_refused_at_the_boundary(client): + """A payload the signature refuses is the caller's 400, not a poison run. + + Args: + client: The test client. + """ + body = json.dumps({ + "workflow": "api.orders", + "handler": "place", + "args": {"order": {"nested": "object"}}, + }) + response = client.post(START_ROUTE, content=body, headers=_auth()) + assert response.status_code == 400, response.text + assert "'order'" in response.json()["error"] diff --git a/tests/units/workflow/test_serde.py b/tests/units/workflow/test_serde.py index 537609c4302..e8890f261c4 100644 --- a/tests/units/workflow/test_serde.py +++ b/tests/units/workflow/test_serde.py @@ -69,3 +69,24 @@ def test_an_unserializable_object_is_refused(): """A value no serializer handles fails here, not at the store.""" with pytest.raises(TypeError): to_run_data({"handle": object()}) + + +def test_decimal_is_refused_rather_than_silently_truncated(): + """Decimal("10.10") must never replay as 10.1. + + The serializer registry hands back a float, losing precision on exactly + the type people reach for to avoid losing precision -- a silent money bug. + Refusing at record time names the fix. + """ + from decimal import Decimal + + with pytest.raises(TypeError, match="Decimal"): + to_run_data({"amount": Decimal("10.10")}) + + +def test_bytes_are_refused_rather_than_becoming_integer_lists(): + """Raw bytes stored as [114, 97, 119] never come back as bytes.""" + with pytest.raises(TypeError, match="base64"): + to_run_data({"blob": b"raw"}) + with pytest.raises(TypeError, match="bytearray"): + to_run_data({"blob": bytearray(b"raw")}) From 5843481f516ed53d5bc7cd02b8874ff04544675a Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 10:36:15 -0700 Subject: [PATCH 101/121] workflows: branches close with their parent Cancelling a rollout left its regional deploys deploying. Two unrelated real-world scenarios -- a CI/CD rollback and a fleet rollout -- hit this as a bug, and they are right: an operator presses cancel to stop the blast radius, and a button that stops only the bookkeeping run has not done the one thing it exists to do. Section 5 had chosen the other side deliberately ("delegation is not ownership"); that default is wrong and this changes it before anyone depends on it. Any terminal transition of a run -- cancelled, failed, timed out, force-finalized, completed -- now requests cancellation of every branch it fanned out that is still running, in the SAME store transaction. Not follow-up: a worker that dies mid-follow-up is exactly the case where the deploys keep going. rx.parallel(..., parent_close="abandon") keeps the old behaviour for work that should genuinely outlive its starter. Each level closes only its own branches. That looks like it could deadlock -- a tier waiting on shards that only get closed once the tier finalizes -- but a run blocked on its own join holds no claim, so it is control-pending the moment it is marked, finalizes, and closes the level beneath it. A three-level test pins that. Also backstops mode="first" race losers, whose cancellation was best-effort follow-up from one worker; a loser is still not fenced at commit while the parent lives on, so the side-effect warning stands. The conformance suite -- the mechanism that is supposed to keep the three stores identical -- was only wired to memory and sqlite. Postgres, the one store whose SQL is hand-written, was not in it. It now joins whenever REFLEX_TEST_POSTGRES is set; all 48 checks pass against a real server, including the two new cascade checks. Overriding the harness store fixture drops the run from 432 cases to 144 by not crossing two independent store parameters. The old all-mode test that asserted branches were left alone now runs under parent_close="abandon", where the bug it actually guards -- a tombstoned join misread as a decided race -- still shows. --- news/workflow-parent-close.bugfix.md | 3 + .../reflex-base/src/reflex_base/workflow.py | 25 +- reflex/workflow/CONTRACT.md | 29 +- reflex/workflow/conformance.py | 67 ++++ reflex/workflow/kernel.py | 10 +- reflex/workflow/postgres.py | 49 ++- reflex/workflow/records.py | 3 + reflex/workflow/store.py | 87 +++++- tests/units/workflow/test_cascade.py | 290 ++++++++++++++++++ tests/units/workflow/test_conformance.py | 51 ++- tests/units/workflow/test_parallel.py | 26 +- 11 files changed, 609 insertions(+), 31 deletions(-) create mode 100644 news/workflow-parent-close.bugfix.md create mode 100644 tests/units/workflow/test_cascade.py diff --git a/news/workflow-parent-close.bugfix.md b/news/workflow-parent-close.bugfix.md new file mode 100644 index 00000000000..89bc94ac606 --- /dev/null +++ b/news/workflow-parent-close.bugfix.md @@ -0,0 +1,3 @@ +Fan-out branches now stop when their parent does. Cancelling, failing, or timing out a run requests cancellation of every branch it fanned out that is still running, in the same store transaction as the terminal transition — an operator cancelling a rollout stops the regional deploys it started, rather than watching them deploy on. Previously children were always abandoned, which two independent real-world scenarios (a CI/CD rollback and a fleet rollout) hit as a bug. The close walks the whole tree: each level closes its own branches, and a branch blocked on its own join holds no claim, so it finalizes at once and closes the level beneath it. `rx.parallel(..., parent_close="abandon")` keeps the old behaviour for work that genuinely should outlive its starter. This also backstops `mode="first"` race losers, whose cancellation was previously best-effort follow-up that died with the worker that sent it. + +The store conformance suite now runs against Postgres whenever `REFLEX_TEST_POSTGRES` is set. It is the mechanism that keeps the three stores behaving identically, and the one store where the SQL is hand-written was not in it. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 60c0d620f3c..0b94fdcb4b4 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -1118,11 +1118,15 @@ class Parallel: then: Handler that receives the list of branch results. mode: ``"all"`` continues once every branch has reported; ``"first"`` continues as soon as one has, and cancels the rest. + parent_close: What becomes of branches still running when the parent + reaches a terminal state -- ``"cancel"`` stops them, ``"abandon"`` + lets them run on. """ branches: tuple[Any, ...] then: Any mode: Literal["all", "first"] = "all" + parent_close: Literal["cancel", "abandon"] = "cancel" def __post_init__(self): """Validate the fan-out. @@ -1140,10 +1144,19 @@ def __post_init__(self): if self.mode not in ("all", "first"): msg = f'parallel() mode must be "all" or "first", got {self.mode!r}.' raise WorkflowDefinitionError(msg) + if self.parent_close not in ("cancel", "abandon"): + msg = ( + 'parallel() parent_close must be "cancel" or "abandon", got ' + f"{self.parent_close!r}." + ) + raise WorkflowDefinitionError(msg) def parallel( - *branches: Any, then: Any, mode: Literal["all", "first"] = "all" + *branches: Any, + then: Any, + mode: Literal["all", "first"] = "all", + parent_close: Literal["cancel", "abandon"] = "cancel", ) -> Parallel: """Run branches concurrently, then continue with all their results. @@ -1162,15 +1175,23 @@ def parallel( Pass ``mode="first"`` to race them instead: the run continues as soon as one branch reports, and the others are cancelled. + Branches stop when the parent does. If the parent is cancelled, fails, or + times out with branches still running, those branches are cancelled too -- + an operator stopping a rollout stops the regional deploys it started. Pass + ``parent_close="abandon"`` when a branch really is delegated work that + should outlive its starter. + Args: branches: Root events to run concurrently. then: Handler to run once the fan-out is satisfied. mode: ``"all"`` to wait for every branch, ``"first"`` to race them. + parent_close: ``"cancel"`` stops branches still running when the + parent finishes; ``"abandon"`` lets them run on. Returns: The control return value. """ - return Parallel(branches=branches, then=then, mode=mode) + return Parallel(branches=branches, then=then, mode=mode, parent_close=parent_close) @dataclasses.dataclass(frozen=True, slots=True) diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 880d1d8e131..b9975405adb 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -180,12 +180,28 @@ Consequences, stated plainly: recovery-budget exhaustion — delivers exactly one arrival to its parent's join slot, atomically with the terminal transition (§1). A join can wait forever only on a child that is still genuinely running. +- **Branches close with their parent.** When a run reaches a terminal state + — cancelled, failed, timed out, force-finalized, or completed — every + branch it fanned out that is still running has cancellation requested *in + the same store transaction as the terminal transition*. Not follow-up: an + operator cancels a rollout to stop the regional deploys, and a worker that + died mid-follow-up would leave them deploying. + - Each level closes only its own branches. A branch blocked on its own join + holds no claim, so it is control-pending the moment it is marked, + finalizes `CANCELLED`, and closes the level beneath it in turn. Depth is + not a loophole and there is no waiting-for-grandchildren deadlock. + - A branch that already finished is left exactly as it finished. + - `rx.parallel(..., parent_close="abandon")` opts a fan-out out, for + delegated work that should genuinely outlive its starter. An operator who + wants an abandoned branch stopped anyway cancels it directly; it is an + ordinary run. - `rx.parallel(..., mode="first")`: the join resolves on the first arrival; the engine then requests cancellation of the losing branches. That request is best-effort follow-up, not part of the winning transaction, and it is - sent by the one worker that saw the winner arrive. Three consequences, - stated plainly because only the first is obvious: - - If that worker dies before sending it, the losers run to completion. + sent by the one worker that saw the winner arrive. Losers are also closed + durably when the parent itself goes terminal (above), which backstops a + worker that died — but only from that moment, so for a parent that runs on + for hours after the race, this remains true meanwhile: - A loser already executing on another worker receives the intent but is not fenced at commit, so it finishes its attempt. - **A loser that runs on therefore performs its side effects.** Its arrival @@ -196,9 +212,7 @@ Consequences, stated plainly: idempotency key, exactly as a retried step must; losing a race is not a guarantee of not having acted. Racing branches whose effects cannot be made idempotent is the wrong shape for `mode="first"`. -- Child runs are ordinary runs; cancelling the parent does not implicitly - cancel children (fan-out is delegation, not ownership). A cancelled - parent's join tombstones; late child arrivals are refused. +- A closed parent's join tombstones; late child arrivals are refused. ## 6. Identity: who is "the same" as whom @@ -354,7 +368,8 @@ its own history event. Every action is legal only from the states listed; anything else is a refused no-op with a reason. -- `cancel(run)` — any nonterminal run. +- `cancel(run)` — any nonterminal run. Closes the branches it fanned out, + per §5. - `resume(run)` — `NEEDS_ATTENTION` only; re-opens the suspended step with a fresh attempt budget. - `retry(run)` — a `FAILED` run: re-opens its failed step with a fresh diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 8b420c55494..685a98276fe 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -1299,6 +1299,71 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: assert [run.run_id for run in quoted] == ["a"] +async def check_finalizing_a_parent_closes_its_branches(store: RunStore) -> None: + """Closing a run marks its branches for cancellation in the same write. + + Best-effort follow-up from the worker that finalized is not enough: an + operator cancels a rollout to stop the regional deploys, and a worker that + dies mid-follow-up would leave them deploying. + """ + await store.admit(make_run(next_ordinal=2), make_step(), _ADMITTED) + for run_id, close in (("kid1", "cancel"), ("kid2", "abandon")): + await store.admit( + make_run( + run_id, parent_run_id="run1", parent_ordinal=1, parent_close=close + ), + make_step(run_id), + _ADMITTED, + ) + assert await store.request_cancel("run1", NOW) + assert await store.finalize_run( + "run1", + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=NOW, + ) + closed = await store.get_run("kid1") + spared = await store.get_run("kid2") + assert closed is not None + assert spared is not None + assert closed.cancel_requested + assert closed.status is RunStatus.CANCELLING + assert not spared.cancel_requested, "abandon must survive its parent" + assert spared.status is RunStatus.PENDING + + +async def check_closing_a_branch_never_revives_a_finished_one( + store: RunStore, +) -> None: + """A branch that already finished is left exactly as it finished.""" + await store.admit(make_run(next_ordinal=2), make_step(), _ADMITTED) + await store.admit( + make_run("kid1", parent_run_id="run1", parent_ordinal=1), + make_step("kid1"), + _ADMITTED, + ) + assert await store.request_cancel("kid1", NOW) + assert await store.finalize_run( + "kid1", + status=RunStatus.COMPLETED, + error=None, + event=HistoryEventType.RUN_COMPLETED, + now=NOW, + ) + assert await store.request_cancel("run1", NOW) + assert await store.finalize_run( + "run1", + status=RunStatus.FAILED, + error={"message": "boom"}, + event=HistoryEventType.RUN_FAILED, + now=NOW, + ) + child = await store.get_run("kid1") + assert child is not None + assert child.status is RunStatus.COMPLETED + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1338,6 +1403,8 @@ async def check_label_filter_handles_awkward_keys(store: RunStore) -> None: check_join_arrivals_count_once, check_finalize_refuses_while_a_step_is_claimed, check_finalize_tombstones_open_slots, + check_finalizing_a_parent_closes_its_branches, + check_closing_a_branch_never_revives_a_finished_one, check_resume_only_reopens_a_suspended_run, check_list_runs_filters_and_orders, check_count_runs_matches_the_listing, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index d965cf5846e..05ebac4ec28 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -1679,7 +1679,11 @@ def _success_completion( ) if isinstance(control, Parallel): children = self._child_records( - claim, control.branches, claim.run.next_ordinal, now + claim, + control.branches, + claim.run.next_ordinal, + now, + control.parent_close, ) then_id = self._resolve_successor(defn, control.then).handler_id join = StepRecord( @@ -2348,6 +2352,7 @@ def _child_records( branches: tuple[Any, ...], join_ordinal: int, now: float, + parent_close: str, ) -> tuple[tuple[RunRecord, StepRecord], ...]: """Build the child runs a fan-out will create. @@ -2360,6 +2365,8 @@ def _child_records( branches: The root events to run concurrently. join_ordinal: The join slot the children report to. now: Current time in epoch seconds. + parent_close: What happens to a branch still running when the + parent reaches a terminal state. Returns: Each child run paired with its root slot. @@ -2407,6 +2414,7 @@ def _child_records( next_ordinal=1, parent_run_id=claim.run.run_id, parent_ordinal=join_ordinal, + parent_close=parent_close, request_key=f"child:{claim.run.run_id}:{join_ordinal}:{index}", deadline=( now + defn.run_timeout if defn.run_timeout is not None else None diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index e6d8f513833..1f728307766 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -81,6 +81,7 @@ flow_key TEXT, parent_run_id TEXT, parent_ordinal INTEGER, + parent_close TEXT NOT NULL DEFAULT 'cancel', request_key TEXT, labels JSONB, deadline DOUBLE PRECISION, @@ -146,6 +147,7 @@ PRIMARY KEY (run_id, wait_key, dedupe_key) ); ALTER TABLE workflow_steps ADD COLUMN IF NOT EXISTS queue TEXT NOT NULL DEFAULT 'default'; +ALTER TABLE workflow_runs ADD COLUMN IF NOT EXISTS parent_close TEXT NOT NULL DEFAULT 'cancel'; CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs (status); CREATE INDEX IF NOT EXISTS idx_workflow_runs_flow ON workflow_runs (workflow_id, flow_key); @@ -234,6 +236,7 @@ def _run_from_row(row: Mapping[str, Any]) -> RunRecord: flow_key=row["flow_key"], parent_run_id=row["parent_run_id"], parent_ordinal=row["parent_ordinal"], + parent_close=row["parent_close"] or "cancel", request_key=row["request_key"], labels=row["labels"], deadline=row["deadline"], @@ -491,10 +494,11 @@ async def _insert_run(self, conn: Connection, run: RunRecord) -> None: await conn.execute( "INSERT INTO workflow_runs (run_id, workflow_id, definition_digest," " status, state, state_version, next_ordinal, result, error," - " flow_key, parent_run_id, parent_ordinal, request_key, labels," + " flow_key, parent_run_id, parent_ordinal, parent_close," + " request_key, labels," " deadline, cancel_requested, created_at, updated_at)" " VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s," - " %s, %s, %s, %s)", + " %s, %s, %s, %s, %s)", ( run.run_id, run.workflow_id, @@ -508,6 +512,7 @@ async def _insert_run(self, conn: Connection, run: RunRecord) -> None: run.flow_key, run.parent_run_id, run.parent_ordinal, + run.parent_close, run.request_key, _json(run.labels), run.deadline, @@ -1061,6 +1066,8 @@ async def commit( ), ) await self._append_events(conn, claim.run.run_id, completion.events, now) + if completion.run_status in TERMINAL_RUN_STATUSES: + await self._close_children(conn, claim.run.run_id, now) if completion.parent_arrival is not None: await self._apply_arrival(conn, *completion.parent_arrival, now) @@ -1249,6 +1256,43 @@ async def admit_children( await self._lock_run(conn, parent) await self._append_events(conn, parent, events, now) + async def _close_children(self, conn: Any, run_id: str, now: float) -> None: + """Request cancellation of branches the closing run fanned out to. + + Called inside the transaction that takes a run terminal, so an + operator cancelling a rollout durably stops the regional deploys it + started -- not best-effort follow-up that dies with the worker. + Grandchildren are not walked here: a marked child is control-pending + the moment it drains (a run blocked on its own join holds no claim), + so it finalizes and closes its own branches in turn. + + Args: + conn: The connection inside an open transaction. + run_id: The run reaching a terminal state. + now: Current time in epoch seconds. + """ + closing = await ( + await conn.execute( + "UPDATE workflow_runs SET cancel_requested = TRUE, status = %s," + " updated_at = %s WHERE parent_run_id = %s" + " AND parent_close <> 'abandon' AND NOT (status = ANY(%s))" + " RETURNING run_id", + ( + RunStatus.CANCELLING.value, + now, + run_id, + [s.value for s in TERMINAL_RUN_STATUSES], + ), + ) + ).fetchall() + for row in closing: + await self._append_events( + conn, + row["run_id"], + ((HistoryEventType.RUN_CANCEL_REQUESTED, {"cause": "parent_close"}),), + now, + ) + async def _apply_arrival( self, conn: Connection, @@ -1655,6 +1699,7 @@ async def finalize_run( ] events.append((event, {} if error is None else dict(error))) await self._append_events(conn, run_id, events, now) + await self._close_children(conn, run_id, now) if parent_arrival is not None: await self._apply_arrival(conn, *parent_arrival, now) return True diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 829c3db307b..415b443a6ca 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -184,6 +184,8 @@ class RunRecord: flow_key: Grouping key for start policies such as singleton, if any. parent_run_id: The run that spawned this one, if any. parent_ordinal: The join slot in the parent this run reports to. + parent_close: What happens to this run when its parent reaches a + terminal state: ``"cancel"`` or ``"abandon"``. request_key: Idempotent admission key, if one was supplied. labels: Server-derived indexing labels. deadline: Absolute run deadline in epoch seconds, if configured. @@ -204,6 +206,7 @@ class RunRecord: flow_key: str | None = None parent_run_id: str | None = None parent_ordinal: int | None = None + parent_close: str = "cancel" request_key: str | None = None labels: dict[str, str] | None = None deadline: float | None = None diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index e28ecb1eb58..0b73e3f7e73 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -1232,6 +1232,8 @@ async def commit( updated_at=now, ) self._append_events(run.run_id, completion.events, now) + if completion.run_status in TERMINAL_RUN_STATUSES: + self._close_children(run.run_id, now) if completion.parent_arrival is not None: parent_id, ordinal, payload, dedupe_key = completion.parent_arrival self._apply_arrival(parent_id, ordinal, payload, dedupe_key, now) @@ -1760,6 +1762,39 @@ async def control_pending(self, now: float) -> tuple[RunRecord, ...]: pending.append(run) return tuple(pending) + def _close_children(self, run_id: str, now: float) -> None: + """Request cancellation of branches the closing run fanned out to. + + Called inside the transaction that takes a run terminal, so an + operator cancelling a rollout durably stops the regional deploys it + started -- not best-effort follow-up that dies with the worker. + Grandchildren are not walked here: a marked child is control-pending + the moment it drains (a run blocked on its own join holds no claim), + so it finalizes and closes its own branches in turn. + + Args: + run_id: The run reaching a terminal state. + now: Current time in epoch seconds. + """ + for child in self._runs.values(): + if ( + child.parent_run_id != run_id + or child.parent_close == "abandon" + or child.status in TERMINAL_RUN_STATUSES + ): + continue + self._runs[child.run_id] = dataclasses.replace( + child, + cancel_requested=True, + status=RunStatus.CANCELLING, + updated_at=now, + ) + self._append_events( + child.run_id, + ((HistoryEventType.RUN_CANCEL_REQUESTED, {"cause": "parent_close"}),), + now, + ) + async def finalize_run( self, run_id: str, @@ -1812,6 +1847,7 @@ async def finalize_run( ) events.append((event, {} if error is None else dict(error))) self._append_events(run_id, events, now) + self._close_children(run_id, now) if parent_arrival is not None: self._apply_arrival(*parent_arrival, now) return True @@ -2312,7 +2348,7 @@ async def next_due( return min(due_times) if due_times else None -SCHEMA_VERSION: Final = 2 +SCHEMA_VERSION: Final = 3 """Stamped into PRAGMA user_version; bump when _SCHEMA or migrations change.""" DATABASE_ENV: Final = "REFLEX_WORKFLOW_DATABASE" @@ -2362,6 +2398,7 @@ def resolve_store(target: str | None = None) -> RunStore: flow_key TEXT, parent_run_id TEXT, parent_ordinal INTEGER, + parent_close TEXT NOT NULL DEFAULT 'cancel', request_key TEXT, labels TEXT, deadline REAL, @@ -2457,6 +2494,13 @@ def resolve_store(target: str | None = None) -> RunStore: ("flow_key", "ALTER TABLE workflow_runs ADD COLUMN flow_key TEXT"), ("parent_run_id", "ALTER TABLE workflow_runs ADD COLUMN parent_run_id TEXT"), ("parent_ordinal", "ALTER TABLE workflow_runs ADD COLUMN parent_ordinal INTEGER"), + ( + "parent_close", + ( + "ALTER TABLE workflow_runs ADD COLUMN parent_close TEXT NOT NULL" + " DEFAULT 'cancel'" + ), + ), ) @@ -2505,6 +2549,7 @@ def _run_from_row(row: sqlite3.Row) -> RunRecord: error=_load(row["error"]), flow_key=row["flow_key"], parent_run_id=row["parent_run_id"], + parent_close=row["parent_close"] or "cancel", parent_ordinal=row["parent_ordinal"], request_key=row["request_key"], labels=_load(row["labels"]), @@ -2813,9 +2858,10 @@ def _insert_run(self, run: RunRecord) -> None: self._db.execute( "INSERT INTO workflow_runs (run_id, workflow_id, definition_digest," " status, state, state_version, next_ordinal, result, error," - " flow_key, parent_run_id, parent_ordinal, request_key, labels," - " deadline, cancel_requested, created_at, updated_at)" - " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + " flow_key, parent_run_id, parent_ordinal, parent_close," + " request_key, labels, deadline, cancel_requested, created_at," + " updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( run.run_id, run.workflow_id, @@ -2829,6 +2875,7 @@ def _insert_run(self, run: RunRecord) -> None: run.flow_key, run.parent_run_id, run.parent_ordinal, + run.parent_close, run.request_key, _dump(run.labels), run.deadline, @@ -2838,6 +2885,35 @@ def _insert_run(self, run: RunRecord) -> None: ), ) + def _close_children_sql(self, run_id: str, now: float) -> None: + """Request cancellation of branches the closing run fanned out to. + + Called inside the transaction that takes a run terminal, so an + operator cancelling a rollout durably stops the regional deploys it + started -- not best-effort follow-up that dies with the worker. + Grandchildren are not walked here: a marked child is control-pending + the moment it drains (a run blocked on its own join holds no claim), + so it finalizes and closes its own branches in turn. + + Args: + run_id: The run reaching a terminal state. + now: Current time in epoch seconds. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + closing = self._db.execute( + "UPDATE workflow_runs SET cancel_requested = 1, status = ?," + " updated_at = ? WHERE parent_run_id = ? AND parent_close <> 'abandon'" + f" AND status NOT IN ({','.join('?' * len(terminal))})" + " RETURNING run_id", + (RunStatus.CANCELLING.value, now, run_id, *terminal), + ).fetchall() + for row in closing: + self._append_events( + row["run_id"], + ((HistoryEventType.RUN_CANCEL_REQUESTED, {"cause": "parent_close"}),), + now, + ) + def _insert_step(self, step: StepRecord) -> None: """Insert a step row inside the current transaction. @@ -3411,6 +3487,8 @@ def work() -> None: ), ) self._append_events(claim.run.run_id, completion.events, now) + if completion.run_status in TERMINAL_RUN_STATUSES: + self._close_children_sql(claim.run.run_id, now) if completion.parent_arrival is not None: self._apply_arrival_sql(*completion.parent_arrival, now) self._db.execute("COMMIT") @@ -4149,6 +4227,7 @@ def work(): ] events.append((event, {} if error is None else dict(error))) self._append_events(run_id, events, now) + self._close_children_sql(run_id, now) if parent_arrival is not None: self._apply_arrival_sql(*parent_arrival, now) self._db.execute("COMMIT") diff --git a/tests/units/workflow/test_cascade.py b/tests/units/workflow/test_cascade.py new file mode 100644 index 00000000000..f527d5b4a19 --- /dev/null +++ b/tests/units/workflow/test_cascade.py @@ -0,0 +1,290 @@ +"""Closing a parent closes the children it fanned out to. + +An operator cancels a rollout to stop the blast radius. If the regional +deploys it spawned keep deploying, the button did not do the one thing it +exists to do. Fan-out children are cancelled with their parent by default; +``parent_close="abandon"`` opts a fan-out out, for the genuine +delegation case where a child should outlive whoever started it. +""" + +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import RunStatus +from reflex.workflow.testing import WorkflowTestHarness + +DEPLOYED: list[str] = [] + + +class Region(rx.State): + """A regional deploy that acts after a soak delay.""" + + __workflow__ = WorkflowConfig(id="cascade.region") + region: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, region: str): + """Soak, then deploy. + + Args: + region: The region to deploy to. + + Returns: + A deferral. + """ + self.region = region + return rx.after("1h", Region.deploy(region)) + + @rx.event(durable=True, effect="non_idempotent_write") + def deploy(self, region: str): + """Perform the deploy. + + Args: + region: The region to deploy to. + + Returns: + Completion. + """ + DEPLOYED.append(region) + return rx.complete(result={"region": region}) + + +def _rollout(**fan_out_kwargs): + """Build a rollout parent fanning out to three regions. + + Args: + fan_out_kwargs: Passed through to ``rx.parallel``. + + Returns: + The parent workflow class. + """ + + class Rollout(rx.State): + __workflow__ = WorkflowConfig(id="cascade.rollout") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Fan out to every region. + + Returns: + The fan-out. + """ + return rx.parallel( + Region.start("us-east"), + Region.start("us-west"), + Region.start("eu"), + then=Rollout.report, + **fan_out_kwargs, + ) + + @rx.event(durable=True, effect="none") + def report(self, results: list): + """Report the rollout. + + Args: + results: One entry per region. + + Returns: + Completion. + """ + return rx.complete(result={"regions": len(results)}) + + return Rollout + + +async def _children(harness: WorkflowTestHarness, parent_id: str): + """List the runs fanned out from a parent. + + Args: + harness: The running harness. + parent_id: The parent run. + + Returns: + The child run records. + """ + runs = await harness.kernel.list_runs() + return [run for run in runs if run.parent_run_id == parent_id] + + +async def test_cancelling_a_parent_stops_the_regions_it_started( + forked_registration_context, +): + """The cancel button stops the blast radius, not just the parent. + + Args: + forked_registration_context: Isolates workflow registration. + """ + DEPLOYED.clear() + rollout = _rollout() + async with WorkflowTestHarness(rollout, Region) as harness: + result = await harness.start(rollout.begin()) + assert result.run_id is not None + assert len(await _children(harness, result.run_id)) == 3 + + await harness.cancel(result.run_id) + await harness.advance("2h") + + assert DEPLOYED == [], ( + f"cancelling the rollout must stop the regions; deployed {DEPLOYED}" + ) + statuses = {run.status for run in await _children(harness, result.run_id)} + assert statuses == {RunStatus.CANCELLED}, statuses + + +async def test_abandon_lets_a_delegated_child_outlive_its_parent( + forked_registration_context, +): + """Delegation is a real shape, so it stays available -- just not silent. + + Args: + forked_registration_context: Isolates workflow registration. + """ + DEPLOYED.clear() + rollout = _rollout(parent_close="abandon") + async with WorkflowTestHarness(rollout, Region) as harness: + result = await harness.start(rollout.begin()) + assert result.run_id is not None + await harness.cancel(result.run_id) + await harness.advance("2h") + + assert sorted(DEPLOYED) == ["eu", "us-east", "us-west"] + statuses = {run.status for run in await _children(harness, result.run_id)} + assert statuses == {RunStatus.COMPLETED}, statuses + + +class Shard(rx.State): + """A grandchild that acts after a delay.""" + + __workflow__ = WorkflowConfig(id="cascade.shard") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, shard: str): + """Soak, then write. + + Args: + shard: The shard to write. + + Returns: + A deferral. + """ + return rx.after("1h", Shard.write(shard)) + + @rx.event(durable=True, effect="non_idempotent_write") + def write(self, shard: str): + """Perform the write. + + Args: + shard: The shard to write. + + Returns: + Completion. + """ + DEPLOYED.append(f"shard:{shard}") + return rx.complete(result={"shard": shard}) + + +class Tier(rx.State): + """A child that fans out again.""" + + __workflow__ = WorkflowConfig(id="cascade.tier") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, tier: str): + """Fan out to two shards. + + Args: + tier: The tier being rolled out. + + Returns: + The fan-out. + """ + return rx.parallel( + Shard.start(f"{tier}-a"), Shard.start(f"{tier}-b"), then=Tier.done + ) + + @rx.event(durable=True, effect="none") + def done(self, results: list): + """Finish the tier. + + Args: + results: One entry per shard. + + Returns: + Completion. + """ + return rx.complete(result={"shards": len(results)}) + + +class Deep(rx.State): + """A rollout three levels deep.""" + + __workflow__ = WorkflowConfig(id="cascade.deep") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Fan out to two tiers. + + Returns: + The fan-out. + """ + return rx.parallel(Tier.start("web"), Tier.start("api"), then=Deep.report) + + @rx.event(durable=True, effect="none") + def report(self, results: list): + """Finish the rollout. + + Args: + results: One entry per tier. + + Returns: + Completion. + """ + return rx.complete(result={"tiers": len(results)}) + + +async def test_the_close_reaches_grandchildren(forked_registration_context): + """Depth is not a loophole; the close walks the whole tree. + + Each level marks only its own branches, and a branch blocked on its own + join holds no claim -- so it is control-pending at once, finalizes, and + closes the level below it. A tier that could not drain until its shards + reported would deadlock the whole scheme. + + Args: + forked_registration_context: Isolates workflow registration. + """ + DEPLOYED.clear() + async with WorkflowTestHarness(Deep, Tier, Shard) as harness: + result = await harness.start(Deep.begin()) + assert result.run_id is not None + await harness.cancel(result.run_id) + await harness.advance("2h") + + assert DEPLOYED == [], f"grandchildren kept writing: {DEPLOYED}" + runs = await harness.kernel.list_runs() + shards = [run for run in runs if run.workflow_id == "cascade.shard"] + assert len(shards) == 4, f"expected four shards, got {len(shards)}" + assert {run.status for run in shards} == {RunStatus.CANCELLED} + + +async def test_a_failed_parent_closes_its_branches(forked_registration_context): + """Cancellation is not the only way a parent stops existing. + + Args: + forked_registration_context: Isolates workflow registration. + """ + DEPLOYED.clear() + rollout = _rollout() + async with WorkflowTestHarness(rollout, Region) as harness: + result = await harness.start(rollout.begin()) + assert result.run_id is not None + assert await harness.kernel.force_finalize( + result.run_id, + status=RunStatus.FAILED, + error={"message": "operator gave up"}, + ) + await harness.advance("2h") + + assert DEPLOYED == [], f"a failed rollout kept deploying: {DEPLOYED}" + statuses = {run.status for run in await _children(harness, result.run_id)} + assert statuses == {RunStatus.CANCELLED}, statuses diff --git a/tests/units/workflow/test_conformance.py b/tests/units/workflow/test_conformance.py index b40d96bc993..455af80dc5a 100644 --- a/tests/units/workflow/test_conformance.py +++ b/tests/units/workflow/test_conformance.py @@ -1,13 +1,41 @@ -"""Run the store conformance suite against every shipped implementation.""" +"""Run the store conformance suite against every shipped implementation. + +Postgres joins the sweep whenever ``REFLEX_TEST_POSTGRES`` points at a live +server. Without it the suite still runs, but only against the two stores that +need no server -- so a Postgres-only divergence in a shared behaviour would +sit undetected until production. Run it with a server before trusting a +change to any store. +""" + +import os +import uuid import pytest from reflex.workflow.conformance import CONFORMANCE_CHECKS from reflex.workflow.store import MemoryRunStore, SqliteRunStore +POSTGRES_URL = os.environ.get("REFLEX_TEST_POSTGRES") or "" + +STORE_KINDS = ["memory", "sqlite", *(["postgres"] if POSTGRES_URL else [])] + + +@pytest.fixture(autouse=True) +def harness_store(): + """Opt out of the shared harness store parameter. + + These checks build their stores directly, so crossing them with the + harness's own store parameter would run every check nine times to test + three things. + + Returns: + The store kind this module reports. + """ + return "memory" -@pytest.fixture(params=["memory", "sqlite"]) -def store(request, tmp_path): + +@pytest.fixture(params=STORE_KINDS) +async def store(request, tmp_path): """A fresh, empty store of each implementation. Args: @@ -19,12 +47,27 @@ def store(request, tmp_path): """ if request.param == "memory": yield MemoryRunStore() - else: + elif request.param == "sqlite": sqlite_store = SqliteRunStore(tmp_path / "workflow.db") yield sqlite_store sqlite_store.close() + else: + from reflex.workflow.postgres import PostgresRunStore + + opened = PostgresRunStore( + POSTGRES_URL, schema=f"wf_conf_{uuid.uuid4().hex}", min_size=0, max_size=4 + ) + yield opened + await opened.close() + opened.drop_schema() @pytest.mark.parametrize("check", CONFORMANCE_CHECKS, ids=lambda check: check.__name__) async def test_store_conforms(store, check): + """Every store answers the same way. + + Args: + store: The store under test. + check: The conformance check to run. + """ await check(store) diff --git a/tests/units/workflow/test_parallel.py b/tests/units/workflow/test_parallel.py index ec139db7d52..7bc77502fca 100644 --- a/tests/units/workflow/test_parallel.py +++ b/tests/units/workflow/test_parallel.py @@ -136,11 +136,12 @@ def start(self, lead: str): raise TransientWorkflowError(msg) -def _router(*branches): +def _router(*branches, **fan_out_kwargs): """Build a parent workflow fanning out to the given branches. Args: branches: The branch classes to fan out to. + fan_out_kwargs: Passed through to ``rx.parallel``. Returns: The parent workflow class. @@ -161,7 +162,9 @@ def begin(self, lead: str): The parallel fan-out. """ return rx.parallel( - *[branch.start(lead) for branch in branches], then=Router.route + *[branch.start(lead) for branch in branches], + then=Router.route, + **fan_out_kwargs, ) @rx.event(durable=True, effect="none") @@ -733,19 +736,20 @@ async def test_unidentifiable_arrivals_keep_their_position( assert [entry["result"] for entry in ordered] == ["a", "b", "c"] -async def test_cancelling_an_all_mode_parent_leaves_its_branches_alone( +async def test_an_abandoned_branch_never_cancels_its_sibling( forked_registration_context, ): - """Delegation is not ownership, so a cancelled parent cancels no children. - - Section 5 says cancelling a parent does not implicitly cancel children. - Cancelling one tombstones its join, and when a branch later finishes its - arrival finds a slot that is no longer blocked -- which is not the same - thing as a race having been decided. Reading it as one made the engine - cancel the sibling of a branch it never raced. + """A tombstoned join is not a decided race, even when nobody cancels. + + Cancelling a parent tombstones its join, and when a branch later finishes + its arrival finds a slot that is no longer blocked -- which is not the + same thing as a race having been decided. Reading it as one made the + engine cancel the sibling of a branch it never raced. Under the default + ``parent_close="cancel"`` both branches stop anyway and that misreading + would hide; ``abandon`` is where it still shows. """ BRANCH_CALLS.clear() - router = _router(Slowish, Slower) + router = _router(Slowish, Slower, parent_close="abandon") async with WorkflowTestHarness(router, Slowish, Slower) as harness: started = await harness.start(router.begin("acme")) assert started.run_id is not None From dc9c79a95fcc3d52e8986938191a4907ac6a4782 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 10:40:47 -0700 Subject: [PATCH 102/121] workflows: operator-action helpers on the test harness retry, skip, force_complete and force_fail, matching what the CLI exposes. Rehearsing repair is most of what a workflow test needs to do, and reaching through harness.kernel for it read as private API -- which also meant a test reaching for a helper that was not there failed with AttributeError, and an xfail swallows that as a pass. That has bitten this suite before. The test that covers them documents three engine behaviours that make the obvious version of it wrong: a failed attempt discards its state, so an attempt counter has to live outside the run; skip only restores a successor that was preallocated, so the fixture needs a real chain; and force_* is refused on a terminal run, so it needs a run that is nonterminal and drained -- one sitting on a timer. --- ...rkflow-harness-operator-helpers.feature.md | 1 + reflex/workflow/testing.py | 76 +++++++++++ tests/units/workflow/test_cascade.py | 6 +- tests/units/workflow/test_testing.py | 120 +++++++++++++++++- 4 files changed, 197 insertions(+), 6 deletions(-) create mode 100644 news/workflow-harness-operator-helpers.feature.md diff --git a/news/workflow-harness-operator-helpers.feature.md b/news/workflow-harness-operator-helpers.feature.md new file mode 100644 index 00000000000..035db9333d3 --- /dev/null +++ b/news/workflow-harness-operator-helpers.feature.md @@ -0,0 +1 @@ +`WorkflowTestHarness` gained `retry`, `skip`, `force_complete`, and `force_fail`, matching the operator actions the CLI exposes. Rehearsing repair is most of what a workflow test needs to do, and reaching through `harness.kernel` for it read as private API. diff --git a/reflex/workflow/testing.py b/reflex/workflow/testing.py index 2c88533b23b..29b2d639362 100644 --- a/reflex/workflow/testing.py +++ b/reflex/workflow/testing.py @@ -22,6 +22,7 @@ ) from reflex.workflow.kernel import WorkflowObserver +from reflex.workflow.records import RunStatus from reflex.workflow.runtime import WorkflowRuntime, _context_runtime from reflex.workflow.store import MemoryRunStore @@ -280,3 +281,78 @@ async def cancel(self, run_id: str) -> bool: cancelled = await self.kernel.cancel(run_id) await self.kernel.run_until_idle() return cancelled + + async def retry(self, run_id: str) -> bool: + """Re-open a failed run at the step that failed, and drain. + + Args: + run_id: The failed run to retry. + + Returns: + True if a failed run was re-opened. + """ + retried = await self.kernel.retry(run_id) + await self.kernel.run_until_idle() + return retried + + async def skip(self, run_id: str) -> bool: + """Give up on a blocking step and let the run carry on, then drain. + + Args: + run_id: The stuck run to unstick. + + Returns: + True if a blocking step was skipped. + """ + skipped = await self.kernel.skip(run_id) + await self.kernel.run_until_idle() + return skipped + + async def force_complete(self, run_id: str, result: Any = None) -> bool: + """Finish a drained run by operator decision, then drain. + + Args: + run_id: The run to complete. + result: The result to record as what it produced. + + Returns: + True if the run was finalized. + """ + return await self._force(run_id, RunStatus.COMPLETED, result=result) + + async def force_fail(self, run_id: str, reason: str) -> bool: + """Fail a drained run by operator decision, then drain. + + Args: + run_id: The run to fail. + reason: The message to record as the failure. + + Returns: + True if the run was finalized. + """ + return await self._force(run_id, RunStatus.FAILED, error={"message": reason}) + + async def _force( + self, + run_id: str, + status: RunStatus, + *, + result: Any = None, + error: dict[str, Any] | None = None, + ) -> bool: + """Force-finalize a run and process what its close unblocks. + + Args: + run_id: The run to finalize. + status: The terminal status to record. + result: Result to record when completing. + error: Error payload to record when failing. + + Returns: + True if the run was finalized. + """ + finalized = await self.kernel.force_finalize( + run_id, status=status, result=result, error=error + ) + await self.kernel.run_until_idle() + return finalized diff --git a/tests/units/workflow/test_cascade.py b/tests/units/workflow/test_cascade.py index f527d5b4a19..97fd1b0932c 100644 --- a/tests/units/workflow/test_cascade.py +++ b/tests/units/workflow/test_cascade.py @@ -278,11 +278,7 @@ async def test_a_failed_parent_closes_its_branches(forked_registration_context): async with WorkflowTestHarness(rollout, Region) as harness: result = await harness.start(rollout.begin()) assert result.run_id is not None - assert await harness.kernel.force_finalize( - result.run_id, - status=RunStatus.FAILED, - error={"message": "operator gave up"}, - ) + assert await harness.force_fail(result.run_id, "operator gave up") await harness.advance("2h") assert DEPLOYED == [], f"a failed rollout kept deploying: {DEPLOYED}" diff --git a/tests/units/workflow/test_testing.py b/tests/units/workflow/test_testing.py index da044a55024..0b201c45b0a 100644 --- a/tests/units/workflow/test_testing.py +++ b/tests/units/workflow/test_testing.py @@ -11,7 +11,7 @@ import pytest from reflex_base.utils.exceptions import WorkflowDefinitionError -from reflex_base.workflow import WorkflowConfig, manual +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual import reflex as rx from reflex.workflow.records import RunStatus @@ -123,3 +123,121 @@ async def test_an_injected_store_is_left_to_its_owner(forked_registration_contex # Still usable afterwards: the run is there for the next harness. assert await store.get_run(result.run_id) is not None + + +STUCK_ATTEMPTS: list[int] = [] + + +class Stuck(rx.State): + """A two-step chain whose middle step always fails.""" + + __workflow__ = WorkflowConfig(id="harness.stuck") + ran_after: bool = False + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Preallocate the chain. + + Returns: + The work, then the step that must survive its failure. + """ + return [Stuck.work, Stuck.after] + + @rx.event( + durable=True, trigger=manual(), effect="none", retry=Retry(max_attempts=1) + ) + def work(self): + """Fail, however many times an operator retries. + + Raises: + TransientWorkflowError: Always. + """ + STUCK_ATTEMPTS.append(1) + msg = "vendor down" + raise TransientWorkflowError(msg) + + @rx.event(durable=True, effect="none") + def after(self): + """Run once the blocking step is past. + + Returns: + Completion. + """ + self.ran_after = True + return rx.complete(result={"ok": True}) + + +class Waiting(rx.State): + """A run that sits on a long timer, nonterminal and drained.""" + + __workflow__ = WorkflowConfig(id="harness.waiting") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Wait a month. + + Returns: + A deferral no test will wait out. + """ + return rx.after("30d", Waiting.later) + + @rx.event(durable=True, effect="none") + def later(self): + """Finish, eventually. + + Returns: + Completion. + """ + return rx.complete(result={"waited": True}) + + +async def test_the_harness_drives_retry_skip_and_force(forked_registration_context): + """Operator repair is most of what a workflow test needs to rehearse. + + Reaching through ``harness.kernel`` for it worked but read as private + API, and a test that reaches for a helper that is not there fails with + AttributeError -- which an xfail will happily swallow as a pass. + + Args: + forked_registration_context: Isolates workflow registration. + """ + STUCK_ATTEMPTS.clear() + async with WorkflowTestHarness(Stuck) as harness: + failed = await harness.start(Stuck.start()) + assert failed.run_id is not None + snapshot = await harness.get_run(failed.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + + assert await harness.retry(failed.run_id) + snapshot = await harness.get_run(failed.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED, "attempt two fails too" + assert len(STUCK_ATTEMPTS) == 2, "retry re-ran the step" + + assert await harness.skip(failed.run_id) + snapshot = await harness.get_run(failed.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state["ran_after"] is True, ( + "skipping restores the successor the failure tombstoned" + ) + + async with WorkflowTestHarness(Waiting) as harness: + run = await harness.start(Waiting.start()) + assert run.run_id is not None + assert await harness.force_complete(run.run_id, {"decided": "by hand"}) + snapshot = await harness.get_run(run.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"decided": "by hand"} + + async with WorkflowTestHarness(Waiting) as harness: + run = await harness.start(Waiting.start()) + assert run.run_id is not None + assert await harness.force_fail(run.run_id, "not worth repairing") + snapshot = await harness.get_run(run.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error is not None + assert snapshot.error["message"] == "not worth repairing" From 0ce143c3386bbbd2158c0a1f2cfd36cbefbcb261 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 10:41:18 -0700 Subject: [PATCH 103/121] workflows: say where release pinning belongs Section 4 documents that the engine does not pin a run to the code that started it, which invites the reading that pinning is a missing feature. It is not: it is a deployment concern, done by routing admissions, and the engine rules hold underneath whatever routing exists. Saying so keeps the contract and the (unbuilt) deploy layer from drifting into two different answers to the same question. --- reflex/workflow/CONTRACT.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index b9975405adb..1a901192810 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -169,6 +169,14 @@ Consequences, stated plainly: steps of one run with different code. Steps are the consistency boundary; the contract makes no promise that one run sees one release. +The engine deliberately does not pin runs to releases, and nothing here +should be read as planning to. Pinning is a *deployment* concern: a hosting +layer that wants one run to see one release does it by routing — admitting +new runs to the new release while old runs finish on the old one — not by +asking the engine to keep old code alive. That routing is not built yet; when +it is, it constrains which workers exist, and every rule above still holds +underneath it. + ## 5. Cancellation, deadlines, and children - `cancel(run_id)` records intent and cancels any in-flight attempt From 53779a3c07d09c9bbab3830ab6bab4d141278879 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 10:49:10 -0700 Subject: [PATCH 104/121] workflows: kill a real worker at real boundaries The suite simulated crashes in-process -- abandon a claim, expire a lease, commit behind a fence. That tests the store's logic and says nothing about what reached the disk, because the process that was supposed to have died is still there to tidy up. "Kill any process at any boundary" was an untested claim. These SIGKILL a worker subprocess at a named boundary and make a fresh process produce the documented outcome: before the effect -> work simply undone, runs once on recovery after an unguarded one -> repeats (section 2, and the reason rx.step exists) after a journal write -> replays, never charges twice after parent finalize -> branches already CANCELLING on disk The last is the one the cascade design turns on: the process that finalized dies with no chance to do anything else, so if branch cancellation were follow-up work the regions would still deploy. Deleting the _close_children_sql call makes exactly that test fail, which is the check that it is measuring what it claims to. Two ways this kind of test lies, both closed: an effect lost to the page cache would make a repeat look exactly-once, so the ledger is fsynced before each kill; and a worker that exits cleanly would pass every assertion that follows, so each scenario asserts it died by signal. Contract section 8 now says which claims rest on simulation and which on real kills, because the difference is the whole point. --- news/workflow-crash-boundaries.feature.md | 1 + reflex/workflow/CONTRACT.md | 24 ++ tests/units/workflow/crash_worker.py | 241 ++++++++++++++++++ tests/units/workflow/test_crash_boundaries.py | 207 +++++++++++++++ 4 files changed, 473 insertions(+) create mode 100644 news/workflow-crash-boundaries.feature.md create mode 100644 tests/units/workflow/crash_worker.py create mode 100644 tests/units/workflow/test_crash_boundaries.py diff --git a/news/workflow-crash-boundaries.feature.md b/news/workflow-crash-boundaries.feature.md new file mode 100644 index 00000000000..8340e08b742 --- /dev/null +++ b/news/workflow-crash-boundaries.feature.md @@ -0,0 +1 @@ +Added a crash-boundary suite that SIGKILLs a real worker subprocess at named boundaries and requires a fresh process to produce the outcome `CONTRACT.md` documents. The rest of the suite simulates crashes in-process, which tests store logic but not what reached the disk. Covered: between claim and handler, after an unguarded effect, after a substep journal write, and immediately after a parent's finalize transaction. Side effects are recorded in an fsynced ledger and each scenario asserts the worker died by signal, so a worker that exits cleanly fails rather than silently passing. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 1a901192810..7471422f41b 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -371,6 +371,30 @@ reason rather than on message text. Nothing on this table is silent: each writes its reason to the run's error and its own history event. +### How this section is held to + +Every rule above is asserted somewhere, but two kinds of test back different +claims and only one of them is evidence about crashes. + +- **In-process simulation** (most of the suite) abandons a claim, expires a + lease, or commits behind a fence. It is a fair test of store logic and no + test at all of what reached the disk, because the process that was supposed + to have died is still there to tidy up. +- **Real kills** (`tests/units/workflow/test_crash_boundaries.py`) SIGKILL a + worker in a separate process at a named boundary and make a fresh process + produce the documented outcome. Effects are recorded in an fsynced ledger, + so an effect that really happened cannot be lost in a way that flatters the + result, and each scenario asserts the worker died by signal rather than + exiting — a scenario whose worker exits cleanly proves nothing while + passing everything. + +Boundaries covered by real kills: between claim and handler (work simply +undone); after an unguarded effect (repeats, §2, which is the cost `rx.step` +exists to remove); after a substep journal write (replays, never repeats); +and immediately after a parent's finalize transaction (branches are already +marked on disk, §5 — the case that distinguishes an in-transaction close from +follow-up). + ## 9. Operator actions Every action is legal only from the states listed; anything else is a refused diff --git a/tests/units/workflow/crash_worker.py b/tests/units/workflow/crash_worker.py new file mode 100644 index 00000000000..e4fd8a6e9c6 --- /dev/null +++ b/tests/units/workflow/crash_worker.py @@ -0,0 +1,241 @@ +"""A worker that dies where it is told to, for the crash-boundary tests. + +Run as a subprocess so the death is real: SIGKILL to a separate process, no +unwinding, no atexit, no flush. Simulating a crash in-process is a fair test +of the store's logic and no test at all of what actually reaches the disk, so +the two together are what the contract's "kill any process at any boundary" +claim rests on. + +Usage: ``crash_worker.py ``. +""" + +import asyncio +import os +import signal +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from reflex_base.workflow import Retry, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import HistoryEventType, RunStatus +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import SqliteRunStore + +LEDGER = Path(os.environ["CRASH_LEDGER"]) +CRASH_AT = os.environ["CRASH_AT"] + + +def record(name: str) -> None: + """Note that a side effect really happened, durably. + + Written and fsynced before any crash can follow it: a ledger entry lost + to the page cache would make a repeated effect look like an exactly-once + one, which is the direction of error that matters here. + + Args: + name: The effect that ran. + """ + with LEDGER.open("a") as handle: + handle.write(f"{name}\n") + handle.flush() + os.fsync(handle.fileno()) + + +def die_at(point: str) -> None: + """End this process immediately if this is the chosen boundary. + + Args: + point: The boundary being passed. + """ + if point == CRASH_AT: + os.kill(os.getpid(), signal.SIGKILL) + + +class Charge(rx.State): + """One step that moves money, with and without a substep journal.""" + + __workflow__ = WorkflowConfig(id="crash.charge") + charged: str = "" + + @rx.event( + durable=True, + trigger=manual(), + effect="non_idempotent_write", + retry=Retry(max_attempts=1), + ) + async def unguarded(self): + """Charge without guarding the call. + + Returns: + Completion. + """ + die_at("after_claim") + record("unguarded") + die_at("after_effect") + self.charged = "unguarded" + return rx.complete(result={"ok": True}) + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=1), + ) + async def guarded(self): + """Charge inside a substep, so the journal can replay it. + + Returns: + Completion. + """ + + def charge_once() -> dict: + """Make the charge. + + Returns: + The charge. + """ + record("guarded") + return {"charge_id": "ch_1"} + + charge = await rx.step("charge", charge_once) + die_at("after_step_record") + self.charged = charge["charge_id"] + return rx.complete(result=charge) + + +class Region(rx.State): + """A branch that acts only after a delay.""" + + __workflow__ = WorkflowConfig(id="crash.region") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, region: str): + """Soak, then deploy. + + Args: + region: The region to deploy. + + Returns: + A deferral. + """ + return rx.after("1h", Region.deploy(region)) + + @rx.event(durable=True, effect="non_idempotent_write") + def deploy(self, region: str): + """Deploy the region. + + Args: + region: The region to deploy. + + Returns: + Completion. + """ + record(f"deploy:{region}") + return rx.complete(result={"region": region}) + + +class Rollout(rx.State): + """A parent that fans out to two regions.""" + + __workflow__ = WorkflowConfig(id="crash.rollout") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Fan out. + + Returns: + The fan-out. + """ + return rx.parallel( + Region.start("us-east"), Region.start("eu"), then=Rollout.report + ) + + @rx.event(durable=True, effect="none") + def report(self, results: list): + """Report the rollout. + + Args: + results: One entry per region. + + Returns: + Completion. + """ + record("report") + return rx.complete(result={"regions": len(results)}) + + +def _write_run_id(run_id: str) -> None: + """Hand the parent run's identity to the next phase. + + Args: + run_id: The run to record. + """ + Path(sys.argv[2] + ".runid").write_text(run_id) + + +def _read_run_id() -> str: + """Read the parent run's identity left by an earlier phase. + + Returns: + The run id. + """ + return Path(sys.argv[2] + ".runid").read_text().strip() + + +async def main() -> None: + """Drive one phase of a crash scenario against a shared SQLite store.""" + db, _, phase = sys.argv[1], sys.argv[2], sys.argv[3] + store = SqliteRunStore(Path(db)) + runtime = WorkflowRuntime(store, lease_duration=1.0) + for workflow_cls in (Charge, Region, Rollout): + runtime.register(workflow_cls) + await runtime.startup(start_worker=False) + kernel = runtime.kernel + + if phase in ("unguarded", "guarded"): + await kernel.start(getattr(Charge, phase)()) + await kernel.recover() + await kernel.run_until_idle() + elif phase == "rollout": + started = await kernel.start(Rollout.begin()) + await kernel.run_until_idle() + die_at("after_fanout") + assert started.run_id is not None + _write_run_id(started.run_id) + elif phase == "cascade": + run_id = _read_run_id() + # Straight at the store, so the kill lands between the finalize + # transaction and literally anything else this process might do. + assert await store.request_cancel(run_id, await _now(store)) + assert await store.finalize_run( + run_id, + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=await _now(store), + ) + die_at("after_finalize") + else: + await kernel.recover() + await kernel.run_until_idle() + store.close() + + +async def _now(store: SqliteRunStore) -> float: + """Read the store's clock, falling back to this process's. + + Args: + store: The store. + + Returns: + Epoch seconds. + """ + import time + + return await store.epoch_time() or time.time() + + +asyncio.run(main()) diff --git a/tests/units/workflow/test_crash_boundaries.py b/tests/units/workflow/test_crash_boundaries.py new file mode 100644 index 00000000000..c45d6084deb --- /dev/null +++ b/tests/units/workflow/test_crash_boundaries.py @@ -0,0 +1,207 @@ +"""Kill a real worker at a real boundary and hold the contract to its word. + +The rest of the suite simulates crashes in-process -- abandoning a claim, +expiring a lease, committing behind a fence. That is a fair test of the +store's logic and no test at all of what actually reached the disk, because +the process that was supposed to have died is still there to tidy up. + +These send SIGKILL to a separate process: no unwinding, no ``atexit``, no +final flush, nothing the kernel could have done "on the way down" because +there is no way down. A fresh process then opens the same database and has to +produce the outcome CONTRACT.md documents. Every scenario asserts against a +durable ledger that is fsynced before each crash, so an effect that really +happened can never be lost in a way that flatters the result. +""" + +import os +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +WORKER = Path(__file__).parent / "crash_worker.py" + +LEASE_LAPSE = 1.4 + + +@pytest.fixture(autouse=True) +def harness_store(): + """Opt out of the shared harness store parameter. + + These drive subprocesses against their own SQLite file, so running each + one per store kind would repeat the same work with the same store. + + Returns: + The store kind this module uses. + """ + return "sqlite" + + +@pytest.fixture +def crash(tmp_path): + """Give a test a database, a ledger, and a way to run the worker. + + Args: + tmp_path: The test's temporary directory. + + Returns: + A ``(phase, crash_at) -> CompletedProcess`` runner, with ``.ledger`` + and ``.db`` paths attached. + """ + db = tmp_path / "crash.db" + ledger = tmp_path / "ledger.txt" + + def run(phase: str, crash_at: str = "none") -> subprocess.CompletedProcess: + """Run one phase of the worker. + + Args: + phase: Which scenario the worker should drive. + crash_at: The boundary at which it should die. + + Returns: + The finished process. + """ + return subprocess.run( + [sys.executable, str(WORKER), str(db), str(ledger), phase], + env={**os.environ, "CRASH_LEDGER": str(ledger), "CRASH_AT": crash_at}, + capture_output=True, + timeout=120, + check=False, + ) + + run.ledger = ledger # pyright: ignore[reportFunctionMemberAccess] + run.db = db # pyright: ignore[reportFunctionMemberAccess] + return run + + +def effects(ledger: Path) -> list[str]: + """Read the durable record of what really ran. + + Args: + ledger: The ledger file. + + Returns: + One entry per side effect that happened, in order. + """ + if not ledger.exists(): + return [] + return ledger.read_text().split() + + +def runs(db: Path) -> dict[str, str]: + """Read run statuses straight out of the database. + + Args: + db: The SQLite file. + + Returns: + Run id to status, read without going through the store. + """ + import sqlite3 + + connection = sqlite3.connect(db) + try: + return { + row[0]: row[1] + for row in connection.execute( + "SELECT run_id, status FROM workflow_runs" + ).fetchall() + } + finally: + connection.close() + + +def assert_killed(finished: subprocess.CompletedProcess) -> None: + """Confirm the worker really was killed rather than exiting. + + A scenario whose worker exited cleanly proves nothing, and would pass + every assertion that follows. + + Args: + finished: The finished process. + """ + assert finished.returncode == -9, ( + f"expected SIGKILL, got {finished.returncode}: {finished.stderr.decode()[-800:]}" + ) + + +def test_a_crash_before_the_effect_costs_nothing(crash): + """Dying between claim and handler leaves the work simply undone. + + Args: + crash: The worker runner. + """ + assert_killed(crash("unguarded", "after_claim")) + assert effects(crash.ledger) == [] + + time.sleep(LEASE_LAPSE) + assert crash("recover").returncode == 0 + assert effects(crash.ledger) == ["unguarded"], "recovery must run it exactly once" + + +def test_an_unguarded_effect_is_at_least_once_as_documented(crash): + """Section 2 promises re-execution, and this is what that costs. + + Not a bug being pinned as behaviour: it is the reason ``rx.step`` exists, + and the number here is what a workflow author is choosing to accept by + calling a provider without one. + + Args: + crash: The worker runner. + """ + assert_killed(crash("unguarded", "after_effect")) + assert effects(crash.ledger) == ["unguarded"] + + time.sleep(LEASE_LAPSE) + assert crash("recover").returncode == 0 + assert effects(crash.ledger) == ["unguarded", "unguarded"], ( + "an unguarded effect repeats after a crash; that is the contract" + ) + + +def test_a_journalled_effect_survives_a_real_kill_exactly_once(crash): + """The substep journal's whole promise, against a real process death. + + The charge is made, the journal records it, and the process is killed + before anything commits. The recovered attempt must replay the recorded + charge rather than make it again, because the money already moved. + + Args: + crash: The worker runner. + """ + assert_killed(crash("guarded", "after_step_record")) + assert effects(crash.ledger) == ["guarded"] + + time.sleep(LEASE_LAPSE) + assert crash("recover").returncode == 0 + assert effects(crash.ledger) == ["guarded"], ( + "the journal must replay the charge, not make a second one" + ) + + +def test_a_parent_closed_and_killed_still_stops_its_branches(crash): + """The close is in the finalize transaction, not follow-up after it. + + This is the scenario the whole design turns on: the process that + finalized the parent dies with no chance to do anything else at all. If + branch cancellation were follow-up work, the regions would still deploy. + + Args: + crash: The worker runner. + """ + assert crash("rollout").returncode == 0 + assert effects(crash.ledger) == [], "regions soak for an hour before deploying" + + assert_killed(crash("cascade", "after_finalize")) + + statuses = sorted(runs(crash.db).values()) + assert statuses == ["CANCELLED", "CANCELLING", "CANCELLING"], ( + f"branches must be marked on disk by the finalize itself, got {statuses}" + ) + + time.sleep(LEASE_LAPSE) + assert crash("recover").returncode == 0 + assert effects(crash.ledger) == [], "a cancelled rollout must never deploy" + assert sorted(runs(crash.db).values()) == ["CANCELLED"] * 3 From e5e80ba6f3e3bec31c2c00c6920457a1a0cf2701 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 11:02:46 -0700 Subject: [PATCH 105/121] workflows: make the contract's vocabulary complete and self-checking Auditing section 1's exit criterion by re-reading the prose was not an audit. Enumerating what the engine can actually show an operator -- run statuses, step statuses, start and delivery dispositions, history events -- and checking each against CONTRACT.md found real holes, in both directions. wait_expired was declared and never emitted. A resolved wait records wait_resolved; an expired one recorded nothing, so "did the approval come through, or did nobody answer?" -- the exact question history exists to answer -- had to be inferred from which handler ran next. Now recorded, with the wait key and the timeout branch that ran. A signal deduplicated by sender key was silent too, which makes "the provider says it delivered" indistinguishable from a delivery that never arrived. Recorded now in the store's own transaction, so it is durable like every other history event, in all three stores. Getting there took two wrong patches: the memory and Postgres deliver() paths sit next to _apply_arrival(), whose duplicate branch is textually identical and means something else entirely -- a duplicate child arrival, not a duplicate signal. The conformance check added here is what caught the Postgres one, on a real server. signal_delivered is deleted rather than emitted: every accepted delivery already records wait_resolved or signal_buffered, and a third event for the same fact is noise. A vocabulary word that never appears makes its own absence uninformative, which is worse than not having it. Section 10 now defines the whole vocabulary, and test_contract_vocabulary.py fails if a member is added without being documented, or documented without anything emitting it. Both directions are mutation-checked. That turns "every failure scenario has one documented outcome" from a claim into something the suite enforces. --- news/workflow-contract-vocabulary.bugfix.md | 3 + reflex/workflow/CONTRACT.md | 65 +++++++++ reflex/workflow/conformance.py | 19 +++ reflex/workflow/kernel.py | 27 ++++ reflex/workflow/postgres.py | 6 + reflex/workflow/records.py | 1 - reflex/workflow/store.py | 17 ++- .../workflow/test_contract_vocabulary.py | 109 ++++++++++++++ tests/units/workflow/test_waits.py | 134 +++++++++++++++++- 9 files changed, 378 insertions(+), 3 deletions(-) create mode 100644 news/workflow-contract-vocabulary.bugfix.md create mode 100644 tests/units/workflow/test_contract_vocabulary.py diff --git a/news/workflow-contract-vocabulary.bugfix.md b/news/workflow-contract-vocabulary.bugfix.md new file mode 100644 index 00000000000..078d3457b6b --- /dev/null +++ b/news/workflow-contract-vocabulary.bugfix.md @@ -0,0 +1,3 @@ +Audited the contract against every value the engine can actually show an operator, and fixed what the audit found. `wait_expired` was declared and never emitted, so a run whose approval timed out looked identical in history to one that was answered — the only trace of the deadline was which handler happened to run next. A signal deduplicated by its sender key was likewise silent, making "the provider says it delivered" impossible to distinguish from a delivery that never arrived; both are now recorded, in the store's transaction, in all three stores. `signal_delivered` was removed: every accepted delivery already records `wait_resolved` or `signal_buffered`, and a declared event nothing writes makes its own absence uninformative. + +`CONTRACT.md` gained section 10, defining the complete observable vocabulary — every run status, step status, start disposition, delivery disposition, and history event. `test_contract_vocabulary.py` fails if a member is added without documenting it, or documented without anything emitting it. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 7471422f41b..f6d213ab241 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -434,3 +434,68 @@ on one key — the singleton promise in §1 governs *admissions*, and an operator re-opening a run is a human override, not an admission. The operator can see what holds the key (`reflex workflows list -w `) and decide; the engine does not silently refuse a repair because a policy would have. + +## 10. The observable vocabulary + +Everything above describes behaviour; this is the complete list of words the +engine uses to describe it. An operator reading a run sees exactly these +values and nothing else, and a value that appears here but is never produced +(or is produced but never appears here) is a defect — `test_contract_vocabulary.py` +fails on either. + +**Run status.** Nonterminal: `PENDING` (admitted, nothing claimed yet), +`RUNNING` (an attempt holds a claim), `RETRYING` (an attempt failed and the +next is scheduled with backoff), `WAITING` (blocked on a signal, a timer, or +a join), `CANCELLING` (cancellation recorded, draining), `NEEDS_ATTENTION` +(suspended for a person; §8 names every reason). Terminal: `COMPLETED`, +`FAILED`, `CANCELLED`, `TIMED_OUT`. + +**Step status.** In flight: `READY` (claimable once it is the frontier), +`BLOCKED` (a wait or join slot not yet satisfied), `CLAIMED` (a worker holds +it under a lease), `RETRY_WAIT` (business retry scheduled), `RECOVERY_WAIT` +(lease lapsed, awaiting re-execution). Terminal: `SUCCEEDED`, `FAILED`, +`TIMED_OUT`, `CANCELLED`, `NEEDS_ATTENTION`, `SKIPPED`. + +**Start disposition** — what admission did with a submission: `started`, +`deduplicated` (§6 request key), `coalesced` (debounce), `skipped` +(singleton), `rejected` (rate limit, with `retry_after`). + +**Delivery disposition** — what the store did with a signal or arrival: +`resolved`, `buffered` (arrived before its wait was armed), `counted` (a join +arrival that is not the last), `duplicate` (repeated sender key), +`expired` (run past its deadline), `unknown_run`, `run_terminal`. + +**History events.** Append-only, one run's whole story: + +| event | means | +|---|---| +| `run_admitted` | the run exists; admission committed | +| `step_scheduled` | a slot was preallocated for future work | +| `attempt_started` | a claimed attempt began executing | +| `attempt_succeeded` | the attempt committed its transition | +| `attempt_failed` | the attempt raised | +| `attempt_timed_out` | the attempt exceeded its `timeout=` | +| `attempt_cancelled` | the attempt was cancelled cooperatively | +| `attempt_abandoned` | the attempt's work was discarded: fenced claim, lost lease, or a commit refused past the run deadline | +| `step_retry_scheduled` | the next business attempt was scheduled with backoff | +| `step_recovered` | a lapsed lease was reclaimed; costs one recovery, not an attempt | +| `step_tombstoned` | a terminal transition closed a slot that will now never run | +| `step_restored` | `retry`/`skip` brought back a slot a failure had tombstoned | +| `step_skipped` | an operator marked a blocking step `SKIPPED` | +| `run_completed`, `run_failed`, `run_timed_out`, `run_cancelled` | the run reached that terminal state | +| `run_cancel_requested` | cancellation intent recorded; `cause: parent_close` when a closing parent did it (§5) | +| `run_needs_attention` | suspended, carrying the `reason` from §8 | +| `run_resumed` | reopened; `origin` distinguishes `resume` from `retry` | +| `child_started` | a fan-out admitted a branch | +| `child_resolved` | a branch's arrival reached its parent's join | +| `wait_armed` | a wait or timer slot was armed | +| `wait_resolved` | a delivery satisfied a wait | +| `wait_expired` | a wait reached its deadline; the `on_timeout` branch runs | +| `signal_buffered` | a signal arrived before its wait was armed | +| `signal_duplicate` | a repeated sender key was ignored | +| `substep_recorded` | an `rx.step` result was journalled | + +Both wait outcomes are recorded, deliberately: "the approval came through" and +"nobody answered in time" lead to different handlers and different +conversations, and a history that showed only which handler ran next would +make an operator infer the difference instead of read it. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 685a98276fe..16f11c679b5 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -1364,6 +1364,24 @@ async def check_closing_a_branch_never_revives_a_finished_one( assert child.status is RunStatus.COMPLETED +async def check_a_duplicate_delivery_is_recorded_in_history(store: RunStore) -> None: + """A no-op delivery still has to be visible to whoever sent it. + + A repeated sender key is correctly ignored, and an ignored delivery that + leaves no record is indistinguishable from one that never arrived -- + which is the question the history exists to answer. + """ + await store.admit( + make_run(), + make_step(status=StepStatus.BLOCKED, wait_key="sig:ping", due_at=0.0), + _ADMITTED, + ) + assert await store.deliver("run1", "sig:ping", "d1", {"v": 1}, NOW) == "resolved" + assert await store.deliver("run1", "sig:ping", "d1", {"v": 1}, NOW) == "duplicate" + kinds = [event.type for event in await store.get_history("run1")] + assert kinds.count(HistoryEventType.SIGNAL_DUPLICATE) == 1, kinds + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1381,6 +1399,7 @@ async def check_closing_a_branch_never_revives_a_finished_one( check_delivery_resolves_a_matching_wait, check_delivery_never_touches_run_state, check_duplicate_deliveries_are_ignored, + check_a_duplicate_delivery_is_recorded_in_history, check_an_early_delivery_is_buffered_then_consumed, check_early_deliveries_queue_in_order, check_children_are_created_with_their_join, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 05ebac4ec28..a80b2c4184b 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -943,6 +943,18 @@ async def signal( ), ), ) + elif disposition == "duplicate": + # Correctly a no-op, but a no-op nobody can see is + # indistinguishable from a delivery that never arrived. + await self._notify_run( + run_id, + ( + ( + HistoryEventType.SIGNAL_DUPLICATE, + {"wait_key": f"sig:{delivery.channel}"}, + ), + ), + ) return disposition async def resume(self, run_id: str) -> bool: @@ -2181,14 +2193,29 @@ async def _execute_claim(self, claim: Claim) -> None: ) return handler = defn.handlers[claim.step.handler_id] + wait_events: tuple[tuple[HistoryEventType, dict[str, Any]], ...] = () if claim.step.status is StepStatus.CLAIMED and claim.step.wait_key is not None: expired = self._expired_wait_handler(defn, claim) if expired is not None: handler = expired + # A resolved wait records WAIT_RESOLVED at delivery; without + # this, an expired one recorded nothing, and the only trace of + # the deadline was which handler happened to run next. + wait_events = ( + ( + HistoryEventType.WAIT_EXPIRED, + { + "wait_key": claim.step.wait_key, + "ordinal": claim.step.ordinal, + "on_timeout": expired.id, + }, + ), + ) steps = await self._store.get_steps(claim.run.run_id) await self._record( claim.run, ( + *wait_events, ( HistoryEventType.ATTEMPT_STARTED, { diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 1f728307766..3d39e6317f9 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1174,6 +1174,12 @@ async def deliver( (run_id, wait_key, dedupe_key), ) if await cursor.fetchone() is not None: + await self._append_events( + conn, + run_id, + ((HistoryEventType.SIGNAL_DUPLICATE, {"wait_key": wait_key}),), + now, + ) return "duplicate" frontier = await self._frontier(conn, run_id) if ( diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 415b443a6ca..db86010f642 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -161,7 +161,6 @@ class HistoryEventType(str, enum.Enum): WAIT_ARMED = "wait_armed" WAIT_RESOLVED = "wait_resolved" WAIT_EXPIRED = "wait_expired" - SIGNAL_DELIVERED = "signal_delivered" SUBSTEP_RECORDED = "substep_recorded" SIGNAL_BUFFERED = "signal_buffered" SIGNAL_DUPLICATE = "signal_duplicate" diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 0b73e3f7e73..3b4acb33127 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -1315,6 +1315,11 @@ async def deliver( return "expired" inbox = self._inbox.setdefault(run_id, {}) if (run_id, wait_key, dedupe_key) in inbox: + self._append_events( + run_id, + ((HistoryEventType.SIGNAL_DUPLICATE, {"wait_key": wait_key}),), + now, + ) return "duplicate" inbox[run_id, wait_key, dedupe_key] = True steps = self._steps[run_id] @@ -3617,7 +3622,17 @@ def work(): (run_id, wait_key, dedupe_key), ).fetchone() if seen is not None: - self._db.execute("ROLLBACK") + self._append_events( + run_id, + ( + ( + HistoryEventType.SIGNAL_DUPLICATE, + {"wait_key": wait_key}, + ), + ), + now, + ) + self._db.execute("COMMIT") return "duplicate" frontier = _frontier(self._load_steps(run_id)) if frontier is not None and _wait_expired(frontier, now): diff --git a/tests/units/workflow/test_contract_vocabulary.py b/tests/units/workflow/test_contract_vocabulary.py new file mode 100644 index 00000000000..565e83bfb1a --- /dev/null +++ b/tests/units/workflow/test_contract_vocabulary.py @@ -0,0 +1,109 @@ +"""Every word the engine can say about a run must be a word the contract defines. + +Section 1's exit criterion is that every failure scenario has one unambiguous +documented outcome. The reasons in section 8 are only part of what an operator +actually reads: a run also has a status, its steps have statuses, a start and a +delivery each come back with a disposition, and its history is written in a +vocabulary of nearly thirty event types. Any of those can be added in a commit +that never touches CONTRACT.md, and nothing would notice. + +So this closes the loop in both directions. A member the contract does not +define is an outcome nobody documented; a member the contract defines but no +code produces is a promise nobody keeps -- and it is the second direction that +found ``wait_expired`` declared and never emitted, which left an expired +approval indistinguishable in history from an answered one. +""" + +from pathlib import Path + +import pytest + +from reflex.workflow.records import ( + HistoryEventType, + RunStatus, + StartDisposition, + StepStatus, +) +from reflex.workflow.store import DeliveryDisposition + +CONTRACT = ( + Path(__file__).parents[3] / "reflex" / "workflow" / "CONTRACT.md" +).read_text() + +ENGINE = tuple( + (Path(__file__).parents[3] / "reflex" / "workflow" / name).read_text() + for name in ("kernel.py", "store.py", "postgres.py", "api.py", "ingress.py") +) + + +@pytest.fixture(autouse=True) +def harness_store(): + """Opt out of the shared harness store parameter. + + This reads files; it has no store to vary. + + Returns: + The store kind this module reports. + """ + return "memory" + + +def _literals(alias) -> tuple[str, ...]: + """Read the members of a ``Literal`` alias. + + Args: + alias: The Literal type alias. + + Returns: + Its string members. + """ + return alias.__args__ + + +VOCABULARY = ( + *((member.value, "run status") for member in RunStatus), + *((member.value, "step status") for member in StepStatus), + *((member.value, "history event") for member in HistoryEventType), + *((value, "start disposition") for value in _literals(StartDisposition)), + *((value, "delivery disposition") for value in _literals(DeliveryDisposition)), +) + + +@pytest.mark.parametrize( + ("value", "kind"), + VOCABULARY, + ids=lambda item: item if isinstance(item, str) else "", +) +def test_every_observable_value_is_documented(value, kind): + """A value an operator can see must be defined in the contract. + + Args: + value: The observable value. + kind: What sort of value it is, for the failure message. + """ + assert f"`{value}`" in CONTRACT, ( + f"The {kind} {value!r} can appear in a run an operator reads, but " + "CONTRACT.md never defines it. Add it to section 10." + ) + + +@pytest.mark.parametrize( + "event", tuple(HistoryEventType), ids=lambda member: member.value +) +def test_every_documented_history_event_is_actually_emitted(event): + """A documented event nothing writes is a promise the engine breaks. + + History is what an operator reads to find out what happened. An event + type that exists in the enum and in the contract but is never written + describes a fact the run will never record -- so its absence from a run's + history means nothing, and reading history gives a false picture. + + Args: + event: The history event type under test. + """ + emitted = any(f"HistoryEventType.{event.name}" in source for source in ENGINE) + assert emitted, ( + f"{event.name} is declared and documented but no engine source emits " + "it. Emit it where it belongs, or delete it -- a vocabulary word that " + "never appears makes its own absence uninformative." + ) diff --git a/tests/units/workflow/test_waits.py b/tests/units/workflow/test_waits.py index 6938bbfeb38..eb96460ee9b 100644 --- a/tests/units/workflow/test_waits.py +++ b/tests/units/workflow/test_waits.py @@ -6,7 +6,7 @@ from reflex_base.workflow import Signal, WorkflowConfig, manual, never, wait_for import reflex as rx -from reflex.workflow.records import RunStatus, StepStatus +from reflex.workflow.records import HistoryEventType, RunStatus, StepStatus from reflex.workflow.store import MemoryRunStore, SqliteRunStore from reflex.workflow.testing import WorkflowTestHarness @@ -488,3 +488,135 @@ def expire(self): "the second stage was answered by the first stage's losing " f"alternative: {snapshot.result}" ) + + +async def test_an_expired_wait_says_so_in_history(forked_registration_context): + """History has to distinguish "nobody answered" from "somebody did". + + A resolved wait records ``wait_resolved``. An expired one recorded + nothing at all, so the only trace of the deadline was that the timeout + branch happened to be the handler that ran next -- an operator asking + "did the approval come through, or did it time out?" had to infer it + from which handler appears, which is exactly what history exists to stop. + + Args: + forked_registration_context: Isolates workflow registration. + """ + flow = _review_flow() + async with WorkflowTestHarness(flow) as harness: + started = await harness.start(flow.start()) + assert started.run_id is not None + await harness.advance("4d") + + history = await harness.kernel.store.get_history(started.run_id) + kinds = [event.type for event in history] + assert HistoryEventType.WAIT_ARMED in kinds + assert HistoryEventType.WAIT_EXPIRED in kinds, ( + f"an expired wait left no trace: {[k.value for k in kinds]}" + ) + assert HistoryEventType.WAIT_RESOLVED not in kinds, ( + "nobody answered, so nothing was resolved" + ) + expiry = next( + event for event in history if event.type is HistoryEventType.WAIT_EXPIRED + ) + assert expiry.data["wait_key"] == "sig:review" + + +async def test_a_resolved_wait_is_not_reported_as_expired( + forked_registration_context, +): + """The other half of the same distinction. + + Args: + forked_registration_context: Isolates workflow registration. + """ + flow = _review_flow() + async with WorkflowTestHarness(flow) as harness: + started = await harness.start(flow.start()) + assert started.run_id is not None + await harness.signal( + started.run_id, flow.review({"approved": True, "by": "ada"}) + ) + + kinds = [ + event.type + for event in await harness.kernel.store.get_history(started.run_id) + ] + assert HistoryEventType.WAIT_RESOLVED in kinds + assert HistoryEventType.WAIT_EXPIRED not in kinds + + +async def test_a_deduplicated_signal_leaves_a_trace(forked_registration_context): + """A sender's retry that changed nothing still has to be visible. + + "The provider says it delivered, so why didn't the run move?" is answered + by the run's history or by nothing at all. A repeated sender key is + correctly a no-op, and a no-op that leaves no record is indistinguishable + from a delivery that never arrived. The run has to still be alive for the + second delivery to be *duplicate* rather than *run_terminal*, which is + why this flow keeps going after it decides. + + Args: + forked_registration_context: Isolates workflow registration. + """ + + class LiveReview(rx.State): + __workflow__ = WorkflowConfig(id="waits.live_review") + decided_by: str = "" + + review = Signal(Decision) + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + """Wait for a decision. + + Returns: + The wait. + """ + return wait_for( + LiveReview.review, + then=LiveReview.decide, + timeout="3d", + on_timeout=LiveReview.expire, + ) + + @rx.event(durable=True, effect="none") + def decide(self, decision: Decision): + """Record the decision and stay alive. + + Args: + decision: The delivered decision. + + Returns: + A long deferral, so the run is still nonterminal. + """ + self.decided_by = decision.by + return rx.after("30d", LiveReview.expire) + + @rx.event(durable=True, effect="none") + def expire(self): + """Finish. + + Returns: + Completion. + """ + return rx.complete(result={"done": True}) + + async with WorkflowTestHarness(LiveReview) as harness: + started = await harness.start(LiveReview.start()) + assert started.run_id is not None + payload = LiveReview.review({"approved": True, "by": "ada"}) + assert await harness.signal(started.run_id, payload, key="hook-1") == "resolved" + assert ( + await harness.signal(started.run_id, payload, key="hook-1") == "duplicate" + ) + + kinds = [ + event.type + for event in await harness.kernel.store.get_history(started.run_id) + ] + assert kinds.count(HistoryEventType.WAIT_RESOLVED) == 1 + assert kinds.count(HistoryEventType.SIGNAL_DUPLICATE) == 1, ( + f"the deduplicated redelivery left no trace: {[k.value for k in kinds]}" + ) From 399277bd8dbdb69efbc2daae7010755edfc1d812 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 11:14:28 -0700 Subject: [PATCH 106/121] workflows: typed run results, and run-id prefixes in the CLI Two phase-3 items, found by walking the thing as a new developer rather than by reading it. RunHandle.result() returned Any. A result crosses the store as JSON, so callers got dicts back whatever the handler passed to rx.complete, and the type checker had nothing to say about it. It is now generic, and result(as_type=Receipt) returns a Receipt -- pyright infers it (checked with reveal_type) and pydantic validates it, so a result that does not fit raises here naming the run instead of becoming an AttributeError in the caller two frames later. Full inference from Orders.place(order) to RunHandle[Receipt] would mean threading generics through Reflex's @rx.event descriptor machinery, which every Reflex app shares. That is a bigger and riskier change than this one and is Alek's call, not mine. The CLI took full run ids only. "workflows dev" prints eight-character prefixes and "workflows list" prints full ids, so reading one and typing into the other -- the normal way these get used together -- gave "No run 'ca40d354' in this database" for an id the tool had just printed. Every command that takes a run id now resolves a prefix the way git does. An exact id is looked up directly and pays nothing; an ambiguous prefix refuses and names the candidates, because acting on the wrong run is worse than being asked to be specific. That also splits an error that used to run three causes together. A run that does not exist now says so; "unknown, already finished, or held by a worker" stays for a run that does exist and still cannot be finalized. --- ...kflow-typed-result-and-prefixes.feature.md | 3 + reflex/workflow/cli.py | 118 +++++++++++++++--- reflex/workflow/handle.py | 56 ++++++++- reflex/workflow/runtime.py | 2 +- tests/units/workflow/test_cli.py | 89 ++++++++++++- tests/units/workflow/test_handle.py | 72 +++++++++++ 6 files changed, 314 insertions(+), 26 deletions(-) create mode 100644 news/workflow-typed-result-and-prefixes.feature.md diff --git a/news/workflow-typed-result-and-prefixes.feature.md b/news/workflow-typed-result-and-prefixes.feature.md new file mode 100644 index 00000000000..714a65ea87e --- /dev/null +++ b/news/workflow-typed-result-and-prefixes.feature.md @@ -0,0 +1,3 @@ +`RunHandle` is generic and `result()` takes `as_type=`: `await handle.result(as_type=Receipt)` returns a `Receipt`, both to the type checker and at runtime. A result crosses the store as plain JSON, so it previously came back as dicts and lists whatever the handler passed to `rx.complete`. The coercion is a real validation — a result that does not fit raises here, naming the run, rather than surfacing as an `AttributeError` somewhere in the caller. + +Every CLI command that takes a run id now accepts an unambiguous prefix, the way git does. `reflex workflows dev` prints eight-character run prefixes while `list` prints full ids, so an operator reading one and typing into the other hit "No run 'ca40d354' in this database" for an id the tool itself had just printed. An exact id is still looked up directly and costs nothing extra; an ambiguous prefix refuses and names the candidates. A run that does not exist now says so precisely, instead of offering three possible causes at once. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 03a780d19fe..1ba357b8500 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -21,7 +21,7 @@ import click from reflex_base.utils import console -from reflex.workflow.records import RunStatus, attempts_made +from reflex.workflow.records import RunQuery, RunStatus, attempts_made if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Iterable @@ -64,10 +64,19 @@ def _operator_action(database: str | None, run_id: str, action: str, **extra): """ import time - applied = _with_store( - database, - lambda store: getattr(store, action)(run_id, time.time(), **extra), - ) + async def apply(store: RunStore) -> Any: + """Resolve the run, then apply the action to it. + + Args: + store: The open run store. + + Returns: + Whether the action applied. + """ + resolved = await _resolve_run_id(store, run_id) + return await getattr(store, action)(resolved, time.time(), **extra) + + applied = _with_store(database, apply) if not applied: console.error( f"Run {run_id!r} is not in a state that allows {action.split('_')[0]!r}." @@ -76,6 +85,9 @@ def _operator_action(database: str | None, run_id: str, action: str, **extra): console.print(f"Applied {action.split('_')[0]} to {run_id}.") +_PREFIX_SCAN_LIMIT = 10_000 + + def _open_store(database: str | None) -> RunStore: """Open the run store the app persists to. @@ -92,6 +104,42 @@ def _open_store(database: str | None) -> RunStore: return resolve_store(database) +async def _resolve_run_id(store: RunStore, run_id: str) -> str: + """Accept an unambiguous id prefix wherever a run id is taken. + + ``list`` prints full ids but ``dev`` prints eight-character prefixes, and + a person who has been reading the second naturally types one. An exact id + is looked up directly and costs nothing extra; only a prefix pays for a + scan. + + Args: + store: The open run store. + run_id: A full run id or a prefix of one. + + Returns: + The full run id. + + Raises: + Exit: If the prefix matches no run, or more than one. + """ + if await store.get_run(run_id) is not None: + return run_id + matches = [ + run.run_id + for run in await store.list_runs(RunQuery(limit=_PREFIX_SCAN_LIMIT)) + if run.run_id.startswith(run_id) + ] + if len(matches) == 1: + return matches[0] + if not matches: + console.error(f"No run {run_id!r} in this database.") + raise click.exceptions.Exit(1) + listed = ", ".join(match[:12] for match in matches[:5]) + more = f" and {len(matches) - 5} more" if len(matches) > 5 else "" + console.error(f"{run_id!r} matches several runs: {listed}{more}.") + raise click.exceptions.Exit(1) + + def _with_store(database: str | None, work: Callable[[RunStore], Awaitable[Any]]): """Open a store, run one unit of work against it, and close it. @@ -1033,10 +1081,11 @@ async def load(store: RunStore): Returns: The run, its steps, and its history events. """ + resolved = await _resolve_run_id(store, run_id) return ( - await store.get_run(run_id), - await store.get_steps(run_id), - await store.get_history(run_id) if history else (), + await store.get_run(resolved), + await store.get_steps(resolved), + await store.get_history(resolved) if history else (), ) run, steps, events = _with_store(database, load) @@ -1220,9 +1269,20 @@ def cancel(database: str | None, run_id: str): """ import time - recorded = _with_store( - database, lambda store: store.request_cancel(run_id, time.time()) - ) + async def request(store: RunStore) -> bool: + """Record the intent against the resolved run. + + Args: + store: The open run store. + + Returns: + Whether intent was recorded. + """ + return await store.request_cancel( + await _resolve_run_id(store, run_id), time.time() + ) + + recorded = _with_store(database, request) if not recorded: console.error(f"Run {run_id!r} is unknown or already finished.") raise click.exceptions.Exit(1) @@ -1280,12 +1340,23 @@ def _finalize( """ from reflex.workflow.kernel import WorkflowKernel - finalized = _with_store( - database, - lambda store: WorkflowKernel([], store).force_finalize( - run_id, status=status, result=result, error=error - ), - ) + async def finish(store: RunStore) -> bool: + """Resolve the run, then finalize it through a kernel. + + Args: + store: The open run store. + + Returns: + Whether the run was finalized. + """ + return await WorkflowKernel([], store).force_finalize( + await _resolve_run_id(store, run_id), + status=status, + result=result, + error=error, + ) + + finalized = _with_store(database, finish) if not finalized: console.error( f"Run {run_id!r} is unknown, already finished, or has a step a " @@ -1339,7 +1410,18 @@ def resume(database: str | None, run_id: str): """Re-open a run suspended for operator attention.""" import time - resumed = _with_store(database, lambda store: store.resume_run(run_id, time.time())) + async def reopen(store: RunStore) -> bool: + """Resolve the run, then re-open it. + + Args: + store: The open run store. + + Returns: + Whether a suspended run was re-opened. + """ + return await store.resume_run(await _resolve_run_id(store, run_id), time.time()) + + resumed = _with_store(database, reopen) if not resumed: console.error(f"Run {run_id!r} is not suspended.") raise click.exceptions.Exit(1) diff --git a/reflex/workflow/handle.py b/reflex/workflow/handle.py index a175f251d5e..58ccadd3906 100644 --- a/reflex/workflow/handle.py +++ b/reflex/workflow/handle.py @@ -14,7 +14,7 @@ from __future__ import annotations import asyncio -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload from reflex_base.utils.exceptions import WorkflowRuntimeError from reflex_base.workflow import DurationLike, parse_duration @@ -28,8 +28,11 @@ DEFAULT_POLL_INTERVAL: float = 0.05 +ResultT = TypeVar("ResultT") +T = TypeVar("T") -class RunHandle: + +class RunHandle(Generic[ResultT]): """One run, and the things a caller does with it. Attributes: @@ -85,9 +88,28 @@ async def status(self) -> RunStatus | None: snapshot = await self.snapshot() return None if snapshot is None else snapshot.status + @overload + async def result( + self, + *, + as_type: type[T], + timeout: DurationLike = "30s", + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> T: ... + + @overload async def result( self, *, + as_type: None = None, + timeout: DurationLike = "30s", + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> ResultT: ... + + async def result( + self, + *, + as_type: type[Any] | None = None, timeout: DurationLike = "30s", poll_interval: float = DEFAULT_POLL_INTERVAL, ) -> Any: @@ -99,16 +121,29 @@ async def result( takes. Compose with ``rx.parallel`` instead, which is what child runs and joins are for. + A result crosses the store as plain JSON data, so it comes back as + dicts and lists whatever the handler passed to ``rx.complete``. Pass + ``as_type`` to get the shape back:: + + receipt = await handle.result(as_type=Receipt) + receipt.total + + That is a real validation, not a cast: a result that does not fit the + declared type raises here, naming the run, rather than becoming an + ``AttributeError`` further along in the caller. + Args: + as_type: Type to validate and coerce the result into. timeout: How long to wait before giving up. poll_interval: Seconds between checks. Returns: - The run's result. + The run's result, coerced to ``as_type`` when one was given. Raises: WorkflowRuntimeError: If the run is unknown, does not finish in - time, or finishes in any state other than completed. + time, finishes in any state other than completed, or produced + a result that does not fit ``as_type``. """ snapshot = await self.wait(timeout=timeout, poll_interval=poll_interval) if snapshot.status is not RunStatus.COMPLETED: @@ -118,7 +153,18 @@ async def result( f"COMPLETED{detail}" ) raise WorkflowRuntimeError(msg) - return snapshot.result + if as_type is None: + return snapshot.result + from pydantic import TypeAdapter, ValidationError + + try: + return TypeAdapter(as_type).validate_python(snapshot.result) + except ValidationError as error: + msg = ( + f"Run {self.run_id} completed with a result that does not fit " + f"{getattr(as_type, '__name__', as_type)}: {error}" + ) + raise WorkflowRuntimeError(msg) from error async def wait( self, diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index 03e85bd36d3..c6e8964c95e 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -414,7 +414,7 @@ async def submit( *, request_key: str | None = None, labels: dict[str, str] | None = None, - ) -> RunHandle: + ) -> RunHandle[Any]: """Start a run and get a handle on it. The same admission as ``start()``, returning the run rather than a diff --git a/tests/units/workflow/test_cli.py b/tests/units/workflow/test_cli.py index 62254ab3721..53022df4ddc 100644 --- a/tests/units/workflow/test_cli.py +++ b/tests/units/workflow/test_cli.py @@ -240,11 +240,33 @@ def test_fail_records_the_operator_s_reason(seeded): def test_finalizing_an_unknown_run_fails(seeded): - """Nothing to finalize is an error, not a silent success.""" + """Nothing to finalize is an error, not a silent success. + + A run that does not exist now says exactly that. The older message + offered three possibilities at once -- unknown, already finished, or + held by a worker -- which is still right for a run that does exist, and + unhelpfully vague for one that does not. + + Args: + seeded: The database and its run ids. + """ database, _, _ = seeded result = _invoke("complete", "-d", database, "no-such-run") assert result.exit_code == 1 - assert "unknown" in result.output + assert "No run" in result.output + + +def test_finalizing_a_finished_run_says_why(seeded): + """A run that exists but cannot be finalized keeps the fuller message. + + Args: + seeded: The database and its run ids. + """ + database, waiting, _ = seeded + assert _invoke("complete", "-d", database, waiting).exit_code == 0 + again = _invoke("complete", "-d", database, waiting) + assert again.exit_code == 1 + assert "already finished" in again.output def test_complete_refuses_a_result_that_is_not_json(seeded): @@ -268,3 +290,66 @@ def test_purge_deletes_only_stale_terminal_runs(seeded): assert purged.exit_code == 0, purged.output assert "Purged 1 run(s)" in purged.output assert _invoke("show", "-d", database, waiting).exit_code == 1 + + +def test_a_run_id_prefix_is_enough(seeded): + """`dev` prints eight-character prefixes, so the CLI has to take them. + + Reading one surface and typing into another is the normal way an operator + uses these commands, and "No run 'ca40d354'" for an id the tool itself + printed is a dead end. + + Args: + seeded: The database and its run ids. + """ + database, waiting_id, _ = seeded + full = _invoke("show", "-d", database, waiting_id, "--json") + assert full.exit_code == 0, full.output + short = _invoke("show", "-d", database, waiting_id[:8], "--json") + assert short.exit_code == 0, short.output + assert json.loads(short.output)["run_id"] == waiting_id + + +def test_an_ambiguous_prefix_refuses_and_names_the_candidates(seeded): + """Acting on the wrong run is worse than being asked to be specific. + + Args: + seeded: The database and its run ids. + """ + database, waiting_id, suspended_id = seeded + shared = "" + for index in range(1, 33): + if waiting_id[:index] != suspended_id[:index]: + break + shared = waiting_id[:index] + if not shared: + pytest.skip("the two seeded run ids share no prefix this time") + result = _invoke("show", "-d", database, shared) + assert result.exit_code == 1 + assert "matches several runs" in result.output + + +def test_an_unknown_prefix_still_says_so(seeded): + """The no-match message must not be lost to the new prefix path. + + Args: + seeded: The database and its run ids. + """ + database, _, _ = seeded + result = _invoke("show", "-d", database, "nosuchrun") + assert result.exit_code == 1 + assert "No run" in result.output + + +def test_operator_actions_take_a_prefix_too(seeded): + """Every command that takes a run id resolves the same way. + + Args: + seeded: The database and its run ids. + """ + database, _, suspended_id = seeded + result = _invoke("resume", "-d", database, suspended_id[:8]) + assert result.exit_code == 0, result.output + run = _load_run(database, suspended_id) + assert run is not None + assert run.status is not RunStatus.NEEDS_ATTENTION diff --git a/tests/units/workflow/test_handle.py b/tests/units/workflow/test_handle.py index 8e363c3b68f..8a3ad61b9ee 100644 --- a/tests/units/workflow/test_handle.py +++ b/tests/units/workflow/test_handle.py @@ -7,6 +7,7 @@ """ import pytest +from pydantic import BaseModel from reflex_base.utils.exceptions import WorkflowRuntimeError from reflex_base.workflow import RateLimit, WorkflowConfig, manual @@ -186,3 +187,74 @@ def go(self): with pytest.raises(WorkflowRuntimeError, match="rejected"): await rx.workflows.submit(Limited.go) _ = harness + + +class Receipt(BaseModel): + """What a completed order produces.""" + + order: str + total: int + + +class Ordering(rx.State): + """A workflow that completes with a structured result.""" + + __workflow__ = WorkflowConfig(id="handle.ordering") + + @rx.event(durable=True, trigger=manual(), effect="none") + def place(self, order: str): + """Complete with a receipt. + + Args: + order: The order identifier. + + Returns: + Completion carrying the receipt. + """ + return rx.complete(result={"order": order, "total": 250}) + + @rx.event(durable=True, trigger=manual(), effect="none") + def wrong(self): + """Complete with something that is not a receipt. + + Returns: + Completion carrying the wrong shape. + """ + return rx.complete(result={"order": "ord_1"}) + + +async def test_a_typed_result_comes_back_as_the_declared_shape( + forked_registration_context, +): + """A result crosses the store as JSON; ``as_type`` restores the shape. + + Args: + forked_registration_context: Isolates workflow registration. + """ + async with WorkflowTestHarness(Ordering) as harness: + handle = await rx.workflows.submit(Ordering.place("ord_1")) + await harness.run_until_idle() + assert await handle.result() == {"order": "ord_1", "total": 250} + + receipt = await handle.result(as_type=Receipt) + assert isinstance(receipt, Receipt) + assert receipt.total == 250 + assert receipt.order == "ord_1" + + +async def test_a_result_that_does_not_fit_names_the_run(forked_registration_context): + """Validation is real, and its error points at the run that produced it. + + A cast would hand the caller a dict that fails with AttributeError + somewhere else entirely, long after the information needed to explain it + has gone. + + Args: + forked_registration_context: Isolates workflow registration. + """ + async with WorkflowTestHarness(Ordering) as harness: + handle = await rx.workflows.submit(Ordering.wrong()) + await harness.run_until_idle() + with pytest.raises(WorkflowRuntimeError, match="does not fit Receipt"): + await handle.result(as_type=Receipt) + assert await handle.result() == {"order": "ord_1"} From 9f02f2746f5d1d05115c002c0ab61863bbd7af18 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 11:20:28 -0700 Subject: [PATCH 107/121] workflows: an OpenTelemetry observer The observer seam was built for this -- MetricsObserver's docstring promises "a metrics endpoint or an OpenTelemetry exporter is a few lines over snapshot()" -- and nothing had cashed it. A span here covers one ATTEMPT, not one run. That is the whole design question and it only has one answer: a span is in-process and time-bounded, and a durable run can wait a day, survive a restart, and execute its steps on different machines. Modelling a run as a span would mean holding one open across processes, which OpenTelemetry cannot do and no backend would draw. Every span carries workflow.run_id instead, so "everything that happened to this run" is an attribute search across however many attempts it took. Steps that never ran produce no spans, because nothing executed to time. Counters reuse MetricsObserver's event-to-name mapping rather than restating it. Two exporters that disagree about one deployment are worse than one exporter, and a test runs both observers over the same runs and asserts they report the same numbers. OpenTelemetry is not a Reflex dependency and this does not make it one: the import is lazy, so importing the module without it works, and constructing the observer raises with the pip line. Both paths are verified. It is in the dev group so CI actually exercises the code rather than skipping the file and shipping it untested. Tested through the real SDK's in-memory exporter and reader, not a mock -- whether a span actually ends, whether a failure survives into its status, and whether the counters agree are not questions a mock can answer. --- news/workflow-opentelemetry.feature.md | 3 + pyproject.toml | 2 + reflex/workflow/otel.py | 188 +++++++++++++++++++++++++ tests/units/workflow/test_otel.py | 165 ++++++++++++++++++++++ uv.lock | 43 ++++++ 5 files changed, 401 insertions(+) create mode 100644 news/workflow-opentelemetry.feature.md create mode 100644 reflex/workflow/otel.py create mode 100644 tests/units/workflow/test_otel.py diff --git a/news/workflow-opentelemetry.feature.md b/news/workflow-opentelemetry.feature.md new file mode 100644 index 00000000000..9c0b915c47e --- /dev/null +++ b/news/workflow-opentelemetry.feature.md @@ -0,0 +1,3 @@ +Added `reflex.workflow.otel.OpenTelemetryObserver`, which forwards workflow activity to OpenTelemetry as a span per attempt plus counters per transition. OpenTelemetry stays optional — Reflex does not install it, importing the module without it is fine, and constructing the observer says exactly what to install. + +A span covers one attempt (claim to commit), not one run. A durable run can wait a day and cross machines, which is not something a span can represent; every span instead carries `workflow.run_id`, so "everything that happened to this run" is an attribute search across however many attempts and processes it took. The counters reuse `MetricsObserver`'s event-to-name mapping, so the two exporters cannot drift into reporting different numbers for the same deployment — a test asserts they agree. diff --git a/pyproject.toml b/pyproject.toml index 6ae0499cea8..aca8d6bc524 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,6 +86,8 @@ dev = [ "hatchling", "libsass", "numpy", + "opentelemetry-api", + "opentelemetry-sdk", "pandas", "pillow", "playwright", diff --git a/reflex/workflow/otel.py b/reflex/workflow/otel.py new file mode 100644 index 00000000000..d5725e707fd --- /dev/null +++ b/reflex/workflow/otel.py @@ -0,0 +1,188 @@ +"""Forward workflow activity to OpenTelemetry. + +Install with ``pip install opentelemetry-api opentelemetry-sdk`` -- this is +never a hard dependency of Reflex, and importing this module without it says +so plainly instead of failing somewhere inside a callback. + +**What is traced, and what deliberately is not.** A span is an in-process, +time-bounded thing. A durable run is neither: it can wait a day, cross a +restart, and execute its steps on different machines. Modelling one run as +one span would mean holding a span open across processes, which OpenTelemetry +cannot do and no backend would render usefully. + +So the span here is *one attempt* -- claim to commit -- which really is +bounded and really does happen in one process. Every span carries +``workflow.run_id``, so "show me everything that happened to this run" is a +search by attribute across however many days and processes it took, rather +than a single trace. Steps that never ran (a run cancelled while waiting) +produce no spans, because nothing executed to time. + +Counters mirror :class:`~reflex.workflow.kernel.MetricsObserver` exactly, +sharing its event-to-name mapping so the two exporters can never drift into +reporting different numbers for the same deployment. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Final + +from reflex.workflow.kernel import MetricsObserver, WorkflowObserver +from reflex.workflow.records import HistoryEventType + +if TYPE_CHECKING: + from collections.abc import Mapping + +INSTRUMENTATION_NAME: Final = "reflex.workflow" + +_ATTEMPT_ENDINGS: Final[dict[HistoryEventType, str]] = { + HistoryEventType.ATTEMPT_SUCCEEDED: "succeeded", + HistoryEventType.ATTEMPT_FAILED: "failed", + HistoryEventType.ATTEMPT_TIMED_OUT: "timed_out", + HistoryEventType.ATTEMPT_CANCELLED: "cancelled", + HistoryEventType.ATTEMPT_ABANDONED: "abandoned", +} + +_OK_ENDINGS: Final = frozenset({"succeeded"}) + +MAX_OPEN_SPANS: Final = 4096 + + +def _require_opentelemetry(): + """Import OpenTelemetry, explaining the install if it is absent. + + Returns: + The ``trace``, ``metrics``, and ``Status``/``StatusCode`` handles. + + Raises: + ImportError: If OpenTelemetry is not installed. + """ + try: + from opentelemetry import metrics, trace + from opentelemetry.trace import Status, StatusCode + except ImportError as error: + msg = ( + "OpenTelemetryObserver needs OpenTelemetry, which Reflex does not " + "install: pip install opentelemetry-api opentelemetry-sdk" + ) + raise ImportError(msg) from error + return trace, metrics, Status, StatusCode + + +class OpenTelemetryObserver(WorkflowObserver): + """Emit a span per attempt and a counter per transition. + + Install it the same way as any observer:: + + app = rx.App(workflow_observer=OpenTelemetryObserver()) + + Callbacks never raise: the kernel swallows observer errors by design, and + an exporter that broke a run would be worse than one that lost a span. + """ + + def __init__(self, *, tracer_provider: Any = None, meter_provider: Any = None): + """Bind to a tracer and a meter. + + Args: + tracer_provider: Provider to take the tracer from; the global one + by default. + meter_provider: Provider to take the meter from; the global one + by default. + """ + trace, metrics, status, status_code = _require_opentelemetry() + self._status = status + self._status_code = status_code + self._tracer = (tracer_provider or trace).get_tracer(INSTRUMENTATION_NAME) + meter = (meter_provider or metrics).get_meter(INSTRUMENTATION_NAME) + self._counters = { + name: meter.create_counter(f"reflex.workflow.{name}") + for name in set(MetricsObserver._COUNTED.values()) # pyright: ignore[reportPrivateUsage] + } + self._open: dict[tuple[str, Any], Any] = {} + + def on_event( + self, + event_type: HistoryEventType, + run_id: str, + workflow_id: str, + data: dict[str, Any], + ) -> None: + """Record one transition as a span edge and a counter increment. + + Args: + event_type: What happened. + run_id: The run it happened to. + workflow_id: That run's workflow identity. + data: Event payload, such as ordinal, handler, attempt, or error. + """ + self._count(event_type, workflow_id) + if event_type is HistoryEventType.ATTEMPT_STARTED: + self._open_span(run_id, workflow_id, data) + return + ending = _ATTEMPT_ENDINGS.get(event_type) + if ending is not None: + self._close_span(run_id, data, ending) + + def _count(self, event_type: HistoryEventType, workflow_id: str) -> None: + """Add one to the counter this event feeds, if any. + + Args: + event_type: What happened. + workflow_id: The workflow it happened in. + """ + name = MetricsObserver._COUNTED.get(event_type) # pyright: ignore[reportPrivateUsage] + counter = self._counters.get(name) if name else None + if counter is not None: + counter.add(1, {"workflow.id": workflow_id}) + + def _open_span( + self, run_id: str, workflow_id: str, data: Mapping[str, Any] + ) -> None: + """Start a span for an attempt that just began. + + Args: + run_id: The run being executed. + workflow_id: That run's workflow identity. + data: The ``attempt_started`` payload. + """ + if len(self._open) >= MAX_OPEN_SPANS: + # An attempt whose ending never arrived -- its worker was killed + # between the two events. Ending it as unset says "we do not know + # how this finished", which is true, and keeps the map bounded. + oldest = next(iter(self._open)) + self._open.pop(oldest).end() + handler = data.get("handler_id", "?") + span = self._tracer.start_span( + f"{workflow_id}.{handler}", + attributes={ + "workflow.id": workflow_id, + "workflow.run_id": run_id, + "workflow.handler_id": handler, + "workflow.step_ordinal": data.get("ordinal", -1), + "workflow.attempt": data.get("attempt", 0), + "workflow.effect": str(data.get("effect", "")), + }, + ) + self._open[run_id, data.get("ordinal")] = span + + def _close_span(self, run_id: str, data: Mapping[str, Any], ending: str) -> None: + """Finish the span for an attempt that just ended. + + Args: + run_id: The run being executed. + data: The ending event's payload. + ending: How the attempt ended. + """ + span = self._open.pop((run_id, data.get("ordinal")), None) + if span is None: + return + span.set_attribute("workflow.attempt_outcome", ending) + reason = data.get("reason") + if reason: + span.set_attribute("workflow.reason", str(reason)) + if ending in _OK_ENDINGS: + span.set_status(self._status(self._status_code.OK)) + else: + span.set_status( + self._status(self._status_code.ERROR, str(reason or ending)) + ) + span.end() diff --git a/tests/units/workflow/test_otel.py b/tests/units/workflow/test_otel.py new file mode 100644 index 00000000000..27c47d89e2b --- /dev/null +++ b/tests/units/workflow/test_otel.py @@ -0,0 +1,165 @@ +"""The OpenTelemetry observer, driven by real runs through a real SDK. + +Exported against the SDK's in-memory exporter and reader rather than a mock, +because the questions worth asking are whether a span actually ends, whether +its status survives a failure, and whether the counters agree with the +non-OpenTelemetry exporter -- none of which a mock can answer. +""" + +import pytest +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.kernel import CompositeObserver, MetricsObserver +from reflex.workflow.testing import WorkflowTestHarness + +pytest.importorskip("opentelemetry", reason="OpenTelemetry is an optional extra") + +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from reflex.workflow.otel import OpenTelemetryObserver + + +class Charge(rx.State): + """One good step and one that fails for good.""" + + __workflow__ = WorkflowConfig(id="otel.charge") + amount: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self, amount: int): + """Succeed. + + Args: + amount: The amount to record. + + Returns: + Completion. + """ + self.amount = amount + return rx.complete(result={"amount": amount}) + + @rx.event( + durable=True, trigger=manual(), effect="none", retry=Retry(max_attempts=1) + ) + def doomed(self): + """Fail. + + Raises: + TransientWorkflowError: Always. + """ + msg = "vendor down" + raise TransientWorkflowError(msg) + + +@pytest.fixture +def otel(): + """Wire an observer to in-memory span and metric collectors. + + Returns: + The observer, the span exporter, and the metric reader. + """ + exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + reader = InMemoryMetricReader() + observer = OpenTelemetryObserver( + tracer_provider=tracer_provider, + meter_provider=MeterProvider(metric_readers=[reader]), + ) + return observer, exporter, reader + + +def _counter(reader, name: str) -> int: + """Total one counter across its attribute sets. + + Args: + reader: The in-memory metric reader. + name: The counter's full name. + + Returns: + The summed value, or zero when the counter never fired. + """ + total = 0 + data = reader.get_metrics_data() + for resource in data.resource_metrics if data else (): + for scope in resource.scope_metrics: + for metric in scope.metrics: + if metric.name == name: + total += sum(point.value for point in metric.data.data_points) + return total + + +async def test_a_successful_attempt_becomes_one_ok_span( + otel, forked_registration_context +): + """The span names the handler and carries the run it belongs to. + + Args: + otel: The observer and its collectors. + forked_registration_context: Isolates workflow registration. + """ + observer, exporter, _ = otel + async with WorkflowTestHarness(Charge, observer=observer) as harness: + started = await harness.start(Charge.go(2500)) + assert started.run_id is not None + + spans = exporter.get_finished_spans() + assert len(spans) == 1, [span.name for span in spans] + span = spans[0] + assert span.name == "otel.charge.go" + assert span.attributes["workflow.run_id"] == started.run_id + assert span.attributes["workflow.handler_id"] == "go" + assert span.attributes["workflow.attempt"] == 1 + assert span.attributes["workflow.attempt_outcome"] == "succeeded" + assert span.status.is_ok + + +async def test_a_failed_attempt_ends_its_span_with_an_error( + otel, forked_registration_context +): + """A span left open would be worse than no span at all. + + Args: + otel: The observer and its collectors. + forked_registration_context: Isolates workflow registration. + """ + observer, exporter, _ = otel + async with WorkflowTestHarness(Charge, observer=observer) as harness: + await harness.start(Charge.doomed()) + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["workflow.attempt_outcome"] == "failed" + assert not spans[0].status.is_ok + + +async def test_the_counters_agree_with_the_other_exporter( + otel, forked_registration_context +): + """Two exporters that disagree about one deployment are worse than one. + + Args: + otel: The observer and its collectors. + forked_registration_context: Isolates workflow registration. + """ + observer, _, reader = otel + metrics_observer = MetricsObserver() + both = CompositeObserver(observer, metrics_observer) + async with WorkflowTestHarness(Charge, observer=both) as harness: + await harness.start(Charge.go(1)) + await harness.start(Charge.go(2)) + await harness.start(Charge.doomed()) + + for name, expected in ( + ("runs_started", 3), + ("runs_completed", 2), + ("runs_failed", 1), + ("attempts", 3), + ): + assert metrics_observer.totals.get(name, 0) == expected, name + assert _counter(reader, f"reflex.workflow.{name}") == expected, name diff --git a/uv.lock b/uv.lock index ff3e75cc6c7..3f2f8326eef 100644 --- a/uv.lock +++ b/uv.lock @@ -2424,6 +2424,45 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/0f/cc6afea3542a5142c5d8fc8211c5e059a8375105d004a41dfa2c7948dbb0/openai-2.53.0-py3-none-any.whl", hash = "sha256:c694ffc747a3c4d1663ef2b07b811315a476164ee5efa3a993967349ebca7618", size = 1659829, upload-time = "2026-08-03T21:41:59.581Z" }, ] +[[package]] +name = "opentelemetry-api" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ee/8b/aa9e2d8b8dfa7c946f7dec5d1f8f6ba8eca062f43509a06bdb5ce93d26c0/opentelemetry_api-1.44.0.tar.gz", hash = "sha256:67647e5e9566edcf421166fdf022b3537f818635daa852b289e34604dc6fb33a", size = 72406, upload-time = "2026-07-16T15:25:32.678Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/6f/a04e900f465ff3221ccc395522503e2d10e79fa21f2723c8e177aae1e0d1/opentelemetry_api-1.44.0-py3-none-any.whl", hash = "sha256:94b98c893a91b88657eaac1e3ba89618cdb85be6918196705354f34728b2cdef", size = 60018, upload-time = "2026-07-16T15:25:11.657Z" }, +] + +[[package]] +name = "opentelemetry-sdk" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5d/77/a6592cbc7c8d9bcc9d6757a9df45e04a7c585e3e6e7a13456da522b21109/opentelemetry_sdk-1.44.0.tar.gz", hash = "sha256:cebe7f65dc12f26ead75c6064de12fd2a9052e5060c0272d402cfa203aae123b", size = 208624, upload-time = "2026-07-16T15:25:46.078Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/23/ff077e61886ee020a17ce9c8b6fa11c601c8d8345b09ea24f605445df62a/opentelemetry_sdk-1.44.0-py3-none-any.whl", hash = "sha256:df081c4c6bcfdb1211e3e86140376792643128a25f8d72d1d27675936e7e96ad", size = 137221, upload-time = "2026-07-16T15:25:29.534Z" }, +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.65b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8f/73/0cbdebcb4cf545fdd328da14f5137e37d0770c3f26185e478b0d15d94f50/opentelemetry_semantic_conventions-0.65b0.tar.gz", hash = "sha256:f9b2b81e9d5b64f11bc952075e7e9c7fb0aab075c7fd1c46d597f1b919852d60", size = 148774, upload-time = "2026-07-16T15:25:46.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/0e/49df70d9b81fb5cbae4bbf2a49d865b09bcbcbc4eb53f5851b1027738d78/opentelemetry_semantic_conventions-0.65b0-py3-none-any.whl", hash = "sha256:1cacde7b0ad306f84c5ef08c3dbe1bbaf20165bba6f8bff43b670e555a086bcb", size = 204645, upload-time = "2026-07-16T15:25:30.688Z" }, +] + [[package]] name = "orjson" version = "3.11.9" @@ -3705,6 +3744,8 @@ dev = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-sdk" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pillow" }, @@ -3787,6 +3828,8 @@ dev = [ { name = "hatchling" }, { name = "libsass" }, { name = "numpy" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-sdk" }, { name = "pandas" }, { name = "pillow" }, { name = "playwright" }, From 20273f61750f8d6d20a20cb7f9858711256d7e1c Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 11:27:18 -0700 Subject: [PATCH 108/121] workflows: audit composition against the graduation dimensions Phase 5's bar is that each feature graduates only once its crash, race, authorization and versioning semantics are tested. Checking that directly, rather than assuming the general tests cover the specific features, turned up two gaps and one pleasant surprise. Authorization: I expected an unverified webhook to be an open run-starter. It is not -- compiling refuses one outright unless someone passes allow_unverified with a non-empty reason, which is better than a warning. But once opted in, doctor said nothing, and doctor exists for "the things whose absence is silent". The refusal protects whoever writes the webhook; whoever deploys it a year later is a different person reading a different surface. It is a note, not an error, because they did opt in -- it just should not be invisible. Versioning: the existing tests all deploy over plain sequential steps. A wait's continuation and a fan-out's join reach their handler by another route, carrying an injected payload or results list, so the compatibility gate applying to them as well was an unverified assumption. It does hold, and both are now pinned: removing the handler suspends the run as unknown_handler instead of raising inside a worker that is serving every other run in the deployment. The join test is mutation-checked -- restoring the removed handler makes it fail -- because a versioning test that passes for the wrong reason looks exactly like one that passes. --- news/workflow-graduation-audit.bugfix.md | 3 + reflex/workflow/cli.py | 11 ++ tests/units/workflow/test_cli_check.py | 57 +++++++ tests/units/workflow/test_versioning.py | 203 ++++++++++++++++++++++- 4 files changed, 273 insertions(+), 1 deletion(-) create mode 100644 news/workflow-graduation-audit.bugfix.md diff --git a/news/workflow-graduation-audit.bugfix.md b/news/workflow-graduation-audit.bugfix.md new file mode 100644 index 00000000000..cc38c3efbf6 --- /dev/null +++ b/news/workflow-graduation-audit.bugfix.md @@ -0,0 +1,3 @@ +`reflex workflows doctor` now names any webhook that accepts unverified deliveries, along with the reason its author gave. Compiling already refuses an unverified webhook unless `allow_unverified=True` is passed with a reason, but that protects whoever writes the code; whoever deploys it later reads the preflight, and a publicly writable run-starter belongs there. + +Added versioning coverage for the composition features. A wait's continuation and a fan-out's join are reached through a different path than a plain successor — with an injected payload or results list — so "the compatibility gate applies there too" was a separate fact from the one the existing tests pinned, and an unchecked one. Both now suspend the run as `unknown_handler` when their handler is removed by a deploy, rather than taking down the worker that was serving every other run. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 1ba357b8500..f8c22b96626 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -482,6 +482,17 @@ def doctor(database: str | None, target: str): f"{secret_env} is unset, so {definition.workflow_id}." f"{handler.name} will refuse every delivery." ) + if isinstance(trigger, WebhookTrigger) and trigger.verify is None: + # Compiling already refused this unless someone opted in, but + # that protects whoever wrote it. Whoever deploys it, possibly + # a year later, is a different person and this is the preflight + # they read. + notes.append( + f"{definition.workflow_id}.{handler.name} accepts unverified " + f"deliveries on topic '{trigger.topic}' -- anyone who knows " + f"the URL can start it. Declared reason: " + f"{trigger.unverified_reason or 'none given'}." + ) if isinstance(trigger, ScheduleTrigger): notes.append( f"{definition.workflow_id}.{handler.name} runs on " diff --git a/tests/units/workflow/test_cli_check.py b/tests/units/workflow/test_cli_check.py index d898ea6b5ef..46586615dd7 100644 --- a/tests/units/workflow/test_cli_check.py +++ b/tests/units/workflow/test_cli_check.py @@ -454,3 +454,60 @@ def test_doctor_notes_schedules_and_unmounted_surfaces( assert result.exit_code == 0, result.output assert "0 9 * * *" in result.output assert "REFLEX_WORKFLOW_API_TOKEN" in result.output + + +UNVERIFIED = ''' +import reflex as rx + + +class Open(rx.State): + __workflow__ = rx.WorkflowConfig(id="doctor.open") + + @rx.event( + durable=True, + effect="none", + trigger=rx.webhook( + "payout", + allow_unverified=True, + unverified_reason="behind an internal load balancer", + ), + ) + def on_hook(self, payload: dict): + """Take an unverified delivery. + + Args: + payload: The delivered body. + + Returns: + Completion. + """ + return rx.complete(result=payload) +''' + + +def test_doctor_names_a_webhook_that_takes_anonymous_deliveries( + tmp_path, forked_registration_context +): + """Opting in protects the author; deploying is done by someone else. + + Compiling already refuses an unverified webhook unless someone passed + allow_unverified with a reason. That is a decision made once, by whoever + wrote it. The person deploying a year later reads this preflight instead, + and an endpoint anyone can post runs into is exactly what it is for. + + Args: + tmp_path: Temporary directory for the module and database. + forked_registration_context: Isolates state registration. + """ + module = tmp_path / "flows_open.py" + module.write_text(UNVERIFIED) + result = CliRunner().invoke( + workflows, ["doctor", str(module), "-d", str(tmp_path / "d.db")] + ) + assert result.exit_code == 0, result.output + # Rich wraps at the terminal width, so the note arrives split across + # lines; the reader sees one sentence and the test should too. + flattened = " ".join(result.output.split()) + assert "unverified deliveries" in flattened + assert "doctor.open.on_hook" in flattened + assert "behind an internal load balancer" in flattened diff --git a/tests/units/workflow/test_versioning.py b/tests/units/workflow/test_versioning.py index a84bbfb05bc..dabdc057202 100644 --- a/tests/units/workflow/test_versioning.py +++ b/tests/units/workflow/test_versioning.py @@ -1,6 +1,12 @@ """Tests for deploying new workflow code while runs are in flight.""" -from reflex_base.workflow import WorkflowConfig, manual, needs_attention +from reflex_base.workflow import ( + Signal, + WorkflowConfig, + manual, + needs_attention, + wait_for, +) import reflex as rx from reflex.workflow.records import RunStatus, StepStatus @@ -284,3 +290,198 @@ async def finish(self, ticket: str = "unset"): assert snapshot is not None assert snapshot.status is RunStatus.COMPLETED assert snapshot.result == "unset" + + +async def test_a_waiting_run_whose_continuation_is_gone_suspends( + forked_registration_context, +): + """A wait reaches its handler by a different path, and it is gated too. + + The compatibility gate is easy to reason about for a plain successor: the + slot names a handler and the claim checks it. A wait's continuation is + reached with an injected payload, and a join's with injected results, so + "the gate applies here as well" is a separate fact rather than the same + one. A worker that crashed on a resolved wait instead of suspending the + run would take out the process for every other run it was serving. + + Args: + forked_registration_context: Isolates workflow registration. + """ + store = MemoryRunStore() + + class Reviewed(rx.State): + __workflow__ = WorkflowConfig(id="versioning.reviewed") + + review = Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Wait for a review. + + Returns: + The wait. + """ + return wait_for( + Reviewed.review, + then=Reviewed.decide, + timeout="3d", + on_timeout=Reviewed.expire, + ) + + @rx.event(durable=True, effect="none") + def decide(self, decision: dict): + """Record the decision. + + Args: + decision: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result=decision) + + @rx.event(durable=True, effect="none") + def expire(self): + """Give up. + + Returns: + Failure. + """ + return rx.fail("no_decision") + + async with WorkflowTestHarness(Reviewed, store=store) as harness: + started = await harness.start(Reviewed.begin()) + assert started.run_id is not None + run_id, resume_at = started.run_id, harness.now + + class Truncated(rx.State): + __workflow__ = WorkflowConfig(id="versioning.reviewed") + + review = Signal(dict) + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Wait for a review. + + Returns: + The wait. + """ + return wait_for( + Truncated.review, + then=Truncated.expire, + timeout="3d", + on_timeout=Truncated.expire, + ) + + @rx.event(durable=True, effect="none") + def expire(self): + """Give up. + + Returns: + Failure. + """ + return rx.fail("no_decision") + + async with WorkflowTestHarness( + Truncated, store=store, start_time=resume_at + 60 + ) as harness: + assert ( + await harness.signal(run_id, Truncated.review({"ok": True})) == "resolved" + ) + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION, snapshot.status + assert snapshot.error is not None + assert snapshot.error["reason"] == "unknown_handler" + assert "decide" in snapshot.error["detail"] + + +async def test_a_join_whose_continuation_is_gone_suspends( + forked_registration_context, +): + """The same fact for a fan-out's results slot. + + The branch soaks so the join is still open across the redeploy; a branch + that finished immediately would resolve the join under the old code and + prove nothing about the new. + + Args: + forked_registration_context: Isolates workflow registration. + """ + store = MemoryRunStore() + + class Branch(rx.State): + __workflow__ = WorkflowConfig(id="versioning.branch") + + @rx.event(durable=True, trigger=manual(), effect="none") + def go(self): + """Soak, then finish. + + Returns: + A deferral. + """ + return rx.after("1h", Branch.done) + + @rx.event(durable=True, effect="none") + def done(self): + """Finish. + + Returns: + Completion. + """ + return rx.complete(result={"done": True}) + + class Fanned(rx.State): + __workflow__ = WorkflowConfig(id="versioning.fanned") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Fan out to one branch. + + Returns: + The fan-out. + """ + return rx.parallel(Branch.go(), then=Fanned.gather) + + @rx.event(durable=True, effect="none") + def gather(self, results: list): + """Collect the branch results. + + Args: + results: One entry per branch. + + Returns: + Completion. + """ + return rx.complete(result={"branches": len(results)}) + + async with WorkflowTestHarness(Fanned, Branch, store=store) as harness: + started = await harness.start(Fanned.begin()) + assert started.run_id is not None + run_id, resume_at = started.run_id, harness.now + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING, "the join is open on the branch" + + class Truncated(rx.State): + __workflow__ = WorkflowConfig(id="versioning.fanned") + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Fan out to one branch. + + Returns: + The fan-out. + """ + return rx.parallel(Branch.go(), then=Truncated.begin) + + async with WorkflowTestHarness( + Truncated, Branch, store=store, start_time=resume_at + ) as harness: + await harness.advance("2h") + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION, snapshot.status + assert snapshot.error is not None + assert snapshot.error["reason"] == "unknown_handler" + assert "gather" in snapshot.error["detail"] From fab1232038e001cc6a2f06351e421d955281420a Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 11:33:09 -0700 Subject: [PATCH 109/121] workflows: stop reflex deploy misdirecting a workflow-only project I had been reporting reflex deploy as blocked in "the hosting repo". That was wrong and I should have checked earlier: the deploy command is in reflex/reflex.py and the hosting CLI is vendored at packages/reflex-hosting-cli. Reading them settles what is actually external and what is not. What is not: the first thing a workflow-only project hits is assert_in_reflex_dir(), failing with "rxconfig.py not found. Move to the root folder of your project, or run reflex init to start a new project." For someone who just ran `reflex workflows init` -- which writes one module and no rxconfig.py on purpose, because there is no frontend to configure -- that instruction means scaffolding the web app they deliberately did not ask for. Deploy now recognises the case, names the modules holding workflows, says why it cannot proceed, and gives the command that does work. What is external, and stays external: hosting zips a backend and then unconditionally zips a frontend (cli.py, the two export_fn calls), and the service has no notion of an app that is workers plus a headless ingress. Making that work is a change to a running service I cannot reach or verify against from here, so this refuses honestly rather than half-implementing a path that would fail further in. The guard reads the .py files in the working directory, so it tolerates one that is not decodable text -- a check that crashed deploy before it started would be worse than the message it replaces. --- news/workflow-deploy-guard.bugfix.md | 1 + reflex/reflex.py | 38 ++++++++++ tests/units/workflow/test_deploy_guard.py | 89 +++++++++++++++++++++++ 3 files changed, 128 insertions(+) create mode 100644 news/workflow-deploy-guard.bugfix.md create mode 100644 tests/units/workflow/test_deploy_guard.py diff --git a/news/workflow-deploy-guard.bugfix.md b/news/workflow-deploy-guard.bugfix.md new file mode 100644 index 00000000000..4372b6313e3 --- /dev/null +++ b/news/workflow-deploy-guard.bugfix.md @@ -0,0 +1 @@ +`reflex deploy` no longer misdirects a workflow-only project. `reflex workflows init` writes one module and deliberately no `rxconfig.py` — there is no frontend to configure — so deploy hit its generic "rxconfig.py not found, run `reflex init` to start a new project" message, which told that reader to scaffold the web app they specifically did not want. A directory holding workflows but no Reflex app is now refused by name, with the reason (hosting cannot deploy a project with no frontend yet) and the command that does work: run `reflex workflows worker` against a database you control. diff --git a/reflex/reflex.py b/reflex/reflex.py index 6295b207711..623ff43cdbf 100644 --- a/reflex/reflex.py +++ b/reflex/reflex.py @@ -997,6 +997,7 @@ def deploy( if interactive: dependency.check_requirements() + _refuse_workflow_only_deploy() prerequisites.assert_in_reflex_dir() # Check if we are set up. @@ -1044,6 +1045,43 @@ def deploy( ) +def _refuse_workflow_only_deploy() -> None: + """Say something true when a workflow-only project reaches deploy. + + ``reflex workflows init`` writes one module and no ``rxconfig.py`` on + purpose -- there is no frontend to configure. Deploy's generic complaint + for a missing config tells the reader to run ``reflex init`` and start a + new project, which for this reader means scaffolding the web app they + deliberately did not ask for. Hosting cannot yet deploy a project with no + frontend, and saying so is better than sending someone down that path. + + Raises: + Exit: When the directory holds workflows but no Reflex app. + """ + from reflex.constants import Config + + if Path(Config.FILE).exists(): + return + holding = [ + candidate.name + for candidate in sorted(Path.cwd().glob("*.py")) + if "__workflow__" in candidate.read_text(errors="ignore") + ] + if not holding: + return + console.error( + f"{', '.join(holding)} define workflows, but this is not a Reflex app " + f"({Config.FILE} is absent) and hosting cannot deploy a project with " + "no frontend yet. Run the workers yourself against your own " + "infrastructure:\n" + " reflex workflows worker \n" + "pointing REFLEX_WORKFLOW_DATABASE at a Postgres you control. To " + "deploy this as a full Reflex app instead, add an app and an " + f"{Config.FILE}." + ) + raise click.exceptions.Exit(1) + + @cli.command() @loglevel_option @click.argument("new_name") diff --git a/tests/units/workflow/test_deploy_guard.py b/tests/units/workflow/test_deploy_guard.py new file mode 100644 index 00000000000..e919d967155 --- /dev/null +++ b/tests/units/workflow/test_deploy_guard.py @@ -0,0 +1,89 @@ +"""Deploy has to say something true to a workflow-only project. + +``reflex workflows init`` writes one module and deliberately no +``rxconfig.py`` -- there is no frontend to configure, which is the whole +point of the workflow-only path. Deploy's generic complaint about a missing +config then tells that reader to run ``reflex init`` and start a new project, +which for them means scaffolding the web app they specifically did not want. +""" + +from pathlib import Path + +import pytest +from click.exceptions import Exit + +from reflex.reflex import _refuse_workflow_only_deploy + +WORKFLOW_MODULE = """ +import reflex as rx +from reflex_base.workflow import WorkflowConfig, manual + + +class Flow(rx.State): + __workflow__ = WorkflowConfig(id="deployguard.flow") + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self): + return rx.complete(result=None) +""" + + +def test_a_workflow_only_project_is_refused_with_something_actionable( + tmp_path, monkeypatch, capsys +): + """Name the module, say why, and give the command that does work. + + Args: + tmp_path: The project directory. + monkeypatch: Used to enter that directory. + capsys: Captures the console output. + """ + (tmp_path / "workflows.py").write_text(WORKFLOW_MODULE) + monkeypatch.chdir(tmp_path) + with pytest.raises(Exit): + _refuse_workflow_only_deploy() + captured = capsys.readouterr() + # console.error goes to stderr; Rich also wraps, so flatten both. + output = " ".join((captured.out + captured.err).split()) + assert "workflows.py" in output + assert "reflex workflows worker" in output + assert "reflex init" not in output, ( + "telling a workflow user to start a new project is the bug being fixed" + ) + + +def test_a_reflex_app_is_left_alone(tmp_path, monkeypatch): + """A real app deploys as it always did, workflows or not. + + Args: + tmp_path: The project directory. + monkeypatch: Used to enter that directory. + """ + (tmp_path / "rxconfig.py").write_text("import reflex as rx\n") + (tmp_path / "workflows.py").write_text(WORKFLOW_MODULE) + monkeypatch.chdir(tmp_path) + _refuse_workflow_only_deploy() + + +def test_a_directory_with_no_workflows_is_left_alone(tmp_path, monkeypatch): + """Someone in the wrong directory still gets the ordinary message. + + Args: + tmp_path: The empty directory. + monkeypatch: Used to enter it. + """ + (tmp_path / "notes.py").write_text("x = 1\n") + monkeypatch.chdir(tmp_path) + _refuse_workflow_only_deploy() + + +def test_an_unreadable_file_does_not_break_the_check(tmp_path, monkeypatch): + """A binary or badly encoded .py must not crash deploy before it starts. + + Args: + tmp_path: The project directory. + monkeypatch: Used to enter it. + """ + Path(tmp_path / "broken.py").write_bytes(b"\xff\xfe\x00binary") + monkeypatch.chdir(tmp_path) + _refuse_workflow_only_deploy() From c9d0948b1784f2bc35dc68af5f076a36188a65df Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 12:54:21 -0700 Subject: [PATCH 110/121] workflows: three Postgres blockers from the review All reproduced first, all fixed, all verified against a real server. 1. Deadline signals lied on Postgres. deliver() selected only status, never deadline, so a delivery to a past-deadline run answered "resolved" where Memory and SQLite answered "expired". The run then finalized TIMED_OUT and dropped the payload -- a person clicking approve was told their decision landed, moments before it was discarded. The conformance suite had no check for it, which is why three stores diverged quietly; it has one now and it failed on Postgres alone before the fix. 2. Cancelling a fan-out deadlocked against its branches. Closing a parent takes the parent row then its children; a branch reporting home takes itself then its parent. Same two rows, opposite orders. A probe measured 130 aborted transactions across 40 rounds. Fixed by ordering, not by retrying: both transactions now take branch rows first, so everything acquires children before self before parent -- deeper before shallower, which cannot cycle. 0 in 150 rounds after. Retrying would have left the latency spike and, where recovery had to converge, a cancelled rollout still running until a lease lapsed, which is the outcome cascade cancellation exists to prevent. 3. Store time moved when the wall clock moved. The offset was added to time.time() on every read, so an NTP step or a resumed snapshot carried the worker with it; it would renew its lease to a moment the store considered past, lose the claim mid-attempt, and let a peer reclaim the step. Time is now carried from the last sync by time.monotonic, which cannot jump. Only real drift remains, and the next recovery corrects it. The test for 3 initially passed while measuring nothing: the kernel decides whether to sync by identity against time.time, and monkeypatching it replaces the object the default argument was bound to, so the sync path silently turned off. It now passes the clock explicitly and asserts the path is live before testing it. --- news/workflow-postgres-blockers.bugfix.md | 7 + reflex/workflow/conformance.py | 17 +++ reflex/workflow/kernel.py | 28 +++- reflex/workflow/postgres.py | 37 +++++- tests/units/workflow/test_clock_authority.py | 72 +++++++++++ tests/units/workflow/test_postgres.py | 127 +++++++++++++++++++ 6 files changed, 280 insertions(+), 8 deletions(-) create mode 100644 news/workflow-postgres-blockers.bugfix.md diff --git a/news/workflow-postgres-blockers.bugfix.md b/news/workflow-postgres-blockers.bugfix.md new file mode 100644 index 00000000000..0405feea342 --- /dev/null +++ b/news/workflow-postgres-blockers.bugfix.md @@ -0,0 +1,7 @@ +Three Postgres-only defects found by an external review, all reproduced before fixing. + +Delivering a signal to a run past its deadline answered `resolved` on Postgres while Memory and SQLite answered `expired`. The run then finalized `TIMED_OUT` and discarded the payload, so a person clicking approve was told their decision landed moments before it was thrown away. `deliver()` never read the deadline column; it does now, and a conformance check covers the case on all three stores. + +Cancelling a fan-out deadlocked against its own branches. Closing a parent takes the parent row and then its children; a branch reporting home takes itself and then its parent — the same two rows in opposite orders, which Postgres detects and aborts. A probe saw 130 aborts across 40 rounds. Both transactions now take branch rows first, so every transaction acquires children before self before parent: deeper rows always before shallower ones, which cannot cycle. The same probe now sees none in 150 rounds. This is lock ordering, not deadlock retries — a retry would have left the latency spike and, when recovery had to converge, a cancelled rollout still running until a lease lapsed. + +A worker derived store time by adding a fixed offset to its wall clock, so an NTP step, a resumed snapshot, or a corrected host moved its clock too. It would then renew its lease to a moment the store considered past, its claim would lapse mid-attempt, and a peer would reclaim the step — two workers on one attempt. Store time is now carried forward from the last sync by `time.monotonic`, which cannot jump, leaving only real drift for the next recovery pass to correct. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 16f11c679b5..e20b1516b8d 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -1382,6 +1382,22 @@ async def check_a_duplicate_delivery_is_recorded_in_history(store: RunStore) -> assert kinds.count(HistoryEventType.SIGNAL_DUPLICATE) == 1, kinds +async def check_a_delivery_to_a_past_deadline_run_is_refused(store: RunStore) -> None: + """A run that can never execute the continuation must not say "resolved". + + Claims exclude past-deadline runs and the sweep is about to finalize this + one TIMED_OUT, so answering "resolved" tells the sender their decision + landed when it is about to be discarded -- and the sender is often a + person clicking approve. + """ + await store.admit( + make_run(deadline=NOW - 1), + make_step(status=StepStatus.BLOCKED, wait_key="sig:ping", due_at=0.0), + _ADMITTED, + ) + assert await store.deliver("run1", "sig:ping", "d1", {"v": 1}, NOW) == "expired" + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1399,6 +1415,7 @@ async def check_a_duplicate_delivery_is_recorded_in_history(store: RunStore) -> check_delivery_resolves_a_matching_wait, check_delivery_never_touches_run_state, check_duplicate_deliveries_are_ignored, + check_a_delivery_to_a_past_deadline_run_is_refused, check_a_duplicate_delivery_is_recorded_in_history, check_an_early_delivery_is_buffered_then_consumed, check_early_deliveries_queue_in_order, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index a80b2c4184b..028c216071d 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -459,7 +459,7 @@ def __init__( # skew is bounded by one round trip plus local drift per recovery # interval. An explicitly injected clock (tests, the dev CLI's # fast-forward) stays authoritative as given and is never synced. - self._store_clock_offset = 0.0 + self._clock_anchor: tuple[float, float] | None = None self._sync_clock_with_store = clock is time.time self._clock = self._store_time if self._sync_clock_with_store else clock self._rng = rng @@ -2848,12 +2848,24 @@ async def _cancel_inflight(self) -> None: self._prune() def _store_time(self) -> float: - """The process clock corrected onto the store's time base. + """The store's clock, carried forward by monotonic elapsed time. + + The wall clock is read once per sync and never between them. NTP + steps, a resumed snapshot, or an operator correcting a drifted host + all move ``time.time`` without warning, and a worker that added a + fixed offset to it moved with it -- renewing its lease to a moment + the store considers past, so its claim lapsed mid-attempt and a peer + reclaimed the step. ``time.monotonic`` cannot jump, so the only error + left is real drift since the last sync, which the next recovery pass + corrects. Returns: Epoch seconds by the store's clock, to within one sync error. """ - return time.time() + self._store_clock_offset + if self._clock_anchor is None: + return time.time() + store_at_sync, monotonic_at_sync = self._clock_anchor + return store_at_sync + (time.monotonic() - monotonic_at_sync) async def _sync_store_clock(self) -> None: """Re-measure the offset between this process and the store's clock. @@ -2863,15 +2875,17 @@ async def _sync_store_clock(self) -> None: """ if not self._sync_clock_with_store: return - before = time.time() + before = time.monotonic() store_now = await self._store.epoch_time() - after = time.time() + after = time.monotonic() if store_now is None: # The process clock is the authority for this store; stop asking. self._sync_clock_with_store = False - self._store_clock_offset = 0.0 + self._clock_anchor = None return - self._store_clock_offset = store_now - (before + after) / 2 + # The store answered somewhere inside the round trip, so credit it to + # the midpoint: the anchor is then off by at most half of one. + self._clock_anchor = (store_now, (before + after) / 2) async def recover(self) -> int: """Renew this kernel's live claims, then reclaim expired ones. diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 3d39e6317f9..90caf539ddc 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1003,6 +1003,8 @@ async def commit( """ pool = await self._open() async with pool.connection() as conn, conn.transaction(): + if completion.run_status in TERMINAL_RUN_STATUSES: + await self._lock_children(conn, claim.run.run_id) await self._lock_run(conn, claim.run.run_id) deadline = await self._check_claim(conn, claim) # Past the deadline the only permitted outcome is TIMED_OUT, and @@ -1160,7 +1162,8 @@ async def deliver( pool = await self._open() async with pool.connection() as conn, conn.transaction(): cursor = await conn.execute( - "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", + "SELECT status, deadline FROM workflow_runs WHERE run_id = %s" + " FOR UPDATE", (run_id,), ) row = await cursor.fetchone() @@ -1168,6 +1171,12 @@ async def deliver( return "unknown_run" if row["status"] in _TERMINAL_RUNS: return "run_terminal" + if row["deadline"] is not None and row["deadline"] <= now: + # Claims exclude past-deadline runs and the sweep is about to + # finalize this one TIMED_OUT, so "resolved" would tell the + # sender -- often a person clicking approve -- that their + # decision landed, moments before it is discarded. + return "expired" cursor = await conn.execute( "SELECT 1 FROM workflow_inbox" " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", @@ -1262,6 +1271,31 @@ async def admit_children( await self._lock_run(conn, parent) await self._append_events(conn, parent, events, now) + async def _lock_children(self, conn: Any, run_id: str) -> None: + """Take the branch rows this transaction will close, before its own. + + Closing a parent locks the parent and then its children; a child + reporting home locks itself and then its parent. Those are the same + two rows in opposite orders, which is an ABBA deadlock -- Postgres + detects it and aborts one side, and under a fan-out being cancelled + that happened constantly. + + Taking the branches first makes every transaction acquire in one + order: children, then self, then parent. Deeper rows are always taken + before shallower ones, so no cycle can form. The ORDER BY matters for + the same reason within one level. + + Args: + conn: The connection inside an open transaction. + run_id: The run whose branches may be closed. + """ + await conn.execute( + "SELECT run_id FROM workflow_runs WHERE parent_run_id = %s" + " AND parent_close <> 'abandon' AND NOT (status = ANY(%s))" + " ORDER BY run_id FOR UPDATE", + (run_id, [s.value for s in TERMINAL_RUN_STATUSES]), + ) + async def _close_children(self, conn: Any, run_id: str, now: float) -> None: """Request cancellation of branches the closing run fanned out to. @@ -1669,6 +1703,7 @@ async def finalize_run( """ pool = await self._open() async with pool.connection() as conn, conn.transaction(): + await self._lock_children(conn, run_id) cursor = await conn.execute( "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", (run_id,), diff --git a/tests/units/workflow/test_clock_authority.py b/tests/units/workflow/test_clock_authority.py index 2aaaee1d5b8..c7c6e7b01c9 100644 --- a/tests/units/workflow/test_clock_authority.py +++ b/tests/units/workflow/test_clock_authority.py @@ -11,6 +11,7 @@ and never synced. """ +import asyncio import time import pytest @@ -121,3 +122,74 @@ async def test_a_store_with_no_clock_of_its_own_stops_being_asked(): await kernel.recover() assert kernel._sync_clock_with_store is False # pyright: ignore[reportPrivateUsage] assert kernel._clock() == pytest.approx(time.time(), abs=1.0) # pyright: ignore[reportPrivateUsage] + + +class _IndependentClockStore(MemoryRunStore): + """A store whose clock is its own, like a database on another machine. + + Deriving it from ``time.time`` would make it jump whenever the worker's + wall clock jumps, which is exactly the thing under test. + """ + + def __init__(self): + """Anchor the store's clock to real elapsed time.""" + super().__init__() + self._base = 1_700_000_000.0 + self._start = time.monotonic() + + async def epoch_time(self) -> float | None: + """Answer with the store's own clock. + + Returns: + Epoch seconds, unaffected by the caller's wall clock. + """ + return self._base + (time.monotonic() - self._start) + + +async def test_a_backward_wall_clock_jump_does_not_move_store_time(monkeypatch): + """Time must not go backwards between syncs, whatever the machine does. + + NTP steps, a hypervisor resuming a snapshot, an operator correcting a + drifted host: all move the wall clock without warning. A worker that + derives store time by adding a fixed offset to that clock moves with it, + and renews its lease to a moment that has, from the store's point of + view, already passed. Its claim lapses while it is still working and a + peer reclaims the step -- two workers on one attempt, which is the one + thing leases exist to prevent. + + Args: + monkeypatch: Used to move the process wall clock. + """ + wall = [1_600_000_000.0] + monkeypatch.setattr(time, "time", lambda: wall[0]) + + # Passed explicitly because the kernel decides whether to sync by identity + # against time.time, and monkeypatching replaces the object the default + # argument was bound to. Without this the sync path silently turns off and + # the test passes while measuring nothing. + kernel = WorkflowKernel([], _IndependentClockStore(), clock=time.time) + assert kernel._sync_clock_with_store is True # pyright: ignore[reportPrivateUsage] + await kernel.recover() + before = kernel._clock() # pyright: ignore[reportPrivateUsage] + + wall[0] -= 45.0 + after = kernel._clock() # pyright: ignore[reportPrivateUsage] + assert after >= before, ( + f"store time went backwards by {before - after:.1f}s when the wall " + "clock did; a lease renewed against it would lapse early" + ) + + +async def test_store_time_still_advances_with_real_elapsed_time(): + """Immunity to jumps must not mean the clock stops. + + A clock that never moved would be just as wrong: leases would never + expire and a crashed worker's step would never be recovered. + """ + kernel = WorkflowKernel([], _IndependentClockStore()) + await kernel.recover() + first = kernel._clock() # pyright: ignore[reportPrivateUsage] + await asyncio.sleep(0.05) + second = kernel._clock() # pyright: ignore[reportPrivateUsage] + assert second > first + assert second - first < 5.0, "advancing far faster than real time is also wrong" diff --git a/tests/units/workflow/test_postgres.py b/tests/units/workflow/test_postgres.py index a79872cbb6d..bff0cada5a2 100644 --- a/tests/units/workflow/test_postgres.py +++ b/tests/units/workflow/test_postgres.py @@ -296,3 +296,130 @@ async def seed_and_close(): cancelled = CliRunner().invoke(workflows, ["cancel", "cli-run", "-d", url]) assert cancelled.exit_code == 0, cancelled.output + + +async def test_closing_a_parent_never_deadlocks_against_its_children(store): + """Cancel a fan-out while its branches report home, repeatedly. + + Closing a parent touches the parent and then its children; a child + reporting home touches itself and then its parent. Same two rows, + opposite orders -- Postgres detects the cycle and aborts one side, and + before the lock ordering was fixed this produced deadlocks on most + rounds. Recovery did eventually converge, but sometimes only after a + lease lapsed, which is half a minute of a cancelled rollout still + running. + + Args: + store: The Postgres store. + """ + rounds, branches = 12, 6 + failures: list[str] = [] + for index in range(rounds): + parent = f"dlp{index}" + kids = [f"dlk{index}_{n}" for n in range(branches)] + await store.admit( + _pg_run(parent, next_ordinal=2), + _pg_step(parent, 1, status=StepStatus.BLOCKED, wait_key="join:1"), + _PG_ADMITTED, + ) + for kid in kids: + await store.admit( + _pg_run(kid, parent_run_id=parent, parent_ordinal=1), + _pg_step(kid), + _PG_ADMITTED, + ) + + async def close_parent(parent=parent): + """Cancel the parent, closing its branches. + + Args: + parent: The parent run. + """ + await store.request_cancel(parent, _PG_NOW) + await store.finalize_run( + parent, + status=RunStatus.CANCELLED, + error=None, + event=HistoryEventType.RUN_CANCELLED, + now=_PG_NOW, + ) + + async def report(kid: str, parent=parent): + """Finish a branch, delivering its arrival. + + Args: + kid: The branch run. + parent: The parent run. + """ + await store.finalize_run( + kid, + status=RunStatus.COMPLETED, + error=None, + event=HistoryEventType.RUN_COMPLETED, + now=_PG_NOW, + parent_arrival=(parent, 1, {"status": "completed"}, kid), + ) + + outcomes = await asyncio.gather( + close_parent(), *(report(kid) for kid in kids), return_exceptions=True + ) + failures.extend( + type(outcome).__name__ + for outcome in outcomes + if isinstance(outcome, BaseException) + ) + assert not failures, f"{len(failures)} transaction(s) aborted: {set(failures)}" + + +_PG_NOW = 1_000_000.0 +_PG_ADMITTED = ((HistoryEventType.RUN_ADMITTED, {}),) + + +def _pg_run(run_id: str, **over) -> RunRecord: + """Build a run record for the deadlock probe. + + Args: + run_id: The run identity. + over: Field overrides. + + Returns: + The record. + """ + fields: dict = { + "run_id": run_id, + "workflow_id": "pg.deadlock", + "definition_digest": "d", + "status": RunStatus.PENDING, + "state": {}, + "state_version": 0, + "next_ordinal": 2, + "created_at": _PG_NOW, + "updated_at": _PG_NOW, + } + fields.update(over) + return RunRecord(**fields) + + +def _pg_step(run_id: str, ordinal: int = 0, **over) -> StepRecord: + """Build a step record for the deadlock probe. + + Args: + run_id: The owning run. + ordinal: The mailbox position. + over: Field overrides. + + Returns: + The record. + """ + fields: dict = { + "run_id": run_id, + "ordinal": ordinal, + "handler_id": "go", + "status": StepStatus.READY, + "args": {}, + "origin": "root", + "created_at": _PG_NOW, + "updated_at": _PG_NOW, + } + fields.update(over) + return StepRecord(**fields) From 439c4259f8f47f84577c9a6dadbe8cf9ce7d9efe Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 13:02:31 -0700 Subject: [PATCH 111/121] workflows: the rest of the review's findings Five smaller defects, all reproduced before fixing. force_complete took the operator's result unchecked, so Memory stored a live Decimal that no other store could hold and SQLite raised a bare "not JSON serializable" from inside json.dumps. Same input, three behaviours, none of them actionable -- and the Memory case only surfaced on the day someone moved to Postgres. It now goes through the same strict serde as a handler's result and all three refuse identically. CLI prefix resolution scanned the newest 10,000 runs and treated a single hit as unique. On a bigger store another run just outside that window could share the prefix, and the next thing the resolver feeds is cancel or complete. There is no prefix query in the store protocol to do better, so it now refuses what it cannot prove and says to pass the full id. Retry.multiplier accepted nan and inf, because every comparison against nan is False and "< 1.0" was the only guard. The backoff then scheduled a step for a moment that never arrives. Cron could not see across a skipped leap century: 2100 is not a leap year, so 2096 to 2104 is eight years and the 1500-day horizon reported no occurrence for a perfectly good expression. Separately, 0 0 30 2 * parsed happily and then never fired, which is indistinguishable from a schedule that is merely waiting -- now a definition error. A weekday restriction still makes it legal, because cron matches day-of-month OR day-of-week and the date has a second path. Occurrences dropped past the catch-up cap now increment a counter on both exporters. They have no run to attach history to, so a log line was the only trace, and noticing that a nightly job silently stopped a week ago is exactly what a counter is for. --- news/workflow-review-fixes.bugfix.md | 11 ++ .../reflex-base/src/reflex_base/workflow.py | 7 +- reflex/workflow/cli.py | 18 +- reflex/workflow/cron.py | 34 +++- reflex/workflow/kernel.py | 56 ++++++ reflex/workflow/otel.py | 16 +- tests/units/workflow/test_review_fixes.py | 163 ++++++++++++++++++ 7 files changed, 295 insertions(+), 10 deletions(-) create mode 100644 news/workflow-review-fixes.bugfix.md create mode 100644 tests/units/workflow/test_review_fixes.py diff --git a/news/workflow-review-fixes.bugfix.md b/news/workflow-review-fixes.bugfix.md new file mode 100644 index 00000000000..90454e03248 --- /dev/null +++ b/news/workflow-review-fixes.bugfix.md @@ -0,0 +1,11 @@ +Fixed the remaining defects from an external review, each reproduced first. + +An operator-supplied result now goes through the same strict serialization as a handler's. Memory kept a live `Decimal` no other store could hold, while SQLite raised a bare "not JSON serializable" from inside `json.dumps` — one input, three behaviours, none of them saying what to do. All three now refuse it identically, naming the fix. + +CLI run-id prefixes are refused rather than guessed when the database holds more runs than a single scan covers. The resolver reads the newest 10,000 runs, so on a larger store a prefix could look unique while another run just outside the window shared it — and `cancel` or `complete` would then act on the wrong run. + +`Retry(multiplier=...)` refuses `nan` and `inf`. Every comparison against `nan` is False, so the `< 1.0` guard let both through, and the backoff they produced scheduled a step for a moment that never arrives. + +The cron search horizon now spans a skipped leap century: 2100 is not a leap year, so the gap from 2096 to 2104 is eight years and `0 0 29 2 *` reported no occurrence at all. A month and day-of-month pairing that no year can satisfy (`0 0 30 2 *`) is now a definition error instead of a schedule that parses and silently never fires — unless a weekday restriction gives the date a second way to match, which cron's OR rule allows. + +Scheduled occurrences dropped past the catch-up cap are counted on `MetricsObserver` and the OpenTelemetry observer. They have no run to carry history, so a log line was their only trace, and "the nightly job silently stopped a week ago" is exactly what a counter is for. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 0b94fdcb4b4..36ce1345871 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -132,8 +132,11 @@ def __post_init__(self): if self.max_attempts < 1: msg = f"Retry.max_attempts must be >= 1, got {self.max_attempts}." raise WorkflowDefinitionError(msg) - if self.multiplier < 1.0: - msg = f"Retry.multiplier must be >= 1.0, got {self.multiplier}." + if not math.isfinite(self.multiplier) or self.multiplier < 1.0: + # Every comparison against nan is False, so "< 1.0" let nan and + # inf straight through, and the backoff they produce is nan or + # inf -- a step scheduled for a time that never arrives. + msg = f"Retry.multiplier must be a finite number >= 1.0, got {self.multiplier}." raise WorkflowDefinitionError(msg) if self.jitter not in ("full", "none"): msg = f'Retry.jitter must be "full" or "none", got {self.jitter!r}.' diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index f8c22b96626..66a153df433 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -124,11 +124,19 @@ async def _resolve_run_id(store: RunStore, run_id: str) -> str: """ if await store.get_run(run_id) is not None: return run_id - matches = [ - run.run_id - for run in await store.list_runs(RunQuery(limit=_PREFIX_SCAN_LIMIT)) - if run.run_id.startswith(run_id) - ] + scanned = await store.list_runs(RunQuery(limit=_PREFIX_SCAN_LIMIT)) + matches = [run.run_id for run in scanned if run.run_id.startswith(run_id)] + if len(scanned) >= _PREFIX_SCAN_LIMIT: + # The newest N runs, not all of them: a prefix that looks unique here + # may match another run just outside the window, and resolving it + # would cancel or complete the wrong one. There is no prefix query in + # the store protocol to do better, so say what is true. + console.error( + f"This database holds more than {_PREFIX_SCAN_LIMIT:,} runs, so " + f"{run_id!r} cannot be shown to match only one. Pass the full run " + "id (reflex workflows list prints them), or purge finished runs." + ) + raise click.exceptions.Exit(1) if len(matches) == 1: return matches[0] if not matches: diff --git a/reflex/workflow/cron.py b/reflex/workflow/cron.py index d2236069c98..3808d66debb 100644 --- a/reflex/workflow/cron.py +++ b/reflex/workflow/cron.py @@ -21,7 +21,12 @@ _FIELD_NAMES: Final = ("minute", "hour", "day of month", "month", "day of week") -MAX_SEARCH_DAYS: Final = 1500 +# The rarest satisfiable date is February 29 across a skipped leap century: +# 2096 to 2104 is eight years, because 2100 is not a leap year. A shorter +# horizon reported "no occurrence" for a schedule that was simply far off. +MAX_SEARCH_DAYS: Final = 3000 + +_LONGEST_MONTH: Final = (0, 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31) def _parse_field(spec: str, index: int) -> frozenset[int]: @@ -115,6 +120,7 @@ def __init__(self, expression: str): self._days_of_week = _parse_field(fields[4], 4) self._dom_restricted = fields[2] != "*" self._dow_restricted = fields[4] != "*" + self._assert_reachable() def _matches_date(self, day: dt.date) -> bool: """Whether a date satisfies the month and day fields. @@ -135,6 +141,32 @@ def _matches_date(self, day: dt.date) -> bool: return dom_hit or dow_hit return dom_hit and dow_hit + def _assert_reachable(self) -> None: + """Refuse a month and day-of-month pairing that no year can satisfy. + + ``0 0 30 2 *`` parses -- every field is in range -- and then never + fires, which looks exactly like a schedule that is merely waiting. + A weekday restriction gives the date a second way to match, so this + only applies when day-of-month is the only selector. + + Raises: + WorkflowDefinitionError: If no date can ever match. + """ + if self._dow_restricted or not self._dom_restricted: + return + if any( + day <= _LONGEST_MONTH[month] + for month in self._months + for day in self._days_of_month + ): + return + msg = ( + f"Cron expression {self.expression!r} can never occur: no day in " + f"{sorted(self._days_of_month)} exists in " + f"{sorted(self._months)}." + ) + raise WorkflowDefinitionError(msg) + def next_after(self, after: float) -> float | None: """Find the first occurrence strictly after a point in time. diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 028c216071d..0065bb9ddce 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -144,6 +144,19 @@ def on_event( data: Event payload, such as ordinal, handler, attempt, or error. """ + def on_schedule_skip(self, schedule_key: str, skipped: int) -> None: + """Handle scheduled occurrences that were dropped rather than run. + + These have no run to carry history, so without this the only trace + of dropped work is a log line. A counter survives the process and can + be alerted on, which is what "we silently stopped doing the nightly + job for a week" needs. + + Args: + schedule_key: The schedule that lost occurrences. + skipped: How many were dropped. + """ + class CompositeObserver(WorkflowObserver): """Fans one transition out to several observers. @@ -188,6 +201,17 @@ def on_event( with contextlib.suppress(Exception): observer.on_event(event_type, run_id, workflow_id, data) + def on_schedule_skip(self, schedule_key: str, skipped: int) -> None: + """Pass dropped occurrences to every observer. + + Args: + schedule_key: The schedule that lost occurrences. + skipped: How many were dropped. + """ + for observer in self.observers: + with contextlib.suppress(Exception): + observer.on_schedule_skip(schedule_key, skipped) + def _branch_index(request_key: str | None) -> int | None: """Recover which branch of a fan-out a child run was. @@ -246,6 +270,21 @@ def __init__(self): self.totals: dict[str, int] = {} self.by_workflow: dict[str, dict[str, int]] = {} + def on_schedule_skip(self, schedule_key: str, skipped: int) -> None: + """Count scheduled occurrences that were dropped rather than run. + + Args: + schedule_key: The schedule that lost occurrences. + skipped: How many were dropped. + """ + self.totals["schedule_occurrences_skipped"] = ( + self.totals.get("schedule_occurrences_skipped", 0) + skipped + ) + by_key = self.by_workflow.setdefault(schedule_key, {}) + by_key["schedule_occurrences_skipped"] = ( + by_key.get("schedule_occurrences_skipped", 0) + skipped + ) + def on_event( self, event_type: HistoryEventType, @@ -1034,6 +1073,12 @@ async def force_finalize( run = await self._store.get_run(run_id) if run is None: return False + # An operator's result is run data like any other. Passing it through + # unchecked let Memory keep a live Decimal that no other store could + # hold, while SQLite raised a bare "not JSON serializable" from inside + # json.dumps -- the same input, three behaviours, none of them saying + # what to do about it. + result = to_run_data({"value": result})["value"] if result is not None else None event = ( HistoryEventType.RUN_COMPLETED if status is RunStatus.COMPLETED @@ -2339,7 +2384,18 @@ async def _admit_due_schedules(self, now: float) -> int: # scheduled work reads as "covered" when it was not; the # operator gets the count and the window, and can start the # missed occurrences by hand if they matter. + dropped = ( + len(schedule.occurrences_between(cursor, now, limit=10_000)) + - MAX_SCHEDULE_CATCHUP + ) occurrences = occurrences[:MAX_SCHEDULE_CATCHUP] + # A log line is the only trace these otherwise leave, and they + # have no run to carry history. A counter survives the process + # and can be alerted on, which is what noticing "the nightly + # job silently stopped for a week" actually needs. + if self._observer is not None: + with contextlib.suppress(Exception): + self._observer.on_schedule_skip(key, max(dropped, 1)) console.warn( f"Schedule {key} missed more than {MAX_SCHEDULE_CATCHUP} " f"occurrences between {cursor:.0f} and {now:.0f}; catching " diff --git a/reflex/workflow/otel.py b/reflex/workflow/otel.py index d5725e707fd..31a3809075e 100644 --- a/reflex/workflow/otel.py +++ b/reflex/workflow/otel.py @@ -93,9 +93,10 @@ def __init__(self, *, tracer_provider: Any = None, meter_provider: Any = None): self._status_code = status_code self._tracer = (tracer_provider or trace).get_tracer(INSTRUMENTATION_NAME) meter = (meter_provider or metrics).get_meter(INSTRUMENTATION_NAME) + names = set(MetricsObserver._COUNTED.values()) # pyright: ignore[reportPrivateUsage] + names.add("schedule_occurrences_skipped") self._counters = { - name: meter.create_counter(f"reflex.workflow.{name}") - for name in set(MetricsObserver._COUNTED.values()) # pyright: ignore[reportPrivateUsage] + name: meter.create_counter(f"reflex.workflow.{name}") for name in names } self._open: dict[tuple[str, Any], Any] = {} @@ -122,6 +123,17 @@ def on_event( if ending is not None: self._close_span(run_id, data, ending) + def on_schedule_skip(self, schedule_key: str, skipped: int) -> None: + """Count scheduled occurrences that were dropped rather than run. + + Args: + schedule_key: The schedule that lost occurrences. + skipped: How many were dropped. + """ + self._counters["schedule_occurrences_skipped"].add( + skipped, {"workflow.schedule_key": schedule_key} + ) + def _count(self, event_type: HistoryEventType, workflow_id: str) -> None: """Add one to the counter this event feeds, if any. diff --git a/tests/units/workflow/test_review_fixes.py b/tests/units/workflow/test_review_fixes.py new file mode 100644 index 00000000000..11017eca4c7 --- /dev/null +++ b/tests/units/workflow/test_review_fixes.py @@ -0,0 +1,163 @@ +"""Fixes for defects an external review found by driving the engine hard. + +Each of these was reproduced before it was fixed, and each is here because +the failure was quiet: a wrong run cancelled, a schedule that never fires, a +backoff that never elapses, an operator result that one store keeps and +another rejects. +""" + +import decimal + +import pytest +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import Retry + +from reflex.workflow.cron import MAX_SEARCH_DAYS, CronSchedule +from reflex.workflow.kernel import WorkflowKernel +from reflex.workflow.records import ( + HistoryEventType, + RunRecord, + RunStatus, + StepRecord, + StepStatus, +) +from reflex.workflow.store import MemoryRunStore, SqliteRunStore + +NOW = 1_000_000.0 + + +def _run(run_id: str = "r1") -> RunRecord: + """Build a run record. + + Args: + run_id: The run identity. + + Returns: + The record. + """ + return RunRecord( + run_id=run_id, + workflow_id="review.flow", + definition_digest="d", + status=RunStatus.PENDING, + state={}, + state_version=0, + next_ordinal=1, + created_at=NOW, + updated_at=NOW, + ) + + +def _step(run_id: str = "r1") -> StepRecord: + """Build a step record that is not due yet. + + Args: + run_id: The owning run. + + Returns: + The record. + """ + return StepRecord( + run_id=run_id, + ordinal=0, + handler_id="go", + status=StepStatus.READY, + args={}, + due_at=NOW + 3600, + origin="root", + created_at=NOW, + updated_at=NOW, + ) + + +@pytest.mark.parametrize("kind", ["memory", "sqlite"]) +async def test_an_operator_result_faces_the_same_serde_everywhere(kind, tmp_path): + """One input must not mean three behaviours across the stores. + + Memory kept a live Decimal no other store could hold; SQLite raised a + bare "not JSON serializable" from inside json.dumps. Neither said what to + do about it, and the Memory case only became a problem on the day someone + moved to Postgres. + + Args: + kind: Which store to build. + tmp_path: Temporary directory for SQLite. + """ + store = MemoryRunStore() if kind == "memory" else SqliteRunStore(tmp_path / "s.db") + await store.admit(_run(), _step(), ((HistoryEventType.RUN_ADMITTED, {}),)) + kernel = WorkflowKernel([], store) + with pytest.raises(TypeError, match="Decimal is not valid run data"): + await kernel.force_finalize( + "r1", status=RunStatus.COMPLETED, result=decimal.Decimal("10.10") + ) + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf")]) +def test_a_non_finite_retry_multiplier_is_refused(bad): + """Every comparison against nan is False, so "< 1.0" let it through. + + The backoff it produced was nan or inf, and a step scheduled for a moment + that never arrives is simply lost. + + Args: + bad: The multiplier to reject. + """ + with pytest.raises(WorkflowDefinitionError, match="finite"): + Retry(max_attempts=3, multiplier=bad) + + +def test_the_cron_horizon_reaches_across_a_skipped_leap_century(): + """2100 is not a leap year, so 2096 to 2104 is eight years apart. + + A horizon shorter than that reported "no occurrence" for a schedule that + was merely far off, which reads identically to a broken expression. + """ + schedule = CronSchedule("0 0 29 2 *") + import datetime as dt + + base = dt.datetime(2096, 3, 1, tzinfo=dt.UTC).timestamp() + found = schedule.next_after(base) + assert found is not None + assert dt.datetime.fromtimestamp(found, dt.UTC).date() == dt.date(2104, 2, 29) + assert MAX_SEARCH_DAYS >= 8 * 366, "the horizon must cover the longest gap" + + +@pytest.mark.parametrize("expression", ["0 0 30 2 *", "0 0 31 4 *", "0 0 31 6 *"]) +def test_a_date_that_cannot_exist_is_refused(expression): + """A schedule that never fires looks exactly like one that is waiting. + + Args: + expression: A cron expression naming an impossible date. + """ + with pytest.raises(WorkflowDefinitionError, match="can never occur"): + CronSchedule(expression) + + +@pytest.mark.parametrize("expression", ["0 0 29 2 *", "0 0 31 1 *", "0 0 31 2 1"]) +def test_rare_but_possible_dates_are_still_accepted(expression): + """February 29 is rare, not impossible, and a weekday gives a second path. + + Under cron's day-of-month/day-of-week OR rule, ``0 0 31 2 1`` still fires + on Mondays in February, so refusing it would be wrong. + + Args: + expression: A cron expression that can occur. + """ + assert CronSchedule(expression).expression == expression + + +def test_dropped_schedule_occurrences_reach_the_metrics(): + """Work the engine decided not to do must not leave only a log line.""" + from reflex.workflow.kernel import CompositeObserver, MetricsObserver + + metrics = MetricsObserver() + CompositeObserver(metrics).on_schedule_skip("nightly", 42) + assert metrics.totals["schedule_occurrences_skipped"] == 42 + assert metrics.by_workflow["nightly"]["schedule_occurrences_skipped"] == 42 + + +def test_the_default_observer_ignores_dropped_occurrences_quietly(): + """The base class must stay a no-op so custom observers keep working.""" + from reflex.workflow.kernel import WorkflowObserver + + assert WorkflowObserver().on_schedule_skip("nightly", 1) is None From 45ba1c70ef730c2fcc9f9b4a6a6372fd5a219196 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Thu, 20 Aug 2026 14:17:20 -0700 Subject: [PATCH 112/121] workflows: second review round Every item reproduced before fixing, verified against a real server. Arrival deadline parity. record_arrival accepted arrivals to a past-deadline parent on all three stores, and then the atomic path disagreed: Postgres refused, Memory and SQLite resolved the join. The same fan-out behaved differently depending on what was behind it, which is worse than either answer alone. All four paths now refuse it as expired, and a conformance check covers them. Schedule seeding. _started_at was captured in __init__, before the first clock sync, so a worker running slow seeded a brand new schedule behind store time and backfilled occurrences from before the deploy. Taken at the end of the first recovery pass instead -- still "when this worker started", now on the store's clock. My first attempt moved it to the first sweep, which is not the same thing and quietly broke restart catch-up; the schedule tests caught it. Empty run id. `cancel "$RUN_ID"` with the variable unset arrives as an empty string, which prefixes every run, and with one run in the database it cancelled it and reported success. That was mine, introduced with prefix support. Refused now. The second Postgres deadlock. recover_orphans locked step rows and then updated their runs; commit takes the run and then the step. The store's stated invariant is run-first and recovery was the single path breaking it, so it now locks the run rows -- which serializes exactly as well, because every writer already obeys that order. 56 aborts in 30 rounds before, 0 in 60 after. Ordering again, not retries. Also: forced-failure error payloads go through strict serde like results; the prefix scan reads one past its cap so a database holding exactly the cap is not treated as truncated; missed occurrences are counted without a ceiling, because the count is what an alert fires on and the bound existed to limit catch-up, not accounting. And Postgres runs in CI. The conformance suite tests every store it can reach, so with no server the Postgres rows skipped and said nothing -- which is how store-specific divergence got as far as review. The workflow suite now runs a second time against a real server on Linux. --- .github/workflows/unit_tests.yml | 22 +++++ news/workflow-review-round-two.bugfix.md | 13 +++ reflex/workflow/cli.py | 13 ++- reflex/workflow/conformance.py | 21 +++++ reflex/workflow/cron.py | 24 ++++++ reflex/workflow/kernel.py | 32 ++++++-- reflex/workflow/postgres.py | 15 +++- reflex/workflow/store.py | 19 ++++- tests/units/workflow/test_clock_authority.py | 21 +++++ tests/units/workflow/test_review_fixes.py | 85 ++++++++++++++++++++ 10 files changed, 254 insertions(+), 11 deletions(-) create mode 100644 news/workflow-review-round-two.bugfix.md diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index d24923c456c..e25892010d4 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -45,6 +45,22 @@ jobs: ports: # Maps port 6379 on service container to the host - 6379:6379 + # The workflow store conformance suite runs against every implementation + # it can reach. Without a server the Postgres rows silently skip, and a + # divergence that only Postgres has -- hand-written SQL, real row locks, + # real deadlock detection -- merges green. + postgres: + image: ${{ matrix.os == 'ubuntu-latest' && 'postgres:16-alpine' || '' }} + env: + POSTGRES_PASSWORD: workflow + POSTGRES_DB: workflow + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 + ports: + - 5432:5432 steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: @@ -66,6 +82,12 @@ jobs: run: | export PYTHONUNBUFFERED=1 uv run pytest tests/units --cov --no-cov-on-fail --cov-report= + - name: Run workflow tests against Postgres + if: ${{ matrix.os == 'ubuntu-latest' }} + run: | + export PYTHONUNBUFFERED=1 + export REFLEX_TEST_POSTGRES=postgresql://postgres:workflow@localhost:5432/workflow + uv run pytest tests/units/workflow --no-cov -q - name: Run unit tests w/ redis if: ${{ matrix.os == 'ubuntu-latest' }} run: | diff --git a/news/workflow-review-round-two.bugfix.md b/news/workflow-review-round-two.bugfix.md new file mode 100644 index 00000000000..0f995d28836 --- /dev/null +++ b/news/workflow-review-round-two.bugfix.md @@ -0,0 +1,13 @@ +A second review round, all reproduced before fixing. + +An arrival to a past-deadline parent is now refused as `expired` on every store. Previously `record_arrival` accepted it everywhere while the atomic path diverged — Postgres refused, Memory and SQLite resolved the join — so the same fan-out behaved differently depending on the store behind it. A conformance check covers both paths. + +A new schedule is seeded from store time rather than worker time. The seed was captured when the kernel was constructed, before the first clock sync, so a worker whose machine ran slow backfilled occurrences from before the deployment existed. It is taken at the end of the first recovery pass instead, which is still "when this worker started" but on the store's clock. + +An empty run id is refused instead of matching every run. `reflex workflows cancel "$RUN_ID"` with the variable unset arrived as an empty string, which is a prefix of everything; with exactly one run in the database it resolved to that run, cancelled it, and reported success. + +A second Postgres deadlock is gone. Recovery locked step rows and then updated their runs, while `commit` locks the run and then the step — the store's own invariant is run-first, and recovery was the one path inverting it. A probe measured 56 aborted transactions across 30 rounds; it now sees none in 60. Fixed by lock ordering, like the first one. + +Forced-failure error payloads face the same strict serialization as results. The prefix scan no longer disables itself when a database holds exactly the scan limit. Missed-occurrence counts are counted rather than sampled, so a long outage is not undercounted by the bound meant for catch-up. + +Postgres now runs in CI. The conformance suite tests every store it can reach, and without a server the Postgres rows skipped silently — which is how store-specific divergences reached review in the first place. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 66a153df433..faac40cad83 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -122,11 +122,20 @@ async def _resolve_run_id(store: RunStore, run_id: str) -> str: Raises: Exit: If the prefix matches no run, or more than one. """ + if not run_id.strip(): + # `reflex workflows cancel "$RUN_ID"` with RUN_ID unset arrives here as + # an empty string, which prefixes every run. With exactly one run in + # the database that resolved and cancelled it, reporting success. + console.error("No run id given. Pass the run to act on.") + raise click.exceptions.Exit(1) if await store.get_run(run_id) is not None: return run_id - scanned = await store.list_runs(RunQuery(limit=_PREFIX_SCAN_LIMIT)) + # One more than the cap: a page that comes back full means there may be + # further runs, while exactly the cap would otherwise read as truncated + # and disable prefixes for a database holding precisely that many. + scanned = await store.list_runs(RunQuery(limit=_PREFIX_SCAN_LIMIT + 1)) matches = [run.run_id for run in scanned if run.run_id.startswith(run_id)] - if len(scanned) >= _PREFIX_SCAN_LIMIT: + if len(scanned) > _PREFIX_SCAN_LIMIT: # The newest N runs, not all of them: a prefix that looks unique here # may match another run just outside the window, and resolving it # would cancel or complete the wrong one. There is no prefix query in diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index e20b1516b8d..bcd4623f292 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -1398,6 +1398,26 @@ async def check_a_delivery_to_a_past_deadline_run_is_refused(store: RunStore) -> assert await store.deliver("run1", "sig:ping", "d1", {"v": 1}, NOW) == "expired" +async def check_an_arrival_to_a_past_deadline_parent_is_refused( + store: RunStore, +) -> None: + """A join that can never run must not be resolved by a late branch. + + The parent is about to finalize TIMED_OUT and its join slot will be + tombstoned, so counting the arrival records a continuation that cannot + happen -- and the stores disagreed about it, which is worse than either + answer: the same fan-out resolved on one store and refused on another. + """ + await store.admit( + make_run(next_ordinal=2, deadline=NOW - 1), + make_step(status=StepStatus.BLOCKED, wait_key="join:0", due_at=0.0), + _ADMITTED, + ) + assert await store.record_arrival( + "run1", 0, {"status": "completed"}, "kid1", NOW + ) == ("expired") + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1437,6 +1457,7 @@ async def check_a_delivery_to_a_past_deadline_run_is_refused(store: RunStore) -> check_force_finalize_records_a_result, check_schedule_cursors_persist, check_join_arrivals_count_once, + check_an_arrival_to_a_past_deadline_parent_is_refused, check_finalize_refuses_while_a_step_is_claimed, check_finalize_tombstones_open_slots, check_finalizing_a_parent_closes_its_branches, diff --git a/reflex/workflow/cron.py b/reflex/workflow/cron.py index 3808d66debb..f48c8ad5e49 100644 --- a/reflex/workflow/cron.py +++ b/reflex/workflow/cron.py @@ -201,6 +201,30 @@ def next_after(self, after: float) -> float | None: return occurrence.timestamp() return None + def count_between(self, after: float, until: float) -> int: + """Count occurrences in a half-open interval, without a ceiling. + + ``occurrences_between`` bounds what it returns so a catch-up cannot + run away. Counting how many were missed is a different question, and + answering it with a bounded list undercounts a long outage exactly + when the number matters most. + + Args: + after: Exclusive lower bound in epoch seconds. + until: Inclusive upper bound in epoch seconds. + + Returns: + How many occurrences fall in the interval. + """ + total = 0 + moment = after + while True: + nxt = self.next_after(moment) + if nxt is None or nxt > until: + return total + total += 1 + moment = nxt + def occurrences_between( self, after: float, until: float, *, limit: int ) -> list[float]: diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 0065bb9ddce..9def63b7ce1 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -542,7 +542,11 @@ def __init__( # previous worker's cursor, and only a schedule this deployment has # never seen starts from now (never backfilling its whole history). self._schedule_cursor: dict[str, float] = {} - self._started_at = clock() + # Filled on first use rather than here: the clock is not synced with + # the store until the first recovery pass, and a worker whose machine + # runs slow would otherwise seed every new schedule behind store time + # and backfill occurrences from before this deployment existed. + self._started_at: float | None = None self._field_adapters: dict[tuple[str, str], TypeAdapter] = {} self._inflight: dict[str, asyncio.Task] = {} self._draining = False @@ -1079,6 +1083,11 @@ async def force_finalize( # json.dumps -- the same input, three behaviours, none of them saying # what to do about it. result = to_run_data({"value": result})["value"] if result is not None else None + # The error payload is stored beside the result and read back the same + # way, so it faces the same rules: an operator reason carrying a + # Decimal or bytes would fail at the store, or worse, only on some + # stores. + error = to_run_data(error) if error is not None else None event = ( HistoryEventType.RUN_COMPLETED if status is RunStatus.COMPLETED @@ -2374,6 +2383,12 @@ async def _admit_due_schedules(self, now: float) -> int: # skip the downtime: an in-memory cursor seeded at startup # treats every missed occurrence as already fired. stored = await self._store.read_schedule_cursor(key) + if self._started_at is None: + # Recovery sets this at startup, right after the clock is + # synced, which is the value that matters. Reaching here + # means a kernel that swept without ever recovering; "now" + # is still the right seed, just an unsynced one. + self._started_at = self._clock() cursor = stored if stored is not None else self._started_at self._schedule_cursor[key] = cursor occurrences = schedule.occurrences_between( @@ -2384,10 +2399,10 @@ async def _admit_due_schedules(self, now: float) -> int: # scheduled work reads as "covered" when it was not; the # operator gets the count and the window, and can start the # missed occurrences by hand if they matter. - dropped = ( - len(schedule.occurrences_between(cursor, now, limit=10_000)) - - MAX_SCHEDULE_CATCHUP - ) + # Counted rather than sampled: a second bounded query would + # undercount a long outage exactly when the number matters + # most, and the count is what an alert fires on. + dropped = schedule.count_between(cursor, now) - MAX_SCHEDULE_CATCHUP occurrences = occurrences[:MAX_SCHEDULE_CATCHUP] # A log line is the only trace these otherwise leave, and they # have no run to carry history. A counter survives the process @@ -2954,6 +2969,13 @@ async def recover(self) -> int: The number of steps recovered. """ await self._sync_store_clock() + if self._started_at is None: + # Now, and not a moment before: this is the seed for a schedule + # this deployment has never seen, and taking it from an unsynced + # clock on a slow machine backfills occurrences from before the + # deployment existed. Recovery runs before the first sweep, so + # the seed is still "when this worker started". + self._started_at = self._clock() await self._renew_leases() now = self._clock() self._next_recovery_at = now + self._recovery_interval diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 90caf539ddc..8c3eb73b877 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1452,7 +1452,8 @@ async def record_arrival( wait_key = f"join:{ordinal}" async with pool.connection() as conn, conn.transaction(): cursor = await conn.execute( - "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", + "SELECT status, deadline FROM workflow_runs WHERE run_id = %s" + " FOR UPDATE", (run_id,), ) run_row = await cursor.fetchone() @@ -1460,6 +1461,10 @@ async def record_arrival( return "unknown_run" if run_row["status"] in _TERMINAL_RUNS: return "run_terminal" + if run_row["deadline"] is not None and run_row["deadline"] <= now: + # The join can never run its continuation: the sweep is about + # to finalize this parent TIMED_OUT and tombstone the slot. + return "expired" cursor = await conn.execute( "SELECT 1 FROM workflow_inbox" " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", @@ -1939,11 +1944,17 @@ async def recover_orphans( exhausted = {"reason": "recovery_budget_exhausted"} async with pool.connection() as conn, conn.transaction(): cursor = await conn.execute( + # Lock the run rows, not the step rows. Every other write + # path takes the run first and the step second -- commit() + # does, through _lock_run -- so locking steps here inverted + # the order and deadlocked against any attempt committing + # late. Holding the run serializes its writers just as well, + # because that is the invariant those writers already obey. "SELECT s.* FROM workflow_steps s" " JOIN workflow_runs r ON r.run_id = s.run_id" " WHERE s.status = %s AND s.lease_expires_at <= %s" " AND NOT (r.status = ANY(%s))" - " ORDER BY s.run_id FOR UPDATE OF s SKIP LOCKED", + " ORDER BY s.run_id FOR UPDATE OF r SKIP LOCKED", (StepStatus.CLAIMED.value, now, _TERMINAL_RUNS), ) rows = await cursor.fetchall() diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 3b4acb33127..19a3ab7a87f 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -1427,6 +1427,11 @@ def _apply_arrival( return "unknown_run" if run.status in TERMINAL_RUN_STATUSES: return "run_terminal" + if run.deadline is not None and run.deadline <= now: + # The join can never run its continuation: the sweep is about to + # finalize this parent TIMED_OUT and tombstone the slot. Counting + # the arrival would record a step that cannot happen. + return "expired" seen = self._inbox.setdefault(run_id, {}) key = (run_id, f"join:{ordinal}", dedupe_key) if key in seen: @@ -3754,12 +3759,16 @@ def _apply_arrival_sql( terminal = tuple(status.value for status in TERMINAL_RUN_STATUSES) wait_key = f"join:{ordinal}" run_row = self._db.execute( - "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + "SELECT status, deadline FROM workflow_runs WHERE run_id = ?", (run_id,) ).fetchone() if run_row is None: return "unknown_run" if run_row["status"] in terminal: return "run_terminal" + if run_row["deadline"] is not None and run_row["deadline"] <= now: + # The join can never run its continuation: the sweep is about to + # finalize this parent TIMED_OUT and tombstone the slot. + return "expired" seen = self._db.execute( "SELECT 1 FROM workflow_inbox" " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", @@ -3852,7 +3861,8 @@ def work(): self._db.execute("BEGIN IMMEDIATE") try: run_row = self._db.execute( - "SELECT status FROM workflow_runs WHERE run_id = ?", (run_id,) + "SELECT status, deadline FROM workflow_runs WHERE run_id = ?", + (run_id,), ).fetchone() if run_row is None: self._db.execute("ROLLBACK") @@ -3860,6 +3870,11 @@ def work(): if run_row["status"] in terminal: self._db.execute("ROLLBACK") return "run_terminal" + if run_row["deadline"] is not None and run_row["deadline"] <= now: + # The join can never run its continuation: the sweep + # is about to finalize this parent TIMED_OUT. + self._db.execute("ROLLBACK") + return "expired" seen = self._db.execute( "SELECT 1 FROM workflow_inbox" " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", diff --git a/tests/units/workflow/test_clock_authority.py b/tests/units/workflow/test_clock_authority.py index c7c6e7b01c9..075cf7706c9 100644 --- a/tests/units/workflow/test_clock_authority.py +++ b/tests/units/workflow/test_clock_authority.py @@ -193,3 +193,24 @@ async def test_store_time_still_advances_with_real_elapsed_time(): second = kernel._clock() # pyright: ignore[reportPrivateUsage] assert second > first assert second - first < 5.0, "advancing far faster than real time is also wrong" + + +async def test_a_new_schedule_is_seeded_from_store_time_not_worker_time(): + """A slow worker must not backfill a schedule from before the deploy. + + The seed for a schedule this deployment has never seen is "now". Taken + from the worker's own clock at construction -- before the first sync -- + "now" on a machine running two minutes slow is two minutes of history, + and the first sweep admits occurrences that were never meant to run. + """ + store = _IndependentClockStore() + kernel = WorkflowKernel([], store) + assert kernel._started_at is None, "the seed must not be taken before syncing" # pyright: ignore[reportPrivateUsage] + + await kernel.recover() + assert kernel._started_at is not None # pyright: ignore[reportPrivateUsage] + store_now = await store.epoch_time() + assert store_now is not None + assert abs(kernel._started_at - store_now) < 1.0, ( # pyright: ignore[reportPrivateUsage] + "the seed must sit on the store's clock, not the worker's" + ) diff --git a/tests/units/workflow/test_review_fixes.py b/tests/units/workflow/test_review_fixes.py index 11017eca4c7..a586934545f 100644 --- a/tests/units/workflow/test_review_fixes.py +++ b/tests/units/workflow/test_review_fixes.py @@ -161,3 +161,88 @@ def test_the_default_observer_ignores_dropped_occurrences_quietly(): from reflex.workflow.kernel import WorkflowObserver assert WorkflowObserver().on_schedule_skip("nightly", 1) is None + + +@pytest.mark.parametrize("empty", ["", " "]) +def test_an_empty_run_id_is_refused_before_it_matches_everything(empty, tmp_path): + """`cancel "$RUN_ID"` with the variable unset must not cancel anything. + + An empty string is a prefix of every run. With exactly one run in the + database it resolved to that run, cancelled it, and reported success -- + the failure mode of an unset shell variable should not be "cancels + production". + + Args: + empty: The blank argument a shell expands to. + tmp_path: Temporary directory for the database. + """ + import asyncio + + from click.exceptions import Exit + + from reflex.workflow.cli import _resolve_run_id + + store = SqliteRunStore(tmp_path / "solo.db") + + async def check() -> None: + """Admit one run, then resolve a blank id against it. + + Raises: + AssertionError: If the blank id resolved to the run. + """ + await store.admit( + _run("only1"), _step("only1"), ((HistoryEventType.RUN_ADMITTED, {}),) + ) + try: + resolved = await _resolve_run_id(store, empty) + except Exit: + return + msg = f"blank id resolved to {resolved!r}" + raise AssertionError(msg) + + asyncio.run(check()) + store.close() + + +async def test_a_forced_failure_reason_faces_strict_serialization(tmp_path): + """The error payload is stored beside the result and read back the same. + + Args: + tmp_path: Temporary directory for the database. + """ + store = SqliteRunStore(tmp_path / "fail.db") + await store.admit(_run(), _step(), ((HistoryEventType.RUN_ADMITTED, {}),)) + kernel = WorkflowKernel([], store) + with pytest.raises(TypeError, match="Decimal is not valid run data"): + await kernel.force_finalize( + "r1", + status=RunStatus.FAILED, + error={"reason": "manual", "amount": decimal.Decimal("1.10")}, + ) + store.close() + + +def test_missed_occurrences_are_counted_without_a_ceiling(): + """A long outage must not be undercounted by a sampling limit. + + The number is what an alert fires on, and "10,000" for an outage that + dropped far more is the kind of wrong that reads as precise. + """ + import datetime as dt + + schedule = CronSchedule("* * * * *") + start = dt.datetime(2026, 1, 1, tzinfo=dt.UTC).timestamp() + # Eleven days of minutes is more than any bounded query returned. + end = start + 11 * 24 * 3600 + assert schedule.count_between(start, end) == 11 * 24 * 60 + + +def test_counting_agrees_with_listing_on_a_small_window(): + """The unbounded count and the bounded list must not disagree.""" + import datetime as dt + + schedule = CronSchedule("0 * * * *") + start = dt.datetime(2026, 1, 1, tzinfo=dt.UTC).timestamp() + end = start + 5 * 3600 + listed = schedule.occurrences_between(start, end, limit=100) + assert schedule.count_between(start, end) == len(listed) From be95a02d6aecb9e6ceaf048fae06f2fd20869786 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Fri, 21 Aug 2026 18:31:43 -0700 Subject: [PATCH 113/121] workflows: terminal transitions close completely, everywhere Three externally-reported defects plus six more an adversarial review found in the fixes themselves, all reproduced before fixing, all pinned by conformance checks that failed on every store first. Recovery-budget exhaustion marked one step and the run FAILED and stopped: preallocated slots stayed open on a dead run and child runs kept working for a parent that no longer existed. It is now the same complete terminal transition failure takes everywhere -- open slots tombstoned, branches closed, the parent's join told. On Postgres the cascade cannot run inside the batch sweep (it holds a page of run rows, and children must be taken before self), so exhausted runs are collected and each fails in its own transaction in canonical order, with a re-check that leaves a renewed or peer-handled claim alone. Skipping the last open slot stamped COMPLETED and stopped: no arrival reached the parent -- BLOCKED 0/1 forever -- and the skipped run's children kept running. Completion by decision is now the same transition as completion by execution. Twelve workers initializing a fresh Postgres schema produced one winner and eleven UniqueViolations: IF NOT EXISTS does not protect the composite-type catalog insert each CREATE TABLE performs. Initialization takes an advisory lock keyed on the schema name and runs all DDL in one transaction, so a worker dying mid-setup leaves nothing behind; the pool is closed instead of leaked when setup fails. The review of these fixes then confirmed six defects in and around them: - The memory exhaustion write rebuilt the run from a pre-pass snapshot, reverting the cancel_requested flag a parent's cascade set on a child exhausting in the same pass. Replaced from the current record; the SQL stores' column updates were already correct. - claim_next locked step-then-run while every arrival path locks run-then-step, and a join slot with a lapsed timeout is claimable -- a reproduced deadlock (the fourth). The claim now locks the run row. - _close_children re-derived its child set after _lock_children pinned it, so a child revived by retry in the window was written parent-first. It now writes exactly the pinned set. - A lease renewal landing between recovery's re-check and its write was acknowledged and discarded in the same instant. Both recovery writes repeat the guard inside the UPDATE and treat zero rows as refusal. - Restore-after-exhaustion brought tombstoned joins and waits back as READY -- running them immediately with missing or partial payloads -- and rewrote delayed slots' due times to now. What was waiting comes back BLOCKED with arrivals and deadlines intact; delays keep their due times. - A second close of the same run appended a duplicate cancel-requested event to each still-cancelling child; the flag is monotonic, so a marked child now has nothing written. Probes: cascade deadlock 0/40, recovery deadlock 0/40, claim-vs-arrival deadlock 0/60 (1/30 before), init race 12/12 x4 (1/12 before). --- news/workflow-terminal-transitions.bugfix.md | 11 + reflex/workflow/CONTRACT.md | 17 +- reflex/workflow/conformance.py | 268 +++++++++++++ reflex/workflow/postgres.py | 374 ++++++++++++++----- reflex/workflow/store.py | 196 ++++++++-- tests/units/workflow/test_postgres.py | 30 ++ 6 files changed, 774 insertions(+), 122 deletions(-) create mode 100644 news/workflow-terminal-transitions.bugfix.md diff --git a/news/workflow-terminal-transitions.bugfix.md b/news/workflow-terminal-transitions.bugfix.md new file mode 100644 index 00000000000..07dceb87674 --- /dev/null +++ b/news/workflow-terminal-transitions.bugfix.md @@ -0,0 +1,11 @@ +Three defects from a third external review round, all reproduced before fixing. + +Exhausting the recovery budget now ends the run the way failure ends it everywhere else. The budget path marked the one exhausted step `FAILED`, stamped the run `FAILED`, and stopped: preallocated successor slots stayed open on a dead run, and child runs kept working for a parent that no longer existed. It now performs the complete terminal transition — open slots tombstoned, branches told to stop, the parent's join delivered one `FAILED` arrival. On Postgres this restructured `recover_orphans` into two phases, because terminal transitions take child locks before their own row and the batch sweep already holds a page of run rows; each exhausted run now fails in its own transaction in the canonical children → self → parent order. + +Skipping the last open slot now completes the run the way completion works. It stamped `COMPLETED` and stopped: no arrival reached a parent joined on the run — which stayed `BLOCKED 0/1` forever — and the skipped run's own children kept running. Completed by an operator's decision and completed by execution are now the same terminal transition, and both new behaviors are pinned by conformance checks on all three stores. + +Twelve workers can now initialize a fresh Postgres schema together. `IF NOT EXISTS` does not make concurrent DDL safe — each `CREATE TABLE` also inserts the table's composite type, and backends that both saw "not exists" raced on `pg_type`: twelve fresh processes produced one winner and eleven `UniqueViolation` crashes, so a fleet's first deploy crash-looped everyone but one. Initialization now takes an advisory lock keyed on the schema name and runs every statement in one transaction, which also means a worker dying mid-setup leaves nothing half-created. + +A fourth Postgres deadlock is gone: `claim_next` locked the frontier step row and then waited on the run row, while every arrival path holds the parent run row and then updates the join step — and a join slot with a lapsed wait deadline is claimable, so a timeout claim could meet the child arrival racing it. The claim now locks the run row, completing the run-first invariant in the one path still inverting it. + +An adversarial review of these fixes found and fixed three more defects in them before they shipped: the memory store's exhaustion write rebuilt the run record from a pre-pass snapshot, silently reverting a `cancel_requested` flag the parent's cascade had just set on a child exhausting in the same pass; a lease renewal landing between recovery's re-check and its write could be acknowledged and immediately discarded, so both recovery writes now repeat the guard inside the UPDATE itself; and restore-after-exhaustion brought tombstoned join and wait slots back as `READY`, which would run them immediately with missing or partial payloads — what was waiting now comes back `BLOCKED` with its arrival count and deadline intact, and delayed slots keep their due times. Postgres schema initialization no longer leaks its connection pool when setup fails, and `_close_children` writes exactly the set of children it locked rather than re-deriving it. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index f6d213ab241..5388e38e762 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -365,7 +365,7 @@ reason rather than on message text. | the next step's handler no longer exists | `unknown_handler` | step and run `NEEDS_ATTENTION`; restore the handler and `resume(run)`, or cancel | | a recorded payload carries arguments the handler no longer accepts | `incompatible_payload` | step and run `NEEDS_ATTENTION`, naming the arguments; restore the parameters or cancel | | a run allocated more steps than `WorkflowConfig.max_steps` | `max_steps_exceeded` | the committing step succeeds, every other open slot is tombstoned, run `FAILED`. The bound is on a runaway loop, so it fails rather than suspending for a person to approve more of the same | -| a step's lease lapsed more times than `max_recoveries` | `recovery_budget_exhausted` | run `FAILED`. Infrastructure recoveries are free of the retry budget precisely so they can be bounded separately; the bound is what stops a poison step cycling forever | +| a step's lease lapsed more times than `max_recoveries` | `recovery_budget_exhausted` | run `FAILED` — the complete terminal transition, identical to any other failure: remaining open slots are tombstoned, children are told to stop, and the parent's join hears one `FAILED` arrival. Infrastructure recoveries are free of the retry budget precisely so they can be bounded separately; the bound is what stops a poison step cycling forever | | an error's `details` cannot be serialized | — | the reason is preserved and the details are replaced by `{"unserializable": repr(...)}`. Losing the payload never turns a failure into a crash | Nothing on this table is silent: each writes its reason to the run's error and @@ -410,10 +410,21 @@ no-op with a reason. - `skip(run)` — a run stopped for attention or failure: marks the blocking step `SKIPPED` (terminal, recorded as a decision rather than an outcome) and lets the run continue at whatever comes next. With nothing left to run, - the run completes with no result rather than sitting pending forever. + the run completes with no result rather than sitting pending forever — and + completing by decision is the same terminal transition as completing by + execution: children are told to stop, and a parent joined on this run + receives one `COMPLETED` arrival (`result: null`) instead of waiting + forever on a run that no longer will. If the run already delivered its + arrival when it first reached a terminal state, the join keeps what it + heard: a run delivers exactly one arrival (§5), and operator repair does + not rewrite a result the parent has already counted. - Both restore the successors the stopping failure tombstoned (`step_restored` in history, fresh budgets), so a preallocated chain's remaining steps — - including its finalizer — still run. Only that failure's casualties come + including its finalizer — still run. What was waiting comes back waiting: + a tombstoned join or wait slot is restored `BLOCKED` with its arrival + count and timeout deadline intact — never `READY`, which would run it + immediately with a missing or partial payload — and a delayed slot keeps + its original due time rather than firing the moment an operator retries. Only that failure's casualties come back: a `CANCELLED` slot in a run these actions accept can have no other source, because run-level cancellation ends in a `CANCELLED` run they refuse and force-finalization leaves them no step to target. diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index bcd4623f292..fb3eb0e3746 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -23,6 +23,7 @@ async def test_my_store_conforms(check): import pytest from reflex.workflow.records import ( + TERMINAL_RUN_STATUSES, HistoryEventType, RunQuery, RunRecord, @@ -1418,6 +1419,268 @@ async def check_an_arrival_to_a_past_deadline_parent_is_refused( ) == ("expired") +async def check_recovery_exhaustion_closes_the_whole_run(store: RunStore) -> None: + """Exhausting the recovery budget must end the run the way failure does. + + The budget path marked the one exhausted step FAILED and stamped the run + FAILED -- and stopped. Preallocated successor slots stayed open, so the + dead run still surfaced wake times and confused retry; child runs kept + working for a parent that no longer existed. A budget exhaustion is a + failure, and failure closes the run completely: open slots tombstoned, + branches told to stop. + """ + await store.admit( + make_run("par0", next_ordinal=1), + make_step( + "par0", + status=StepStatus.BLOCKED, + wait_key="join:0", + join_expected=1, + origin="join", + due_at=0.0, + ), + _ADMITTED, + ) + await store.admit( + make_run(parent_run_id="par0", parent_ordinal=0), make_step(), _ADMITTED + ) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + assert claim.run.run_id == "run1", "par0 has no claimable slot" + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.PENDING, + state={"n": 1}, + new_steps=( + make_step(ordinal=1, due_at=NOW), + make_step(ordinal=2, due_at=NOW + 3600), + ), + next_ordinal=3, + events=((HistoryEventType.ATTEMPT_SUCCEEDED, {}),), + children=( + ( + make_run("kid", parent_run_id="run1", parent_ordinal=2), + make_step("kid", due_at=NOW + 7200), + ), + ), + ), + NOW, + ) + claim = await store.claim_next(NOW, lease_duration=0.0) + assert claim is not None + assert claim.run.run_id == "run1" + + _, failed = await store.recover_orphans(NOW + 1, max_recoveries=0) + assert "run1" in failed + + run = await store.get_run("run1") + assert run is not None + assert run.status is RunStatus.FAILED + steps = await store.get_steps("run1") + assert steps[1].status is StepStatus.FAILED + assert steps[2].status is StepStatus.CANCELLED, ( + f"the preallocated slot must be tombstoned CANCELLED -- the exact " + f"status retry restores -- not left {steps[2].status}" + ) + kid = await store.get_run("kid") + assert kid is not None + assert kid.cancel_requested, "the child kept working for a dead parent" + par_steps = await store.get_steps("par0") + assert par_steps[0].join_arrived == 1, "the parent never heard the child exhausted" + assert par_steps[0].args["__results__"][0]["status"] == RunStatus.FAILED.value + # No leftover-claim sweep here: with every slot's status pinned exactly + # above, a frontier-based claim has nothing left to find, and a loop + # that cannot iterate reads as coverage it does not provide. + + +async def check_skipping_the_last_step_completes_like_a_completion( + store: RunStore, +) -> None: + """A skip that finishes the run must finish it everywhere it matters. + + Skipping the only open slot stamped the run COMPLETED and stopped there: + no arrival reached the parent's join, which stayed BLOCKED 0/1 forever, + and the skipped run's own children kept running. "Completed by an + operator's decision" and "completed" must be the same terminal + transition. + """ + await store.admit( + make_run("par", next_ordinal=1), + make_step( + "par", + status=StepStatus.BLOCKED, + wait_key="join:0", + join_expected=1, + origin="join", + due_at=0.0, + ), + _ADMITTED, + ) + await store.admit( + make_run( + "kid", + parent_run_id="par", + parent_ordinal=0, + status=RunStatus.FAILED, + error={"reason": "boom"}, + ), + make_step("kid", status=StepStatus.FAILED, error={"reason": "boom"}), + _ADMITTED, + ) + await store.admit(make_run("gk", parent_run_id="kid"), make_step("gk"), _ADMITTED) + + assert await store.skip_step("kid", NOW) + + kid = await store.get_run("kid") + assert kid is not None + assert kid.status is RunStatus.COMPLETED + par_steps = await store.get_steps("par") + assert par_steps[0].join_arrived == 1, "the parent never heard the child finished" + assert par_steps[0].status is StepStatus.READY + assert par_steps[0].args["__results__"][0]["status"] == RunStatus.COMPLETED.value + gk = await store.get_run("gk") + assert gk is not None + assert gk.cancel_requested, "the grandchild kept working under a closed branch" + + +async def check_a_cascade_flag_survives_the_childs_own_exhaustion( + store: RunStore, +) -> None: + """One crashed worker takes out a parent and its child in the same pass. + + The parent exhausts first and its cascade durably requests the child's + cancellation; the child's own exhaustion then fails the child. Rebuilding + the child's record from a pre-pass snapshot reverted the flag the cascade + had just set, and a later retry could revive a branch whose parent is + dead with nothing ever re-requesting its cancellation. The failure may + keep the child FAILED, but the request must survive it. + + Run ids are chosen so every store processes the parent first ("a-" sorts + and inserts before "z-"), which is the ordering that exercises the + cascade-then-exhaust path. + """ + await store.admit(make_run("a-par"), make_step("a-par"), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=0.0) + assert claim is not None + assert claim.run.run_id == "a-par" + await store.admit( + make_run("z-kid", parent_run_id="a-par"), make_step("z-kid"), _ADMITTED + ) + claim = await store.claim_next(NOW, lease_duration=0.0) + assert claim is not None + assert claim.run.run_id == "z-kid" + + _, failed = await store.recover_orphans(NOW + 1, max_recoveries=0) + assert "a-par" in failed + + kid = await store.get_run("z-kid") + assert kid is not None + assert kid.status in TERMINAL_RUN_STATUSES + assert kid.cancel_requested, ( + "the parent's cascade requested cancellation and the child's own " + "exhaustion erased the request" + ) + + +async def check_a_second_close_does_not_repeat_the_cancel_request( + store: RunStore, +) -> None: + """History records the decision to stop a branch once, not per close. + + A failed run whose children were already told to stop can reach a second + terminal transition -- an operator skipping its failed step completes it. + The second close used to append another cancel-requested event to every + still-CANCELLING child, so history read as if the decision were made + twice. The flag is durable and monotonic; a marked child has nothing + left to write. + """ + await store.admit( + make_run("x", status=RunStatus.FAILED, error={"reason": "boom"}), + make_step("x", status=StepStatus.FAILED, error={"reason": "boom"}), + _ADMITTED, + ) + await store.admit(make_run("kid", parent_run_id="x"), make_step("kid"), _ADMITTED) + assert await store.request_cancel("kid", NOW) + assert await store.skip_step("x", NOW + 1) + + kid = await store.get_run("kid") + assert kid is not None + assert kid.cancel_requested + events = [ + event + for event in await store.get_history("kid") + if event.type is HistoryEventType.RUN_CANCEL_REQUESTED + ] + assert len(events) == 1, f"the stop decision was recorded {len(events)} times" + + +async def check_retry_after_exhaustion_restores_waits_as_waits( + store: RunStore, +) -> None: + """A restored join must wait again, not run with missing inputs. + + Exhaustion tombstones every open slot, including a BLOCKED join holding + partial arrivals and a delayed slot holding a future due time. Retry + restores the chain -- and restoring a join as READY would run its + handler immediately with a partial result set, while rewriting a delayed + slot's due time to "now" would erase the delay. What was waiting comes + back waiting, with its arrival count and its deadline intact. + """ + await store.admit(make_run(), make_step(), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + await store.commit( + claim, + StepCompletion( + step_status=StepStatus.SUCCEEDED, + run_status=RunStatus.PENDING, + state={"n": 1}, + new_steps=( + make_step(ordinal=1, due_at=NOW), + make_step( + ordinal=2, + status=StepStatus.BLOCKED, + wait_key="join:2", + join_expected=2, + join_arrived=1, + origin="join", + due_at=NOW + 9000, + ), + make_step(ordinal=3, due_at=NOW + 3600), + ), + next_ordinal=4, + events=((HistoryEventType.ATTEMPT_SUCCEEDED, {}),), + ), + NOW, + ) + claim = await store.claim_next(NOW, lease_duration=0.0) + assert claim is not None + assert claim.step.ordinal == 1 + + _, failed = await store.recover_orphans(NOW + 1, max_recoveries=0) + assert "run1" in failed + steps = await store.get_steps("run1") + assert steps[2].status is StepStatus.CANCELLED + assert steps[3].status is StepStatus.CANCELLED + + assert await store.retry_run("run1", NOW + 10) + steps = await store.get_steps("run1") + assert steps[2].status is StepStatus.BLOCKED, ( + f"the join came back {steps[2].status}, and READY would run it with " + "one of two results" + ) + assert steps[2].join_arrived == 1, "the arrival already counted was lost" + assert steps[2].due_at == pytest.approx(NOW + 9000), ( + "the join's timeout deadline was rewritten" + ) + assert steps[3].status is StepStatus.READY + assert steps[3].due_at == pytest.approx(NOW + 3600), ( + "the delayed slot's due time was rewritten; its delay is erased" + ) + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1459,6 +1722,11 @@ async def check_an_arrival_to_a_past_deadline_parent_is_refused( check_join_arrivals_count_once, check_an_arrival_to_a_past_deadline_parent_is_refused, check_finalize_refuses_while_a_step_is_claimed, + check_recovery_exhaustion_closes_the_whole_run, + check_a_cascade_flag_survives_the_childs_own_exhaustion, + check_skipping_the_last_step_completes_like_a_completion, + check_a_second_close_does_not_repeat_the_cancel_request, + check_retry_after_exhaustion_restores_waits_as_waits, check_finalize_tombstones_open_slots, check_finalizing_a_parent_closes_its_branches, check_closing_a_branch_never_revives_a_finished_one, diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 8c3eb73b877..1cc0964a69c 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -383,16 +383,45 @@ async def configure(conn: AsyncConnection) -> None: open=False, ) await pool.open(wait=True) - async with pool.connection() as conn: - if schema is not None: - await conn.execute( - SQL("CREATE SCHEMA IF NOT EXISTS {}").format(Identifier(schema)) - ) - await conn.execute(_set_search_path(schema)) - await conn.execute(_SCHEMA) + try: + await self._initialize_schema(pool, schema) + except BaseException: + # The pool is already open; abandoning it here would leak its + # connections and every retry would leak another pool's worth. + await pool.close() + raise self._pool = pool return pool + async def _initialize_schema(self, pool: Any, schema: str | None) -> None: + """Create this store's schema and tables, safely under concurrency. + + Args: + pool: The open connection pool. + schema: The schema to create tables in, or None for the search + path's default. + """ + async with pool.connection() as conn, conn.transaction(): + # IF NOT EXISTS does not make concurrent DDL safe: each + # CREATE TABLE also inserts the table's composite type, and + # two backends that both saw "not exists" race on pg_type -- + # twelve fresh workers produced one winner and eleven + # UniqueViolations. The advisory lock serializes the + # initializers so the losers' IF NOT EXISTS genuinely sees + # the winner's objects, and running every statement in one + # transaction means a worker that dies mid-setup leaves + # nothing half-created behind. + await conn.execute( + "SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", + (f"reflex-workflow-ddl:{schema or ''}",), + ) + if schema is not None: + await conn.execute( + SQL("CREATE SCHEMA IF NOT EXISTS {}").format(Identifier(schema)) + ) + await conn.execute(_set_search_path(schema)) + await conn.execute(_SCHEMA) + @property def schema(self) -> str | None: """The schema this store's tables live in, if one was named. @@ -897,7 +926,16 @@ async def claim_next( f" WHERE {_RUNNABLE_PREDICATE} AND {_FRONTIER_PREDICATE}" f" AND {_CLAIMABLE_PREDICATE} AND {_QUEUE_PREDICATE}" " ORDER BY r.created_at, s.run_id" - " FOR UPDATE OF s SKIP LOCKED LIMIT 1", + # OF r, not OF s: every arrival path holds the parent run row + # and then updates the join step, and a join slot with a + # lapsed wait deadline is claimable -- locking the step first + # here was the last remaining inversion of the run-first + # invariant, and it deadlocked a timeout claim against the + # child arrival racing it. Holding the run row serializes the + # step writers just as well, because every one of them takes + # the run row first. SKIP LOCKED at run granularity also + # matches the serial mailbox: one claim per run. + " FOR UPDATE OF r SKIP LOCKED LIMIT 1", params, ) candidate = await cursor.fetchone() @@ -1003,8 +1041,9 @@ async def commit( """ pool = await self._open() async with pool.connection() as conn, conn.transaction(): + locked_children: list[str] = [] if completion.run_status in TERMINAL_RUN_STATUSES: - await self._lock_children(conn, claim.run.run_id) + locked_children = await self._lock_children(conn, claim.run.run_id) await self._lock_run(conn, claim.run.run_id) deadline = await self._check_claim(conn, claim) # Past the deadline the only permitted outcome is TIMED_OUT, and @@ -1069,7 +1108,7 @@ async def commit( ) await self._append_events(conn, claim.run.run_id, completion.events, now) if completion.run_status in TERMINAL_RUN_STATUSES: - await self._close_children(conn, claim.run.run_id, now) + await self._close_children(conn, now, locked_children) if completion.parent_arrival is not None: await self._apply_arrival(conn, *completion.parent_arrival, now) @@ -1271,7 +1310,7 @@ async def admit_children( await self._lock_run(conn, parent) await self._append_events(conn, parent, events, now) - async def _lock_children(self, conn: Any, run_id: str) -> None: + async def _lock_children(self, conn: Any, run_id: str) -> list[str]: """Take the branch rows this transaction will close, before its own. Closing a parent locks the parent and then its children; a child @@ -1285,18 +1324,32 @@ async def _lock_children(self, conn: Any, run_id: str) -> None: before shallower ones, so no cycle can form. The ORDER BY matters for the same reason within one level. + Children already told to cancel are excluded: the flag is durable + and monotonic, so there is nothing left to write on them -- and + skipping them keeps a second close from appending a duplicate + cancel-requested event to their history. + Args: conn: The connection inside an open transaction. run_id: The run whose branches may be closed. + + Returns: + The child run ids this transaction holds, for _close_children -- + which must write exactly this set, because a child revived + between the lock and the write would otherwise be written + parent-before-child, the inversion this ordering exists to + prevent. """ - await conn.execute( + cursor = await conn.execute( "SELECT run_id FROM workflow_runs WHERE parent_run_id = %s" - " AND parent_close <> 'abandon' AND NOT (status = ANY(%s))" + " AND parent_close <> 'abandon' AND NOT cancel_requested" + " AND NOT (status = ANY(%s))" " ORDER BY run_id FOR UPDATE", (run_id, [s.value for s in TERMINAL_RUN_STATUSES]), ) + return [row["run_id"] for row in await cursor.fetchall()] - async def _close_children(self, conn: Any, run_id: str, now: float) -> None: + async def _close_children(self, conn: Any, now: float, children: list[str]) -> None: """Request cancellation of branches the closing run fanned out to. Called inside the transaction that takes a run terminal, so an @@ -1308,19 +1361,30 @@ async def _close_children(self, conn: Any, run_id: str, now: float) -> None: Args: conn: The connection inside an open transaction. - run_id: The run reaching a terminal state. now: Current time in epoch seconds. + children: The child run ids _lock_children pinned earlier in this + transaction. """ + if not children: + return closing = await ( await conn.execute( + # Exactly the set _lock_children pinned, never re-derived: a + # child revived between the lock and this write would match a + # fresh predicate without ever having been locked, and this + # transaction would then take its row while holding the + # parent -- the shallower-before-deeper inversion the lock + # ordering exists to prevent. The revived child is the + # operator's decision and keeps running; it was terminal when + # this close began. "UPDATE workflow_runs SET cancel_requested = TRUE, status = %s," - " updated_at = %s WHERE parent_run_id = %s" - " AND parent_close <> 'abandon' AND NOT (status = ANY(%s))" + " updated_at = %s WHERE run_id = ANY(%s)" + " AND NOT (status = ANY(%s))" " RETURNING run_id", ( RunStatus.CANCELLING.value, now, - run_id, + children, [s.value for s in TERMINAL_RUN_STATUSES], ), ) @@ -1708,7 +1772,7 @@ async def finalize_run( """ pool = await self._open() async with pool.connection() as conn, conn.transaction(): - await self._lock_children(conn, run_id) + locked_children = await self._lock_children(conn, run_id) cursor = await conn.execute( "SELECT status FROM workflow_runs WHERE run_id = %s FOR UPDATE", (run_id,), @@ -1745,7 +1809,7 @@ async def finalize_run( ] events.append((event, {} if error is None else dict(error))) await self._append_events(conn, run_id, events, now) - await self._close_children(conn, run_id, now) + await self._close_children(conn, now, locked_children) if parent_arrival is not None: await self._apply_arrival(conn, *parent_arrival, now) return True @@ -1804,12 +1868,18 @@ async def _restore_tombstoned(conn: Any, run_id: str, now: float) -> list[int]: The restored ordinals, in order. """ cursor = await conn.execute( - "UPDATE workflow_steps SET status = %s, attempts = 0, due_at = %s," + # Waits and joins come back BLOCKED with their arrival counts and + # deadlines intact, never READY -- restored-as-READY they would + # run immediately with a missing or partial payload. Plain slots + # keep their own due_at, so a restored delay still waits out its + # delay instead of firing the moment an operator retries. + "UPDATE workflow_steps SET status = CASE WHEN wait_key IS NULL" + " THEN %s ELSE %s END, attempts = 0," " lease_expires_at = 0, error = NULL, updated_at = %s" " WHERE run_id = %s AND status = %s RETURNING ordinal", ( StepStatus.READY.value, - now, + StepStatus.BLOCKED.value, now, run_id, StepStatus.CANCELLED.value, @@ -1876,6 +1946,12 @@ async def skip_step(self, run_id: str, now: float) -> bool: """ pool = await self._open() async with pool.connection() as conn, conn.transaction(): + # A skip that removes the last open slot completes the run, and + # terminal transitions take child locks before their own row. + # Whether this skip completes is not known until the slot count + # below, so the branches are taken unconditionally -- the order + # is what matters, not the need. + locked_children = await self._lock_children(conn, run_id) await self._lock_run(conn, run_id) cursor = await conn.execute( "SELECT s.ordinal AS ordinal FROM workflow_steps s" @@ -1926,6 +2002,31 @@ async def skip_step(self, run_id: str, now: float) -> bool: if not open_left: events.append((HistoryEventType.RUN_COMPLETED, {})) await self._append_events(conn, run_id, tuple(events), now) + if not open_left: + # Completed by an operator's decision is still completed: + # branches are told to stop, and a parent joined on this run + # hears it finished instead of waiting forever. + await self._close_children(conn, now, locked_children) + cursor = await conn.execute( + "SELECT parent_run_id, parent_ordinal FROM workflow_runs" + " WHERE run_id = %s", + (run_id,), + ) + parent = await cursor.fetchone() + if parent is not None and parent["parent_run_id"] is not None: + await self._apply_arrival( + conn, + parent["parent_run_id"], + parent["parent_ordinal"], + { + "run_id": run_id, + "status": RunStatus.COMPLETED.value, + "result": None, + "error": None, + }, + run_id, + now, + ) return True async def recover_orphans( @@ -1941,7 +2042,9 @@ async def recover_orphans( How many steps were transitioned, and the runs failed outright. """ pool = await self._open() - exhausted = {"reason": "recovery_budget_exhausted"} + recovered = 0 + failed: list[str] = [] + overdrawn: list[StepRecord] = [] async with pool.connection() as conn, conn.transaction(): cursor = await conn.execute( # Lock the run rows, not the step rows. Every other write @@ -1958,84 +2061,169 @@ async def recover_orphans( (StepStatus.CLAIMED.value, now, _TERMINAL_RUNS), ) rows = await cursor.fetchall() - recovered = 0 - failed: list[str] = [] for row in rows: step = _step_from_row(row) - recovered += 1 if step.recoveries + 1 > max_recoveries: - await conn.execute( - "UPDATE workflow_steps SET status = %s, recoveries = %s," - " lease_expires_at = 0, error = %s, updated_at = %s" - " WHERE run_id = %s AND ordinal = %s", - ( - StepStatus.FAILED.value, - step.recoveries + 1, - _json(exhausted), - now, - step.run_id, - step.ordinal, - ), - ) - await conn.execute( - "UPDATE workflow_runs SET status = %s, error = %s," - " updated_at = %s WHERE run_id = %s", - ( - RunStatus.FAILED.value, - _json(exhausted), - now, - step.run_id, - ), - ) - failed.append(step.run_id) - cursor = await conn.execute( - "SELECT parent_run_id, parent_ordinal FROM workflow_runs" - " WHERE run_id = %s", - (step.run_id,), - ) - parent = await cursor.fetchone() - if parent is not None and parent["parent_run_id"] is not None: - await self._apply_arrival( - conn, - parent["parent_run_id"], - parent["parent_ordinal"], - { - "run_id": step.run_id, - "status": RunStatus.FAILED.value, - "result": None, - "error": dict(exhausted), - }, - step.run_id, - now, - ) - await self._append_events( - conn, - step.run_id, - ((HistoryEventType.RUN_FAILED, dict(exhausted)),), + # Exhaustion is a terminal transition, and terminal + # transitions take child locks before their own row. This + # transaction already holds a batch of run rows, so + # taking children now would acquire shallower before + # deeper -- the inversion both deadlock fixes removed. + # The step stays CLAIMED with its lapsed lease (nothing + # can claim it) and fails in its own transaction below, + # in the canonical order. + overdrawn.append(step) + continue + cursor = await conn.execute( + # Guarded like the exhaustion write below: a renewal can + # land between the batch SELECT and this UPDATE because + # renew_lease takes no run-row lock. + "UPDATE workflow_steps SET status = %s, recoveries = %s," + " due_at = %s, lease_expires_at = 0, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s AND status = %s" + " AND lease_expires_at <= %s", + ( + StepStatus.RECOVERY_WAIT.value, + step.recoveries + 1, + now, now, - ) - else: - await conn.execute( - "UPDATE workflow_steps SET status = %s, recoveries = %s," - " due_at = %s, lease_expires_at = 0, updated_at = %s" - " WHERE run_id = %s AND ordinal = %s", - ( - StepStatus.RECOVERY_WAIT.value, - step.recoveries + 1, - now, - now, - step.run_id, - step.ordinal, - ), - ) - await self._append_events( - conn, step.run_id, - ((HistoryEventType.STEP_RECOVERED, {"ordinal": step.ordinal}),), + step.ordinal, + StepStatus.CLAIMED.value, now, - ) + ), + ) + if cursor.rowcount == 0: + continue + recovered += 1 + await self._append_events( + conn, + step.run_id, + ((HistoryEventType.STEP_RECOVERED, {"ordinal": step.ordinal}),), + now, + ) + for step in overdrawn: + if await self._fail_exhausted(step, now): + recovered += 1 + failed.append(step.run_id) return recovered, tuple(failed) + async def _fail_exhausted(self, step: StepRecord, now: float) -> bool: + """Fail a run whose step outlived its recovery budget, completely. + + The budget path used to mark the one step and the run FAILED and stop + there: preallocated slots stayed open on a dead run, and child runs + kept working for a parent that no longer existed. This is the same + terminal transition failure takes everywhere else -- open slots + tombstoned, branches closed, the parent told -- in the same lock + order: children, then self, then parent. + + Args: + step: The exhausted step, as phase one saw it. + now: Current time in epoch seconds. + + Returns: + True if this call performed the transition. + """ + exhausted = {"reason": "recovery_budget_exhausted"} + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + locked_children = await self._lock_children(conn, step.run_id) + cursor = await conn.execute( + "SELECT status, parent_run_id, parent_ordinal FROM workflow_runs" + " WHERE run_id = %s FOR UPDATE", + (step.run_id,), + ) + run_row = await cursor.fetchone() + if run_row is None or run_row["status"] in _TERMINAL_RUNS: + return False + cursor = await conn.execute( + "SELECT status, lease_expires_at, recoveries FROM workflow_steps" + " WHERE run_id = %s AND ordinal = %s", + (step.run_id, step.ordinal), + ) + step_row = await cursor.fetchone() + if ( + step_row is None + or step_row["status"] != StepStatus.CLAIMED.value + or step_row["lease_expires_at"] > now + ): + # Renewed, recovered, or failed by a peer between phases; the + # attempt is someone else's to account for. + return False + cursor = await conn.execute( + # The guard repeats the re-check inside the write itself: + # renew_lease is the one writer that takes no run-row lock, + # so a renewal can land in the round trip between the SELECT + # above and this UPDATE -- and an unguarded write would + # acknowledge the worker's lease and fail its run in the + # same instant, discarding in-flight work the store just + # promised another lease_duration to. + "UPDATE workflow_steps SET status = %s, recoveries = %s," + " lease_expires_at = 0, error = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = %s AND status = %s" + " AND lease_expires_at <= %s", + ( + StepStatus.FAILED.value, + step_row["recoveries"] + 1, + _json(exhausted), + now, + step.run_id, + step.ordinal, + StepStatus.CLAIMED.value, + now, + ), + ) + if cursor.rowcount == 0: + # Renewed in the window; the attempt is someone else's to + # account for after all. + return False + cursor = await conn.execute( + "SELECT ordinal FROM workflow_steps WHERE run_id = %s" + " AND ordinal != %s AND NOT (status = ANY(%s)) ORDER BY ordinal", + (step.run_id, step.ordinal, _TERMINAL_STEPS), + ) + tombstoned = [row["ordinal"] for row in await cursor.fetchall()] + if tombstoned: + await conn.execute( + "UPDATE workflow_steps SET status = %s, updated_at = %s" + " WHERE run_id = %s AND ordinal = ANY(%s)", + (StepStatus.CANCELLED.value, now, step.run_id, tombstoned), + ) + await conn.execute( + "UPDATE workflow_runs SET status = %s, error = %s, updated_at = %s" + " WHERE run_id = %s", + (RunStatus.FAILED.value, _json(exhausted), now, step.run_id), + ) + await self._append_events( + conn, + step.run_id, + ( + *( + (HistoryEventType.STEP_TOMBSTONED, {"ordinal": ordinal}) + for ordinal in tombstoned + ), + (HistoryEventType.RUN_FAILED, dict(exhausted)), + ), + now, + ) + await self._close_children(conn, now, locked_children) + if run_row["parent_run_id"] is not None: + await self._apply_arrival( + conn, + run_row["parent_run_id"], + run_row["parent_ordinal"], + { + "run_id": step.run_id, + "status": RunStatus.FAILED.value, + "result": None, + "error": dict(exhausted), + }, + step.run_id, + now, + ) + return True + async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: """List runs matching a query, newest first. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 19a3ab7a87f..08ec5b3b4bf 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -1790,8 +1790,14 @@ def _close_children(self, run_id: str, now: float) -> None: if ( child.parent_run_id != run_id or child.parent_close == "abandon" + or child.cancel_requested or child.status in TERMINAL_RUN_STATUSES ): + # cancel_requested is durable and monotonic, so an + # already-marked child has nothing left to write -- and + # skipping it keeps a second close (a failed run later + # skip-completed by an operator) from appending a duplicate + # cancel-requested event to its history. continue self._runs[child.run_id] = dataclasses.replace( child, @@ -1922,6 +1928,27 @@ async def skip_step(self, run_id: str, now: float) -> bool: updated_at=now, ) self._append_events(run_id, events, now) + if not open_left: + # Completed by an operator's decision is still + # completed: branches are told to stop, and a parent + # joined on this run hears it finished instead of + # waiting forever on a run that no longer will. + self._close_children(run_id, now) + if run.parent_run_id is not None and ( + run.parent_ordinal is not None + ): + self._apply_arrival( + run.parent_run_id, + run.parent_ordinal, + { + "run_id": run_id, + "status": RunStatus.COMPLETED.value, + "result": None, + "error": None, + }, + run_id, + now, + ) return True return False @@ -2003,11 +2030,17 @@ def _restore_tombstoned(steps: list[StepRecord], now: float) -> list[int]: restored: list[int] = [] for index, step in enumerate(steps): if step.status is StepStatus.CANCELLED: + # A wait or join comes back as what it was -- BLOCKED, with + # its arrival count and deadline intact -- never as READY: + # restored-as-READY it would run immediately with a missing + # or partial payload. Plain slots keep their own due_at too, + # so a restored delay still waits out its delay instead of + # firing the moment an operator retries. + waiting = step.wait_key is not None steps[index] = dataclasses.replace( step, - status=StepStatus.READY, + status=StepStatus.BLOCKED if waiting else StepStatus.READY, attempts=0, - due_at=now, lease_expires_at=0.0, error=None, updated_at=now, @@ -2079,12 +2112,49 @@ async def recover_orphans( error={"reason": "recovery_budget_exhausted"}, updated_at=now, ) + # Replaced from the CURRENT record, not the loop's + # snapshot: when a parent and its child both exhaust + # in one pass, the parent's cascade has already set + # cancel_requested on this child, and a snapshot + # rebuild would silently revert it -- letting a later + # retry revive a branch whose parent is dead. The SQL + # stores update columns in place and keep the flag; + # this is the same semantics. self._runs[run.run_id] = dataclasses.replace( - run, + self._runs[run.run_id], status=RunStatus.FAILED, error={"reason": "recovery_budget_exhausted"}, updated_at=now, ) + # Exhaustion is a failure, and failure closes the run + # completely: open slots tombstoned so a dead run + # surfaces no wake times, branches told to stop so + # children do not keep working for a parent that no + # longer exists. + tombstoned = [] + for other in list(steps): + if ( + other.ordinal != step.ordinal + and other.status not in TERMINAL_STEP_STATUSES + ): + steps[other.ordinal] = dataclasses.replace( + other, + status=StepStatus.CANCELLED, + updated_at=now, + ) + tombstoned.append(other.ordinal) + self._append_events( + run.run_id, + tuple( + ( + HistoryEventType.STEP_TOMBSTONED, + {"ordinal": ordinal}, + ) + for ordinal in tombstoned + ), + now, + ) + self._close_children(run.run_id, now) failed.append(run.run_id) if run.parent_run_id is not None and ( run.parent_ordinal is not None @@ -2111,25 +2181,25 @@ async def recover_orphans( ), now, ) - else: - steps[step.ordinal] = dataclasses.replace( - step, - status=StepStatus.RECOVERY_WAIT, - recoveries=step.recoveries + 1, - due_at=now, - lease_expires_at=0.0, - updated_at=now, - ) - self._append_events( - run.run_id, + break + steps[step.ordinal] = dataclasses.replace( + step, + status=StepStatus.RECOVERY_WAIT, + recoveries=step.recoveries + 1, + due_at=now, + lease_expires_at=0.0, + updated_at=now, + ) + self._append_events( + run.run_id, + ( ( - ( - HistoryEventType.STEP_RECOVERED, - {"ordinal": step.ordinal}, - ), + HistoryEventType.STEP_RECOVERED, + {"ordinal": step.ordinal}, ), - now, - ) + ), + now, + ) return recovered, tuple(failed) async def list_runs(self, query: RunQuery) -> tuple[RunRecord, ...]: @@ -2913,6 +2983,7 @@ def _close_children_sql(self, run_id: str, now: float) -> None: closing = self._db.execute( "UPDATE workflow_runs SET cancel_requested = 1, status = ?," " updated_at = ? WHERE parent_run_id = ? AND parent_close <> 'abandon'" + " AND NOT cancel_requested" f" AND status NOT IN ({','.join('?' * len(terminal))})" " RETURNING run_id", (RunStatus.CANCELLING.value, now, run_id, *terminal), @@ -4330,13 +4401,19 @@ def work() -> bool: # "from there" is a lie unless they come back. Nothing # independently cancelled can be in a run these # actions accept (see MemoryRunStore._restore_tombstoned). + # Waits and joins come back BLOCKED with their + # arrival counts and deadlines intact, never READY -- + # restored-as-READY they would run immediately with a + # missing or partial payload. Plain slots keep their + # own due_at, so a restored delay still waits. self._db.execute( - "UPDATE workflow_steps SET status = ?, attempts = 0," - " due_at = ?, lease_expires_at = 0, error = NULL," + "UPDATE workflow_steps SET status = CASE WHEN" + " wait_key IS NULL THEN ? ELSE ? END, attempts = 0," + " lease_expires_at = 0, error = NULL," " updated_at = ? WHERE run_id = ? AND status = ?", ( StepStatus.READY.value, - now, + StepStatus.BLOCKED.value, now, run_id, StepStatus.CANCELLED.value, @@ -4370,6 +4447,30 @@ def work() -> bool: if not open_left: events.append((HistoryEventType.RUN_COMPLETED, {})) self._append_events(run_id, events, now) + if not open_left: + # Completed by an operator's decision is still + # completed: branches are told to stop, and a parent + # joined on this run hears it finished instead of + # waiting forever. + self._close_children_sql(run_id, now) + parent = self._db.execute( + "SELECT parent_run_id, parent_ordinal FROM" + " workflow_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if parent is not None and (parent["parent_run_id"] is not None): + self._apply_arrival_sql( + parent["parent_run_id"], + parent["parent_ordinal"], + { + "run_id": run_id, + "status": RunStatus.COMPLETED.value, + "result": None, + "error": None, + }, + run_id, + now, + ) self._db.execute("COMMIT") except BaseException: self._db.execute("ROLLBACK") @@ -4434,13 +4535,19 @@ def work() -> bool: # "from there" is a lie unless they come back. Nothing # independently cancelled can be in a run these # actions accept (see MemoryRunStore._restore_tombstoned). + # Waits and joins come back BLOCKED with their + # arrival counts and deadlines intact, never READY -- + # restored-as-READY they would run immediately with a + # missing or partial payload. Plain slots keep their + # own due_at, so a restored delay still waits. self._db.execute( - "UPDATE workflow_steps SET status = ?, attempts = 0," - " due_at = ?, lease_expires_at = 0, error = NULL," + "UPDATE workflow_steps SET status = CASE WHEN" + " wait_key IS NULL THEN ? ELSE ? END, attempts = 0," + " lease_expires_at = 0, error = NULL," " updated_at = ? WHERE run_id = ? AND status = ?", ( StepStatus.READY.value, - now, + StepStatus.BLOCKED.value, now, run_id, StepStatus.CANCELLED.value, @@ -4580,6 +4687,43 @@ def work(): step.run_id, ), ) + # Exhaustion is a failure, and failure closes the + # run completely: open slots tombstoned, branches + # told to stop. + step_terminal = tuple( + s.value for s in TERMINAL_STEP_STATUSES + ) + open_rows = self._db.execute( + "SELECT ordinal FROM workflow_steps" + " WHERE run_id = ? AND ordinal != ? AND status" + " NOT IN" + f" ({','.join('?' * len(step_terminal))})", + (step.run_id, step.ordinal, *step_terminal), + ).fetchall() + for open_row in open_rows: + self._db.execute( + "UPDATE workflow_steps SET status = ?," + " updated_at = ? WHERE run_id = ?" + " AND ordinal = ?", + ( + StepStatus.CANCELLED.value, + now, + step.run_id, + open_row["ordinal"], + ), + ) + self._append_events( + step.run_id, + tuple( + ( + HistoryEventType.STEP_TOMBSTONED, + {"ordinal": open_row["ordinal"]}, + ) + for open_row in open_rows + ), + now, + ) + self._close_children_sql(step.run_id, now) failed.append(step.run_id) parent = self._db.execute( "SELECT parent_run_id, parent_ordinal FROM" diff --git a/tests/units/workflow/test_postgres.py b/tests/units/workflow/test_postgres.py index bff0cada5a2..334d63cebb4 100644 --- a/tests/units/workflow/test_postgres.py +++ b/tests/units/workflow/test_postgres.py @@ -423,3 +423,33 @@ def _pg_step(run_id: str, ordinal: int = 0, **over) -> StepRecord: } fields.update(over) return StepRecord(**fields) + + +async def test_twelve_workers_can_initialize_a_fresh_schema_together(): + """First deploy of a fleet: every worker races the same CREATE statements. + + IF NOT EXISTS does not make concurrent DDL safe -- each CREATE TABLE also + inserts the table's composite type, and two backends that both saw "not + exists" race on pg_type. Twelve fresh workers produced one winner and + eleven UniqueViolations, so a fleet's first deploy crash-looped everyone + but one. The advisory lock serializes initializers; this drives twelve + concurrent stores at one brand-new schema and requires them all to come + up. + """ + from reflex.workflow.postgres import PostgresRunStore + + schema = f"initrace_{uuid.uuid4().hex[:12]}" + stores = [ + PostgresRunStore(POSTGRES_URL, schema=schema, min_size=0, max_size=2) + for _ in range(12) + ] + try: + outcomes = await asyncio.gather( + *(store.epoch_time() for store in stores), return_exceptions=True + ) + errors = [o for o in outcomes if isinstance(o, BaseException)] + assert not errors, f"{len(errors)} of 12 initializers failed: {errors[:2]}" + finally: + for store in stores: + await store.close() + stores[0].drop_schema() From 71cdc92389e94ab820072de50f4248d9d06dd1af Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Mon, 24 Aug 2026 12:38:14 -0700 Subject: [PATCH 114/121] workflows: one validation semantics at every boundary Ticket 1 of the standalone-GA plan: Python starts, HTTP starts, webhooks, signal deliveries, and worker dispatch now validate arguments through one module (reflex.workflow.validation), with one rule for where errors land -- boundaries refuse before anything is written, dispatch suspends without consuming attempts. What each boundary got: - Webhooks validated the declared model and then forwarded the RAW payload, throwing the validation away: no coercions, no defaults. The canonical form goes onward now. An object that cannot fill a multi-parameter root is a 400 naming the arguments -- absent fields used to arrive as None and fail inside the handler, on a run that already existed. And a single parameter annotated with a type receives its same-named field of an object payload rather than the whole event dict: def on_paid(self, id: str) received the entire payload as `id`, silently -- one of this repo's own tests pinned that corruption as expected behavior. - Python starts refuse mistyped or missing arguments at the call site, where the stack frame shows which caller is wrong, instead of admitting a run whose first dispatch suspends with the same message. - Signal channels are compiled onto the definition. Delivering to a channel the workflow does not declare raises at the sender, naming the declared channels -- a typo'd channel used to buffer forever, silently. A channel's declared model is enforced on every route in, including ChannelDelivery built without Signal.__call__ (approval redemptions), and the canonical form is what goes onward. Duplicate channel names refuse to compile. - A schedule root with required parameters refuses to compile: a schedule fires with no caller, so every occurrence would admit and immediately suspend. Authoring errors surface at compile, not at 2am. - Dispatch extends the incompatible_payload gate to types: a recorded value the redeployed code's hints no longer fit suspends the run NEEDS_ATTENTION with the argument named and zero attempts consumed -- retrying cannot change what the code declares. The new conformance check for payload-less signals found a Postgres divergence in passing: workflow_inbox.payload and workflow_substeps .payload are NOT NULL, and both inserts passed Python None as SQL NULL, so a signal with no payload -- an "approved" ping -- and a substep whose call returned nothing crashed on Postgres alone. None is the JSON value null, not an absent column; both inserts now say so, and the check pins all three stores. --- news/workflow-boundary-validation.md | 5 + reflex/workflow/CONTRACT.md | 21 +- reflex/workflow/api.py | 32 +- reflex/workflow/conformance.py | 29 ++ reflex/workflow/definition.py | 47 +++ reflex/workflow/ingress.py | 47 ++- reflex/workflow/kernel.py | 72 +++- reflex/workflow/postgres.py | 13 +- reflex/workflow/store.py | 6 +- reflex/workflow/validation.py | 121 ++++++ tests/units/workflow/test_ingress.py | 5 +- tests/units/workflow/test_validation.py | 518 ++++++++++++++++++++++++ 12 files changed, 869 insertions(+), 47 deletions(-) create mode 100644 news/workflow-boundary-validation.md create mode 100644 reflex/workflow/validation.py create mode 100644 tests/units/workflow/test_validation.py diff --git a/news/workflow-boundary-validation.md b/news/workflow-boundary-validation.md new file mode 100644 index 00000000000..f74f2abfd38 --- /dev/null +++ b/news/workflow-boundary-validation.md @@ -0,0 +1,5 @@ +Every boundary that accepts run data now validates it the same way, through one module (`reflex.workflow.validation`). A boundary refuses before anything is written: an invalid webhook or HTTP payload gets a 400 and creates zero runs, an invalid Python start or signal delivery raises at the call site. Dispatch — which judges payloads recorded before a redeploy changed the code — suspends the run `NEEDS_ATTENTION` instead, naming the argument, and never consumes retry attempts. + +A webhook's declared `model=` is no longer validate-and-discard: what goes onward is the canonical form, with coercions applied and defaults filled. A single root parameter annotated with a type receives its field of an object payload rather than the whole event dict — `def on_paid(self, id: str)` used to receive the entire payload as `id`, silently. Multi-parameter roots refuse objects that cannot fill them instead of defaulting absent fields to None. + +Signal channels are compiled onto the workflow definition: delivering to a channel the workflow does not declare is rejected at the sender (a typo used to buffer forever), duplicate channel names refuse to compile, and a channel's declared model is enforced on every route in — including deliveries built without `Signal.__call__`, such as approval-token redemptions. A schedule root with required parameters now refuses to compile, since every occurrence would admit and immediately suspend. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 5388e38e762..be7f21e2c1d 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -71,6 +71,25 @@ Substep results (`rx.step`) are the deliberate exception: each records in its **own** transaction the moment the callable returns, because their purpose is to survive a crash that prevents the attempt from ever committing. +### One validation semantics at every boundary + +Python starts, HTTP starts, webhooks, signal deliveries, and worker dispatch +all validate arguments the same way (`reflex.workflow.validation`): required +parameters must bind, and supplied values must satisfy the declared types. +A **boundary** refuses before anything is written — an invalid webhook or +HTTP payload is a 400 and zero runs; an invalid Python start or signal is an +exception at the call site. **Dispatch** — which judges payloads recorded +before a redeploy changed the code — suspends the run `NEEDS_ATTENTION` +instead (§8, `incompatible_payload`) and never consumes retry attempts. +A payload model declared on a webhook (`model=...`) or channel +(`rx.Signal(Model)`) is enforced on every route in, including deliveries +built without `Signal.__call__`, and what goes onward is the **canonical** +form — coercions applied, defaults filled — not the raw input. A single +root parameter receives the whole payload when the payload satisfies its +declared type, and otherwise its same-named field of an object payload; an +unknown channel is rejected at the sender when the workflow is registered +in the sending process. + ## 2. When handlers re-execute A handler runs more than once in exactly two situations, both bounded: @@ -363,7 +382,7 @@ reason rather than on message text. |---|---|---| | the run's workflow class is not registered in this process | `unknown_workflow` | step and run `NEEDS_ATTENTION`; re-register and `resume(run)`. Not a failure: a worker that does not serve a workflow must not decide that workflow's fate | | the next step's handler no longer exists | `unknown_handler` | step and run `NEEDS_ATTENTION`; restore the handler and `resume(run)`, or cancel | -| a recorded payload carries arguments the handler no longer accepts | `incompatible_payload` | step and run `NEEDS_ATTENTION`, naming the arguments; restore the parameters or cancel | +| a recorded payload carries arguments the handler no longer accepts — names it no longer declares, required names it cannot fill, or values its declared types no longer fit | `incompatible_payload` | step and run `NEEDS_ATTENTION`, naming the arguments; restore the parameters (or their types) or cancel. Never consumes retry attempts: retrying cannot change what the code declares | | a run allocated more steps than `WorkflowConfig.max_steps` | `max_steps_exceeded` | the committing step succeeds, every other open slot is tombstoned, run `FAILED`. The bound is on a runaway loop, so it fails rather than suspending for a person to approve more of the same | | a step's lease lapsed more times than `max_recoveries` | `recovery_budget_exhausted` | run `FAILED` — the complete terminal transition, identical to any other failure: remaining open slots are tombstoned, children are told to stop, and the parent's join hears one `FAILED` arrival. Infrastructure recoveries are free of the retry budget precisely so they can be bounded separately; the bound is what stops a poison step cycling forever | | an error's `details` cannot be serialized | — | the reason is preserved and the details are replaced by `{"unserializable": repr(...)}`. Losing the payload never turns a failure into a crash | diff --git a/reflex/workflow/api.py b/reflex/workflow/api.py index 50c6badea0e..a4e20735a36 100644 --- a/reflex/workflow/api.py +++ b/reflex/workflow/api.py @@ -25,6 +25,7 @@ from reflex.workflow.definition import unbound_params from reflex.workflow.records import attempts_made +from reflex.workflow.validation import mistyped_args if TYPE_CHECKING: from collections.abc import Callable, Coroutine @@ -67,35 +68,6 @@ def _authorized(request: Request, token: str) -> bool: return hmac.compare_digest(presented, token) -def _mistyped_args(handler: Any, args: dict[str, Any]) -> list[str]: - """Check supplied arguments against the handler's declared types. - - Args: - handler: The resolved handler definition. - args: The caller-supplied arguments. - - Returns: - One message per argument that cannot validate, empty when all fit. - """ - from pydantic import TypeAdapter, ValidationError - - problems: list[str] = [] - for name, value in args.items(): - hint = handler.type_hints.get(name) - if hint is None: - continue - try: - TypeAdapter(hint).validate_python(value) - except ValidationError: - problems.append( - f"{name!r} does not validate as {getattr(hint, '__name__', hint)}" - ) - except Exception: - # An exotic hint pydantic cannot adapt is not the caller's fault. - pass - return problems - - def start_endpoint( runtime: WorkflowRuntime, token: str ) -> Callable[[Request], Coroutine[Any, Any, JSONResponse]]: @@ -170,7 +142,7 @@ async def endpoint(request: Request) -> JSONResponse: {"error": "args must be a JSON object"}, status_code=400 ) args = raw_args or {} - mistyped = _mistyped_args(handler, args) + mistyped = mistyped_args(handler, args) if mistyped: # Admitting a payload the handler's signature refuses creates a # run whose first attempt can only raise; the caller gets a 202 diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index fb3eb0e3746..b32d9f54673 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -1681,6 +1681,34 @@ async def check_retry_after_exhaustion_restores_waits_as_waits( ) +async def check_none_is_a_legal_payload_everywhere(store: RunStore) -> None: + """None is the JSON value null, not an absent column. + + A signal with no payload -- an "approved" ping -- and a journaled + substep whose call returned nothing are both everyday shapes, and one + store refusing what the others accept is a divergence someone only + finds after migrating. + """ + await store.admit( + make_run(), + make_step(status=StepStatus.BLOCKED, wait_key="sig:ping", due_at=0.0), + _ADMITTED, + ) + assert await store.deliver("run1", "sig:ping", "d1", None, NOW) == "resolved" + steps = await store.get_steps("run1") + assert steps[0].args["__payload__"] is None + + await store.admit(make_run("sub1"), make_step("sub1"), _ADMITTED) + claim = await store.claim_next(NOW, lease_duration=LEASE) + while claim is not None and claim.run.run_id != "sub1": + # run1's resolved continuation is claimable too, and stores order + # their frontiers differently. + claim = await store.claim_next(NOW, lease_duration=LEASE) + assert claim is not None + assert await store.record_substep("sub1", 0, claim.step.epoch, "notify", None, NOW) + assert await store.get_substeps("sub1", 0) == {"notify": None} + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1706,6 +1734,7 @@ async def check_retry_after_exhaustion_restores_waits_as_waits( check_list_children_finds_a_joins_branches, check_claims_respect_queue_boundaries, check_substeps_record_once_and_fence_stale_writers, + check_none_is_a_legal_payload_everywhere, check_recovery_respects_a_live_lease, check_a_terminal_run_refuses_further_control, check_finalize_delivers_a_childs_arrival, diff --git a/reflex/workflow/definition.py b/reflex/workflow/definition.py index 241a3d55b28..0676c163a5c 100644 --- a/reflex/workflow/definition.py +++ b/reflex/workflow/definition.py @@ -24,6 +24,7 @@ RateLimit, Retry, ScheduleTrigger, + Signal, Singleton, Throttle, Trigger, @@ -117,6 +118,7 @@ class WorkflowDefinition: handler_ids_by_name: Map from Python method name to handler id. roots: Handler ids that declare a trigger and may start runs. fields: Run-state field schemas in declaration order. + channels: Signal channels the class declares, keyed by channel name. """ workflow_id: str @@ -129,6 +131,7 @@ class WorkflowDefinition: handler_ids_by_name: Mapping[str, str] roots: tuple[str, ...] fields: tuple[FieldSchema, ...] + channels: Mapping[str, Signal] def _error(workflow_cls: type, msg: str) -> WorkflowDefinitionError: @@ -675,6 +678,35 @@ def _compute_digest( return hashlib.sha256(payload.encode()).hexdigest() +def channels_of(state_cls: type) -> dict[str, Signal]: + """Collect the signal channels a workflow class declares. + + Args: + state_cls: The workflow class. + + Returns: + Declared channels keyed by channel name. + + Raises: + WorkflowDefinitionError: If two declarations share one channel name. + """ + channels: dict[str, Signal] = {} + for klass in reversed(state_cls.__mro__): + for value in vars(klass).values(): + if not isinstance(value, Signal): + continue + existing = channels.get(value.name) + if existing is not None and existing is not value: + msg = ( + f"Workflow {state_cls.__name__}: two rx.Signal declarations " + f"share the channel name {value.name!r}; a delivery could " + "not say which one it means." + ) + raise WorkflowDefinitionError(msg) + channels[value.name] = value + return channels + + def unbound_params(handler: HandlerDefinition, supplied: set[str]) -> set[str]: """Parameters a handler requires that a recorded payload cannot fill. @@ -731,6 +763,20 @@ def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: durable_names = frozenset(defn.name for defn in handlers.values()) for defn in handlers.values(): _validate_handler_body(workflow_cls, defn, durable_names) + for defn in handlers.values(): + if isinstance(defn.trigger, ScheduleTrigger): + unfillable = sorted(unbound_params(defn, set())) + if unfillable: + # A schedule fires with no caller to supply arguments, so a + # required parameter here means every occurrence would admit + # and immediately suspend. That is an authoring error, and + # authoring errors surface at compile, not at 2am. + raise _error( + workflow_cls, + f"schedule root {defn.name!r} requires parameters " + f"{unfillable}, which a schedule occurrence cannot " + "supply; give them defaults.", + ) roots = tuple( defn.id for defn in sorted(handlers.values(), key=lambda d: d.id) @@ -757,4 +803,5 @@ def compile_workflow(workflow_cls: type[BaseState]) -> WorkflowDefinition: handler_ids_by_name={defn.name: defn.id for defn in handlers.values()}, roots=roots, fields=fields, + channels=channels_of(workflow_cls), ) diff --git a/reflex/workflow/ingress.py b/reflex/workflow/ingress.py index 168d965047e..e1aa3b3a4fd 100644 --- a/reflex/workflow/ingress.py +++ b/reflex/workflow/ingress.py @@ -12,11 +12,13 @@ import json from typing import TYPE_CHECKING, Any -from pydantic import TypeAdapter, ValidationError +from pydantic import ValidationError from reflex_base.utils import console from starlette.requests import Request from starlette.responses import JSONResponse +from reflex.workflow.validation import canonical_payload, missing_args, mistyped_args + if TYPE_CHECKING: from collections.abc import Callable, Coroutine, Mapping @@ -185,18 +187,35 @@ def _legacy_dedupe_keys(trigger: WebhookTrigger, payload: Any) -> tuple[str, ... def _root_args(handler: HandlerDefinition, payload: Any) -> dict[str, Any]: """Map a decoded payload onto the root handler's parameters. + One parameter can mean two things: "hand me the event object" or "hand + me this one field". The declared type decides -- if the whole payload + satisfies the parameter's hint it is passed whole, and otherwise a dict + payload carrying the parameter's name is unpacked to that field. Before + types were consulted, ``def on_paid(self, id: str)`` received the entire + event dict as ``id`` and nothing ever said so. + Args: handler: The root handler definition. payload: The decoded request payload. Returns: - The keyword arguments to start the root with. + The keyword arguments to start the root with; absent object fields + are omitted rather than filled with None, so the boundary's + missing-argument check can see them. """ if not handler.params: return {} if len(handler.params) == 1: - return {handler.params[0]: payload} - return {name: payload.get(name) for name in handler.params} + name = handler.params[0] + if ( + mistyped_args(handler, {name: payload}) + and isinstance(payload, dict) + and name in payload + and not mistyped_args(handler, {name: payload[name]}) + ): + return {name: payload[name]} + return {name: payload} + return {name: payload[name] for name in handler.params if name in payload} def webhook_endpoint( @@ -255,7 +274,10 @@ async def endpoint(request: Request) -> JSONResponse: if route.trigger.model is not None: try: - TypeAdapter(route.trigger.model).validate_python(payload) + # The canonical form -- coercions applied, defaults filled -- + # is what goes onward. Validating and then passing the raw + # payload threw the validation away. + payload = canonical_payload(route.trigger.model, payload) except ValidationError: return JSONResponse( {"error": "payload does not match the declared model"}, @@ -287,6 +309,21 @@ async def endpoint(request: Request) -> JSONResponse: status_code=400, ) args = _root_args(route.handler, payload) + absent = missing_args(route.handler, args) + wrong = mistyped_args(route.handler, args) + if absent or wrong: + # Refused before any run exists. Admitting would produce a run + # that suspends on its first step for a reason the provider is + # never told; a 400 naming the arguments is retryable after the + # sender fixes their payload, and creates nothing until then. + faults = [ + *(f"missing required argument {name!r}" for name in absent), + *wrong, + ] + return JSONResponse( + {"error": f"payload does not fit the handler: {'; '.join(faults)}"}, + status_code=400, + ) result = await runtime.kernel.start( spec(**args) if args else spec, request_key=request_key, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 9def63b7ce1..bf75c283903 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -20,7 +20,7 @@ import uuid from typing import TYPE_CHECKING, Any, Final -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from reflex_base.event.processor.base_state_processor import _transform_event_payload from reflex_base.utils import console from reflex_base.utils.exceptions import WorkflowDefinitionError, WorkflowRuntimeError @@ -68,6 +68,7 @@ StepCompletion, _child_admission_events, ) +from reflex.workflow.validation import canonical_payload, mistyped_args if TYPE_CHECKING: from collections.abc import Callable, Iterable, Mapping @@ -712,6 +713,21 @@ async def start( does not match the admitting ingress. """ defn, handler, payload = self._resolve_target(target) + unbound = sorted(unbound_params(handler, set(payload))) + problems = mistyped_args(handler, payload) + if unbound or problems: + # Refused before anything is written: admitting this run would + # only postpone the same message to its first dispatch, minus the + # stack frame that shows the caller which call site is wrong. + faults = [ + *(f"missing required argument {name!r}" for name in unbound), + *problems, + ] + msg = ( + f"Cannot start {handler.id!r} of {defn.workflow_id!r}: " + f"{'; '.join(faults)}." + ) + raise WorkflowDefinitionError(msg) declared = getattr(handler.trigger, "kind", None) # A handler without a trigger is a mid-flow step, not a root; no # ingress -- and no test privilege -- makes it startable. @@ -957,12 +973,48 @@ async def signal( Returns: What the store did with the delivery. + + Raises: + WorkflowDefinitionError: If the run's workflow is registered here + and does not declare the channel, or the payload does not + satisfy the channel's declared model. """ + payload = delivery.payload + run = await self._store.get_run(run_id) + defn = self._definitions.get(run.workflow_id) if run is not None else None + if defn is not None: + channel = defn.channels.get(delivery.channel) + if channel is None: + # A typo'd channel would buffer forever: the store cannot + # know the name is wrong, so the sender must hear it here, + # from the process that knows what the workflow declares. + declared = sorted(defn.channels) or [""] + msg = ( + f"Workflow {defn.workflow_id!r} declares no channel " + f"{delivery.channel!r}; declared channels: " + f"{', '.join(declared)}." + ) + raise WorkflowDefinitionError(msg) + if channel.model is not None: + # Every route into a channel validates the same way -- + # including deliveries built without Signal.__call__, such + # as approval-token redemptions -- and what goes onward is + # the canonical form the model promises, not the raw input. + try: + payload = canonical_payload(channel.model, payload) + except ValidationError as error: + msg = ( + f"Channel {delivery.channel!r} of " + f"{defn.workflow_id!r} expects " + f"{channel.model.__name__}: {error.error_count()} " + "validation error(s)." + ) + raise WorkflowDefinitionError(msg) from error disposition = await self._store.deliver( run_id, f"sig:{delivery.channel}", key or uuid.uuid4().hex, - to_run_data({"value": delivery.payload})["value"], + to_run_data({"value": payload})["value"], self._clock(), ) if disposition == "resolved": @@ -2195,6 +2247,22 @@ def _incompatible_reason( "or cancel the run." ), } + wrong = mistyped_args(handler, claim.step.args) + if wrong: + # A recorded value that no longer fits the parameter's type is a + # redeploy problem exactly like a renamed parameter: the payload + # was valid when it was recorded and the code changed underneath + # it. Dispatching anyway raises from inside the handler and + # burns retry attempts on a state no retry can change. + return { + "reason": "incompatible_payload", + "handler_id": handler.id, + "detail": ( + f"Recorded payload no longer fits handler {handler.id!r}: " + f"{'; '.join(wrong)}. Restore the parameter types, or " + "cancel the run." + ), + } return None @staticmethod diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 1cc0964a69c..e97ee8a59f9 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -1183,7 +1183,7 @@ async def deliver( run_id: str, wait_key: str, dedupe_key: str, - payload: dict[str, Any], + payload: Any, now: float, ) -> DeliveryDisposition: """Deliver a payload to a run, resolving its wait or buffering it. @@ -1249,7 +1249,7 @@ async def deliver( wait_key, dedupe_key, await self._next_inbox_seq(conn, run_id), - _json(payload), + Jsonb(payload), "CONSUMED" if resolves else "PENDING", now, ), @@ -1457,7 +1457,7 @@ async def _apply_arrival( wait_key, dedupe_key, await self._next_inbox_seq(conn, run_id), - _json(payload), + Jsonb(payload), now, ), ) @@ -1553,7 +1553,7 @@ async def record_arrival( wait_key, dedupe_key, await self._next_inbox_seq(conn, run_id), - _json(payload), + Jsonb(payload), now, ), ) @@ -2375,7 +2375,10 @@ async def record_substep( "INSERT INTO workflow_substeps" " (run_id, ordinal, key, payload, created_at)" " VALUES (%s, %s, %s, %s, %s) ON CONFLICT DO NOTHING", - (run_id, ordinal, key, _json(payload), now), + # Jsonb, not _json: a substep that recorded None recorded + # the JSON value null, and the column is NOT NULL because + # every journal entry has a payload. + (run_id, ordinal, key, Jsonb(payload), now), ) if cursor.rowcount: await self._append_events( diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 08ec5b3b4bf..f055fa5ea54 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -333,7 +333,7 @@ async def deliver( run_id: str, wait_key: str, dedupe_key: str, - payload: dict[str, Any], + payload: Any, now: float, ) -> DeliveryDisposition: """Deliver a payload to a run, resolving its wait or buffering it. @@ -1286,7 +1286,7 @@ async def deliver( run_id: str, wait_key: str, dedupe_key: str, - payload: dict[str, Any], + payload: Any, now: float, ) -> DeliveryDisposition: """Deliver a payload to a run, resolving its wait or buffering it. @@ -3651,7 +3651,7 @@ async def deliver( run_id: str, wait_key: str, dedupe_key: str, - payload: dict[str, Any], + payload: Any, now: float, ) -> DeliveryDisposition: """Deliver a payload to a run, resolving its wait or buffering it. diff --git a/reflex/workflow/validation.py b/reflex/workflow/validation.py new file mode 100644 index 00000000000..a7bfb772ba7 --- /dev/null +++ b/reflex/workflow/validation.py @@ -0,0 +1,121 @@ +"""One validation semantics for every boundary that accepts run data. + +Python starts, HTTP starts, webhooks, signal deliveries, and worker dispatch +all answer the same two questions -- are these the arguments the handler +declares, and do the values fit its types -- and they used to answer them +in different places with different strictness. A payload one boundary +admitted could then suspend at dispatch, or worse, run wrong. Every boundary +now asks this module, so "does this payload fit" has exactly one answer. + +The rule for where errors land: a boundary refuses *before* anything is +written, so an invalid payload creates nothing; dispatch -- which judges +payloads recorded before a redeploy changed the code -- suspends the run for +an operator instead, and never consumes retry attempts, because retrying +cannot change what the code declares. +""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Any + +from reflex.workflow.definition import channels_of, unbound_params + +if TYPE_CHECKING: + from reflex.workflow.definition import HandlerDefinition + +__all__ = [ + "canonical_payload", + "channels_of", + "missing_args", + "mistyped_args", +] + + +@functools.lru_cache(maxsize=1024) +def _adapter(hint: Any): + """Build (once) the pydantic adapter for a type hint. + + Args: + hint: The handler parameter's declared type. + + Returns: + The cached ``TypeAdapter``, or None for a hint pydantic cannot adapt. + """ + from pydantic import TypeAdapter + + try: + return TypeAdapter(hint) + except Exception: + # An exotic hint is the author's business, never the caller's fault. + return None + + +def mistyped_args(handler: HandlerDefinition, args: dict[str, Any]) -> list[str]: + """Check supplied arguments against the handler's declared types. + + Args: + handler: The resolved handler definition. + args: The caller-supplied arguments. + + Returns: + One message per argument that cannot validate, empty when all fit. + """ + from pydantic import ValidationError + + problems: list[str] = [] + for name, value in args.items(): + if name.startswith("__"): + continue + hint = handler.type_hints.get(name) + adapter = _adapter(hint) if hint is not None else None + if adapter is None: + continue + try: + adapter.validate_python(value) + except ValidationError: + problems.append( + f"{name!r} does not validate as {getattr(hint, '__name__', hint)}" + ) + except Exception: + pass + return problems + + +def missing_args(handler: HandlerDefinition, args: dict[str, Any]) -> list[str]: + """Required parameters the supplied arguments leave unbound. + + Args: + handler: The resolved handler definition. + args: The caller-supplied arguments. + + Returns: + The unbound required parameter names, sorted. + """ + return sorted(unbound_params(handler, set(args))) + + +def canonical_payload(model: type, payload: Any) -> Any: + """Validate a payload against a model and return its canonical form. + + Validating and then passing the *raw* payload onward throws the + validation away: coercions, defaults, and alias resolution never happen, + so the handler receives something subtly different from what the model + promised. The canonical form is the validated object dumped back to + JSON-compatible data -- what the model says the payload *is*. + + Args: + model: The declared payload model. + payload: The decoded payload to validate. + + Returns: + The validated payload in JSON-canonical form. + + Raises: + pydantic.ValidationError: If the payload does not satisfy the model. + """ + from pydantic import TypeAdapter + + adapter = TypeAdapter(model) + validated = adapter.validate_python(payload) + return adapter.dump_python(validated, mode="json") diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py index 458a1124138..f687ed27193 100644 --- a/tests/units/workflow/test_ingress.py +++ b/tests/units/workflow/test_ingress.py @@ -651,5 +651,8 @@ async def test_github_form_encoded_deliveries_are_understood( await runtime.kernel.run_until_idle() snapshot = await runtime.kernel.get_run(response.json()["run_id"]) assert snapshot is not None - assert snapshot.result == {"sha": "abc123"}, "the wrapped JSON is the payload" + assert snapshot.result == "abc123", ( + "the wrapped JSON is the payload, and a parameter annotated str " + "receives its field, not the enclosing object" + ) await runtime.shutdown() diff --git a/tests/units/workflow/test_validation.py b/tests/units/workflow/test_validation.py new file mode 100644 index 00000000000..92a95d69fa7 --- /dev/null +++ b/tests/units/workflow/test_validation.py @@ -0,0 +1,518 @@ +"""One validation semantics at every boundary that accepts run data. + +The acceptance bar, per boundary: an invalid payload is refused before +anything exists (webhook and HTTP -> 400 with zero runs; Python starts and +signals -> an exception at the call site), and dispatch -- judging payloads +recorded before a redeploy -- suspends without consuming retry attempts. +""" + +import hashlib +import hmac as hmac_mod +import json + +import pytest +from pydantic import BaseModel +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import Signal, WorkflowConfig, hmac_signature, manual, webhook +from starlette.applications import Starlette +from starlette.routing import Route +from starlette.testclient import TestClient + +import reflex as rx +from reflex.workflow.definition import channels_of, compile_workflow +from reflex.workflow.ingress import WEBHOOK_ROUTE, webhook_endpoint +from reflex.workflow.records import RunQuery, RunStatus +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.store import MemoryRunStore +from reflex.workflow.testing import WorkflowTestHarness + +SECRET = "whsec_validation" + + +class Shipment(BaseModel): + """A provider event with a coercible field and a defaulted one.""" + + order_id: str + parcels: int + carrier: str = "ups" + + +def _sign(body: bytes) -> str: + """Sign a body the way the HMAC verifier expects. + + Args: + body: The raw request body. + + Returns: + The hex signature. + """ + return hmac_mod.new(SECRET.encode(), body, hashlib.sha256).hexdigest() + + +def _shipping_flow(): + """Build a webhook workflow whose root declares a payload model. + + Returns: + The workflow class. + """ + + class Shipping(rx.State): + __workflow__ = WorkflowConfig(id="validation.shipping") + seen: str = "" + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "shipped", + model=Shipment, + verify=hmac_signature( + secret_env="VALIDATION_WEBHOOK_SECRET", header="X-Signature" + ), + ), + ) + def on_shipped(self, event: dict): + """Record the canonical event. + + Args: + event: The validated payload. + + Returns: + Completion. + """ + return rx.complete(result=event) + + return Shipping + + +async def _webhook_client(monkeypatch, flow): + """Stand up a runtime and a webhook test client for one workflow. + + Args: + monkeypatch: Used to install the webhook secret. + flow: The workflow class to register. + + Returns: + The runtime and the entered test client. + """ + monkeypatch.setenv("VALIDATION_WEBHOOK_SECRET", SECRET) + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(flow) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + return runtime, TestClient(app) + + +async def test_an_invalid_webhook_payload_is_refused_with_zero_runs( + monkeypatch, forked_registration_context +): + """The acceptance bar itself: 400 out, nothing admitted. + + Args: + monkeypatch: Used to install the webhook secret. + forked_registration_context: Isolated state registry. + """ + runtime, client = await _webhook_client(monkeypatch, _shipping_flow()) + with client: + body = json.dumps({"order_id": "o1", "parcels": "not-a-number"}).encode() + response = client.post( + "/_workflow/webhook/shipped", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert response.status_code == 400, response.text + assert await runtime.kernel._store.count_runs(RunQuery()) == 0 # pyright: ignore[reportPrivateUsage] + await runtime.shutdown() + + +async def test_the_validated_canonical_payload_is_what_goes_onward( + monkeypatch, forked_registration_context +): + """Validation must change what the handler receives, not just gatekeep. + + "5" coerces to 5 and the absent carrier fills with its default; a + boundary that validated and then forwarded the raw payload dropped both, + so the handler saw something subtly different from what the model + promised. + + Args: + monkeypatch: Used to install the webhook secret. + forked_registration_context: Isolated state registry. + """ + runtime, client = await _webhook_client(monkeypatch, _shipping_flow()) + with client: + body = json.dumps({"order_id": "o1", "parcels": "5"}).encode() + response = client.post( + "/_workflow/webhook/shipped", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert response.status_code == 202, response.text + run_id = response.json()["run_id"] + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(run_id) + assert snapshot is not None + assert snapshot.result == {"order_id": "o1", "parcels": 5, "carrier": "ups"} + await runtime.shutdown() + + +async def test_a_missing_webhook_argument_is_refused_not_defaulted_to_none( + monkeypatch, forked_registration_context +): + """A multi-parameter root refuses an object that cannot fill it. + + Absent fields used to arrive as None and fail inside the handler, on a + run that already existed and burned attempts on an unfixable payload. + + Args: + monkeypatch: Used to install the webhook secret. + forked_registration_context: Isolated state registry. + """ + + class TwoArg(rx.State): + __workflow__ = WorkflowConfig(id="validation.twoarg") + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "pair", + verify=hmac_signature( + secret_env="VALIDATION_WEBHOOK_SECRET", header="X-Signature" + ), + ), + ) + def on_pair(self, left: str, right: str): + """Take two named fields. + + Args: + left: One field. + right: The other. + + Returns: + Completion. + """ + return rx.complete(result=[left, right]) + + runtime, client = await _webhook_client(monkeypatch, TwoArg) + with client: + body = json.dumps({"left": "only"}).encode() + response = client.post( + "/_workflow/webhook/pair", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert response.status_code == 400, response.text + assert "right" in response.json()["error"] + assert await runtime.kernel._store.count_runs(RunQuery()) == 0 # pyright: ignore[reportPrivateUsage] + await runtime.shutdown() + + +def test_a_single_typed_parameter_receives_its_field_not_the_object( + forked_registration_context, +): + """``def go(self, order_id: str)`` gets the string, never the dict. + + Before types were consulted, the whole event object landed in the lone + parameter and nothing ever said so. + + Args: + forked_registration_context: Isolated state registry. + """ + from reflex.workflow.ingress import _root_args + + class OneField(rx.State): + __workflow__ = WorkflowConfig(id="validation.onefield") + + @rx.event(durable=True, effect="none", trigger=manual()) + def go(self, order_id: str): + """Take one typed field. + + Args: + order_id: The order. + """ + + defn = compile_workflow(OneField) + handler = next(iter(defn.handlers.values())) + assert _root_args(handler, {"order_id": "o1", "noise": 1}) == {"order_id": "o1"} + assert _root_args(handler, "o1") == {"order_id": "o1"} + + +async def test_a_python_start_with_bad_arguments_creates_nothing( + forked_registration_context, +): + """The Python boundary refuses at the call site, like every other. + + Args: + forked_registration_context: Isolated state registry. + """ + + class Typed(rx.State): + __workflow__ = WorkflowConfig(id="validation.typed") + + @rx.event(durable=True, effect="none", trigger=manual()) + def begin(self, count: int): + """Start with a typed argument. + + Args: + count: A number. + + Returns: + Completion. + """ + return rx.complete(result=count) + + store = MemoryRunStore() + async with WorkflowTestHarness(Typed, store=store) as harness: + with pytest.raises(WorkflowDefinitionError, match="count"): + await harness.start( + Typed.begin("not-a-number") # pyright: ignore[reportArgumentType] + ) + with pytest.raises(WorkflowDefinitionError, match="missing required"): + await harness.start(Typed.begin()) + assert await store.count_runs(RunQuery()) == 0 + result = await harness.start(Typed.begin(3)) + assert result.disposition == "started" + + +async def test_an_unknown_channel_is_rejected_at_the_sender( + forked_registration_context, +): + """A typo'd channel must fail the sender, not buffer forever. + + Args: + forked_registration_context: Isolated state registry. + """ + from reflex_base.workflow import ChannelDelivery + + class Waits(rx.State): + __workflow__ = WorkflowConfig(id="validation.waits") + approved = Signal() + + @rx.event(durable=True, effect="none", trigger=manual()) + def begin(self): + """Wait for approval. + + Returns: + The wait. + """ + return rx.wait_for(Waits.approved, then=Waits.done, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def done(self, decision): + """Finish. + + Args: + decision: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result=decision) + + async with WorkflowTestHarness(Waits) as harness: + result = await harness.start(Waits.begin()) + assert result.run_id is not None + await harness.run_until_idle() + with pytest.raises(WorkflowDefinitionError, match="approved"): + await harness.kernel.signal( + result.run_id, ChannelDelivery(channel="aproved", payload=None) + ) + assert ( + await harness.kernel.signal(result.run_id, Waits.approved(None)) + == "resolved" + ) + + +async def test_channel_payloads_are_validated_on_every_route_in( + forked_registration_context, +): + """A raw ChannelDelivery cannot smuggle past the declared model. + + Signal.__call__ validates, but approvals and future HTTP senders build + ChannelDelivery directly; the kernel is where every route converges, so + the kernel enforces the model and forwards the canonical form. + + Args: + forked_registration_context: Isolated state registry. + """ + from reflex_base.workflow import ChannelDelivery + + class Modeled(rx.State): + __workflow__ = WorkflowConfig(id="validation.modeled") + shipped = Signal(Shipment) + + @rx.event(durable=True, effect="none", trigger=manual()) + def begin(self): + """Wait for the shipment. + + Returns: + The wait. + """ + return rx.wait_for(Modeled.shipped, then=Modeled.done, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def done(self, shipment): + """Finish with the delivered payload. + + Args: + shipment: The canonical shipment. + + Returns: + Completion. + """ + return rx.complete(result=shipment) + + async with WorkflowTestHarness(Modeled) as harness: + result = await harness.start(Modeled.begin()) + assert result.run_id is not None + await harness.run_until_idle() + with pytest.raises(WorkflowDefinitionError, match="Shipment"): + await harness.kernel.signal( + result.run_id, + ChannelDelivery(channel="shipped", payload={"parcels": "x"}), + ) + disposition = await harness.kernel.signal( + result.run_id, + ChannelDelivery( + channel="shipped", payload={"order_id": "o1", "parcels": "2"} + ), + ) + assert disposition == "resolved" + await harness.run_until_idle() + snapshot = await harness.get_run(result.run_id) + assert snapshot is not None + assert snapshot.result == {"order_id": "o1", "parcels": 2, "carrier": "ups"} + + +def test_two_declarations_sharing_a_channel_name_refuse_to_compile( + forked_registration_context, +): + """A delivery could not say which declaration it means. + + Args: + forked_registration_context: Isolated state registry. + """ + + class Ambiguous(rx.State): + __workflow__ = WorkflowConfig(id="validation.ambiguous") + first = Signal(name="decision") + second = Signal(name="decision") + + @rx.event(durable=True, effect="none", trigger=manual()) + def begin(self): + """Start.""" + + with pytest.raises(WorkflowDefinitionError, match="decision"): + channels_of(Ambiguous) + + +def test_a_schedule_root_with_required_parameters_refuses_to_compile( + forked_registration_context, +): + """A schedule fires with no caller; every occurrence would suspend. + + Args: + forked_registration_context: Isolated state registry. + """ + from reflex_base.workflow import schedule + + class Nightly(rx.State): + __workflow__ = WorkflowConfig(id="validation.nightly") + + @rx.event(durable=True, effect="none", trigger=schedule("0 3 * * *")) + def run_report(self, region: str): + """Demand an argument no schedule can supply. + + Args: + region: Unfillable. + """ + + with pytest.raises(WorkflowDefinitionError, match="region"): + compile_workflow(Nightly) + + +async def test_a_retyped_parameter_suspends_without_consuming_attempts( + forked_registration_context, +): + """Schema incompatibility is a redeploy problem, never a retry burn. + + The payload was valid when recorded; the code changed underneath it. + Retrying cannot change what the code declares, so the run suspends with + the argument named and its attempt budget untouched. + + Args: + forked_registration_context: Isolated state registry. + """ + store = MemoryRunStore() + + def _first(): + class Retyped(rx.State): + __workflow__ = WorkflowConfig(id="validation.retyped") + + @rx.event(durable=True, effect="none", trigger=manual()) + def begin(self): + """Schedule the finish with an int. + + Returns: + The deferral. + """ + return rx.after("1h", Retyped.finish(5)) + + @rx.event(durable=True, effect="read") + def finish(self, count: int): + """Finish with a count. + + Args: + count: The recorded number. + """ + + return Retyped + + first_cls = _first() + async with WorkflowTestHarness(first_cls, store=store) as harness: + result = await harness.start(first_cls.begin()) + assert result.run_id is not None + run_id, resume_at = result.run_id, harness.now + + class Retyped(rx.State): + __workflow__ = WorkflowConfig(id="validation.retyped") + + @rx.event(durable=True, effect="none", trigger=manual()) + def begin(self): + """Unchanged root. + + Returns: + The deferral. + """ + return rx.after( + "1h", + Retyped.finish(5), # pyright: ignore[reportArgumentType] + ) + + @rx.event(durable=True, effect="read") + def finish(self, count: Shipment): + """Now demand a model where an int was recorded. + + Args: + count: The retyped parameter. + """ + + async with WorkflowTestHarness( + Retyped, store=store, start_time=resume_at + 3600 + ) as harness: + await harness.run_until_idle() + snapshot = await harness.get_run(run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.NEEDS_ATTENTION + assert snapshot.error is not None + assert snapshot.error["reason"] == "incompatible_payload" + assert "count" in snapshot.error["detail"] + steps = await store.get_steps(run_id) + pending = [s for s in steps if s.handler_id.endswith("finish")] + assert pending + assert pending[0].attempts == 0, ( + "suspension must not consume the attempt budget" + ) From 994a33be487a46ac6a72aea0565695fbab3215e3 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Mon, 24 Aug 2026 13:43:20 -0700 Subject: [PATCH 115/121] workflows: reflex workflows serve -- the standalone service Ticket 2 of the standalone-GA plan. One process, no frontend, no rx.App: webhook and approval ingress, the run HTTP API, health and readiness probes, Prometheus metrics, an OpenAPI document, and the worker loop. --ingress-only and --worker-only split the halves for separate scaling against the same store; both keep the probes. The API is POST /runs, GET /runs (workflow/status/label filters), GET /runs/{id}, POST /runs/{id}/signals/{channel}, and POST /runs/{id}/cancel|retry|resume. Authorization is scoped bearer tokens: REFLEX_WORKFLOW_API_TOKEN grants everything as before, and REFLEX_WORKFLOW_API_TOKEN_READ, _START, _SIGNAL, _OPERATE each grant exactly one scope -- no valid token is 401, a valid token without the route's scope is 403, so a dashboard's credential cannot cancel runs and a relay's cannot read them. The existing api.py endpoints take an authorizer callable now (a bare token still works), so embedded and standalone modes share one implementation. HTTP signals go through kernel.signal, so channel validation, model canonicalization, and dispositions are identical to Python's by construction -- resolved/buffered/duplicate are 202, unknown_run is 404, run_terminal and expired are 409, an unknown channel or refused payload is the kernel's 400. Webhooks and approval links keep their embedded-mode paths byte-for-byte, so a Stripe URL configured against an rx.App keeps working when the deployment moves to serve; the plan's acceptance scenario -- a signed Stripe delivery with no frontend anywhere -- passes, along with a forged-signature 401 and an id-keyed redelivery answering deduplicated. Shutdown is the graceful sequence: uvicorn stops accepting, then the lifespan gives running attempts the drain budget, then the service closes its own store. That last step was missing first and Postgres made it visible: a pool's maintenance tasks do not exit for a bare loop-close cancellation, which leaked the pool in production and hung any embedding that waits for the loop's tasks -- the test suite's TestClient portal join reproduced it deterministically. The store now closes on its own loop while it still runs. Verified end to end by a test that boots the real server in a child process, drives the API over real sockets, sends SIGTERM, and asserts the drain completed -- uvicorn then re-raises the signal it captured, so dying by SIGTERM after "Application shutdown complete" is the graceful exit, and the assertion knows that. The SQLite variant of the same teardown was harsher: with the store now closed at lifespan end, close() could run while a work item abandoned by a cancelled awaiter was still mid-statement on its worker thread -- a segfault in the sqlite3 C layer, not an exception. close() now takes the store's lock, so it waits out whatever is executing. Full suite against real Postgres: 1,776 passed in 69s, no hangs, no dumps. --- news/workflow-serve.md | 5 + reflex/workflow/CONTRACT.md | 22 + reflex/workflow/api.py | 47 +- reflex/workflow/cli.py | 128 ++++++ reflex/workflow/serve.py | 584 +++++++++++++++++++++++++ reflex/workflow/store.py | 12 +- tests/units/workflow/test_cli_serve.py | 182 ++++++++ tests/units/workflow/test_serve.py | 385 ++++++++++++++++ 8 files changed, 1354 insertions(+), 11 deletions(-) create mode 100644 news/workflow-serve.md create mode 100644 reflex/workflow/serve.py create mode 100644 tests/units/workflow/test_cli_serve.py create mode 100644 tests/units/workflow/test_serve.py diff --git a/news/workflow-serve.md b/news/workflow-serve.md new file mode 100644 index 00000000000..d3afef948bd --- /dev/null +++ b/news/workflow-serve.md @@ -0,0 +1,5 @@ +`reflex workflows serve module.py` runs a workflow deployment as one standalone process — no frontend, no `rx.App`: webhook and approval ingress, the run HTTP API, the worker loop, `/healthz`, `/readyz`, Prometheus `/metrics`, and an OpenAPI document. `--ingress-only` and `--worker-only` split the halves for separate scaling against the same store, each keeping its probes. On SIGTERM the server stops accepting requests and gives running attempts the drain budget to commit, verified by a test that boots the real server and kills it with a real signal. + +The run API is `POST /runs`, `GET /runs` (filters: workflow, status, labels), `GET /runs/{id}`, `POST /runs/{id}/signals/{channel}`, and `POST /runs/{id}/cancel|retry|resume`. Authorization is scoped: `REFLEX_WORKFLOW_API_TOKEN` grants everything as before, while `REFLEX_WORKFLOW_API_TOKEN_READ`, `_START`, `_SIGNAL`, and `_OPERATE` each grant exactly one scope — a missing token is 401, a valid token without the route's scope is 403. HTTP signals go through the same kernel path as Python ones, so channel validation, payload canonicalization, and dispositions are identical by construction; webhooks keep their provider-signature authentication and their embedded-mode paths byte-for-byte, so a Stripe URL configured against an `rx.App` keeps working when the deployment moves to `serve`. + +The standalone service also closes its own store at shutdown, on its own loop. That fixed two teardown defects the test suite surfaced: a Postgres pool whose maintenance tasks never exit for a bare loop-close (leaking the pool and hanging embeddings that wait for loop tasks), and a SQLite connection closed while an abandoned work item was still mid-statement on its worker thread — a C-level segfault. `SqliteRunStore.close()` now takes the store's lock so it waits out in-flight statements. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index be7f21e2c1d..632a555ae94 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -312,6 +312,28 @@ Checked in this order at start: - Multiple workers share one Postgres store via `SKIP LOCKED` claims; SQLite is a one-process store (calls off-loop, contention bounded); memory is for tests. All three answer the same conformance suite. +### The standalone service + +`reflex workflows serve module.py` is the deployment shape for a workflow +with no frontend: webhook and approval ingress, the run HTTP API +(`POST /runs`, `GET /runs`, `GET /runs/{id}`, `POST /runs/{id}/signals/ +{channel}`, `POST /runs/{id}/cancel|retry|resume`), `/healthz`, `/readyz`, +`/metrics`, `/openapi.json`, and the worker loop in one process. +`--ingress-only` and `--worker-only` split the halves for separate scaling +against the same store; both keep the probes and metrics. Shutdown is the +graceful sequence: the server stops accepting, then running attempts get +the drain budget (§ drain) to commit. + +API authorization is scoped bearer tokens: `REFLEX_WORKFLOW_API_TOKEN` +grants every scope; `REFLEX_WORKFLOW_API_TOKEN_READ`, `_START`, `_SIGNAL`, +and `_OPERATE` grant exactly one each, so a dashboard's credential cannot +cancel runs and a relay's credential cannot read them. No valid token is +401; a valid token without the route's scope is 403. Webhooks and approval +links do not use these tokens — they authenticate with provider signatures +and signed link tokens respectively. HTTP signal deliveries pass through +the same kernel path as Python ones, so dispositions, channel validation, +and payload canonicalization are identical by construction. + ### Tenancy and who runs the workers Managed and customer-hosted are a deployment split, not a semantic one. A diff --git a/reflex/workflow/api.py b/reflex/workflow/api.py index a4e20735a36..aca99412b11 100644 --- a/reflex/workflow/api.py +++ b/reflex/workflow/api.py @@ -51,6 +51,32 @@ def api_token() -> str | None: return os.environ.get(TOKEN_ENV) or None +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine + + AuthorizeFn = Callable[[Request], JSONResponse | None] +else: + AuthorizeFn = object + + +def _refusal(request: Request, token: str | AuthorizeFn) -> JSONResponse | None: + """Authorize a request against a token or a scope authorizer. + + Args: + request: The incoming request. + token: The single bearer token, or an authorizer returning a refusal + response (401/403) or None to admit. + + Returns: + The refusal to send, or None when the request may proceed. + """ + if callable(token): + return token(request) + if _authorized(request, token): + return None + return JSONResponse({"error": "unauthorized"}, status_code=401) + + def _authorized(request: Request, token: str) -> bool: """Check a request's bearer token in constant time. @@ -69,7 +95,7 @@ def _authorized(request: Request, token: str) -> bool: def start_endpoint( - runtime: WorkflowRuntime, token: str + runtime: WorkflowRuntime, token: str | AuthorizeFn ) -> Callable[[Request], Coroutine[Any, Any, JSONResponse]]: """Build the endpoint that starts a run. @@ -90,8 +116,9 @@ async def endpoint(request: Request) -> JSONResponse: Returns: The admission result, or an error. """ - if not _authorized(request, token): - return JSONResponse({"error": "unauthorized"}, status_code=401) + refused = _refusal(request, token) + if refused is not None: + return refused body = await request.body() if len(body) > MAX_BODY_BYTES: return JSONResponse({"error": "payload too large"}, status_code=413) @@ -181,7 +208,7 @@ async def endpoint(request: Request) -> JSONResponse: def run_endpoint( - runtime: WorkflowRuntime, token: str + runtime: WorkflowRuntime, token: str | AuthorizeFn ) -> Callable[[Request], Coroutine[Any, Any, JSONResponse]]: """Build the endpoint that reads one run. @@ -202,8 +229,9 @@ async def endpoint(request: Request) -> JSONResponse: Returns: The run projection, or an error. """ - if not _authorized(request, token): - return JSONResponse({"error": "unauthorized"}, status_code=401) + refused = _refusal(request, token) + if refused is not None: + return refused run_id = request.path_params.get("run_id", "") snapshot = await runtime.kernel.get_run(run_id) if snapshot is None: @@ -275,7 +303,7 @@ def render_prometheus(snapshot: dict[str, Any]) -> str: def metrics_endpoint( - runtime: WorkflowRuntime, token: str + runtime: WorkflowRuntime, token: str | AuthorizeFn ) -> Callable[[Request], Coroutine[Any, Any, Response]]: """Build the endpoint that exposes this process's counters. @@ -302,8 +330,9 @@ async def endpoint(request: Request) -> Response: # noqa: RUF029 """ # Reading in-process counters needs no await; the signature is a # Starlette endpoint's, not a claim that this does I/O. - if not _authorized(request, token): - return JSONResponse({"error": "unauthorized"}, status_code=401) + refused = _refusal(request, token) + if refused is not None: + return refused return PlainTextResponse( render_prometheus(runtime.metrics.snapshot()), media_type="text/plain; version=0.0.4; charset=utf-8", diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index faac40cad83..44dfea9fd20 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -1286,6 +1286,134 @@ def _load_module(target: str): return importlib.import_module(target) +def _run_server(app: Any, host: str, port: int) -> None: + """Serve an ASGI app; separated so tests can intercept it. + + Args: + app: The ASGI application. + host: The interface to bind. + port: The port to bind. + """ + import uvicorn + + uvicorn.run(app, host=host, port=port, log_level="info", lifespan="on") + + +@workflows.command() +@database_option +@click.argument("target") +@click.option("--host", default="0.0.0.0", help="Interface to bind.") +@click.option("--port", default=8000, type=int, help="Port to bind.") +@click.option( + "--queue", + "queues", + multiple=True, + help="Execute only these queues. Repeatable; default serves every queue.", +) +@click.option( + "--ingress-only", + is_flag=True, + help="Accept webhooks and API calls but execute nothing.", +) +@click.option( + "--worker-only", + is_flag=True, + help="Execute steps but accept nothing; keeps /healthz, /readyz, /metrics.", +) +@click.option( + "--drain", + default=None, + help=( + "How long shutdown lets running attempts finish. Defaults to " + "REFLEX_WORKFLOW_DRAIN, or 30s." + ), +) +def serve( + database: str | None, + target: str, + host: str, + port: int, + queues: tuple[str, ...], + ingress_only: bool, + worker_only: bool, + drain: str | None, +): + """Serve TARGET's workflows as a standalone service: ingress, API, worker. + + One process, no frontend, no rx.App: webhook and approval ingress, the + run HTTP API (POST /runs, GET /runs, signals, operator actions), health + and readiness probes, Prometheus metrics, an OpenAPI document, and the + worker loop. Scale the halves separately with --ingress-only and + --worker-only against the same database. + + The API authenticates with bearer tokens: REFLEX_WORKFLOW_API_TOKEN grants + everything; REFLEX_WORKFLOW_API_TOKEN_READ, _START, _SIGNAL, and _OPERATE + each grant one scope. Webhooks authenticate with provider signatures. + + On SIGTERM the server stops accepting requests, then gives running + attempts --drain to commit, so a rolling deploy hands over cleanly. + """ + from reflex_base.utils.exceptions import WorkflowDefinitionError + from reflex_base.workflow import parse_duration + + from reflex.workflow.runtime import WorkflowRuntime, configured_drain + from reflex.workflow.serve import build_app + from reflex.workflow.store import resolve_store + + if ingress_only and worker_only: + console.error("--ingress-only and --worker-only exclude each other.") + raise click.exceptions.Exit(1) + if drain is None: + drain_seconds = configured_drain() + else: + try: + drain_seconds = parse_duration(drain) + except Exception as err: + console.error(f"--drain {drain!r} is not a duration: {err}") + raise click.exceptions.Exit(1) from None + + try: + module = _load_module(target) + except Exception as err: + console.error(f"Could not load {target!r}: {err}") + raise click.exceptions.Exit(1) from None + classes = [ + value + for value in vars(module).values() + if isinstance(value, type) and "__workflow__" in vars(value) + ] + if not classes: + console.error( + f"No workflow classes in {target!r}. A workflow is an rx.State " + "subclass with __workflow__ = rx.WorkflowConfig(id=...)." + ) + raise click.exceptions.Exit(1) + + runtime = WorkflowRuntime(resolve_store(database), queues=queues or None) + try: + for workflow_cls in classes: + runtime.register(workflow_cls) + except WorkflowDefinitionError as err: + console.error(f"Cannot serve {target!r}: {err}") + raise click.exceptions.Exit(1) from None + app = build_app( + runtime, + worker=not ingress_only, + ingress=not worker_only, + drain=drain_seconds, + ) + served = ", ".join(sorted(d.workflow_id for d in runtime.definitions)) + mode = ( + "ingress only" + if ingress_only + else "worker only" + if worker_only + else "ingress + worker" + ) + console.print(f"Serving {served} on {host}:{port} ({mode}).") + _run_server(app, host, port) + + @workflows.command() @database_option @click.argument("run_id") diff --git a/reflex/workflow/serve.py b/reflex/workflow/serve.py new file mode 100644 index 00000000000..42dcc7ccb56 --- /dev/null +++ b/reflex/workflow/serve.py @@ -0,0 +1,584 @@ +"""The standalone workflow service: ingress, API, and worker in one process. + +``reflex workflows serve workflows.py`` runs everything a deployed workflow +needs with no frontend and no ``rx.App``: webhook and approval ingress, the +run HTTP API, the worker loop, and the probes an orchestrator points at. +``--ingress-only`` and ``--worker-only`` split the same process for separate +scaling; both halves keep ``/healthz``, ``/readyz``, and ``/metrics``. + +Authorization is scoped. ``REFLEX_WORKFLOW_API_TOKEN`` grants everything, as it +always has; ``REFLEX_WORKFLOW_API_TOKEN_READ``, ``_START``, ``_SIGNAL``, and +``_OPERATE`` each grant exactly one scope, so the credential a dashboard +holds cannot cancel runs and the credential a webhook relay holds cannot +read them. A request with no valid token is a 401; a valid token without +the route's scope is a 403. +""" + +from __future__ import annotations + +import hmac +import json +import os +from typing import TYPE_CHECKING, Any, Final + +from reflex_base.utils import console +from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.workflow import ChannelDelivery +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +from reflex.workflow.api import ( + MAX_BODY_BYTES, + TOKEN_ENV, + metrics_endpoint, + run_endpoint, + start_endpoint, +) +from reflex.workflow.records import RunQuery, RunStatus +from reflex.workflow.runtime import _close_store + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Callable + + from starlette.requests import Request + + from reflex.workflow.runtime import WorkflowRuntime + +SCOPES: Final = ("read", "start", "signal", "operate") +SCOPE_TOKEN_ENVS: Final = {scope: f"{TOKEN_ENV}_{scope.upper()}" for scope in SCOPES} + + +def _bearer(request: Request) -> str: + """Extract the bearer token a request presents. + + Args: + request: The incoming request. + + Returns: + The presented token, or an empty string. + """ + header = request.headers.get("authorization", "") + scheme, _, credential = header.partition(" ") + return credential.strip() if scheme.lower() == "bearer" else "" + + +class ScopedTokens: + """The service's token-to-scope mapping, read from the environment. + + Attributes: + grants: Token to granted scopes. + """ + + def __init__(self): + """Read every configured token.""" + self.grants: dict[str, frozenset[str]] = {} + universal = os.environ.get(TOKEN_ENV) + if universal: + self.grants[universal] = frozenset(SCOPES) + for scope, env in SCOPE_TOKEN_ENVS.items(): + token = os.environ.get(env) + if token: + merged = self.grants.get(token, frozenset()) | {scope} + self.grants[token] = merged + + def __bool__(self) -> bool: + """Whether any token is configured. + + Returns: + True when at least one token exists. + """ + return bool(self.grants) + + def scopes_for(self, request: Request) -> frozenset[str] | None: + """Resolve the scopes a request's token grants. + + Args: + request: The incoming request. + + Returns: + The granted scopes, or None when no configured token matches. + """ + presented = _bearer(request) + if not presented: + return None + for token, scopes in self.grants.items(): + # Compared in constant time, every candidate every time, so the + # comparison count does not leak which token was close. + if hmac.compare_digest(presented, token): + return scopes + return None + + def require(self, scope: str) -> Callable[[Request], JSONResponse | None]: + """Build an authorizer demanding one scope. + + Args: + scope: The scope the route requires. + + Returns: + An authorizer returning a refusal response or None to admit. + """ + + def authorize(request: Request) -> JSONResponse | None: + """Refuse a request without the scope. + + Args: + request: The incoming request. + + Returns: + The refusal, or None to admit. + """ + granted = self.scopes_for(request) + if granted is None: + return JSONResponse({"error": "unauthorized"}, status_code=401) + if scope not in granted: + return JSONResponse( + {"error": f"token lacks the {scope!r} scope"}, + status_code=403, + ) + return None + + return authorize + + +async def _read_json(request: Request) -> tuple[Any, JSONResponse | None]: + """Read and decode a JSON request body within the size cap. + + Args: + request: The incoming request. + + Returns: + The decoded payload and None, or None and the refusal to send. + """ + body = await request.body() + if len(body) > MAX_BODY_BYTES: + return None, JSONResponse({"error": "payload too large"}, status_code=413) + try: + return (json.loads(body) if body else None), None + except json.JSONDecodeError: + return None, JSONResponse({"error": "payload is not JSON"}, status_code=400) + + +def list_runs_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens): + """Build the run-listing endpoint. + + Args: + runtime: The runtime owning the runs. + tokens: The service's token scopes. + + Returns: + The endpoint callable. + """ + authorize = tokens.require("read") + + async def endpoint(request: Request) -> JSONResponse: + """List runs, newest first, filtered by query parameters. + + Args: + request: The incoming request. + + Returns: + The run summaries. + """ + refused = authorize(request) + if refused is not None: + return refused + params = request.query_params + raw_statuses = params.getlist("status") + known = {status.value: status for status in RunStatus} + unknown = [raw for raw in raw_statuses if raw.upper() not in known] + if unknown: + return JSONResponse( + {"error": f"unknown status {unknown[0]!r}"}, status_code=400 + ) + statuses = [known[raw.upper()] for raw in raw_statuses] + labels = { + name.removeprefix("label."): value + for name, value in params.items() + if name.startswith("label.") + } + try: + limit = min(int(params.get("limit", "50")), 500) + except ValueError: + return JSONResponse({"error": "limit must be an integer"}, 400) + runs = await runtime.kernel._store.list_runs( # pyright: ignore[reportPrivateUsage] + RunQuery( + workflow_id=params.get("workflow"), + statuses=tuple(statuses), + labels=labels or None, + limit=limit, + ) + ) + return JSONResponse({ + "runs": [ + { + "run_id": run.run_id, + "workflow": run.workflow_id, + "status": run.status.value, + "labels": run.labels or {}, + "created_at": run.created_at, + "updated_at": run.updated_at, + } + for run in runs + ] + }) + + return endpoint + + +def signal_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens): + """Build the endpoint that delivers a signal to a run's channel. + + Args: + runtime: The runtime owning the runs. + tokens: The service's token scopes. + + Returns: + The endpoint callable. + """ + authorize = tokens.require("signal") + + async def endpoint(request: Request) -> JSONResponse: + """Deliver the request body to the named channel. + + Args: + request: The incoming request. + + Returns: + The delivery disposition, or an error. + """ + refused = authorize(request) + if refused is not None: + return refused + payload, bad = await _read_json(request) + if bad is not None: + return bad + run_id = request.path_params["run_id"] + channel = request.path_params["channel"] + key = request.headers.get("idempotency-key") or request.query_params.get("key") + try: + disposition = await runtime.kernel.signal( + run_id, ChannelDelivery(channel=channel, payload=payload), key=key + ) + except WorkflowDefinitionError as error: + # An unknown channel or a payload the channel's model refuses is + # the sender's bug; the kernel's message names the fix. + return JSONResponse({"error": str(error)}, status_code=400) + status = { + "unknown_run": 404, + "run_terminal": 409, + "expired": 409, + }.get(disposition, 202) + return JSONResponse({"disposition": disposition}, status_code=status) + + return endpoint + + +def operator_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens, action: str): + """Build one operator action endpoint. + + Args: + runtime: The runtime owning the runs. + tokens: The service's token scopes. + action: One of ``cancel``, ``retry``, or ``resume``. + + Returns: + The endpoint callable. + """ + authorize = tokens.require("operate") + + async def endpoint(request: Request) -> JSONResponse: + """Apply the operator action to the addressed run. + + Args: + request: The incoming request. + + Returns: + Whether the action applied, or an error. + """ + refused = authorize(request) + if refused is not None: + return refused + run_id = request.path_params["run_id"] + if await runtime.kernel.get_run(run_id) is None: + return JSONResponse({"error": "unknown run"}, status_code=404) + applied = await getattr(runtime.kernel, action)(run_id) + if not applied: + # The run exists but is not in a state this action accepts -- + # retrying a healthy run, resuming one that is not suspended. + return JSONResponse( + {"error": f"run does not accept {action} in its current state"}, + status_code=409, + ) + return JSONResponse({"applied": True}, status_code=202) + + return endpoint + + +def health_endpoint(): + """Build the liveness probe. + + Returns: + The endpoint callable. + """ + + async def endpoint(request: Request) -> JSONResponse: # noqa: RUF029 + """Answer that the process is up. + + Args: + request: The incoming request. + + Returns: + 200 always; a dead process does not answer. + """ + return JSONResponse({"status": "ok"}) + + return endpoint + + +def ready_endpoint(runtime: WorkflowRuntime): + """Build the readiness probe. + + Args: + runtime: The runtime whose store must be reachable. + + Returns: + The endpoint callable. + """ + + async def endpoint(request: Request) -> JSONResponse: + """Answer whether this process can do useful work right now. + + Args: + request: The incoming request. + + Returns: + 200 with the store reachable, 503 otherwise. + """ + try: + await runtime.kernel._store.epoch_time() # pyright: ignore[reportPrivateUsage] + except Exception as error: + return JSONResponse( + {"status": "unready", "store": str(error)}, status_code=503 + ) + return JSONResponse({"status": "ready"}) + + return endpoint + + +def openapi_endpoint(runtime: WorkflowRuntime): + """Build the OpenAPI document endpoint. + + Args: + runtime: The runtime whose workflows the document describes. + + Returns: + The endpoint callable. + """ + + async def endpoint(request: Request) -> JSONResponse: # noqa: RUF029 + """Serve the API description. + + Args: + request: The incoming request. + + Returns: + The OpenAPI 3.1 document. + """ + workflows = sorted(defn.workflow_id for defn in runtime.definitions) + run_ref = {"$ref": "#/components/schemas/Disposition"} + document = { + "openapi": "3.1.0", + "info": { + "title": "Reflex Workflows", + "version": "1", + "description": ( + f"Workflow service for: {', '.join(workflows) or 'none'}" + ), + }, + "components": { + "securitySchemes": {"bearer": {"type": "http", "scheme": "bearer"}}, + "schemas": { + "Disposition": { + "type": "object", + "properties": { + "disposition": {"type": "string"}, + "run_id": {"type": ["string", "null"]}, + }, + } + }, + }, + "security": [{"bearer": []}], + "paths": { + "/runs": { + "post": { + "summary": "Start a run (scope: start)", + "responses": {"202": {"description": "Admitted"}}, + "requestBody": { + "content": { + "application/json": { + "schema": { + "type": "object", + "required": ["workflow", "handler"], + "properties": { + "workflow": {"type": "string"}, + "handler": {"type": "string"}, + "args": {"type": "object"}, + "request_key": {"type": "string"}, + "labels": {"type": "object"}, + }, + } + } + } + }, + }, + "get": { + "summary": "List runs (scope: read)", + "responses": {"200": {"description": "Run summaries"}}, + }, + }, + "/runs/{run_id}": { + "get": { + "summary": "Read one run (scope: read)", + "responses": { + "200": { + "description": "The run", + "content": {"application/json": {"schema": run_ref}}, + } + }, + } + }, + "/runs/{run_id}/signals/{channel}": { + "post": { + "summary": "Deliver a signal (scope: signal)", + "responses": {"202": {"description": "Delivered"}}, + } + }, + **{ + f"/runs/{{run_id}}/{action}": { + "post": { + "summary": f"{action.title()} a run (scope: operate)", + "responses": {"202": {"description": "Applied"}}, + } + } + for action in ("cancel", "retry", "resume") + }, + "/healthz": {"get": {"summary": "Liveness", "security": []}}, + "/readyz": {"get": {"summary": "Readiness", "security": []}}, + "/metrics": {"get": {"summary": "Prometheus metrics (scope: read)"}}, + }, + } + return JSONResponse(document) + + return endpoint + + +def build_app( + runtime: WorkflowRuntime, + *, + worker: bool = True, + ingress: bool = True, + drain: float | str = "25s", + tokens: ScopedTokens | None = None, +) -> Starlette: + """Compose the standalone service application. + + Args: + runtime: The runtime to serve. + worker: Whether this process executes steps. + ingress: Whether this process accepts webhooks and API calls. + drain: How long shutdown gives in-flight attempts to commit. + tokens: Token scopes; read from the environment when omitted. + + Returns: + The ASGI application, with lifespan wired to the runtime. + """ + from contextlib import asynccontextmanager + + tokens = tokens if tokens is not None else ScopedTokens() + + @asynccontextmanager + async def lifespan(app: Starlette) -> AsyncIterator[None]: + """Start the runtime with the app and drain it on shutdown. + + The server stops accepting requests before this exits, so the order + is exactly the graceful sequence: ingress closes, then the worker + gets the drain budget to commit what it holds. + + Args: + app: The application being served. + + Yields: + Nothing; the service runs while suspended here. + """ + await runtime.startup(start_worker=worker) + try: + yield + finally: + await runtime.shutdown(drain=drain) + # The service owns its store, so the store dies with the service + # -- on this loop, while it still runs. A Postgres pool's + # maintenance tasks do not exit for a bare loop-close + # cancellation, which leaks the pool in production and hangs any + # embedding that waits for the loop's tasks to finish. + await _close_store(runtime.store) + + routes = [ + Route("/healthz", health_endpoint(), methods=["GET"]), + Route("/readyz", ready_endpoint(runtime), methods=["GET"]), + Route( + "/metrics", + metrics_endpoint(runtime, tokens.require("read")), + methods=["GET"], + ), + ] + if ingress: + from reflex.workflow.approvals import APPROVAL_ROUTE, approval_endpoint + from reflex.workflow.ingress import ( + WEBHOOK_ROUTE, + collect_webhook_routes, + webhook_endpoint, + ) + + if not tokens: + console.warn( + f"No API token configured ({TOKEN_ENV} or scoped variants); " + "the run API is refusing every request. Webhooks still work: " + "they authenticate with their own signatures." + ) + routes += [ + Route("/openapi.json", openapi_endpoint(runtime), methods=["GET"]), + Route( + "/runs", + start_endpoint(runtime, tokens.require("start")), + methods=["POST"], + ), + Route("/runs", list_runs_endpoint(runtime, tokens), methods=["GET"]), + Route( + "/runs/{run_id}", + run_endpoint(runtime, tokens.require("read")), + methods=["GET"], + ), + Route( + "/runs/{run_id}/signals/{channel}", + signal_endpoint(runtime, tokens), + methods=["POST"], + ), + *( + Route( + f"/runs/{{run_id}}/{action}", + operator_endpoint(runtime, tokens, action), + methods=["POST"], + ) + for action in ("cancel", "retry", "resume") + ), + # The embedded-mode paths, kept byte-for-byte: a Stripe URL or a + # minted approval link configured against an rx.App keeps working + # when the deployment moves to the standalone service. + Route(APPROVAL_ROUTE, approval_endpoint(runtime), methods=["GET", "POST"]), + ] + if collect_webhook_routes(runtime.definitions): + routes.append( + Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"]) + ) + return Starlette(routes=routes, lifespan=lifespan) diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index f055fa5ea54..ae5923bc142 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -2900,8 +2900,16 @@ def _migrate(self) -> None: raise def close(self) -> None: - """Close the backing database connection.""" - self._db.close() + """Close the backing database connection. + + Under the store's lock: every operation runs its SQL holding it, and + an operation abandoned by a cancelled awaiter is still executing on + its worker thread. Closing the connection out from under that thread + is a segfault in the sqlite3 C layer, not an exception -- taking the + lock makes close wait out whatever is mid-statement. + """ + with self._lock: + self._db.close() def _append_events( self, diff --git a/tests/units/workflow/test_cli_serve.py b/tests/units/workflow/test_cli_serve.py new file mode 100644 index 00000000000..9c49189908c --- /dev/null +++ b/tests/units/workflow/test_cli_serve.py @@ -0,0 +1,182 @@ +"""The serve command boots a real server and hands over cleanly on SIGTERM. + +These drive the actual console entry in a child process -- uvicorn, real +sockets, a real signal -- because "stops accepting and drains" is a claim +about process behavior that an in-process TestClient cannot make. +""" + +import json +import signal +import socket +import subprocess +import sys +import time +import urllib.error +import urllib.request + +MODULE = ''' +import reflex as rx + + +class Pinger(rx.State): + __workflow__ = rx.WorkflowConfig(id="serve.pinger") + + @rx.event(durable=True, trigger=rx.manual(), effect="none") + def go(self, name: str): + """Complete immediately. + + Args: + name: Who pinged. + + Returns: + Completion. + """ + return rx.complete(result={"hello": name}) +''' + +RUNNER = "from reflex.workflow.cli import workflows; workflows()" + + +def _free_port() -> int: + """Reserve an ephemeral port. + + Returns: + A port number that was free at reservation time. + """ + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _get(url: str) -> tuple[int, dict]: + """GET a URL. + + Args: + url: The URL. + + Returns: + Status code and decoded body. + """ + with urllib.request.urlopen(url, timeout=2) as response: + return response.status, json.loads(response.read()) + + +def test_serve_boots_starts_a_run_and_drains_on_sigterm(tmp_path): + """Boot, probe, start a run over HTTP, SIGTERM, exit clean. + + Args: + tmp_path: Working directory for the module and database. + """ + module = tmp_path / "pinger.py" + module.write_text(MODULE) + port = _free_port() + process = subprocess.Popen( + [ + sys.executable, + "-c", + RUNNER, + "serve", + str(module), + "--host", + "127.0.0.1", + "--port", + str(port), + "-d", + str(tmp_path / "serve.db"), + "--drain", + "5s", + ], + env={ + **__import__("os").environ, + "REFLEX_WORKFLOW_API_TOKEN": "tk_serve_cli", + }, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + base = f"http://127.0.0.1:{port}" + try: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + try: + status, body = _get(f"{base}/healthz") + if status == 200: + break + except (urllib.error.URLError, ConnectionError): + time.sleep(0.2) + else: + msg = "server never became healthy" + raise AssertionError(msg) + status, body = _get(f"{base}/readyz") + assert status == 200, body + + request = urllib.request.Request( + f"{base}/runs", + data=json.dumps({ + "workflow": "serve.pinger", + "handler": "go", + "args": {"name": "cli"}, + }).encode(), + headers={ + "Authorization": "Bearer tk_serve_cli", + "Content-Type": "application/json", + }, + method="POST", + ) + with urllib.request.urlopen(request, timeout=5) as response: + admitted = json.loads(response.read()) + assert response.status == 202 + run_id = admitted["run_id"] + + snapshot: dict = {} + deadline = time.monotonic() + 15 + while time.monotonic() < deadline: + request = urllib.request.Request( + f"{base}/runs/{run_id}", + headers={"Authorization": "Bearer tk_serve_cli"}, + ) + with urllib.request.urlopen(request, timeout=5) as response: + snapshot = json.loads(response.read()) + if snapshot["status"] == "COMPLETED": + break + time.sleep(0.1) + assert snapshot["status"] == "COMPLETED", snapshot + assert snapshot["result"] == {"hello": "cli"} + finally: + process.send_signal(signal.SIGTERM) + output, _ = process.communicate(timeout=30) + # Uvicorn drains and then re-raises the captured SIGTERM so the parent + # sees the true cause of death; -SIGTERM after a completed shutdown IS + # the graceful exit. What must never appear is a shutdown that started + # and did not finish. + assert process.returncode in (0, -signal.SIGTERM), output + assert "Application shutdown complete" in output, output + + +def test_serve_refuses_contradictory_modes(tmp_path): + """--ingress-only and --worker-only exclude each other. + + Args: + tmp_path: Working directory for the module. + """ + module = tmp_path / "pinger.py" + module.write_text(MODULE) + result = subprocess.run( + [ + sys.executable, + "-c", + RUNNER, + "serve", + str(module), + "--ingress-only", + "--worker-only", + "-d", + str(tmp_path / "x.db"), + ], + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode != 0 + assert "exclude each other" in result.stdout + result.stderr diff --git a/tests/units/workflow/test_serve.py b/tests/units/workflow/test_serve.py new file mode 100644 index 00000000000..9842c1c52ca --- /dev/null +++ b/tests/units/workflow/test_serve.py @@ -0,0 +1,385 @@ +"""The standalone service: scoped auth, the run API, probes, and lifecycle. + +The acceptance bar from the plan: ``workflows serve`` receives a signed +Stripe webhook with no ``rx.App`` and no frontend; shutdown drains before +losing anything; Python and HTTP signaling agree on dispositions. +""" + +import hmac as hmac_mod +import json +import time + +import pytest +from reflex_base.workflow import ( + Signal, + WorkflowConfig, + manual, + stripe_signature, + webhook, +) +from starlette.testclient import TestClient + +import reflex as rx +from reflex.workflow import testing +from reflex.workflow.runtime import WorkflowRuntime +from reflex.workflow.serve import SCOPES, ScopedTokens, build_app + +STRIPE_SECRET = "whsec_serve_test" + + +class Orders(rx.State): + """A workflow with a manual root and a signal channel.""" + + __workflow__ = WorkflowConfig(id="serve.orders") + note: str = "" + + @rx.event(durable=True, effect="none", trigger=manual()) + def place(self, order_id: str): + """Start an order and wait for its shipment. + + Args: + order_id: The order. + + Returns: + The wait. + """ + self.note = order_id + return rx.wait_for(Orders.shipped, then=Orders.close, timeout=rx.never) + + shipped = Signal() + + @rx.event(durable=True, effect="none") + def close(self, payload): + """Finish with the shipment payload. + + Args: + payload: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result={"order": self.note, "shipment": payload}) + + +def _tokens(**grants: str) -> ScopedTokens: + """Build token scopes without touching the environment. + + Args: + grants: token=scope pairs, scope "all" meaning every scope. + + Returns: + The scope mapping. + """ + tokens = ScopedTokens.__new__(ScopedTokens) + tokens.grants = { + token: frozenset(SCOPES) if scope == "all" else frozenset({scope}) + for token, scope in grants.items() + } + return tokens + + +def _auth(token: str) -> dict[str, str]: + """Bearer header for a token. + + Args: + token: The token to present. + + Returns: + The header mapping. + """ + return {"Authorization": f"Bearer {token}"} + + +@pytest.fixture +def service(forked_registration_context): + """A served runtime over the parametrized store, with scoped tokens. + + Args: + forked_registration_context: Isolated state registry. + + Yields: + The entered test client. + """ + runtime = WorkflowRuntime(testing.MemoryRunStore()) + runtime.register(Orders) + app = build_app( + runtime, + worker=True, + drain=0, + tokens=_tokens( + tk_all="all", + tk_read="read", + tk_start="start", + tk_signal="signal", + tk_operate="operate", + ), + ) + with TestClient(app) as client: + yield client + + +def _start_order(client: TestClient, token: str = "tk_start") -> str: + """Start one order run over HTTP. + + Args: + client: The service client. + token: The token to start with. + + Returns: + The admitted run id. + """ + response = client.post( + "/runs", + json={ + "workflow": "serve.orders", + "handler": "place", + "args": {"order_id": "o1"}, + }, + headers=_auth(token), + ) + assert response.status_code == 202, response.text + return response.json()["run_id"] + + +def test_every_scope_gates_exactly_its_routes(service): + """Read cannot write, start cannot read, and nothing works tokenless. + + Args: + service: The served client. + """ + body = {"workflow": "serve.orders", "handler": "place", "args": {"order_id": "x"}} + assert service.post("/runs", json=body).status_code == 401 + assert service.get("/runs").status_code == 401 + assert service.post("/runs", json=body, headers=_auth("tk_read")).status_code == 403 + assert service.get("/runs", headers=_auth("tk_start")).status_code == 403 + assert service.get("/metrics", headers=_auth("tk_signal")).status_code == 403 + assert service.get("/metrics", headers=_auth("tk_read")).status_code == 200 + run_id = _start_order(service) + assert ( + service.post(f"/runs/{run_id}/cancel", headers=_auth("tk_signal")).status_code + == 403 + ) + assert ( + service.post( + f"/runs/{run_id}/signals/shipped", + json={}, + headers=_auth("tk_operate"), + ).status_code + == 403 + ) + assert service.get(f"/runs/{run_id}", headers=_auth("tk_all")).status_code == 200 + + +def test_the_run_lifecycle_works_end_to_end_over_http(service): + """Start, list, read, signal, and read the result, all over the API. + + Args: + service: The served client. + """ + run_id = _start_order(service) + listed = service.get( + "/runs", params={"workflow": "serve.orders"}, headers=_auth("tk_read") + ) + assert listed.status_code == 200 + assert any(run["run_id"] == run_id for run in listed.json()["runs"]) + + delivered = service.post( + f"/runs/{run_id}/signals/shipped", + json={"parcel": "P-1"}, + headers={**_auth("tk_signal"), "Idempotency-Key": "evt_1"}, + ) + assert delivered.status_code == 202 + assert delivered.json()["disposition"] in ("resolved", "buffered") + duplicate = service.post( + f"/runs/{run_id}/signals/shipped", + json={"parcel": "P-1"}, + headers={**_auth("tk_signal"), "Idempotency-Key": "evt_1"}, + ) + assert duplicate.json()["disposition"] == "duplicate" + + snapshot: dict = {} + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + snapshot = service.get(f"/runs/{run_id}", headers=_auth("tk_read")).json() + if snapshot["status"] == "COMPLETED": + break + time.sleep(0.05) + assert snapshot["status"] == "COMPLETED", snapshot + assert snapshot["result"] == {"order": "o1", "shipment": {"parcel": "P-1"}} + + +def test_http_signals_share_python_signal_semantics(service): + """The HTTP boundary maps kernel dispositions, never invents its own. + + Args: + service: The served client. + """ + missing = service.post( + "/runs/nope/signals/shipped", json={}, headers=_auth("tk_signal") + ) + assert missing.status_code == 404 + assert missing.json()["disposition"] == "unknown_run" + + run_id = _start_order(service) + unknown_channel = service.post( + f"/runs/{run_id}/signals/shiped", json={}, headers=_auth("tk_signal") + ) + assert unknown_channel.status_code == 400 + assert "shiped" in unknown_channel.json()["error"] + + +def test_operator_actions_answer_404_and_409_precisely(service): + """Unknown runs and wrong-state runs are different failures. + + Args: + service: The served client. + """ + assert ( + service.post("/runs/nope/cancel", headers=_auth("tk_operate")).status_code + == 404 + ) + run_id = _start_order(service) + retried = service.post(f"/runs/{run_id}/retry", headers=_auth("tk_operate")) + assert retried.status_code == 409, "a healthy run does not accept retry" + cancelled = service.post(f"/runs/{run_id}/cancel", headers=_auth("tk_operate")) + assert cancelled.status_code == 202 + + +def test_probes_and_openapi_need_no_token(service): + """An orchestrator holds no credentials; probes must answer anyway. + + Args: + service: The served client. + """ + assert service.get("/healthz").json() == {"status": "ok"} + assert service.get("/readyz").json() == {"status": "ready"} + document = service.get("/openapi.json").json() + assert document["openapi"].startswith("3.") + assert "/runs/{run_id}/signals/{channel}" in document["paths"] + + +def test_readyz_reports_an_unreachable_store(forked_registration_context): + """Ready means "can do useful work", and that means the store answers. + + Args: + forked_registration_context: Isolated state registry. + """ + runtime = WorkflowRuntime(testing.MemoryRunStore()) + runtime.register(Orders) + app = build_app(runtime, worker=False, drain=0, tokens=_tokens(t="all")) + with TestClient(app) as client: + + async def broken(): # noqa: RUF029 + msg = "store is down" + raise ConnectionError(msg) + + runtime.kernel._store.epoch_time = broken # pyright: ignore[reportAttributeAccessIssue] + response = client.get("/readyz") + assert response.status_code == 503 + assert "store is down" in response.json()["store"] + + +def test_worker_only_mode_serves_probes_and_nothing_else( + forked_registration_context, +): + """A pure worker still answers its orchestrator, but accepts no work. + + Args: + forked_registration_context: Isolated state registry. + """ + runtime = WorkflowRuntime(testing.MemoryRunStore()) + runtime.register(Orders) + app = build_app( + runtime, worker=True, ingress=False, drain=0, tokens=_tokens(t="all") + ) + with TestClient(app) as client: + assert client.get("/healthz").status_code == 200 + assert client.get("/metrics", headers=_auth("t")).status_code == 200 + assert client.post("/runs", json={}, headers=_auth("t")).status_code == 404 + assert client.get("/openapi.json").status_code == 404 + + +def test_shutdown_drains_the_runtime(forked_registration_context): + """Closing the server hands the drain budget to the runtime. + + Args: + forked_registration_context: Isolated state registry. + """ + runtime = WorkflowRuntime(testing.MemoryRunStore()) + runtime.register(Orders) + seen: list = [] + original = runtime.shutdown + + async def spying_shutdown(drain=0): + seen.append(drain) + await original(drain=drain) + + runtime.shutdown = spying_shutdown # pyright: ignore[reportAttributeAccessIssue] + app = build_app(runtime, worker=True, drain="7s", tokens=_tokens(t="all")) + with TestClient(app): + pass + assert seen == ["7s"], "shutdown must receive exactly the configured drain" + + +def test_a_signed_stripe_webhook_lands_with_no_rx_app( + monkeypatch, forked_registration_context +): + """The plan's acceptance scenario, minus uvicorn: serve takes a real + provider delivery with no frontend anywhere. + + Args: + monkeypatch: Used to install the Stripe secret. + forked_registration_context: Isolated state registry. + """ + monkeypatch.setenv("SERVE_STRIPE_SECRET", STRIPE_SECRET) + + class Billing(rx.State): + __workflow__ = WorkflowConfig(id="serve.billing") + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "invoice_paid", + dedupe_by="id", + verify=stripe_signature(secret_env="SERVE_STRIPE_SECRET"), + ), + ) + def on_paid(self, id: str): + """Record the invoice. + + Args: + id: The invoice identifier. + + Returns: + Completion. + """ + return rx.complete(result=id) + + runtime = WorkflowRuntime(testing.MemoryRunStore()) + runtime.register(Billing) + app = build_app(runtime, worker=True, drain=0, tokens=_tokens(t="all")) + body = json.dumps({"id": "inv_9"}).encode() + timestamp = int(time.time()) + digest = hmac_mod.new( + STRIPE_SECRET.encode(), f"{timestamp}.".encode() + body, "sha256" + ).hexdigest() + with TestClient(app) as client: + accepted = client.post( + "/_workflow/webhook/invoice_paid", + content=body, + headers={"Stripe-Signature": f"t={timestamp},v1={digest}"}, + ) + assert accepted.status_code == 202, accepted.text + forged = client.post( + "/_workflow/webhook/invoice_paid", + content=body, + headers={"Stripe-Signature": f"t={timestamp},v1={'0' * 64}"}, + ) + assert forged.status_code == 401 + redelivered = client.post( + "/_workflow/webhook/invoice_paid", + content=body, + headers={"Stripe-Signature": f"t={timestamp},v1={digest}"}, + ) + assert redelivered.json()["disposition"] == "deduplicated" From 47d369a4d5dd09dfb59fa9ea807175d4b7438fd3 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Mon, 24 Aug 2026 13:49:49 -0700 Subject: [PATCH 116/121] workflows: business-key addressing for runs and signals Ticket 3 of the standalone-GA plan, and the doorstep of correlated webhook delivery: the request key is already a durable unique index -- it is what makes webhook redelivery idempotent -- so it now doubles as a run's business address. Python: rx.workflows.get_by_key(Order, "order_123") returns a handle; rx.workflows.signal_by_key(Order, "order_123", Order.shipped(p), key=event_id) delivers to the run the key admitted, answering "unknown_key" when nothing did -- a new DeliveryDisposition, distinct from unknown_run because the caller's next move differs: an unknown key often means "not yet", and Phase 2's durable inbox will buffer on exactly that answer. HTTP: GET /workflows/{id}/keys/{request_key} (read scope) and POST /workflows/{id}/keys/{request_key}/signals/{channel} (signal scope) on the standalone service, unknown_key mapping to 404. The business key addresses; the sender's idempotency key deduplicates -- the same two-identity split the webhook ingress already lives by. Both routes go through kernel.signal, so channel validation, model canonicalization, and dispositions stay identical to Python's by construction. Contract gains the addressing note; OpenAPI documents both paths. --- news/workflow-by-key.md | 1 + reflex/workflow/CONTRACT.md | 13 +++ reflex/workflow/kernel.py | 70 +++++++++++++ reflex/workflow/runtime.py | 51 ++++++++++ reflex/workflow/serve.py | 125 +++++++++++++++++++++++- reflex/workflow/store.py | 1 + tests/units/workflow/test_serve.py | 55 +++++++++++ tests/units/workflow/test_validation.py | 58 +++++++++++ 8 files changed, 373 insertions(+), 1 deletion(-) create mode 100644 news/workflow-by-key.md diff --git a/news/workflow-by-key.md b/news/workflow-by-key.md new file mode 100644 index 00000000000..46ae6d87d42 --- /dev/null +++ b/news/workflow-by-key.md @@ -0,0 +1 @@ +Runs are addressable by business key. The request key was already a durable unique index — it is what makes webhook redelivery idempotent — so it now doubles as the run's business address: `rx.workflows.get_by_key(Order, "order_123")` returns a handle and `rx.workflows.signal_by_key(Order, "order_123", Order.shipped(payload), key=event_id)` delivers to the run that key admitted, with `"unknown_key"` when nothing did. Over HTTP the same pair is `GET /workflows/{workflow}/keys/{request_key}` (read scope) and `POST /workflows/{workflow}/keys/{request_key}/signals/{channel}` (signal scope), mapping `unknown_key` to 404. The business key addresses; the sender's idempotency key still deduplicates. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 632a555ae94..cce5f90d71c 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -260,6 +260,19 @@ Checked in this order at start: refuses its signal rather than letting it resolve a later wait on the same channel. +### Business-key addressing + +The request key is a durable unique index per workflow (`§6`), so it doubles +as a run's business address: `rx.workflows.get_by_key(Order, "order_123")` +and `rx.workflows.signal_by_key(Order, "order_123", Order.shipped(p), +key=event_id)` reach the run the key admitted, without the caller ever +storing the engine's run id. Over HTTP the same pair is +`GET /workflows/{id}/keys/{request_key}` and `POST .../signals/{channel}`. +A key that admitted nothing answers `unknown_key` (HTTP 404) — the caller +decides whether that means "not yet" (buffer upstream, or start the run) or +"never". Signal dedupe stays the sender's idempotency key; the business key +addresses, the event id deduplicates. + ## 7. Workers - A **claim** takes the run's frontier step (lowest unresolved ordinal) — diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index bf75c283903..8d04ad38d9e 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -957,6 +957,76 @@ async def cancel(self, run_id: str) -> bool: self._wakeup.set() return recorded + def _workflow_id_of(self, workflow: Any) -> str: + """Resolve a workflow argument to its stable identity. + + Args: + workflow: A registered workflow class, or a workflow id string. + + Returns: + The workflow id. + + Raises: + WorkflowRuntimeError: If the workflow is not registered here. + """ + if isinstance(workflow, str): + if workflow in self._definitions: + return workflow + else: + defn = self._definitions_by_cls.get(workflow) + if defn is not None: + return defn.workflow_id + known = ", ".join(sorted(self._definitions)) or "" + msg = ( + f"Workflow {workflow!r} is not registered with this runtime; " + f"registered: {known}." + ) + raise WorkflowRuntimeError(msg) + + async def find_by_key(self, workflow: Any, request_key: str) -> str | None: + """Find the run a business key admitted, if any. + + The request key is already a durable unique index -- it is what makes + webhook redelivery idempotent -- so it doubles as the business + address of a run: ``order_123`` finds the order's run without anyone + having threaded the engine's run id through their own tables. + + Args: + workflow: The registered workflow class or its id. + request_key: The admission key the run was started under. + + Returns: + The run id, or None when the key admitted nothing. + """ + return await self._store.find_by_request_key( + self._workflow_id_of(workflow), request_key + ) + + async def signal_by_key( + self, + workflow: Any, + request_key: str, + delivery: ChannelDelivery, + *, + key: str | None = None, + ) -> DeliveryDisposition: + """Deliver a signal to the run a business key admitted. + + Args: + workflow: The registered workflow class or its id. + request_key: The admission key the run was started under. + delivery: The addressed payload, e.g. ``Order.shipped(payload)``. + key: Sender idempotency key; a repeated key is a no-op. + + Returns: + What the store did with the delivery, or ``"unknown_key"`` when + the key admitted nothing. + """ + run_id = await self.find_by_key(workflow, request_key) + if run_id is None: + return "unknown_key" + return await self.signal(run_id, delivery, key=key) + async def signal( self, run_id: str, diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index c6e8964c95e..ae350cf89b7 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -448,6 +448,57 @@ async def submit( raise WorkflowRuntimeError(msg) return RunHandle(result.run_id, result.disposition) + @staticmethod + async def get_by_key(workflow: Any, request_key: str) -> RunHandle[Any] | None: + """Find the run a business key admitted, as a handle. + + The request key is already a durable unique index -- it is what makes + redelivery idempotent -- so it doubles as the business address of a + run: ``order_123`` finds the order's run without anyone having stored + the engine's run id:: + + handle = await rx.workflows.get_by_key(Order, "order_123") + + Args: + workflow: The registered workflow class or its id. + request_key: The admission key the run was started under. + + Returns: + A handle on the run, or None when the key admitted nothing. + """ + run_id = await get_runtime().kernel.find_by_key(workflow, request_key) + return None if run_id is None else RunHandle(run_id, "found") + + @staticmethod + async def signal_by_key( + workflow: Any, + request_key: str, + delivery: Any, + *, + key: str | None = None, + ) -> str: + """Deliver a signal to the run a business key admitted. + + Usage:: + + await rx.workflows.signal_by_key( + Order, "order_123", Order.shipped(payload), key=event_id + ) + + Args: + workflow: The registered workflow class or its id. + request_key: The admission key the run was started under. + delivery: The addressed payload, e.g. ``Order.shipped(payload)``. + key: Sender idempotency key; a repeated key is a no-op. + + Returns: + The delivery disposition, ``"unknown_key"`` when the key admitted + nothing. + """ + return await get_runtime().kernel.signal_by_key( + workflow, request_key, delivery, key=key + ) + @staticmethod async def cancel(run_id: str) -> bool: """Request cancellation of a run. diff --git a/reflex/workflow/serve.py b/reflex/workflow/serve.py index 42dcc7ccb56..88d8f319156 100644 --- a/reflex/workflow/serve.py +++ b/reflex/workflow/serve.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any, Final from reflex_base.utils import console -from reflex_base.utils.exceptions import WorkflowDefinitionError +from reflex_base.utils.exceptions import WorkflowDefinitionError, WorkflowRuntimeError from reflex_base.workflow import ChannelDelivery from starlette.applications import Starlette from starlette.responses import JSONResponse @@ -274,6 +274,104 @@ async def endpoint(request: Request) -> JSONResponse: return endpoint +def key_read_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens): + """Build the endpoint that reads a run by its business key. + + Args: + runtime: The runtime owning the runs. + tokens: The service's token scopes. + + Returns: + The endpoint callable. + """ + authorize = tokens.require("read") + + async def endpoint(request: Request) -> JSONResponse: + """Resolve the key and answer with the run's identity and status. + + Args: + request: The incoming request. + + Returns: + The run reference, or 404. + """ + refused = authorize(request) + if refused is not None: + return refused + workflow_id = request.path_params["workflow_id"] + request_key = request.path_params["request_key"] + try: + run_id = await runtime.kernel.find_by_key(workflow_id, request_key) + except WorkflowRuntimeError as error: + return JSONResponse({"error": str(error)}, status_code=404) + if run_id is None: + return JSONResponse({"error": "unknown key"}, status_code=404) + snapshot = await runtime.kernel.get_run(run_id) + if snapshot is None: + return JSONResponse({"error": "unknown key"}, status_code=404) + return JSONResponse({ + "run_id": snapshot.run_id, + "workflow": snapshot.workflow_id, + "status": snapshot.status.value, + "result": snapshot.result, + "error": snapshot.error, + }) + + return endpoint + + +def key_signal_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens): + """Build the endpoint that signals a run by its business key. + + Args: + runtime: The runtime owning the runs. + tokens: The service's token scopes. + + Returns: + The endpoint callable. + """ + authorize = tokens.require("signal") + + async def endpoint(request: Request) -> JSONResponse: + """Deliver the request body to the keyed run's channel. + + Args: + request: The incoming request. + + Returns: + The delivery disposition, or an error. + """ + refused = authorize(request) + if refused is not None: + return refused + payload, bad = await _read_json(request) + if bad is not None: + return bad + key = request.headers.get("idempotency-key") or request.query_params.get("key") + try: + disposition = await runtime.kernel.signal_by_key( + request.path_params["workflow_id"], + request.path_params["request_key"], + ChannelDelivery( + channel=request.path_params["channel"], payload=payload + ), + key=key, + ) + except WorkflowRuntimeError as error: + return JSONResponse({"error": str(error)}, status_code=404) + except WorkflowDefinitionError as error: + return JSONResponse({"error": str(error)}, status_code=400) + status = { + "unknown_key": 404, + "unknown_run": 404, + "run_terminal": 409, + "expired": 409, + }.get(disposition, 202) + return JSONResponse({"disposition": disposition}, status_code=status) + + return endpoint + + def operator_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens, action: str): """Build one operator action endpoint. @@ -463,6 +561,18 @@ async def endpoint(request: Request) -> JSONResponse: # noqa: RUF029 } for action in ("cancel", "retry", "resume") }, + "/workflows/{workflow_id}/keys/{request_key}": { + "get": { + "summary": "Read a run by business key (scope: read)", + "responses": {"200": {"description": "The run"}}, + } + }, + "/workflows/{workflow_id}/keys/{request_key}/signals/{channel}": { + "post": { + "summary": ("Deliver a signal by business key (scope: signal)"), + "responses": {"202": {"description": "Delivered"}}, + } + }, "/healthz": {"get": {"summary": "Liveness", "security": []}}, "/readyz": {"get": {"summary": "Readiness", "security": []}}, "/metrics": {"get": {"summary": "Prometheus metrics (scope: read)"}}, @@ -572,6 +682,19 @@ async def lifespan(app: Starlette) -> AsyncIterator[None]: ) for action in ("cancel", "retry", "resume") ), + # Business-key addressing over the durable request-key index: the + # caller that knows "order_123" reaches the order's run without + # ever having stored the engine's run id. + Route( + "/workflows/{workflow_id}/keys/{request_key}", + key_read_endpoint(runtime, tokens), + methods=["GET"], + ), + Route( + "/workflows/{workflow_id}/keys/{request_key}/signals/{channel}", + key_signal_endpoint(runtime, tokens), + methods=["POST"], + ), # The embedded-mode paths, kept byte-for-byte: a Stripe URL or a # minted approval link configured against an rx.App keeps working # when the deployment moves to the standalone service. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index ae5923bc142..0562f8785a9 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -56,6 +56,7 @@ "duplicate", "unknown_run", "run_terminal", + "unknown_key", ] diff --git a/tests/units/workflow/test_serve.py b/tests/units/workflow/test_serve.py index 9842c1c52ca..69df51a94a1 100644 --- a/tests/units/workflow/test_serve.py +++ b/tests/units/workflow/test_serve.py @@ -383,3 +383,58 @@ def on_paid(self, id: str): headers={"Stripe-Signature": f"t={timestamp},v1={digest}"}, ) assert redelivered.json()["disposition"] == "deduplicated" + + +def test_business_keys_address_runs_without_run_ids(service): + """`order_123` reaches the order's run; nobody stored a run id. + + Args: + service: The served client. + """ + started = service.post( + "/runs", + json={ + "workflow": "serve.orders", + "handler": "place", + "args": {"order_id": "o1"}, + "request_key": "order_123", + }, + headers=_auth("tk_start"), + ) + assert started.status_code == 202 + run_id = started.json()["run_id"] + + found = service.get( + "/workflows/serve.orders/keys/order_123", headers=_auth("tk_read") + ) + assert found.status_code == 200 + assert found.json()["run_id"] == run_id + + missing = service.get( + "/workflows/serve.orders/keys/order_999", headers=_auth("tk_read") + ) + assert missing.status_code == 404 + + delivered = service.post( + "/workflows/serve.orders/keys/order_123/signals/shipped", + json={"parcel": "P-9"}, + headers={**_auth("tk_signal"), "Idempotency-Key": "evt_9"}, + ) + assert delivered.status_code == 202 + redelivered = service.post( + "/workflows/serve.orders/keys/order_123/signals/shipped", + json={"parcel": "P-9"}, + headers={**_auth("tk_signal"), "Idempotency-Key": "evt_9"}, + ) + assert redelivered.json()["disposition"] == "duplicate" + unkeyed = service.post( + "/workflows/serve.orders/keys/order_999/signals/shipped", + json={}, + headers=_auth("tk_signal"), + ) + assert unkeyed.status_code == 404 + assert unkeyed.json()["disposition"] == "unknown_key" + unknown_workflow = service.get( + "/workflows/serve.nope/keys/order_123", headers=_auth("tk_read") + ) + assert unknown_workflow.status_code == 404 diff --git a/tests/units/workflow/test_validation.py b/tests/units/workflow/test_validation.py index 92a95d69fa7..d7f21e34aa8 100644 --- a/tests/units/workflow/test_validation.py +++ b/tests/units/workflow/test_validation.py @@ -516,3 +516,61 @@ def finish(self, count: Shipment): assert pending[0].attempts == 0, ( "suspension must not consume the attempt budget" ) + + +async def test_python_by_key_lookup_and_signal(forked_registration_context): + """The Python forms of business-key addressing mirror the HTTP ones. + + Args: + forked_registration_context: Isolated state registry. + """ + + class Keyed(rx.State): + __workflow__ = WorkflowConfig(id="validation.keyed") + go = Signal() + + @rx.event(durable=True, effect="none", trigger=manual()) + def begin(self): + """Wait for the go signal. + + Returns: + The wait. + """ + return rx.wait_for(Keyed.go, then=Keyed.done, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def done(self, payload): + """Finish. + + Args: + payload: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result=payload) + + async with WorkflowTestHarness(Keyed) as harness: + result = await harness.start(Keyed.begin(), request_key="order_7") + assert result.run_id is not None + await harness.run_until_idle() + + assert await harness.kernel.find_by_key(Keyed, "order_7") == result.run_id + assert await harness.kernel.find_by_key("validation.keyed", "order_7") == ( + result.run_id + ) + assert await harness.kernel.find_by_key(Keyed, "order_8") is None + assert ( + await harness.kernel.signal_by_key( + Keyed, "order_8", Keyed.go(None), key="e1" + ) + == "unknown_key" + ) + assert ( + await harness.kernel.signal_by_key( + Keyed, "order_7", Keyed.go(None), key="e1" + ) + == "resolved" + ) + with pytest.raises(Exception, match="not registered"): + await harness.kernel.find_by_key("validation.nope", "order_7") From 16349de8ddc4edb2c4407ce62b227f8fc3a57d8d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Mon, 24 Aug 2026 17:43:19 -0700 Subject: [PATCH 117/121] workflows: the channel inbox -- durable correlated delivery Phase 2's foundation, on all three stores: a webhook event addressed to a signal channel by business key is durable from the moment it is acknowledged, delivered exactly once whenever its run exists, and visible when it cannot be. ingest_channel_delivery makes one transaction of the whole decision. The row keyed by (workflow, channel, correlation_key, event id) is written first, so the ack and the record are one fact: a provider redelivery and a crash-after-ack replay both collapse into "duplicate". If the correlation key already admitted a run, the payload is delivered in the same transaction -- through the same per-run inbox and dedupe as every other signal. If not, the row waits PENDING. Admission flushes parked mail inside the admitting transaction, on both doors -- plain admit and policy admit_flow -- so a crash cannot separate "the run exists" from "its early mail reached it". On Postgres the channel rows are locked before the run row on every path (ingest, admit, admit_flow, replay), one order, no cycle. A delivery nothing can take is a dead letter, never a silent drop: a terminal or past-deadline run at ingest, or a PENDING row unclaimed past its TTL (sweep_parked). Dead letters carry their reason, list via list_parked, and replay via replay_parked with the same idempotency -- replaying a delivered row is "duplicate", never a second signal. Store surface: ingest_channel_delivery, list_parked, replay_parked, sweep_parked; new dispositions "parked" and "dead_letter" (documented in the contract vocabulary, which this repo's guard test enforces); SQLite schema v4 adds workflow_channel_inbox; Postgres adds the same table to its advisory-locked DDL. deliver() on SQLite and Postgres is factored into an in-transaction form both admission and ingest share, with refusal branches that write nothing so the caller's transaction stays committable. Six conformance checks pin it everywhere, including the acceptance flow: park before the run exists, redeliver three times, admit the run -- the wait resolves with exactly one payload, and the late redelivery is still "duplicate". Ingress wiring (rx.Signal(trigger=rx.webhook(..., correlate_by=...)) routing) and the SIGKILL crash test ride the next commit; the base API for it (correlate_by on webhook(), Signal.trigger validation) is already in. --- news/workflow-channel-inbox.md | 1 + .../reflex-base/src/reflex_base/workflow.py | 44 +- reflex/workflow/CONTRACT.md | 10 +- reflex/workflow/conformance.py | 139 +++ reflex/workflow/postgres.py | 503 +++++++-- reflex/workflow/records.py | 45 + reflex/workflow/store.py | 954 +++++++++++++++--- 7 files changed, 1474 insertions(+), 222 deletions(-) create mode 100644 news/workflow-channel-inbox.md diff --git a/news/workflow-channel-inbox.md b/news/workflow-channel-inbox.md new file mode 100644 index 00000000000..47886afc19c --- /dev/null +++ b/news/workflow-channel-inbox.md @@ -0,0 +1 @@ +The durable core of correlated webhook delivery: a channel inbox on every store. `ingest_channel_delivery` accepts a provider event addressed to a workflow channel by business key and makes one transaction of the whole decision — the row keyed by the provider's event id is written first, so acknowledging the provider and recording the event are one fact, and a crash after the ack replays as `duplicate`, never a second signal. A delivery whose run already exists lands immediately; one that arrives early waits `PENDING` and is flushed **inside the admitting transaction** — through either admission door, plain or policy — so a crash cannot separate "the run exists" from "its early mail reached it". A delivery nothing can take (terminal run, past deadline, or unclaimed past TTL) becomes a visible dead letter with a reason, listable and replayable with the same event-id idempotency. Six conformance checks pin the semantics on Memory, SQLite, and Postgres, including the acceptance flow: park before the run, redeliver three times, admit — exactly one signal. diff --git a/packages/reflex-base/src/reflex_base/workflow.py b/packages/reflex-base/src/reflex_base/workflow.py index 36ce1345871..2c8bae065ed 100644 --- a/packages/reflex-base/src/reflex_base/workflow.py +++ b/packages/reflex-base/src/reflex_base/workflow.py @@ -267,6 +267,11 @@ class WebhookTrigger(Trigger): request genuinely came from the provider. dedupe_by: Payload field used as the ingress deduplication key, so a provider redelivering an event does not start a second run. + correlate_by: Payload field carrying the business key of the run this + delivery belongs to. Only meaningful on a channel trigger: the + value is matched against runs' request keys, so + ``correlate_by="order_id"`` routes a shipment event to the order + run started under ``request_key="order_123"``. allow_unverified: Acknowledge that this endpoint accepts anonymous traffic. Only valid with a non-empty ``unverified_reason``. unverified_reason: Why anonymous traffic is acceptable here. @@ -278,6 +283,7 @@ class WebhookTrigger(Trigger): model: type | None = None verify: WebhookVerifier | None = None dedupe_by: str | None = None + correlate_by: str | None = None allow_unverified: bool = False unverified_reason: str = "" @@ -354,10 +360,11 @@ def webhook( model: type | None = None, verify: WebhookVerifier | None = None, dedupe_by: str | None = None, + correlate_by: str | None = None, allow_unverified: bool = False, unverified_reason: str = "", ) -> WebhookTrigger: - """Create a webhook trigger for a workflow root handler. + """Create a webhook trigger for a root handler or a signal channel. Args: topic: Stable provider event topic. @@ -365,6 +372,9 @@ def webhook( verify: Callable given the raw body and headers that returns whether the request genuinely came from the provider. dedupe_by: Payload field used as the ingress deduplication key. + correlate_by: Payload field carrying the business key of the run the + delivery belongs to; used by channel triggers to route the event + to the run started under that request key. allow_unverified: Acknowledge that this endpoint accepts anonymous traffic. unverified_reason: Why anonymous traffic is acceptable here. @@ -376,6 +386,7 @@ def webhook( model=model, verify=verify, dedupe_by=dedupe_by, + correlate_by=correlate_by, allow_unverified=allow_unverified, unverified_reason=unverified_reason, ) @@ -980,15 +991,44 @@ class Onboarding(rx.State): name: The channel name, defaulting to the attribute name. """ - def __init__(self, model: type | None = None, *, name: str | None = None): + def __init__( + self, + model: type | None = None, + *, + name: str | None = None, + trigger: WebhookTrigger | None = None, + ): """Declare a channel. Args: model: The payload model deliveries must satisfy. name: Explicit channel name; defaults to the attribute name. + trigger: A webhook that delivers into this channel. It must name + both identities a correlated delivery needs: ``correlate_by`` + for which run, ``dedupe_by`` for which event. + + Raises: + WorkflowDefinitionError: If a trigger is given without + ``correlate_by`` and ``dedupe_by``, or is not a webhook. """ self.model = model self.name = name or "" + self.trigger = trigger + if trigger is not None: + if not isinstance(trigger, WebhookTrigger): + msg = ( + "rx.Signal(trigger=...) takes rx.webhook(...); a channel " + "is delivered into, so no other trigger kind can feed it." + ) + raise WorkflowDefinitionError(msg) + if trigger.correlate_by is None or trigger.dedupe_by is None: + msg = ( + f"Channel webhook {trigger.topic!r} needs correlate_by " + "(which run this event belongs to) and dedupe_by (which " + "event this delivery is); without both, a delivery " + "cannot be routed exactly once." + ) + raise WorkflowDefinitionError(msg) def __set_name__(self, owner: type, name: str) -> None: """Adopt the attribute name as the channel name. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index cce5f90d71c..e07a750ab13 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -528,7 +528,15 @@ it under a lease), `RETRY_WAIT` (business retry scheduled), `RECOVERY_WAIT` **Delivery disposition** — what the store did with a signal or arrival: `resolved`, `buffered` (arrived before its wait was armed), `counted` (a join arrival that is not the last), `duplicate` (repeated sender key), -`expired` (run past its deadline), `unknown_run`, `run_terminal`. +`expired` (run past its deadline), `unknown_run`, `run_terminal`, +`unknown_key` (a business key that admitted nothing), `parked` (a correlated +webhook delivery accepted before its run exists — durable in the channel +inbox, flushed inside the admitting transaction when the run arrives), and +`dead_letter` (a correlated delivery nothing can take: its run is terminal or +past deadline, or the parked delivery went unclaimed past its TTL; visible +via the channel inbox and replayable by an operator with the same event-id +idempotency, so a replay of a delivered row is a `duplicate`, never a second +signal). **History events.** Append-only, one run's whole story: diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index b32d9f54673..526d2eb8316 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -25,6 +25,7 @@ async def test_my_store_conforms(check): from reflex.workflow.records import ( TERMINAL_RUN_STATUSES, HistoryEventType, + ParkedStatus, RunQuery, RunRecord, RunStatus, @@ -1709,6 +1710,139 @@ async def check_none_is_a_legal_payload_everywhere(store: RunStore) -> None: assert await store.get_substeps("sub1", 0) == {"notify": None} +async def check_a_delivery_before_its_run_lands_exactly_once( + store: RunStore, +) -> None: + """The Phase 2 acceptance flow, minus the crash: park, redeliver, admit. + + A shipment event arrives before the order workflow exists, the provider + sends it three times, and the run starts later. Exactly one signal must + reach the run -- the channel-inbox row keyed by the provider's event id + is what collapses redelivery and crash-after-ack replays into one fact. + """ + first = await store.ingest_channel_delivery( + "conformance.flow", "shipped", "order_1", "evt_1", {"parcel": "P1"}, NOW + ) + assert first == "parked", "no run exists yet; the delivery must wait" + for _ in range(2): + again = await store.ingest_channel_delivery( + "conformance.flow", "shipped", "order_1", "evt_1", {"parcel": "P1"}, NOW + ) + assert again == "duplicate", "a redelivery is the same event" + + await store.admit( + make_run(request_key="order_1"), + make_step(status=StepStatus.BLOCKED, wait_key="sig:shipped", due_at=0.0), + _ADMITTED, + ) + steps = await store.get_steps("run1") + assert steps[0].status is StepStatus.READY, "the parked payload resolved it" + assert steps[0].args["__payload__"] == {"parcel": "P1"} + + late = await store.ingest_channel_delivery( + "conformance.flow", "shipped", "order_1", "evt_1", {"parcel": "P1"}, NOW + 1 + ) + assert late == "duplicate", "the event id survives delivery" + rows = await store.list_parked(workflow_id="conformance.flow") + assert len(rows) == 1 + assert rows[0].status is ParkedStatus.DELIVERED + assert rows[0].run_id == "run1" + + +async def check_a_delivery_to_a_live_run_lands_immediately(store: RunStore) -> None: + """With the run already waiting, ingest is an ordinary delivery. + + The channel-inbox row is still written -- it is the event-id dedupe for + every later redelivery -- but the payload goes straight through. + """ + await store.admit( + make_run(request_key="order_2"), + make_step(status=StepStatus.BLOCKED, wait_key="sig:shipped", due_at=0.0), + _ADMITTED, + ) + landed = await store.ingest_channel_delivery( + "conformance.flow", "shipped", "order_2", "evt_2", {"n": 2}, NOW + ) + assert landed == "resolved" + again = await store.ingest_channel_delivery( + "conformance.flow", "shipped", "order_2", "evt_2", {"n": 2}, NOW + ) + assert again == "duplicate" + steps = await store.get_steps("run1") + assert steps[0].args["__payload__"] == {"n": 2} + + +async def check_a_dead_letter_is_visible_and_replayable(store: RunStore) -> None: + """A delivery nothing can take becomes an operator's problem, loudly. + + The run is terminal, so the payload can never land -- but silence would + read as delivered. The row dies visibly, and after the operator revives + the run, replay routes the same row with the same idempotency. + """ + await store.admit( + make_run(request_key="order_3", status=RunStatus.FAILED), + make_step(status=StepStatus.FAILED, error={"reason": "boom"}), + _ADMITTED, + ) + dead = await store.ingest_channel_delivery( + "conformance.flow", "shipped", "order_3", "evt_3", {"n": 3}, NOW + ) + assert dead == "dead_letter" + rows = await store.list_parked(status=ParkedStatus.DEAD) + assert len(rows) == 1 + assert rows[0].reason == "run_terminal" + + assert await store.retry_run("run1", NOW + 1) + replayed = await store.replay_parked(rows[0].parked_id, NOW + 2) + assert replayed == "buffered", "the revived run has no wait open yet" + rows = await store.list_parked(workflow_id="conformance.flow") + assert rows[0].status is ParkedStatus.DELIVERED + assert await store.replay_parked(rows[0].parked_id, NOW + 3) == "duplicate", ( + "replaying a delivered row must never signal twice" + ) + assert await store.replay_parked("missing", NOW) == "unknown_key" + + +async def check_unclaimed_deliveries_become_dead_letters(store: RunStore) -> None: + """A parked delivery whose run never arrives surfaces, not lingers.""" + await store.ingest_channel_delivery( + "conformance.flow", "shipped", "order_4", "evt_4", {"n": 4}, NOW + ) + assert await store.sweep_parked(NOW + 100, ttl=3600) == 0, "not yet unclaimed" + assert await store.sweep_parked(NOW + 4000, ttl=3600) == 1 + rows = await store.list_parked(status=ParkedStatus.DEAD) + assert len(rows) == 1 + assert rows[0].reason == "unclaimed" + assert await store.sweep_parked(NOW + 5000, ttl=3600) == 0, "dead rows stay dead" + + +async def check_policy_admission_also_flushes_parked_mail(store: RunStore) -> None: + """A run admitted through a start policy still receives its early mail. + + Policy admission is a second door into existence; a delivery parked + before the run must not depend on which door the run came through. + """ + from reflex.workflow.store import FlowGate + + parked = await store.ingest_channel_delivery( + "conformance.flow", "shipped", "order_5", "evt_5", {"n": 5}, NOW + ) + assert parked == "parked" + admission = await store.admit_flow( + make_run(request_key="order_5", flow_key="order_5"), + make_step(status=StepStatus.BLOCKED, wait_key="sig:shipped", due_at=0.0), + _ADMITTED, + FlowGate(), + NOW, + ) + assert admission.disposition == "started" + steps = await store.get_steps(admission.run_id or "run1") + assert steps[0].status is StepStatus.READY + assert steps[0].args["__payload__"] == {"n": 5} + rows = await store.list_parked(workflow_id="conformance.flow") + assert rows[0].status is ParkedStatus.DELIVERED + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1726,6 +1860,11 @@ async def check_none_is_a_legal_payload_everywhere(store: RunStore) -> None: check_delivery_resolves_a_matching_wait, check_delivery_never_touches_run_state, check_duplicate_deliveries_are_ignored, + check_a_delivery_before_its_run_lands_exactly_once, + check_a_delivery_to_a_live_run_lands_immediately, + check_a_dead_letter_is_visible_and_replayable, + check_unclaimed_deliveries_become_dead_letters, + check_policy_admission_also_flushes_parked_mail, check_a_delivery_to_a_past_deadline_run_is_refused, check_a_duplicate_delivery_is_recorded_in_history, check_an_early_delivery_is_buffered_then_consumed, diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index e97ee8a59f9..85f6fbf8b2d 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -16,6 +16,7 @@ import asyncio import dataclasses +import uuid from typing import TYPE_CHECKING, Any, Final from reflex_base.utils.exceptions import WorkflowRuntimeError @@ -27,6 +28,8 @@ TERMINAL_STEP_STATUSES, HistoryEvent, HistoryEventType, + ParkedDelivery, + ParkedStatus, RunRecord, RunStatus, StepRecord, @@ -136,6 +139,22 @@ created_at DOUBLE PRECISION NOT NULL, PRIMARY KEY (run_id, ordinal, key) ); +CREATE TABLE IF NOT EXISTS workflow_channel_inbox ( + parked_id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + channel TEXT NOT NULL, + correlation_key TEXT NOT NULL, + dedupe_key TEXT NOT NULL, + payload JSONB NOT NULL, + status TEXT NOT NULL, + reason TEXT, + run_id TEXT, + created_at DOUBLE PRECISION NOT NULL, + updated_at DOUBLE PRECISION NOT NULL, + UNIQUE (workflow_id, channel, correlation_key, dedupe_key) +); +CREATE INDEX IF NOT EXISTS idx_workflow_channel_inbox_route + ON workflow_channel_inbox (workflow_id, correlation_key, status); CREATE TABLE IF NOT EXISTS workflow_inbox ( run_id TEXT NOT NULL, wait_key TEXT NOT NULL, @@ -214,6 +233,30 @@ def _json(value: Any) -> Any: return None if value is None else Jsonb(value) +def _parked_from_row(row: Mapping[str, Any]) -> ParkedDelivery: + """Build a parked-delivery record from a database row. + + Args: + row: The ``workflow_channel_inbox`` row. + + Returns: + The record. + """ + return ParkedDelivery( + parked_id=row["parked_id"], + workflow_id=row["workflow_id"], + channel=row["channel"], + correlation_key=row["correlation_key"], + dedupe_key=row["dedupe_key"], + payload=row["payload"], + status=ParkedStatus(row["status"]), + reason=row["reason"], + run_id=row["run_id"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + def _run_from_row(row: Mapping[str, Any]) -> RunRecord: """Build a run record from a database row. @@ -676,6 +719,9 @@ async def admit( pool = await self._open() async with pool.connection() as conn, conn.transaction(): if run.request_key is not None: + # Channel-inbox rows before the run: ingest locks its row and + # then the run, so admission takes them in the same order. + await self._lock_parked_conn(conn, run.workflow_id, run.request_key) cursor = await conn.execute( "INSERT INTO workflow_dedupe (workflow_id, request_key, run_id)" " VALUES (%s, %s, %s) ON CONFLICT DO NOTHING RETURNING run_id", @@ -693,6 +739,17 @@ async def admit( await self._insert_run(conn, run) await self._insert_step(conn, root_step) await self._append_events(conn, run.run_id, events, run.created_at) + if run.request_key is not None: + # Deliveries that arrived before this run did, flushed inside + # the admitting transaction: a crash cannot separate "the run + # exists" from "its early mail reached it". + await self._flush_parked_conn( + conn, + run.workflow_id, + run.request_key, + run.run_id, + run.created_at, + ) return True, run.run_id async def admit_flow( @@ -825,9 +882,23 @@ async def admit_flow( if cursor.rowcount: return FlowAdmission("coalesced", active[0]["run_id"]) due_at = now + gate.debounce + if run.request_key is not None: + # Channel-inbox rows before the run row, matching ingest's + # order. + await self._lock_parked_conn(conn, run.workflow_id, run.request_key) await self._insert_run(conn, run) await self._insert_step(conn, dataclasses.replace(root_step, due_at=due_at)) await self._append_events(conn, run.run_id, events, run.created_at) + if run.request_key is not None: + # Policy admission is still admission: early mail flushes on + # this door exactly as on the plain one. + await self._flush_parked_conn( + conn, + run.workflow_id, + run.request_key, + run.run_id, + run.created_at, + ) return FlowAdmission("started", run.run_id, cancelled=tuple(cancelled)) async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> int: @@ -1200,92 +1271,376 @@ async def deliver( """ pool = await self._open() async with pool.connection() as conn, conn.transaction(): - cursor = await conn.execute( - "SELECT status, deadline FROM workflow_runs WHERE run_id = %s" - " FOR UPDATE", - (run_id,), + return await self._deliver_with( + conn, run_id, wait_key, dedupe_key, payload, now ) - row = await cursor.fetchone() - if row is None: - return "unknown_run" - if row["status"] in _TERMINAL_RUNS: - return "run_terminal" - if row["deadline"] is not None and row["deadline"] <= now: - # Claims exclude past-deadline runs and the sweep is about to - # finalize this one TIMED_OUT, so "resolved" would tell the - # sender -- often a person clicking approve -- that their - # decision landed, moments before it is discarded. - return "expired" - cursor = await conn.execute( - "SELECT 1 FROM workflow_inbox" - " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", - (run_id, wait_key, dedupe_key), + + async def _deliver_with( + self, + conn: Any, + run_id: str, + wait_key: str, + dedupe_key: str, + payload: Any, + now: float, + ) -> DeliveryDisposition: + """Deliver inside the caller's open transaction. + + Takes the run row here, so a caller holding channel-inbox rows keeps + the canonical channel-before-run lock order. Refusal branches write + nothing, so the caller's transaction stays committable. + + Args: + conn: The connection inside an open transaction. + run_id: The receiving run. + wait_key: The address the waiting slot declared. + dedupe_key: Sender-supplied identity, making redelivery a no-op. + payload: JSON-compatible payload to hand the resuming handler. + now: Current time in epoch seconds. + + Returns: + What the store did with the delivery. + """ + cursor = await conn.execute( + "SELECT status, deadline FROM workflow_runs WHERE run_id = %s FOR UPDATE", + (run_id,), + ) + row = await cursor.fetchone() + if row is None: + return "unknown_run" + if row["status"] in _TERMINAL_RUNS: + return "run_terminal" + if row["deadline"] is not None and row["deadline"] <= now: + # Claims exclude past-deadline runs and the sweep is about to + # finalize this one TIMED_OUT, so "resolved" would tell the + # sender -- often a person clicking approve -- that their + # decision landed, moments before it is discarded. + return "expired" + cursor = await conn.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = %s AND wait_key = %s AND dedupe_key = %s", + (run_id, wait_key, dedupe_key), + ) + if await cursor.fetchone() is not None: + await self._append_events( + conn, + run_id, + ((HistoryEventType.SIGNAL_DUPLICATE, {"wait_key": wait_key}),), + now, ) - if await cursor.fetchone() is not None: - await self._append_events( - conn, - run_id, - ((HistoryEventType.SIGNAL_DUPLICATE, {"wait_key": wait_key}),), + return "duplicate" + frontier = await self._frontier(conn, run_id) + if ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and 0.0 < frontier.due_at <= now + ): + return "expired" + resolves = ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and frontier.wait_key == wait_key + ) + await conn.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at) VALUES (%s, %s, %s, %s, %s, %s, %s)", + ( + run_id, + wait_key, + dedupe_key, + await self._next_inbox_seq(conn, run_id), + Jsonb(payload), + "CONSUMED" if resolves else "PENDING", + now, + ), + ) + if resolves and frontier is not None: + await conn.execute( + "UPDATE workflow_steps SET status = %s, due_at = %s, args = %s," + " updated_at = %s WHERE run_id = %s AND ordinal = %s", + ( + StepStatus.READY.value, now, - ) - return "duplicate" - frontier = await self._frontier(conn, run_id) - if ( - frontier is not None - and frontier.status is StepStatus.BLOCKED - and 0.0 < frontier.due_at <= now - ): - return "expired" - resolves = ( - frontier is not None - and frontier.status is StepStatus.BLOCKED - and frontier.wait_key == wait_key + _json({**frontier.args, "__payload__": payload}), + now, + run_id, + frontier.ordinal, + ), + ) + await self._append_events( + conn, + run_id, + ( + ( + HistoryEventType.WAIT_RESOLVED, + {"ordinal": frontier.ordinal, "wait_key": wait_key}, + ), + ), + now, + ) + else: + await self._append_events( + conn, + run_id, + ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), + now, ) + return "resolved" if resolves else "buffered" + + async def _flush_parked_conn( + self, conn: Any, workflow_id: str, request_key: str, run_id: str, now: float + ) -> None: + """Deliver PENDING channel-inbox rows to a freshly admitted run. + + The caller must have locked these rows (``_lock_parked_conn``) + before creating the run, keeping the canonical channel-before-run + order. + + Args: + conn: The connection inside the admitting transaction. + workflow_id: The workflow identity. + request_key: The admission key, matched against correlation keys. + run_id: The run that now exists. + now: Current time in epoch seconds. + """ + cursor = await conn.execute( + "SELECT parked_id, channel, dedupe_key, payload FROM" + " workflow_channel_inbox WHERE workflow_id = %s AND" + " correlation_key = %s AND status = %s ORDER BY created_at", + (workflow_id, request_key, ParkedStatus.PENDING.value), + ) + for row in await cursor.fetchall(): + disposition = await self._deliver_with( + conn, + run_id, + f"sig:{row['channel']}", + row["dedupe_key"], + row["payload"], + now, + ) + delivered = disposition in ("resolved", "buffered", "duplicate") await conn.execute( - "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," - " payload, status, created_at) VALUES (%s, %s, %s, %s, %s, %s, %s)", + "UPDATE workflow_channel_inbox SET status = %s, reason = %s," + " run_id = %s, updated_at = %s WHERE parked_id = %s", ( - run_id, - wait_key, + ParkedStatus.DELIVERED.value + if delivered + else ParkedStatus.DEAD.value, + None if delivered else disposition, + run_id if delivered else None, + now, + row["parked_id"], + ), + ) + + async def _lock_parked_conn( + self, conn: Any, workflow_id: str, request_key: str + ) -> None: + """Take the PENDING channel-inbox rows an admission will flush. + + Before the run row exists: ingest locks its channel row and then the + run, so admission must also take channel rows first or the two paths + meet on the same rows in opposite orders. + + Args: + conn: The connection inside the admitting transaction. + workflow_id: The workflow identity. + request_key: The admission key. + """ + await conn.execute( + "SELECT parked_id FROM workflow_channel_inbox WHERE workflow_id = %s" + " AND correlation_key = %s AND status = %s ORDER BY parked_id" + " FOR UPDATE", + (workflow_id, request_key, ParkedStatus.PENDING.value), + ) + + async def _route_parked_conn( + self, conn: Any, parked_id: str, now: float + ) -> DeliveryDisposition: + """Route one PENDING channel-inbox row inside an open transaction. + + Args: + conn: The connection, already holding the row. + parked_id: The row to route. + now: Current time in epoch seconds. + + Returns: + The routing outcome. + """ + cursor = await conn.execute( + "SELECT workflow_id, channel, correlation_key, dedupe_key, payload" + " FROM workflow_channel_inbox WHERE parked_id = %s", + (parked_id,), + ) + row = await cursor.fetchone() + cursor = await conn.execute( + "SELECT run_id FROM workflow_dedupe WHERE workflow_id = %s" + " AND request_key = %s", + (row["workflow_id"], row["correlation_key"]), + ) + target = await cursor.fetchone() + if target is None: + return "parked" + disposition = await self._deliver_with( + conn, + target["run_id"], + f"sig:{row['channel']}", + row["dedupe_key"], + row["payload"], + now, + ) + delivered = disposition in ("resolved", "buffered", "duplicate") + await conn.execute( + "UPDATE workflow_channel_inbox SET status = %s, reason = %s," + " run_id = %s, updated_at = %s WHERE parked_id = %s", + ( + ParkedStatus.DELIVERED.value if delivered else ParkedStatus.DEAD.value, + None if delivered else disposition, + target["run_id"] if delivered else None, + now, + parked_id, + ), + ) + return disposition if delivered else "dead_letter" + + async def ingest_channel_delivery( + self, + workflow_id: str, + channel: str, + correlation_key: str, + dedupe_key: str, + payload: Any, + now: float, + ) -> DeliveryDisposition: + """Durably accept a correlated provider event, exactly once. + + Args: + workflow_id: The workflow whose channel the event addresses. + channel: The channel name. + correlation_key: The business key naming the target run. + dedupe_key: The provider's event identity. + payload: The canonical event payload. + now: Current time in epoch seconds. + + Returns: + The routing outcome. + """ + pool = await self._open() + parked_id = uuid.uuid4().hex + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "INSERT INTO workflow_channel_inbox (parked_id, workflow_id," + " channel, correlation_key, dedupe_key, payload, status," + " created_at, updated_at)" + " VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)" + " ON CONFLICT DO NOTHING RETURNING parked_id", + ( + parked_id, + workflow_id, + channel, + correlation_key, dedupe_key, - await self._next_inbox_seq(conn, run_id), Jsonb(payload), - "CONSUMED" if resolves else "PENDING", + ParkedStatus.PENDING.value, + now, now, ), ) - if resolves and frontier is not None: - await conn.execute( - "UPDATE workflow_steps SET status = %s, due_at = %s, args = %s," - " updated_at = %s WHERE run_id = %s AND ordinal = %s", - ( - StepStatus.READY.value, - now, - _json({**frontier.args, "__payload__": payload}), - now, - run_id, - frontier.ordinal, - ), - ) - await self._append_events( - conn, - run_id, - ( - ( - HistoryEventType.WAIT_RESOLVED, - {"ordinal": frontier.ordinal, "wait_key": wait_key}, - ), - ), - now, - ) - else: - await self._append_events( - conn, - run_id, - ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), + if await cursor.fetchone() is None: + # The event id is the identity: a provider redelivery and a + # crash-after-ack replay both land here. + return "duplicate" + return await self._route_parked_conn(conn, parked_id, now) + + async def list_parked( + self, + *, + workflow_id: str | None = None, + status: ParkedStatus | None = None, + limit: int = 100, + ) -> tuple[ParkedDelivery, ...]: + """List channel-inbox deliveries, newest first. + + Args: + workflow_id: Restrict to one workflow. + status: Restrict to one lifecycle state. + limit: Maximum rows. + + Returns: + The matching deliveries. + """ + clauses, params = [], [] + if workflow_id is not None: + clauses.append("workflow_id = %s") + params.append(workflow_id) + if status is not None: + clauses.append("status = %s") + params.append(status.value) + where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + f"SELECT * FROM workflow_channel_inbox{where}" + " ORDER BY created_at DESC LIMIT %s", + (*params, limit), + ) + return tuple(_parked_from_row(row) for row in await cursor.fetchall()) + + async def replay_parked(self, parked_id: str, now: float) -> DeliveryDisposition: + """Re-attempt routing of a parked or dead delivery. + + Args: + parked_id: The delivery to replay. + now: Current time in epoch seconds. + + Returns: + The routing outcome, or ``unknown_key`` if no such delivery. + """ + pool = await self._open() + async with pool.connection() as conn, conn.transaction(): + cursor = await conn.execute( + "SELECT status FROM workflow_channel_inbox WHERE parked_id = %s" + " FOR UPDATE", + (parked_id,), + ) + row = await cursor.fetchone() + if row is None: + return "unknown_key" + if row["status"] == ParkedStatus.DELIVERED.value: + # Replaying what already reached its run must never signal + # twice. + return "duplicate" + await conn.execute( + "UPDATE workflow_channel_inbox SET status = %s, reason = NULL," + " updated_at = %s WHERE parked_id = %s", + (ParkedStatus.PENDING.value, now, parked_id), + ) + return await self._route_parked_conn(conn, parked_id, now) + + async def sweep_parked(self, now: float, ttl: float) -> int: + """Turn PENDING deliveries older than a ttl into DEAD letters. + + Args: + now: Current time in epoch seconds. + ttl: Age in seconds beyond which PENDING is unclaimed. + + Returns: + How many deliveries became dead letters. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "UPDATE workflow_channel_inbox SET status = %s," + " reason = 'unclaimed', updated_at = %s" + " WHERE status = %s AND created_at < %s", + ( + ParkedStatus.DEAD.value, now, - ) - return "resolved" if resolves else "buffered" + ParkedStatus.PENDING.value, + now - ttl, + ), + ) + return cursor.rowcount async def admit_children( self, diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index db86010f642..4d288ed5f5f 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -312,6 +312,51 @@ class StartResult: retry_after: float | None = None +class ParkedStatus(str, enum.Enum): + """Lifecycle of a correlated webhook delivery in the channel inbox.""" + + PENDING = "PENDING" + DELIVERED = "DELIVERED" + DEAD = "DEAD" + + +@dataclasses.dataclass(frozen=True, slots=True) +class ParkedDelivery: + """One correlated provider event, durable from the moment it was acked. + + The channel inbox is what makes webhook-to-signal delivery exactly-once: + the row's identity is the provider's event id, so redelivery and + crash-after-ack replays collapse into it, and a delivery that arrives + before its run exists waits here instead of being dropped. + + Attributes: + parked_id: Stable identity of this delivery record. + workflow_id: The workflow whose channel the event addresses. + channel: The channel name. + correlation_key: The business key naming the target run. + dedupe_key: The provider's event identity. + payload: The canonical event payload. + status: Where the delivery is in its lifecycle. + reason: Why a DEAD delivery died, e.g. ``run_terminal``, ``expired``, + or ``unclaimed``. + run_id: The run the delivery reached, once DELIVERED. + created_at: When the delivery was first acknowledged. + updated_at: Last transition time. + """ + + parked_id: str + workflow_id: str + channel: str + correlation_key: str + dedupe_key: str + payload: Any + status: ParkedStatus + reason: str | None + run_id: str | None + created_at: float + updated_at: float + + @dataclasses.dataclass(frozen=True, slots=True) class RunQuery: """Filters for listing runs in an operator surface. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 0562f8785a9..db57de2eaae 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -23,6 +23,7 @@ import json import sqlite3 import threading +import uuid from pathlib import Path from typing import TYPE_CHECKING, Any, Final, Literal, Protocol @@ -35,6 +36,8 @@ TERMINAL_STEP_STATUSES, HistoryEvent, HistoryEventType, + ParkedDelivery, + ParkedStatus, RunQuery, RunRecord, RunStatus, @@ -45,7 +48,7 @@ ) if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Iterable, Mapping DeliveryDisposition = Literal[ @@ -57,6 +60,8 @@ "unknown_run", "run_terminal", "unknown_key", + "parked", + "dead_letter", ] @@ -378,6 +383,93 @@ async def admit_children( """ ... + async def ingest_channel_delivery( + self, + workflow_id: str, + channel: str, + correlation_key: str, + dedupe_key: str, + payload: Any, + now: float, + ) -> DeliveryDisposition: + """Durably accept a correlated provider event, exactly once. + + The row keyed by the provider's event id is written first, in the + same transaction as any delivery, so acknowledging the provider and + recording the event are one fact: a crash after the ack replays as + ``duplicate``, never as a second signal. If the correlation key + already admitted a run, the payload is delivered to it here; if not, + the row waits PENDING for the run to be admitted; if the run is + terminal or past its deadline, the row becomes a DEAD letter an + operator can see and replay. + + Args: + workflow_id: The workflow whose channel the event addresses. + channel: The channel name. + correlation_key: The business key naming the target run. + dedupe_key: The provider's event identity. + payload: The canonical event payload. + now: Current time in epoch seconds. + + Returns: + ``resolved``/``buffered`` when delivered, ``duplicate`` for a + redelivery, ``parked`` when no run exists yet, ``dead_letter`` + when the run can no longer take it. + """ + ... + + async def list_parked( + self, + *, + workflow_id: str | None = None, + status: ParkedStatus | None = None, + limit: int = 100, + ) -> tuple[ParkedDelivery, ...]: + """List channel-inbox deliveries, newest first. + + Args: + workflow_id: Restrict to one workflow. + status: Restrict to one lifecycle state. + limit: Maximum rows. + + Returns: + The matching deliveries. + """ + ... + + async def replay_parked(self, parked_id: str, now: float) -> DeliveryDisposition: + """Re-attempt routing of a parked or dead delivery. + + The operator's answer to a dead letter whose cause is fixed: the row + goes through the same routing as ingest, with the same idempotency, + so replaying a delivery that already reached its run is a + ``duplicate``, never a second signal. + + Args: + parked_id: The delivery to replay. + now: Current time in epoch seconds. + + Returns: + The routing outcome, or ``unknown_key`` if no such delivery. + """ + ... + + async def sweep_parked(self, now: float, ttl: float) -> int: + """Turn PENDING deliveries older than a ttl into DEAD letters. + + A delivery whose run never arrived must eventually become visible as + a problem rather than waiting forever: ``unclaimed`` dead letters are + what an operator alerts on. + + Args: + now: Current time in epoch seconds. + ttl: Age in seconds beyond which PENDING is unclaimed. + + Returns: + How many deliveries became dead letters. + """ + ... + async def record_arrival( self, run_id: str, @@ -981,6 +1073,7 @@ def __init__(self): self._steps: dict[str, list[StepRecord]] = {} self._substeps: dict[tuple[str, int], dict[str, Any]] = {} self._schedule_cursors: dict[str, float] = {} + self._parked: list[ParkedDelivery] = [] self._history: dict[str, list[HistoryEvent]] = {} self._dedupe: dict[tuple[str, str], str] = {} self._inbox: dict[str, dict[tuple[str, str, str], bool]] = {} @@ -1041,8 +1134,210 @@ async def admit( self._runs[run.run_id] = run self._steps[run.run_id] = [root_step] self._append_events(run.run_id, events, run.created_at) + if run.request_key is not None: + # Deliveries that arrived before this run did: flushed inside + # the admitting transaction, so a crash cannot separate "the + # run exists" from "its early mail reached it". + self._flush_parked_locked( + run.workflow_id, run.request_key, run.run_id, run.created_at + ) return True, run.run_id + def _flush_parked_locked( + self, workflow_id: str, request_key: str, run_id: str, now: float + ) -> None: + """Deliver PENDING channel-inbox rows to a freshly admitted run. + + Args: + workflow_id: The workflow identity. + request_key: The admission key, matched against correlation keys. + run_id: The run that now exists. + now: Current time in epoch seconds. + """ + for index, parked in enumerate(self._parked): + if ( + parked.status is not ParkedStatus.PENDING + or parked.workflow_id != workflow_id + or parked.correlation_key != request_key + ): + continue + disposition = self._deliver_locked( + run_id, + f"sig:{parked.channel}", + parked.dedupe_key, + parked.payload, + now, + ) + if disposition in ("resolved", "buffered", "duplicate"): + self._parked[index] = dataclasses.replace( + parked, + status=ParkedStatus.DELIVERED, + run_id=run_id, + updated_at=now, + ) + else: + self._parked[index] = dataclasses.replace( + parked, + status=ParkedStatus.DEAD, + reason=disposition, + updated_at=now, + ) + + async def ingest_channel_delivery( + self, + workflow_id: str, + channel: str, + correlation_key: str, + dedupe_key: str, + payload: Any, + now: float, + ) -> DeliveryDisposition: + """Durably accept a correlated provider event, exactly once. + + Args: + workflow_id: The workflow whose channel the event addresses. + channel: The channel name. + correlation_key: The business key naming the target run. + dedupe_key: The provider's event identity. + payload: The canonical event payload. + now: Current time in epoch seconds. + + Returns: + The routing outcome. + """ + async with self._lock: + for parked in self._parked: + if ( + parked.workflow_id == workflow_id + and parked.channel == channel + and parked.correlation_key == correlation_key + and parked.dedupe_key == dedupe_key + ): + # The event id is the identity: a provider redelivery and + # a crash-after-ack replay both land here, whatever state + # the earlier row reached. + return "duplicate" + record = ParkedDelivery( + parked_id=uuid.uuid4().hex, + workflow_id=workflow_id, + channel=channel, + correlation_key=correlation_key, + dedupe_key=dedupe_key, + payload=payload, + status=ParkedStatus.PENDING, + reason=None, + run_id=None, + created_at=now, + updated_at=now, + ) + self._parked.append(record) + return self._route_parked_locked(len(self._parked) - 1, now) + + def _route_parked_locked(self, index: int, now: float) -> DeliveryDisposition: + """Route one PENDING channel-inbox row to its run, if it exists yet. + + Args: + index: The row's position. + now: Current time in epoch seconds. + + Returns: + The routing outcome. + """ + parked = self._parked[index] + run_id = self._dedupe.get((parked.workflow_id, parked.correlation_key)) + if run_id is None: + return "parked" + disposition = self._deliver_locked( + run_id, f"sig:{parked.channel}", parked.dedupe_key, parked.payload, now + ) + if disposition in ("resolved", "buffered", "duplicate"): + self._parked[index] = dataclasses.replace( + parked, + status=ParkedStatus.DELIVERED, + run_id=run_id, + updated_at=now, + ) + return disposition if disposition != "duplicate" else "duplicate" + self._parked[index] = dataclasses.replace( + parked, status=ParkedStatus.DEAD, reason=disposition, updated_at=now + ) + return "dead_letter" + + async def list_parked( + self, + *, + workflow_id: str | None = None, + status: ParkedStatus | None = None, + limit: int = 100, + ) -> tuple[ParkedDelivery, ...]: + """List channel-inbox deliveries, newest first. + + Args: + workflow_id: Restrict to one workflow. + status: Restrict to one lifecycle state. + limit: Maximum rows. + + Returns: + The matching deliveries. + """ + async with self._lock: + rows = [ + parked + for parked in self._parked + if (workflow_id is None or parked.workflow_id == workflow_id) + and (status is None or parked.status is status) + ] + rows.sort(key=lambda parked: parked.created_at, reverse=True) + return tuple(rows[:limit]) + + async def replay_parked(self, parked_id: str, now: float) -> DeliveryDisposition: + """Re-attempt routing of a parked or dead delivery. + + Args: + parked_id: The delivery to replay. + now: Current time in epoch seconds. + + Returns: + The routing outcome, or ``unknown_key`` if no such delivery. + """ + async with self._lock: + for index, parked in enumerate(self._parked): + if parked.parked_id != parked_id: + continue + if parked.status is ParkedStatus.DELIVERED: + return "duplicate" + self._parked[index] = dataclasses.replace( + parked, status=ParkedStatus.PENDING, reason=None, updated_at=now + ) + return self._route_parked_locked(index, now) + return "unknown_key" + + async def sweep_parked(self, now: float, ttl: float) -> int: + """Turn PENDING deliveries older than a ttl into DEAD letters. + + Args: + now: Current time in epoch seconds. + ttl: Age in seconds beyond which PENDING is unclaimed. + + Returns: + How many deliveries became dead letters. + """ + async with self._lock: + swept = 0 + for index, parked in enumerate(self._parked): + if ( + parked.status is ParkedStatus.PENDING + and now - parked.created_at > ttl + ): + self._parked[index] = dataclasses.replace( + parked, + status=ParkedStatus.DEAD, + reason="unclaimed", + updated_at=now, + ) + swept += 1 + return swept + async def claim_next( self, now: float, @@ -1303,62 +1598,82 @@ async def deliver( What the store did with the delivery. """ async with self._lock: - run = self._runs.get(run_id) - if run is None: - return "unknown_run" - if run.status in TERMINAL_RUN_STATUSES: - return "run_terminal" - if run.deadline is not None and run.deadline <= now: - # The run can never execute a continuation: claims exclude - # past-deadline runs and the sweep will finalize TIMED_OUT. - # Answering "resolved" here would tell the sender their - # decision was recorded when it is about to be discarded. - return "expired" - inbox = self._inbox.setdefault(run_id, {}) - if (run_id, wait_key, dedupe_key) in inbox: - self._append_events( - run_id, - ((HistoryEventType.SIGNAL_DUPLICATE, {"wait_key": wait_key}),), - now, - ) - return "duplicate" - inbox[run_id, wait_key, dedupe_key] = True - steps = self._steps[run_id] - frontier = _frontier(steps) - if frontier is not None and _wait_expired(frontier, now): - return "expired" - if ( - frontier is not None - and frontier.status is StepStatus.BLOCKED - and frontier.wait_key == wait_key - ): - steps[frontier.ordinal] = dataclasses.replace( - frontier, - status=StepStatus.READY, - due_at=now, - args={**frontier.args, "__payload__": payload}, - updated_at=now, - ) - self._append_events( - run_id, - ( - ( - HistoryEventType.WAIT_RESOLVED, - {"ordinal": frontier.ordinal, "wait_key": wait_key}, - ), - ), - now, - ) - return "resolved" - self._pending.setdefault(run_id, {}).setdefault(wait_key, []).append( - payload + return self._deliver_locked(run_id, wait_key, dedupe_key, payload, now) + + def _deliver_locked( + self, + run_id: str, + wait_key: str, + dedupe_key: str, + payload: Any, + now: float, + ) -> DeliveryDisposition: + """Deliver with the store lock already held. + + Args: + run_id: The receiving run. + wait_key: The address the waiting slot declared. + dedupe_key: Sender-supplied identity, making redelivery a no-op. + payload: JSON-compatible payload to hand the resuming handler. + now: Current time in epoch seconds. + + Returns: + What the store did with the delivery. + """ + run = self._runs.get(run_id) + if run is None: + return "unknown_run" + if run.status in TERMINAL_RUN_STATUSES: + return "run_terminal" + if run.deadline is not None and run.deadline <= now: + # The run can never execute a continuation: claims exclude + # past-deadline runs and the sweep will finalize TIMED_OUT. + # Answering "resolved" here would tell the sender their + # decision was recorded when it is about to be discarded. + return "expired" + inbox = self._inbox.setdefault(run_id, {}) + if (run_id, wait_key, dedupe_key) in inbox: + self._append_events( + run_id, + ((HistoryEventType.SIGNAL_DUPLICATE, {"wait_key": wait_key}),), + now, + ) + return "duplicate" + inbox[run_id, wait_key, dedupe_key] = True + steps = self._steps[run_id] + frontier = _frontier(steps) + if frontier is not None and _wait_expired(frontier, now): + return "expired" + if ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and frontier.wait_key == wait_key + ): + steps[frontier.ordinal] = dataclasses.replace( + frontier, + status=StepStatus.READY, + due_at=now, + args={**frontier.args, "__payload__": payload}, + updated_at=now, ) self._append_events( run_id, - ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), + ( + ( + HistoryEventType.WAIT_RESOLVED, + {"ordinal": frontier.ordinal, "wait_key": wait_key}, + ), + ), now, ) - return "buffered" + return "resolved" + self._pending.setdefault(run_id, {}).setdefault(wait_key, []).append(payload) + self._append_events( + run_id, + ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), + now, + ) + return "buffered" async def admit_children( self, @@ -1585,6 +1900,12 @@ async def admit_flow( self._runs[run.run_id] = run self._steps[run.run_id] = [dataclasses.replace(root_step, due_at=due_at)] self._append_events(run.run_id, events, run.created_at) + if run.request_key is not None: + # Policy admission is still admission: early mail flushes on + # this door exactly as on the plain one. + self._flush_parked_locked( + run.workflow_id, run.request_key, run.run_id, run.created_at + ) return FlowAdmission("started", run.run_id, cancelled=tuple(cancelled)) async def purge_runs(self, before: float, *, workflow_id: str | None = None) -> int: @@ -2429,7 +2750,7 @@ async def next_due( return min(due_times) if due_times else None -SCHEMA_VERSION: Final = 3 +SCHEMA_VERSION: Final = 4 """Stamped into PRAGMA user_version; bump when _SCHEMA or migrations change.""" DATABASE_ENV: Final = "REFLEX_WORKFLOW_DATABASE" @@ -2549,6 +2870,22 @@ def resolve_store(target: str | None = None) -> RunStore: ON workflow_steps (status, due_at, queue); CREATE INDEX IF NOT EXISTS idx_workflow_inbox_pending ON workflow_inbox (run_id, wait_key, status, seq); +CREATE TABLE IF NOT EXISTS workflow_channel_inbox ( + parked_id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + channel TEXT NOT NULL, + correlation_key TEXT NOT NULL, + dedupe_key TEXT NOT NULL, + payload TEXT NOT NULL, + status TEXT NOT NULL, + reason TEXT, + run_id TEXT, + created_at REAL NOT NULL, + updated_at REAL NOT NULL, + UNIQUE (workflow_id, channel, correlation_key, dedupe_key) +); +CREATE INDEX IF NOT EXISTS idx_workflow_channel_inbox_route + ON workflow_channel_inbox (workflow_id, correlation_key, status); """ _STEP_MIGRATIONS: Final = ( @@ -2720,6 +3057,30 @@ def _child_admission_events( ) +def _parked_from_row(row: Mapping[str, Any]) -> ParkedDelivery: + """Build a parked-delivery record from a database row. + + Args: + row: The ``workflow_channel_inbox`` row. + + Returns: + The record. + """ + return ParkedDelivery( + parked_id=row["parked_id"], + workflow_id=row["workflow_id"], + channel=row["channel"], + correlation_key=row["correlation_key"], + dedupe_key=row["dedupe_key"], + payload=json.loads(row["payload"]), + status=ParkedStatus(row["status"]), + reason=row["reason"], + run_id=row["run_id"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + def _sqlite_frontier_query( select: str, now: float, @@ -3099,6 +3460,17 @@ def work(): self._insert_run(run) self._insert_step(root_step) self._append_events(run.run_id, events, run.created_at) + if run.request_key is not None: + # Deliveries that arrived before this run did, flushed + # inside the admitting transaction: a crash cannot + # separate "the run exists" from "its early mail + # reached it". + self._flush_parked_in_txn( + run.workflow_id, + run.request_key, + run.run_id, + run.created_at, + ) self._db.execute("COMMIT") except BaseException: self._db.execute("ROLLBACK") @@ -3107,6 +3479,275 @@ def work(): return await asyncio.to_thread(work) + def _flush_parked_in_txn( + self, workflow_id: str, request_key: str, run_id: str, now: float + ) -> None: + """Deliver PENDING channel-inbox rows to a freshly admitted run. + + Args: + workflow_id: The workflow identity. + request_key: The admission key, matched against correlation keys. + run_id: The run that now exists. + now: Current time in epoch seconds. + """ + rows = self._db.execute( + "SELECT parked_id, channel, dedupe_key, payload FROM" + " workflow_channel_inbox WHERE workflow_id = ? AND" + " correlation_key = ? AND status = ? ORDER BY created_at", + (workflow_id, request_key, ParkedStatus.PENDING.value), + ).fetchall() + for row in rows: + disposition = self._deliver_in_txn( + run_id, + f"sig:{row['channel']}", + row["dedupe_key"], + json.loads(row["payload"]), + now, + ) + delivered = disposition in ("resolved", "buffered", "duplicate") + self._db.execute( + "UPDATE workflow_channel_inbox SET status = ?, reason = ?," + " run_id = ?, updated_at = ? WHERE parked_id = ?", + ( + ParkedStatus.DELIVERED.value + if delivered + else ParkedStatus.DEAD.value, + None if delivered else disposition, + run_id if delivered else None, + now, + row["parked_id"], + ), + ) + + def _route_parked_in_txn(self, parked_id: str, now: float) -> DeliveryDisposition: + """Route one PENDING channel-inbox row inside an open transaction. + + Args: + parked_id: The row to route. + now: Current time in epoch seconds. + + Returns: + The routing outcome. + """ + row = self._db.execute( + "SELECT workflow_id, channel, correlation_key, dedupe_key, payload" + " FROM workflow_channel_inbox WHERE parked_id = ?", + (parked_id,), + ).fetchone() + target = self._db.execute( + "SELECT run_id FROM workflow_dedupe WHERE workflow_id = ?" + " AND request_key = ?", + (row["workflow_id"], row["correlation_key"]), + ).fetchone() + if target is None: + return "parked" + disposition = self._deliver_in_txn( + target["run_id"], + f"sig:{row['channel']}", + row["dedupe_key"], + json.loads(row["payload"]), + now, + ) + delivered = disposition in ("resolved", "buffered", "duplicate") + self._db.execute( + "UPDATE workflow_channel_inbox SET status = ?, reason = ?, run_id = ?," + " updated_at = ? WHERE parked_id = ?", + ( + ParkedStatus.DELIVERED.value if delivered else ParkedStatus.DEAD.value, + None if delivered else disposition, + target["run_id"] if delivered else None, + now, + parked_id, + ), + ) + return disposition if delivered else "dead_letter" + + async def ingest_channel_delivery( + self, + workflow_id: str, + channel: str, + correlation_key: str, + dedupe_key: str, + payload: Any, + now: float, + ) -> DeliveryDisposition: + """Durably accept a correlated provider event, exactly once. + + Args: + workflow_id: The workflow whose channel the event addresses. + channel: The channel name. + correlation_key: The business key naming the target run. + dedupe_key: The provider's event identity. + payload: The canonical event payload. + now: Current time in epoch seconds. + + Returns: + The routing outcome. + """ + + def work() -> DeliveryDisposition: + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + parked_id = uuid.uuid4().hex + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + inserted = self._db.execute( + "INSERT INTO workflow_channel_inbox (parked_id," + " workflow_id, channel, correlation_key, dedupe_key," + " payload, status, created_at, updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)" + " ON CONFLICT DO NOTHING", + ( + parked_id, + workflow_id, + channel, + correlation_key, + dedupe_key, + json.dumps(payload), + ParkedStatus.PENDING.value, + now, + now, + ), + ) + if inserted.rowcount == 0: + # The event id is the identity: a provider redelivery + # and a crash-after-ack replay both land here. + self._db.execute("COMMIT") + return "duplicate" + disposition = self._route_parked_in_txn(parked_id, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return disposition + + return await asyncio.to_thread(work) + + async def list_parked( + self, + *, + workflow_id: str | None = None, + status: ParkedStatus | None = None, + limit: int = 100, + ) -> tuple[ParkedDelivery, ...]: + """List channel-inbox deliveries, newest first. + + Args: + workflow_id: Restrict to one workflow. + status: Restrict to one lifecycle state. + limit: Maximum rows. + + Returns: + The matching deliveries. + """ + + def work() -> tuple[ParkedDelivery, ...]: + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + clauses, params = [], [] + if workflow_id is not None: + clauses.append("workflow_id = ?") + params.append(workflow_id) + if status is not None: + clauses.append("status = ?") + params.append(status.value) + where = f" WHERE {' AND '.join(clauses)}" if clauses else "" + with self._lock: + rows = self._db.execute( + f"SELECT * FROM workflow_channel_inbox{where}" + " ORDER BY created_at DESC LIMIT ?", + (*params, limit), + ).fetchall() + return tuple(_parked_from_row(row) for row in rows) + + return await asyncio.to_thread(work) + + async def replay_parked(self, parked_id: str, now: float) -> DeliveryDisposition: + """Re-attempt routing of a parked or dead delivery. + + Args: + parked_id: The delivery to replay. + now: Current time in epoch seconds. + + Returns: + The routing outcome, or ``unknown_key`` if no such delivery. + """ + + def work() -> DeliveryDisposition: + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + row = self._db.execute( + "SELECT status FROM workflow_channel_inbox WHERE parked_id = ?", + (parked_id,), + ).fetchone() + if row is None: + self._db.execute("COMMIT") + return "unknown_key" + if row["status"] == ParkedStatus.DELIVERED.value: + # Replaying what already reached its run must never + # signal twice. + self._db.execute("COMMIT") + return "duplicate" + self._db.execute( + "UPDATE workflow_channel_inbox SET status = ?," + " reason = NULL, updated_at = ? WHERE parked_id = ?", + (ParkedStatus.PENDING.value, now, parked_id), + ) + disposition = self._route_parked_in_txn(parked_id, now) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + return disposition + + return await asyncio.to_thread(work) + + async def sweep_parked(self, now: float, ttl: float) -> int: + """Turn PENDING deliveries older than a ttl into DEAD letters. + + Args: + now: Current time in epoch seconds. + ttl: Age in seconds beyond which PENDING is unclaimed. + + Returns: + How many deliveries became dead letters. + """ + + def work() -> int: + """Run the operation on the worker thread. + + Returns: + The operation's result. + """ + with self._lock: + cursor = self._db.execute( + "UPDATE workflow_channel_inbox SET status = ?," + " reason = 'unclaimed', updated_at = ?" + " WHERE status = ? AND created_at < ?", + ( + ParkedStatus.DEAD.value, + now, + ParkedStatus.PENDING.value, + now - ttl, + ), + ) + return cursor.rowcount + + return await asyncio.to_thread(work) + async def admit_flow( self, run: RunRecord, @@ -3227,6 +3868,15 @@ def work() -> FlowAdmission: self._insert_run(run) self._insert_step(dataclasses.replace(root_step, due_at=due_at)) self._append_events(run.run_id, events, run.created_at) + if run.request_key is not None: + # Policy admission is still admission: early mail + # flushes on this door exactly as on the plain one. + self._flush_parked_in_txn( + run.workflow_id, + run.request_key, + run.run_id, + run.created_at, + ) self._db.execute("COMMIT") except BaseException: self._db.execute("ROLLBACK") @@ -3682,109 +4332,123 @@ def work(): Returns: The operation's result. """ - terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) with self._lock: self._db.execute("BEGIN IMMEDIATE") try: - row = self._db.execute( - "SELECT status, deadline FROM workflow_runs WHERE run_id = ?", - (run_id,), - ).fetchone() - if row is None: - self._db.execute("ROLLBACK") - return "unknown_run" - if row["status"] in terminal: - self._db.execute("ROLLBACK") - return "run_terminal" - if row["deadline"] is not None and row["deadline"] <= now: - # A past-deadline run can never execute the - # continuation; saying "resolved" would be a lie. - self._db.execute("ROLLBACK") - return "expired" - seen = self._db.execute( - "SELECT 1 FROM workflow_inbox" - " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", - (run_id, wait_key, dedupe_key), - ).fetchone() - if seen is not None: - self._append_events( - run_id, - ( - ( - HistoryEventType.SIGNAL_DUPLICATE, - {"wait_key": wait_key}, - ), - ), - now, - ) - self._db.execute("COMMIT") - return "duplicate" - frontier = _frontier(self._load_steps(run_id)) - if frontier is not None and _wait_expired(frontier, now): - self._db.execute("ROLLBACK") - return "expired" - resolves = ( - frontier is not None - and frontier.status is StepStatus.BLOCKED - and frontier.wait_key == wait_key + disposition = self._deliver_in_txn( + run_id, wait_key, dedupe_key, payload, now ) - self._db.execute( - "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," - " payload, status, created_at)" - " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" - " workflow_inbox WHERE run_id = ?), ?, ?, ?)", - ( - run_id, - wait_key, - dedupe_key, - run_id, - json.dumps(payload), - "CONSUMED" if resolves else "PENDING", - now, - ), - ) - if resolves and frontier is not None: - self._db.execute( - "UPDATE workflow_steps SET status = ?, due_at = ?, args = ?," - " updated_at = ? WHERE run_id = ? AND ordinal = ?", - ( - StepStatus.READY.value, - now, - json.dumps({**frontier.args, "__payload__": payload}), - now, - run_id, - frontier.ordinal, - ), - ) - self._append_events( - run_id, - ( - ( - HistoryEventType.WAIT_RESOLVED, - {"ordinal": frontier.ordinal, "wait_key": wait_key}, - ), - ), - now, - ) - else: - self._append_events( - run_id, - ( - ( - HistoryEventType.SIGNAL_BUFFERED, - {"wait_key": wait_key}, - ), - ), - now, - ) self._db.execute("COMMIT") except BaseException: self._db.execute("ROLLBACK") raise - return "resolved" if resolves else "buffered" + return disposition return await asyncio.to_thread(work) + def _deliver_in_txn( + self, + run_id: str, + wait_key: str, + dedupe_key: str, + payload: Any, + now: float, + ) -> DeliveryDisposition: + """Deliver inside the caller's open transaction. + + Refusal branches write nothing, so the caller's transaction stays + committable whatever this returns -- which is what lets admission + flush parked deliveries in its own transaction. + + Args: + run_id: The receiving run. + wait_key: The address the waiting slot declared. + dedupe_key: Sender-supplied identity, making redelivery a no-op. + payload: JSON-compatible payload to hand the resuming handler. + now: Current time in epoch seconds. + + Returns: + What the store did with the delivery. + """ + terminal = tuple(s.value for s in TERMINAL_RUN_STATUSES) + row = self._db.execute( + "SELECT status, deadline FROM workflow_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + return "unknown_run" + if row["status"] in terminal: + return "run_terminal" + if row["deadline"] is not None and row["deadline"] <= now: + # A past-deadline run can never execute the continuation; saying + # "resolved" would be a lie. + return "expired" + seen = self._db.execute( + "SELECT 1 FROM workflow_inbox" + " WHERE run_id = ? AND wait_key = ? AND dedupe_key = ?", + (run_id, wait_key, dedupe_key), + ).fetchone() + if seen is not None: + self._append_events( + run_id, + ((HistoryEventType.SIGNAL_DUPLICATE, {"wait_key": wait_key}),), + now, + ) + return "duplicate" + frontier = _frontier(self._load_steps(run_id)) + if frontier is not None and _wait_expired(frontier, now): + return "expired" + resolves = ( + frontier is not None + and frontier.status is StepStatus.BLOCKED + and frontier.wait_key == wait_key + ) + self._db.execute( + "INSERT INTO workflow_inbox (run_id, wait_key, dedupe_key, seq," + " payload, status, created_at)" + " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1 FROM" + " workflow_inbox WHERE run_id = ?), ?, ?, ?)", + ( + run_id, + wait_key, + dedupe_key, + run_id, + json.dumps(payload), + "CONSUMED" if resolves else "PENDING", + now, + ), + ) + if resolves and frontier is not None: + self._db.execute( + "UPDATE workflow_steps SET status = ?, due_at = ?, args = ?," + " updated_at = ? WHERE run_id = ? AND ordinal = ?", + ( + StepStatus.READY.value, + now, + json.dumps({**frontier.args, "__payload__": payload}), + now, + run_id, + frontier.ordinal, + ), + ) + self._append_events( + run_id, + ( + ( + HistoryEventType.WAIT_RESOLVED, + {"ordinal": frontier.ordinal, "wait_key": wait_key}, + ), + ), + now, + ) + else: + self._append_events( + run_id, + ((HistoryEventType.SIGNAL_BUFFERED, {"wait_key": wait_key}),), + now, + ) + return "resolved" if resolves else "buffered" + async def admit_children( self, runs: tuple[tuple[RunRecord, StepRecord], ...], From ab179d59b0b91d0e2ec4bb92bfa0734f1731ae12 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Mon, 24 Aug 2026 17:54:14 -0700 Subject: [PATCH 118/121] workflows: correlated webhook-to-signal delivery, end to end Phase 2 complete on top of the channel inbox. A provider event now reaches a waiting run with no glue code: class Order(rx.State): shipped = rx.Signal( Shipment, trigger=rx.webhook( "shippo.shipped", verify=shippo_verifier, dedupe_by="event_id", correlate_by="order_id", ), ) Root webhooks start runs; channel webhooks locate them. correlate_by names the payload field carrying the business key, matched against request keys; dedupe_by names the provider's event identity. A channel trigger requires both -- without them a delivery cannot be routed exactly once -- and one topic identifies exactly one target, root or channel, enforced at route collection. The webhook endpoint verifies the signature, canonicalizes against the channel's model, extracts both identities (a missing one is a 400 before anything exists -- accepting it would mint a dead letter for a sender error a 400 would have fixed), and hands the payload to kernel.ingest_channel, which checks the channel is declared and wakes the worker on a resolved delivery. Every durable outcome acknowledges 202: once the row is committed the provider must stop retrying, whether the payload landed, parked, deduplicated, or died visibly. Recovery sweeps PENDING deliveries unclaimed past 30 days into "unclaimed" dead letters and says so. Operators get the loop on both surfaces: reflex workflows deadletters [--status ...|--all|--replay ID] and GET /deadletters + POST /deadletters/{id}/replay on the standalone service (read/operate scopes). The plan's acceptance scenario now runs with a real SIGKILL: the shipment webhook arrives before the order workflow exists, the process is killed immediately after acknowledging it, the provider redelivers twice from fresh processes, the order workflow starts later -- exactly one signal arrives, held by the fsynced effect ledger, by the run's inbox read straight from the database, and by the single DELIVERED channel row. The same flow minus the kill runs over real HTTP in the ingress suite, and the dead-letter loop runs over the standalone service's own webhook route. --- news/workflow-channel-inbox.md | 2 + reflex/workflow/CONTRACT.md | 28 ++++ reflex/workflow/cli.py | 69 ++++++++ reflex/workflow/ingress.py | 156 +++++++++++++++--- reflex/workflow/kernel.py | 59 +++++++ reflex/workflow/serve.py | 112 +++++++++++++ tests/units/workflow/crash_worker.py | 57 ++++++- tests/units/workflow/test_crash_boundaries.py | 48 ++++++ tests/units/workflow/test_ingress.py | 151 +++++++++++++++++ tests/units/workflow/test_serve.py | 85 ++++++++++ 10 files changed, 744 insertions(+), 23 deletions(-) diff --git a/news/workflow-channel-inbox.md b/news/workflow-channel-inbox.md index 47886afc19c..f8a90faac92 100644 --- a/news/workflow-channel-inbox.md +++ b/news/workflow-channel-inbox.md @@ -1 +1,3 @@ The durable core of correlated webhook delivery: a channel inbox on every store. `ingest_channel_delivery` accepts a provider event addressed to a workflow channel by business key and makes one transaction of the whole decision — the row keyed by the provider's event id is written first, so acknowledging the provider and recording the event are one fact, and a crash after the ack replays as `duplicate`, never a second signal. A delivery whose run already exists lands immediately; one that arrives early waits `PENDING` and is flushed **inside the admitting transaction** — through either admission door, plain or policy — so a crash cannot separate "the run exists" from "its early mail reached it". A delivery nothing can take (terminal run, past deadline, or unclaimed past TTL) becomes a visible dead letter with a reason, listable and replayable with the same event-id idempotency. Six conformance checks pin the semantics on Memory, SQLite, and Postgres, including the acceptance flow: park before the run, redeliver three times, admit — exactly one signal. + +The full path is wired end to end: `rx.Signal(Model, trigger=rx.webhook(topic, verify=..., dedupe_by="event_id", correlate_by="order_id"))` declares a channel fed directly by a provider. Root webhooks start runs; channel webhooks locate them by business key, with one topic identifying exactly one target. The webhook endpoint verifies, canonicalizes, extracts both identities (400 when either is missing — before anything exists), and acknowledges with the routing disposition once the row is durable. Recovery sweeps unclaimed deliveries into dead letters after 30 days; `reflex workflows deadletters [--replay ID]` and `GET /deadletters` + `POST /deadletters/{id}/replay` give operators the loop. The acceptance scenario holds under a real SIGKILL: the shipment webhook arrives before the order workflow, the process dies immediately after acknowledging it, the provider redelivers twice, the run starts later — exactly one signal arrives, proven by the fsynced ledger and the database read directly. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index e07a750ab13..d13d4c5836e 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -260,6 +260,34 @@ Checked in this order at start: refuses its signal rather than letting it resolve a later wait on the same channel. +### Correlated webhook delivery + +A signal channel can be fed directly by a provider: +`rx.Signal(Model, trigger=rx.webhook(topic, verify=..., dedupe_by="event_id", +correlate_by="order_id"))`. Root webhooks start runs; channel webhooks locate +them: `correlate_by` names the payload field carrying the business key, +matched against runs' request keys (§6), and `dedupe_by` names the +provider's event identity. Both are required on a channel trigger, and one +topic identifies exactly one target — root or channel, never both. + +The delivery is durable from the moment it is acknowledged. Ingest writes +the channel-inbox row keyed by the event id and routes it in the **same +transaction**: to the run when the correlation key admitted one, to +`PENDING` when none exists yet, to a `DEAD` letter when the run is terminal +or past its deadline. Admission — through either door, plain or policy — +flushes `PENDING` rows for its request key inside the admitting +transaction, so a crash cannot separate "the run exists" from "its early +mail reached it". A `PENDING` row unclaimed past the TTL (30 days) becomes +a `DEAD` letter with reason `unclaimed`, swept by recovery and warned about. + +Dead letters are never silent: they list with their reason +(`reflex workflows deadletters`, `GET /deadletters`) and replay +(`--replay`, `POST /deadletters/{id}/replay`) through the same routing with +the same event-id idempotency — replaying a delivered row is a `duplicate`, +never a second signal. The acceptance shape, held by a real SIGKILL test: +delivery before the run, killed after the ack, redelivered twice, run +started later — exactly one signal arrives. + ### Business-key addressing The request key is a durable unique index per workflow (`§6`), so it doubles diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 44dfea9fd20..6e582044f7d 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -1414,6 +1414,75 @@ def serve( _run_server(app, host, port) +@workflows.command() +@database_option +@click.option( + "--status", + type=click.Choice(["pending", "delivered", "dead"]), + default=None, + help="Show only deliveries in this state; default shows dead letters.", +) +@click.option("--all", "show_all", is_flag=True, help="Show every delivery state.") +@click.option( + "--replay", + "replay_id", + default=None, + help="Re-attempt routing of one delivery by its id.", +) +def deadletters( + database: str | None, + status: str | None, + show_all: bool, + replay_id: str | None, +): + """Inspect and replay correlated webhook deliveries. + + A delivery that arrived before its run waits PENDING; one nothing can + take is DEAD with a reason. Replay routes a delivery again with the same + event-id idempotency, so replaying one that already landed is a no-op. + """ + from reflex.workflow.records import ParkedStatus + + async def act(store: RunStore): + """Run the inspection or replay. + + Args: + store: The open run store. + + Returns: + The rows to render, or the replay disposition. + """ + if replay_id is not None: + return await store.replay_parked(replay_id, time.time()) + chosen = ( + None + if show_all + else ParkedStatus(status.upper()) + if status + else ParkedStatus.DEAD + ) + return await store.list_parked(status=chosen) + + import time + + result = _with_store(database, act) + if replay_id is not None: + console.print(f"Replay: {result}") + raise click.exceptions.Exit( + 0 if result in ("resolved", "buffered", "duplicate") else 1 + ) + if not result: + console.print("No matching deliveries.") + return + for row in result: + target = row.run_id or f"key={row.correlation_key!r}" + detail = f" reason={row.reason}" if row.reason else "" + console.print( + f"{row.parked_id} {row.status.value:<9} {row.workflow_id}." + f"{row.channel} {target}{detail}" + ) + + @workflows.command() @database_option @click.argument("run_id") diff --git a/reflex/workflow/ingress.py b/reflex/workflow/ingress.py index e1aa3b3a4fd..d270ddc7cbf 100644 --- a/reflex/workflow/ingress.py +++ b/reflex/workflow/ingress.py @@ -22,7 +22,7 @@ if TYPE_CHECKING: from collections.abc import Callable, Coroutine, Mapping - from reflex_base.workflow import WebhookTrigger + from reflex_base.workflow import Signal, WebhookTrigger from reflex.workflow.definition import HandlerDefinition, WorkflowDefinition from reflex.workflow.runtime import WorkflowRuntime @@ -33,32 +33,38 @@ class WebhookRoute: - """One workflow root reachable over HTTP. + """One webhook topic reachable over HTTP: a root to start or a channel + to deliver into. Attributes: - definition: The workflow definition owning the root. - handler: The root handler started by this topic. + definition: The workflow definition owning the target. + handler: The root handler started by this topic, for a root route. trigger: The webhook trigger declaring the topic and its verification. + channel: The signal channel this topic delivers into, for a channel + route. """ - __slots__ = ("definition", "handler", "trigger") + __slots__ = ("channel", "definition", "handler", "trigger") def __init__( self, definition: WorkflowDefinition, - handler: HandlerDefinition, + handler: HandlerDefinition | None, trigger: WebhookTrigger, + channel: Signal | None = None, ): """Initialize the route. Args: - definition: The workflow definition owning the root. - handler: The root handler started by this topic. + definition: The workflow definition owning the target. + handler: The root handler started by this topic, or None. trigger: The webhook trigger declaring the topic. + channel: The signal channel this topic delivers into, or None. """ self.definition = definition self.handler = handler self.trigger = trigger + self.channel = channel def collect_webhook_routes( @@ -80,22 +86,52 @@ def collect_webhook_routes( from reflex_base.workflow import WebhookTrigger routes: dict[str, WebhookRoute] = {} + + def claim(topic: str, route: WebhookRoute, target: str) -> None: + """Claim a topic, refusing a second claimant. + + Args: + topic: The provider topic. + route: The route claiming it. + target: Human name of the claimant, for the error. + + Raises: + WorkflowDefinitionError: If the topic is already claimed. + """ + existing = routes.get(topic) + if existing is not None: + held_by = ( + f"{existing.definition.workflow_id}.{existing.handler.id}" + if existing.handler is not None + else f"{existing.definition.workflow_id}." + f"{existing.channel.name if existing.channel else '?'}" + ) + msg = ( + f"Webhook topic {topic!r} is claimed by both {held_by} and " + f"{target}; a topic must identify exactly one target." + ) + raise WorkflowDefinitionError(msg) + routes[topic] = route + for definition in definitions: for handler_id in definition.roots: handler = definition.handlers[handler_id] trigger = handler.trigger if not isinstance(trigger, WebhookTrigger): continue - existing = routes.get(trigger.topic) - if existing is not None: - msg = ( - f"Webhook topic {trigger.topic!r} is claimed by both " - f"{existing.definition.workflow_id}.{existing.handler.id} and " - f"{definition.workflow_id}.{handler.id}; a topic must " - "identify exactly one root." - ) - raise WorkflowDefinitionError(msg) - routes[trigger.topic] = WebhookRoute(definition, handler, trigger) + claim( + trigger.topic, + WebhookRoute(definition, handler, trigger), + f"{definition.workflow_id}.{handler.id}", + ) + for channel in definition.channels.values(): + if channel.trigger is None: + continue + claim( + channel.trigger.topic, + WebhookRoute(definition, None, channel.trigger, channel), + f"{definition.workflow_id}.{channel.name}", + ) return routes @@ -117,7 +153,22 @@ def _identity_value( The identity, or None when the declared source is absent. """ assert trigger.dedupe_by is not None - source = trigger.dedupe_by + return _extract_identity(trigger.dedupe_by, payload, headers) + + +def _extract_identity( + source: str, payload: Any, headers: Mapping[str, str] +) -> str | None: + """Extract one identity value from a payload field or a header. + + Args: + source: A payload field path, or ``"header:Name"``. + payload: The decoded request payload. + headers: The request headers. + + Returns: + The identity, or None when the declared source is absent. + """ if source.startswith("header:"): name = source[len("header:") :] value = headers.get(name.lower()) or headers.get(name) @@ -218,6 +269,62 @@ def _root_args(handler: HandlerDefinition, payload: Any) -> dict[str, Any]: return {name: payload[name] for name in handler.params if name in payload} +async def _ingest_channel( + runtime: WorkflowRuntime, + route: WebhookRoute, + payload: Any, + headers: Mapping[str, str], +) -> JSONResponse: + """Route a verified channel delivery into the durable channel inbox. + + Args: + runtime: The workflow runtime. + route: The channel route the topic resolved to. + payload: The canonical payload. + headers: The request headers, for header-sourced identities. + + Returns: + The acknowledgement. Every durable outcome is a 202: once the row is + committed the provider must stop retrying, whether the payload + landed, parked, deduplicated, or died visibly for an operator. + """ + assert route.channel is not None + trigger = route.trigger + dedupe = _extract_identity(trigger.dedupe_by or "", payload, headers) + if dedupe is None: + return JSONResponse( + { + "error": ( + f"delivery carries no {trigger.dedupe_by!r}, which this " + "channel deduplicates by" + ) + }, + status_code=400, + ) + correlation = _extract_identity(trigger.correlate_by or "", payload, headers) + if correlation is None: + # Without the business key there is no run to route to and no key to + # park under; accepting it would create a dead letter for a sender + # error a 400 would have fixed. + return JSONResponse( + { + "error": ( + f"delivery carries no {trigger.correlate_by!r}, which this " + "channel correlates by" + ) + }, + status_code=400, + ) + disposition = await runtime.kernel.ingest_channel( + route.definition.workflow_id, + route.channel.name, + str(correlation), + str(dedupe), + payload, + ) + return JSONResponse({"disposition": disposition}, status_code=202) + + def webhook_endpoint( runtime: WorkflowRuntime, ) -> Callable[[Request], Coroutine[Any, Any, JSONResponse]]: @@ -272,18 +379,25 @@ async def endpoint(request: Request) -> JSONResponse: except ValueError: return JSONResponse({"error": "payload is not JSON"}, status_code=400) - if route.trigger.model is not None: + model = route.trigger.model or ( + route.channel.model if route.channel is not None else None + ) + if model is not None: try: # The canonical form -- coercions applied, defaults filled -- # is what goes onward. Validating and then passing the raw # payload threw the validation away. - payload = canonical_payload(route.trigger.model, payload) + payload = canonical_payload(model, payload) except ValidationError: return JSONResponse( {"error": "payload does not match the declared model"}, status_code=400, ) + if route.channel is not None: + return await _ingest_channel(runtime, route, payload, headers) + + assert route.handler is not None spec = getattr(route.definition.state_cls, route.handler.name) if len(route.handler.params) > 1 and not isinstance(payload, dict): # Several named parameters can only be filled from an object. diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index 8d04ad38d9e..f0fd11fd72c 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -85,6 +85,13 @@ RECOVERY_INTERVAL_FRACTION = 1 / 2 MAX_SCHEDULE_CATCHUP = 10 +PARKED_DELIVERY_TTL = 30 * 86_400.0 +"""Seconds a parked channel delivery waits for its run before dead-lettering. + +Thirty days matches the longest provider retry horizons with room to spare: +a delivery still unclaimed after a month is a correlation nobody is coming +for, and an operator should see it rather than the table growing forever. +""" class _HandlerCancelledError(Exception): @@ -1002,6 +1009,50 @@ async def find_by_key(self, workflow: Any, request_key: str) -> str | None: self._workflow_id_of(workflow), request_key ) + async def ingest_channel( + self, + workflow_id: str, + channel_name: str, + correlation_key: str, + dedupe_key: str, + payload: Any, + ) -> DeliveryDisposition: + """Durably accept a correlated webhook delivery for a channel. + + Args: + workflow_id: The workflow whose channel the event addresses. + channel_name: The channel name. + correlation_key: The business key naming the target run. + dedupe_key: The provider's event identity. + payload: The canonical event payload. + + Returns: + The routing outcome. + + Raises: + WorkflowDefinitionError: If the workflow is registered here and + does not declare the channel. + """ + defn = self._definitions.get(workflow_id) + if defn is not None and channel_name not in defn.channels: + declared = sorted(defn.channels) or [""] + msg = ( + f"Workflow {workflow_id!r} declares no channel " + f"{channel_name!r}; declared channels: {', '.join(declared)}." + ) + raise WorkflowDefinitionError(msg) + disposition = await self._store.ingest_channel_delivery( + workflow_id, + channel_name, + correlation_key, + dedupe_key, + to_run_data({"value": payload})["value"], + self._clock(), + ) + if disposition == "resolved": + self._wakeup.set() + return disposition + async def signal_by_key( self, workflow: Any, @@ -3116,6 +3167,14 @@ async def recover(self) -> int: self._started_at = self._clock() await self._renew_leases() now = self._clock() + swept = await self._store.sweep_parked(now, PARKED_DELIVERY_TTL) + if swept: + console.warn( + f"{swept} parked channel deliver{'y' if swept == 1 else 'ies'} " + "went unclaimed past the TTL and became dead letters; " + "list them with the store's list_parked and replay any that " + "matter." + ) self._next_recovery_at = now + self._recovery_interval recovered, failed = await self._store.recover_orphans(now, self._max_recoveries) for run_id in failed: diff --git a/reflex/workflow/serve.py b/reflex/workflow/serve.py index 88d8f319156..a58719a4ae4 100644 --- a/reflex/workflow/serve.py +++ b/reflex/workflow/serve.py @@ -372,6 +372,96 @@ async def endpoint(request: Request) -> JSONResponse: return endpoint +def deadletters_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens): + """Build the endpoint that lists correlated webhook deliveries. + + Args: + runtime: The runtime owning the store. + tokens: The service's token scopes. + + Returns: + The endpoint callable. + """ + authorize = tokens.require("read") + + async def endpoint(request: Request) -> JSONResponse: + """List deliveries, dead letters by default. + + Args: + request: The incoming request. + + Returns: + The delivery rows. + """ + refused = authorize(request) + if refused is not None: + return refused + from reflex.workflow.records import ParkedStatus + + raw = request.query_params.get("status", "dead") + if raw == "all": + chosen = None + else: + try: + chosen = ParkedStatus(raw.upper()) + except ValueError: + return JSONResponse( + {"error": f"unknown status {raw!r}"}, status_code=400 + ) + rows = await runtime.kernel._store.list_parked(status=chosen) # pyright: ignore[reportPrivateUsage] + return JSONResponse({ + "deliveries": [ + { + "parked_id": row.parked_id, + "workflow": row.workflow_id, + "channel": row.channel, + "correlation_key": row.correlation_key, + "status": row.status.value, + "reason": row.reason, + "run_id": row.run_id, + "created_at": row.created_at, + } + for row in rows + ] + }) + + return endpoint + + +def deadletter_replay_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens): + """Build the endpoint that replays one delivery. + + Args: + runtime: The runtime owning the store. + tokens: The service's token scopes. + + Returns: + The endpoint callable. + """ + authorize = tokens.require("operate") + + async def endpoint(request: Request) -> JSONResponse: + """Route the delivery again, with the same event-id idempotency. + + Args: + request: The incoming request. + + Returns: + The routing outcome. + """ + refused = authorize(request) + if refused is not None: + return refused + disposition = await runtime.kernel._store.replay_parked( # pyright: ignore[reportPrivateUsage] + request.path_params["parked_id"], + runtime.kernel._clock(), # pyright: ignore[reportPrivateUsage] + ) + status = {"unknown_key": 404, "dead_letter": 409}.get(disposition, 202) + return JSONResponse({"disposition": disposition}, status_code=status) + + return endpoint + + def operator_endpoint(runtime: WorkflowRuntime, tokens: ScopedTokens, action: str): """Build one operator action endpoint. @@ -573,6 +663,18 @@ async def endpoint(request: Request) -> JSONResponse: # noqa: RUF029 "responses": {"202": {"description": "Delivered"}}, } }, + "/deadletters": { + "get": { + "summary": ("List correlated webhook deliveries (scope: read)"), + "responses": {"200": {"description": "The deliveries"}}, + } + }, + "/deadletters/{parked_id}/replay": { + "post": { + "summary": "Replay a delivery (scope: operate)", + "responses": {"202": {"description": "Routed"}}, + } + }, "/healthz": {"get": {"summary": "Liveness", "security": []}}, "/readyz": {"get": {"summary": "Readiness", "security": []}}, "/metrics": {"get": {"summary": "Prometheus metrics (scope: read)"}}, @@ -695,6 +797,16 @@ async def lifespan(app: Starlette) -> AsyncIterator[None]: key_signal_endpoint(runtime, tokens), methods=["POST"], ), + Route( + "/deadletters", + deadletters_endpoint(runtime, tokens), + methods=["GET"], + ), + Route( + "/deadletters/{parked_id}/replay", + deadletter_replay_endpoint(runtime, tokens), + methods=["POST"], + ), # The embedded-mode paths, kept byte-for-byte: a Stripe URL or a # minted approval link configured against an rx.App keeps working # when the deployment moves to the standalone service. diff --git a/tests/units/workflow/crash_worker.py b/tests/units/workflow/crash_worker.py index e4fd8a6e9c6..4c2fe7cf239 100644 --- a/tests/units/workflow/crash_worker.py +++ b/tests/units/workflow/crash_worker.py @@ -185,17 +185,70 @@ def _read_run_id() -> str: return Path(sys.argv[2] + ".runid").read_text().strip() +class Order(rx.State): + """A run that waits for a correlated shipment event.""" + + __workflow__ = WorkflowConfig(id="crash.order") + + shipped = rx.Signal() + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Wait for the shipment. + + Returns: + The wait. + """ + return rx.wait_for(Order.shipped, then=Order.close, timeout=rx.never) + + @rx.event(durable=True, effect="none", retry=Retry(max_attempts=1)) + def close(self, shipment): + """Handle the shipment; the ledger is the exactly-once evidence. + + Args: + shipment: The delivered payload. + + Returns: + Completion. + """ + record("shipped-handled") + return rx.complete(result=shipment) + + async def main() -> None: """Drive one phase of a crash scenario against a shared SQLite store.""" db, _, phase = sys.argv[1], sys.argv[2], sys.argv[3] store = SqliteRunStore(Path(db)) runtime = WorkflowRuntime(store, lease_duration=1.0) - for workflow_cls in (Charge, Region, Rollout): + for workflow_cls in (Charge, Region, Rollout, Order): runtime.register(workflow_cls) await runtime.startup(start_worker=False) kernel = runtime.kernel - if phase in ("unguarded", "guarded"): + if phase == "ingest_shipment": + # The provider's first delivery: durable, acked, then the process is + # killed with nothing else done -- the crash-after-ack window. + disposition = await kernel.ingest_channel( + "crash.order", "shipped", "order_1", "evt_1", {"parcel": "P-1"} + ) + assert disposition == "parked", disposition + record("acked") + die_at("after_ack") + elif phase == "redeliver": + # The provider retries twice from a fresh process; both must collapse + # into the durable row the crashed process left behind. + for _ in range(2): + disposition = await kernel.ingest_channel( + "crash.order", "shipped", "order_1", "evt_1", {"parcel": "P-1"} + ) + assert disposition == "duplicate", disposition + record("redelivered") + elif phase == "start_order": + started = await kernel.start(Order.begin(), request_key="order_1") + assert started.run_id is not None + _write_run_id(started.run_id) + await kernel.run_until_idle() + elif phase in ("unguarded", "guarded"): await kernel.start(getattr(Charge, phase)()) await kernel.recover() await kernel.run_until_idle() diff --git a/tests/units/workflow/test_crash_boundaries.py b/tests/units/workflow/test_crash_boundaries.py index c45d6084deb..ed8812a4813 100644 --- a/tests/units/workflow/test_crash_boundaries.py +++ b/tests/units/workflow/test_crash_boundaries.py @@ -205,3 +205,51 @@ def test_a_parent_closed_and_killed_still_stops_its_branches(crash): assert crash("recover").returncode == 0 assert effects(crash.ledger) == [], "a cancelled rollout must never deploy" assert sorted(runs(crash.db).values()) == ["CANCELLED"] * 3 + + +def test_a_shipment_acked_then_crashed_lands_exactly_once(crash): + """The correlated-delivery acceptance, with a real kill in the window. + + The shipment webhook arrives before the order workflow exists and the + process is SIGKILLed immediately after acknowledging it. The provider + then redelivers twice. The order workflow starts later. Exactly one + signal must reach it -- the ledger records the handler running, and the + run's inbox holds exactly one row for the channel, read straight from + the database. + + Args: + crash: The subprocess runner. + """ + first = crash("ingest_shipment", "after_ack") + assert first.returncode == -9, ( + f"expected SIGKILL, got {first.returncode}: {first.stderr.decode()[-800:]}" + ) + assert effects(crash.ledger) == ["acked"], "the ack must be durable pre-crash" + + redelivered = crash("redeliver") + assert redelivered.returncode == 0, redelivered.stderr.decode()[-800:] + started = crash("start_order") + assert started.returncode == 0, started.stderr.decode()[-800:] + + ledger = effects(crash.ledger) + assert ledger.count("shipped-handled") == 1, ledger + assert ledger.count("redelivered") == 2, ledger + + import sqlite3 + + connection = sqlite3.connect(crash.db) + try: + (inbox_rows,) = connection.execute( + "SELECT COUNT(*) FROM workflow_inbox WHERE wait_key = 'sig:shipped'" + ).fetchone() + (channel_rows,) = connection.execute( + "SELECT COUNT(*) FROM workflow_channel_inbox" + ).fetchone() + (delivered_rows,) = connection.execute( + "SELECT COUNT(*) FROM workflow_channel_inbox WHERE status = 'DELIVERED'" + ).fetchone() + finally: + connection.close() + assert inbox_rows == 1, "exactly one signal reached the run" + assert channel_rows == 1, "three deliveries are one durable event" + assert delivered_rows == 1 diff --git a/tests/units/workflow/test_ingress.py b/tests/units/workflow/test_ingress.py index f687ed27193..9722cbd9602 100644 --- a/tests/units/workflow/test_ingress.py +++ b/tests/units/workflow/test_ingress.py @@ -656,3 +656,154 @@ async def test_github_form_encoded_deliveries_are_understood( "receives its field, not the enclosing object" ) await runtime.shutdown() + + +class Fulfil(rx.State): + """A workflow whose channel is fed by a correlated provider webhook.""" + + __workflow__ = WorkflowConfig(id="ingress.fulfil") + + shipped = rx.Signal( + trigger=webhook( + "carrier_shipped", + dedupe_by="event_id", + correlate_by="order_id", + verify=hmac_signature( + secret_env="STRIPE_WEBHOOK_SECRET", header="X-Signature" + ), + ) + ) + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Wait for the shipment. + + Returns: + The wait. + """ + return rx.wait_for(Fulfil.shipped, then=Fulfil.close, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def close(self, shipment): + """Finish with the shipment. + + Args: + shipment: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result=shipment) + + +async def test_a_correlated_webhook_reaches_its_run_exactly_once( + monkeypatch, forked_registration_context +): + """Early delivery, three sends, late run: one signal, over real HTTP. + + Args: + monkeypatch: Used to install the webhook secret. + forked_registration_context: Isolated state registry. + """ + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + runtime = WorkflowRuntime(MemoryRunStore()) + runtime.register(Fulfil) + await runtime.startup(start_worker=False) + app = Starlette( + routes=[Route(WEBHOOK_ROUTE, webhook_endpoint(runtime), methods=["POST"])] + ) + body = json.dumps({ + "event_id": "evt_7", + "order_id": "ord_7", + "parcel": "P-7", + }).encode() + with TestClient(app) as client: + first = client.post( + "/_workflow/webhook/carrier_shipped", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert first.status_code == 202, first.text + assert first.json()["disposition"] == "parked" + for _ in range(2): + again = client.post( + "/_workflow/webhook/carrier_shipped", + content=body, + headers={"x-signature": _sign(body)}, + ) + assert again.status_code == 202 + assert again.json()["disposition"] == "duplicate" + + keyless = json.dumps({"event_id": "evt_8", "parcel": "P-8"}).encode() + refused = client.post( + "/_workflow/webhook/carrier_shipped", + content=keyless, + headers={"x-signature": _sign(keyless)}, + ) + assert refused.status_code == 400 + assert "order_id" in refused.json()["error"] + + started = await runtime.kernel.start(Fulfil.begin(), request_key="ord_7") + assert started.run_id is not None + await runtime.kernel.run_until_idle() + snapshot = await runtime.kernel.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == { + "event_id": "evt_7", + "order_id": "ord_7", + "parcel": "P-7", + } + parked = await runtime.kernel._store.list_parked() # pyright: ignore[reportPrivateUsage] + assert len(parked) == 1 + assert parked[0].status.value == "DELIVERED" + await runtime.shutdown() + + +def test_a_channel_topic_cannot_collide_with_a_root_topic( + monkeypatch, forked_registration_context +): + """One topic, one target: a root and a channel cannot share it. + + Args: + monkeypatch: Used to install the webhook secret. + forked_registration_context: Isolated state registry. + """ + from reflex.workflow.definition import compile_workflow + from reflex.workflow.ingress import collect_webhook_routes + + monkeypatch.setenv("STRIPE_WEBHOOK_SECRET", SECRET) + + class Both(rx.State): + __workflow__ = WorkflowConfig(id="ingress.both") + + colliding = rx.Signal( + trigger=webhook( + "shipped", + dedupe_by="id", + correlate_by="order", + verify=hmac_signature( + secret_env="STRIPE_WEBHOOK_SECRET", header="X-Signature" + ), + ) + ) + + @rx.event( + durable=True, + effect="none", + trigger=webhook( + "shipped", + verify=hmac_signature( + secret_env="STRIPE_WEBHOOK_SECRET", header="X-Signature" + ), + ), + ) + def on_shipped(self, event: dict): + """Claim the same topic as the channel. + + Args: + event: The payload. + """ + + with pytest.raises(WorkflowDefinitionError, match="shipped"): + collect_webhook_routes((compile_workflow(Both),)) diff --git a/tests/units/workflow/test_serve.py b/tests/units/workflow/test_serve.py index 69df51a94a1..470b827611e 100644 --- a/tests/units/workflow/test_serve.py +++ b/tests/units/workflow/test_serve.py @@ -438,3 +438,88 @@ def test_business_keys_address_runs_without_run_ids(service): "/workflows/serve.nope/keys/order_123", headers=_auth("tk_read") ) assert unknown_workflow.status_code == 404 + + +def test_dead_letters_are_visible_and_replayable_over_http( + forked_registration_context, +): + """The operator's dead-letter loop over serve: park, list, replay. + + Args: + forked_registration_context: Isolated state registry. + """ + + class Freight(rx.State): + __workflow__ = WorkflowConfig(id="serve.freight") + + arrived = Signal( + trigger=webhook( + "freight_arrived", + dedupe_by="event_id", + correlate_by="shipment_id", + allow_unverified=True, + unverified_reason="test-only channel", + ) + ) + + @rx.event(durable=True, effect="none", trigger=manual()) + def begin(self): + """Wait for arrival. + + Returns: + The wait. + """ + return rx.wait_for(Freight.arrived, then=Freight.close, timeout=rx.never) + + @rx.event(durable=True, effect="none") + def close(self, event): + """Finish. + + Args: + event: The delivered payload. + + Returns: + Completion. + """ + return rx.complete(result=event) + + runtime = WorkflowRuntime(testing.MemoryRunStore()) + runtime.register(Freight) + app = build_app( + runtime, + worker=False, + drain=0, + tokens=_tokens(tk_read="read", tk_operate="operate"), + ) + with TestClient(app) as client: + assert ( + client.get("/deadletters?status=all", headers=_auth("tk_read")).json()[ + "deliveries" + ] + == [] + ) + parked = client.post( + "/_workflow/webhook/freight_arrived", + json={"event_id": "evt_d", "shipment_id": "ship_9"}, + ) + assert parked.status_code == 202, parked.text + assert parked.json()["disposition"] == "parked" + + rows = client.get( + "/deadletters?status=pending", headers=_auth("tk_read") + ).json()["deliveries"] + assert len(rows) == 1 + assert rows[0]["correlation_key"] == "ship_9" + parked_id = rows[0]["parked_id"] + + forbidden = client.post( + f"/deadletters/{parked_id}/replay", headers=_auth("tk_read") + ) + assert forbidden.status_code == 403 + replayed = client.post( + f"/deadletters/{parked_id}/replay", headers=_auth("tk_operate") + ) + assert replayed.status_code == 202 + assert replayed.json()["disposition"] == "parked", "still no run to take it" + missing = client.post("/deadletters/nope/replay", headers=_auth("tk_operate")) + assert missing.status_code == 404 From 5b5ed5074de9bc5fde0fea5a3d65ffcbdb46f29f Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Mon, 24 Aug 2026 18:00:58 -0700 Subject: [PATCH 119/121] workflows: competitor-pattern parity suite Ten canonical patterns from five ecosystems, recreated as executable Reflex workflows from the projects' own official examples and run as tests on every store: Temporal's money transfer, Restate's travel-saga compensation, Inngest's onboarding wait-for-event / fan-out / checkpointing, Prefect's API-sourced ETL, Celery's chain and chord, Airflow's TaskFlow ETL, and DBOS's transactional outbox. 48 test rows across Memory, SQLite, and Postgres, all green. Beyond regression cover, this is the parity ledger: when a competitor's pattern needs contortions here, the contortion is written down where CI runs it. Gaps the exercise surfaced, tracked for the roadmap: per-step retry / timeout / queue overrides on rx.step; pub/sub fan-out atop the new correlated-event inbox; a first-class saga/compensation stack; result piping and lighter parallel branches; and transaction-coupled database steps, where DBOS's same-transaction outbox guarantee has no current equivalent. --- .../workflow/test_examples_airflow_dbos.py | 171 ++++++ tests/units/workflow/test_examples_inngest.py | 547 ++++++++++++++++++ .../workflow/test_examples_prefect_celery.py | 357 ++++++++++++ .../test_examples_temporal_restate.py | 528 +++++++++++++++++ 4 files changed, 1603 insertions(+) create mode 100644 tests/units/workflow/test_examples_airflow_dbos.py create mode 100644 tests/units/workflow/test_examples_inngest.py create mode 100644 tests/units/workflow/test_examples_prefect_celery.py create mode 100644 tests/units/workflow/test_examples_temporal_restate.py diff --git a/tests/units/workflow/test_examples_airflow_dbos.py b/tests/units/workflow/test_examples_airflow_dbos.py new file mode 100644 index 00000000000..697f0418e23 --- /dev/null +++ b/tests/units/workflow/test_examples_airflow_dbos.py @@ -0,0 +1,171 @@ +"""Executable Reflex translations of canonical Airflow and DBOS examples. + +Official sources: + +* https://airflow.apache.org/docs/apache-airflow/stable/tutorial_taskflow_api.html +* https://docs.dbos.dev/python/examples/outbox +""" + +import pytest +from reflex_base.workflow import ( + Retry, + TransientWorkflowError, + WorkflowConfig, + complete, + manual, +) + +import reflex as rx +from reflex.workflow.records import RunStatus +from reflex.workflow.testing import WorkflowTestHarness + +AIRFLOW_LOADS: list[float] = [] + + +class AirflowTaskFlowETL(rx.State): + """Airflow's TaskFlow tutorial: extract, transform, then load order data.""" + + __workflow__ = WorkflowConfig(id="examples.airflow_taskflow_etl") + + orders: dict[str, float] = {} + total_order_value: float = 0.0 + + @rx.event(durable=True, trigger=manual(), effect="read") + def extract(self): + """Extract the tutorial's hard-coded orders. + + Returns: + The transform step. + """ + self.orders = {"1001": 301.27, "1002": 433.21, "1003": 502.22} + return AirflowTaskFlowETL.transform + + @rx.event(durable=True, effect="none") + def transform(self): + """Compute the same aggregate as the TaskFlow transform task. + + Returns: + The load step. + """ + self.total_order_value = sum(self.orders.values()) + return AirflowTaskFlowETL.load + + @rx.event(durable=True, effect="idempotent_write") + def load(self): + """Stand in for loading the aggregate into an analytics sink. + + Returns: + Completion containing the aggregate. + """ + AIRFLOW_LOADS.append(self.total_order_value) + return complete(result={"total_order_value": self.total_order_value}) + + +async def test_airflow_taskflow_etl_translates_to_durable_handlers( + forked_registration_context, +): + """The TaskFlow tutorial maps directly to three persisted transitions.""" + AIRFLOW_LOADS.clear() + + async with WorkflowTestHarness(AirflowTaskFlowETL) as harness: + started = await harness.start(AirflowTaskFlowETL.extract) + assert started.run_id is not None + + run = await harness.get_run(started.run_id) + assert run is not None + assert run.status is RunStatus.COMPLETED + assert run.result["total_order_value"] == pytest.approx(1236.70) + assert [pytest.approx(1236.70)] == AIRFLOW_LOADS + assert [step.handler_id for step in run.steps] == [ + "extract", + "transform", + "load", + ] + + +DBOS_INSERTS: list[str] = [] +DBOS_NOTIFICATIONS: list[str] = [] +DBOS_NOTIFICATION_ATTEMPTS: list[str] = [] + + +def _insert_order(customer: str, item: str, quantity: int) -> str: + """Simulate DBOS's transactional order insert. + + Returns: + The inserted order ID. + """ + order_id = f"order-{customer}-{item}-{quantity}" + DBOS_INSERTS.append(order_id) + return order_id + + +def _send_order_notification(order_id: str) -> None: + """Fail the first delivery, like a temporarily unavailable broker.""" + DBOS_NOTIFICATION_ATTEMPTS.append(order_id) + if len(DBOS_NOTIFICATION_ATTEMPTS) == 1: + msg = "notification broker unavailable" + raise TransientWorkflowError(msg) + DBOS_NOTIFICATIONS.append(order_id) + + +class DBOSTransactionalOutbox(rx.State): + """Best available translation of DBOS's transactional-outbox workflow.""" + + __workflow__ = WorkflowConfig(id="examples.dbos_transactional_outbox") + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=2, initial_delay="1s", jitter="none"), + ) + async def place_order(self, customer: str, item: str, quantity: int): + """Checkpoint the insert before attempting the notification. + + This recreates DBOS's authoring shape, but not its atomic guarantee: + ``rx.step`` records after the callable returns, so the real insert + must still be idempotent across a crash between commit and recording. + + Returns: + Completion containing the order ID. + """ + order_id = await rx.step( + "insert_order", _insert_order, customer, item, quantity + ) + await rx.step("send_order_notification", _send_order_notification, order_id) + return complete(result={"order_id": order_id, "notified": True}) + + +async def test_dbos_outbox_shape_replays_the_insert_after_notification_failure( + forked_registration_context, +): + """A recorded Reflex substep is not repeated by a later handler retry.""" + DBOS_INSERTS.clear() + DBOS_NOTIFICATIONS.clear() + DBOS_NOTIFICATION_ATTEMPTS.clear() + + async with WorkflowTestHarness(DBOSTransactionalOutbox) as harness: + started = await harness.start( + DBOSTransactionalOutbox.place_order("alice", "book", 2) + ) + assert started.run_id is not None + + retrying = await harness.get_run(started.run_id) + assert retrying is not None + assert retrying.status is RunStatus.RETRYING + assert len(DBOS_INSERTS) == 1 + + await harness.advance("1s") + completed = await harness.get_run(started.run_id) + assert completed is not None + assert completed.status is RunStatus.COMPLETED + assert completed.result == { + "order_id": "order-alice-book-2", + "notified": True, + } + assert DBOS_INSERTS == ["order-alice-book-2"] + assert DBOS_NOTIFICATION_ATTEMPTS == [ + "order-alice-book-2", + "order-alice-book-2", + ] + assert DBOS_NOTIFICATIONS == ["order-alice-book-2"] diff --git a/tests/units/workflow/test_examples_inngest.py b/tests/units/workflow/test_examples_inngest.py new file mode 100644 index 00000000000..5c75b3570db --- /dev/null +++ b/tests/units/workflow/test_examples_inngest.py @@ -0,0 +1,547 @@ +"""Executable Reflex translations of canonical Inngest workflow examples. + +Official sources: + +* https://www.inngest.com/docs/features/inngest-functions/steps-workflows/wait-for-event +* https://www.inngest.com/docs/guides/fan-out-jobs +* https://www.inngest.com/docs/learn/inngest-steps +""" + +from pydantic import BaseModel +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import HistoryEventType, RunStatus, StepStatus +from reflex.workflow.testing import WorkflowTestHarness + +ONBOARDING_EMAILS: list[tuple[str, str]] = [] +FANOUT_EFFECTS: list[tuple[str, str]] = [] +CRM_ATTEMPTS: list[str] = [] +CHECKPOINT_CALLS: list[str] = [] + + +class OnboardingCompleted(BaseModel): + """The event that releases an onboarding drip campaign.""" + + user_id: str + + +def _send_onboarding_email(email: str, template: str) -> dict[str, str]: + """Simulate an email provider call. + + Args: + email: The recipient. + template: The email template. + + Returns: + The provider response. + """ + ONBOARDING_EMAILS.append((email, template)) + return {"email": email, "template": template} + + +class OnboardingDrip(rx.State): + """Welcome a user, then send tips or a timeout nudge.""" + + __workflow__ = WorkflowConfig(id="examples.inngest.onboarding") + + completed = rx.Signal(OnboardingCompleted) + user_id: str = "" + email: str = "" + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def account_created(self, user_id: str, email: str): + """Send the welcome email and wait up to three days. + + Args: + user_id: The user's business identity and run request key. + email: The user's email address. + + Returns: + A durable wait for onboarding completion. + """ + self.user_id = user_id + self.email = email + await rx.step("welcome-email", _send_onboarding_email, email, "welcome") + return rx.wait_for( + OnboardingDrip.completed, + then=OnboardingDrip.send_tips, + timeout="3d", + on_timeout=OnboardingDrip.send_nudge, + ) + + @rx.event(durable=True, effect="idempotent_write") + async def send_tips(self, completion: OnboardingCompleted): + """Send the successful-onboarding tips email. + + Args: + completion: The completion event addressed to this user's run. + + Returns: + Completed campaign result, or a correlation failure. + """ + if completion.user_id != self.user_id: + return rx.fail( + "correlation_mismatch", + details={ + "expected_user_id": self.user_id, + "received_user_id": completion.user_id, + }, + ) + await rx.step("tips-email", _send_onboarding_email, self.email, "tips") + return rx.complete(result={"path": "completed", "user_id": completion.user_id}) + + @rx.event(durable=True, effect="idempotent_write") + async def send_nudge(self): + """Send the timeout nudge. + + Returns: + Completed campaign result. + """ + await rx.step("nudge-email", _send_onboarding_email, self.email, "nudge") + return rx.complete(result={"path": "timeout", "user_id": self.user_id}) + + +def _record_signup_effect(service: str, user_id: str) -> dict[str, str]: + """Record one successful downstream signup integration. + + Args: + service: The downstream service name. + user_id: The signed-up user. + + Returns: + The service result. + """ + FANOUT_EFFECTS.append((service, user_id)) + return {"service": service, "user_id": user_id} + + +def _create_crm_contact(user_id: str) -> dict[str, str]: + """Fail CRM once to demonstrate branch-level independence. + + Args: + user_id: The signed-up user. + + Returns: + The CRM result after recovery. + + Raises: + TransientWorkflowError: On the first provider attempt. + """ + CRM_ATTEMPTS.append(user_id) + if len(CRM_ATTEMPTS) == 1: + msg = "CRM temporarily unavailable" + raise TransientWorkflowError(msg) + return _record_signup_effect("crm", user_id) + + +class SignupWelcome(rx.State): + """The welcome-email branch of signup fan-out.""" + + __workflow__ = WorkflowConfig(id="examples.inngest.signup.welcome") + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def send(self, user_id: str, email: str): + """Send the welcome email. + + Args: + user_id: The signed-up user. + email: The user's email address. + + Returns: + The branch completion. + """ + del email + result = await rx.step( + "welcome-email", _record_signup_effect, "welcome", user_id + ) + return rx.complete(result=result) + + +class StripeTrial(rx.State): + """The Stripe-trial branch of signup fan-out.""" + + __workflow__ = WorkflowConfig(id="examples.inngest.signup.stripe") + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def start(self, user_id: str, email: str): + """Create the user's Stripe trial. + + Args: + user_id: The signed-up user. + email: The user's email address. + + Returns: + The branch completion. + """ + del email + result = await rx.step("stripe-trial", _record_signup_effect, "stripe", user_id) + return rx.complete(result=result) + + +class CrmContact(rx.State): + """The CRM branch of signup fan-out.""" + + __workflow__ = WorkflowConfig(id="examples.inngest.signup.crm") + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=2, initial_delay="1s", jitter="none"), + ) + async def add(self, user_id: str, email: str): + """Create the CRM contact, retrying a transient outage. + + Args: + user_id: The signed-up user. + email: The user's email address. + + Returns: + The branch completion. + """ + del email + result = await rx.step("crm-contact", _create_crm_contact, user_id) + return rx.complete(result=result) + + +class MailingList(rx.State): + """The mailing-list branch of signup fan-out.""" + + __workflow__ = WorkflowConfig(id="examples.inngest.signup.mailing_list") + + @rx.event(durable=True, trigger=manual(), effect="idempotent_write") + async def subscribe(self, user_id: str, email: str): + """Subscribe the user to the mailing list. + + Args: + user_id: The signed-up user. + email: The user's email address. + + Returns: + The branch completion. + """ + del email + result = await rx.step( + "mailing-list", _record_signup_effect, "mailing-list", user_id + ) + return rx.complete(result=result) + + +class SignupFanout(rx.State): + """Run four signup integrations independently, then join them.""" + + __workflow__ = WorkflowConfig(id="examples.inngest.signup") + services: list[str] = [] + + @rx.event(durable=True, trigger=manual(), effect="none") + def signup(self, user_id: str, email: str): + """Fan one signup out to all downstream services. + + Args: + user_id: The signed-up user. + email: The user's email address. + + Returns: + A parallel branch join. + """ + return rx.parallel( + SignupWelcome.send(user_id, email), + StripeTrial.start(user_id, email), + CrmContact.add(user_id, email), + MailingList.subscribe(user_id, email), + then=SignupFanout.join, + ) + + @rx.event(durable=True, effect="none") + def join(self, results: list[dict]): + """Collect all four independent branch results. + + Args: + results: Child outcomes in declaration order. + + Returns: + Completion with every service name. + """ + self.services = [entry["result"]["service"] for entry in results] + return rx.complete(result={"services": self.services}) + + +def _create_external_customer(user_id: str) -> dict[str, str]: + """Simulate a non-repeatable external customer creation. + + Args: + user_id: The user being provisioned. + + Returns: + The external customer identity. + """ + CHECKPOINT_CALLS.append("create-customer") + return {"customer_id": f"cus_{user_id}"} + + +def _finalize_profile(customer_id: str) -> dict[str, str]: + """Fail once after customer creation, then finalize the profile. + + Args: + customer_id: The already-created external customer. + + Returns: + The finalized profile. + + Raises: + TransientWorkflowError: On the first attempt. + """ + CHECKPOINT_CALLS.append("finalize-profile") + if CHECKPOINT_CALLS.count("finalize-profile") == 1: + msg = "profile service temporarily unavailable" + raise TransientWorkflowError(msg) + return {"customer_id": customer_id, "status": "ready"} + + +class CheckpointedSignup(rx.State): + """Demonstrate Inngest ``step.run``-style replay with ``rx.step``.""" + + __workflow__ = WorkflowConfig(id="examples.inngest.checkpoint") + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=2, initial_delay="1s", jitter="none"), + ) + async def provision(self, user_id: str): + """Create a customer once even when a later step retries. + + Args: + user_id: The user being provisioned. + + Returns: + Completion with the finalized profile. + """ + customer = await rx.step("create-customer", _create_external_customer, user_id) + profile = await rx.step( + "finalize-profile", _finalize_profile, customer["customer_id"] + ) + return rx.complete(result=profile) + + +async def test_onboarding_completion_sends_tips_before_timeout( + forked_registration_context, +): + """A matching business-key signal chooses tips, never the nudge.""" + ONBOARDING_EMAILS.clear() + async with WorkflowTestHarness(OnboardingDrip) as harness: + started = await harness.start( + OnboardingDrip.account_created("user-complete", "done@example.com"), + request_key="user-complete", + ) + assert started.run_id is not None + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.WAITING + assert ONBOARDING_EMAILS == [("done@example.com", "welcome")] + + await harness.advance("2d") + assert ( + await rx.workflows.signal_by_key( + OnboardingDrip, + "user-complete", + OnboardingDrip.completed(OnboardingCompleted(user_id="user-complete")), + key="onboarding-completed:user-complete", + ) + == "resolved" + ) + await harness.run_until_idle() + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"path": "completed", "user_id": "user-complete"} + assert ONBOARDING_EMAILS == [ + ("done@example.com", "welcome"), + ("done@example.com", "tips"), + ] + + await harness.advance("2d") + assert ("done@example.com", "nudge") not in ONBOARDING_EMAILS + + +async def test_onboarding_timeout_sends_nudge_not_tips( + forked_registration_context, +): + """No completion within three days deterministically chooses the nudge.""" + ONBOARDING_EMAILS.clear() + async with WorkflowTestHarness(OnboardingDrip) as harness: + started = await harness.start( + OnboardingDrip.account_created("user-timeout", "late@example.com"), + request_key="user-timeout", + ) + assert started.run_id is not None + + await harness.advance("3d") + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"path": "timeout", "user_id": "user-timeout"} + assert ONBOARDING_EMAILS == [ + ("late@example.com", "welcome"), + ("late@example.com", "nudge"), + ] + + +async def test_onboarding_signal_before_wait_is_buffered( + forked_registration_context, +): + """An admitted run keeps a completion that beats welcome-email execution.""" + ONBOARDING_EMAILS.clear() + async with WorkflowTestHarness(OnboardingDrip) as harness: + started = await harness.start_only( + OnboardingDrip.account_created("user-early", "early@example.com"), + request_key="user-early", + ) + assert started.run_id is not None + assert ( + await rx.workflows.signal_by_key( + OnboardingDrip, + "user-early", + OnboardingDrip.completed(OnboardingCompleted(user_id="user-early")), + key="onboarding-completed:user-early", + ) + == "buffered" + ) + + await harness.run_until_idle() + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == {"path": "completed", "user_id": "user-early"} + assert ONBOARDING_EMAILS == [ + ("early@example.com", "welcome"), + ("early@example.com", "tips"), + ] + + +async def test_onboarding_rejects_mismatched_signal_payload( + forked_registration_context, +): + """Addressing a run by key cannot silently substitute another user.""" + ONBOARDING_EMAILS.clear() + async with WorkflowTestHarness(OnboardingDrip) as harness: + started = await harness.start( + OnboardingDrip.account_created("user-a", "a@example.com"), + request_key="user-a", + ) + assert started.run_id is not None + + assert ( + await rx.workflows.signal_by_key( + OnboardingDrip, + "user-a", + OnboardingDrip.completed(OnboardingCompleted(user_id="user-b")), + key="onboarding-completed:user-b", + ) + == "resolved" + ) + await harness.run_until_idle() + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.error == { + "reason": "correlation_mismatch", + "details": { + "expected_user_id": "user-a", + "received_user_id": "user-b", + }, + } + assert ONBOARDING_EMAILS == [("a@example.com", "welcome")] + + +async def test_signup_fanout_is_independent_and_joins( + forked_registration_context, +): + """Three integrations finish while CRM retries; the parent waits for all.""" + FANOUT_EFFECTS.clear() + CRM_ATTEMPTS.clear() + workflows = (SignupFanout, SignupWelcome, StripeTrial, CrmContact, MailingList) + async with WorkflowTestHarness(*workflows) as harness: + started = await harness.start( + SignupFanout.signup("user-42", "user42@example.com") + ) + assert started.run_id is not None + + parent = await harness.get_run(started.run_id) + assert parent is not None + assert parent.status is RunStatus.WAITING + assert parent.steps[1].status is StepStatus.BLOCKED + assert parent.steps[1].join_expected == 4 + assert parent.steps[1].join_arrived == 3 + assert sorted(FANOUT_EFFECTS) == [ + ("mailing-list", "user-42"), + ("stripe", "user-42"), + ("welcome", "user-42"), + ] + assert CRM_ATTEMPTS == ["user-42"] + + runs = await harness.kernel.list_runs() + children = [run for run in runs if run.parent_run_id == started.run_id] + assert len(children) == 4 + assert sum(run.status is RunStatus.COMPLETED for run in children) == 3 + + await harness.advance("1s") + + parent = await harness.get_run(started.run_id) + assert parent is not None + assert parent.status is RunStatus.COMPLETED + assert parent.result == { + "services": ["welcome", "stripe", "crm", "mailing-list"] + } + assert sorted(FANOUT_EFFECTS) == [ + ("crm", "user-42"), + ("mailing-list", "user-42"), + ("stripe", "user-42"), + ("welcome", "user-42"), + ] + assert CRM_ATTEMPTS == ["user-42", "user-42"] + + runs = await harness.kernel.list_runs() + children = [run for run in runs if run.parent_run_id == started.run_id] + assert len(children) == 4 + assert all(run.status is RunStatus.COMPLETED for run in children) + + +async def test_completed_step_is_not_repeated_after_later_failure( + forked_registration_context, +): + """A later retry replays the customer result instead of creating another.""" + CHECKPOINT_CALLS.clear() + async with WorkflowTestHarness(CheckpointedSignup) as harness: + started = await harness.start(CheckpointedSignup.provision("user-7")) + assert started.run_id is not None + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is not RunStatus.COMPLETED + assert CHECKPOINT_CALLS == ["create-customer", "finalize-profile"] + + await harness.advance("1s") + + snapshot = await harness.get_run(started.run_id) + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.result == { + "customer_id": "cus_user-7", + "status": "ready", + } + assert CHECKPOINT_CALLS.count("create-customer") == 1 + assert CHECKPOINT_CALLS.count("finalize-profile") == 2 + + history = await harness.kernel.store.get_history(started.run_id) + recorded = [ + event.data["key"] + for event in history + if event.type is HistoryEventType.SUBSTEP_RECORDED + ] + assert recorded == ["create-customer", "finalize-profile"] diff --git a/tests/units/workflow/test_examples_prefect_celery.py b/tests/units/workflow/test_examples_prefect_celery.py new file mode 100644 index 00000000000..7200c6a87b7 --- /dev/null +++ b/tests/units/workflow/test_examples_prefect_celery.py @@ -0,0 +1,357 @@ +"""Canonical Prefect and Celery examples expressed as Reflex workflows. + +The examples intentionally follow the official documentation shapes closely: + +* Prefect API-sourced ETL: + https://docs.prefect.io/v3/examples/run-api-sourced-etl +* Celery chains, groups, and chords: + https://docs.celeryq.dev/en/stable/userguide/canvas.html + +External API and database calls are simulated so the durable orchestration is +tested deterministically against every configured workflow store. +""" + +from reflex_base.workflow import Retry, TransientWorkflowError, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import RunStatus, StepStatus +from reflex.workflow.testing import WorkflowTestHarness + +Article = dict[str, str | int] + +FETCH_ATTEMPTS: dict[int, int] = {} +FETCH_FAILURES_REMAINING: dict[int, int] = {} +FETCH_ALWAYS_FAIL: set[int] = set() +LOADED_BATCHES: list[list[Article]] = [] + + +def _reset_prefect_io() -> None: + """Reset the deterministic API and database doubles.""" + FETCH_ATTEMPTS.clear() + FETCH_FAILURES_REMAINING.clear() + FETCH_ALWAYS_FAIL.clear() + LOADED_BATCHES.clear() + + +def _simulated_api_page(page: int) -> list[Article]: + """Return the deterministic response for one API page. + + Args: + page: The requested page number. + + Returns: + Two article records. + + Raises: + TransientWorkflowError: When the configured API double is unavailable. + """ + FETCH_ATTEMPTS[page] = FETCH_ATTEMPTS.get(page, 0) + 1 + failures_remaining = FETCH_FAILURES_REMAINING.get(page, 0) + if page in FETCH_ALWAYS_FAIL or failures_remaining: + if failures_remaining: + FETCH_FAILURES_REMAINING[page] = failures_remaining - 1 + msg = f"article API returned 503 for page {page}" + raise TransientWorkflowError(msg) + return [ + { + "id": page * 10 + offset, + "title": f"article-{page}-{offset}", + "published_at": f"2026-08-{page + offset:02d}", + "url": f"https://example.test/articles/{page}-{offset}", + } + for offset in (1, 2) + ] + + +class PrefectFetchPage(rx.State): + """One independently retryable API-page task.""" + + __workflow__ = WorkflowConfig(id="examples.prefect.fetch_page") + + @rx.event( + durable=True, + trigger=manual(), + effect="read", + retry=Retry(max_attempts=2, initial_delay="1s", jitter="none"), + ) + def fetch(self, page: int): + """Fetch and return one page. + + Args: + page: The requested page number. + + Returns: + Completion carrying that page's records. + """ + return rx.complete(result=_simulated_api_page(page)) + + +class PrefectApiEtl(rx.State): + """Fan out extraction, transform the join, then durably load it.""" + + __workflow__ = WorkflowConfig(id="examples.prefect.api_etl") + articles: list[Article] = [] + page_count: int = 0 + + @rx.event(durable=True, trigger=manual(), effect="none") + def etl(self, pages: list[int]): + """Start one durable child run per API page. + + Args: + pages: Page numbers to extract. + + Returns: + A fan-out joined by ``combine_pages``. + """ + self.page_count = len(pages) + return rx.parallel( + *(PrefectFetchPage.fetch(page) for page in pages), + then=PrefectApiEtl.combine_pages, + ) + + @rx.event(durable=True, effect="none") + def combine_pages(self, results: list): + """Flatten successful pages or fail the ETL before transformation. + + Args: + results: Ordered child-run outcomes from the fan-out. + + Returns: + The transform step, or terminal failure. + """ + failures = [result for result in results if result["status"] != "COMPLETED"] + if failures: + return rx.fail(reason=f"{len(failures)} API page(s) failed") + records = [article for result in results for article in result["result"]] + return PrefectApiEtl.transform(records) + + @rx.event(durable=True, effect="none") + def transform(self, records: list[Article]): + """Normalize the joined API response into the load schema. + + Args: + records: Flattened raw articles. + + Returns: + The durable load step. + """ + self.articles = [ + { + "id": record["id"], + "title": str(record["title"]).upper(), + "published_at": record["published_at"], + "url": record["url"], + } + for record in records + ] + return PrefectApiEtl.load + + @rx.event(durable=True, effect="idempotent_write") + def load(self): + """Write the transformed batch to the simulated database. + + Returns: + Completion carrying the load summary. + """ + LOADED_BATCHES.append([dict(article) for article in self.articles]) + return rx.complete( + result={"loaded": len(self.articles), "pages": self.page_count} + ) + + +class CeleryChain(rx.State): + """The official ``add(2, 2) | add(4) | add(8)`` Canvas chain.""" + + __workflow__ = WorkflowConfig(id="examples.celery.chain") + value: int = 0 + intermediate_results: list[int] = [] + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, left: int, right: int): + """Run the first addition and allocate the rest of the chain. + + Args: + left: First addend. + right: Second addend. + + Returns: + The remaining additions and final result step. + """ + self.value = left + right + self.intermediate_results = [self.value] + return [ + CeleryChain.add_previous(4), + CeleryChain.add_previous(8), + CeleryChain.finish, + ] + + @rx.event(durable=True, effect="none") + def add_previous(self, addend: int): + """Add to the durable result of the preceding handler. + + Args: + addend: Value from the next Celery partial signature. + """ + self.value += addend + self.intermediate_results = [*self.intermediate_results, self.value] + + @rx.event(durable=True, effect="none") + def finish(self): + """Return the chain's final value. + + Returns: + Completion carrying the final sum. + """ + return rx.complete(result=self.value) + + +class CeleryAdd(rx.State): + """The task placed in the Celery-style group header.""" + + __workflow__ = WorkflowConfig(id="examples.celery.add") + + @rx.event(durable=True, trigger=manual(), effect="none") + def add(self, left: int, right: int): + """Add two values. + + Args: + left: First addend. + right: Second addend. + + Returns: + Completion carrying their sum. + """ + return rx.complete(result=left + right) + + +class CeleryChord(rx.State): + """A parallel group whose ordered results feed a sum callback.""" + + __workflow__ = WorkflowConfig(id="examples.celery.chord") + group_results: list[int] = [] + + @rx.event(durable=True, trigger=manual(), effect="none") + def start(self, size: int): + """Create the equivalent of ``group(add.s(i, i)) | tsum.s()``. + + Args: + size: Number of additions in the group. + + Returns: + A parallel fan-out joined by the sum callback. + """ + return rx.parallel( + *(CeleryAdd.add(index, index) for index in range(size)), + then=CeleryChord.sum_results, + ) + + @rx.event(durable=True, effect="none") + def sum_results(self, results: list): + """Fail like a chord or sum every successful group result. + + Args: + results: Ordered child-run outcomes from the group. + + Returns: + Terminal failure or completion carrying the total. + """ + failures = [result for result in results if result["status"] != "COMPLETED"] + if failures: + return rx.fail(reason=f"{len(failures)} chord task(s) failed") + self.group_results = [result["result"] for result in results] + return rx.complete(result=sum(self.group_results)) + + +async def test_prefect_etl_retries_one_page_then_loads(): + """Each page retries independently before transform and load resume.""" + _reset_prefect_io() + FETCH_FAILURES_REMAINING[2] = 1 + + async with WorkflowTestHarness(PrefectApiEtl, PrefectFetchPage) as harness: + started = await harness.start(PrefectApiEtl.etl([1, 2, 3])) + assert started.run_id is not None + + waiting = await harness.get_run(started.run_id) + assert waiting is not None + assert waiting.status is RunStatus.WAITING + assert waiting.steps[1].status is StepStatus.BLOCKED + assert waiting.steps[1].join_arrived == 2 + assert FETCH_ATTEMPTS == {1: 1, 2: 1, 3: 1} + assert not LOADED_BATCHES + + await harness.advance("1s") + + completed = await harness.get_run(started.run_id) + assert completed is not None + assert completed.status is RunStatus.COMPLETED + assert completed.result == {"loaded": 6, "pages": 3} + assert FETCH_ATTEMPTS == {1: 1, 2: 2, 3: 1} + assert len(LOADED_BATCHES) == 1 + assert [article["id"] for article in LOADED_BATCHES[0]] == [ + 11, + 12, + 21, + 22, + 31, + 32, + ] + assert all( + str(article["title"]).startswith("ARTICLE-") + for article in LOADED_BATCHES[0] + ) + + +async def test_prefect_etl_does_not_load_after_page_retry_exhaustion(): + """A terminal page failure fails the parent without calling the loader.""" + _reset_prefect_io() + FETCH_ALWAYS_FAIL.add(2) + + async with WorkflowTestHarness(PrefectApiEtl, PrefectFetchPage) as harness: + started = await harness.start(PrefectApiEtl.etl([1, 2, 3])) + assert started.run_id is not None + await harness.advance("1s") + + failed = await harness.get_run(started.run_id) + assert failed is not None + assert failed.status is RunStatus.FAILED + assert FETCH_ATTEMPTS == {1: 1, 2: 2, 3: 1} + assert not LOADED_BATCHES + + +async def test_celery_canvas_chain_threads_each_intermediate_result(): + """The official three-addition chain produces 4, then 8, then 16.""" + async with WorkflowTestHarness(CeleryChain) as harness: + started = await harness.start(CeleryChain.start(2, 2)) + assert started.run_id is not None + + completed = await harness.get_run(started.run_id) + assert completed is not None + assert completed.status is RunStatus.COMPLETED + assert completed.result == 16 + assert completed.state["intermediate_results"] == [4, 8, 16] + assert [step.status for step in completed.steps] == [ + StepStatus.SUCCEEDED, + StepStatus.SUCCEEDED, + StepStatus.SUCCEEDED, + StepStatus.SUCCEEDED, + ] + + +async def test_celery_group_chord_adds_in_parallel_then_sums(): + """Ten ``add(i, i)`` child runs join into the documented total of 90.""" + async with WorkflowTestHarness(CeleryChord, CeleryAdd) as harness: + started = await harness.start(CeleryChord.start(10)) + assert started.run_id is not None + + completed = await harness.get_run(started.run_id) + assert completed is not None + assert completed.status is RunStatus.COMPLETED + assert completed.result == 90 + assert completed.state["group_results"] == list(range(0, 20, 2)) + assert completed.steps[1].join_expected == 10 + assert completed.steps[1].join_arrived == 10 + + runs = await harness.kernel.list_runs() + children = [run for run in runs if run.parent_run_id == started.run_id] + assert len(children) == 10 + assert all(child.status is RunStatus.COMPLETED for child in children) diff --git a/tests/units/workflow/test_examples_temporal_restate.py b/tests/units/workflow/test_examples_temporal_restate.py new file mode 100644 index 00000000000..cc6da98e7ce --- /dev/null +++ b/tests/units/workflow/test_examples_temporal_restate.py @@ -0,0 +1,528 @@ +"""Competitor saga examples expressed with Reflex Workflows. + +Official sources: + +* https://github.com/temporalio/money-transfer-project-template-python +* https://docs.restate.dev/guides/sagas + +The examples intentionally keep provider calls as ordinary Python functions. +``rx.step`` is the durable activity boundary and the provider-facing key from +``rx.current_run`` closes the small crash window before a step result records. +""" + +from __future__ import annotations + +from typing import Any + +from reflex_base.workflow import Retry, WorkflowConfig, manual + +import reflex as rx +from reflex.workflow.records import RunStatus +from reflex.workflow.testing import WorkflowTestHarness + +MONEY_EFFECTS: list[tuple[str, str]] = [] +TRAVEL_EFFECTS: list[tuple[str, str]] = [] + + +class DepositRejected(Exception): + """The target bank rejected a deposit permanently.""" + + +class HotelUnavailable(Exception): + """The hotel cannot satisfy this booking request.""" + + +def _money_effect( + operation: str, + account: str, + amount: int, + reference_id: str, + idempotency_key: str, +) -> dict[str, Any]: + """Simulate an idempotent banking API call. + + Args: + operation: Banking operation being performed. + account: Account affected by the operation. + amount: Amount transferred in cents. + reference_id: Business transfer identifier. + idempotency_key: Provider-facing request identity. + + Returns: + A JSON-compatible provider receipt. + """ + MONEY_EFFECTS.append((f"{operation}:{account}", idempotency_key)) + return { + "operation": operation, + "account": account, + "amount": amount, + "reference_id": reference_id, + "idempotency_key": idempotency_key, + } + + +def _withdraw( + account: str, amount: int, reference_id: str, idempotency_key: str +) -> dict[str, Any]: + """Withdraw money from the source account. + + Args: + account: Source account. + amount: Amount in cents. + reference_id: Business transfer identifier. + idempotency_key: Provider-facing request identity. + + Returns: + The withdrawal receipt. + """ + return _money_effect("withdraw", account, amount, reference_id, idempotency_key) + + +def _deposit( + account: str, + amount: int, + reference_id: str, + idempotency_key: str, + reject: bool, +) -> dict[str, Any]: + """Deposit money, optionally simulating a permanent rejection. + + Args: + account: Target account. + amount: Amount in cents. + reference_id: Business transfer identifier. + idempotency_key: Provider-facing request identity. + reject: Whether the target bank rejects this deposit. + + Returns: + The deposit receipt. + + Raises: + DepositRejected: When ``reject`` is true. + """ + receipt = _money_effect("deposit", account, amount, reference_id, idempotency_key) + if reject: + msg = f"target account {account} rejected transfer {reference_id}" + raise DepositRejected(msg) + return receipt + + +def _refund( + account: str, amount: int, reference_id: str, idempotency_key: str +) -> dict[str, Any]: + """Compensate a successful withdrawal. + + Args: + account: Source account receiving its money back. + amount: Amount in cents. + reference_id: Business transfer identifier. + idempotency_key: Provider-facing request identity. + + Returns: + The refund receipt. + """ + return _money_effect("refund", account, amount, reference_id, idempotency_key) + + +class TemporalMoneyTransfer(rx.State): + """Temporal's withdraw/deposit/refund tutorial as a Reflex workflow.""" + + __workflow__ = WorkflowConfig(id="examples.temporal.money_transfer") + + reference_id: str = "" + phase: str = "new" + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=3, initial_delay="1s", jitter="none"), + ) + async def transfer( + self, + source_account: str, + target_account: str, + amount: int, + reference_id: str, + reject_deposit: bool = False, + ): + """Withdraw, deposit, and refund the withdrawal if deposit fails. + + Args: + source_account: Account money leaves. + target_account: Account money enters. + amount: Amount in cents. + reference_id: Stable business transfer identifier. + reject_deposit: Simulate a terminal target-bank rejection. + + Returns: + Completion on success, or a compensated business failure. + """ + context = rx.current_run() + if context is None: + msg = "money transfer must run inside a durable attempt" + raise RuntimeError(msg) + + self.reference_id = reference_id + self.phase = "withdrawing" + withdrawal = await rx.step( + "withdraw", + _withdraw, + source_account, + amount, + reference_id, + context.idempotency_key(scope="withdraw"), + ) + + try: + self.phase = "depositing" + deposit = await rx.step( + "deposit", + _deposit, + target_account, + amount, + reference_id, + context.idempotency_key(scope="deposit"), + reject_deposit, + ) + except DepositRejected as error: + self.phase = "refunding" + refund = await rx.step( + "refund", + _refund, + source_account, + amount, + reference_id, + context.idempotency_key(scope="refund"), + ) + self.phase = "refunded" + return rx.fail( + "deposit_rejected", + details={"message": str(error), "refund": refund}, + ) + + self.phase = "completed" + return rx.complete(result={"withdrawal": withdrawal, "deposit": deposit}) + + +def _travel_effect( + operation: str, customer_id: str, idempotency_key: str +) -> dict[str, str]: + """Simulate an idempotent travel-provider operation. + + Args: + operation: Reservation or compensation being performed. + customer_id: Customer owning the itinerary. + idempotency_key: Provider-facing request identity. + + Returns: + A JSON-compatible provider receipt. + """ + TRAVEL_EFFECTS.append((operation, idempotency_key)) + return { + "operation": operation, + "customer_id": customer_id, + "reservation_id": f"{operation}-{customer_id}", + "idempotency_key": idempotency_key, + } + + +def _book_flight(customer_id: str, idempotency_key: str) -> dict[str, str]: + """Reserve the flight leg. + + Returns: + The flight receipt. + """ + return _travel_effect("book_flight", customer_id, idempotency_key) + + +def _book_car(customer_id: str, idempotency_key: str) -> dict[str, str]: + """Reserve the rental car. + + Returns: + The car receipt. + """ + return _travel_effect("book_car", customer_id, idempotency_key) + + +def _book_hotel( + customer_id: str, idempotency_key: str, unavailable: bool +) -> dict[str, str]: + """Reserve the hotel, optionally simulating a terminal failure. + + Args: + customer_id: Customer owning the itinerary. + idempotency_key: Provider-facing request identity. + unavailable: Whether the hotel is fully booked. + + Returns: + The hotel receipt. + + Raises: + HotelUnavailable: When ``unavailable`` is true. + """ + receipt = _travel_effect("book_hotel", customer_id, idempotency_key) + if unavailable: + msg = f"hotel is unavailable for customer {customer_id}" + raise HotelUnavailable(msg) + return receipt + + +def _cancel_hotel(customer_id: str, idempotency_key: str) -> dict[str, str]: + """Compensate the hotel reservation. + + Returns: + The cancellation receipt. + """ + return _travel_effect("cancel_hotel", customer_id, idempotency_key) + + +def _cancel_car(customer_id: str, idempotency_key: str) -> dict[str, str]: + """Compensate the car reservation. + + Returns: + The cancellation receipt. + """ + return _travel_effect("cancel_car", customer_id, idempotency_key) + + +def _cancel_flight(customer_id: str, idempotency_key: str) -> dict[str, str]: + """Compensate the flight reservation. + + Returns: + The cancellation receipt. + """ + return _travel_effect("cancel_flight", customer_id, idempotency_key) + + +class RestateTravelBooking(rx.State): + """Restate's flight/car/hotel saga as a Reflex workflow.""" + + __workflow__ = WorkflowConfig(id="examples.restate.travel_booking") + + customer_id: str = "" + phase: str = "new" + + @rx.event( + durable=True, + trigger=manual(), + effect="idempotent_write", + retry=Retry(max_attempts=3, initial_delay="1s", jitter="none"), + ) + async def book(self, customer_id: str, hotel_unavailable: bool = False): + """Book flight, car, and hotel, compensating in reverse on failure. + + Args: + customer_id: Customer owning the itinerary. + hotel_unavailable: Simulate a terminal hotel failure. + + Returns: + Completion on success, or a compensated business failure. + """ + context = rx.current_run() + if context is None: + msg = "travel booking must run inside a durable attempt" + raise RuntimeError(msg) + + self.customer_id = customer_id + self.phase = "booking_flight" + flight = await rx.step( + "book-flight", + _book_flight, + customer_id, + context.idempotency_key(scope="book-flight"), + ) + self.phase = "booking_car" + car = await rx.step( + "book-car", + _book_car, + customer_id, + context.idempotency_key(scope="book-car"), + ) + + try: + self.phase = "booking_hotel" + hotel = await rx.step( + "book-hotel", + _book_hotel, + customer_id, + context.idempotency_key(scope="book-hotel"), + hotel_unavailable, + ) + except HotelUnavailable as error: + # Restate registers each compensation before its booking attempt. + # Even an ambiguous hotel failure is therefore cancelled first, + # followed by the earlier successful reservations in reverse. + self.phase = "cancelling_hotel" + cancelled_hotel = await rx.step( + "cancel-hotel", + _cancel_hotel, + customer_id, + context.idempotency_key(scope="cancel-hotel"), + ) + self.phase = "cancelling_car" + cancelled_car = await rx.step( + "cancel-car", + _cancel_car, + customer_id, + context.idempotency_key(scope="cancel-car"), + ) + self.phase = "cancelling_flight" + cancelled_flight = await rx.step( + "cancel-flight", + _cancel_flight, + customer_id, + context.idempotency_key(scope="cancel-flight"), + ) + self.phase = "compensated" + return rx.fail( + "hotel_unavailable", + details={ + "message": str(error), + "compensated": [ + cancelled_hotel["operation"], + cancelled_car["operation"], + cancelled_flight["operation"], + ], + }, + ) + + self.phase = "completed" + return rx.complete(result={"flight": flight, "car": car, "hotel": hotel}) + + +def _operations(effects: list[tuple[str, str]]) -> list[str]: + """Return the provider operation names in execution order. + + Args: + effects: Recorded provider calls. + + Returns: + Operation names without their idempotency keys. + """ + return [operation for operation, _ in effects] + + +def _assert_scoped_keys(effects: list[tuple[str, str]]) -> None: + """Assert every provider call received a distinct stable key scope. + + Args: + effects: Recorded provider calls. + """ + keys = [key for _, key in effects] + assert all(len(key) == 32 for key in keys) + assert len(keys) == len(set(keys)) + + +async def test_temporal_money_transfer_happy_path( + forked_registration_context, +): + """A successful withdrawal and deposit complete the transfer.""" + MONEY_EFFECTS.clear() + async with WorkflowTestHarness(TemporalMoneyTransfer) as harness: + started = await harness.start( + TemporalMoneyTransfer.transfer( # pyright: ignore[reportCallIssue] + "alice", "bob", 5000, "tx-1" + ) + ) + assert started.run_id is not None + snapshot = await harness.get_run(started.run_id) + + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state["phase"] == "completed" + assert snapshot.result["withdrawal"]["account"] == "alice" + assert snapshot.result["deposit"]["account"] == "bob" + assert _operations(MONEY_EFFECTS) == ["withdraw:alice", "deposit:bob"] + _assert_scoped_keys(MONEY_EFFECTS) + + +async def test_temporal_money_transfer_refunds_after_deposit_failure( + forked_registration_context, +): + """A rejected deposit durably refunds the successful withdrawal.""" + MONEY_EFFECTS.clear() + async with WorkflowTestHarness(TemporalMoneyTransfer) as harness: + started = await harness.start( + TemporalMoneyTransfer.transfer( # pyright: ignore[reportCallIssue] + "alice", "bob", 5000, "tx-2", reject_deposit=True + ) + ) + assert started.run_id is not None + snapshot = await harness.get_run(started.run_id) + + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.state["phase"] == "refunded" + assert snapshot.error is not None + assert snapshot.error["reason"] == "deposit_rejected" + assert snapshot.error["details"]["refund"]["account"] == "alice" + assert _operations(MONEY_EFFECTS) == [ + "withdraw:alice", + "deposit:bob", + "refund:alice", + ] + _assert_scoped_keys(MONEY_EFFECTS) + + +async def test_restate_travel_booking_happy_path( + forked_registration_context, +): + """Flight, car, and hotel all remain booked on success.""" + TRAVEL_EFFECTS.clear() + async with WorkflowTestHarness(RestateTravelBooking) as harness: + started = await harness.start( + RestateTravelBooking.book( # pyright: ignore[reportCallIssue] + "customer-1" + ) + ) + assert started.run_id is not None + snapshot = await harness.get_run(started.run_id) + + assert snapshot is not None + assert snapshot.status is RunStatus.COMPLETED + assert snapshot.state["phase"] == "completed" + assert snapshot.result["flight"]["operation"] == "book_flight" + assert snapshot.result["car"]["operation"] == "book_car" + assert snapshot.result["hotel"]["operation"] == "book_hotel" + assert _operations(TRAVEL_EFFECTS) == [ + "book_flight", + "book_car", + "book_hotel", + ] + _assert_scoped_keys(TRAVEL_EFFECTS) + + +async def test_restate_travel_booking_compensates_in_reverse_order( + forked_registration_context, +): + """A terminal hotel failure cancels hotel, car, and then flight.""" + TRAVEL_EFFECTS.clear() + async with WorkflowTestHarness(RestateTravelBooking) as harness: + started = await harness.start( + RestateTravelBooking.book( # pyright: ignore[reportCallIssue] + "customer-2", hotel_unavailable=True + ) + ) + assert started.run_id is not None + snapshot = await harness.get_run(started.run_id) + + assert snapshot is not None + assert snapshot.status is RunStatus.FAILED + assert snapshot.state["phase"] == "compensated" + assert snapshot.error is not None + assert snapshot.error["reason"] == "hotel_unavailable" + assert snapshot.error["details"]["compensated"] == [ + "cancel_hotel", + "cancel_car", + "cancel_flight", + ] + assert _operations(TRAVEL_EFFECTS) == [ + "book_flight", + "book_car", + "book_hotel", + "cancel_hotel", + "cancel_car", + "cancel_flight", + ] + _assert_scoped_keys(TRAVEL_EFFECTS) From d464bb90f2f357e22953737f7b4fb95ed301e426 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Mon, 24 Aug 2026 18:45:10 -0700 Subject: [PATCH 120/121] workflows: release-pinned execution and the worker fleet The engine half of workflow-only deploys, per the roadmap's Phase 3 release-safety spec and the standing design decision: runs pin to the release that admitted them, and a deploy is a set of reads, not a ceremony. Two identities per run now, doing different jobs. The definition digest (structural, existing) decides whether code CAN run a payload -- mismatches suspend at dispatch. The new release id -- REFLEX_RELEASE_ID or WorkflowRuntime(release=...) -- decides whether code MAY: every run and every fan-out child stamps the release that admitted it, and claim_next on all three stores skips runs pinned elsewhere, so a run drains on the code that recorded its payloads and never silently mixes two releases. Pinning binds only when both sides declare: unpinned runs are anyone's, and a worker with no release (dev, tests) serves everything -- which is also why the entire existing suite runs unchanged. Workers register a durable identity -- id, release, queues, capacity -- at startup (after the first recovery, so the timestamp is store-synced), heartbeat on the lease-renewal cadence rather than a second timer, and deregister on clean shutdown; a crashed worker stays listed with a stale heartbeat, which is exactly what a fleet page should show. reflex workflows fleet renders the registry and per-release active-run counts, and --can-retire RELEASE is the deploy gate: nonzero while any active run is pinned to that release, because stopping its workers early strands those runs until leases lapse. RunQuery gains a release_id filter (all three stores + the memory matcher), RunSnapshot and the HTTP read surfaces report each run's release, and SQLite reaches schema v5 (release_id column + the workflow_workers table; Postgres adds both additively in its advisory-locked DDL). Pinned by three conformance checks on every store -- routing, registry round-trip, retirement counts -- and a rolling N/N-1 kernel test: a run sleeps for a day on v1, v2 deploys and takes new admissions, v2 never claims v1's sleeping run, v1 drains it, and the retirement count reaches zero. One test lesson recorded in the diff: worker heartbeats are store-clock timestamps, and two clock syncs can differ by milliseconds either way, so freshness assertions carry the sync's own jitter tolerance. --- news/workflow-release-routing.md | 1 + reflex/workflow/CONTRACT.md | 25 +- reflex/workflow/api.py | 1 + reflex/workflow/cli.py | 71 ++++++ reflex/workflow/conformance.py | 87 +++++++ reflex/workflow/kernel.py | 41 +++- reflex/workflow/postgres.py | 106 ++++++++- reflex/workflow/records.py | 32 +++ reflex/workflow/runtime.py | 5 + reflex/workflow/serve.py | 1 + reflex/workflow/store.py | 235 ++++++++++++++++++- tests/units/workflow/test_release_routing.py | 182 ++++++++++++++ 12 files changed, 779 insertions(+), 8 deletions(-) create mode 100644 news/workflow-release-routing.md create mode 100644 tests/units/workflow/test_release_routing.py diff --git a/news/workflow-release-routing.md b/news/workflow-release-routing.md new file mode 100644 index 00000000000..cdf8701ccb4 --- /dev/null +++ b/news/workflow-release-routing.md @@ -0,0 +1 @@ +Release-pinned execution, the engine half of workflow-only deploys. Runs (children included) stamp the release that admitted them — `REFLEX_RELEASE_ID` or `WorkflowRuntime(release=...)` — and a worker of a different release never claims them: the run drains on the code that recorded its payloads, so one run never silently mixes two releases. Pinning binds only when both sides declare, so dev and test workers keep serving everything. Workers register their release, queues, and capacity at startup, heartbeat on the lease-renewal cadence, and deregister on clean shutdown; `reflex workflows fleet` shows the registry with heartbeat ages, and `--can-retire RELEASE` is the deploy gate — nonzero while any active run is still pinned to that release. `RunQuery` gains a `release_id` filter, read surfaces report each run's release, and conformance pins the routing, the registry, and the retirement count on all three stores, with a rolling N/N−1 test proving a sleeping run resumes on the release that admitted it while the new release takes new work. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index d13d4c5836e..5eab60b489f 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -196,7 +196,30 @@ asking the engine to keep old code alive. That routing is not built yet; when it is, it constrains which workers exist, and every rule above still holds underneath it. -## 5. Cancellation, deadlines, and children +### Release-pinned execution + +Two identities per run, doing different jobs. The **definition digest** +(structural) decides whether code *can* run a payload — mismatches suspend +at dispatch (§8). The **release id** (identity of the deployed artifact, +from `REFLEX_RELEASE_ID` or `WorkflowRuntime(release=...)`) decides whether +code *may*: a run pins to the release that admitted it, children included, +and a worker of a different release never claims it — the run drains on the +code that recorded its payloads, so one run never silently mixes two +releases. A run or worker with no declared release is unconstrained (dev, +tests, pre-release deployments): pinning binds only when both sides declare. + +Workers register their identity — release, queues, capacity — at startup, +heartbeat on the lease-renewal cadence, and deregister on clean shutdown; +a crashed worker stays listed with a stale heartbeat, which is what a +fleet page should show (`reflex workflows fleet`, `RunStore.list_workers`). +Rolling deploys are then reads, not ceremonies: new admissions carry the +new release the moment its workers start; the old release's workers drain +what they own; rollback is starting old-release workers again; and the +retirement gate is a count — `reflex workflows fleet --can-retire R` exits +nonzero while any active run is pinned to R, because stopping R's workers +early strands those runs until their leases lapse. + +## 5. Cancellation, deadlines, and children## 5. Cancellation, deadlines, and children - `cancel(run_id)` records intent and cancels any in-flight attempt cooperatively. The run finalizes `CANCELLED` only once no step is claimed diff --git a/reflex/workflow/api.py b/reflex/workflow/api.py index aca99412b11..681457e1da4 100644 --- a/reflex/workflow/api.py +++ b/reflex/workflow/api.py @@ -240,6 +240,7 @@ async def endpoint(request: Request) -> JSONResponse: "run_id": snapshot.run_id, "workflow": snapshot.workflow_id, "status": snapshot.status.value, + "release": snapshot.release_id, "result": snapshot.result, "error": snapshot.error, "steps": [ diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 6e582044f7d..35d6c792004 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -1414,6 +1414,77 @@ def serve( _run_server(app, host, port) +@workflows.command() +@database_option +@click.option( + "--can-retire", + "retire_release", + default=None, + help=( + "Exit 0 when no active run is pinned to this release, 1 otherwise; " + "the gate a deploy runs before stopping that release's workers." + ), +) +def fleet(database: str | None, retire_release: str | None): + """Show registered workers and the runs each release still owns. + + A worker that stopped cleanly disappears; one that crashed stays listed + with a stale heartbeat, which is exactly what this page should show. + """ + import time + + from reflex.workflow.records import TERMINAL_RUN_STATUSES, RunQuery, RunStatus + + active = tuple(s for s in RunStatus if s not in TERMINAL_RUN_STATUSES) + + async def read(store: RunStore): + """Collect workers and per-release active counts. + + Args: + store: The open run store. + + Returns: + The workers, the release counts, and the gate answer. + """ + workers = await store.list_workers() + releases = sorted({w.release_id for w in workers if w.release_id is not None}) + if retire_release is not None and retire_release not in releases: + releases.append(retire_release) + counts = { + release: await store.count_runs( + RunQuery(release_id=release, statuses=active) + ) + for release in releases + } + return workers, counts + + workers, counts = _with_store(database, read) + if retire_release is not None: + held = counts.get(retire_release, 0) + if held: + console.error( + f"Release {retire_release!r} still owns {held} active " + f"run{'s' if held != 1 else ''}; retiring its workers now " + "would strand them until their leases lapse." + ) + raise click.exceptions.Exit(1) + console.print(f"Release {retire_release!r} owns no active runs.") + return + if not workers: + console.print("No workers registered.") + now = time.time() + for worker in workers: + age = now - worker.heartbeat_at + queues = ", ".join(worker.queues) or "all queues" + console.print( + f"{worker.worker_id[:12]} release={worker.release_id or '-'} " + f"{queues} capacity={worker.capacity} " + f"heartbeat {age:.0f}s ago" + ) + for release, held in counts.items(): + console.print(f"release {release}: {held} active run(s)") + + @workflows.command() @database_option @click.option( diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 526d2eb8316..77834617e07 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -31,6 +31,7 @@ async def test_my_store_conforms(check): RunStatus, StepRecord, StepStatus, + WorkerRecord, ) from reflex.workflow.store import StaleClaimError, StepCompletion @@ -1843,6 +1844,89 @@ async def check_policy_admission_also_flushes_parked_mail(store: RunStore) -> No assert rows[0].status is ParkedStatus.DELIVERED +async def check_a_pinned_run_drains_only_on_its_release(store: RunStore) -> None: + """Release routing: a run never mixes two releases' code. + + A v2 worker must not claim a run v1 admitted -- the run drains on the + release whose code recorded its payloads. An unpinned run (admitted + before releases were declared) is claimable by anyone, and a worker with + no declared release (dev, tests) claims anything: pinning constrains + only when both sides declare. + """ + await store.admit(make_run(release_id="v1"), make_step(due_at=0.0), _ADMITTED) + assert await store.claim_next(NOW, release="v2") is None, ( + "a v2 worker claimed a v1-pinned run" + ) + claim = await store.claim_next(NOW, release="v1") + assert claim is not None + assert claim.run.release_id == "v1" + await store.release_claim(claim, status=StepStatus.READY, events=(), now=NOW) + + claim = await store.claim_next(NOW, release=None) + assert claim is not None, "an undeclared worker serves every release" + await store.release_claim(claim, status=StepStatus.READY, events=(), now=NOW) + + await store.admit(make_run("free1"), make_step("free1", due_at=0.0), _ADMITTED) + claim = await store.claim_next(NOW, release="v2") + assert claim is not None + assert claim.run.run_id == "free1", "an unpinned run is anyone's to run" + + +async def check_worker_registry_roundtrip(store: RunStore) -> None: + """The fleet surface: register, heartbeat, list, deregister.""" + await store.register_worker( + WorkerRecord( + worker_id="w1", + release_id="v1", + queues=("default",), + capacity=8, + started_at=NOW, + heartbeat_at=NOW, + ) + ) + await store.register_worker( + WorkerRecord( + worker_id="w2", + release_id=None, + queues=(), + capacity=4, + started_at=NOW + 1, + heartbeat_at=NOW + 1, + ) + ) + workers = await store.list_workers() + assert [worker.worker_id for worker in workers] == ["w2", "w1"] + assert workers[1].queues == ("default",) + + await store.heartbeat_worker("w1", NOW + 60) + workers = await store.list_workers() + beat = next(worker for worker in workers if worker.worker_id == "w1") + assert beat.heartbeat_at == pytest.approx(NOW + 60) + + await store.deregister_worker("w2") + workers = await store.list_workers() + assert [worker.worker_id for worker in workers] == ["w1"] + + +async def check_release_counts_answer_the_retirement_question( + store: RunStore, +) -> None: + """Count what still runs the release being replaced; zero means retire.""" + await store.admit(make_run(release_id="v1"), make_step(due_at=0.0), _ADMITTED) + await store.admit( + make_run("done1", release_id="v1", status=RunStatus.COMPLETED), + make_step("done1", status=StepStatus.SUCCEEDED), + _ADMITTED, + ) + active = [status for status in RunStatus if status not in TERMINAL_RUN_STATUSES] + assert ( + await store.count_runs(RunQuery(release_id="v1", statuses=tuple(active))) == 1 + ) + assert ( + await store.count_runs(RunQuery(release_id="v2", statuses=tuple(active))) == 0 + ), "nothing pins to a release that admitted nothing" + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1865,6 +1949,9 @@ async def check_policy_admission_also_flushes_parked_mail(store: RunStore) -> No check_a_dead_letter_is_visible_and_replayable, check_unclaimed_deliveries_become_dead_letters, check_policy_admission_also_flushes_parked_mail, + check_a_pinned_run_drains_only_on_its_release, + check_worker_registry_roundtrip, + check_release_counts_answer_the_retirement_question, check_a_delivery_to_a_past_deadline_run_is_refused, check_a_duplicate_delivery_is_recorded_in_history, check_an_early_delivery_is_buffered_then_consumed, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index f0fd11fd72c..c9830e2ffa3 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -14,6 +14,7 @@ import contextlib import dataclasses import operator +import os import random import time import traceback @@ -55,6 +56,7 @@ StartResult, StepRecord, StepStatus, + WorkerRecord, ) from reflex.workflow.serde import to_run_data from reflex.workflow.steps import SubstepJournal, bind_journal, unbind_journal @@ -462,6 +464,7 @@ def __init__( observer: WorkflowObserver | None = None, max_concurrency: int = DEFAULT_MAX_CONCURRENCY, queues: Iterable[str] | None = None, + release: str | None = None, ): """Initialize the kernel. @@ -483,6 +486,10 @@ def __init__( max_concurrency: How many attempts this kernel runs at once. Each belongs to a different run, so a run's own steps stay serial. queues: Queues this kernel's worker serves; None serves them all. + release: The deployed artifact identity this worker runs, read + from ``REFLEX_RELEASE_ID`` when omitted. Runs admitted here + pin to it, and this worker never claims a run pinned to a + different release. Steps land on the queue their handler declared, "default" otherwise, so a deployment can dedicate processes to slow or sensitive work. @@ -563,6 +570,9 @@ def __init__( self._worker_id = uuid.uuid4().hex self._observer = observer self._queues = tuple(queues) if queues is not None else None + self._release = ( + release if release is not None else os.environ.get("REFLEX_RELEASE_ID") + ) or None self._wakeup = asyncio.Event() self._closing = False self._worker: asyncio.Task | None = None @@ -873,6 +883,7 @@ def _admission_records( flow_key=flow_key, request_key=request_key, labels=labels, + release_id=self._release, deadline=(now + defn.run_timeout) if defn.run_timeout is not None else None, created_at=now, updated_at=now, @@ -1334,6 +1345,7 @@ async def get_run(self, run_id: str) -> RunSnapshot | None: result=run.result, error=run.error, steps=steps, + release_id=run.release_id, ) def _adapter(self, defn: WorkflowDefinition, field_name: str) -> TypeAdapter: @@ -2250,6 +2262,11 @@ async def _renew_leases(self) -> None: """ for lease in list(self._leases.values()): await self._renew(lease) + if self._worker is not None: + # The same cadence that proves claims alive proves the worker + # alive; a heartbeat needing its own timer would drift from the + # one signal operators actually watch. + await self._store.heartbeat_worker(self._worker_id, self._clock()) async def _cancel_requested(self, run_id: str) -> bool: """Whether a run carries cancellation intent. @@ -2698,6 +2715,7 @@ def _child_records( status=RunStatus.PENDING, state={field.name: field.default for field in defn.fields}, state_version=0, + release_id=self._release, next_ordinal=1, parent_run_id=claim.run.run_id, parent_ordinal=join_ordinal, @@ -3050,7 +3068,10 @@ async def _fill_slots(self, now: float) -> list[asyncio.Task]: started: list[asyncio.Task] = [] while len(self._inflight) < self._max_concurrency: claim = await self._store.claim_next( - now, lease_duration=self._lease_duration, queues=self._queues + now, + lease_duration=self._lease_duration, + queues=self._queues, + release=self._release, ) if claim is None: break @@ -3275,6 +3296,20 @@ async def start_worker(self) -> None: return self._closing = False await self.recover() + now = self._clock() + # Registered after the first recovery so the clock is store-synced; + # the fleet surface is how a deploy gate can see who runs which + # release, at what capacity, and how recently they proved alive. + await self._store.register_worker( + WorkerRecord( + worker_id=self._worker_id, + release_id=self._release, + queues=self._queues or (), + capacity=self._max_concurrency, + started_at=now, + heartbeat_at=now, + ) + ) self._worker = asyncio.create_task(self._worker_loop()) async def aclose(self, drain: float = 0.0) -> None: @@ -3304,6 +3339,10 @@ async def aclose(self, drain: float = 0.0) -> None: with contextlib.suppress(asyncio.CancelledError): await self._worker self._worker = None + # A clean shutdown removes the registration; a crash leaves it with a + # stale heartbeat, which is exactly what a fleet page should show. + with contextlib.suppress(Exception): + await self._store.deregister_worker(self._worker_id) if self._inflight and self._draining: await asyncio.wait(set(self._inflight.values()), timeout=drain) self._draining = False diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index 85f6fbf8b2d..dfb9a33ffb8 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -34,6 +34,7 @@ RunStatus, StepRecord, StepStatus, + WorkerRecord, ) from reflex.workflow.store import ( Claim, @@ -89,6 +90,7 @@ labels JSONB, deadline DOUBLE PRECISION, cancel_requested BOOLEAN NOT NULL DEFAULT FALSE, + release_id TEXT, created_at DOUBLE PRECISION NOT NULL, updated_at DOUBLE PRECISION NOT NULL ); @@ -139,6 +141,15 @@ created_at DOUBLE PRECISION NOT NULL, PRIMARY KEY (run_id, ordinal, key) ); +CREATE TABLE IF NOT EXISTS workflow_workers ( + worker_id TEXT PRIMARY KEY, + release_id TEXT, + queues JSONB NOT NULL, + capacity INTEGER NOT NULL, + started_at DOUBLE PRECISION NOT NULL, + heartbeat_at DOUBLE PRECISION NOT NULL +); +ALTER TABLE workflow_runs ADD COLUMN IF NOT EXISTS release_id TEXT; CREATE TABLE IF NOT EXISTS workflow_channel_inbox ( parked_id TEXT PRIMARY KEY, workflow_id TEXT NOT NULL, @@ -203,6 +214,8 @@ "NOT (r.status = ANY(%(terminal_runs)s)) AND r.status <> 'NEEDS_ATTENTION'" " AND NOT r.cancel_requested" " AND (r.deadline IS NULL OR r.deadline > %(now)s)" + " AND (r.release_id IS NULL OR %(release)s::text IS NULL" + " OR r.release_id = %(release)s)" ) @@ -282,6 +295,7 @@ def _run_from_row(row: Mapping[str, Any]) -> RunRecord: parent_close=row["parent_close"] or "cancel", request_key=row["request_key"], labels=row["labels"], + release_id=row["release_id"], deadline=row["deadline"], cancel_requested=row["cancel_requested"], created_at=row["created_at"], @@ -340,6 +354,9 @@ def _run_filters(query: RunQuery) -> tuple[str, tuple[Any, ...]]: if query.definition_digest is not None: clauses.append("definition_digest = %s") params.append(query.definition_digest) + if query.release_id is not None: + clauses.append("release_id = %s") + params.append(query.release_id) if query.statuses: clauses.append("status = ANY(%s)") params.append([status.value for status in query.statuses]) @@ -568,9 +585,9 @@ async def _insert_run(self, conn: Connection, run: RunRecord) -> None: " status, state, state_version, next_ordinal, result, error," " flow_key, parent_run_id, parent_ordinal, parent_close," " request_key, labels," - " deadline, cancel_requested, created_at, updated_at)" + " deadline, cancel_requested, release_id, created_at, updated_at)" " VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s," - " %s, %s, %s, %s, %s)", + " %s, %s, %s, %s, %s, %s)", ( run.run_id, run.workflow_id, @@ -589,6 +606,7 @@ async def _insert_run(self, conn: Connection, run: RunRecord) -> None: _json(run.labels), run.deadline, run.cancel_requested, + run.release_id, run.created_at, run.updated_at, ), @@ -966,6 +984,7 @@ async def claim_next( *, lease_duration: float = DEFAULT_LEASE_DURATION, queues: tuple[str, ...] | None = None, + release: str | None = None, ) -> Claim | None: """Claim the due frontier step of some runnable run. @@ -978,6 +997,9 @@ async def claim_next( lease_duration: Seconds of renewal silence tolerated before the claim is treated as orphaned. queues: Queues this worker serves; None serves every queue. + release: The claiming worker's release identity. A run pinned to + a different release is skipped: it drains on the release that + admitted it, so one run never mixes two releases' code. Returns: A fenced claim, or None when nothing is claimable right now. @@ -989,6 +1011,7 @@ async def claim_next( "terminal_steps": _TERMINAL_STEPS, "claimable": _CLAIMABLE_STEPS, "queues": list(queues) if queues is not None else None, + "release": release, } async with pool.connection() as conn, conn.transaction(): cursor = await conn.execute( @@ -1503,6 +1526,82 @@ async def _route_parked_conn( ) return disposition if delivered else "dead_letter" + async def register_worker(self, worker: WorkerRecord) -> None: + """Record (or refresh) a worker's registration. + + Args: + worker: The worker's identity, release, queues, and capacity. + """ + pool = await self._open() + async with pool.connection() as conn: + await conn.execute( + "INSERT INTO workflow_workers (worker_id, release_id, queues," + " capacity, started_at, heartbeat_at)" + " VALUES (%s, %s, %s, %s, %s, %s)" + " ON CONFLICT (worker_id) DO UPDATE SET release_id =" + " excluded.release_id, queues = excluded.queues," + " capacity = excluded.capacity," + " heartbeat_at = excluded.heartbeat_at", + ( + worker.worker_id, + worker.release_id, + Jsonb(list(worker.queues)), + worker.capacity, + worker.started_at, + worker.heartbeat_at, + ), + ) + + async def heartbeat_worker(self, worker_id: str, now: float) -> None: + """Refresh a worker's sign of life. + + Args: + worker_id: The worker. + now: Current time in epoch seconds. + """ + pool = await self._open() + async with pool.connection() as conn: + await conn.execute( + "UPDATE workflow_workers SET heartbeat_at = %s WHERE worker_id = %s", + (now, worker_id), + ) + + async def deregister_worker(self, worker_id: str) -> None: + """Remove a worker that shut down cleanly. + + Args: + worker_id: The worker. + """ + pool = await self._open() + async with pool.connection() as conn: + await conn.execute( + "DELETE FROM workflow_workers WHERE worker_id = %s", + (worker_id,), + ) + + async def list_workers(self) -> tuple[WorkerRecord, ...]: + """List registered workers, most recently started first. + + Returns: + The registrations. + """ + pool = await self._open() + async with pool.connection() as conn: + cursor = await conn.execute( + "SELECT * FROM workflow_workers ORDER BY started_at DESC" + ) + return tuple( + WorkerRecord( + worker_id=row["worker_id"], + release_id=row["release_id"], + queues=tuple(row["queues"]), + capacity=row["capacity"], + started_at=row["started_at"], + heartbeat_at=row["heartbeat_at"], + ) + for row in await cursor.fetchall() + ) + async def ingest_channel_delivery( self, workflow_id: str, @@ -2846,6 +2945,9 @@ async def next_due( "terminal_steps": _TERMINAL_STEPS, "claimable": _CLAIMABLE_STEPS, "queues": list(queues) if queues is not None else None, + # next_due bounds the sleep for every runnable run; release + # pinning shapes who claims, not when the fleet wakes. + "release": None, } async with pool.connection() as conn: cursor = await conn.execute( diff --git a/reflex/workflow/records.py b/reflex/workflow/records.py index 4d288ed5f5f..03d2de87b57 100644 --- a/reflex/workflow/records.py +++ b/reflex/workflow/records.py @@ -187,6 +187,10 @@ class RunRecord: terminal state: ``"cancel"`` or ``"abandon"``. request_key: Idempotent admission key, if one was supplied. labels: Server-derived indexing labels. + release_id: Immutable identity of the deployed artifact that admitted + this run. Runs pin to their admitting release: a worker of a + different release does not claim them, so one run never silently + mixes two releases' code. deadline: Absolute run deadline in epoch seconds, if configured. cancel_requested: Whether cancellation intent has been recorded. created_at: Admission time in epoch seconds. @@ -208,6 +212,7 @@ class RunRecord: parent_close: str = "cancel" request_key: str | None = None labels: dict[str, str] | None = None + release_id: str | None = None deadline: float | None = None cancel_requested: bool = False created_at: float = 0.0 @@ -312,6 +317,27 @@ class StartResult: retry_after: float | None = None +@dataclasses.dataclass(frozen=True, slots=True) +class WorkerRecord: + """One live worker's registration, as the fleet surface reads it. + + Attributes: + worker_id: The worker's unique identity. + release_id: The deployed artifact this worker runs, if declared. + queues: The queues this worker serves; empty means every queue. + capacity: Concurrent attempts this worker runs at most. + started_at: When the worker registered, in epoch seconds. + heartbeat_at: The worker's last sign of life, in epoch seconds. + """ + + worker_id: str + release_id: str | None + queues: tuple[str, ...] + capacity: int + started_at: float + heartbeat_at: float + + class ParkedStatus(str, enum.Enum): """Lifecycle of a correlated webhook delivery in the channel inbox.""" @@ -367,6 +393,9 @@ class RunQuery: definition. This is what answers "is anything still running the release I am replacing", which a deploy gate and an operator watching a rollout both need. + release_id: Restrict to runs pinned to one release — with + non-terminal ``statuses``, the "can this release's workers + retire" question. statuses: Restrict to these run statuses; empty means any. labels: Require every one of these server-derived label values. created_before: Pagination cursor, as the ``(created_at, run_id)`` of @@ -378,6 +407,7 @@ class RunQuery: workflow_id: str | None = None definition_digest: str | None = None + release_id: str | None = None statuses: tuple[RunStatus, ...] = () labels: Mapping[str, str] | None = None created_before: tuple[float, str] | None = None @@ -396,6 +426,7 @@ class RunSnapshot: state_version: Committed state version. result: Run result, if completed with one. error: Terminal or suspension error payload. + release_id: The release that admitted the run and drains it, if any. steps: All mailbox slots in ordinal order. """ @@ -407,3 +438,4 @@ class RunSnapshot: result: Any error: dict[str, Any] | None steps: tuple[StepRecord, ...] + release_id: str | None = None diff --git a/reflex/workflow/runtime.py b/reflex/workflow/runtime.py index ae350cf89b7..07d83c1d9f4 100644 --- a/reflex/workflow/runtime.py +++ b/reflex/workflow/runtime.py @@ -74,6 +74,7 @@ def __init__( max_recoveries: int = DEFAULT_MAX_RECOVERIES, max_concurrency: int = DEFAULT_MAX_CONCURRENCY, queues: Iterable[str] | None = None, + release: str | None = None, ): """Initialize the runtime. @@ -91,6 +92,8 @@ def __init__( max_recoveries: Infrastructure recovery budget per logical step. max_concurrency: How many attempts run at once. queues: Queues this process's worker serves; None serves all. + release: The deployed artifact identity this runtime's worker + runs, read from REFLEX_RELEASE_ID by the kernel when omitted. """ self._store = store self._clock = clock @@ -111,6 +114,7 @@ def __init__( self._max_recoveries = max_recoveries self._max_concurrency = max_concurrency self._queues = tuple(queues) if queues is not None else None + self._release = release self._definitions: dict[str, WorkflowDefinition] = {} self._classes: dict[type, str] = {} self._kernel: WorkflowKernel | None = None @@ -222,6 +226,7 @@ async def startup(self, *, start_worker: bool = True) -> None: max_recoveries=self._max_recoveries, max_concurrency=self._max_concurrency, queues=self._queues, + release=self._release, ) if start_worker: await self._kernel.start_worker() diff --git a/reflex/workflow/serve.py b/reflex/workflow/serve.py index a58719a4ae4..e9ad151193b 100644 --- a/reflex/workflow/serve.py +++ b/reflex/workflow/serve.py @@ -216,6 +216,7 @@ async def endpoint(request: Request) -> JSONResponse: "workflow": run.workflow_id, "status": run.status.value, "labels": run.labels or {}, + "release": run.release_id, "created_at": run.created_at, "updated_at": run.updated_at, } diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index db57de2eaae..7476753a993 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -43,6 +43,7 @@ RunStatus, StepRecord, StepStatus, + WorkerRecord, step_claimable_at, step_wake_at, ) @@ -242,6 +243,7 @@ async def claim_next( *, lease_duration: float = DEFAULT_LEASE_DURATION, queues: tuple[str, ...] | None = None, + release: str | None = None, ) -> Claim | None: """Claim the due frontier step of some runnable run. @@ -256,6 +258,9 @@ async def claim_next( queues: Queues this worker serves; None serves every queue. A run whose frontier sits on an unserved queue is skipped whole, because claiming a later slot would break its ordering. + release: The claiming worker's release identity. A run pinned + to a different release is skipped: it drains on the release + that admitted it, so one run never mixes two releases' code. Returns: A fenced claim, or None when nothing is claimable right now. @@ -383,6 +388,43 @@ async def admit_children( """ ... + async def register_worker(self, worker: WorkerRecord) -> None: + """Record (or refresh) a worker's registration. + + Args: + worker: The worker's identity, release, queues, and capacity. + """ + ... + + async def heartbeat_worker(self, worker_id: str, now: float) -> None: + """Refresh a worker's sign of life. + + Args: + worker_id: The worker. + now: Current time in epoch seconds. + """ + ... + + async def deregister_worker(self, worker_id: str) -> None: + """Remove a worker that shut down cleanly. + + Args: + worker_id: The worker. + """ + ... + + async def list_workers(self) -> tuple[WorkerRecord, ...]: + """List registered workers, most recently started first. + + A worker that died without deregistering stays listed with a stale + heartbeat; staleness is the reader's judgement, because only the + reader knows what cadence it expects. + + Returns: + The registrations. + """ + ... + async def ingest_channel_delivery( self, workflow_id: str, @@ -1008,6 +1050,8 @@ def _matches_query(run: RunRecord, query: RunQuery) -> bool: and run.definition_digest != query.definition_digest ): return False + if query.release_id is not None and run.release_id != query.release_id: + return False if query.statuses and run.status not in query.statuses: return False if query.created_before is not None and (run.created_at, run.run_id) >= ( @@ -1074,6 +1118,7 @@ def __init__(self): self._substeps: dict[tuple[str, int], dict[str, Any]] = {} self._schedule_cursors: dict[str, float] = {} self._parked: list[ParkedDelivery] = [] + self._workers: dict[str, WorkerRecord] = {} self._history: dict[str, list[HistoryEvent]] = {} self._dedupe: dict[tuple[str, str], str] = {} self._inbox: dict[str, dict[tuple[str, str, str], bool]] = {} @@ -1183,6 +1228,51 @@ def _flush_parked_locked( updated_at=now, ) + async def register_worker(self, worker: WorkerRecord) -> None: + """Record (or refresh) a worker's registration. + + Args: + worker: The worker's identity, release, queues, and capacity. + """ + async with self._lock: + self._workers[worker.worker_id] = worker + + async def heartbeat_worker(self, worker_id: str, now: float) -> None: + """Refresh a worker's sign of life. + + Args: + worker_id: The worker. + now: Current time in epoch seconds. + """ + async with self._lock: + worker = self._workers.get(worker_id) + if worker is not None: + self._workers[worker_id] = dataclasses.replace(worker, heartbeat_at=now) + + async def deregister_worker(self, worker_id: str) -> None: + """Remove a worker that shut down cleanly. + + Args: + worker_id: The worker. + """ + async with self._lock: + self._workers.pop(worker_id, None) + + async def list_workers(self) -> tuple[WorkerRecord, ...]: + """List registered workers, most recently started first. + + Returns: + The registrations. + """ + async with self._lock: + return tuple( + sorted( + self._workers.values(), + key=lambda worker: worker.started_at, + reverse=True, + ) + ) + async def ingest_channel_delivery( self, workflow_id: str, @@ -1344,6 +1434,7 @@ async def claim_next( *, lease_duration: float = DEFAULT_LEASE_DURATION, queues: tuple[str, ...] | None = None, + release: str | None = None, ) -> Claim | None: """Claim the due frontier step of some runnable run. @@ -1354,6 +1445,9 @@ async def claim_next( queues: Queues this worker serves; None serves every queue. A run whose frontier sits on an unserved queue is skipped whole, because claiming a later slot would break its ordering. + release: The claiming worker's release identity. A run pinned to + a different release is skipped: it drains on the release that + admitted it, so one run never mixes two releases' code. Returns: A fenced claim, or None when nothing is claimable right now. @@ -1366,6 +1460,14 @@ async def claim_next( frontier = _frontier(steps) if frontier is None or not step_claimable_at(frontier, now): continue + if ( + release is not None + and run.release_id is not None + and run.release_id != release + ): + # Pinned to another release: it drains on the workers + # that admitted it, never on this one. + continue if queues is not None and frontier.queue not in queues: continue claimed = dataclasses.replace( @@ -2750,7 +2852,7 @@ async def next_due( return min(due_times) if due_times else None -SCHEMA_VERSION: Final = 4 +SCHEMA_VERSION: Final = 5 """Stamped into PRAGMA user_version; bump when _SCHEMA or migrations change.""" DATABASE_ENV: Final = "REFLEX_WORKFLOW_DATABASE" @@ -2805,9 +2907,18 @@ def resolve_store(target: str | None = None) -> RunStore: labels TEXT, deadline REAL, cancel_requested INTEGER NOT NULL DEFAULT 0, + release_id TEXT, created_at REAL NOT NULL, updated_at REAL NOT NULL ); +CREATE TABLE IF NOT EXISTS workflow_workers ( + worker_id TEXT PRIMARY KEY, + release_id TEXT, + queues TEXT NOT NULL, + capacity INTEGER NOT NULL, + started_at REAL NOT NULL, + heartbeat_at REAL NOT NULL +); CREATE TABLE IF NOT EXISTS workflow_steps ( run_id TEXT NOT NULL, ordinal INTEGER NOT NULL, @@ -2909,6 +3020,7 @@ def resolve_store(target: str | None = None) -> RunStore: ) _RUN_MIGRATIONS: Final = ( + ("release_id", "ALTER TABLE workflow_runs ADD COLUMN release_id TEXT"), ("flow_key", "ALTER TABLE workflow_runs ADD COLUMN flow_key TEXT"), ("parent_run_id", "ALTER TABLE workflow_runs ADD COLUMN parent_run_id TEXT"), ("parent_ordinal", "ALTER TABLE workflow_runs ADD COLUMN parent_ordinal INTEGER"), @@ -2971,6 +3083,7 @@ def _run_from_row(row: sqlite3.Row) -> RunRecord: parent_ordinal=row["parent_ordinal"], request_key=row["request_key"], labels=_load(row["labels"]), + release_id=row["release_id"], deadline=row["deadline"], cancel_requested=bool(row["cancel_requested"]), created_at=row["created_at"], @@ -3089,6 +3202,7 @@ def _sqlite_frontier_query( due_only: bool, order: str, limit: int, + release: str | None = None, ) -> tuple[str, tuple[Any, ...]]: """Build the query for claimable-or-waking frontier steps. @@ -3108,6 +3222,8 @@ def _sqlite_frontier_query( clock event alone can make claimable, however far out. order: ORDER BY expression. limit: Maximum rows. + release: The claiming worker's release; a run pinned to a different + release is skipped, never claimed. Returns: The SQL and its parameters. @@ -3129,6 +3245,10 @@ def _sqlite_frontier_query( waking_params = (*claimable, StepStatus.BLOCKED.value) queue_sql = f" AND s.queue IN ({marks(queues)})" if queues is not None else "" queue_params = tuple(queues) if queues is not None else () + release_sql = ( + " AND (r.release_id IS NULL OR r.release_id = ?)" if release is not None else "" + ) + release_params = (release,) if release is not None else () sql = ( f"SELECT {select} FROM workflow_steps s" " JOIN workflow_runs r ON r.run_id = s.run_id" @@ -3136,6 +3256,7 @@ def _sqlite_frontier_query( f" AND r.status NOT IN ({marks(terminal_runs)})" " AND r.status != ? AND r.cancel_requested = 0" " AND (r.deadline IS NULL OR r.deadline > ?)" + f"{release_sql}" " AND NOT EXISTS (SELECT 1 FROM workflow_steps x" " WHERE x.run_id = s.run_id AND x.ordinal < s.ordinal" f" AND x.status NOT IN ({marks(terminal_steps)}))" @@ -3146,6 +3267,7 @@ def _sqlite_frontier_query( *terminal_runs, RunStatus.NEEDS_ATTENTION.value, now, + *release_params, *terminal_steps, *queue_params, ) @@ -3172,6 +3294,9 @@ def _sqlite_run_filters(query: RunQuery) -> tuple[str, tuple[Any, ...]]: if query.definition_digest is not None: clauses.append("definition_digest = ?") params.append(query.definition_digest) + if query.release_id is not None: + clauses.append("release_id = ?") + params.append(query.release_id) if query.statuses: placeholders = ",".join("?" * len(query.statuses)) clauses.append(f"status IN ({placeholders})") @@ -3309,9 +3434,9 @@ def _insert_run(self, run: RunRecord) -> None: "INSERT INTO workflow_runs (run_id, workflow_id, definition_digest," " status, state, state_version, next_ordinal, result, error," " flow_key, parent_run_id, parent_ordinal, parent_close," - " request_key, labels, deadline, cancel_requested, created_at," - " updated_at)" - " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + " request_key, labels, deadline, cancel_requested, release_id," + " created_at, updated_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( run.run_id, run.workflow_id, @@ -3330,6 +3455,7 @@ def _insert_run(self, run: RunRecord) -> None: _dump(run.labels), run.deadline, int(run.cancel_requested), + run.release_id, run.created_at, run.updated_at, ), @@ -3562,6 +3688,102 @@ def _route_parked_in_txn(self, parked_id: str, now: float) -> DeliveryDispositio ) return disposition if delivered else "dead_letter" + async def register_worker(self, worker: WorkerRecord) -> None: + """Record (or refresh) a worker's registration. + + Args: + worker: The worker's identity, release, queues, and capacity. + """ + + def work() -> None: + """Run the operation on the worker thread.""" + with self._lock: + self._db.execute( + "INSERT INTO workflow_workers (worker_id, release_id," + " queues, capacity, started_at, heartbeat_at)" + " VALUES (?, ?, ?, ?, ?, ?)" + " ON CONFLICT (worker_id) DO UPDATE SET release_id =" + " excluded.release_id, queues = excluded.queues," + " capacity = excluded.capacity," + " heartbeat_at = excluded.heartbeat_at", + ( + worker.worker_id, + worker.release_id, + json.dumps(list(worker.queues)), + worker.capacity, + worker.started_at, + worker.heartbeat_at, + ), + ) + + await asyncio.to_thread(work) + + async def heartbeat_worker(self, worker_id: str, now: float) -> None: + """Refresh a worker's sign of life. + + Args: + worker_id: The worker. + now: Current time in epoch seconds. + """ + + def work() -> None: + """Run the operation on the worker thread.""" + with self._lock: + self._db.execute( + "UPDATE workflow_workers SET heartbeat_at = ? WHERE worker_id = ?", + (now, worker_id), + ) + + await asyncio.to_thread(work) + + async def deregister_worker(self, worker_id: str) -> None: + """Remove a worker that shut down cleanly. + + Args: + worker_id: The worker. + """ + + def work() -> None: + """Run the operation on the worker thread.""" + with self._lock: + self._db.execute( + "DELETE FROM workflow_workers WHERE worker_id = ?", + (worker_id,), + ) + + await asyncio.to_thread(work) + + async def list_workers(self) -> tuple[WorkerRecord, ...]: + """List registered workers, most recently started first. + + Returns: + The registrations. + """ + + def work() -> tuple[WorkerRecord, ...]: + """Run the operation on the worker thread. + + Returns: + The registrations. + """ + with self._lock: + rows = self._db.execute( + "SELECT * FROM workflow_workers ORDER BY started_at DESC" + ).fetchall() + return tuple( + WorkerRecord( + worker_id=row["worker_id"], + release_id=row["release_id"], + queues=tuple(json.loads(row["queues"])), + capacity=row["capacity"], + started_at=row["started_at"], + heartbeat_at=row["heartbeat_at"], + ) + for row in rows + ) + + return await asyncio.to_thread(work) + async def ingest_channel_delivery( self, workflow_id: str, @@ -3956,6 +4178,7 @@ async def claim_next( *, lease_duration: float = DEFAULT_LEASE_DURATION, queues: tuple[str, ...] | None = None, + release: str | None = None, ) -> Claim | None: """Claim the due frontier step of some runnable run. @@ -3966,6 +4189,9 @@ async def claim_next( queues: Queues this worker serves; None serves every queue. A run whose frontier sits on an unserved queue is skipped whole, because claiming a later slot would break its ordering. + release: The claiming worker's release identity. A run pinned to + a different release is skipped: it drains on the release that + admitted it, so one run never mixes two releases' code. Returns: A fenced claim, or None when nothing is claimable right now. @@ -3993,6 +4219,7 @@ def work(): due_only=True, order="s.due_at, s.run_id", limit=1, + release=release, ) row = self._db.execute(sql, params).fetchone() if row is not None: diff --git a/tests/units/workflow/test_release_routing.py b/tests/units/workflow/test_release_routing.py new file mode 100644 index 00000000000..b5d24c19cf4 --- /dev/null +++ b/tests/units/workflow/test_release_routing.py @@ -0,0 +1,182 @@ +"""Release-pinned execution across a rolling deploy. + +The acceptance shape: a workflow sleeps for a day, a new release deploys, +new runs use the new release, and the sleeping run resumes on the release +that admitted it. Rollback is the same property read the other way, and the +retirement gate is a count. +""" + +import asyncio + +import pytest +from reflex_base.workflow import WorkflowConfig, manual + +import reflex as rx +from reflex.workflow import testing +from reflex.workflow.definition import compile_workflow +from reflex.workflow.kernel import WorkflowKernel +from reflex.workflow.records import TERMINAL_RUN_STATUSES, RunQuery, RunStatus + + +def _flow(): + """Build a workflow that sleeps between two steps. + + Returns: + The workflow class. + """ + + class Deploying(rx.State): + __workflow__ = WorkflowConfig(id="release.deploying") + note: str = "" + + @rx.event(durable=True, trigger=manual(), effect="none") + def begin(self): + """Sleep a day before finishing. + + Returns: + The deferral. + """ + return rx.after("1d", Deploying.finish) + + @rx.event(durable=True, effect="none") + def finish(self): + """Complete. + + Returns: + Completion. + """ + return rx.complete(result="done") + + return Deploying + + +ACTIVE = tuple(s for s in RunStatus if s not in TERMINAL_RUN_STATUSES) + + +async def test_a_sleeping_run_resumes_on_the_release_that_admitted_it( + forked_registration_context, +): + """v2 never claims v1's sleeping run; v1 drains it; the gate sees both. + + Args: + forked_registration_context: Isolated state registry. + """ + store = testing.MemoryRunStore() + flow = _flow() + clock = [1_000_000.0] + + v1 = WorkflowKernel( + [compile_workflow(flow)], store, release="v1", clock=lambda: clock[0] + ) + started = await v1.start(flow.begin()) + assert started.run_id is not None + await v1.run_until_idle() + sleeping = await v1.get_run(started.run_id) + assert sleeping is not None + assert sleeping.release_id == "v1" + assert sleeping.status is RunStatus.WAITING + + # The new release arrives; its worker takes new admissions only. + v2 = WorkflowKernel( + [compile_workflow(flow)], store, release="v2", clock=lambda: clock[0] + ) + fresh = await v2.start(flow.begin()) + assert fresh.run_id is not None + fresh_run = await v2.get_run(fresh.run_id) + assert fresh_run is not None + assert fresh_run.release_id == "v2" + + clock[0] += 90_000.0 + assert await store.claim_next(clock[0], release="v2") is not None, ( + "v2 claims its own due run" + ) + remaining = await store.claim_next(clock[0], release="v2") + assert remaining is None, "v2 must never claim the run v1 admitted" + + # The retirement gate: v1 cannot retire while its run is active. + assert await store.count_runs(RunQuery(release_id="v1", statuses=ACTIVE)) == 1 + await v1.run_until_idle() + drained = await v1.get_run(started.run_id) + assert drained is not None + assert drained.status is RunStatus.COMPLETED + assert await store.count_runs(RunQuery(release_id="v1", statuses=ACTIVE)) == 0, ( + "with its runs drained, v1's workers may retire" + ) + + +async def test_workers_register_their_release_and_deregister_cleanly( + forked_registration_context, +): + """The fleet surface shows who runs what, and a clean stop removes it. + + Args: + forked_registration_context: Isolated state registry. + """ + store = testing.MemoryRunStore() + flow = _flow() + kernel = WorkflowKernel( + [compile_workflow(flow)], + store, + release="v7", + queues=("billing",), + max_concurrency=3, + ) + await kernel.start_worker() + workers = await store.list_workers() + assert len(workers) == 1 + assert workers[0].release_id == "v7" + assert workers[0].queues == ("billing",) + assert workers[0].capacity == 3 + + before = workers[0].heartbeat_at + await asyncio.sleep(0.05) + await kernel.recover() + workers = await store.list_workers() + # Each recovery re-measures the store-clock offset, and two midpoint + # estimates can differ by a few milliseconds in either direction; the + # heartbeat provably refreshed if it moved at all beyond that jitter. + assert workers[0].heartbeat_at > before - 1.0 + assert workers[0].heartbeat_at != before + + await kernel.aclose() + assert await store.list_workers() == () + + +async def test_an_unreleased_dev_kernel_serves_pinned_runs( + forked_registration_context, +): + """A worker with no declared release is dev tooling; it runs anything. + + Args: + forked_registration_context: Isolated state registry. + """ + store = testing.MemoryRunStore() + flow = _flow() + v1 = WorkflowKernel([compile_workflow(flow)], store, release="v1") + started = await v1.start(flow.begin()) + assert started.run_id is not None + + dev = WorkflowKernel([compile_workflow(flow)], store) + assert dev._release is None # pyright: ignore[reportPrivateUsage] + claim = await store.claim_next(2_000_000_000.0, release=None) + assert claim is not None, "the undeclared worker claims the pinned run" + + +def test_release_defaults_from_the_environment( + monkeypatch, forked_registration_context +): + """REFLEX_RELEASE_ID is the deploy surface's way in. + + Args: + monkeypatch: Used to set the environment. + forked_registration_context: Isolated state registry. + """ + monkeypatch.setenv("REFLEX_RELEASE_ID", "rel-2026-08-24") + kernel = WorkflowKernel([], testing.MemoryRunStore()) + assert kernel._release == "rel-2026-08-24" # pyright: ignore[reportPrivateUsage] + monkeypatch.setenv("REFLEX_RELEASE_ID", "") + bare = WorkflowKernel([], testing.MemoryRunStore()) + assert bare._release is None # pyright: ignore[reportPrivateUsage] + + +_ = pytest From 08897b5c8363b65f76d22c54ea2ffa0114e6093d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Mon, 24 Aug 2026 19:28:25 -0700 Subject: [PATCH 121/121] workflows: operator mutations record who asked and why The audit half of the operator console's acceptance criterion ("every mutation shows who performed it and why"), landed ahead of the UI so the console renders a record that already exists. Attribution rides the history. The operator store operations -- request_cancel, retry_run, skip_step, resume_run, finalize_run -- take an optional attribution mapping and merge it into the operator-facing history event, inside the same transaction as the mutation itself. An audit that lives outside the history would drift from it; this one cannot, because the record that answers "what happened" is the record that answers "who did this". The kernel's operator methods (cancel, retry, skip, resume, force_finalize) take actor and reason and build the payload. The CLI stamps the invoking user (REFLEX_ACTOR overrides, getpass fallback) and grows --reason on cancel, retry, skip, resume, and complete; fail already demanded a reason and now records it as attribution too. The HTTP API records the caller's X-Actor claim -- tokens are anonymous, so the header is a claim, recorded as given, and an authenticating proxy can stamp it; "api" beats naming nobody -- plus the body's reason on cancel, retry, and resume. Pinned by a conformance check exercising all five operations on all three stores, a serve test proving X-Actor and the body reason land in history, and a CLI subprocess test proving --reason and REFLEX_ACTOR do. Contract 9 states the rule. --- news/workflow-operator-attribution.md | 1 + reflex/workflow/CONTRACT.md | 8 ++ reflex/workflow/cli.py | 66 +++++++-- reflex/workflow/conformance.py | 72 ++++++++++ reflex/workflow/kernel.py | 89 ++++++++++-- reflex/workflow/postgres.py | 69 ++++++++-- reflex/workflow/serve.py | 12 +- reflex/workflow/store.py | 186 +++++++++++++++++++++++--- tests/units/workflow/test_cli.py | 92 +++++++++++++ tests/units/workflow/test_serve.py | 40 ++++++ 10 files changed, 582 insertions(+), 53 deletions(-) create mode 100644 news/workflow-operator-attribution.md diff --git a/news/workflow-operator-attribution.md b/news/workflow-operator-attribution.md new file mode 100644 index 00000000000..670ac34dc73 --- /dev/null +++ b/news/workflow-operator-attribution.md @@ -0,0 +1 @@ +Operator mutations record who asked and why. `cancel`, `retry`, `skip`, `resume`, `complete`, and `fail` — on the CLI (`--reason`, actor from the invoking user or `REFLEX_ACTOR`) and the HTTP API (`X-Actor` header, body `reason`) — carry attribution into the same history events, in the same store transactions, as the mutations they describe, so the record that answers "what happened" also answers "who did this" and cannot drift from it. Pinned by a conformance check on all three stores plus end-to-end tests through the CLI and the standalone service. This is the audit half of the operator console's acceptance criterion, landed ahead of the UI. diff --git a/reflex/workflow/CONTRACT.md b/reflex/workflow/CONTRACT.md index 5eab60b489f..92a1a179a69 100644 --- a/reflex/workflow/CONTRACT.md +++ b/reflex/workflow/CONTRACT.md @@ -502,6 +502,14 @@ follow-up). ## 9. Operator actions +Every operator mutation records **who asked and why** when the surface +knows: the CLI stamps the invoking user (`REFLEX_ACTOR` overrides) and its +`--reason`; the HTTP API records the caller's `X-Actor` claim (default +`api` — tokens are anonymous, and an authenticating proxy can stamp the +header) and the body's `reason`. Attribution rides the same history events, +in the same transactions, as the mutations they describe — the record that +answers "what happened" answers "who did this", and cannot drift from it. + Every action is legal only from the states listed; anything else is a refused no-op with a reason. diff --git a/reflex/workflow/cli.py b/reflex/workflow/cli.py index 35d6c792004..175add48f2e 100644 --- a/reflex/workflow/cli.py +++ b/reflex/workflow/cli.py @@ -50,6 +50,24 @@ def webhook_root_names(definitions: Iterable[WorkflowDefinition]) -> list[str]: ) +def _cli_attribution(reason: str | None) -> dict[str, str]: + """Who is running this command, and why, for the run's history. + + Args: + reason: The operator's stated reason, if any. + + Returns: + The attribution mapping. + """ + import getpass + + actor = os.environ.get("REFLEX_ACTOR") or getpass.getuser() + payload = {"actor": actor} + if reason: + payload["reason"] = reason + return payload + + def _operator_action(database: str | None, run_id: str, action: str, **extra): """Apply one operator action to a run, reporting what happened. @@ -1557,7 +1575,8 @@ async def act(store: RunStore): @workflows.command() @database_option @click.argument("run_id") -def cancel(database: str | None, run_id: str): +@click.option("--reason", default=None, help="Why, recorded in the run's history.") +def cancel(database: str | None, run_id: str, reason: str | None): """Request cancellation of a run. The running worker finalizes it; if no worker is running, it is cancelled @@ -1575,7 +1594,9 @@ async def request(store: RunStore) -> bool: Whether intent was recorded. """ return await store.request_cancel( - await _resolve_run_id(store, run_id), time.time() + await _resolve_run_id(store, run_id), + time.time(), + _cli_attribution(reason), ) recorded = _with_store(database, request) @@ -1588,25 +1609,31 @@ async def request(store: RunStore) -> bool: @workflows.command() @database_option @click.argument("run_id") -def retry(database: str | None, run_id: str): +@click.option("--reason", default=None, help="Why, recorded in the run's history.") +def retry(database: str | None, run_id: str, reason: str | None): """Re-open a failed run at the step that failed. The step runs again with a fresh attempt budget; the original failure stays in the run's history. """ - _operator_action(database, run_id, "retry_run") + _operator_action( + database, run_id, "retry_run", attribution=_cli_attribution(reason) + ) @workflows.command() @database_option @click.argument("run_id") -def skip(database: str | None, run_id: str): +@click.option("--reason", default=None, help="Why, recorded in the run's history.") +def skip(database: str | None, run_id: str, reason: str | None): """Skip the step blocking a stopped run and let it continue. For a step that cannot succeed and is not worth failing the run over. It is recorded as an operator decision, not as an outcome. """ - _operator_action(database, run_id, "skip_step") + _operator_action( + database, run_id, "skip_step", attribution=_cli_attribution(reason) + ) def _finalize( @@ -1616,6 +1643,7 @@ def _finalize( *, result: Any = None, error: dict[str, Any] | None = None, + reason: str | None = None, ): """End a run by operator decision. @@ -1630,6 +1658,7 @@ def _finalize( status: The terminal status to record. result: Result to record when completing. error: Error payload to record when failing. + reason: Why, recorded in the run's history. Raises: Exit: When the run is unknown, already finished, or has a claimed step. @@ -1650,6 +1679,8 @@ async def finish(store: RunStore) -> bool: status=status, result=result, error=error, + actor=_cli_attribution(None)["actor"], + reason=reason, ) finalized = _with_store(database, finish) @@ -1665,10 +1696,16 @@ async def finish(store: RunStore) -> bool: @workflows.command() @database_option @click.argument("run_id") +@click.option("--reason", default=None, help="Why, recorded in the run's history.") @click.option( "--result", "result_json", default=None, help="Result to record, as JSON." ) -def complete(database: str | None, run_id: str, result_json: str | None): +def complete( + database: str | None, + run_id: str, + result_json: str | None, + reason: str | None, +): """End a run as completed by operator decision. For a run no code path will finish: a wait nobody will answer, a branch @@ -1683,7 +1720,7 @@ def complete(database: str | None, run_id: str, result_json: str | None): except json.JSONDecodeError as err: console.error(f"--result is not JSON: {err}") raise click.exceptions.Exit(1) from None - _finalize(database, run_id, RunStatus.COMPLETED, result=result) + _finalize(database, run_id, RunStatus.COMPLETED, result=result, reason=reason) @workflows.command() @@ -1696,13 +1733,16 @@ def fail(database: str | None, run_id: str, reason: str): The reason is recorded on the run, so the history says a person decided this rather than leaving a failure with no explanation. """ - _finalize(database, run_id, RunStatus.FAILED, error={"reason": reason}) + _finalize( + database, run_id, RunStatus.FAILED, error={"reason": reason}, reason=reason + ) @workflows.command() @database_option @click.argument("run_id") -def resume(database: str | None, run_id: str): +@click.option("--reason", default=None, help="Why, recorded in the run's history.") +def resume(database: str | None, run_id: str, reason: str | None): """Re-open a run suspended for operator attention.""" import time @@ -1715,7 +1755,11 @@ async def reopen(store: RunStore) -> bool: Returns: Whether a suspended run was re-opened. """ - return await store.resume_run(await _resolve_run_id(store, run_id), time.time()) + return await store.resume_run( + await _resolve_run_id(store, run_id), + time.time(), + _cli_attribution(reason), + ) resumed = _with_store(database, reopen) if not resumed: diff --git a/reflex/workflow/conformance.py b/reflex/workflow/conformance.py index 77834617e07..8837291ec16 100644 --- a/reflex/workflow/conformance.py +++ b/reflex/workflow/conformance.py @@ -1927,6 +1927,77 @@ async def check_release_counts_answer_the_retirement_question( ), "nothing pins to a release that admitted nothing" +async def check_operator_actions_carry_attribution(store: RunStore) -> None: + """Every operator mutation records who asked and why, in the run's story. + + An audit that lives outside the history would drift from it; attribution + rides the same events, in the same transactions, so "who did this" is + answered by the record that already answers "what happened". + """ + who = {"actor": "alex", "reason": "customer asked"} + + await store.admit(make_run(), make_step(due_at=0.0), _ADMITTED) + assert await store.request_cancel("run1", NOW, who) + events = await store.get_history("run1") + cancel_event = next( + event for event in events if event.type is HistoryEventType.RUN_CANCEL_REQUESTED + ) + assert cancel_event.data["actor"] == "alex" + assert cancel_event.data["reason"] == "customer asked" + + await store.admit( + make_run("fail1", status=RunStatus.FAILED, error={"reason": "boom"}), + make_step("fail1", status=StepStatus.FAILED, error={"reason": "boom"}), + _ADMITTED, + ) + assert await store.retry_run("fail1", NOW, who) + events = await store.get_history("fail1") + resumed = next( + event for event in events if event.type is HistoryEventType.RUN_RESUMED + ) + assert resumed.data["actor"] == "alex" + + await store.admit( + make_run("skip1", status=RunStatus.FAILED, error={"reason": "boom"}), + make_step("skip1", status=StepStatus.FAILED, error={"reason": "boom"}), + _ADMITTED, + ) + assert await store.skip_step("skip1", NOW, who) + events = await store.get_history("skip1") + skipped = next( + event for event in events if event.type is HistoryEventType.STEP_SKIPPED + ) + assert skipped.data["reason"] == "customer asked" + + await store.admit( + make_run("sus1", status=RunStatus.NEEDS_ATTENTION), + make_step("sus1", status=StepStatus.NEEDS_ATTENTION), + _ADMITTED, + ) + assert await store.resume_run("sus1", NOW, who) + events = await store.get_history("sus1") + reopened = next( + event for event in events if event.type is HistoryEventType.RUN_RESUMED + ) + assert reopened.data["actor"] == "alex" + + await store.admit(make_run("fin1"), make_step("fin1", due_at=0.0), _ADMITTED) + assert await store.finalize_run( + "fin1", + status=RunStatus.FAILED, + error={"reason": "manual"}, + event=HistoryEventType.RUN_FAILED, + now=NOW, + attribution=who, + ) + events = await store.get_history("fin1") + failed = next( + event for event in events if event.type is HistoryEventType.RUN_FAILED + ) + assert failed.data["actor"] == "alex" + assert failed.data["reason"] == "customer asked" + + CONFORMANCE_CHECKS: tuple[Callable[[RunStore], Awaitable[None]], ...] = ( check_admit_creates_a_run, check_reads_do_not_alias_stored_state, @@ -1952,6 +2023,7 @@ async def check_release_counts_answer_the_retirement_question( check_a_pinned_run_drains_only_on_its_release, check_worker_registry_roundtrip, check_release_counts_answer_the_retirement_question, + check_operator_actions_carry_attribution, check_a_delivery_to_a_past_deadline_run_is_refused, check_a_duplicate_delivery_is_recorded_in_history, check_an_early_delivery_is_buffered_then_consumed, diff --git a/reflex/workflow/kernel.py b/reflex/workflow/kernel.py index c9830e2ffa3..8e3c3db9e39 100644 --- a/reflex/workflow/kernel.py +++ b/reflex/workflow/kernel.py @@ -952,7 +952,29 @@ async def _admit( self._wakeup.set() return StartResult(disposition="started", run_id=authoritative_run_id) - async def cancel(self, run_id: str) -> bool: + @staticmethod + def _attribution(actor: str | None, reason: str | None) -> dict[str, str] | None: + """Build the who-and-why payload an operator event carries. + + Args: + actor: Who asked, if known. + reason: Why, if given. + + Returns: + The attribution mapping, or None when neither is known. + """ + payload = { + key: value for key, value in (("actor", actor), ("reason", reason)) if value + } + return payload or None + + async def cancel( + self, + run_id: str, + *, + actor: str | None = None, + reason: str | None = None, + ) -> bool: """Request cancellation of a run. The in-flight attempt, if any, is cancelled cooperatively; the run is @@ -960,14 +982,18 @@ async def cancel(self, run_id: str) -> bool: Args: run_id: The run to cancel. + actor: Who asked, recorded in the run's history. + reason: Why, recorded alongside. Returns: True if intent was recorded on a nonterminal run. """ - recorded = await self._store.request_cancel(run_id, self._clock()) + attribution = self._attribution(actor, reason) + recorded = await self._store.request_cancel(run_id, self._clock(), attribution) if recorded: await self._notify_run( - run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),) + run_id, + ((HistoryEventType.RUN_CANCEL_REQUESTED, dict(attribution or {})),), ) task = self._inflight.get(run_id) if task is not None: @@ -1184,31 +1210,59 @@ async def signal( ) return disposition - async def resume(self, run_id: str) -> bool: + async def resume( + self, + run_id: str, + *, + actor: str | None = None, + reason: str | None = None, + ) -> bool: """Re-open a run that is suspended for operator attention. Args: run_id: The run to resume. + actor: Who asked, recorded in the run's history. + reason: Why, recorded alongside. Returns: True if a suspended run was re-opened. """ - resumed = await self._store.resume_run(run_id, self._clock()) + resumed = await self._store.resume_run( + run_id, self._clock(), self._attribution(actor, reason) + ) if resumed: - await self._notify_run(run_id, ((HistoryEventType.RUN_RESUMED, {}),)) + await self._notify_run( + run_id, + ( + ( + HistoryEventType.RUN_RESUMED, + dict(self._attribution(actor, reason) or {}), + ), + ), + ) self._wakeup.set() return resumed - async def retry(self, run_id: str) -> bool: + async def retry( + self, + run_id: str, + *, + actor: str | None = None, + reason: str | None = None, + ) -> bool: """Re-open a failed run at the step that failed. Args: run_id: The run to retry. + actor: Who asked, recorded in the run's history. + reason: Why, recorded alongside. Returns: True if a failed run was re-opened. """ - retried = await self._store.retry_run(run_id, self._clock()) + retried = await self._store.retry_run( + run_id, self._clock(), self._attribution(actor, reason) + ) if retried: await self._notify_run( run_id, ((HistoryEventType.RUN_RESUMED, {"origin": "retry"}),) @@ -1216,16 +1270,26 @@ async def retry(self, run_id: str) -> bool: self._wakeup.set() return retried - async def skip(self, run_id: str) -> bool: + async def skip( + self, + run_id: str, + *, + actor: str | None = None, + reason: str | None = None, + ) -> bool: """Skip the step blocking a stopped run and let it continue. Args: run_id: The run to unstick. + actor: Who asked, recorded in the run's history. + reason: Why, recorded alongside. Returns: True if a blocking step was skipped. """ - skipped = await self._store.skip_step(run_id, self._clock()) + skipped = await self._store.skip_step( + run_id, self._clock(), self._attribution(actor, reason) + ) if skipped: await self._notify_run( run_id, ((HistoryEventType.STEP_SKIPPED, {"origin": "operator"}),) @@ -1240,6 +1304,8 @@ async def force_finalize( status: RunStatus, result: Any = None, error: dict[str, Any] | None = None, + actor: str | None = None, + reason: str | None = None, ) -> bool: """End a run by operator decision, tombstoning what it had open. @@ -1250,6 +1316,8 @@ async def force_finalize( Args: run_id: The run to finalize. + actor: Who asked, recorded in the run's history. + reason: Why, recorded alongside. status: The terminal status to record. result: Result to record when completing. error: Error payload to record when failing. @@ -1285,6 +1353,7 @@ async def force_finalize( now=now, result=result, parent_arrival=self._arrival_for(run, status, result, error), + attribution=self._attribution(actor, reason), ) if finalized: self._notify(run, ((event, {"origin": "operator"}),)) diff --git a/reflex/workflow/postgres.py b/reflex/workflow/postgres.py index dfb9a33ffb8..5d8b8f6637c 100644 --- a/reflex/workflow/postgres.py +++ b/reflex/workflow/postgres.py @@ -2153,12 +2153,20 @@ async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: ) return cursor.rowcount > 0 - async def request_cancel(self, run_id: str, now: float) -> bool: + async def request_cancel( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Record cancellation intent on a run. Args: run_id: The run to cancel. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if intent was recorded on a nonterminal run. @@ -2173,7 +2181,10 @@ async def request_cancel(self, run_id: str, now: float) -> bool: if cursor.rowcount == 0: return False await self._append_events( - conn, run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now + conn, + run_id, + ((HistoryEventType.RUN_CANCEL_REQUESTED, dict(attribution or {})),), + now, ) return True @@ -2208,6 +2219,7 @@ async def finalize_run( now: float, result: Any = None, parent_arrival: tuple[str, int, dict[str, Any], str] | None = None, + attribution: Mapping[str, str] | None = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -2217,6 +2229,9 @@ async def finalize_run( error: Error payload recorded on the run. event: The terminal history event type. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". result: Result to record, for an operator forcing completion. parent_arrival: When this run is a child, the arrival to deliver to its parent's join, applied in this same transaction. @@ -2261,19 +2276,30 @@ async def finalize_run( (HistoryEventType.STEP_TOMBSTONED, {"ordinal": open_row["ordinal"]}) for open_row in open_rows ] - events.append((event, {} if error is None else dict(error))) + events.append(( + event, + {**({} if error is None else dict(error)), **(attribution or {})}, + )) await self._append_events(conn, run_id, events, now) await self._close_children(conn, now, locked_children) if parent_arrival is not None: await self._apply_arrival(conn, *parent_arrival, now) return True - async def resume_run(self, run_id: str, now: float) -> bool: + async def resume_run( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Re-open a suspended run so its frontier step runs again. Args: run_id: The run to resume. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a suspended run was re-opened. @@ -2300,7 +2326,10 @@ async def resume_run(self, run_id: str, now: float) -> bool: ), ) await self._append_events( - conn, run_id, ((HistoryEventType.RUN_RESUMED, {}),), now + conn, + run_id, + ((HistoryEventType.RUN_RESUMED, dict(attribution or {})),), + now, ) return True @@ -2341,12 +2370,20 @@ async def _restore_tombstoned(conn: Any, run_id: str, now: float) -> list[int]: ) return sorted(row["ordinal"] for row in await cursor.fetchall()) - async def retry_run(self, run_id: str, now: float) -> bool: + async def retry_run( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Re-open a failed run at the step that failed. Args: run_id: The run to retry. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a failed run was re-opened. @@ -2378,7 +2415,10 @@ async def retry_run(self, run_id: str, now: float) -> bool: conn, run_id, ( - (HistoryEventType.RUN_RESUMED, {"origin": "retry"}), + ( + HistoryEventType.RUN_RESUMED, + {"origin": "retry", **(attribution or {})}, + ), *( (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) for ordinal in restored @@ -2388,12 +2428,20 @@ async def retry_run(self, run_id: str, now: float) -> bool: ) return True - async def skip_step(self, run_id: str, now: float) -> bool: + async def skip_step( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Give up on a stuck step and let the run carry on past it. Args: run_id: The run to unstick. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a blocking step was skipped. @@ -2447,7 +2495,10 @@ async def skip_step(self, run_id: str, now: float) -> bool: ), ) events = [ - (HistoryEventType.STEP_SKIPPED, {"ordinal": row["ordinal"]}), + ( + HistoryEventType.STEP_SKIPPED, + {"ordinal": row["ordinal"], **(attribution or {})}, + ), *( (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) for ordinal in restored diff --git a/reflex/workflow/serve.py b/reflex/workflow/serve.py index e9ad151193b..88956b28555 100644 --- a/reflex/workflow/serve.py +++ b/reflex/workflow/serve.py @@ -491,7 +491,17 @@ async def endpoint(request: Request) -> JSONResponse: run_id = request.path_params["run_id"] if await runtime.kernel.get_run(run_id) is None: return JSONResponse({"error": "unknown run"}, status_code=404) - applied = await getattr(runtime.kernel, action)(run_id) + payload, bad = await _read_json(request) + if bad is not None: + return bad + reason = payload.get("reason") if isinstance(payload, dict) else None + # Tokens are anonymous; X-Actor is the caller's claim of identity, + # recorded as given. An audit that names "api" beats one that names + # nobody, and a proxy that authenticates people can stamp the header. + actor = request.headers.get("x-actor") or "api" + applied = await getattr(runtime.kernel, action)( + run_id, actor=actor, reason=reason + ) if not applied: # The run exists but is not in a state this action accepts -- # retrying a healthy run, resuming one that is not suspended. diff --git a/reflex/workflow/store.py b/reflex/workflow/store.py index 7476753a993..c9db6036390 100644 --- a/reflex/workflow/store.py +++ b/reflex/workflow/store.py @@ -676,12 +676,20 @@ async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: """ ... - async def request_cancel(self, run_id: str, now: float) -> bool: + async def request_cancel( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Record cancellation intent on a run. Args: run_id: The run to cancel. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if intent was recorded on a nonterminal run. @@ -712,6 +720,7 @@ async def finalize_run( now: float, result: Any = None, parent_arrival: tuple[str, int, dict[str, Any], str] | None = None, + attribution: Mapping[str, str] | None = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -721,6 +730,9 @@ async def finalize_run( error: Error payload recorded on the run. event: The terminal history event type. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". result: Result to record, for an operator forcing completion. parent_arrival: When this run is a child, the arrival to deliver to its parent's join, applied in this same transaction. @@ -731,7 +743,12 @@ async def finalize_run( """ ... - async def skip_step(self, run_id: str, now: float) -> bool: + async def skip_step( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Give up on a stuck step and let the run carry on past it. The operator's answer to a step that cannot succeed and is not worth @@ -744,13 +761,21 @@ async def skip_step(self, run_id: str, now: float) -> bool: Args: run_id: The run to unstick. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a blocking step was skipped. """ ... - async def retry_run(self, run_id: str, now: float) -> bool: + async def retry_run( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Re-open a failed run at the step that failed. The operator's answer to a run that failed for a reason now fixed: the @@ -760,13 +785,21 @@ async def retry_run(self, run_id: str, now: float) -> bool: Args: run_id: The run to retry. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a failed run was re-opened. """ ... - async def resume_run(self, run_id: str, now: float) -> bool: + async def resume_run( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Re-open a suspended run so its frontier step runs again. Suspension is an operator state, not an outcome: the run waits for a @@ -777,6 +810,9 @@ async def resume_run(self, run_id: str, now: float) -> bool: Args: run_id: The run to resume. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a suspended run was re-opened. @@ -2144,12 +2180,20 @@ async def defer_root(self, run_id: str, due_at: float, now: float) -> bool: steps[0] = dataclasses.replace(steps[0], due_at=due_at, updated_at=now) return True - async def request_cancel(self, run_id: str, now: float) -> bool: + async def request_cancel( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Record cancellation intent on a run. Args: run_id: The run to cancel. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if intent was recorded on a nonterminal run. @@ -2165,7 +2209,9 @@ async def request_cancel(self, run_id: str, now: float) -> bool: updated_at=now, ) self._append_events( - run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now + run_id, + ((HistoryEventType.RUN_CANCEL_REQUESTED, dict(attribution or {})),), + now, ) return True @@ -2245,6 +2291,7 @@ async def finalize_run( now: float, result: Any = None, parent_arrival: tuple[str, int, dict[str, Any], str] | None = None, + attribution: Mapping[str, str] | None = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -2254,6 +2301,9 @@ async def finalize_run( error: Error payload recorded on the run. event: The terminal history event type. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". result: Result to record, for an operator forcing completion. parent_arrival: When this run is a child, the arrival to deliver to its parent's join, applied in this same transaction. @@ -2285,14 +2335,22 @@ async def finalize_run( result=result if result is not None else run.result, updated_at=now, ) - events.append((event, {} if error is None else dict(error))) + events.append(( + event, + {**({} if error is None else dict(error)), **(attribution or {})}, + )) self._append_events(run_id, events, now) self._close_children(run_id, now) if parent_arrival is not None: self._apply_arrival(*parent_arrival, now) return True - async def skip_step(self, run_id: str, now: float) -> bool: + async def skip_step( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Give up on a stuck step and let the run carry on past it. The operator's answer to a step that cannot succeed and is not worth @@ -2305,6 +2363,9 @@ async def skip_step(self, run_id: str, now: float) -> bool: Args: run_id: The run to unstick. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a blocking step was skipped. @@ -2337,7 +2398,10 @@ async def skip_step(self, run_id: str, now: float) -> bool: other.status not in TERMINAL_STEP_STATUSES for other in steps ) events = [ - (HistoryEventType.STEP_SKIPPED, {"ordinal": step.ordinal}), + ( + HistoryEventType.STEP_SKIPPED, + {"ordinal": step.ordinal, **(attribution or {})}, + ), *( (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) for ordinal in restored @@ -2376,7 +2440,12 @@ async def skip_step(self, run_id: str, now: float) -> bool: return True return False - async def retry_run(self, run_id: str, now: float) -> bool: + async def retry_run( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Re-open a failed run at the step that failed. The operator's answer to a run that failed for a reason now fixed: the @@ -2386,6 +2455,9 @@ async def retry_run(self, run_id: str, now: float) -> bool: Args: run_id: The run to retry. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a failed run was re-opened. @@ -2418,7 +2490,10 @@ async def retry_run(self, run_id: str, now: float) -> bool: self._append_events( run_id, ( - (HistoryEventType.RUN_RESUMED, {"origin": "retry"}), + ( + HistoryEventType.RUN_RESUMED, + {"origin": "retry", **(attribution or {})}, + ), *( (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) for ordinal in restored @@ -2472,12 +2547,20 @@ def _restore_tombstoned(steps: list[StepRecord], now: float) -> list[int]: restored.append(step.ordinal) return restored - async def resume_run(self, run_id: str, now: float) -> bool: + async def resume_run( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Re-open a suspended run so its frontier step runs again. Args: run_id: The run to resume. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a suspended run was re-opened. @@ -2501,7 +2584,9 @@ async def resume_run(self, run_id: str, now: float) -> bool: self._runs[run_id] = dataclasses.replace( run, status=RunStatus.PENDING, error=None, updated_at=now ) - self._append_events(run_id, ((HistoryEventType.RUN_RESUMED, {}),), now) + self._append_events( + run_id, ((HistoryEventType.RUN_RESUMED, dict(attribution or {})),), now + ) return True async def recover_orphans( @@ -5082,12 +5167,20 @@ def work(): return await asyncio.to_thread(work) - async def request_cancel(self, run_id: str, now: float) -> bool: + async def request_cancel( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Record cancellation intent on a run. Args: run_id: The run to cancel. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if intent was recorded on a nonterminal run. @@ -5113,7 +5206,14 @@ def work(): self._db.execute("ROLLBACK") return False self._append_events( - run_id, ((HistoryEventType.RUN_CANCEL_REQUESTED, {}),), now + run_id, + ( + ( + HistoryEventType.RUN_CANCEL_REQUESTED, + dict(attribution or {}), + ), + ), + now, ) self._db.execute("COMMIT") except BaseException: @@ -5163,6 +5263,7 @@ async def finalize_run( now: float, result: Any = None, parent_arrival: tuple[str, int, dict[str, Any], str] | None = None, + attribution: Mapping[str, str] | None = None, ) -> bool: """Terminate a drained run and tombstone its unresolved slots. @@ -5172,6 +5273,9 @@ async def finalize_run( error: Error payload recorded on the run. event: The terminal history event type. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". result: Result to record, for an operator forcing completion. parent_arrival: When this run is a child, the arrival to deliver to its parent's join, applied in this same transaction. @@ -5226,7 +5330,13 @@ def work(): (HistoryEventType.STEP_TOMBSTONED, {"ordinal": row["ordinal"]}) for row in open_rows ] - events.append((event, {} if error is None else dict(error))) + events.append(( + event, + { + **({} if error is None else dict(error)), + **(attribution or {}), + }, + )) self._append_events(run_id, events, now) self._close_children_sql(run_id, now) if parent_arrival is not None: @@ -5239,7 +5349,12 @@ def work(): return await asyncio.to_thread(work) - async def skip_step(self, run_id: str, now: float) -> bool: + async def skip_step( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Give up on a stuck step and let the run carry on past it. The operator's answer to a step that cannot succeed and is not worth @@ -5252,6 +5367,9 @@ async def skip_step(self, run_id: str, now: float) -> bool: Args: run_id: The run to unstick. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a blocking step was skipped. @@ -5338,7 +5456,10 @@ def work() -> bool: ), ) events = [ - (HistoryEventType.STEP_SKIPPED, {"ordinal": row["ordinal"]}), + ( + HistoryEventType.STEP_SKIPPED, + {"ordinal": row["ordinal"], **(attribution or {})}, + ), *( (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) for ordinal in restored @@ -5379,7 +5500,12 @@ def work() -> bool: return await asyncio.to_thread(work) - async def retry_run(self, run_id: str, now: float) -> bool: + async def retry_run( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Re-open a failed run at the step that failed. The operator's answer to a run that failed for a reason now fixed: the @@ -5389,6 +5515,9 @@ async def retry_run(self, run_id: str, now: float) -> bool: Args: run_id: The run to retry. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a failed run was re-opened. @@ -5456,7 +5585,10 @@ def work() -> bool: self._append_events( run_id, ( - (HistoryEventType.RUN_RESUMED, {"origin": "retry"}), + ( + HistoryEventType.RUN_RESUMED, + {"origin": "retry", **(attribution or {})}, + ), *( (HistoryEventType.STEP_RESTORED, {"ordinal": ordinal}) for ordinal in restored @@ -5472,12 +5604,20 @@ def work() -> bool: return await asyncio.to_thread(work) - async def resume_run(self, run_id: str, now: float) -> bool: + async def resume_run( + self, + run_id: str, + now: float, + attribution: Mapping[str, str] | None = None, + ) -> bool: """Re-open a suspended run so its frontier step runs again. Args: run_id: The run to resume. now: Current time in epoch seconds. + attribution: Who asked and why, e.g. ``{"actor": ..., "reason": + ...}``, merged into the operator-facing history event so the + run's own story answers "who did this". Returns: True if a suspended run was re-opened. @@ -5518,7 +5658,9 @@ def work(): ), ) self._append_events( - run_id, ((HistoryEventType.RUN_RESUMED, {}),), now + run_id, + ((HistoryEventType.RUN_RESUMED, dict(attribution or {})),), + now, ) self._db.execute("COMMIT") except BaseException: diff --git a/tests/units/workflow/test_cli.py b/tests/units/workflow/test_cli.py index 53022df4ddc..32ee06096fc 100644 --- a/tests/units/workflow/test_cli.py +++ b/tests/units/workflow/test_cli.py @@ -353,3 +353,95 @@ def test_operator_actions_take_a_prefix_too(seeded): run = _load_run(database, suspended_id) assert run is not None assert run.status is not RunStatus.NEEDS_ATTENTION + + +def test_cancel_records_the_operator_and_their_reason(tmp_path, monkeypatch): + """`--reason` and the invoking user land in the run's history. + + Args: + tmp_path: Working directory for the database. + monkeypatch: Used to pin the actor. + """ + import asyncio + import json as jsonlib + import sqlite3 + import subprocess + import sys + import time + + from reflex.workflow.records import ( + HistoryEventType, + RunRecord, + RunStatus, + StepRecord, + StepStatus, + ) + from reflex.workflow.store import SqliteRunStore + + db = tmp_path / "attr.db" + now = time.time() + + async def seed() -> None: + """Admit one long-waiting run.""" + store = SqliteRunStore(db) + await store.admit( + RunRecord( + run_id="auditrun1", + workflow_id="cli.audit", + definition_digest="d", + status=RunStatus.WAITING, + state={}, + state_version=1, + next_ordinal=2, + created_at=now, + updated_at=now, + ), + StepRecord( + run_id="auditrun1", + ordinal=0, + handler_id="go", + status=StepStatus.READY, + args={}, + due_at=now + 86_400, + origin="root", + created_at=now, + updated_at=now, + ), + ((HistoryEventType.RUN_ADMITTED, {}),), + ) + store.close() + + asyncio.run(seed()) + result = subprocess.run( + [ + sys.executable, + "-c", + "from reflex.workflow.cli import workflows; workflows()", + "cancel", + "auditrun1", + "--reason", + "fat-fingered order", + "-d", + str(db), + ], + env={**__import__("os").environ, "REFLEX_ACTOR": "alek"}, + capture_output=True, + text=True, + timeout=120, + check=False, + ) + assert result.returncode == 0, result.stderr[-500:] + + connection = sqlite3.connect(db) + try: + rows = [ + jsonlib.loads(row[0]) + for row in connection.execute( + "SELECT data FROM workflow_history WHERE run_id = 'auditrun1'" + " AND type = ?", + (HistoryEventType.RUN_CANCEL_REQUESTED.value,), + ).fetchall() + ] + finally: + connection.close() + assert rows == [{"actor": "alek", "reason": "fat-fingered order"}] diff --git a/tests/units/workflow/test_serve.py b/tests/units/workflow/test_serve.py index 470b827611e..b9cc838ced5 100644 --- a/tests/units/workflow/test_serve.py +++ b/tests/units/workflow/test_serve.py @@ -523,3 +523,43 @@ def close(self, event): assert replayed.json()["disposition"] == "parked", "still no run to take it" missing = client.post("/deadletters/nope/replay", headers=_auth("tk_operate")) assert missing.status_code == 404 + + +def test_operator_actions_record_who_and_why(forked_registration_context): + """X-Actor and the body's reason land in the run's history. + + Args: + forked_registration_context: Isolated state registry. + """ + from reflex.workflow.records import HistoryEventType + + runtime = WorkflowRuntime(testing.MemoryRunStore()) + runtime.register(Orders) + app = build_app(runtime, worker=False, drain=0, tokens=_tokens(tk="all")) + with TestClient(app) as client: + started = client.post( + "/runs", + json={ + "workflow": "serve.orders", + "handler": "place", + "args": {"order_id": "o-audit"}, + }, + headers=_auth("tk"), + ) + run_id = started.json()["run_id"] + cancelled = client.post( + f"/runs/{run_id}/cancel", + json={"reason": "duplicate order"}, + headers={**_auth("tk"), "X-Actor": "ops@example.com"}, + ) + assert cancelled.status_code == 202 + assert client.portal is not None + events = client.portal.call( # pyright: ignore[reportAttributeAccessIssue] + runtime.kernel._store.get_history, # pyright: ignore[reportPrivateUsage] + run_id, + ) + cancel_event = next( + event for event in events if event.type is HistoryEventType.RUN_CANCEL_REQUESTED + ) + assert cancel_event.data["actor"] == "ops@example.com" + assert cancel_event.data["reason"] == "duplicate order"