diff --git a/src/mcp/shared/exceptions.py b/src/mcp/shared/exceptions.py index f153ea319..ec6dc529f 100644 --- a/src/mcp/shared/exceptions.py +++ b/src/mcp/shared/exceptions.py @@ -104,3 +104,19 @@ 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, ...]: + """Make the class pickle-safe. + + Without this, the default ``Exception.__reduce_ex__`` reconstructs by + calling ``cls(*self.args)`` — and ``self.args`` was set by + :class:`MCPError`'s ``__init__`` to ``(code, message, data)``, which + does not match this subclass's constructor signature + ``(elicitations, message=None)``. The mismatch raises + ``TypeError: takes from 2 to 3 positional arguments but 4 were given`` + on every unpickle (closes #2431 for this subclass). + + We reconstruct from the high-level fields so the round-trip preserves + the typed ``elicitations`` list rather than the wire-format data dict. + """ + return (self.__class__, (self._elicitations, self.message)) diff --git a/tests/shared/test_exceptions.py b/tests/shared/test_exceptions.py index 9a7466264..bb5330e22 100644 --- a/tests/shared/test_exceptions.py +++ b/tests/shared/test_exceptions.py @@ -162,3 +162,57 @@ def test_url_elicitation_required_error_exception_message() -> None: # The exception's string representation should match the message assert str(error) == "URL elicitation required" + + +def test_url_elicitation_required_error_pickle_round_trip() -> None: + """UrlElicitationRequiredError must survive pickle round-trips. + + Without __reduce__, the default Exception unpickle path tries to + reconstruct via cls(*self.args), where args is the parent MCPError's + (code, message, data) tuple — wrong arity for this subclass's + (elicitations, message=None) signature. Regression for #2431. + """ + import pickle + + elicitations = [ + ElicitRequestURLParams( + mode="url", + message="Auth required", + url="https://example.com/auth1", + elicitation_id="test-1", + ), + ElicitRequestURLParams( + mode="url", + message="More auth", + url="https://example.com/auth2", + elicitation_id="test-2", + ), + ] + err = UrlElicitationRequiredError(elicitations, message="Two auths needed") + + restored = pickle.loads(pickle.dumps(err)) + + assert isinstance(restored, UrlElicitationRequiredError) + assert str(restored) == "Two auths needed" + assert len(restored.elicitations) == 2 + assert restored.elicitations[0].elicitation_id == "test-1" + assert restored.elicitations[1].elicitation_id == "test-2" + # error.code + error.data are reconstructed by re-running __init__, + # so they should match the original. + assert restored.error.code == URL_ELICITATION_REQUIRED + + +def test_mcp_error_pickle_round_trip() -> None: + """MCPError already pickles via the default Exception mechanism because + its constructor accepts (code, message, data) — same shape as args. + Pinning that contract here so a future __init__ refactor can't silently + break it. + """ + import pickle + + err = MCPError(code=-32603, message="Internal error", data={"context": "x"}) + restored = pickle.loads(pickle.dumps(err)) + assert isinstance(restored, MCPError) + assert restored.code == -32603 + assert restored.message == "Internal error" + assert restored.error.data == {"context": "x"}