Retry poll rejected while stream is draining for continue-as-new

The poll update's validator rejected new polls during detach-for-CAN with an
untyped RuntimeError, which subscribe() did not classify and re-raised — ending
the subscription with an error during a routine rollover.

Give the validator the well-known StreamDraining ApplicationError type and have
subscribe() back off and retry on it, so the poll lands on the successor run
once the rollover completes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: brianstrauch

temporalio/contrib/workflow_streams/_client.py

@@ -556,6 +556,14 @@ async def subscribe(
                     # base_offset).
                     offset = 0
                     continue
+                if cause_type == "StreamDraining":
+                    # Workflow is detaching for continue-as-new. Back off and
+                    # retry; the poll lands on the successor run once the
+                    # rollover completes.
+                    cooldown_secs = poll_cooldown.total_seconds()
+                    if cooldown_secs > 0:
+                        await asyncio.sleep(cooldown_secs)
+                    continue
                 if cause_type == "AcceptedUpdateCompletedWorkflow":
                     # Workflow returned (or continued-as-new) before
                     # this poll's update completed. Either follow the

temporalio/contrib/workflow_streams/_stream.py

@@ -460,9 +460,18 @@ async def _on_poll(self, payload: PollInput) -> PollResult:
         )
 
     def _validate_poll(self, _payload: PollInput) -> None:
-        """Reject new polls when pollers are detached for continue-as-new."""
+        """Reject new polls when pollers are detached for continue-as-new.
+
+        Uses the well-known ``StreamDraining`` type so a subscriber recognizes
+        the rollover-in-progress and retries until its poll lands on the
+        successor run, rather than surfacing the rejection as an error.
+        """
         if self._detaching:
-            raise RuntimeError("Workflow pollers are detached for continue-as-new")
+            raise ApplicationError(
+                "Workflow pollers are detached for continue-as-new",
+                type="StreamDraining",
+                non_retryable=True,
+            )
 
     def _on_offset(self) -> int:
         """Return the current global offset (base_offset + log length)."""

tests/contrib/workflow_streams/test_workflow_streams.py

@@ -2066,6 +2066,118 @@ async def test_follow_continue_as_new_describes_polled_run(client: Client) -> No
         await new_handle.signal(ContinueAsNewHelperWorkflow.close)
 
 
+@workflow.defn
+class DrainingGateWorkflow:
+    """CAN workflow that detaches pollers and then *holds* in the draining state
+    until released, so a subscriber deterministically hits the draining poll
+    rejection before the rollover completes."""
+
+    @workflow.init
+    def __init__(self, input: CANWorkflowInputTyped) -> None:
+        self.stream = WorkflowStream(prior_state=input.stream_state)
+        self._should_continue = False
+        self._release = False
+        self._closed = False
+
+    @workflow.signal
+    def close(self) -> None:
+        self._closed = True
+
+    @workflow.signal
+    def trigger_continue(self) -> None:
+        self._should_continue = True
+
+    @workflow.signal
+    def release(self) -> None:
+        self._release = True
+
+    @workflow.run
+    async def run(self, _input: CANWorkflowInputTyped) -> None:
+        del _input
+        await workflow.wait_condition(lambda: self._should_continue or self._closed)
+        if self._closed:
+            return
+        # Detach but stay open until released, so new polls are rejected with
+        # StreamDraining for a deterministic window.
+        self.stream.detach_pollers()
+        await workflow.wait_condition(lambda: self._release)
+        await workflow.wait_condition(workflow.all_handlers_finished)
+        workflow.continue_as_new(
+            args=[CANWorkflowInputTyped(stream_state=self.stream.get_state())]
+        )
+
+
+@pytest.mark.asyncio
+async def test_subscribe_retries_while_draining(client: Client) -> None:
+    """A poll rejected because the stream is draining for continue-as-new must
+    be retried, not surfaced as an error: the subscription stays alive through
+    the rollover and resumes on the successor run."""
+    async with new_worker(client, DrainingGateWorkflow) as worker:
+        handle = await client.start_workflow(
+            DrainingGateWorkflow.run,
+            CANWorkflowInputTyped(),
+            id=f"workflow-stream-draining-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+        await handle.signal(
+            "__temporal_workflow_stream_publish",
+            PublishInput(
+                items=[PublishEntry(topic="events", data=_wire_bytes(b"item-0"))],
+                publisher_id="pub",
+                sequence=1,
+            ),
+        )
+
+        stream = WorkflowStreamClient.create(client, handle.id)
+        received: list[WorkflowStreamItem] = []
+
+        async def consume() -> None:
+            async for item in stream.subscribe(
+                from_offset=0, poll_cooldown=timedelta(0), result_type=bytes
+            ):
+                received.append(item)
+
+        async def received_count() -> int:
+            return len(received)
+
+        task = asyncio.create_task(consume())
+        new_handle = client.get_workflow_handle(handle.id)
+        try:
+            await assert_eq_eventually(1, received_count)
+
+            # Detach; the subscriber's polls are now rejected with StreamDraining.
+            await handle.signal(DrainingGateWorkflow.trigger_continue)
+            # The subscription must keep retrying, not error out.
+            await asyncio.sleep(1.0)
+            assert not task.done(), "draining rejection must not end the subscription"
+
+            # Release: the workflow continues-as-new; the subscription resumes on
+            # the successor run and receives an item published there.
+            await handle.signal(DrainingGateWorkflow.release)
+            await assert_eq_eventually(
+                True, lambda: _is_different_run(handle, new_handle)
+            )
+            await new_handle.signal(
+                "__temporal_workflow_stream_publish",
+                PublishInput(
+                    items=[PublishEntry(topic="events", data=_wire_bytes(b"item-1"))],
+                    publisher_id="pub",
+                    sequence=2,
+                ),
+            )
+
+            await assert_eq_eventually(2, received_count)
+            assert [i.data for i in received] == [b"item-0", b"item-1"]
+            assert [i.offset for i in received] == [0, 1]
+        finally:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+            await new_handle.signal(DrainingGateWorkflow.close)
+
+
 # ---------------------------------------------------------------------------
 # Cross-workflow workflow stream (Scenario 1)
 # ---------------------------------------------------------------------------