mcp_compressor_core/
sdk.rs

1use std::collections::{BTreeMap, HashMap};
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use async_trait::async_trait;
6
7use serde_json::{json, Value};
8
9use crate::client_gen::cli::CliGenerator;
10use crate::client_gen::generator::{
11    artifact_map, write_artifacts, ClientGenerator, GeneratedArtifact, GeneratorConfig,
12};
13use crate::client_gen::python::PythonGenerator;
14use crate::client_gen::typescript::TypeScriptGenerator;
15use crate::compression::engine::Tool;
16use crate::compression::CompressionLevel;
17use crate::ffi::{normalize_sdk_servers, FfiSdkServerConfig, FfiSdkServersConfig};
18use crate::proxy::{RunningToolProxy, ToolProxyServer};
19use crate::server::{BackendAuthMode, BackendServerConfig};
20use crate::server::{CompressedServer, CompressedServerConfig, ProxyTransformMode};
21use crate::Error;
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum CompressorMode {
25    CompressedTools,
26    Cli,
27    JustBash,
28}
29
30impl From<CompressorMode> for ProxyTransformMode {
31    fn from(value: CompressorMode) -> Self {
32        match value {
33            CompressorMode::CompressedTools => Self::CompressedTools,
34            CompressorMode::Cli => Self::Cli,
35            CompressorMode::JustBash => Self::JustBash,
36        }
37    }
38}
39
40type HeaderProvider = Arc<dyn Fn() -> Result<BTreeMap<String, String>, Error> + Send + Sync>;
41
42#[derive(Clone)]
43pub struct ServerConfig {
44    inner: FfiSdkServerConfig,
45    auth_provider: Option<HeaderProvider>,
46    oauth_app_name: Option<String>,
47}
48
49impl ServerConfig {
50    pub fn command(command: impl Into<String>) -> Self {
51        Self {
52            inner: FfiSdkServerConfig::Structured {
53                command: Some(command.into()),
54                url: None,
55                args: Vec::new(),
56                headers: BTreeMap::new(),
57                oauth_app_name: None,
58            },
59            auth_provider: None,
60            oauth_app_name: None,
61        }
62    }
63
64    pub fn url(url: impl Into<String>) -> Self {
65        Self {
66            inner: FfiSdkServerConfig::Structured {
67                command: None,
68                url: Some(url.into()),
69                args: Vec::new(),
70                headers: BTreeMap::new(),
71                oauth_app_name: None,
72            },
73            auth_provider: None,
74            oauth_app_name: None,
75        }
76    }
77
78    pub fn arg(mut self, arg: impl Into<String>) -> Self {
79        if let FfiSdkServerConfig::Structured { args, .. } = &mut self.inner {
80            args.push(arg.into());
81        }
82        self
83    }
84
85    pub fn args(mut self, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
86        if let FfiSdkServerConfig::Structured { args: stored, .. } = &mut self.inner {
87            stored.extend(args.into_iter().map(Into::into));
88        }
89        self
90    }
91
92    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
93        if let FfiSdkServerConfig::Structured { headers, .. } = &mut self.inner {
94            headers.insert(name.into(), value.into());
95        }
96        self
97    }
98
99    pub fn auth_provider(
100        mut self,
101        provider: impl Fn() -> Result<BTreeMap<String, String>, Error> + Send + Sync + 'static,
102    ) -> Self {
103        self.auth_provider = Some(Arc::new(provider));
104        self
105    }
106
107    pub fn oauth_app_name(mut self, app_name: impl Into<String>) -> Self {
108        self.oauth_app_name = Some(app_name.into());
109        self
110    }
111
112    fn materialize(mut self) -> (FfiSdkServerConfig, Option<HeaderProvider>) {
113        if let (FfiSdkServerConfig::Structured { oauth_app_name, .. }, Some(app_name)) =
114            (&mut self.inner, self.oauth_app_name.take())
115        {
116            *oauth_app_name = Some(app_name);
117        }
118        (self.inner, self.auth_provider.take())
119    }
120}
121
122#[derive(Clone)]
123pub struct CompressorClientBuilder {
124    servers: BTreeMap<String, ServerConfig>,
125    compression_level: CompressionLevel,
126    server_name: Option<String>,
127    include_tools: Vec<String>,
128    exclude_tools: Vec<String>,
129    toonify: bool,
130    mode: CompressorMode,
131}
132
133impl Default for CompressorClientBuilder {
134    fn default() -> Self {
135        Self {
136            servers: BTreeMap::new(),
137            compression_level: CompressionLevel::Max,
138            server_name: None,
139            include_tools: Vec::new(),
140            exclude_tools: Vec::new(),
141            toonify: false,
142            mode: CompressorMode::CompressedTools,
143        }
144    }
145}
146
147impl CompressorClientBuilder {
148    pub fn server(mut self, name: impl Into<String>, config: ServerConfig) -> Self {
149        self.servers.insert(name.into(), config);
150        self
151    }
152
153    pub fn compression_level(mut self, level: CompressionLevel) -> Self {
154        self.compression_level = level;
155        self
156    }
157
158    pub fn server_name(mut self, server_name: impl Into<String>) -> Self {
159        self.server_name = Some(server_name.into());
160        self
161    }
162
163    pub fn include_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
164        self.include_tools = tools.into_iter().map(Into::into).collect();
165        self
166    }
167
168    pub fn exclude_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
169        self.exclude_tools = tools.into_iter().map(Into::into).collect();
170        self
171    }
172
173    pub fn toonify(mut self, enabled: bool) -> Self {
174        self.toonify = enabled;
175        self
176    }
177
178    pub fn mode(mut self, mode: CompressorMode) -> Self {
179        self.mode = mode;
180        self
181    }
182
183    pub fn build(self) -> CompressorClient {
184        CompressorClient { builder: self }
185    }
186}
187
188#[derive(Clone)]
189pub struct CompressorClient {
190    builder: CompressorClientBuilder,
191}
192
193impl CompressorClient {
194    pub fn builder() -> CompressorClientBuilder {
195        CompressorClientBuilder::default()
196    }
197
198    pub async fn connect(&self) -> Result<CompressorProxy, Error> {
199        let materialized = self
200            .builder
201            .servers
202            .clone()
203            .into_iter()
204            .map(|(name, config)| {
205                let (config, provider) = config.materialize();
206                (name, config, provider)
207            })
208            .collect::<Vec<_>>();
209        let providers = materialized
210            .iter()
211            .filter_map(|(name, _, provider)| {
212                provider.clone().map(|provider| (name.clone(), provider))
213            })
214            .collect::<BTreeMap<_, _>>();
215        let ffi_configs = materialized
216            .into_iter()
217            .map(|(name, config, _)| (name, config))
218            .collect::<Vec<_>>();
219        let backends = normalize_sdk_servers(FfiSdkServersConfig::from_iter(ffi_configs))?;
220        let backends = backends
221            .into_iter()
222            .map(|backend| {
223                let name = backend.name.clone();
224                let mut backend = BackendServerConfig::from(backend);
225                if let Some(provider) = providers.get(&name) {
226                    backend = backend
227                        .with_header_provider(Arc::clone(provider))
228                        .with_auth_mode(BackendAuthMode::ExplicitHeaders);
229                }
230                backend
231            })
232            .collect::<Vec<_>>();
233        let server = CompressedServer::connect_multi_stdio(
234            CompressedServerConfig {
235                level: self.builder.compression_level.clone(),
236                server_name: self.builder.server_name.clone(),
237                include_tools: self.builder.include_tools.clone(),
238                exclude_tools: self.builder.exclude_tools.clone(),
239                toonify: self.builder.toonify,
240                transform_mode: self.builder.mode.into(),
241                ..CompressedServerConfig::default()
242            },
243            backends,
244        )
245        .await?;
246        CompressorProxy::start(server).await
247    }
248}
249
250pub struct CompressorProxy {
251    default_server: Option<String>,
252    frontend_tools: Vec<Tool>,
253    backend_tools: Vec<Tool>,
254    backend_tools_by_server: Vec<(String, Tool)>,
255    just_bash_providers: Vec<crate::server::JustBashProviderSpec>,
256    proxy: RunningToolProxy,
257}
258
259impl CompressorProxy {
260    async fn start(server: CompressedServer) -> Result<Self, Error> {
261        let default_server = server.default_server_name().map(str::to_string);
262        let frontend_tools = server.list_frontend_tools().await?;
263        let backend_tools = server.backend_tools();
264        let backend_tools_by_server = server.backend_tools_by_server();
265        let just_bash_providers = server.just_bash_provider_specs();
266        let proxy = ToolProxyServer::start(server).await?;
267        Ok(Self {
268            default_server,
269            frontend_tools,
270            backend_tools,
271            backend_tools_by_server,
272            just_bash_providers,
273            proxy,
274        })
275    }
276
277    pub fn bridge_url(&self) -> &str {
278        self.proxy.bridge_url()
279    }
280
281    pub fn token(&self) -> &str {
282        self.proxy.token_value()
283    }
284
285    pub fn tools(&self) -> &[Tool] {
286        &self.frontend_tools
287    }
288
289    pub fn backend_tools(&self) -> &[Tool] {
290        &self.backend_tools
291    }
292
293    pub fn just_bash_providers(&self) -> &[crate::server::JustBashProviderSpec] {
294        &self.just_bash_providers
295    }
296
297    pub fn schema(&self, tool_name: &str) -> Result<Value, Error> {
298        self.schema_on(self.default_server.as_deref(), tool_name)
299    }
300
301    pub fn schema_on(&self, server: Option<&str>, tool_name: &str) -> Result<Value, Error> {
302        let matches = self
303            .backend_tools_by_server
304            .iter()
305            .filter(|(server_name, tool)| {
306                tool.name == tool_name && server.map(|server| server == server_name).unwrap_or(true)
307            })
308            .collect::<Vec<_>>();
309        match matches.as_slice() {
310            [(_, tool)] => Ok(tool.input_schema.clone()),
311            [] => Err(Error::ToolNotFound(tool_name.to_string())),
312            _ => Err(Error::Config(
313                "Multiple backend tools matched; specify a server".to_string(),
314            )),
315        }
316    }
317
318    pub async fn invoke(&self, tool_name: &str, input: Value) -> Result<String, Error> {
319        self.invoke_on(self.default_server.as_deref(), tool_name, input)
320            .await
321    }
322
323    pub async fn invoke_on(
324        &self,
325        server: Option<&str>,
326        tool_name: &str,
327        input: Value,
328    ) -> Result<String, Error> {
329        let wrapper = self.invoke_wrapper(server)?;
330        self.invoke_wrapper_tool(
331            &wrapper,
332            json!({
333                "tool_name": tool_name,
334                "tool_input": input,
335            }),
336        )
337        .await
338    }
339
340    async fn invoke_wrapper_tool(&self, wrapper: &str, input: Value) -> Result<String, Error> {
341        let exec_url = self.proxy.exec_url();
342        ensure_loopback_proxy_url(&exec_url)?;
343        let client = reqwest::Client::new();
344        let response = client
345            // codeql[rust/cleartext-transmission]
346            // This session token is sent only to the local loopback proxy after
347            // ensure_loopback_proxy_url rejects non-loopback and non-http URLs.
348            .post(exec_url)
349            .header("Authorization", format!("Bearer {}", self.token()))
350            .json(&json!({
351                "tool": wrapper,
352                "input": input
353            }))
354            .send()
355            .await
356            .map_err(|error| Error::Config(format!("proxy request failed: {error}")))?;
357        let status = response.status();
358        let text = response
359            .text()
360            .await
361            .map_err(|error| Error::Config(format!("proxy response failed: {error}")))?;
362        if status.is_success() {
363            Ok(text)
364        } else {
365            Err(Error::Config(format!(
366                "proxy request failed with {status}: {text}"
367            )))
368        }
369    }
370
371    pub fn executable_tools(&self) -> BTreeMap<String, Box<dyn ExecutableTool + '_>> {
372        self.frontend_tools
373            .iter()
374            .map(|tool| {
375                (
376                    tool.name.clone(),
377                    Box::new(ProxyExecutableTool { proxy: self, tool: tool.clone() })
378                        as Box<dyn ExecutableTool>,
379                )
380            })
381            .collect()
382    }
383
384    pub fn write_cli_client(
385        &self,
386        output_dir: impl AsRef<Path>,
387        name: Option<&str>,
388    ) -> Result<GeneratedClient, Error> {
389        self.write_client(GeneratedClientKind::Cli, output_dir, name)
390    }
391
392    pub fn generate_code_client(
393        &self,
394        language: CodeLanguage,
395        output_dir: impl AsRef<Path>,
396        name: Option<&str>,
397    ) -> Result<GeneratedClient, Error> {
398        let kind = match language {
399            CodeLanguage::Python => GeneratedClientKind::Python,
400            CodeLanguage::TypeScript => GeneratedClientKind::TypeScript,
401        };
402        self.generate_client(kind, output_dir, name)
403    }
404
405    pub fn write_code_client(
406        &self,
407        language: CodeLanguage,
408        output_dir: impl AsRef<Path>,
409        name: Option<&str>,
410    ) -> Result<GeneratedClient, Error> {
411        let generated = self.generate_code_client(language, output_dir, name)?;
412        generated.write_to_disk()
413    }
414
415    pub fn generate_client(
416        &self,
417        kind: GeneratedClientKind,
418        output_dir: impl AsRef<Path>,
419        name: Option<&str>,
420    ) -> Result<GeneratedClient, Error> {
421        let generator_config = self.generator_config(output_dir, name);
422        let artifacts = match kind {
423            GeneratedClientKind::Cli => CliGenerator.render(&generator_config),
424            GeneratedClientKind::Python => PythonGenerator.render(&generator_config),
425            GeneratedClientKind::TypeScript => TypeScriptGenerator.render(&generator_config),
426        }?;
427        let file_contents = artifact_map(&artifacts);
428        let files = artifacts
429            .iter()
430            .map(|artifact| generator_config.output_dir.join(&artifact.file_name))
431            .collect();
432        let environment = kind.environment(&generator_config);
433        Ok(GeneratedClient {
434            kind,
435            output_dir: generator_config.output_dir,
436            files,
437            file_contents,
438            artifacts,
439            environment,
440        })
441    }
442
443    pub fn write_client(
444        &self,
445        kind: GeneratedClientKind,
446        output_dir: impl AsRef<Path>,
447        name: Option<&str>,
448    ) -> Result<GeneratedClient, Error> {
449        let generated = self.generate_client(kind, output_dir, name)?;
450        generated.write_to_disk()
451    }
452
453    fn generator_config(&self, output_dir: impl AsRef<Path>, name: Option<&str>) -> GeneratorConfig {
454        GeneratorConfig {
455            cli_name: name
456                .or(self.default_server.as_deref())
457                .unwrap_or("mcp")
458                .to_string(),
459            bridge_url: self.bridge_url().to_string(),
460            token: self.token().to_string(),
461            tools: self.backend_tools.clone(),
462            session_pid: 0,
463            output_dir: output_dir.as_ref().to_path_buf(),
464        }
465    }
466
467    fn invoke_wrapper(&self, server: Option<&str>) -> Result<String, Error> {
468        let suffix = "_invoke_tool";
469        let matches = self
470            .frontend_tools
471            .iter()
472            .filter(|tool| tool.name.ends_with(suffix))
473            .filter(|tool| {
474                server
475                    .map(|name| tool.name == format!("{name}{suffix}"))
476                    .unwrap_or(true)
477            })
478            .map(|tool| tool.name.clone())
479            .collect::<Vec<_>>();
480        match matches.as_slice() {
481            [name] => Ok(name.clone()),
482            [] => Err(Error::Config(format!(
483                "No compressed invoke wrapper found for server {}",
484                server.unwrap_or("<default>")
485            ))),
486            _ => Err(Error::Config(
487                "Multiple compressed invoke wrappers found; specify a server".to_string(),
488            )),
489        }
490    }
491}
492
493fn ensure_loopback_proxy_url(url: &str) -> Result<(), Error> {
494    let parsed = reqwest::Url::parse(url)
495        .map_err(|error| Error::Config(format!("invalid proxy URL {url:?}: {error}")))?;
496    if parsed.scheme() != "http" {
497        return Err(Error::Config(format!(
498            "refusing to send proxy session token to non-http URL: {url}"
499        )));
500    }
501    let Some(host) = parsed.host_str() else {
502        return Err(Error::Config(format!("proxy URL has no host: {url}")));
503    };
504    let is_loopback_ip = host
505        .parse::<std::net::IpAddr>()
506        .map(|addr| addr.is_loopback())
507        .unwrap_or(false);
508    if host != "localhost" && !is_loopback_ip {
509        return Err(Error::Config(format!(
510            "refusing to send proxy session token to non-loopback URL: {url}"
511        )));
512    }
513    Ok(())
514}
515
516#[async_trait]
517pub trait ExecutableTool: Send + Sync {
518    fn name(&self) -> &str;
519    fn description(&self) -> Option<&str>;
520    fn input_schema(&self) -> &Value;
521    async fn execute(&self, input: Value) -> Result<String, Error>;
522}
523
524struct ProxyExecutableTool<'a> {
525    proxy: &'a CompressorProxy,
526    tool: Tool,
527}
528
529#[async_trait]
530impl ExecutableTool for ProxyExecutableTool<'_> {
531    fn name(&self) -> &str {
532        &self.tool.name
533    }
534
535    fn description(&self) -> Option<&str> {
536        self.tool.description.as_deref()
537    }
538
539    fn input_schema(&self) -> &Value {
540        &self.tool.input_schema
541    }
542
543    async fn execute(&self, input: Value) -> Result<String, Error> {
544        self.proxy.invoke_wrapper_tool(&self.tool.name, input).await
545    }
546}
547
548#[derive(Debug, Clone, Copy, PartialEq, Eq)]
549pub enum CodeLanguage {
550    Python,
551    TypeScript,
552}
553
554#[derive(Debug, Clone, PartialEq, Eq)]
555pub struct GeneratedClient {
556    pub kind: GeneratedClientKind,
557    pub output_dir: PathBuf,
558    /// Paths where artifacts would be written, or were written after `write_to_disk`.
559    pub files: Vec<PathBuf>,
560    /// Generated artifact contents keyed by file name.
561    pub file_contents: BTreeMap<String, String>,
562    pub environment: HashMap<String, String>,
563    artifacts: Vec<GeneratedArtifact>,
564}
565
566impl GeneratedClient {
567    pub fn write_to_disk(mut self) -> Result<Self, Error> {
568        self.files = write_artifacts(&self.artifacts, &self.output_dir)?;
569        Ok(self)
570    }
571}
572
573#[derive(Debug, Clone, Copy, PartialEq, Eq)]
574pub enum GeneratedClientKind {
575    Cli,
576    Python,
577    TypeScript,
578}
579
580impl GeneratedClientKind {
581    fn environment(self, config: &GeneratorConfig) -> HashMap<String, String> {
582        match self {
583            GeneratedClientKind::Python => HashMap::from([(
584                "PYTHONPATH".to_string(),
585                config.output_dir.to_string_lossy().to_string(),
586            )]),
587            GeneratedClientKind::Cli | GeneratedClientKind::TypeScript => HashMap::new(),
588        }
589    }
590}
591
592#[cfg(test)]
593mod tests {
594    use serde_json::json;
595
596    use super::*;
597
598    fn fixture_path(name: &str) -> String {
599        format!("{}/tests/fixtures/{name}", env!("CARGO_MANIFEST_DIR"))
600    }
601
602    fn python_command() -> String {
603        std::env::var("PYTHON").unwrap_or_else(|_| "python3".to_string())
604    }
605
606    #[test]
607    fn server_config_oauth_app_name_is_preserved_for_transport_layer() {
608        let config = ServerConfig::url("https://example.test/mcp")
609            .oauth_app_name("Rovo Dev")
610            .materialize()
611            .0;
612
613        match config {
614            FfiSdkServerConfig::Structured { oauth_app_name, .. } => {
615                assert_eq!(oauth_app_name.as_deref(), Some("Rovo Dev"));
616            }
617            FfiSdkServerConfig::CommandOrUrl(_) => panic!("expected structured config"),
618        }
619    }
620
621    #[test]
622    fn loopback_proxy_url_guard_allows_only_local_http_urls() {
623        assert!(ensure_loopback_proxy_url("http://127.0.0.1:1234/exec").is_ok());
624        assert!(ensure_loopback_proxy_url("http://localhost:1234/exec").is_ok());
625        assert!(ensure_loopback_proxy_url("https://127.0.0.1:1234/exec").is_err());
626        assert!(ensure_loopback_proxy_url("http://example.com:1234/exec").is_err());
627    }
628
629    #[test]
630    fn server_config_auth_provider_is_preserved_for_transport_layer() {
631        let (config, provider) = ServerConfig::url("https://example.test/mcp")
632            .header("X-Static", "yes")
633            .auth_provider(|| {
634                Ok(BTreeMap::from([(
635                    "Authorization".to_string(),
636                    "Bearer dynamic".to_string(),
637                )]))
638            })
639            .materialize();
640
641        let backends = normalize_sdk_servers(FfiSdkServersConfig::from_iter([(
642            "remote".to_string(),
643            config,
644        )]))
645        .unwrap();
646
647        assert_eq!(backends[0].command_or_url, "https://example.test/mcp");
648        assert_eq!(
649            backends[0].args,
650            ["-H", "X-Static=yes", "--auth", "explicit-headers"]
651        );
652        assert!(provider.is_some());
653    }
654
655    #[tokio::test]
656    async fn compressor_client_invokes_single_server_without_compressor_subprocess() {
657        let client = CompressorClient::builder()
658            .server(
659                "alpha",
660                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
661            )
662            .compression_level(CompressionLevel::Max)
663            .build();
664        let proxy = client.connect().await.unwrap();
665        assert!(proxy
666            .tools()
667            .iter()
668            .any(|tool| tool.name == "alpha_invoke_tool"));
669        let result = proxy
670            .invoke("echo", json!({ "message": "rust-sdk" }))
671            .await
672            .unwrap();
673        assert_eq!(result, "alpha:rust-sdk");
674
675        let executable = proxy.executable_tools();
676        let invoke = executable.get("alpha_invoke_tool").unwrap();
677        let executable_result = invoke
678            .execute(json!({
679                "tool_name": "echo",
680                "tool_input": { "message": "executable-rust" }
681            }))
682            .await
683            .unwrap();
684        assert_eq!(executable_result, "alpha:executable-rust");
685    }
686
687    #[tokio::test]
688    async fn compressor_client_routes_multiple_servers() {
689        let client = CompressorClient::builder()
690            .server(
691                "alpha",
692                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
693            )
694            .server(
695                "beta",
696                ServerConfig::command(python_command()).arg(fixture_path("beta_server.py")),
697            )
698            .compression_level(CompressionLevel::Max)
699            .build();
700        let proxy = client.connect().await.unwrap();
701        let alpha = proxy
702            .invoke_on(Some("alpha"), "add", json!({ "a": 2, "b": 3 }))
703            .await
704            .unwrap();
705        let beta = proxy
706            .invoke_on(Some("beta"), "multiply", json!({ "a": 4, "b": 5 }))
707            .await
708            .unwrap();
709        assert_eq!(alpha, "5");
710        assert_eq!(beta, "20");
711    }
712
713    #[tokio::test]
714    async fn compressor_client_writes_generated_clients() {
715        let client = CompressorClient::builder()
716            .server(
717                "alpha",
718                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
719            )
720            .compression_level(CompressionLevel::Max)
721            .build();
722        let proxy = client.connect().await.unwrap();
723        let tempdir = tempfile::tempdir().unwrap();
724        let generated = proxy
725            .write_code_client(CodeLanguage::Python, tempdir.path(), Some("alpha"))
726            .unwrap();
727        assert_eq!(generated.kind, GeneratedClientKind::Python);
728        assert!(generated.files.iter().any(|path| path.ends_with("alpha.py")));
729        assert_eq!(
730            generated.environment.get("PYTHONPATH"),
731            Some(&tempdir.path().to_string_lossy().to_string())
732        );
733
734        let cli = proxy.write_cli_client(tempdir.path(), Some("alpha")).unwrap();
735        assert_eq!(cli.kind, GeneratedClientKind::Cli);
736    }
737
738    #[tokio::test]
739    async fn compressor_client_exposes_cli_and_just_bash_modes() {
740        let cli = CompressorClient::builder()
741            .server(
742                "alpha",
743                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
744            )
745            .mode(CompressorMode::Cli)
746            .build()
747            .connect()
748            .await
749            .unwrap();
750        assert!(cli.tools().iter().any(|tool| tool.name == "alpha_help"));
751
752        let bash = CompressorClient::builder()
753            .server(
754                "alpha",
755                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
756            )
757            .mode(CompressorMode::JustBash)
758            .build()
759            .connect()
760            .await
761            .unwrap();
762        assert!(bash.tools().iter().any(|tool| tool.name == "bash_tool"));
763        assert!(bash.tools().iter().any(|tool| tool.name == "alpha_help"));
764        let provider = bash
765            .just_bash_providers()
766            .iter()
767            .find(|provider| provider.provider_name == "alpha")
768            .unwrap();
769        assert_eq!(provider.help_tool_name, "alpha_help");
770        assert!(provider.tools.iter().any(|command| {
771            command.command_name == "echo"
772                && command.backend_tool_name == "echo"
773                && command.invoke_tool_name == "alpha_invoke_tool"
774        }));
775    }
776}
mcp_compressor_core/sdk.rs

mcp_compressor_core/
sdk.rs
