Support overriding randomness source for ids

Make message id generation use Random and add use_random and
use_random_async @contextmanagers for oerriding it.

The main goal here is to support durable execution frameworks where we
need to be able to provide a deterministic randomness source.

Author: msullivan

examples/temporal-direct/main.py

@@ -29,7 +29,9 @@
 import dataclasses
 import datetime
 import json
+import os
 import sys
+import threading
 import uuid
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -42,9 +44,14 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
 
+# ai itself is safe to import in a workflow... but its transitive deps
+# httpx and modelsdotdev are *not*
 with temporalio.workflow.unsafe.imports_passed_through():
-    import ai
+    import httpx  # noqa: F401
+    import modelsdotdev  # noqa: F401
+    import pydantic_core  # noqa: F401
 
+import ai
 
 MODEL_ID = "gateway:anthropic/claude-sonnet-4.6"
 
@@ -94,6 +101,28 @@ async def get_population_activity(city: str) -> int:
 }
 
 
+# ── Stream-event sink ────────────────────────────────────────────
+#
+# The durable result of the workflow is the final list of Messages.
+# Live token-by-token streaming is a separate, ephemeral concern: the
+# LLM activity emits every model-stream event to an out-of-band sink
+# as it drains the stream. The sink is a plain file named in
+# DURABILITY_EVENT_LOG (/dev/stdout is a good choice for observing);
+# each line is one event serialized as JSON
+# (``event.model_dump_json()``), so a consumer can rebuild the event.
+
+_EVENT_LOG_ENV = "DURABILITY_EVENT_LOG"
+_event_lock = threading.Lock()
+
+
+def _emit_event(event: ai.events.Event) -> None:
+    path = os.environ.get(_EVENT_LOG_ENV)
+    if not path:
+        return
+    with _event_lock, open(path, "a") as f:
+        f.write(event.model_dump_json() + "\n")
+
+
 @dataclasses.dataclass
 class LLMParams:
     model_id: str
@@ -114,8 +143,8 @@ async def llm_call_activity(params: LLMParams) -> LLMResult:
     tools = [ai.Tool.model_validate(t) for t in params.tool_schemas]
 
     async with ai.stream(model, messages, tools=tools) as s:
-        async for _event in s:
-            pass
+        async for event in s:
+            _emit_event(event)
         if s.message is None:
             raise RuntimeError("LLM stream ended without a final message")
         return LLMResult(message=s.message.model_dump())
@@ -206,8 +235,15 @@ async def _call() -> ai.events.ToolCallResult:
 
 @temporalio.workflow.defn
 class WeatherWorkflow:
+    def __init__(self) -> None:
+        self._messages: list[ai.messages.Message] = []
+
+    # Draw message/part ids from the workflow's deterministic RNG so they
+    # are stable across replay. ``workflow.random`` is passed as a factory
+    # (it's only valid inside the workflow) and resolved on each call.
     @temporalio.workflow.run
