refactor(server): pass workflow transitions as commit_transition params

WorkflowTransition added a DTO between every caller and the registry while
carrying no behavior of its own. Replace it with keyword-only params on
commit_transition so each call site reads as a direct delta. While there,
_persist_terminal_locked now collects only terminal tasks, warns when a
non-terminal task reaches the path, and skips the commit for a workflow with
no terminal move.

Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

Author: kaiitunnz

src/server/registries/workflow.py

@@ -1,6 +1,5 @@
 import json
 from collections.abc import Sequence
-from dataclasses import dataclass
 from enum import StrEnum
 from typing import Any
 
@@ -60,24 +59,6 @@ class WorkflowSched(BaseModel):
     epoch_frontier: int = 0
 
 
-@dataclass(frozen=True, slots=True)
-class WorkflowTransition:
-    """A durable workflow state delta committed as one atomic Redis transaction.
-
-    ``records`` are upserted; each status field moves its task ids into that
-    status's set membership; ``sched`` snapshots the schedule when present.
-    """
-
-    workflow_id: str
-    records: Sequence[PersistedTask] = ()
-    dispatched: Sequence[str] = ()
-    pending: Sequence[str] = ()
-    done: Sequence[str] = ()
-    failed: Sequence[str] = ()
-    cancelled: Sequence[str] = ()
-    sched: WorkflowSched | None = None
-
-
 class WorkflowStatus(StrEnum):
     PENDING = "PENDING"
     DISPATCHED = "DISPATCHED"
@@ -278,36 +259,47 @@ async def get_workflow_async(self, workflow_id: str) -> Workflow | None:
             remaining_tasks,
         )
 
-    def commit_transition(self, transition: WorkflowTransition) -> None:
+    def commit_transition(
+        self,
+        workflow_id: str,
+        *,
+        records: Sequence[PersistedTask] = (),
+        dispatched: Sequence[str] = (),
+        pending: Sequence[str] = (),
+        done: Sequence[str] = (),
+        failed: Sequence[str] = (),
+        cancelled: Sequence[str] = (),
+        sched: WorkflowSched | None = None,
+    ) -> None:
         """Apply a workflow state delta as one atomic control-Redis transaction.
 
-        Task records, status-set membership moves, the workflow's ``updated_at``,
-        and the optional schedule snapshot commit together or not at all, so a
-        crash mid-persist can never leave durable state half-applied.
+        ``records`` are upserted; ``dispatched`` / ``pending`` / ``done`` /
+        ``failed`` / ``cancelled`` move their task ids into the matching status-set
+        membership; ``sched`` snapshots the schedule when present. The records,
+        membership moves, the workflow's ``updated_at``, and the schedule snapshot
+        commit together or not at all, so a crash mid-persist can never leave
+        durable state half-applied.
         """
-        wf = transition.workflow_id
-        leaving = (*transition.done, *transition.failed, *transition.cancelled)
-        touched_membership = bool(
-            transition.dispatched or transition.pending or leaving
-        )
+        terminal = (*done, *failed, *cancelled)
+        touched_membership = bool(dispatched or pending or terminal)
         with self._rds.sync.control_pipeline() as pipe:
-            for item in transition.records:
+            for item in records:
                 pipe.set(task_state_key(item.record.task_id), item.model_dump_json())
-            if transition.dispatched:
-                pipe.sadd(workflow_dispatched_tasks_key(wf), *transition.dispatched)
-            if transition.pending:
-                pipe.srem(workflow_dispatched_tasks_key(wf), *transition.pending)
-            if leaving:
-                pipe.srem(workflow_tasks_key(wf), *leaving)
-                pipe.srem(workflow_dispatched_tasks_key(wf), *leaving)
-            if transition.failed:
-                pipe.sadd(workflow_failed_tasks_key(wf), *transition.failed)
-            if transition.cancelled:
-                pipe.sadd(workflow_cancelled_tasks_key(wf), *transition.cancelled)
-            if touched_membership or transition.sched is not None:
-                pipe.hset(workflow_key(wf), mapping=_workflow_update())
-            if transition.sched is not None:
-                pipe.set(workflow_sched_key(wf), transition.sched.model_dump_json())
+            if dispatched:
+                pipe.sadd(workflow_dispatched_tasks_key(workflow_id), *dispatched)
+            if pending:
+                pipe.srem(workflow_dispatched_tasks_key(workflow_id), *pending)
+            if terminal:
+                pipe.srem(workflow_tasks_key(workflow_id), *terminal)
+                pipe.srem(workflow_dispatched_tasks_key(workflow_id), *terminal)
+            if failed:
+                pipe.sadd(workflow_failed_tasks_key(workflow_id), *failed)
+            if cancelled:
+                pipe.sadd(workflow_cancelled_tasks_key(workflow_id), *cancelled)
+            if touched_membership or sched is not None:
+                pipe.hset(workflow_key(workflow_id), mapping=_workflow_update())
+            if sched is not None:
+                pipe.set(workflow_sched_key(workflow_id), sched.model_dump_json())
             pipe.execute()
 
     # ---- Durable task state (for restart rehydration) ----------------- #

