From 07cb9adfd0993bc36f70c4dd2b962e485f78fc7c Mon Sep 17 00:00:00 2001 From: Brian Strauch Date: Tue, 20 Oct 2026 14:28:05 -0700 Subject: [PATCH 1/2] Preserve Workflow Streams after workflow closure --- changelog/fixed/sleepy-otter-stream.md | 1 + temporalio/contrib/workflow_streams/README.md | 16 +- .../contrib/workflow_streams/_client.py | 189 +++---- .../contrib/workflow_streams/_stream.py | 36 +- temporalio/contrib/workflow_streams/_types.py | 9 +- .../workflow_streams/test_closed_streams.py | 518 ++++++++++++++++++ .../workflow_streams/test_workflow_streams.py | 109 ---- 7 files changed, 635 insertions(+), 243 deletions(-) create mode 100644 changelog/fixed/sleepy-otter-stream.md create mode 100644 tests/contrib/workflow_streams/test_closed_streams.py diff --git a/changelog/fixed/sleepy-otter-stream.md b/changelog/fixed/sleepy-otter-stream.md new file mode 100644 index 0000000000..2fdfc99d3d --- /dev/null +++ b/changelog/fixed/sleepy-otter-stream.md @@ -0,0 +1 @@ +Preserve retained Workflow Stream items after workflow completion, failure, cancellation, termination, or timeout by draining them through read-only queries before subscriptions finish. diff --git a/temporalio/contrib/workflow_streams/README.md b/temporalio/contrib/workflow_streams/README.md index 2fb4f94855..e8ca461e22 100644 --- a/temporalio/contrib/workflow_streams/README.md +++ b/temporalio/contrib/workflow_streams/README.md @@ -13,7 +13,8 @@ scales with durable batches rather than tokens. Under the hood the stream is built directly on Temporal's existing message-passing primitives: Signals carry publishes, Updates serve -long-poll subscriptions, and a Query exposes the current global offset. +long-poll subscriptions, and Queries expose retained items and the current +global offset. The library packages the boilerplate that turns those primitives into a usable stream: batching to amortize per-event overhead, deduplication for exactly-once delivery, topic filtering, and continue-as-new helpers @@ -31,3 +32,16 @@ snippets throughout. For runnable end-to-end examples, see the [Workflow Streams samples](https://github.com/temporalio/samples-python/tree/main/workflow-streams). + +Subscriptions to completed, failed, canceled, terminated, or timed-out runs +drain their remaining retained items through the +`__temporal_workflow_stream_poll` Query, then finish. Each Query targets the +confirmed terminal run, including when the Workflow ID is reused. Queries +return immediately, paginate within the same response size limit as Updates, +and start at the retained base when the requested offset has been truncated. + +Upgrade workers before enabling upgraded clients. Existing clients work with +upgraded workers; upgraded clients surface missing-query errors from older +workers. Closed-stream reads require retained history and an available +compatible worker. Query rejection policies and server/query errors remain +visible; truncated items and expired history cannot be recovered. diff --git a/temporalio/contrib/workflow_streams/_client.py b/temporalio/contrib/workflow_streams/_client.py index 8ed6cb5b51..7ad2a1dfe0 100644 --- a/temporalio/contrib/workflow_streams/_client.py +++ b/temporalio/contrib/workflow_streams/_client.py @@ -8,7 +8,7 @@ sync payload converter. The **codec chain** (e.g. encryption, compression) is **not** run per item — it runs once at the envelope level when Temporal's SDK encodes the ``__temporal_workflow_stream_publish`` -signal args and the ``__temporal_workflow_stream_poll`` update result. +signal args and the ``__temporal_workflow_stream_poll`` update/query result. Running the codec per item as well would double-encrypt / double-compress, because the envelope path covers the items again. The per-item ``Payload`` still carries the encoding metadata (``encoding: json/plain``, @@ -161,10 +161,6 @@ def __init__( self._pending_seq: int = 0 self._pending_since: float | None = None self._topic_types: dict[str, type[Any]] = {} - # Run id the most recent poll's update was admitted to. Captured before - # waiting for the outcome so a mid-poll continue-as-new can be detected by - # describing that specific run. None until the first poll is admitted. - self._polled_run_id: str | None = None @classmethod def create( @@ -529,7 +525,11 @@ async def subscribe( """Async iterator that polls for new items. Automatically follows continue-as-new chains when the client - was created via :py:meth:`create`. + was created via :py:meth:`create`. After the workflow completes, + fails, is canceled, terminated, or times out, queries deliver its + remaining retained items before the iterator finishes. These reads + require retained history and an available worker with the poll query + handler. Query failures, including configured query rejection, propagate. Args: topics: Topic filter. A single topic name, a list of topic @@ -570,65 +570,68 @@ async def subscribe( else: topic_filter = topics offset = from_offset + poll_handle = self._handle + terminal_handle: WorkflowHandle[Any, Any] | None = None while True: - try: - # Wait only for ACCEPTED so the handle (and the run id it was - # admitted to) is available before we block on the outcome; if - # the run continues-as-new mid-poll, result() fails but we still - # know which run to inspect. - handle = await self._handle.start_update( + if terminal_handle is not None: + result = await terminal_handle.query( "__temporal_workflow_stream_poll", PollInput(topics=topic_filter, from_offset=offset), - wait_for_stage=WorkflowUpdateStage.ACCEPTED, result_type=PollResult, ) - self._polled_run_id = handle.workflow_run_id - result: PollResult = await handle.result() - except asyncio.CancelledError: - raise - except WorkflowUpdateFailedError as e: - cause_type = getattr(e.cause, "type", None) - 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 == 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. - 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 - # chain or exit cleanly. - if await self._follow_continue_as_new(): + else: + # Each attempt tracks its own accepted run, including when the + # same stream client is shared by concurrent subscriptions. + polled_run_id: str | None = None + try: + handle = await poll_handle.start_update( + "__temporal_workflow_stream_poll", + PollInput(topics=topic_filter, from_offset=offset), + wait_for_stage=WorkflowUpdateStage.ACCEPTED, + result_type=PollResult, + ) + polled_run_id = handle.workflow_run_id + result = await handle.result() + except asyncio.CancelledError: + raise + except ( + WorkflowUpdateFailedError, + WorkflowUpdateRPCTimeoutOrCancelledError, + RPCError, + ) as err: + if isinstance(err, WorkflowUpdateFailedError): + cause_type = getattr(err.cause, "type", None) + if cause_type == TRUNCATED_OFFSET_ERROR_TYPE: + offset = 0 + continue + if cause_type == STREAM_DRAINING_ERROR_TYPE: + cooldown_secs = poll_cooldown.total_seconds() + if cooldown_secs > 0: + await asyncio.sleep(cooldown_secs) + continue + if isinstance(err.__cause__, asyncio.CancelledError): + raise err.__cause__ + desc = await self._describe_polled_run(poll_handle, polled_run_id) + if ( + desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW + and self._client is not None + ): + poll_handle = self._client.get_workflow_handle( + self._workflow_id + ) + continue + if desc.run_id and desc.status in ( + WorkflowExecutionStatus.COMPLETED, + WorkflowExecutionStatus.FAILED, + WorkflowExecutionStatus.CANCELED, + WorkflowExecutionStatus.TERMINATED, + WorkflowExecutionStatus.TIMED_OUT, + ): + terminal_handle = poll_handle._client.get_workflow_handle( + self._workflow_id, run_id=desc.run_id + ) continue - return - raise - except WorkflowUpdateRPCTimeoutOrCancelledError as e: - if isinstance(e.__cause__, asyncio.CancelledError): - raise e.__cause__ - if await self._follow_continue_as_new(): - continue - return - except RPCError as e: - # Workflow may have completed between polls; subscribe - # exits cleanly on terminal status so callers don't - # have to wrap the iterator in error handling for the - # normal end-of-stream case. - if e.status != RPCStatusCode.NOT_FOUND: raise - if await self._follow_continue_as_new(): - continue - if await self._workflow_in_terminal_state(): - return - raise converter = self._payload_converter() for wire_item in result.items: payload = _decode_payload(wire_item.data) @@ -643,61 +646,23 @@ async def subscribe( offset=wire_item.offset, ) offset = result.next_offset - cooldown_secs = poll_cooldown.total_seconds() - if not result.more_ready and cooldown_secs > 0: - await asyncio.sleep(cooldown_secs) - - 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 - continue-as-new be detected: a rolled-over run is closed with status - CONTINUED_AS_NEW, whereas the latest run would report RUNNING. Falls - back to the latest run when no run id has been captured yet, or when no - client is available to target a specific run. - """ - if self._client is not None: - return await self._client.get_workflow_handle( - self._workflow_id, run_id=self._polled_run_id + if terminal_handle is not None: + if not result.more_ready: + return + elif not result.more_ready: + cooldown_secs = poll_cooldown.total_seconds() + if cooldown_secs > 0: + await asyncio.sleep(cooldown_secs) + + async def _describe_polled_run( + self, handle: WorkflowHandle[Any, Any], polled_run_id: str | None + ) -> WorkflowExecutionDescription: + """Describe the accepted run, or the targeted execution before acceptance.""" + if polled_run_id is not None: + return await handle._client.get_workflow_handle( + self._workflow_id, run_id=polled_run_id ).describe() - return await self._handle.describe() - - async def _follow_continue_as_new(self) -> bool: - """Check if the polled run continued-as-new and re-target the handle. - - Returns True if the handle was updated (caller should retry). The - successor run id is not needed — re-targeting to an unpinned handle - makes the next poll address the latest (successor) run. - """ - if self._client is None: - return False - try: - desc = await self._describe_polled_run() - except Exception: - return False - if desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW: - self._handle = self._client.get_workflow_handle(self._workflow_id) - return True - return False - - async def _workflow_in_terminal_state(self) -> bool: - """Return True if the polled run has reached a terminal state. - - Used by ``subscribe()`` to distinguish "workflow finished — - stream is done" from "wrong workflow id" when a poll RPC - returns NOT_FOUND. - """ - try: - desc = await self._describe_polled_run() - except Exception: - return False - return desc.status in ( - WorkflowExecutionStatus.COMPLETED, - WorkflowExecutionStatus.FAILED, - WorkflowExecutionStatus.CANCELED, - WorkflowExecutionStatus.TERMINATED, - WorkflowExecutionStatus.TIMED_OUT, - ) + return await handle.describe() async def get_offset(self) -> int: """Query the current global offset (base_offset + log length).""" diff --git a/temporalio/contrib/workflow_streams/_stream.py b/temporalio/contrib/workflow_streams/_stream.py index ae8608c3b9..e2da3a130b 100644 --- a/temporalio/contrib/workflow_streams/_stream.py +++ b/temporalio/contrib/workflow_streams/_stream.py @@ -52,6 +52,7 @@ _PUBLISH_SIGNAL = "__temporal_workflow_stream_publish" _POLL_UPDATE = "__temporal_workflow_stream_poll" +_POLL_QUERY = "__temporal_workflow_stream_poll" _OFFSET_QUERY = "__temporal_workflow_stream_offset" _MAX_POLL_RESPONSE_BYTES = 1_000_000 @@ -84,6 +85,7 @@ class WorkflowStream: - ``__temporal_workflow_stream_publish`` signal — external publish with dedup - ``__temporal_workflow_stream_poll`` update — long-poll subscription + - ``__temporal_workflow_stream_poll`` query — immediate retained-item read - ``__temporal_workflow_stream_offset`` query — current log length Note: @@ -158,6 +160,7 @@ def __init__(self, prior_state: WorkflowStreamState | None = None) -> None: workflow.set_update_handler( _POLL_UPDATE, self._on_poll, validator=self._validate_poll ) + workflow.set_query_handler(_POLL_QUERY, self._collect_batch) workflow.set_query_handler(_OFFSET_QUERY, self._on_offset) def _publish_to_topic(self, topic: str, value: Any) -> None: @@ -337,9 +340,9 @@ async def continue_as_new( def truncate(self, up_to_offset: int) -> None: """Discard log entries before ``up_to_offset``. - After truncation, polls requesting an offset before the new - base will receive an ApplicationError. All global offsets - remain monotonic. + After truncation, poll updates requesting a nonzero offset before + the new base receive an ApplicationError. Poll queries read from + the retained base instead. All global offsets remain monotonic. Raises ApplicationError (not ValueError) when ``up_to_offset`` is past the end of the log so that callers invoking this from @@ -407,21 +410,18 @@ async def _on_poll(self, payload: PollInput) -> PollResult: or self._detaching ), ) - log_offset = payload.from_offset - self._base_offset - if log_offset < 0: - if payload.from_offset == 0: - # "From the beginning" — start at whatever is available. - log_offset = 0 - else: - # Subscriber had a specific position that's been - # truncated. ApplicationError fails this update (client - # gets the error) without crashing the workflow task — - # avoids a poison pill during replay. - raise ApplicationError( - f"Requested offset {payload.from_offset} has been truncated. " - f"Current base offset is {self._base_offset}.", - type=TRUNCATED_OFFSET_ERROR_TYPE, - ) + if payload.from_offset != 0 and payload.from_offset < self._base_offset: + # ApplicationError fails this update without failing the workflow task. + raise ApplicationError( + f"Requested offset {payload.from_offset} has been truncated. " + f"Current base offset is {self._base_offset}.", + type=TRUNCATED_OFFSET_ERROR_TYPE, + ) + return self._collect_batch(payload) + + def _collect_batch(self, payload: PollInput) -> PollResult: + """Read retained items immediately without changing stream state.""" + log_offset = max(0, payload.from_offset - self._base_offset) all_new = self._log[log_offset:] if payload.topics: topic_set = set(payload.topics) diff --git a/temporalio/contrib/workflow_streams/_types.py b/temporalio/contrib/workflow_streams/_types.py index a58cf75aeb..9bfcbe467b 100644 --- a/temporalio/contrib/workflow_streams/_types.py +++ b/temporalio/contrib/workflow_streams/_types.py @@ -5,7 +5,7 @@ ``Payload`` by the payload converter at publish time, and the resulting bytes/metadata are preserved per item so subscribers can decode with ``subscribe(result_type=T)``. The codec chain (e.g. encryption, compression) -applies once at the outer signal/update envelope level — not separately to each +applies once at the outer signal/update/query envelope level — not separately to each embedded item — so codec behavior is symmetric between workflow-side and client-side publishing. @@ -106,7 +106,10 @@ class PublishInput: @dataclass class PollInput: - """Update payload: request to poll for new items. + """Update/query payload: request to poll for new items. + + Updates wait for items and reject truncated nonzero offsets. Queries + return immediately and clamp offsets to the retained base. .. warning:: This class is experimental and may change in future versions. @@ -127,7 +130,7 @@ class _WorkflowStreamWireItem: @dataclass class PollResult: - """Update response: items matching the poll request. + """Update/query response: items matching the poll request. .. warning:: This class is experimental and may change in future versions. diff --git a/tests/contrib/workflow_streams/test_closed_streams.py b/tests/contrib/workflow_streams/test_closed_streams.py new file mode 100644 index 0000000000..d05224c333 --- /dev/null +++ b/tests/contrib/workflow_streams/test_closed_streams.py @@ -0,0 +1,518 @@ +"""Closed-run stream reads and subscription run isolation.""" + +from __future__ import annotations + +import asyncio +import json +import uuid +from dataclasses import dataclass +from datetime import timedelta +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from temporalio import workflow +from temporalio.client import ( + Client, + WorkflowExecutionStatus, + WorkflowHandle, + WorkflowQueryFailedError, + WorkflowQueryRejectedError, + WorkflowUpdateFailedError, +) +from temporalio.common import QueryRejectCondition, RawValue +from temporalio.contrib.workflow_streams import ( + PollInput, + PollResult, + WorkflowStream, + WorkflowStreamClient, +) +from temporalio.contrib.workflow_streams._types import ( + _decode_payload, + _WorkflowStreamWireItem, +) +from temporalio.converter import DataConverter +from temporalio.exceptions import ApplicationError +from temporalio.service import RPCError, RPCStatusCode +from tests.helpers import assert_eq_eventually, new_worker + +POLL = "__temporal_workflow_stream_poll" +WIRE_ABC = "ChgKCGVuY29kaW5nEgxiaW5hcnkvcGxhaW4SA2FiYw==" +TERMINAL_STATUSES = [ + WorkflowExecutionStatus.COMPLETED, + WorkflowExecutionStatus.FAILED, + WorkflowExecutionStatus.CANCELED, + WorkflowExecutionStatus.TERMINATED, + WorkflowExecutionStatus.TIMED_OUT, +] + + +def mock_stream( + status: WorkflowExecutionStatus = WorkflowExecutionStatus.COMPLETED, +) -> tuple[WorkflowStreamClient, Any, Any, Any]: + client = MagicMock(spec=Client) + client.data_converter = DataConverter.default + poll = MagicMock(spec=WorkflowHandle) + poll.id = "workflow" + poll._client = client + poll.start_update = AsyncMock( + side_effect=RPCError("closed", RPCStatusCode.NOT_FOUND, b"") + ) + poll.describe = AsyncMock( + return_value=SimpleNamespace(status=status, run_id="closed-run") + ) + read = MagicMock(spec=WorkflowHandle) + read.describe = poll.describe + read.query = AsyncMock(return_value=PollResult()) + client.get_workflow_handle.return_value = read + return WorkflowStreamClient(poll, client=client), client, poll, read + + +@pytest.mark.parametrize("status", TERMINAL_STATUSES) +async def test_workflow_streams_terminal_pages(status: WorkflowExecutionStatus) -> None: + stream, client, poll, read = mock_stream(status) + read.query.side_effect = [ + PollResult([_WorkflowStreamWireItem("events", WIRE_ABC, 3)], 4, True), + PollResult([_WorkflowStreamWireItem("events", WIRE_ABC, 5)], 6, False), + ] + items = [item async for item in stream.subscribe("events", 3, result_type=bytes)] + assert [(item.offset, item.data) for item in items] == [(3, b"abc"), (5, b"abc")] + client.get_workflow_handle.assert_called_once_with("workflow", run_id="closed-run") + assert poll.start_update.await_count == 1 + assert [call.args for call in read.query.await_args_list] == [ + (POLL, PollInput(["events"], 3)), + (POLL, PollInput(["events"], 4)), + ] + + +async def test_workflow_streams_direct_handle_keeps_default_converter() -> None: + _, client, poll, read = mock_stream() + read.query.return_value = PollResult( + [_WorkflowStreamWireItem("events", WIRE_ABC, 0)], 1 + ) + stream = WorkflowStreamClient(poll) + items = [item async for item in stream.subscribe(result_type=RawValue)] + assert items[0].data.payload.data == b"abc" + client.get_workflow_handle.assert_called_once_with("workflow", run_id="closed-run") + + +@pytest.mark.parametrize( + "failure", ["running", "unknown", "missing_run", "describe", "query"] +) +async def test_workflow_streams_unconfirmed_closure_and_query_errors( + failure: str, +) -> None: + stream, _, poll, read = mock_stream() + expected: Exception + if failure in ("running", "unknown", "missing_run"): + if failure == "missing_run": + poll.describe.return_value.run_id = "" + else: + poll.describe.return_value.status = ( + WorkflowExecutionStatus.RUNNING if failure == "running" else None + ) + expected = poll.start_update.side_effect + elif failure == "describe": + expected = RPCError("missing", RPCStatusCode.NOT_FOUND, b"") + poll.describe.side_effect = expected + else: + expected = WorkflowQueryFailedError("missing handler or query rejected") + read.query.side_effect = expected + with pytest.raises(type(expected)) as err: + _ = [item async for item in stream.subscribe()] + assert err.value is expected + + +async def test_workflow_streams_run_tracking_is_local_to_each_attempt() -> None: + stream, client, poll, _ = mock_stream() + admitted: asyncio.Queue[str] = asyncio.Queue() + outcomes: dict[str, asyncio.Future[PollResult]] = {} + reads: dict[str, Any] = {} + + async def start_update(*_args: Any, **_kwargs: Any) -> Any: + run_id = f"run-{len(outcomes)}" + future: asyncio.Future[PollResult] = asyncio.get_running_loop().create_future() + outcomes[run_id] = future + admitted.put_nowait(run_id) + return SimpleNamespace(workflow_run_id=run_id, result=lambda: future) + + def pinned_handle(workflow_id: str, *, run_id: str | None = None) -> Any: + assert workflow_id == "workflow" + assert run_id is not None + if run_id not in reads: + read = MagicMock(spec=WorkflowHandle) + read.describe = AsyncMock( + return_value=SimpleNamespace( + status=WorkflowExecutionStatus.COMPLETED, run_id=run_id + ) + ) + read.query = AsyncMock( + return_value=PollResult( + [_WorkflowStreamWireItem(run_id, WIRE_ABC, 0)], 1 + ) + ) + reads[run_id] = read + return reads[run_id] + + poll.start_update.side_effect = start_update + client.get_workflow_handle.side_effect = pinned_handle + + async def consume() -> list[str]: + return [item.topic async for item in stream.subscribe()] + + first = asyncio.create_task(consume()) + assert await admitted.get() == "run-0" + second = asyncio.create_task(consume()) + assert await admitted.get() == "run-1" + for run_id in ["run-1", "run-0"]: + outcomes[run_id].set_exception( + WorkflowUpdateFailedError( + ApplicationError("closed", type="AcceptedUpdateCompletedWorkflow") + ) + ) + assert await first == ["run-0"] + assert await second == ["run-1"] + + +async def test_workflow_streams_continue_as_new_resets_accepted_run() -> None: + stream, client, poll, read = mock_stream() + accepted = SimpleNamespace( + workflow_run_id="old-run", + result=AsyncMock( + side_effect=WorkflowUpdateFailedError( + ApplicationError("closed", type="AcceptedUpdateCompletedWorkflow") + ) + ), + ) + poll.start_update.side_effect = None + poll.start_update.return_value = accepted + old = MagicMock(spec=WorkflowHandle) + old.describe = AsyncMock( + return_value=SimpleNamespace( + status=WorkflowExecutionStatus.CONTINUED_AS_NEW, run_id="old-run" + ) + ) + successor = MagicMock(spec=WorkflowHandle) + successor._client = client + successor.start_update = AsyncMock( + side_effect=RPCError("closed", RPCStatusCode.NOT_FOUND, b"") + ) + successor.describe = AsyncMock( + return_value=SimpleNamespace( + status=WorkflowExecutionStatus.COMPLETED, run_id="successor-run" + ) + ) + read.query.return_value = PollResult( + [_WorkflowStreamWireItem("events", WIRE_ABC, 0)], 1 + ) + + def get_handle(workflow_id: str, *, run_id: str | None = None) -> Any: + assert workflow_id == "workflow" + return old if run_id == "old-run" else successor if run_id is None else read + + client.get_workflow_handle.side_effect = get_handle + assert [item.data async for item in stream.subscribe()] == [b"abc"] + old.describe.assert_awaited_once() + successor.describe.assert_awaited_once() + assert client.get_workflow_handle.call_args_list[-1].kwargs == { + "run_id": "successor-run" + } + + +async def test_workflow_streams_shared_wire_vector() -> None: + converter = DataConverter.default + request = '{"topics":["events"],"from_offset":3}' + response = ( + '{"items":[{"topic":"events","data":"' + + WIRE_ABC + + '","offset":3}],"next_offset":4,"more_ready":false}' + ) + for value, expected in [ + (PollInput(["events"], 3), request), + (PollResult([_WorkflowStreamWireItem("events", WIRE_ABC, 3)], 4), response), + ]: + payloads = await converter.encode([value]) + assert json.loads(payloads[0].data) == json.loads(expected) + decoded = (await converter.decode(payloads, [type(value)]))[0] + assert decoded == value + assert _decode_payload(WIRE_ABC).data == b"abc" + assert _decode_payload(WIRE_ABC).metadata == {"encoding": b"binary/plain"} + + +@dataclass +class ClosedStreamInput: + count: int = 3 + size: int = 1 + truncate: int = 0 + fail: bool = False + tail: bool = False + + +@workflow.defn +class ClosedStreamWorkflow: + @workflow.init + def __init__(self, arg: ClosedStreamInput) -> None: + self.stream = WorkflowStream() + self.closed = False + + @workflow.signal + def close(self) -> None: + self.closed = True + + @workflow.signal + def detach(self) -> None: + self.stream.detach_pollers() + + @workflow.run + async def run(self, arg: ClosedStreamInput) -> None: + for i in range(arg.count): + self.stream.topic("events" if i % 2 == 0 else "other").publish( + bytes([i]) * arg.size + ) + self.stream.truncate(arg.truncate) + await workflow.wait_condition(lambda: self.closed) + if arg.tail: + self.stream.topic("events").publish(b"tail") + if arg.fail: + raise ApplicationError("closed stream failure") + + +@pytest.mark.parametrize("status", TERMINAL_STATUSES) +async def test_workflow_streams_read_after_terminal_cold_replay( + client: Client, status: WorkflowExecutionStatus +) -> None: + async with new_worker( + client, ClosedStreamWorkflow, max_cached_workflows=0 + ) as worker: + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(fail=status == WorkflowExecutionStatus.FAILED), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + run_timeout=timedelta(seconds=2) + if status == WorkflowExecutionStatus.TIMED_OUT + else None, + ) + stream = WorkflowStreamClient(handle) + await assert_eq_eventually(3, stream.get_offset) + if status == WorkflowExecutionStatus.CANCELED: + await handle.cancel() + elif status == WorkflowExecutionStatus.TERMINATED: + await handle.terminate() + elif status != WorkflowExecutionStatus.TIMED_OUT: + await handle.signal(ClosedStreamWorkflow.close) + + async def execution_status() -> WorkflowExecutionStatus: + status = (await handle.describe()).status + assert status is not None + return status + + await assert_eq_eventually(status, execution_status) + before = await handle.fetch_history() + items = [item async for item in stream.subscribe(result_type=bytes)] + assert [item.data for item in items] == [b"\x00", b"\x01", b"\x02"] + assert [item.offset for item in items] == [0, 1, 2] + assert (await handle.fetch_history()).events == before.events + + +@pytest.mark.parametrize( + "arg,topics,offset,expected", + [ + (ClosedStreamInput(count=0), None, 0, []), + (ClosedStreamInput(), ["missing"], 0, []), + (ClosedStreamInput(), None, 1, [1, 2]), + (ClosedStreamInput(), None, 3, []), + (ClosedStreamInput(), None, 10, []), + (ClosedStreamInput(count=6, truncate=3), None, 1, [3, 4, 5]), + (ClosedStreamInput(count=8, size=200_000), None, 0, list(range(8))), + (ClosedStreamInput(count=8, size=200_000), ["events"], 0, [0, 2, 4, 6]), + ], +) +async def test_workflow_streams_closed_retention_and_pagination( + client: Client, + arg: ClosedStreamInput, + topics: list[str] | None, + offset: int, + expected: list[int], +) -> None: + async with new_worker(client, ClosedStreamWorkflow) as worker: + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + arg, + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.signal(ClosedStreamWorkflow.close) + await handle.result() + stream = WorkflowStreamClient.create(client, handle.id) + items = [ + item async for item in stream.subscribe(topics, offset, result_type=bytes) + ] + assert [item.offset for item in items] == expected + assert [item.data for item in items] == [ + bytes([i]) * arg.size for i in expected + ] + + +@pytest.mark.filterwarnings( + "ignore::temporalio.workflow.UnfinishedUpdateHandlersWarning" +) +async def test_workflow_streams_accepted_poll_drains_unread_tail( + client: Client, +) -> None: + async with new_worker(client, ClosedStreamWorkflow) as worker: + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(tail=True), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + stream = WorkflowStreamClient(handle) + await assert_eq_eventually(3, stream.get_offset) + accepted = asyncio.Event() + original_start = handle.start_update + + async def start(*args: Any, **kwargs: Any) -> Any: + update = await original_start(*args, **kwargs) + accepted.set() + return update + + async def consume() -> list[bytes]: + return [ + item.data + async for item in stream.subscribe( + from_offset=3, result_type=bytes, poll_cooldown=timedelta(0) + ) + ] + + with patch.object(handle, "start_update", side_effect=start): + task = asyncio.create_task(consume()) + await accepted.wait() + await handle.signal(ClosedStreamWorkflow.close) + assert await task == [b"tail"] + + +async def test_workflow_streams_query_is_immediate_while_draining( + client: Client, +) -> None: + async with new_worker(client, ClosedStreamWorkflow) as worker: + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(count=6, truncate=3), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.signal(ClosedStreamWorkflow.detach) + result = await handle.query(POLL, PollInput([], 1), result_type=PollResult) + assert [item.offset for item in result.items] == [3, 4, 5] + empty = await handle.query( + POLL, PollInput(["missing"], 6), result_type=PollResult + ) + assert empty == PollResult([], 6) + await handle.signal(ClosedStreamWorkflow.close) + await handle.result() + + +async def test_workflow_streams_closed_codec_and_raw_conversion(client: Client) -> None: + from tests.test_converter import SimpleCodec + + config = client.config() + config["data_converter"] = DataConverter(payload_codec=SimpleCodec()) + codec_client = Client(**config) + async with new_worker( + codec_client, ClosedStreamWorkflow, max_cached_workflows=0 + ) as worker: + handle = await codec_client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.signal(ClosedStreamWorkflow.close) + await handle.result() + stream = WorkflowStreamClient.create(codec_client, handle.id) + typed = [item.data async for item in stream.subscribe(result_type=bytes)] + raw = [ + item.data.payload.data + async for item in stream.subscribe(result_type=RawValue) + ] + assert typed == raw == [b"\x00", b"\x01", b"\x02"] + + +async def test_workflow_streams_closed_explicit_run_after_workflow_id_reuse( + client: Client, +) -> None: + async with new_worker(client, ClosedStreamWorkflow) as worker: + workflow_id = f"closed-stream-{uuid.uuid4()}" + old = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(), + id=workflow_id, + task_queue=worker.task_queue, + ) + await old.signal(ClosedStreamWorkflow.close) + await old.result() + new = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(count=1), + id=workflow_id, + task_queue=worker.task_queue, + ) + pinned = WorkflowStreamClient( + client.get_workflow_handle(workflow_id, run_id=old.result_run_id) + ) + + async def consume() -> list[int]: + return [item.offset async for item in pinned.subscribe()] + + assert await asyncio.gather(consume(), consume()) == [[0, 1, 2], [0, 1, 2]] + await new.signal(ClosedStreamWorkflow.close) + await new.result() + latest = WorkflowStreamClient.create(client, workflow_id) + assert [item.offset async for item in latest.subscribe()] == [0] + + +@workflow.defn +class LegacyClosedStreamWorkflow: + @workflow.run + async def run(self) -> None: + pass + + +async def test_workflow_streams_missing_query_and_configured_query_rejection( + client: Client, +) -> None: + async with new_worker( + client, ClosedStreamWorkflow, LegacyClosedStreamWorkflow + ) as worker: + legacy = await client.start_workflow( + LegacyClosedStreamWorkflow.run, + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await legacy.result() + with pytest.raises(WorkflowQueryFailedError): + _ = [item async for item in WorkflowStreamClient(legacy).subscribe()] + + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.signal(ClosedStreamWorkflow.close) + await handle.result() + config = client.config() + config["default_workflow_query_reject_condition"] = ( + QueryRejectCondition.NOT_OPEN + ) + stream = WorkflowStreamClient.create(Client(**config), handle.id) + with pytest.raises(WorkflowQueryRejectedError): + _ = [item async for item in stream.subscribe()] + + missing = WorkflowStreamClient.create(client, f"missing-{uuid.uuid4()}") + with pytest.raises(RPCError) as err: + _ = [item async for item in missing.subscribe()] + assert err.value.status == RPCStatusCode.NOT_FOUND diff --git a/tests/contrib/workflow_streams/test_workflow_streams.py b/tests/contrib/workflow_streams/test_workflow_streams.py index 79d0fa5409..65dbc454fa 100644 --- a/tests/contrib/workflow_streams/test_workflow_streams.py +++ b/tests/contrib/workflow_streams/test_workflow_streams.py @@ -30,7 +30,6 @@ from temporalio import activity, nexus, workflow from temporalio.client import ( Client, - WorkflowExecutionStatus, WorkflowHandle, WorkflowUpdateFailedError, WorkflowUpdateStage, @@ -2185,114 +2184,6 @@ async def test_continue_as_new_helper(client: Client) -> None: await new_handle.signal(ContinueAsNewHelperWorkflow.close) -@pytest.mark.asyncio -async def test_follow_continue_as_new_describes_polled_run(client: Client) -> None: - """Regression test for continue-as-new detection. - - ``_follow_continue_as_new`` must describe the *specific run the poll was - admitted to* — a rolled-over run is closed with status CONTINUED_AS_NEW, - whereas the latest (successor) run reports RUNNING. The previous - implementation described the latest run, so the check never fired and a poll - failure during a rollover stopped the subscription instead of following it. - - Driving the exact poll-failure race deterministically is impractical (the - workflow drains in-flight polls before continuing-as-new), so this asserts - the helper's decision directly against a real post-rollover run. - """ - async with new_worker(client, ContinueAsNewHelperWorkflow) as worker: - handle = await client.start_workflow( - ContinueAsNewHelperWorkflow.run, - CANWorkflowInputTyped(), - id=f"workflow-stream-can-follow-{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, - ), - ) - old_run_id = handle.result_run_id - - await handle.signal(ContinueAsNewHelperWorkflow.trigger_continue) - new_handle = client.get_workflow_handle(handle.id) - await assert_eq_eventually(True, lambda: _is_different_run(handle, new_handle)) - - # The fix's premise: the polled (old) run reports CONTINUED_AS_NEW; the - # latest run reports RUNNING. - old_desc = await client.get_workflow_handle( - handle.id, run_id=old_run_id - ).describe() - assert old_desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW - latest_desc = await client.get_workflow_handle(handle.id).describe() - assert latest_desc.status == WorkflowExecutionStatus.RUNNING - - # The client follows the rollover when it describes the polled run, but - # the previous latest-run behavior (polled_run_id unset) would not. - following = WorkflowStreamClient.create(client, handle.id) - following._polled_run_id = old_run_id - assert await following._follow_continue_as_new() is True - - latest_only = WorkflowStreamClient.create(client, handle.id) - latest_only._polled_run_id = None # describes the latest run, as the bug did - assert await latest_only._follow_continue_as_new() is False - - 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 From dcf32a9668a0a0f439e24087a4725ba224fdde69 Mon Sep 17 00:00:00 2001 From: Brian Strauch Date: Wed, 28 Oct 2026 09:32:13 -0700 Subject: [PATCH 2/2] Address Workflow Streams review feedback --- changelog/fixed/sleepy-otter-stream.md | 1 + .../contrib/workflow_streams/_client.py | 11 +- .../contrib/workflow_streams/_stream.py | 10 +- .../workflow_streams/test_closed_streams.py | 116 ++++++++++++++++-- 4 files changed, 124 insertions(+), 14 deletions(-) diff --git a/changelog/fixed/sleepy-otter-stream.md b/changelog/fixed/sleepy-otter-stream.md index 2fdfc99d3d..82b0957a70 100644 --- a/changelog/fixed/sleepy-otter-stream.md +++ b/changelog/fixed/sleepy-otter-stream.md @@ -1 +1,2 @@ Preserve retained Workflow Stream items after workflow completion, failure, cancellation, termination, or timeout by draining them through read-only queries before subscriptions finish. +Preserve the original Workflow Streams poll error when the follow-up execution description fails with an RPC error. diff --git a/temporalio/contrib/workflow_streams/_client.py b/temporalio/contrib/workflow_streams/_client.py index 7ad2a1dfe0..791151e6bc 100644 --- a/temporalio/contrib/workflow_streams/_client.py +++ b/temporalio/contrib/workflow_streams/_client.py @@ -530,6 +530,8 @@ async def subscribe( remaining retained items before the iterator finishes. These reads require retained history and an available worker with the poll query handler. Query failures, including configured query rejection, propagate. + If describing the execution fails with an RPC error, the original poll + error propagates. Args: topics: Topic filter. A single topic name, a list of topic @@ -611,7 +613,12 @@ async def subscribe( continue if isinstance(err.__cause__, asyncio.CancelledError): raise err.__cause__ - desc = await self._describe_polled_run(poll_handle, polled_run_id) + try: + desc = await self._describe_polled_run( + poll_handle, polled_run_id + ) + except RPCError: + raise err if ( desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW and self._client is not None @@ -627,6 +634,8 @@ async def subscribe( WorkflowExecutionStatus.TERMINATED, WorkflowExecutionStatus.TIMED_OUT, ): + # Use the handle's client so direct-handle subscriptions + # can pin terminal reads without an explicit client. terminal_handle = poll_handle._client.get_workflow_handle( self._workflow_id, run_id=desc.run_id ) diff --git a/temporalio/contrib/workflow_streams/_stream.py b/temporalio/contrib/workflow_streams/_stream.py index e2da3a130b..b018f86a48 100644 --- a/temporalio/contrib/workflow_streams/_stream.py +++ b/temporalio/contrib/workflow_streams/_stream.py @@ -51,8 +51,8 @@ ) _PUBLISH_SIGNAL = "__temporal_workflow_stream_publish" -_POLL_UPDATE = "__temporal_workflow_stream_poll" -_POLL_QUERY = "__temporal_workflow_stream_poll" +# Updates and queries share a name because they implement the same poll protocol. +_POLL = "__temporal_workflow_stream_poll" _OFFSET_QUERY = "__temporal_workflow_stream_offset" _MAX_POLL_RESPONSE_BYTES = 1_000_000 @@ -157,10 +157,8 @@ def __init__(self, prior_state: WorkflowStreamState | None = None) -> None: self._topic_types: dict[str, type[Any]] = {} workflow.set_signal_handler(_PUBLISH_SIGNAL, self._on_publish) - workflow.set_update_handler( - _POLL_UPDATE, self._on_poll, validator=self._validate_poll - ) - workflow.set_query_handler(_POLL_QUERY, self._collect_batch) + workflow.set_update_handler(_POLL, self._on_poll, validator=self._validate_poll) + workflow.set_query_handler(_POLL, self._collect_batch) workflow.set_query_handler(_OFFSET_QUERY, self._on_offset) def _publish_to_topic(self, topic: str, value: Any) -> None: diff --git a/tests/contrib/workflow_streams/test_closed_streams.py b/tests/contrib/workflow_streams/test_closed_streams.py index d05224c333..d032bc2e12 100644 --- a/tests/contrib/workflow_streams/test_closed_streams.py +++ b/tests/contrib/workflow_streams/test_closed_streams.py @@ -21,6 +21,7 @@ WorkflowQueryFailedError, WorkflowQueryRejectedError, WorkflowUpdateFailedError, + WorkflowUpdateRPCTimeoutOrCancelledError, ) from temporalio.common import QueryRejectCondition, RawValue from temporalio.contrib.workflow_streams import ( @@ -98,6 +99,62 @@ async def test_workflow_streams_direct_handle_keeps_default_converter() -> None: client.get_workflow_handle.assert_called_once_with("workflow", run_id="closed-run") +async def test_workflow_streams_update_to_query_resumes_at_next_offset() -> None: + stream, client, poll, read = mock_stream() + poll.start_update.side_effect = [ + SimpleNamespace( + workflow_run_id="closed-run", + result=AsyncMock( + return_value=PollResult( + [_WorkflowStreamWireItem("events", WIRE_ABC, i) for i in (0, 1)], + 2, + True, + ) + ), + ), + SimpleNamespace( + workflow_run_id="closed-run", + result=AsyncMock( + return_value=PollResult( + [_WorkflowStreamWireItem("events", WIRE_ABC, 2)], 3 + ) + ), + ), + SimpleNamespace( + workflow_run_id="closed-run", + result=AsyncMock( + side_effect=WorkflowUpdateFailedError( + ApplicationError("closed", type="AcceptedUpdateCompletedWorkflow") + ) + ), + ), + ] + read.query.side_effect = [ + PollResult([_WorkflowStreamWireItem("events", WIRE_ABC, 3)], 4, True), + PollResult([_WorkflowStreamWireItem("events", WIRE_ABC, 4)], 5), + ] + items = [ + item + async for item in stream.subscribe( + "events", result_type=bytes, poll_cooldown=timedelta(0) + ) + ] + assert [(item.offset, item.data) for item in items] == [ + (i, b"abc") for i in range(5) + ] + assert [call.args for call in poll.start_update.await_args_list] == [ + (POLL, PollInput(["events"], offset)) for offset in (0, 2, 3) + ] + assert [call.args for call in read.query.await_args_list] == [ + (POLL, PollInput(["events"], offset)) for offset in (3, 4) + ] + read.describe.assert_awaited_once() + assert client.get_workflow_handle.call_count == 2 + for call in client.get_workflow_handle.call_args_list: + assert call.args == ("workflow",) + assert call.kwargs == {"run_id": "closed-run"} + + @pytest.mark.parametrize( "failure", ["running", "unknown", "missing_run", "describe", "query"] ) @@ -115,8 +172,8 @@ async def test_workflow_streams_unconfirmed_closure_and_query_errors( ) expected = poll.start_update.side_effect elif failure == "describe": - expected = RPCError("missing", RPCStatusCode.NOT_FOUND, b"") - poll.describe.side_effect = expected + poll.describe.side_effect = RPCError("missing", RPCStatusCode.NOT_FOUND, b"") + expected = poll.start_update.side_effect else: expected = WorkflowQueryFailedError("missing handler or query rejected") read.query.side_effect = expected @@ -125,6 +182,45 @@ async def test_workflow_streams_unconfirmed_closure_and_query_errors( assert err.value is expected +@pytest.mark.parametrize("accepted", [False, True]) +@pytest.mark.parametrize("failure", ["rpc", "update", "timeout"]) +async def test_workflow_streams_describe_failure_preserves_poll_error( + accepted: bool, failure: str +) -> None: + stream, _, poll, read = mock_stream() + expected: Exception + if failure == "rpc": + expected = RPCError("poll failed", RPCStatusCode.UNAVAILABLE, b"") + elif failure == "update": + expected = WorkflowUpdateFailedError(ApplicationError("poll failed")) + else: + expected = WorkflowUpdateRPCTimeoutOrCancelledError() + cause = expected.__cause__ + poll.describe.side_effect = RPCError( + "describe failed", RPCStatusCode.NOT_FOUND, b"" + ) + if accepted: + poll.start_update.side_effect = None + poll.start_update.return_value = SimpleNamespace( + workflow_run_id="closed-run", result=AsyncMock(side_effect=expected) + ) + else: + poll.start_update.side_effect = expected + with pytest.raises(type(expected)) as err: + _ = [item async for item in stream.subscribe()] + assert err.value is expected + assert err.value.__cause__ is cause + read.query.assert_not_awaited() + + +async def test_workflow_streams_describe_cancellation_propagates() -> None: + stream, _, poll, read = mock_stream() + poll.describe.side_effect = asyncio.CancelledError + with pytest.raises(asyncio.CancelledError): + _ = [item async for item in stream.subscribe()] + read.query.assert_not_awaited() + + async def test_workflow_streams_run_tracking_is_local_to_each_attempt() -> None: stream, client, poll, _ = mock_stream() admitted: asyncio.Queue[str] = asyncio.Queue() @@ -376,14 +472,15 @@ async def test_workflow_streams_accepted_poll_drains_unread_tail( async def start(*args: Any, **kwargs: Any) -> Any: update = await original_start(*args, **kwargs) - accepted.set() + if args[1].from_offset == 3: + accepted.set() return update - async def consume() -> list[bytes]: + async def consume() -> list[tuple[int, bytes]]: return [ - item.data + (item.offset, item.data) async for item in stream.subscribe( - from_offset=3, result_type=bytes, poll_cooldown=timedelta(0) + result_type=bytes, poll_cooldown=timedelta(0) ) ] @@ -391,7 +488,12 @@ async def consume() -> list[bytes]: task = asyncio.create_task(consume()) await accepted.wait() await handle.signal(ClosedStreamWorkflow.close) - assert await task == [b"tail"] + assert await task == [ + (0, b"\x00"), + (1, b"\x01"), + (2, b"\x02"), + (3, b"tail"), + ] async def test_workflow_streams_query_is_immediate_while_draining(