Narrow message_handler's parameter to notifications and exceptions (#3168)

tests/client/test_client_caching.py

@@ -42,7 +42,7 @@ from mcp_types import (
 )
 from mcp_types.version import LATEST_MODERN_VERSION
 
-from mcp.client import Client
+from mcp.client import Client, IncomingMessage
 from mcp.client._transport import TransportStreams
 from mcp.client.caching import (
     CacheConfig,
@@ -57,13 +57,10 @@ from mcp.server.caching import CacheHint
 from mcp.shared.exceptions import MCPError
 from mcp.shared.memory import MessageStream, create_client_server_memory_streams
 from mcp.shared.message import SessionMessage
-from mcp.shared.session import RequestResponder
 from tests.interaction._connect import BASE_URL, mounted_app
 
 pytestmark = pytest.mark.anyio
 
-IncomingMessage = RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception
-
 
 def _coordinator(client: Client) -> ClientResponseCache:
     cache = client._response_cache

tests/client/test_logging_callback.py

@@ -1,6 +1,5 @@
 from typing import Literal
 
-import mcp_types as types
 import pytest
 from mcp_types import (
     LoggingMessageNotificationParams,
@@ -8,8 +7,8 @@ from mcp_types import (
 )
 
 from mcp import Client
+from mcp.client import IncomingMessage
 from mcp.server.mcpserver import Context, MCPServer
-from mcp.shared.session import RequestResponder
 
 
 class LoggingCollector:
@@ -55,9 +54,7 @@ async def test_logging_callback():
         return True
 
     # Create a message handler to catch exceptions
-    async def message_handler(
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
+    async def message_handler(message: IncomingMessage) -> None:
         if isinstance(message, Exception):  # pragma: no cover
             raise message
 

tests/client/test_notification_response.py

@@ -16,8 +16,8 @@ from starlette.responses import JSONResponse, Response
 from starlette.routing import Route
 
 from mcp import ClientSession, MCPError
+from mcp.client import IncomingMessage
 from mcp.client.streamable_http import streamable_http_client
-from mcp.shared.session import RequestResponder
 
 pytestmark = pytest.mark.anyio
 
@@ -82,9 +82,7 @@ async def test_non_compliant_notification_response() -> None:
     """
     returned_exception = None
 
-    async def message_handler(  # pragma: no cover
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
+    async def message_handler(message: IncomingMessage) -> None:  # pragma: no cover
         nonlocal returned_exception
         if isinstance(message, Exception):
             returned_exception = message

tests/client/test_session.py

@@ -37,7 +37,7 @@ from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, LATEST_HANDSHAKE_VERS
 from pydantic import FileUrl, ValidationError
 
 from mcp import MCPError
-from mcp.client import ClientRequestContext
+from mcp.client import ClientRequestContext, IncomingMessage
 from mcp.client.client import Client
 from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
 from mcp.client.subscriptions import ToolsListChanged, listen
@@ -45,7 +45,6 @@ from mcp.server import Server, ServerRequestContext
 from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair
 from mcp.shared.dispatcher import CallOptions, DispatchContext, OnNotify, OnNotifyIntercept, OnRequest
 from mcp.shared.message import SessionMessage
-from mcp.shared.session import RequestResponder
 from mcp.shared.subscriptions import SUBSCRIPTION_ID_META_KEY
 from mcp.shared.transport_context import TransportContext
 
@@ -123,9 +122,7 @@ async def test_client_session_initialize():
             )
 
     # Create a message handler to catch exceptions
-    async def message_handler(  # pragma: no cover
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
+    async def message_handler(message: IncomingMessage) -> None:  # pragma: no cover
         if isinstance(message, Exception):
             raise message
 
@@ -1228,10 +1225,8 @@ async def test_raising_notification_callbacks_over_direct_dispatch_cost_only_tha
     async def logging_callback(params: types.LoggingMessageNotificationParams) -> None:
         raise ValueError("logging callback boom")
 
-    async def message_handler(
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
-        assert not isinstance(message, RequestResponder | Exception)
+    async def message_handler(message: IncomingMessage) -> None:
+        assert not isinstance(message, Exception)
         teed.append(message)
         raise ValueError("message handler boom")