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 {
378                        proxy: self,
379                        tool: tool.clone(),
380                    }) as Box<dyn ExecutableTool>,
381                )
382            })
383            .collect()
384    }
385
386    pub fn write_cli_client(
387        &self,
388        output_dir: impl AsRef<Path>,
389        name: Option<&str>,
390    ) -> Result<GeneratedClient, Error> {
391        self.write_client(GeneratedClientKind::Cli, output_dir, name)
392    }
393
394    pub fn generate_code_client(
395        &self,
396        language: CodeLanguage,
397        output_dir: impl AsRef<Path>,
398        name: Option<&str>,
399    ) -> Result<GeneratedClient, Error> {
400        let kind = match language {
401            CodeLanguage::Python => GeneratedClientKind::Python,
402            CodeLanguage::TypeScript => GeneratedClientKind::TypeScript,
403        };
404        self.generate_client(kind, output_dir, name)
405    }
406
407    pub fn write_code_client(
408        &self,
409        language: CodeLanguage,
410        output_dir: impl AsRef<Path>,
411        name: Option<&str>,
412    ) -> Result<GeneratedClient, Error> {
413        let generated = self.generate_code_client(language, output_dir, name)?;
414        generated.write_to_disk()
415    }
416
417    pub fn generate_client(
418        &self,
419        kind: GeneratedClientKind,
420        output_dir: impl AsRef<Path>,
421        name: Option<&str>,
422    ) -> Result<GeneratedClient, Error> {
423        let generator_config = self.generator_config(output_dir, name);
424        let artifacts = match kind {
425            GeneratedClientKind::Cli => CliGenerator.render(&generator_config),
426            GeneratedClientKind::Python => PythonGenerator.render(&generator_config),
427            GeneratedClientKind::TypeScript => TypeScriptGenerator.render(&generator_config),
428        }?;
429        let file_contents = artifact_map(&artifacts);
430        let files = artifacts
431            .iter()
432            .map(|artifact| generator_config.output_dir.join(&artifact.file_name))
433            .collect();
434        let environment = kind.environment(&generator_config);
435        Ok(GeneratedClient {
436            kind,
437            output_dir: generator_config.output_dir,
438            files,
439            file_contents,
440            artifacts,
441            environment,
442        })
443    }
444
445    pub fn write_client(
446        &self,
447        kind: GeneratedClientKind,
448        output_dir: impl AsRef<Path>,
449        name: Option<&str>,
450    ) -> Result<GeneratedClient, Error> {
451        let generated = self.generate_client(kind, output_dir, name)?;
452        generated.write_to_disk()
453    }
454
455    fn generator_config(
456        &self,
457        output_dir: impl AsRef<Path>,
458        name: Option<&str>,
459    ) -> GeneratorConfig {
460        GeneratorConfig {
461            cli_name: name
462                .or(self.default_server.as_deref())
463                .unwrap_or("mcp")
464                .to_string(),
465            bridge_url: self.bridge_url().to_string(),
466            token: self.token().to_string(),
467            tools: self.backend_tools.clone(),
468            session_pid: 0,
469            output_dir: output_dir.as_ref().to_path_buf(),
470            extra_cli_bridges: Vec::new(),
471        }
472    }
473
474    fn invoke_wrapper(&self, server: Option<&str>) -> Result<String, Error> {
475        let suffix = "_invoke_tool";
476        let matches = self
477            .frontend_tools
478            .iter()
479            .filter(|tool| tool.name.ends_with(suffix))
480            .filter(|tool| {
481                server
482                    .map(|name| tool.name == format!("{name}{suffix}"))
483                    .unwrap_or(true)
484            })
485            .map(|tool| tool.name.clone())
486            .collect::<Vec<_>>();
487        match matches.as_slice() {
488            [name] => Ok(name.clone()),
489            [] => Err(Error::Config(format!(
490                "No compressed invoke wrapper found for server {}",
491                server.unwrap_or("<default>")
492            ))),
493            _ => Err(Error::Config(
494                "Multiple compressed invoke wrappers found; specify a server".to_string(),
495            )),
496        }
497    }
498}
499
500fn ensure_loopback_proxy_url(url: &str) -> Result<(), Error> {
501    let parsed = reqwest::Url::parse(url)
502        .map_err(|error| Error::Config(format!("invalid proxy URL {url:?}: {error}")))?;
503    if parsed.scheme() != "http" {
504        return Err(Error::Config(format!(
505            "refusing to send proxy session token to non-http URL: {url}"
506        )));
507    }
508    let Some(host) = parsed.host_str() else {
509        return Err(Error::Config(format!("proxy URL has no host: {url}")));
510    };
511    let is_loopback_ip = host
512        .parse::<std::net::IpAddr>()
513        .map(|addr| addr.is_loopback())
514        .unwrap_or(false);
515    if host != "localhost" && !is_loopback_ip {
516        return Err(Error::Config(format!(
517            "refusing to send proxy session token to non-loopback URL: {url}"
518        )));
519    }
520    Ok(())
521}
522
523#[async_trait]
524pub trait ExecutableTool: Send + Sync {
525    fn name(&self) -> &str;
526    fn description(&self) -> Option<&str>;
527    fn input_schema(&self) -> &Value;
528    async fn execute(&self, input: Value) -> Result<String, Error>;
529}
530
531struct ProxyExecutableTool<'a> {
532    proxy: &'a CompressorProxy,
533    tool: Tool,
534}
535
536#[async_trait]
537impl ExecutableTool for ProxyExecutableTool<'_> {
538    fn name(&self) -> &str {
539        &self.tool.name
540    }
541
542    fn description(&self) -> Option<&str> {
543        self.tool.description.as_deref()
544    }
545
546    fn input_schema(&self) -> &Value {
547        &self.tool.input_schema
548    }
549
550    async fn execute(&self, input: Value) -> Result<String, Error> {
551        self.proxy.invoke_wrapper_tool(&self.tool.name, input).await
552    }
553}
554
555#[derive(Debug, Clone, Copy, PartialEq, Eq)]
556pub enum CodeLanguage {
557    Python,
558    TypeScript,
559}
560
561#[derive(Debug, Clone, PartialEq, Eq)]
562pub struct GeneratedClient {
563    pub kind: GeneratedClientKind,
564    pub output_dir: PathBuf,
565    /// Paths where artifacts would be written, or were written after `write_to_disk`.
566    pub files: Vec<PathBuf>,
567    /// Generated artifact contents keyed by file name.
568    pub file_contents: BTreeMap<String, String>,
569    pub environment: HashMap<String, String>,
570    artifacts: Vec<GeneratedArtifact>,
571}
572
573impl GeneratedClient {
574    pub fn write_to_disk(mut self) -> Result<Self, Error> {
575        self.files = write_artifacts(&self.artifacts, &self.output_dir)?;
576        Ok(self)
577    }
578}
579
580#[derive(Debug, Clone, Copy, PartialEq, Eq)]
581pub enum GeneratedClientKind {
582    Cli,
583    Python,
584    TypeScript,
585}
586
587impl GeneratedClientKind {
588    fn environment(self, config: &GeneratorConfig) -> HashMap<String, String> {
589        match self {
590            GeneratedClientKind::Python => HashMap::from([(
591                "PYTHONPATH".to_string(),
592                config.output_dir.to_string_lossy().to_string(),
593            )]),
594            GeneratedClientKind::Cli | GeneratedClientKind::TypeScript => HashMap::new(),
595        }
596    }
597}
598
599#[cfg(test)]
600mod tests {
601    use serde_json::json;
602
603    use super::*;
604
605    fn fixture_path(name: &str) -> String {
606        format!("{}/tests/fixtures/{name}", env!("CARGO_MANIFEST_DIR"))
607    }
608
609    fn python_command() -> String {
610        std::env::var("PYTHON").unwrap_or_else(|_| "python3".to_string())
611    }
612
613    #[test]
614    fn server_config_oauth_app_name_is_preserved_for_transport_layer() {
615        let config = ServerConfig::url("https://example.test/mcp")
616            .oauth_app_name("Rovo Dev")
617            .materialize()
618            .0;
619
620        match config {
621            FfiSdkServerConfig::Structured { oauth_app_name, .. } => {
622                assert_eq!(oauth_app_name.as_deref(), Some("Rovo Dev"));
623            }
624            FfiSdkServerConfig::CommandOrUrl(_) => panic!("expected structured config"),
625        }
626    }
627
628    #[test]
629    fn loopback_proxy_url_guard_allows_only_local_http_urls() {
630        assert!(ensure_loopback_proxy_url("http://127.0.0.1:1234/exec").is_ok());
631        assert!(ensure_loopback_proxy_url("http://localhost:1234/exec").is_ok());
632        assert!(ensure_loopback_proxy_url("https://127.0.0.1:1234/exec").is_err());
633        assert!(ensure_loopback_proxy_url("http://example.com:1234/exec").is_err());
634    }
635
636    #[test]
637    fn server_config_auth_provider_is_preserved_for_transport_layer() {
638        let (config, provider) = ServerConfig::url("https://example.test/mcp")
639            .header("X-Static", "yes")
640            .auth_provider(|| {
641                Ok(BTreeMap::from([(
642                    "Authorization".to_string(),
643                    "Bearer dynamic".to_string(),
644                )]))
645            })
646            .materialize();
647
648        let backends = normalize_sdk_servers(FfiSdkServersConfig::from_iter([(
649            "remote".to_string(),
650            config,
651        )]))
652        .unwrap();
653
654        assert_eq!(backends[0].command_or_url, "https://example.test/mcp");
655        assert_eq!(
656            backends[0].args,
657            ["-H", "X-Static=yes", "--auth", "explicit-headers"]
658        );
659        assert!(provider.is_some());
660    }
661
662    #[tokio::test]
663    async fn compressor_client_invokes_single_server_without_compressor_subprocess() {
664        let client = CompressorClient::builder()
665            .server(
666                "alpha",
667                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
668            )
669            .compression_level(CompressionLevel::Max)
670            .build();
671        let proxy = client.connect().await.unwrap();
672        assert!(proxy
673            .tools()
674            .iter()
675            .any(|tool| tool.name == "alpha_invoke_tool"));
676        let result = proxy
677            .invoke("echo", json!({ "message": "rust-sdk" }))
678            .await
679            .unwrap();
680        assert_eq!(result, "alpha:rust-sdk");
681
682        let executable = proxy.executable_tools();
683        let invoke = executable.get("alpha_invoke_tool").unwrap();
684        let executable_result = invoke
685            .execute(json!({
686                "tool_name": "echo",
687                "tool_input": { "message": "executable-rust" }
688            }))
689            .await
690            .unwrap();
691        assert_eq!(executable_result, "alpha:executable-rust");
692    }
693
694    #[tokio::test]
695    async fn compressor_client_routes_multiple_servers() {
696        let client = CompressorClient::builder()
697            .server(
698                "alpha",
699                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
700            )
701            .server(
702                "beta",
703                ServerConfig::command(python_command()).arg(fixture_path("beta_server.py")),
704            )
705            .compression_level(CompressionLevel::Max)
706            .build();
707        let proxy = client.connect().await.unwrap();
708        let alpha = proxy
709            .invoke_on(Some("alpha"), "add", json!({ "a": 2, "b": 3 }))
710            .await
711            .unwrap();
712        let beta = proxy
713            .invoke_on(Some("beta"), "multiply", json!({ "a": 4, "b": 5 }))
714            .await
715            .unwrap();
716        assert_eq!(alpha, "5");
717        assert_eq!(beta, "20");
718    }
719
720    #[tokio::test]
721    async fn compressor_client_writes_generated_clients() {
722        let client = CompressorClient::builder()
723            .server(
724                "alpha",
725                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
726            )
727            .compression_level(CompressionLevel::Max)
728            .build();
729        let proxy = client.connect().await.unwrap();
730        let tempdir = tempfile::tempdir().unwrap();
731        let generated = proxy
732            .write_code_client(CodeLanguage::Python, tempdir.path(), Some("alpha"))
733            .unwrap();
734        assert_eq!(generated.kind, GeneratedClientKind::Python);
735        assert!(generated
736            .files
737            .iter()
738            .any(|path| path.ends_with("alpha.py")));
739        assert_eq!(
740            generated.environment.get("PYTHONPATH"),
741            Some(&tempdir.path().to_string_lossy().to_string())
742        );
743
744        let cli = proxy
745            .write_cli_client(tempdir.path(), Some("alpha"))
746            .unwrap();
747        assert_eq!(cli.kind, GeneratedClientKind::Cli);
748    }
749
750    #[tokio::test]
751    async fn compressor_client_exposes_cli_and_just_bash_modes() {
752        let cli = CompressorClient::builder()
753            .server(
754                "alpha",
755                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
756            )
757            .mode(CompressorMode::Cli)
758            .build()
759            .connect()
760            .await
761            .unwrap();
762        assert!(cli.tools().iter().any(|tool| tool.name == "alpha_help"));
763
764        let bash = CompressorClient::builder()
765            .server(
766                "alpha",
767                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
768            )
769            .mode(CompressorMode::JustBash)
770            .build()
771            .connect()
772            .await
773            .unwrap();
774        assert!(bash.tools().iter().any(|tool| tool.name == "bash_tool"));
775        assert!(bash.tools().iter().any(|tool| tool.name == "alpha_help"));
776        let provider = bash
777            .just_bash_providers()
778            .iter()
779            .find(|provider| provider.provider_name == "alpha")
780            .unwrap();
781        assert_eq!(provider.help_tool_name, "alpha_help");
782        assert!(provider.tools.iter().any(|command| {
783            command.command_name == "echo"
784                && command.backend_tool_name == "echo"
785                && command.invoke_tool_name == "alpha_invoke_tool"
786        }));
787    }
788}
mcp_compressor_core/sdk.rs

mcp_compressor_core/
sdk.rs
