fix(examples): Prevent transport from closing due to premature client drop (#324)

Author: kahlive

examples/simple-chat-client/src/bin/simple_chat.rs

@@ -74,7 +74,7 @@ async fn main() -> Result<()> {
             if config.mcp.is_some() {
                 let mcp_clients = config.create_mcp_clients().await?;
 
-                for (name, client) in mcp_clients {
+                for (name, client) in mcp_clients.iter() {
                     println!("load MCP tool: {}", name);
                     let server = client.peer().clone();
                     let tools = get_mcp_tools(server).await?;
@@ -84,6 +84,7 @@ async fn main() -> Result<()> {
                         tool_set.add_tool(tool);
                     }
                 }
+                tool_set.set_clients(mcp_clients);
             }
 
             // create chat session

examples/simple-chat-client/src/tool.rs

@@ -3,8 +3,9 @@ use std::{collections::HashMap, sync::Arc};
 use anyhow::Result;
 use async_trait::async_trait;
 use rmcp::{
+    RoleClient,
     model::{CallToolRequestParam, CallToolResult, Tool as McpTool},
-    service::ServerSink,
+    service::{RunningService, ServerSink},
 };
 use serde_json::Value;
 
@@ -70,9 +71,14 @@ impl Tool for McpToolAdapter {
 #[derive(Default)]
 pub struct ToolSet {
     tools: HashMap<String, Arc<dyn Tool>>,
+    clients: HashMap<String, RunningService<RoleClient, ()>>,
 }
 
 impl ToolSet {
+    pub fn set_clients(&mut self, clients: HashMap<String, RunningService<RoleClient, ()>>) {
+        self.clients = clients;
+    }
+
     pub fn add_tool<T: Tool + 'static>(&mut self, tool: T) {
         self.tools.insert(tool.name(), Arc::new(tool));
     }