diff --git a/CHANGELOG.md b/CHANGELOG.md index 3dc178710..8d180fd7b 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/temporalio/contrib/workflow_streams/_client.py b/temporalio/contrib/workflow_streams/_client.py index 605bf3f03..6cfec7bb1 100644 --- a/temporalio/contrib/workflow_streams/_client.py +++ b/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 diff --git a/tests/contrib/workflow_streams/test_workflow_streams.py b/tests/contrib/workflow_streams/test_workflow_streams.py index e7cedd038..e3014b34e 100644 --- a/tests/contrib/workflow_streams/test_workflow_streams.py +++ b/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,10 +1168,8 @@ 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" @@ -1123,6 +1177,25 @@ async def subscribe_and_collect() -> None: 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."""