From c68482a6fb62c960952b1efe080297efb86bb3a2 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 14:58:15 -0700 Subject: [PATCH] Added the Redis and in-memory Stream backends. Both call the `temporalio-streams` runtime that the bridge links. CI installs `redis-server` on Linux and macOS, so the Redis lane runs there. --- .github/actions/install-redis/action.yml | 17 ++ .github/workflows/ci.yml | 5 + temporalio/client/_client.py | 5 +- temporalio/common/_streams/_extension.py | 145 ++++++++++ temporalio/contrib/stream/__init__.py | 90 ++++++ temporalio/contrib/stream/testing.py | 56 ++++ temporalio/worker/_replayer.py | 4 +- tests/contrib/streams/conftest.py | 77 ++++++ tests/contrib/streams/test_stream_client.py | 286 +++++++++++++++++++- 9 files changed, 680 insertions(+), 5 deletions(-) create mode 100644 .github/actions/install-redis/action.yml create mode 100644 temporalio/common/_streams/_extension.py create mode 100644 temporalio/contrib/stream/__init__.py create mode 100644 temporalio/contrib/stream/testing.py create mode 100644 tests/contrib/streams/conftest.py diff --git a/.github/actions/install-redis/action.yml b/.github/actions/install-redis/action.yml new file mode 100644 index 0000000000..edcc8e752d --- /dev/null +++ b/.github/actions/install-redis/action.yml @@ -0,0 +1,17 @@ +name: Install Redis +description: Install redis-server for the Stream tests on Linux and macOS, and make their Redis lane fail without it. +runs: + using: composite + steps: + - if: runner.os == 'Linux' + shell: bash + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends redis-server + - if: runner.os == 'macOS' + shell: bash + run: brew install redis + # Redis has no official Windows build, so the Redis lane skips there. + - if: runner.os != 'Windows' + shell: bash + run: echo "TEMPORAL_STREAMS_REQUIRE_REDIS=1" >> "$GITHUB_ENV" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f498f08f09..7e179d471c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -60,6 +60,7 @@ jobs: if: ${{ matrix.clippyLinter }} - run: poe build-develop - run: poe lint + - uses: ./.github/actions/install-redis - run: mkdir junit-xml - run: poe test ${{matrix.pytestExtraArgs}} -s --junit-xml=junit-xml/${{ matrix.python }}--${{ matrix.os }}.xml timeout-minutes: 15 @@ -158,6 +159,7 @@ jobs: - run: poe build-develop-fips - name: Confirm the FIPS build is loaded run: uv run python -c "from temporalio.bridge import temporal_sdk_bridge; assert temporal_sdk_bridge.FIPS, 'not a FIPS build'" + - uses: ./.github/actions/install-redis - run: mkdir junit-xml - name: Run tests against the FIPS build run: | @@ -237,6 +239,7 @@ jobs: - name: Check generation unchanged run: | [[ -z $(git status --porcelain temporalio tests) ]] || (git diff temporalio tests; echo "Protos changed"; exit 1) + - uses: ./.github/actions/install-redis - name: Test with protobuf 3.x run: poe test -s env: @@ -270,6 +273,7 @@ jobs: - run: uv sync --all-extras - run: poe build-develop - run: poe lint + - uses: ./.github/actions/install-redis - run: mkdir junit-xml - run: poe test -s --junit-xml=junit-xml/latest-deps.xml timeout-minutes: 15 @@ -331,6 +335,7 @@ jobs: env: TEMPORAL_CLIENT_CLOUD_API_KEY: ${{ secrets.TEMPORAL_CLIENT_CLOUD_API_KEY }} TEMPORAL_CLIENT_CLOUD_API_VERSION: v0.19.1 + - uses: ./.github/actions/install-redis - run: mkdir junit-xml - run: poe test -s --workflow-environment envconfig --junit-xml=junit-xml/cloud.xml timeout-minutes: 15 diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index 569b17aefc..e5ba9d9a28 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -238,8 +238,9 @@ async def connect( payload_limits: Warning thresholds for outbound payload/memo sizes. Over-threshold fields are logged but still sent. Set a threshold to 0 to disable it. header_codec_behavior: Encoding behavior for headers sent by the client. - stream_backends: The Stream backends of this client and its Workers. Names must be - unique. With more than one backend, exactly one must have ``default=True``. See + stream_backends: The Stream backends of this client and its Workers, such as + :py:class:`temporalio.contrib.stream.RedisStreamBackend`. Names must be unique. + With more than one backend, exactly one must have ``default=True``. See :py:class:`temporalio.common.StreamBackend`. .. warning:: diff --git a/temporalio/common/_streams/_extension.py b/temporalio/common/_streams/_extension.py new file mode 100644 index 0000000000..dbb68c82d1 --- /dev/null +++ b/temporalio/common/_streams/_extension.py @@ -0,0 +1,145 @@ +"""The Stream backends that call the temporalio-streams runtime in the bridge.""" + +from __future__ import annotations + +import logging +import weakref +from collections.abc import Sequence +from typing import Any + +import temporalio.bridge.streams +import temporalio.client +import temporalio.converter +import temporalio.runtime +import temporalio.service +import temporalio.worker +from temporalio.common._streams._backend import BackendCalls, _ExtensionCalls +from temporalio.common._streams._errors import StreamConfigError +from temporalio.common._streams._plugin import StreamBackend +from temporalio.common._streams.proto import v1 as ext_pb + +logger = logging.getLogger(__name__) + + +class ExtensionStreamBackend(StreamBackend): + """Base class of the Redis and memory backends, which call the temporalio-streams runtime. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + name: str, + memory: bool, + config: ext_pb.RedisStreamConfig | ext_pb.MemoryStreamConfig, + runtime: temporalio.runtime.Runtime | None, + backend_name: str, + default: bool, + ) -> None: + """Create the backend over the bridge's Stream runtime, on Redis or in memory. + + ``backend_name`` is the name that Stream handles store for this backend. ``default`` + makes it the backend of a create with no backend name. + """ + config.backend = backend_name + self._config = config.SerializeToString() + # On Redis, the clients of one Core runtime share one runtime object and its connections. + # The memory backend has one runtime object, so all its clients share its Streams. + self._backends: weakref.WeakKeyDictionary[ + temporalio.runtime.Runtime, _ExtensionCalls + ] = weakref.WeakKeyDictionary() + self._memory: _ExtensionCalls | None = None + if memory: + self._memory = self._new_backend( + runtime or temporalio.runtime.Runtime.default(), True + ) + super().__init__( + name, + backend_name=backend_name, + default=default, + ) + + def _new_backend( + self, runtime: temporalio.runtime.Runtime, memory: bool = False + ) -> _ExtensionCalls: + if self._memory is not None and not memory: + return self._memory + backend = self._backends.get(runtime) + if backend is None: + backend = _ExtensionCalls( + temporalio.bridge.streams.StreamRuntime( + runtime._core_runtime, memory, self._config + ) + ) + self._backends[runtime] = backend + return backend + + def _configure_client( + self, config: temporalio.client.ClientConfig + ) -> temporalio.client.ClientConfig: + config = super()._configure_client(config) + backend = self._new_backend(_runtime_of(config["service_client"])) + config["interceptors"] = [ + *config.get("interceptors", []), + _StreamInterceptor(self, backend, config["data_converter"]), + ] + return config + + def _backend_for(self, client: temporalio.client.Client) -> BackendCalls: + # The Stream runtime that this backend gave the client. + interceptor = self._interceptor_of( + client.config(active_config=True).get("interceptors", []) + ) + if interceptor is None or interceptor.backend is None: + raise StreamConfigError( + f"The client has no {type(self).__name__}. Register it with " + "stream_backends=." + ) + return interceptor.backend + + def _interceptor_of(self, interceptors: Sequence[Any]) -> _StreamInterceptor | None: + for interceptor in interceptors: + if ( + isinstance(interceptor, _StreamInterceptor) + and interceptor.plugin is self + ): + return interceptor + return None + + async def shutdown(self) -> None: + """Close every open writer, including the client's. + + Call it when the process stops writing to Streams. + """ + failed = 0 + for backend in list(self._backends.values()): + failed += await backend.close_writers() + if failed: + logger.warning("Stream shutdown left %d failed writers", failed) + + +def _runtime_of( + service: temporalio.service.ServiceClient, +) -> temporalio.runtime.Runtime: + # Find the runtime of a client's connection the same way the Worker does. + service = getattr(service, "worker_service_client", service) + runtime = getattr(service, "_runtime", None) or getattr( + getattr(service, "config", None), "runtime", None + ) + return runtime or temporalio.runtime.Runtime.default() + + +class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor): + # There is one for each client of the backend, with that client's Stream runtime. The + # client's Workers get it from the client. + def __init__( + self, + plugin: ExtensionStreamBackend, + backend: _ExtensionCalls | None, + converter: temporalio.converter.DataConverter | None, + ) -> None: + super().__init__() + self.plugin = plugin + self.backend = backend + self.converter = converter or temporalio.converter.DataConverter.default diff --git a/temporalio/contrib/stream/__init__.py b/temporalio/contrib/stream/__init__.py new file mode 100644 index 0000000000..4e8c7fcd20 --- /dev/null +++ b/temporalio/contrib/stream/__init__.py @@ -0,0 +1,90 @@ +"""The Redis Stream backend. + +Register it on the client. A client creates a Stream with +:py:meth:`temporalio.client.Client.create_stream`. ``backend=`` picks a backend that isn't the +default:: + + from temporalio.contrib.stream import RedisStreamBackend + + client = await Client.connect( + "localhost:7233", + stream_backends=[RedisStreamBackend("redis://localhost", default=True)], + ) + +.. warning:: + This package is experimental. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import timedelta + +from temporalio.common._streams._extension import ExtensionStreamBackend +from temporalio.common._streams.proto import v1 as ext_pb + + +class RedisStreamBackend(ExtensionStreamBackend): + """Sends Stream calls to Redis, through the temporalio-streams runtime in the bridge. + + Register it on the client with ``stream_backends=``. A Worker uses its client's backends, + so one backend serves Workflow code, Activities and the client. The clients of the backend + that share a :py:class:`temporalio.runtime.Runtime` share one runtime object and its + connections, and so do their Workers. It registers the backend name ``redis``, and each + Stream handle that it creates stores that name. A Worker serves only one backend. A client + can have more than one, and ``backend=`` picks one. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + url: str | Sequence[str], + *, + cluster: bool = False, + key_prefix: str | None = None, + retention: timedelta | None = None, + safety_ttl: timedelta | None = None, + response_timeout: timedelta | None = None, + name: str = "redis", + default: bool = False, + ) -> None: + """Create the backend. + + The backend connects to Redis at the first Stream call on a runtime, and the clients of + that runtime share the connection. + + Args: + url: A Redis URL, or the seed URLs of a cluster. Credentials and TLS come from the + URL. + cluster: Whether the URLs are of a Redis cluster. + key_prefix: The prefix of every key. Defaults to ``temporal-streams``. + retention: The default Retention of a new Stream. Defaults to 7 days. + safety_ttl: How long an open Workflow-owned Stream is kept after its last write. + Defaults to 90 days. + response_timeout: The longest time that one Redis call can take. Defaults to 5 s. + name: The backend name that Stream handles store, so every process can route a + handle that it receives. + default: Make this the backend of a create with no backend name. With more than + one registered backend, exactly one must have it, or the client raises + ``StreamConfigError``. + """ + config = ext_pb.RedisStreamConfig( + urls=[url] if isinstance(url, str) else list(url), + cluster=cluster, + key_prefix=key_prefix or "", + ) + for field, value in ( + ("retention", retention), + ("safety_ttl", safety_ttl), + ("response_timeout", response_timeout), + ): + if value is not None: + getattr(config, field).FromTimedelta(value) + super().__init__( + "temporalio.RedisStreamBackend", False, config, None, name, default + ) + + +__all__ = ["RedisStreamBackend"] diff --git a/temporalio/contrib/stream/testing.py b/temporalio/contrib/stream/testing.py new file mode 100644 index 0000000000..8d393ae59c --- /dev/null +++ b/temporalio/contrib/stream/testing.py @@ -0,0 +1,56 @@ +"""A Stream backend for tests, with no Redis. + +.. warning:: + This module is experimental. +""" + +from __future__ import annotations + +from datetime import timedelta + +import temporalio.runtime +from temporalio.common._streams._extension import ExtensionStreamBackend +from temporalio.common._streams.proto import v1 as ext_pb + + +class MemoryStreamBackend(ExtensionStreamBackend): + """Keeps Streams in this process, with the same rules as Redis, for tests. + + Register it on the client, as with ``RedisStreamBackend``. Every Worker and client that + shares one backend shares its Streams. The Streams are lost when the process ends. It + registers the backend name ``memory``. It has Redis's rules, so tests can swap the Store + without changing Workflow code. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + *, + retention: timedelta | None = None, + safety_ttl: timedelta | None = None, + runtime: temporalio.runtime.Runtime | None = None, + name: str = "memory", + default: bool = False, + ) -> None: + """Create the backend. + + Args: + retention: The default Retention of a new Stream. Defaults to 7 days. + safety_ttl: How long an open Workflow-owned Stream is kept after its last write. + Defaults to 90 days. + runtime: The runtime of the backend's memory Store. Defaults to the default + runtime. + name: The backend name that Stream handles store. + default: Make this the backend of a create with no backend name. With more than + one registered backend, exactly one must have it. + """ + config = ext_pb.MemoryStreamConfig() + if retention is not None: + config.retention.FromTimedelta(retention) + if safety_ttl is not None: + config.safety_ttl.FromTimedelta(safety_ttl) + super().__init__( + "temporalio.MemoryStreamBackend", True, config, runtime, name, default + ) diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 80dcf8c0da..e4ae5e9d94 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -68,8 +68,8 @@ def __init__( most of the arguments. Most of the same arguments need to be passed to the replayer that were passed to the worker when the workflow originally ran. Pass ``stream_backends`` with the same backends as the Worker's - client. A replay makes no Stream call, so the Store doesn't need to be - reachable. + client, such as :py:class:`temporalio.contrib.stream.RedisStreamBackend`. + A replay makes no Stream call, so the Store doesn't need to be reachable. Note, unlike the worker, for the replayer the workflow_task_executor will default to a new thread pool executor with no max_workers set that diff --git a/tests/contrib/streams/conftest.py b/tests/contrib/streams/conftest.py new file mode 100644 index 0000000000..594a706ae7 --- /dev/null +++ b/tests/contrib/streams/conftest.py @@ -0,0 +1,77 @@ +import os +import shutil +import socket +import subprocess +import time +import uuid +from collections.abc import AsyncIterator, Iterator + +import pytest + +from temporalio.client import Client +from temporalio.common import StreamBackend +from temporalio.contrib.stream import RedisStreamBackend +from temporalio.contrib.stream.testing import MemoryStreamBackend + +# Ports that other local services own. +_RESERVED = {7333, 7861, 7931, 6431, 7101, 7102, 7103} + + +def _free_port() -> int: + while True: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + if port not in _RESERVED: + return port + + +# One scratch server serves every test module, and each test uses its own key prefix. To use a +# known port, set TEMPORAL_STREAMS_REDIS_URL to your own server. +@pytest.fixture(scope="session") +def redis_url() -> Iterator[str]: + given = os.environ.get("TEMPORAL_STREAMS_REDIS_URL") + if given: + yield given + return + server = shutil.which("redis-server") + if server is None: + # CI sets this where it installs Redis, so the Redis lane can't skip there unseen. + if os.environ.get("TEMPORAL_STREAMS_REQUIRE_REDIS"): + pytest.fail("redis-server is not installed") + pytest.skip("redis-server is not installed") + port = _free_port() + process = subprocess.Popen( + [server, "--port", str(port), "--save", "", "--appendonly", "no"], + stdout=subprocess.DEVNULL, + ) + try: + for _ in range(100): + try: + socket.create_connection(("127.0.0.1", port), timeout=0.1).close() + break + except OSError: + time.sleep(0.05) + yield f"redis://127.0.0.1:{port}" + finally: + process.terminate() + process.wait() + + +@pytest.fixture(params=["memory", "redis"]) +async def stream_client( + client: Client, request: pytest.FixtureRequest +) -> AsyncIterator[Client]: + # Every live test runs on the memory Store and on Redis, since both must have the same + # rules. The backend is the default, so a create with no backend name uses it. + backend: StreamBackend + if request.param == "memory": + backend = MemoryStreamBackend(default=True) + else: + redis_url = request.getfixturevalue("redis_url") + backend = RedisStreamBackend( + redis_url, key_prefix=f"t-{uuid.uuid4()}", default=True + ) + config = client.config() + config["stream_backends"] = [backend] + yield Client(**config) diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index be315b7654..ce1d1bf3e9 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -1,11 +1,18 @@ +import asyncio +import dataclasses import uuid from dataclasses import dataclass +from datetime import timedelta +from typing import Any import pytest from temporalio import common, workflow +from temporalio.api.enums.v1 import StreamErrorReason from temporalio.client import Client -from temporalio.converter import DataConverter +from temporalio.contrib.stream.testing import MemoryStreamBackend +from temporalio.converter import DataConverter, DefaultFailureConverter +from temporalio.exceptions import ApplicationError from tests.contrib.streams.fake_extension import fake_client from tests.helpers import new_worker @@ -82,3 +89,280 @@ async def test_query_result_binds_its_client_stream(client: Client): assert [item.data for item in items if isinstance(item, common.Record)] == [ "from the query result" ] + + +async def test_async_with_closes_a_client_stream(stream_client: Client): + async with stream_client.create_stream(Tokens, id=f"block-{uuid.uuid4()}") as done: + await done.publish("a") + failed = stream_client.stream(Tokens, id=f"block-{uuid.uuid4()}") + with pytest.raises(ValueError): + async with stream_client.create_stream(Tokens, id=failed.id): + raise ValueError("turn failed") + async with await stream_client.create_stream( + Tokens, id=f"block-{uuid.uuid4()}" + ) as early: + await early.close(result="early") + cancelled = stream_client.stream(Tokens, id=f"block-{uuid.uuid4()}") + with pytest.raises(asyncio.CancelledError): + async with stream_client.create_stream(Tokens, id=cancelled.id): + raise asyncio.CancelledError() + for stream, result, failure, status in ( + (done, None, None, common.CloseStatus.COMPLETED), + (failed, None, "turn failed", common.CloseStatus.FAILED), + (early, "early", None, common.CloseStatus.COMPLETED), + (cancelled, None, "Cancelled", common.CloseStatus.CANCELED), + ): + items = [item async for item in stream.subscribe()] + closes = [item for item in items if isinstance(item, common.Closed)] + assert len(closes) == 1 and closes[0].result == result + assert closes[0].status == status + assert (failure is None) == (closes[0].failure is None) + assert failure is None or failure in str(closes[0].failure) + + +async def test_create_gets_a_stream_that_exists(stream_client: Client): + stream_id = f"again-{uuid.uuid4()}" + first = await stream_client.create_stream( + Tokens, id=stream_id, retention=timedelta(days=1) + ) + await first.publish("a") + # The Stream exists, so the create ignores the other Retention and returns it. + again = await stream_client.create_stream( + Tokens, id=stream_id, retention=timedelta(days=2) + ) + assert again == first + assert await again.get_offset() == await first.get_offset() > 0 + await first.close() + with pytest.raises(common.StreamClosedError) as closed: + await stream_client.create_stream(Tokens, id=stream_id) + assert closed.value.reason == StreamErrorReason.STREAM_ERROR_REASON_CLOSED + + +async def test_publish_to_a_stream_that_nobody_created_raises(stream_client: Client): + stream = stream_client.stream(Tokens, id=f"missing-{uuid.uuid4()}") + with pytest.raises(common.StreamNotFoundError) as missing: + await stream.publish("lost") + assert missing.value.reason == StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND + + +async def test_next_offset_resumes_after_the_last_item(stream_client: Client): + stream = await stream_client.create_stream(Tokens, id=f"resume-{uuid.uuid4()}") + last = await stream.publish("a", "b") + tokens = stream.subscribe(wait_new_records=False) + assert tokens.next_offset == 0 + assert [item.data async for item in tokens if isinstance(item, common.Record)] == [ + "a", + "b", + ] + assert tokens.next_offset == last + 1 + new = stream.subscribe(offset=common.END, wait_new_records=False) + assert new.next_offset is None + assert [item async for item in new] == [] + assert new.next_offset == last + 1 + await stream.publish("c") + resumed = stream.subscribe(offset=tokens.next_offset, wait_new_records=False) + assert [item.data async for item in resumed if isinstance(item, common.Record)] == [ + "c" + ] + + +class _EncodesCommonAttributes(DefaultFailureConverter): + def __init__(self) -> None: + super().__init__(encode_common_attributes=True) + + +async def test_closed_failure_decodes_with_the_clients_converter(stream_client: Client): + # The failure converter puts the message in the encoded attributes, which only the payload + # converter reads. + config = stream_client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, failure_converter_class=_EncodesCommonAttributes + ) + client = Client(**config) + stream = await client.create_stream(Tokens, id=f"failed-{uuid.uuid4()}") + await stream.close( + failure=ApplicationError("tool failed", "detail", {"k": 1}, type="Tool") + ) + [closed] = [item async for item in stream.subscribe()] + assert isinstance(closed, common.Closed) + assert isinstance(closed.failure, ApplicationError) + assert closed.failure.message == "tool failed" and closed.failure.type == "Tool" + assert closed.failure.details == ("detail", {"k": 1}) + + +async def test_last_reads_the_newest_records_then_follows(stream_client: Client): + stream = await stream_client.create_stream(Tokens, id=f"last-{uuid.uuid4()}") + await stream.publish("a", "b", "c") + # A chat view shows the last few Records first, then waits for new ones. + tail = stream.subscribe(last=2) + assert tail.next_offset is None + shown = [await anext(tail), await anext(tail)] + await stream.publish("d") + shown.append(await anext(tail)) + await tail.aclose() + assert [item.data for item in shown if isinstance(item, common.Record)] == [ + "b", + "c", + "d", + ] + newest = stream.subscribe(last=1, wait_new_records=False) + assert [item.data async for item in newest if isinstance(item, common.Record)] == [ + "d" + ] + with pytest.raises(ValueError, match="no offset"): + stream.subscribe(last=2, offset=0) + + +@common.defn(name="turns") +class TurnDicts(common.StreamDefinition[dict[str, Any]]): ... + + +async def test_payload_type_of_the_definition_decodes_each_value( + stream_client: Client, +): + stream = await stream_client.create_stream(Turns, id=f"typed-{uuid.uuid4()}") + await stream.publish(Turn("model")) + + async def read(definition: type[common.StreamDefinition[Any]]) -> list[object]: + items = stream_client.stream(definition, id=stream.id).subscribe( + wait_new_records=False + ) + return [item.data async for item in items if isinstance(item, common.Record)] + + assert await read(Turns) == [Turn("model")] + assert await read(TurnDicts) == [{"kind": "model", "text": ""}] + + +@common.defn(name="turns") +class LooseTurns(common.StreamDefinition[str]): ... + + +async def test_record_that_fails_to_decode_raises_and_resumes_past_it( + stream_client: Client, +): + stream_id = f"poison-{uuid.uuid4()}" + writer = await stream_client.create_stream(LooseTurns, id=stream_id) + bad = await writer.publish("not a turn") + reader = stream_client.stream(Turns, id=stream_id) + await reader.publish(Turn("after")) + read = reader.subscribe(wait_new_records=False) + with pytest.raises(common.StreamError, match=f"offset {bad}"): + await anext(read) + assert read.next_offset == bad + 1 + resumed = reader.subscribe(offset=read.next_offset, wait_new_records=False) + assert [item.data async for item in resumed if isinstance(item, common.Record)] == [ + Turn("after") + ] + + +async def test_async_with_refuses_a_workflow_stream_before_the_block( + stream_client: Client, +): + # The close at the block end would raise and hide the block's own error. + owned: common.ClientStream[Turn] = common.ClientStream( + id="agent-1:run-1:turns", workflow_id="agent-1", first_execution_run_id="run-1" + ) + ran = False + with pytest.raises(common.StreamNotOwnerError): + async with stream_client.stream(Turns, stream=owned): + ran = True + assert not ran + + +async def test_record_over_the_store_limit_raises_too_large(stream_client: Client): + # The client's writer splits a batch at the Store's limit, so only one Record can be over it. + stream = await stream_client.create_stream(Tokens, id=f"large-{uuid.uuid4()}") + with pytest.raises(common.StreamTooLargeError) as refused: + await stream.publish("x" * 2_000_000) + assert refused.value.reason == StreamErrorReason.STREAM_ERROR_REASON_TOO_LARGE + assert refused.value.non_retryable + assert await stream.get_offset() == 0 + + +async def test_read_keeps_no_offsets_of_client_records(stream_client: Client): + # Nothing withdraws a client's Records, so a long UI read of them keeps no state. + stream = await stream_client.create_stream(Tokens, id=f"chat-{uuid.uuid4()}") + await stream.publish("a", "b", "c") + read = stream.subscribe(wait_new_records=False) + assert [item.data async for item in read if isinstance(item, common.Record)] == [ + "a", + "b", + "c", + ] + assert read._shown == {} + + +async def test_create_of_an_expired_stream_raises(stream_client: Client): + # An open client Stream expires one Retention after its last write. + stream_id = f"expired-{uuid.uuid4()}" + retention = timedelta(milliseconds=200) + stream = await stream_client.create_stream( + Tokens, id=stream_id, retention=retention + ) + await stream.publish("old") + await asyncio.sleep(0.5) + with pytest.raises(common.StreamExpiredError): + await stream_client.create_stream(Tokens, id=stream_id, retention=retention) + + +async def test_named_producer_fences_its_earlier_epoch(stream_client: Client): + stream = await stream_client.create_stream(Tokens, id=f"named-{uuid.uuid4()}") + # A server that restarts writes as the same name with a higher epoch, so the zombie that it + # replaces can't write after it. + before = stream.producer("ingest", epoch=1) + await before.publish("before restart") + after = stream.producer("ingest", epoch=2) + await after.publish("after restart") + with pytest.raises(common.StreamFencedError): + await before.publish("zombie") + # The earlier epoch's Records are real messages, so nothing withdraws them. + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert [type(item) for item in items] == [common.Record, common.Record] + with pytest.raises(ValueError): + stream.producer("ingest", epoch=0) + + +async def test_clients_of_one_runtime_share_the_backend_calls(stream_client: Client): + # A service that makes a client for each tenant keeps one set of Redis connections. + other = Client(**stream_client.config()) + [backend] = stream_client.config(active_config=True)["stream_backends"] + assert backend._backend_for(other) is backend._backend_for(stream_client) + stream = await other.create_stream(Tokens, id=f"shared-{uuid.uuid4()}") + await stream.publish("from the other client") + again = stream_client.stream(Tokens, id=stream.id) + items = [item async for item in again.subscribe(wait_new_records=False)] + assert [item.data for item in items if isinstance(item, common.Record)] == [ + "from the other client" + ] + + +async def test_default_is_native_with_no_default_backend(client: Client): + config = client.config() + config["stream_backends"] = [MemoryStreamBackend()] + one = Client(**config) + with pytest.raises(common.StreamConfigError, match="default is native"): + one.stream(Tokens, id="d-2") + with pytest.raises(common.StreamConfigError, match="default is native"): + await one.create_stream(Tokens, id="d-2") + assert one.stream(Tokens, id="d-2", backend="memory").backend == "memory" + assert (await one.create_stream(Tokens, id="d-2", backend="memory")).backend == ( + "memory" + ) + + +async def test_stream_takes_the_handle_backend_or_the_default(client: Client): + config = client.config() + config["stream_backends"] = [ + MemoryStreamBackend(name="a", default=True), + MemoryStreamBackend(name="b"), + ] + two = Client(**config) + created = await two.create_stream(Tokens, id="moving", backend="b") + await created.publish("on b") + assert await two.stream(Tokens, stream=created).get_offset() > 0 + # A handle with no backend name uses the default backend. + assert two.stream(Tokens, stream=common.ClientStream(id="moving")).backend == "a" + with pytest.raises(ValueError): + two.stream(Tokens, stream=created, backend="a") + with pytest.raises(TypeError): + two.stream(Tokens)