mcp_proxy_args/
rewrite.rs

1use std::collections::{HashMap, HashSet};
2use std::path::PathBuf;
3
4use thiserror::Error;
5
6use crate::cache::{CacheError, write_content_addressed};
7use crate::flags::{
8    IMPORT_INITIALIZE, IMPORT_INITIALIZE_FILE, IMPORT_TOOLS, IMPORT_TOOLS_FILE, long_flag,
9};
10use crate::load::{LoadError, RawFallbackImportSpec, try_load_fallback};
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct RewriteOptions {
14    pub cache_dir: PathBuf,
15    pub file_prefix: String,
16}
17
18impl RewriteOptions {
19    pub fn new(cache_dir: impl Into<PathBuf>) -> Self {
20        Self {
21            cache_dir: cache_dir.into(),
22            file_prefix: "mcp-proxy-import".to_string(),
23        }
24    }
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct RewriteResult {
29    pub args: Vec<String>,
30    pub created_files: Vec<PathBuf>,
31    pub changed: bool,
32}
33
34#[derive(Debug, Error)]
35pub enum RewriteError {
36    #[error("flag --{flag} is missing its value")]
37    MissingValue { flag: &'static str },
38    #[error("flag --{flag} was provided more than once")]
39    DuplicateFlag { flag: &'static str },
40    #[error("fallback import must include both initialize and tools metadata")]
41    IncompletePair,
42    #[error("both inline and file sources were provided for {kind}")]
43    ConflictingSources { kind: &'static str },
44    #[error("cache path is not valid UTF-8: {path}")]
45    NonUtf8CachePath { path: PathBuf },
46    #[error(transparent)]
47    Load(#[from] LoadError),
48    #[error(transparent)]
49    Cache(#[from] CacheError),
50}
51
52#[derive(Debug, Clone)]
53struct ParsedValue {
54    flag_index: usize,
55    value_index: Option<usize>,
56    value: String,
57}
58
59#[derive(Debug, Default)]
60struct ParsedImports {
61    initialize_inline: Option<ParsedValue>,
62    initialize_file: Option<ParsedValue>,
63    tools_inline: Option<ParsedValue>,
64    tools_file: Option<ParsedValue>,
65}
66
67impl ParsedImports {
68    fn insert(
69        slot: &mut Option<ParsedValue>,
70        value: ParsedValue,
71        flag: &'static str,
72    ) -> Result<(), RewriteError> {
73        if slot.is_some() {
74            return Err(RewriteError::DuplicateFlag { flag });
75        }
76        *slot = Some(value);
77        Ok(())
78    }
79
80    fn validate(&self) -> Result<(), RewriteError> {
81        if self.initialize_inline.is_some() && self.initialize_file.is_some() {
82            return Err(RewriteError::ConflictingSources { kind: "initialize" });
83        }
84        if self.tools_inline.is_some() && self.tools_file.is_some() {
85            return Err(RewriteError::ConflictingSources { kind: "tools" });
86        }
87
88        let has_initialize = self.initialize_inline.is_some() || self.initialize_file.is_some();
89        let has_tools = self.tools_inline.is_some() || self.tools_file.is_some();
90        if has_initialize != has_tools {
91            return Err(RewriteError::IncompletePair);
92        }
93        Ok(())
94    }
95}
96
97fn executable_name(command: &str) -> Option<&str> {
98    command
99        .rsplit(['/', '\\'])
100        .next()
101        .filter(|name| !name.is_empty())
102}
103
104pub fn is_mcp_proxy_convert(command: &str, args: &[String]) -> bool {
105    if !matches!(
106        executable_name(command),
107        Some("mcp-proxy" | "mcp-proxy.exe")
108    ) {
109        return false;
110    }
111
112    for arg in args {
113        if arg == "--" {
114            return false;
115        }
116        if matches!(arg.as_str(), "-v" | "--verbose" | "-q" | "--quiet")
117            || (arg.starts_with('-')
118                && !arg.starts_with("--")
119                && arg.len() > 1
120                && arg[1..].chars().all(|ch| matches!(ch, 'v' | 'q')))
121        {
122            continue;
123        }
124        return arg == "convert";
125    }
126    false
127}
128
129fn recognized_flag(arg: &str) -> Option<(&'static str, Option<&str>)> {
130    for name in [
131        IMPORT_INITIALIZE,
132        IMPORT_INITIALIZE_FILE,
133        IMPORT_TOOLS,
134        IMPORT_TOOLS_FILE,
135    ] {
136        let full = long_flag(name);
137        if arg == full {
138            return Some((name, None));
139        }
140        if let Some(value) = arg.strip_prefix(&format!("{full}=")) {
141            return Some((name, Some(value)));
142        }
143    }
144    None
145}
146
147fn parse_imports(args: &[String]) -> Result<ParsedImports, RewriteError> {
148    let mut parsed = ParsedImports::default();
149    let mut index = 0;
150
151    while index < args.len() {
152        if args[index] == "--" {
153            break;
154        }
155        let Some((flag, inline_value)) = recognized_flag(&args[index]) else {
156            index += 1;
157            continue;
158        };
159
160        let (value, value_index) = if let Some(value) = inline_value {
161            (value.to_string(), None)
162        } else {
163            let next = index + 1;
164            let value = args
165                .get(next)
166                .ok_or(RewriteError::MissingValue { flag })?
167                .clone();
168            if recognized_flag(&value).is_some() || value == "--" {
169                return Err(RewriteError::MissingValue { flag });
170            }
171            (value, Some(next))
172        };
173        let parsed_value = ParsedValue {
174            flag_index: index,
175            value_index,
176            value,
177        };
178
179        match flag {
180            IMPORT_INITIALIZE => ParsedImports::insert(
181                &mut parsed.initialize_inline,
182                parsed_value,
183                IMPORT_INITIALIZE,
184            )?,
185            IMPORT_INITIALIZE_FILE => ParsedImports::insert(
186                &mut parsed.initialize_file,
187                parsed_value,
188                IMPORT_INITIALIZE_FILE,
189            )?,
190            IMPORT_TOOLS => {
191                ParsedImports::insert(&mut parsed.tools_inline, parsed_value, IMPORT_TOOLS)?
192            }
193            IMPORT_TOOLS_FILE => {
194                ParsedImports::insert(&mut parsed.tools_file, parsed_value, IMPORT_TOOLS_FILE)?
195            }
196            _ => {}
197        }
198
199        index = value_index.map_or(index + 1, |value_index| value_index + 1);
200    }
201
202    parsed.validate()?;
203    Ok(parsed)
204}
205
206pub fn rewrite_convert_import_args_to_files(
207    command: &str,
208    args: &[String],
209    options: &RewriteOptions,
210) -> Result<RewriteResult, RewriteError> {
211    if !is_mcp_proxy_convert(command, args) {
212        return Ok(RewriteResult {
213            args: args.to_vec(),
214            created_files: Vec::new(),
215            changed: false,
216        });
217    }
218
219    let parsed = parse_imports(args)?;
220    let has_any = parsed.initialize_inline.is_some()
221        || parsed.initialize_file.is_some()
222        || parsed.tools_inline.is_some()
223        || parsed.tools_file.is_some();
224    if !has_any || (parsed.initialize_inline.is_none() && parsed.tools_inline.is_none()) {
225        validate_sources(&parsed)?;
226        return Ok(RewriteResult {
227            args: args.to_vec(),
228            created_files: Vec::new(),
229            changed: false,
230        });
231    }
232
233    validate_sources(&parsed)?;
234    if options.cache_dir.to_str().is_none() {
235        return Err(RewriteError::NonUtf8CachePath {
236            path: options.cache_dir.clone(),
237        });
238    }
239
240    let mut replacements: HashMap<usize, Vec<String>> = HashMap::new();
241    let mut skipped_indexes = HashSet::new();
242    let mut created_files = Vec::new();
243
244    for (kind, file_flag, parsed_value) in [
245        (
246            "initialize",
247            IMPORT_INITIALIZE_FILE,
248            parsed.initialize_inline.as_ref(),
249        ),
250        ("tools", IMPORT_TOOLS_FILE, parsed.tools_inline.as_ref()),
251    ] {
252        let Some(parsed_value) = parsed_value else {
253            continue;
254        };
255        let (path, created) = write_content_addressed(
256            &options.cache_dir,
257            &options.file_prefix,
258            kind,
259            parsed_value.value.as_bytes(),
260        )?;
261        if created {
262            created_files.push(path.clone());
263        }
264        let path_arg = path
265            .to_str()
266            .ok_or_else(|| RewriteError::NonUtf8CachePath { path: path.clone() })?
267            .to_string();
268        replacements.insert(
269            parsed_value.flag_index,
270            vec![long_flag(file_flag), path_arg],
271        );
272        if let Some(value_index) = parsed_value.value_index {
273            skipped_indexes.insert(value_index);
274        }
275    }
276
277    let mut rewritten = Vec::with_capacity(args.len());
278    for (index, arg) in args.iter().enumerate() {
279        if skipped_indexes.contains(&index) {
280            continue;
281        }
282        if let Some(values) = replacements.get(&index) {
283            rewritten.extend(values.iter().cloned());
284        } else {
285            rewritten.push(arg.clone());
286        }
287    }
288
289    Ok(RewriteResult {
290        args: rewritten,
291        created_files,
292        changed: true,
293    })
294}
295
296fn validate_sources(parsed: &ParsedImports) -> Result<(), RewriteError> {
297    let spec = RawFallbackImportSpec {
298        initialize_inline: parsed
299            .initialize_inline
300            .as_ref()
301            .map(|value| value.value.clone()),
302        initialize_file: parsed
303            .initialize_file
304            .as_ref()
305            .map(|value| PathBuf::from(&value.value)),
306        tools_inline: parsed
307            .tools_inline
308            .as_ref()
309            .map(|value| value.value.clone()),
310        tools_file: parsed
311            .tools_file
312            .as_ref()
313            .map(|value| PathBuf::from(&value.value)),
314    };
315    try_load_fallback(&spec)?;
316    Ok(())
317}
mcp_proxy_args/rewrite.rs

mcp_proxy_args/
rewrite.rs
