[Brook S7] Added Workflow Streams, which write through Local Activities.

temporalio/contrib/stream/__init__.py

@@ -1,8 +1,8 @@
 """The Redis Stream backend.
 
-Register it on the client. A client creates a Stream with
-:py:attr:`temporalio.client.Client.create_stream`. ``backend=`` picks a backend that isn't the
-default::
+Register it on the client. Workflow code and clients create Streams with
+:py:data:`temporalio.workflow.create_stream` and :py:attr:`temporalio.client.Client.create_stream`.
+``backend=`` picks a backend that isn't the default::
 
     from temporalio.contrib.stream import RedisStreamBackend
 
@@ -11,6 +11,10 @@
         stream_backends=[RedisStreamBackend("redis://localhost", default=True)],
     )
 
+A Workflow's close is a regular Activity on the Workflow's task queue. So a Worker on that queue
+that polls Activities must run the close, and must register the Redis backend. A Worker with
+``no_remote_activities=True`` logs a warning, since the close waits for another Worker.
+
 .. warning::
     This package is experimental.
 """
@@ -34,6 +38,11 @@ class RedisStreamBackend(ExtensionStreamBackend):
     each Stream ref that it creates keeps that name. A Worker serves only one backend. A client
     can have more than one, and ``backend=`` picks one.
 
+    A Workflow's close is a regular Activity on the Workflow's task queue, which runs after the
+    Workflow Task commits. So a Worker on that queue that polls Activities must have this
+    backend. A Worker with ``no_remote_activities=True`` logs a warning, since the close waits
+    for another Worker.
+
     .. warning::
         This API is experimental.
     """

temporalio/contrib/stream/testing.py

@@ -19,7 +19,9 @@ class MemoryStreamBackend(ExtensionStreamBackend):
     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 its ``name``, ``memory`` by default. It has Redis's rules, so tests can swap the
-    Store without changing Workflow code.
+    Store without changing Workflow code. As on Redis, a Workflow's close is a regular Activity
+    on the Workflow's task queue, so a Worker on that queue that polls Activities must have the
+    backend.
 
     .. warning::
         This API is experimental.

temporalio/nexus/system/_support.py

@@ -456,39 +456,62 @@ def wake(_: object, done: asyncio.Future[None] = done) -> None:
 def request_id() -> str:
     """The id of one create call, which its retries send again.
 
-    In a Workflow it comes from ``workflow.uuid4()``, so a replay sends the same id.
+    Outside a Workflow it's a random UUID. In a Workflow it's empty, since the Workflow create
+    sets an id that is the same in each run of the chain.
     """
-    from temporalio.workflow import in_workflow, uuid4
+    from temporalio.workflow import in_workflow
 
-    return str(uuid4() if in_workflow() else uuid.uuid4())
+    # A workflow.uuid4() here would move the user's workflow.random() sequence, and the
+    # Workflow create replaces the id anyway.
+    return "" if in_workflow() else str(uuid.uuid4())
 
 
 # The Stream sources. Each one checks where it runs and gives the value of that context, or
-# raises where it has none to give.
-
-
-def _no_workflow_streams(source: str) -> RuntimeError:
-    # This SDK doesn't run Workflow Stream calls yet, so a generated Workflow call fails here.
-    return RuntimeError(f"Workflow code can't call the Stream service yet ({source})")
+# raises where it has none to give. The Streams API imports this module, so the sources import
+# the API only when they run.
 
 
 def workflow_writer() -> typing.Any:
-    """The writer of the Workflow run's writes."""
-    raise _no_workflow_streams("workflow_writer")
+    """The writer of the Workflow run's writes, with the previous run in a retry run."""
+    from temporalio.streams import _ids
+    from temporalio.streams.system import models
+    from temporalio.workflow import info
+
+    return models.StreamWriter(
+        producer=_ids.workflow_producer(info().run_id),
+        failed_tasks=[],
+        withdraw_run=withdraw_run(),
+    )
 
 
 def withdraw_run() -> str | None:
