chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[responses API] Add list_tools_for_servers and threading server_keys in routers (#16540)

Browse Source
This commit is contained in:
Chang Su
2026-01-05 22:42:22 -08:00
committed by GitHub
parent 1d7ad4afcc
commit 9bf76c11a9
15 changed files with 265 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/mcp/manager.rs
+34 -1
View File
@@ -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<Tool> {
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(),
sgl-model-gateway/src/routers/grpc/common/responses/utils.rs
+18 -15
View File
@@ -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<McpManager>,
tools: Option<&[ResponseTool]>,
) -> Result<bool, Response> {
) -> Result<(bool, Vec<String>), 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
sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs
+3 -5
View File
@@ -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<McpManager>,
mcp_tools: &[mcp::Tool],
) {
// Build mcp_list_tools item
let tools = mcp_manager.list_tools();
let tools = mcp_tools;
let tools_info: Vec<McpToolInfo> = tools
.iter()
.map(|t| McpToolInfo {
sgl-model-gateway/src/routers/grpc/harmony/responses/context.rs
+5 -1
View File
@@ -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<McpManager>,
/// Server keys for MCP tools requested in this context
pub requested_servers: Arc<StdRwLock<Vec<String>>>,
/// Response storage for loading conversation history
pub response_storage: Arc<dyn ResponseStorage>,
@@ -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,
sgl-model-gateway/src/routers/grpc/harmony/responses/non_streaming.rs
+15 -6
View File
@@ -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(),
sgl-model-gateway/src/routers/grpc/harmony/responses/streaming.rs
+13 -4
View File
@@ -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();
sgl-model-gateway/src/routers/grpc/regular/responses/common.rs
+2 -1
View File
@@ -175,8 +175,9 @@ pub(super) fn generate_mcp_id(prefix: &str) -> String {
pub(super) fn build_mcp_list_tools_item(
mcp: &Arc<McpManager>,
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<McpToolInfo> = tools
.iter()
.map(|t| McpToolInfo {
sgl-model-gateway/src/routers/grpc/regular/responses/context.rs
+8 -1
View File
@@ -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<McpManager>,
/// Server keys for MCP tools requested in this context
pub requested_servers: Arc<StdRwLock<Vec<String>>>,
/// Background task handles for cancellation support
pub background_tasks: Arc<RwLock<HashMap<String, BackgroundTaskInfo>>>,
}
@@ -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())),
}
}
sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs
+8 -2
View File
@@ -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");
sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs
+15 -3
View File
@@ -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
sgl-model-gateway/src/routers/mcp_utils.rs
+79 -52
View File
@@ -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<String>,
}
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<McpManager>,
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<McpManager>, Vec<String>)> {
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
}
}
sgl-model-gateway/src/routers/openai/context.rs
+1
View File
@@ -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;
sgl-model-gateway/src/routers/openai/responses/mcp.rs
+18 -8
View File
@@ -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<mcp::McpManager>,
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<mcp::McpManager>,
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<mcp::McpManager>,
original_body: &ResponsesRequest,
server_keys: &[String],
) -> Result<Value, String> {
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<mcp::McpManager>, server_label: &str) -> Value {
let tools = mcp.list_tools();
pub(super) fn build_mcp_list_tools_item(
mcp: &Arc<mcp::McpManager>,
server_label: &str,
server_keys: &[String],
) -> Value {
let tools = mcp.list_tools_for_servers(server_keys);
let tools_json: Vec<Value> = tools
.iter()
.map(|t| {
sgl-model-gateway/src/routers/openai/responses/non_streaming.rs
+13 -6
View File
@@ -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(),
sgl-model-gateway/src/routers/openai/responses/streaming.rs
+33 -8
View File
@@ -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<crate::mcp::McpManager>,
server_keys: Vec<String>,
) -> 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::<Result<Bytes, io::Error>>();
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
}