Produce ChatModelStream objects from MessagesTransformer

Replace the passthrough (chunk, metadata) tuple projection with one
that yields a ChatModelStream per LLM call, routed by run_id. Handle
both v2 protocol-event payloads (message-start/chunk/message-finish)
and whole AIMessage payloads from on_chain_end (replayed via
message_to_events). Wire _bind_pump from GraphRunStream so nested
sync streams share the caller-driven pump.

libs/langgraph/langgraph/stream/run_stream.py

@@ -60,6 +60,11 @@ class GraphRunStream:
 
         Sync iteration is caller-driven, so a cursor that catches up to
         the buffer's tail needs a way to ask the graph for more events.
+
+        Also calls `_bind_pump` on any transformer that exposes it, so
+        that transformers producing ChatModelStream objects (e.g.
+        MessagesTransformer) can wire the pull callback on each stream as
+        it's created.
         """
         mux._events._request_more = self._pump_next
         for value in mux.extensions.values():
@@ -67,6 +72,9 @@ class GraphRunStream:
                 value._request_more = self._pump_next
             elif isinstance(value, StreamChannel):
                 value._log._request_more = self._pump_next
+        for transformer in mux._transformers:
+            if hasattr(transformer, "_bind_pump"):
+                transformer._bind_pump(self._pump_next)
 
     def _pump_next(self) -> bool:
         """Pull one event from the graph and push it through the mux.

libs/langgraph/langgraph/stream/transformers.py

@@ -1,10 +1,21 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
+
+from langchain_core.language_models._compat_bridge import message_to_events
+from langchain_core.language_models.chat_model_stream import (
+    AsyncChatModelStream,
+    ChatModelStream,
+)
+from langchain_core.messages import AIMessageChunk, BaseMessage
+from langchain_protocol.protocol import MessagesData
 
 from langgraph.stream._event_log import EventLog
 from langgraph.stream._types import ProtocolEvent, StreamTransformer
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
 
 class ValuesTransformer(StreamTransformer):
     """Capture values events as an iterable of state snapshots.
@@ -51,34 +62,142 @@ class ValuesTransformer(StreamTransformer):
 
 
 class MessagesTransformer(StreamTransformer):
-    """Pass through raw (chunk, metadata) tuples from messages events.
+    """Capture messages events as ChatModelStream objects.
 
-    This is the same shape as today's `stream_mode="messages"` output.
-    A follow-on PR will replace this with a richer transformer that
-    produces ChatModelStream objects using the protocol handler.
+    The messages projection yields one `ChatModelStream` (or
+    `AsyncChatModelStream`) per LLM call. Consumers iterate
+    `run.messages` to get stream handles, then use each handle's typed
+    projections (`.text`, `.reasoning`, `.tool_calls`, `.usage`,
+    `.output`) for per-message content.
 
-    Only root-namespace messages events are captured; tokens emitted
-    from subgraphs are dropped from the `messages` projection. Consumers
-    that need subgraph tokens should iterate the raw event stream or
-    register a custom transformer.
+    Two input shapes are handled (via `params["data"] = (payload,
+    metadata)` from `StreamMessagesHandler`):
 
-    Native transformer — projection keys are exposed as direct
-    attributes on the run stream (e.g. `run.messages`).
+    1. Protocol event (dict with `"event"` key) — emitted by
+       `stream_v2()` / `astream_v2()` via the `on_stream_event`
+       callback. Routed to an existing `ChatModelStream` by
+       `metadata["run_id"]`. A `message-start` event creates a new
+       stream; `message-finish` closes it.
+    2. Whole `AIMessage` — emitted from `on_chain_end` when a node
+       returns a finalized message. Replayed as a synthetic protocol
+       event lifecycle via `message_to_events`, then the
+       already-complete stream is pushed to the log.
+
+    V1 `AIMessageChunk` tuples (from `on_llm_new_token`) are not
+    streamed into this projection: chat models that want to populate
+    `run.messages` with content-block streaming must use
+    `stream_v2()` / `astream_v2()`. Models called via the legacy
+    `stream()` method still surface their final `AIMessage` via
+    `on_chain_end` when a node returns it as state.
+
+    Only root-namespace events are captured; tokens from subgraphs are
+    dropped. Consumers that need subgraph tokens should iterate the raw
+    event stream or register a custom transformer.
+
+    Native transformer — the `messages` projection is exposed as a
+    direct attribute on the run stream.
     """
 
     _native = True
 
     def __init__(self) -> None:
-        self._log: EventLog[tuple[Any, dict[str, Any]]] = EventLog()
+        self._log: EventLog[ChatModelStream] = EventLog()
+        # Correlate protocol events back to a ChatModelStream by run_id
+        # (attached to the event's metadata by StreamMessagesHandler).
+        self._by_run: dict[str, ChatModelStream] = {}
+        self._pump_fn: Callable[[], bool] | None = None
 
     def init(self) -> dict[str, Any]:
         return {"messages": self._log}
 
