From dea56d85878a22ee8e33ed8ba97845252ca546a6 Mon Sep 17 00:00:00 2001
From: Mukunda Katta <mukunda.vjcs6@gmail.com>
Date: Thu, 7 May 2026 10:56:33 -0700
Subject: [PATCH] fix(shared): make UrlElicitationRequiredError pickle-safe
 (refs #2431)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The default Exception unpickle path reconstructs by calling
cls(*self.args), where args was set by the MCPError parent __init__ to
(code, message, data) — three positionals. That doesn't match
UrlElicitationRequiredError's signature (elicitations, message=None=2)
so every unpickle raises:

    TypeError: UrlElicitationRequiredError.__init__() takes from 2 to 3
    positional arguments but 4 were given

Reproduces with stdlib pickle and cloudpickle alike.

Add __reduce__ that reconstructs from the typed (elicitations, message)
pair so the round-trip preserves the high-level fields rather than the
wire-format data dict.

  - regression test for the round-trip on a 2-elicitation error
  - companion test that pins MCPError's (already-working) round-trip so
    a future __init__ refactor can't silently break it

Issue #2431 originally describes a related but slightly different
constructor shape (`McpError(error: ErrorData)`) that no longer exists
on main. The pickle-safety problem in the title still holds — just on
the URL-elicitation subclass — and that's what this fix addresses.
---
 src/mcp/shared/exceptions.py    | 16 ++++++++++
 tests/shared/test_exceptions.py | 54 +++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)

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"}