Backport startup diagnostics to stable Python 2.5.8

docs/reference/worker.md

@@ -42,4 +42,16 @@ The worker-plane cleanup operation is separate from
 operator management action for retiring or recovering worker records and is not
 used by normal `Worker` shutdown.
 
+`run_until(timeout=...)` includes Server discovery and registration in its
+workflow wait budget. A timeout before startup finishes identifies the worker
+and workflow in the error. A timeout after startup reports that the workflow
+has not reached a terminal state. Neither timeout cancels the durable run.
+A replacement worker can continue that same run.
+
+When `stop()` is called during startup, it waits for registration to finish
+before starting the shared drain deadline. Server can record a registration
+before its reply reaches the SDK. A worker that never receives that reply has
+not begun polling and becomes stale under the Server's heartbeat policy.
+Use the existing workflow handle and queue when bringing up a replacement.
+
 ::: durable_workflow.worker

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "durable-workflow"
-version = "2.5.7"
+version = "2.5.8"
 description = "Python client and worker SDK for Durable Workflow Cloud and self-hosted Server"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -71,8 +71,8 @@ durable-workflow-replay-conformance = "durable_workflow.replay_conformance:main"
 durable-workflow-workflow-updates-conformance = "durable_workflow.workflow_updates_conformance:main"
 
 [tool.durable-workflow]
-product-train = "2.5.7"
-registry-version = "2.5.7"
+product-train = "2.5.8"
+registry-version = "2.5.8"
 supported-server-versions = "2.5.16"
 worker-protocol-version = "1.20"
 control-plane-version = "2"

src/durable_workflow/worker.py

@@ -4244,6 +4244,7 @@ async def run_until(
         """
         background_tasks: list[asyncio.Task[Any]] = []
         deadline = asyncio.get_running_loop().time() + timeout
+        startup_completed = False
 
         self._begin_run()
         try:
@@ -4252,6 +4253,7 @@ async def run_until(
             finally:
                 self._registration_done.set()
 
+            startup_completed = True
             if self._stop.is_set():
                 raise asyncio.CancelledError
             background_tasks.append(asyncio.create_task(self._heartbeat_loop()))
@@ -4273,6 +4275,11 @@ async def run_until(
                 run_until_loop, timeout=max(0.0, deadline - asyncio.get_running_loop().time()),
             )
         except asyncio.TimeoutError as error:
+            if not startup_completed:
+                raise TimeoutError(
+                    f"worker {self.worker_id} startup did not finish within the "
+                    f"{timeout}s workflow wait budget for {workflow_id}"
+                ) from error
             raise TimeoutError(f"workflow {workflow_id} not terminal after {timeout}s") from error
         finally:
             primary_error = sys.exc_info()[1]

tests/test_worker.py

@@ -5101,6 +5101,34 @@ async def test_run_until_returns_terminal_description(self, mock_client: AsyncMo
         assert mock_client.describe_workflow.await_count == 2
         assert mock_client.poll_workflow_task.call_args.kwargs["history_page_size"] == 500
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("blocked_method", ["get_cluster_info", "register_worker"])
+    async def test_run_until_identifies_startup_timeout_before_polling(
+        self, mock_client: AsyncMock, blocked_method: str
+    ) -> None:
+        started = asyncio.Event()
+
+        async def blocked(*args: object, **kwargs: object) -> None:
+            started.set()
+            await asyncio.Event().wait()
+
+        setattr(mock_client, blocked_method, AsyncMock(side_effect=blocked))
+        worker = Worker(mock_client, task_queue="q1", workflows=[TestWorkflow])
+
+        with pytest.raises(TimeoutError) as failure:
+            await worker.run_until(workflow_id="wf-startup-timeout", timeout=0.05)
+
+        assert started.is_set()
+        assert "startup" in str(failure.value)
+        assert worker.worker_id in str(failure.value)
+        assert "wf-startup-timeout" in str(failure.value)
+        assert isinstance(failure.value.__cause__, asyncio.TimeoutError)
+        assert worker._stop.is_set()
+        assert worker._registration_done.is_set()
+        mock_client.poll_workflow_task.assert_not_awaited()
+        mock_client.poll_activity_task.assert_not_awaited()
+        mock_client.deregister_worker_registration.assert_not_awaited()
+
     @pytest.mark.asyncio
     async def test_run_until_times_out_and_stops_worker(self, mock_client: AsyncMock) -> None:
         mock_client.describe_workflow = AsyncMock(
@@ -5119,9 +5147,11 @@ async def test_run_until_times_out_and_stops_worker(self, mock_client: AsyncMock
             poll_timeout=0.01,
         )
 
-        with pytest.raises(TimeoutError, match="wf-timeout"):
+        with pytest.raises(TimeoutError, match="wf-timeout") as failure:
             await worker.run_until(workflow_id="wf-timeout", timeout=0.02, poll_interval=0.01)
 
+        assert "not terminal" in str(failure.value)
+        assert "startup" not in str(failure.value)
         assert worker._stop.is_set()
 
     @pytest.mark.asyncio