Skip to main content

wraith_runtime/
mcp.rs

1use crate::config::{McpServerConfig, ScopedMcpServerConfig};
2
3const CLAUDEAI_SERVER_PREFIX: &str = "claude.ai ";
4const CCR_PROXY_PATH_MARKERS: [&str; 2] = ["/v2/session_ingress/shttp/mcp/", "/v2/ccr-sessions/"];
5
6#[must_use]
7pub fn normalize_name_for_mcp(name: &str) -> String {
8    let mut normalized = name
9        .chars()
10        .map(|ch| match ch {
11            'a'..='z' | 'A'..='Z' | '0'..='9' | '_' | '-' => ch,
12            _ => '_',
13        })
14        .collect::<String>();
15
16    if name.starts_with(CLAUDEAI_SERVER_PREFIX) {
17        normalized = collapse_underscores(&normalized)
18            .trim_matches('_')
19            .to_string();
20    }
21
22    normalized
23}
24
25#[must_use]
26pub fn mcp_tool_prefix(server_name: &str) -> String {
27    format!("mcp__{}__", normalize_name_for_mcp(server_name))
28}
29
30#[must_use]
31pub fn mcp_tool_name(server_name: &str, tool_name: &str) -> String {
32    format!(
33        "{}{}",
34        mcp_tool_prefix(server_name),
35        normalize_name_for_mcp(tool_name)
36    )
37}
38
39#[must_use]
40pub fn unwrap_ccr_proxy_url(url: &str) -> String {
41    if !CCR_PROXY_PATH_MARKERS
42        .iter()
43        .any(|marker| url.contains(marker))
44    {
45        return url.to_string();
46    }
47
48    let Some(query_start) = url.find('?') else {
49        return url.to_string();
50    };
51    let query = &url[query_start + 1..];
52    for pair in query.split('&') {
53        let mut parts = pair.splitn(2, '=');
54        if matches!(parts.next(), Some("mcp_url")) {
55            if let Some(value) = parts.next() {
56                return percent_decode(value);
57            }
58        }
59    }
60
61    url.to_string()
62}
63
64#[must_use]
65pub fn mcp_server_signature(config: &McpServerConfig) -> Option<String> {
66    match config {
67        McpServerConfig::Stdio(config) => {
68            let mut command = vec![config.command.clone()];
69            command.extend(config.args.clone());
70            Some(format!("stdio:{}", render_command_signature(&command)))
71        }
72        McpServerConfig::Sse(config) | McpServerConfig::Http(config) => {
73            Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url)))
74        }
75        McpServerConfig::Ws(config) => Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url))),
76        McpServerConfig::ManagedProxy(config) => {
77            Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url)))
78        }
79        McpServerConfig::Sdk(_) => None,
80    }
81}
82
83#[must_use]
84pub fn scoped_mcp_config_hash(config: &ScopedMcpServerConfig) -> String {
85    let rendered = match &config.config {
86        McpServerConfig::Stdio(stdio) => format!(
87            "stdio|{}|{}|{}",
88            stdio.command,
89            render_command_signature(&stdio.args),
90            render_env_signature(&stdio.env)
91        ),
92        McpServerConfig::Sse(remote) => format!(
93            "sse|{}|{}|{}|{}",
94            remote.url,
95            render_env_signature(&remote.headers),
96            remote.headers_helper.as_deref().unwrap_or(""),
97            render_oauth_signature(remote.oauth.as_ref())
98        ),
99        McpServerConfig::Http(remote) => format!(
100            "http|{}|{}|{}|{}",
101            remote.url,
102            render_env_signature(&remote.headers),
103            remote.headers_helper.as_deref().unwrap_or(""),
104            render_oauth_signature(remote.oauth.as_ref())
105        ),
106        McpServerConfig::Ws(ws) => format!(
107            "ws|{}|{}|{}",
108            ws.url,
109            render_env_signature(&ws.headers),
110            ws.headers_helper.as_deref().unwrap_or("")
111        ),
112        McpServerConfig::Sdk(sdk) => format!("sdk|{}", sdk.name),
113        McpServerConfig::ManagedProxy(proxy) => {
114            format!("claudeai-proxy|{}|{}", proxy.url, proxy.id)
115        }
116    };
117    stable_hex_hash(&rendered)
118}
119
120fn render_command_signature(command: &[String]) -> String {
121    let escaped = command
122        .iter()
123        .map(|part| part.replace('\\', "\\\\").replace('|', "\\|"))
124        .collect::<Vec<_>>();
125    format!("[{}]", escaped.join("|"))
126}
127
128fn render_env_signature(map: &std::collections::BTreeMap<String, String>) -> String {
129    map.iter()
130        .map(|(key, value)| format!("{key}={value}"))
131        .collect::<Vec<_>>()
132        .join(";")
133}
134
135fn render_oauth_signature(oauth: Option<&crate::config::McpOAuthConfig>) -> String {
136    oauth.map_or_else(String::new, |oauth| {
137        format!(
138            "{}|{}|{}|{}",
139            oauth.client_id.as_deref().unwrap_or(""),
140            oauth
141                .callback_port
142                .map_or_else(String::new, |port| port.to_string()),
143            oauth.auth_server_metadata_url.as_deref().unwrap_or(""),
144            oauth.xaa.map_or_else(String::new, |flag| flag.to_string())
145        )
146    })
147}
148
149fn stable_hex_hash(value: &str) -> String {
150    let mut hash = 0xcbf2_9ce4_8422_2325_u64;
151    for byte in value.as_bytes() {
152        hash ^= u64::from(*byte);
153        hash = hash.wrapping_mul(0x0100_0000_01b3);
154    }
155    format!("{hash:016x}")
156}
157
158fn collapse_underscores(value: &str) -> String {
159    let mut collapsed = String::with_capacity(value.len());
160    let mut last_was_underscore = false;
161    for ch in value.chars() {
162        if ch == '_' {
163            if !last_was_underscore {
164                collapsed.push(ch);
165            }
166            last_was_underscore = true;
167        } else {
168            collapsed.push(ch);
169            last_was_underscore = false;
170        }
171    }
172    collapsed
173}
174
175fn percent_decode(value: &str) -> String {
176    let bytes = value.as_bytes();
177    let mut decoded = Vec::with_capacity(bytes.len());
178    let mut index = 0;
179    while index < bytes.len() {
180        match bytes[index] {
181            b'%' if index + 2 < bytes.len() => {
182                let hex = &value[index + 1..index + 3];
183                if let Ok(byte) = u8::from_str_radix(hex, 16) {
184                    decoded.push(byte);
185                    index += 3;
186                    continue;
187                }
188                decoded.push(bytes[index]);
189                index += 1;
190            }
191            b'+' => {
192                decoded.push(b' ');
193                index += 1;
194            }
195            byte => {
196                decoded.push(byte);
197                index += 1;
198            }
199        }
200    }
201    String::from_utf8_lossy(&decoded).into_owned()
202}
203
204#[cfg(test)]
205mod tests {
206    use std::collections::BTreeMap;
207
208    use crate::config::{
209        ConfigSource, McpRemoteServerConfig, McpServerConfig, McpStdioServerConfig,
210        McpWebSocketServerConfig, ScopedMcpServerConfig,
211    };
212
213    use super::{
214        mcp_server_signature, mcp_tool_name, normalize_name_for_mcp, scoped_mcp_config_hash,
215        unwrap_ccr_proxy_url,
216    };
217
218    #[test]
219    fn normalizes_server_names_for_mcp_tooling() {
220        assert_eq!(normalize_name_for_mcp("github.com"), "github_com");
221        assert_eq!(normalize_name_for_mcp("tool name!"), "tool_name_");
222        assert_eq!(
223            normalize_name_for_mcp("claude.ai Example   Server!!"),
224            "claude_ai_Example_Server"
225        );
226        assert_eq!(
227            mcp_tool_name("claude.ai Example Server", "weather tool"),
228            "mcp__claude_ai_Example_Server__weather_tool"
229        );
230    }
231
232    #[test]
233    fn unwraps_ccr_proxy_urls_for_signature_matching() {
234        let wrapped = "https://api.anthropic.com/v2/session_ingress/shttp/mcp/123?mcp_url=https%3A%2F%2Fvendor.example%2Fmcp&other=1";
235        assert_eq!(unwrap_ccr_proxy_url(wrapped), "https://vendor.example/mcp");
236        assert_eq!(
237            unwrap_ccr_proxy_url("https://vendor.example/mcp"),
238            "https://vendor.example/mcp"
239        );
240    }
241
242    #[test]
243    fn computes_signatures_for_stdio_and_remote_servers() {
244        let stdio = McpServerConfig::Stdio(McpStdioServerConfig {
245            command: "uvx".to_string(),
246            args: vec!["mcp-server".to_string()],
247            env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]),
248        });
249        assert_eq!(
250            mcp_server_signature(&stdio),
251            Some("stdio:[uvx|mcp-server]".to_string())
252        );
253
254        let remote = McpServerConfig::Ws(McpWebSocketServerConfig {
255            url: "https://api.anthropic.com/v2/ccr-sessions/1?mcp_url=wss%3A%2F%2Fvendor.example%2Fmcp".to_string(),
256            headers: BTreeMap::new(),
257            headers_helper: None,
258        });
259        assert_eq!(
260            mcp_server_signature(&remote),
261            Some("url:wss://vendor.example/mcp".to_string())
262        );
263    }
264
265    #[test]
266    fn scoped_hash_ignores_scope_but_tracks_config_content() {
267        let base_config = McpServerConfig::Http(McpRemoteServerConfig {
268            url: "https://vendor.example/mcp".to_string(),
269            headers: BTreeMap::from([("Authorization".to_string(), "Bearer token".to_string())]),
270            headers_helper: Some("helper.sh".to_string()),
271            oauth: None,
272        });
273        let user = ScopedMcpServerConfig {
274            scope: ConfigSource::User,
275            config: base_config.clone(),
276        };
277        let local = ScopedMcpServerConfig {
278            scope: ConfigSource::Local,
279            config: base_config,
280        };
281        assert_eq!(
282            scoped_mcp_config_hash(&user),
283            scoped_mcp_config_hash(&local)
284        );
285
286        let changed = ScopedMcpServerConfig {
287            scope: ConfigSource::Local,
288            config: McpServerConfig::Http(McpRemoteServerConfig {
289                url: "https://vendor.example/v2/mcp".to_string(),
290                headers: BTreeMap::new(),
291                headers_helper: None,
292                oauth: None,
293            }),
294        };
295        assert_ne!(
296            scoped_mcp_config_hash(&user),
297            scoped_mcp_config_hash(&changed)
298        );
299    }
300}