-    async def run(self, user_query: str) -> str:
+    @ai.messages.use_random_async(temporalio.workflow.random)
+    async def run(self, user_query: str) -> list[dict[str, Any]]:
         model = ai.get_model(MODEL_ID)
         messages: list[ai.messages.Message] = [
             ai.system_message(
@@ -216,12 +252,16 @@ async def run(self, user_query: str) -> str:
             ai.user_message(user_query),
         ]
 
-        final_text = ""
         async with weather_agent.run(model, messages) as stream:
-            async for event in stream:
-                if isinstance(event, ai.events.StreamEnd):
-                    final_text = event.message.text
-        return final_text
+            async for _event in stream:
+                pass
+        self._messages = stream.messages
+        return [m.model_dump() for m in self._messages]
+
+    @temporalio.workflow.query
+    def messages(self) -> list[dict[str, Any]]:
+        """The final messages, as rebuilt by this replay."""
+        return [m.model_dump() for m in self._messages]
 
 
 # ── Entry point ──────────────────────────────────────────────────
@@ -246,13 +286,14 @@ async def main(user_query: str) -> None:
         print(f"Workflow: {workflow_id}")
         print(f"Query:    {user_query}\n")
 
-        result = await client.execute_workflow(
+        messages = await client.execute_workflow(
             WeatherWorkflow.run,
             user_query,
             id=workflow_id,
             task_queue=TASK_QUEUE,
         )
-        print(result)
+        final = ai.messages.Message.model_validate(messages[-1])
+        print(f"\nResult:   {final.text}")
 
 
 if __name__ == "__main__":

examples/temporal-direct/test_durability.py

@@ -48,11 +48,14 @@
 from typing import Any
 
 import main as ex
+import pydantic
 import temporalio.client
 import temporalio.testing
 import temporalio.worker
 from _durability_worker import LOGGED_ACTIVITIES  # noqa: PLC2701
 
+import ai
+
 # ── Helpers ──────────────────────────────────────────────────────
 
 
@@ -71,6 +74,22 @@ def read_activity_log(log_file: pathlib.Path) -> Counter[str]:
     return Counter(log_file.read_text().splitlines())
 
 
+_EVENT_ADAPTER: pydantic.TypeAdapter[ai.events.Event] = pydantic.TypeAdapter(
+    ai.events.DiscriminatedEvent
+)
+
+
+def read_streamed_events(event_log: pathlib.Path) -> list[ai.events.Event]:
+    """Deserialize the events the LLM activity wrote to the streaming sink."""
+    if not event_log.exists():
+        return []
+    return [
+        _EVENT_ADAPTER.validate_json(line)
+        for line in event_log.read_text().splitlines()
+        if line
+    ]
+
+
 QUERY = "What's the weather and population of New York and Los Angeles?"
 
 
@@ -92,11 +111,12 @@ async def test_happy_path(
             task_queue=ex.TASK_QUEUE,
         )
 
-    assert result, "expected non-empty text"
+    assert result, "expected a non-empty conversation"
+    final_text = ai.messages.Message.model_validate(result[-1]).text
     assert (
-        "8,336,817" in result or "8336817" in result
-    ), f"expected NYC population in result, got: {result!r}"
-    print(f"  ✓ workflow {wid} produced {len(result)} chars")
+        "8,336,817" in final_text or "8336817" in final_text
+    ), f"expected NYC population in result, got: {final_text!r}"
+    print(f"  ✓ workflow {wid} produced {len(result)} messages")
     print(f"  ✓ activity calls: {dict(read_activity_log(log_file))}")
     return wid
 
@@ -118,6 +138,80 @@ async def test_replay_determinism(
     print(f"  ✓ replay clean for {workflow_id} ({len(history.events)} events)")
 
 
+# ── Test: workflow-minted ids are stable across replay ───────────
+
+
+async def test_workflow_minted_ids_are_deterministic(
+    client: temporalio.client.Client, workflow_id: str
+) -> None:
+    print("\n── test_workflow_minted_ids_are_deterministic ─────")
+
+    def ids(messages: list[dict[str, Any]]) -> list[str]:
+        return [
+            id_
+            for m in messages
+            for id_ in (m["id"], *(p["id"] for p in m["parts"]))
+        ]
+
+    handle = client.get_workflow_handle(workflow_id)
+    original = ids(await handle.result())
+
+    # Querying the closed workflow forces a worker to replay the history
+    # and re-run ``run`` to answer -- so these messages are rebuilt by the
+    # replay, not the original execution.
+    async with make_worker(client):
+        replayed = ids(await handle.query(ex.WeatherWorkflow.messages))
+
+    # Any id generated in workflow code (system/user/tool messages and the
+    # tool-result parts inside them) would be re-minted on replay and
+    # diverge here. Ids sourced from a cached activity result survive
+    # identically regardless.
+    assert original == replayed, (
+        "workflow re-minted ids on replay -- nondeterministic id "
+        "generation in workflow code:\n"
+        f"  original: {original}\n"
+        f"  replayed: {replayed}"
+    )
+    print(f"  ✓ all {len(original)} message/part ids stable across replay")
+
+
+# ── Test: streamed ids match the final returned messages ─────────
+
+
+async def test_stream_ids_match_final_messages(
+    client: temporalio.client.Client, event_log: pathlib.Path
+) -> None:
+    print("\n── test_stream_ids_match_final_messages ───────────")
+    event_log.write_text("")
+
+    async with make_worker(client):
+        wid = f"ids-{uuid.uuid4().hex[:8]}"
+        messages = await client.execute_workflow(
+            ex.WeatherWorkflow.run,
+            QUERY,
+            id=wid,
+            task_queue=ex.TASK_QUEUE,
+        )
+
+    streamed = read_streamed_events(event_log)
+    streamed_ids = {e.message.id for e in streamed}
+    final_ids = {m["id"] for m in messages}
+    assert streamed, "no events reached the streaming sink"
+
+    # Every assistant turn streamed to the sink must be identifiable in
+    # the durable result by the same message id. Today the workflow's
+    # Stream re-mints the id when it reassembles the message, so the id a
+    # client saw streaming live is absent from the final messages.
+    missing = streamed_ids - final_ids
+    assert not missing, (
+        f"message ids streamed to the sink are absent from the final "
+        f"returned messages: {sorted(missing)}\n"
+        f"  streamed: {sorted(streamed_ids)}\n"
+        f"  final:    {sorted(final_ids)}"
+    )
+    print(f"  ✓ all {len(streamed_ids)} streamed id(s) present in final")
+
+
 # ── Test 3: activity caching across a worker restart ─────────────
 
 
@@ -234,11 +328,15 @@ def _unraisablehook(unraisable: Any) -> None:
 
 
 async def main() -> None:
-    log_file = pathlib.Path(tempfile.mkdtemp()) / "activity_log.txt"
+    tmp = pathlib.Path(tempfile.mkdtemp())
+    log_file = tmp / "activity_log.txt"
+    event_log = tmp / "event_log.txt"
     # In-process workers (worker2 below, plus the happy-path worker)
     # share this module's ``LOGGED_ACTIVITIES``, which look at this env
     # var to decide whether to log. Set it for the whole run.
     os.environ["DURABILITY_ACTIVITY_LOG"] = str(log_file)
+    # The LLM activity emits every stream event to this sink.
+    os.environ[ex._EVENT_LOG_ENV] = str(event_log)
 
     print("Starting embedded Temporal dev server...")
     async with (
@@ -248,6 +346,8 @@ async def main() -> None:
 
         wid = await test_happy_path(client, log_file)
         await test_replay_determinism(client, wid)
+        await test_workflow_minted_ids_are_deterministic(client, wid)
+        await test_stream_ids_match_final_messages(client, event_log)
         await test_activity_caching(env, client, log_file)
 
     print("\nAll durability checks passed.")