update check_client_capability

Author: ochafik

src/mcp/server/session.py

@@ -120,6 +120,12 @@ def check_client_capability(self, capability: types.ClientCapabilities) -> bool:
         if capability.sampling is not None:
             if client_caps.sampling is None:
                 return False
+            if capability.sampling.context is not None:
+                if client_caps.sampling.context is None:
+                    return False
+            if capability.sampling.tools is not None:
+                if client_caps.sampling.tools is None:
+                    return False
 
         if capability.elicitation is not None:
             if client_caps.elicitation is None:
@@ -234,6 +240,7 @@ async def create_message(
             max_tokens: Maximum number of tokens to generate.
             system_prompt: Optional system prompt.
             include_context: Optional context inclusion setting.
+                Requires client to have sampling.context capability.
             temperature: Optional sampling temperature.
             stop_sequences: Optional stop sequences.
             metadata: Optional metadata to pass through to the LLM provider.
@@ -251,6 +258,20 @@ async def create_message(
             McpError: If tool_use or tool_result blocks are misused when tools are provided.
         """
 
+        if tools is not None or tool_choice is not None:
+            has_tools_cap = self.check_client_capability( \
+                types.ClientCapabilities(sampling=types.SamplingCapability(tools=types.SamplingToolsCapability()))
+            )
+            if not has_tools_cap:
+                from mcp.shared.exceptions import McpError
+
+                raise McpError(
+                    types.ErrorData(
+                        code=types.INVALID_PARAMS,
+                        message="Client does not support sampling tools capability",
+                    )
+                )
+
         if messages and tools:
             last_content = messages[-1].content_as_list
             has_tool_results = any(c.type == "tool_result" for c in last_content)

tests/server/test_session.py

@@ -309,6 +309,15 @@ async def test_create_message_tool_result_validation():
                 capabilities=ServerCapabilities(),
             ),
         ) as session:
+            # Set up client params with sampling.tools capability for the test
+            session._client_params = types.InitializeRequestParams(
+                protocolVersion=types.LATEST_PROTOCOL_VERSION,
+                capabilities=types.ClientCapabilities(
+                    sampling=types.SamplingCapability(tools=types.SamplingToolsCapability())
+                ),
+                clientInfo=types.Implementation(name="test", version="1.0"),
+            )
+
             tool = types.Tool(name="test_tool", inputSchema={"type": "object"})
             text = types.TextContent(type="text", text="hello")
             tool_use = types.ToolUseContent(type="tool_use", id="call_1", name="test_tool", input={})