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")
 

tests/interaction/README.md

@@ -40,7 +40,7 @@ flows — with a single subprocess test for stdio.
 ```text
 tests/interaction/
   _requirements.py      the requirements manifest (see below)
-  _helpers.py           shared type aliases + the wire-recording transport
+  _helpers.py           the wire-recording transport
   _connect.py           the transport-parametrized connection factories
   conftest.py           the connect fixture (the transport matrix)
   test_coverage.py      enforces the manifest ↔ test contract

tests/interaction/_helpers.py

@@ -1,28 +1,16 @@
 """Shared helpers for the interaction suite.
 
-Keep this module small: it exists only for (a) types that every test would otherwise have to
-assemble from the SDK's internals to annotate a client callback, and (b) the recording transport
-used by the wire-level tests. Server fixtures and assertion helpers belong in the test that uses
-them.
+Keep this module small: it exists only for the recording transport used by the wire-level
+tests. Server fixtures and assertion helpers belong in the test that uses them.
 """
 
 from types import TracebackType
 
 import anyio
-from mcp_types import ClientResult, ServerNotification, ServerRequest
 from typing_extensions import Self
 
 from mcp.client._transport import ReadStream, Transport, TransportStreams, WriteStream
 from mcp.shared.message import SessionMessage
-from mcp.shared.session import RequestResponder
-
-# TODO: this union is the parameter type of every client message handler (MessageHandlerFnT),
-# but the SDK does not export a name for it -- writing a correctly-typed handler requires
-# importing RequestResponder from mcp.shared.session and assembling the union by hand. It
-# should be a named, exported alias next to MessageHandlerFnT (like ClientRequestContext is
-# for the request callbacks), at which point this alias can be deleted.
-IncomingMessage = RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception
-"""Everything a client message handler can receive."""
 
 
 class _RecordingReadStream:

tests/interaction/lowlevel/test_cancellation.py

@@ -29,13 +29,12 @@ from mcp_types import (
 )
 
 from mcp import MCPError
-from mcp.client import ClientRequestContext, ClientSession
+from mcp.client import ClientRequestContext, ClientSession, IncomingMessage
 from mcp.server import Server, ServerRequestContext
 from mcp.shared.memory import MessageStream, create_client_server_memory_streams
 from mcp.shared.message import SessionMessage
 from tests._stamp import Unstamp
 from tests.interaction._connect import Connect
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/interaction/lowlevel/test_elicitation.py

@@ -28,12 +28,11 @@ from mcp_types import (
 )
 
 from mcp import MCPError, UrlElicitationRequiredError
-from mcp.client import ClientRequestContext, ClientSession
+from mcp.client import ClientRequestContext, ClientSession, IncomingMessage
 from mcp.server import Server, ServerRequestContext
 from mcp.shared.memory import MessageStream, create_client_server_memory_streams
 from mcp.shared.message import SessionMessage
 from tests.interaction._connect import Connect
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/interaction/lowlevel/test_flows.py

@@ -29,12 +29,11 @@ from mcp_types import (
 )
 
 from mcp import MCPError, UrlElicitationRequiredError
-from mcp.client import ClientRequestContext
+from mcp.client import ClientRequestContext, IncomingMessage
 from mcp.server import Server, ServerRequestContext
 from mcp.server.session import ServerSession
 from tests._stamp import Unstamp
 from tests.interaction._connect import Connect
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/interaction/lowlevel/test_list_changed.py

@@ -26,9 +26,9 @@ from mcp_types import (
     ToolListChangedNotification,
 )
 
+from mcp.client import IncomingMessage
 from mcp.server import Server, ServerRequestContext
 from tests.interaction._connect import Connect
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/interaction/lowlevel/test_progress.py

@@ -15,12 +15,12 @@ import pytest
 from inline_snapshot import snapshot
 from mcp_types import CallToolResult, ProgressNotification, ProgressNotificationParams, ProgressToken, TextContent
 
+from mcp.client import IncomingMessage
 from mcp.server import Server, ServerRequestContext
 from mcp.server.session import ServerSession
-from mcp.shared.session import ProgressFnT
+from mcp.shared.dispatcher import ProgressFnT
 from tests._stamp import Unstamp
 from tests.interaction._connect import Connect
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/interaction/lowlevel/test_resources.py

@@ -26,10 +26,10 @@ from mcp_types import (
 )
 
 from mcp import MCPError
+from mcp.client import IncomingMessage
 from mcp.server import Server, ServerRequestContext
 from tests._stamp import Unstamp
 from tests.interaction._connect import Connect
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/interaction/mcpserver/test_context.py

@@ -17,12 +17,11 @@ from mcp_types import (
 from pydantic import BaseModel
 
 from mcp import MCPError
-from mcp.client import ClientRequestContext
+from mcp.client import ClientRequestContext, IncomingMessage
 from mcp.server.elicitation import AcceptedElicitation
 from mcp.server.mcpserver import Context, MCPServer
 from tests._stamp import Unstamp
 from tests.interaction._connect import Connect
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/interaction/mcpserver/test_tools.py

@@ -17,12 +17,12 @@ from mcp_types import (
 from pydantic import BaseModel, Field
 
 from mcp import MCPError
+from mcp.client import IncomingMessage
 from mcp.server.mcpserver import Context, MCPServer
 from mcp.server.mcpserver.exceptions import ToolError
 from mcp.shared.exceptions import UrlElicitationRequiredError
 from tests._stamp import Unstamp
 from tests.interaction._connect import Connect
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/interaction/transports/test_streamable_http.py

@@ -23,12 +23,11 @@ from mcp_types import (
 )
 from pydantic import BaseModel
 
-from mcp.client import ClientRequestContext
+from mcp.client import ClientRequestContext, IncomingMessage
 from mcp.server.elicitation import AcceptedElicitation
 from mcp.server.mcpserver import Context, MCPServer
 from mcp.shared.exceptions import MCPError
 from tests.interaction._connect import connect_over_streamable_http
-from tests.interaction._helpers import IncomingMessage
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio

tests/server/mcpserver/test_integration.py

@@ -14,7 +14,6 @@ import json
 import pytest
 from inline_snapshot import snapshot
 from mcp_types import (
-    ClientResult,
     CreateMessageRequestParams,
     CreateMessageResult,
     ElicitRequestParams,
@@ -30,7 +29,6 @@ from mcp_types import (
     ResourceListChangedNotification,
     ResourceTemplateReference,
     ServerNotification,
-    ServerRequest,
     TextContent,
     TextResourceContents,
     ToolListChangedNotification,
@@ -48,8 +46,7 @@ from examples.snippets.servers import (
     structured_output,
     tool_progress,
 )
-from mcp.client import Client, ClientRequestContext
-from mcp.shared.session import RequestResponder
+from mcp.client import Client, ClientRequestContext, IncomingMessage
 
 pytestmark = pytest.mark.anyio
 
@@ -63,9 +60,7 @@ class NotificationCollector:
         self.resource_notifications: list[NotificationParams | None] = []
         self.tool_notifications: list[NotificationParams | None] = []
 
-    async def handle_generic_notification(
-        self, message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception
-    ) -> None:
+    async def handle_generic_notification(self, message: IncomingMessage) -> None:
         """Handle any server notification and route to appropriate handler."""
         if isinstance(message, ServerNotification):  # pragma: no branch
             if isinstance(message, ProgressNotification):
@@ -180,7 +175,7 @@ async def test_tool_progress() -> None:
     """Test tool progress reporting."""
     collector = NotificationCollector()
 
-    async def message_handler(message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception):
+    async def message_handler(message: IncomingMessage):
         await collector.handle_generic_notification(message)
         if isinstance(message, Exception):  # pragma: no cover
             raise message
@@ -259,7 +254,7 @@ async def test_notifications() -> None:
     """Test notifications and logging functionality."""
     collector = NotificationCollector()
 
-    async def message_handler(message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception):
+    async def message_handler(message: IncomingMessage):
         await collector.handle_generic_notification(message)
         if isinstance(message, Exception):  # pragma: no cover
             raise message

tests/shared/test_streamable_http.py

@@ -43,7 +43,7 @@ from starlette.routing import Mount
 from starlette.types import Message, Scope
 
 from mcp import MCPError
-from mcp.client import ClientRequestContext
+from mcp.client import ClientRequestContext, IncomingMessage
 from mcp.client.session import ClientSession
 from mcp.client.streamable_http import StreamableHTTPTransport, streamable_http_client
 from mcp.server import Server, ServerRequestContext
@@ -64,7 +64,6 @@ from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared._compat import resync_tracer
 from mcp.shared._context_streams import create_context_streams
 from mcp.shared.message import ClientMessageMetadata, ServerMessageMetadata, SessionMessage
-from mcp.shared.session import RequestResponder
 from tests.interaction.transports import StreamingASGITransport
 
 # Test constants
@@ -968,9 +967,7 @@ async def test_streamable_http_client_get_stream(basic_app: Starlette) -> None:
     notifications_received: list[types.ServerNotification] = []
 
     # Define message handler to capture notifications
-    async def message_handler(  # pragma: no branch
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
+    async def message_handler(message: IncomingMessage) -> None:  # pragma: no branch
         if isinstance(message, types.ServerNotification):  # pragma: no branch
             notifications_received.append(message)
 
@@ -1128,9 +1125,7 @@ async def test_streamable_http_client_resumption(event_app: tuple[SimpleEventSto
     first_notification_received = anyio.Event()
     resumption_token_received = anyio.Event()
 
-    async def message_handler(  # pragma: no branch
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
+    async def message_handler(message: IncomingMessage) -> None:  # pragma: no branch
         if isinstance(message, types.ServerNotification):  # pragma: no branch
             captured_notifications.append(message)
             # Look for our first notification
@@ -1798,14 +1793,11 @@ async def test_streamable_http_client_auto_reconnects(
     _, app = event_app
     captured_notifications: list[str] = []
 
-    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 branch
             return  # pragma: no cover
-        if isinstance(message, types.ServerNotification):  # pragma: no branch
-            if isinstance(message, types.LoggingMessageNotification):  # pragma: no branch
-                captured_notifications.append(str(message.params.data))
+        if isinstance(message, types.LoggingMessageNotification):  # pragma: no branch
+            captured_notifications.append(str(message.params.data))
 
     async with (
         make_client(app) as http_client,
@@ -1866,14 +1858,11 @@ async def test_streamable_http_sse_polling_full_cycle(
     _, app = event_app
     all_notifications: list[str] = []
 
-    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 branch
             return  # pragma: no cover
-        if isinstance(message, types.ServerNotification):  # pragma: no branch
-            if isinstance(message, types.LoggingMessageNotification):  # pragma: no branch
-                all_notifications.append(str(message.params.data))
+        if isinstance(message, types.LoggingMessageNotification):  # pragma: no branch
+            all_notifications.append(str(message.params.data))
 
     async with (
         make_client(app) as http_client,
@@ -1909,14 +1898,11 @@ async def test_streamable_http_events_replayed_after_disconnect(
     _, app = event_app
     notification_data: list[str] = []
 
-    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 branch
             return  # pragma: no cover
-        if isinstance(message, types.ServerNotification):  # pragma: no branch
-            if isinstance(message, types.LoggingMessageNotification):  # pragma: no branch
-                notification_data.append(str(message.params.data))
+        if isinstance(message, types.LoggingMessageNotification):  # pragma: no branch
+            notification_data.append(str(message.params.data))
 
     async with (
         make_client(app) as http_client,
@@ -2042,14 +2028,11 @@ async def test_standalone_get_stream_reconnection(event_app: tuple[SimpleEventSt
     _, app = event_app
     received_notifications: list[str] = []
 
-    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):
             return  # pragma: no cover
-        if isinstance(message, types.ServerNotification):  # pragma: no branch
-            if isinstance(message, types.ResourceUpdatedNotification):  # pragma: no branch
-                received_notifications.append(str(message.params.uri))
+        if isinstance(message, types.ResourceUpdatedNotification):  # pragma: no branch
+            received_notifications.append(str(message.params.uri))
 
     async with (
         make_client(app) as http_client,
@@ -2183,9 +2166,7 @@ async def test_standalone_stream_teardown_mid_listen_is_not_an_error(caplog: pyt
     app = Starlette(routes=[Mount("/mcp", app=session_manager.handle_request)])
     notified = anyio.Event()
 
-    async def message_handler(
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
+    async def message_handler(message: IncomingMessage) -> None:
         # Only the standalone-stream notification is teed to the handler here.
         assert isinstance(message, types.ResourceUpdatedNotification)
         notified.set()