Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions src/mcp/shared/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,20 @@
from mcp.types import URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError


def _restore_mcp_error(exc_type: type[MCPError], error: ErrorData) -> MCPError:
"""Reconstruct a pickled MCPError or subclass from ErrorData."""
if exc_type is UrlElicitationRequiredError:
return exc_type.from_error(error)

if hasattr(exc_type, "from_error_data"):
return exc_type.from_error_data(error)

restored = exc_type.__new__(exc_type)
Exception.__init__(restored, error.code, error.message, error.data)
restored.error = error
return restored


class MCPError(Exception):
"""Exception type raised when an error arrives over an MCP connection."""

Expand Down Expand Up @@ -40,6 +54,9 @@ def from_error_data(cls, error: ErrorData) -> MCPError:
def __str__(self) -> str:
return self.message

def __reduce__(self) -> tuple[Any, tuple[type[MCPError], ErrorData]]:
return (_restore_mcp_error, (type(self), self.error))


class StatelessModeNotSupported(RuntimeError):
"""Raised when attempting to use a method that is not supported in stateless mode.
Expand Down
40 changes: 40 additions & 0 deletions tests/shared/test_exceptions.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
"""Tests for MCP exception classes."""

import pickle

import pytest

from mcp.shared.exceptions import MCPError, UrlElicitationRequiredError
Expand Down Expand Up @@ -162,3 +164,41 @@ 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_mcp_error_pickle_roundtrip() -> None:
"""Test that MCPError survives a normal pickle round-trip."""
original = MCPError(
code=-32600,
message="Authentication Required",
data={"scope": "files.read"},
)

restored = pickle.loads(pickle.dumps(original))

assert isinstance(restored, MCPError)
assert restored.code == -32600
assert restored.message == "Authentication Required"
assert restored.data == {"scope": "files.read"}
assert str(restored) == "Authentication Required"


def test_url_elicitation_required_error_pickle_roundtrip() -> None:
"""Test that specialized MCPError subclasses survive pickle too."""
original = UrlElicitationRequiredError(
[
ElicitRequestURLParams(
mode="url",
message="Auth required",
url="https://example.com/auth",
elicitation_id="test-123",
)
]
)

restored = pickle.loads(pickle.dumps(original))

assert isinstance(restored, UrlElicitationRequiredError)
assert restored.elicitations[0].elicitation_id == "test-123"
assert restored.elicitations[0].url == "https://example.com/auth"
assert restored.message == "URL elicitation required"
Loading