From fe31e624b2ceb1484c3ea811b7994edd2789a856 Mon Sep 17 00:00:00 2001 From: ulofiai Date: Fri, 7 Aug 2026 09:41:44 +0800 Subject: [PATCH 1/2] fix(shared): make UrlElicitationRequiredError pickle-safe Signed-off-by: ulofiai --- src/mcp/shared/exceptions.py | 4 ++++ tests/shared/test_exceptions.py | 37 +++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/src/mcp/shared/exceptions.py b/src/mcp/shared/exceptions.py index c2a7fd44e7..2d914650f9 100644 --- a/src/mcp/shared/exceptions.py +++ b/src/mcp/shared/exceptions.py @@ -107,6 +107,10 @@ def elicitations(self) -> list[ElicitRequestURLParams]: """The list of URL elicitations required before the request can proceed.""" return self._elicitations + def __reduce__(self) -> tuple[type, tuple[list[ElicitRequestURLParams], str]]: + """Support pickling by reconstructing with the original constructor signature.""" + return (self.__class__, (self._elicitations, self.message)) + @classmethod def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError: """Reconstruct from an ErrorData received over the wire.""" diff --git a/tests/shared/test_exceptions.py b/tests/shared/test_exceptions.py index 9da4f606d3..c3281963c4 100644 --- a/tests/shared/test_exceptions.py +++ b/tests/shared/test_exceptions.py @@ -1,5 +1,7 @@ """Tests for MCP exception classes.""" +import pickle + import pytest from mcp_types import URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError @@ -173,3 +175,38 @@ 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.error == original.error + + +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.error == original.error + assert restored.elicitations == original.elicitations From dfb911989b1b16670c42bdffc222dfcf1782ddb1 Mon Sep 17 00:00:00 2001 From: ulofiai Date: Fri, 7 Aug 2026 10:02:36 +0800 Subject: [PATCH 2/2] fix(shared): make MCPError subclasses pickle-safe Signed-off-by: ulofiai --- src/mcp/shared/exceptions.py | 14 ++++++++++---- tests/shared/test_exceptions.py | 17 ++++++++++++++++- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/src/mcp/shared/exceptions.py b/src/mcp/shared/exceptions.py index 2d914650f9..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. @@ -107,10 +117,6 @@ def elicitations(self) -> list[ElicitRequestURLParams]: """The list of URL elicitations required before the request can proceed.""" return self._elicitations - def __reduce__(self) -> tuple[type, tuple[list[ElicitRequestURLParams], str]]: - """Support pickling by reconstructing with the original constructor signature.""" - return (self.__class__, (self._elicitations, self.message)) - @classmethod def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError: """Reconstruct from an ErrorData received over the wire.""" diff --git a/tests/shared/test_exceptions.py b/tests/shared/test_exceptions.py index c3281963c4..aa11ae515c 100644 --- a/tests/shared/test_exceptions.py +++ b/tests/shared/test_exceptions.py @@ -5,7 +5,7 @@ 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: @@ -184,7 +184,21 @@ def test_mcp_error_pickle_roundtrip() -> None: 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: @@ -208,5 +222,6 @@ def test_url_elicitation_required_error_pickle_roundtrip() -> None: 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