diff --git a/temporalio/contrib/streams/nexus/__init__.py b/temporalio/contrib/streams/nexus/__init__.py index 386f3a24cd..25f4e55f29 100644 --- a/temporalio/contrib/streams/nexus/__init__.py +++ b/temporalio/contrib/streams/nexus/__init__.py @@ -29,6 +29,17 @@ authorizer and the namespace's payload codec, the parts only the SDK has, and the stream runtime runs every rule of the contract on the store of its Worker's client. + +A Nexus operation can also hand its caller a stream. +:class:`StreamOperationHandler` is such an operation: its start attaches the +caller to the stream's notifier on the server and puts the +:class:`StreamRef` in the operation token (:func:`stream_ref_from_token`). +With :meth:`temporalio.contrib.streams.StreamStorePlugin.notify_on_append` +on, the default, the stream runtime tells the notifier after every append, which the caller +sees as operation progress. +:meth:`temporalio.contrib.streams.StreamStorePlugin.close_stream` completes +the operation, or :func:`temporalio.contrib.streams.close_workflow_stream` +from the owning Workflow. """ from temporalio.contrib.streams import StreamRef @@ -49,6 +60,10 @@ StreamAccess, TemporalStreamsHandler, ) +from temporalio.contrib.streams.nexus._operation import ( + StreamOperationHandler, + stream_ref_from_token, +) __all__ = [ "AppendInput", @@ -59,8 +74,10 @@ "RecordWire", "StreamAccess", "StreamCursor", + "StreamOperationHandler", "StreamRef", "TemporalStreams", "TemporalStreamsHandler", "TemporalStreamsHttpClient", + "stream_ref_from_token", ] diff --git a/temporalio/contrib/streams/nexus/_operation.py b/temporalio/contrib/streams/nexus/_operation.py new file mode 100644 index 0000000000..8caf917589 --- /dev/null +++ b/temporalio/contrib/streams/nexus/_operation.py @@ -0,0 +1,175 @@ +"""A Nexus operation whose start hands its caller a stream. + +The operation stays open while the stream runs. Its start attaches the +caller's callback to the stream's notifier on the server, so the notifier +tells the caller when the stream moves (Nexus operation progress) and +completes the operation when the stream closes. The caller learns which +stream it is from the operation token. The stream runtime resolves the +stream's run chain, mints the token and attaches the callback. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Generic, TypeVar + +import nexusrpc +import nexusrpc.handler + +import temporalio.nexus +from temporalio.bridge.proto.streams import ( + AttachCallerRequest, + DetachCallerRequest, + StreamOwnerKind, +) +from temporalio.bridge.streams_generated import StreamNexusService +from temporalio.client import Client +from temporalio.contrib.streams._plugin import call, store_for_client +from temporalio.contrib.streams._ref import StreamRef, stream_address +from temporalio.contrib.streams.nexus._generated.models import StreamOperationToken +from temporalio.contrib.streams.nexus._handler import _handler_error + +__all__ = ["StreamOperationHandler", "stream_ref_from_token"] + +InputT = TypeVar("InputT") +OutputT = TypeVar("OutputT") + + +def stream_ref_from_token(token: str) -> StreamRef: + """The stream a stream-returning operation's token names. + + This is how a caller learns the stream: the start answers with the + operation token, and the token carries the stream reference. + + Raises: + ValueError: ``token`` did not come from :class:`StreamOperationHandler`. + """ + return StreamOperationToken.from_token(token).ref + + +async def _caller_calls(client: Client) -> StreamNexusService: + """The stream runtime's caller calls, on the operation Worker's client's store.""" + try: + return await store_for_client(client)._nexus_service_for(client) + except ValueError as error: + # The handler's configuration is at fault, so the caller's request + # isn't refused, and a retry may land once the Worker is fixed. + raise nexusrpc.HandlerError( + f"the stream operation's Worker can't reach a stream store: {error}", + type=nexusrpc.HandlerErrorType.INTERNAL, + ) from None + + +FIRST_WAKE_HEADER = "temporal-stream-first-wake-after" +"""The start header that asks for the callback's first wake after a counter. + +A follower that reads only when woken sends it, so its idle callback hears +of the first record. A reader that reads first sends none. +""" + + +class StreamOperationHandler( + nexusrpc.handler.OperationHandler[InputT, OutputT], Generic[InputT, OutputT] +): + """A Nexus operation whose start hands its caller a stream. + + ``open_stream`` names the stream for a start: it creates or finds the + stream and returns its reference. The start then attaches the caller's + callback to that stream's notifier and answers asynchronously, with the + reference in the operation token. The callback is idle: the server tells + the caller when the stream moves only after a read that found nothing + asked for a wake, when the store has + :meth:`temporalio.contrib.streams.StreamStorePlugin.notify_on_append` on + (the default), and completes the operation with the close result when the stream closes. + A start with the ``temporal-stream-first-wake-after`` header, which + :func:`follow_stream` sends, asks for the first wake at the attach. + A cancel detaches the caller. + + The operation's output type is the close result's type. The records + themselves are read through the stream service, not carried by the + operation. + + The client of the Worker that runs the operation must carry a stream + store plugin, as the stream service's handler's does, since the stream + runtime attaches the caller through it. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + open_stream: Callable[ + [nexusrpc.handler.StartOperationContext, InputT], Awaitable[StreamRef] + ], + ) -> None: + """Hand out the stream ``open_stream`` names for each start.""" + self._open_stream = open_stream + + async def start( + self, ctx: nexusrpc.handler.StartOperationContext, input: InputT + ) -> nexusrpc.handler.StartOperationResultAsync: + """Attach the caller to the stream and answer with its reference. + + Raises: + nexusrpc.HandlerError: The caller sent no callback, so it could + never hear about the stream (``BAD_REQUEST``), or the server + didn't take the attach. + """ + if not ctx.callback_url: + # Refused before open_stream runs, so a start that can't work + # creates nothing. + raise nexusrpc.HandlerError( + "a stream-returning operation needs a caller callback; start it " + "asynchronously", + type=nexusrpc.HandlerErrorType.BAD_REQUEST, + retryable_override=False, + ) + client = temporalio.nexus.client() + service = await _caller_calls(client) + ref = await self._open_stream(ctx, input) + first_wake = ctx.headers.get(FIRST_WAKE_HEADER) + stream = stream_address(client.namespace, ref) + if ref.kind not in ("workflow", "activity", "workflow_activity"): + # The stream runtime refuses an owner kind this release doesn't open. + stream.owner_kind = StreamOwnerKind.STREAM_OWNER_KIND_UNSPECIFIED + request = AttachCallerRequest( + stream=stream, + request_id=ctx.request_id, + callback_url=ctx.callback_url, + callback_headers=dict(ctx.callback_headers), + ) + if first_wake is not None: + try: + request.wake_after_counter = int(first_wake) + except ValueError: + raise nexusrpc.HandlerError( + f"{FIRST_WAKE_HEADER} must be a counter, got {first_wake!r}", + type=nexusrpc.HandlerErrorType.BAD_REQUEST, + retryable_override=False, + ) from None + response = await call(service.attach_caller(request)) + if response.HasField("failure"): + raise _handler_error(response.failure) + return nexusrpc.handler.StartOperationResultAsync( + token=response.operation_token + ) + + async def cancel( + self, ctx: nexusrpc.handler.CancelOperationContext, token: str + ) -> None: + """Detach the caller: the stream goes on, the operation ends. + + Raises: + nexusrpc.HandlerError: ``token`` is not one this operation issued + (``BAD_REQUEST``), or the server didn't take the detach. + """ + client = temporalio.nexus.client() + service = await _caller_calls(client) + response = await call( + service.detach_caller( + DetachCallerRequest(namespace=client.namespace, operation_token=token) + ) + ) + if response.HasField("failure"): + raise _handler_error(response.failure) diff --git a/tests/contrib/streams/test_nexus_stream_operation.py b/tests/contrib/streams/test_nexus_stream_operation.py new file mode 100644 index 0000000000..5b4460cfcb --- /dev/null +++ b/tests/contrib/streams/test_nexus_stream_operation.py @@ -0,0 +1,617 @@ +"""A Nexus operation whose start hands its caller a stream. + +The handler attaches the caller's callback to the stream's notifier on the +server and puts the stream reference in the operation token. Core notifies +the notifier after each append, a Workflow's flushes included, and closing +the stream completes the operation with the close result. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +import nexusrpc +import nexusrpc.handler +import pytest + +import temporalio.nexus +from temporalio import workflow +from temporalio.api.enums.v1 import StreamOwnerKind +from temporalio.api.stream.v1 import StreamReference +from temporalio.api.workflowservice.v1 import ( + DescribeStreamNotifierRequest, + DescribeStreamNotifierResponse, +) +from temporalio.client import Client +from temporalio.contrib.streams import ( + RecordKind, + StreamRef, + StreamStorePlugin, + StreamUnsupportedError, + close_workflow_stream, + stream_handle, + stream_producer, +) +from temporalio.contrib.streams.memory import MemoryStreams +from temporalio.contrib.streams.nexus import ( + StreamOperationHandler, + stream_ref_from_token, +) +from temporalio.contrib.streams.nexus._operation import FIRST_WAKE_HEADER +from temporalio.service import RPCError, RPCStatusCode +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Replayer, Worker +from temporalio.worker._nexus import _NexusTaskCancellation +from tests.contrib.streams._support import call_workflow_service, streams_crate_nexus +from tests.helpers.nexus import make_nexus_endpoint_name + +REF = StreamRef.for_workflow("owner-1", topic="tokens") + + +def start_context(callback_url: str | None = "temporal://system") -> Any: + return nexusrpc.handler.StartOperationContext( + service="ChatService", + operation="chat", + headers={}, + task_cancellation=_NexusTaskCancellation(), + request_id="attach-1", + callback_url=callback_url, + callback_headers={"nexus-callback-token": "abc"}, + ) + + +def cancel_context() -> Any: + return nexusrpc.handler.CancelOperationContext( + service="ChatService", + operation="chat", + headers={}, + task_cancellation=_NexusTaskCancellation(), + ) + + +# Core attaches and detaches the caller and mints the token. Its tests in +# crates/streams/src/nexus/operation.rs cover the attach, the token, a pinned +# ref and the chain a cancel detaches from. These cover what stays in lang. + + +async def test_a_start_without_a_callback_is_refused_before_it_opens_the_stream() -> ( + None +): + opened: list[str] = [] + + async def open_stream(_ctx: Any, prompt: str) -> StreamRef: + opened.append(prompt) + return REF + + with pytest.raises(nexusrpc.HandlerError) as raised: + await StreamOperationHandler(open_stream).start( + start_context(callback_url=None), "hi" + ) + + assert raised.value.type == nexusrpc.HandlerErrorType.BAD_REQUEST + assert not opened + + +def _raised(error: nexusrpc.HandlerError) -> tuple[str, bool | None, str | None]: + """The type, retryability and stream condition class a caller sees.""" + failure = error.original_failure + details: Mapping[str, Any] = (failure.details if failure else None) or {} + cause = details.get("cause", {}).get("applicationFailureInfo", {}) + return error.type.name, error.retryable_override, cause.get("type") + + +async def test_a_start_on_a_worker_without_a_store_fails_before_it_opens_the_stream( + client: Client, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("temporalio.nexus.client", lambda: client) + opened: list[str] = [] + + async def open_stream(_ctx: Any, prompt: str) -> StreamRef: + opened.append(prompt) + return REF + + with pytest.raises(nexusrpc.HandlerError) as raised: + await StreamOperationHandler(open_stream).start(start_context(), "hi") + + # The handler's configuration is at fault, not the caller's request. + assert raised.value.type == nexusrpc.HandlerErrorType.INTERNAL + assert not opened + + +async def test_a_start_for_an_owner_that_does_not_exist_is_not_found( + client: Client, monkeypatch: pytest.MonkeyPatch +) -> None: + store_client = await connect(client, MemoryStreams()) + monkeypatch.setattr("temporalio.nexus.client", lambda: store_client) + + async def open_missing(_ctx: Any, _prompt: str) -> StreamRef: + return StreamRef.for_workflow(f"missing-{uuid.uuid4()}", topic="tokens") + + with pytest.raises(nexusrpc.HandlerError) as raised: + await StreamOperationHandler(open_missing).start(start_context(), "hi") + + # Core's table, as the stream service's operations answer with it. + assert _raised(raised.value) == ("NOT_FOUND", False, "StreamNotFoundError") + + +async def test_a_start_for_an_activity_that_does_not_exist_is_not_found( + client: Client, monkeypatch: pytest.MonkeyPatch +) -> None: + # The ref reaches Core as an Activity owner, which Core checks with a + # describe of the Activity. + store_client = await connect(client, MemoryStreams()) + monkeypatch.setattr("temporalio.nexus.client", lambda: store_client) + + async def open_missing(_ctx: Any, _prompt: str) -> StreamRef: + return StreamRef.for_activity( + f"missing-{uuid.uuid4()}", run_id=str(uuid.uuid4()), topic="tokens" + ) + + with pytest.raises(nexusrpc.HandlerError) as raised: + await StreamOperationHandler(open_missing).start(start_context(), "hi") + + assert _raised(raised.value) == ("NOT_FOUND", False, "StreamNotFoundError") + + +async def test_a_cancel_with_a_token_that_names_no_stream_is_a_bad_request( + client: Client, monkeypatch: pytest.MonkeyPatch +) -> None: + store_client = await connect(client, MemoryStreams()) + monkeypatch.setattr("temporalio.nexus.client", lambda: store_client) + + async def open_stream(_ctx: Any, _prompt: str) -> StreamRef: + return REF + + with pytest.raises(nexusrpc.HandlerError) as raised: + await StreamOperationHandler(open_stream).cancel( + cancel_context(), "not-a-token" + ) + + assert _raised(raised.value) == ("BAD_REQUEST", False, "ValueError") + + +def test_a_token_reads_as_the_shared_vectors_say() -> None: + vectors = json.loads( + (streams_crate_nexus() / "operation_token_vectors.json").read_text() + ) + for vector in vectors["vectors"] + vectors["decode_only"]: + ref = vector["ref"] + assert stream_ref_from_token(vector["token"]) == StreamRef( + kind=ref["kind"], + workflow_id=ref["workflow_id"], + chain_run_id=ref.get("chain_run_id"), + topic=ref["topic"], + activity_id=ref.get("activity_id"), + ) + for token in vectors["refused"]: + try: + decoded = stream_ref_from_token(token) + except ValueError: + continue + # A ref of a kind this release lacks decodes, and opening it refuses it. + with pytest.raises(StreamUnsupportedError): + decoded._require_supported() # type: ignore[reportPrivateUsage] + + +# Live: a caller Workflow sees the stream's progress and the close result. +# These need a server with Nexus progress and the stream notifier, and skip on +# any other server. + + +@nexusrpc.service +class ChatService: + chat: nexusrpc.Operation[str, str] + + +@dataclass +class Observed: + counters: list[int] + result: str + + +@workflow.defn(name="StreamOperationCaller") +class StreamOperationCaller: + def __init__(self) -> None: + self._counters: list[int] = [] + + @workflow.run + async def run(self, endpoint: str) -> Observed: + nexus_client = workflow.create_nexus_client( + service=ChatService, endpoint=endpoint + ) + # This caller never reads, so it asks for the first wake at the start, + # as a follower does. After that delivery its callback is idle. + handle = await nexus_client.start_operation( + ChatService.chat, "hello", headers={FIRST_WAKE_HEADER: "0"} + ) + last = 0 + while (progress := await handle.progress(after_counter=last)) is not None: + self._counters.append(progress.counter) + last = progress.counter + return Observed(counters=list(self._counters), result=await handle) + + @workflow.query + def counters(self) -> list[int]: + return list(self._counters) + + +@workflow.defn(name="OwnerUntilDone") +class OwnerUntilDone: + """Owns a stream until signaled done.""" + + def __init__(self) -> None: + self._done = False + + @workflow.run + async def run(self) -> None: + await workflow.wait_condition(lambda: self._done) + + @workflow.signal + def done(self) -> None: + self._done = True + + +async def connect(client: Client, store: StreamStorePlugin) -> Client: + """A client of the test server that carries ``store``.""" + return await Client.connect( + client.service_client.config.target_host, + namespace=client.namespace, + plugins=[store], + ) + + +async def _skip_without_notifier(client: Client, ref: StreamRef) -> None: + try: + await call_workflow_service( + client, + "DescribeStreamNotifier", + DescribeStreamNotifierRequest( + namespace=client.namespace, + stream_ref=StreamReference( + owner_kind=StreamOwnerKind.STREAM_OWNER_KIND_WORKFLOW, + workflow_id=ref.workflow_id, + topic=ref.topic, + ), + ), + DescribeStreamNotifierResponse, + ) + except RPCError as err: + if err.status == RPCStatusCode.UNIMPLEMENTED: + pytest.skip(f"server has no stream notifier: {err.message}") + # A server that keys the notifier by run chain refuses the probe's + # empty run id, which still shows it has the notifier. + if err.status not in ( + RPCStatusCode.NOT_FOUND, + RPCStatusCode.INVALID_ARGUMENT, + ): + raise + + +async def _wait_for(predicate: Any, attempts: int = 100) -> Any: + for _ in range(attempts): + value = await predicate() + if value: + return value + await asyncio.sleep(0.1) + return None + + +def _store_for(backing: str) -> StreamStorePlugin: + if backing == "memory": + return MemoryStreams().notify_on_append() + url = os.environ.get("STREAMS_REDIS_URL") + if not url: + pytest.skip("set STREAMS_REDIS_URL to run the Redis cases") + from temporalio.contrib.streams.redis import RedisStreams + + return RedisStreams(url, key_prefix=f"test-{uuid.uuid4().hex}").notify_on_append() + + +async def _values(client: Client, ref: StreamRef) -> list[Any]: + return [ + record.value + async for record in stream_handle(client, ref).read() + if record.kind is RecordKind.DATA + ] + + +async def test_a_caller_sees_appends_as_progress_and_the_close_as_the_result( + client: Client, env: WorkflowEnvironment +) -> None: + ref = StreamRef.for_workflow(f"owner-{uuid.uuid4()}", topic="tokens") + await _skip_without_notifier(client, ref) + task_queue = f"stream-operation-{uuid.uuid4()}" + endpoint = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint, task_queue) + store = MemoryStreams().notify_on_append() + streams_client = await connect(client, store) + started = asyncio.Event() + + async def open_owned_stream(_ctx: Any, _prompt: str) -> StreamRef: + # The notifier is keyed by the owner's run chain, so the owner runs. + await streams_client.start_workflow( + OwnerUntilDone.run, id=ref.workflow_id, task_queue=task_queue + ) + started.set() + return ref + + @nexusrpc.handler.service_handler(service=ChatService) + class ChatServiceHandler: + @nexusrpc.handler.operation_handler + def chat(self) -> nexusrpc.handler.OperationHandler[str, str]: + return StreamOperationHandler(open_owned_stream) + + async with Worker( + streams_client, + task_queue=task_queue, + workflows=[StreamOperationCaller, OwnerUntilDone], + nexus_service_handlers=[ChatServiceHandler()], + ): + handle = await streams_client.start_workflow( + StreamOperationCaller.run, + endpoint, + id=f"stream-operation-caller-{uuid.uuid4()}", + task_queue=task_queue, + ) + await asyncio.wait_for(started.wait(), 10) + # Progress that reaches the caller before its started event is dropped, + # so the first append waits for the operation to be started. + + async def operation_started() -> bool: + history = await handle.fetch_history() + return any( + event.HasField("nexus_operation_started_event_attributes") + for event in history.events + ) + + assert await _wait_for(operation_started) + producer = stream_handle(streams_client, ref).producer( + producer_id="writer", attempt=1 + ) + await producer.append("a") + + async def first_progress() -> list[int]: + return await handle.query(StreamOperationCaller.counters) + + assert await _wait_for(first_progress), "the caller never saw progress" + await producer.append("b", "c") + await store.close_stream(streams_client, ref, "done") + await store.flush_notifications() + + observed = await asyncio.wait_for(handle.result(), 20) + await streams_client.get_workflow_handle(ref.workflow_id).signal( + OwnerUntilDone.done + ) + + assert observed.result == "done" + assert observed.counters, observed + assert observed.counters == sorted(set(observed.counters)), observed + + +TOKENS_TOPIC = "tokens" + + +@workflow.defn(name="StreamOwnerThatPublishes") +class StreamOwnerThatPublishes: + """Publishes its tokens on its own stream, then waits to be told to end.""" + + def __init__(self) -> None: + self._done = False + + @workflow.run + async def run(self, tokens: list[str]) -> None: + writer = stream_producer(TOKENS_TOPIC) + for token in tokens: + writer.append(token) + await workflow.sleep(timedelta(milliseconds=200)) + await writer.finish() + await workflow.wait_condition(lambda: self._done) + + @workflow.signal + def done(self) -> None: + self._done = True + + +async def test_a_workflows_own_publishes_reach_the_caller_as_progress( + client: Client, env: WorkflowEnvironment +) -> None: + owner_id = f"owner-{uuid.uuid4()}" + ref = StreamRef.for_workflow(owner_id, topic=TOKENS_TOPIC) + await _skip_without_notifier(client, ref) + task_queue = f"stream-owner-{uuid.uuid4()}" + endpoint = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint, task_queue) + store = MemoryStreams().notify_on_append() + streams_client = await connect(client, store) + + async def start_owner(_ctx: Any, _prompt: str) -> StreamRef: + await temporalio.nexus.client().start_workflow( + StreamOwnerThatPublishes.run, + ["a", "b", "c"], + id=owner_id, + task_queue=task_queue, + ) + return ref + + @nexusrpc.handler.service_handler(service=ChatService) + class ChatServiceHandler: + @nexusrpc.handler.operation_handler + def chat(self) -> nexusrpc.handler.OperationHandler[str, str]: + return StreamOperationHandler(start_owner) + + async with Worker( + streams_client, + task_queue=task_queue, + workflows=[StreamOperationCaller, StreamOwnerThatPublishes], + nexus_service_handlers=[ChatServiceHandler()], + ): + caller = await streams_client.start_workflow( + StreamOperationCaller.run, + endpoint, + id=f"stream-owner-caller-{uuid.uuid4()}", + task_queue=task_queue, + ) + + # A Workflow's records are notified when the activation that + # published them completes, and the caller's first wake hears of them. + async def first_progress() -> list[int] | None: + counters = await caller.query(StreamOperationCaller.counters) + return counters or None + + assert await _wait_for(first_progress), "the caller saw no progress" + await streams_client.get_workflow_handle(owner_id).signal( + StreamOwnerThatPublishes.done + ) + await store.close_stream(streams_client, ref, "3 tokens") + observed = await asyncio.wait_for(caller.result(), 20) + + assert observed.result == "3 tokens" + assert observed.counters == sorted(set(observed.counters)), observed + + +@workflow.defn(name="StreamOwnerThatCloses") +class StreamOwnerThatCloses: + """Publishes its tokens on its own stream, then closes it from Workflow code.""" + + @workflow.run + async def run(self, tokens: list[str]) -> None: + writer = stream_producer(TOKENS_TOPIC) + for token in tokens: + writer.append(token) + await workflow.sleep(timedelta(milliseconds=200)) + await close_workflow_stream(TOKENS_TOPIC, result=f"{len(tokens)} tokens") + # A second close of the topic keeps the first outcome. + await close_workflow_stream(TOKENS_TOPIC, result="again") + + +async def test_a_workflow_closes_its_own_stream( + client: Client, env: WorkflowEnvironment +) -> None: + owner_id = f"owner-{uuid.uuid4()}" + ref = StreamRef.for_workflow(owner_id, topic=TOKENS_TOPIC) + await _skip_without_notifier(client, ref) + task_queue = f"stream-closer-{uuid.uuid4()}" + endpoint = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint, task_queue) + store = MemoryStreams().notify_on_append() + streams_client = await connect(client, store) + + async def start_owner(_ctx: Any, _prompt: str) -> StreamRef: + await temporalio.nexus.client().start_workflow( + StreamOwnerThatCloses.run, + ["a", "b", "c"], + id=owner_id, + task_queue=task_queue, + ) + return ref + + @nexusrpc.handler.service_handler(service=ChatService) + class ChatServiceHandler: + @nexusrpc.handler.operation_handler + def chat(self) -> nexusrpc.handler.OperationHandler[str, str]: + return StreamOperationHandler(start_owner) + + async with Worker( + streams_client, + task_queue=task_queue, + workflows=[StreamOperationCaller, StreamOwnerThatCloses], + nexus_service_handlers=[ChatServiceHandler()], + ): + observed = await asyncio.wait_for( + streams_client.execute_workflow( + StreamOperationCaller.run, + endpoint, + id=f"stream-closer-caller-{uuid.uuid4()}", + task_queue=task_queue, + ), + 30, + ) + owner_result = await streams_client.get_workflow_handle(owner_id).result() + + assert observed.result == "3 tokens" + assert owner_result is None + + +@workflow.defn(name="OwnerThatClosesInOneTask") +class OwnerThatClosesInOneTask: + """Publishes its tokens and closes its stream with no wait between, then + keeps running, so only the close can end a read.""" + + def __init__(self) -> None: + self._done = False + + @workflow.run + async def run(self, count: int) -> None: + writer = stream_producer(TOKENS_TOPIC) + for n in range(count): + writer.append(f"t{n}") + await close_workflow_stream(TOKENS_TOPIC, result=f"{count} tokens") + await workflow.wait_condition(lambda: self._done) + + @workflow.signal + def done(self) -> None: + self._done = True + + +@pytest.mark.parametrize("backing", ["memory", "redis"]) +async def test_a_close_right_after_the_last_publishes_comes_after_them( + client: Client, env: WorkflowEnvironment, backing: str +) -> None: + store = _store_for(backing) + owner_id = f"owner-{uuid.uuid4()}" + ref = StreamRef.for_workflow(owner_id, topic=TOKENS_TOPIC) + await _skip_without_notifier(client, ref) + task_queue = f"stream-closer-{uuid.uuid4()}" + endpoint = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint, task_queue) + streams_client = await connect(client, store) + + async def start_owner(_ctx: Any, _prompt: str) -> StreamRef: + await temporalio.nexus.client().start_workflow( + OwnerThatClosesInOneTask.run, 50, id=owner_id, task_queue=task_queue + ) + return ref + + @nexusrpc.handler.service_handler(service=ChatService) + class ChatServiceHandler: + @nexusrpc.handler.operation_handler + def chat(self) -> nexusrpc.handler.OperationHandler[str, str]: + return StreamOperationHandler(start_owner) + + async with Worker( + streams_client, + task_queue=task_queue, + workflows=[StreamOperationCaller, OwnerThatClosesInOneTask], + nexus_service_handlers=[ChatServiceHandler()], + ): + observed = await asyncio.wait_for( + streams_client.execute_workflow( + StreamOperationCaller.run, + endpoint, + id=f"stream-closer-caller-{uuid.uuid4()}", + task_queue=task_queue, + ), + 30, + ) + # The operation completed, so the store is closed and holds every + # record: a read gets them all and ends while the owner still runs. + values = await asyncio.wait_for(_values(streams_client, ref), 10) + owner = streams_client.get_workflow_handle(owner_id) + await owner.signal(OwnerThatClosesInOneTask.done) + await owner.result() + history = await owner.fetch_history() + + assert observed.result == "50 tokens" + assert values == [f"t{n}" for n in range(50)] + # The flushes and the close are Local Activity markers, so a replay needs + # no store. + await Replayer( + workflows=[OwnerThatClosesInOneTask], plugins=[MemoryStreams()] + ).replay_workflow(history)