diff --git a/src/mcp/shared/exceptions.py b/src/mcp/shared/exceptions.py index c2a7fd44e7..150a075ad0 100644 --- a/src/mcp/shared/exceptions.py +++ b/src/mcp/shared/exceptions.py @@ -51,6 +51,16 @@ def from_error_data(cls, error: ErrorData) -> MCPError: def __str__(self) -> str: return self.message + def __reduce__(self) -> tuple[Any, tuple[type[MCPError], tuple[Any, ...]], dict[str, Any]]: + return (_restore_mcp_error, (type(self), self.args), self.__dict__) + + +def _restore_mcp_error(error_type: type[MCPError], args: tuple[Any, ...]) -> MCPError: + """Reconstruct an MCPError without invoking a subclass constructor.""" + restored = MCPError.__new__(error_type) + Exception.__init__(restored, *args) + return restored + class NoBackChannelError(MCPError): """Raised when a server-initiated request has no channel that can deliver it. diff --git a/tests/shared/test_exceptions.py b/tests/shared/test_exceptions.py index 9da4f606d3..aa11ae515c 100644 --- a/tests/shared/test_exceptions.py +++ b/tests/shared/test_exceptions.py @@ -1,9 +1,11 @@ """Tests for MCP exception classes.""" +import pickle + import pytest from mcp_types import URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError -from mcp.shared.exceptions import MCPError, UrlElicitationRequiredError +from mcp.shared.exceptions import MCPError, NoBackChannelError, UrlElicitationRequiredError def test_url_elicitation_required_error_create_with_single_elicitation() -> None: @@ -173,3 +175,53 @@ def test_from_jsonrpc_error_preserves_code_message_and_data() -> None: ) error = MCPError.from_jsonrpc_error(wire) assert error.error == ErrorData(code=URL_ELICITATION_REQUIRED, message="go elsewhere", data={"hint": "y"}) + + +def test_mcp_error_pickle_roundtrip() -> None: + """MCPError preserves its structured payload across a pickle round-trip.""" + original = MCPError(code=-32600, message="Invalid request", data={"detail": "bad"}) + + restored = pickle.loads(pickle.dumps(original)) + + assert type(restored) is MCPError + assert restored.args == original.args + assert restored.error == original.error + assert str(restored) == str(original) + + +def test_no_back_channel_error_pickle_roundtrip_preserves_method() -> None: + """NoBackChannelError preserves its method and structured payload when pickled.""" + original = NoBackChannelError("sampling/createMessage") + + restored = pickle.loads(pickle.dumps(original)) + + assert isinstance(restored, NoBackChannelError) + assert restored.args == original.args + assert restored.error == original.error + assert restored.method == original.method + + +def test_url_elicitation_required_error_pickle_roundtrip() -> None: + """UrlElicitationRequiredError preserves its typed state when pickled.""" + elicitations = [ + ElicitRequestURLParams( + mode="url", + message="First authorization", + url="https://example.com/auth/first", + elicitation_id="auth-1", + ), + ElicitRequestURLParams( + mode="url", + message="Second authorization", + url="https://example.com/auth/second", + elicitation_id="auth-2", + ), + ] + original = UrlElicitationRequiredError(elicitations, message="Authorization required") + + restored = pickle.loads(pickle.dumps(original)) + + assert type(restored) is UrlElicitationRequiredError + assert restored.args == original.args + assert restored.error == original.error + assert restored.elicitations == original.elicitations