1use std::sync::Arc;
 2
 3use collections::HashMap;
 4use gpui::{AppContext, Global, ReadGlobal};
 5use parking_lot::RwLock;
 6
 7struct GlobalContextServerRegistry(Arc<ContextServerRegistry>);
 8
 9impl Global for GlobalContextServerRegistry {}
10
11pub struct ContextServerRegistry {
12    command_registry: RwLock<HashMap<String, Vec<Arc<str>>>>,
13    tool_registry: RwLock<HashMap<String, Vec<Arc<str>>>>,
14}
15
16impl ContextServerRegistry {
17    pub fn global(cx: &AppContext) -> Arc<Self> {
18        GlobalContextServerRegistry::global(cx).0.clone()
19    }
20
21    pub fn register(cx: &mut AppContext) {
22        cx.set_global(GlobalContextServerRegistry(Arc::new(
23            ContextServerRegistry {
24                command_registry: RwLock::new(HashMap::default()),
25                tool_registry: RwLock::new(HashMap::default()),
26            },
27        )))
28    }
29
30    pub fn register_command(&self, server_id: String, command_name: &str) {
31        let mut registry = self.command_registry.write();
32        registry
33            .entry(server_id)
34            .or_default()
35            .push(command_name.into());
36    }
37
38    pub fn unregister_command(&self, server_id: &str, command_name: &str) {
39        let mut registry = self.command_registry.write();
40        if let Some(commands) = registry.get_mut(server_id) {
41            commands.retain(|name| name.as_ref() != command_name);
42        }
43    }
44
45    pub fn get_commands(&self, server_id: &str) -> Option<Vec<Arc<str>>> {
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();
67        registry.get(server_id).cloned()
68    }
69}