Skip to main content

sac/mcp/
registry.rs

1use super::*;
2
3#[derive(Clone)]
4pub struct McpRegistry {
5    tools: Arc<HashMap<String, Arc<McpToolBinding>>>,
6}
7
8#[derive(Clone)]
9struct McpToolBinding {
10    tool_name: String,
11    definition: ToolDefinition,
12    server: Arc<McpServer>,
13}
14
15struct McpServer {
16    _service: Arc<McpService>,
17    tool_call_timeout: Duration,
18}
19
20#[derive(Clone)]
21pub(super) struct NacMcpClientHandler {
22    root_uri: String,
23    root_name: String,
24}
25
26impl McpRegistry {
27    pub async fn load(cwd: &Path, sandbox: Option<&SandboxSession>) -> Result<Option<Arc<Self>>> {
28        let Some(path) = default_config_path() else {
29            return Ok(None);
30        };
31        if !path.exists() {
32            return Ok(None);
33        }
34
35        let raw = match std::fs::read_to_string(&path) {
36            Ok(raw) => raw,
37            Err(error) => {
38                eprintln!(
39                    "MCP config at '{}' could not be read; MCP will be disabled: {:#}",
40                    path.display(),
41                    error
42                );
43                return Ok(None);
44            }
45        };
46        let config: McpConfigFile = match toml::from_str(&raw) {
47            Ok(config) => config,
48            Err(error) => {
49                eprintln!(
50                    "MCP config at '{}' is invalid; MCP will be disabled: {:#}",
51                    path.display(),
52                    error
53                );
54                return Ok(None);
55            }
56        };
57
58        let root_uri = if sandbox.is_some() {
59            "file:///workspace".to_string()
60        } else {
61            Url::from_directory_path(cwd)
62                .map_err(|_| anyhow!("failed to build file:// root for {}", cwd.display()))?
63                .to_string()
64        };
65        let root_name = if sandbox.is_some() {
66            "workspace".to_string()
67        } else {
68            cwd.file_name()
69                .and_then(|value| value.to_str())
70                .unwrap_or("workspace")
71                .to_string()
72        };
73
74        let handler = NacMcpClientHandler {
75            root_uri,
76            root_name,
77        };
78
79        let mut tools = HashMap::new();
80        let mut seen_names = HashMap::<String, usize>::new();
81
82        for (server_name, server_config) in config.mcp_servers {
83            if !server_config.enabled {
84                continue;
85            }
86
87            let service = match timeout(
88                MCP_CONNECT_TIMEOUT,
89                connect_server(&server_name, &server_config, &handler, sandbox),
90            )
91            .await
92            {
93                Ok(Ok(service)) => Arc::new(service),
94                Ok(Err(error)) => {
95                    eprintln!(
96                        "MCP server '{}' is unavailable and will be skipped: {:#}",
97                        server_name, error
98                    );
99                    continue;
100                }
101                Err(_) => {
102                    eprintln!(
103                        "MCP server '{}' timed out during connect after {}s and will be skipped",
104                        server_name,
105                        MCP_CONNECT_TIMEOUT.as_secs()
106                    );
107                    continue;
108                }
109            };
110
111            let listed_tools = match timeout(MCP_TOOL_INVENTORY_TIMEOUT, service.list_all_tools())
112                .await
113            {
114                Ok(Ok(tools)) => tools,
115                Ok(Err(error)) => {
116                    eprintln!(
117                        "MCP server '{}' could not list tools and will be skipped: {:#}",
118                        server_name, error
119                    );
120                    continue;
121                }
122                Err(_) => {
123                    eprintln!(
124                        "MCP server '{}' timed out while listing tools after {}s and will be skipped",
125                        server_name,
126                        MCP_TOOL_INVENTORY_TIMEOUT.as_secs()
127                    );
128                    continue;
129                }
130            };
131
132            let tool_call_timeout = server_config
133                .tool_call_timeout_secs
134                .map(Duration::from_secs)
135                .unwrap_or(MCP_TOOL_CALL_TIMEOUT);
136            let server = Arc::new(McpServer {
137                _service: service.clone(),
138                tool_call_timeout,
139            });
140            for tool in listed_tools {
141                let qualified_name = allocate_tool_name(&server_name, &tool.name, &mut seen_names);
142                let definition = tool_definition(&qualified_name, &server_name, &tool);
143                tools.insert(
144                    qualified_name,
145                    Arc::new(McpToolBinding {
146                        tool_name: tool.name.to_string(),
147                        definition,
148                        server: server.clone(),
149                    }),
150                );
151            }
152        }
153
154        if tools.is_empty() {
155            return Ok(None);
156        }
157
158        Ok(Some(Arc::new(Self {
159            tools: Arc::new(tools),
160        })))
161    }
162
163    pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
164        let mut definitions: Vec<ToolDefinition> = self
165            .tools
166            .values()
167            .map(|binding| binding.definition.clone())
168            .collect();
169        definitions.sort_by(|left, right| left.function.name.cmp(&right.function.name));
170        definitions
171    }
172
173    pub async fn call_tool(&self, name: &str, args: Value) -> ToolResult {
174        let Some(binding) = self.tools.get(name) else {
175            return ToolResult {
176                content: format!("Error: unknown MCP tool '{}'", name),
177                is_error: true,
178            };
179        };
180
181        let arguments = match args {
182            Value::Object(map) => Some(map),
183            Value::Null => None,
184            _ => {
185                return ToolResult {
186                    content: format!("Error: MCP tool '{}' requires object arguments", name),
187                    is_error: true,
188                }
189            }
190        };
191
192        let mut params = CallToolRequestParams::new(binding.tool_name.clone());
193        if let Some(arguments) = arguments {
194            params = params.with_arguments(arguments);
195        }
196        let tool_timeout = binding.server.tool_call_timeout;
197        match timeout(
198            tool_timeout,
199            binding.server._service.call_tool(params),
200        )
201        .await
202        {
203            Ok(Ok(result)) => flatten_tool_result(result),
204            Ok(Err(error)) => ToolResult {
205                content: format!("Error calling MCP tool '{}': {}", name, error),
206                is_error: true,
207            },
208            Err(_) => ToolResult {
209                content: format!(
210                    "Error calling MCP tool '{}': timed out after {}s",
211                    name,
212                    tool_timeout.as_secs()
213                ),
214                is_error: true,
215            },
216        }
217    }
218
219    pub fn has_tool(&self, name: &str) -> bool {
220        self.tools.contains_key(name)
221    }
222}
223
224impl ClientHandler for NacMcpClientHandler {
225    fn get_info(&self) -> ClientInfo {
226        ClientInfo::new(
227            serde_json::from_value(serde_json::json!({
228                "roots": {
229                    "listChanged": true
230                }
231            }))
232            .expect("valid MCP client capabilities"),
233            Implementation::new("sac", env!("CARGO_PKG_VERSION")),
234        )
235    }
236
237    async fn list_roots(
238        &self,
239        _request_context: rmcp::service::RequestContext<RoleClient>,
240    ) -> std::result::Result<ListRootsResult, rmcp::model::ErrorData> {
241        Ok(ListRootsResult::new(vec![
242            Root::new(self.root_uri.clone()).with_name(self.root_name.clone())
243        ]))
244    }
245}
246
247pub(super) fn tool_definition(full_name: &str, server_name: &str, tool: &Tool) -> ToolDefinition {
248    let description = tool
249        .description
250        .as_ref()
251        .map(|value| value.to_string())
252        .unwrap_or_else(|| format!("MCP tool '{}' from server '{}'", tool.name, server_name));
253    ToolDefinition {
254        def_type: "function".to_string(),
255        function: FunctionDef {
256            name: full_name.to_string(),
257            description,
258            parameters: tool.schema_as_json_value(),
259        },
260    }
261}