From 96739ebc8061ead02df8173586266539c07e4b45 Mon Sep 17 00:00:00 2001 From: Nishchay Mahor Date: Wed, 27 May 2026 20:57:03 -0700 Subject: [PATCH] fix(shared): make McpError and UrlElicitationRequiredError pickle-safe on v1.x `McpError.__init__(error)` calls `super().__init__(error.message)`, which sets `Exception.args` to `(message_str,)`. The default pickle path then reconstructs by calling `cls(*self.args)`, so unpickling a `McpError` invokes `McpError(message_str)` and crashes with `AttributeError: 'str' object has no attribute 'message'`. `UrlElicitationRequiredError` inherits the same `args` and additionally has a different signature `(elicitations, message=None)`, so its unpickle path also binds the message string to the `elicitations` parameter and errors out. Add `__reduce__` to both classes so pickle reconstructs from the typed `ErrorData` / `_elicitations` fields and the round-trip preserves the high-level payload rather than the raw wire-format args. Existing call sites and the `from_error` constructor are unchanged. A partial fix already landed on `main` for `McpError` (#XXXX) and #2555 is in review for `UrlElicitationRequiredError` on `main`. This change covers the v1.x release branch where both classes are still broken. Github-Issue: #2431 --- src/mcp/shared/exceptions.py | 17 +++++++ tests/shared/test_exceptions.py | 78 +++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/src/mcp/shared/exceptions.py b/src/mcp/shared/exceptions.py index 4943114912..0992ab61cd 100644 --- a/src/mcp/shared/exceptions.py +++ b/src/mcp/shared/exceptions.py @@ -17,6 +17,14 @@ def __init__(self, error: ErrorData): super().__init__(error.message) self.error = error + def __reduce__(self) -> tuple[Any, ...]: + # `Exception.__init__(error.message)` stores a plain string in + # ``self.args``, so the default pickle path reconstructs by calling + # ``McpError(message_str)`` and crashes because ``__init__`` expects + # ``ErrorData``. Reconstruct from ``self.error`` instead so the typed + # payload survives a pickle round-trip. + return (self.__class__, (self.error,)) + class UrlElicitationRequiredError(McpError): """ @@ -69,3 +77,12 @@ def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError: raw_elicitations = cast(list[dict[str, Any]], data.get("elicitations", [])) elicitations = [ElicitRequestURLParams.model_validate(e) for e in raw_elicitations] return cls(elicitations, error.message) + + def __reduce__(self) -> tuple[Any, ...]: + # ``McpError.__init__`` stores the message string in ``self.args``, so + # the default pickle path reconstructs by calling + # ``UrlElicitationRequiredError(message_str)`` — the string ends up + # bound to the ``elicitations`` parameter and unpickling crashes. + # Reconstruct from the typed elicitations and message so the round-trip + # preserves the high-level fields rather than the wire-format dict. + return (self.__class__, (self._elicitations, self.error.message)) diff --git a/tests/shared/test_exceptions.py b/tests/shared/test_exceptions.py index 8845dfe781..d0cb538267 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.shared.exceptions import McpError, UrlElicitationRequiredError @@ -157,3 +159,79 @@ def test_exception_message(self) -> None: # The exception's string representation should match the message assert str(error) == "URL elicitation required" + + def test_pickle_roundtrip_preserves_elicitations_and_message(self) -> None: + """Pickling a single-elicitation error reconstructs the typed payload.""" + original = UrlElicitationRequiredError( + [ + ElicitRequestURLParams( + mode="url", + message="Auth required", + url="https://example.com/auth", + elicitationId="test-123", + ) + ], + message="Custom auth message", + ) + + restored = pickle.loads(pickle.dumps(original)) + + assert isinstance(restored, UrlElicitationRequiredError) + assert isinstance(restored, McpError) + assert restored.error.code == URL_ELICITATION_REQUIRED + assert restored.error.message == "Custom auth message" + assert str(restored) == "Custom auth message" + assert len(restored.elicitations) == 1 + assert restored.elicitations[0].elicitationId == "test-123" + assert restored.elicitations[0].url == "https://example.com/auth" + + def test_pickle_roundtrip_preserves_multiple_elicitations(self) -> None: + """Pickling a multi-elicitation error keeps the typed list intact.""" + original = UrlElicitationRequiredError( + [ + ElicitRequestURLParams( + mode="url", + message="Auth 1", + url="https://example.com/auth1", + elicitationId="test-1", + ), + ElicitRequestURLParams( + mode="url", + message="Auth 2", + url="https://example.com/auth2", + elicitationId="test-2", + ), + ] + ) + + restored = pickle.loads(pickle.dumps(original)) + + assert restored.error.message == "URL elicitations required" + assert [e.elicitationId for e in restored.elicitations] == ["test-1", "test-2"] + assert all(isinstance(e, ElicitRequestURLParams) for e in restored.elicitations) + + +class TestMcpErrorPickle: + """Pickle round-trip coverage for McpError.""" + + def test_pickle_roundtrip_preserves_error_data(self) -> None: + """The ErrorData payload should survive a pickle round-trip intact.""" + original = McpError(ErrorData(code=-32600, message="Invalid Request", data={"path": "/foo"})) + + restored = pickle.loads(pickle.dumps(original)) + + assert isinstance(restored, McpError) + assert restored.error.code == -32600 + assert restored.error.message == "Invalid Request" + assert restored.error.data == {"path": "/foo"} + assert str(restored) == "Invalid Request" + + def test_pickle_roundtrip_without_data_field(self) -> None: + """An ErrorData with no `data` field should round-trip cleanly.""" + original = McpError(ErrorData(code=-32601, message="Method not found")) + + restored = pickle.loads(pickle.dumps(original)) + + assert restored.error.code == -32601 + assert restored.error.message == "Method not found" + assert restored.error.data is None