Skip to content

Commit 96739eb

Browse files
committed
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
1 parent 1abcca2 commit 96739eb

2 files changed

Lines changed: 95 additions & 0 deletions

File tree

src/mcp/shared/exceptions.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,14 @@ def __init__(self, error: ErrorData):
1717
super().__init__(error.message)
1818
self.error = error
1919

20+
def __reduce__(self) -> tuple[Any, ...]:
21+
# `Exception.__init__(error.message)` stores a plain string in
22+
# ``self.args``, so the default pickle path reconstructs by calling
23+
# ``McpError(message_str)`` and crashes because ``__init__`` expects
24+
# ``ErrorData``. Reconstruct from ``self.error`` instead so the typed
25+
# payload survives a pickle round-trip.
26+
return (self.__class__, (self.error,))
27+
2028

2129
class UrlElicitationRequiredError(McpError):
2230
"""
@@ -69,3 +77,12 @@ def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError:
6977
raw_elicitations = cast(list[dict[str, Any]], data.get("elicitations", []))
7078
elicitations = [ElicitRequestURLParams.model_validate(e) for e in raw_elicitations]
7179
return cls(elicitations, error.message)
80+
81+
def __reduce__(self) -> tuple[Any, ...]:
82+
# ``McpError.__init__`` stores the message string in ``self.args``, so
83+
# the default pickle path reconstructs by calling
84+
# ``UrlElicitationRequiredError(message_str)`` — the string ends up
85+
# bound to the ``elicitations`` parameter and unpickling crashes.
86+
# Reconstruct from the typed elicitations and message so the round-trip
87+
# preserves the high-level fields rather than the wire-format dict.
88+
return (self.__class__, (self._elicitations, self.error.message))

tests/shared/test_exceptions.py

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
"""Tests for MCP exception classes."""
22

3+
import pickle
4+
35
import pytest
46

57
from mcp.shared.exceptions import McpError, UrlElicitationRequiredError
@@ -157,3 +159,79 @@ def test_exception_message(self) -> None:
157159

158160
# The exception's string representation should match the message
159161
assert str(error) == "URL elicitation required"
162+
163+
def test_pickle_roundtrip_preserves_elicitations_and_message(self) -> None:
164+
"""Pickling a single-elicitation error reconstructs the typed payload."""
165+
original = UrlElicitationRequiredError(
166+
[
167+
ElicitRequestURLParams(
168+
mode="url",
169+
message="Auth required",
170+
url="https://example.com/auth",
171+
elicitationId="test-123",
172+
)
173+
],
174+
message="Custom auth message",
175+
)
176+
177+
restored = pickle.loads(pickle.dumps(original))
178+
179+
assert isinstance(restored, UrlElicitationRequiredError)
180+
assert isinstance(restored, McpError)
181+
assert restored.error.code == URL_ELICITATION_REQUIRED
182+
assert restored.error.message == "Custom auth message"
183+
assert str(restored) == "Custom auth message"
184+
assert len(restored.elicitations) == 1
185+
assert restored.elicitations[0].elicitationId == "test-123"
186+
assert restored.elicitations[0].url == "https://example.com/auth"
187+
188+
def test_pickle_roundtrip_preserves_multiple_elicitations(self) -> None:
189+
"""Pickling a multi-elicitation error keeps the typed list intact."""
190+
original = UrlElicitationRequiredError(
191+
[
192+
ElicitRequestURLParams(
193+
mode="url",
194+
message="Auth 1",
195+
url="https://example.com/auth1",
196+
elicitationId="test-1",
197+
),
198+
ElicitRequestURLParams(
199+
mode="url",
200+
message="Auth 2",
201+
url="https://example.com/auth2",
202+
elicitationId="test-2",
203+
),
204+
]
205+
)
206+
207+
restored = pickle.loads(pickle.dumps(original))
208+
209+
assert restored.error.message == "URL elicitations required"
210+
assert [e.elicitationId for e in restored.elicitations] == ["test-1", "test-2"]
211+
assert all(isinstance(e, ElicitRequestURLParams) for e in restored.elicitations)
212+
213+
214+
class TestMcpErrorPickle:
215+
"""Pickle round-trip coverage for McpError."""
216+
217+
def test_pickle_roundtrip_preserves_error_data(self) -> None:
218+
"""The ErrorData payload should survive a pickle round-trip intact."""
219+
original = McpError(ErrorData(code=-32600, message="Invalid Request", data={"path": "/foo"}))
220+
221+
restored = pickle.loads(pickle.dumps(original))
222+
223+
assert isinstance(restored, McpError)
224+
assert restored.error.code == -32600
225+
assert restored.error.message == "Invalid Request"
226+
assert restored.error.data == {"path": "/foo"}
227+
assert str(restored) == "Invalid Request"
228+
229+
def test_pickle_roundtrip_without_data_field(self) -> None:
230+
"""An ErrorData with no `data` field should round-trip cleanly."""
231+
original = McpError(ErrorData(code=-32601, message="Method not found"))
232+
233+
restored = pickle.loads(pickle.dumps(original))
234+
235+
assert restored.error.code == -32601
236+
assert restored.error.message == "Method not found"
237+
assert restored.error.data is None

0 commit comments

Comments
 (0)