Merge branch 'main' into pydantic2.11-support

Author: medaminezghal

.pre-commit-config.yaml

@@ -7,15 +7,29 @@ repos:
       - id: prettier
         types_or: [yaml, json5]
 
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.8.1
+  - repo: local
     hooks:
       - id: ruff-format
+        name: Ruff Format
+        entry: uv run ruff
+        args: [format]
+        language: system
+        types: [python]
+        pass_filenames: false
       - id: ruff
-        args: [--fix, --exit-non-zero-on-fix]
-
-  - repo: local
-    hooks:
+        name: Ruff
+        entry: uv run ruff
+        args: ["check", "--fix", "--exit-non-zero-on-fix"]
+        types: [python]
+        language: system
+        pass_filenames: false
+      - id: pyright
+        name: pyright
+        entry: uv run pyright
+        args: [src]
+        language: system
+        types: [python]
+        pass_filenames: false
       - id: uv-lock-check
         name: Check uv.lock is up to date
         entry: uv lock --check

src/mcp/client/__main__.py

@@ -38,9 +38,13 @@ async def message_handler(
 async def run_session(
     read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
     write_stream: MemoryObjectSendStream[JSONRPCMessage],
+    client_info: types.Implementation | None = None,
 ):
     async with ClientSession(
-        read_stream, write_stream, message_handler=message_handler
+        read_stream,
+        write_stream,
+        message_handler=message_handler,
+        client_info=client_info,
     ) as session:
         logger.info("Initializing session")
         await session.initialize()

src/mcp/client/session.py

@@ -10,6 +10,8 @@
 from mcp.shared.session import BaseSession, RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
+DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
+
 
 class SamplingFnT(Protocol):
     async def __call__(
@@ -97,6 +99,7 @@ def __init__(
         list_roots_callback: ListRootsFnT | None = None,
         logging_callback: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
+        client_info: types.Implementation | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -105,6 +108,7 @@ def __init__(
             types.ServerNotification,
             read_timeout_seconds=read_timeout_seconds,
         )
+        self._client_info = client_info or DEFAULT_CLIENT_INFO
         self._sampling_callback = sampling_callback or _default_sampling_callback
         self._list_roots_callback = list_roots_callback or _default_list_roots_callback
         self._logging_callback = logging_callback or _default_logging_callback
@@ -130,7 +134,7 @@ async def initialize(self) -> types.InitializeResult:
                             experimental=None,
                             roots=roots,
                         ),
-                        clientInfo=types.Implementation(name="mcp", version="0.1.0"),
+                        clientInfo=self._client_info,
                     ),
                 )
             ),

src/mcp/server/fastmcp/tools/base.py

@@ -2,7 +2,7 @@
 
 import inspect
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, get_origin
 
 from pydantic import BaseModel, Field
 