+    def _bind_pump(self, fn: Callable[[], bool]) -> None:
+        """Wire the sync pull callback. Called by GraphRunStream._wire_request_more."""
+        self._pump_fn = fn
+
+    def _make_stream(
+        self,
+        *,
+        namespace: list[str],
+        node: str | None,
+        message_id: str | None,
+    ) -> ChatModelStream:
+        """Create a ChatModelStream (sync) or AsyncChatModelStream (async)."""
+        if self._pump_fn is not None:
+            stream: ChatModelStream = ChatModelStream(
+                namespace=namespace,
+                node=node,
+                message_id=message_id,
+            )
+            stream.set_request_more(self._pump_fn)
+        else:
+            stream = AsyncChatModelStream(
+                namespace=namespace,
+                node=node,
+                message_id=message_id,
+            )
+        return stream
+
     def process(self, event: ProtocolEvent) -> bool:
         if event["method"] != "messages":
             return True
         params = event["params"]
         if params["namespace"]:
             return True
-        self._log.push(params["data"])
+
+        payload, metadata = params["data"]
+        node: str | None = metadata.get("langgraph_node")
+        run_id = str(metadata.get("run_id", "")) if metadata else ""
+
+        if isinstance(payload, dict) and "event" in payload:
+            self._route_protocol_event(
+                cast("MessagesData", payload), run_id=run_id, node=node
+            )
+        elif isinstance(payload, BaseMessage) and not isinstance(
+            payload, AIMessageChunk
+        ):
+            self._route_whole_message(payload, node=node)
+        # Legacy AIMessageChunk tuples (from on_llm_new_token) are ignored;
+        # v1 streaming callers must switch to stream_v2() to populate this
+        # projection.
+
         return True
+
+    def _route_protocol_event(
+        self,
+        event: MessagesData,
+        *,
+        run_id: str,
+        node: str | None,
+    ) -> None:
+        event_type = event.get("event")
+        if event_type == "message-start":
+            message_id = event.get("message_id")
+            stream = self._make_stream(
+                namespace=[],
+                node=node,
+                message_id=str(message_id) if message_id is not None else None,
+            )
+            self._by_run[run_id] = stream
+            self._log.push(stream)
+            stream.dispatch(event)
+        elif run_id in self._by_run:
+            stream = self._by_run[run_id]
+            stream.dispatch(event)
+            if event_type == "message-finish":
+                del self._by_run[run_id]
+
+    def _route_whole_message(self, message: BaseMessage, *, node: str | None) -> None:
+        stream = self._make_stream(namespace=[], node=node, message_id=message.id)
+        for evt in message_to_events(message, message_id=message.id):
+            stream.dispatch(evt)
+        self._log.push(stream)
+
+    def finalize(self) -> None:
+        """Clear any routing state — streams close themselves via `message-finish`."""
+        self._by_run.clear()
+
+    def fail(self, err: BaseException) -> None:
+        """Propagate run error to any streams still open when the graph fails."""
+        for stream in list(self._by_run.values()):
+            stream.fail(err)
+        self._by_run.clear()

libs/langgraph/tests/test_streaming_handler.py

@@ -807,22 +807,40 @@ class TestValuesTransformer:
 
 class TestMessagesTransformer:
     def test_captures_root_messages(self) -> None:
+        """Protocol-event lifecycle produces a ChatModelStream in the log."""
         t = MessagesTransformer()
         t.init()
         t._log._bind(is_async=False)
+        t._bind_pump(lambda: False)
 
-        t.process(_event("messages", ("chunk", {"meta": True})))
+        meta = {"langgraph_node": "llm", "run_id": "run-1"}
+        for evt in (
+            {"event": "message-start", "role": "ai", "message_id": "run-1"},
+            {"event": "message-finish", "reason": "stop"},
+        ):
+            t.process(_event("messages", (evt, meta)))
         t._log.close()
         items = list(t._log)
         assert len(items) == 1
-        assert items[0] == ("chunk", {"meta": True})
+        # Items in the messages log are ChatModelStream objects, not raw
+        # tuples — the content-block-centric projection.
+        assert hasattr(items[0], "dispatch")
+        assert items[0].message_id == "run-1"
 
     def test_ignores_non_root_namespace(self) -> None:
         t = MessagesTransformer()
         t.init()
         t._log._bind(is_async=False)
+        t._bind_pump(lambda: False)
 
-        t.process(_event("messages", ("chunk", {}), namespace=["sub"]))
+        meta = {"langgraph_node": "llm", "run_id": "run-1"}
+        t.process(
+            _event(
+                "messages",
+                ({"event": "message-start", "message_id": "run-1"}, meta),
+                namespace=["sub"],
+            )
+        )
         t._log.close()
         assert list(t._log) == []