tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

fix: allow concurrent tool execution within the same MCP extension (#7202)

Browse Source 
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-02-13 18:10:46 +05:30
committed by GitHub
parent 7c253e6b10
commit ad3bea3685
3 changed files with 45 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/scenario_tests/scenario_runner.rs
+1 -2
View File
@@ -142,7 +142,6 @@ where
F: Fn(&ScenarioResult) -> Result<()>, F: Fn(&ScenarioResult) -> Result<()>,
{ {
use goose::config::ExtensionConfig; use goose::config::ExtensionConfig;
use tokio::sync::Mutex;
goose::agents::moim::SKIP.with(|f| f.set(true)); goose::agents::moim::SKIP.with(|f| f.set(true));
@@ -229,7 +228,7 @@ where
bundled: None, bundled: None,
available_tools: vec![], available_tools: vec![],
}, },
Arc::new(Mutex::new(Box::new(mock_client))), Arc::new(mock_client),
None, None,
None, None,
) )
crates/goose/src/agents/extension_manager.rs
+26 -68
View File
@@ -50,7 +50,7 @@ use rmcp::transport::auth::AuthClient;
use schemars::_private::NoSerialize; use schemars::_private::NoSerialize;
use serde_json::Value; use serde_json::Value;
type McpClientBox = Arc<Mutex<Box<dyn McpClientTrait>>>; type McpClientBox = Arc<dyn McpClientTrait>;
static RE_ENV_BRACES: Lazy<regex::Regex> = static RE_ENV_BRACES: Lazy<regex::Regex> =
Lazy::new(|| regex::Regex::new(r"\$\{\s*([A-Za-z_][A-Za-z0-9_]*)\s*\}").expect("valid regex")); Lazy::new(|| regex::Regex::new(r"\$\{\s*([A-Za-z_][A-Za-z0-9_]*)\s*\}").expect("valid regex"));
@@ -691,7 +691,7 @@ impl ExtensionManager {
let mut extensions = self.extensions.lock().await; let mut extensions = self.extensions.lock().await;
extensions.insert( extensions.insert(
sanitized_name, sanitized_name,
Extension::new(config, Arc::new(Mutex::new(client)), server_info, temp_dir), Extension::new(config, Arc::from(client), server_info, temp_dir),
); );
drop(extensions); drop(extensions);
self.invalidate_tools_cache_and_bump_version().await; self.invalidate_tools_cache_and_bump_version().await;
@@ -862,8 +862,7 @@ impl ExtensionManager {
let ext_name = name.clone(); let ext_name = name.clone();
async move { async move {
let mut tools = Vec::new(); let mut tools = Vec::new();
let client_guard = client.lock().await; let mut client_tools = match client
let mut client_tools = match client_guard
.list_tools(session_id, None, cancel_token.clone()) .list_tools(session_id, None, cancel_token.clone())
.await .await
{ {
@@ -913,7 +912,7 @@ impl ExtensionManager {
break; break;
} }
client_tools = match client_guard client_tools = match client
.list_tools(session_id, client_tools.next_cursor, cancel_token.clone()) .list_tools(session_id, client_tools.next_cursor, cancel_token.clone())
.await .await
{ {
@@ -1065,8 +1064,7 @@ impl ExtensionManager {
.await .await
.ok_or(ErrorData::new(ErrorCode::INVALID_PARAMS, error_msg, None))?; .ok_or(ErrorData::new(ErrorCode::INVALID_PARAMS, error_msg, None))?;
let client_guard = client.lock().await; client
client_guard
.read_resource(session_id, uri, cancellation_token) .read_resource(session_id, uri, cancellation_token)
.await .await
.map_err(|_| { .map_err(|_| {
@@ -1093,9 +1091,7 @@ impl ExtensionManager {
}; };
for (extension_name, client) in extensions_to_check { for (extension_name, client) in extensions_to_check {
let client_guard = client.lock().await; match client
match client_guard
.list_resources(session_id, None, CancellationToken::default()) .list_resources(session_id, None, CancellationToken::default())
.await .await
{ {
@@ -1132,8 +1128,7 @@ impl ExtensionManager {
) )
})?; })?;
let client_guard = client.lock().await; client
client_guard
.list_resources(session_id, None, cancellation_token) .list_resources(session_id, None, cancellation_token)
.await .await
.map_err(|e| { .map_err(|e| {
@@ -1307,7 +1302,7 @@ impl ExtensionManager {
let arguments = tool_call.arguments.clone(); let arguments = tool_call.arguments.clone();
let client = resolved.client.clone(); let client = resolved.client.clone();
let notifications_receiver = client.lock().await.subscribe().await; let notifications_receiver = client.subscribe().await;
let session_id = session_id.to_string(); let session_id = session_id.to_string();
let actual_tool_name = resolved.actual_tool_name; let actual_tool_name = resolved.actual_tool_name;
let working_dir_str = working_dir.map(|p| p.to_string_lossy().to_string()); let working_dir_str = working_dir.map(|p| p.to_string_lossy().to_string());
@@ -1319,8 +1314,7 @@ impl ExtensionManager {
session_id, session_id,
working_dir_str working_dir_str
); );
let client_guard = client.lock().await; client
client_guard
.call_tool( .call_tool(
&session_id, &session_id,
&actual_tool_name, &actual_tool_name,
@@ -1360,8 +1354,7 @@ impl ExtensionManager {
) )
})?; })?;
let client_guard = client.lock().await; client
client_guard
.list_prompts(session_id, None, cancellation_token) .list_prompts(session_id, None, cancellation_token)
.await .await
.map_err(|e| { .map_err(|e| {
@@ -1435,8 +1428,7 @@ impl ExtensionManager {
.await .await
.ok_or_else(|| anyhow::anyhow!("Extension {} not found", extension_name))?; .ok_or_else(|| anyhow::anyhow!("Extension {} not found", extension_name))?;
let client_guard = client.lock().await; client
client_guard
.get_prompt(session_id, name, arguments, cancellation_token) .get_prompt(session_id, name, arguments, cancellation_token)
.await .await
.map_err(|e| anyhow::anyhow!("Failed to get prompt: {}", e)) .map_err(|e| anyhow::anyhow!("Failed to get prompt: {}", e))
@@ -1540,8 +1532,7 @@ impl ExtensionManager {
}; };
for (name, client) in platform_clients { for (name, client) in platform_clients {
let client_guard = client.lock().await; if let Some(moim_content) = client.get_moim(session_id).await {
if let Some(moim_content) = client_guard.get_moim(session_id).await {
tracing::debug!("MOIM content from {}: {} chars", name, moim_content.len()); tracing::debug!("MOIM content from {}: {} chars", name, moim_content.len());
content.push('\n'); content.push('\n');
content.push_str(&moim_content); content.push_str(&moim_content);
@@ -1707,24 +1698,15 @@ mod tests {
// Add some mock clients using the helper method // Add some mock clients using the helper method
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("test_client".to_string(), Arc::new(MockClient {}))
"test_client".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("__cli__ent__".to_string(), Arc::new(MockClient {}))
"__cli__ent__".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("client 🚀".to_string(), Arc::new(MockClient {}))
"client 🚀".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
let tool_call = CallToolRequestParams { let tool_call = CallToolRequestParams {
@@ -1852,7 +1834,7 @@ mod tests {
extension_manager extension_manager
.add_mock_extension_with_tools( .add_mock_extension_with_tools(
"test_extension".to_string(), "test_extension".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))), Arc::new(MockClient {}),
available_tools, available_tools,
) )
.await; .await;
@@ -1882,7 +1864,7 @@ mod tests {
extension_manager extension_manager
.add_mock_extension_with_tools( .add_mock_extension_with_tools(
"test_extension".to_string(), "test_extension".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))), Arc::new(MockClient {}),
vec![], // Empty available_tools means all tools are available by default vec![], // Empty available_tools means all tools are available by default
) )
.await; .await;
@@ -1914,7 +1896,7 @@ mod tests {
extension_manager extension_manager
.add_mock_extension_with_tools( .add_mock_extension_with_tools(
"test_extension".to_string(), "test_extension".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))), Arc::new(MockClient {}),
available_tools, available_tools,
) )
.await; .await;
@@ -2018,10 +2000,7 @@ mod tests {
ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("ext_a".to_string(), Arc::new(MockClient {}))
"ext_a".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
let tools_after_first = extension_manager let tools_after_first = extension_manager
@@ -2036,10 +2015,7 @@ mod tests {
assert!(!tool_names.iter().any(|n| n.starts_with("ext_b__"))); assert!(!tool_names.iter().any(|n| n.starts_with("ext_b__")));
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("ext_b".to_string(), Arc::new(MockClient {}))
"ext_b".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
let tools_after_second = extension_manager let tools_after_second = extension_manager
@@ -2061,16 +2037,10 @@ mod tests {
ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("ext_a".to_string(), Arc::new(MockClient {}))
"ext_a".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("ext_b".to_string(), Arc::new(MockClient {}))
"ext_b".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
let tools_before = extension_manager let tools_before = extension_manager
@@ -2099,16 +2069,10 @@ mod tests {
ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("ext_a".to_string(), Arc::new(MockClient {}))
"ext_a".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("ext_b".to_string(), Arc::new(MockClient {}))
"ext_b".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
let tools = extension_manager let tools = extension_manager
@@ -2128,16 +2092,10 @@ mod tests {
ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("ext_a".to_string(), Arc::new(MockClient {}))
"ext_a".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
extension_manager extension_manager
.add_mock_extension( .add_mock_extension("ext_b".to_string(), Arc::new(MockClient {}))
"ext_b".to_string(),
Arc::new(Mutex::new(Box::new(MockClient {}))),
)
.await; .await;
let tools = extension_manager let tools = extension_manager
crates/goose/src/agents/mcp_client.rs
+18 -33
View File
@@ -106,8 +106,10 @@ pub trait McpClientTrait: Send + Sync {
pub struct GooseClient { pub struct GooseClient {
notification_handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>, notification_handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>,
provider: SharedProvider, provider: SharedProvider,
// Single-slot because calls are serialized per MCP client. /// Fallback session_id for server-initiated callbacks (e.g. sampling/createMessage)
current_session_id: Arc<Mutex<Option<String>>>, /// that don't include the session_id in their MCP extensions metadata.
/// Set once on first request; never cleared (the id is invariant per McpClient).
session_id: Mutex<Option<String>>,
} }
impl GooseClient { impl GooseClient {
@@ -118,23 +120,21 @@ impl GooseClient {
GooseClient { GooseClient {
notification_handlers: handlers, notification_handlers: handlers,
provider, provider,
current_session_id: Arc::new(Mutex::new(None)), session_id: Mutex::new(None),
} }
} }
async fn set_current_session_id(&self, session_id: &str) { async fn set_session_id(&self, session_id: &str) {
let mut slot = self.current_session_id.lock().await; let mut slot = self.session_id.lock().await;
assert!(
slot.as_deref().is_none_or(|s| s == session_id),
"McpClient received requests from different sessions"
);
*slot = Some(session_id.to_string()); *slot = Some(session_id.to_string());
} }
async fn clear_current_session_id(&self) {
let mut slot = self.current_session_id.lock().await;
*slot = None;
}
async fn current_session_id(&self) -> Option<String> { async fn current_session_id(&self) -> Option<String> {
let slot = self.current_session_id.lock().await; self.session_id.lock().await.clone()
slot.clone()
} }
async fn resolve_session_id(&self, extensions: &Extensions) -> Option<String> { async fn resolve_session_id(&self, extensions: &Extensions) -> Option<String> {
@@ -416,31 +416,17 @@ impl McpClient {
cancel_token: CancellationToken, cancel_token: CancellationToken,
) -> Result<ServerResult, Error> { ) -> Result<ServerResult, Error> {
let request = inject_session_context_into_request(request, Some(session_id), working_dir); let request = inject_session_context_into_request(request, Some(session_id), working_dir);
// ExtensionManager serializes calls per MCP connection, so one current_session_id slot // The inner mutex is held only for the send; the actual response wait
// is sufficient for mapping callbacks to the active request session. // happens outside the lock so concurrent calls can overlap.
let handle = { let handle = {
let client = self.client.lock().await; let client = self.client.lock().await;
client.service().set_current_session_id(session_id).await; client.service().set_session_id(session_id).await;
client client
.send_cancellable_request(request, PeerRequestOptions::no_options()) .send_cancellable_request(request, PeerRequestOptions::no_options())
.await .await
}; }?;
let handle = match handle { await_response(handle, self.timeout, &cancel_token).await
Ok(handle) => handle,
Err(err) => {
let client = self.client.lock().await;
client.service().clear_current_session_id().await;
return Err(err);
}
};
let result = await_response(handle, self.timeout, &cancel_token).await;
let client = self.client.lock().await;
client.service().clear_current_session_id().await;
result
} }
} }
@@ -857,8 +843,7 @@ mod tests {
runtime.block_on(async { runtime.block_on(async {
let client = new_client(); let client = new_client();
if let Some(session_id) = current_session { if let Some(session_id) = current_session {
let mut slot = client.current_session_id.lock().await; client.set_session_id(session_id).await;
*slot = Some(session_id.to_string());
} }
let extensions = let extensions =