Prevent workflow stream subscriptions from hiding cancellation

CHANGELOG.md

@@ -37,6 +37,8 @@ to include examples, links to docs, or any other relevant information.
 
 ### Fixed
 
+- `WorkflowStreamClient.subscribe` now propagates task cancellation instead of
+  ending the subscription normally.
 - `StrandsPlugin` now disables Botocore retries for its default Bedrock model so
   model request retries are handled exclusively by Temporal.
 - `temporalio.contrib.openai_agents` now honors the `retry-after-ms` and

temporalio/contrib/workflow_streams/_client.py

@@ -551,7 +551,7 @@ async def subscribe(
                 self._polled_run_id = handle.workflow_run_id
                 result: PollResult = await handle.result()
             except asyncio.CancelledError:
-                return
+                raise
             except WorkflowUpdateFailedError as e:
                 cause_type = getattr(e.cause, "type", None)
                 if cause_type == TRUNCATED_OFFSET_ERROR_TYPE:
@@ -577,7 +577,9 @@ async def subscribe(
                         continue
                     return
                 raise
-            except WorkflowUpdateRPCTimeoutOrCancelledError:
+            except WorkflowUpdateRPCTimeoutOrCancelledError as e:
+                if isinstance(e.__cause__, asyncio.CancelledError):
+                    raise e.__cause__
                 if await self._follow_continue_as_new():
                     continue
                 return

tests/contrib/workflow_streams/test_workflow_streams.py

@@ -48,7 +48,7 @@
 )
 from temporalio.contrib.workflow_streams._types import _encode_payload
 from temporalio.converter import DataConverter
-from temporalio.exceptions import ApplicationError
+from temporalio.exceptions import ActivityError, ApplicationError
 from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
@@ -88,6 +88,37 @@ async def run(self) -> None:
         await workflow.wait_condition(lambda: self._closed)
 
 
+@workflow.defn
+class CancelSubscriptionWorkflow:
+    @workflow.init
+    def __init__(self) -> None:
+        self.stream = WorkflowStream()
+        self._cancel_requested = False
+
+    @workflow.signal
+    def cancel_subscription(self) -> None:
+        self._cancel_requested = True
+
+    @workflow.run
+    async def run(self) -> str:
+        self.stream.topic("events", type=bytes).publish(b"seed")
+        handle = workflow.start_activity(
+            "subscribe_until_cancelled",
+            start_to_close_timeout=timedelta(seconds=30),
+            heartbeat_timeout=timedelta(seconds=1),
+            cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED,
+        )
+        await workflow.wait_condition(lambda: self._cancel_requested)
+        handle.cancel()
+        try:
+            result = await handle
+        except ActivityError as err:
+            result = type(err.cause).__name__
+        self.stream.detach_pollers()
+        await workflow.wait_condition(workflow.all_handlers_finished)
+        return result
+
+
 @workflow.defn
 class ActivityPublishWorkflow:
     @workflow.init
@@ -360,6 +391,31 @@ async def publish_items(count: int) -> None:
             client.topic("events", type=bytes).publish(f"item-{i}".encode())
 
 
+class CancellableSubscriber:
+    def __init__(self) -> None:
+        self.started = asyncio.Event()
+
+    @activity.defn(name="subscribe_until_cancelled")
+    async def subscribe(self) -> str:
+        async def heartbeat() -> None:
+            while True:
+                activity.heartbeat()
+                await asyncio.sleep(0.1)
+
+        heartbeat_task = asyncio.create_task(heartbeat())
+        try:
+            stream = WorkflowStreamClient.from_within_activity()
+            try:
+                async for _ in stream.subscribe(result_type=bytes):
+                    self.started.set()
+            except asyncio.CancelledError:
+                return "subscription-cancelled"
+            return "subscription-ended"
+        finally:
+            heartbeat_task.cancel()
+            await asyncio.gather(heartbeat_task, return_exceptions=True)
+
+
 @activity.defn(name="publish_multi_topic")
 async def publish_multi_topic(count: int) -> None:
     topics = ["a", "b", "c"]
@@ -1076,7 +1132,7 @@ async def test_priority_flush(client: Client) -> None:
 @pytest.mark.asyncio
 async def test_iterator_cancellation(client: Client) -> None:
     """Cancelling a subscription iterator after it has yielded an item
-    completes cleanly."""
+    propagates cancellation."""
     async with new_worker(
         client,
         BasicWorkflowStreamWorkflow,
@@ -1112,17 +1168,34 @@ async def subscribe_and_collect() -> None:
         async with _async_timeout(5):
             await first_item.wait()
         task.cancel()
-        try:
+        with pytest.raises(asyncio.CancelledError):
             await task
-        except asyncio.CancelledError:
-            pass
 
         assert len(items) == 1
         assert items[0].data == b"seed"
 
         await handle.signal(BasicWorkflowStreamWorkflow.close)
 
 
+@pytest.mark.asyncio
+async def test_activity_subscription_propagates_cancellation(client: Client) -> None:
+    subscriber = CancellableSubscriber()
+    async with new_worker(
+        client,
+        CancelSubscriptionWorkflow,
+        activities=[subscriber.subscribe],
+    ) as worker:
+        handle = await client.start_workflow(
+            CancelSubscriptionWorkflow.run,
+            id=f"workflow-stream-activity-cancel-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+        async with _async_timeout(5):
+            await subscriber.started.wait()
+        await handle.signal(CancelSubscriptionWorkflow.cancel_subscription)
+        assert await handle.result() == "subscription-cancelled"
+
+
 @pytest.mark.asyncio
 async def test_context_manager_flushes_on_exit(client: Client) -> None:
     """Context manager exit flushes all buffered items."""