From b0584be3e3d62fac7e114410d56e489604b3806c Mon Sep 17 00:00:00 2001 From: DABH Date: Tue, 14 Jul 2026 00:39:18 -0500 Subject: [PATCH] Add a public workflow-side read API to workflow streams WorkflowStream gains read() plus base_offset / end_offset properties so workflow code can consume its own stream's entries deterministically, e.g. items a streaming activity publishes back to the workflow that scheduled it. read() mirrors the poll handler's semantics (topic filtering, global per-item offsets, from_offset=0 meaning the start of whatever is retained, TruncatedOffset for non-zero offsets that fell behind truncation) and subscribe's result_type decoding, and end_offset composes with workflow.wait_condition to consume entries as they land. Additive only; no behavior changes to the publish/poll/query paths. --- CHANGELOG.md | 7 + .../contrib/workflow_streams/_stream.py | 155 +++++++++++++ .../contrib/workflow_streams/_topic_handle.py | 4 +- temporalio/contrib/workflow_streams/_types.py | 13 +- .../workflow_streams/test_workflow_streams.py | 204 ++++++++++++++++++ 5 files changed, 376 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a1a92a95..117a292cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,13 @@ to include examples, links to docs, or any other relevant information. ### Added +- Workflow Streams: `WorkflowStream` now has a public workflow-side read API — `read()` plus the + `base_offset` / `end_offset` properties — so workflow code can deterministically consume entries + from its own stream (for example, items a streaming activity publishes back to the workflow that + scheduled it). `read()` supports topic filtering and the same `result_type` decoding as + `subscribe`, and `end_offset` composes with `workflow.wait_condition` to consume entries as they + land. + ### Changed ### Deprecated diff --git a/temporalio/contrib/workflow_streams/_stream.py b/temporalio/contrib/workflow_streams/_stream.py index ae8608c3b..743e344ff 100644 --- a/temporalio/contrib/workflow_streams/_stream.py +++ b/temporalio/contrib/workflow_streams/_stream.py @@ -86,6 +86,10 @@ class WorkflowStream: - ``__temporal_workflow_stream_poll`` update — long-poll subscription - ``__temporal_workflow_stream_offset`` query — current log length + Workflow code can also read the log back directly via + :meth:`read`, with :attr:`base_offset` / :attr:`end_offset` + exposing the retained offset range. + Note: Because the publish handler is registered dynamically from ``__init__``, on the activation where the stream is @@ -234,6 +238,157 @@ def topic( self._topic_types[name] = bound return WorkflowTopicHandle(self, name, bound) + @property + def base_offset(self) -> int: + """Global offset of the oldest entry still retained in the log. + + Starts at 0, advances when :meth:`truncate` discards a prefix, + and carries across continue-as-new via ``prior_state``. Entries + at global offsets before this value have been discarded; reads + and polls that request an older (non-zero) offset fail with a + ``TruncatedOffset`` error. + """ + return self._base_offset + + @property + def end_offset(self) -> int: + """Global offset one past the newest entry in the log. + + Equal to :attr:`base_offset` plus the number of retained + entries; the next published item lands at this offset. This is + the same value the ``__temporal_workflow_stream_offset`` query + reports. It is monotonic — :meth:`truncate` moves + :attr:`base_offset`, not the end — which makes it suitable for + :func:`temporalio.workflow.wait_condition` predicates; see + :meth:`read`. + """ + return self._base_offset + len(self._log) + + @overload + def read( + self, + topics: str | list[str] | None = ..., + from_offset: int = ..., + *, + result_type: type[T], + ) -> list[WorkflowStreamItem[T]]: ... + @overload + def read( + self, + topics: str | list[str] | None = ..., + from_offset: int = ..., + *, + result_type: None = None, + ) -> list[WorkflowStreamItem[Any]]: ... + + def read( + self, + topics: str | list[str] | None = None, + from_offset: int = 0, + *, + result_type: type | None = None, + ) -> list[WorkflowStreamItem[Any]]: + """Return retained entries at global offsets >= ``from_offset``. + + The workflow-side read path: a deterministic, non-blocking + snapshot of the in-memory log with no I/O and no history + events, safe to call from workflow code. External consumers + should keep using :meth:`WorkflowStreamClient.subscribe`; this + method is for workflows that consume their own stream — e.g. + items an activity publishes back to the workflow that + scheduled it. + + To consume entries as they land, combine with + :attr:`end_offset` and + :func:`temporalio.workflow.wait_condition`: + + .. code-block:: python + + cursor = self.stream.end_offset + while not done: + await workflow.wait_condition( + lambda: self.stream.end_offset > cursor + ) + for item in self.stream.read(from_offset=cursor): + ... # item.offset is the item's global offset + cursor = self.stream.end_offset + + Args: + topics: Topic filter. A single topic name, a list of topic + names, or None. None or empty list means all topics. + from_offset: Global offset to start reading from. ``0`` + means "from the beginning of whatever is retained" + (i.e. from :attr:`base_offset`), matching poll + semantics. Offsets at or past :attr:`end_offset` + return an empty list. + result_type: Optional target type. Each returned + :class:`WorkflowStreamItem` has its ``data`` decoded + via the workflow's payload converter. When omitted, + the converter's default ``Any`` decoding is used. Pass + ``result_type=temporalio.common.RawValue`` for an + opaque ``RawValue`` wrapping the original ``Payload`` + — useful for heterogeneous topics where the caller + dispatches on ``Payload.metadata`` or wants per-entry + control over decode errors. + + Returns: + Matching entries in log order, each with its global + ``offset`` populated. + + Raises: + RuntimeError: If ``result_type`` is ``Payload`` — the + payload converter has no Payload decode path; use + ``result_type=RawValue`` instead. + ApplicationError: If ``from_offset`` is non-zero and below + :attr:`base_offset` (type ``TruncatedOffset``) — the + requested entries have been truncated. Callers that + may lag behind :meth:`truncate` can compare against + :attr:`base_offset` first. + """ + if result_type is Payload: + raise RuntimeError( + "Cannot read with result_type=Payload: the payload " + "converter has no Payload decode path. Omit result_type " + "for default decoding, or pass result_type=RawValue to " + "receive a RawValue wrapping the raw Payload." + ) + log_offset = from_offset - self._base_offset + if log_offset < 0: + if from_offset == 0: + # "From the beginning" — start at whatever is retained. + log_offset = 0 + else: + raise ApplicationError( + f"Requested offset {from_offset} has been truncated. " + f"Current base offset is {self._base_offset}.", + type=TRUNCATED_OFFSET_ERROR_TYPE, + ) + topic_set: set[str] | None + if topics is None: + topic_set = None + elif isinstance(topics, str): + topic_set = {topics} + else: + topic_set = set(topics) or None + converter = workflow.payload_converter() + items: list[WorkflowStreamItem[Any]] = [] + for i, entry in enumerate(self._log[log_offset:]): + if topic_set is not None and entry.topic not in topic_set: + continue + data: Any = ( + converter.from_payload(entry.data) + if result_type is None + else converter.from_payload(entry.data, result_type) + ) + items.append( + WorkflowStreamItem( + topic=entry.topic, + data=data, + offset=self._base_offset + log_offset + i, + ) + ) + return items + def get_state( self, *, publisher_ttl: timedelta = timedelta(seconds=900) ) -> WorkflowStreamState: diff --git a/temporalio/contrib/workflow_streams/_topic_handle.py b/temporalio/contrib/workflow_streams/_topic_handle.py index 3b94e226f..872395e59 100644 --- a/temporalio/contrib/workflow_streams/_topic_handle.py +++ b/temporalio/contrib/workflow_streams/_topic_handle.py @@ -123,7 +123,9 @@ class WorkflowTopicHandle(Generic[T]): This class is experimental and may change in future versions. Constructed via :meth:`WorkflowStream.topic`. Has no - ``subscribe`` — workflows do not consume their own stream. + ``subscribe`` — long-poll subscription is the external consumption + path. Workflow code that consumes its own stream reads the log + directly via :meth:`WorkflowStream.read`. """ def __init__( diff --git a/temporalio/contrib/workflow_streams/_types.py b/temporalio/contrib/workflow_streams/_types.py index a58cf75ae..82981357d 100644 --- a/temporalio/contrib/workflow_streams/_types.py +++ b/temporalio/contrib/workflow_streams/_types.py @@ -57,15 +57,16 @@ class WorkflowStreamItem(Generic[T]): This class is experimental and may change in future versions. The ``data`` field carries the decoded value produced by - :meth:`WorkflowStreamClient.subscribe`. The generic parameter ``T`` - matches the ``result_type`` passed to ``subscribe``: an instance of - ``T`` when ``result_type=T``, the converter's default ``Any`` - decoding when ``result_type`` is omitted, or a + :meth:`WorkflowStreamClient.subscribe` and + :meth:`WorkflowStream.read`. The generic parameter ``T`` matches + the ``result_type`` passed to either: an instance of ``T`` when + ``result_type=T``, the converter's default ``Any`` decoding when + ``result_type`` is omitted, or a :class:`temporalio.common.RawValue` wrapping the original ``Payload`` when ``result_type=RawValue``. - The ``offset`` field is populated at poll time from the item's - position in the global log. + The ``offset`` field is populated at poll/read time from the + item's position in the global log. """ topic: str diff --git a/tests/contrib/workflow_streams/test_workflow_streams.py b/tests/contrib/workflow_streams/test_workflow_streams.py index 12026ff1a..efefe0bc0 100644 --- a/tests/contrib/workflow_streams/test_workflow_streams.py +++ b/tests/contrib/workflow_streams/test_workflow_streams.py @@ -2698,6 +2698,210 @@ async def test_subscribe_iterates_through_more_ready(client: Client) -> None: await handle.signal(BasicWorkflowStreamWorkflow.close) +# --------------------------------------------------------------------------- +# Workflow-side read API +# --------------------------------------------------------------------------- + + +@dataclass +class WorkflowReadResult: + values: list[str] + offsets: list[int] + base_offset: int + end_offset: int + + +@workflow.defn +class WorkflowReadConsumerWorkflow: + """Consumes its own stream from workflow code while an activity publishes. + + The read loop is the documented pattern from ``WorkflowStream.read``: + wait for ``end_offset`` to advance, drain with ``read``, repeat. The + test runs the worker with ``max_cached_workflows=0`` so every + activation replays the loop from history, proving the read path is + deterministic. + """ + + @workflow.init + def __init__(self, count: int) -> None: + self.stream = WorkflowStream() + + @workflow.run + async def run(self, count: int) -> WorkflowReadResult: + activity_handle = workflow.start_activity( + "publish_items", + count, + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=10), + ) + values: list[str] = [] + offsets: list[int] = [] + cursor = self.stream.end_offset + while len(values) < count: + await workflow.wait_condition(lambda: self.stream.end_offset > cursor) + for item in self.stream.read("events", cursor, result_type=bytes): + values.append(item.data.decode()) + offsets.append(item.offset) + cursor = self.stream.end_offset + await activity_handle + return WorkflowReadResult( + values=values, + offsets=offsets, + base_offset=self.stream.base_offset, + end_offset=self.stream.end_offset, + ) + + +@pytest.mark.asyncio +async def test_workflow_reads_activity_published_items(client: Client) -> None: + """A workflow consumes items an activity publishes to its own stream via + the workflow-side read API, and the loop survives replay + (``max_cached_workflows=0``).""" + count = 10 + async with new_worker( + client, + WorkflowReadConsumerWorkflow, + activities=[publish_items], + max_cached_workflows=0, + ) as worker: + result = await client.execute_workflow( + WorkflowReadConsumerWorkflow.run, + count, + id=f"workflow-stream-wf-read-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result.values == [f"item-{i}" for i in range(count)] + assert result.offsets == list(range(count)) + assert result.base_offset == 0 + assert result.end_offset == count + + +@dataclass +class ReadProbeResult: + all_topics: list[str] + all_values: list[str] + all_offsets: list[int] + a_offsets: list[int] + suffix_values: list[str] + empty_read_len: int + base_before: int + end_before: int + default_values: list[str] + raw_values_ok: bool + payload_reject_error: str + base_after_truncate: int + end_after_truncate: int + from_zero_offsets_after_truncate: list[int] + truncated_error_type: str + + +@workflow.defn +class ReadProbeWorkflow: + """Exercises ``WorkflowStream.read`` entirely inside workflow code. + + Covers global offsets with and without topic filtering, suffix and + empty reads, ``base_offset`` / ``end_offset`` before and after + ``truncate``, default / ``RawValue`` decoding, the ``Payload`` + ``result_type`` rejection, and the ``TruncatedOffset`` error for a + non-zero ``from_offset`` that fell behind truncation. Publishing is + workflow-side, so the whole scenario is synchronous and + deterministic — no external round trips. + """ + + @workflow.init + def __init__(self) -> None: + self.stream = WorkflowStream() + + @workflow.run + async def run(self) -> ReadProbeResult: + from temporalio.api.common.v1 import Payload + + topic_a = self.stream.topic("a", type=str) + topic_b = self.stream.topic("b", type=str) + for i in range(3): + topic_a.publish(f"a-{i}") # global offsets 0, 2, 4 + topic_b.publish(f"b-{i}") # global offsets 1, 3, 5 + + all_items = self.stream.read(result_type=str) + a_items = self.stream.read("a", result_type=str) + suffix_items = self.stream.read(from_offset=4, result_type=str) + empty_items = self.stream.read( + from_offset=self.stream.end_offset, result_type=str + ) + default_items = self.stream.read(["a", "b"], 4) + raw_items = self.stream.read("a", result_type=RawValue) + try: + self.stream.read(result_type=Payload) + except RuntimeError as e: + payload_reject_error = str(e) + else: + payload_reject_error = "" + base_before = self.stream.base_offset + end_before = self.stream.end_offset + + self.stream.truncate(4) + try: + self.stream.read(from_offset=2) + except ApplicationError as e: + truncated_error_type = e.type or "" + else: + truncated_error_type = "" + from_zero_items = self.stream.read(result_type=str) + + return ReadProbeResult( + all_topics=[item.topic for item in all_items], + all_values=[item.data for item in all_items], + all_offsets=[item.offset for item in all_items], + a_offsets=[item.offset for item in a_items], + suffix_values=[item.data for item in suffix_items], + empty_read_len=len(empty_items), + base_before=base_before, + end_before=end_before, + default_values=[item.data for item in default_items], + raw_values_ok=len(raw_items) == 3 + and all(isinstance(item.data, RawValue) for item in raw_items), + payload_reject_error=payload_reject_error, + base_after_truncate=self.stream.base_offset, + end_after_truncate=self.stream.end_offset, + from_zero_offsets_after_truncate=[item.offset for item in from_zero_items], + truncated_error_type=truncated_error_type, + ) + + +@pytest.mark.asyncio +async def test_workflow_read_offsets_filters_decode_and_truncation( + client: Client, +) -> None: + """``read`` returns global offsets (also under topic filters), honors + ``from_offset`` and decode modes, and is truncation-aware: after + ``truncate``, ``from_offset=0`` reads from ``base_offset`` while a + non-zero truncated offset raises ``TruncatedOffset``.""" + async with new_worker(client, ReadProbeWorkflow) as worker: + result = await client.execute_workflow( + ReadProbeWorkflow.run, + id=f"workflow-stream-read-probe-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result.all_topics == ["a", "b", "a", "b", "a", "b"] + assert result.all_values == ["a-0", "b-0", "a-1", "b-1", "a-2", "b-2"] + assert result.all_offsets == [0, 1, 2, 3, 4, 5] + # Filtered reads keep global (not per-topic) offsets. + assert result.a_offsets == [0, 2, 4] + assert result.suffix_values == ["a-2", "b-2"] + assert result.empty_read_len == 0 + assert result.base_before == 0 + assert result.end_before == 6 + # Default decoding of str payloads yields the strings back. + assert result.default_values == ["a-2", "b-2"] + assert result.raw_values_ok is True + assert "result_type=Payload" in result.payload_reject_error + # truncate(4) moves base_offset, not the end. + assert result.base_after_truncate == 4 + assert result.end_after_truncate == 6 + assert result.from_zero_offsets_after_truncate == [4, 5] + assert result.truncated_error_type == "TruncatedOffset" + + @pytest.mark.asyncio async def test_cross_namespace_nexus_stream( client: Client, env: WorkflowEnvironment