Merge pull request 'fix: chat hangs after error in mcp sse_reader' (#371) from chat-hangs into main

Reviewed-on: https://git.biggo.com/Funmula/dive-mcp-host/pulls/371

Author: LouisFunmula

dive_mcp_host/host/errors.py

@@ -39,7 +39,9 @@ def __init__(
         self.query = query
         self.state_values = state_values
         self.error = error
-        super().__init__(f"Error in query, {error}")
+        super().__init__(
+            f"Error in query, type: {type(self.error).__name__}, content: {error}"
+        )
 
 
 class GraphNotCompiledError(MCPHostError):

dive_mcp_host/host/tools/mcp_server.py

@@ -39,6 +39,7 @@
 
     from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
     from mcp.shared.message import SessionMessage
+    from mcp.shared.session import RequestResponder
 
     from dive_mcp_host.host.conf import ServerConfig
 
@@ -141,6 +142,28 @@ def __init__(
         else:
             raise InvalidMcpServerError(self.config.name, "Invalid server config")
 
+    async def _message_handler(
+        self,
+        message: RequestResponder[types.ServerRequest, types.ClientResult]
+        | types.ServerNotification
+        | Exception,
+    ) -> None:
+        """Used for handling mcp special responses.
+
+        Such as:
+        - Exception (Literal python exception)
+        - ProgressResult (ServerNotification) ... etc
+        """
+        logger.info(
+            "handling message for %s, type: %s, content: %s",
+            self.name,
+            type(message).__name__,
+            message,
+        )
+
+        if isinstance(message, Exception):
+            raise message
+
     async def _init_tool_info(self, session: ClientSession) -> None:
         """Initialize the session."""
         async with asyncio.timeout(10):
@@ -331,7 +354,9 @@ async def _stdio_client_watcher(self) -> None:  # noqa: C901, PLR0915
                         ),
                         errlog=self._stderr_log_proxy,
                     ) as (stream_read, stream_send, pid),
-                    ClientSession(stream_read, stream_send) as session,
+                    ClientSession(
+                        stream_read, stream_send, message_handler=self._message_handler
+                    ) as session,
                 ):
                     self._session = session
                     self._pid = pid
@@ -363,6 +388,7 @@ async def _stdio_client_watcher(self) -> None:  # noqa: C901, PLR0915
                 httpx.ConnectError,
                 httpx.InvalidURL,
                 httpx.TooManyRedirects,
+                httpx.ConnectTimeout,
             ) as eg:
                 err_msg = (
                     f"Client initialization error for {self.name}: {eg.exceptions}"
@@ -550,7 +576,6 @@ def _http_get_client(
                     key: value.get_secret_value()
                     for key, value in self.config.headers.items()
                 },
-                sse_read_timeout=0.1,
             )
         if self.config.transport == "websocket":
             return websocket_client(
@@ -564,7 +589,7 @@ async def _http_init_client(self) -> None:
         """Initialize the HTTP client."""
         async with (
             self._http_get_client() as streams,
-            ClientSession(*streams) as session,
+            ClientSession(*streams, message_handler=self._message_handler) as session,
         ):
             await self._init_tool_info(session)
 
@@ -639,7 +664,9 @@ async def session_ctx() -> AsyncGenerator[ClientSession, None]:
             """
             async with (
                 self._http_get_client() as streams,
-                ClientSession(*streams) as session,
+                ClientSession(
+                    *streams, message_handler=self._message_handler
+                ) as session,
                 self._session_wrapper(),
             ):
                 await session.initialize()
@@ -750,7 +777,9 @@ async def session_ctx() -> AsyncGenerator[ClientSession, None]:
             """
             async with (
                 self._http_get_client() as streams,
-                ClientSession(*streams) as session,
+                ClientSession(
+                    *streams, message_handler=self._message_handler
+                ) as session,
                 self._session_wrapper(
                     restart_client=lambda e: isinstance(e, httpx.ConnectError)
                 ),