adjust types after the spec revision

Author: ihrpr

src/mcp/client/session.py

@@ -143,8 +143,11 @@ async def initialize(self) -> types.InitializeResult:
             if self._sampling_callback is not _default_sampling_callback
             else None
         )
-        # TODO: change this to a more specific type
-        elicitation = types.ElicitationCapability()
+        elicitation = (
+            types.ElicitationCapability()
+            if self._elicitation_callback is not _default_elicitation_callback
+            else None
+        )
         roots = (
             # TODO: Should this be based on whether we
             # _will_ send notifications, or only whether

src/mcp/server/fastmcp/server.py

@@ -10,11 +10,11 @@
     asynccontextmanager,
 )
 from itertools import chain
-from typing import Any, Generic, Literal
+from typing import Any, Generic, Literal, TypeVar
 
 import anyio
 import pydantic_core
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError
 from pydantic.networks import AnyUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from starlette.applications import Starlette
@@ -67,6 +67,8 @@
 
 logger = get_logger(__name__)
 
+ElicitedModelT = TypeVar("ElicitedModelT", bound=BaseModel)
+
 
 class Settings(BaseSettings, Generic[LifespanResultT]):
     """FastMCP server settings.
@@ -1005,35 +1007,48 @@ async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContent
     async def elicit(
         self,
         message: str,
-        requestedSchema: dict[str, Any],
-    ) -> dict[str, Any]:
+        schema: type[ElicitedModelT],
+    ) -> ElicitedModelT:
         """Elicit information from the client/user.
 
         This method can be used to interactively ask for additional information from the
-        client within a tool's execution.
-        The client might display the message to the user and collect a response
-        according to the provided schema. Or in case a client is an agent, it might
-        decide how to handle the elicitation -- either by asking the user or
-        automatically generating a response.
+        client within a tool's execution. The client might display the message to the
+        user and collect a response according to the provided schema. Or in case a
+        client
+        is an agent, it might decide how to handle the elicitation -- either by asking
+        the user or automatically generating a response.
 
         Args:
-            message: The message to present to the user
-            requestedSchema: JSON Schema defining the expected response structure
+            schema: A Pydantic model class defining the expected response structure
+            message: Optional message to present to the user. If not provided, will use
+                    a default message based on the schema
 
         Returns:
-            The user's response as a dict matching the request schema structure
+            An instance of the schema type with the user's response
 
         Raises:
-            ValueError: If elicitation is not supported by the client or fails
+            Exception: If the user declines or cancels the elicitation
+            ValidationError: If the response doesn't match the schema
         """
 
+        json_schema = schema.model_json_schema()
+
         result = await self.request_context.session.elicit(
             message=message,
-            requestedSchema=requestedSchema,
+            requestedSchema=json_schema,
             related_request_id=self.request_id,
         )
 
-        return result.content
+        if result.action == "accept" and result.content:
+            # Validate and parse the content using the schema
+            try:
+                return schema.model_validate(result.content)
+            except ValidationError as e:
+                raise ValueError(f"Invalid response: {e}")
+        elif result.action == "decline":
+            raise Exception("User declined to provide information")
+        else:
+            raise Exception("User cancelled the request")
 
     async def log(
         self,

src/mcp/server/session.py

@@ -277,14 +277,14 @@ async def list_roots(self) -> types.ListRootsResult:
     async def elicit(
         self,
         message: str,
-        requestedSchema: dict[str, Any],
+        requestedSchema: types.ElicitRequestedSchema,
         related_request_id: types.RequestId | None = None,
     ) -> types.ElicitResult:
         """Send an elicitation/create request.
 
         Args:
             message: The message to present to the user
-            requestedSchema: JSON Schema defining the expected response structure
+            requestedSchema: Schema defining the expected response structure
 
         Returns:
             The client's response

src/mcp/shared/session.py

@@ -369,7 +369,8 @@ async def _receive_loop(self) -> None:
                             request=validated_request,
                             session=self,
                             on_complete=lambda r: self._in_flight.pop(
-                                r.request_id, None),
+                                r.request_id, None
+                            ),
                             message_metadata=message.metadata,
                         )
                         self._in_flight[responder.request_id] = responder
@@ -394,7 +395,8 @@ async def _receive_loop(self) -> None:
                             ),
                         )
                         session_message = SessionMessage(
-                            message=JSONRPCMessage(error_response))
+                            message=JSONRPCMessage(error_response)
+                        )
                         await self._write_stream.send(session_message)
 
                 elif isinstance(message.message.root, JSONRPCNotification):

src/mcp/types.py

