[MCP] Introduce an experimental official rust sdk based mcp client (#4252)
The [official Rust
SDK](57fc428c57)
has come a long way since we first started our mcp client implementation
5 months ago and, today, it is much more complete than our own
stdio-only implementation.
This PR introduces a new config flag `experimental_use_rmcp_client`
which will use a new mcp client powered by the sdk instead of our own.
To keep this PR simple, I've only implemented the same stdio MCP
functionality that we had but will expand on it with future PRs.
---------
Co-authored-by: pakrym-oai <pakrym@openai.com>
This commit is contained in:
142
codex-rs/rmcp-client/src/bin/rmcp_test_server.rs
Normal file
142
codex-rs/rmcp-client/src/bin/rmcp_test_server.rs
Normal file
@@ -0,0 +1,142 @@
|
||||
use std::borrow::Cow;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use rmcp::ErrorData as McpError;
|
||||
use rmcp::ServiceExt;
|
||||
use rmcp::handler::server::ServerHandler;
|
||||
use rmcp::model::CallToolRequestParam;
|
||||
use rmcp::model::CallToolResult;
|
||||
use rmcp::model::JsonObject;
|
||||
use rmcp::model::ListToolsResult;
|
||||
use rmcp::model::PaginatedRequestParam;
|
||||
use rmcp::model::ServerCapabilities;
|
||||
use rmcp::model::ServerInfo;
|
||||
use rmcp::model::Tool;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
use tokio::task;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestToolServer {
|
||||
tools: Arc<Vec<Tool>>,
|
||||
}
|
||||
pub fn stdio() -> (tokio::io::Stdin, tokio::io::Stdout) {
|
||||
(tokio::io::stdin(), tokio::io::stdout())
|
||||
}
|
||||
impl TestToolServer {
|
||||
fn new() -> Self {
|
||||
let tools = vec![Self::echo_tool()];
|
||||
Self {
|
||||
tools: Arc::new(tools),
|
||||
}
|
||||
}
|
||||
|
||||
fn echo_tool() -> Tool {
|
||||
#[expect(clippy::expect_used)]
|
||||
let schema: JsonObject = serde_json::from_value(json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"message": { "type": "string" },
|
||||
"env_var": { "type": "string" }
|
||||
},
|
||||
"required": ["message"],
|
||||
"additionalProperties": false
|
||||
}))
|
||||
.expect("echo tool schema should deserialize");
|
||||
|
||||
Tool::new(
|
||||
Cow::Borrowed("echo"),
|
||||
Cow::Borrowed("Echo back the provided message and include environment data."),
|
||||
Arc::new(schema),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct EchoArgs {
|
||||
message: String,
|
||||
#[allow(dead_code)]
|
||||
env_var: Option<String>,
|
||||
}
|
||||
|
||||
impl ServerHandler for TestToolServer {
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo {
|
||||
capabilities: ServerCapabilities::builder()
|
||||
.enable_tools()
|
||||
.enable_tool_list_changed()
|
||||
.build(),
|
||||
..ServerInfo::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn list_tools(
|
||||
&self,
|
||||
_request: Option<PaginatedRequestParam>,
|
||||
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
|
||||
) -> impl std::future::Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
|
||||
let tools = self.tools.clone();
|
||||
async move {
|
||||
Ok(ListToolsResult {
|
||||
tools: (*tools).clone(),
|
||||
next_cursor: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
request: CallToolRequestParam,
|
||||
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
|
||||
) -> Result<CallToolResult, McpError> {
|
||||
match request.name.as_ref() {
|
||||
"echo" => {
|
||||
let args: EchoArgs = match request.arguments {
|
||||
Some(arguments) => serde_json::from_value(serde_json::Value::Object(
|
||||
arguments.into_iter().collect(),
|
||||
))
|
||||
.map_err(|err| McpError::invalid_params(err.to_string(), None))?,
|
||||
None => {
|
||||
return Err(McpError::invalid_params(
|
||||
"missing arguments for echo tool",
|
||||
None,
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let env_snapshot: HashMap<String, String> = std::env::vars().collect();
|
||||
let structured_content = json!({
|
||||
"echo": args.message,
|
||||
"env": env_snapshot.get("MCP_TEST_VALUE"),
|
||||
});
|
||||
|
||||
Ok(CallToolResult {
|
||||
content: Vec::new(),
|
||||
structured_content: Some(structured_content),
|
||||
is_error: Some(false),
|
||||
meta: None,
|
||||
})
|
||||
}
|
||||
other => Err(McpError::invalid_params(
|
||||
format!("unknown tool: {other}"),
|
||||
None,
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
eprintln!("starting rmcp test server");
|
||||
// Run the server with STDIO transport. If the client disconnects we simply
|
||||
// bubble up the error so the process exits.
|
||||
let service = TestToolServer::new();
|
||||
let running = service.serve(stdio()).await?;
|
||||
|
||||
// Wait for the client to finish interacting with the server.
|
||||
running.waiting().await?;
|
||||
// Drain background tasks to ensure clean shutdown.
|
||||
task::yield_now().await;
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user