Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add ability to register tools in McpServer (#35068)

Browse source 
Makes it easier to add tools to a server by implementing a trait

Release Notes:

- N/A
This commit is contained in:
Agus Zubiaga 2025-07-24 23:19:20 -03:00 • committed by GitHub
parent b446d66be7
commit 15c9da4ea4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 408 additions and 244 deletions
Download patch file Download diff file Expand all files Collapse all files

11
crates/agent_servers/src/claude.rs
View file

@ -2,6 +2,7 @@ mod mcp_server;
pub mod tools;
use collections::HashMap;
use context_server::listener::McpServerTool;
use project::Project;
use settings::SettingsStore;
use smol::process::Child;
@ -332,10 +333,16 @@ async fn spawn_claude(
&format!(
"mcp__{}__{}",
mcp_server::SERVER_NAME,
mcp_server::PERMISSION_TOOL
mcp_server::PermissionTool::NAME,
),
"--allowedTools",
"mcp__zed__Read,mcp__zed__Edit",
&format!(
"mcp__{}__{},mcp__{}__{}",
mcp_server::SERVER_NAME,
mcp_server::EditTool::NAME,
mcp_server::SERVER_NAME,
mcp_server::ReadTool::NAME
),
"--disallowedTools",
"Read,Edit",
])

418
crates/agent_servers/src/claude/mcp_server.rs
View file