@@ -1174,16 +1174,16 @@ class ClientNotification(
     pass
 
 
+# Type for elicitation schema - a JSON Schema dict
+ElicitRequestedSchema: TypeAlias = dict[str, Any]
+"""Schema for elicitation requests."""
+
+
 class ElicitRequestParams(RequestParams):
     """Parameters for elicitation requests."""
 
     message: str
-    """The message to present to the user."""
-
-    requestedSchema: dict[str, Any]
-    """
-    A JSON Schema object defining the expected structure of the response.
-    """
+    requestedSchema: ElicitRequestedSchema
     model_config = ConfigDict(extra="allow")
 
 
@@ -1195,10 +1195,21 @@ class ElicitRequest(Request[ElicitRequestParams, Literal["elicitation/create"]])
 
 
 class ElicitResult(Result):
-    """The client's response to an elicitation/create request from the server."""
+    """The client's response to an elicitation request."""
 
-    content: dict[str, Any]
-    """The response from the client, matching the structure of requestedSchema."""
+    action: Literal["accept", "decline", "cancel"]
+    """
+    The user action in response to the elicitation.
+    - "accept": User submitted the form/confirmed the action
+    - "decline": User explicitly declined the action
+    - "cancel": User dismissed without making an explicit choice
+    """
+
+    content: dict[str, str | int | float | bool | None] | None = None
+    """
+    The submitted form data, only present when action is "accept".
+    Contains values matching the requested schema.
+    """
 
 
 class ClientResult(

tests/issues/test_malformed_input.py

@@ -1,4 +1,4 @@
-# Claude Debug  
+# Claude Debug
 """Test for HackerOne vulnerability report #3156202 - malformed input DOS."""
 
 import anyio
@@ -38,7 +38,7 @@ async def test_malformed_initialize_request_does_not_crash_server():
             method="initialize",
             # params=None  # Missing required params field
         )
-        
+
         # Wrap in session message
         request_message = SessionMessage(message=JSONRPCMessage(malformed_request))
 
@@ -54,22 +54,22 @@ async def test_malformed_initialize_request_does_not_crash_server():
         ):
             # Send the malformed request
             await read_send_stream.send(request_message)
-            
+
             # Give the session time to process the request
             await anyio.sleep(0.1)
-            
+
             # Check that we received an error response instead of a crash
             try:
                 response_message = write_receive_stream.receive_nowait()
                 response = response_message.message.root
-                
+
                 # Verify it's a proper JSON-RPC error response
                 assert isinstance(response, JSONRPCError)
                 assert response.jsonrpc == "2.0"
                 assert response.id == "f20fe86132ed4cd197f89a7134de5685"
                 assert response.error.code == INVALID_PARAMS
                 assert "Invalid request parameters" in response.error.message
-                
+
                 # Verify the session is still alive and can handle more requests
                 # Send another malformed request to confirm server stability
                 another_malformed_request = JSONRPCRequest(
@@ -81,18 +81,18 @@ async def test_malformed_initialize_request_does_not_crash_server():
                 another_request_message = SessionMessage(
                     message=JSONRPCMessage(another_malformed_request)
                 )
-                
+
                 await read_send_stream.send(another_request_message)
                 await anyio.sleep(0.1)
-                
+
                 # Should get another error response, not a crash
                 second_response_message = write_receive_stream.receive_nowait()
                 second_response = second_response_message.message.root
-                
+
                 assert isinstance(second_response, JSONRPCError)
                 assert second_response.id == "test_id_2"
                 assert second_response.error.code == INVALID_PARAMS
-                
+
             except anyio.WouldBlock:
                 pytest.fail("No response received - server likely crashed")
     finally:
@@ -140,14 +140,14 @@ async def test_multiple_concurrent_malformed_requests():
                     message=JSONRPCMessage(malformed_request)
                 )
                 malformed_requests.append(request_message)
-            
+
             # Send all requests
             for request in malformed_requests:
                 await read_send_stream.send(request)
-            
+
             # Give time to process
             await anyio.sleep(0.2)
-            
+
             # Verify we get error responses for all requests
             error_responses = []
             try:
@@ -156,10 +156,10 @@ async def test_multiple_concurrent_malformed_requests():
                     error_responses.append(response_message.message.root)
             except anyio.WouldBlock:
                 pass  # No more messages
-            
+
             # Should have received 10 error responses
             assert len(error_responses) == 10
-            
+
             for i, response in enumerate(error_responses):
                 assert isinstance(response, JSONRPCError)
                 assert response.id == f"malformed_{i}"
@@ -169,4 +169,4 @@ async def test_multiple_concurrent_malformed_requests():
         await read_send_stream.aclose()
         await write_send_stream.aclose()
         await read_receive_stream.aclose()
-        await write_receive_stream.aclose()
+        await write_receive_stream.aclose()