add callback for logging message notification (#314)

src/mcp/client/session.py

@@ -24,6 +24,13 @@ class ListRootsFnT(Protocol):
     ) -> types.ListRootsResult | types.ErrorData: ...
 
 
+class LoggingFnT(Protocol):
+    async def __call__(
+        self,
+        params: types.LoggingMessageNotificationParams,
+    ) -> None: ...
+
+
 async def _default_sampling_callback(
     context: RequestContext["ClientSession", Any],
     params: types.CreateMessageRequestParams,
@@ -43,6 +50,12 @@ async def _default_list_roots_callback(
     )
 
 
+async def _default_logging_callback(
+    params: types.LoggingMessageNotificationParams,
+) -> None:
+    pass
+
+
 ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(
     types.ClientResult | types.ErrorData
 )
@@ -64,6 +77,7 @@ class ClientSession(
         read_timeout_seconds: timedelta | None = None,
         sampling_callback: SamplingFnT | None = None,
         list_roots_callback: ListRootsFnT | None = None,
+        logging_callback: LoggingFnT | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -74,6 +88,7 @@ class ClientSession(
         )
         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
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability()
@@ -321,3 +336,13 @@ class ClientSession(
                     return await responder.respond(
                         types.ClientResult(root=types.EmptyResult())
                     )
+
+    async def _received_notification(
+        self, notification: types.ServerNotification
+    ) -> None:
+        """Handle notifications from the server."""
+        match notification.root:
+            case types.LoggingMessageNotification(params=params):
+                await self._logging_callback(params)
+            case _:
+                pass