Add Stream.replay_message for id-preserving message replaying (#185)

This allows Stream to preserve the message id. I made the old thing,
replay_message_events, private. I might move it into Stream also in a
follow-up but didn't want to make this one too noisy.

Author: msullivan

examples/temporal-direct/main.py

@@ -153,11 +153,10 @@ async def loop(
 
             # 2. Wrap the complete message in a synthetic stream so we can
             #    drive the rest of the loop with ToolRunner — same shape as
-            #    the default loop. ``replay_message_events`` is the framework
-            #    helper that decomposes a complete ``Message`` back into the
-            #    events a streaming adapter would have produced.
+            #    the default loop. ``Stream.replay_message`` replays events
+            #    from llm_msg.
             async with (
-                ai.Stream(ai.events.replay_message_events(llm_msg)) as stream,
+                ai.Stream.replay_message(llm_msg) as stream,
                 ai.ToolRunner() as tr,
             ):
                 async for event in ai.util.merge(stream, tr.events()):

src/ai/models/core/api.py

@@ -135,6 +135,53 @@ def __init__(
         # rather than look like a normal end of turn.
         self._ended = False
 
+    @classmethod
+    def replay_message(
+        cls,
+        message: types.messages.Message,
+        *,
+        output_type: type[StreamOutputT] | None = None,
+    ) -> Stream[StreamOutputT]:
+        """Synthesize stream events for ``msg``.
+
+        Use when you have a complete ``Message`` from a non-streaming source —
+        e.g., the result of a Temporal activity, a cached LLM response, or an
+        offline test fixture — and want to feed it through code that consumes
+        an async event stream (``ai.Stream``, ``ai.ToolRunner``, custom loops
+        that mirror the default loop's shape, etc.)::
+
+            async with ai.Stream.replay_message(msg) as stream:
+                async with ai.ToolRunner() as tr:
+                    async for event in ai.util.merge(stream, tr.events()):
+                        ...
+
+        Each part is emitted as the start/delta/end triple a streaming adapter
+        would have produced, in part order, bracketed by ``StreamStart`` and
+        ``StreamEnd``. The full body of text/reasoning/tool-args is sent as a
+        single delta — the granularity of the model's original chunking is
+        not recoverable from a complete message.
+
+        Each part's ``provider_metadata`` (and the message's) rides on its end
+        event, mirroring the real adapters, so a rebuilt turn keeps it --
+        reasoning signatures included, which must survive to replay the turn
+        back to the provider.
+
+        Parts with no model-layer event analog — ``ToolResultPart``,
+        ``HookPart`` — are skipped silently; they are agent-layer concerns
+        and never appear on the model stream.
+
+        ``stream.message`` keeps ``message``'s id; the parts are rebuilt
+        from the stream.
+        """
+        seed = types.messages.Message(
+            id=message.id, role=message.role, parts=[]
+        )
+        return cls(
+            types.events._replay_message_events(message),
+            seed_message=seed,
+            output_type=output_type,
+        )
+
     async def aclose(self) -> None:
         await self._gen.aclose()
 

src/ai/types/events.py

@@ -183,47 +183,24 @@ class HookResolution(BaseEvent):
 ]
 
 
-async def replay_message_events(
+async def _replay_message_events(
     msg: messages.Message,
 ) -> AsyncGenerator[Event]:
-    """Synthesize stream events for ``msg``.
-
-    Use when you have a complete ``Message`` from a non-streaming source —
-    e.g., the result of a Temporal activity, a cached LLM response, or an
-    offline test fixture — and want to feed it through code that consumes
-    an async event stream (``ai.Stream``, ``ai.ToolRunner``, custom loops
-    that mirror the default loop's shape, etc.)::
-
-        async with ai.Stream(ai.events.replay_message_events(msg)) as stream:
-            async with ai.ToolRunner() as tr:
-                async for event in ai.util.merge(stream, tr.events()):
-                    ...
-
-    Each part is emitted as the start/delta/end triple a streaming adapter
-    would have produced, in part order, bracketed by ``StreamStart`` and
-    ``StreamEnd``. The full body of text/reasoning/tool-args is sent as a
-    single delta — the granularity of the model's original chunking is
-    not recoverable from a complete message.
-
-    Parts with no model-layer event analog — ``ToolResultPart``,
-    ``HookPart`` — are skipped silently; they are agent-layer concerns
-    and never appear on the model stream.
-    """
+    """Synthesize stream events for ``msg``."""
+    # See Stream.replay_message
     yield StreamStart()
     for part in msg.parts:
         if isinstance(part, messages.TextPart):
             yield TextStart(block_id=part.id)
             if part.text:
                 yield TextDelta(block_id=part.id, chunk=part.text)
-            yield TextEnd(block_id=part.id)
+            yield TextEnd(
+                block_id=part.id, provider_metadata=part.provider_metadata
+            )
         elif isinstance(part, messages.ReasoningPart):
             yield ReasoningStart(block_id=part.id)
             if part.text:
                 yield ReasoningDelta(block_id=part.id, chunk=part.text)
-            # Carry the signature (and any other reasoning metadata) on the
-            # end event, mirroring how the real adapters emit it -- otherwise
-            # a replayed-then-rebuilt turn loses its signature and can't be
-            # replayed to the provider.
             yield ReasoningEnd(
                 block_id=part.id,
                 provider_metadata=part.provider_metadata,
@@ -238,7 +215,11 @@ async def replay_message_events(
                     tool_call_id=part.tool_call_id,
                     chunk=part.tool_args,
                 )
-            yield ToolEnd(tool_call_id=part.tool_call_id, tool_call=part)
+            yield ToolEnd(
+                tool_call_id=part.tool_call_id,
+                tool_call=part,
+                provider_metadata=part.provider_metadata,
+            )
         elif isinstance(part, messages.BuiltinToolCallPart):
             yield BuiltinToolStart(
                 tool_call_id=part.tool_call_id,
@@ -249,7 +230,11 @@ async def replay_message_events(
                     tool_call_id=part.tool_call_id,
                     chunk=part.tool_args,
                 )
-            yield BuiltinToolEnd(tool_call_id=part.tool_call_id, tool_call=part)
+            yield BuiltinToolEnd(
+                tool_call_id=part.tool_call_id,
+                tool_call=part,
+                provider_metadata=part.provider_metadata,
+            )
         elif isinstance(part, messages.BuiltinToolReturnPart):
             yield BuiltinToolResult(tool_call_id=part.tool_call_id, result=part)
         elif isinstance(part, messages.FilePart):
@@ -258,8 +243,9 @@ async def replay_message_events(
                 data=part.data,
                 media_type=part.media_type,
                 filename=part.filename,
+                provider_metadata=part.provider_metadata,
             )
-    yield StreamEnd()
+    yield StreamEnd(provider_metadata=msg.provider_metadata)
 
 
 # ---------------------------------------------------------------------------

tests/models/core/test_api.py

@@ -464,6 +464,43 @@ async def _spy_stream(
         assert not any(isinstance(e, events_.TextDelta) for e in events)
 
 
+async def test_replay_message_preserves_id() -> None:
+    """``Stream.replay_message`` keeps the message id and rebuilds parts."""
+    msg = messages_.Message(
+        role="assistant",
+        parts=[
+            messages_.TextPart(id="t1", text="calling tools"),
+            messages_.ToolCallPart(
+                tool_call_id="tc-1",
+                tool_name="weather",
+                tool_args='{"city":"SF"}',
+            ),
+        ],
+    )
+
+    async with ai.Stream.replay_message(msg) as stream:
+        events: list[events_.Event] = [event async for event in stream]
+
+    # The rebuilt message keeps the original id (a fresh Message, not the
+    # original object) and the parts are reconstructed from the events --
+    # no duplication, same part ids.
+    assert stream.message is not msg
+    assert stream.message.id == msg.id
+    assert len(stream.message.parts) == 2
+    assert stream.text == "calling tools"
+    assert [tc.tool_call_id for tc in stream.tool_calls] == ["tc-1"]
+
+    # The full event set is re-emitted (and is visible to consumers --
+    # nothing is flagged ``replay``), enough to drive a tool runner.
+    assert events
+    assert not any(e.replay for e in events)
+    tool_ends = [e for e in events if isinstance(e, events_.ToolEnd)]
+    assert len(tool_ends) == 1
+    assert tool_ends[0].tool_call.tool_call_id == "tc-1"
+    assert any(isinstance(e, events_.TextDelta) for e in events)
+    assert any(isinstance(e, events_.ToolStart) for e in events)
+
+
 def test_tool_end_replay_flag_excluded_from_json() -> None:
     """The replay flag is internal — it must not appear in serialized output."""
     ev = events_.ToolEnd(

tests/types/test_events.py

@@ -23,7 +23,7 @@ async def test_reasoning_signature_survives_replay(self) -> None:
         )
 
         async with models.Stream(
-            events.replay_message_events(original)
+            events._replay_message_events(original)
         ) as stream:
             async for _ in stream:
                 pass
@@ -54,10 +54,44 @@ async def test_reasoning_signature_on_end_event(self) -> None:
 
         reasoning_ends = [
             e
-            async for e in events.replay_message_events(msg)
+            async for e in events._replay_message_events(msg)
             if isinstance(e, events.ReasoningEnd)
         ]
         assert len(reasoning_ends) == 1
         assert reasoning_ends[0].provider_metadata == {
             "anthropic": {"signature": "sig"}
         }
+
+    async def test_provider_metadata_survives_replay(self) -> None:
+        """provider_metadata on every part, and the message itself, round-
+        trips through the aggregator -- not just reasoning signatures."""
+        original = messages.Message(
+            role="assistant",
+            parts=[
+                messages.TextPart(text="hi", provider_metadata={"p": {"t": 1}}),
+                messages.ToolCallPart(
+                    tool_call_id="tc-1",
+                    tool_name="weather",
+                    tool_args="{}",
+                    provider_metadata={"p": {"tc": 2}},
+                ),
+            ],
+            provider_metadata={"p": {"msg": 3}},
+        )
+
+        async with models.Stream(
+            events._replay_message_events(original)
+        ) as stream:
+            async for _ in stream:
+                pass
+
+        rebuilt = stream.message
+        assert rebuilt.provider_metadata == {"p": {"msg": 3}}
+        text = next(
+            p for p in rebuilt.parts if isinstance(p, messages.TextPart)
+        )
+        assert text.provider_metadata == {"p": {"t": 1}}
+        tool = next(
+            p for p in rebuilt.parts if isinstance(p, messages.ToolCallPart)
+        )
+        assert tool.provider_metadata == {"p": {"tc": 2}}