From ce7826c315c7d0c218dc55aa580361aa5fa82766 Mon Sep 17 00:00:00 2001
From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com>
Date: Wed, 27 May 2026 12:02:10 +0530
Subject: [PATCH] fix: guard wait() against transient FAILED during saga
 startup

wait()'s poll safety-net could see FAILED before the saga
transitioned the run to COMPENSATING.  When the tracker event
has not yet fired, defer to the event instead of returning
the possibly-transient terminal state.

Also retry _mark_run_compensating on transient "database is
locked" errors so SQLite contention does not silently leave
the run stuck in FAILED.
---
 py_src/taskito/workflows/run.py               | 21 +++++++++++++++++++
 py_src/taskito/workflows/saga/orchestrator.py | 20 +++++++++++++-----
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/py_src/taskito/workflows/run.py b/py_src/taskito/workflows/run.py
index 2269938b..09b3002d 100644
--- a/py_src/taskito/workflows/run.py
+++ b/py_src/taskito/workflows/run.py
@@ -68,6 +68,27 @@ def wait(self, timeout: float | None = None, poll_interval: float = 0.1) -> Work
             while True:
                 snapshot = self.status()
                 if snapshot.state.is_terminal():
+                    # FAILED and COMPLETED_WITH_FAILURES can transition to
+                    # COMPENSATING when a saga is about to start.  Between
+                    # mark_workflow_node_result (sets FAILED in Rust) and
+                    # start_compensation (sets COMPENSATING from Python)
+                    # there is a window where the poll sees the transient
+                    # terminal state.  When the tracker owns an event for
+                    # this run, defer to the event — it is set only after
+                    # the saga decision is made — instead of returning the
+                    # possibly-transient state.
+                    if (
+                        snapshot.state
+                        in (WorkflowState.FAILED, WorkflowState.COMPLETED_WITH_FAILURES)
+                        and event is not None
+                        and not event.is_set()
+                    ):
+                        grace = poll_interval
+                        if deadline is not None:
+                            grace = min(grace, max(0.0, deadline - time.monotonic()))
+                        if grace > 0:
+                            event.wait(timeout=grace)
+                            continue
                     return snapshot
 
                 remaining: float
diff --git a/py_src/taskito/workflows/saga/orchestrator.py b/py_src/taskito/workflows/saga/orchestrator.py
index d6ef8aa1..3872d7f7 100644
--- a/py_src/taskito/workflows/saga/orchestrator.py
+++ b/py_src/taskito/workflows/saga/orchestrator.py
@@ -601,11 +601,21 @@ def _record_node_outcome(
         )
 
     def _mark_run_compensating(self, run_id: str) -> None:
-        """Move the run from ``Running`` to ``Compensating`` in storage."""
-        try:
-            self._queue._inner.set_workflow_run_compensating(run_id)
-        except Exception:
-            logger.exception("saga: failed to mark run %s compensating", run_id)
+        """Move the run to ``Compensating`` in storage.
+
+        Retries briefly on transient SQLite "database is locked" errors
+        so that contention from a concurrent result-handler write does
+        not silently leave the run in ``Failed``.
+        """
+        for attempt in range(5):
+            try:
+                self._queue._inner.set_workflow_run_compensating(run_id)
+                return
+            except Exception as exc:
+                if "database is locked" not in str(exc) or attempt == 4:
+                    logger.exception("saga: failed to mark run %s compensating", run_id)
+                    return
+                time.sleep(0.05 * (attempt + 1))
 
     def _finalize_locked(self, run_id: str, *, succeeded: bool) -> None:
         """Cleanly finalize a saga. Must be called with ``self._lock``."""