fix(lib): refactor MCP client to use a global Tokio runtime for improved resource management

Author: Gioee

src/lib.rs

@@ -15,6 +15,9 @@ use rmcp::{ServiceExt, RoleClient};
 use rmcp::model::{ClientInfo, ClientCapabilities, Implementation};
 use serde_json;
 
+// Global runtime - one runtime per process that stays alive
+static GLOBAL_RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
+
 // Global client instance - one client per process
 static GLOBAL_CLIENT: OnceLock<Mutex<Option<McpClient>>> = OnceLock::new();
 
@@ -237,26 +240,29 @@ type RunningClient = rmcp::service::RunningService<RoleClient, ClientInfo>;
 
 /// Opaque handle for MCP client
 pub struct McpClient {
-    runtime: tokio::runtime::Runtime,
     service: Arc<TokioMutex<Option<RunningClient>>>,
     server_url: Mutex<Option<String>>,
 }
 
+/// Get or create the global Tokio runtime
+fn get_runtime() -> &'static tokio::runtime::Runtime {
+    GLOBAL_RUNTIME.get_or_init(|| {
+        tokio::runtime::Runtime::new().expect("Failed to create Tokio runtime")
+    })
+}
+
 /// Create a new MCP client
 /// Returns NULL on error
 #[no_mangle]
 pub extern "C" fn mcp_client_new() -> *mut McpClient {
-    match tokio::runtime::Runtime::new() {
-        Ok(runtime) => {
-            let client = Box::new(McpClient {
-                runtime,
-                service: Arc::new(TokioMutex::new(None)),
-                server_url: Mutex::new(None),
-            });
-            Box::into_raw(client)
-        }
-        Err(_) => ptr::null_mut(),
-    }
+    // Ensure runtime is initialized
+    let _ = get_runtime();
+
+    let client = Box::new(McpClient {
+        service: Arc::new(TokioMutex::new(None)),
+        server_url: Mutex::new(None),
+    });
+    Box::into_raw(client)
 }
 
 /// Free an MCP client
@@ -323,15 +329,11 @@ pub extern "C" fn mcp_connect(
         }
     };
 
-    // Create a new McpClient with runtime
+    // Ensure global runtime is initialized
+    let _ = get_runtime();
+
+    // Create a new McpClient
     let new_client = McpClient {
-        runtime: match tokio::runtime::Runtime::new() {
-            Ok(r) => r,
-            Err(e) => {
-                let error = format!(r#"{{"error": "Failed to create runtime: {}"}}"#, e);
-                return CString::new(error).unwrap_or_default().into_raw();
-            }
-        },
         service: Arc::new(TokioMutex::new(None)),
         server_url: Mutex::new(None),
     };
@@ -340,7 +342,7 @@ pub extern "C" fn mcp_connect(
 
     let (result, maybe_service) = if use_sse {
         // Use SSE transport (legacy) with optional custom headers
-        new_client.runtime.block_on(async {
+        get_runtime().block_on(async {
             // Create HTTP client with optional custom headers
             let mut client_builder = reqwest::Client::builder();
             if let Some(ref headers_map) = headers_map {
@@ -426,7 +428,7 @@ pub extern "C" fn mcp_connect(
         })
     } else {
         // Use streamable HTTP transport (default) with optional custom headers
-        new_client.runtime.block_on(async {
+        get_runtime().block_on(async {
             // For Streamable HTTP, we need to extract the Authorization header specifically
             // since it has a dedicated field, and we'll use a custom HTTP client for other headers
             let auth_header_value = headers_map.as_ref().and_then(|m| m.get("Authorization")).map(|s| s.clone());
@@ -525,7 +527,7 @@ pub extern "C" fn mcp_connect(
 
     // Store service and URL if connection succeeded
     if let Some((service, url)) = maybe_service {
-        new_client.runtime.block_on(async {
+        get_runtime().block_on(async {
             *new_client.service.lock().await = Some(service);
         });
         *new_client.server_url.lock().unwrap() = Some(url);
@@ -604,7 +606,7 @@ pub extern "C" fn mcp_list_tools_json(_client_ptr: *mut McpClient) -> *mut c_cha
         }
     };
 
-    let result = client.runtime.block_on(async {
+    let result = get_runtime().block_on(async {
         let service_guard = client.service.lock().await;
         let service = match service_guard.as_ref() {
             Some(s) => s,
@@ -698,7 +700,7 @@ pub extern "C" fn mcp_call_tool_json(
         }
     };
 
-    let result = client.runtime.block_on(async {
+    let result = get_runtime().block_on(async {
         let service_guard = client.service.lock().await;
         let service = match service_guard.as_ref() {
             Some(s) => s,
@@ -787,10 +789,8 @@ pub extern "C" fn mcp_list_tools_init() -> usize {
         if let Some(client) = client_opt.as_ref() {
             // Clone the Arc to share the service across async boundaries
             let service_arc = client.service.clone();
-            let runtime_handle = client.runtime.handle().clone();
-
-            // Spawn async task directly on the runtime (like the official rmcp examples)
-            runtime_handle.spawn(async move {
+            // Spawn async task directly on the global runtime (like the official rmcp examples)
+            get_runtime().spawn(async move {
                 let service_guard = service_arc.lock().await;
                 if let Some(service) = service_guard.as_ref() {
                     match service.list_tools(None).await {
@@ -872,10 +872,8 @@ pub extern "C" fn mcp_call_tool_init(tool_name: *const c_char, arguments: *const
         let client_opt = client_mutex.lock().unwrap();
         if let Some(client) = client_opt.as_ref() {
             let service_arc = client.service.clone();
-            let runtime_handle = client.runtime.handle().clone();
-
-            // Spawn async task directly on the runtime (like the official rmcp examples)
-            runtime_handle.spawn(async move {
+            // Spawn async task directly on the global runtime (like the official rmcp examples)
+            get_runtime().spawn(async move {
                 let service_guard = service_arc.lock().await;
                 if let Some(service) = service_guard.as_ref() {
                     // Parse arguments
@@ -975,7 +973,7 @@ pub extern "C" fn mcp_stream_wait(stream_id: usize, timeout_ms: u64) -> *mut Str
         let mut channels = STREAM_CHANNELS.lock().unwrap();
         if let Some(rx) = channels.get_mut(&stream_id) {
             // We need to use block_on for the async recv operation
-            client.runtime.block_on(async {
+            get_runtime().block_on(async {
                 let timeout = tokio::time::Duration::from_millis(timeout_ms);
                 match tokio::time::timeout(timeout, rx.recv()).await {
                     Ok(Some(chunk)) => {