From 0f2f66fc8f0565a09ad7a562536b4eb77b259743 Mon Sep 17 00:00:00 2001
From: Nick Hollon <nick.hollon@langchain.dev>
Date: Fri, 17 Apr 2026 16:22:50 -0400
Subject: [PATCH] refactor(langgraph): scope on_stream_event to
 StreamMessagesHandlerV2

Moves `on_stream_event` off the v1 `StreamMessagesHandler` base and
onto the v2 subclass. Content-block events are a v2-only concept, so
forwarding them only when the v2 handler is attached keeps the
messages channel's shape predictable for v1 callers: plain
`graph.stream(stream_mode="messages")` now ignores `on_stream_event`
entirely, even if a node explicitly calls `model.stream_v2()` on a
v1-flagged run. Dedupe of the returned AIMessage still works in that
case via `_find_and_emit_messages` / `on_chain_end`.

Also makes the v2 handler's `on_llm_new_token` override an explicit
pass-through with a comment rather than `return None`, so it reads as
an intentional no-op instead of a missing return value.
---
 libs/langgraph/langgraph/pregel/_messages.py | 84 +++++++++++---------
 1 file changed, 45 insertions(+), 39 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/_messages.py b/libs/langgraph/langgraph/pregel/_messages.py
index 56eb829e3..1c196add6 100644
--- a/libs/langgraph/langgraph/pregel/_messages.py
+++ b/libs/langgraph/langgraph/pregel/_messages.py
@@ -155,37 +155,6 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
                     stream_metadata["tags"] = filtered_tags
             self.metadata[run_id] = (ns, stream_metadata)
 
-    def on_stream_event(
-        self,
-        event: dict[str, Any],
-        *,
-        run_id: UUID,
-        parent_run_id: UUID | None = None,
-        tags: list[str] | None = None,
-        **kwargs: Any,
-    ) -> Any:
-        """Forward a protocol event from `stream_v2` as a messages stream part.
-
-        Fires once per `MessagesData` event (`message-start`, per-block
-        `content-block-*`, `message-finish`). The transformer layer
-        correlates events back to a single `ChatModelStream` via
-        `metadata["run_id"]` — attached here so the v1
-        `stream_mode="messages"` output (which emits
-        `(AIMessageChunk, metadata)` via `on_llm_new_token`) keeps its
-        original metadata shape.
-        """
-        if meta := self.metadata.get(run_id):
-            # Record message_id on message-start so on_chain_end's
-            # dedupe skips the finalized AIMessage the node returns
-            # (otherwise the messages projection double-counts: once
-            # from streaming, once from the chain output).
-            if event.get("event") == "message-start":
-                msg_id = event.get("message_id")
-                if msg_id:
-                    self.seen.add(msg_id)
-            v2_meta = {**meta[1], "run_id": str(run_id)}
-            self.stream((meta[0], "messages", (event, v2_meta)))
-
     def on_llm_new_token(
         self,
         token: str,
@@ -320,13 +289,50 @@ class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler
         tags: list[str] | None = None,
         **kwargs: Any,
     ) -> Any:
-        """Suppress v1 chunk emission on the messages channel.
+        """Intentional no-op — v1 chunks are not used on v2-flagged runs.
 
-        The v2 marker already steers `invoke` to the event generator,
-        so `on_llm_new_token` should not fire under normal routing.
-        This override guards against any caller (e.g. a node that
-        calls `model.stream()` directly, which still fires the v1
-        callback) leaking AIMessageChunks onto a v2-flagged messages
-        stream.
+        The v2 marker already steers `invoke` to the event generator, so
+        `on_llm_new_token` should not fire under normal routing. This
+        override stays a pass-through (no call to `super()`) to make
+        the intent explicit and to guard against any caller (e.g. a
+        node that calls `model.stream()` directly, which still fires
+        the v1 callback) leaking AIMessageChunks onto a v2-flagged
+        messages stream.
         """
-        return None
+        # Intentionally empty: v2 handler does not forward v1 chunks.
+
+    def on_stream_event(
+        self,
+        event: dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Forward a protocol event from `stream_v2` as a messages stream part.
+
+        Fires once per `MessagesData` event (`message-start`, per-block
+        `content-block-*`, `message-finish`). The transformer layer
+        correlates events back to a single `ChatModelStream` via
+        `metadata["run_id"]` — attached here so the v1
+        `stream_mode="messages"` output (which emits
+        `(AIMessageChunk, metadata)` via `on_llm_new_token`) keeps its
+        original metadata shape.
+
+        Lives on the v2 handler rather than the v1 base: content-block
+        events are a v2-only concept, and forwarding them only when the
+        v2 handler is attached keeps the message channel's shape
+        predictable for v1 callers.
+        """
+        if meta := self.metadata.get(run_id):
+            # Record message_id on message-start so on_chain_end's
+            # dedupe skips the finalized AIMessage the node returns
+            # (otherwise the messages projection double-counts: once
+            # from streaming, once from the chain output).
+            if event.get("event") == "message-start":
+                msg_id = event.get("message_id")
+                if msg_id:
+                    self.seen.add(msg_id)
+            v2_meta = {**meta[1], "run_id": str(run_id)}
+            self.stream((meta[0], "messages", (event, v2_meta)))