From cfcaefd09b47e1a1629dac36dd30b29363e78f69 Mon Sep 17 00:00:00 2001
From: Yujong Lee <yujong@berri.ai>
Date: Sat, 5 Sep 2026 13:29:53 -0700
Subject: [PATCH] refactor(native): use endpoint binding for WebSocket dispatch

---
 litellm/rust_bridge/responses_websocket.py |  6 ++--
 litellm/rust_bridge/runtime.py             | 41 ----------------------
 2 files changed, 3 insertions(+), 44 deletions(-)

diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py
index 8d0b7c263d0..d82b3b94669 100644
--- a/litellm/rust_bridge/responses_websocket.py
+++ b/litellm/rust_bridge/responses_websocket.py
@@ -14,16 +14,16 @@ from litellm.rust_bridge.protocols import (
     RustResponsesWebSocketConnection,
 )
 from litellm.rust_bridge.runtime import (
-    AsyncEndpointDispatch,
     BridgeErrorContext,
+    EndpointBinding,
     async_none,
     identity,
 )
 from litellm.rust_bridge.timeouts import timeout_to_seconds
 
-_RESPONSES_WEBSOCKET: Final[AsyncEndpointDispatch[RustResponsesWebSocketConnection]] = AsyncEndpointDispatch.native(
+_RESPONSES_WEBSOCKET: Final[EndpointBinding[RustResponsesWebSocketConnection]] = EndpointBinding.native(
     route="responses_websocket",
-    asynchronous=lambda native: native.ResponsesWebSocketConnection,
+    select=lambda native: native.ResponsesWebSocketConnection,
     enabled=rust_enabled,
 )
 
diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py
index eef349b22b5..355fa1b17eb 100644
--- a/litellm/rust_bridge/runtime.py
+++ b/litellm/rust_bridge/runtime.py
@@ -449,47 +449,6 @@ class EndpointDispatch(Generic[SyncBindingT, AsyncBindingT]):
         )
 
 
-@dataclass(frozen=True, slots=True)
-class AsyncEndpointDispatch(Generic[AsyncBindingT]):
-    asynchronous: EndpointBinding[AsyncBindingT]
-
-    @staticmethod
-    def native(
-        *,
-        route: str,
-        asynchronous: Callable[[NativeModule], SelectedAsyncT],
-        enabled: RustEnablement,
-    ) -> AsyncEndpointDispatch[SelectedAsyncT]:
-        return AsyncEndpointDispatch(
-            asynchronous=EndpointBinding.native(route=route, select=asynchronous, enabled=enabled)
-        )
-
-    def override(self, value: AsyncBindingT | None) -> None:
-        self.asynchronous.override(value)
-
-    def reset(self) -> None:
-        self.asynchronous.reset()
-
-    async def ainvoke(
-        self,
-        *,
-        prepare: Callable[[], RequestT],
-        call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]],
-        fallback: Callable[[], Awaitable[ResultT]],
-        adapt: Callable[[NativeT], ResultT],
-        error_context: BridgeErrorContext,
-        eligible: bool = True,
-    ) -> ResultT:
-        return await self.asynchronous.ainvoke(
-            prepare=prepare,
-            call=call,
-            fallback=fallback,
-            adapt=adapt,
-            error_context=error_context,
-            eligible=eligible,
-        )
-
-
 def _error_message(error: BaseException) -> str:
     reason: Final[object] = error.args[0] if error.args else str(error)
     return reason if isinstance(reason, str) else str(reason)