@@ -53,7 +53,9 @@ def from_function(
         if context_kwarg is None:
             sig = inspect.signature(fn)
             for param_name, param in sig.parameters.items():
-                if param.annotation is Context:
+                if get_origin(param.annotation) is not None:
+                    continue
+                if issubclass(param.annotation, Context):
                     context_kwarg = param_name
                     break
 

src/mcp/shared/memory.py

@@ -10,6 +10,7 @@
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
+import mcp.types as types
 from mcp.client.session import (
     ClientSession,
     ListRootsFnT,
@@ -65,6 +66,7 @@ async def create_connected_server_and_client_session(
     list_roots_callback: ListRootsFnT | None = None,
     logging_callback: LoggingFnT | None = None,
     message_handler: MessageHandlerFnT | None = None,
+    client_info: types.Implementation | None = None,
     raise_exceptions: bool = False,
 ) -> AsyncGenerator[ClientSession, None]:
     """Creates a ClientSession that is connected to a running MCP server."""
@@ -95,6 +97,7 @@ async def create_connected_server_and_client_session(
                     list_roots_callback=list_roots_callback,
                     logging_callback=logging_callback,
                     message_handler=message_handler,
+                    client_info=client_info,
                 ) as client_session:
                     await client_session.initialize()
                     yield client_session

tests/client/test_session.py

@@ -2,7 +2,7 @@
 import pytest
 
 import mcp.types as types
-from mcp.client.session import ClientSession
+from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
 from mcp.shared.session import RequestResponder
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
@@ -111,3 +111,131 @@ async def message_handler(
     # Check that the client sent the initialized notification
     assert initialized_notification
     assert isinstance(initialized_notification.root, InitializedNotification)
+
+
+@pytest.mark.anyio
+async def test_client_session_custom_client_info():
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
+        JSONRPCMessage
+    ](1)
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
+        JSONRPCMessage
+    ](1)
+
+    custom_client_info = Implementation(name="test-client", version="1.2.3")
+    received_client_info = None
+
+    async def mock_server():
+        nonlocal received_client_info
+
+        jsonrpc_request = await client_to_server_receive.receive()
+        assert isinstance(jsonrpc_request.root, JSONRPCRequest)
+        request = ClientRequest.model_validate(
+            jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
+        )
+        assert isinstance(request.root, InitializeRequest)
+        received_client_info = request.root.params.clientInfo
+
+        result = ServerResult(
+            InitializeResult(
+                protocolVersion=LATEST_PROTOCOL_VERSION,
+                capabilities=ServerCapabilities(),
+                serverInfo=Implementation(name="mock-server", version="0.1.0"),
+            )
+        )
+
+        async with server_to_client_send:
+            await server_to_client_send.send(
+                JSONRPCMessage(
+                    JSONRPCResponse(
+                        jsonrpc="2.0",
+                        id=jsonrpc_request.root.id,
+                        result=result.model_dump(
+                            by_alias=True, mode="json", exclude_none=True
+                        ),
+                    )
+                )
+            )
+            # Receive initialized notification
+            await client_to_server_receive.receive()
+
+    async with (
+        ClientSession(
+            server_to_client_receive,
+            client_to_server_send,
+            client_info=custom_client_info,
+        ) as session,
+        anyio.create_task_group() as tg,
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        tg.start_soon(mock_server)
+        await session.initialize()
+
+    # Assert that the custom client info was sent
+    assert received_client_info == custom_client_info
+
+
+@pytest.mark.anyio
+async def test_client_session_default_client_info():
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
+        JSONRPCMessage
+    ](1)
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
+        JSONRPCMessage
+    ](1)
+
+    received_client_info = None
+
+    async def mock_server():
+        nonlocal received_client_info
+
+        jsonrpc_request = await client_to_server_receive.receive()
+        assert isinstance(jsonrpc_request.root, JSONRPCRequest)
+        request = ClientRequest.model_validate(
+            jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
+        )
+        assert isinstance(request.root, InitializeRequest)
+        received_client_info = request.root.params.clientInfo
+
+        result = ServerResult(
+            InitializeResult(
+                protocolVersion=LATEST_PROTOCOL_VERSION,
+                capabilities=ServerCapabilities(),
+                serverInfo=Implementation(name="mock-server", version="0.1.0"),
+            )
+        )
+
+        async with server_to_client_send:
+            await server_to_client_send.send(
+                JSONRPCMessage(
+                    JSONRPCResponse(
+                        jsonrpc="2.0",
+                        id=jsonrpc_request.root.id,
+                        result=result.model_dump(
+                            by_alias=True, mode="json", exclude_none=True
+                        ),
+                    )
+                )
+            )
+            # Receive initialized notification
+            await client_to_server_receive.receive()
+
+    async with (
+        ClientSession(
+            server_to_client_receive,
+            client_to_server_send,
+        ) as session,
+        anyio.create_task_group() as tg,
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        tg.start_soon(mock_server)
+        await session.initialize()
+
+    # Assert that the default client info was sent
+    assert received_client_info == DEFAULT_CLIENT_INFO

tests/server/fastmcp/test_tool_manager.py

@@ -4,8 +4,11 @@
 import pytest
 from pydantic import BaseModel
 
+from mcp.server.fastmcp import Context, FastMCP
 from mcp.server.fastmcp.exceptions import ToolError
 from mcp.server.fastmcp.tools import ToolManager
+from mcp.server.session import ServerSessionT
+from mcp.shared.context import LifespanContextT
 
 
 class TestAddTools:
@@ -194,8 +197,6 @@ def concat_strs(vals: list[str] | str) -> str:
 
     @pytest.mark.anyio
     async def test_call_tool_with_complex_model(self):
-        from mcp.server.fastmcp import Context
-
         class MyShrimpTank(BaseModel):
             class Shrimp(BaseModel):
                 name: str
@@ -223,8 +224,6 @@ def name_shrimp(tank: MyShrimpTank, ctx: Context) -> list[str]:
 class TestToolSchema:
     @pytest.mark.anyio
     async def test_context_arg_excluded_from_schema(self):
-        from mcp.server.fastmcp import Context
-
         def something(a: int, ctx: Context) -> int:
             return a
 
@@ -241,7 +240,6 @@ class TestContextHandling:
     def test_context_parameter_detection(self):
         """Test that context parameters are properly detected in
         Tool.from_function()."""
-        from mcp.server.fastmcp import Context
 
         def tool_with_context(x: int, ctx: Context) -> str:
             return str(x)
@@ -256,10 +254,17 @@ def tool_without_context(x: int) -> str:
         tool = manager.add_tool(tool_without_context)
         assert tool.context_kwarg is None
 
+        def tool_with_parametrized_context(
+            x: int, ctx: Context[ServerSessionT, LifespanContextT]
+        ) -> str:
+            return str(x)
+
+        tool = manager.add_tool(tool_with_parametrized_context)
+        assert tool.context_kwarg == "ctx"
+
     @pytest.mark.anyio
     async def test_context_injection(self):
         """Test that context is properly injected during tool execution."""
-        from mcp.server.fastmcp import Context, FastMCP
 
         def tool_with_context(x: int, ctx: Context) -> str:
             assert isinstance(ctx, Context)
@@ -276,7 +281,6 @@ def tool_with_context(x: int, ctx: Context) -> str:
     @pytest.mark.anyio
     async def test_context_injection_async(self):
         """Test that context is properly injected in async tools."""
-        from mcp.server.fastmcp import Context, FastMCP
 
         async def async_tool(x: int, ctx: Context) -> str:
             assert isinstance(ctx, Context)
@@ -293,7 +297,6 @@ async def async_tool(x: int, ctx: Context) -> str:
     @pytest.mark.anyio
     async def test_context_optional(self):
         """Test that context is optional when calling tools."""
-        from mcp.server.fastmcp import Context
 
         def tool_with_context(x: int, ctx: Context | None = None) -> str:
             return str(x)
@@ -307,7 +310,6 @@ def tool_with_context(x: int, ctx: Context | None = None) -> str:
     @pytest.mark.anyio
     async def test_context_error_handling(self):
         """Test error handling when context injection fails."""
-        from mcp.server.fastmcp import Context, FastMCP
 
         def tool_with_context(x: int, ctx: Context) -> str:
             raise ValueError("Test error")