feat(sdk-py): wire websocket stream selection (#7832)

libs/sdk-py/langgraph_sdk/_async/stream.py

@@ -23,7 +23,12 @@ from langchain_core.language_models.chat_model_stream import AsyncChatModelStrea
 from langchain_protocol import Event, SubscribeParams
 
 from langgraph_sdk._async.http import HttpClient
-from langgraph_sdk.stream.transport import EventStreamHandle, ProtocolSseTransport
+from langgraph_sdk.stream.transport import (
+    AsyncProtocolTransport,
+    EventStreamHandle,
+    ProtocolSseTransport,
+    ProtocolWebSocketTransport,
+)
 
 
 class InterruptPayload(TypedDict):
@@ -1153,6 +1158,7 @@ class AsyncThreadStream:
         max_queue_size: int = 1024,
         run_start_timeout: float | None = None,
         explicit_thread_id: bool = False,
+        transport_kind: Literal["sse", "websocket"] = "sse",
     ) -> None:
         self._http = http
         self._headers = dict(headers or {})
@@ -1161,8 +1167,9 @@ class AsyncThreadStream:
         self._max_queue_size = max_queue_size
         self._run_start_timeout = run_start_timeout
         self._explicit_thread_id = explicit_thread_id
+        self._transport_kind = transport_kind
         self._closed = False
-        self._transport: ProtocolSseTransport | None = None
+        self._transport: AsyncProtocolTransport | None = None
         self._open_handles: list[EventStreamHandle] = []
         self._next_command_id = 1
         self._next_subscription_id = 1
@@ -1221,7 +1228,12 @@ class AsyncThreadStream:
     async def __aenter__(self) -> AsyncThreadStream:
         if self._closed:
             raise RuntimeError("AsyncThreadStream is closed and cannot be re-entered.")
