fix: add StatelessModeNotSupported exception and improve tests (#1828)

src/mcp/server/session.py

@@ -49,6 +49,7 @@ import mcp.types as types
 from mcp.server.experimental.session_features import ExperimentalServerSessionFeatures
 from mcp.server.models import InitializationOptions
 from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
+from mcp.shared.exceptions import StatelessModeNotSupported
 from mcp.shared.experimental.tasks.capabilities import check_tasks_capability
 from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
@@ -157,26 +158,6 @@ class ServerSession(
 
         return True
 
-    def _require_stateful_mode(self, feature_name: str) -> None:
-        """Raise an error if trying to use a feature that requires stateful mode.
-
-        Server-to-client requests (sampling, elicitation, list_roots) are not
-        supported in stateless HTTP mode because there is no persistent connection
-        for bidirectional communication.
-
-        Args:
-            feature_name: Name of the feature being used (for error message)
-
-        Raises:
-            RuntimeError: If the session is in stateless mode
-        """
-        if self._stateless:
-            raise RuntimeError(
-                f"Cannot use {feature_name} in stateless HTTP mode. "
-                "Stateless mode does not support server-to-client requests. "
-                "Use stateful mode (stateless_http=False) to enable this feature."
-            )
-
     async def _receive_loop(self) -> None:
         async with self._incoming_message_stream_writer:
             await super()._receive_loop()
@@ -332,9 +313,10 @@ class ServerSession(
         Raises:
             McpError: If tools are provided but client doesn't support them.
             ValueError: If tool_use or tool_result message structure is invalid.
-            RuntimeError: If called in stateless HTTP mode.
+            StatelessModeNotSupported: If called in stateless HTTP mode.
         """
-        self._require_stateful_mode("sampling")
+        if self._stateless:
+            raise StatelessModeNotSupported(method="sampling")
         client_caps = self._client_params.capabilities if self._client_params else None
         validate_sampling_tools(client_caps, tools, tool_choice)
         validate_tool_use_result_messages(messages)
@@ -372,7 +354,8 @@ class ServerSession(
 
     async def list_roots(self) -> types.ListRootsResult:
         """Send a roots/list request."""
-        self._require_stateful_mode("list_roots")
+        if self._stateless:
+            raise StatelessModeNotSupported(method="list_roots")
         return await self.send_request(
             types.ServerRequest(types.ListRootsRequest()),
             types.ListRootsResult,
@@ -417,9 +400,10 @@ class ServerSession(
             The client's response with form data
 
         Raises:
-            RuntimeError: If called in stateless HTTP mode.
+            StatelessModeNotSupported: If called in stateless HTTP mode.
         """
-        self._require_stateful_mode("elicitation")
+        if self._stateless:
+            raise StatelessModeNotSupported(method="elicitation")
         return await self.send_request(
             types.ServerRequest(
                 types.ElicitRequest(
@@ -455,9 +439,10 @@ class ServerSession(
             The client's response indicating acceptance, decline, or cancellation
 
         Raises:
-            RuntimeError: If called in stateless HTTP mode.
+            StatelessModeNotSupported: If called in stateless HTTP mode.
         """
-        self._require_stateful_mode("elicitation")
+        if self._stateless:
+            raise StatelessModeNotSupported(method="elicitation")
         return await self.send_request(
             types.ServerRequest(
                 types.ElicitRequest(

src/mcp/shared/exceptions.py

@@ -18,6 +18,24 @@ class McpError(Exception):
         self.error = error
 
 
+class StatelessModeNotSupported(RuntimeError):
+    """
+    Raised when attempting to use a method that is not supported in stateless mode.
+
+    Server-to-client requests (sampling, elicitation, list_roots) are not
+    supported in stateless HTTP mode because there is no persistent connection
+    for bidirectional communication.
+    """
+
+    def __init__(self, method: str):
+        super().__init__(
+            f"Cannot use {method} in stateless HTTP mode. "
+            "Stateless mode does not support server-to-client requests. "
+            "Use stateful mode (stateless_http=False) to enable this feature."
+        )
+        self.method = method
+
+
 class UrlElicitationRequiredError(McpError):
     """
     Specialized error for when a tool requires URL mode elicitation(s) before proceeding.

tests/server/test_stateless_mode.py

@@ -7,47 +7,32 @@ that appropriate errors are raised when attempting to use unsupported features.
 See: https://github.com/modelcontextprotocol/python-sdk/issues/1097
 """
 
+from collections.abc import AsyncGenerator
+from typing import Any
+
 import anyio
 import pytest
 
 import mcp.types as types
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
+from mcp.shared.exceptions import StatelessModeNotSupported
 from mcp.shared.message import SessionMessage
 from mcp.types import ServerCapabilities
 
 
-def create_test_streams():
+@pytest.fixture
-    """Create memory streams for testing."""
+async def stateless_session() -> AsyncGenerator[ServerSession, None]:
+    """Create a stateless ServerSession for testing."""
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
-    return (
-        server_to_client_send,
-        server_to_client_receive,
-        client_to_server_send,
-        client_to_server_receive,
-    )
 
-
+    init_options = InitializationOptions(
-def create_init_options():
-    """Create default initialization options for testing."""
-    return InitializationOptions(
         server_name="test",
         server_version="0.1.0",
         capabilities=ServerCapabilities(),
     )
 
-
-@pytest.mark.anyio
-async def test_list_roots_fails_in_stateless_mode():
-    """Test that list_roots raises RuntimeError in stateless mode."""
-    (
-        server_to_client_send,
-        server_to_client_receive,
-        client_to_server_send,
-        client_to_server_receive,
-    ) = create_test_streams()
-
     async with (
         client_to_server_send,
         client_to_server_receive,
@@ -57,187 +42,100 @@ async def test_list_roots_fails_in_stateless_mode():
         async with ServerSession(
             client_to_server_receive,
             server_to_client_send,
-            create_init_options(),
+            init_options,
             stateless=True,
         ) as session:
-            with pytest.raises(RuntimeError) as exc_info:
+            yield session
-                await session.list_roots()
-
-            assert "stateless HTTP mode" in str(exc_info.value)
-            assert "list_roots" in str(exc_info.value)
 
 
 @pytest.mark.anyio
-async def test_create_message_fails_in_stateless_mode():
+async def test_list_roots_fails_in_stateless_mode(stateless_session: ServerSession):
-    """Test that create_message raises RuntimeError in stateless mode."""
+    """Test that list_roots raises StatelessModeNotSupported in stateless mode."""
-    (
+    with pytest.raises(StatelessModeNotSupported, match="list_roots"):
-        server_to_client_send,
+        await stateless_session.list_roots()
-        server_to_client_receive,
-        client_to_server_send,
-        client_to_server_receive,
-    ) = create_test_streams()
 
-    async with (
+
-        client_to_server_send,
+@pytest.mark.anyio
-        client_to_server_receive,
+async def test_create_message_fails_in_stateless_mode(stateless_session: ServerSession):
-        server_to_client_send,
+    """Test that create_message raises StatelessModeNotSupported in stateless mode."""
-        server_to_client_receive,
+    with pytest.raises(StatelessModeNotSupported, match="sampling"):
-    ):
+        await stateless_session.create_message(
-        async with ServerSession(
+            messages=[
-            client_to_server_receive,
+                types.SamplingMessage(
-            server_to_client_send,
+                    role="user",
-            create_init_options(),
+                    content=types.TextContent(type="text", text="hello"),
-            stateless=True,
-        ) as session:
-            with pytest.raises(RuntimeError) as exc_info:
-                await session.create_message(
-                    messages=[
-                        types.SamplingMessage(
-                            role="user",
-                            content=types.TextContent(type="text", text="hello"),
-                        )
-                    ],
-                    max_tokens=100,
                 )
-
+            ],
-            assert "stateless HTTP mode" in str(exc_info.value)
+            max_tokens=100,
-            assert "sampling" in str(exc_info.value)
+        )
 
 
 @pytest.mark.anyio
-async def test_elicit_form_fails_in_stateless_mode():
+async def test_elicit_form_fails_in_stateless_mode(stateless_session: ServerSession):
-    """Test that elicit_form raises RuntimeError in stateless mode."""
+    """Test that elicit_form raises StatelessModeNotSupported in stateless mode."""
-    (
+    with pytest.raises(StatelessModeNotSupported, match="elicitation"):
-        server_to_client_send,
+        await stateless_session.elicit_form(
-        server_to_client_receive,
+            message="Please provide input",
-        client_to_server_send,
+            requestedSchema={"type": "object", "properties": {}},
-        client_to_server_receive,
+        )
-    ) = create_test_streams()
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-    ):
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            create_init_options(),
-            stateless=True,
-        ) as session:
-            with pytest.raises(RuntimeError) as exc_info:
-                await session.elicit_form(
-                    message="Please provide input",
-                    requestedSchema={"type": "object", "properties": {}},
-                )
-
-            assert "stateless HTTP mode" in str(exc_info.value)
-            assert "elicitation" in str(exc_info.value)
 
 
 @pytest.mark.anyio
-async def test_elicit_url_fails_in_stateless_mode():
+async def test_elicit_url_fails_in_stateless_mode(stateless_session: ServerSession):
-    """Test that elicit_url raises RuntimeError in stateless mode."""
+    """Test that elicit_url raises StatelessModeNotSupported in stateless mode."""
-    (
+    with pytest.raises(StatelessModeNotSupported, match="elicitation"):
-        server_to_client_send,
+        await stateless_session.elicit_url(
-        server_to_client_receive,
+            message="Please authenticate",
-        client_to_server_send,
+            url="https://example.com/auth",
-        client_to_server_receive,
+            elicitation_id="test-123",
-    ) = create_test_streams()
+        )
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-    ):
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            create_init_options(),
-            stateless=True,
-        ) as session:
-            with pytest.raises(RuntimeError) as exc_info:
-                await session.elicit_url(
-                    message="Please authenticate",
-                    url="https://example.com/auth",
-                    elicitation_id="test-123",
-                )
-
-            assert "stateless HTTP mode" in str(exc_info.value)
-            assert "elicitation" in str(exc_info.value)
 
 
 @pytest.mark.anyio
-async def test_elicit_deprecated_fails_in_stateless_mode():
+async def test_elicit_deprecated_fails_in_stateless_mode(stateless_session: ServerSession):
     """Test that the deprecated elicit method also fails in stateless mode."""
-    (
+    with pytest.raises(StatelessModeNotSupported, match="elicitation"):
-        server_to_client_send,
+        await stateless_session.elicit(
-        server_to_client_receive,
+            message="Please provide input",
-        client_to_server_send,
+            requestedSchema={"type": "object", "properties": {}},
-        client_to_server_receive,
+        )
-    ) = create_test_streams()
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-    ):
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            create_init_options(),
-            stateless=True,
-        ) as session:
-            with pytest.raises(RuntimeError) as exc_info:
-                await session.elicit(
-                    message="Please provide input",
-                    requestedSchema={"type": "object", "properties": {}},
-                )
-
-            assert "stateless HTTP mode" in str(exc_info.value)
-            assert "elicitation" in str(exc_info.value)
 
 
 @pytest.mark.anyio
-async def test_require_stateful_mode_does_not_raise_in_stateful_mode():
+async def test_stateless_error_message_is_actionable(stateless_session: ServerSession):
-    """Test that _require_stateful_mode does not raise in stateful mode."""
-    (
-        server_to_client_send,
-        server_to_client_receive,
-        client_to_server_send,
-        client_to_server_receive,
-    ) = create_test_streams()
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-    ):
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            create_init_options(),
-            stateless=False,  # Stateful mode
-        ) as session:
-            # These should not raise - the check passes in stateful mode
-            session._require_stateful_mode("list_roots")
-            session._require_stateful_mode("sampling")
-            session._require_stateful_mode("elicitation")
-
-
-@pytest.mark.anyio
-async def test_stateless_error_message_is_actionable():
     """Test that the error message provides actionable guidance."""
-    (
+    with pytest.raises(StatelessModeNotSupported) as exc_info:
-        server_to_client_send,
+        await stateless_session.list_roots()
-        server_to_client_receive,
+
-        client_to_server_send,
+    error_message = str(exc_info.value)
-        client_to_server_receive,
+    # Should mention it's stateless mode
-    ) = create_test_streams()
+    assert "stateless HTTP mode" in error_message
+    # Should explain why it doesn't work
+    assert "server-to-client requests" in error_message
+    # Should tell user how to fix it
+    assert "stateless_http=False" in error_message
+
+
+@pytest.mark.anyio
+async def test_exception_has_method_attribute(stateless_session: ServerSession):
+    """Test that the exception has a method attribute for programmatic access."""
+    with pytest.raises(StatelessModeNotSupported) as exc_info:
+        await stateless_session.list_roots()
+
+    assert exc_info.value.method == "list_roots"
+
+
+@pytest.fixture
+async def stateful_session() -> AsyncGenerator[ServerSession, None]:
+    """Create a stateful ServerSession for testing."""
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
+
+    init_options = InitializationOptions(
+        server_name="test",
+        server_version="0.1.0",
+        capabilities=ServerCapabilities(),
+    )
 
     async with (
         client_to_server_send,
@@ -248,16 +146,32 @@ async def test_stateless_error_message_is_actionable():
         async with ServerSession(
             client_to_server_receive,
             server_to_client_send,
-            create_init_options(),
+            init_options,
-            stateless=True,
+            stateless=False,
         ) as session:
-            with pytest.raises(RuntimeError) as exc_info:
+            yield session
-                await session.list_roots()
 
-            error_message = str(exc_info.value)
+
-            # Should mention it's stateless mode
+@pytest.mark.anyio
-            assert "stateless HTTP mode" in error_message
+async def test_stateful_mode_does_not_raise_stateless_error(
-            # Should explain why it doesn't work
+    stateful_session: ServerSession, monkeypatch: pytest.MonkeyPatch
-            assert "server-to-client requests" in error_message
+):
-            # Should tell user how to fix it
+    """Test that StatelessModeNotSupported is not raised in stateful mode.
-            assert "stateless_http=False" in error_message
+
+    We mock send_request to avoid blocking on I/O while still verifying
+    that the stateless check passes.
+    """
+    send_request_called = False
+
+    async def mock_send_request(*_: Any, **__: Any) -> types.ListRootsResult:
+        nonlocal send_request_called
+        send_request_called = True
+        return types.ListRootsResult(roots=[])
+
+    monkeypatch.setattr(stateful_session, "send_request", mock_send_request)
+
+    # This should NOT raise StatelessModeNotSupported
+    result = await stateful_session.list_roots()
+
+    assert send_request_called
+    assert isinstance(result, types.ListRootsResult)