-    """The previous run whose Records a retry run's create withdraws."""
-    from temporalio.workflow import in_workflow
+    """The previous run whose Records a retry run's first write withdraws."""
+    from temporalio.workflow import in_workflow, info
 
-    if in_workflow():
-        raise _no_workflow_streams("withdraw_run")
-    # Code outside a Workflow isn't a retry run.
-    return None
+    if not in_workflow():
+        # Code outside a Workflow isn't a retry run.
+        return None
+    run = info()
+    return run.continued_run_id if run.attempt > 1 and run.continued_run_id else None
 
 
 def client_writer() -> typing.Any:
-    """The writer of a write outside a Workflow."""
-    # Each client has its own Producer, and a generated call doesn't say which client makes it.
-    # So the SDK's Streams send their writer, and a call without one fails.
-    raise RuntimeError("A Stream write outside a Workflow needs its Producer")
+    """The writer of a write outside a Workflow, which is the Activity attempt's writer."""
+    import temporalio.api.stream.v1 as stream_v1
+    from temporalio.streams import _activity
+    from temporalio.streams.system import models
+
+    streams = _activity.current.get()
+    if streams is None:
+        # Each client has its own Producer, and a generated call doesn't say which client makes
+        # it. So the SDK's Streams send their writer, and a call without one fails.
+        raise RuntimeError("A Stream write outside a Workflow needs its Producer")
+    writer = stream_v1.StreamWriter(
+        producer=streams.producer,
+        # Only a Local Activity has the run's task facts.
+        previous_completed_task=streams.previous_completed_task,
+        failed_tasks=streams.failed_tasks,
+    )
+    return models._StreamWriterTransferTypeConverter().from_transfer_type(
+        writer, models.StreamWriter
+    )

temporalio/streams/_activity.py

@@ -0,0 +1,60 @@
+"""The Stream state of one Activity attempt."""
+
+from __future__ import annotations
+
+import asyncio
+import contextvars
+from dataclasses import dataclass, field
+
+import temporalio.api.stream.v1 as stream_v1
+import temporalio.converter
+from temporalio.streams._backend import BackendCalls
+from temporalio.streams._writer import RecordWriter
+
+
+@dataclass
+class ActivityStreams:
+    """What a Stream backend gives the Streams of one Activity attempt."""
+
+    backend: BackendCalls
+    converter: temporalio.converter.DataConverter
+    namespace: str
+    producer: stream_v1.Producer
+    loop: asyncio.AbstractEventLoop
+    # The attempt's Workflow ID, which the close Activity checks against a bound Stream's owner.
+    workflow_id: str | None = None
+    # A Local Activity's writes use these to withdraw the run's orphaned Records.
+    previous_completed_task: stream_v1.WorkflowTaskEpoch | None = None
+    failed_tasks: tuple[int, ...] = ()
+    writers: dict[str, RecordWriter] = field(default_factory=dict)
+
+    def writer(self, stream_id: str) -> RecordWriter:
+        """The attempt's writer for a Stream, one per Stream."""
+        writer = self.writers.get(stream_id)
+        if writer is None:
+            writer = RecordWriter(
+                writers=self.backend.writers(self.namespace, self.converter),
+                stream_id=stream_id,
+                producer=self.producer,
+                free=self.backend.free_writer,
+                previous_completed_task=self.previous_completed_task,
+                failed_tasks=self.failed_tasks,
+            )
+            self.writers[stream_id] = writer
+        return writer
+
+    async def close(self) -> None:
+        """Close every writer of the attempt."""
+        errors: list[BaseException] = []
+        for writer in self.writers.values():
+            try:
+                await writer.close()
+            except Exception as err:
+                errors.append(err)
+        if errors:
+            raise errors[0]
+
+
+current: contextvars.ContextVar[ActivityStreams | None] = contextvars.ContextVar(
+    "__temporal_activity_streams", default=None
+)