diff --git a/crates/proxy/src/config/mod.rs b/crates/proxy/src/config/mod.rs index 6b58717..50d91f0 100644 --- a/crates/proxy/src/config/mod.rs +++ b/crates/proxy/src/config/mod.rs @@ -501,6 +501,9 @@ pub struct LoadResult { /// Resolved master_key from LiteLLM general_settings, if present. /// Caller should apply as PROXY_API_KEYS if that var is not already set. pub litellm_master_key: Option, + /// Tool-related config from a simple YAML config file. None when the config + /// was loaded from env vars, TOML, or LiteLLM format (which has no tool sections). + pub tool_config: Option, } impl MultiConfig { @@ -533,6 +536,7 @@ impl MultiConfig { multi_config: parsed.multi_config, model_router: Some(Arc::new(std::sync::RwLock::new(parsed.router))), litellm_master_key: None, + tool_config: Some(parsed.tool_config), }; } @@ -574,18 +578,21 @@ impl MultiConfig { multi_config: mc, model_router: Some(Arc::new(std::sync::RwLock::new(parsed.router))), litellm_master_key: parsed.master_key, + tool_config: None, // LiteLLM format has no tool sections }; } LoadResult { multi_config: Self::from_toml_file(&path), model_router: None, litellm_master_key: None, + tool_config: None, } } else { LoadResult { multi_config: Self::from_legacy_env(), model_router: None, litellm_master_key: None, + tool_config: None, } } } diff --git a/crates/proxy/src/config/simple.rs b/crates/proxy/src/config/simple.rs index 485f54a..e85f0f0 100644 --- a/crates/proxy/src/config/simple.rs +++ b/crates/proxy/src/config/simple.rs @@ -137,10 +137,32 @@ pub struct McpServerConfig { // Parsed result + public parser // --------------------------------------------------------------------------- +/// Tool-related config extracted from SimpleConfig, passed up to main.rs +/// so it can build ToolEngineState without re-parsing the config file. +#[derive(Debug)] +pub struct ToolStartupConfig { + pub tool_execution: Option, + pub builtin_tools: Option>, + pub mcp_servers: Option>, +} + +impl ToolStartupConfig { + /// Returns true when at least one tool-related section was present in the config. + /// Used to decide whether to construct a ToolEngineState at all. + pub fn has_any(&self) -> bool { + self.tool_execution.is_some() + || self.builtin_tools.is_some() + || self.mcp_servers.is_some() + } +} + /// Result from parsing a simple YAML config file. pub struct SimpleParsed { pub multi_config: MultiConfig, pub router: ModelRouter, + /// Tool-related sections extracted from the config. None-valued when no tool + /// sections were present (callers should check `has_any()` before using). + pub tool_config: ToolStartupConfig, } /// Parse a simple YAML config string and produce a `MultiConfig + ModelRouter`. @@ -293,6 +315,11 @@ pub fn parse_simple_yaml(yaml: &str) -> SimpleParsed { SimpleParsed { multi_config: multi, router, + tool_config: ToolStartupConfig { + tool_execution: config.tool_execution, + builtin_tools: config.builtin_tools, + mcp_servers: config.mcp_servers, + }, } } diff --git a/crates/proxy/src/main.rs b/crates/proxy/src/main.rs index 1467a3b..abcd385 100644 --- a/crates/proxy/src/main.rs +++ b/crates/proxy/src/main.rs @@ -1,4 +1,4 @@ -use anyllm_proxy::{admin, config, server::routes}; +use anyllm_proxy::{admin, config, server::routes, tools}; use std::sync::Arc; use tracing_subscriber::prelude::*; @@ -157,6 +157,80 @@ async fn main() { } } + // Build tool engine state from config, if tool sections were present. + // Only constructed when at least one of tool_execution / builtin_tools / mcp_servers + // is present in the config file, to avoid overhead when tools are unused. + let tool_engine_state: Option> = + if let Some(tc) = load_result.tool_config.filter(|tc| tc.has_any()) { + let simple_config_shell = config::simple::SimpleConfig { + routing_strategy: None, + listen_port: None, + log_bodies: None, + models: vec![], + tool_execution: tc.tool_execution, + builtin_tools: tc.builtin_tools, + mcp_servers: tc.mcp_servers, + }; + let (policy, loop_config) = simple_config_shell.build_tool_config(); + + let mut registry = tools::ToolRegistry::new(); + // Register built-in tools (gated behind the dangerous-builtin-tools feature). + anyllm_proxy::tools::builtin::register_all(&mut registry); + + // Build MCP manager and discover tools from configured servers. + let mcp_manager = if let Some(ref servers) = simple_config_shell.mcp_servers { + let manager = Arc::new(tools::McpServerManager::new()); + for server_cfg in servers { + match tools::McpServerManager::discover_tools(&server_cfg.url).await { + Ok(discovered) => { + tracing::info!( + server = %server_cfg.name, + url = %server_cfg.url, + tools = discovered.len(), + "MCP server connected and tools discovered" + ); + manager.register_server_blocking( + &server_cfg.name, + &server_cfg.url, + discovered, + ); + } + Err(e) => { + tracing::warn!( + server = %server_cfg.name, + url = %server_cfg.url, + error = %e, + "MCP server unreachable at startup; tools from this server will be unavailable" + ); + } + } + } + // Register all discovered MCP tools into the registry. + tools::mcp::register_mcp_tools(&manager, &mut registry); + Some(manager) + } else { + None + }; + + tracing::info!( + registered_tools = registry.list_names().len(), + mcp_servers = mcp_manager + .as_ref() + .map(|m| m.list_servers_blocking().len()) + .unwrap_or(0), + "tool execution engine initialized" + ); + + Some(Arc::new(routes::ToolEngineState { + registry: Arc::new(registry), + policy: Arc::new(policy), + loop_config, + mcp_manager, + })) + } else { + None + }; + // Admin web UI is opt-in: pass --webui or --admin to enable. // DISABLE_ADMIN=1 overrides the flag (useful in container/scripted environments). let flag_set = args.iter().any(|a| a == "--webui" || a == "--admin"); @@ -335,7 +409,9 @@ async fn main() { virtual_keys, hmac_secret, model_router: model_router.clone(), - mcp_manager: None, + mcp_manager: tool_engine_state + .as_ref() + .and_then(|s| s.mcp_manager.clone()), }; // Admin token: use env var or generate random UUID written to a file. @@ -447,12 +523,12 @@ async fn main() { None }; - // Build proxy router with optional shared admin state. + // Build proxy router with optional shared admin state and tool engine. let app = routes::app_multi_with_shared( multi_config, admin_parts.as_ref().map(|(s, _, _)| s.clone()), model_router, - None, // Tool engine: wired in when config-driven setup is implemented + tool_engine_state, ); // --- Start servers ---