headers in handshake

Author: andrewmjc

src/mcp/client/websocket.py

@@ -2,6 +2,7 @@
 import logging
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
+from typing import Any
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
@@ -17,6 +18,7 @@
 @asynccontextmanager
 async def websocket_client(
     url: str,
+    headers: dict[str, Any] | None = None,
 ) -> AsyncGenerator[
     tuple[
         MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
@@ -48,7 +50,7 @@ async def websocket_client(
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
 
     # Connect using websockets, requesting the "mcp" subprotocol
-    async with ws_connect(url, subprotocols=[Subprotocol("mcp")]) as ws:
+    async with ws_connect(url, subprotocols=[Subprotocol("mcp")], additional_headers=headers) as ws:
 
         async def ws_reader():
             """