fix(wac): return the checkpointed value from step(), not the live object (#10367)

* fix(wac): return the checkpointed value from step(), not the live object Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * docs(wac): regenerate system prompts and narrow the round-trip claim Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * style(wac): condense the round-trip comments and fix the fallback note Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(sdk): type step() as the JSON round trip of its body's result Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(sdk): apply the JSON round trip to task() and the standalone paths Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(sdk): encode bigint, keep unknown as unknown, align dropped-key results Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(wac): null out results whose key JSON.stringify would drop Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(wac): normalize only the top-level result, keeping nested keys as they were Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(wac): normalize a child task's result so a deployed job cannot fail to parse Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * docs(sdk): pin non-finite number behavior in Jsonified and its tests Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(sdk): admit undefined for keys whose value JSON.stringify may omit Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(sdk): make a key JSON.stringify may omit optional, not just nullable Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(sdk): treat a class-valued property as dropped, like any other function Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-18 08:01:26 +00:00 · 2026-07-28 09:46:23 +02:00
parent b0c7e09173
commit 044ce39e5f
17 changed files with 695 additions and 65 deletions
@@ -1,11 +1,47 @@
 """Tests for the Workflow-as-Code SDK."""

 import asyncio
+import json
 import pytest

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


+class _StubInlineClient:
+    """Stands in for the httpx client the inline fast path POSTs with.
+
+    Decodes each request body, so a test sees exactly the JSON that reaches
+    ``/jobs/wac/inline_checkpoint`` — and therefore what a replay reads back.
+    """
+
+    def __init__(self):
+        self.posted = []
+
+    async def post(self, url, content=None):
+        self.posted.append(json.loads(content))
+
+        class _Response:
+            def raise_for_status(self):
+                pass
+
+        return _Response()
+
+    async def aclose(self):
+        pass
+
+
+def _set_inline_fast_path_env(monkeypatch):
+    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)
+
+
@task
 async def extract_data(url: str):
    pass  # body unused in workflow context
@@ -974,31 +1010,14 @@ class TestRaisingInlineStepIsCheckpointed:
        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)
+        _set_inline_fast_path_env(monkeypatch)

-        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
+        stub = _StubInlineClient()
+        posted = stub.posted

        async def run():
            ctx = WorkflowCtx({})
-            ctx._inline_http_client = _StubClient()
+            ctx._inline_http_client = stub
            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.
@@ -1031,6 +1050,52 @@ class TestRaisingInlineStepIsCheckpointed:
        assert r["result"] == 10


+class TestInlineStepRoundParity:
+    """The round that runs a ``step()`` body must see what a replay sees.
+
+    The fast path returns the value it checkpointed, not the in-memory one:
+    a workflow branching on a datetime attribute or a tuple would otherwise
+    take one path on the round that ran the body and another on every replay,
+    which can change which tasks get dispatched, not just crash later.
+    """
+
+    CASES = [
+        ("dt", lambda: datetime(2026, 1, 1, tzinfo=timezone.utc), "2026-01-01 00:00:00+00:00"),
+        ("pair", lambda: (1, 2), [1, 2]),
+        ("intkeys", lambda: {1: "a"}, {"1": "a"}),
+    ]
+
+    def test_outside_a_workflow_returns_the_same_shape(self):
+        """No checkpoint, no replay — but a local run must not hand back a shape
+        a deployed one never produces, or testing a workflow locally proves
+        nothing. The async task path is the sharp edge: the wrapper is sync, so
+        the value has to be round-tripped after the await, not before."""
+
+        @task
+        async def make_pair():
+            return (1, datetime(2026, 1, 1, tzinfo=timezone.utc))
+
+        assert asyncio.run(step("pair", lambda: (1, 2))) == [1, 2]
+        assert asyncio.run(make_pair()) == [1, "2026-01-01 00:00:00+00:00"]
+
+    def test_live_round_matches_checkpoint_and_replay(self, monkeypatch):
+        _set_inline_fast_path_env(monkeypatch)
+
+        async def run():
+            for key, fn, expected in self.CASES:
+                stub = _StubInlineClient()
+                ctx = WorkflowCtx({})
+                ctx._inline_http_client = stub
+                live = await ctx._run_inline_step(key, fn)
+                checkpointed = stub.posted[0]["result"]
+                assert checkpointed == expected
+                assert live == expected and type(live) is type(expected)
+                replayed = WorkflowCtx({"completed_steps": {key: checkpointed}})
+                assert await replayed._run_inline_step(key, fn) == live
+
+        asyncio.run(run())
+
+
 # =====================================================================
 # TASK OPTIONS TESTS
 # =====================================================================
