From 9e949e58773ff0ac5f34efed7c4af1aa0fb9dedc Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sun, 26 Oct 2025 01:47:12 -0700 Subject: [PATCH] [router] centralize mcp tool args handling (#12155) --- sgl-router/src/mcp/error.rs | 3 + sgl-router/src/mcp/manager.rs | 22 ++- sgl-router/src/mcp/mod.rs | 2 + sgl-router/src/mcp/tool_args.rs | 104 +++++++++++++ .../src/routers/grpc/responses/tool_loop.rs | 145 ++++++++---------- sgl-router/src/routers/openai/mcp.rs | 69 +++++---- 6 files changed, 229 insertions(+), 116 deletions(-) create mode 100644 sgl-router/src/mcp/tool_args.rs diff --git a/sgl-router/src/mcp/error.rs b/sgl-router/src/mcp/error.rs index 03b8b4cd1..31ffdb8a2 100644 --- a/sgl-router/src/mcp/error.rs +++ b/sgl-router/src/mcp/error.rs @@ -31,6 +31,9 @@ pub enum McpError { #[error("Prompt not found: {0}")] PromptNotFound(String), + #[error("Invalid arguments: {0}")] + InvalidArguments(String), + #[error(transparent)] Sdk(#[from] Box), diff --git a/sgl-router/src/mcp/manager.rs b/sgl-router/src/mcp/manager.rs index 4bbd18533..1f34c4c01 100644 --- a/sgl-router/src/mcp/manager.rs +++ b/sgl-router/src/mcp/manager.rs @@ -37,7 +37,9 @@ use crate::mcp::{ connection_pool::McpConnectionPool, error::{McpError, McpResult}, inventory::ToolInventory, + tool_args::ToolArgs, }; + /// Type alias for MCP client type McpClient = RunningService; @@ -221,18 +223,28 @@ impl McpManager { .collect() } - /// Call a tool by name + /// Call a tool by name with automatic type coercion + /// + /// Accepts either JSON string or parsed Map as arguments. + /// Automatically converts string numbers to actual numbers based on tool schema. pub async fn call_tool( &self, tool_name: &str, - args: Option>, + args: impl Into, ) -> McpResult { - // Get server that owns this tool - let (server_name, _tool_info) = self + // Get tool info for schema and server + let (server_name, tool_info) = self .inventory .get_tool(tool_name) .ok_or_else(|| McpError::ToolNotFound(tool_name.to_string()))?; + // Convert args with type coercion based on schema + let tool_schema = tool_info.parameters.as_ref(); + let args_map = args + .into() + .into_map(tool_schema) + .map_err(McpError::InvalidArguments)?; + // Get client for that server let client = self .get_client(&server_name) @@ -242,7 +254,7 @@ impl McpManager { // Call the tool let request = CallToolRequestParam { name: Cow::Owned(tool_name.to_string()), - arguments: args, + arguments: args_map, }; client diff --git a/sgl-router/src/mcp/mod.rs b/sgl-router/src/mcp/mod.rs index 1257c44cb..d48cdbfc3 100644 --- a/sgl-router/src/mcp/mod.rs +++ b/sgl-router/src/mcp/mod.rs @@ -14,6 +14,7 @@ pub mod inventory; pub mod manager; pub mod oauth; pub mod proxy; +pub mod tool_args; // Re-export the main types for convenience pub use config::{ @@ -25,3 +26,4 @@ pub use error::{McpError, McpResult}; pub use inventory::ToolInventory; pub use manager::{McpManager, McpManagerStats}; pub use proxy::{create_http_client, resolve_proxy_config}; +pub use tool_args::ToolArgs; diff --git a/sgl-router/src/mcp/tool_args.rs b/sgl-router/src/mcp/tool_args.rs new file mode 100644 index 000000000..288182e60 --- /dev/null +++ b/sgl-router/src/mcp/tool_args.rs @@ -0,0 +1,104 @@ +//! Tool arguments handling and type coercion +//! +//! This module provides utilities for handling MCP tool arguments, +//! supporting both JSON strings and parsed Maps with automatic type coercion. + +use serde_json::Map; + +/// Tool arguments input - supports both JSON strings and parsed Maps +pub enum ToolArgs { + /// JSON string that needs parsing + JsonString(String), + /// Already parsed map + Map(Option>), +} + +impl ToolArgs { + /// Convert to Map with type coercion based on tool schema + pub(crate) fn into_map( + self, + tool_schema: Option<&serde_json::Value>, + ) -> Result>, String> { + match self { + ToolArgs::JsonString(s) => { + if s.is_empty() || s == "{}" { + return Ok(None); + } + let mut value: serde_json::Value = + serde_json::from_str(&s).map_err(|e| format!("parse tool args: {}", e))?; + Self::coerce_types(&mut value, tool_schema)?; + let result = match value { + serde_json::Value::Object(m) => Some(m), + _ => None, + }; + Ok(result) + } + ToolArgs::Map(map) => { + if let Some(m) = map { + let mut value = serde_json::Value::Object(m); + Self::coerce_types(&mut value, tool_schema)?; + let result = match value { + serde_json::Value::Object(m) => Some(m), + _ => None, + }; + Ok(result) + } else { + Ok(None) + } + } + } + } + + /// Coerce string numbers to actual numbers based on schema + /// LLMs often output numbers as strings, so we need to convert them + fn coerce_types( + value: &mut serde_json::Value, + tool_schema: Option<&serde_json::Value>, + ) -> Result<(), String> { + let Some(params) = tool_schema else { + return Ok(()); + }; + let Some(props) = params.get("properties").and_then(|p| p.as_object()) else { + return Ok(()); + }; + let Some(args) = value.as_object_mut() else { + return Ok(()); + }; + + for (key, val) in args.iter_mut() { + let should_be_number = props + .get(key) + .and_then(|s| s.get("type")) + .and_then(|t| t.as_str()) + .is_some_and(|t| matches!(t, "number" | "integer")); + + if should_be_number { + if let Some(s) = val.as_str() { + if let Ok(num) = s.parse::() { + *val = serde_json::json!(num); + } + } + } + } + Ok(()) + } +} + +// Implement From traits for convenient conversion +impl From for ToolArgs { + fn from(s: String) -> Self { + ToolArgs::JsonString(s) + } +} + +impl From<&str> for ToolArgs { + fn from(s: &str) -> Self { + ToolArgs::JsonString(s.to_string()) + } +} + +impl From>> for ToolArgs { + fn from(map: Option>) -> Self { + ToolArgs::Map(map) + } +} diff --git a/sgl-router/src/routers/grpc/responses/tool_loop.rs b/sgl-router/src/routers/grpc/responses/tool_loop.rs index 8555aa166..3f2181d37 100644 --- a/sgl-router/src/routers/grpc/responses/tool_loop.rs +++ b/sgl-router/src/routers/grpc/responses/tool_loop.rs @@ -99,63 +99,6 @@ fn extract_all_tool_calls_from_chat( } } -/// Execute an MCP tool call -async fn execute_mcp_call( - mcp_mgr: &Arc, - tool_name: &str, - args_json_str: &str, -) -> Result { - // Parse arguments JSON string to Value - let mut args_value: serde_json::Value = - serde_json::from_str::(args_json_str) - .map_err(|e| format!("parse tool args: {}", e))?; - - // Get tool info to access schema for type coercion - let tool_info = mcp_mgr - .get_tool(tool_name) - .ok_or_else(|| format!("tool not found: {}", tool_name))?; - - // Coerce string numbers to actual numbers based on schema (LLMs often output numbers as strings) - if let Some(params) = &tool_info.parameters { - let properties = params.get("properties").and_then(|p| p.as_object()); - let args_obj = args_value.as_object_mut(); - - if let (Some(props), Some(args)) = (properties, args_obj) { - for (key, val) in args.iter_mut() { - let should_be_number = props - .get(key) - .and_then(|s| s.get("type")) - .and_then(|t| t.as_str()) - .is_some_and(|t| matches!(t, "number" | "integer")); - - if should_be_number { - if let Some(s) = val.as_str() { - if let Ok(num) = s.parse::() { - *val = json!(num); - } - } - } - } - } - } - - let args_obj = args_value.as_object().cloned(); - - debug!( - "Calling MCP tool '{}' with args: {}", - tool_name, args_json_str - ); - - let result = mcp_mgr - .call_tool(tool_name, args_obj) - .await - .map_err(|e| format!("tool call failed: {}", e))?; - - let output_str = serde_json::to_string(&result) - .map_err(|e| format!("Failed to serialize tool result: {}", e))?; - Ok(output_str) -} - /// State for tracking multi-turn tool calling loop struct ToolLoopState { iteration: usize, @@ -381,18 +324,32 @@ pub(super) async fn execute_tool_loop( // Increment after check state.total_calls += 1; - // Execute the MCP tool - let (output_str, success, error) = - match execute_mcp_call(&mcp_manager, &tool_name, &args_json_str).await { + // Execute the MCP tool - manager handles parsing and type coercion + debug!( + "Calling MCP tool '{}' with args: {}", + tool_name, args_json_str + ); + let (output_str, success, error) = match mcp_manager + .call_tool(tool_name.as_str(), args_json_str.as_str()) + .await + { + Ok(result) => match serde_json::to_string(&result) { Ok(output) => (output, true, None), - Err(err) => { - warn!("Tool execution failed: {}", err); - let error_msg = err.clone(); - // Return error as output, let model decide how to proceed - let error_json = json!({ "error": err }).to_string(); - (error_json, false, Some(error_msg)) + Err(e) => { + let err = format!("Failed to serialize tool result: {}", e); + warn!("{}", err); + let error_json = json!({ "error": &err }).to_string(); + (error_json, false, Some(err)) } - }; + }, + Err(err) => { + let err_str = format!("tool call failed: {}", err); + warn!("Tool execution failed: {}", err_str); + // Return error as output, let model decide how to proceed + let error_json = json!({ "error": &err_str }).to_string(); + (error_json, false, Some(err_str)) + } + }; // Record the call in state state.record_call( @@ -796,9 +753,16 @@ async fn execute_tool_loop_streaming_internal( emitter.emit_mcp_call_arguments_done(output_index, &item_id, &args_json_str); emitter.send_event(&event, &tx)?; - // Execute the MCP tool - let (output_str, success, error) = - match execute_mcp_call(&mcp_manager, &tool_name, &args_json_str).await { + // Execute the MCP tool - manager handles parsing and type coercion + debug!( + "Calling MCP tool '{}' with args: {}", + tool_name, args_json_str + ); + let (output_str, success, error) = match mcp_manager + .call_tool(tool_name.as_str(), args_json_str.as_str()) + .await + { + Ok(result) => match serde_json::to_string(&result) { Ok(output) => { // Emit mcp_call.completed let event = emitter.emit_mcp_call_completed(output_index, &item_id); @@ -822,8 +786,9 @@ async fn execute_tool_loop_streaming_internal( emitter.complete_output_item(output_index); (output, true, None) } - Err(err) => { - warn!("Tool execution failed: {}", err); + Err(e) => { + let err = format!("Failed to serialize tool result: {}", e); + warn!("{}", err); // Emit mcp_call.failed let event = emitter.emit_mcp_call_failed(output_index, &item_id, &err); emitter.send_event(&event, &tx)?; @@ -836,7 +801,7 @@ async fn execute_tool_loop_streaming_internal( "server_label": state.server_label, "status": "failed", "arguments": args_json_str, - "error": err + "error": &err }); // Emit output_item.done @@ -844,11 +809,37 @@ async fn execute_tool_loop_streaming_internal( emitter.send_event(&event, &tx)?; emitter.complete_output_item(output_index); - let error_msg = err.clone(); - let error_json = json!({ "error": err }).to_string(); - (error_json, false, Some(error_msg)) + let error_json = json!({ "error": &err }).to_string(); + (error_json, false, Some(err)) } - }; + }, + Err(err) => { + let err_str = format!("tool call failed: {}", err); + warn!("Tool execution failed: {}", err_str); + // Emit mcp_call.failed + let event = emitter.emit_mcp_call_failed(output_index, &item_id, &err_str); + emitter.send_event(&event, &tx)?; + + // Build failed item + let item_done = json!({ + "id": item_id, + "type": "mcp_call", + "name": tool_name, + "server_label": state.server_label, + "status": "failed", + "arguments": args_json_str, + "error": &err_str + }); + + // Emit output_item.done + let event = emitter.emit_output_item_done(output_index, &item_done); + emitter.send_event(&event, &tx)?; + + emitter.complete_output_item(output_index); + let error_json = json!({ "error": &err_str }).to_string(); + (error_json, false, Some(err_str)) + } + }; // Record the call in state state.record_call( diff --git a/sgl-router/src/routers/openai/mcp.rs b/sgl-router/src/routers/openai/mcp.rs index ef860655d..4504de1bb 100644 --- a/sgl-router/src/routers/openai/mcp.rs +++ b/sgl-router/src/routers/openai/mcp.rs @@ -14,7 +14,7 @@ use axum::http::HeaderMap; use bytes::Bytes; use serde_json::{json, to_value, Value}; use tokio::sync::mpsc; -use tracing::{info, warn}; +use tracing::{debug, info, warn}; use super::utils::{event_types, generate_id}; use crate::{ @@ -191,31 +191,6 @@ pub async fn ensure_request_mcp_client( // Tool Execution // ============================================================================ -/// Execute an MCP tool call -pub(super) async fn execute_mcp_call( - mcp_mgr: &Arc, - tool_name: &str, - args_json_str: &str, -) -> Result<(String, String), String> { - let args_value: Value = - serde_json::from_str(args_json_str).map_err(|e| format!("parse tool args: {}", e))?; - let args_obj = args_value.as_object().cloned(); - - let server_name = mcp_mgr - .get_tool(tool_name) - .map(|t| t.server) - .ok_or_else(|| format!("tool not found: {}", tool_name))?; - - let result = mcp_mgr - .call_tool(tool_name, args_obj) - .await - .map_err(|e| format!("tool call failed: {}", e))?; - - let output_str = serde_json::to_string(&result) - .map_err(|e| format!("Failed to serialize tool result: {}", e))?; - Ok((server_name, output_str)) -} - /// Execute detected tool calls and send completion events to client /// Returns false if client disconnected during execution pub(super) async fn execute_streaming_tool_calls( @@ -249,12 +224,26 @@ pub(super) async fn execute_streaming_tool_calls( &call.arguments_buffer }; - let call_result = execute_mcp_call(active_mcp, &call.name, args_str).await; + // Call tool directly - manager handles parsing and type coercion + debug!("Calling MCP tool '{}' with args: {}", call.name, args_str); + let call_result = active_mcp.call_tool(&call.name, args_str).await; let (output_str, success, error_msg) = match call_result { - Ok((_, output)) => (output, true, None), + Ok(result) => match serde_json::to_string(&result) { + Ok(output) => (output, true, None), + Err(e) => { + let err = format!("Failed to serialize tool result: {}", e); + warn!("{}", err); + (json!({ "error": &err }).to_string(), false, Some(err)) + } + }, Err(err) => { - warn!("Tool execution failed during streaming: {}", err); - (json!({ "error": &err }).to_string(), false, Some(err)) + let err_str = format!("tool call failed: {}", err); + warn!("Tool execution failed during streaming: {}", err_str); + ( + json!({ "error": &err_str }).to_string(), + false, + Some(err_str), + ) } }; @@ -674,15 +663,27 @@ pub(super) async fn execute_tool_loop( ); } - // Execute tool - let call_result = execute_mcp_call(active_mcp, &tool_name, &args_json_str).await; + // Execute tool - manager handles parsing and type coercion + debug!( + "Calling MCP tool '{}' with args: {}", + tool_name, args_json_str + ); + let call_result = active_mcp + .call_tool(&tool_name, args_json_str.as_str()) + .await; let output_str = match call_result { - Ok((_, output)) => output, + Ok(result) => match serde_json::to_string(&result) { + Ok(output) => output, + Err(e) => { + warn!("Failed to serialize tool result: {}", e); + json!({ "error": format!("Serialization error: {}", e) }).to_string() + } + }, Err(err) => { warn!("Tool execution failed: {}", err); // Return error as output, let model decide how to proceed - json!({ "error": err }).to_string() + json!({ "error": format!("tool call failed: {}", err) }).to_string() } };