Fix mypy errors, tighten init() return type, move import to top level

libs/langgraph/langgraph/stream/_mux.py

@@ -117,9 +117,7 @@ class StreamMux:
                 # Ensure the channel's log matches the mux's mode.
                 if value._is_async != self._is_async:
                     value._is_async = self._is_async
-                    value._log = (
-                        AsyncEventLog() if self._is_async else EventLog()
-                    )
+                    value._log = AsyncEventLog() if self._is_async else EventLog()
                 self._channels.append(value)
                 channel_name = value.name
 

libs/langgraph/langgraph/stream/_types.py

@@ -45,7 +45,7 @@ class StreamTransformer(ABC):
     """
 
     @abstractmethod
-    def init(self) -> Any:
+    def init(self) -> dict[str, Any]:
         """Return the projection dict.
 
         Keys become entries in `run.extensions`. If the transformer has

libs/langgraph/langgraph/stream/run_stream.py

@@ -4,7 +4,8 @@ import asyncio
 from collections.abc import AsyncIterator, Iterator
 from typing import Any
 
-from langgraph.stream._event_log import EventLog
+from langgraph.stream._convert import convert_to_protocol_event
+from langgraph.stream._event_log import AsyncEventLog, EventLog
 from langgraph.stream._mux import StreamMux
 from langgraph.stream._types import ProtocolEvent
 from langgraph.stream.stream_channel import StreamChannel
@@ -43,18 +44,14 @@ class GraphRunStream:
         # when its cursor catches up to the buffer.
         self._wire_request_more(mux, extensions)
 
-    def _wire_request_more(
-        self, mux: StreamMux, extensions: dict[str, Any]
-    ) -> None:
+    def _wire_request_more(self, mux: StreamMux, extensions: dict[str, Any]) -> None:
         """Set _request_more on all sync EventLogs so iteration drives the graph."""
         if isinstance(mux._events, EventLog):
             mux._events._request_more = self._pump_next
         for value in extensions.values():
             if isinstance(value, EventLog):
                 value._request_more = self._pump_next
-            elif isinstance(value, StreamChannel) and isinstance(
-                value._log, EventLog
-            ):
+            elif isinstance(value, StreamChannel) and isinstance(value._log, EventLog):
                 value._log._request_more = self._pump_next
 
     def _pump_next(self) -> bool:
@@ -74,8 +71,6 @@ class GraphRunStream:
             self._mux.fail(e)
             self._exhausted = True
             return False
-        from langgraph.stream._convert import convert_to_protocol_event
-
         self._mux.push(convert_to_protocol_event(part))
         return True
 
@@ -106,6 +101,7 @@ class GraphRunStream:
 
     def __iter__(self) -> Iterator[ProtocolEvent]:
         """Iterate all protocol events from the mux's main event log."""
+        assert isinstance(self._mux._events, EventLog)
         return iter(self._mux._events)
 
 
@@ -172,4 +168,5 @@ class AsyncGraphRunStream:
 
     def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
         """Iterate all protocol events from the mux's main event log."""
+        assert isinstance(self._mux._events, AsyncEventLog)
         return self._mux._events.__aiter__()

libs/langgraph/langgraph/stream/stream_channel.py

@@ -68,7 +68,6 @@ class StreamChannel(Generic[T]):
     def __aiter__(self) -> AsyncIterator[T]:
         if not isinstance(self._log, AsyncEventLog):
             raise RuntimeError(
-                "Cannot use async iteration on a sync StreamChannel. "
-                "Use 'for' instead."
+                "Cannot use async iteration on a sync StreamChannel. Use 'for' instead."
             )
         return self._log.__aiter__()

libs/langgraph/tests/test_streaming_handler.py

@@ -418,7 +418,6 @@ class TestStreamingHandlerSyncErrors:
             list(run)
 
 
-
 class TestStreamingHandlerSyncInterrupt:
     def test_interrupted(self) -> None:
         graph = _build_interrupt_graph()
@@ -773,6 +772,7 @@ class TestStreamMuxResilience:
     def test_close_continues_after_finalize_error(self) -> None:
         """If a transformer's finalize() raises, the main event log and
         remaining transformers should still be closed/finalized."""
+
         class BrokenFinalizer(StreamTransformer):
             def init(self) -> dict[str, Any]:
                 return {}
@@ -812,6 +812,7 @@ class TestStreamMuxResilience:
     def test_fail_continues_after_transformer_error(self) -> None:
         """If a transformer's fail() raises, the main event log and
         remaining transformers should still be failed."""
+
         class BrokenFailer(StreamTransformer):
             def init(self) -> dict[str, Any]:
                 return {}
@@ -848,6 +849,7 @@ class TestStreamMuxResilience:
 
     def test_close_still_closes_channels_after_finalize_error(self) -> None:
         """Channels should be closed even if a transformer's finalize raises."""
+
         class BrokenWithChannel(StreamTransformer):
             def __init__(self) -> None:
                 self._channel: StreamChannel[str] = StreamChannel("ch")