From a613fac21ee7768f2e76392dde2d8c991cbb97e9 Mon Sep 17 00:00:00 2001
From: Durable Workflow <support@durable-workflow.com>
Date: Mon, 12 Oct 2026 05:45:01 +0000
Subject: [PATCH] Identify worker startup timeouts within the workflow wait
 budget

---
 docs/reference/worker.md       | 12 ++++++++++++
 src/durable_workflow/worker.py |  7 +++++++
 tests/test_worker.py           | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/docs/reference/worker.md b/docs/reference/worker.md
index 56230f3..964cea3 100644
--- a/docs/reference/worker.md
+++ b/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
diff --git a/src/durable_workflow/worker.py b/src/durable_workflow/worker.py
index dfee727..90faa31 100644
--- a/src/durable_workflow/worker.py
+++ b/src/durable_workflow/worker.py
@@ -4123,6 +4123,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:
@@ -4131,6 +4132,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()))
@@ -4152,6 +4154,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]
diff --git a/tests/test_worker.py b/tests/test_worker.py
index 43b6cbc..eb67847 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -5060,6 +5060,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(
@@ -5078,9 +5106,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