Skip to content

Commit 8dc5231

Browse files
fix: make UrlElicitationRequiredError pickle-safe
Add __reduce__ method so pickle reconstructs with the correct (elicitations, message) arguments instead of the (code, message, data) tuple stored in Exception.args. Closes #2431 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 5cbd259 commit 8dc5231

File tree

2 files changed

+64
-0
lines changed

2 files changed

+64
-0
lines changed

src/mcp/shared/exceptions.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,10 @@ def elicitations(self) -> list[ElicitRequestURLParams]:
9494
"""The list of URL elicitations required before the request can proceed."""
9595
return self._elicitations
9696

97+
def __reduce__(self) -> tuple[type, tuple[list[ElicitRequestURLParams], str]]:
98+
"""Support pickling by reconstructing with the original __init__ signature."""
99+
return (self.__class__, (self._elicitations, self.error.message))
100+
97101
@classmethod
98102
def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError:
99103
"""Reconstruct from an ErrorData received over the wire."""

tests/shared/test_exceptions.py

Lines changed: 60 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
@@ -162,3 +164,61 @@ def test_url_elicitation_required_error_exception_message() -> None:
162164

163165
# The exception's string representation should match the message
164166
assert str(error) == "URL elicitation required"
167+
168+
169+
def test_mcp_error_pickle_roundtrip() -> None:
170+
"""Test that MCPError survives a pickle round-trip."""
171+
original = MCPError(code=-32600, message="Invalid request", data={"detail": "bad"})
172+
173+
restored = pickle.loads(pickle.dumps(original))
174+
175+
assert type(restored) is MCPError
176+
assert restored.error.code == original.error.code
177+
assert restored.error.message == original.error.message
178+
assert restored.error.data == original.error.data
179+
180+
181+
def test_url_elicitation_required_error_pickle_roundtrip() -> None:
182+
"""Test that UrlElicitationRequiredError survives a pickle round-trip."""
183+
elicitations = [
184+
ElicitRequestURLParams(
185+
mode="url",
186+
message="Auth required",
187+
url="https://example.com/auth",
188+
elicitation_id="test-123",
189+
),
190+
]
191+
original = UrlElicitationRequiredError(elicitations, message="Please authenticate")
192+
193+
restored = pickle.loads(pickle.dumps(original))
194+
195+
assert type(restored) is UrlElicitationRequiredError
196+
assert restored.error.code == URL_ELICITATION_REQUIRED
197+
assert restored.error.message == "Please authenticate"
198+
assert len(restored.elicitations) == 1
199+
assert restored.elicitations[0].elicitation_id == "test-123"
200+
assert restored.elicitations[0].url == "https://example.com/auth"
201+
202+
203+
def test_url_elicitation_required_error_pickle_default_message() -> None:
204+
"""Test pickle round-trip preserves the auto-generated default message."""
205+
elicitations = [
206+
ElicitRequestURLParams(
207+
mode="url",
208+
message="Auth",
209+
url="https://example.com/auth",
210+
elicitation_id="e1",
211+
),
212+
ElicitRequestURLParams(
213+
mode="url",
214+
message="Auth2",
215+
url="https://example.com/auth2",
216+
elicitation_id="e2",
217+
),
218+
]
219+
original = UrlElicitationRequiredError(elicitations)
220+
221+
restored = pickle.loads(pickle.dumps(original))
222+
223+
assert restored.error.message == "URL elicitations required"
224+
assert len(restored.elicitations) == 2

0 commit comments

Comments
 (0)