[Brook] Ended reads of a StreamRef with the owning run.

scripts/_proto/temporal/api/streamservice/v1/stream_service.proto

@@ -13,6 +13,10 @@ import "temporal/api/failure/v1/message.proto";
 // Identifies one Stream. The id is opaque to callers.
 message StreamRef {
     string id = 1;
+    // The Workflow that owns the Stream. Set on a Workflow's ref, empty on a standalone,
+    // client or Activity ref. The Store keys on id only and never reads these.
+    string workflow_id = 2;
+    string first_execution_run_id = 3;
 }
 
 enum ProducerKind {

scripts/_wit/stream-service.wit

@@ -36,6 +36,10 @@ interface stream-service {
   record stream-ref {
     /// @nexus.doc "The Stream id. 1 to 4096 bytes after encoding."
     id: string,
+    /// @nexus.doc "The Workflow that owns the Stream. Empty on a standalone, client or Activity ref."
+    workflow-id: option<string>,
+    /// @nexus.doc "The first run id of the owning run chain. Empty when workflow-id is."
+    first-execution-run-id: option<string>,
   }
 
   /// @nexus.proto "temporal.api.streamservice.v1.ProducerKind"

temporalio/api/streamservice/v1/stream_service_pb2.pyi

@@ -142,14 +142,31 @@ class StreamRef(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
 
     ID_FIELD_NUMBER: builtins.int
+    WORKFLOW_ID_FIELD_NUMBER: builtins.int
+    FIRST_EXECUTION_RUN_ID_FIELD_NUMBER: builtins.int
     id: builtins.str
+    workflow_id: builtins.str
+    """The Workflow that owns the Stream. Set on a Workflow's ref, empty on a standalone,
+    client or Activity ref. The Store keys on id only and never reads these.
+    """
+    first_execution_run_id: builtins.str
     def __init__(
         self,
         *,
         id: builtins.str = ...,
+        workflow_id: builtins.str = ...,
+        first_execution_run_id: builtins.str = ...,
     ) -> None: ...
     def ClearField(
-        self, field_name: typing_extensions.Literal["id", b"id"]
+        self,
+        field_name: typing_extensions.Literal[
+            "first_execution_run_id",
+            b"first_execution_run_id",
+            "id",
+            b"id",
+            "workflow_id",
+            b"workflow_id",
+        ],
     ) -> None: ...
 
 global___StreamRef = StreamRef

temporalio/bridge/Cargo.toml

@@ -94,7 +94,7 @@ url = "2.2"
 # crates, so the bridge sees one type of each.
 [dependencies.temporalio-streams]
 git = "https://github.com/moetemp/sdk-streams"
-rev = "1c8c01f5ce87ca4ebb3fba155430920ff0a95c5c"
+rev = "813f9d129da2e90f86e8542e9a1ad766e2e7bf4f"
 default-features = false
 optional = true
 

temporalio/contrib/streams/_client.py

@@ -138,6 +138,7 @@ async def find_id() -> str:
             raise ValueError(
                 "The server gave no first run id, so the Stream id is unknown"
             )
+        stream._owner = (handle.id, first)
         return stream_id_of(first)
 
     def stream_id_of(first_run_id: str) -> str:
@@ -150,21 +151,25 @@ def stream_id_of(first_run_id: str) -> str:
         if handle.first_execution_run_id
         else find_id
     )
-    return ClientStream(
+    stream = ClientStream(
         handle._client,
         backend_of_client(handle._client),
         stream_id,
         batch_options(batch_interval, max_batch_size, max_retry_duration),
         wait_for_create=wait_for_create,
         handle=handle if end_with_run else None,
     )
+    if handle.first_execution_run_id:
+        stream._owner = (handle.id, handle.first_execution_run_id)
+    return stream
 
 
 def for_ref(
     client: temporalio.client.Client,
     ref: StreamRef,
     *,
     wait_for_create: timedelta | None = None,
+    end_with_run: bool = False,
     batch_interval: timedelta | None = None,
     max_batch_size: int | None = None,
     max_retry_duration: timedelta | None = None,
@@ -178,17 +183,33 @@ def for_ref(
         ref: The reference of the Stream.
         wait_for_create: How long a read waits for a Stream that doesn't exist yet. Defaults
             to no wait.
+        end_with_run: End a read when the owning Workflow's run is no longer running, as
+            :py:func:`for_handle` does. Only a Workflow's Stream has an owner.
         batch_interval: How often this client's writes flush. Defaults to 100 ms.
         max_batch_size: At most this many Records in one batch.
         max_retry_duration: How long a flush retries a Store failure.
+
+    Raises:
+        ValueError: ``end_with_run`` is set, and the ref has no owning Workflow.
     """
-    return ClientStream(
+    handle = None
+    if end_with_run:
+        if not ref.workflow_id or not ref.first_execution_run_id:
+            raise ValueError("end_with_run needs a ref of a Workflow's Stream")
+        handle = client.get_workflow_handle(
+            ref.workflow_id, first_execution_run_id=ref.first_execution_run_id
+        )
+    stream = ClientStream(
         client,
         backend_of_client(client),
         ref.id,
         batch_options(batch_interval, max_batch_size, max_retry_duration),
         wait_for_create=wait_for_create,
+        handle=handle,
     )
+    if ref.workflow_id and ref.first_execution_run_id:
+        stream._owner = (ref.workflow_id, ref.first_execution_run_id)
+    return stream
 
 
 async def create_stream(

temporalio/contrib/streams/_outside.py

@@ -41,6 +41,8 @@ def __init__(
         self._namespace = namespace
         # A client that has only a Workflow id learns the first run id on its first call.
         self._id = stream_id if isinstance(stream_id, str) else ""
+        # The owning Workflow of a Workflow's Stream, which its ref carries.
+        self._owner: tuple[str, str] | None = None
         self._find_id = None if isinstance(stream_id, str) else stream_id
         self._writer_factory = writer
         self._writer: StreamWriter | None = None
@@ -61,7 +63,13 @@ def ref(self) -> StreamRef:
             raise RuntimeError(
                 "The Stream id is not known yet. Call get_offset() first."
             )
-        return StreamRef(self._id)
+        if self._owner is None:
+            return StreamRef(self._id)
+        return StreamRef(
+            self._id,
+            workflow_id=self._owner[0],
+            first_execution_run_id=self._owner[1],
+        )
 
     async def _stream_id(self) -> str:
         if not self._id:

temporalio/contrib/streams/_reader.py

@@ -18,8 +18,11 @@
 
 # The service waits at most 60 s per read, so a longer wait is a loop of reads.
 _MAX_WAIT = timedelta(seconds=60)
-# How often a reader that ends with the run checks the run while it waits for the create.
+# How often a reader that ends with the run checks the run, while it waits for the create
+# or for new Records.
 _RUN_CHECK = timedelta(seconds=2)
+# The shortest idle time a read asks for, so a read never asks for none.
+_MIN_IDLE = timedelta(milliseconds=1)
 # How long one loop of the generated long poll may run before the facade looks again.
 _POLL_DEADLINE = timedelta(minutes=5)
 
@@ -51,7 +54,15 @@ async def read(
     # A reader that ends with the run waits for the create in short steps, so it sees the
     # run end soon.
     create_step = _RUN_CHECK if run_ended is not None else _MAX_WAIT
+    # The idle time counts from the last Record across reads, since a read that ends with the
+    # run polls in short steps.
+    idle_since = loop.time()
     while True:
+        idle_left = (
+            max(idle_timeout - timedelta(seconds=loop.time() - idle_since), _MIN_IDLE)
+            if idle_timeout is not None
+            else None
+        )
         create_left = (
             timedelta(seconds=max(create_until - loop.time(), 0))
             if create_until is not None
@@ -63,14 +74,18 @@ async def read(
             start=None if offset is not None else ReadStart.READ_START_BEGINNING,
             topics=names,
             topic_prefixes=list(topic_prefixes),
-            idle_timeout=idle_timeout,
+            idle_timeout=idle_left,
             create_wait=min(create_left, create_step) if create_left else None,
         )
         try:
             with _convert.raw_bodies(converter.failure_converter):
                 if wait_new_records:
+                    # A reader that ends with the run looks at the run between short polls.
                     answer = await service.read_records_until_records(
-                        request, deadline=_POLL_DEADLINE
+                        request,
+                        deadline=_RUN_CHECK
+                        if run_ended is not None
+                        else _POLL_DEADLINE,
                     )
                 else:
                     answer = await service.read_records(request)
@@ -90,6 +105,7 @@ async def read(
         if answer.next_offset is not None:
             offset = answer.next_offset
         if answer.records:
+            idle_since = loop.time()
             continue
         if answer.done or not wait_new_records:
             return

temporalio/contrib/streams/_stream.py

@@ -8,12 +8,42 @@
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
+import temporalio.converter
+
 if TYPE_CHECKING:
     from temporalio.contrib.streams._reader import Item
 
 T = TypeVar("T")
 
 
+class _StreamRefJson(
+    temporalio.converter.TransferTypeConverter["StreamRef", dict[str, str]]
+):
+    # The JSON form is the proto JSON of temporal.api.streamservice.v1.StreamRef, so every SDK
+    # reads a ref that another wrote. Snake case keys from older refs still decode.
+    transfer_type: type[dict[str, str]] | None = dict
+
+    def to_transfer_type(self, value: StreamRef) -> dict[str, str]:
+        json = {"id": value.id}
+        if value.workflow_id:
+            json["workflowId"] = value.workflow_id
+        if value.first_execution_run_id:
+            json["firstExecutionRunId"] = value.first_execution_run_id
+        return json
+
+    def from_transfer_type(
+        self, value: dict[str, str], type_hint: type[StreamRef]
+    ) -> StreamRef:
+        return StreamRef(
+            value["id"],
+            workflow_id=value.get("workflowId") or value.get("workflow_id") or None,
+            first_execution_run_id=value.get("firstExecutionRunId")
+            or value.get("first_execution_run_id")
+            or None,
+        )
+
+
+@temporalio.converter.transfer_type_convertible(_StreamRefJson)
 @dataclass(frozen=True)
 class StreamRef:
     """Names one Stream. Pass it to Workflows and Activities as an argument.
@@ -25,6 +55,14 @@ class StreamRef:
     id: str
     """The Stream id. It is opaque to callers."""
 
+    workflow_id: str | None = None
+    """For a Workflow's Stream, the id of the Workflow that owns it. None for a standalone
+    Stream."""
+
+    first_execution_run_id: str | None = None
+    """For a Workflow's Stream, the first run id of the owning run chain. None for a
+    standalone Stream."""
+
 
 class _Writable(abc.ABC):
     # A Stream kind that Topics publish to.

temporalio/contrib/streams/_workflow.py

@@ -80,7 +80,7 @@ class WorkflowStream(_Writable):
 
     def __init__(
         self,
-        stream_id: str,
+        ref: StreamRef,
         *,
         create: bool,
         retention: timedelta | None,
@@ -89,7 +89,8 @@ def __init__(
         """Create a Stream handle. Use :py:func:`workflow_stream` or
         :py:func:`create_workflow_stream` instead.
         """
-        self._id = stream_id
+        self._ref = ref
+        self._id = ref.id
         self._create = create
         self._retention = retention
         self._auto_close = auto_close
@@ -100,7 +101,14 @@ def __init__(
     @property
     def ref(self) -> StreamRef:
         """Reference to pass to Activities, child Workflows and clients."""
-        return StreamRef(self._id)
+        return self._ref
+
+    def _model_ref(self) -> models.StreamRef:
+        return models.StreamRef(
+            id=self._id,
+            workflow_id=self._ref.workflow_id,
+            first_execution_run_id=self._ref.first_execution_run_id,
+        )
 
     @property
     def auto_close(self) -> bool:
@@ -143,7 +151,7 @@ def _publish(self, topic: str, value: Any, force_flush: bool) -> None:
         self._ensure_created()
         self._used = True
         append_records(
-            stream=models.StreamRef(id=self._id),
+            stream=self._model_ref(),
             records=[models.AppendRecord(topic=topic, body=value)],
         )
 
@@ -158,9 +166,7 @@ def close(
         self._check("close")
         self._ensure_created()
         self._closed = True
-        close_stream(
-            stream=models.StreamRef(id=self._id), result=result, failure=failure
-        )
+        close_stream(stream=self._model_ref(), result=result, failure=failure)
 
     def __enter__(self) -> WorkflowStream:
         """Use the Stream in a ``with`` block, which closes it at the end."""
@@ -202,18 +208,18 @@ def workflow_stream(
     state = _state()
     if ref is None:
         info = temporalio.workflow.info()
-        stream_id = _ids.default_stream_id(
-            info.workflow_id, info.first_execution_run_id
+        ref = StreamRef(
+            _ids.default_stream_id(info.workflow_id, info.first_execution_run_id),
+            workflow_id=info.workflow_id,
+            first_execution_run_id=info.first_execution_run_id,
         )
         create, close = True, True if auto_close is None else auto_close
     else:
-        stream_id = ref.id
         create, close = False, bool(auto_close)
+    stream_id = ref.id
     stream = state.streams.get(stream_id)
     if stream is None:
-        stream = WorkflowStream(
-            stream_id, create=create, retention=None, auto_close=close
-        )
+        stream = WorkflowStream(ref, create=create, retention=None, auto_close=close)
         state.streams[stream_id] = stream
     elif auto_close is not None:
         stream._auto_close = auto_close
@@ -237,9 +243,12 @@ def create_workflow_stream(
     stream_id = _ids.scoped_stream_id(
         info.workflow_id, info.first_execution_run_id, name
     )
-    stream = WorkflowStream(
-        stream_id, create=True, retention=retention, auto_close=False
+    ref = StreamRef(
+        stream_id,
+        workflow_id=info.workflow_id,
+        first_execution_run_id=info.first_execution_run_id,
     )
+    stream = WorkflowStream(ref, create=True, retention=retention, auto_close=False)
     stream._check("create")
     stream._ensure_created()
     _state().streams[stream_id] = stream