fix: run a workflow-as-code task called from inside a task in place (#11699)

* fix: run a workflow-as-code task called from inside a task in place

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: hand an async task body an awaitable for a sync nested task

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: mark the task-body scope per nested task body

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs: note the task-body behaviour of step and sleep

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

python-client/wmill/tests/test_workflow.py

@@ -7,7 +7,7 @@ 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
+from wmill.client import WorkflowCtx, _StepSuspend, TaskError, workflow, task, task_script, step, sleep, parallel, wait_for_approval, _run_workflow, _run_workflow_async
 
 
 class _FakeItems(dict):
@@ -752,6 +752,84 @@ class TestChildModeParallel:
         assert result["result"] == 8
 
 
+class TestSdkCallsInsideTaskBody:
+    """A task body runs as one child job, so the SDK calls it makes cannot be
+    dispatched as steps of their own."""
+
+    def _run_child(self, wf, key):
+        async def bounded():
+            return await asyncio.wait_for(
+                _run_workflow_async(wf, {"completed_steps": {}, "_executing_key": key}, {}),
+                timeout=2,
+            )
+
+        return asyncio.run(bounded())
+
+    def test_nested_task_and_step_run_in_the_tasks_own_job(self):
+        @task
+        async def inner(x: str):
+            return f"inner:{x}"
+
+        @task
+        def sync_leaf(x: str):
+            return f"sync:{x}"
+
+        # A sync body calls its tasks without awaiting, whatever called it.
+        @task
+        def sync_inner(x: str):
+            return sync_leaf(x).upper().lower()
+
+        @task
+        async def outer(x: str):
+            via_task = await inner(x)
+            via_sync = await sync_inner(x)
+            via_step = await step("s", lambda: "stepped")
+            return f"{via_task}+{via_sync}+{via_step}"
+
+        @workflow
+        async def wf():
+            return await outer("a")
+
+        result = self._run_child(wf, "outer")
+        assert result == {"type": "complete", "result": "inner:a+sync:a+stepped"}
+
+    def test_sibling_task_of_the_workflow_body_is_not_run(self):
+        # The workflow body keeps replaying beside the task body, and its own
+        # calls must stay parked: running them would execute a sibling task in
+        # this job.
+        ran = []
+
+        @task
+        async def inner(x: str):
+            ran.append(x)
+            return x
+
+        @task
+        async def outer(x: str):
+            await asyncio.sleep(0)
+            return await inner("nested")
+
+        @workflow
+        async def wf():
+            return await asyncio.gather(outer("a"), inner("sibling"))
+
+        result = self._run_child(wf, "outer")
+        assert result == {"type": "complete", "result": "nested"}
+        assert ran == ["nested"]
+
+    def test_step_that_would_need_a_dispatch_fails_instead_of_hanging(self):
+        @task
+        async def dispatching():
+            return await task_script("f/other/script")()
+
+        @workflow
+        async def wf():
+            return await dispatching()
+
+        with pytest.raises(RuntimeError, match="cannot be called from inside a task"):
+            self._run_child(wf, "dispatching")
+
+
 class TestKeyDeterminism:
     def test_keys_consistent_across_replays(self):
         r1 = _run_workflow(det_wf, {}, {"n": 3})

python-client/wmill/wmill/client.py

@@ -2766,6 +2766,7 @@ def parse_sql_client_name(name: str) -> tuple[str, Optional[str]]:
 
 import asyncio as _asyncio
 import contextvars as _contextvars
+import inspect as _inspect
 import sys as _sys
 import traceback as _traceback
 
@@ -2993,6 +2994,64 @@ _workflow_ctx: _contextvars.ContextVar["WorkflowCtx"] = _contextvars.ContextVar(
     "_workflow_ctx"
 )
 
+# Set while a task body runs in its child job: ``"async"`` or ``"sync"`` after
+# the kind of function the body is. A flag on the ctx cannot stand in: under
+# ``asyncio.gather`` the workflow body keeps replaying beside the task body, and
+# its calls must still park rather than run a sibling task's body in this job.
+_in_task_body: _contextvars.ContextVar[Optional[str]] = _contextvars.ContextVar(
+    "_in_task_body", default=None
+)
+
+
+def _dispatch_ctx() -> "WorkflowCtx | None":
+    """The ctx SDK calls dispatch through, or None where they run in place:
+    outside a workflow, and inside a task body, which is one job and has no
+    rounds of its own to dispatch or checkpoint a step in."""
+    if _in_task_body.get():
+        return None
+    return _workflow_ctx.get(None)
+
+
+async def _resolved(value):
+    return value
+
+
+def _call_in_task_body(func, caller: str, args, kwargs):
+    """Run a task called from inside a task body in place, round-tripping its
+    result as a dispatched one would be.
+
+    The workflow body awaits every task, sync ones included, so an async caller
+    gets an awaitable back either way; a sync caller gets the value. The scope
+    is re-marked for the nested body, which may be of the other kind."""
+
+    async def run_async(coro):
+        token = _in_task_body.set("async")
+        try:
+            return _json_round_trip(await coro)
+        finally:
+            _in_task_body.reset(token)
+
+    if _inspect.iscoroutinefunction(func):
+        return run_async(func(*args, **kwargs))
+    token = _in_task_body.set("sync")
+    try:
+        result = func(*args, **kwargs)
+    finally:
+        _in_task_body.reset(token)
+    if _asyncio.iscoroutine(result):
+        return run_async(result)
+    result = _json_round_trip(result)
+    return _resolved(result) if caller == "async" else result
+
+
+def _assert_not_inside_task_body(what: str) -> None:
+    if _in_task_body.get():
+        raise RuntimeError(
+            f"{what} cannot be called from inside a task: a task runs as a single "
+            "job and cannot suspend or dispatch steps. Call it from the @workflow "
+            "body instead."
+        )
+
 
 class WorkflowCtx:
     """Internal context for workflow replay/suspension.
@@ -3117,15 +3176,20 @@ class WorkflowCtx:
         raise _StepSuspend({"mode": "sleep", "key": key, "seconds": seconds, "steps": []})
 
     async def _resolved(self, value):
-        return value
+        return await _resolved(value)
 
     async def _execute_directly(self, func, **kwargs):
+        token = _in_task_body.set(
+            "async" if _inspect.iscoroutinefunction(func) else "sync"
+        )
         try:
             result = func(**kwargs)
             if _asyncio.iscoroutine(result):
                 result = await result
         except Exception as exc:
             raise _StepFailure(exc) from exc
+        finally:
+            _in_task_body.reset(token)
         raise _StepSuspend({"mode": "step_complete", "steps": [], "result": result})
 
     async def _never_resolve(self):
@@ -3367,7 +3431,9 @@ def task(
 
     - **v2 (inside @workflow)**: dispatches as a checkpoint step.
     - **v1 (WM_JOB_ID set, no @workflow)**: dispatches via HTTP API.
-    - **Standalone**: executes the function body directly.
+    - **Standalone, or called from inside another v2 task**: executes the
+      function body directly. A task runs as a single job, so the tasks it
+      calls run in that same job, without their own options.
 
     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
@@ -3447,12 +3513,16 @@ def task(
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
             # WAC v2: inside a @workflow context
-            ctx = _workflow_ctx.get(None)
+            ctx = _dispatch_ctx()
             if ctx is not None:
                 script = task_path if task_path else task_name
                 merged = _merge_args(args, kwargs)
                 return ctx._next_step(task_name, script, func, _task_options=_task_opts, **merged)
 
+            caller = _in_task_body.get()
+            if caller is not None:
+                return _call_in_task_body(func, caller, args, kwargs)
+
             # WAC v1: running inside a Windmill job but not in a @workflow
             if (
                 os.environ.get("WM_JOB_ID") is not None
@@ -3531,6 +3601,7 @@ def task_script(
     _opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s, "retry": _checked_retry(retry)}.items() if v is not None} or None
 
     def wrapper(**kwargs):
+        _assert_not_inside_task_body(f'task_script("{path}")')
         ctx = _workflow_ctx.get(None)
         if ctx is not None:
             return ctx._next_step(name, path, dispatch_type="script", _task_options=_opts, **kwargs)
@@ -3570,6 +3641,7 @@ def task_flow(
     _opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s, "retry": _checked_retry(retry)}.items() if v is not None} or None
 
     def wrapper(**kwargs):
+        _assert_not_inside_task_body(f'task_flow("{path}")')
         ctx = _workflow_ctx.get(None)
         if ctx is not None:
             return ctx._next_step(name, path, dispatch_type="flow", _task_options=_opts, **kwargs)
@@ -3604,8 +3676,11 @@ async def step(name: str, fn):
     ``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.
+
+    Called from inside a task body, ``fn`` runs in place and nothing is
+    checkpointed: a task is a single job and is never replayed.
     """
-    ctx: WorkflowCtx | None = _workflow_ctx.get(None)
+    ctx = _dispatch_ctx()
     if ctx is not None:
         return await ctx._run_inline_step(name, fn)
     result = fn()
@@ -3620,9 +3695,9 @@ async def sleep(seconds: int):
     """Server-side sleep — suspend the workflow for the given duration without holding a worker.
 
     Inside a @workflow, the parent job suspends and auto-resumes after ``seconds``.
-    Outside a workflow, falls back to ``asyncio.sleep``.
+    Outside a workflow, or inside a task body, falls back to ``asyncio.sleep``.
     """
-    ctx: WorkflowCtx | None = _workflow_ctx.get(None)
+    ctx = _dispatch_ctx()
     if ctx is not None:
         return await ctx._sleep(seconds)
     await _asyncio.sleep(seconds)
@@ -3660,6 +3735,7 @@ async def wait_for_approval(
         await step("notify", lambda: send_email(urls["resume"], urls["cancel"]))
         result = await wait_for_approval(key="manager", timeout=3600)
     """
+    _assert_not_inside_task_body("wait_for_approval()")
     ctx: WorkflowCtx | None = _workflow_ctx.get(None)
     if ctx is not None:
         return await ctx._wait_for_approval(