8a96ea25c4
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 --------- Co-authored-by: Marshall Bowers <elliott.codes@gmail.com>
69 lines
2.2 KiB
Rust
69 lines
2.2 KiB
Rust
use std::sync::Arc;
|
|
|
|
use collections::HashMap;
|
|
use gpui::{AppContext, Global, ReadGlobal};
|
|
use parking_lot::RwLock;
|
|
|
|
struct GlobalContextServerRegistry(Arc<ContextServerRegistry>);
|
|
|
|
impl Global for GlobalContextServerRegistry {}
|
|
|
|
pub struct ContextServerRegistry {
|
|
command_registry: RwLock<HashMap<String, Vec<Arc<str>>>>,
|
|
tool_registry: RwLock<HashMap<String, Vec<Arc<str>>>>,
|
|
}
|
|
|
|
impl ContextServerRegistry {
|
|
pub fn global(cx: &AppContext) -> Arc<Self> {
|
|
GlobalContextServerRegistry::global(cx).0.clone()
|
|
}
|
|
|
|
pub fn register(cx: &mut AppContext) {
|
|
cx.set_global(GlobalContextServerRegistry(Arc::new(
|
|
ContextServerRegistry {
|
|
command_registry: RwLock::new(HashMap::default()),
|
|
tool_registry: RwLock::new(HashMap::default()),
|
|
},
|
|
)))
|
|
}
|
|
|
|
pub fn register_command(&self, server_id: String, command_name: &str) {
|
|
let mut registry = self.command_registry.write();
|
|
registry
|
|
.entry(server_id)
|
|
.or_default()
|
|
.push(command_name.into());
|
|
}
|
|
|
|
pub fn unregister_command(&self, server_id: &str, command_name: &str) {
|
|
let mut registry = self.command_registry.write();
|
|
if let Some(commands) = registry.get_mut(server_id) {
|
|
commands.retain(|name| name.as_ref() != command_name);
|
|
}
|
|
}
|
|
|
|
pub fn get_commands(&self, server_id: &str) -> Option<Vec<Arc<str>>> {
|
|
let registry = self.command_registry.read();
|
|
registry.get(server_id).cloned()
|
|
}
|
|
|
|
pub fn register_tool(&self, server_id: String, tool_name: &str) {
|
|
let mut registry = self.tool_registry.write();
|
|
registry
|
|
.entry(server_id)
|
|
.or_default()
|
|
.push(tool_name.into());
|
|
}
|
|
|
|
pub fn unregister_tool(&self, server_id: &str, tool_name: &str) {
|
|
let mut registry = self.tool_registry.write();
|
|
if let Some(tools) = registry.get_mut(server_id) {
|
|
tools.retain(|name| name.as_ref() != tool_name);
|
|
}
|
|
}
|
|
|
|
pub fn get_tools(&self, server_id: &str) -> Option<Vec<Arc<str>>> {
|
|
let registry = self.tool_registry.read();
|
|
registry.get(server_id).cloned()
|
|
}
|
|
}
|