Address review feedback on workflow_streams continue-as-new fix

- Hoist StreamDraining/TruncatedOffset error types to constants in
  _types.py so the raise sites (_stream.py) and handling sites
  (_client.py) cannot diverge.
- Annotate _describe_polled_run return type.
- Drop dead non_retryable=True flags: all three ApplicationErrors are
  raised from update validators/handlers, where retry policy does not
  apply, so the flag was a no-op.
- Document poll_cooldown semantics and warn against timedelta(0).
- Add test verifying the subscribe flow captures the polled run id.

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

Author: brianstrauch

temporalio/contrib/workflow_streams/_client.py

@@ -31,6 +31,7 @@
 from temporalio.api.common.v1 import Payload
 from temporalio.client import (
     Client,
+    WorkflowExecutionDescription,
     WorkflowExecutionStatus,
     WorkflowHandle,
     WorkflowUpdateFailedError,
@@ -42,6 +43,8 @@
 
 from ._topic_handle import TopicHandle
 from ._types import (
+    STREAM_DRAINING_ERROR_TYPE,
+    TRUNCATED_OFFSET_ERROR_TYPE,
     PollInput,
     PollResult,
     PublishEntry,
@@ -509,9 +512,11 @@ async def subscribe(
                 ``Payload`` — useful for heterogeneous topics where
                 the caller dispatches on ``Payload.metadata`` or wants
                 to forward the bytes without decoding.
-            poll_cooldown: Minimum interval between polls to avoid
-                overwhelming the workflow when items arrive faster
-                than the poll round-trip. Defaults to 100ms.
+            poll_cooldown: Minimum interval between polls when caught
+                up (backlogs always drain at full speed). Defaults to
+                100ms. Avoid ``timedelta(0)``: an idle subscriber
+                busy-loops, and each poll grows workflow history toward
+                its limit. Use 0 only in tests.
 
         Yields:
             :class:`WorkflowStreamItem` for each matching item.
@@ -549,14 +554,14 @@ async def subscribe(
                 return
             except WorkflowUpdateFailedError as e:
                 cause_type = getattr(e.cause, "type", None)
-                if cause_type == "TruncatedOffset":
+                if cause_type == TRUNCATED_OFFSET_ERROR_TYPE:
                     # Subscriber fell behind truncation. Retry from
                     # offset 0 which the stream treats as "from the
                     # beginning of whatever exists" (i.e., from
                     # base_offset).
                     offset = 0
                     continue
-                if cause_type == "StreamDraining":
+                if cause_type == STREAM_DRAINING_ERROR_TYPE:
                     # Workflow is detaching for continue-as-new. Back off and
                     # retry; the poll lands on the successor run once the
                     # rollover completes.
@@ -606,7 +611,7 @@ async def subscribe(
             if not result.more_ready and cooldown_secs > 0:
                 await asyncio.sleep(cooldown_secs)
 
-    async def _describe_polled_run(self):
+    async def _describe_polled_run(self) -> WorkflowExecutionDescription:
         """Describe the specific run the most recent poll was admitted to.
 
         Describing that run (rather than the latest) is what lets a

temporalio/contrib/workflow_streams/_stream.py

@@ -37,6 +37,8 @@
 
 from ._topic_handle import WorkflowTopicHandle
 from ._types import (
+    STREAM_DRAINING_ERROR_TYPE,
+    TRUNCATED_OFFSET_ERROR_TYPE,
     PollInput,
     PollResult,
     PublisherState,
@@ -357,7 +359,6 @@ def truncate(self, up_to_offset: int) -> None:
                 f"Cannot truncate to offset {up_to_offset}: "
                 f"valid range is [{self._base_offset}, {self._base_offset + len(self._log)})",
                 type="TruncateOutOfRange",
-                non_retryable=True,
             )
         self._log = self._log[log_index:]
         self._base_offset = up_to_offset
@@ -419,8 +420,7 @@ async def _on_poll(self, payload: PollInput) -> PollResult:
                 raise ApplicationError(
                     f"Requested offset {payload.from_offset} has been truncated. "
                     f"Current base offset is {self._base_offset}.",
-                    type="TruncatedOffset",
-                    non_retryable=True,
+                    type=TRUNCATED_OFFSET_ERROR_TYPE,
                 )
         all_new = self._log[log_offset:]
         if payload.topics:
@@ -469,8 +469,7 @@ def _validate_poll(self, _payload: PollInput) -> None:
         if self._detaching:
             raise ApplicationError(
                 "Workflow pollers are detached for continue-as-new",
-                type="StreamDraining",
-                non_retryable=True,
+                type=STREAM_DRAINING_ERROR_TYPE,
             )
 
     def _on_offset(self) -> int:

temporalio/contrib/workflow_streams/_types.py

@@ -26,6 +26,13 @@
 
 T = TypeVar("T")
 
+# Well-known ``ApplicationError.type`` values the stream workflow uses to reject
+# polls, and which ``WorkflowStreamClient.subscribe`` recognizes to drive retry
+# behavior. Defined here so the raise sites (``_stream.py``) and the handling
+# sites (``_client.py``) cannot diverge.
+STREAM_DRAINING_ERROR_TYPE = "StreamDraining"
+TRUNCATED_OFFSET_ERROR_TYPE = "TruncatedOffset"
+
 
 # basedpyright flags _-prefixed module-level functions as unused even when
 # sibling modules import them (_stream.py, _client.py). Vanilla pyright does

tests/contrib/workflow_streams/test_workflow_streams.py

@@ -2066,6 +2066,57 @@ async def test_follow_continue_as_new_describes_polled_run(client: Client) -> No
         await new_handle.signal(ContinueAsNewHelperWorkflow.close)
 
 
+@pytest.mark.asyncio
+async def test_subscribe_captures_polled_run_id(client: Client) -> None:
+    """The ``subscribe`` loop must record the run each poll was admitted to.
+
+    ``_follow_continue_as_new`` relies on ``_polled_run_id`` to describe the
+    *polled* run rather than the latest one (see
+    ``test_follow_continue_as_new_describes_polled_run``). This verifies the
+    real subscribe flow populates that field, rather than setting it by hand.
+    """
+    async with new_worker(client, BasicWorkflowStreamWorkflow) as worker:
+        handle = await client.start_workflow(
+            BasicWorkflowStreamWorkflow.run,
+            id=f"workflow-stream-polled-run-{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())
+        try:
+            await assert_eq_eventually(1, received_count)
+            # The poll was admitted to the run we started; subscribe must have
+            # captured exactly that run id, not left it unset.
+            assert stream._polled_run_id == handle.result_run_id
+        finally:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+            await handle.signal(BasicWorkflowStreamWorkflow.close)
+
+
 @workflow.defn
 class DrainingGateWorkflow:
     """CAN workflow that detaches pollers and then *holds* in the draining state