catch and send errors on stream

Author: gyang-xai

src/mcp/client/streamable_http.py

@@ -251,6 +251,13 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
                     await event_source.response.aclose()
                     break
 
+    async def _send_error_response(self, ctx: RequestContext, error: Exception) -> None:
+        """Send an error response to the client."""
+        error_data = ErrorData(code=32000, message=str(error))
+        jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=ctx.session_message.message.root.id, error=error_data)
+        session_message = SessionMessage(message=JSONRPCMessage(jsonrpc_error))
+        await ctx.read_stream_writer.send(session_message)
+
     async def _handle_post_request(self, ctx: RequestContext) -> None:
         """Handle a POST request with response processing."""
         headers = self._prepare_request_headers(ctx.headers)
@@ -403,10 +410,13 @@ async def post_writer(
                     )
 
                     async def handle_request_async():
-                        if is_resumption:
-                            await self._handle_resumption_request(ctx)
-                        else:
-                            await self._handle_post_request(ctx)
+                        try:
+                            if is_resumption:
+                                await self._handle_resumption_request(ctx)
+                            else:
+                                await self._handle_post_request(ctx)
+                        except Exception as e:
+                            await self._send_error_response(ctx, e)
 
                     # If this is a request, start a new task to handle it
                     if isinstance(message.root, JSONRPCRequest):
@@ -475,6 +485,9 @@ async def streamablehttp_client(
     read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
 
+    transport._write_stream = write_stream_reader
+    transport._read_stream_writer = read_stream_writer
+
     async with anyio.create_task_group() as tg:
         try:
             logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")