Merge branch 'main' into feature/add-enable-disable-methods-tools

Author: davila7

examples/clients/simple-chatbot/mcp_simple_chatbot/main.py

@@ -245,7 +245,7 @@ def get_response(self, messages: list[dict[str, str]]) -> str:
         }
         payload = {
             "messages": messages,
-            "model": "llama-3.2-90b-vision-preview",
+            "model": "meta-llama/llama-4-scout-17b-16e-instruct",
             "temperature": 0.7,
             "max_tokens": 4096,
             "top_p": 1,
@@ -284,12 +284,9 @@ def __init__(self, servers: list[Server], llm_client: LLMClient) -> None:
 
     async def cleanup_servers(self) -> None:
         """Clean up all servers properly."""
-        cleanup_tasks = [
-            asyncio.create_task(server.cleanup()) for server in self.servers
-        ]
-        if cleanup_tasks:
+        for server in reversed(self.servers):
             try:
-                await asyncio.gather(*cleanup_tasks, return_exceptions=True)
+                await server.cleanup()
             except Exception as e:
                 logging.warning(f"Warning during final cleanup: {e}")
 

src/mcp/client/session_group.py

@@ -77,10 +77,10 @@ class ClientSessionGroup:
     the client and can be accessed via the session.
 
     Example Usage:
-        name_fn = lambda name, server_info: f"{(server_info.name)}-{name}"
+        name_fn = lambda name, server_info: f"{(server_info.name)}_{name}"
         async with ClientSessionGroup(component_name_hook=name_fn) as group:
             for server_params in server_params:
-                group.connect_to_server(server_param)
+                await group.connect_to_server(server_param)
             ...
 
     """
@@ -145,14 +145,15 @@ async def __aexit__(
     ) -> bool | None:
         """Closes session exit stacks and main exit stack upon completion."""
 
+        # Only close the main exit stack if we created it
+        if self._owns_exit_stack:
+            await self._exit_stack.aclose()
+
         # Concurrently close session stacks.
         async with anyio.create_task_group() as tg:
             for exit_stack in self._session_exit_stacks.values():
                 tg.start_soon(exit_stack.aclose)
 
-        # Only close the main exit stack if we created it
-        if self._owns_exit_stack:
-            await self._exit_stack.aclose()
 
     @property
     def sessions(self) -> list[mcp.ClientSession]:

src/mcp/client/sse.py

@@ -10,7 +10,7 @@
 from httpx_sse import aconnect_sse
 
 import mcp.types as types
-from mcp.shared._httpx_utils import create_mcp_http_client
+from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
 from mcp.shared.message import SessionMessage
 
 logger = logging.getLogger(__name__)
@@ -26,6 +26,7 @@ async def sse_client(
     headers: dict[str, Any] | None = None,
     timeout: float = 5,
     sse_read_timeout: float = 60 * 5,
+    httpx_client_factory: McpHttpClientFactory = create_mcp_http_client,
     auth: httpx.Auth | None = None,
 ):
     """
@@ -53,7 +54,7 @@ async def sse_client(
     async with anyio.create_task_group() as tg:
         try:
             logger.info(f"Connecting to SSE endpoint: {remove_request_params(url)}")
-            async with create_mcp_http_client(headers=headers, auth=auth) as client:
+            async with httpx_client_factory(headers=headers, auth=auth) as client:
                 async with aconnect_sse(
                     client,
                     "GET",

src/mcp/client/streamable_http.py

@@ -19,7 +19,7 @@
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
 
-from mcp.shared._httpx_utils import create_mcp_http_client
+from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
 from mcp.shared.message import ClientMessageMetadata, SessionMessage
 from mcp.types import (
     ErrorData,
@@ -430,6 +430,7 @@ async def streamablehttp_client(
     timeout: timedelta = timedelta(seconds=30),
     sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
     terminate_on_close: bool = True,
+    httpx_client_factory: McpHttpClientFactory = create_mcp_http_client,
     auth: httpx.Auth | None = None,
 ) -> AsyncGenerator[
     tuple[
@@ -464,7 +465,7 @@ async def streamablehttp_client(
         try:
             logger.info(f"Connecting to StreamableHTTP endpoint: {url}")
 
-            async with create_mcp_http_client(
+            async with httpx_client_factory(
                 headers=transport.request_headers,
                 timeout=httpx.Timeout(
                     transport.timeout.seconds, read=transport.sse_read_timeout.seconds

src/mcp/shared/_httpx_utils.py

@@ -1,12 +1,21 @@
 """Utilities for creating standardized httpx AsyncClient instances."""
 
-from typing import Any
+from typing import Any, Protocol
 
 import httpx
 
 __all__ = ["create_mcp_http_client"]
 
 
+class McpHttpClientFactory(Protocol):
+    def __call__(
+        self,
+        headers: dict[str, str] | None = None,
+        timeout: httpx.Timeout | None = None,
+        auth: httpx.Auth | None = None,
+    ) -> httpx.AsyncClient: ...
+
+
 def create_mcp_http_client(
     headers: dict[str, str] | None = None,
     timeout: httpx.Timeout | None = None,