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
298 return;298 return;
299 };299 };
300 300
301 if let Some(prompts) = protocol.list_prompts().await.log_err() {301 if protocol.capable(context_servers::protocol::ServerCapability::Prompts) {
302 for prompt in prompts302 if let Some(prompts) = protocol.list_prompts().await.log_err() {
303 .into_iter()303 for prompt in prompts
304 .filter(context_server_command::acceptable_prompt)304 .into_iter()
305 {305 .filter(context_server_command::acceptable_prompt)
306 log::info!(306 {
307 "registering context server command: {:?}",307 log::info!(
308 prompt.name308 "registering context server command: {:?}",
309 );309 prompt.name
310 context_server_registry.register_command(310 );
311 server.id.clone(),311 context_server_registry.register_command(
312 prompt.name.as_str(),312 server.id.clone(),
313 );313 prompt.name.as_str(),
314 slash_command_registry.register_command(314 );
315 context_server_command::ContextServerSlashCommand::new(315 slash_command_registry.register_command(
316 &server, prompt,316 context_server_command::ContextServerSlashCommand::new(
317 ),317 &server, prompt,
318 true,318 ),
319 );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 }
320 }359 }
321 }360 }
322 })361 })
@@ -334,6 +373,14 @@ fn register_context_server_handlers(cx: &mut AppContext) { · 11 unmodified lines
334 context_server_registry.unregister_command(&server_id, &command_name);373 context_server_registry.unregister_command(&server_id, &command_name);
335 }374 }
336 }375 }
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 }
337 }384 }
338 },385 },
339 )386 )
crates/assistant/src/tools.rs+1−0
@@ -1 +1,2 @@
1pub mod context_server_tool;
1pub mod now_tool;2pub mod now_tool;
crates/assistant/src/tools/context_server_tool.rs+82−0
@@ -0,0 +1,82 @@
1use anyhow::{anyhow, bail};
2use assistant_tool::Tool;
3use context_servers::manager::ContextServerManager;
4use context_servers::types;
5use gpui::Task;
6
7pub struct ContextServerTool {
8 server_id: String,
9 tool: types::Tool,
10}
11
12impl 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
21impl 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
180 180
181 Ok(completion)181 Ok(completion)
182 }182 }
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 }
183}216}
184 217
185impl InitializedContextServerProtocol {218impl InitializedContextServerProtocol {
crates/context_servers/src/registry.rs+27−5
@@ -9,7 +9,8 @@ struct GlobalContextServerRegistry(Arc<ContextServerRegistry>); · 8 unmodified lines
9impl Global for GlobalContextServerRegistry {}9impl Global for GlobalContextServerRegistry {}
10 10
11pub struct ContextServerRegistry {11pub 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>>>>,
13}14}
14 15
15impl ContextServerRegistry {16impl ContextServerRegistry {
@@ -20,13 +21,14 @@ impl ContextServerRegistry { · 4 unmodified lines
20 pub fn register(cx: &mut AppContext) {21 pub fn register(cx: &mut AppContext) {
21 cx.set_global(GlobalContextServerRegistry(Arc::new(22 cx.set_global(GlobalContextServerRegistry(Arc::new(
22 ContextServerRegistry {23 ContextServerRegistry {
23 registry: RwLock::new(HashMap::default()),24 command_registry: RwLock::new(HashMap::default()),
25 tool_registry: RwLock::new(HashMap::default()),
24 },26 },
25 )))27 )))
26 }28 }
27 29
28 pub fn register_command(&self, server_id: String, command_name: &str) {30 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();
30 registry32 registry
31 .entry(server_id)33 .entry(server_id)
32 .or_default()34 .or_default()
@@ -34,14 +36,34 @@ impl ContextServerRegistry { · 1 unmodified line
34 }36 }
35 37
36 pub fn unregister_command(&self, server_id: &str, command_name: &str) {38 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();
38 if let Some(commands) = registry.get_mut(server_id) {40 if let Some(commands) = registry.get_mut(server_id) {
39 commands.retain(|name| name.as_ref() != command_name);41 commands.retain(|name| name.as_ref() != command_name);
40 }42 }
41 }43 }
42 44
43 pub fn get_commands(&self, server_id: &str) -> Option<Vec<Arc<str>>> {45 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();
45 registry.get(server_id).cloned()67 registry.get(server_id).cloned()
46 }68 }
47}69}
crates/context_servers/src/types.rs+18−0
@@ -16,6 +16,8 @@ pub enum RequestType { · 15 unmodified lines
16 PromptsList,16 PromptsList,
17 CompletionComplete,17 CompletionComplete,
18 Ping,18 Ping,
19 ListTools,
20 ListResourceTemplates,
19}21}
20 22
21impl RequestType {23impl RequestType {
@@ -32,6 +34,8 @@ impl RequestType { · 10 unmodified lines
32 RequestType::PromptsList => "prompts/list",34 RequestType::PromptsList => "prompts/list",
33 RequestType::CompletionComplete => "completion/complete",35 RequestType::CompletionComplete => "completion/complete",
34 RequestType::Ping => "ping",36 RequestType::Ping => "ping",
37 RequestType::ListTools => "tools/list",
38 RequestType::ListResourceTemplates => "resources/templates/list",
35 }39 }
36 }40 }
37}41}
@@ -402,3 +406,17 @@ pub struct Completion { · 364 unmodified lines
402 pub values: Vec<String>,406 pub values: Vec<String>,
403 pub total: CompletionTotal,407 pub total: CompletionTotal,
404}408}
409
410#[derive(Debug, Deserialize)]
411#[serde(rename_all = "camelCase")]
412pub struct CallToolResponse {
413 pub tool_result: serde_json::Value,
414}
415
416#[derive(Debug, Deserialize)]
417#[serde(rename_all = "camelCase")]
418pub 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
505 LanguageModelToolUse {505 LanguageModelToolUse {
506 id: tool_use.id,506 id: tool_use.id,
507 name: tool_use.name,507 name: tool_use.name,
508 input: serde_json::Value::from_str(508 input: if tool_use.input_json.is_empty() {
509 &tool_use.input_json,509 serde_json::Value::Null
510 )510 } else {
511 .map_err(|err| anyhow!(err))?,511 serde_json::Value::from_str(
512 &tool_use.input_json,
513 )
514 .map_err(|err| anyhow!(err))?
515 },
512 },516 },
513 ))517 ))
514 })),518 })),