fix: isolate async invoke processing context

statemachine/engines/base.py

@@ -39,8 +39,8 @@
 
 
 @contextmanager
-def _processing_step(engine: "BaseEngine") -> Iterator[None]:
-    """Mark the current context as running a step of ``engine``, until the block ends."""
+def _processing_step(engine: "BaseEngine | None") -> Iterator[None]:
+    """Set the current processing engine, or clear it, until the block ends."""
     token = _running_engine.set(engine)
     try:
         yield

statemachine/invoke.py

@@ -451,8 +451,12 @@ def _spawn_one_async(self, callback: "CallbackWrapper", **kwargs):
         self._active[ctx.invokeid] = invocation
         self._debug("%s invoke spawn async: %s on state %s", self._log_id, ctx.invokeid, state.id)
 
+        from .engines.base import _processing_step
+
         loop = asyncio.get_running_loop()
-        task = loop.create_task(self._run_async_handler(callback, handler, ctx, invocation))
+        # Invoke handlers run outside the macrostep that spawned them.
+        with _processing_step(None):
+            task = loop.create_task(self._run_async_handler(callback, handler, ctx, invocation))
         invocation.task = task
 
     async def _run_async_handler(

tests/test_invoke_context.py

@@ -0,0 +1,109 @@
+import asyncio
+from contextvars import ContextVar
+
+import pytest
+from statemachine.engines.base import _processing_step
+from statemachine.engines.base import _running_engine
+
+from statemachine import State
+from statemachine import StateChart
+
+
+@pytest.mark.parametrize("sm_runner", ["async"], indirect=True)
+async def test_async_invoke_waits_for_raised_event(sm_runner):
+    start = asyncio.Event()
+    attempted = asyncio.Event()
+    returned = asyncio.Event()
+    busy_entered = asyncio.Event()
+    release_busy = asyncio.Event()
+    release_worker = asyncio.Event()
+    request_id = ContextVar("request_id", default=None)
+    trace = []
+    contexts = []
+
+    async def worker(machine):
+        contexts.append(request_id.get())
+        await start.wait()
+        attempted.set()
+        await machine.raise_("tick")
+        contexts.append(request_id.get())
+        trace.append("invoke_returned")
+        returned.set()
+        await release_worker.wait()
+
+    class Machine(StateChart):
+        active = State(initial=True, invoke=worker)
+        busy = active.to(None)
+        tick = active.to(None)
+
+        async def on_busy(self):
+            busy_entered.set()
+            await release_busy.wait()
+            trace.append("busy_completed")
+
+        def on_tick(self):
+            trace.append("tick")
+
+    token = request_id.set("request")
+    machine = await sm_runner.start(Machine)
+    busy_task = asyncio.create_task(machine.send("busy"))
+    try:
+        await asyncio.wait_for(busy_entered.wait(), timeout=1)
+        start.set()
+        await asyncio.wait_for(attempted.wait(), timeout=1)
+        premature = returned.is_set()
+        release_busy.set()
+        await asyncio.wait_for(busy_task, timeout=1)
+        await asyncio.wait_for(returned.wait(), timeout=1)
+
+        assert not premature
+        assert trace == ["busy_completed", "tick", "invoke_returned"]
+        assert contexts == ["request", "request"]
+        assert request_id.get() == "request"
+        assert _running_engine.get() is None
+    finally:
+        release_busy.set()
+        release_worker.set()
+        tasks = [inv.task for inv in machine._engine._invoke_manager._active.values() if inv.task]
+        machine._engine.stop()
+        await asyncio.gather(busy_task, *tasks, return_exceptions=True)
+        request_id.reset(token)
+
+
+@pytest.mark.parametrize("sm_runner", ["async"], indirect=True)
+async def test_async_invoke_restores_context_when_task_factory_raises(sm_runner):
+    release = asyncio.Event()
+    observed = []
+
+    async def worker():
+        await release.wait()
+
+    class Machine(StateChart):
+        active = State(initial=True, invoke=worker)
+        tick = active.to(None)
+
+    machine = await sm_runner.start(Machine)
+    manager = machine._engine._invoke_manager
+    loop = asyncio.get_running_loop()
+    original_factory = loop.get_task_factory()
+
+    def failing_factory(loop, coroutine, **kwargs):
+        observed.append(_running_engine.get())
+        coroutine.close()
+        raise RuntimeError("cannot create invoke task")
+
+    try:
+        manager.mark_for_invoke(Machine.active)
+        loop.set_task_factory(failing_factory)
+        with _processing_step(machine._engine):
+            with pytest.raises(RuntimeError, match="cannot create invoke task"):
+                await manager.spawn_pending_async()
+            assert _running_engine.get() is machine._engine
+        assert _running_engine.get() is None
+        assert observed == [None]
+    finally:
+        loop.set_task_factory(original_factory)
+        release.set()
+        tasks = [inv.task for inv in manager._active.values() if inv.task]
+        machine._engine.stop()
+        await asyncio.gather(*tasks, return_exceptions=True)