Skip to content

Commit e49ef71

Browse files
committed
fix: make URL elicitation errors pickle-safe
1 parent 6705402 commit e49ef71

2 files changed

Lines changed: 28 additions & 0 deletions

File tree

src/mcp/shared/exceptions.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
from collections.abc import Callable
34
from typing import Any, cast
45

56
from mcp_types import INVALID_REQUEST, URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError
@@ -117,3 +118,7 @@ def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError:
117118
raw_elicitations = cast(list[dict[str, Any]], data.get("elicitations", []))
118119
elicitations = [ElicitRequestURLParams.model_validate(e) for e in raw_elicitations]
119120
return cls(elicitations, error.message)
121+
122+
def __reduce__(self) -> tuple[Callable[[ErrorData], UrlElicitationRequiredError], tuple[ErrorData]]:
123+
"""Reconstruct the exception from its wire-compatible error payload."""
124+
return (self.from_error, (self.error,))

tests/shared/test_exceptions.py

Lines changed: 23 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
from mcp_types import URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError
57

@@ -117,6 +119,27 @@ def test_url_elicitation_required_error_serialization_roundtrip() -> None:
117119
assert reconstructed.elicitations[0].message == original.elicitations[0].message
118120

119121

122+
def test_url_elicitation_required_error_pickle_roundtrip() -> None:
123+
"""Pickle reconstruction preserves the specialized exception payload."""
124+
original = UrlElicitationRequiredError(
125+
[
126+
ElicitRequestURLParams(
127+
mode="url",
128+
message="Auth required",
129+
url="https://example.com/auth",
130+
elicitation_id="test-123",
131+
)
132+
],
133+
message="Please authenticate",
134+
)
135+
136+
restored = pickle.loads(pickle.dumps(original))
137+
138+
assert isinstance(restored, UrlElicitationRequiredError)
139+
assert restored.error == original.error
140+
assert restored.elicitations == original.elicitations
141+
142+
120143
def test_url_elicitation_required_error_data_contains_elicitations() -> None:
121144
"""Test that error data contains properly serialized elicitations."""
122145
elicitation = ElicitRequestURLParams(

0 commit comments

Comments
 (0)