From fd4257300e6e3819631292fe3bd65375eee3ea2b Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 27 May 2026 16:48:22 -0400 Subject: [PATCH] feat(sdk-py): wire websocket stream selection (#7832) --- libs/sdk-py/langgraph_sdk/_async/stream.py | 18 +++++++-- libs/sdk-py/langgraph_sdk/_async/threads.py | 8 +++- libs/sdk-py/langgraph_sdk/_sync/stream.py | 15 +++++-- libs/sdk-py/langgraph_sdk/_sync/threads.py | 6 +++ .../streaming/test_sync_thread_stream.py | 24 +++++++++++ .../tests/streaming/test_thread_stream.py | 40 +++++++++++++++++++ libs/sdk-py/tests/test_api_parity.py | 4 +- 7 files changed, 107 insertions(+), 8 deletions(-) diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py index d8110de26..70a86d0fe 100644 --- a/libs/sdk-py/langgraph_sdk/_async/stream.py +++ b/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, diff --git a/libs/sdk-py/langgraph_sdk/_async/threads.py b/libs/sdk-py/langgraph_sdk/_async/threads.py index 2b1db1077..7b2930047 100644 --- a/libs/sdk-py/langgraph_sdk/_async/threads.py +++ b/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( diff --git a/libs/sdk-py/langgraph_sdk/_sync/stream.py b/libs/sdk-py/langgraph_sdk/_sync/stream.py index 7c1b07505..0aaccdf47 100644 --- a/libs/sdk-py/langgraph_sdk/_sync/stream.py +++ b/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, diff --git a/libs/sdk-py/langgraph_sdk/_sync/threads.py b/libs/sdk-py/langgraph_sdk/_sync/threads.py index 557a2a751..ad1139a00 100644 --- a/libs/sdk-py/langgraph_sdk/_sync/threads.py +++ b/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( diff --git a/libs/sdk-py/tests/streaming/test_sync_thread_stream.py b/libs/sdk-py/tests/streaming/test_sync_thread_stream.py index 75a4ba450..63be80953 100644 --- a/libs/sdk-py/tests/streaming/test_sync_thread_stream.py +++ b/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] + ) diff --git a/libs/sdk-py/tests/streaming/test_thread_stream.py b/libs/sdk-py/tests/streaming/test_thread_stream.py index b430d2670..43500c852 100644 --- a/libs/sdk-py/tests/streaming/test_thread_stream.py +++ b/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] + ) diff --git a/libs/sdk-py/tests/test_api_parity.py b/libs/sdk-py/tests/test_api_parity.py index 5227679cd..b67f8bfe0 100644 --- a/libs/sdk-py/tests/test_api_parity.py +++ b/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