@@ -2702,6 +2702,13 @@ class TaskError(Exception):
        self.result = result


+def _json_round_trip(value):
+    """Put a value through the checkpoint's encoding without checkpointing it, so
+    the paths that never persist anything still hand back the shape the ones that
+    do would. ``default=str`` matches the worker wrapper's encoder."""
+    return json.loads(json.dumps(value, default=str))
+
+
 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``."""
@@ -2936,7 +2943,22 @@ class WorkflowCtx:
        _token = os.environ.get("WM_TOKEN")
        if _fast_path_enabled and _job_id and _workspace and _base and _token:
            _fast_path_ok = False
+            _replay_result = None
            try:
+                # ``default=str`` is the encoder the worker wrapper uses on the
+                # suspend path, so both arms checkpoint the same value — and a
+                # datetime or set takes the fast path instead of silently
+                # degrading to a suspend round.
+                _payload = _json_mod.dumps(
+                    {
+                        "key": key,
+                        "result": result,
+                        "started_at": started_at,
+                        "duration_ms": duration_ms,
+                    },
+                    default=str,
+                )
+                _replay_result = _json_mod.loads(_payload)["result"]
                if self._inline_lock is None:
                    self._inline_lock = _asyncio.Lock()
                # Lock wraps only the POST, not fn() above — concurrent
@@ -2953,12 +2975,7 @@ class WorkflowCtx:
                        )
                    _resp = await self._inline_http_client.post(
                        f"{_base}/api/w/{_workspace}/jobs/wac/inline_checkpoint/{_job_id}",
-                        json={
-                            "key": key,
-                            "result": result,
-                            "started_at": started_at,
-                            "duration_ms": duration_ms,
-                        },
+                        content=_payload,
                    )
                    _resp.raise_for_status()
                _fast_path_ok = True
@@ -2976,7 +2993,11 @@ class WorkflowCtx:
                # 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
+                # Return the round trip of what was checkpointed, never the
+                # in-memory value: handing back the live object would let the
+                # round that ran the body branch on a type — tuple, datetime —
+                # that no replay of it ever sees.
+                return _replay_result

        raise _StepSuspend({
            "mode": "inline_checkpoint",
@@ -3009,6 +3030,10 @@ def task(
    - **v1 (WM_JOB_ID set, no @workflow)**: dispatches via HTTP API.
    - **Standalone**: executes the function body directly.

+    A task runs as its own job, so its result is always encoded as JSON and
+    decoded back before the caller sees it: a ``datetime`` comes back as a
+    string, a tuple as a list.
+
    Usage::

        @task
@@ -3081,8 +3106,19 @@ def task(
                print(f"Task {func.__name__} ({child_job_id}) completed")
                return job_result

-            # Standalone — execute directly
-            return func(*args, **kwargs)
+            # Standalone — execute directly, but round-trip the result: a task's
+            # value crosses JSON in every other path, so a local run must agree.
+            # This wrapper is sync, so an ``async def`` task hands back a
+            # coroutine here — round-tripping that would serialize the coroutine
+            # object itself.
+            result = func(*args, **kwargs)
+            if _asyncio.iscoroutine(result):
+
+                async def _round_trip_awaited():
+                    return _json_round_trip(await result)
+
+                return _round_trip_awaited()
+            return _json_round_trip(result)

        wrapper._is_task = True
        wrapper._task_path = task_path
@@ -3186,6 +3222,10 @@ async def step(name: str, fn):
    On replay the cached value is returned without re-executing ``fn``.
    Use for lightweight deterministic operations (timestamps, random IDs,
    config reads) that should not incur the overhead of a child job.
+
+    ``fn``'s result is encoded as JSON and decoded back before it is returned,
+    so the round that runs the body sees the same types every replay sees:
+    a ``datetime`` comes back as a string, a tuple as a list.
    """
    ctx: WorkflowCtx | None = _workflow_ctx.get(None)
    if ctx is not None:
@@ -3193,7 +3233,9 @@ async def step(name: str, fn):
    result = fn()
    if _asyncio.iscoroutine(result):
        result = await result
-    return result
+    # Outside a workflow nothing is checkpointed, but round-trip anyway: running
+    # the script locally must not hand back a shape a deployed run never sees.
+    return _json_round_trip(result)


 async def sleep(seconds: int):