diff --git a/src/mcp/shared/exceptions.py b/src/mcp/shared/exceptions.py index 4943114912..97089c2706 100644 --- a/src/mcp/shared/exceptions.py +++ b/src/mcp/shared/exceptions.py @@ -17,6 +17,10 @@ def __init__(self, error: ErrorData): super().__init__(error.message) self.error = error + def __reduce__(self) -> tuple[type[McpError], tuple[ErrorData]]: + """Reconstruct the exception from its complete wire error payload.""" + return type(self), (self.error,) + class UrlElicitationRequiredError(McpError): """ @@ -69,3 +73,7 @@ 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, tuple[ErrorData]]: + """Reconstruct the specialized exception from its wire error payload.""" + return self.from_error, (self.error,) diff --git a/tests/shared/test_exceptions.py b/tests/shared/test_exceptions.py index 8845dfe781..0889ba4fa8 100644 --- a/tests/shared/test_exceptions.py +++ b/tests/shared/test_exceptions.py @@ -1,11 +1,24 @@ """Tests for MCP exception classes.""" +import pickle + import pytest from mcp.shared.exceptions import McpError, UrlElicitationRequiredError from mcp.types import URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData +def test_mcp_error_pickle_roundtrip() -> None: + """Test that McpError preserves its wire payload through pickle.""" + original = McpError(ErrorData(code=-32600, message="Authentication Required", data={"retry": True})) + + reconstructed = pickle.loads(pickle.dumps(original)) + + assert isinstance(reconstructed, McpError) + assert reconstructed.error == original.error + assert str(reconstructed) == str(original) + + class TestUrlElicitationRequiredError: """Tests for UrlElicitationRequiredError exception class.""" @@ -157,3 +170,24 @@ def test_exception_message(self) -> None: # The exception's string representation should match the message assert str(error) == "URL elicitation required" + + def test_pickle_roundtrip(self) -> None: + """Test that URL elicitation errors preserve their payload through pickle.""" + original = UrlElicitationRequiredError( + [ + ElicitRequestURLParams( + mode="url", + message="Auth required", + url="https://example.com/auth", + elicitationId="test-123", + ) + ], + message="Custom message", + ) + + reconstructed = pickle.loads(pickle.dumps(original)) + + assert isinstance(reconstructed, UrlElicitationRequiredError) + assert reconstructed.error == original.error + assert reconstructed.elicitations == original.elicitations + assert str(reconstructed) == "Custom message"