Adds support for ElicitCompleteNotification callbacks

Author: kylestratis

src/mcp/client/session.py

@@ -33,6 +33,10 @@ async def __call__(
         params: types.ElicitRequestParams,
     ) -> types.ElicitResult | types.ErrorData: ...  # pragma: no branch
 
+class ElicitCompleteFnT(Protocol):
+    async def __call__(
+            self, params: types.ElicitCompleteNotificationParams,
+    ) -> None: ... #pragma: no branch
 
 class ListRootsFnT(Protocol):
     async def __call__(
@@ -111,6 +115,11 @@ async def _default_elicitation_callback(
         message="Elicitation not supported",
     )
 
+async def _default_elicit_complete_callback(
+        params: types.ElicitCompleteNotificationParams
+) -> None:
+    pass
+
 
 async def _default_list_roots_callback(
     context: RequestContext["ClientSession", Any],
@@ -172,6 +181,7 @@ def __init__(
         read_timeout_seconds: timedelta | None = None,
         sampling_callback: SamplingFnT | None = None,
         elicitation_callback: ElicitationFnT | None = None,
+        elicit_complete_callback: ElicitCompleteFnT | None = None,
         list_roots_callback: ListRootsFnT | None = None,
         logging_callback: LoggingFnT | None = None,
         progress_notification_callback: ProgressNotificationFnT | None = None,
@@ -192,6 +202,7 @@ def __init__(
         self._client_info = client_info or DEFAULT_CLIENT_INFO
         self._sampling_callback = sampling_callback or _default_sampling_callback
         self._elicitation_callback = elicitation_callback or _default_elicitation_callback
+        self._elicit_complete_callback = elicit_complete_callback or _default_elicit_complete_callback
         self._list_roots_callback = list_roots_callback or _default_list_roots_callback
         self._logging_callback = logging_callback or _default_logging_callback
         self._progress_notification_callback = progress_notification_callback or _default_progress_callback
@@ -638,7 +649,7 @@ async def _received_notification(self, notification: types.ServerNotification) -
                 # Handle elicitation completion notification
                 # Clients MAY use this to retry requests or update UI
                 # The notification contains the elicitationId of the completed elicitation
-                pass
+                await self._elicit_complete_callback(params)
             case _:  # pragma: no cover
                 # CancelledNotification is handled separately in shared/session.py
                 # and should never reach this point. This case is defensive.

src/mcp/shared/memory.py

@@ -15,6 +15,7 @@
 import mcp.types as types
 from mcp.client.session import (
     ClientSession,
+    ElicitCompleteFnT,
     ElicitationFnT,
     ListRootsFnT,
     LoggingFnT,
@@ -74,6 +75,7 @@ async def create_connected_server_and_client_session(
     client_info: types.Implementation | None = None,
     raise_exceptions: bool = False,
     elicitation_callback: ElicitationFnT | None = None,
+    elicit_complete_callback: ElicitCompleteFnT | None = None,
 ) -> AsyncGenerator[ClientSession, None]:
     """Creates a ClientSession that is connected to a running MCP server."""
 
@@ -113,6 +115,7 @@ async def create_connected_server_and_client_session(
                     message_handler=message_handler,
                     client_info=client_info,
                     elicitation_callback=elicitation_callback,
+                    elicit_complete_callback=elicit_complete_callback,
                 ) as client_session:
                     await client_session.initialize()
                     yield client_session

tests/client/test_notification_callbacks.py

@@ -82,6 +82,18 @@ async def __call__(self) -> None:
         self.notification_count += 1
 
 
+class ElicitCompleteCollector:
+    """Collector for ElicitCompleteNotification events."""
+
+    def __init__(self) -> None:
+        """Initialize the collector."""
+        self.notifications: list[types.ElicitCompleteNotificationParams] = []
+
+    async def __call__(self, params: types.ElicitCompleteNotificationParams) -> None:
+        """Collect an elicit complete notification."""
+        self.notifications.append(params)
+
+
 @pytest.fixture
 def progress_collector() -> ProgressNotificationCollector:
     """Create a progress notification collector."""
@@ -112,6 +124,12 @@ def prompt_list_changed_collector() -> PromptListChangedCollector:
     return PromptListChangedCollector()
 
 
+@pytest.fixture
+def elicit_complete_collector() -> ElicitCompleteCollector:
+    """Create an elicit complete collector."""
+    return ElicitCompleteCollector()
+
+
 @pytest.mark.anyio
 async def test_progress_notification_callback(progress_collector: ProgressNotificationCollector) -> None:
     """Test that progress notifications are correctly received by the callback."""
@@ -298,6 +316,41 @@ async def message_handler(
         assert prompt_list_changed_collector.notification_count == 1
 
 
+@pytest.mark.anyio
+async def test_elicit_complete_callback(elicit_complete_collector: ElicitCompleteCollector) -> None:
+    """Test that elicit complete notifications are correctly received by the callback."""
+    from mcp.server.fastmcp import FastMCP
+
+    server = FastMCP("test")
+
+    @server.tool("send_elicit_complete")
+    async def send_elicit_complete_tool(elicitation_id: str) -> bool:
+        """Send an elicit complete notification to the client."""
+        await server.get_context().session.send_elicit_complete(elicitation_id)
+        return True
+
+    async def message_handler(
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:
+        """Handle exceptions from the session."""
+        if isinstance(message, Exception):  # pragma: no cover
+            raise message
+
+    async with create_session(
+        server._mcp_server,
+        elicit_complete_callback=elicit_complete_collector,
+        message_handler=message_handler,
+    ) as client_session:
+        # Trigger elicit complete notification
+        result = await client_session.call_tool("send_elicit_complete", {"elicitation_id": "test-elicit-123"})
+        assert result.isError is False
+
+        # Verify the notification was received
+        assert len(elicit_complete_collector.notifications) == 1
+        notification = elicit_complete_collector.notifications[0]
+        assert notification.elicitationId == "test-elicit-123"
+
+
 @pytest.mark.anyio
 @pytest.mark.parametrize(
     "notification_type,callback_param,collector_fixture,tool_name,tool_args,verification",
@@ -350,6 +403,17 @@ async def message_handler(
             {},
             lambda c: c.notification_count == 1,  # type: ignore[attr-defined]
         ),
+        (
+            "elicit_complete",
+            "elicit_complete_callback",
+            "elicit_complete_collector",
+            "send_elicit_complete",
+            {"elicitation_id": "param-test-elicit-456"},
+            lambda c: (  # type: ignore[misc]
+                len(c.notifications) == 1  # type: ignore[attr-defined]
+                and c.notifications[0].elicitationId == "param-test-elicit-456"  # type: ignore[attr-defined]
+            ),
+        ),
     ],
 )
 async def test_notification_callback_parametrized(
@@ -407,6 +471,12 @@ async def change_prompt_list_tool() -> bool:
         await server.get_context().session.send_prompt_list_changed()
         return True
 
+    @server.tool("send_elicit_complete")
+    async def send_elicit_complete_tool(elicitation_id: str) -> bool:
+        """Send an elicit complete notification to the client."""
+        await server.get_context().session.send_elicit_complete(elicitation_id)
+        return True
+
     async def message_handler(
         message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
     ) -> None:
@@ -478,6 +548,12 @@ async def send_prompt_list_changed_tool() -> bool:
         await server.get_context().session.send_prompt_list_changed()
         return True
 
+    @server.tool("send_elicit_complete")
+    async def send_elicit_complete_tool(elicitation_id: str) -> bool:
+        """Send an elicit complete notification."""
+        await server.get_context().session.send_elicit_complete(elicitation_id)
+        return True
+
     # Create session WITHOUT custom callbacks - all will use defaults
     async with create_session(server._mcp_server) as client_session:
         # Test progress notification with default callback
@@ -507,6 +583,10 @@ async def send_prompt_list_changed_tool() -> bool:
         result5 = await client_session.call_tool("send_prompt_list_changed", {})
         assert result5.isError is False
 
+        # Test elicit complete with default callback
+        result6 = await client_session.call_tool("send_elicit_complete", {"elicitation_id": "test-123"})
+        assert result6.isError is False
+
 
 @pytest.mark.anyio
 async def test_progress_tool_without_progress_token() -> None: