diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index f8a51d9b7..7b78b0c96 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -12,3 +12,6 @@ # manage repo-wide concerns. /temporalio/contrib/workflow_streams/ @temporalio/ai-sdk @temporalio/sdk /tests/contrib/workflow_streams/ @temporalio/ai-sdk @temporalio/sdk +/temporalio/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/meandering-river.md b/changelog/added/meandering-river.md new file mode 100644 index 000000000..ec7aca925 --- /dev/null +++ b/changelog/added/meandering-river.md @@ -0,0 +1,2 @@ +Add experimental Streams. A Stream is an ordered log of Records that Workflows, Activities and clients write, and that clients and Activities read. `temporalio.streams` has the shared types, such as the `StreamRef` data, `temporalio.workflow.WorkflowStream` makes the calls in Workflow code, and `temporalio.client.Stream` makes them elsewhere. Register a backend with `stream_backends=` on `Client.connect`, `Replayer` or `WorkflowEnvironment`, then use `workflow.create_stream[T]`, `Client.create_stream[T]` or `Client.get_stream[T]`, and `to_stream` for a ref that came as data. +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/streams/__init__.py b/temporalio/streams/__init__.py index a4c2a6a6e..a7eee1595 100644 --- a/temporalio/streams/__init__.py +++ b/temporalio/streams/__init__.py @@ -1,5 +1,34 @@ """The Streams API, which Workflows, Activities and clients share. +A Stream is an ordered log of Records that Workflows, Activities and clients write, and that +clients and Activities read. A Stream has one payload type, and a call takes it as a type +argument, as in ``client.get_stream[Turn]("chat-42")``. With no type argument, the payload type +is ``Any``, and a read gives JSON values. + +Code creates a Stream before its first publish. :py:data:`temporalio.workflow.create_stream` and +:py:attr:`temporalio.client.Client.create_stream` fail with :py:class:`StreamAlreadyExistsError` if +the id exists, and ``create_or_get_stream`` gets the Stream instead. In Workflow code, each run of a +chain sends the same create for an id, so a retry run or a cron run gets the Stream that an earlier +run made. Workflow code gets a :py:class:`temporalio.workflow.WorkflowStream`, and other code a +:py:class:`temporalio.client.Stream`. A client gets any Stream by its id with +:py:attr:`temporalio.client.Client.get_stream`, with no call. + +A Stream's :py:class:`StreamRef` is data with no calls. Code passes it as an argument to a +Workflow, an Activity, a child Workflow or the next run, or hands it out through a query. +:py:func:`temporalio.workflow.to_stream` and :py:meth:`temporalio.client.Client.to_stream` give +the Stream of a ref. In an Activity, a Stream from ``activity.client()`` writes as the attempt. + +A read gives :py:class:`Record` and :py:class:`Superseded` items, and the ``kind`` of each item +narrows its type. A close carries no value. A read of a closed Stream ends once it gives the last +item. + +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 ref keeps 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. """ diff --git a/tests/contrib/streams/examples/README.md b/tests/contrib/streams/examples/README.md new file mode 100644 index 000000000..bc2dea4f4 --- /dev/null +++ b/tests/contrib/streams/examples/README.md @@ -0,0 +1,15 @@ +# Examples for the 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 Stream by its id | [`ui_backend_reads.py`](ui_backend_reads.py) | +| A UI backend reads a Stream by its id, its resume notes | [`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 000000000..e69de29bb 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 000000000..0eac15b36 --- /dev/null +++ b/tests/contrib/streams/examples/activity_streams_tokens.py @@ -0,0 +1,137 @@ +"""Streams API proposal, "An Activity streams tokens to the Workflow's Stream". + +The Workflow passes its Stream's ref to the Activity as a field of the input. An async Activity +gets the Stream from ``activity.client()``, so it writes as the attempt, and batches the model's +deltas with ``BufferedPublisher``. A sync Activity has no ``activity.client()``, so it keeps the +Worker's client. It wraps its Stream in ``SyncStream``, whose calls block its thread, so it works +in a thread pool executor. +""" + +from __future__ import annotations + +import concurrent.futures +import uuid +from collections.abc import AsyncIterator +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import activity, streams, workflow +from temporalio.client import BufferedPublisher, Client, SyncStream +from temporalio.worker import Worker + +Message = str + + +@dataclass +class Reply: + text: str + + +class FakeResponse: + """A model's streamed response, one delta for each word.""" + + def __init__(self, messages: list[Message]) -> None: + self.words = messages[-1].split() + + async def __aiter__(self) -> AsyncIterator[str]: + for word in self.words: + yield word + + def reply(self) -> Reply: + return Reply(" ".join(self.words).upper()) + + +class FakeLLM: + """Stands in for a model client.""" + + def stream(self, messages: list[Message]) -> FakeResponse: + return FakeResponse(messages) + + +llm = FakeLLM() + + +@dataclass +class CallModelInput: + messages: list[Message] + tokens: streams.StreamRef[str] # the ref that the Workflow passes + + +@dataclass +class SummarizeInput: + text: str + tokens: streams.StreamRef[str] + + +@activity.defn +async def call_model(input: CallModelInput) -> Reply: + stream = activity.client().to_stream(input.tokens) # writes as this attempt + response = llm.stream(input.messages) + async with BufferedPublisher(stream) as tokens: # batches deltas + async for delta in response: + tokens.publish(delta) + return response.reply() + + +class Summarizer: + def __init__(self, client: Client) -> None: + self._client = client # the Worker's client + + @activity.defn + def summarize(self, input: SummarizeInput) -> int: + tokens = SyncStream(self._client.to_stream(input.tokens)) + # The sync publish blocks until the Record is written, and returns its offset. + return tokens.publish(f"summary: {input.text}") + + +@workflow.defn +class Model: + def __init__(self) -> None: + self._tokens: workflow.WorkflowStream[str] | None = None + + @workflow.query + def stream(self) -> streams.StreamRef[str] | None: + return self._tokens.ref if self._tokens else None + + @workflow.run + async def run(self, prompt: str) -> int: + self._tokens = tokens = await workflow.create_stream[str]( + f"{workflow.info().workflow_id}-tokens" + ) + reply = await workflow.execute_activity( + call_model, + CallModelInput([prompt], tokens.ref), + start_to_close_timeout=timedelta(minutes=5), + ) + summary = await workflow.execute_activity_method( + Summarizer.summarize, + SummarizeInput(reply.text, tokens.ref), + start_to_close_timeout=timedelta(minutes=5), + ) + await tokens.close() + return summary + + +async def main( + client: Client, task_queue: str +) -> tuple[int, list[streams.Record[str]]]: + with concurrent.futures.ThreadPoolExecutor() as executor: + async with Worker( + client, + task_queue=task_queue, + workflows=[Model], + activities=[call_model, Summarizer(client).summarize], + activity_executor=executor, + ): + handle = await client.start_workflow( + Model.run, + "one two three", + id=f"model-{uuid.uuid4()}", + task_queue=task_queue, + ) + summary = await handle.result() + ref = await handle.query(Model.stream) + assert ref is not None + stream = client.to_stream(ref) + records = [r async for r in stream.subscribe() if isinstance(r, streams.Record)] + return summary, 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 000000000..1b1a7fdec --- /dev/null +++ b/tests/contrib/streams/examples/agent_publishes_turns.py @@ -0,0 +1,136 @@ +"""Streams API proposal, "An agent Workflow publishes its turns". + +The Workflow creates its two Streams, turns and token deltas, under ids of its choice, and hands +out the turns Stream's ref through a ``stream`` query. It publishes each turn with an awaited +``publish``, and passes the tokens ref to the model Activity in its one input object. A run +after continue-as-new gets the earlier run's refs and creates nothing. 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, streams, workflow +from temporalio.client import BufferedPublisher, Client +from temporalio.worker import Worker + +MINUTE = timedelta(minutes=1) + +Message = str + + +@dataclass +class Turn: + kind: str + text: str | None = None + + +@dataclass +class Reply: + text: str + tool_calls: list[str] = field(default_factory=list) + + +@dataclass +class CallModelInput: + messages: list[Message] + tokens: streams.StreamRef[str] # the ref that the Workflow passes + + +@activity.defn +async def call_model(input: CallModelInput) -> Reply: + stream = activity.client().to_stream(input.tokens) # writes as this attempt + async with BufferedPublisher(stream) as tokens: # batches deltas + for delta in input.messages[-1].split(): + tokens.publish(delta) + if len(input.messages) == 1: + return Reply("", tool_calls=["search"]) + return Reply("the answer") + + +@activity.defn +async def run_tool(call: str) -> Message: + return f"{call} result" + + +@dataclass +class AgentInput: + session_id: str + messages: list[Message] + # The run before Continue-as-New sets the refs. + turns: streams.StreamRef[Turn] | None = None + tokens: streams.StreamRef[str] | None = None + + +@workflow.defn +class Agent: + def __init__(self) -> None: + self._turns: workflow.WorkflowStream[Turn] | None = None # run creates it + + @workflow.query + def stream(self) -> streams.StreamRef[Turn] | None: # one way to hand it out + return self._turns.ref if self._turns else None # None before the create + + @workflow.run + async def run(self, input: AgentInput) -> str: + messages = input.messages + # A Workflow creates only standalone Streams, under an id of its choice. A run after + # Continue-as-New gets the earlier run's refs and creates nothing. + if input.turns and input.tokens: + self._turns = turns = workflow.to_stream(input.turns) + tokens = workflow.to_stream(input.tokens) + else: + self._turns = turns = await workflow.create_stream[Turn]( + f"{input.session_id}-turns" + ) # the default Retention + # Token deltas matter only to live readers, so they get a short Retention. + tokens = await workflow.create_stream[str]( + f"{input.session_id}-tokens", retention=timedelta(hours=1) + ) + while True: + if workflow.info().is_continue_as_new_suggested(): + workflow.continue_as_new( # forward the refs, so the next run keeps the Streams + AgentInput(input.session_id, messages, turns.ref, tokens.ref) + ) + await turns.publish(Turn("model")) # one Local Activity + reply = await workflow.execute_activity( + call_model, + CallModelInput(messages, tokens.ref), + start_to_close_timeout=MINUTE, + ) + if not reply.tool_calls: + await turns.publish(Turn("done", reply.text)) + await turns.close() # the done Turn carries the reply + 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=MINUTE + ) + ) + + +async def main(client: Client, task_queue: str) -> list[streams.Record[Turn]]: + async with Worker( + client, + task_queue=task_queue, + workflows=[Agent], + activities=[call_model, run_tool], + ): + session_id = f"session-{uuid.uuid4()}" + handle = await client.start_workflow( + Agent.run, + AgentInput(session_id, ["find it"]), + id=session_id, + task_queue=task_queue, + ) + await handle.result() + # The Workflow hands out its turns ref, and the query's client reads the Stream. + ref = await handle.query(Agent.stream) + assert ref is not None + stream = client.to_stream(ref, payload_type=Turn) + return [r async for r in stream.subscribe() if isinstance(r, streams.Record)] 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 000000000..796366c9f --- /dev/null +++ b/tests/contrib/streams/examples/client_creates_stream.py @@ -0,0 +1,74 @@ +"""Streams API proposal, "A client creates a Stream and passes it to a Workflow". + +The client creates the Stream with its settings, passes its ref to the Workflow as data, reads +it, and closes it at the end of ``async with``. The create fails if the id exists, and +``create_or_get_stream`` would get the Stream instead. The Workflow gets the ref as data, turns +it into a Stream to write, and passes the same ref to an Activity. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import activity, streams, workflow +from temporalio.client import Client +from temporalio.worker import Worker + + +@dataclass +class Turn: + kind: str + text: str | None = None + + +SECONDS_30 = timedelta(seconds=30) + + +@activity.defn +async def notify_user(ref: streams.StreamRef[Turn]) -> str: + # A real Activity sends the Stream id to the user's device, which reads it by id. + return ref.id + + +@workflow.defn +class Chat: + @workflow.run + async def run(self, ref: streams.StreamRef[Turn]) -> None: + # The Workflow only writes. The client closes the Stream. + turns = workflow.to_stream(ref) + turns.publish(Turn("model", "hi")) + await workflow.execute_activity( + notify_user, ref, start_to_close_timeout=SECONDS_30 + ) + await turns.publish(Turn("last", "bye")) + + +def is_last_turn(r: streams.Record[Turn]) -> bool: + return r.data.kind == "last" + + +async def main(client: Client, task_queue: str) -> list[Turn]: + async with Worker( + client, task_queue=task_queue, workflows=[Chat], activities=[notify_user] + ): + chat_id = f"chat-{uuid.uuid4()}" + turns: list[Turn] = [] + render = turns.append + # create_stream raises if the id exists, and create_or_get_stream gets it instead. + stream = await client.create_stream[Turn](chat_id, retention=timedelta(days=1)) + async with stream: # optional, the block end closes it + await client.start_workflow( + Chat.run, + stream.ref, # the ref is data, and stream.proto_ref goes in a proto message + id=chat_id, + task_queue=task_queue, + ) + async for r in stream.subscribe(): + if r.kind != streams.ItemKind.RECORD: + continue + render(r.data) + if is_last_turn(r): + 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 000000000..e63724432 --- /dev/null +++ b/tests/contrib/streams/examples/resume_from_offset.py @@ -0,0 +1,99 @@ +"""Streams API proposal, the resume notes of "A UI backend reads a Stream by its id". + +The chat server creates the Stream and passes its ref to the Workflow. The Reader gets the +Stream by its id, 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 ``streams.END``. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass + +from temporalio import streams, workflow +from temporalio.client import Client +from temporalio.worker import Worker + + +@dataclass +class Turn: + text: str + + +@workflow.defn +class Chat: + def __init__(self) -> None: + self.more: list[str] = [] + + @workflow.run + async def run(self, ref: streams.StreamRef[Turn], first: list[str]) -> None: + turns = workflow.to_stream(ref) + await turns.publish(*map(Turn, first)) + await workflow.wait_condition(lambda: bool(self.more)) + 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) -> None: + if offset is not None: + self.offsets[key] = offset + + +async def read_some( + client: Client, chat_id: str, store: OffsetStore, count: int +) -> list[str]: + stream = client.get_stream[Turn](chat_id) + seen = [] + read = stream.subscribe(offset=store.load_offset("turns")) + async for r in read: + if r.kind != streams.ItemKind.RECORD: + continue + seen.append(r.data.text) + store.save_offset("turns", read.next_offset) # the next read starts after it + if len(seen) == count: + break + await read.aclose() + return seen + + +async def read_new(client: Client, chat_id: str) -> list[str]: + # A view that shows only what comes next starts after the newest Record. + stream = client.get_stream[Turn](chat_id) + return [ + r.data.text + async for r in stream.subscribe(offset=streams.END, wait_new_records=False) + if isinstance(r, streams.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]): + chat_id = f"chat-{uuid.uuid4()}" + async with await client.create_stream[Turn](chat_id) as turns: + handle = await client.start_workflow( + Chat.run, + args=[turns.ref, ["a", "b"]], + id=chat_id, + task_queue=task_queue, + ) + store = OffsetStore() + before = await read_some(client, chat_id, store, 2) + new = await read_new(client, chat_id) + await handle.signal(Chat.add, ["c", "d"]) + after = await read_some(client, chat_id, 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 000000000..4c6e9b2f4 --- /dev/null +++ b/tests/contrib/streams/examples/test_examples.py @@ -0,0 +1,83 @@ +"""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 streams +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")] + + +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): + records = await agent_publishes_turns.main(stream_client, task_queue()) + Turn = agent_publishes_turns.Turn + # The read ends at the close, after the done Turn that carries the reply. + assert [record.data for record in records] == [ + Turn("model"), + Turn("model"), + Turn("done", "the answer"), + ] + + +async def test_activity_streams_tokens(stream_client: Client): + summary, 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 summary == 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(streams.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 000000000..978976ac1 --- /dev/null +++ b/tests/contrib/streams/examples/ui_backend_reads.py @@ -0,0 +1,96 @@ +"""Streams API proposal, "A UI backend reads a Stream by its id". + +A chat server creates the Stream and passes its ref to the agent Workflow, which publishes each +turn. The UI backend gets the Stream by its id, with no call, and renders each item as it +arrives. It may read before the creator has created the Stream, so the read waits for the +create. The read ends when the chat server closes the Stream. +""" + +from __future__ import annotations + +import asyncio +import uuid +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import activity, streams, workflow +from temporalio.client import Client +from temporalio.worker import Worker + + +@dataclass +class Turn: + kind: str + text: str | None = None + + +@activity.defn +async def call_model(prompt: str) -> str: + return prompt.upper() + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, ref: streams.StreamRef[Turn], prompt: str) -> str: + turns = workflow.to_stream(ref) + 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)) + return reply + + +async def chat_server(client: Client, chat_id: str, task_queue: str) -> None: + # The block end closes the Stream after the agent's done Turn. + async with await client.create_stream[Turn](chat_id) as turns: + await client.execute_workflow( + Agent.run, + args=[turns.ref, "hello there"], + id=chat_id, + task_queue=task_queue, + ) + + +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] = {} + + def render(self, turn: Turn, offset: int) -> None: + self.shown[offset] = turn + + def unrender(self, offsets: tuple[int, ...]) -> None: + for offset in offsets: + del self.shown[offset] + + +# Where the view goes on. A first view has none, so it reads from the start. +saved: list[int] = [] + + +def saved_offset() -> int | None: + return saved[-1] if saved else None + + +async def main(client: Client, task_queue: str) -> Screen: + async with Worker( + client, task_queue=task_queue, workflows=[Agent], activities=[call_model] + ): + chat_id = f"chat-{uuid.uuid4()}" + server = asyncio.create_task(chat_server(client, chat_id, task_queue)) + screen = Screen() + stream = client.get_stream[Turn](chat_id) # by id, with no call + async for r in stream.subscribe( + offset=saved_offset(), # wherever the app keeps it, or streams.END + create_wait=timedelta(seconds=10), # the creator may not have created it + ): + match r.kind: # the kind narrows the item's type + case streams.ItemKind.SUPERSEDED: + screen.unrender(r.withdrawn_offsets) + case _: + screen.render(r.data, r.offset) # a Turn + await server + 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 000000000..400b7cdb4 --- /dev/null +++ b/tests/contrib/streams/examples/worker_setup.py @@ -0,0 +1,84 @@ +"""Streams 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 streams, workflow +from temporalio.client import Client +from temporalio.contrib.stream import RedisStreamBackend +from temporalio.contrib.stream.testing import MemoryStreamBackend +from temporalio.worker import Worker + + +@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[str]( + f"{workflow.info().workflow_id}-tokens" + ) + self.tokens.publish_nowait(f"hello {name}") + + @workflow.query + def stream(self) -> streams.StreamRef[str]: + return self.tokens.ref + + @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: + # This SDK 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 = client.to_stream(await handle.query(Greeter.stream)) + return [r.data async for r in stream.subscribe() if isinstance(r, streams.Record)] + + +async def create_on_second_backend(client: Client) -> str | None: + stream = await client.create_stream[str](f"chat-{uuid.uuid4()}", backend="redis-b") + return stream.ref.backend diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index 218dac27a..582e43730 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -1,13 +1,44 @@ +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, streams, workflow -from temporalio.client import Client, Stream +from temporalio.api.common.v1 import Payload +from temporalio.api.enums.v1 import EventType, StreamErrorReason +from temporalio.client import ( + BufferedPublisher, + Client, + Interceptor, + Stream, + WorkflowFailureError, + WorkflowHandle, +) +from temporalio.common import RawValue, RetryPolicy +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.streams._extension import ExtensionStreamBackend, _WorkflowInbound from temporalio.streams.system import models 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 stream_status from tests.helpers import new_worker @@ -179,3 +210,630 @@ async def test_later_write_withdraws_a_skipped_local_activity(stream_client: Cli superseded = [item for item in items if isinstance(item, streams.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"] + + +@dataclass +class Turn: + kind: str + text: str | None = None + + +def own_id(name: str = "tokens") -> str: + # A Workflow picks the id of each Stream that it creates. These tests use its Workflow ID. + return f"{workflow.info().workflow_id}-{name}" + + +def stream_of( + client: Client, handle: WorkflowHandle[Any, Any], name: str = "tokens" +) -> Stream[Any]: + return client.get_stream(f"{handle.id}-{name}") + + +async def closed(client: Client, handle: WorkflowHandle[Any, Any], name: str) -> bool: + status = await stream_status(client, f"{handle.id}-{name}") + return status == models.StreamStatus.StreamStatusClosed + + +@activity.defn +async def call_model(input: CallModelInput) -> str: + stream = activity.client().to_stream(input.stream) + for word in input.prompt.split(): + await 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[Turn](own_id("turns")) + tokens = await workflow.create_stream[str](own_id("tokens")) + await turns.publish(Turn("model")) + reply = await workflow.execute_activity( + call_model, + CallModelInput(prompt, tokens.ref), + start_to_close_timeout=timedelta(seconds=30), + ) + await turns.publish(Turn("done", reply)) + await turns.close() + await tokens.close() + return reply + + +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 = stream_of(stream_client, handle) + 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, streams.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, ref: streams.StreamRef[str] | None) -> str: + if ref is None: + stream = await workflow.create_stream[str](own_id()) + await stream.publish("first run") + workflow.continue_as_new(stream.ref) + await workflow.to_stream(ref).close() + 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, + None, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == "second run" + items = await read_all(stream_of(stream_client, handle)) + assert await closed(stream_client, handle, "tokens") + assert [i.data for i in records(items)] == ["first run"] + + +@workflow.defn +class CancelledBlock: + @workflow.run + async def run(self) -> None: + async with await workflow.create_stream[str](own_id("scoped")) as turn: + await turn.publish("started") + await workflow.wait_condition(lambda: False) + + +async def test_cancelled_block_closes_the_stream(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 = stream_of(stream_client, handle, "scoped") + reader = turn.subscribe(create_wait=timedelta(seconds=20)) + assert isinstance(await anext(reader), streams.Record) + await reader.aclose() + await handle.cancel() + items = await read_all(turn) + assert await closed(stream_client, handle, "scoped") + assert [i.data for i in records(items)] == ["started"] + + +@workflow.defn +class PlainWriter: + @workflow.run + async def run(self) -> str: + stream = await workflow.create_stream[str](own_id("plain")) + await stream.publish("plain") + async with await workflow.create_stream[str](own_id("scoped")) as scoped: + await scoped.publish("scoped") + await stream.close() + 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(stream_of(stream_client, handle, "plain")) + scoped = await read_all(stream_of(stream_client, handle, "scoped")) + assert await closed(stream_client, handle, "plain") + assert await closed(stream_client, handle, "scoped") + assert [i.data for i in records(items)] == ["plain"] + assert [i.data for i in records(scoped)] == ["scoped"] + + +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 + + +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 = codec_client.get_stream[Turn](f"{handle.id}-turns") + tokens = codec_client.get_stream[str](f"{handle.id}-tokens") + turn_items = await read_all(turns) + decoded = await read_all(tokens) + # RawValue as the payload type reads the raw Payloads. + raw_tokens = codec_client.get_stream[RawValue](tokens.ref.id) + raw = [item async for item in raw_tokens.subscribe(wait_new_records=False)] + assert [item.data for item in records(turn_items)] == [ + Turn("model"), + Turn("done", "A B"), + ] + assert [item.data for item in records(decoded)] == ["a", "b"] + first = raw[0] + assert isinstance(first, streams.Record) + assert first.data.payload.metadata["encoding"] == b"json/plain" + + +@workflow.defn +class WaitsForSignal: + def __init__(self) -> None: + self.go = False + + @workflow.run + async def run(self) -> str: + stream = await workflow.create_stream[str](own_id()) + await stream.publish("before restart") + async with await workflow.create_stream[str](own_id("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 = stream_of(stream_client, handle) + reader = first.subscribe(create_wait=timedelta(seconds=20)) + assert isinstance(await anext(reader), streams.Record) + await reader.aclose() + # The new Worker has no cache, so it replays the run. A replay sends no create, which would + # fail on the Streams that exist, and no second close. + 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(stream_of(stream_client, handle, "scoped")) + assert await closed(stream_client, handle, "tokens") + assert await closed(stream_client, handle, "scoped") + assert [i.data for i in records(items)] == ["before restart", "after restart"] + assert [i.data for i in records(scoped)] == ["scoped"] + + +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[str](f"shutdown-{uuid.uuid4()}") + async with 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, streams.Record) + ] + assert written == ["buffered"] + + +@workflow.defn +class WaitsForever: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream[str](own_id()) + 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 = stream_of(stream_client, handle).subscribe( + create_wait=timedelta(seconds=20) + ) + assert isinstance(await anext(reader), streams.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 + # Each retry run sends the create of the first run, so it gets the same Stream. + stream = await workflow.create_stream[str](own_id()) + 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 = stream_of(stream_client, handle) + items = [item async for item in stream.subscribe(wait_new_records=False)] + # The end of the chain doesn't close the Stream. + assert not await closed(stream_client, handle, "tokens") + written = records(items) + superseded = [item for item in items if isinstance(item, streams.Superseded)] + 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[str](f"timer-{uuid.uuid4()}") + first = await stream.publish("one") + seen: list[streams.Record[Any]] = [] + + async def read_until_quiet() -> None: + async for item in stream.subscribe(): + assert isinstance(item, streams.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, streams.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 + # ref 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[str](stream_id) + on_memory = await stream_client.create_stream[str](stream_id, backend="memory") + assert (on_redis.ref.backend, on_memory.ref.backend) == ("redis", "memory") + await on_redis.publish("redis") + await on_memory.publish("memory") + for backend, expected in (("redis", "redis"), ("memory", "memory")): + again = stream_client.get_stream[str](stream_id, backend=backend) + data = [ + item.data + async for item in again.subscribe(wait_new_records=False) + if isinstance(item, streams.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[str](own_id()) + await stream.publish("token") + await stream.close() + 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 = stream_of(client, handle) + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert await closed(client, handle, "tokens") + assert [r.data for r in records(items)] == ["token"] + 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, ref: streams.StreamRef[str]) -> str: + try: + await workflow.to_stream(ref).publish("from workflow") + except streams.StreamError as err: + return StreamErrorReason.Name(StreamErrorReason.ValueType(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[str](f"closed-{uuid.uuid4()}") + await stream.close() + async with new_worker(stream_client, WritesToStream) as worker: + assert ( + await stream_client.execute_workflow( + WritesToStream.run, + stream.ref, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + == "STREAM_ERROR_REASON_CLOSED" + ) + # The close wrote no item, so a read of the closed Stream ends at once with none. + assert [item async for item in stream.subscribe()] == [] + + +@workflow.defn +class ContinuesThenRetries: + @workflow.run + async def run(self, ref: streams.StreamRef[str] | None) -> str: + if ref is None: + created = await workflow.create_stream[str](own_id()) + await created.publish("first") + workflow.continue_as_new(created.ref) + stream = workflow.to_stream(ref) + info = workflow.info() + await stream.publish(f"attempt-{info.attempt}") + if info.attempt == 1: + raise ApplicationError("the continued run fails once") + await stream.close() + 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(stream_of(stream_client, handle)) + assert await closed(stream_client, 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, streams.Superseded)] + assert [r.data for r in written if superseded[0].withdraws(r)] == ["attempt-1"]