diff --git a/sgl-model-gateway/src/mcp/manager.rs b/sgl-model-gateway/src/mcp/manager.rs index 295e2fcf0..802fcf4bf 100644 --- a/sgl-model-gateway/src/mcp/manager.rs +++ b/sgl-model-gateway/src/mcp/manager.rs @@ -161,6 +161,39 @@ impl McpManager { .collect() } + /// List tools only from specific servers plus all static servers + /// + /// This method filters tools to only include: + /// 1. Tools from static servers (always visible) + /// 2. Tools from the specified dynamic servers + /// + /// This provides request-scoped tool isolation while maintaining + /// global visibility for static servers. + pub fn list_tools_for_servers(&self, server_keys: &[String]) -> Vec { + self.inventory + .list_tools() + .into_iter() + .filter(|(_tool_name, server_key, _tool_info)| { + // Include if: + // 1. It's a static server (check by name in static_clients) + // 2. It's in the requested servers list + self.is_static_server_by_key(server_key) || server_keys.contains(server_key) + }) + .map(|(_tool_name, _server_key, tool_info)| tool_info) + .collect() + } + + /// Check if a server key belongs to a static server + /// + /// Static servers can be identified by checking if their name + /// exists in the static_clients map. We need to handle the fact + /// that static servers use name as key while dynamic use URL. + fn is_static_server_by_key(&self, server_key: &str) -> bool { + // For static servers, the server_key in inventory is the server name + // Check if this key exists in static_clients + self.static_clients.contains_key(server_key) + } + /// Call a tool by name with automatic type coercion /// /// Accepts either JSON string or parsed Map as arguments. @@ -736,7 +769,7 @@ impl McpManager { } /// Generate a unique key for a server config - fn server_key(config: &McpServerConfig) -> String { + pub fn server_key(config: &McpServerConfig) -> String { // Extract URL from transport or use name match &config.transport { McpTransport::Streamable { url, .. } => url.clone(), diff --git a/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs b/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs index 1f8aa3b28..303b9f533 100644 --- a/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs +++ b/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs @@ -22,11 +22,12 @@ use crate::{ /// Ensure MCP connection succeeds if MCP tools are declared /// /// Checks if request declares MCP tools, and if so, validates that -/// the MCP client can be created and connected. +/// the MCP clients can be created and connected. +/// Returns Ok((has_mcp_tools, server_keys)) on success. pub(crate) async fn ensure_mcp_connection( mcp_manager: &Arc, tools: Option<&[ResponseTool]>, -) -> Result { +) -> Result<(bool, Vec), Response> { let has_mcp_tools = tools .map(|t| { t.iter() @@ -36,23 +37,25 @@ pub(crate) async fn ensure_mcp_connection( if has_mcp_tools { if let Some(tools) = tools { - if ensure_request_mcp_client(mcp_manager, tools) - .await - .is_none() - { - error!( - function = "ensure_mcp_connection", - "Failed to connect to MCP server" - ); - return Err(error::failed_dependency( - "connect_mcp_server_failed", - "Failed to connect to MCP server. Check server_url and authorization.", - )); + match ensure_request_mcp_client(mcp_manager, tools).await { + Some((_manager, server_keys)) => { + return Ok((true, server_keys)); + } + None => { + error!( + function = "ensure_mcp_connection", + "Failed to connect to MCP servers" + ); + return Err(error::failed_dependency( + "connect_mcp_server_failed", + "Failed to connect to MCP servers. Check server_url and authorization.", + )); + } } } } - Ok(has_mcp_tools) + Ok((false, Vec::new())) } /// Validate that workers are available for the requested model diff --git a/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs b/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs index e143c37d5..6b7ab4b74 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs @@ -1,7 +1,5 @@ //! Shared helpers and state tracking for Harmony Responses -use std::sync::Arc; - use axum::response::Response; use serde_json::{from_value, json, to_string, Value}; use tracing::{debug, error, warn}; @@ -10,7 +8,7 @@ use uuid::Uuid; use super::{context::HarmonyResponsesContext, execution::ToolResult}; use crate::{ data_connector::ResponseId, - mcp::McpManager, + mcp, protocols::{ common::{ToolCall, ToolChoice, ToolChoiceValue}, responses::{ @@ -217,10 +215,10 @@ pub(super) fn build_next_request_with_tools( pub(super) fn inject_mcp_metadata( response: &mut ResponsesResponse, tracking: &McpCallTracking, - mcp_manager: &Arc, + mcp_tools: &[mcp::Tool], ) { // Build mcp_list_tools item - let tools = mcp_manager.list_tools(); + let tools = mcp_tools; let tools_info: Vec = tools .iter() .map(|t| McpToolInfo { diff --git a/sgl-model-gateway/src/routers/grpc/harmony/responses/context.rs b/sgl-model-gateway/src/routers/grpc/harmony/responses/context.rs index a35fab7c0..9dee1e367 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/responses/context.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/responses/context.rs @@ -1,6 +1,6 @@ //! Context for Harmony Responses execution -use std::sync::Arc; +use std::sync::{Arc, RwLock as StdRwLock}; use crate::{ data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}, @@ -23,6 +23,9 @@ pub(crate) struct HarmonyResponsesContext { /// MCP manager for tool execution pub mcp_manager: Arc, + /// Server keys for MCP tools requested in this context + pub requested_servers: Arc>>, + /// Response storage for loading conversation history pub response_storage: Arc, @@ -47,6 +50,7 @@ impl HarmonyResponsesContext { pipeline, components, mcp_manager, + requested_servers: Arc::new(StdRwLock::new(Vec::new())), response_storage, conversation_storage, conversation_item_storage, diff --git a/sgl-model-gateway/src/routers/grpc/harmony/responses/non_streaming.rs b/sgl-model-gateway/src/routers/grpc/harmony/responses/non_streaming.rs index c3e0f1e13..b70d75455 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/responses/non_streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/responses/non_streaming.rs @@ -57,9 +57,15 @@ pub(crate) async fn serve_harmony_responses( let current_request = load_previous_messages(ctx, request).await?; // Check MCP connection and get whether MCP tools are present - let has_mcp_tools = + let (has_mcp_tools, server_keys) = ensure_mcp_connection(&ctx.mcp_manager, current_request.tools.as_deref()).await?; + // Set the server keys in the context + { + let mut servers = ctx.requested_servers.write().unwrap(); + *servers = server_keys; + } + let response = if has_mcp_tools { execute_with_mcp_loop(ctx, current_request).await? } else { @@ -96,8 +102,11 @@ async fn execute_with_mcp_loop( // Extract user's max_tool_calls limit (if set) let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize); - // Add static MCP tools from inventory to the request - let mcp_tools = ctx.mcp_manager.list_tools(); + // Add filtered MCP tools (static + requested dynamic) to the request + let mcp_tools = { + let servers = ctx.requested_servers.read().unwrap(); + ctx.mcp_manager.list_tools_for_servers(&servers) + }; if !mcp_tools.is_empty() { let mcp_response_tools = convert_mcp_tools_to_response_tools(&mcp_tools); @@ -216,7 +225,7 @@ async fn execute_with_mcp_loop( // Inject MCP metadata if any calls were executed if mcp_tracking.total_calls() > 0 { - inject_mcp_metadata(&mut response, &mcp_tracking, &ctx.mcp_manager); + inject_mcp_metadata(&mut response, &mcp_tracking, &mcp_tools); } return Ok(response); @@ -258,7 +267,7 @@ async fn execute_with_mcp_loop( // Inject MCP metadata for all executed calls if mcp_tracking.total_calls() > 0 { - inject_mcp_metadata(&mut response, &mcp_tracking, &ctx.mcp_manager); + inject_mcp_metadata(&mut response, &mcp_tracking, &mcp_tools); } return Ok(response); @@ -291,7 +300,7 @@ async fn execute_with_mcp_loop( ); // Inject MCP metadata into final response - inject_mcp_metadata(&mut response, &mcp_tracking, &ctx.mcp_manager); + inject_mcp_metadata(&mut response, &mcp_tracking, &mcp_tools); debug!( mcp_calls = mcp_tracking.total_calls(), diff --git a/sgl-model-gateway/src/routers/grpc/harmony/responses/streaming.rs b/sgl-model-gateway/src/routers/grpc/harmony/responses/streaming.rs index bd326c241..d3e299851 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/responses/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/responses/streaming.rs @@ -47,12 +47,18 @@ pub(crate) async fn serve_harmony_responses_stream( }; // Check MCP connection BEFORE starting stream and get whether MCP tools are present - let has_mcp_tools = + let (has_mcp_tools, server_keys) = match ensure_mcp_connection(&ctx.mcp_manager, current_request.tools.as_deref()).await { - Ok(has_mcp) => has_mcp, + Ok(result) => result, Err(response) => return response, }; + // Set the server keys in the context + { + let mut servers = ctx.requested_servers.write().unwrap(); + *servers = server_keys; + } + // Create SSE channel let (tx, rx) = mpsc::unbounded_channel(); @@ -124,8 +130,11 @@ async fn execute_mcp_tool_loop_streaming( // Extract user's max_tool_calls limit (if set) let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize); - // Add static MCP tools from inventory - let mcp_tools = ctx.mcp_manager.list_tools(); + // Add filtered MCP tools (static + requested dynamic) to the request + let mcp_tools = { + let servers = ctx.requested_servers.read().unwrap(); + ctx.mcp_manager.list_tools_for_servers(&servers) + }; if !mcp_tools.is_empty() { let mcp_response_tools = convert_mcp_tools_to_response_tools(&mcp_tools); let mut all_tools = current_request.tools.clone().unwrap_or_default(); diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs index ecbc99c03..eef6957f5 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs @@ -175,8 +175,9 @@ pub(super) fn generate_mcp_id(prefix: &str) -> String { pub(super) fn build_mcp_list_tools_item( mcp: &Arc, server_label: &str, + server_keys: &[String], ) -> ResponseOutputItem { - let tools = mcp.list_tools(); + let tools = mcp.list_tools_for_servers(server_keys); let tools_info: Vec = tools .iter() .map(|t| McpToolInfo { diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/context.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/context.rs index fc2eb1e3e..c598662a3 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/context.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/context.rs @@ -3,7 +3,10 @@ //! Bundles all dependencies needed by responses handlers to avoid passing //! 10+ parameters to every function. -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::HashMap, + sync::{Arc, RwLock as StdRwLock}, +}; use tokio::{sync::RwLock, task::JoinHandle}; @@ -50,6 +53,9 @@ pub(crate) struct ResponsesContext { /// MCP manager for tool support pub mcp_manager: Arc, + /// Server keys for MCP tools requested in this context + pub requested_servers: Arc>>, + /// Background task handles for cancellation support pub background_tasks: Arc>>, } @@ -71,6 +77,7 @@ impl ResponsesContext { conversation_storage, conversation_item_storage, mcp_manager, + requested_servers: Arc::new(StdRwLock::new(Vec::new())), background_tasks: Arc::new(RwLock::new(HashMap::new())), } } diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs index 30cbd9628..c908139a3 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs @@ -110,12 +110,18 @@ async fn route_responses_streaming( }; // 2. Check MCP connection and get whether MCP tools are present - let has_mcp_tools = + let (has_mcp_tools, server_keys) = match ensure_mcp_connection(&ctx.mcp_manager, request.tools.as_deref()).await { - Ok(has_mcp) => has_mcp, + Ok(result) => result, Err(response) => return response, }; + // Set the server keys in the context + { + let mut servers = ctx.requested_servers.write().unwrap(); + *servers = server_keys; + } + if has_mcp_tools { debug!("MCP tools detected in streaming mode, using streaming tool loop"); diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs index d8a83c204..9bc01ad57 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs @@ -48,7 +48,14 @@ pub(super) async fn route_responses_internal( let modified_request = load_conversation_history(ctx, &request).await?; // 2. Check MCP connection and get whether MCP tools are present - let has_mcp_tools = ensure_mcp_connection(&ctx.mcp_manager, request.tools.as_deref()).await?; + let (has_mcp_tools, server_keys) = + ensure_mcp_connection(&ctx.mcp_manager, request.tools.as_deref()).await?; + + // Set the server keys in the context + { + let mut servers = ctx.requested_servers.write().unwrap(); + *servers = server_keys; + } let responses_response = if has_mcp_tools { debug!("MCP tools detected, using tool loop"); @@ -167,7 +174,10 @@ pub(super) async fn execute_tool_loop( ); // Get MCP tools and convert to chat format (do this once before loop) - let mcp_tools = ctx.mcp_manager.list_tools(); + let mcp_tools = { + let servers = ctx.requested_servers.read().unwrap(); + ctx.mcp_manager.list_tools_for_servers(&servers) + }; let mcp_chat_tools = convert_mcp_tools_to_chat_tools(&mcp_tools); trace!( "Converted {} MCP tools to chat format", @@ -399,7 +409,9 @@ pub(super) async fn execute_tool_loop( // Inject MCP metadata into output if state.total_calls > 0 { // Prepend mcp_list_tools item - let mcp_list_tools = build_mcp_list_tools_item(&ctx.mcp_manager, &server_label); + let servers = ctx.requested_servers.read().unwrap(); + let mcp_list_tools = + build_mcp_list_tools_item(&ctx.mcp_manager, &server_label, &servers); responses_response.output.insert(0, mcp_list_tools); // Append all mcp_call items at the end diff --git a/sgl-model-gateway/src/routers/mcp_utils.rs b/sgl-model-gateway/src/routers/mcp_utils.rs index e388febd5..ddba6f40c 100644 --- a/sgl-model-gateway/src/routers/mcp_utils.rs +++ b/sgl-model-gateway/src/routers/mcp_utils.rs @@ -34,12 +34,16 @@ pub struct McpLoopConfig { /// Maximum iterations as safety limit (default: DEFAULT_MAX_ITERATIONS). /// Prevents infinite loops when max_tool_calls is not set by user. pub max_iterations: usize, + /// Server keys for filtering MCP tools. + /// Contains keys for dynamic servers that were connected for this request. + pub server_keys: Vec, } impl Default for McpLoopConfig { fn default() -> Self { Self { max_iterations: DEFAULT_MAX_ITERATIONS, + server_keys: Vec::new(), } } } @@ -70,67 +74,90 @@ pub fn extract_server_label(tools: Option<&[ResponseTool]>, default_label: &str) // MCP Connection // ============================================================================ -/// Ensure MCP client is connected for request-level MCP tools. +/// Ensure MCP clients are connected for all request-level MCP tools. /// -/// This function extracts MCP server configuration from request tools (server_url, authorization) -/// and ensures a client connection is established via the connection pool. +/// This function extracts MCP server configurations from ALL request tools (server_url, authorization) +/// and ensures client connections are established via the connection pool. /// -/// Returns `Some(())` if a dynamic MCP tool was found and client was created/retrieved, -/// `None` if no MCP tools with server_url were found or connection failed. +/// Returns `Some((manager, server_keys))` if MCP tools were found and clients created, +/// `None` if no MCP tools with server_url were found. pub async fn ensure_request_mcp_client( mcp_manager: &Arc, tools: &[ResponseTool], -) -> Option<()> { - // Find an MCP tool with a server_url - let tool = tools - .iter() - .find(|t| matches!(t.r#type, ResponseToolType::Mcp) && t.server_url.is_some())?; +) -> Option<(Arc, Vec)> { + let mut server_keys = Vec::new(); + let mut has_mcp_tools = false; - let server_url = tool.server_url.as_ref()?.trim().to_string(); + // Process all MCP tools + for tool in tools { + if matches!(tool.r#type, ResponseToolType::Mcp) && tool.server_url.is_some() { + has_mcp_tools = true; + let Some(server_url) = tool.server_url.as_ref().map(|s| s.trim().to_string()) else { + continue; + }; - // Validate URL scheme - if !(server_url.starts_with("http://") || server_url.starts_with("https://")) { - warn!( - "Ignoring MCP server_url with unsupported scheme: {}", - server_url - ); - return None; + // Validate URL scheme + if !(server_url.starts_with("http://") || server_url.starts_with("https://")) { + warn!( + "Ignoring MCP server_url with unsupported scheme: {}", + server_url + ); + continue; + } + + // Extract server label and auth token + let name = tool + .server_label + .clone() + .unwrap_or_else(|| "request-mcp".to_string()); + let token = tool.authorization.clone(); + + // Determine transport type based on URL pattern + let transport = if server_url.contains("/sse") { + McpTransport::Sse { + url: server_url.clone(), + token, + } + } else { + McpTransport::Streamable { + url: server_url.clone(), + token, + } + }; + + // Create server config + let server_config = McpServerConfig { + name, + transport, + proxy: None, + required: false, + }; + + // Get the server key for tracking + let server_key = McpManager::server_key(&server_config); + + // Use get_or_create_client to establish connection + match mcp_manager.get_or_create_client(server_config).await { + Ok(_client) => { + // Track this server for filtering + if !server_keys.contains(&server_key) { + server_keys.push(server_key); + } + } + Err(err) => { + warn!( + "Failed to get/create MCP connection for {}: {}", + server_key, err + ); + // Continue processing other tools + } + } + } } - // Extract server label and auth token - let name = tool - .server_label - .clone() - .unwrap_or_else(|| "request-mcp".to_string()); - let token = tool.authorization.clone(); - - // Determine transport type based on URL pattern - let transport = if server_url.contains("/sse") { - McpTransport::Sse { - url: server_url.clone(), - token, - } + if has_mcp_tools && !server_keys.is_empty() { + Some((mcp_manager.clone(), server_keys)) } else { - McpTransport::Streamable { - url: server_url.clone(), - token, - } - }; - - // Create server config - let server_config = McpServerConfig { - name, - transport, - proxy: None, - required: false, - }; - - // Use get_or_create_client to establish connection - match mcp_manager.get_or_create_client(server_config).await { - Ok(_client) => Some(()), - Err(err) => { - warn!("Failed to get/create MCP connection: {}", err); - None - } + None } } diff --git a/sgl-model-gateway/src/routers/openai/context.rs b/sgl-model-gateway/src/routers/openai/context.rs index 462557fe6..6f3dbebe6 100644 --- a/sgl-model-gateway/src/routers/openai/context.rs +++ b/sgl-model-gateway/src/routers/openai/context.rs @@ -238,6 +238,7 @@ pub struct StreamingEventContext<'a> { pub server_label: &'a str, pub original_request: &'a ResponsesRequest, pub previous_response_id: Option<&'a str>, + pub server_keys: &'a [String], } pub type StreamingRequest = OwnedStreamingContext; diff --git a/sgl-model-gateway/src/routers/openai/responses/mcp.rs b/sgl-model-gateway/src/routers/openai/responses/mcp.rs index 4281f9934..8f2ae7a62 100644 --- a/sgl-model-gateway/src/routers/openai/responses/mcp.rs +++ b/sgl-model-gateway/src/routers/openai/responses/mcp.rs @@ -202,6 +202,7 @@ pub(super) async fn execute_streaming_tool_calls( pub(super) fn prepare_mcp_payload_for_streaming( payload: &mut Value, active_mcp: &Arc, + server_keys: &[String], ) { if let Some(obj) = payload.as_object_mut() { // Remove any non-function tools from outgoing payload @@ -217,7 +218,7 @@ pub(super) fn prepare_mcp_payload_for_streaming( } // Build function tools for all discovered MCP tools - let tools = active_mcp.list_tools(); + let tools = active_mcp.list_tools_for_servers(server_keys); let mut tools_json = Vec::with_capacity(tools.len()); for t in tools { let parameters = Value::Object((*t.input_schema).clone()); @@ -310,8 +311,9 @@ pub(super) fn send_mcp_list_tools_events( server_label: &str, output_index: usize, sequence_number: &mut u64, + server_keys: &[String], ) -> bool { - let tools_item_full = build_mcp_list_tools_item(mcp, server_label); + let tools_item_full = build_mcp_list_tools_item(mcp, server_label, server_keys); let item_id = tools_item_full .get("id") .and_then(|v| v.as_str()) @@ -464,13 +466,14 @@ pub(super) fn inject_mcp_metadata_streaming( state: &ToolLoopState, mcp: &Arc, server_label: &str, + server_keys: &[String], ) { if let Some(output_array) = response.get_mut("output").and_then(|v| v.as_array_mut()) { output_array.retain(|item| { item.get("type").and_then(|t| t.as_str()) != Some(ItemType::MCP_LIST_TOOLS) }); - let list_tools_item = build_mcp_list_tools_item(mcp, server_label); + let list_tools_item = build_mcp_list_tools_item(mcp, server_label, server_keys); output_array.insert(0, list_tools_item); let mcp_call_items = @@ -482,7 +485,7 @@ pub(super) fn inject_mcp_metadata_streaming( } } else if let Some(obj) = response.as_object_mut() { let mut output_items = Vec::new(); - output_items.push(build_mcp_list_tools_item(mcp, server_label)); + output_items.push(build_mcp_list_tools_item(mcp, server_label, server_keys)); output_items.extend(build_executed_mcp_call_items( &state.conversation_history, server_label, @@ -584,6 +587,7 @@ pub(super) async fn execute_tool_loop( "max_tool_calls", active_mcp, original_body, + &config.server_keys, ); } @@ -634,7 +638,8 @@ pub(super) async fn execute_tool_loop( let server_label = extract_server_label(original_body.tools.as_deref(), "mcp"); // Build mcp_list_tools item - let list_tools_item = build_mcp_list_tools_item(active_mcp, &server_label); + let list_tools_item = + build_mcp_list_tools_item(active_mcp, &server_label, &config.server_keys); // Insert at beginning of output array if let Some(output_array) = response_json @@ -668,6 +673,7 @@ pub(super) fn build_incomplete_response( reason: &str, active_mcp: &Arc, original_body: &ResponsesRequest, + server_keys: &[String], ) -> Result { let obj = response .as_object_mut() @@ -712,7 +718,7 @@ pub(super) fn build_incomplete_response( // Add mcp_list_tools and executed mcp_call items at the beginning if state.total_calls > 0 || !mcp_call_items.is_empty() { - let list_tools_item = build_mcp_list_tools_item(active_mcp, &server_label); + let list_tools_item = build_mcp_list_tools_item(active_mcp, &server_label, server_keys); output_array.insert(0, list_tools_item); // Add mcp_call items for executed calls using helper @@ -758,8 +764,12 @@ pub(super) fn build_incomplete_response( // ============================================================================ /// Build a mcp_list_tools output item -pub(super) fn build_mcp_list_tools_item(mcp: &Arc, server_label: &str) -> Value { - let tools = mcp.list_tools(); +pub(super) fn build_mcp_list_tools_item( + mcp: &Arc, + server_label: &str, + server_keys: &[String], +) -> Value { + let tools = mcp.list_tools_for_servers(server_keys); let tools_json: Vec = tools .iter() .map(|t| { diff --git a/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs b/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs index 77e2e3c39..f02d9f3e2 100644 --- a/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs +++ b/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs @@ -50,11 +50,15 @@ pub async fn handle_non_streaming_response(mut ctx: RequestContext) -> Response } }; - if let Some(ref tools) = original_body.tools { - ensure_request_mcp_client(mcp_manager, tools.as_slice()).await; - } + let server_keys = match original_body.tools.as_ref() { + Some(tools) => match ensure_request_mcp_client(mcp_manager, tools.as_slice()).await { + Some((_manager, keys)) => keys, + None => Vec::new(), + }, + None => Vec::new(), + }; - let active_mcp = if mcp_manager.list_tools().is_empty() { + let active_mcp = if mcp_manager.list_tools_for_servers(&server_keys).is_empty() { None } else { Some(mcp_manager) @@ -63,8 +67,11 @@ pub async fn handle_non_streaming_response(mut ctx: RequestContext) -> Response let mut response_json: Value; if let Some(mcp) = active_mcp { - let config = McpLoopConfig::default(); - prepare_mcp_payload_for_streaming(&mut payload, mcp); + let config = McpLoopConfig { + server_keys: server_keys.clone(), + ..McpLoopConfig::default() + }; + prepare_mcp_payload_for_streaming(&mut payload, mcp, &server_keys); match execute_tool_loop( ctx.components.client(), diff --git a/sgl-model-gateway/src/routers/openai/responses/streaming.rs b/sgl-model-gateway/src/routers/openai/responses/streaming.rs index 2b012cbdf..2176d5cf4 100644 --- a/sgl-model-gateway/src/routers/openai/responses/streaming.rs +++ b/sgl-model-gateway/src/routers/openai/responses/streaming.rs @@ -442,7 +442,13 @@ pub(super) fn send_final_response_event( } if let Some(mcp) = active_mcp { - inject_mcp_metadata_streaming(&mut final_response, state, mcp, ctx.server_label); + inject_mcp_metadata_streaming( + &mut final_response, + state, + mcp, + ctx.server_label, + ctx.server_keys, + ); } mask_tools_as_mcp(&mut final_response, ctx.original_request); @@ -632,10 +638,11 @@ pub(super) async fn handle_streaming_with_tool_interception( headers: Option<&HeaderMap>, req: StreamingRequest, active_mcp: &Arc, + server_keys: Vec, ) -> Response { // Transform MCP tools to function tools in payload let mut payload = req.payload; - prepare_mcp_payload_for_streaming(&mut payload, active_mcp); + prepare_mcp_payload_for_streaming(&mut payload, active_mcp, &server_keys); let (tx, rx) = mpsc::unbounded_channel::>(); let should_store = req.original_body.store.unwrap_or(false); @@ -650,11 +657,15 @@ pub(super) async fn handle_streaming_with_tool_interception( let headers_opt = headers.cloned(); let payload_clone = payload.clone(); let active_mcp_clone = Arc::clone(active_mcp); + let server_keys_clone = server_keys.clone(); // Spawn the streaming loop task tokio::spawn(async move { let mut state = ToolLoopState::new(original_request.input.clone()); - let loop_config = McpLoopConfig::default(); + let loop_config = McpLoopConfig { + server_keys: server_keys_clone.clone(), + ..McpLoopConfig::default() + }; let max_tool_calls = original_request.max_tool_calls.map(|n| n as usize); let tools_json = payload_clone.get("tools").cloned().unwrap_or(json!([])); let base_payload = payload_clone.clone(); @@ -680,6 +691,7 @@ pub(super) async fn handle_streaming_with_tool_interception( server_label, original_request: &original_request, previous_response_id: previous_response_id.as_deref(), + server_keys: &server_keys_clone, }; loop { @@ -789,6 +801,7 @@ pub(super) async fn handle_streaming_with_tool_interception( server_label, list_tools_index, &mut sequence_number, + &server_keys_clone, ) { // Client disconnected return; @@ -868,6 +881,7 @@ pub(super) async fn handle_streaming_with_tool_interception( &state, &active_mcp_clone, server_label, + &server_keys_clone, ); mask_tools_as_mcp(&mut response_json, &original_request); @@ -977,11 +991,15 @@ pub async fn handle_streaming_response(ctx: RequestContext) -> Response { let original_body = ctx.responses_request(); let mcp_manager = ctx.components.mcp_manager().expect("MCP manager required"); - if let Some(ref tools) = original_body.tools { - ensure_request_mcp_client(mcp_manager, tools.as_slice()).await; - } + let server_keys = match original_body.tools.as_ref() { + Some(tools) => match ensure_request_mcp_client(mcp_manager, tools.as_slice()).await { + Some((_manager, keys)) => keys, + None => Vec::new(), + }, + None => Vec::new(), + }; - let active_mcp = if mcp_manager.list_tools().is_empty() { + let active_mcp = if mcp_manager.list_tools_for_servers(&server_keys).is_empty() { None } else { Some(mcp_manager.clone()) @@ -1003,5 +1021,12 @@ pub async fn handle_streaming_response(ctx: RequestContext) -> Response { let active_mcp = active_mcp.unwrap(); // MCP is active - transform tools and set up interception - handle_streaming_with_tool_interception(&client, headers.as_ref(), req, &active_mcp).await + handle_streaming_with_tool_interception( + &client, + headers.as_ref(), + req, + &active_mcp, + server_keys, + ) + .await }