Skip to content

Implement MCP sampling protocol for dynamic tool discovery #2525

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
8 changes: 6 additions & 2 deletions crates/chat-cli/src/cli/chat/tools/custom_tool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@ use crate::mcp_client::{
Client as McpClient,
ClientConfig as McpClientConfig,
JsonRpcResponse,
JsonRpcStdioTransport,
MessageContent,
Messenger,
PromptGet,
Expand All @@ -55,6 +54,9 @@ pub struct CustomToolConfig {
/// A boolean flag to denote whether or not to load this mcp server
#[serde(default)]
pub disabled: bool,
/// Enable MCP sampling support for this server
#[serde(default)]
pub sampling_enabled: bool,
/// A flag to denote whether this is a server from the legacy mcp.json
#[serde(skip)]
pub is_from_legacy_mcp_json: bool,
Expand Down Expand Up @@ -103,6 +105,7 @@ impl CustomToolClient {
env,
timeout,
disabled: _,
sampling_enabled,
..
} = config;

Expand All @@ -122,8 +125,9 @@ impl CustomToolClient {
"version": "1.0.0"
}),
env: processed_env,
sampling_enabled,
};
let client = McpClient::<JsonRpcStdioTransport>::from_config(mcp_client_config)?;
let client = McpClient::<StdioTransport>::from_config(mcp_client_config, Some(std::sync::Arc::new(os.client.clone())))?;
Ok(CustomToolClient::Stdio {
server_name,
client,
Expand Down
Loading