From 0a937990682024b98eb02867392665f998d7409e Mon Sep 17 00:00:00 2001 From: David Soria Parra Date: Thu, 23 Jan 2025 19:08:23 +0000 Subject: [PATCH] refactor: extract request and notification handling into separate methods --- src/mcp/server/lowlevel/server.py | 122 +++++++++++++++--------------- 1 file changed, 63 insertions(+), 59 deletions(-) diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 32ea279..c3f2abf 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -433,73 +433,77 @@ async def run( match message: case RequestResponder(request=types.ClientRequest(root=req)): - logger.info( - f"Processing request of type {type(req).__name__}" + await self._handle_request( + message, req, session, raise_exceptions ) - if type(req) in self.request_handlers: - handler = self.request_handlers[type(req)] - logger.debug( - f"Dispatching request of type {type(req).__name__}" - ) - - token = None - try: - # Set our global state that can be retrieved via - # app.get_request_context() - token = request_ctx.set( - RequestContext( - message.request_id, - message.request_meta, - session, - ) - ) - response = await handler(req) - except McpError as err: - response = err.error - except Exception as err: - if raise_exceptions: - raise err - response = types.ErrorData( - code=0, message=str(err), data=None - ) - finally: - # Reset the global state after we are done - if token is not None: - request_ctx.reset(token) - - await message.respond(response) - else: - await message.respond( - types.ErrorData( - code=types.METHOD_NOT_FOUND, - message="Method not found", - ) - ) - - logger.debug("Response sent") case types.ClientNotification(root=notify): - if type(notify) in self.notification_handlers: - assert type(notify) in self.notification_handlers - - handler = self.notification_handlers[type(notify)] - logger.debug( - f"Dispatching notification of type " - f"{type(notify).__name__}" - ) - - try: - await handler(notify) - except Exception as err: - logger.error( - f"Uncaught exception in notification handler: " - f"{err}" - ) + await self._handle_notification(notify) for warning in w: logger.info( f"Warning: {warning.category.__name__}: {warning.message}" ) + async def _handle_request( + self, + message: RequestResponder, + req: Any, + session: ServerSession, + raise_exceptions: bool, + ): + logger.info(f"Processing request of type {type(req).__name__}") + if type(req) in self.request_handlers: + handler = self.request_handlers[type(req)] + logger.debug(f"Dispatching request of type {type(req).__name__}") + + token = None + try: + # Set our global state that can be retrieved via + # app.get_request_context() + token = request_ctx.set( + RequestContext( + message.request_id, + message.request_meta, + session, + ) + ) + response = await handler(req) + except McpError as err: + response = err.error + except Exception as err: + if raise_exceptions: + raise err + response = types.ErrorData(code=0, message=str(err), data=None) + finally: + # Reset the global state after we are done + if token is not None: + request_ctx.reset(token) + + await message.respond(response) + else: + await message.respond( + types.ErrorData( + code=types.METHOD_NOT_FOUND, + message="Method not found", + ) + ) + + logger.debug("Response sent") + + async def _handle_notification(self, notify: Any): + if type(notify) in self.notification_handlers: + assert type(notify) in self.notification_handlers + + handler = self.notification_handlers[type(notify)] + logger.debug( + f"Dispatching notification of type " f"{type(notify).__name__}" + ) + + try: + await handler(notify) + except Exception as err: + logger.error(f"Uncaught exception in notification handler: " f"{err}") + async def _ping_handler(request: types.PingRequest) -> types.ServerResult: return types.ServerResult(types.EmptyResult())