From 0bf1696a9a86cb4302c7792926a2172caf1b1973 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 15:08:43 -0700 Subject: [PATCH] Added the Stream examples and the end-to-end Redis tests. The examples run on the memory and Redis lanes, as the tests do. The `temporalio.common` module doc and the changelog describe Streams. --- .github/CODEOWNERS | 2 + changelog/added/babbling-brook.md | 2 + temporalio/common/__init__.py | 23 +- tests/contrib/streams/examples/README.md | 15 + tests/contrib/streams/examples/__init__.py | 0 .../examples/activity_streams_tokens.py | 98 +++ .../streams/examples/agent_publishes_turns.py | 116 +++ .../streams/examples/client_creates_stream.py | 56 ++ .../streams/examples/resume_from_offset.py | 107 +++ .../contrib/streams/examples/test_examples.py | 81 +++ .../streams/examples/ui_backend_reads.py | 90 +++ .../contrib/streams/examples/worker_setup.py | 88 +++ tests/contrib/streams/test_redis.py | 680 +++++++++++++++++- 13 files changed, 1355 insertions(+), 3 deletions(-) create mode 100644 changelog/added/babbling-brook.md create mode 100644 tests/contrib/streams/examples/README.md create mode 100644 tests/contrib/streams/examples/__init__.py create mode 100644 tests/contrib/streams/examples/activity_streams_tokens.py create mode 100644 tests/contrib/streams/examples/agent_publishes_turns.py create mode 100644 tests/contrib/streams/examples/client_creates_stream.py create mode 100644 tests/contrib/streams/examples/resume_from_offset.py create mode 100644 tests/contrib/streams/examples/test_examples.py create mode 100644 tests/contrib/streams/examples/ui_backend_reads.py create mode 100644 tests/contrib/streams/examples/worker_setup.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index f8a51d9b71..a698708d70 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -12,3 +12,5 @@ # manage repo-wide concerns. /temporalio/contrib/workflow_streams/ @temporalio/ai-sdk @temporalio/sdk /tests/contrib/workflow_streams/ @temporalio/ai-sdk @temporalio/sdk +/temporalio/contrib/stream/ @temporalio/ai-sdk @temporalio/sdk +/tests/contrib/streams/ @temporalio/ai-sdk @temporalio/sdk diff --git a/changelog/added/babbling-brook.md b/changelog/added/babbling-brook.md new file mode 100644 index 0000000000..44d9cc6b22 --- /dev/null +++ b/changelog/added/babbling-brook.md @@ -0,0 +1,2 @@ +Add experimental Streams to `temporalio.common`. A Stream is an ordered log of Records that Workflows, Activities and clients write, and that clients and Activities read. Register a backend with `stream_backends=` on `Client.connect`, `Replayer` or `WorkflowEnvironment`, then use `workflow.create_stream`, `Client.create_stream` or `Client.stream`. +Add a Redis Stream backend in `temporalio.contrib.stream`, and an in-memory one for tests in `temporalio.contrib.stream.testing`. The Redis backend needs a build with the bridge's `streams` feature, which the official builds have. diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index dd78538070..3d326fc1f6 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -1,4 +1,25 @@ -"""Common code used in the Temporal SDK.""" +"""Common code used in the Temporal SDK. + +This module also has the Streams API. A Stream is an ordered log of Records that Workflows, +Activities and clients write and read. A Stream definition is a subclass of +:py:class:`StreamDefinition` with ``@common.defn``. It gives the Stream name and its one payload +type. Code creates a Stream before its first publish, and a create is create-or-get. Workflow +code creates one with :py:func:`temporalio.workflow.create_stream`, and a client with +:py:meth:`temporalio.client.Client.create_stream`. A client binds any Stream with +:py:meth:`temporalio.client.Client.stream`. A Stream is data. A Workflow can pass it as an +argument to an Activity, a child Workflow or the next run, or hand it out through a query. The +parameter's type sets the handle class for that context, such as :py:class:`ActivityStream`. + +A :py:class:`StreamBackend` makes the Stream calls and registers a backend name. Examples are +:py:class:`temporalio.contrib.stream.RedisStreamBackend` (``redis``) and +:py:class:`temporalio.contrib.stream.testing.MemoryStreamBackend` (``memory``). Register +backends with ``Client.connect(..., stream_backends=[...])``. A handle stores the name of its +backend. With more than one registered backend, exactly one must have ``default=True``. A +create with no ``backend=`` uses that backend. A Worker serves only one backend. + +.. warning:: + The Streams API is experimental. +""" from __future__ import annotations diff --git a/tests/contrib/streams/examples/README.md b/tests/contrib/streams/examples/README.md new file mode 100644 index 0000000000..fdfe69c477 --- /dev/null +++ b/tests/contrib/streams/examples/README.md @@ -0,0 +1,15 @@ +# Examples for the Brook streams API proposal + +Each file is a complete program for one section of the proposal. `test_examples.py` runs each one +end to end on the memory Store and on Redis. + +| Proposal section | File | +|---|---| +| A UI backend reads a Workflow's Stream | [`ui_backend_reads.py`](ui_backend_reads.py) | +| A client resumes a Stream from an offset | [`resume_from_offset.py`](resume_from_offset.py) | +| A client creates a Stream and passes it to a Workflow | [`client_creates_stream.py`](client_creates_stream.py) | +| An agent Workflow publishes its turns | [`agent_publishes_turns.py`](agent_publishes_turns.py) | +| An Activity streams tokens to the Workflow's Stream | [`activity_streams_tokens.py`](activity_streams_tokens.py) | +| Worker setup | [`worker_setup.py`](worker_setup.py) | + +The proposal's Nexus examples aren't in this directory. diff --git a/tests/contrib/streams/examples/__init__.py b/tests/contrib/streams/examples/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/contrib/streams/examples/activity_streams_tokens.py b/tests/contrib/streams/examples/activity_streams_tokens.py new file mode 100644 index 0000000000..22e2fee28d --- /dev/null +++ b/tests/contrib/streams/examples/activity_streams_tokens.py @@ -0,0 +1,98 @@ +"""Brook API proposal, "An Activity streams tokens to the Workflow's Stream". + +An async Activity publishes each token with an awaited ``publish``, which returns the offset, +and batches the rest with ``BufferedPublisher``. A sync Activity takes ``SyncActivityStream``, +which works in a thread pool executor. +""" + +from __future__ import annotations + +import concurrent.futures +import uuid +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import activity, common, workflow +from temporalio.client import Client +from temporalio.worker import Worker + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +@dataclass +class CallModelInput: + prompt: str + tokens: common.ActivityStream[str] # the Stream that the Workflow passes + + +@dataclass +class SummarizeInput: + text: str + tokens: common.SyncActivityStream[str] + + +@activity.defn +async def call_model(input: CallModelInput) -> int: + words = input.prompt.split() + # Each awaited publish has written its Record when it returns. + last = await input.tokens.publish(words[0]) + # A fast model sends many small tokens, so the helper batches them. + async with common.BufferedPublisher( + input.tokens, max_delay=timedelta(milliseconds=100) + ) as fast: + for word in words[1:]: + fast.publish(word) + return last + + +@activity.defn +def summarize(input: SummarizeInput) -> int: + return input.tokens.publish(f"summary: {input.text}") + + +@workflow.defn +class Model: + @workflow.query + def stream(self) -> common.WorkflowStream[str]: + return workflow.stream_handle(Tokens) + + @workflow.run + async def run(self, prompt: str) -> list[int]: + tokens = await workflow.create_stream(Tokens) + first = await workflow.execute_activity( + call_model, + CallModelInput(prompt, common.ActivityStream.from_handle(tokens)), + start_to_close_timeout=timedelta(minutes=5), + ) + summary = await workflow.execute_activity( + summarize, + SummarizeInput(prompt, common.SyncActivityStream.from_handle(tokens)), + start_to_close_timeout=timedelta(minutes=5), + ) + await tokens.close() + return [first, summary] + + +async def main( + client: Client, task_queue: str +) -> tuple[list[int], list[common.Record[str]]]: + with concurrent.futures.ThreadPoolExecutor() as executor: + async with Worker( + client, + task_queue=task_queue, + workflows=[Model], + activities=[call_model, summarize], + activity_executor=executor, + ): + handle = await client.start_workflow( + Model.run, + "one two three", + id=f"model-{uuid.uuid4()}", + task_queue=task_queue, + ) + offsets = await handle.result() + stream = await handle.query("stream", result_type=common.ClientStream[str]) + records = [r async for r in stream.subscribe() if isinstance(r, common.Record)] + return offsets, records diff --git a/tests/contrib/streams/examples/agent_publishes_turns.py b/tests/contrib/streams/examples/agent_publishes_turns.py new file mode 100644 index 0000000000..ff86579473 --- /dev/null +++ b/tests/contrib/streams/examples/agent_publishes_turns.py @@ -0,0 +1,116 @@ +"""Brook API proposal, "An agent Workflow publishes its turns". + +The Workflow creates its two Streams, turns and token deltas, and hands out the turns Stream +through a ``stream`` query. It publishes each turn with an awaited ``publish``, and passes the +tokens Stream to the model Activity in its one input object. A create is create-or-get, so the +run after continue-as-new gets the same Streams. It closes both Streams before it returns, since +a run's end doesn't close a Stream. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass, field +from datetime import timedelta + +from temporalio import activity, common, workflow +from temporalio.client import Client +from temporalio.worker import Worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +@dataclass +class Reply: + text: str + tool_calls: list[str] = field(default_factory=list) + + +@dataclass +class CallModelInput: + messages: list[str] + tokens: common.ActivityStream[str] + + +@activity.defn +async def call_model(input: CallModelInput) -> Reply: + await input.tokens.publish(*input.messages[-1].split()) + if len(input.messages) == 1: + return Reply("", tool_calls=["search"]) + return Reply("the answer") + + +@activity.defn +async def run_tool(call: str) -> str: + return f"{call} result" + + +@workflow.defn +class Agent: + @workflow.query + def stream(self) -> common.WorkflowStream[Turn]: + # The handle makes no call, so a reader can get it before the create. + return workflow.stream_handle(Turns) + + @workflow.run + async def run(self, messages: list[str]) -> str: + # Create-or-get, so a retry or the run after continue-as-new gets the same Streams. + turns = await workflow.create_stream(Turns, retention=timedelta(days=1)) + tokens = await workflow.create_stream(Tokens, retention=timedelta(days=1)) + while True: + if workflow.info().is_continue_as_new_suggested(): + workflow.continue_as_new(messages) + # The publish is one Local Activity. History keeps its offset, not the Turn. + await turns.publish(Turn("model")) + reply = await workflow.execute_activity( + call_model, + CallModelInput(messages, common.ActivityStream.from_handle(tokens)), + start_to_close_timeout=timedelta(minutes=5), + ) + if not reply.tool_calls: + await turns.publish(Turn("done", reply.text)) + # A run's end doesn't close the Streams, so the Workflow closes them. + await turns.close(result=reply.text) + await tokens.close() + return reply.text + for call in reply.tool_calls: + messages.append( + await workflow.execute_activity( + run_tool, call, start_to_close_timeout=timedelta(minutes=5) + ) + ) + + +async def main(client: Client, task_queue: str) -> list[common.Record | common.Closed]: + async with Worker( + client, + task_queue=task_queue, + workflows=[Agent], + activities=[call_model, run_tool], + ): + handle = await client.start_workflow( + Agent.run, + ["find it"], + id=f"agent-{uuid.uuid4()}", + task_queue=task_queue, + ) + await handle.result() + # The Workflow hands out its turns Stream, and the client reads it as a ClientStream. + stream = await handle.query("stream", result_type=common.ClientStream[Turn]) + return [ + item + async for item in stream.subscribe() + if isinstance(item, (common.Record, common.Closed)) + ] diff --git a/tests/contrib/streams/examples/client_creates_stream.py b/tests/contrib/streams/examples/client_creates_stream.py new file mode 100644 index 0000000000..2b0813cf94 --- /dev/null +++ b/tests/contrib/streams/examples/client_creates_stream.py @@ -0,0 +1,56 @@ +"""Brook API proposal, "A client creates a Stream and passes it to a Workflow". + +The client creates the Stream with its settings, passes it to the Workflow as data, reads it, +and closes it, since the creator owns it. A create is create-or-get by Stream id. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import common, workflow +from temporalio.client import Client +from temporalio.worker import Worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@workflow.defn +class Chat: + @workflow.run + async def run(self, stream: common.WorkflowStream[Turn]) -> None: + # The Workflow only writes. The client closes the Stream. + await stream.publish(Turn("model", "hi"), Turn("last", "bye")) + + +async def main(client: Client, task_queue: str) -> list[Turn]: + async with Worker(client, task_queue=task_queue, workflows=[Chat]): + turns = [] + # The block closes the Stream at its end, or with the error that ends it. + async with client.create_stream( + Turns, id=f"chat-{uuid.uuid4()}", retention=timedelta(days=1) + ) as stream: + await client.start_workflow( + Chat.run, + # A Stream is data, so the Workflow gets it as an argument. + common.WorkflowStream.from_handle(stream), + id=f"chat-{uuid.uuid4()}", + task_queue=task_queue, + ) + async for r in stream.subscribe(): + if not isinstance(r, common.Record): + continue + turns.append(r.data) + if r.data.kind == "last": + break + return turns diff --git a/tests/contrib/streams/examples/resume_from_offset.py b/tests/contrib/streams/examples/resume_from_offset.py new file mode 100644 index 0000000000..e27e1829fd --- /dev/null +++ b/tests/contrib/streams/examples/resume_from_offset.py @@ -0,0 +1,107 @@ +"""Brook API proposal, "A client resumes a Stream from an offset". + +The Workflow hands out its ``turns`` Stream through a ``stream`` query. The Reader saves the +read's ``next_offset`` after each Record, stops, and later resumes from it with ``offset``. A +view that wants only new Records starts at ``common.END``. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import common, workflow +from temporalio.client import Client, WorkflowHandle +from temporalio.worker import Worker + + +@dataclass +class Turn: + text: str + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@workflow.defn +class Chat: + def __init__(self) -> None: + self.more: list[str] = [] + self.done = False + + @workflow.query + def stream(self) -> common.WorkflowStream[Turn]: + return workflow.stream_handle(Turns) + + @workflow.run + async def run(self, first: list[str]) -> None: + turns = await workflow.create_stream(Turns) + await turns.publish(*map(Turn, first)) + await workflow.wait_condition(lambda: bool(self.more) or self.done) + await turns.publish(*map(Turn, self.more)) + + @workflow.signal + def add(self, texts: list[str]) -> None: + self.more = texts + + +class OffsetStore: + """Where the Reader keeps its place, such as a database row.""" + + def __init__(self) -> None: + self.offsets: dict[str, int] = {} + + def load_offset(self, key: str) -> int | None: + return self.offsets.get(key) + + def save_offset(self, key: str, offset: int) -> None: + self.offsets[key] = offset + + +async def read_some( + handle: WorkflowHandle[Chat, None], store: OffsetStore, count: int +) -> list[str]: + stream = await handle.query("stream", result_type=common.ClientStream[Turn]) + seen = [] + read = stream.subscribe( + offset=store.load_offset("turns"), create_wait=timedelta(seconds=10) + ) + async for r in read: + if isinstance(r, common.Closed): + break + if isinstance(r, common.Record): + seen.append(r.data.text) + # The read's next offset is after this Record, so the next read starts there. + store.save_offset("turns", read.next_offset or 0) + if len(seen) == count: + break + await read.aclose() + return seen + + +async def read_new(handle: WorkflowHandle[Chat, None]) -> list[str]: + # A view that shows only what comes next starts after the newest Record. + stream = await handle.query("stream", result_type=common.ClientStream[Turn]) + return [ + r.data.text + async for r in stream.subscribe(offset=common.END, wait_new_records=False) + if isinstance(r, common.Record) + ] + + +async def main( + client: Client, task_queue: str +) -> tuple[list[str], list[str], list[str]]: + async with Worker(client, task_queue=task_queue, workflows=[Chat]): + handle = await client.start_workflow( + Chat.run, ["a", "b"], id=f"chat-{uuid.uuid4()}", task_queue=task_queue + ) + store = OffsetStore() + before = await read_some(handle, store, 2) + new = await read_new(handle) + await handle.signal(Chat.add, ["c", "d"]) + after = await read_some(handle, store, 2) + await handle.result() + return before, new, after diff --git a/tests/contrib/streams/examples/test_examples.py b/tests/contrib/streams/examples/test_examples.py new file mode 100644 index 0000000000..ba7c7c35ec --- /dev/null +++ b/tests/contrib/streams/examples/test_examples.py @@ -0,0 +1,81 @@ +"""Runs each example end to end on the memory Store and on Redis, so a broken one fails.""" + +import uuid + +import pytest + +from temporalio import common +from temporalio.client import Client +from tests.contrib.streams.examples import ( + activity_streams_tokens, + agent_publishes_turns, + client_creates_stream, + resume_from_offset, + ui_backend_reads, + worker_setup, +) + + +def task_queue() -> str: + return f"tq-{uuid.uuid4()}" + + +async def test_ui_backend_reads(stream_client: Client): + screen = await ui_backend_reads.main(stream_client, task_queue()) + Turn = ui_backend_reads.Turn + assert list(screen.shown.values()) == [Turn("model"), Turn("done", "HELLO THERE")] + assert screen.result == "HELLO THERE" + + +async def test_resume_from_offset(stream_client: Client): + assert await resume_from_offset.main(stream_client, task_queue()) == ( + ["a", "b"], + [], + ["c", "d"], + ) + + +async def test_client_creates_stream(stream_client: Client): + turns = await client_creates_stream.main(stream_client, task_queue()) + assert [turn.kind for turn in turns] == ["model", "last"] + + +async def test_agent_publishes_turns(stream_client: Client): + items = await agent_publishes_turns.main(stream_client, task_queue()) + data = [item.data for item in items if isinstance(item, common.Record)] + Turn = agent_publishes_turns.Turn + assert data == [Turn("model"), Turn("model"), Turn("done", "the answer")] + assert isinstance(items[-1], common.Closed) and items[-1].result == "the answer" + + +async def test_activity_streams_tokens(stream_client: Client): + offsets, records = await activity_streams_tokens.main(stream_client, task_queue()) + assert [record.data for record in records] == [ + "one", + "two", + "three", + "summary: one two three", + ] + assert offsets == [records[0].offset, records[-1].offset] + + +async def test_worker_setup_on_memory(client: Client): + target = client.service_client.config.target_host + test_client = await worker_setup.connect_for_tests(target) + assert await worker_setup.main(test_client, task_queue()) == [ + "hello ada", + "bye ada", + ] + + +async def test_worker_setup_on_redis(client: Client, redis_url: str): + target = client.service_client.config.target_host + redis_client = await worker_setup.connect_with_redis(target, redis_url) + assert await worker_setup.main(redis_client, task_queue()) == [ + "hello ada", + "bye ada", + ] + migrating = await worker_setup.connect_while_migrating(target, redis_url, redis_url) + assert await worker_setup.create_on_second_backend(migrating) == "redis-b" + with pytest.raises(common.StreamConfigError, match="client with one backend"): + worker_setup.run_worker(migrating, task_queue()) diff --git a/tests/contrib/streams/examples/ui_backend_reads.py b/tests/contrib/streams/examples/ui_backend_reads.py new file mode 100644 index 0000000000..068addf6ff --- /dev/null +++ b/tests/contrib/streams/examples/ui_backend_reads.py @@ -0,0 +1,90 @@ +"""Brook API proposal, "A UI backend reads a Workflow's Stream". + +The agent Workflow hands out its ``turns`` Stream through a ``stream`` query. The UI backend +starts it, gets the Stream from the query, and renders each item as it arrives. A Stream is +data, so the client reads it as a ``ClientStream``. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +from temporalio import activity, common, workflow +from temporalio.client import Client +from temporalio.worker import Worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@activity.defn +async def call_model(prompt: str) -> str: + return prompt.upper() + + +@workflow.defn +class Agent: + @workflow.query + def stream(self) -> common.WorkflowStream[Turn]: + # The handle makes no call, so a reader can get it before the create. + return workflow.stream_handle(Turns) + + @workflow.run + async def run(self, prompt: str) -> str: + turns = await workflow.create_stream(Turns) + await turns.publish(Turn("model")) + reply = await workflow.execute_activity( + call_model, prompt, start_to_close_timeout=timedelta(minutes=5) + ) + await turns.publish(Turn("done", reply)) + # A run's end doesn't close the Stream, so the Workflow closes it before it returns. + await turns.close(result=reply) + return reply + + +class Screen: + """A fake UI. It keeps what it shows by offset, so a withdrawal can remove items.""" + + def __init__(self) -> None: + self.shown: dict[int, Turn] = {} + self.result: Any = None + + def render(self, record: common.Record[Turn]) -> None: + self.shown[record.offset] = record.data + + def unrender(self, offsets: tuple[int, ...]) -> None: + for offset in offsets: + del self.shown[offset] + + +async def main(client: Client, task_queue: str) -> Screen: + async with Worker( + client, task_queue=task_queue, workflows=[Agent], activities=[call_model] + ): + handle = await client.start_workflow( + Agent.run, "hello there", id=f"agent-{uuid.uuid4()}", task_queue=task_queue + ) + # The Workflow hands out its Stream. + stream = await handle.query("stream", result_type=common.ClientStream[Turn]) + screen = Screen() + # The Workflow may not have created the Stream yet, so the read waits for it. + async for r in stream.subscribe(create_wait=timedelta(seconds=10)): + match r: + case common.Closed(): + screen.result = r.result if r.failure is None else r.failure + case common.Superseded(): + # The offsets of the Records that this read showed and the item withdraws. + screen.unrender(r.withdrawn_offsets) + case _: + screen.render(r) + return screen diff --git a/tests/contrib/streams/examples/worker_setup.py b/tests/contrib/streams/examples/worker_setup.py new file mode 100644 index 0000000000..a83f17d128 --- /dev/null +++ b/tests/contrib/streams/examples/worker_setup.py @@ -0,0 +1,88 @@ +"""Brook API proposal, "Worker setup". + +Register the Stream backends on the client with ``stream_backends=``. The Worker uses its +client's backend, and serves only one. Each backend registers a name. A create with no backend +name uses the backend with ``default=True``, and ``backend=`` picks another. Tests use +``MemoryStreamBackend`` without changing Workflow code. Sync Workflow code, such as +``@workflow.init``, uses the ``_nowait`` calls. +""" + +from __future__ import annotations + +import uuid + +from temporalio import common, workflow +from temporalio.client import Client +from temporalio.contrib.stream import RedisStreamBackend +from temporalio.contrib.stream.testing import MemoryStreamBackend +from temporalio.worker import Worker + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +@workflow.defn +class Greeter: + @workflow.init + def __init__(self, name: str) -> None: + # The init can't await, so it creates and publishes with the _nowait calls. + self.tokens = workflow.create_stream_nowait(Tokens) + self.tokens.publish_nowait(f"hello {name}") + + @workflow.query + def stream(self) -> common.WorkflowStream[str]: + return self.tokens + + @workflow.run + async def run(self, name: str) -> str: + await self.tokens.publish(f"bye {name}") + await self.tokens.close() + return name + + +async def connect_with_redis(target_host: str, redis_url: str) -> Client: + # Brook has no native backend, so its one Redis backend is the default. + return await Client.connect( + target_host, stream_backends=[RedisStreamBackend(redis_url, default=True)] + ) + + +async def connect_for_tests(target_host: str) -> Client: + return await Client.connect( + target_host, stream_backends=[MemoryStreamBackend(default=True)] + ) + + +async def connect_while_migrating(target_host: str, url_a: str, url_b: str) -> Client: + # Two backends during a migration. A create with no backend= uses the default one. A + # Worker serves only one backend, so a Worker gets a client with one backend. + return await Client.connect( + target_host, + stream_backends=[ + RedisStreamBackend(url_a, name="redis-a", default=True), + RedisStreamBackend(url_b, name="redis-b"), + ], + ) + + +def run_worker(client: Client, task_queue: str) -> Worker: + # The Worker uses its client's one backend, so it needs no Stream setup of its own. + return Worker(client, task_queue=task_queue, workflows=[Greeter]) + + +async def main(client: Client, task_queue: str) -> list[str]: + async with run_worker(client, task_queue): + handle = await client.start_workflow( + Greeter.run, "ada", id=f"greeter-{uuid.uuid4()}", task_queue=task_queue + ) + await handle.result() + stream = await handle.query("stream", result_type=common.ClientStream[str]) + return [r.data async for r in stream.subscribe() if isinstance(r, common.Record)] + + +async def create_on_second_backend(client: Client) -> str | None: + stream = await client.create_stream( + Tokens, id=f"chat-{uuid.uuid4()}", backend="redis-b" + ) + return stream.backend diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index 02d249ad82..cc7a056799 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -1,21 +1,62 @@ +import asyncio +import dataclasses +import gc import uuid +from collections.abc import Generator, Sequence from dataclasses import dataclass from datetime import timedelta -from typing import Any +from types import SimpleNamespace +from typing import Any, cast + +import pytest from temporalio import activity, common, workflow -from temporalio.client import Client +from temporalio.api.common.v1 import Payload +from temporalio.api.enums.v1 import EventType, StreamErrorReason +from temporalio.client import Client, Interceptor, WorkflowFailureError, WorkflowHandle +from temporalio.common import RawValue, RetryPolicy +from temporalio.common._streams._extension import ( + ExtensionStreamBackend, + _WorkflowInbound, +) +from temporalio.contrib.stream import RedisStreamBackend +from temporalio.contrib.stream.testing import MemoryStreamBackend +from temporalio.converter import DataConverter, PayloadCodec +from temporalio.exceptions import ApplicationError from temporalio.worker import ( + ExecuteWorkflowInput, + Replayer, UnsandboxedWorkflowRunner, + WorkflowInboundInterceptor, + WorkflowInterceptorClassInput, +) +from temporalio.worker import Interceptor as WorkerInterceptor +from temporalio.worker.workflow_sandbox import ( + SandboxedWorkflowRunner, + SandboxRestrictions, ) from tests.contrib.streams._support import workflow_stream from tests.helpers import new_worker +@dataclass +class Turn: + kind: str + text: str = "" + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + @common.defn(name="tokens") class Tokens(common.StreamDefinition[str]): ... +@common.defn(name="scoped") +class Scoped(common.StreamDefinition[str]): ... + + @dataclass class CallModelInput: prompt: str @@ -30,6 +71,31 @@ def records(items: list[Any]) -> list[common.Record[Any]]: return [item for item in items if isinstance(item, common.Record)] +@activity.defn +async def call_model(input: CallModelInput) -> str: + for word in input.prompt.split(): + await input.stream.publish(word) + return input.prompt.upper() + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + turns = await workflow.create_stream(Turns) + tokens = await workflow.create_stream(Tokens) + await turns.publish(Turn("model")) + reply = await workflow.execute_activity( + call_model, + CallModelInput(prompt, common.ActivityStream.from_handle(tokens)), + start_to_close_timeout=timedelta(seconds=30), + ) + await turns.publish(Turn("done", reply)) + await turns.close(result=reply) + await tokens.close() + return reply + + attempts: dict[str, int] = {} @@ -73,6 +139,163 @@ async def test_failed_workflow_task_is_withdrawn(stream_client: Client): assert isinstance(items[-1], common.Closed) +async def test_resume_after_a_superseded_item_keeps_its_record(stream_client: Client): + key = str(uuid.uuid4()) + async with new_worker( + stream_client, + PublishThenFailTask, + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await stream_client.start_workflow( + PublishThenFailTask.run, + key, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + stream = workflow_stream(handle, Tokens) + items = [item async for item in stream.subscribe(wait_new_records=False)] + # The retried task's batch has its Superseded item and its first Record in one entry. + record_offsets = {r.offset for r in records(items)} + [shared] = [ + i + for i in items + if isinstance(i, common.Superseded) and i.offset in record_offsets + ] + reader = stream.subscribe(wait_new_records=False) + async for item in reader: + if item == shared: + break + resume_at = reader.next_offset + await reader.aclose() + assert resume_at == shared.offset + resumed = [ + item + async for item in stream.subscribe(offset=resume_at, wait_new_records=False) + ] + assert [r.data for r in records(resumed) if r.offset == shared.offset] == ["after"] + + +@activity.defn +async def noop() -> None: + pass + + +@workflow.defn +class RandomAroundActivity: + @workflow.run + async def run(self) -> bool: + before = workflow.random().getstate() + handle = workflow.start_activity( + noop, start_to_close_timeout=timedelta(seconds=30) + ) + same = workflow.random().getstate() == before + await handle + return same + + +async def test_backend_leaves_the_workflow_random_alone(stream_client: Client): + async with new_worker( + stream_client, RandomAroundActivity, activities=[noop] + ) as worker: + assert await stream_client.execute_workflow( + RandomAroundActivity.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + + +@workflow.defn +class WritesThenContinues: + @workflow.run + async def run(self, write: bool) -> str: + if write: + stream = await workflow.create_stream(Tokens) + await stream.publish("first run") + workflow.continue_as_new(False) + await workflow.stream_handle(Tokens).close(result="second run") + return "second run" + + +async def test_later_run_closes_a_stream_it_did_not_write(stream_client: Client): + async with new_worker(stream_client, WritesThenContinues) as worker: + handle = await stream_client.start_workflow( + WritesThenContinues.run, + True, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == "second run" + items = await read_all(workflow_stream(handle, Tokens)) + assert [i.data for i in records(items)] == ["first run"] + closed = items[-1] + assert isinstance(closed, common.Closed) and closed.result == "second run" + + +@workflow.defn +class CancelledBlock: + @workflow.run + async def run(self) -> None: + async with await workflow.create_stream(Scoped) as turn: + await turn.publish("started") + await workflow.wait_condition(lambda: False) + + +async def test_cancelled_block_closes_with_a_failure(stream_client: Client): + async with new_worker(stream_client, CancelledBlock) as worker: + handle = await stream_client.start_workflow( + CancelledBlock.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + turn = workflow_stream(handle, Scoped) + reader = turn.subscribe(create_wait=timedelta(seconds=20)) + assert isinstance(await anext(reader), common.Record) + await reader.aclose() + await handle.cancel() + items = await read_all(turn) + closed = items[-1] + assert isinstance(closed, common.Closed) + assert closed.failure is not None + assert closed.status == common.CloseStatus.CANCELED + + +@common.defn(name="plain") +class Plain(common.StreamDefinition[str]): ... + + +@workflow.defn +class PlainWriter: + @workflow.run + async def run(self) -> str: + stream = await workflow.create_stream(Plain) + await stream.publish("plain") + async with await workflow.create_stream(Scoped) as scoped: + await scoped.publish("scoped") + await stream.close(result="done") + return "done" + + +async def test_temporal_only_passthrough_keeps_records(stream_client: Client): + # Only the minimum modules and temporalio pass through. Values are plain, since a + # dataclass from a sandbox with its own dataclasses module can't convert outside it. + runner = SandboxedWorkflowRunner( + restrictions=dataclasses.replace( + SandboxRestrictions.default, + passthrough_modules=SandboxRestrictions.passthrough_modules_with_temporal, + ) + ) + async with new_worker(stream_client, PlainWriter, workflow_runner=runner) as worker: + handle = await stream_client.start_workflow( + PlainWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + items = await read_all(workflow_stream(handle, Plain)) + scoped = await read_all(workflow_stream(handle, Scoped)) + assert [i.data for i in records(items)] == ["plain"] + assert isinstance(items[-1], common.Closed) and items[-1].result == "done" + assert [i.data for i in records(scoped)] == ["scoped"] + assert isinstance(scoped[-1], common.Closed) + + @activity.defn async def local_tokens(input: CallModelInput) -> None: await input.stream.publish(*input.prompt.split()) @@ -165,3 +388,456 @@ async def test_later_write_withdraws_a_skipped_local_activity(stream_client: Cli 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) + + +class _Rot13Codec(PayloadCodec): + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + return [ + Payload( + metadata={"encoding": b"binary/rot"}, data=p.SerializeToString()[::-1] + ) + for p in payloads + ] + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + out = [] + for p in payloads: + if p.metadata.get("encoding") != b"binary/rot": + out.append(p) + continue + inner = Payload() + inner.ParseFromString(p.data[::-1]) + out.append(inner) + return out + + +@common.defn(name="tokens") +class RawTokens(common.StreamDefinition[RawValue]): ... + + +async def test_records_go_through_the_payload_codec(stream_client: Client): + config = stream_client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, payload_codec=_Rot13Codec() + ) + codec_client = Client(**config) + async with new_worker(codec_client, Agent, activities=[call_model]) as worker: + handle = await codec_client.start_workflow( + Agent.run, "a b", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + turns = await read_all(workflow_stream(handle, Turns)) + tokens = workflow_stream(handle, Tokens) + decoded = await read_all(tokens) + # A definition with RawValue as its payload type reads the raw Payloads. + raw_tokens = codec_client.stream(RawTokens, id=tokens.id) + raw = [item async for item in raw_tokens.subscribe(wait_new_records=False)] + assert [item.data for item in records(turns)] == [ + Turn("model"), + Turn("done", "A B"), + ] + assert [item.data for item in records(decoded)] == ["a", "b"] + first = raw[0] + assert isinstance(first, common.Record) + assert first.data.payload.metadata["encoding"] == b"json/plain" + + +async def test_client_close_goes_through_the_payload_codec(stream_client: Client): + config = stream_client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, payload_codec=_Rot13Codec() + ) + codec_client = Client(**config) + done = await codec_client.create_stream(Tokens, id=f"done-{uuid.uuid4()}") + await done.close(result=Turn("done")) + failed = await codec_client.create_stream(Tokens, id=f"failed-{uuid.uuid4()}") + await failed.close(failure=ApplicationError("bad turn")) + [closed] = [item async for item in done.subscribe()] + assert isinstance(closed, common.Closed) + assert closed.result == {"kind": "done", "text": ""} + [closed] = [item async for item in failed.subscribe()] + assert isinstance(closed, common.Closed) and str(closed.failure) == "bad turn" + + +@workflow.defn +class WaitsForSignal: + def __init__(self) -> None: + self.go = False + + @workflow.run + async def run(self) -> str: + stream = await workflow.create_stream(Tokens) + await stream.publish("before restart") + async with await workflow.create_stream(Scoped) as scoped: + await scoped.publish("scoped") + await workflow.wait_condition(lambda: self.go) + await stream.publish("after restart") + await stream.close() + return "done" + + @workflow.signal + def proceed(self) -> None: + self.go = True + + +async def test_new_worker_replays_and_writes_each_record_once(stream_client: Client): + task_queue = f"tq-{uuid.uuid4()}" + async with new_worker(stream_client, WaitsForSignal, task_queue=task_queue): + handle = await stream_client.start_workflow( + WaitsForSignal.run, id=f"wf-{uuid.uuid4()}", task_queue=task_queue + ) + first = workflow_stream(handle, Tokens) + reader = first.subscribe(create_wait=timedelta(seconds=20)) + assert isinstance(await anext(reader), common.Record) + await reader.aclose() + # The new Worker has no cache, so it replays the run and sends the creates and the close + # of the scoped Stream again. + async with new_worker(stream_client, WaitsForSignal, task_queue=task_queue): + await handle.signal(WaitsForSignal.proceed) + assert await handle.result() == "done" + items = await read_all(first) + scoped = await read_all(workflow_stream(handle, Scoped)) + assert [i.data for i in records(items)] == ["before restart", "after restart"] + assert isinstance(items[-1], common.Closed) + assert [i.data for i in records(scoped)] == ["scoped"] + assert [type(i) for i in scoped].count(common.Closed) == 1 + + +async def test_backend_replays_a_history_made_without_it( + client: Client, stream_client: Client +): + async with new_worker(client, RandomAroundActivity, activities=[noop]) as worker: + handle = await client.start_workflow( + RandomAroundActivity.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() + history = await handle.fetch_history() + replayer = Replayer( + workflows=[RandomAroundActivity], + stream_backends=stream_client.config()["stream_backends"], + ) + await replayer.replay_workflow(history) + + +async def test_backend_shutdown_after_client_writes(stream_client: Client): + [backend] = stream_client.config()["stream_backends"] + assert isinstance(backend, ExtensionStreamBackend) + stream = await stream_client.create_stream(Tokens, id=f"shutdown-{uuid.uuid4()}") + async with common.BufferedPublisher(stream) as tokens: + tokens.publish("buffered") + await backend.shutdown() + written = [ + item.data + async for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + assert written == ["buffered"] + + +@workflow.defn +class WaitsForever: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream(Tokens) + await stream.publish("waiting") + await workflow.wait_condition(lambda: False) + + +class _Suspends: + def __await__(self) -> Generator[None, None, None]: + yield + + +class _NeverEnds(WorkflowInboundInterceptor): + def __init__(self) -> None: # pyright: ignore[reportMissingSuperCall] + pass + + async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any: + await _Suspends() + + +def test_closing_an_evicted_run_touches_no_runtime(): + # Python closes an evicted run's coroutine with GeneratorExit, outside the Workflow's loop. + # The interceptor must let it pass without a call to the runtime. + inbound = _WorkflowInbound(_NeverEnds()) + run = inbound.execute_workflow( + cast(ExecuteWorkflowInput, cast(object, SimpleNamespace(headers={}, args=[]))) + ) + run.send(None) + run.close() + + +@pytest.mark.parametrize("cache, safe", [(0, True), (10, True), (10, False)]) +async def test_eviction_prints_nothing( + stream_client: Client, capfd: pytest.CaptureFixture[str], cache: int, safe: bool +): + # With no cache, each Workflow Task evicts the run. With a cache, the Worker's shutdown + # evicts it. Either way the run's coroutine closes. + async with new_worker( + stream_client, + WaitsForever, + max_cached_workflows=cache, + disable_safe_workflow_eviction=not safe, + ) as w: + handle = await stream_client.start_workflow( + WaitsForever.run, id=f"wf-{uuid.uuid4()}", task_queue=w.task_queue + ) + reader = workflow_stream(handle, Tokens).subscribe( + create_wait=timedelta(seconds=20) + ) + assert isinstance(await anext(reader), common.Record) + await reader.aclose() + await handle.terminate() + gc.collect() + captured = capfd.readouterr() + assert "Exception ignored" not in captured.err + captured.out + + +@workflow.defn +class AlwaysFails: + @workflow.run + async def run(self, non_retryable: bool) -> None: + attempt = workflow.info().attempt + stream = await workflow.create_stream(Tokens) + await stream.publish(f"run-{attempt}") + raise ApplicationError(f"attempt {attempt} fails", non_retryable=non_retryable) + + +@pytest.mark.parametrize("non_retryable", [False, True]) +async def test_final_failure_under_a_retry_policy( + stream_client: Client, non_retryable: bool +): + async with new_worker(stream_client, AlwaysFails) as worker: + handle = await stream_client.start_workflow( + AlwaysFails.run, + non_retryable, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + retry_policy=RetryPolicy( + initial_interval=timedelta(milliseconds=10), maximum_attempts=3 + ), + ) + with pytest.raises(WorkflowFailureError): + await handle.result() + stream = workflow_stream(handle, Tokens) + items = [item async for item in stream.subscribe(wait_new_records=False)] + written = records(items) + superseded = [item for item in items if isinstance(item, common.Superseded)] + # The end of the chain doesn't close the Stream. + assert not any(isinstance(item, common.Closed) for item in items) + if non_retryable: + # A non-retryable failure ends the chain on attempt 1. + assert [r.data for r in written] == ["run-1"] + assert superseded == [] + return + # Each retry run withdraws the Records of the run before it. + assert [r.data for r in written] == ["run-1", "run-2", "run-3"] + assert len(superseded) == 2 + assert [r.data for r in written if superseded[0].withdraws(r)] == ["run-1"] + assert [r.data for r in written if superseded[1].withdraws(r)] == ["run-2"] + assert not any(s.withdraws(written[-1]) for s in superseded) + + +async def test_reader_stops_on_its_own_timer_and_resumes(stream_client: Client): + stream = await stream_client.create_stream(Tokens, id=f"timer-{uuid.uuid4()}") + first = await stream.publish("one") + seen: list[common.Record[Any]] = [] + + async def read_until_quiet() -> None: + async for item in stream.subscribe(): + assert isinstance(item, common.Record) + seen.append(item) + + # A read has no idle timeout, so a reader that wants one uses its own timer. + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(read_until_quiet(), timeout=0.5) + assert [item.offset for item in seen] == [first] + await stream.publish("two") + resumed = [ + item.data + async for item in stream.subscribe( + offset=seen[-1].offset + 1, wait_new_records=False + ) + if isinstance(item, common.Record) + ] + assert resumed == ["two"] + + +async def test_two_backends_on_one_client(client: Client, redis_url: str): + # A client can have two Stream backends, such as during a migration to another backend. Each + # handle stores its backend name, so a call goes to the backend that has the Stream. + config = client.config() + config["stream_backends"] = [ + RedisStreamBackend(redis_url, key_prefix=f"t-{uuid.uuid4()}", default=True), + MemoryStreamBackend(), + ] + stream_client = Client(**config) + stream_id = f"two-{uuid.uuid4()}" + on_redis = await stream_client.create_stream(Tokens, id=stream_id) + on_memory = await stream_client.create_stream( + Tokens, id=stream_id, backend="memory" + ) + assert (on_redis.backend, on_memory.backend) == ("redis", "memory") + await on_redis.publish("redis") + await on_memory.publish("memory") + for backend, expected in (("redis", "redis"), ("memory", "memory")): + again = stream_client.stream(Tokens, id=stream_id, backend=backend) + data = [ + item.data + async for item in again.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + assert data == [expected] + + +failed_last_task: set[str] = set() + + +class _FailsTheLastTaskOnce(Interceptor, WorkerInterceptor): + # It's on the client before the Stream backend's interceptor, so it's the outer Workflow + # interceptor. It fails the task after the backend drained the run's writes. + def workflow_interceptor_class( + self, input: WorkflowInterceptorClassInput + ) -> type[WorkflowInboundInterceptor]: + return _FailsTheLastTaskOnceInbound + + +class _FailsTheLastTaskOnceInbound(WorkflowInboundInterceptor): + async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any: + result = await super().execute_workflow(input) + workflow_id = workflow.info().workflow_id + if workflow_id not in failed_last_task: + failed_last_task.add(workflow_id) + raise RuntimeError("fail the last Workflow Task once") + return result + + +@workflow.defn +class PublishesClosesAndReturns: + @workflow.run + async def run(self) -> str: + stream = await workflow.create_stream(Tokens) + await stream.publish("token") + await stream.close(result="done") + return "done" + + +async def close_activity_events( + handle: WorkflowHandle[Any, Any], +) -> list[EventType.ValueType]: + # The events of the close Activity, from its schedule on. + history = await handle.fetch_history() + scheduled = { + event.event_id + for event in history.events + if event.activity_task_scheduled_event_attributes.activity_type.name + == "__temporal_streams_close" + } + return [ + event.event_type + for event in history.events + if event.event_id in scheduled + or event.activity_task_started_event_attributes.scheduled_event_id in scheduled + or event.activity_task_completed_event_attributes.scheduled_event_id + in scheduled + ] + + +@pytest.mark.parametrize("cache", [0, 10]) +async def test_close_survives_a_failed_last_task(stream_client: Client, cache: int): + # The close is an Activity, committed with the Workflow Task that scheduled it. The last + # task fails after it, and its retry replays the close and writes nothing more. + config = stream_client.config() + config["interceptors"] = [_FailsTheLastTaskOnce()] + client = Client(**config) + async with new_worker( + client, + PublishesClosesAndReturns, + max_cached_workflows=cache, + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await client.start_workflow( + PublishesClosesAndReturns.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == "done" + stream = workflow_stream(handle, Tokens) + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert [r.data for r in records(items)] == ["token"] + closes = [item for item in items if isinstance(item, common.Closed)] + assert len(closes) == 1 and items[-1] is closes[0] + assert closes[0].result == "done" + assert closes[0].status == common.CloseStatus.COMPLETED + assert await close_activity_events(handle) == [ + EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED, + EventType.EVENT_TYPE_ACTIVITY_TASK_STARTED, + EventType.EVENT_TYPE_ACTIVITY_TASK_COMPLETED, + ] + + +@workflow.defn +class WritesToStream: + @workflow.run + async def run(self, stream: common.WorkflowStream[str]) -> str: + try: + await stream.publish("from workflow") + except common.StreamError as err: + return StreamErrorReason.Name(err.reason) + return "wrote" + + +async def test_write_to_a_closed_stream_raises_at_the_await(stream_client: Client): + stream = await stream_client.create_stream(Tokens, id=f"closed-{uuid.uuid4()}") + await stream.close() + async with new_worker(stream_client, WritesToStream) as worker: + assert ( + await stream_client.execute_workflow( + WritesToStream.run, + common.WorkflowStream.from_handle(stream), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + == "STREAM_ERROR_REASON_CLOSED" + ) + items = [item async for item in stream.subscribe()] + assert len(items) == 1 and isinstance(items[0], common.Closed) + + +@workflow.defn +class ContinuesThenRetries: + @workflow.run + async def run(self, stream: common.WorkflowStream[str] | None) -> str: + if stream is None: + stream = await workflow.create_stream(Tokens) + await stream.publish("first") + workflow.continue_as_new(stream) + info = workflow.info() + await stream.publish(f"attempt-{info.attempt}") + if info.attempt == 1: + raise ApplicationError("the continued run fails once") + await stream.close(result="done") + return "done" + + +async def test_retry_of_a_continued_run_writes_the_same_stream(stream_client: Client): + async with new_worker(stream_client, ContinuesThenRetries) as worker: + handle = await stream_client.start_workflow( + ContinuesThenRetries.run, + None, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + retry_policy=RetryPolicy( + initial_interval=timedelta(milliseconds=10), maximum_attempts=2 + ), + ) + assert await handle.result() == "done" + items = await read_all(workflow_stream(handle, Tokens)) + written = records(items) + assert [r.data for r in written] == ["first", "attempt-1", "attempt-2"] + superseded = [item for item in items if isinstance(item, common.Superseded)] + assert [r.data for r in written if superseded[0].withdraws(r)] == ["attempt-1"] + closes = [item for item in items if isinstance(item, common.Closed)] + assert len(closes) == 1 and closes[0].result == "done"