Improve error handling and adjust height

Signed-off-by: Aivin V. Solatorio <[email protected]>

mcp_remote_client.py

@@ -69,54 +69,70 @@ class MCPClientWrapper:
         self.tools = []
 
     async def connect(self, server_path_or_url: str) -> str:
-        # If there's an existing session, close it
-        if self.exit_stack:
-            await self.exit_stack.aclose()
+        try:
+            # If there's an existing session, close it
+            if self.exit_stack:
+                await self.exit_stack.aclose()
 
-        self.exit_stack = AsyncExitStack()
+            self.exit_stack = AsyncExitStack()
 
-        if server_path_or_url.endswith(".py"):
-            command = "python"
+            if server_path_or_url.endswith(".py"):
+                command = "python"
 
-            server_params = StdioServerParameters(
-                command=command,
-                args=[server_path_or_url],
-                env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"},
-            )
+                server_params = StdioServerParameters(
+                    command=command,
+                    args=[server_path_or_url],
+                    env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"},
+                )
 
-            # Launch MCP subprocess and bind streams on the current running loop
-            stdio_transport = await self.exit_stack.enter_async_context(
-                stdio_client(server_params)
-            )
-            self.stdio, self.write = stdio_transport
-        else:
-            sse_transport = await self.exit_stack.enter_async_context(
-                sse_client(
-                    server_path_or_url,
-                    headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
+                print(
+                    f"Starting MCP server with command: {command} {server_path_or_url}"
+                )
+                # Launch MCP subprocess and bind streams on the current running loop
+                stdio_transport = await self.exit_stack.enter_async_context(
+                    stdio_client(server_params)
                 )
+                self.stdio, self.write = stdio_transport
+            else:
+                print(f"Connecting to MCP server at: {server_path_or_url}")
+                sse_transport = await self.exit_stack.enter_async_context(
+                    sse_client(
+                        server_path_or_url,
+                        headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
+                    )
+                )
+                self.stdio, self.write = sse_transport
+
+            print("Creating MCP client session...")
+            # Create ClientSession on this same loop
+            self.session = await self.exit_stack.enter_async_context(
+                ClientSession(self.stdio, self.write)
             )
-            self.stdio, self.write = sse_transport
+            await self.session.initialize()
+            print("MCP session initialized successfully")
 
-        # Create ClientSession on this same loop
-        self.session = await self.exit_stack.enter_async_context(
-            ClientSession(self.stdio, self.write)
-        )
-        await self.session.initialize()
-
-        response = await self.session.list_tools()
-        self.tools = [
-            {
-                "name": tool.name,
-                "description": tool.description,
-                "input_schema": tool.inputSchema,
-            }
-            for tool in response.tools
-        ]
-
-        print("Available tools:", self.tools)
-        tool_names = [tool["name"] for tool in self.tools]
-        return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
+            response = await self.session.list_tools()
+            self.tools = [
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "input_schema": tool.inputSchema,
+                }
+                for tool in response.tools
+            ]
+
+            print("Available tools:", self.tools)
+            tool_names = [tool["name"] for tool in self.tools]
+            return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
+        except Exception as e:
+            error_msg = f"Failed to connect to MCP server: {str(e)}"
+            print(error_msg)
+            # Clean up on error
+            if self.exit_stack:
+                await self.exit_stack.aclose()
+                self.exit_stack = None
+            self.session = None
+            return error_msg
 
     async def process_message(
         self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
@@ -239,7 +255,32 @@ class MCPClientWrapper:
                 yield [result_messages[-1]]
 
                 print(f"Calling tool: {tool_name} with args: {tool_args}")
-                result = await self.session.call_tool(tool_name, tool_args)
+                try:
+                    # Check if session is still valid
+                    if not self.session or not self.stdio or not self.write:
+                        raise Exception(
+                            "MCP session is not connected or has been closed"
+                        )
+
+                    result = await self.session.call_tool(tool_name, tool_args)
+                except Exception as e:
+                    error_msg = f"Error calling tool {tool_name}: {str(e)}"
+                    print(error_msg)
+                    result_messages.append(
+                        {
+                            "role": "assistant",
+                            "content": f"Sorry, I encountered an error while calling the tool: {error_msg}. Please try again.",
+                            "metadata": {
+                                "title": f"Tool Error for {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
+                                "status": "error",
+                                "id": f"error_{tool_name}",
+                            },
+                        }
+                    )
+                    partial_messages.append(result_messages[-1])
+                    yield [result_messages[-1]]
+                    partial_messages = []
+                    continue
 
                 if result_messages and "metadata" in result_messages[-2]:
                     result_messages[-2]["metadata"]["status"] = "done"
@@ -360,7 +401,7 @@ def gradio_interface(
 
         chatbot = gr.Chatbot(
             value=[],
-            height=600,
+            height="75vh",
             type="messages",
             show_copy_button=True,
             avatar_images=("img/small-user.png", "img/small-robot.png"),
@@ -375,9 +416,11 @@ def gradio_interface(
             )
             clear_btn = gr.Button("Clear Chat", scale=1)
 
-        connect_btn.click(client.connect, inputs=server_path, outputs=status)
+        # connect_btn.click(client.connect, inputs=server_path, outputs=status)
         # Automatically call client.connect(...) as soon as the interface loads
-        demo.load(fn=client.connect, inputs=server_path, outputs=status)
+        demo.load(
+            fn=client.connect, inputs=server_path, outputs=status, show_progress="full"
+        )
 
         msg.submit(client.process_message, [msg, chatbot], [chatbot, msg])
         clear_btn.click(lambda: [], None, chatbot)