mcp_proxy/
admin_tools.rs

1//! MCP admin tools for proxy introspection.
2//!
3//! Registers tools under the `proxy/` namespace that allow any MCP client
4//! to query proxy status. Uses `ChannelTransport` to add an in-process
5//! backend to the proxy.
6
7use std::sync::Arc;
8
9use schemars::JsonSchema;
10use serde::{Deserialize, Serialize};
11use tower_mcp::client::ChannelTransport;
12use tower_mcp::proxy::{AddBackendError, McpProxy};
13use tower_mcp::{CallToolResult, McpRouter, NoParams, SessionHandle, ToolBuilder};
14
15use crate::admin::AdminState;
16use crate::config::ProxyConfig;
17
18/// Shared state accessible to admin tool handlers.
19#[derive(Clone)]
20struct AdminToolState {
21    admin_state: AdminState,
22    session_handle: SessionHandle,
23    config_snapshot: Arc<String>,
24    proxy: McpProxy,
25}
26
27#[derive(Serialize)]
28struct BackendInfo {
29    namespace: String,
30    healthy: bool,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    last_checked_at: Option<String>,
33    consecutive_failures: u32,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    error: Option<String>,
36    #[serde(skip_serializing_if = "Option::is_none")]
37    transport: Option<String>,
38}
39
40#[derive(Serialize)]
41struct BackendsResult {
42    proxy_name: String,
43    proxy_version: String,
44    backend_count: usize,
45    backends: Vec<BackendInfo>,
46}
47
48#[derive(Serialize)]
49struct SessionResult {
50    active_sessions: usize,
51}
52
53/// Register admin tools as an in-process backend on the proxy.
54///
55/// Tools are added under the `proxy/` namespace:
56/// - `proxy/list_backends` -- list backends with health status
57/// - `proxy/health_check` -- cached health check results
58/// - `proxy/session_count` -- active session count
59/// - `proxy/add_backend` -- dynamically add an HTTP backend
60/// - `proxy/config` -- dump current config (TOML)
61/// - `proxy/call_tool` -- (search mode only) invoke any backend tool by name
62pub async fn register_admin_tools(
63    proxy: &McpProxy,
64    admin_state: AdminState,
65    session_handle: SessionHandle,
66    config: &ProxyConfig,
67    discovery_tools: Option<Vec<tower_mcp::Tool>>,
68) -> Result<(), AddBackendError> {
69    let config_toml =
70        toml::to_string_pretty(config).unwrap_or_else(|e| format!("error serializing: {e}"));
71
72    let search_mode = config.proxy.tool_exposure == crate::config::ToolExposure::Search;
73
74    let state = AdminToolState {
75        admin_state,
76        session_handle,
77        config_snapshot: Arc::new(config_toml),
78        proxy: proxy.clone(),
79    };
80
81    let router = build_admin_router(state, discovery_tools, search_mode);
82    let transport = ChannelTransport::new(router);
83
84    proxy.add_backend("proxy", transport).await
85}
86
87fn build_admin_router(
88    state: AdminToolState,
89    discovery_tools: Option<Vec<tower_mcp::Tool>>,
90    search_mode: bool,
91) -> McpRouter {
92    let state_for_backends = state.clone();
93    let list_backends = ToolBuilder::new("list_backends")
94        .description("List all proxy backends with health status")
95        .handler(move |_: NoParams| {
96            let s = state_for_backends.clone();
97            async move {
98                let health = s.admin_state.health().await;
99                let backends: Vec<BackendInfo> = health
100                    .iter()
101                    .map(|b| BackendInfo {
102                        namespace: b.namespace.clone(),
103                        healthy: b.healthy,
104                        last_checked_at: b.last_checked_at.map(|t| t.to_rfc3339()),
105                        consecutive_failures: b.consecutive_failures,
106                        error: b.error.clone(),
107                        transport: b.transport.clone(),
108                    })
109                    .collect();
110
111                let result = BackendsResult {
112                    proxy_name: s.admin_state.proxy_name().to_string(),
113                    proxy_version: s.admin_state.proxy_version().to_string(),
114                    backend_count: s.admin_state.backend_count(),
115                    backends,
116                };
117
118                Ok(CallToolResult::text(
119                    serde_json::to_string_pretty(&result).unwrap(),
120                ))
121            }
122        })
123        .build();
124
125    let state_for_sessions = state.clone();
126    let session_count = ToolBuilder::new("session_count")
127        .description("Get the number of active MCP sessions")
128        .handler(move |_: NoParams| {
129            let s = state_for_sessions.clone();
130            async move {
131                let count = s.session_handle.session_count().await;
132                let result = SessionResult {
133                    active_sessions: count,
134                };
135                Ok(CallToolResult::text(
136                    serde_json::to_string_pretty(&result).unwrap(),
137                ))
138            }
139        })
140        .build();
141
142    let config_snapshot = Arc::clone(&state.config_snapshot);
143    let config_tool = ToolBuilder::new("config")
144        .description("Dump the current proxy configuration")
145        .handler(move |_: NoParams| {
146            let config = Arc::clone(&config_snapshot);
147            async move { Ok(CallToolResult::text((*config).clone())) }
148        })
149        .build();
150
151    let state_for_health = state.clone();
152    let health_check = ToolBuilder::new("health_check")
153        .description("Get cached health check results for all backends")
154        .handler(move |_: NoParams| {
155            let s = state_for_health.clone();
156            async move {
157                let health = s.admin_state.health().await;
158                let backends: Vec<BackendInfo> = health
159                    .iter()
160                    .map(|b| BackendInfo {
161                        namespace: b.namespace.clone(),
162                        healthy: b.healthy,
163                        last_checked_at: b.last_checked_at.map(|t| t.to_rfc3339()),
164                        consecutive_failures: b.consecutive_failures,
165                        error: b.error.clone(),
166                        transport: b.transport.clone(),
167                    })
168                    .collect();
169                let healthy_count = backends.iter().filter(|b| b.healthy).count();
170                let total = backends.len();
171                let result = HealthCheckResult {
172                    status: if healthy_count == total {
173                        "healthy"
174                    } else {
175                        "degraded"
176                    }
177                    .to_string(),
178                    healthy_count,
179                    total_count: total,
180                    backends,
181                };
182                Ok(CallToolResult::text(
183                    serde_json::to_string_pretty(&result).unwrap(),
184                ))
185            }
186        })
187        .build();
188
189    let state_for_add = state.clone();
190    let add_backend = ToolBuilder::new("add_backend")
191        .description("Dynamically add an HTTP backend to the proxy")
192        .handler(move |input: AddBackendInput| {
193            let s = state_for_add.clone();
194            async move {
195                let transport = tower_mcp::client::HttpClientTransport::new(&input.url);
196                match s.proxy.add_backend(&input.name, transport).await {
197                    Ok(()) => Ok(CallToolResult::text(format!(
198                        "Backend '{}' added successfully at {}",
199                        input.name, input.url
200                    ))),
201                    Err(e) => Ok(CallToolResult::text(format!(
202                        "Failed to add backend '{}': {e}",
203                        input.name
204                    ))),
205                }
206            }
207        })
208        .build();
209
210    let mut router = McpRouter::new()
211        .server_info("mcp-proxy-admin", "0.1.0")
212        .tool(list_backends)
213        .tool(health_check)
214        .tool(session_count)
215        .tool(add_backend)
216        .tool(config_tool);
217
218    if search_mode {
219        let state_for_call = state.clone();
220        let call_tool = ToolBuilder::new("call_tool")
221            .description(
222                "Invoke any backend tool by its fully-qualified name. Use proxy/search_tools \
223                 to discover available tools, then call them through this tool.",
224            )
225            .handler(move |input: CallToolInput| {
226                let s = state_for_call.clone();
227                async move {
228                    use tower::Service;
229                    use tower_mcp::protocol::{CallToolParams, McpRequest, McpResponse, RequestId};
230                    use tower_mcp::router::{Extensions, RouterRequest};
231
232                    let req = RouterRequest {
233                        id: RequestId::Number(0),
234                        inner: McpRequest::CallTool(CallToolParams {
235                            name: input.name.clone(),
236                            arguments: input.arguments.unwrap_or_default().into(),
237                            meta: None,
238                            task: None,
239                        }),
240                        extensions: Extensions::new(),
241                    };
242
243                    let mut proxy = s.proxy.clone();
244                    match proxy.call(req).await {
245                        Ok(resp) => match resp.inner {
246                            Ok(McpResponse::CallTool(result)) => Ok(result),
247                            Ok(_) => Ok(CallToolResult::text(format!(
248                                "Unexpected response type for tool '{}'",
249                                input.name
250                            ))),
251                            Err(e) => Ok(CallToolResult::text(format!(
252                                "Error calling '{}': {}",
253                                input.name, e.message
254                            ))),
255                        },
256                        Err(_) => Ok(CallToolResult::text(format!(
257                            "Internal error calling '{}'",
258                            input.name
259                        ))),
260                    }
261                }
262            })
263            .build();
264        router = router.tool(call_tool);
265    }
266
267    if let Some(tools) = discovery_tools {
268        for tool in tools {
269            router = router.tool(tool);
270        }
271    }
272
273    router
274}
275
276#[derive(Serialize)]
277struct HealthCheckResult {
278    status: String,
279    healthy_count: usize,
280    total_count: usize,
281    backends: Vec<BackendInfo>,
282}
283
284#[derive(Debug, Deserialize, JsonSchema)]
285struct AddBackendInput {
286    /// Name/namespace for the new backend
287    name: String,
288    /// URL of the HTTP MCP server
289    url: String,
290}
291
292/// Input for the `proxy/call_tool` meta-tool (search mode only).
293#[derive(Debug, Deserialize, JsonSchema)]
294struct CallToolInput {
295    /// Fully-qualified tool name (e.g. "math/add", "files/read_file")
296    name: String,
297    /// Arguments to pass to the tool
298    arguments: Option<serde_json::Map<String, serde_json::Value>>,
299}
mcp_proxy/admin_tools.rs

mcp_proxy/
admin_tools.rs
