[responses API] Add list_tools_for_servers and threading server_keys in routers (#16540)
This commit is contained in:
@@ -161,6 +161,39 @@ impl McpManager {
|
|||||||
.collect()
|
.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
|
/// Call a tool by name with automatic type coercion
|
||||||
///
|
///
|
||||||
/// Accepts either JSON string or parsed Map as arguments.
|
/// Accepts either JSON string or parsed Map as arguments.
|
||||||
@@ -736,7 +769,7 @@ impl McpManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Generate a unique key for a server config
|
/// 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
|
// Extract URL from transport or use name
|
||||||
match &config.transport {
|
match &config.transport {
|
||||||
McpTransport::Streamable { url, .. } => url.clone(),
|
McpTransport::Streamable { url, .. } => url.clone(),
|
||||||
|
|||||||
@@ -22,11 +22,12 @@ use crate::{
|
|||||||
/// Ensure MCP connection succeeds if MCP tools are declared
|
/// Ensure MCP connection succeeds if MCP tools are declared
|
||||||
///
|
///
|
||||||
/// Checks if request declares MCP tools, and if so, validates that
|
/// 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(
|
pub(crate) async fn ensure_mcp_connection(
|
||||||
mcp_manager: &Arc<McpManager>,
|
mcp_manager: &Arc<McpManager>,
|
||||||
tools: Option<&[ResponseTool]>,
|
tools: Option<&[ResponseTool]>,
|
||||||
) -> Result<bool, Response> {
|
) -> Result<(bool, Vec<String>), Response> {
|
||||||
let has_mcp_tools = tools
|
let has_mcp_tools = tools
|
||||||
.map(|t| {
|
.map(|t| {
|
||||||
t.iter()
|
t.iter()
|
||||||
@@ -36,23 +37,25 @@ pub(crate) async fn ensure_mcp_connection(
|
|||||||
|
|
||||||
if has_mcp_tools {
|
if has_mcp_tools {
|
||||||
if let Some(tools) = tools {
|
if let Some(tools) = tools {
|
||||||
if ensure_request_mcp_client(mcp_manager, tools)
|
match ensure_request_mcp_client(mcp_manager, tools).await {
|
||||||
.await
|
Some((_manager, server_keys)) => {
|
||||||
.is_none()
|
return Ok((true, server_keys));
|
||||||
{
|
}
|
||||||
|
None => {
|
||||||
error!(
|
error!(
|
||||||
function = "ensure_mcp_connection",
|
function = "ensure_mcp_connection",
|
||||||
"Failed to connect to MCP server"
|
"Failed to connect to MCP servers"
|
||||||
);
|
);
|
||||||
return Err(error::failed_dependency(
|
return Err(error::failed_dependency(
|
||||||
"connect_mcp_server_failed",
|
"connect_mcp_server_failed",
|
||||||
"Failed to connect to MCP server. Check server_url and authorization.",
|
"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
|
/// Validate that workers are available for the requested model
|
||||||
|
|||||||
@@ -1,7 +1,5 @@
|
|||||||
//! Shared helpers and state tracking for Harmony Responses
|
//! Shared helpers and state tracking for Harmony Responses
|
||||||
|
|
||||||
use std::sync::Arc;
|
|
||||||
|
|
||||||
use axum::response::Response;
|
use axum::response::Response;
|
||||||
use serde_json::{from_value, json, to_string, Value};
|
use serde_json::{from_value, json, to_string, Value};
|
||||||
use tracing::{debug, error, warn};
|
use tracing::{debug, error, warn};
|
||||||
@@ -10,7 +8,7 @@ use uuid::Uuid;
|
|||||||
use super::{context::HarmonyResponsesContext, execution::ToolResult};
|
use super::{context::HarmonyResponsesContext, execution::ToolResult};
|
||||||
use crate::{
|
use crate::{
|
||||||
data_connector::ResponseId,
|
data_connector::ResponseId,
|
||||||
mcp::McpManager,
|
mcp,
|
||||||
protocols::{
|
protocols::{
|
||||||
common::{ToolCall, ToolChoice, ToolChoiceValue},
|
common::{ToolCall, ToolChoice, ToolChoiceValue},
|
||||||
responses::{
|
responses::{
|
||||||
@@ -217,10 +215,10 @@ pub(super) fn build_next_request_with_tools(
|
|||||||
pub(super) fn inject_mcp_metadata(
|
pub(super) fn inject_mcp_metadata(
|
||||||
response: &mut ResponsesResponse,
|
response: &mut ResponsesResponse,
|
||||||
tracking: &McpCallTracking,
|
tracking: &McpCallTracking,
|
||||||
mcp_manager: &Arc<McpManager>,
|
mcp_tools: &[mcp::Tool],
|
||||||
) {
|
) {
|
||||||
// Build mcp_list_tools item
|
// Build mcp_list_tools item
|
||||||
let tools = mcp_manager.list_tools();
|
let tools = mcp_tools;
|
||||||
let tools_info: Vec<McpToolInfo> = tools
|
let tools_info: Vec<McpToolInfo> = tools
|
||||||
.iter()
|
.iter()
|
||||||
.map(|t| McpToolInfo {
|
.map(|t| McpToolInfo {
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
//! Context for Harmony Responses execution
|
//! Context for Harmony Responses execution
|
||||||
|
|
||||||
use std::sync::Arc;
|
use std::sync::{Arc, RwLock as StdRwLock};
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage},
|
data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage},
|
||||||
@@ -23,6 +23,9 @@ pub(crate) struct HarmonyResponsesContext {
|
|||||||
/// MCP manager for tool execution
|
/// MCP manager for tool execution
|
||||||
pub mcp_manager: Arc<McpManager>,
|
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
|
/// Response storage for loading conversation history
|
||||||
pub response_storage: Arc<dyn ResponseStorage>,
|
pub response_storage: Arc<dyn ResponseStorage>,
|
||||||
|
|
||||||
@@ -47,6 +50,7 @@ impl HarmonyResponsesContext {
|
|||||||
pipeline,
|
pipeline,
|
||||||
components,
|
components,
|
||||||
mcp_manager,
|
mcp_manager,
|
||||||
|
requested_servers: Arc::new(StdRwLock::new(Vec::new())),
|
||||||
response_storage,
|
response_storage,
|
||||||
conversation_storage,
|
conversation_storage,
|
||||||
conversation_item_storage,
|
conversation_item_storage,
|
||||||
|
|||||||
@@ -57,9 +57,15 @@ pub(crate) async fn serve_harmony_responses(
|
|||||||
let current_request = load_previous_messages(ctx, request).await?;
|
let current_request = load_previous_messages(ctx, request).await?;
|
||||||
|
|
||||||
// Check MCP connection and get whether MCP tools are present
|
// 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?;
|
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 {
|
let response = if has_mcp_tools {
|
||||||
execute_with_mcp_loop(ctx, current_request).await?
|
execute_with_mcp_loop(ctx, current_request).await?
|
||||||
} else {
|
} else {
|
||||||
@@ -96,8 +102,11 @@ async fn execute_with_mcp_loop(
|
|||||||
// Extract user's max_tool_calls limit (if set)
|
// Extract user's max_tool_calls limit (if set)
|
||||||
let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize);
|
let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize);
|
||||||
|
|
||||||
// Add static MCP tools from inventory to the request
|
// Add filtered MCP tools (static + requested dynamic) to the request
|
||||||
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)
|
||||||
|
};
|
||||||
if !mcp_tools.is_empty() {
|
if !mcp_tools.is_empty() {
|
||||||
let mcp_response_tools = convert_mcp_tools_to_response_tools(&mcp_tools);
|
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
|
// Inject MCP metadata if any calls were executed
|
||||||
if mcp_tracking.total_calls() > 0 {
|
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);
|
return Ok(response);
|
||||||
@@ -258,7 +267,7 @@ async fn execute_with_mcp_loop(
|
|||||||
|
|
||||||
// Inject MCP metadata for all executed calls
|
// Inject MCP metadata for all executed calls
|
||||||
if mcp_tracking.total_calls() > 0 {
|
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);
|
return Ok(response);
|
||||||
@@ -291,7 +300,7 @@ async fn execute_with_mcp_loop(
|
|||||||
);
|
);
|
||||||
|
|
||||||
// Inject MCP metadata into final response
|
// 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!(
|
debug!(
|
||||||
mcp_calls = mcp_tracking.total_calls(),
|
mcp_calls = mcp_tracking.total_calls(),
|
||||||
|
|||||||
@@ -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
|
// 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 {
|
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,
|
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
|
// Create SSE channel
|
||||||
let (tx, rx) = mpsc::unbounded_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)
|
// Extract user's max_tool_calls limit (if set)
|
||||||
let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize);
|
let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize);
|
||||||
|
|
||||||
// Add static MCP tools from inventory
|
// Add filtered MCP tools (static + requested dynamic) to the request
|
||||||
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)
|
||||||
|
};
|
||||||
if !mcp_tools.is_empty() {
|
if !mcp_tools.is_empty() {
|
||||||
let mcp_response_tools = convert_mcp_tools_to_response_tools(&mcp_tools);
|
let mcp_response_tools = convert_mcp_tools_to_response_tools(&mcp_tools);
|
||||||
let mut all_tools = current_request.tools.clone().unwrap_or_default();
|
let mut all_tools = current_request.tools.clone().unwrap_or_default();
|
||||||
|
|||||||
@@ -175,8 +175,9 @@ pub(super) fn generate_mcp_id(prefix: &str) -> String {
|
|||||||
pub(super) fn build_mcp_list_tools_item(
|
pub(super) fn build_mcp_list_tools_item(
|
||||||
mcp: &Arc<McpManager>,
|
mcp: &Arc<McpManager>,
|
||||||
server_label: &str,
|
server_label: &str,
|
||||||
|
server_keys: &[String],
|
||||||
) -> ResponseOutputItem {
|
) -> ResponseOutputItem {
|
||||||
let tools = mcp.list_tools();
|
let tools = mcp.list_tools_for_servers(server_keys);
|
||||||
let tools_info: Vec<McpToolInfo> = tools
|
let tools_info: Vec<McpToolInfo> = tools
|
||||||
.iter()
|
.iter()
|
||||||
.map(|t| McpToolInfo {
|
.map(|t| McpToolInfo {
|
||||||
|
|||||||
@@ -3,7 +3,10 @@
|
|||||||
//! Bundles all dependencies needed by responses handlers to avoid passing
|
//! Bundles all dependencies needed by responses handlers to avoid passing
|
||||||
//! 10+ parameters to every function.
|
//! 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};
|
use tokio::{sync::RwLock, task::JoinHandle};
|
||||||
|
|
||||||
@@ -50,6 +53,9 @@ pub(crate) struct ResponsesContext {
|
|||||||
/// MCP manager for tool support
|
/// MCP manager for tool support
|
||||||
pub mcp_manager: Arc<McpManager>,
|
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
|
/// Background task handles for cancellation support
|
||||||
pub background_tasks: Arc<RwLock<HashMap<String, BackgroundTaskInfo>>>,
|
pub background_tasks: Arc<RwLock<HashMap<String, BackgroundTaskInfo>>>,
|
||||||
}
|
}
|
||||||
@@ -71,6 +77,7 @@ impl ResponsesContext {
|
|||||||
conversation_storage,
|
conversation_storage,
|
||||||
conversation_item_storage,
|
conversation_item_storage,
|
||||||
mcp_manager,
|
mcp_manager,
|
||||||
|
requested_servers: Arc::new(StdRwLock::new(Vec::new())),
|
||||||
background_tasks: Arc::new(RwLock::new(HashMap::new())),
|
background_tasks: Arc::new(RwLock::new(HashMap::new())),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -110,12 +110,18 @@ async fn route_responses_streaming(
|
|||||||
};
|
};
|
||||||
|
|
||||||
// 2. Check MCP connection and get whether MCP tools are present
|
// 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 {
|
match ensure_mcp_connection(&ctx.mcp_manager, request.tools.as_deref()).await {
|
||||||
Ok(has_mcp) => has_mcp,
|
Ok(result) => result,
|
||||||
Err(response) => return response,
|
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 {
|
if has_mcp_tools {
|
||||||
debug!("MCP tools detected in streaming mode, using streaming tool loop");
|
debug!("MCP tools detected in streaming mode, using streaming tool loop");
|
||||||
|
|
||||||
|
|||||||
@@ -48,7 +48,14 @@ pub(super) async fn route_responses_internal(
|
|||||||
let modified_request = load_conversation_history(ctx, &request).await?;
|
let modified_request = load_conversation_history(ctx, &request).await?;
|
||||||
|
|
||||||
// 2. Check MCP connection and get whether MCP tools are present
|
// 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 {
|
let responses_response = if has_mcp_tools {
|
||||||
debug!("MCP tools detected, using tool loop");
|
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)
|
// 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);
|
let mcp_chat_tools = convert_mcp_tools_to_chat_tools(&mcp_tools);
|
||||||
trace!(
|
trace!(
|
||||||
"Converted {} MCP tools to chat format",
|
"Converted {} MCP tools to chat format",
|
||||||
@@ -399,7 +409,9 @@ pub(super) async fn execute_tool_loop(
|
|||||||
// Inject MCP metadata into output
|
// Inject MCP metadata into output
|
||||||
if state.total_calls > 0 {
|
if state.total_calls > 0 {
|
||||||
// Prepend mcp_list_tools item
|
// 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);
|
responses_response.output.insert(0, mcp_list_tools);
|
||||||
|
|
||||||
// Append all mcp_call items at the end
|
// Append all mcp_call items at the end
|
||||||
|
|||||||
@@ -34,12 +34,16 @@ pub struct McpLoopConfig {
|
|||||||
/// Maximum iterations as safety limit (default: DEFAULT_MAX_ITERATIONS).
|
/// Maximum iterations as safety limit (default: DEFAULT_MAX_ITERATIONS).
|
||||||
/// Prevents infinite loops when max_tool_calls is not set by user.
|
/// Prevents infinite loops when max_tool_calls is not set by user.
|
||||||
pub max_iterations: usize,
|
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 {
|
impl Default for McpLoopConfig {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
max_iterations: DEFAULT_MAX_ITERATIONS,
|
max_iterations: DEFAULT_MAX_ITERATIONS,
|
||||||
|
server_keys: Vec::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -70,23 +74,27 @@ pub fn extract_server_label(tools: Option<&[ResponseTool]>, default_label: &str)
|
|||||||
// MCP Connection
|
// 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)
|
/// This function extracts MCP server configurations from ALL request tools (server_url, authorization)
|
||||||
/// and ensures a client connection is established via the connection pool.
|
/// and ensures client connections are established via the connection pool.
|
||||||
///
|
///
|
||||||
/// Returns `Some(())` if a dynamic MCP tool was found and client was created/retrieved,
|
/// Returns `Some((manager, server_keys))` if MCP tools were found and clients created,
|
||||||
/// `None` if no MCP tools with server_url were found or connection failed.
|
/// `None` if no MCP tools with server_url were found.
|
||||||
pub async fn ensure_request_mcp_client(
|
pub async fn ensure_request_mcp_client(
|
||||||
mcp_manager: &Arc<McpManager>,
|
mcp_manager: &Arc<McpManager>,
|
||||||
tools: &[ResponseTool],
|
tools: &[ResponseTool],
|
||||||
) -> Option<()> {
|
) -> Option<(Arc<McpManager>, Vec<String>)> {
|
||||||
// Find an MCP tool with a server_url
|
let mut server_keys = Vec::new();
|
||||||
let tool = tools
|
let mut has_mcp_tools = false;
|
||||||
.iter()
|
|
||||||
.find(|t| matches!(t.r#type, ResponseToolType::Mcp) && t.server_url.is_some())?;
|
|
||||||
|
|
||||||
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
|
// Validate URL scheme
|
||||||
if !(server_url.starts_with("http://") || server_url.starts_with("https://")) {
|
if !(server_url.starts_with("http://") || server_url.starts_with("https://")) {
|
||||||
@@ -94,7 +102,7 @@ pub async fn ensure_request_mcp_client(
|
|||||||
"Ignoring MCP server_url with unsupported scheme: {}",
|
"Ignoring MCP server_url with unsupported scheme: {}",
|
||||||
server_url
|
server_url
|
||||||
);
|
);
|
||||||
return None;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Extract server label and auth token
|
// Extract server label and auth token
|
||||||
@@ -125,12 +133,31 @@ pub async fn ensure_request_mcp_client(
|
|||||||
required: false,
|
required: false,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Get the server key for tracking
|
||||||
|
let server_key = McpManager::server_key(&server_config);
|
||||||
|
|
||||||
// Use get_or_create_client to establish connection
|
// Use get_or_create_client to establish connection
|
||||||
match mcp_manager.get_or_create_client(server_config).await {
|
match mcp_manager.get_or_create_client(server_config).await {
|
||||||
Ok(_client) => Some(()),
|
Ok(_client) => {
|
||||||
|
// Track this server for filtering
|
||||||
|
if !server_keys.contains(&server_key) {
|
||||||
|
server_keys.push(server_key);
|
||||||
|
}
|
||||||
|
}
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
warn!("Failed to get/create MCP connection: {}", err);
|
warn!(
|
||||||
|
"Failed to get/create MCP connection for {}: {}",
|
||||||
|
server_key, err
|
||||||
|
);
|
||||||
|
// Continue processing other tools
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if has_mcp_tools && !server_keys.is_empty() {
|
||||||
|
Some((mcp_manager.clone(), server_keys))
|
||||||
|
} else {
|
||||||
None
|
None
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|||||||
@@ -238,6 +238,7 @@ pub struct StreamingEventContext<'a> {
|
|||||||
pub server_label: &'a str,
|
pub server_label: &'a str,
|
||||||
pub original_request: &'a ResponsesRequest,
|
pub original_request: &'a ResponsesRequest,
|
||||||
pub previous_response_id: Option<&'a str>,
|
pub previous_response_id: Option<&'a str>,
|
||||||
|
pub server_keys: &'a [String],
|
||||||
}
|
}
|
||||||
|
|
||||||
pub type StreamingRequest = OwnedStreamingContext;
|
pub type StreamingRequest = OwnedStreamingContext;
|
||||||
|
|||||||
@@ -202,6 +202,7 @@ pub(super) async fn execute_streaming_tool_calls(
|
|||||||
pub(super) fn prepare_mcp_payload_for_streaming(
|
pub(super) fn prepare_mcp_payload_for_streaming(
|
||||||
payload: &mut Value,
|
payload: &mut Value,
|
||||||
active_mcp: &Arc<mcp::McpManager>,
|
active_mcp: &Arc<mcp::McpManager>,
|
||||||
|
server_keys: &[String],
|
||||||
) {
|
) {
|
||||||
if let Some(obj) = payload.as_object_mut() {
|
if let Some(obj) = payload.as_object_mut() {
|
||||||
// Remove any non-function tools from outgoing payload
|
// 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
|
// 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());
|
let mut tools_json = Vec::with_capacity(tools.len());
|
||||||
for t in tools {
|
for t in tools {
|
||||||
let parameters = Value::Object((*t.input_schema).clone());
|
let parameters = Value::Object((*t.input_schema).clone());
|
||||||
@@ -310,8 +311,9 @@ pub(super) fn send_mcp_list_tools_events(
|
|||||||
server_label: &str,
|
server_label: &str,
|
||||||
output_index: usize,
|
output_index: usize,
|
||||||
sequence_number: &mut u64,
|
sequence_number: &mut u64,
|
||||||
|
server_keys: &[String],
|
||||||
) -> bool {
|
) -> 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
|
let item_id = tools_item_full
|
||||||
.get("id")
|
.get("id")
|
||||||
.and_then(|v| v.as_str())
|
.and_then(|v| v.as_str())
|
||||||
@@ -464,13 +466,14 @@ pub(super) fn inject_mcp_metadata_streaming(
|
|||||||
state: &ToolLoopState,
|
state: &ToolLoopState,
|
||||||
mcp: &Arc<mcp::McpManager>,
|
mcp: &Arc<mcp::McpManager>,
|
||||||
server_label: &str,
|
server_label: &str,
|
||||||
|
server_keys: &[String],
|
||||||
) {
|
) {
|
||||||
if let Some(output_array) = response.get_mut("output").and_then(|v| v.as_array_mut()) {
|
if let Some(output_array) = response.get_mut("output").and_then(|v| v.as_array_mut()) {
|
||||||
output_array.retain(|item| {
|
output_array.retain(|item| {
|
||||||
item.get("type").and_then(|t| t.as_str()) != Some(ItemType::MCP_LIST_TOOLS)
|
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);
|
output_array.insert(0, list_tools_item);
|
||||||
|
|
||||||
let mcp_call_items =
|
let mcp_call_items =
|
||||||
@@ -482,7 +485,7 @@ pub(super) fn inject_mcp_metadata_streaming(
|
|||||||
}
|
}
|
||||||
} else if let Some(obj) = response.as_object_mut() {
|
} else if let Some(obj) = response.as_object_mut() {
|
||||||
let mut output_items = Vec::new();
|
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(
|
output_items.extend(build_executed_mcp_call_items(
|
||||||
&state.conversation_history,
|
&state.conversation_history,
|
||||||
server_label,
|
server_label,
|
||||||
@@ -584,6 +587,7 @@ pub(super) async fn execute_tool_loop(
|
|||||||
"max_tool_calls",
|
"max_tool_calls",
|
||||||
active_mcp,
|
active_mcp,
|
||||||
original_body,
|
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");
|
let server_label = extract_server_label(original_body.tools.as_deref(), "mcp");
|
||||||
|
|
||||||
// Build mcp_list_tools item
|
// 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
|
// Insert at beginning of output array
|
||||||
if let Some(output_array) = response_json
|
if let Some(output_array) = response_json
|
||||||
@@ -668,6 +673,7 @@ pub(super) fn build_incomplete_response(
|
|||||||
reason: &str,
|
reason: &str,
|
||||||
active_mcp: &Arc<mcp::McpManager>,
|
active_mcp: &Arc<mcp::McpManager>,
|
||||||
original_body: &ResponsesRequest,
|
original_body: &ResponsesRequest,
|
||||||
|
server_keys: &[String],
|
||||||
) -> Result<Value, String> {
|
) -> Result<Value, String> {
|
||||||
let obj = response
|
let obj = response
|
||||||
.as_object_mut()
|
.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
|
// Add mcp_list_tools and executed mcp_call items at the beginning
|
||||||
if state.total_calls > 0 || !mcp_call_items.is_empty() {
|
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);
|
output_array.insert(0, list_tools_item);
|
||||||
|
|
||||||
// Add mcp_call items for executed calls using helper
|
// 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
|
/// Build a mcp_list_tools output item
|
||||||
pub(super) fn build_mcp_list_tools_item(mcp: &Arc<mcp::McpManager>, server_label: &str) -> Value {
|
pub(super) fn build_mcp_list_tools_item(
|
||||||
let tools = mcp.list_tools();
|
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
|
let tools_json: Vec<Value> = tools
|
||||||
.iter()
|
.iter()
|
||||||
.map(|t| {
|
.map(|t| {
|
||||||
|
|||||||
@@ -50,11 +50,15 @@ pub async fn handle_non_streaming_response(mut ctx: RequestContext) -> Response
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
if let Some(ref tools) = original_body.tools {
|
let server_keys = match original_body.tools.as_ref() {
|
||||||
ensure_request_mcp_client(mcp_manager, tools.as_slice()).await;
|
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
|
None
|
||||||
} else {
|
} else {
|
||||||
Some(mcp_manager)
|
Some(mcp_manager)
|
||||||
@@ -63,8 +67,11 @@ pub async fn handle_non_streaming_response(mut ctx: RequestContext) -> Response
|
|||||||
let mut response_json: Value;
|
let mut response_json: Value;
|
||||||
|
|
||||||
if let Some(mcp) = active_mcp {
|
if let Some(mcp) = active_mcp {
|
||||||
let config = McpLoopConfig::default();
|
let config = McpLoopConfig {
|
||||||
prepare_mcp_payload_for_streaming(&mut payload, mcp);
|
server_keys: server_keys.clone(),
|
||||||
|
..McpLoopConfig::default()
|
||||||
|
};
|
||||||
|
prepare_mcp_payload_for_streaming(&mut payload, mcp, &server_keys);
|
||||||
|
|
||||||
match execute_tool_loop(
|
match execute_tool_loop(
|
||||||
ctx.components.client(),
|
ctx.components.client(),
|
||||||
|
|||||||
@@ -442,7 +442,13 @@ pub(super) fn send_final_response_event(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if let Some(mcp) = active_mcp {
|
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);
|
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>,
|
headers: Option<&HeaderMap>,
|
||||||
req: StreamingRequest,
|
req: StreamingRequest,
|
||||||
active_mcp: &Arc<crate::mcp::McpManager>,
|
active_mcp: &Arc<crate::mcp::McpManager>,
|
||||||
|
server_keys: Vec<String>,
|
||||||
) -> Response {
|
) -> Response {
|
||||||
// Transform MCP tools to function tools in payload
|
// Transform MCP tools to function tools in payload
|
||||||
let mut payload = req.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 (tx, rx) = mpsc::unbounded_channel::<Result<Bytes, io::Error>>();
|
||||||
let should_store = req.original_body.store.unwrap_or(false);
|
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 headers_opt = headers.cloned();
|
||||||
let payload_clone = payload.clone();
|
let payload_clone = payload.clone();
|
||||||
let active_mcp_clone = Arc::clone(active_mcp);
|
let active_mcp_clone = Arc::clone(active_mcp);
|
||||||
|
let server_keys_clone = server_keys.clone();
|
||||||
|
|
||||||
// Spawn the streaming loop task
|
// Spawn the streaming loop task
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let mut state = ToolLoopState::new(original_request.input.clone());
|
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 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 tools_json = payload_clone.get("tools").cloned().unwrap_or(json!([]));
|
||||||
let base_payload = payload_clone.clone();
|
let base_payload = payload_clone.clone();
|
||||||
@@ -680,6 +691,7 @@ pub(super) async fn handle_streaming_with_tool_interception(
|
|||||||
server_label,
|
server_label,
|
||||||
original_request: &original_request,
|
original_request: &original_request,
|
||||||
previous_response_id: previous_response_id.as_deref(),
|
previous_response_id: previous_response_id.as_deref(),
|
||||||
|
server_keys: &server_keys_clone,
|
||||||
};
|
};
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
@@ -789,6 +801,7 @@ pub(super) async fn handle_streaming_with_tool_interception(
|
|||||||
server_label,
|
server_label,
|
||||||
list_tools_index,
|
list_tools_index,
|
||||||
&mut sequence_number,
|
&mut sequence_number,
|
||||||
|
&server_keys_clone,
|
||||||
) {
|
) {
|
||||||
// Client disconnected
|
// Client disconnected
|
||||||
return;
|
return;
|
||||||
@@ -868,6 +881,7 @@ pub(super) async fn handle_streaming_with_tool_interception(
|
|||||||
&state,
|
&state,
|
||||||
&active_mcp_clone,
|
&active_mcp_clone,
|
||||||
server_label,
|
server_label,
|
||||||
|
&server_keys_clone,
|
||||||
);
|
);
|
||||||
|
|
||||||
mask_tools_as_mcp(&mut response_json, &original_request);
|
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 original_body = ctx.responses_request();
|
||||||
let mcp_manager = ctx.components.mcp_manager().expect("MCP manager required");
|
let mcp_manager = ctx.components.mcp_manager().expect("MCP manager required");
|
||||||
|
|
||||||
if let Some(ref tools) = original_body.tools {
|
let server_keys = match original_body.tools.as_ref() {
|
||||||
ensure_request_mcp_client(mcp_manager, tools.as_slice()).await;
|
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
|
None
|
||||||
} else {
|
} else {
|
||||||
Some(mcp_manager.clone())
|
Some(mcp_manager.clone())
|
||||||
@@ -1003,5 +1021,12 @@ pub async fn handle_streaming_response(ctx: RequestContext) -> Response {
|
|||||||
let active_mcp = active_mcp.unwrap();
|
let active_mcp = active_mcp.unwrap();
|
||||||
|
|
||||||
// MCP is active - transform tools and set up interception
|
// 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
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user