History

context_servers: Support tools (#19548)

This PR depends on #19547 This PR adds support for tools from context servers. Context servers are free to expose tools that Zed can pass to models. When called by the model, Zed forwards the request to context servers. This allows for some interesting techniques. Context servers can easily expose tools such as querying local databases, reading or writing local files, reading resources over authenticated APIs (e.g. kubernetes, asana, etc). This is currently experimental. Things to discuss * I want to still add a confirm dialog asking people if a server is allows to use the tool. Should do this or just use the tool and assume trustworthyness of context servers? * Can we add tool use behind a local setting flag? Release Notes: - N/A ---------

David Soria Parra committed 2 years ago

Commit 8a96ea21 parent cdddb4d7 files · +235 −28Browse files

Unified · Split

7 files
  1. crates/assistant/src/assistant.rs +66 −19
  2. crates/assistant/src/tools.rs +1 −0
  3. crates/assistant/src/tools/context_server_tool.rs +82 −0
  4. crates/context_servers/src/protocol.rs +33 −0
  5. crates/context_servers/src/registry.rs +27 −5
  6. crates/context_servers/src/types.rs +18 −0
  7. crates/language_model/src/provider/anthropic.rs +8 −4
crates/assistant/src/assistant.rs+66−19
@@ -298,25 +298,64 @@ fn register_context_server_handlers(cx: &mut AppContext) { · 297 unmodified lines
298298 return;
299299 };
300300
301− if let Some(prompts) = protocol.list_prompts().await.log_err() {
302− for prompt in prompts
303− .into_iter()
304− .filter(context_server_command::acceptable_prompt)
305− {
306− log::info!(
307− "registering context server command: {:?}",
308− prompt.name
309− );
310− context_server_registry.register_command(
311− server.id.clone(),
312− prompt.name.as_str(),
313− );
314− slash_command_registry.register_command(
315− context_server_command::ContextServerSlashCommand::new(
316− &server, prompt,
317− ),
318− true,
319− );
301+ if protocol.capable(context_servers::protocol::ServerCapability::Prompts) {
302+ if let Some(prompts) = protocol.list_prompts().await.log_err() {
303+ for prompt in prompts
304+ .into_iter()
305+ .filter(context_server_command::acceptable_prompt)
306+ {
307+ log::info!(
308+ "registering context server command: {:?}",
309+ prompt.name
310+ );
311+ context_server_registry.register_command(
312+ server.id.clone(),
313+ prompt.name.as_str(),
314+ );
315+ slash_command_registry.register_command(
316+ context_server_command::ContextServerSlashCommand::new(
317+ &server, prompt,
318+ ),
319+ true,
320+ );
321+ }
322+ }
323+ }
324+ })
325+ .detach();
326+ }
327+ },
328+ );
329+
330+ cx.update_model(
331+ &manager,
332+ |manager: &mut context_servers::manager::ContextServerManager, cx| {
333+ let tool_registry = ToolRegistry::global(cx);
334+ let context_server_registry = ContextServerRegistry::global(cx);
335+ if let Some(server) = manager.get_server(server_id) {
336+ cx.spawn(|_, _| async move {
337+ let Some(protocol) = server.client.read().clone() else {
338+ return;
339+ };
340+
341+ if protocol.capable(context_servers::protocol::ServerCapability::Tools) {
342+ if let Some(tools) = protocol.list_tools().await.log_err() {
343+ for tool in tools.tools {
344+ log::info!(
345+ "registering context server tool: {:?}",
346+ tool.name
347+ );
348+ context_server_registry.register_tool(
349+ server.id.clone(),
350+ tool.name.as_str(),
351+ );
352+ tool_registry.register_tool(
353+ tools::context_server_tool::ContextServerTool::new(
354+ server.id.clone(),
355+ tool
356+ ),
357+ );
358+ }
320359 }
321360 }
322361 })
@@ -334,6 +373,14 @@ fn register_context_server_handlers(cx: &mut AppContext) { · 11 unmodified lines
334373 context_server_registry.unregister_command(&server_id, &command_name);
335374 }
336375 }
376+
377+ if let Some(tools) = context_server_registry.get_tools(server_id) {
378+ let tool_registry = ToolRegistry::global(cx);
379+ for tool_name in tools {
380+ tool_registry.unregister_tool_by_name(&tool_name);
381+ context_server_registry.unregister_tool(&server_id, &tool_name);
382+ }
383+ }
337384 }
338385 },
339386 )
crates/assistant/src/tools.rs+1−0
@@ -1 +1,2 @@
1+pub mod context_server_tool;
12 pub mod now_tool;
crates/assistant/src/tools/context_server_tool.rs+82−0
@@ -0,0 +1,82 @@
1+use anyhow::{anyhow, bail};
2+use assistant_tool::Tool;
3+use context_servers::manager::ContextServerManager;
4+use context_servers::types;
5+use gpui::Task;
6+
7+pub struct ContextServerTool {
8+ server_id: String,
9+ tool: types::Tool,
10+}
11+
12+impl ContextServerTool {
13+ pub fn new(server_id: impl Into<String>, tool: types::Tool) -> Self {
14+ Self {
15+ server_id: server_id.into(),
16+ tool,
17+ }
18+ }
19+}
20+
21+impl Tool for ContextServerTool {
22+ fn name(&self) -> String {
23+ self.tool.name.clone()
24+ }
25+
26+ fn description(&self) -> String {
27+ self.tool.description.clone().unwrap_or_default()
28+ }
29+
30+ fn input_schema(&self) -> serde_json::Value {
31+ match &self.tool.input_schema {
32+ serde_json::Value::Null => {
33+ serde_json::json!({ "type": "object", "properties": [] })
34+ }
35+ serde_json::Value::Object(map) if map.is_empty() => {
36+ serde_json::json!({ "type": "object", "properties": [] })
37+ }
38+ _ => self.tool.input_schema.clone(),
39+ }
40+ }
41+
42+ fn run(
43+ self: std::sync::Arc<Self>,
44+ input: serde_json::Value,
45+ _workspace: gpui::WeakView<workspace::Workspace>,
46+ cx: &mut ui::WindowContext,
47+ ) -> gpui::Task<gpui::Result<String>> {
48+ let manager = ContextServerManager::global(cx);
49+ let manager = manager.read(cx);
50+ if let Some(server) = manager.get_server(&self.server_id) {
51+ cx.foreground_executor().spawn({
52+ let tool_name = self.tool.name.clone();
53+ async move {
54+ let Some(protocol) = server.client.read().clone() else {
55+ bail!("Context server not initialized");
56+ };
57+
58+ let arguments = if let serde_json::Value::Object(map) = input {
59+ Some(map.into_iter().collect())
60+ } else {
61+ None
62+ };
63+
64+ log::trace!(
65+ "Running tool: {} with arguments: {:?}",
66+ tool_name,
67+ arguments
68+ );
69+ let response = protocol.run_tool(tool_name, arguments).await?;
70+
71+ let tool_result = match response.tool_result {
72+ serde_json::Value::String(s) => s,
73+ _ => serde_json::to_string(&response.tool_result)?,
74+ };
75+ Ok(tool_result)
76+ }
77+ })
78+ } else {
79+ Task::ready(Err(anyhow!("Context server not found")))
80+ }
81+ }
82+}
crates/context_servers/src/protocol.rs+33−0
@@ -180,6 +180,39 @@ impl InitializedContextServerProtocol { · 179 unmodified lines
180180
181181 Ok(completion)
182182 }
183+
184+ /// List MCP tools.
185+ pub async fn list_tools(&self) -> Result<types::ListToolsResponse> {
186+ self.check_capability(ServerCapability::Tools)?;
187+
188+ let response = self
189+ .inner
190+ .request::<types::ListToolsResponse>(types::RequestType::ListTools.as_str(), ())
191+ .await?;
192+
193+ Ok(response)
194+ }
195+
196+ /// Executes a tool with the given arguments
197+ pub async fn run_tool<P: AsRef<str>>(
198+ &self,
199+ tool: P,
200+ arguments: Option<HashMap<String, serde_json::Value>>,
201+ ) -> Result<types::CallToolResponse> {
202+ self.check_capability(ServerCapability::Tools)?;
203+
204+ let params = types::CallToolParams {
205+ name: tool.as_ref().to_string(),
206+ arguments,
207+ };
208+
209+ let response: types::CallToolResponse = self
210+ .inner
211+ .request(types::RequestType::CallTool.as_str(), params)
212+ .await?;
213+
214+ Ok(response)
215+ }
183216 }
184217
185218 impl InitializedContextServerProtocol {
crates/context_servers/src/registry.rs+27−5
@@ -9,7 +9,8 @@ struct GlobalContextServerRegistry(Arc<ContextServerRegistry>); · 8 unmodified lines
99 impl Global for GlobalContextServerRegistry {}
1010
1111 pub struct ContextServerRegistry {
12− registry: RwLock<HashMap<String, Vec<Arc<str>>>>,
12+ command_registry: RwLock<HashMap<String, Vec<Arc<str>>>>,
13+ tool_registry: RwLock<HashMap<String, Vec<Arc<str>>>>,
1314 }
1415
1516 impl ContextServerRegistry {
@@ -20,13 +21,14 @@ impl ContextServerRegistry { · 4 unmodified lines
2021 pub fn register(cx: &mut AppContext) {
2122 cx.set_global(GlobalContextServerRegistry(Arc::new(
2223 ContextServerRegistry {
23− registry: RwLock::new(HashMap::default()),
24+ command_registry: RwLock::new(HashMap::default()),
25+ tool_registry: RwLock::new(HashMap::default()),
2426 },
2527 )))
2628 }
2729
2830 pub fn register_command(&self, server_id: String, command_name: &str) {
29− let mut registry = self.registry.write();
31+ let mut registry = self.command_registry.write();
3032 registry
3133 .entry(server_id)
3234 .or_default()
@@ -34,14 +36,34 @@ impl ContextServerRegistry { · 1 unmodified line
3436 }
3537
3638 pub fn unregister_command(&self, server_id: &str, command_name: &str) {
37− let mut registry = self.registry.write();
39+ let mut registry = self.command_registry.write();
3840 if let Some(commands) = registry.get_mut(server_id) {
3941 commands.retain(|name| name.as_ref() != command_name);
4042 }
4143 }
4244
4345 pub fn get_commands(&self, server_id: &str) -> Option<Vec<Arc<str>>> {
44− let registry = self.registry.read();
46+ let registry = self.command_registry.read();
47+ registry.get(server_id).cloned()
48+ }
49+
50+ pub fn register_tool(&self, server_id: String, tool_name: &str) {
51+ let mut registry = self.tool_registry.write();
52+ registry
53+ .entry(server_id)
54+ .or_default()
55+ .push(tool_name.into());
56+ }
57+
58+ pub fn unregister_tool(&self, server_id: &str, tool_name: &str) {
59+ let mut registry = self.tool_registry.write();
60+ if let Some(tools) = registry.get_mut(server_id) {
61+ tools.retain(|name| name.as_ref() != tool_name);
62+ }
63+ }
64+
65+ pub fn get_tools(&self, server_id: &str) -> Option<Vec<Arc<str>>> {
66+ let registry = self.tool_registry.read();
4567 registry.get(server_id).cloned()
4668 }
4769 }
crates/context_servers/src/types.rs+18−0
@@ -16,6 +16,8 @@ pub enum RequestType { · 15 unmodified lines
1616 PromptsList,
1717 CompletionComplete,
1818 Ping,
19+ ListTools,
20+ ListResourceTemplates,
1921 }
2022
2123 impl RequestType {
@@ -32,6 +34,8 @@ impl RequestType { · 10 unmodified lines
3234 RequestType::PromptsList => "prompts/list",
3335 RequestType::CompletionComplete => "completion/complete",
3436 RequestType::Ping => "ping",
37+ RequestType::ListTools => "tools/list",
38+ RequestType::ListResourceTemplates => "resources/templates/list",
3539 }
3640 }
3741 }
@@ -402,3 +406,17 @@ pub struct Completion { · 364 unmodified lines
402406 pub values: Vec<String>,
403407 pub total: CompletionTotal,
404408 }
409+
410+#[derive(Debug, Deserialize)]
411+#[serde(rename_all = "camelCase")]
412+pub struct CallToolResponse {
413+ pub tool_result: serde_json::Value,
414+}
415+
416+#[derive(Debug, Deserialize)]
417+#[serde(rename_all = "camelCase")]
418+pub struct ListToolsResponse {
419+ pub tools: Vec<Tool>,
420+ #[serde(skip_serializing_if = "Option::is_none")]
421+ pub next_cursor: Option<String>,
422+}
crates/language_model/src/provider/anthropic.rs+8−4
@@ -505,10 +505,14 @@ pub fn map_to_language_model_completion_events( · 504 unmodified lines
505505 LanguageModelToolUse {
506506 id: tool_use.id,
507507 name: tool_use.name,
508− input: serde_json::Value::from_str(
509− &tool_use.input_json,
510− )
511− .map_err(|err| anyhow!(err))?,
508+ input: if tool_use.input_json.is_empty() {
509+ serde_json::Value::Null
510+ } else {
511+ serde_json::Value::from_str(
512+ &tool_use.input_json,
513+ )
514+ .map_err(|err| anyhow!(err))?
515+ },
512516 },
513517 ))
514518 })),