Skip to main content

mcp_compressor_core/ffi/
session.rs

1use crate::proxy::{RunningToolProxy, ToolProxyServer};
2use crate::server::{CompressedServer, CompressedServerConfig, ProxyTransformMode};
3use crate::Error;
4
5use super::dto::{
6    FfiBackendConfig, FfiCompressedSessionConfig, FfiCompressedSessionInfo, FfiSdkServerConfig,
7    FfiSdkServersConfig, FfiTool,
8};
9
10pub fn normalize_sdk_servers(servers: FfiSdkServersConfig) -> Result<Vec<FfiBackendConfig>, Error> {
11    servers
12        .into_iter()
13        .map(|(name, config)| normalize_sdk_server(name, config))
14        .collect()
15}
16
17fn normalize_sdk_server(
18    name: String,
19    config: FfiSdkServerConfig,
20) -> Result<FfiBackendConfig, Error> {
21    match config {
22        FfiSdkServerConfig::CommandOrUrl(command_or_url) => Ok(FfiBackendConfig {
23            name,
24            command_or_url,
25            args: Vec::new(),
26            oauth_app_name: None,
27        }),
28        FfiSdkServerConfig::Structured {
29            command,
30            url,
31            mut args,
32            headers,
33            oauth_app_name,
34        } => {
35            let command_or_url = url
36                .or(command)
37                .ok_or_else(|| Error::Config(format!("server {name} must define command or url")))?;
38            if !headers.is_empty() {
39                let mut header_args = Vec::new();
40                for (key, value) in headers {
41                    header_args.push("-H".to_string());
42                    header_args.push(format!("{key}={value}"));
43                }
44                if !args.iter().any(|arg| arg == "--auth") {
45                    header_args.push("--auth".to_string());
46                    header_args.push("explicit-headers".to_string());
47                }
48                header_args.extend(args);
49                args = header_args;
50            }
51            Ok(FfiBackendConfig {
52                name,
53                command_or_url,
54                args,
55                oauth_app_name,
56            })
57        }
58    }
59}
60
61pub struct FfiCompressedSession {
62    info: FfiCompressedSessionInfo,
63    _proxy: RunningToolProxy,
64}
65
66impl FfiCompressedSession {
67    pub fn bridge_url(&self) -> &str {
68        &self.info.bridge_url
69    }
70
71    pub fn token(&self) -> &str {
72        &self.info.token
73    }
74
75    pub fn info(&self) -> FfiCompressedSessionInfo {
76        self.info.clone()
77    }
78
79    pub fn close(self) {}
80}
81
82fn parse_ffi_transform_mode(value: Option<&str>) -> Result<ProxyTransformMode, Error> {
83    match value.unwrap_or("compressed-tools") {
84        "compressed-tools" | "compressed" | "normal" => Ok(ProxyTransformMode::CompressedTools),
85        "cli" | "cli-mode" => Ok(ProxyTransformMode::Cli),
86        "just-bash" | "just_bash" => Ok(ProxyTransformMode::JustBash),
87        other => Err(Error::Config(format!("invalid transform mode: {other}"))),
88    }
89}
90
91async fn compressed_session_from_server(
92    server: CompressedServer,
93) -> Result<FfiCompressedSession, Error> {
94    let frontend_tools = server
95        .list_frontend_tools()
96        .await?
97        .into_iter()
98        .map(FfiTool::from)
99        .collect();
100    let backend_tools = server.backend_tools().into_iter().map(FfiTool::from).collect();
101    let backend_tools_by_server = server
102        .backend_tools_by_server()
103        .into_iter()
104        .map(|(server_name, tool)| super::dto::FfiBackendTool {
105            server_name,
106            tool: FfiTool::from(tool),
107        })
108        .collect();
109    let just_bash_providers = server
110        .just_bash_provider_specs()
111        .into_iter()
112        .map(Into::into)
113        .collect();
114    let proxy = ToolProxyServer::start(server).await?;
115    Ok(FfiCompressedSession {
116        info: FfiCompressedSessionInfo {
117            bridge_url: proxy.bridge_url().to_string(),
118            token: proxy.token_value().to_string(),
119            frontend_tools,
120            backend_tools,
121            backend_tools_by_server,
122            just_bash_providers,
123        },
124        _proxy: proxy,
125    })
126}
127
128pub async fn start_compressed_session(
129    config: FfiCompressedSessionConfig,
130    backends: Vec<FfiBackendConfig>,
131) -> Result<FfiCompressedSession, Error> {
132    start_compressed_session_with_backend_configs(
133        config,
134        backends.into_iter().map(Into::into).collect(),
135    )
136    .await
137}
138
139pub async fn start_compressed_session_with_backend_configs(
140    config: FfiCompressedSessionConfig,
141    backends: Vec<crate::server::BackendServerConfig>,
142) -> Result<FfiCompressedSession, Error> {
143    let server = CompressedServer::connect_multi_stdio(
144        CompressedServerConfig {
145            level: config.compression_level.parse()?,
146            server_name: config.server_name,
147            include_tools: config.include_tools,
148            exclude_tools: config.exclude_tools,
149            toonify: config.toonify,
150            transform_mode: parse_ffi_transform_mode(config.transform_mode.as_deref())?,
151            ..CompressedServerConfig::default()
152        },
153        backends,
154    )
155    .await?;
156    compressed_session_from_server(server).await
157}
158
159pub async fn start_compressed_session_from_mcp_config(
160    config: FfiCompressedSessionConfig,
161    mcp_config_json: &str,
162) -> Result<FfiCompressedSession, Error> {
163    let server = CompressedServer::connect_mcp_config_json(
164        CompressedServerConfig {
165            level: config.compression_level.parse()?,
166            server_name: config.server_name,
167            include_tools: config.include_tools,
168            exclude_tools: config.exclude_tools,
169            toonify: config.toonify,
170            transform_mode: parse_ffi_transform_mode(config.transform_mode.as_deref())?,
171            ..CompressedServerConfig::default()
172        },
173        mcp_config_json,
174    )
175    .await?;
176    compressed_session_from_server(server).await
177}