From 86d50a10c445c1c3b2f11dd90280ecbcd8989174 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 15:05:00 -0700 Subject: [PATCH] Added Activity Stream handles and `BufferedPublisher`. An Activity writes as its attempt, so a retry withdraws the Records of the earlier attempts. An Activity's `BufferedPublisher` flushes at the end of the attempt. --- temporalio/common/__init__.py | 14 + temporalio/common/_streams/__init__.py | 10 + .../common/_streams/_activity_stream.py | 295 +++++++ temporalio/common/_streams/_buffered.py | 176 ++++ temporalio/common/_streams/_stream_client.py | 7 +- .../contrib/streams/test_activity_handles.py | 801 ++++++++++++++++++ tests/contrib/streams/test_pydantic.py | 71 ++ tests/contrib/streams/test_redis.py | 103 ++- tests/contrib/streams/test_stream_client.py | 14 + 9 files changed, 1487 insertions(+), 4 deletions(-) create mode 100644 temporalio/common/_streams/_activity_stream.py create mode 100644 temporalio/common/_streams/_buffered.py create mode 100644 tests/contrib/streams/test_activity_handles.py create mode 100644 tests/contrib/streams/test_pydantic.py diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index ba622ecb5..dd7853807 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -34,6 +34,12 @@ import temporalio.types if TYPE_CHECKING: + from temporalio.common._streams._activity_stream import ( + ActivityStream, + SyncActivityStream, + SyncSubscription, + ) + from temporalio.common._streams._buffered import BufferedPublisher from temporalio.common._streams._defn import StreamDefinition, defn from temporalio.common._streams._errors import ( StreamClosedError, @@ -70,7 +76,9 @@ __all__ = [ "ActivityIDConflictPolicy", "ActivityIDReusePolicy", + "ActivityStream", "AutoUpgradeVersioningOverride", + "BufferedPublisher", "ClientStream", "CloseStatus", "Closed", @@ -125,6 +133,8 @@ "StreamUnsupportedError", "Subscription", "Superseded", + "SyncActivityStream", + "SyncSubscription", "TypedSearchAttributes", "VersioningBehavior", "VersioningOverride", @@ -1591,6 +1601,8 @@ def wait_sync(self, timeout: float | None = None) -> None: # so this module imports the Stream names on first use. _STREAM_NAMES = frozenset( [ + "ActivityStream", + "BufferedPublisher", "ClientStream", "CloseStatus", "Closed", @@ -1616,6 +1628,8 @@ def wait_sync(self, timeout: float | None = None) -> None: "StreamUnsupportedError", "Subscription", "Superseded", + "SyncActivityStream", + "SyncSubscription", "WorkflowStream", "defn", ] diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 4c828426f..ace0c5ba7 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -1,5 +1,11 @@ """The Streams API, which :py:mod:`temporalio.common` exports.""" +from temporalio.common._streams._activity_stream import ( + ActivityStream, + SyncActivityStream, + SyncSubscription, +) +from temporalio.common._streams._buffered import BufferedPublisher from temporalio.common._streams._defn import StreamDefinition, defn from temporalio.common._streams._errors import ( StreamClosedError, @@ -31,6 +37,8 @@ __all__ = [ "END", + "ActivityStream", + "BufferedPublisher", "ClientStream", "CloseStatus", "Closed", @@ -55,6 +63,8 @@ "StreamUnsupportedError", "Subscription", "Superseded", + "SyncActivityStream", + "SyncSubscription", "WorkflowStream", "defn", ] diff --git a/temporalio/common/_streams/_activity_stream.py b/temporalio/common/_streams/_activity_stream.py new file mode 100644 index 000000000..0013f3a50 --- /dev/null +++ b/temporalio/common/_streams/_activity_stream.py @@ -0,0 +1,295 @@ +"""Stream handles in Activity code.""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +from collections.abc import Coroutine, Generator +from datetime import timedelta +from typing import Any, Generic, TypeVar + +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.common._streams import _activity, _errors +from temporalio.common._streams._access import Access, read_spec, subscription +from temporalio.common._streams._handle import StreamHandle +from temporalio.common._streams._items import Closed, End, Record, Superseded +from temporalio.common._streams._reader import Subscription + +T = TypeVar("T") +_R = TypeVar("_R") + +# How long a blocked sync call waits at a time. +_WAIT_STEP = 0.1 + + +def _start( + loop: asyncio.AbstractEventLoop, + call: Coroutine[Any, Any, _R], + future: concurrent.futures.Future[_R], +) -> None: + # Runs in the Worker's loop. The Activity thread can cancel the future before this runs. + if future.cancelled(): + call.close() + return + task = loop.create_task(call) + + def task_done(task: asyncio.Task[_R]) -> None: + if task.cancelled(): + future.cancel() + elif future.set_running_or_notify_cancel(): + error = task.exception() + if error is None: + future.set_result(task.result()) + else: + future.set_exception(error) + + def future_done(future: concurrent.futures.Future[_R]) -> None: + if future.cancelled() and not loop.is_closed(): + loop.call_soon_threadsafe(task.cancel) + + task.add_done_callback(task_done) + future.add_done_callback(future_done) + + +def _bound(handle: StreamHandle[Any]) -> _activity.ActivityStreams: + streams = _activity.current.get() + if streams is None: + raise _errors.StreamConfigError( + "No Stream backend runs this Activity. Register one on the Worker's client " + "with stream_backends=." + ) + if handle.backend and handle.backend != streams.backend_name: + raise _errors.StreamConfigError( + f"The Stream {handle.id} is of the backend {handle.backend}. The Worker has only " + f"{streams.backend_name}." + ) + if not handle.backend and not streams.default: + raise _errors.no_default_backend( + "Worker", [streams.backend_name], "a handle that has its backend name" + ) + return streams + + +def _not_owner(handle: StreamHandle[Any]) -> _errors.StreamNotOwnerError: + return _errors.StreamNotOwnerError( + f"Can't close the Stream {handle.id} from an Activity. A Workflow's Stream closes " + "only from Workflow code, and a client's only from a client.", + reason=StreamErrorReason.STREAM_ERROR_REASON_NOT_OWNER, + ) + + +def _access(handle: StreamHandle[Any]) -> Access: + streams = _bound(handle) + stream_id = handle.id + return Access( + backend=streams.backend, + converter=streams.converter, + namespace=streams.namespace, + stream_id=stream_id, + writer=lambda: streams.writer(stream_id), + ) + + +class ActivityStream(StreamHandle[T]): + """A Stream in an async Activity. The Activity can write and read it. + + An Activity gets a Stream only as an argument, such as a field of its input. Each publish + writes before it returns. Use :py:class:`BufferedPublisher` to batch. + + .. warning:: + This API is experimental. + """ + + def _access(self) -> Access: + return _access(self) + + async def publish(self, *values: T) -> int: + """Write one Record for each value, in order, and return the last Record's offset. + + A list is one value. + + Raises: + StreamNotFoundError: Nobody created the Stream. + """ + return await self._access().publish(values) + + def subscribe( + self, + *, + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + ) -> Subscription[T]: + """Read the Stream's items in order. See :py:meth:`ClientStream.subscribe`.""" + read = read_spec( + self._payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + access = self._access() + return subscription(lambda position: access.items(read, position), read) + + async def get_offset(self) -> int: + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + + Offsets are sparse and above zero, so ``offset=0`` reads from the start. + """ + return await self._access().get_offset() + + async def close( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Raise :py:class:`StreamNotOwnerError`, since an Activity can't close a Stream. + + A Workflow's Stream closes only from Workflow code, and a client's Stream only from a + client. The rule includes a Local Activity of the run chain that created the Stream. A + Local Activity runs before its Workflow Task commits. If the task then failed, the close + would stay, and the task's retry couldn't create the Stream again. + """ + del result, failure + raise _not_owner(self) + + +class SyncSubscription(Generic[T]): + """The items of one read of a Stream in a sync Activity, in order. Iterate it with ``for``. + + Each iteration blocks the Activity's thread. See :py:class:`Subscription`. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, stream: SyncActivityStream[Any], subscription: Subscription[T] + ) -> None: + """Create a read. Use a handle's ``subscribe`` instead.""" + self._subscription = subscription + self._steps = stream._items(subscription) + + @property + def next_offset(self) -> int | None: + """The offset to resume from. See :py:attr:`Subscription.next_offset`.""" + return self._subscription.next_offset + + def __iter__(self) -> SyncSubscription[T]: + """The read is its own iterator.""" + return self + + def __next__(self) -> Record[T] | Superseded | Closed: + """The next item.""" + return next(self._steps) + + def close(self) -> None: + """End the read early. It makes no more calls.""" + self._steps.close() + + +class SyncActivityStream(StreamHandle[T]): + """A Stream in a sync Activity. Its calls block the Activity's thread. + + The calls run on the Worker's event loop, so the Activity thread only waits for them. + + .. warning:: + This API is experimental. + """ + + def _access(self) -> Access: + return _access(self) + + def _run(self, call: Coroutine[Any, Any, _R]) -> _R: + loop = _bound(self).loop + try: + on_loop = asyncio.get_running_loop() is loop + except RuntimeError: + on_loop = False + if on_loop: + call.close() + raise RuntimeError( + "A SyncActivityStream blocks, so use ActivityStream in an async Activity" + ) + # The future exists before the call is scheduled, so a cancel raised in this thread at + # any point after the schedule reaches the except below, which cancels the call. + future: concurrent.futures.Future[_R] = concurrent.futures.Future() + try: + try: + loop.call_soon_threadsafe(_start, loop, call, future) + except RuntimeError: + # The Worker's loop is closed, such as at shutdown, so the call never starts. + call.close() + raise + while True: + # Short waits let the Worker raise an exception in this thread, such as a cancel. + try: + return future.result(timeout=_WAIT_STEP) + except concurrent.futures.TimeoutError: + continue + except BaseException: + # Cancel the call too, so a long poll doesn't keep running in the Worker's loop after + # the Activity thread is canceled, such as from a heartbeat. + future.cancel() + raise + + def publish(self, *values: T) -> int: + """Write one Record for each value, in order, and return the last Record's offset. + + A list is one value. + + Raises: + StreamNotFoundError: Nobody created the Stream. + """ + return self._run(self._access().publish(values)) + + def subscribe( + self, + *, + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + ) -> SyncSubscription[T]: + """Read the Stream's items in order. See :py:meth:`ClientStream.subscribe`.""" + read = read_spec( + self._payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + access = self._access() + return SyncSubscription( + self, subscription(lambda position: access.items(read, position), read) + ) + + def _items( + self, items: Subscription[T] + ) -> Generator[Record[T] | Superseded | Closed, None, None]: + try: + while True: + try: + item = self._run(_next(items)) + except StopAsyncIteration: + return + yield item + finally: + self._run(items.aclose()) + + def get_offset(self) -> int: + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + + Offsets are sparse and above zero, so ``offset=0`` reads from the start. + """ + return self._run(self._access().get_offset()) + + def close( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Raise :py:class:`StreamNotOwnerError`. See :py:meth:`ActivityStream.close`.""" + del result, failure + raise _not_owner(self) + + +async def _next(items: Subscription[T]) -> Record[T] | Superseded | Closed: + return await items.__anext__() diff --git a/temporalio/common/_streams/_buffered.py b/temporalio/common/_streams/_buffered.py new file mode 100644 index 000000000..fdf65c637 --- /dev/null +++ b/temporalio/common/_streams/_buffered.py @@ -0,0 +1,176 @@ +"""A publisher that batches the Records of one Stream.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta +from types import TracebackType +from typing import Generic, TypeVar + +from temporalio.common._streams import _activity +from temporalio.common._streams._activity_stream import ActivityStream +from temporalio.common._streams._stream_client import ClientStream + +T = TypeVar("T") + +# Records in one batch of the crate's writer. A full batch is written at once. +_BATCH = 1000 + + +class BufferedPublisher(Generic[T]): + """Batches the publishes of one Stream. + + Use it in an Activity or a client that writes often. :py:meth:`publish` returns at once. + The publisher writes the buffered values as one batch at most ``max_delay`` after the first + of them, or at once when they fill a batch. One write is in flight at a time, and values + published during it wait for the next one. Use it in ``async with``, which writes the + remaining values at the end of the block. In an Activity, the end of the attempt also + writes the remaining values. + + ``await flush()`` is the backpressure point. A caller that publishes faster than the Store + accepts Records must await it from time to time, so the buffer doesn't grow without bound. + + After a batch fails, the publisher writes nothing more, so a reader sees no gap. Each later + call raises the error. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + target: ActivityStream[T] | ClientStream[T], + *, + max_delay: timedelta = timedelta(milliseconds=100), + ) -> None: + """Create a publisher. + + Args: + target: An ``ActivityStream`` or a ``ClientStream``. + max_delay: The longest time that a value waits in the buffer. + """ + self._stream = target + self._max_delay = max_delay + self._buffer: list[T] = [] + # When the oldest buffered value arrived, so the next batch is at most max_delay later. + self._first_at = 0.0 + self._timer: asyncio.TimerHandle | None = None + self._writing: asyncio.Task[None] | None = None + self._published = 0 + # Count of the values that a write took, written or not. + self._settled = 0 + self._error: BaseException | None = None + self._last_offset = 0 + self._unwritten = 0 + self._reported = False + streams = _activity.current.get() + if streams is not None: + streams.publishers.append(self) + + def publish(self, *values: T) -> None: + """Add one Record for each value to the buffer, and return at once. + + A list is one value. + + Raises: + StreamError: A batch failed to write, so the publisher takes no more values. + """ + self._raise_error() + if not values: + raise TypeError("publish takes one or more values") + loop = asyncio.get_running_loop() + if not self._buffer: + self._first_at = loop.time() + self._buffer.extend(values) + self._published += len(values) + self._schedule() + + async def flush(self) -> int: + """Write every value published before the call, and return the last written offset. + + Returns 0 if the publisher hasn't written anything. + """ + target = self._published + while self._settled < target: + if self._writing is None: + self._start() + assert self._writing is not None + await asyncio.shield(self._writing) + self._raise_error() + return self._last_offset + + def _raise_error(self) -> None: + if self._error is not None: + self._reported = True + raise self._error + + async def _flush_at_attempt_end(self) -> None: + # Fail the attempt for unwritten values, unless the Activity already got the error and + # continued. + reported = self._reported + try: + await self.flush() + except Exception: + if not reported: + raise + + def _schedule(self) -> None: + if self._writing is not None or not self._buffer: + return + if len(self._buffer) >= _BATCH: + self._start() + elif self._timer is None: + self._timer = asyncio.get_running_loop().call_at( + self._first_at + self._max_delay.total_seconds(), self._start + ) + + def _start(self) -> None: + if self._timer is not None: + self._timer.cancel() + self._timer = None + if self._writing is not None or not self._buffer: + return + batch, self._buffer = self._buffer, [] + self._writing = asyncio.get_running_loop().create_task(self._write(batch)) + + async def _write(self, batch: list[T]) -> None: + try: + if self._error is None: + self._last_offset = await self._send(batch) + else: + # A failed batch fails the later ones too, so a reader sees no gap. + self._unwritten += len(batch) + except Exception as err: + self._unwritten += len(batch) + self._error = err + finally: + self._settled += len(batch) + self._writing = None + # Schedule the values published during the write, for when their delay ends or they + # fill a batch. + self._schedule() + + async def _send(self, batch: list[T]) -> int: + return await self._stream._access().publish(batch) + + @property + def unwritten(self) -> int: + """How many published values a failed batch left unwritten.""" + return self._unwritten + + async def __aenter__(self) -> BufferedPublisher[T]: + """Use the publisher in ``async with``.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Write the buffered values. A write error does not hide an error of the block.""" + try: + await self.flush() + except Exception: + if exc is None: + raise diff --git a/temporalio/common/_streams/_stream_client.py b/temporalio/common/_streams/_stream_client.py index 210f9bf3e..eeb533e9b 100644 --- a/temporalio/common/_streams/_stream_client.py +++ b/temporalio/common/_streams/_stream_client.py @@ -90,9 +90,10 @@ class ClientStream(StreamHandle[T]): :py:meth:`temporalio.client.Client.stream`. A ``ClientStream`` that a query of a client returns, such as a Workflow's ``stream`` query, uses that client. - Each publish writes before it returns. All handles of one client and Stream write as one - Producer. :py:meth:`producer` returns a handle that writes as a named Producer. Use a handle - in ``async with`` to close the Stream at the end of the block. + Each publish writes before it returns. Use :py:class:`BufferedPublisher` to batch. All + handles of one client and Stream write as one Producer. :py:meth:`producer` returns a + handle that writes as a named Producer. Use a handle in ``async with`` to close the Stream + at the end of the block. .. warning:: This API is experimental. diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py new file mode 100644 index 000000000..55fff5d36 --- /dev/null +++ b/tests/contrib/streams/test_activity_handles.py @@ -0,0 +1,801 @@ +import asyncio +import contextlib +import contextvars +import ctypes +import gc +import inspect +import sys +import threading +import time +import uuid +from collections.abc import Coroutine, Iterator, Sequence +from dataclasses import dataclass +from datetime import timedelta +from types import SimpleNamespace +from typing import Any, cast + +import pytest + +from temporalio import activity, common, workflow +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.client import Client, WorkflowHandle +from temporalio.common import RetryPolicy +from temporalio.common._streams import _activity +from temporalio.exceptions import CancelledError +from tests.contrib.streams._support import workflow_stream +from tests.helpers import new_worker + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +@common.defn(name="other") +class Other(common.StreamDefinition[str]): ... + + +@dataclass +class CallModelInput: + prompt: str + stream: common.ActivityStream[str] + + +@dataclass +class SyncInput: + prompt: str + stream: common.SyncActivityStream[str] + + +async def read( + handle: WorkflowHandle[Any, Any], + definition: type[common.StreamDefinition[Any]] = Tokens, +) -> list[object]: + stream = workflow_stream(handle, definition) + return [item async for item in stream.subscribe(wait_new_records=False)] + + +def records(items: Sequence[object]) -> list[common.Record[Any]]: + return [item for item in items if isinstance(item, common.Record)] + + +@activity.defn +async def call_model(input: CallModelInput) -> list[int]: + assert input.stream._payload_type is str + offsets = [await input.stream.publish(word) for word in input.prompt.split()] + async with common.BufferedPublisher( + input.stream, max_delay=timedelta(seconds=10) + ) as buffered: + buffered.publish("model-a") + buffered.publish("model-b") + offsets.append(await input.stream.get_offset()) + return offsets + + +@activity.defn +def call_model_sync(input: SyncInput) -> int: + last = input.stream.publish(*input.prompt.split()) + assert input.stream.get_offset() == last + 1 + return last + + +@activity.defn +async def read_back(stream: common.ActivityStream[str]) -> list[Any]: + return [ + item.data + async for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> list[Any]: + stream = await workflow.create_stream(Tokens) + offsets = await workflow.execute_activity( + call_model, + CallModelInput(prompt, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + ) + last = await workflow.execute_activity( + call_model_sync, + SyncInput("sync words", common.SyncActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + ) + read = await workflow.execute_activity( + read_back, + common.ActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close() + return [offsets, last, read] + + +async def test_activity_publishes_and_returns_offsets(stream_client: Client): + async with new_worker( + stream_client, + Agent, + activities=[call_model, call_model_sync, read_back], + activity_executor=_executor(), + ) as worker: + handle = await stream_client.start_workflow( + Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + offsets, last, read_tokens = await handle.result() + items = await read(handle) + data = records(items) + written = ["hi", "there", "model-a", "model-b", "sync", "words"] + assert [record.data for record in data] == written + assert offsets[:2] == [data[0].offset, data[1].offset] + assert offsets[2] == data[3].offset + 1 + assert last == data[5].offset + assert read_tokens == written + assert isinstance(items[-1], common.Closed) + + +def _executor() -> Any: + import concurrent.futures + + return concurrent.futures.ThreadPoolExecutor(max_workers=2) + + +attempts: dict[str, int] = {} + + +@activity.defn +async def flaky(input: CallModelInput) -> None: + attempts[input.prompt] = attempts.get(input.prompt, 0) + 1 + if attempts[input.prompt] == 1: + await input.stream.publish("attempt-1") + raise RuntimeError("first attempt fails") + + +@workflow.defn +class Retrying: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + await stream.publish("start") + await workflow.execute_activity( + flaky, + CallModelInput(key, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def test_retry_withdraws_a_received_stream(stream_client: Client): + key = str(uuid.uuid4()) + async with new_worker(stream_client, Retrying, activities=[flaky]) as worker: + handle = await stream_client.start_workflow( + Retrying.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + items = await read(handle) + superseded = [item for item in items if isinstance(item, common.Superseded)] + # The second attempt writes nothing, and its start still withdraws the first one's Record. + assert len(superseded) == 1 + withdrawn = [r.data for r in records(items) if superseded[0].withdraws(r)] + assert withdrawn == ["attempt-1"] + + +@dataclass +class FinishedStreams: + key: str + closed: common.ActivityStream[str] + expired: common.ActivityStream[str] + + +@activity.defn +async def reads_after_a_failed_attempt(input: FinishedStreams) -> list[Any]: + attempts[input.key] = attempts.get(input.key, 0) + 1 + if attempts[input.key] == 1: + raise RuntimeError("first attempt fails") + items = input.closed.subscribe(wait_new_records=False) + return [item.data async for item in items if isinstance(item, common.Record)] + + +@workflow.defn +class ReadsFinishedStreams: + @workflow.run + async def run(self, key: str, expired: common.WorkflowStream[str]) -> list[Any]: + stream = await workflow.create_stream(Tokens) + await stream.publish("turn") + await stream.close() + return await workflow.execute_activity( + reads_after_a_failed_attempt, + FinishedStreams( + key, + common.ActivityStream.from_handle(stream), + common.ActivityStream.from_handle(expired), + ), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def test_retry_goes_on_with_a_closed_or_expired_stream(stream_client: Client): + # On a Stream that accepts no write, the second attempt's open withdraws nothing, and the + # attempt only reads. + expired = await stream_client.create_stream( + Tokens, id=f"expired-{uuid.uuid4()}", retention=timedelta(milliseconds=200) + ) + await expired.publish("old") + await asyncio.sleep(0.5) + key = str(uuid.uuid4()) + async with new_worker( + stream_client, ReadsFinishedStreams, activities=[reads_after_a_failed_attempt] + ) as worker: + result = await stream_client.execute_workflow( + ReadsFinishedStreams.run, + args=[key, common.WorkflowStream.from_handle(expired)], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == ["turn"] and attempts[key] == 2 + + +# Set by a test after its readers showed the first attempt's Record. +shown: dict[str, asyncio.Event] = {} + + +@activity.defn +async def fails_after_it_is_shown(input: CallModelInput) -> None: + attempts[input.prompt] = attempts.get(input.prompt, 0) + 1 + tokens = input.stream + if attempts[input.prompt] == 1: + await tokens.publish("attempt-1") + await shown[input.prompt].wait() + raise RuntimeError("first attempt fails") + await tokens.publish("attempt-2") + + +@workflow.defn +class ShownThenWithdrawn: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + await workflow.execute_activity( + fails_after_it_is_shown, + CallModelInput(key, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + await stream.close() + + +async def test_reader_withdraws_a_shown_record_after_a_reconnect( + stream_client: Client, +): + key = str(uuid.uuid4()) + shown[key] = asyncio.Event() + async with new_worker( + stream_client, ShownThenWithdrawn, activities=[fails_after_it_is_shown] + ) as worker: + handle = await stream_client.start_workflow( + ShownThenWithdrawn.run, + key, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + stream = workflow_stream(handle, Tokens) + live = stream.subscribe(create_wait=timedelta(seconds=20)) + reconnects = stream.subscribe(create_wait=timedelta(seconds=20)) + first = await anext(live) + before = await anext(reconnects) + # The reader closes its connection and keeps only the Records that it showed and its + # resume offset, as a UI does across a page reload. + resume_at = reconnects.next_offset + await reconnects.aclose() + shown[key].set() + live_items = [first, *[item async for item in live]] + resumed = [item async for item in stream.subscribe(offset=resume_at)] + assert isinstance(first, common.Record) and first.data == "attempt-1" + assert isinstance(before, common.Record) and before.offset == first.offset + [live_superseded] = [i for i in live_items if isinstance(i, common.Superseded)] + assert live_superseded.withdrawn_offsets == (first.offset,) + # A new read got no Record of the failed attempt, so the item lists no offset. The kept + # Record is still withdrawn. + [superseded] = [i for i in resumed if isinstance(i, common.Superseded)] + assert superseded.withdrawn_offsets == () + assert superseded.withdraws(before) + assert [r.data for r in records(resumed)] == ["attempt-2"] + assert isinstance(resumed[-1], common.Closed) + + +@activity.defn +async def wrong_backend(stream: common.ActivityStream[str]) -> str: + try: + await stream.publish("x") + except common.StreamConfigError: + return "config" + return "wrote" + + +@activity.defn +async def close_it(stream: common.ActivityStream[str]) -> str: + await stream.publish("last") + try: + await stream.close(result="closed by activity") + except common.StreamNotOwnerError as err: + return StreamErrorReason.Name(err.reason) + return "closed" + + +@workflow.defn +class Closing: + @workflow.run + async def run(self) -> list[str]: + stream = await workflow.create_stream(Tokens) + backend = await workflow.execute_activity( + wrong_backend, + common.ActivityStream[str](id=stream.id, backend="elsewhere"), + start_to_close_timeout=timedelta(seconds=30), + ) + owner = await workflow.execute_activity( + close_it, + common.ActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + # A Local Activity runs before its Workflow Task commits, so it can't close it either. + local = await workflow.execute_local_activity( + close_it, + common.ActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + # A handle with only the id names no Workflow, and no Activity can close a Stream. + by_id = await workflow.execute_local_activity( + close_it, + common.ActivityStream[str](id=stream.id), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close(result="closed by the run") + return [backend, owner, local, by_id] + + +async def test_activity_checks_the_backend_and_the_owner(stream_client: Client): + async with new_worker( + stream_client, Closing, activities=[wrong_backend, close_it] + ) as worker: + handle = await stream_client.start_workflow( + Closing.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + # An Activity writes to the Workflow's Stream but can't close it. + assert await handle.result() == [ + "config", + *["STREAM_ERROR_REASON_NOT_OWNER"] * 3, + ] + items = await read(handle) + assert [r.data for r in records(items)] == ["last", "last", "last"] + assert isinstance(items[-1], common.Closed) + assert items[-1].result == "closed by the run" + + +class _Batches(common.ActivityStream[int]): + # A Stream whose writes go to a list, so the batches show with no backend. + def __init__(self, fail: bool = False) -> None: + super().__init__(id="batches") + self.batches: list[tuple[int, ...]] = [] + self.fail = fail + self.in_flight = 0 + self.most_in_flight = 0 + + def _access(self) -> Any: + return SimpleNamespace(publish=self._write) + + async def _write(self, values: Sequence[int]) -> int: + if self.fail: + raise common.StreamStorageError("store down") + self.in_flight += 1 + self.most_in_flight = max(self.most_in_flight, self.in_flight) + await asyncio.sleep(0.01) + self.in_flight -= 1 + self.batches.append(tuple(values)) + return sum(len(batch) for batch in self.batches) + + +async def test_buffered_publisher_batches_by_delay(): + stream = _Batches() + publisher = common.BufferedPublisher(stream, max_delay=timedelta(milliseconds=20)) + publisher.publish(1) + publisher.publish(2, 3) + await asyncio.sleep(0.1) + assert stream.batches == [(1, 2, 3)] + publisher.publish(4) + assert await publisher.flush() == 4 + assert stream.batches == [(1, 2, 3), (4,)] + assert await publisher.flush() == 4 + + +async def test_buffered_publisher_writes_a_full_batch_with_one_write_in_flight(): + stream = _Batches() + publisher = common.BufferedPublisher(stream, max_delay=timedelta(minutes=10)) + publisher.publish(*range(1000)) + publisher.publish(*range(1000, 2500)) + await asyncio.sleep(0.2) + # A full batch is written without waiting for the delay. The values published during its + # write are in the next batch. + assert [len(batch) for batch in stream.batches] == [1000, 1500] + assert stream.most_in_flight == 1 + assert await publisher.flush() == 2500 + + +async def test_buffered_publisher_raises_a_failed_batch(): + stream = _Batches(fail=True) + with pytest.raises(common.StreamStorageError): + async with common.BufferedPublisher(stream) as publisher: + publisher.publish(1) + with pytest.raises(TypeError): + common.BufferedPublisher(stream).publish() + + +@activity.defn +async def local_step(stream: common.ActivityStream[str]) -> list[Any]: + # Each Workflow write is an awaited Local Activity, so the run's create and Records are in + # the Store before a later Local Activity starts. + seen = [ + item.data + async for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + await stream.publish("local") + return seen + + +@workflow.defn +class LocalOrder: + @workflow.run + async def run(self) -> list[Any]: + stream = await workflow.create_stream(Tokens) + await stream.publish("before") + seen = await workflow.execute_local_activity( + local_step, + common.ActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.publish("after") + await stream.close() + return seen + + +async def test_workflow_writes_go_out_before_a_local_activity(stream_client: Client): + async with new_worker(stream_client, LocalOrder, activities=[local_step]) as worker: + handle = await stream_client.start_workflow( + LocalOrder.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == ["before"] + items = await read(handle) + written = records(items) + assert [record.data for record in written] == ["before", "local", "after"] + producer = written[1].producer + assert producer is not None and producer.id.startswith("la:") + assert isinstance(items[-1], common.Closed) + + +async def test_buffered_publisher_stays_failed(): + stream = _Batches(fail=True) + publisher = common.BufferedPublisher(stream) + publisher.publish(1, 2) + with pytest.raises(common.StreamStorageError): + await publisher.flush() + stream.fail = False + # A later batch would leave a gap, so the publisher refuses it and keeps the error. + with pytest.raises(common.StreamStorageError): + publisher.publish(3) + with pytest.raises(common.StreamStorageError): + await publisher.flush() + assert stream.batches == [] and publisher.unwritten == 2 + + +@activity.defn +async def never_flushes(input: CallModelInput) -> None: + publisher = common.BufferedPublisher(input.stream, max_delay=timedelta(minutes=10)) + publisher.publish(*input.prompt.split()) + + +@dataclass +class ManyStreams: + key: str + targets: list[common.ActivityStream[str]] + maybe: common.ActivityStream[str] | None = None + + +@activity.defn +async def flaky_many(input: ManyStreams) -> None: + attempts[input.key] = attempts.get(input.key, 0) + 1 + assert input.maybe is not None + if attempts[input.key] == 1: + for stream in [*input.targets, input.maybe]: + await stream.publish("attempt-1") + raise RuntimeError("first attempt fails") + + +@activity.defn +def sync_read(stream: common.SyncActivityStream[str]) -> list[Any]: + return [ + item.data + for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + + +@workflow.defn +class Leftovers: + @workflow.run + async def run(self, key: str) -> list[Any]: + stream = await workflow.create_stream(Tokens) + other = await workflow.create_stream(Other) + await workflow.execute_activity( + never_flushes, + CallModelInput("a b", common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + ) + await workflow.execute_activity( + flaky_many, + ManyStreams( + key, + [common.ActivityStream.from_handle(stream)], + common.ActivityStream.from_handle(other), + ), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + return await workflow.execute_activity( + sync_read, + common.SyncActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + + +async def test_attempt_end_flushes_and_retry_finds_nested_streams( + stream_client: Client, +): + key = str(uuid.uuid4()) + async with new_worker( + stream_client, + Leftovers, + activities=[never_flushes, flaky_many, sync_read], + activity_executor=_executor(), + ) as worker: + handle = await stream_client.start_workflow( + Leftovers.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + # The publisher was never flushed, so the end of the attempt wrote its values. + assert await handle.result() == ["a", "b", "attempt-1"] + items = await read(handle) + other = await read(handle, Other) + for read_items in (items, other): + superseded = [i for i in read_items if isinstance(i, common.Superseded)] + withdrawn = [r.data for r in records(read_items) if superseded[0].withdraws(r)] + assert withdrawn == ["attempt-1"] + + +@contextlib.contextmanager +def _loop_in_thread() -> Iterator[asyncio.AbstractEventLoop]: + # Each test closes its loop. The finalizer of a loop left open can run in a later caller + # thread and swallow the cancel that a test raises there. + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + try: + yield loop + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(5) + loop.close() + + +def _long_poll( + started: threading.Event, stopped: threading.Event +) -> Coroutine[Any, Any, None]: + async def long_poll() -> None: + started.set() + try: + await asyncio.sleep(60) + finally: + stopped.set() + + return long_poll() + + +def _caller( + loop: asyncio.AbstractEventLoop, + call: Coroutine[Any, Any, None], + errors: list[BaseException], +) -> threading.Thread: + stream: common.SyncActivityStream[Any] = common.SyncActivityStream(id="s") + token = _activity.current.set( + cast(Any, SimpleNamespace(loop=loop, backend_name="", default=True)) + ) + + def run() -> None: + try: + stream._run(call) + except BaseException as err: + errors.append(err) + + caller = threading.Thread(target=contextvars.copy_context().run, args=(run,)) + _activity.current.reset(token) + return caller + + +def _waits_for_the_call(caller: threading.Thread) -> bool: + frame = sys._current_frames().get(caller.ident or 0) + while frame is not None: + if ( + frame.f_code.co_name == "result" + and "concurrent" in frame.f_code.co_filename + ): + return True + frame = frame.f_back + return False + + +def test_sync_call_stops_when_the_thread_is_cancelled(): + started, stopped = threading.Event(), threading.Event() + errors: list[BaseException] = [] + with _loop_in_thread() as loop: + caller = _caller(loop, _long_poll(started, stopped), errors) + # A garbage collection in the caller thread runs finalizers there, and a finalizer + # swallows an exception that is raised in its thread. + gc.collect() + gc.disable() + try: + caller.start() + assert started.wait(5) + # Raise only while the caller waits on the call's future, so the test checks _run and + # not where the scheduler happens to stop the caller thread. + deadline = time.monotonic() + 5 + while not _waits_for_the_call(caller) and time.monotonic() < deadline: + time.sleep(0.001) + assert _waits_for_the_call(caller) + # The Worker cancels a sync Activity by raising in its thread. + assert caller.ident is not None + ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_ulong(caller.ident), ctypes.py_object(CancelledError) + ) + caller.join(5) + finally: + gc.enable() + assert stopped.wait(5) + assert errors and isinstance(errors[0], CancelledError) + + +def test_sync_call_on_a_closed_loop_closes_its_coroutine(): + errors: list[BaseException] = [] + loop = asyncio.new_event_loop() + loop.close() + call = _long_poll(threading.Event(), threading.Event()) + caller = _caller(loop, call, errors) + caller.start() + caller.join(5) + assert errors and isinstance(errors[0], RuntimeError) + # A closed coroutine can't warn at shutdown that nobody awaited it. + assert inspect.getcoroutinestate(call) == inspect.CORO_CLOSED + + +def test_sync_call_stops_when_the_cancel_lands_right_after_the_schedule(): + started, stopped = threading.Event(), threading.Event() + errors: list[BaseException] = [] + with _loop_in_thread() as loop: + schedule = loop.call_soon_threadsafe + cancelled: list[bool] = [] + + def schedule_then_cancel(*args: Any) -> Any: + handle = schedule(*args) + # The Worker's cancel lands in the Activity thread as soon as the call is scheduled. + if threading.current_thread() is caller and not cancelled: + cancelled.append(True) + assert started.wait(5) + raise CancelledError() + return handle + + setattr(loop, "call_soon_threadsafe", schedule_then_cancel) + caller = _caller(loop, _long_poll(started, stopped), errors) + caller.start() + caller.join(5) + assert stopped.wait(5) + assert errors and isinstance(errors[0], CancelledError) + + +@activity.defn +async def plain(name: str) -> str: + return f"hello {name}" + + +@workflow.defn +class NeverTouchesStreams: + @workflow.run + async def run(self) -> str: + return await workflow.execute_activity( + plain, "ada", start_to_close_timeout=timedelta(seconds=10) + ) + + +@workflow.defn +class ReceivesAStream: + @workflow.run + async def run(self, stream: common.WorkflowStream[str]) -> str: + # The Activity gets no Stream, but the run received one. + return await workflow.execute_activity( + plain, stream.id, start_to_close_timeout=timedelta(seconds=10) + ) + + +@workflow.defn +class RetriesAfterContinueAsNew: + @workflow.run + async def run(self, key: str, stream: common.WorkflowStream[str] | None) -> None: + if stream is None: + created = await workflow.create_stream(Tokens) + workflow.continue_as_new(args=[key, created]) + await workflow.execute_activity( + flaky, + CallModelInput(key, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def keyed(handle: WorkflowHandle[Any, Any], activity_type: str) -> list[bool]: + # Whether each schedule of the Activity has the Producer key. + history = await handle.fetch_history() + return [ + "__temporal_stream_key" in attributes.header.fields + for event in history.events + if ( + attributes := event.activity_task_scheduled_event_attributes + ).activity_type.name + == activity_type + ] + + +async def test_only_a_run_that_uses_streams_keys_its_activities(stream_client: Client): + # The key is in History, so a run with no Stream adds none. A run that created, bound or + # received a Stream adds it to each Activity, so an Activity that writes keeps its Producer id. + key = str(uuid.uuid4()) + async with new_worker( + stream_client, + NeverTouchesStreams, + ReceivesAStream, + Retrying, + RetriesAfterContinueAsNew, + activities=[plain, flaky], + ) as worker: + none = await stream_client.start_workflow( + NeverTouchesStreams.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await none.result() == "hello ada" + assert await keyed(none, "plain") == [False] + creates = await stream_client.start_workflow( + Retrying.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await creates.result() + assert await keyed(creates, "flaky") == [True] + receives = await stream_client.start_workflow( + ReceivesAStream.run, + common.WorkflowStream.from_handle(workflow_stream(creates, Tokens)), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await receives.result() + assert await keyed(receives, "plain") == [True] + # The next run gets the Stream as its argument. Its retried Activity withdraws the first + # attempt's Records, as in the run that created the Stream. + continued_key = str(uuid.uuid4()) + continued = await stream_client.start_workflow( + RetriesAfterContinueAsNew.run, + args=[continued_key, None], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await continued.result() + assert await keyed(continued, "flaky") == [True] + items = await read(continued) + superseded = [item for item in items if isinstance(item, common.Superseded)] + assert len(superseded) == 1 + assert [r.data for r in records(items) if superseded[0].withdraws(r)] == [ + "attempt-1" + ] diff --git a/tests/contrib/streams/test_pydantic.py b/tests/contrib/streams/test_pydantic.py new file mode 100644 index 000000000..d7fadc069 --- /dev/null +++ b/tests/contrib/streams/test_pydantic.py @@ -0,0 +1,71 @@ +import uuid +from datetime import timedelta + +import pydantic + +from temporalio import activity, common, workflow +from temporalio.client import Client +from temporalio.contrib.pydantic import pydantic_data_converter +from temporalio.contrib.stream.testing import MemoryStreamBackend +from tests.helpers import new_worker + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +class CallModelInput(pydantic.BaseModel): + prompt: str + stream: common.ActivityStream[str] + + +class TurnResult(pydantic.BaseModel): + text: str + stream: common.ClientStream[str] + + +@activity.defn +async def call_model(input: CallModelInput) -> int: + assert input.stream._payload_type is str + return await input.stream.publish(*input.prompt.split()) + + +@workflow.defn +class TurnWorkflow: + @workflow.run + async def run(self, prompt: str) -> TurnResult: + stream = await workflow.create_stream(Tokens) + await workflow.execute_activity( + call_model, + CallModelInput( + prompt=prompt, stream=common.ActivityStream.from_handle(stream) + ), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close(result=prompt.upper()) + return TurnResult( + text=prompt.upper(), stream=common.ClientStream.from_handle(stream) + ) + + +async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): + config = client.config() + config["stream_backends"] = [MemoryStreamBackend(default=True)] + config["data_converter"] = pydantic_data_converter + client = Client(**config) + async with new_worker(client, TurnWorkflow, activities=[call_model]) as worker: + handle = await client.start_workflow( + TurnWorkflow.run, + "hi there", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + result = await handle.result() + assert result.text == "HI THERE" + assert result.stream.workflow_id == handle.id + assert result.stream.first_execution_run_id == handle.first_execution_run_id + assert result.stream.backend == "memory" + stream = client.stream(Tokens, id=result.stream.id) + items = [item async for item in stream.subscribe()] + assert [i.data for i in items if isinstance(i, common.Record)] == ["hi", "there"] + assert isinstance(items[-1], common.Closed) and items[-1].result is not None diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index b1ec38ae2..02d249ad8 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -1,8 +1,9 @@ import uuid +from dataclasses import dataclass from datetime import timedelta from typing import Any -from temporalio import common, workflow +from temporalio import activity, common, workflow from temporalio.client import Client from temporalio.worker import ( UnsandboxedWorkflowRunner, @@ -15,6 +16,12 @@ class Tokens(common.StreamDefinition[str]): ... +@dataclass +class CallModelInput: + prompt: str + stream: common.ActivityStream[str] + + async def read_all(stream: common.ClientStream[Any]) -> list[Any]: return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] @@ -64,3 +71,97 @@ async def test_failed_workflow_task_is_withdrawn(stream_client: Client): withdrawn = [r for r in written if any(i.withdraws(r) for i in superseded)] assert withdrawn == [written[1]] assert isinstance(items[-1], common.Closed) + + +@activity.defn +async def local_tokens(input: CallModelInput) -> None: + await input.stream.publish(*input.prompt.split()) + + +@workflow.defn +class LocalThenFailTask: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + first = attempts.setdefault(key, 0) == 0 + await workflow.execute_local_activity( + local_tokens, + CallModelInput( + "0 1" if first else "2 3", common.ActivityStream.from_handle(stream) + ), + start_to_close_timeout=timedelta(seconds=30), + ) + if first: + attempts[key] = 1 + raise RuntimeError("fail this Workflow Task once") + await stream.close() + + +async def test_rerun_local_activity_withdraws_its_failed_task(stream_client: Client): + async with new_worker( + stream_client, + LocalThenFailTask, + activities=[local_tokens], + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await stream_client.start_workflow( + LocalThenFailTask.run, + str(uuid.uuid4()), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = await read_all(workflow_stream(handle, Tokens)) + written = records(items) + superseded = [item for item in items if isinstance(item, common.Superseded)] + visible = [r for r in written if not any(i.withdraws(r) for i in superseded)] + assert [r.data for r in visible] == ["2", "3"] + # The rerun keeps the Producer id and uses the retry task's epoch. + first, second = written[0].producer, visible[0].producer + assert first is not None and second is not None and first.id == second.id + assert first.epoch is not None and second.epoch is not None + assert (first.epoch.task_attempt, first.epoch.task_started_event_id) < ( + second.epoch.task_attempt, + second.epoch.task_started_event_id, + ) + + +@workflow.defn +class LocalOnlyInFailedTask: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + if attempts.setdefault(key, 0) == 0: + attempts[key] = 1 + await workflow.execute_local_activity( + local_tokens, + CallModelInput("0 1", common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + ) + raise RuntimeError("fail this Workflow Task once") + await stream.publish("other") + await stream.close() + + +async def test_later_write_withdraws_a_skipped_local_activity(stream_client: Client): + # The retry task skips the Local Activity, so only its own write can withdraw "0 1". The + # write also uses the skipped call's key. + async with new_worker( + stream_client, + LocalOnlyInFailedTask, + activities=[local_tokens], + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await stream_client.start_workflow( + LocalOnlyInFailedTask.run, + str(uuid.uuid4()), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = await read_all(workflow_stream(handle, Tokens)) + written = records(items) + superseded = [item for item in items if isinstance(item, common.Superseded)] + visible = [r for r in written if not any(s.withdraws(r) for s in superseded)] + assert [r.data for r in visible] == ["other"] + assert isinstance(items[-1], common.Closed) diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 7d3125174..aaf818a22 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -308,6 +308,20 @@ async def read(definition: type[common.StreamDefinition[Any]]) -> list[object]: assert await read(TurnDicts) == [{"kind": "model", "text": ""}] +async def test_buffered_publisher_keeps_the_publish_order(stream_client: Client): + stream = await stream_client.create_stream(Turns, id=f"buffered-{uuid.uuid4()}") + async with common.BufferedPublisher(stream) as publisher: + publisher.publish(Turn("user")) + publisher.publish(Turn("a"), Turn("b")) + publisher.publish(Turn("model")) + written = [ + item.data + async for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + assert written == [Turn("user"), Turn("a"), Turn("b"), Turn("model")] + + @common.defn(name="turns") class LooseTurns(common.StreamDefinition[str]): ...