address comments

Author: Kludex

src/mcp/client/client.py

@@ -95,8 +95,8 @@ async def main():
     elicitation_callback: ElicitationFnT | None = None
     """Callback for handling elicitation requests."""
 
-    _session: ClientSession | None = None
-    _exit_stack: AsyncExitStack | None = None
+    _session: ClientSession | None = field(init=False, default=None)
+    _exit_stack: AsyncExitStack | None = field(init=False, default=None)
     _transport: Transport = field(init=False)
 
     def __post_init__(self) -> None:
@@ -113,10 +113,7 @@ async def __aenter__(self) -> Client:
             raise RuntimeError("Client is already entered; cannot reenter")
 
         async with AsyncExitStack() as exit_stack:
-            # Transports may return additional values (e.g., streamable_http_client returns a 3-tuple)
-            # We only need the first two elements (read_stream, write_stream)
-            streams = await exit_stack.enter_async_context(self._transport)
-            read_stream, write_stream = streams[0], streams[1]
+            read_stream, write_stream = await exit_stack.enter_async_context(self._transport)
 
             self._session = await exit_stack.enter_async_context(
                 ClientSession(

src/mcp/client/streamable_http.py

@@ -32,10 +32,10 @@
 logger = logging.getLogger(__name__)
 
 
+# TODO(Marcelo): Put the TransportStreams in a module under shared, so we can import here.
 SessionMessageOrError = SessionMessage | Exception
 StreamWriter = MemoryObjectSendStream[SessionMessageOrError]
 StreamReader = MemoryObjectReceiveStream[SessionMessage]
-GetSessionIdCallback = Callable[[], str | None]
 
 MCP_SESSION_ID = "mcp-session-id"
 MCP_PROTOCOL_VERSION = "mcp-protocol-version"