fix(wac): checkpoint step errors so a caught exception does not hang replay (#10348)

* fix(wac): checkpoint step errors so a caught exception does not hang replay Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix(wac): honour a step suspend the workflow body caught and swallowed Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix(wac): park every suspend, not only those from a failing step Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * chore(wac): keep the generated bun wrapper comment-free Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix(wac): park the child task-completion suspend and align error identity Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * test(wac): pin the TaskError identity of replayed step and task failures Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-19 08:01:25 +00:00 · 2026-07-27 12:38:32 +02:00
parent 78e115bee5
commit be5e3bbfc4
5 changed files with 418 additions and 28 deletions
@@ -3,7 +3,7 @@
 import asyncio
 import pytest

-from wmill.client import WorkflowCtx, _StepSuspend, TaskError, workflow, task, step, sleep, parallel, wait_for_approval, _run_workflow
+from wmill.client import WorkflowCtx, _StepSuspend, TaskError, workflow, task, step, sleep, parallel, wait_for_approval, _run_workflow, _run_workflow_async


@task
@@ -914,6 +914,105 @@ class TestErrorPropagation:
        assert "step failed" in r["result"]["caught"]


+class TestRaisingInlineStepIsCheckpointed:
+    """A ``step()`` whose body raises must still land in ``completed_steps``.
+
+    Otherwise a workflow that catches the exception and later dispatches a task
+    replays with ``_executing_key`` set, reaches the unrecorded key, and parks
+    on the never-resolving future forever.
+    """
+
+    MARKER = {
+        "__wmill_error": True,
+        "message": "boom",
+        "step_key": "risky",
+        "result": {"error": "boom", "type": "ValueError"},
+    }
+
+    @staticmethod
+    def _boom():
+        raise ValueError("boom")
+
+    @classmethod
+    def _wf(cls):
+        @workflow
+        async def wf(x: int):
+            try:
+                await step("risky", cls._boom)
+            except Exception:
+                pass
+            return await double(x=x)
+
+        return wf
+
+    def test_first_run_emits_error_checkpoint(self):
+        r = _run_workflow(self._wf(), {}, {"x": 5})
+        assert r["type"] == "inline_checkpoint"
+        assert r["key"] == "risky"
+        assert r["result"] == self.MARKER
+
+    def test_fast_path_posts_error_and_raises_the_replay_exception(self, monkeypatch):
+        """The default path: the checkpoint is POSTed and the workflow body gets
+        the same ``TaskError`` a replay rebuilds from the marker — raising the
+        original ``ValueError`` here would make ``except ValueError:`` catch on
+        this run and miss on the next one."""
+        for var, val in (
+            ("WM_JOB_ID", "job-1"),
+            ("WM_WORKSPACE", "admins"),
+            ("BASE_INTERNAL_URL", "http://localhost:8000"),
+            ("WM_TOKEN", "tok"),
+        ):
+            monkeypatch.setenv(var, val)
+
+        posted = []
+
+        class _StubResponse:
+            def raise_for_status(self):
+                pass
+
+        class _StubClient:
+            async def post(self, url, json=None):
+                posted.append(json)
+                return _StubResponse()
+
+            async def aclose(self):
+                pass
+
+        async def run():
+            ctx = WorkflowCtx({})
+            ctx._inline_http_client = _StubClient()
+            with pytest.raises(TaskError, match="boom") as live:
+                await ctx._run_inline_step("risky", self._boom)
+            # ...and the replay of that very checkpoint raises the same thing.
+            replayed = WorkflowCtx({"completed_steps": {"risky": self.MARKER}})
+            with pytest.raises(TaskError, match="boom") as replay:
+                await replayed._run_inline_step("risky", self._boom)
+            assert type(live.value) is type(replay.value)
+            assert live.value.args == replay.value.args
+            assert live.value.result == replay.value.result == self.MARKER["result"]
+            assert isinstance(live.value.__cause__, ValueError)
+
+        asyncio.run(run())
+        assert len(posted) == 1
+        assert posted[0]["key"] == "risky"
+        assert posted[0]["result"] == self.MARKER
+
+    def test_replay_reraises_and_does_not_hang(self):
+        checkpoint = {
+            "completed_steps": {"risky": self.MARKER},
+            "_executing_key": "double",
+        }
+
+        async def run():
+            return await asyncio.wait_for(
+                _run_workflow_async(self._wf(), checkpoint, {"x": 5}), timeout=5
+            )
+
+        r = asyncio.run(run())
+        assert r["type"] == "complete"
+        assert r["result"] == 10
+
+
 # =====================================================================
 # TASK OPTIONS TESTS
 # =====================================================================
@@ -2689,6 +2689,29 @@ class TaskError(Exception):
        self.result = result


+def _step_error_marker(key: str, exc: BaseException) -> dict:
+    """Serialize a failed ``step()`` body into the ``__wmill_error`` marker that
+    task failures also use, so it can be stored in ``completed_steps``."""
+    return {
+        "__wmill_error": True,
+        "message": str(exc),
+        "step_key": key,
+        "result": {"error": str(exc), "type": type(exc).__name__},
+    }
+
+
+def _step_error_from_marker(marker: dict, name: str) -> TaskError:
+    """Rebuild the exception a failed step raises. Both the run that produced the
+    failure and every later replay go through here, so a workflow's ``except``
+    clauses see the same type either way."""
+    return TaskError(
+        marker.get("message", f"Step '{name}' failed"),
+        step_key=marker.get("step_key", ""),
+        child_job_id=marker.get("child_job_id", ""),
+        result=marker.get("result"),
+    )
+
+
 _workflow_ctx: _contextvars.ContextVar["WorkflowCtx"] = _contextvars.ContextVar(
    "_workflow_ctx"
 )
@@ -2858,12 +2881,7 @@ class WorkflowCtx:
        if key in self._completed:
            val = self._completed[key]
            if isinstance(val, dict) and val.get("__wmill_error"):
-                raise TaskError(
-                    val.get("message", f"Step '{name}' failed"),
-                    step_key=val.get("step_key", ""),
-                    child_job_id=val.get("child_job_id", ""),
-                    result=val.get("result"),
-                )
+                raise _step_error_from_marker(val, name)
            return val

        if self._executing_key is not None:
@@ -2873,9 +2891,18 @@ class WorkflowCtx:
        started_at = _dt.now(_tz.utc).isoformat()
        print(f"WM_WAC_STEP: {_json_mod.dumps({'key': key, 'started_at': started_at})}")
        t0 = _time_mod.monotonic()
-        result = fn()
-        if _asyncio.iscoroutine(result):
-            result = await result
+        # A raised step still has to reach ``completed_steps``, or a replay with
+        # ``_executing_key`` set finds nothing recorded and parks forever on the
+        # ``_asyncio.Future()`` above. ``_StepSuspend`` and ``CancelledError`` are
+        # ``BaseException``, so they pass through untouched.
+        step_error: Optional[Exception] = None
+        try:
+            result = fn()
+            if _asyncio.iscoroutine(result):
+                result = await result
+        except Exception as _exc:
+            step_error = _exc
+            result = _step_error_marker(key, _exc)
        duration_ms = int((_time_mod.monotonic() - t0) * 1000)

        # Fast path: POST the delta to the new per-job API endpoint and return
@@ -2892,6 +2919,7 @@ class WorkflowCtx:
        _base = os.environ.get("BASE_INTERNAL_URL")
        _token = os.environ.get("WM_TOKEN")
        if _fast_path_enabled and _job_id and _workspace and _base and _token:
+            _fast_path_ok = False
            try:
                if self._inline_lock is None:
                    self._inline_lock = _asyncio.Lock()
@@ -2917,7 +2945,7 @@ class WorkflowCtx:
                        },
                    )
                    _resp.raise_for_status()
-                return result
+                _fast_path_ok = True
            except Exception as _e:
                logger.info(
                    "WAC v2 inline fast path failed for key %s, falling back to suspend: %s",
@@ -2925,6 +2953,14 @@ class WorkflowCtx:
                    _e,
                )
                # fall through to the legacy suspend path
+            if _fast_path_ok:
+                # Raise what a replay would rebuild from the marker, never the
+                # original: a replay cannot reconstruct the original type, so
+                # raising it here would make ``except ValueError:`` catch on this
+                # run and miss on the next. ``__cause__`` is for tracebacks only.
+                if step_error is not None:
+                    raise _step_error_from_marker(result, name) from step_error
+                return result

        raise _StepSuspend({
            "mode": "inline_checkpoint",