Skip to main content

systemprompt_cli/commands/shared/
mcp_tools.rs

1//! MCP tool probing shared by `admin agents tools` and `plugins mcp tools`.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use anyhow::{Context, Result};
7use rmcp::ServiceExt;
8use rmcp::model::{ClientCapabilities, ClientInfo, Implementation};
9use rmcp::transport::streamable_http_client::{
10    StreamableHttpClientTransport, StreamableHttpClientTransportConfig,
11};
12use std::time::Duration;
13use systemprompt_identifiers::{AgentName, ContextId, SessionId, SessionToken, TraceId};
14use systemprompt_mcp::services::client::HttpClientWithContext;
15use systemprompt_models::execution::context::RequestContext;
16use tokio::time::timeout;
17use tracing::debug;
18
19fn probe_context(server_name: &str) -> RequestContext {
20    RequestContext::new(
21        SessionId::new(format!("cli-{server_name}")),
22        TraceId::generate(),
23        ContextId::derived_from_cli_probe(server_name),
24        AgentName::system(),
25    )
26}
27
28#[derive(Debug)]
29pub struct ToolInfo {
30    pub name: String,
31    pub description: Option<String>,
32    pub parameters_count: usize,
33    pub input_schema: Option<serde_json::Value>,
34    pub output_schema: Option<serde_json::Value>,
35}
36
37pub async fn list_tools_unauthenticated(
38    server_name: &str,
39    port: u16,
40    timeout_secs: u64,
41) -> Result<Vec<ToolInfo>> {
42    let url = format!("http://127.0.0.1:{}/mcp", port);
43    let config = StreamableHttpClientTransportConfig::with_uri(url.as_str());
44    let transport = StreamableHttpClientTransport::with_client(
45        HttpClientWithContext::new(probe_context(server_name)),
46        config,
47    );
48
49    let client_info = ClientInfo::new(
50        ClientCapabilities::default(),
51        Implementation::new(format!("systemprompt-cli-{}", server_name), "1.0.0"),
52    );
53
54    let client = timeout(
55        Duration::from_secs(timeout_secs),
56        client_info.serve(transport),
57    )
58    .await
59    .context("Connection timeout")?
60    .context("Failed to connect to MCP server")?;
61
62    let tools_response = client
63        .list_tools(None)
64        .await
65        .context("Failed to list tools")?;
66
67    let tools: Vec<ToolInfo> = tools_response
68        .tools
69        .into_iter()
70        .map(convert_tool_info)
71        .collect();
72
73    client.cancel().await?;
74    Ok(tools)
75}
76
77pub async fn list_tools_authenticated(
78    server_name: &str,
79    port: u16,
80    token: &SessionToken,
81    timeout_secs: u64,
82) -> Result<Vec<ToolInfo>> {
83    let url = format!("http://127.0.0.1:{}/mcp", port);
84
85    let config = StreamableHttpClientTransportConfig::with_uri(url.as_str())
86        .auth_header(token.as_str().to_owned());
87    let context = probe_context(server_name).with_auth_token(token.as_str());
88    let transport =
89        StreamableHttpClientTransport::with_client(HttpClientWithContext::new(context), config);
90
91    let client_info = ClientInfo::new(
92        ClientCapabilities::default(),
93        Implementation::new(format!("systemprompt-cli-{}", server_name), "1.0.0"),
94    );
95
96    let client = timeout(
97        Duration::from_secs(timeout_secs),
98        client_info.serve(transport),
99    )
100    .await
101    .context("Connection timeout")?
102    .context("Failed to connect to MCP server")?;
103
104    let tools_response = client
105        .list_tools(None)
106        .await
107        .context("Failed to list tools")?;
108
109    let tools: Vec<ToolInfo> = tools_response
110        .tools
111        .into_iter()
112        .map(convert_tool_info)
113        .collect();
114
115    client.cancel().await?;
116    Ok(tools)
117}
118
119fn convert_tool_info(tool: rmcp::model::Tool) -> ToolInfo {
120    let input_schema = serde_json::to_value(&tool.input_schema)
121        .inspect_err(|e| debug!("Failed to serialize input schema: {}", e))
122        .ok();
123    let output_schema = tool.output_schema.and_then(|s| {
124        serde_json::to_value(s.as_ref())
125            .inspect_err(|e| debug!("Failed to serialize output schema: {}", e))
126            .ok()
127    });
128    let parameters_count = input_schema
129        .as_ref()
130        .and_then(|s| s.get("properties"))
131        .and_then(|p| p.as_object())
132        .map_or(0, serde_json::Map::len);
133
134    ToolInfo {
135        name: tool.name.to_string(),
136        description: tool.description.map(|d| d.to_string()),
137        parameters_count,
138        input_schema,
139        output_schema,
140    }
141}