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 = match serde_json::to_value(&tool.input_schema) {
121        Ok(schema) => Some(schema),
122        Err(e) => {
123            debug!(tool = %tool.name, error = %e, "Failed to serialize input schema");
124            None
125        },
126    };
127    let output_schema =
128        tool.output_schema
129            .as_ref()
130            .and_then(|s| match serde_json::to_value(s.as_ref()) {
131                Ok(schema) => Some(schema),
132                Err(e) => {
133                    debug!(tool = %tool.name, error = %e, "Failed to serialize output schema");
134                    None
135                },
136            });
137    let parameters_count = input_schema
138        .as_ref()
139        .and_then(|s| s.get("properties"))
140        .and_then(|p| p.as_object())
141        .map_or(0, serde_json::Map::len);
142
143    ToolInfo {
144        name: tool.name.to_string(),
145        description: tool.description.map(|d| d.to_string()),
146        parameters_count,
147        input_schema,
148        output_schema,
149    }
150}