-        self._transport = ProtocolSseTransport(
+        transport_cls = (
+            ProtocolWebSocketTransport
+            if self._transport_kind == "websocket"
+            else ProtocolSseTransport
+        )
+        self._transport = transport_cls(
             client=self._http.client,
             thread_id=self.thread_id,
             headers=self._headers,

libs/sdk-py/langgraph_sdk/_async/threads.py

@@ -743,6 +743,7 @@ class ThreadsClient:
         assistant_id: str,
         headers: Mapping[str, str] | None = None,
         run_start_timeout: float | None = None,
+        transport: Literal["sse", "websocket"] = "sse",
     ) -> AsyncThreadStream:
         """Open a v3 thread-centric streaming session.
 
@@ -758,15 +759,19 @@ class ThreadsClient:
             thread_id: optional explicit thread identifier. Defaults to a
                 fresh UUIDv4.
             assistant_id: assistant the run will use. Required.
-            headers: optional headers forwarded on every command and SSE
+            headers: optional headers forwarded on every command and event
                 request for this stream session.
             run_start_timeout: optional seconds to wait for an in-flight
                 `run.start` before subscribing operations raise
                 `asyncio.TimeoutError`. Defaults to `None` (wait forever).
+            transport: event transport to use — `"sse"` (default) or
+                `"websocket"`.
 
         Returns:
             An `AsyncThreadStream` to use as an async context manager.
         """
+        if transport not in ("sse", "websocket"):
+            raise ValueError("transport must be 'sse' or 'websocket'.")
         return AsyncThreadStream(
             http=self.http,
             thread_id=thread_id if thread_id is not None else str(uuid.uuid4()),
@@ -774,6 +779,7 @@ class ThreadsClient:
             headers=headers,
             run_start_timeout=run_start_timeout,
             explicit_thread_id=thread_id is not None,
+            transport_kind=transport,
         )
 
     async def join_stream(

libs/sdk-py/langgraph_sdk/_sync/stream.py

@@ -23,9 +23,11 @@ from langchain_protocol import Event, SubscribeParams
 
 from langgraph_sdk._sync.http import SyncHttpClient
 from langgraph_sdk.stream.sync_controller import SyncStreamController, _SyncSubscription
-from langgraph_sdk.stream.transport.sync_http import (
+from langgraph_sdk.stream.transport import (
     SyncEventStreamHandle,
     SyncProtocolSseTransport,
+    SyncProtocolTransport,
+    SyncProtocolWebSocketTransport,
 )
 
 
@@ -1146,6 +1148,7 @@ class SyncThreadStream:
         headers: Mapping[str, str] | None = None,
         run_start_timeout: float | None = None,
         explicit_thread_id: bool = False,
+        transport_kind: Literal["sse", "websocket"] = "sse",
     ) -> None:
         self._http = http
         self._headers = dict(headers or {})
@@ -1153,8 +1156,9 @@ class SyncThreadStream:
         self.assistant_id = assistant_id
         self._run_start_timeout = run_start_timeout
         self._explicit_thread_id = explicit_thread_id
+        self._transport_kind = transport_kind
         self._closed = False
-        self._transport: SyncProtocolSseTransport | None = None
+        self._transport: SyncProtocolTransport | None = None
         self._controller: SyncStreamController | None = None
         self._command_id_lock = threading.Lock()
         self._next_command_id = 1
@@ -1179,7 +1183,12 @@ class SyncThreadStream:
     def __enter__(self) -> SyncThreadStream:
         if self._closed:
             raise RuntimeError("SyncThreadStream is closed and cannot be re-entered.")
-        self._transport = SyncProtocolSseTransport(
+        transport_cls = (
+            SyncProtocolWebSocketTransport
+            if self._transport_kind == "websocket"
+            else SyncProtocolSseTransport
+        )
+        self._transport = transport_cls(
             client=self._http.client,
             thread_id=self.thread_id,
             headers=self._headers,

libs/sdk-py/langgraph_sdk/_sync/threads.py

@@ -731,6 +731,7 @@ class SyncThreadsClient:
         assistant_id: str,
         headers: Mapping[str, str] | None = None,
         run_start_timeout: float | None = None,
+        transport: Literal["sse", "websocket"] = "sse",
     ) -> SyncThreadStream:
         """Open a v3 thread-centric streaming session.
 
@@ -740,10 +741,14 @@ class SyncThreadsClient:
             assistant_id: assistant the run will use. Required.
             headers: optional headers forwarded on every command and SSE
                 request for this stream session.
+            transport: event transport to use, `"sse"` (default) or
+                `"websocket"`.
 
         Returns:
             A `SyncThreadStream` to use as a context manager.
         """
+        if transport not in ("sse", "websocket"):
+            raise ValueError("transport must be 'sse' or 'websocket'.")
         return SyncThreadStream(
             http=self.http,
             thread_id=thread_id if thread_id is not None else str(uuid.uuid4()),
@@ -751,6 +756,7 @@ class SyncThreadsClient:
             headers=headers,
             run_start_timeout=run_start_timeout,
             explicit_thread_id=thread_id is not None,
+            transport_kind=transport,
         )
 
     def join_stream(

libs/sdk-py/tests/streaming/test_sync_thread_stream.py

@@ -416,3 +416,27 @@ def test_sync_lifecycle_watcher_reconnects_with_since_after_transport_drop():
     assert terminal.status == "completed"
     assert terminal.error is None
     assert fake.stream_request_bodies[1]["since"] == 1
+
+
+def test_sync_threads_stream_accepts_websocket_transport_option():
+    with httpx.Client(base_url="http://test") as raw:
+        threads = SyncThreadsClient(SyncHttpClient(raw))
+        stream = threads.stream(
+            thread_id="t-1",
+            assistant_id="agent",
+            transport="websocket",
+        )
+    assert stream._transport_kind == "websocket"
+
+
+def test_sync_threads_stream_rejects_unknown_transport_option():
+    import pytest
+
+    with httpx.Client(base_url="http://test") as raw:
+        threads = SyncThreadsClient(SyncHttpClient(raw))
+        with pytest.raises(ValueError, match="transport"):
+            threads.stream(
+                thread_id="t-1",
+                assistant_id="agent",
+                transport="bogus",  # ty: ignore[invalid-argument-type]
+            )

libs/sdk-py/tests/streaming/test_thread_stream.py

@@ -11,6 +11,10 @@ import pytest
 from langgraph_sdk._async.http import HttpClient
 from langgraph_sdk._async.stream import AsyncThreadStream
 from langgraph_sdk._async.threads import ThreadsClient
+from langgraph_sdk.stream.transport import (
+    ProtocolSseTransport,
+    ProtocolWebSocketTransport,
+)
 from streaming._events import (
     lifecycle_completed_event,
     lifecycle_event,
@@ -171,6 +175,19 @@ async def test_aenter_constructs_transport_with_thread_id():
             assert stream._transport.thread_id == "t-1"
 
 
+async def test_aenter_selects_websocket_transport():
+    async with httpx.AsyncClient(base_url="http://test") as raw:
+        from langgraph_sdk._async.http import HttpClient
+        from langgraph_sdk._async.threads import ThreadsClient
+
+        threads = ThreadsClient(HttpClient(raw))
+        stream = threads.stream(
+            thread_id="t-1", assistant_id="agent", transport="websocket"
+        )
+        async with stream:
+            assert isinstance(stream._transport, ProtocolWebSocketTransport)
+
+
 async def test_aexit_closes_transport():
     fake = FakeServer()
     transport = httpx.ASGITransport(app=fake.app)
@@ -183,6 +200,7 @@ async def test_aexit_closes_transport():
         async with stream:
             inner_transport = stream._transport
         assert inner_transport is not None
+        assert isinstance(inner_transport, ProtocolSseTransport)
         assert inner_transport._closed is True
 
 
@@ -806,3 +824,25 @@ async def test_output_with_timeout_returns_new_awaitable_not_self():
             assert bounded is not thread.output
             assert bounded._timeout == 0.5
             assert thread.output._timeout is None
+
+
+async def test_threads_stream_accepts_websocket_transport_option():
+    async with httpx.AsyncClient(base_url="http://test") as raw:
+        threads = ThreadsClient(HttpClient(raw))
+        stream = threads.stream(
+            thread_id="t-1",
+            assistant_id="agent",
+            transport="websocket",
+        )
+    assert stream._transport_kind == "websocket"
+
+
+async def test_threads_stream_rejects_unknown_transport_option():
+    async with httpx.AsyncClient(base_url="http://test") as raw:
+        threads = ThreadsClient(HttpClient(raw))
+        with pytest.raises(ValueError, match="transport"):
+            threads.stream(
+                thread_id="t-1",
+                assistant_id="agent",
+                transport="bogus",  # ty: ignore[invalid-argument-type]
+            )

libs/sdk-py/tests/test_api_parity.py

@@ -45,7 +45,9 @@ def _normalize_return_annotation(ann: object) -> str:
     s = re.sub(r"Generator\[([^,\]]+)(?:,[^\]]*)?\]", r"Iterator[\1]", s)
     s = re.sub(r"AsyncIterator\[(.+)\]", r"Iterator[\1]", s)
     s = re.sub(r"AsyncIterable\[(.+)\]", r"Iterable[\1]", s)
-    s = s.replace("AsyncThreadStream", "SyncThreadStream")
+    # Normalize Async/Sync class prefixes so AsyncFoo and SyncFoo both compare as Foo.
+    s = re.sub(r"\bAsync([A-Z])", r"\1", s)
+    s = re.sub(r"\bSync([A-Z])", r"\1", s)
     return s