src/server/task/runtime.py

@@ -13,12 +13,7 @@
 
 from ..hooks import SUPPLIER_RESOLVERS
 from ..registries.worker import Worker, WorkerRegistry
-from ..registries.workflow import (
-    PersistedTask,
-    WorkflowRegistry,
-    WorkflowSched,
-    WorkflowTransition,
-)
+from ..registries.workflow import PersistedTask, WorkflowRegistry, WorkflowSched
 from ..utils.time import parse_iso_ts
 from .models import (
     TERMINAL_TASK_STATUSES,
@@ -393,19 +388,16 @@ def _persist_locked(self, *task_ids: str) -> None:
         """Commit task records (no membership change) atomically, per workflow."""
         by_workflow: dict[str, list[str]] = defaultdict(list)
         for task_id in dict.fromkeys(task_ids):
-            record = self._tasks.get(task_id)
-            if record is not None:
+            if record := self._tasks.get(task_id):
                 by_workflow[record.workflow_id].append(task_id)
         for workflow_id, ids in by_workflow.items():
             self._workflow_registry.commit_transition(
-                WorkflowTransition(
-                    workflow_id=workflow_id, records=self._records_locked(*ids)
-                )
+                workflow_id, records=self._records_locked(*ids)
             )
 
     def _persist_terminal_locked(self, *task_ids: str, sched: bool = True) -> None:
         """Commit each task's final state — its record and its done/failed/cancelled
-        set membership (by current status) — plus the workflow schedule, as one atomic
+        set membership (by current status) — and the workflow schedule, as one atomic
         transaction per workflow and the single last step of a transition.
 
         Committing only after all in-memory mutations means a failed or crashed write
@@ -416,36 +408,36 @@ def _persist_terminal_locked(self, *task_ids: str, sched: bool = True) -> None:
         ordered tasks carry ``position_in_epoch`` (so the ready-queue helpers never hit
         their guards).
         """
-        records: dict[str, list[str]] = defaultdict(list)
         moves: dict[str, tuple[list[str], list[str], list[str]]] = defaultdict(
             lambda: ([], [], [])
         )
         for task_id in dict.fromkeys(task_ids):
             record = self._tasks.get(task_id)
             if record is None:
                 continue
-            workflow_id = record.workflow_id
-            records[workflow_id].append(task_id)
-            done, failed, cancelled = moves[workflow_id]
             match record.status:
                 case TaskStatus.DONE:
-                    done.append(task_id)
+                    moves[record.workflow_id][0].append(task_id)
                 case TaskStatus.FAILED:
-                    failed.append(task_id)
+                    moves[record.workflow_id][1].append(task_id)
                 case TaskStatus.CANCELLED:
-                    cancelled.append(task_id)
-        for workflow_id, ids in records.items():
-            done, failed, cancelled = moves[workflow_id]
-            self._workflow_registry.commit_transition(
-                WorkflowTransition(
-                    workflow_id=workflow_id,
+                    moves[record.workflow_id][2].append(task_id)
+                case _:
+                    self._logger.warning(
+                        "Non-terminal task %s (%s) skipped in terminal persist",
+                        task_id,
+                        record.status,
+                    )
+        for workflow_id, (done, failed, cancelled) in moves.items():
+            if ids := done + failed + cancelled:
+                self._workflow_registry.commit_transition(
+                    workflow_id,
                     records=self._records_locked(*ids),
                     done=done,
                     failed=failed,
                     cancelled=cancelled,
                     sched=self._sched_locked(workflow_id) if sched else None,
                 )
-            )
 
     def _repersist_terminal_workflow_locked(self, workflow_id: str) -> None:
         """Re-commit the workflow's already-terminal tasks and schedule state.
@@ -467,9 +459,7 @@ def _repersist_terminal_workflow_locked(self, workflow_id: str) -> None:
             self._persist_terminal_locked(*terminal_ids)
         else:
             self._workflow_registry.commit_transition(
-                WorkflowTransition(
-                    workflow_id=workflow_id, sched=self._sched_locked(workflow_id)
-                )
+                workflow_id, sched=self._sched_locked(workflow_id)
             )
 
     # ------------------------------------------------------------------ #
@@ -698,11 +688,9 @@ def mark_pending(self, task_id: str, *, increment_retry: bool = False) -> None:
                 except Exception:
                     record.attempts = (record.attempts or 0) + 1
             self._workflow_registry.commit_transition(
-                WorkflowTransition(
-                    workflow_id=record.workflow_id,
-                    records=self._records_locked(task_id),
-                    pending=[task_id],
-                )
+                record.workflow_id,
+                records=self._records_locked(task_id),
+                pending=[task_id],
             )
 
     def requeue(self, task_id: str, *, front: bool = False) -> bool:
@@ -775,11 +763,9 @@ def _plan_merge_locked(
                 sibling_record.assigned_worker = None
                 sibling_record.merge_slice = None
         self._workflow_registry.commit_transition(
-            WorkflowTransition(
-                workflow_id=record.workflow_id,
-                records=self._records_locked(task_id, *siblings),
-                dispatched=siblings,
-            )
+            record.workflow_id,
+            records=self._records_locked(task_id, *siblings),
+            dispatched=siblings,
         )
 
         return siblings
@@ -935,11 +921,9 @@ def mark_dispatched(self, task_id: str, worker: Worker) -> None:
             self._remove_from_ready_locked(task_id)
             self._merge_bucket_remove(task_id)
             self._workflow_registry.commit_transition(
-                WorkflowTransition(
-                    workflow_id=record.workflow_id,
-                    records=self._records_locked(task_id),
-                    dispatched=[task_id],
-                )
+                record.workflow_id,
+                records=self._records_locked(task_id),
+                dispatched=[task_id],
             )
 
     def mark_started(
@@ -962,11 +946,9 @@ def mark_started(
             if worker_id:
                 record.assigned_worker = worker_id
             self._workflow_registry.commit_transition(
-                WorkflowTransition(
-                    workflow_id=record.workflow_id,
-                    records=self._records_locked(task_id),
-                    dispatched=[task_id],
-                )
+                record.workflow_id,
+                records=self._records_locked(task_id),
+                dispatched=[task_id],
             )
 
     def mark_updated(self, task_id: str, payload: dict[str, Any]) -> None:
@@ -1232,12 +1214,10 @@ def cancel_workflow(self, workflow_id: str, reason: str = "cancelled") -> list[s
             for task_id, _ in workflow_tasks:
                 self._task_epoch_index.pop(task_id, None)
             self._workflow_registry.commit_transition(
-                WorkflowTransition(
-                    workflow_id=workflow_id,
-                    records=self._records_locked(*touched),
-                    cancelled=cancelled,
-                    sched=self._sched_locked(workflow_id),
-                )
+                workflow_id,
+                records=self._records_locked(*touched),
+                cancelled=cancelled,
+                sched=self._sched_locked(workflow_id),
             )
 
         for interrupt in interrupts:

tests/server/dispatcher/helpers.py

@@ -31,7 +31,7 @@ class WorkflowRegistryStub:
     async def register_workflow_async(self, workflow_id: str, tasks: list[Any]) -> None:
         return None
 
-    def commit_transition(self, transition: Any) -> None: ...
+    def commit_transition(self, workflow_id: str, **kwargs: Any) -> None: ...
     async def save_task_states_async(self, items: Any) -> None: ...
     async def save_workflow_sched_async(
         self, wid: str, in_epoch_order: bool, frontier: int

tests/server/task/test_runtime_epoch_order.py

@@ -13,7 +13,7 @@ class _WorkflowRegistryStub:
     async def register_workflow_async(self, workflow_id: str, tasks: list[Any]) -> None:
         return None
 
-    def commit_transition(self, transition: Any) -> None:
+    def commit_transition(self, workflow_id: str, **kwargs: Any) -> None:
         return None
 
     async def save_task_states_async(self, items: Any) -> None: