From 3047c8c0cb79c7fa5fa9924f813156a0e847b2ad Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 19:28:08 -0700 Subject: [PATCH 1/4] Added StreamClient with create, get_handle and get_workflow_stream. --- temporalio/contrib/streams/__init__.py | 8 +- temporalio/contrib/streams/_backend.py | 32 ++ temporalio/contrib/streams/_buffered.py | 5 +- temporalio/contrib/streams/_stream_client.py | 319 ++++++++++++++++++ .../contrib/streams/test_activity_handles.py | 3 +- tests/contrib/streams/test_redis.py | 2 +- tests/contrib/streams/test_stream_client.py | 174 ++++++++++ 7 files changed, 537 insertions(+), 6 deletions(-) create mode 100644 temporalio/contrib/streams/_stream_client.py create mode 100644 tests/contrib/streams/test_stream_client.py diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 9f8ff7d6c..bddd7cc88 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -20,7 +20,6 @@ from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams._buffered import BufferedPublisher from temporalio.contrib.streams._client import ( - ClientStream, create_stream, for_handle, for_ref, @@ -48,6 +47,11 @@ from temporalio.contrib.streams._items import Closed, Idle, Record, Superseded from temporalio.contrib.streams._outside import OutsideStream from temporalio.contrib.streams._stream import StreamRef, Topic +from temporalio.contrib.streams._stream_client import ( + ClientStream, + ClientTopic, + StreamClient, +) from temporalio.contrib.streams._workflow import ( WorkflowStream, create_workflow_stream, @@ -59,6 +63,7 @@ "ActivityTopic", "BufferedPublisher", "ClientStream", + "ClientTopic", "Closed", "Idle", "OutsideStream", @@ -66,6 +71,7 @@ "Record", "StreamBackend", "StreamBufferFullError", + "StreamClient", "StreamDefinition", "StreamClosedError", "StreamError", diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/contrib/streams/_backend.py index 0ebab40bd..a26e614be 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -174,3 +174,35 @@ def backend_of_client(client: temporalio.client.Client) -> StreamBackend: raise _errors.StreamTransportMissingError( "The client has no Stream plugin, such as RedisStreamPlugin" ) + + +def client_backend( + client: temporalio.client.Client, name: str | None +) -> tuple[str, StreamBackend]: + """The name and the backend of a client's Stream plugin. None picks the first plugin. + + Raises: + StreamTransportMissingError: The client has no Stream plugin. + StreamUnsupportedError: The client has no Stream plugin of that name. + """ + backends = [ + (getattr(plugin, "backend_name", ""), backend) + for plugin in client.config()["plugins"] + if isinstance( + backend := getattr(plugin, "_temporal_stream_backend", None), StreamBackend + ) + ] + if not backends: + raise _errors.StreamTransportMissingError( + "The client has no Stream plugin, such as RedisStreamPlugin" + ) + if not name: + return backends[0] + for backend_name, backend in backends: + if backend_name == name: + return backend_name, backend + raise _errors.StreamUnsupportedError( + f"The client has no Stream backend {name}. It has " + f"{', '.join(backend_name for backend_name, _ in backends)}.", + reason="stream_unsupported", + ) diff --git a/temporalio/contrib/streams/_buffered.py b/temporalio/contrib/streams/_buffered.py index b0a25ad47..9d44a03af 100644 --- a/temporalio/contrib/streams/_buffered.py +++ b/temporalio/contrib/streams/_buffered.py @@ -7,7 +7,8 @@ from types import TracebackType from typing import Any, Generic, TypeVar -from temporalio.contrib.streams._access import _AsyncTopic +from temporalio.contrib.streams._activity_stream import ActivityTopic +from temporalio.contrib.streams._stream_client import ClientTopic T = TypeVar("T") @@ -25,7 +26,7 @@ class BufferedPublisher(Generic[T]): def __init__( self, - topic: _AsyncTopic[T], + topic: ActivityTopic[T] | ClientTopic[T], *, max_delay: timedelta = timedelta(milliseconds=100), ) -> None: diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py new file mode 100644 index 000000000..a4f815198 --- /dev/null +++ b/temporalio/contrib/streams/_stream_client.py @@ -0,0 +1,319 @@ +"""Streams in client code.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence +from datetime import timedelta +from typing import Any, TypeVar, overload + +import temporalio.api.streamservice.v1 as stream_pb +import temporalio.client +from temporalio.contrib.streams import _errors, _ids +from temporalio.contrib.streams._access import Access, _AsyncTopic, read_spec +from temporalio.contrib.streams._backend import StreamBackend, client_backend +from temporalio.contrib.streams._client import _client_producer +from temporalio.contrib.streams._defn import TopicDefinition, definition_of +from temporalio.contrib.streams._handle import StreamHandle +from temporalio.contrib.streams._reader import Item +from temporalio.contrib.streams._writer import BatchOptions, StreamWriter + +E = TypeVar("E") +T = TypeVar("T") + + +class ClientTopic(_AsyncTopic[T]): + """One Topic of a :py:class:`ClientStream`. All Topics of a Stream share one order. + + .. warning:: + This API is experimental. + """ + + +class ClientStream(StreamHandle[E]): + """A Stream that a client writes, reads and closes. Get one from :py:class:`StreamClient`. + + Each publish writes before it returns. Use :py:class:`BufferedPublisher` to batch. Each + handle writes as its own Producer. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + *, + id: str, + workflow_id: str | None = None, + first_execution_run_id: str | None = None, + backend: str | None = None, + definition: type[E] | None = None, + ) -> None: + """Create a handle. Use :py:class:`StreamClient` instead.""" + super().__init__( + id=id, + workflow_id=workflow_id, + first_execution_run_id=first_execution_run_id, + backend=backend, + definition=definition, + ) + self._client: temporalio.client.Client | None = None + self._stream_backend: StreamBackend | None = None + self._find_first_run: Callable[[], Awaitable[str]] | None = None + self._writer: StreamWriter | None = None + # The Stream name of a Workflow's Stream whose first run id is not known yet. + self._stream_name = "" + + def _bind( + self, + client: temporalio.client.Client, + backend: StreamBackend, + find_first_run: Callable[[], Awaitable[str]] | None = None, + ) -> ClientStream[E]: + self._client = client + self._stream_backend = backend + self._find_first_run = find_first_run + return self + + @property + def id(self) -> str: + """The Stream id. + + Raises: + RuntimeError: The Workflow handle had no first run id, and no call has found it + yet. Call :py:meth:`get_offset` first. + """ + if not self._id: + raise RuntimeError( + "The Stream id is not known yet. Call get_offset() first." + ) + return self._id + + async def _stream_id(self) -> str: + if not self._id: + assert self._find_first_run is not None and self._workflow_id is not None + self._first_execution_run_id = await self._find_first_run() + self._find_first_run = None + self._id = _ids.scoped_stream_id( + self._workflow_id, self._first_execution_run_id, self._stream_name + ) + return self._id + + def _access(self) -> Access: + client, backend = self._client, self._stream_backend + if client is None or backend is None: + raise _errors.StreamTransportMissingError( + "This ClientStream came from a payload, so it has no client. " + "Use StreamClient.get_handle." + ) + + def writer() -> StreamWriter: + if self._writer is None: + self._writer = StreamWriter( + backend=backend, + converter=client.data_converter, + namespace=client.namespace, + stream_id=self._stream_id, + producer=_client_producer(client), + batch=BatchOptions(), + create_first=False, + loop=asyncio.get_running_loop(), + ) + return self._writer + + return Access( + backend=backend, + converter=client.data_converter, + namespace=client.namespace, + stream_id=self._stream_id, + writer=writer, + ) + + @overload + def topic(self, topic: TopicDefinition[T]) -> ClientTopic[T]: ... + + @overload + def topic(self, topic: str) -> ClientTopic[Any]: ... + + @overload + def topic(self, topic: str, *, type: type[T]) -> ClientTopic[T]: ... + + def topic( + self, topic: str | TopicDefinition[Any], *, type: Any = None + ) -> ClientTopic[Any]: + """A Topic of this Stream. + + Args: + topic: A Topic of the Stream's definition, such as ``Events.turns``, or a name. + A name of the definition takes its type. + type: The type of a Topic outside the definition, or of a Stream with no + definition. + """ + name, topic_type = self._topic_spec(topic, type) + return ClientTopic(self, name, topic_type) + + def subscribe( + self, + *topics: str | TopicDefinition[Any] | ClientTopic[Any], + from_offset: int | None = None, + wait_new_records: bool = True, + prefixes: Sequence[str] = (), + create_wait: timedelta | None = None, + ) -> AsyncGenerator[Item, None]: + """Read the Stream's items in order. + + The read ends after a ``Closed`` item, or at the newest Record when + ``wait_new_records`` is False. A reader that wants to stop after a quiet time sets its + own timer, and resumes later with ``from_offset``. A ``Superseded`` item withdraws + earlier Records of a failed attempt, and it passes every Topic filter. + + Args: + topics: Read these Topics. With no Topics and no prefixes, read every Topic. + from_offset: Start at this offset, inclusive. Defaults to the first Record. + wait_new_records: Wait for new Records at the end of the Stream. + prefixes: Also read the Topics whose names start with one of these. + create_wait: How long to wait for a Stream that nobody created yet, such as the + Stream of a Workflow that just started. Defaults to no wait. + """ + names, types = read_spec(self, topics) + return self._access().read( + topics=names, + prefixes=prefixes, + types=types, + from_offset=from_offset, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + + async def get_offset(self) -> int: + """The next offset of the Stream, the newest offset plus one.""" + return await self._access().get_offset() + + async def close( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Write this handle's Records, then close the Stream. + + Readers get the Records before it, then a ``Closed`` item with ``result`` or + ``failure``. + """ + await self._access().close(result, failure) + + +class StreamClient: + """Creates and gets Streams over a Temporal client with a Stream plugin. + + .. warning:: + This API is experimental. + """ + + def __init__(self, client: temporalio.client.Client) -> None: + """Create a Stream client. It makes no call.""" + self._client = client + + async def create( + self, + definition: type[E], + *, + id: str, + retention: timedelta | None = None, + backend: str | None = None, + ) -> ClientStream[E]: + """Create a Stream that this client owns, so only a client closes it. + + It is idempotent for one id and one Retention. + + Args: + definition: The Stream definition class. + id: The Stream id, as a Workflow takes a Workflow id. + retention: How long the Records stay after the close. Defaults to the backend's. + backend: The name of the backend. Defaults to the client's first Stream plugin. + + Raises: + StreamUnsupportedError: The client has no backend of that name. No call is made. + StreamClosedError: A Stream of that id was created and closed before. + StreamError: The Stream exists with another Retention, with the reason + ``stream_already_exists``. + """ + definition_of(definition) + backend_name, stream_backend = client_backend(self._client, backend) + request = stream_pb.CreateStreamRequest( + namespace=self._client.namespace, stream_id=id + ) + if retention is not None: + request.retention.FromTimedelta(retention) + response = await stream_backend.call_async( + "CreateStream", request, stream_pb.CreateStreamResponse + ) + if response.status == stream_pb.StreamStatus.STREAM_STATUS_CLOSED: + raise _errors.StreamClosedError( + f"The Stream {id} was created and closed before", + reason="stream_closed", + ) + stream = ClientStream( + id=id, backend=backend_name or None, definition=definition + ) + return stream._bind(self._client, stream_backend) + + def get_handle( + self, definition: type[E], *, id: str, backend: str | None = None + ) -> ClientStream[E]: + """A Stream by its id, such as one that this or another client created. No call. + + Args: + definition: The Stream definition class. + id: The Stream id. + backend: The name of the backend. Defaults to the client's first Stream plugin. + """ + definition_of(definition) + backend_name, stream_backend = client_backend(self._client, backend) + stream = ClientStream( + id=id, backend=backend_name or None, definition=definition + ) + return stream._bind(self._client, stream_backend) + + def get_workflow_stream( + self, + handle: temporalio.client.WorkflowHandle[Any, Any], + definition: type[E], + *, + name: str | None = None, + backend: str | None = None, + ) -> ClientStream[E]: + """A Stream that a Workflow created, by the Workflow handle and the definition. + + It builds the id from the Workflow id and the first run id, so it makes no call. A + handle with no first run id makes one Describe call on the Stream's first use. + + Args: + handle: A handle of the Workflow, best the one that ``start_workflow`` returned. + definition: The Stream definition class. + name: The Stream name that the Workflow gave. Defaults to the definition's name. + backend: The name of the backend. Defaults to the client's first Stream plugin. + """ + stream_name = name or definition_of(definition).name + backend_name, stream_backend = client_backend(self._client, backend) + first_run = handle.first_execution_run_id + stream = ClientStream( + id=_ids.scoped_stream_id(handle.id, first_run, stream_name) + if first_run + else "", + workflow_id=handle.id, + first_execution_run_id=first_run, + backend=backend_name or None, + definition=definition, + ) + if first_run: + return stream._bind(self._client, stream_backend) + + async def find_first_run() -> str: + description = await handle.describe() + first = description.raw_description.workflow_execution_info.first_run_id + if not first: + raise ValueError( + "The server gave no first run id, so the Stream id is unknown" + ) + return first + + stream._stream_name = stream_name + return stream._bind(self._client, stream_backend, find_first_run) diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py index 73b06cd5e..75107b852 100644 --- a/tests/contrib/streams/test_activity_handles.py +++ b/tests/contrib/streams/test_activity_handles.py @@ -11,7 +11,6 @@ from temporalio.common import RetryPolicy from temporalio.contrib import streams from temporalio.contrib.streams import workflow as stream_workflow -from temporalio.contrib.streams._access import _AsyncTopic from tests.helpers import new_worker @@ -227,7 +226,7 @@ async def test_activity_closes_and_checks_the_backend(stream_client: Client): @dataclass -class _Batches(_AsyncTopic[int]): +class _Batches(streams.ActivityTopic[int]): batches: list[tuple[int, ...]] = field(default_factory=list) fail: bool = False diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index 37cf4ad18..18c81ca1c 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -64,7 +64,7 @@ async def run(self, prompt: str) -> str: return reply -async def read_all(stream: streams.ClientStream) -> list[streams.Record | object]: +async def read_all(stream: streams.OutsideStream) -> list[streams.Record | object]: return [item async for item in stream.subscribe(idle_timeout=timedelta(seconds=20))] diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py new file mode 100644 index 000000000..40f4cb0a9 --- /dev/null +++ b/tests/contrib/streams/test_stream_client.py @@ -0,0 +1,174 @@ +import uuid +from dataclasses import dataclass +from datetime import timedelta + +import pytest + +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib import streams +from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.converter import DataConverter +from temporalio.plugin import SimplePlugin +from tests.contrib.streams.fake_extension import FakeExtension +from tests.helpers import new_worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@streams.defn(name="events") +class Events: + turns = streams.topic(Turn) + tokens = streams.topic(str) + + +def fake_client(client: Client) -> tuple[Client, FakeExtension]: + extension = FakeExtension() + plugin = SimplePlugin("fake-streams") + setattr(plugin, "_temporal_stream_backend", streams.StreamBackend(extension)) + setattr(plugin, "backend_name", "fake") + config = client.config() + config["plugins"] = [plugin] + return Client(**config), extension + + +async def test_unknown_backend_raises_before_any_call(client: Client): + client, extension = fake_client(client) + stream_client = streams.StreamClient(client) + with pytest.raises(streams.StreamUnsupportedError, match="It has fake"): + await stream_client.create(Events, id="chat-1", backend="redis") + with pytest.raises(streams.StreamUnsupportedError): + stream_client.get_handle(Events, id="chat-1", backend="redis") + assert not extension.streams and not extension.namespaces + stream = await stream_client.create(Events, id="chat-1") + assert stream.backend == "fake" + + +async def test_client_without_plugin_has_no_transport(client: Client): + with pytest.raises(streams.StreamTransportMissingError): + streams.StreamClient(client).get_handle(Events, id="chat-1") + + +async def test_decoded_client_stream_has_no_client(client: Client): + client, _ = fake_client(client) + stream = await streams.StreamClient(client).create(Events, id="chat-2") + converter = DataConverter.default.payload_converter + decoded = converter.from_payload( + converter.to_payload(stream), streams.ClientStream[Events] + ) + assert decoded == stream + with pytest.raises(streams.StreamTransportMissingError, match="get_handle"): + await decoded.get_offset() + + +async def test_handle_without_first_run_describes_once(client: Client): + client, _ = fake_client(client) + handle = client.get_workflow_handle("agent-1") + calls = [] + + async def describe() -> object: + calls.append(1) + + class Info: + first_run_id = "first" + + class Raw: + workflow_execution_info = Info() + + class Description: + raw_description = Raw() + + return Description() + + setattr(handle, "describe", describe) + stream_client = streams.StreamClient(client) + await stream_client.create(Events, id="agent-1:first:n") + stream = stream_client.get_workflow_stream(handle, Events, name="n") + with pytest.raises(RuntimeError, match="not known yet"): + stream.id + assert await stream.get_offset() == 0 + assert stream.id == "agent-1:first:n" + assert stream.first_execution_run_id == "first" + assert await stream.get_offset() == 0 + assert calls == [1] + + +@workflow.defn +class WriteToStream: + @workflow.run + async def run(self, stream: stream_workflow.WorkflowStream[Events]) -> None: + await stream.topic(Events.turns).publish(Turn("workflow")) + + +async def test_client_stream_takes_workflow_writes(stream_client: Client): + client = streams.StreamClient(stream_client) + stream = await client.create(Events, id=f"chat-{uuid.uuid4()}") + assert await stream.get_offset() == 0 + first = await stream.topic(Events.turns).publish(Turn("client")) + last = await stream.topic(Events.tokens).publish("a", "b") + assert last > first + async with new_worker(stream_client, WriteToStream) as worker: + await stream_client.execute_workflow( + WriteToStream.run, + stream_workflow.WorkflowStream.from_handle(stream), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await stream.close(result="bye") + items = [item async for item in stream.subscribe()] + data = [item.data for item in items if isinstance(item, streams.Record)] + assert data == [Turn("client"), "a", "b", Turn("workflow")] + assert isinstance(items[-1], streams.Closed) and items[-1].result == "bye" + + # Another handle of the same Stream reads it, from an offset and by Topic. + again = client.get_handle(Events, id=stream.id) + resumed = [ + item.data + async for item in again.subscribe(Events.tokens, from_offset=first + 1) + if isinstance(item, streams.Record) + ] + assert resumed == ["a", "b"] + by_prefix = [ + item.data + async for item in again.subscribe(prefixes=["tok"]) + if isinstance(item, streams.Record) + ] + assert by_prefix == ["a", "b"] + only_turns = [item async for item in again.topic(Events.turns).subscribe()] + assert [type(item) for item in only_turns] == [ + streams.Record, + streams.Record, + streams.Closed, + ] + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + await workflow.sleep(timedelta(milliseconds=200)) + stream = await stream_workflow.create_stream(Events) + await stream.topic(Events.tokens).publish(*prompt.split()) + return prompt.upper() + + +async def test_client_reads_a_workflow_stream(stream_client: Client): + async with new_worker(stream_client, Agent) as worker: + handle = await stream_client.start_workflow( + Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + stream = streams.StreamClient(stream_client).get_workflow_stream(handle, Events) + # The Workflow creates its Stream later, so the read waits for the create. + items = [ + item async for item in stream.subscribe(create_wait=timedelta(seconds=20)) + ] + assert [item.data for item in items if isinstance(item, streams.Record)] == [ + "hi", + "there", + ] + assert isinstance(items[-1], streams.Closed) and items[-1].result == "HI THERE" + assert stream.workflow_id == handle.id From 70b40e5c338a8d5501d6a1f7390453e075dacbd2 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 19:50:31 -0700 Subject: [PATCH 2/4] Took auto_close on StreamClient.create. --- temporalio/contrib/streams/_stream_client.py | 11 ++++++++--- tests/contrib/streams/test_stream_client.py | 14 ++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index a4f815198..297feefbe 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -217,23 +217,26 @@ async def create( *, id: str, retention: timedelta | None = None, + auto_close: timedelta | None = None, backend: str | None = None, ) -> ClientStream[E]: """Create a Stream that this client owns, so only a client closes it. - It is idempotent for one id and one Retention. + It is idempotent for one id, one Retention and one ``auto_close``. Args: definition: The Stream definition class. id: The Stream id, as a Workflow takes a Workflow id. retention: How long the Records stay after the close. Defaults to the backend's. + auto_close: Close the Stream with no result when no Record was appended for this + long, so a Stream that nobody closes does not stay open. Defaults to never. backend: The name of the backend. Defaults to the client's first Stream plugin. Raises: StreamUnsupportedError: The client has no backend of that name. No call is made. StreamClosedError: A Stream of that id was created and closed before. - StreamError: The Stream exists with another Retention, with the reason - ``stream_already_exists``. + StreamError: The Stream exists with another Retention or ``auto_close``, with + the reason ``stream_already_exists``. """ definition_of(definition) backend_name, stream_backend = client_backend(self._client, backend) @@ -242,6 +245,8 @@ async def create( ) if retention is not None: request.retention.FromTimedelta(retention) + if auto_close is not None: + request.auto_close.FromTimedelta(auto_close) response = await stream_backend.call_async( "CreateStream", request, stream_pb.CreateStreamResponse ) diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 40f4cb0a9..cee9ac147 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -1,3 +1,4 @@ +import asyncio import uuid from dataclasses import dataclass from datetime import timedelta @@ -172,3 +173,16 @@ async def test_client_reads_a_workflow_stream(stream_client: Client): ] assert isinstance(items[-1], streams.Closed) and items[-1].result == "HI THERE" assert stream.workflow_id == handle.id + + +async def test_auto_close_of_a_client_stream(stream_client: Client): + stream = await streams.StreamClient(stream_client).create( + Events, id=f"quiet-{uuid.uuid4()}", auto_close=timedelta(seconds=1) + ) + await stream.topic(Events.tokens).publish("only") + await asyncio.sleep(1.5) + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert [item.data for item in items if isinstance(item, streams.Record)] == ["only"] + assert isinstance(items[-1], streams.Closed) and items[-1].result is None + with pytest.raises(streams.StreamClosedError): + await stream.topic(Events.tokens).publish("late") From c2dd5132c8ad163a2025d572011f610561dfdb19 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 20:20:31 -0700 Subject: [PATCH 3/4] Shared one Describe call and named the client's default backend. --- temporalio/contrib/streams/_backend.py | 16 +++++--- temporalio/contrib/streams/_stream_client.py | 25 +++++++++--- tests/contrib/streams/test_stream_client.py | 42 ++++++++++++++++++++ 3 files changed, 72 insertions(+), 11 deletions(-) diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/contrib/streams/_backend.py index 53afc7d3f..577d45e7b 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -208,14 +208,19 @@ def backend_of_client(client: temporalio.client.Client) -> StreamBackend: def client_backend( client: temporalio.client.Client, name: str | None ) -> tuple[str, StreamBackend]: - """The name and the backend of a client's Stream plugin. None picks the first plugin. + """The name and the backend of a client's Stream plugin. None picks the plugin marked + as the default, or else the first one. Raises: StreamTransportMissingError: The client has no Stream plugin. StreamUnsupportedError: The client has no Stream plugin of that name. """ backends = [ - (getattr(plugin, "backend_name", ""), backend) + ( + getattr(plugin, "backend_name", ""), + backend, + getattr(plugin, "default", False), + ) for plugin in client.config()["plugins"] if isinstance( backend := getattr(plugin, "_temporal_stream_backend", None), StreamBackend @@ -226,12 +231,13 @@ def client_backend( "The client has no Stream plugin, such as RedisStreamPlugin" ) if not name: - return backends[0] - for backend_name, backend in backends: + defaults = [entry for entry in backends if entry[2]] or backends + return defaults[0][0], defaults[0][1] + for backend_name, backend, _ in backends: if backend_name == name: return backend_name, backend raise _errors.StreamUnsupportedError( f"The client has no Stream backend {name}. It has " - f"{', '.join(backend_name for backend_name, _ in backends)}.", + f"{', '.join(entry[0] for entry in backends)}.", reason="stream_unsupported", ) diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index 297feefbe..c276beba3 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -63,6 +63,7 @@ def __init__( self._writer: StreamWriter | None = None # The Stream name of a Workflow's Stream whose first run id is not known yet. self._stream_name = "" + self._finding: asyncio.Future[str] | None = None def _bind( self, @@ -91,14 +92,26 @@ def id(self) -> str: async def _stream_id(self) -> str: if not self._id: - assert self._find_first_run is not None and self._workflow_id is not None - self._first_execution_run_id = await self._find_first_run() - self._find_first_run = None - self._id = _ids.scoped_stream_id( - self._workflow_id, self._first_execution_run_id, self._stream_name - ) + # Calls that start together share one Describe call. + if self._finding is None: + assert self._find_first_run is not None + self._finding = asyncio.ensure_future(self._find_first_run()) + first_run = await asyncio.shield(self._finding) + if not self._id: + assert self._workflow_id is not None + self._first_execution_run_id = first_run + self._id = _ids.scoped_stream_id( + self._workflow_id, first_run, self._stream_name + ) return self._id + def _fields(self) -> dict[str, str]: + fields = super()._fields() + if not self._id: + # The Stream name tells apart the Streams of one Workflow before the id is known. + fields["name"] = self._stream_name + return fields + def _access(self) -> Access: client, backend = self._client, self._stream_backend if client is None or backend is None: diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index cee9ac147..38a0dad0d 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -2,6 +2,8 @@ import uuid from dataclasses import dataclass from datetime import timedelta +from types import SimpleNamespace +from typing import Any import pytest @@ -9,6 +11,7 @@ from temporalio.client import Client from temporalio.contrib import streams from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.contrib.streams.testing import MemoryStreamPlugin from temporalio.converter import DataConverter from temporalio.plugin import SimplePlugin from tests.contrib.streams.fake_extension import FakeExtension @@ -186,3 +189,42 @@ async def test_auto_close_of_a_client_stream(stream_client: Client): assert isinstance(items[-1], streams.Closed) and items[-1].result is None with pytest.raises(streams.StreamClosedError): await stream.topic(Events.tokens).publish("late") + + +async def test_unresolved_workflow_stream_compares_and_prints(client: Client): + client, _ = fake_client(client) + stream_client = streams.StreamClient(client) + handle = client.get_workflow_handle("agent-2") + first = stream_client.get_workflow_stream(handle, Events) + same = stream_client.get_workflow_stream(handle, Events) + other = stream_client.get_workflow_stream(handle, Events, name="other") + key = hash(first) + assert first == same and first != other + assert "agent-2" in repr(first) and "events" in repr(first) + describes = [] + + async def describe() -> Any: + describes.append(1) + await asyncio.sleep(0.05) + info = SimpleNamespace(first_run_id="run-2") + return SimpleNamespace( + raw_description=SimpleNamespace(workflow_execution_info=info) + ) + + setattr(handle, "describe", describe) + await stream_client.create(Events, id="agent-2:run-2:events") + # Calls that start together share one Describe call. + assert await asyncio.gather(first.get_offset(), first.get_offset()) == [0, 0] + assert describes == [1] + assert first.id == "agent-2:run-2:events" + assert hash(first) == key + + +async def test_default_plugin_is_the_default_backend(client: Client): + config = client.config() + config["plugins"] = [ + MemoryStreamPlugin(), + MemoryStreamPlugin(name="b", default=True), + ] + stream = await streams.StreamClient(Client(**config)).create(Events, id="d-1") + assert stream.backend == "b" From 018ad9d90620ce47cc938093ec54a7a0b82c2a11 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 21:18:31 -0700 Subject: [PATCH 4/4] Renamed get_workflow_stream to workflow_stream. --- temporalio/contrib/streams/_stream_client.py | 2 +- tests/contrib/streams/test_stream_client.py | 15 ++++++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index c276beba3..3fffa2a0d 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -290,7 +290,7 @@ def get_handle( ) return stream._bind(self._client, stream_backend) - def get_workflow_stream( + def workflow_stream( self, handle: temporalio.client.WorkflowHandle[Any, Any], definition: type[E], diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 38a0dad0d..6cffb7ec0 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -91,7 +91,7 @@ class Description: setattr(handle, "describe", describe) stream_client = streams.StreamClient(client) await stream_client.create(Events, id="agent-1:first:n") - stream = stream_client.get_workflow_stream(handle, Events, name="n") + stream = stream_client.workflow_stream(handle, Events, name="n") with pytest.raises(RuntimeError, match="not known yet"): stream.id assert await stream.get_offset() == 0 @@ -165,7 +165,7 @@ async def test_client_reads_a_workflow_stream(stream_client: Client): handle = await stream_client.start_workflow( Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - stream = streams.StreamClient(stream_client).get_workflow_stream(handle, Events) + stream = streams.StreamClient(stream_client).workflow_stream(handle, Events) # The Workflow creates its Stream later, so the read waits for the create. items = [ item async for item in stream.subscribe(create_wait=timedelta(seconds=20)) @@ -195,9 +195,9 @@ async def test_unresolved_workflow_stream_compares_and_prints(client: Client): client, _ = fake_client(client) stream_client = streams.StreamClient(client) handle = client.get_workflow_handle("agent-2") - first = stream_client.get_workflow_stream(handle, Events) - same = stream_client.get_workflow_stream(handle, Events) - other = stream_client.get_workflow_stream(handle, Events, name="other") + first = stream_client.workflow_stream(handle, Events) + same = stream_client.workflow_stream(handle, Events) + other = stream_client.workflow_stream(handle, Events, name="other") key = hash(first) assert first == same and first != other assert "agent-2" in repr(first) and "events" in repr(first) @@ -228,3 +228,8 @@ async def test_default_plugin_is_the_default_backend(client: Client): ] stream = await streams.StreamClient(Client(**config)).create(Events, id="d-1") assert stream.backend == "b" + + +def test_workflow_stream_has_no_get_prefix(): + assert hasattr(streams.StreamClient, "workflow_stream") + assert not hasattr(streams.StreamClient, "get_workflow_stream")