aljaz/mcp-python-sdk
1
0
Fork 0
mirror of https://github.com/aljazceru/mcp-python-sdk.git synced 2025-12-19 14:54:24 +01:00
Code Issues Packages Projects Releases Wiki Activity

fix: improve error handling and request cancellation for issue #88

Browse Source
This commit is contained in:
David Soria Parra
2025-02-04 13:58:44 +00:00
parent 827e494df4
commit 08cfbe522a
2 changed files with 41 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
src/mcp/shared/session.py
View File

@@ -1,3 +1,4 @@
import logging
from contextlib import AbstractAsyncContextManager
from datetime import timedelta
from typing import Any, Callable, Generic, TypeVar
@@ -273,19 +274,28 @@ class BaseSession(
await self._incoming_message_stream_writer.send(responder)
elif isinstance(message.root, JSONRPCNotification):
notification = self._receive_notification_type.model_validate(
message.root.model_dump(
by_alias=True, mode="json", exclude_none=True
try:
notification = self._receive_notification_type.model_validate(
message.root.model_dump(
by_alias=True, mode="json", exclude_none=True
)
)
# Handle cancellation notifications
if isinstance(notification.root, CancelledNotification):
cancelled_id = notification.root.params.requestId
if cancelled_id in self._in_flight:
await self._in_flight[cancelled_id].cancel()
else:
await self._received_notification(notification)
await self._incoming_message_stream_writer.send(
notification
)
except Exception as e:
# For other validation errors, log and continue
logging.warning(
f"Failed to validate notification: {e}. "
f"Message was: {message.root}"
)
)
# Handle cancellation notifications
if isinstance(notification.root, CancelledNotification):
cancelled_id = notification.root.params.requestId
if cancelled_id in self._in_flight:
await self._in_flight[cancelled_id].cancel()
else:
await self._received_notification(notification)
await self._incoming_message_stream_writer.send(notification)
else: # Response or error
stream = self._response_streams.pop(message.root.id, None)
if stream: