diff --git a/crates/mcp/src/core/orchestrator.rs b/crates/mcp/src/core/orchestrator.rs index 60d6bc4d0b..ef92c3d671 100644 --- a/crates/mcp/src/core/orchestrator.rs +++ b/crates/mcp/src/core/orchestrator.rs @@ -461,6 +461,14 @@ impl McpOrchestrator { /// Apply tool configurations from server config (aliases, response formats, arg mappings). fn apply_tool_configs(&self, config: &McpServerConfig) { + self.apply_tool_configs_for_server_key(&config.name, config); + } + + /// Apply tool configurations for a specific inventory server key. + /// + /// Static servers use `config.name` as the server key; dynamic servers use + /// URL-based server keys from the connection pool. + fn apply_tool_configs_for_server_key(&self, server_key: &str, config: &McpServerConfig) { let Some(tools) = &config.tools else { return; }; @@ -469,11 +477,11 @@ impl McpOrchestrator { // Check if the tool exists if !self .tool_inventory - .has_tool_qualified(&config.name, tool_name) + .has_tool_qualified(server_key, tool_name) { warn!( "Tool config for '{}:{}' but tool not found on server", - config.name, tool_name + server_key, tool_name ); continue; } @@ -498,30 +506,30 @@ impl McpOrchestrator { if let Some(alias_name) = &tool_config.alias { if let Err(e) = self.register_alias( alias_name, - &config.name, + server_key, tool_name, arg_mapping, response_format.clone(), ) { warn!( "Failed to register alias '{}' for '{}:{}': {}", - alias_name, config.name, tool_name, e + alias_name, server_key, tool_name, e ); } else { info!( "Registered alias '{}' → '{}:{}' with format {:?}", - alias_name, config.name, tool_name, response_format + alias_name, server_key, tool_name, response_format ); } } else if response_format != ResponseFormat::Passthrough { // No alias, but has custom response format - update the entry directly - if let Some(mut entry) = self.tool_inventory.get_entry(&config.name, tool_name) { + if let Some(mut entry) = self.tool_inventory.get_entry(server_key, tool_name) { entry.response_format = response_format.clone(); entry.arg_mapping.clone_from(&arg_mapping); self.tool_inventory.insert_entry(entry); info!( "Set response format {:?} for '{}:{}'", - response_format, config.name, tool_name + response_format, server_key, tool_name ); } } @@ -534,6 +542,14 @@ impl McpOrchestrator { /// corresponding tool should use the response format associated with the builtin type /// (e.g., WebSearchPreview -> WebSearchCall) unless explicitly overridden in the tools config. fn apply_builtin_response_format(&self, config: &McpServerConfig) { + self.apply_builtin_response_format_for_server_key(&config.name, config); + } + + fn apply_builtin_response_format_for_server_key( + &self, + server_key: &str, + config: &McpServerConfig, + ) { let Some(builtin_type) = &config.builtin_type else { return; }; @@ -548,7 +564,7 @@ impl McpOrchestrator { if has_explicit_config { debug!( - server = %config.name, + server = %server_key, tool = %tool_name, "Builtin tool has explicit config, skipping auto-apply of response_format" ); @@ -559,10 +575,10 @@ impl McpOrchestrator { let updated = self .tool_inventory - .update_entry(&config.name, tool_name, |entry| { + .update_entry(server_key, tool_name, |entry| { if entry.response_format != response_format { info!( - server = %config.name, + server = %server_key, tool = %tool_name, builtin_type = %builtin_type, format = ?response_format, @@ -574,7 +590,7 @@ impl McpOrchestrator { if !updated { warn!( - server = %config.name, + server = %server_key, tool = %tool_name, builtin_type = %builtin_type, "Builtin tool not found on server" @@ -832,6 +848,7 @@ impl McpOrchestrator { /// - `server_key` - Internal identifier for the server (used for `call_tool`) /// - `tool_name` - The MCP tool to call on that server /// - `response_format` - The format to use for response transformation + /// - `server_config` - The matched MCP server config /// /// Returns `None` if no server is configured for this built-in type. /// @@ -839,7 +856,7 @@ impl McpOrchestrator { /// /// ```ignore /// // Check if web_search_preview is configured - /// if let Some((server_key, tool_name, format)) = + /// if let Some((server_key, tool_name, format, _server_cfg)) = /// orchestrator.find_builtin_server(BuiltinToolType::WebSearchPreview) /// { /// // Route to MCP server @@ -857,7 +874,7 @@ impl McpOrchestrator { pub fn find_builtin_server( &self, builtin_type: BuiltinToolType, - ) -> Option<(String, String, ResponseFormat)> { + ) -> Option<(String, String, ResponseFormat, McpServerConfig)> { // Helper to extract builtin info from a server config let extract_builtin = |server_config: &McpServerConfig| { if let (Some(cfg_type), Some(tool_name)) = ( @@ -877,6 +894,7 @@ impl McpOrchestrator { server_config.name.clone(), tool_name.clone(), response_format, + server_config.clone(), )); } } @@ -1585,6 +1603,9 @@ impl McpOrchestrator { .with_category(ToolCategory::Dynamic); inventory_clone.insert_entry(entry); } + + self.apply_tool_configs_for_server_key(&server_key, &config); + self.apply_builtin_response_format_for_server_key(&server_key, &config); } Err(e) => warn!("Failed to list tools from '{}': {}", server_key, e), } @@ -2475,7 +2496,7 @@ mod tests { let result = orchestrator.find_builtin_server(BuiltinToolType::WebSearchPreview); assert!(result.is_some()); - let (server_key, tool_name, response_format) = result.unwrap(); + let (server_key, tool_name, response_format, _) = result.unwrap(); assert_eq!(server_key, "brave"); assert_eq!(tool_name, "brave_web_search"); assert_eq!(response_format, ResponseFormat::WebSearchCall); @@ -2543,7 +2564,7 @@ mod tests { let result = orchestrator.find_builtin_server(BuiltinToolType::WebSearchPreview); assert!(result.is_some()); - let (server_key, tool_name, response_format) = result.unwrap(); + let (server_key, tool_name, response_format, _) = result.unwrap(); assert_eq!(server_key, "custom-search"); assert_eq!(tool_name, "my_search"); // Should use the custom Passthrough format, not the default WebSearchCall diff --git a/model_gateway/src/routers/anthropic/router.rs b/model_gateway/src/routers/anthropic/router.rs index e39d2a3e7c..9677198e95 100644 --- a/model_gateway/src/routers/anthropic/router.rs +++ b/model_gateway/src/routers/anthropic/router.rs @@ -108,6 +108,8 @@ impl RouterTrait for AnthropicRouter { url: Some(server.url.clone()), authorization: server.authorization_token.clone(), headers: HashMap::new(), + builtin_type: None, + builtin_tool_name: None, allowed_tools: toolset_allowed.get(&server.name).and_then(|v| v.clone()), }) .collect(); diff --git a/model_gateway/src/routers/grpc/common/responses/utils.rs b/model_gateway/src/routers/grpc/common/responses/utils.rs index 0eee7a9b5e..ffa09665e5 100644 --- a/model_gateway/src/routers/grpc/common/responses/utils.rs +++ b/model_gateway/src/routers/grpc/common/responses/utils.rs @@ -58,7 +58,7 @@ pub(crate) async fn ensure_mcp_connection( } if let Some(tools) = tools { - match ensure_request_mcp_client(mcp_orchestrator, tools).await { + match ensure_request_mcp_client(mcp_orchestrator, tools, None).await { Some(mcp_servers) => { return Ok((true, mcp_servers)); } diff --git a/model_gateway/src/routers/mcp_utils.rs b/model_gateway/src/routers/mcp_utils.rs index 2ec75d0b96..66f1ca4a3f 100644 --- a/model_gateway/src/routers/mcp_utils.rs +++ b/model_gateway/src/routers/mcp_utils.rs @@ -2,6 +2,7 @@ use std::{collections::HashMap, sync::Arc}; +use axum::http::HeaderMap; use openai_protocol::responses::ResponseTool; use smg_mcp::{ BuiltinToolType, McpOrchestrator, McpServerBinding, McpServerConfig, McpTransport, @@ -21,6 +22,9 @@ pub struct McpServerInput { pub url: Option, pub authorization: Option, pub headers: HashMap, + /// Built-in routing metadata to preserve response format for dynamic connections. + pub builtin_type: Option, + pub builtin_tool_name: Option, /// Optional per-server tool allowlist. pub allowed_tools: Option>, } @@ -78,8 +82,8 @@ pub async fn connect_mcp_servers( proxy: None, required: false, tools: None, - builtin_type: None, - builtin_tool_name: None, + builtin_type: input.builtin_type.clone(), + builtin_tool_name: input.builtin_tool_name.clone(), }; let server_key = McpOrchestrator::server_key(&server_config); @@ -158,7 +162,7 @@ pub fn collect_builtin_routing( _ => continue, }; - if let Some((server_name, tool_name, response_format)) = + if let Some((server_name, tool_name, response_format, _)) = mcp_orchestrator.find_builtin_server(builtin_type) { debug!( @@ -216,7 +220,8 @@ pub async fn ensure_mcp_servers( // Add builtin tool routing servers for &builtin_type in builtin_types { - if let Some((server_name, tool_name, _)) = orchestrator.find_builtin_server(builtin_type) { + if let Some((server_name, tool_name, _, _)) = orchestrator.find_builtin_server(builtin_type) + { debug!( builtin_type = ?builtin_type, server = %server_name, @@ -252,6 +257,7 @@ pub async fn ensure_mcp_servers( pub async fn ensure_request_mcp_client( mcp_orchestrator: &Arc, tools: &[ResponseTool], + request_headers: Option<&HeaderMap>, ) -> Option> { let inputs: Vec = tools .iter() @@ -261,6 +267,8 @@ pub async fn ensure_request_mcp_client( url: mcp.server_url.clone(), authorization: mcp.authorization.clone(), headers: mcp.headers.clone().unwrap_or_default(), + builtin_type: None, + builtin_tool_name: None, allowed_tools: mcp.allowed_tools.clone(), }), _ => None, @@ -269,13 +277,96 @@ pub async fn ensure_request_mcp_client( let builtin_types = extract_builtin_types(tools); + let mut inputs = inputs; + // Built-in tools (e.g. image_generation/web_search_preview) are also routed + // through MCP when configured in mcp.local.yaml, so we must add dynamic + // inputs here to carry forwarded transport headers from the request. + if !builtin_types.is_empty() { + for builtin_type in &builtin_types { + let Some((server_name, _, _, server_cfg)) = + mcp_orchestrator.find_builtin_server(*builtin_type) + else { + continue; + }; + + if let Some((url, token, headers)) = + resolve_forwarded_headers(&server_cfg, request_headers) + { + inputs.push(McpServerInput { + label: server_name, + url: Some(url), + authorization: token, + headers, + builtin_type: server_cfg.builtin_type.clone(), + builtin_tool_name: server_cfg.builtin_tool_name.clone(), + allowed_tools: None, + }); + } + } + } + ensure_mcp_servers(mcp_orchestrator, &inputs, &builtin_types).await } +/// Resolve forwarded headers from a static MCP server config. +/// +/// In `mcp.local.yaml` transport.headers, an empty value means: +/// "forward the same header from the incoming request". +/// +/// Returns dynamic connection settings only when at least one forwarded-header +/// reference is resolved from request headers, and the transport is HTTP-based. +/// +/// Behavior note: +/// - Empty configured header value (`""`) is treated as a per-request placeholder. +/// - Non-empty configured header value is copied as a static header value. +/// - Dynamic routing is enabled only if at least one placeholder header is +/// successfully resolved from `request_headers`; otherwise callers should +/// fall back to static server routing. +fn resolve_forwarded_headers( + server_cfg: &McpServerConfig, + request_headers: Option<&HeaderMap>, +) -> Option<(String, Option, HashMap)> { + // Header forwarding only applies to HTTP transports (SSE/Streamable). + // Stdio has no HTTP header channel, so it cannot participate here. + let (url, token, configured_headers) = match &server_cfg.transport { + McpTransport::Sse { + url, + token, + headers, + } + | McpTransport::Streamable { + url, + token, + headers, + } => (url.clone(), token.clone(), headers), + McpTransport::Stdio { .. } => return None, + }; + + let mut merged_headers = HashMap::new(); + let mut has_resolved_forward_ref = false; + + for (name, value) in configured_headers { + if value.is_empty() { + if let Some(header_value) = request_headers + .and_then(|h| h.get(name.as_str())) + .and_then(|v| v.to_str().ok()) + { + merged_headers.insert(name.clone(), header_value.to_string()); + has_resolved_forward_ref = true; + } + } else { + merged_headers.insert(name.clone(), value.clone()); + } + } + + has_resolved_forward_ref.then_some((url, token, merged_headers)) +} + #[cfg(test)] mod tests { use std::{collections::HashMap, sync::Arc}; + use axum::http::{HeaderMap, HeaderValue}; use openai_protocol::{ common::Function, responses::{ @@ -567,7 +658,7 @@ mod tests { WebSearchPreviewTool::default(), )]; - let result = ensure_request_mcp_client(&orchestrator, &tools).await; + let result = ensure_request_mcp_client(&orchestrator, &tools, None).await; // Should return Some because built-in routing is configured assert!(result.is_some()); @@ -590,7 +681,7 @@ mod tests { WebSearchPreviewTool::default(), )]; - let result = ensure_request_mcp_client(&orchestrator, &tools).await; + let result = ensure_request_mcp_client(&orchestrator, &tools, None).await; // Should return None because no MCP or built-in routing is available assert!(result.is_none()); @@ -610,7 +701,7 @@ mod tests { }, })]; - let result = ensure_request_mcp_client(&orchestrator, &tools).await; + let result = ensure_request_mcp_client(&orchestrator, &tools, None).await; // Should return None - function tools don't need MCP processing assert!(result.is_none()); @@ -634,7 +725,7 @@ mod tests { ResponseTool::WebSearchPreview(WebSearchPreviewTool::default()), ]; - let result = ensure_request_mcp_client(&orchestrator, &tools).await; + let result = ensure_request_mcp_client(&orchestrator, &tools, None).await; // Should return Some because web_search_preview has built-in routing assert!(result.is_some()); @@ -643,4 +734,66 @@ mod tests { assert_eq!(mcp_servers.len(), 1); assert_eq!(mcp_servers[0].label, "search-server"); } + + #[test] + fn test_resolve_forwarded_headers_returns_none_when_no_forward_header_value_present() { + let server_cfg = McpServerConfig { + name: "image-server".to_string(), + transport: McpTransport::Streamable { + url: "http://localhost:9999/mcp".to_string(), + token: None, + headers: HashMap::from([ + ("opc-compartment-id".to_string(), "".to_string()), + ("x-static".to_string(), "fixed".to_string()), + ]), + }, + proxy: None, + required: false, + tools: None, + builtin_type: Some(BuiltinToolType::ImageGeneration), + builtin_tool_name: Some("generate_image".to_string()), + }; + + let headers = HeaderMap::new(); + let resolved = resolve_forwarded_headers(&server_cfg, Some(&headers)); + assert!(resolved.is_none()); + } + + #[test] + fn test_resolve_forwarded_headers_returns_some_when_forward_header_value_present() { + let server_cfg = McpServerConfig { + name: "image-server".to_string(), + transport: McpTransport::Streamable { + url: "http://localhost:9999/mcp".to_string(), + token: None, + headers: HashMap::from([ + ("opc-compartment-id".to_string(), "".to_string()), + ("x-static".to_string(), "fixed".to_string()), + ]), + }, + proxy: None, + required: false, + tools: None, + builtin_type: Some(BuiltinToolType::ImageGeneration), + builtin_tool_name: Some("generate_image".to_string()), + }; + + let mut headers = HeaderMap::new(); + headers.insert( + "opc-compartment-id", + HeaderValue::from_str("ocid1.compartment.oc1..example").unwrap(), + ); + + let resolved = resolve_forwarded_headers(&server_cfg, Some(&headers)); + assert!(resolved.is_some()); + let (_, _, merged_headers) = resolved.unwrap(); + assert_eq!( + merged_headers.get("opc-compartment-id").map(String::as_str), + Some("ocid1.compartment.oc1..example") + ); + assert_eq!( + merged_headers.get("x-static").map(String::as_str), + Some("fixed") + ); + } } diff --git a/model_gateway/src/routers/openai/responses/non_streaming.rs b/model_gateway/src/routers/openai/responses/non_streaming.rs index de842afe0f..fa2d1c88b5 100644 --- a/model_gateway/src/routers/openai/responses/non_streaming.rs +++ b/model_gateway/src/routers/openai/responses/non_streaming.rs @@ -59,7 +59,7 @@ pub async fn handle_non_streaming_response(mut ctx: RequestContext) -> Response // Check for MCP tools and create session if needed let mcp_servers = if let Some(tools) = original_body.tools.as_deref() { - ensure_request_mcp_client(mcp_orchestrator, tools).await + ensure_request_mcp_client(mcp_orchestrator, tools, ctx.headers()).await } else { None }; diff --git a/model_gateway/src/routers/openai/responses/streaming.rs b/model_gateway/src/routers/openai/responses/streaming.rs index 7142f9202e..cd0f474cef 100644 --- a/model_gateway/src/routers/openai/responses/streaming.rs +++ b/model_gateway/src/routers/openai/responses/streaming.rs @@ -1039,7 +1039,7 @@ pub async fn handle_streaming_response(ctx: RequestContext) -> Response { // Check for MCP tools and create request context if needed let mcp_servers = if let Some(tools) = original_body.tools.as_deref() { - ensure_request_mcp_client(&mcp_orchestrator, tools).await + ensure_request_mcp_client(&mcp_orchestrator, tools, headers.as_ref()).await } else { None };