1use crate::common::types::{GeneratedCode, GeneratedFile};
34use crate::common::typescript::{
35 disambiguate_identifier, extract_properties, sanitize_ts_identifier, to_camel_case,
36};
37use crate::progressive::types::{
38 BridgeContext, CategoryInfo, IndexContext, PropertyInfo, ToolCategorization, ToolContext,
39 ToolSummary,
40};
41use crate::template_engine::TemplateEngine;
42use mcp_execution_core::metadata::{
43 METADATA_FILE_NAME, METADATA_SCHEMA_VERSION, ParameterMetadata, ServerMetadata, ToolMetadata,
44};
45use mcp_execution_core::{Error, Result};
46use mcp_execution_introspector::{ServerInfo, ToolInfo};
47use std::collections::{HashMap, HashSet};
48
49#[derive(Debug)]
66pub struct ProgressiveGenerator<'a> {
67 engine: TemplateEngine<'a>,
68}
69
70impl<'a> ProgressiveGenerator<'a> {
71 pub fn new() -> Result<Self> {
89 let engine = TemplateEngine::new()?;
90 Ok(Self { engine })
91 }
92
93 pub fn generate(&self, server_info: &ServerInfo) -> Result<GeneratedCode> {
148 tracing::info!(
149 "Generating progressive loading code for server: {}",
150 server_info.name
151 );
152
153 let mut code = GeneratedCode::new();
154 let server_id = server_info.id.as_str();
155 let typescript_names = resolve_typescript_names(&server_info.tools);
156 let mut tool_metadata = Vec::with_capacity(server_info.tools.len());
157
158 for (idx, tool) in server_info.tools.iter().enumerate() {
160 let typescript_name = typescript_names.get(idx).cloned().unwrap_or_default();
161 let tool_context =
162 self.create_tool_context(server_id, tool, None, typescript_name.clone())?;
163 let tool_code = self.engine.render("progressive/tool", &tool_context)?;
164
165 code.add_file(GeneratedFile {
166 path: format!("{}.ts", tool_context.typescript_name),
167 content: tool_code,
168 });
169
170 tracing::debug!("Generated tool file: {}.ts", tool_context.typescript_name);
171
172 tool_metadata.push(self.create_tool_metadata(tool, None, typescript_name)?);
173 }
174
175 let index_context = self.create_index_context(server_info, None, &typescript_names)?;
177 let index_code = self.engine.render("progressive/index", &index_context)?;
178
179 code.add_file(GeneratedFile {
180 path: "index.ts".to_string(),
181 content: index_code,
182 });
183
184 tracing::debug!("Generated index.ts");
185
186 let bridge_context = BridgeContext::default();
188 let bridge_code = self
189 .engine
190 .render("progressive/runtime-bridge", &bridge_context)?;
191
192 code.add_file(GeneratedFile {
193 path: "_runtime/mcp-bridge.ts".to_string(),
194 content: bridge_code,
195 });
196
197 tracing::debug!("Generated _runtime/mcp-bridge.ts");
198
199 code.add_file(GeneratedFile {
201 path: "package.json".to_string(),
202 content: "{\"type\":\"module\"}\n".to_string(),
203 });
204
205 tracing::debug!("Generated package.json");
206
207 code.add_file(Self::create_metadata_file(server_info, tool_metadata)?);
209
210 tracing::debug!("Generated {}", METADATA_FILE_NAME);
211
212 tracing::info!(
213 "Successfully generated {} files for {} (progressive loading)",
214 code.file_count(),
215 server_info.name
216 );
217
218 Ok(code)
219 }
220
221 pub fn generate_with_categories(
275 &self,
276 server_info: &ServerInfo,
277 categorizations: &HashMap<String, ToolCategorization>,
278 ) -> Result<GeneratedCode> {
279 tracing::info!(
280 "Generating progressive loading code with categorizations for server: {}",
281 server_info.name
282 );
283
284 let mut code = GeneratedCode::new();
285 let server_id = server_info.id.as_str();
286 let typescript_names = resolve_typescript_names(&server_info.tools);
287 let mut tool_metadata = Vec::with_capacity(server_info.tools.len());
288
289 for (idx, tool) in server_info.tools.iter().enumerate() {
291 let tool_name = tool.name.as_str();
292 let categorization = categorizations.get(tool_name);
293 let typescript_name = typescript_names.get(idx).cloned().unwrap_or_default();
294 let tool_context =
295 self.create_tool_context(server_id, tool, categorization, typescript_name.clone())?;
296 let tool_code = self.engine.render("progressive/tool", &tool_context)?;
297
298 code.add_file(GeneratedFile {
299 path: format!("{}.ts", tool_context.typescript_name),
300 content: tool_code,
301 });
302
303 tracing::debug!(
304 "Generated tool file: {}.ts (category: {:?})",
305 tool_context.typescript_name,
306 categorization.map(|c| &c.category)
307 );
308
309 tool_metadata.push(self.create_tool_metadata(tool, categorization, typescript_name)?);
310 }
311
312 let index_context =
314 self.create_index_context(server_info, Some(categorizations), &typescript_names)?;
315 let index_code = self.engine.render("progressive/index", &index_context)?;
316
317 code.add_file(GeneratedFile {
318 path: "index.ts".to_string(),
319 content: index_code,
320 });
321
322 tracing::debug!(
323 "Generated index.ts with {} categorizations",
324 categorizations.len()
325 );
326
327 let bridge_context = BridgeContext::default();
329 let bridge_code = self
330 .engine
331 .render("progressive/runtime-bridge", &bridge_context)?;
332
333 code.add_file(GeneratedFile {
334 path: "_runtime/mcp-bridge.ts".to_string(),
335 content: bridge_code,
336 });
337
338 tracing::debug!("Generated _runtime/mcp-bridge.ts");
339
340 code.add_file(GeneratedFile {
342 path: "package.json".to_string(),
343 content: "{\"type\":\"module\"}\n".to_string(),
344 });
345
346 tracing::debug!("Generated package.json");
347
348 code.add_file(Self::create_metadata_file(server_info, tool_metadata)?);
350
351 tracing::debug!("Generated {}", METADATA_FILE_NAME);
352
353 tracing::info!(
354 "Successfully generated {} files for {} with categorizations (progressive loading)",
355 code.file_count(),
356 server_info.name
357 );
358
359 Ok(code)
360 }
361
362 fn create_tool_context(
374 &self,
375 server_id: &str,
376 tool: &mcp_execution_introspector::ToolInfo,
377 categorization: Option<&ToolCategorization>,
378 typescript_name: String,
379 ) -> Result<ToolContext> {
380 let properties = self.extract_property_infos(&tool.input_schema)?;
382
383 let description = sanitize_jsdoc(&tool.description, 256);
384 let short_description = Some(categorization.map_or_else(
387 || description.clone(),
388 |c| sanitize_jsdoc(&c.short_description, 256),
389 ));
390
391 Ok(ToolContext {
392 server_id: sanitize_jsdoc(server_id, 256),
393 name: sanitize_jsdoc(tool.name.as_str(), 256),
394 name_literal: sanitize_ts_string_literal(tool.name.as_str()),
395 server_id_literal: sanitize_ts_string_literal(server_id),
396 typescript_name,
397 description,
398 input_schema: sanitize_schema_jsdoc_descriptions(tool.input_schema.clone()),
399 properties,
400 category: categorization.map(|c| sanitize_jsdoc(&c.category, 128)),
401 keywords: categorization.map(|c| sanitize_jsdoc(&c.keywords, 256)),
402 short_description,
403 })
404 }
405
406 fn create_index_context(
412 &self,
413 server_info: &ServerInfo,
414 categorizations: Option<&HashMap<String, ToolCategorization>>,
415 typescript_names: &[String],
416 ) -> Result<IndexContext> {
417 let tools: Vec<ToolSummary> = server_info
418 .tools
419 .iter()
420 .enumerate()
421 .map(|(idx, tool)| {
422 let tool_name = tool.name.as_str();
423 let cat = categorizations.and_then(|c| c.get(tool_name));
424 ToolSummary {
425 typescript_name: typescript_names.get(idx).cloned().unwrap_or_default(),
426 description: sanitize_jsdoc(&tool.description, 256),
427 category: cat.map(|c| sanitize_jsdoc(&c.category, 128)),
428 keywords: cat.map(|c| sanitize_jsdoc(&c.keywords, 256)),
429 short_description: cat.map(|c| sanitize_jsdoc(&c.short_description, 256)),
430 }
431 })
432 .collect();
433
434 let category_groups = categorizations.map(|_| {
436 let mut groups: HashMap<String, Vec<ToolSummary>> = HashMap::new();
437
438 for tool in &tools {
439 let cat_name = tool
440 .category
441 .clone()
442 .unwrap_or_else(|| "uncategorized".to_string());
443 groups.entry(cat_name).or_default().push(tool.clone());
444 }
445
446 let mut result: Vec<CategoryInfo> = groups
447 .into_iter()
448 .map(|(name, tools)| CategoryInfo { name, tools })
449 .collect();
450
451 result.sort_by(|a, b| {
453 if a.name == "uncategorized" {
454 std::cmp::Ordering::Greater
455 } else if b.name == "uncategorized" {
456 std::cmp::Ordering::Less
457 } else {
458 a.name.cmp(&b.name)
459 }
460 });
461
462 result
463 });
464
465 Ok(IndexContext {
466 server_name: sanitize_jsdoc(&server_info.name, 256),
467 server_version: sanitize_jsdoc(&server_info.version, 64),
468 tool_count: server_info.tools.len(),
469 tools,
470 categories: category_groups,
471 })
472 }
473
474 fn extract_property_infos(&self, schema: &serde_json::Value) -> Result<Vec<PropertyInfo>> {
486 Ok(self
487 .extract_property_data(schema)?
488 .into_iter()
489 .map(|(info, _raw_description)| info)
490 .collect())
491 }
492
493 fn extract_property_data(
507 &self,
508 schema: &serde_json::Value,
509 ) -> Result<Vec<(PropertyInfo, Option<String>)>> {
510 let raw_properties = extract_properties(schema);
511
512 let mut properties = Vec::new();
513 let mut used_names = HashSet::new();
514 for prop in raw_properties {
515 let raw_name = prop["name"]
516 .as_str()
517 .ok_or_else(|| Error::ValidationError {
518 field: "name".to_string(),
519 reason: "Property name is not a string".to_string(),
520 })?
521 .to_string();
522
523 let typescript_type = prop["type"]
524 .as_str()
525 .ok_or_else(|| Error::ValidationError {
526 field: "type".to_string(),
527 reason: "Property type is not a string".to_string(),
528 })?
529 .to_string();
530
531 let required = prop["required"].as_bool().unwrap_or(false);
532
533 let raw_description = if let Some(obj) = schema.as_object() {
536 obj.get("properties")
537 .and_then(|props| props.as_object())
538 .and_then(|props| props.get(&raw_name))
539 .and_then(|prop_schema| prop_schema.as_object())
540 .and_then(|obj| obj.get("description"))
541 .and_then(|desc| desc.as_str())
542 .map(str::to_string)
543 } else {
544 None
545 };
546 let description = raw_description
547 .as_deref()
548 .map(|desc| sanitize_jsdoc(desc, 256));
549
550 let base_name = sanitize_ts_identifier(&raw_name);
551 properties.push((
552 PropertyInfo {
553 name: disambiguate_identifier(&base_name, &mut used_names),
554 typescript_type,
555 description,
556 required,
557 },
558 raw_description,
559 ));
560 }
561
562 Ok(properties)
563 }
564
565 fn create_tool_metadata(
581 &self,
582 tool: &ToolInfo,
583 categorization: Option<&ToolCategorization>,
584 typescript_name: String,
585 ) -> Result<ToolMetadata> {
586 let properties = self.extract_property_data(&tool.input_schema)?;
587
588 let description = (!tool.description.is_empty()).then(|| tool.description.clone());
589 let category = categorization.map(|c| c.category.clone());
590 let keywords = categorization.map_or_else(Vec::new, |c| {
591 c.keywords
592 .split(',')
593 .map(str::trim)
594 .filter(|s| !s.is_empty())
595 .map(str::to_string)
596 .collect()
597 });
598
599 Ok(ToolMetadata {
600 name: tool.name.as_str().to_string(),
601 typescript_name,
602 category,
603 keywords,
604 description,
605 parameters: properties
606 .into_iter()
607 .map(|(p, raw_description)| ParameterMetadata {
608 name: p.name,
609 typescript_type: p.typescript_type,
610 required: p.required,
611 description: raw_description,
612 })
613 .collect(),
614 })
615 }
616
617 fn create_metadata_file(
625 server_info: &ServerInfo,
626 tools: Vec<ToolMetadata>,
627 ) -> Result<GeneratedFile> {
628 let meta = ServerMetadata {
629 schema_version: METADATA_SCHEMA_VERSION,
630 server_id: server_info.id.as_str().to_string(),
631 server_name: server_info.name.clone(),
632 server_version: server_info.version.clone(),
633 tools,
634 };
635
636 let content =
637 serde_json::to_string_pretty(&meta).map_err(|e| Error::SerializationError {
638 message: format!("failed to serialize {METADATA_FILE_NAME}"),
639 source: Some(e),
640 })?;
641
642 Ok(GeneratedFile {
643 path: METADATA_FILE_NAME.to_string(),
644 content,
645 })
646 }
647}
648
649fn sanitize_jsdoc(s: &str, max_len: usize) -> String {
654 let sanitized = s.replace("*/", "*\\/").replace(['\r', '\n'], " ");
655 if sanitized.chars().count() > max_len {
656 sanitized.chars().take(max_len).collect()
657 } else {
658 sanitized
659 }
660}
661
662fn sanitize_ts_string_literal(s: &str) -> String {
668 s.replace('\\', "\\\\")
669 .replace('\'', "\\'")
670 .replace('\r', "\\r")
671 .replace('\n', "\\n")
672}
673
674const RESERVED_WORDS: &[&str] = &[
680 "arguments",
681 "await",
682 "break",
683 "case",
684 "catch",
685 "class",
686 "const",
687 "continue",
688 "debugger",
689 "default",
690 "delete",
691 "do",
692 "else",
693 "enum",
694 "eval",
695 "export",
696 "extends",
697 "false",
698 "finally",
699 "for",
700 "function",
701 "if",
702 "implements",
703 "import",
704 "in",
705 "instanceof",
706 "interface",
707 "let",
708 "new",
709 "null",
710 "package",
711 "private",
712 "protected",
713 "public",
714 "return",
715 "static",
716 "super",
717 "switch",
718 "this",
719 "throw",
720 "true",
721 "try",
722 "typeof",
723 "var",
724 "void",
725 "while",
726 "with",
727 "yield",
728];
729
730fn resolve_typescript_names(tools: &[ToolInfo]) -> Vec<String> {
749 let mut used: HashSet<String> = RESERVED_WORDS.iter().map(|&s| s.to_string()).collect();
750 let mut resolved = Vec::with_capacity(tools.len());
751
752 for tool in tools {
753 let base = sanitize_ts_identifier(&to_camel_case(tool.name.as_str()));
754 resolved.push(disambiguate_identifier(&base, &mut used));
755 }
756
757 resolved
758}
759
760fn sanitize_schema_jsdoc_descriptions(mut value: serde_json::Value) -> serde_json::Value {
761 sanitize_schema_jsdoc_value(&mut value);
762 value
763}
764
765fn sanitize_schema_jsdoc_value(value: &mut serde_json::Value) {
766 match value {
767 serde_json::Value::Object(map) => {
768 for (key, child) in map.iter_mut() {
769 if key == "description" {
770 if let Some(description) = child.as_str() {
771 *child = serde_json::Value::String(sanitize_jsdoc(description, 256));
772 } else {
773 *child = serde_json::Value::Null;
774 }
775 } else {
776 sanitize_schema_jsdoc_value(child);
777 }
778 }
779 }
780 serde_json::Value::Array(values) => {
781 for child in values {
782 sanitize_schema_jsdoc_value(child);
783 }
784 }
785 _ => {}
786 }
787}
788
789#[cfg(test)]
790mod tests {
791 use super::*;
792 use mcp_execution_core::{ServerId, ToolName};
793 use mcp_execution_introspector::{ServerCapabilities, ToolInfo};
794 use serde_json::json;
795
796 fn create_test_server_info() -> ServerInfo {
797 ServerInfo {
798 id: ServerId::new("test-server"),
799 name: "Test Server".to_string(),
800 version: "1.0.0".to_string(),
801 tools: vec![
802 ToolInfo {
803 name: ToolName::new("create_issue"),
804 description: "Creates a new issue".to_string(),
805 input_schema: json!({
806 "type": "object",
807 "properties": {
808 "title": {
809 "type": "string",
810 "description": "Issue title"
811 },
812 "body": {
813 "type": "string",
814 "description": "Issue body"
815 }
816 },
817 "required": ["title"]
818 }),
819 output_schema: None,
820 },
821 ToolInfo {
822 name: ToolName::new("update_issue"),
823 description: "Updates an existing issue".to_string(),
824 input_schema: json!({
825 "type": "object",
826 "properties": {
827 "id": {
828 "type": "number"
829 }
830 },
831 "required": ["id"]
832 }),
833 output_schema: None,
834 },
835 ],
836 capabilities: ServerCapabilities {
837 supports_tools: true,
838 supports_resources: false,
839 supports_prompts: false,
840 },
841 }
842 }
843
844 #[test]
845 fn test_progressive_generator_new() {
846 let generator = ProgressiveGenerator::new();
847 assert!(generator.is_ok());
848 }
849
850 #[test]
851 fn test_generate_progressive_files() {
852 let generator = ProgressiveGenerator::new().unwrap();
853 let server_info = create_test_server_info();
854
855 let code = generator.generate(&server_info).unwrap();
856
857 assert_eq!(code.file_count(), 6);
864
865 let tool_files: Vec<_> = code.files.iter().map(|f| f.path.as_str()).collect();
867
868 assert!(tool_files.contains(&"createIssue.ts"));
869 assert!(tool_files.contains(&"updateIssue.ts"));
870 assert!(tool_files.contains(&"index.ts"));
871 assert!(tool_files.contains(&"_runtime/mcp-bridge.ts"));
872 assert!(tool_files.contains(&"package.json"));
873 assert!(tool_files.contains(&"_meta.json"));
874 }
875
876 #[test]
877 fn test_generate_meta_json_preserves_parameter_descriptions() {
878 let generator = ProgressiveGenerator::new().unwrap();
882 let server_info = create_test_server_info();
883
884 let code = generator.generate(&server_info).unwrap();
885 let meta_file = code.files.iter().find(|f| f.path == "_meta.json").unwrap();
886 let meta: ServerMetadata = serde_json::from_str(&meta_file.content).unwrap();
887
888 assert_eq!(meta.schema_version, METADATA_SCHEMA_VERSION);
889 assert_eq!(meta.server_id, "test-server");
890 assert_eq!(meta.server_name, "Test Server");
891 assert_eq!(meta.server_version, "1.0.0");
892 assert_eq!(meta.tools.len(), 2);
893
894 let create_issue = meta
895 .tools
896 .iter()
897 .find(|t| t.name == "create_issue")
898 .unwrap();
899 assert_eq!(create_issue.typescript_name, "createIssue");
900 let title = create_issue
901 .parameters
902 .iter()
903 .find(|p| p.name == "title")
904 .unwrap();
905 assert_eq!(title.description, Some("Issue title".to_string()));
906 assert!(title.required);
907 }
908
909 #[test]
910 fn test_generate_with_categories_meta_json_includes_categorization() {
911 let generator = ProgressiveGenerator::new().unwrap();
912 let server_info = create_test_server_info();
913
914 let mut categorizations = HashMap::new();
915 categorizations.insert(
916 "create_issue".to_string(),
917 ToolCategorization {
918 category: "issues".to_string(),
919 keywords: "create, issue , new".to_string(),
920 short_description: "Create a new issue".to_string(),
921 },
922 );
923
924 let code = generator
925 .generate_with_categories(&server_info, &categorizations)
926 .unwrap();
927 let meta_file = code.files.iter().find(|f| f.path == "_meta.json").unwrap();
928 let meta: ServerMetadata = serde_json::from_str(&meta_file.content).unwrap();
929
930 let create_issue = meta
931 .tools
932 .iter()
933 .find(|t| t.name == "create_issue")
934 .unwrap();
935 assert_eq!(create_issue.category, Some("issues".to_string()));
936 assert_eq!(
937 create_issue.keywords,
938 vec!["create".to_string(), "issue".to_string(), "new".to_string()]
939 );
940
941 let update_issue = meta
942 .tools
943 .iter()
944 .find(|t| t.name == "update_issue")
945 .unwrap();
946 assert!(update_issue.category.is_none());
947 assert!(update_issue.keywords.is_empty());
948 }
949
950 #[test]
951 fn test_generate_meta_json_parameter_description_is_raw_not_jsdoc_sanitized() {
952 let raw_description = format!(
957 "Matches C-style /* */ comment blocks.\nSecond line follows. {}",
958 "x".repeat(300)
959 );
960 assert!(raw_description.contains("*/"));
961 assert!(raw_description.contains('\n'));
962 assert!(raw_description.chars().count() > 256);
963
964 let server_info = ServerInfo {
965 id: ServerId::new("test-server"),
966 name: "Test Server".to_string(),
967 version: "1.0.0".to_string(),
968 tools: vec![ToolInfo {
969 name: ToolName::new("send_message"),
970 description: "Sends a message".to_string(),
971 input_schema: json!({
972 "type": "object",
973 "properties": {
974 "notes": {
975 "type": "string",
976 "description": raw_description
977 }
978 },
979 "required": []
980 }),
981 output_schema: None,
982 }],
983 capabilities: ServerCapabilities {
984 supports_tools: true,
985 supports_resources: false,
986 supports_prompts: false,
987 },
988 };
989
990 let generator = ProgressiveGenerator::new().unwrap();
991 let code = generator.generate(&server_info).unwrap();
992
993 let meta_file = code.files.iter().find(|f| f.path == "_meta.json").unwrap();
995 let meta: ServerMetadata = serde_json::from_str(&meta_file.content).unwrap();
996 let send_message = meta
997 .tools
998 .iter()
999 .find(|t| t.name == "send_message")
1000 .unwrap();
1001 let notes = send_message
1002 .parameters
1003 .iter()
1004 .find(|p| p.name == "notes")
1005 .unwrap();
1006 assert_eq!(notes.description, Some(raw_description.clone()));
1007
1008 let ts_file = code
1011 .files
1012 .iter()
1013 .find(|f| f.path == "sendMessage.ts")
1014 .unwrap();
1015 assert!(
1016 !ts_file.content.contains(raw_description.as_str()),
1017 "the .ts file must not contain the raw, un-sanitized description verbatim"
1018 );
1019 assert!(
1020 ts_file.content.contains("*\\/"),
1021 "the .ts file must escape '*/' to avoid closing the JSDoc comment early"
1022 );
1023 assert!(
1024 !ts_file
1025 .content
1026 .contains("Matches C-style /* */ comment blocks.\nSecond"),
1027 "the .ts file must flatten newlines within the description to spaces"
1028 );
1029 }
1030
1031 #[test]
1032 fn test_create_tool_context() {
1033 let generator = ProgressiveGenerator::new().unwrap();
1034 let tool = ToolInfo {
1035 name: ToolName::new("send_message"),
1036 description: "Sends a message".to_string(),
1037 input_schema: json!({
1038 "type": "object",
1039 "properties": {
1040 "text": {"type": "string"}
1041 },
1042 "required": ["text"]
1043 }),
1044 output_schema: None,
1045 };
1046
1047 let categorization = ToolCategorization {
1048 category: "messaging".to_string(),
1049 keywords: "send,message,chat".to_string(),
1050 short_description: "Send a message".to_string(),
1051 };
1052 let context = generator
1053 .create_tool_context(
1054 "test-server",
1055 &tool,
1056 Some(&categorization),
1057 "sendMessage".to_string(),
1058 )
1059 .unwrap();
1060
1061 assert_eq!(context.server_id, "test-server");
1062 assert_eq!(context.name, "send_message");
1063 assert_eq!(context.name_literal, "send_message");
1064 assert_eq!(context.server_id_literal, "test-server");
1065 assert_eq!(context.typescript_name, "sendMessage");
1066 assert_eq!(context.description, "Sends a message");
1067 assert_eq!(context.properties.len(), 1);
1068 assert_eq!(context.properties[0].name, "text");
1069 assert_eq!(context.category, Some("messaging".to_string()));
1070 assert_eq!(context.keywords, Some("send,message,chat".to_string()));
1071 assert_eq!(
1072 context.short_description,
1073 Some("Send a message".to_string())
1074 );
1075 }
1076
1077 #[test]
1078 fn test_create_tool_context_without_categorization_falls_back_to_description() {
1079 let generator = ProgressiveGenerator::new().unwrap();
1080 let tool = ToolInfo {
1081 name: ToolName::new("format_document"),
1082 description: "Format document with language-specific rules".to_string(),
1083 input_schema: json!({
1084 "type": "object",
1085 "properties": {
1086 "text": {"type": "string"}
1087 },
1088 "required": ["text"]
1089 }),
1090 output_schema: None,
1091 };
1092
1093 let context = generator
1094 .create_tool_context("test-server", &tool, None, "formatDocument".to_string())
1095 .unwrap();
1096
1097 assert_eq!(
1098 context.short_description,
1099 Some("Format document with language-specific rules".to_string())
1100 );
1101
1102 let rendered = generator
1104 .engine
1105 .render("progressive/tool", &context)
1106 .unwrap();
1107 assert!(rendered.contains("@description Format document with language-specific rules"));
1108 }
1109
1110 #[test]
1111 fn test_create_tool_context_input_schema_is_sanitized() {
1112 let generator = ProgressiveGenerator::new().unwrap();
1113 let tool = ToolInfo {
1114 name: ToolName::new("send_message"),
1115 description: "Sends a message".to_string(),
1116 input_schema: json!({
1117 "type": "object",
1118 "description": "Schema */ injected\nnext",
1119 "properties": {
1120 "text": {"type": "string"}
1121 },
1122 "required": ["text"]
1123 }),
1124 output_schema: None,
1125 };
1126
1127 let context = generator
1128 .create_tool_context("test-server", &tool, None, "sendMessage".to_string())
1129 .unwrap();
1130
1131 let expected = sanitize_schema_jsdoc_descriptions(tool.input_schema);
1132 assert_eq!(context.input_schema, expected);
1133 assert_eq!(
1134 context.input_schema["description"],
1135 json!("Schema *\\/ injected next")
1136 );
1137 }
1138
1139 #[test]
1140 fn test_create_index_context() {
1141 let generator = ProgressiveGenerator::new().unwrap();
1142 let server_info = create_test_server_info();
1143 let typescript_names = resolve_typescript_names(&server_info.tools);
1144
1145 let context = generator
1146 .create_index_context(&server_info, None, &typescript_names)
1147 .unwrap();
1148
1149 assert_eq!(context.server_name, "Test Server");
1150 assert_eq!(context.server_version, "1.0.0");
1151 assert_eq!(context.tool_count, 2);
1152 assert_eq!(context.tools.len(), 2);
1153 assert_eq!(context.tools[0].typescript_name, "createIssue");
1154 assert!(context.categories.is_none());
1155 }
1156
1157 #[test]
1158 fn test_extract_property_infos() {
1159 let generator = ProgressiveGenerator::new().unwrap();
1160 let schema = json!({
1161 "type": "object",
1162 "properties": {
1163 "name": {
1164 "type": "string",
1165 "description": "User name"
1166 },
1167 "age": {
1168 "type": "number"
1169 }
1170 },
1171 "required": ["name"]
1172 });
1173
1174 let props = generator.extract_property_infos(&schema).unwrap();
1175
1176 assert_eq!(props.len(), 2);
1177
1178 let name_prop = props.iter().find(|p| p.name == "name").unwrap();
1180 assert_eq!(name_prop.typescript_type, "string");
1181 assert_eq!(name_prop.description, Some("User name".to_string()));
1182 assert!(name_prop.required);
1183
1184 let age_prop = props.iter().find(|p| p.name == "age").unwrap();
1186 assert_eq!(age_prop.typescript_type, "number");
1187 assert!(!age_prop.required);
1188 }
1189
1190 #[test]
1191 fn test_extract_property_infos_sanitizes_malicious_property_name() {
1192 let generator = ProgressiveGenerator::new().unwrap();
1193 let schema = json!({
1194 "type": "object",
1195 "properties": {
1196 "x: string }; export const pwned = 1; interface J {": {
1197 "type": "string",
1198 "description": "Evil property"
1199 }
1200 },
1201 "required": []
1202 });
1203
1204 let props = generator.extract_property_infos(&schema).unwrap();
1205
1206 assert_eq!(props.len(), 1);
1207 assert!(!props[0].name.contains(['{', '}', ';', ':', ' ']));
1208 assert_eq!(props[0].description, Some("Evil property".to_string()));
1211 }
1212
1213 #[test]
1214 fn test_extract_property_infos_disambiguates_colliding_sibling_names() {
1215 let generator = ProgressiveGenerator::new().unwrap();
1219 let schema = json!({
1220 "type": "object",
1221 "properties": {
1222 "a-b": {"type": "string"},
1223 "a.b": {"type": "number"}
1224 },
1225 "required": []
1226 });
1227
1228 let props = generator.extract_property_infos(&schema).unwrap();
1229 let mut names: Vec<&str> = props.iter().map(|p| p.name.as_str()).collect();
1230 names.sort_unstable();
1231
1232 assert_eq!(names, vec!["a_b", "a_b_2"]);
1233 }
1234
1235 #[test]
1236 fn test_extract_property_infos_disambiguates_three_way_collision() {
1237 let generator = ProgressiveGenerator::new().unwrap();
1238 let schema = json!({
1239 "type": "object",
1240 "properties": {
1241 "a-b": {"type": "string"},
1242 "a.b": {"type": "number"},
1243 "a b": {"type": "boolean"}
1244 },
1245 "required": []
1246 });
1247
1248 let props = generator.extract_property_infos(&schema).unwrap();
1249 let mut names: Vec<&str> = props.iter().map(|p| p.name.as_str()).collect();
1250 names.sort_unstable();
1251
1252 assert_eq!(names, vec!["a_b", "a_b_2", "a_b_3"]);
1253 }
1254
1255 #[test]
1256 fn test_generate_disambiguates_colliding_top_level_params() {
1257 let generator = ProgressiveGenerator::new().unwrap();
1258 let mut server_info = create_test_server_info();
1259 server_info.tools[0].input_schema = json!({
1260 "type": "object",
1261 "properties": {
1262 "a-b": {"type": "string"},
1263 "a.b": {"type": "number"}
1264 },
1265 "required": []
1266 });
1267
1268 let code = generator.generate(&server_info).unwrap();
1269 let tool = code
1270 .files
1271 .iter()
1272 .find(|f| f.path == "createIssue.ts")
1273 .unwrap();
1274
1275 assert_eq!(
1276 tool.content.matches("a_b:").count() + tool.content.matches("a_b?:").count(),
1277 1,
1278 "field 'a_b' must appear exactly once in the Params interface: {}",
1279 tool.content
1280 );
1281 assert_eq!(
1282 tool.content.matches("a_b_2:").count() + tool.content.matches("a_b_2?:").count(),
1283 1,
1284 "disambiguated field 'a_b_2' must appear exactly once in the Params interface: {}",
1285 tool.content
1286 );
1287 }
1288
1289 #[test]
1290 fn test_generate_sanitizes_property_name_injection() {
1291 let generator = ProgressiveGenerator::new().unwrap();
1292 let mut server_info = create_test_server_info();
1293 server_info.tools[0].input_schema = json!({
1294 "type": "object",
1295 "properties": {
1296 "x: string }; export const pwned = evil(); interface J {": {"type": "string"}
1297 },
1298 "required": []
1299 });
1300
1301 let code = generator.generate(&server_info).unwrap();
1302 let tool = code
1303 .files
1304 .iter()
1305 .find(|f| f.path == "createIssue.ts")
1306 .unwrap();
1307
1308 assert!(
1309 !tool.content.contains("export const pwned"),
1310 "raw property name must not inject a top-level statement: {}",
1311 tool.content
1312 );
1313 }
1314
1315 #[test]
1316 fn test_sanitize_jsdoc_strips_comment_terminator() {
1317 assert_eq!(sanitize_jsdoc("Foo */ bar", 256), "Foo *\\/ bar");
1318 }
1319
1320 #[test]
1321 fn test_sanitize_jsdoc_replaces_newlines() {
1322 assert_eq!(
1323 sanitize_jsdoc("line1\nline2\r\nline3", 256),
1324 "line1 line2 line3"
1325 );
1326 }
1327
1328 #[test]
1329 fn test_sanitize_jsdoc_truncates() {
1330 let long = "a".repeat(300);
1331 assert_eq!(sanitize_jsdoc(&long, 256).chars().count(), 256);
1332 }
1333
1334 #[test]
1335 fn test_sanitize_jsdoc_passthrough() {
1336 assert_eq!(sanitize_jsdoc("Normal string", 256), "Normal string");
1337 }
1338
1339 #[test]
1340 fn test_sanitize_ts_string_literal_escapes_quote_and_backslash() {
1341 assert_eq!(
1342 sanitize_ts_string_literal(r"it's a \test"),
1343 r"it\'s a \\test"
1344 );
1345 }
1346
1347 #[test]
1348 fn test_sanitize_ts_string_literal_escape_order_prevents_double_escaping() {
1349 assert_eq!(sanitize_ts_string_literal("\\'"), r"\\\'");
1352 }
1353
1354 #[test]
1355 fn test_sanitize_ts_string_literal_escapes_newlines() {
1356 assert_eq!(
1357 sanitize_ts_string_literal("line1\nline2\rline3"),
1358 "line1\\nline2\\rline3"
1359 );
1360 }
1361
1362 #[test]
1367 fn test_sanitize_ts_identifier_passthrough_valid() {
1368 assert_eq!(sanitize_ts_identifier("sendMessage_1"), "sendMessage_1");
1369 }
1370
1371 #[test]
1372 fn test_generate_sanitizes_call_site_string_literal_injection() {
1373 let generator = ProgressiveGenerator::new().unwrap();
1374 let mut server_info = create_test_server_info();
1375 server_info.tools[0].name = ToolName::new("create_issue'); alert('pwned");
1376
1377 let code = generator.generate(&server_info).unwrap();
1378 let tool = code
1379 .files
1380 .iter()
1381 .find(|f| {
1382 std::path::Path::new(&f.path)
1383 .extension()
1384 .is_some_and(|ext| ext.eq_ignore_ascii_case("ts"))
1385 && f.path != "index.ts"
1386 })
1387 .unwrap();
1388
1389 assert!(
1390 !tool.content.contains("'); alert('pwned"),
1391 "raw quote must not break out of the callMCPTool string literal: {}",
1392 tool.content
1393 );
1394 }
1395
1396 #[test]
1397 fn test_resolve_typescript_names_disambiguates_collisions() {
1398 let tools = vec![
1399 ToolInfo {
1400 name: ToolName::new("foo-bar"),
1401 description: String::new(),
1402 input_schema: json!({}),
1403 output_schema: None,
1404 },
1405 ToolInfo {
1406 name: ToolName::new("foo.bar"),
1407 description: String::new(),
1408 input_schema: json!({}),
1409 output_schema: None,
1410 },
1411 ToolInfo {
1412 name: ToolName::new("foo bar"),
1413 description: String::new(),
1414 input_schema: json!({}),
1415 output_schema: None,
1416 },
1417 ];
1418
1419 let resolved = resolve_typescript_names(&tools);
1420 let mut names: Vec<&String> = resolved.iter().collect();
1421 names.sort();
1422
1423 assert_eq!(resolved.len(), 3);
1425 let unique: HashSet<&String> = names.iter().copied().collect();
1426 assert_eq!(
1427 unique.len(),
1428 3,
1429 "collisions must be disambiguated: {names:?}"
1430 );
1431 assert_eq!(resolved[0], "foo_bar");
1432 }
1433
1434 #[test]
1435 fn test_resolve_typescript_names_disambiguates_identical_raw_names() {
1436 let tools = vec![
1440 ToolInfo {
1441 name: ToolName::new("dup"),
1442 description: "First".to_string(),
1443 input_schema: json!({}),
1444 output_schema: None,
1445 },
1446 ToolInfo {
1447 name: ToolName::new("dup"),
1448 description: "Second".to_string(),
1449 input_schema: json!({}),
1450 output_schema: None,
1451 },
1452 ];
1453
1454 let resolved = resolve_typescript_names(&tools);
1455
1456 assert_eq!(resolved, vec!["dup".to_string(), "dup_2".to_string()]);
1457 }
1458
1459 #[test]
1460 fn test_resolve_typescript_names_disambiguates_three_way_identical_raw_names() {
1461 let tools: Vec<ToolInfo> = (0..3)
1462 .map(|_| ToolInfo {
1463 name: ToolName::new("dup"),
1464 description: String::new(),
1465 input_schema: json!({}),
1466 output_schema: None,
1467 })
1468 .collect();
1469
1470 let resolved = resolve_typescript_names(&tools);
1471
1472 assert_eq!(
1473 resolved,
1474 vec!["dup".to_string(), "dup_2".to_string(), "dup_3".to_string()]
1475 );
1476 }
1477
1478 #[test]
1479 fn test_resolve_typescript_names_disambiguates_reserved_words() {
1480 let reserved_tool_names = [
1481 "delete",
1482 "typeof",
1483 "class",
1484 "new",
1485 "import",
1486 "export",
1487 "in",
1488 "instanceof",
1489 "void",
1490 "enum",
1491 "eval",
1492 "arguments",
1493 ];
1494
1495 for name in reserved_tool_names {
1496 let tools = vec![ToolInfo {
1497 name: ToolName::new(name),
1498 description: String::new(),
1499 input_schema: json!({}),
1500 output_schema: None,
1501 }];
1502
1503 let resolved = resolve_typescript_names(&tools);
1504 let typescript_name = &resolved[0];
1505
1506 assert_ne!(
1507 typescript_name, name,
1508 "reserved word {name} must be disambiguated"
1509 );
1510 assert!(
1511 !RESERVED_WORDS.contains(&typescript_name.as_str()),
1512 "resolved name {typescript_name} for tool {name} must not be a reserved word"
1513 );
1514 }
1515 }
1516
1517 #[test]
1518 fn test_resolve_typescript_names_reserved_word_avoids_existing_collision() {
1519 let tools = vec![
1520 ToolInfo {
1521 name: ToolName::new("class"),
1522 description: String::new(),
1523 input_schema: json!({}),
1524 output_schema: None,
1525 },
1526 ToolInfo {
1534 name: ToolName::new("class-2"),
1535 description: String::new(),
1536 input_schema: json!({}),
1537 output_schema: None,
1538 },
1539 ];
1540
1541 let resolved = resolve_typescript_names(&tools);
1542
1543 assert_ne!(
1544 resolved[0], resolved[1],
1545 "a reserved-word tool's fallback name must not collide with an unrelated tool that already claims it"
1546 );
1547 assert!(!RESERVED_WORDS.contains(&resolved[0].as_str()));
1548 }
1549
1550 #[test]
1551 fn test_generate_sanitizes_reserved_word_tool_name() {
1552 let generator = ProgressiveGenerator::new().unwrap();
1553 let mut server_info = create_test_server_info();
1554 server_info.tools = vec![ToolInfo {
1555 name: ToolName::new("delete"),
1556 description: "Delete something".to_string(),
1557 input_schema: json!({}),
1558 output_schema: None,
1559 }];
1560
1561 let code = generator.generate(&server_info).unwrap();
1562 let tool_file = code.files.iter().find(|f| f.path == "delete_2.ts").unwrap();
1563
1564 assert!(!tool_file.content.contains("export async function delete("));
1565 assert!(
1566 tool_file
1567 .content
1568 .contains("export async function delete_2(")
1569 );
1570 }
1571
1572 #[test]
1573 fn test_generate_disambiguates_colliding_tool_names() {
1574 let generator = ProgressiveGenerator::new().unwrap();
1575 let mut server_info = create_test_server_info();
1576 server_info.tools = vec![
1577 ToolInfo {
1578 name: ToolName::new("foo-bar"),
1579 description: "First".to_string(),
1580 input_schema: json!({}),
1581 output_schema: None,
1582 },
1583 ToolInfo {
1584 name: ToolName::new("foo.bar"),
1585 description: "Second".to_string(),
1586 input_schema: json!({}),
1587 output_schema: None,
1588 },
1589 ];
1590
1591 let code = generator.generate(&server_info).unwrap();
1592
1593 let tool_files: Vec<&str> = code
1595 .files
1596 .iter()
1597 .filter(|f| f.path == "foo_bar.ts" || f.path == "foo_bar_2.ts")
1598 .map(|f| f.path.as_str())
1599 .collect();
1600 assert_eq!(
1601 tool_files.len(),
1602 2,
1603 "colliding names must not overwrite each other's file: {tool_files:?}"
1604 );
1605
1606 let index = code.files.iter().find(|f| f.path == "index.ts").unwrap();
1607 assert_eq!(
1608 index.content.matches("export { foo_bar,").count(),
1609 1,
1610 "index.ts must export the first tool's identifier exactly once"
1611 );
1612 assert_eq!(
1613 index.content.matches("export { foo_bar_2,").count(),
1614 1,
1615 "index.ts must export the disambiguated second identifier exactly once"
1616 );
1617 }
1618
1619 #[test]
1620 fn test_generate_disambiguates_identical_raw_tool_names() {
1621 let generator = ProgressiveGenerator::new().unwrap();
1625 let mut server_info = create_test_server_info();
1626 server_info.tools = vec![
1627 ToolInfo {
1628 name: ToolName::new("dup"),
1629 description: "First".to_string(),
1630 input_schema: json!({}),
1631 output_schema: None,
1632 },
1633 ToolInfo {
1634 name: ToolName::new("dup"),
1635 description: "Second".to_string(),
1636 input_schema: json!({}),
1637 output_schema: None,
1638 },
1639 ];
1640
1641 let code = generator.generate(&server_info).unwrap();
1642
1643 let dup_files: Vec<&str> = code
1644 .files
1645 .iter()
1646 .filter(|f| f.path == "dup.ts" || f.path == "dup_2.ts")
1647 .map(|f| f.path.as_str())
1648 .collect();
1649 assert_eq!(
1650 dup_files.len(),
1651 2,
1652 "identical raw tool names must not overwrite each other's file: {dup_files:?}"
1653 );
1654
1655 let index = code.files.iter().find(|f| f.path == "index.ts").unwrap();
1656 assert_eq!(
1657 index.content.matches("export { dup,").count(),
1658 1,
1659 "index.ts must export the first tool's identifier exactly once"
1660 );
1661 assert_eq!(
1662 index.content.matches("export { dup_2,").count(),
1663 1,
1664 "index.ts must export the disambiguated second identifier exactly once"
1665 );
1666 }
1667
1668 #[test]
1669 fn test_sanitize_schema_jsdoc_drops_non_string_descriptions() {
1670 let sanitized = sanitize_schema_jsdoc_descriptions(json!({
1671 "type": "object",
1672 "description": {"text": "Schema */ injected\nnext"},
1673 "properties": {
1674 "title": {
1675 "type": "string",
1676 "description": ["Title */ injected\nnext"]
1677 }
1678 }
1679 }));
1680
1681 assert!(sanitized["description"].is_null());
1682 assert!(sanitized["properties"]["title"]["description"].is_null());
1683 }
1684
1685 #[test]
1686 fn test_sanitize_schema_jsdoc_recurses_into_array_items() {
1687 let sanitized = sanitize_schema_jsdoc_descriptions(json!({
1688 "type": "object",
1689 "properties": {
1690 "tags": {
1691 "type": "array",
1692 "items": [
1693 {
1694 "type": "string",
1695 "description": "Tag */ injected\nnext"
1696 }
1697 ]
1698 }
1699 }
1700 }));
1701
1702 let description = sanitized["properties"]["tags"]["items"][0]["description"]
1703 .as_str()
1704 .unwrap();
1705
1706 assert_eq!(description, "Tag *\\/ injected next");
1707 }
1708
1709 #[test]
1710 fn test_sanitize_jsdoc_truncation_boundary_injection() {
1711 let max_len = 256;
1712 let payload = format!("{}*/{}", "a".repeat(max_len - 1), "trailer");
1716
1717 let sanitized = sanitize_jsdoc(&payload, max_len);
1718
1719 assert!(
1720 !sanitized.contains("*/"),
1721 "truncation must not re-open the JSDoc comment: {sanitized}"
1722 );
1723 assert_eq!(sanitized.chars().count(), max_len);
1724 }
1725
1726 #[test]
1727 fn test_generate_sanitizes_jsdoc_injection() {
1728 let generator = ProgressiveGenerator::new().unwrap();
1729 let mut server_info = create_test_server_info();
1730 server_info.name = "Evil */ injection".to_string();
1731 server_info.version = "1.0\n<script>".to_string();
1732
1733 let code = generator.generate(&server_info).unwrap();
1734 let index = code.files.iter().find(|f| f.path == "index.ts").unwrap();
1735
1736 assert!(
1738 !index.content.contains("Evil */ injection"),
1739 "Server name should be sanitized in JSDoc"
1740 );
1741 assert!(
1742 !index.content.contains("1.0\n<script>"),
1743 "Server version should have newlines stripped"
1744 );
1745 }
1746
1747 #[test]
1748 fn test_generate_sanitizes_schema_and_category_jsdoc_injection() {
1749 let generator = ProgressiveGenerator::new().unwrap();
1750 let mut server_info = create_test_server_info();
1751 server_info.tools[0].input_schema = json!({
1752 "type": "object",
1753 "description": "Schema */ injected\nnext",
1754 "properties": {
1755 "title": {
1756 "type": "string",
1757 "description": "Title */ injected\nnext"
1758 }
1759 },
1760 "required": ["title"]
1761 });
1762
1763 let mut categorizations = HashMap::new();
1764 categorizations.insert(
1765 "create_issue".to_string(),
1766 ToolCategorization {
1767 category: "issues */ injected\nnext".to_string(),
1768 keywords: "create,*/ injected\nnext".to_string(),
1769 short_description: "Create */ injected\nnext".to_string(),
1770 },
1771 );
1772
1773 let code = generator
1774 .generate_with_categories(&server_info, &categorizations)
1775 .unwrap();
1776 let tool = code
1777 .files
1778 .iter()
1779 .find(|f| f.path == "createIssue.ts")
1780 .unwrap();
1781
1782 for raw in [
1783 "Schema */ injected",
1784 "Title */ injected",
1785 "issues */ injected",
1786 "create,*/ injected",
1787 "Create */ injected",
1788 ] {
1789 assert!(
1790 !tool.content.contains(raw),
1791 "generated JSDoc should not contain raw injection text: {raw}"
1792 );
1793 }
1794
1795 assert!(tool.content.contains("Schema *\\/ injected next"));
1796 assert!(tool.content.contains("Title *\\/ injected next"));
1797 assert!(tool.content.contains("issues *\\/ injected next"));
1798 assert!(tool.content.contains("create,*\\/ injected next"));
1799 assert!(tool.content.contains("Create *\\/ injected next"));
1800 }
1801}