Spaces:
Running
Running
| from typing import TypeAlias | |
| import mcp.types | |
| from mcp.client.session import MessageHandlerFnT | |
| from mcp.shared.session import RequestResponder | |
| Message: TypeAlias = ( | |
| RequestResponder[mcp.types.ServerRequest, mcp.types.ClientResult] | |
| | mcp.types.ServerNotification | |
| | Exception | |
| ) | |
| MessageHandlerT: TypeAlias = MessageHandlerFnT | |
| class MessageHandler: | |
| """ | |
| This class is used to handle MCP messages sent to the client. It is used to handle all messages, | |
| requests, notifications, and exceptions. Users can override any of the hooks | |
| """ | |
| async def __call__( | |
| self, | |
| message: RequestResponder[mcp.types.ServerRequest, mcp.types.ClientResult] | |
| | mcp.types.ServerNotification | |
| | Exception, | |
| ) -> None: | |
| return await self.dispatch(message) | |
| async def dispatch(self, message: Message) -> None: | |
| # handle all messages | |
| await self.on_message(message) | |
| match message: | |
| # requests | |
| case RequestResponder(): | |
| # handle all requests | |
| await self.on_request(message) | |
| # handle specific requests | |
| match message.request.root: | |
| case mcp.types.PingRequest(): | |
| await self.on_ping(message.request.root) | |
| case mcp.types.ListRootsRequest(): | |
| await self.on_list_roots(message.request.root) | |
| case mcp.types.CreateMessageRequest(): | |
| await self.on_create_message(message.request.root) | |
| # notifications | |
| case mcp.types.ServerNotification(): | |
| # handle all notifications | |
| await self.on_notification(message) | |
| # handle specific notifications | |
| match message.root: | |
| case mcp.types.CancelledNotification(): | |
| await self.on_cancelled(message.root) | |
| case mcp.types.ProgressNotification(): | |
| await self.on_progress(message.root) | |
| case mcp.types.LoggingMessageNotification(): | |
| await self.on_logging_message(message.root) | |
| case mcp.types.ToolListChangedNotification(): | |
| await self.on_tool_list_changed(message.root) | |
| case mcp.types.ResourceListChangedNotification(): | |
| await self.on_resource_list_changed(message.root) | |
| case mcp.types.PromptListChangedNotification(): | |
| await self.on_prompt_list_changed(message.root) | |
| case mcp.types.ResourceUpdatedNotification(): | |
| await self.on_resource_updated(message.root) | |
| case Exception(): | |
| await self.on_exception(message) | |
| async def on_message(self, message: Message) -> None: | |
| pass | |
| async def on_request( | |
| self, message: RequestResponder[mcp.types.ServerRequest, mcp.types.ClientResult] | |
| ) -> None: | |
| pass | |
| async def on_ping(self, message: mcp.types.PingRequest) -> None: | |
| pass | |
| async def on_list_roots(self, message: mcp.types.ListRootsRequest) -> None: | |
| pass | |
| async def on_create_message(self, message: mcp.types.CreateMessageRequest) -> None: | |
| pass | |
| async def on_notification(self, message: mcp.types.ServerNotification) -> None: | |
| pass | |
| async def on_exception(self, message: Exception) -> None: | |
| pass | |
| async def on_progress(self, message: mcp.types.ProgressNotification) -> None: | |
| pass | |
| async def on_logging_message( | |
| self, message: mcp.types.LoggingMessageNotification | |
| ) -> None: | |
| pass | |
| async def on_tool_list_changed( | |
| self, message: mcp.types.ToolListChangedNotification | |
| ) -> None: | |
| pass | |
| async def on_resource_list_changed( | |
| self, message: mcp.types.ResourceListChangedNotification | |
| ) -> None: | |
| pass | |
| async def on_prompt_list_changed( | |
| self, message: mcp.types.PromptListChangedNotification | |
| ) -> None: | |
| pass | |
| async def on_resource_updated( | |
| self, message: mcp.types.ResourceUpdatedNotification | |
| ) -> None: | |
| pass | |
| async def on_cancelled(self, message: mcp.types.CancelledNotification) -> None: | |
| pass | |