@ -1,49 +1,24 @@
use std::path::PathBuf;
use crate::claude::tools::{ClaudeTool, EditToolParams, ReadToolParams};
use acp_thread::AcpThread;
use agent_client_protocol as acp;
use anyhow::{Context, Result};
use collections::HashMap;
use context_server::listener::{McpServerTool, ToolResponse};
use context_server::types::{
CallToolParams, CallToolResponse, Implementation, InitializeParams, InitializeResponse,
ListToolsResponse, ProtocolVersion, ServerCapabilities, Tool, ToolAnnotations,
ToolResponseContent, ToolsCapabilities, requests,
Implementation, InitializeParams, InitializeResponse, ProtocolVersion, ServerCapabilities,
ToolAnnotations, ToolResponseContent, ToolsCapabilities, requests,
};
use gpui::{App, AsyncApp, Entity, Task, WeakEntity};
use gpui::{App, AsyncApp, Task, WeakEntity};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use crate::claude::tools::{ClaudeTool, EditToolParams, ReadToolParams};
pub struct ClaudeZedMcpServer {
server: context_server::listener::McpServer,
}
pub const SERVER_NAME: &str = "zed";
pub const READ_TOOL: &str = "Read";
pub const EDIT_TOOL: &str = "Edit";
pub const PERMISSION_TOOL: &str = "Confirmation";
#[derive(Deserialize, JsonSchema, Debug)]
struct PermissionToolParams {
tool_name: String,
input: serde_json::Value,
tool_use_id: Option<String>,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct PermissionToolResponse {
behavior: PermissionToolBehavior,
updated_input: serde_json::Value,
}
#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
enum PermissionToolBehavior {
Allow,
Deny,
}
impl ClaudeZedMcpServer {
pub async fn new(
@ -52,9 +27,15 @@ impl ClaudeZedMcpServer {
) -> Result<Self> {
let mut mcp_server = context_server::listener::McpServer::new(cx).await?;
mcp_server.handle_request::<requests::Initialize>(Self::handle_initialize);
mcp_server.handle_request::<requests::ListTools>(Self::handle_list_tools);
mcp_server.handle_request::<requests::CallTool>(move |request, cx| {
Self::handle_call_tool(request, thread_rx.clone(), cx)
mcp_server.add_tool(PermissionTool {
thread_rx: thread_rx.clone(),
});
mcp_server.add_tool(ReadTool {
thread_rx: thread_rx.clone(),
});
mcp_server.add_tool(EditTool {
thread_rx: thread_rx.clone(),
});
Ok(Self { server: mcp_server })
@ -96,193 +77,6 @@ impl ClaudeZedMcpServer {
})
})
}
fn handle_list_tools(_: (), cx: &App) -> Task<Result<ListToolsResponse>> {
cx.foreground_executor().spawn(async move {
Ok(ListToolsResponse {
tools: vec![
Tool {
name: PERMISSION_TOOL.into(),
input_schema: schemars::schema_for!(PermissionToolParams).into(),
description: None,
annotations: None,
},
Tool {
name: READ_TOOL.into(),
input_schema: schemars::schema_for!(ReadToolParams).into(),
description: Some("Read the contents of a file. In sessions with mcp__zed__Read always use it instead of Read as it contains the most up-to-date contents.".to_string()),
annotations: Some(ToolAnnotations {
title: Some("Read file".to_string()),
read_only_hint: Some(true),
destructive_hint: Some(false),
open_world_hint: Some(false),
// if time passes the contents might change, but it's not going to do anything different
// true or false seem too strong, let's try a none.
idempotent_hint: None,
}),
},
Tool {
name: EDIT_TOOL.into(),
input_schema: schemars::schema_for!(EditToolParams).into(),
description: Some("Edits a file. In sessions with mcp__zed__Edit always use it instead of Edit as it will show the diff to the user better.".to_string()),
annotations: Some(ToolAnnotations {
title: Some("Edit file".to_string()),
read_only_hint: Some(false),
destructive_hint: Some(false),
open_world_hint: Some(false),
idempotent_hint: Some(false),
}),
},
],
next_cursor: None,
meta: None,
})
})
}
fn handle_call_tool(
request: CallToolParams,
mut thread_rx: watch::Receiver<WeakEntity<AcpThread>>,
cx: &App,
) -> Task<Result<CallToolResponse>> {
cx.spawn(async move |cx| {
let Some(thread) = thread_rx.recv().await?.upgrade() else {
anyhow::bail!("Thread closed");
};
if request.name.as_str() == PERMISSION_TOOL {
let input =
serde_json::from_value(request.arguments.context("Arguments required")?)?;
let result = Self::handle_permissions_tool_call(input, thread, cx).await?;
Ok(CallToolResponse {
content: vec![ToolResponseContent::Text {
text: serde_json::to_string(&result)?,
}],
is_error: None,
meta: None,
})
} else if request.name.as_str() == READ_TOOL {
let input =
serde_json::from_value(request.arguments.context("Arguments required")?)?;
let content = Self::handle_read_tool_call(input, thread, cx).await?;
Ok(CallToolResponse {
content,
is_error: None,
meta: None,
})
} else if request.name.as_str() == EDIT_TOOL {
let input =
serde_json::from_value(request.arguments.context("Arguments required")?)?;
Self::handle_edit_tool_call(input, thread, cx).await?;
Ok(CallToolResponse {
content: vec![],
is_error: None,
meta: None,
})
} else {
anyhow::bail!("Unsupported tool");
}
})
}
fn handle_read_tool_call(
ReadToolParams {
abs_path,
offset,
limit,
}: ReadToolParams,
thread: Entity<AcpThread>,
cx: &AsyncApp,
) -> Task<Result<Vec<ToolResponseContent>>> {
cx.spawn(async move |cx| {
let content = thread
.update(cx, |thread, cx| {
thread.read_text_file(abs_path, offset, limit, false, cx)
})?
.await?;
Ok(vec![ToolResponseContent::Text { text: content }])
})
}
fn handle_edit_tool_call(
params: EditToolParams,
thread: Entity<AcpThread>,
cx: &AsyncApp,
) -> Task<Result<()>> {
cx.spawn(async move |cx| {
let content = thread
.update(cx, |threads, cx| {
threads.read_text_file(params.abs_path.clone(), None, None, true, cx)
})?
.await?;
let new_content = content.replace(&params.old_text, &params.new_text);
if new_content == content {
return Err(anyhow::anyhow!("The old_text was not found in the content"));
}
thread
.update(cx, |threads, cx| {
threads.write_text_file(params.abs_path, new_content, cx)
})?
.await?;
Ok(())
})
}
fn handle_permissions_tool_call(
params: PermissionToolParams,
thread: Entity<AcpThread>,
cx: &AsyncApp,
) -> Task<Result<PermissionToolResponse>> {
cx.spawn(async move |cx| {
let claude_tool = ClaudeTool::infer(&params.tool_name, params.input.clone());
let tool_call_id =
acp::ToolCallId(params.tool_use_id.context("Tool ID required")?.into());
let allow_option_id = acp::PermissionOptionId("allow".into());
let reject_option_id = acp::PermissionOptionId("reject".into());
let chosen_option = thread
.update(cx, |thread, cx| {
thread.request_tool_call_permission(
claude_tool.as_acp(tool_call_id),
vec![
acp::PermissionOption {
id: allow_option_id.clone(),
label: "Allow".into(),
kind: acp::PermissionOptionKind::AllowOnce,
},
acp::PermissionOption {
id: reject_option_id,
label: "Reject".into(),
kind: acp::PermissionOptionKind::RejectOnce,
},
],
cx,
)
})?
.await?;
if chosen_option == allow_option_id {
Ok(PermissionToolResponse {
behavior: PermissionToolBehavior::Allow,
updated_input: params.input,
})
} else {
Ok(PermissionToolResponse {
behavior: PermissionToolBehavior::Deny,
updated_input: params.input,
})
}
})
}
}
#[derive(Serialize)]
@ -299,3 +93,187 @@ pub struct McpServerConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub env: Option<HashMap<String, String>>,
}
// Tools
#[derive(Clone)]
pub struct PermissionTool {
thread_rx: watch::Receiver<WeakEntity<AcpThread>>,
}
#[derive(Deserialize, JsonSchema, Debug)]
pub struct PermissionToolParams {
tool_name: String,
input: serde_json::Value,
tool_use_id: Option<String>,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PermissionToolResponse {
behavior: PermissionToolBehavior,
updated_input: serde_json::Value,
}
#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
enum PermissionToolBehavior {
Allow,
Deny,
}
impl McpServerTool for PermissionTool {
type Input = PermissionToolParams;
const NAME: &'static str = "Confirmation";
fn description(&self) -> &'static str {
"Request permission for tool calls"
}
async fn run(&self, input: Self::Input, cx: &mut AsyncApp) -> Result<ToolResponse> {
let mut thread_rx = self.thread_rx.clone();
let Some(thread) = thread_rx.recv().await?.upgrade() else {
anyhow::bail!("Thread closed");
};
let claude_tool = ClaudeTool::infer(&input.tool_name, input.input.clone());
let tool_call_id = acp::ToolCallId(input.tool_use_id.context("Tool ID required")?.into());
let allow_option_id = acp::PermissionOptionId("allow".into());
let reject_option_id = acp::PermissionOptionId("reject".into());
let chosen_option = thread
.update(cx, |thread, cx| {
thread.request_tool_call_permission(
claude_tool.as_acp(tool_call_id),
vec![
acp::PermissionOption {
id: allow_option_id.clone(),
label: "Allow".into(),
kind: acp::PermissionOptionKind::AllowOnce,
},
acp::PermissionOption {
id: reject_option_id.clone(),
label: "Reject".into(),
kind: acp::PermissionOptionKind::RejectOnce,
},
],
cx,
)
})?
.await?;
let response = if chosen_option == allow_option_id {
PermissionToolResponse {
behavior: PermissionToolBehavior::Allow,
updated_input: input.input,
}
} else {
PermissionToolResponse {
behavior: PermissionToolBehavior::Deny,
updated_input: input.input,
}
};
Ok(ToolResponse {
content: vec![ToolResponseContent::Text {
text: serde_json::to_string(&response)?,
}],
structured_content: None,
})
}
}
#[derive(Clone)]
pub struct ReadTool {
thread_rx: watch::Receiver<WeakEntity<AcpThread>>,
}
impl McpServerTool for ReadTool {
type Input = ReadToolParams;
const NAME: &'static str = "Read";
fn description(&self) -> &'static str {
"Read the contents of a file. In sessions with mcp__zed__Read always use it instead of Read as it contains the most up-to-date contents."
}
fn annotations(&self) -> ToolAnnotations {
ToolAnnotations {
title: Some("Read file".to_string()),
read_only_hint: Some(true),
destructive_hint: Some(false),
open_world_hint: Some(false),
idempotent_hint: None,
}
}
async fn run(&self, input: Self::Input, cx: &mut AsyncApp) -> Result<ToolResponse> {
let mut thread_rx = self.thread_rx.clone();
let Some(thread) = thread_rx.recv().await?.upgrade() else {
anyhow::bail!("Thread closed");
};
let content = thread
.update(cx, |thread, cx| {
thread.read_text_file(input.abs_path, input.offset, input.limit, false, cx)
})?
.await?;
Ok(ToolResponse {
content: vec![ToolResponseContent::Text { text: content }],
structured_content: None,
})
}
}
#[derive(Clone)]
pub struct EditTool {
thread_rx: watch::Receiver<WeakEntity<AcpThread>>,
}
impl McpServerTool for EditTool {
type Input = EditToolParams;
const NAME: &'static str = "Edit";
fn description(&self) -> &'static str {
"Edits a file. In sessions with mcp__zed__Edit always use it instead of Edit as it will show the diff to the user better."
}
fn annotations(&self) -> ToolAnnotations {
ToolAnnotations {
title: Some("Edit file".to_string()),
read_only_hint: Some(false),
destructive_hint: Some(false),
open_world_hint: Some(false),
idempotent_hint: Some(false),
}
}
async fn run(&self, input: Self::Input, cx: &mut AsyncApp) -> Result<ToolResponse> {
let mut thread_rx = self.thread_rx.clone();
let Some(thread) = thread_rx.recv().await?.upgrade() else {
anyhow::bail!("Thread closed");
};
let content = thread
.update(cx, |thread, cx| {
thread.read_text_file(input.abs_path.clone(), None, None, true, cx)
})?
.await?;
let new_content = content.replace(&input.old_text, &input.new_text);
if new_content == content {
return Err(anyhow::anyhow!("The old_text was not found in the content"));
}
thread
.update(cx, |thread, cx| {
thread.write_text_file(input.abs_path, new_content, cx)
})?
.await?;
Ok(ToolResponse {
content: vec![],
structured_content: None,
})
}
}