Improve call to action when disconnected to the mcp server. Try closing the client when terminating the app

mcp_remote_client.py

@@ -142,6 +142,7 @@ class MCPClientWrapper:
 
     async def disconnect(self):
         if self.exit_stack:
+            print("Disconnecting from MCP server...")
             await self.exit_stack.aclose()
             self.exit_stack = None
             self.session = None
@@ -154,7 +155,7 @@ class MCPClientWrapper:
                 {"role": "user", "content": message},
                 {
                     "role": "assistant",
-                    "content": "Please connect to an MCP server first.",
+                    "content": "Please connect to an MCP server first by reloading the page.",
                 },
             ]
             yield messages, gr.Textbox(value="")
@@ -423,60 +424,68 @@ def gradio_interface(
 
     # Disable auto-dark mode by setting theme to None
     with gr.Blocks(title="WDI MCP Client", css=custom_css, theme=None) as demo:
-        gr.Markdown("# Development Data Chat")
-        # gr.Markdown("Connect to the WDI MCP server and chat with the assistant")
-
-        with gr.Accordion(
-            "Connect to the WDI MCP server and chat with the assistant",
-            open=False,
-            visible=server_path_or_url.endswith(".py"),
-        ):
-            with gr.Row(equal_height=True):
-                with gr.Column(scale=4):
-                    server_path = gr.Textbox(
-                        label="Server Script Path",
-                        placeholder="Enter path to server script (e.g., wdi_mcp_server.py)",
-                        value=server_path_or_url,
-                    )
-                with gr.Column(scale=1):
-                    connect_btn = gr.Button("Connect")
-
-            status = gr.Textbox(label="Connection Status", interactive=False)
-
-        chatbot = gr.Chatbot(
-            value=[],
-            height="81vh",
-            type="messages",
-            show_copy_button=False,
-            avatar_images=("img/small-user.png", "img/small-robot.png"),
-            autoscroll=True,
-            layout="panel",
-            placeholder="Ask development data questions!",
-        )
-
-        with gr.Row(equal_height=True):
-            msg = gr.Textbox(
-                label=None,
-                placeholder="Ask about what indicators are available for a specific topic (e.g., What's the definition of GDP?)",
-                scale=4,
-                show_label=False,
+        try:
+            gr.Markdown("# Development Data Chat")
+            # gr.Markdown("Connect to the WDI MCP server and chat with the assistant")
+
+            with gr.Accordion(
+                "Connect to the WDI MCP server and chat with the assistant",
+                open=False,
+                visible=server_path_or_url.endswith(".py"),
+            ):
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=4):
+                        server_path = gr.Textbox(
+                            label="Server Script Path",
+                            placeholder="Enter path to server script (e.g., wdi_mcp_server.py)",
+                            value=server_path_or_url,
+                        )
+                    with gr.Column(scale=1):
+                        connect_btn = gr.Button("Connect")
+
+                status = gr.Textbox(label="Connection Status", interactive=False)
+
+            chatbot = gr.Chatbot(
+                value=[],
+                height="81vh",
+                type="messages",
+                show_copy_button=False,
+                avatar_images=("img/small-user.png", "img/small-robot.png"),
+                autoscroll=True,
+                layout="panel",
+                placeholder="Ask development data questions!",
             )
-            # clear_btn = gr.Button("Clear Chat", scale=1)
 
-        # 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, show_progress="full"
-        )
+            with gr.Row(equal_height=True):
+                msg = gr.Textbox(
+                    label=None,
+                    placeholder="Ask about what indicators are available for a specific topic (e.g., What's the definition of GDP?)",
+                    scale=4,
+                    show_label=False,
+                )
+                # clear_btn = gr.Button("Clear Chat", scale=1)
+
+            # 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,
+                show_progress="full",
+            )
 
-        msg.submit(
-            client.process_message,
-            [msg, chatbot],
-            [chatbot, msg],
-            concurrency_limit=10,
-        )
-        # clear_btn.click(lambda: [], None, chatbot)
+            msg.submit(
+                client.process_message,
+                [msg, chatbot],
+                [chatbot, msg],
+                concurrency_limit=10,
+            )
+            # clear_btn.click(lambda: [], None, chatbot)
 
+        except KeyboardInterrupt:
+            print("Keyboard interrupt received. Disconnecting from MCP server...")
+            asyncio.run(client.disconnect())
+            raise KeyboardInterrupt
         # demo.unload(client.disconnect)
 
     return demo