1use crate::common::types::{GeneratedCode, GeneratedFile};
34use crate::common::typescript::{
35 MAX_SCHEMA_RECURSION_DEPTH, disambiguate_identifier, extract_properties,
36 sanitize_ts_identifier, to_camel_case,
37};
38use crate::progressive::types::{
39 BridgeContext, CategoryInfo, IndexContext, PropertyInfo, ToolCategorization, ToolContext,
40 ToolSummary,
41};
42use crate::template_engine::TemplateEngine;
43use mcp_execution_core::ResourceKind;
44use mcp_execution_core::metadata::{
45 METADATA_FILE_NAME, METADATA_SCHEMA_VERSION, ParameterMetadata, ServerMetadata, ToolMetadata,
46};
47use mcp_execution_core::{Error, Result};
48use mcp_execution_introspector::{ServerInfo, ToolInfo};
49use std::collections::{HashMap, HashSet};
50
51const FIXED_FILE_COUNT: usize = 5;
55
56pub const MAX_GENERATED_FILES: usize =
76 mcp_execution_introspector::MAX_TOOL_COUNT + FIXED_FILE_COUNT;
77
78pub const MAX_GENERATED_BYTES: usize = 2
97 * mcp_execution_introspector::MAX_TOOL_COUNT
98 * (mcp_execution_introspector::MAX_TOOL_NAME_LEN
99 + mcp_execution_introspector::MAX_TOOL_DESCRIPTION_LEN
100 + mcp_execution_introspector::MAX_SCHEMA_SIZE_BYTES);
101
102const PACKAGE_JSON: &str = "{\"type\":\"module\",\"devDependencies\":{\"@types/node\":\"^22\"}}\n";
111
112const TSCONFIG_JSON: &str = r#"{
152 "compilerOptions": {
153 "target": "ES2022",
154 "module": "NodeNext",
155 "moduleResolution": "NodeNext",
156 "strict": true,
157 "noEmit": true,
158 "allowImportingTsExtensions": true,
159 "skipLibCheck": true,
160 "types": ["node"]
161 },
162 "include": ["**/*.ts"]
163}
164"#;
165
166#[derive(Debug)]
183pub struct ProgressiveGenerator<'a> {
184 engine: TemplateEngine<'a>,
185}
186
187impl ProgressiveGenerator<'_> {
188 pub fn new() -> Result<Self> {
206 let engine = TemplateEngine::new()?;
207 Ok(Self { engine })
208 }
209
210 pub fn generate(&self, server_info: &ServerInfo) -> Result<GeneratedCode> {
270 self.generate_with_categories(server_info, &HashMap::new())
271 }
272
273 #[tracing::instrument(
327 skip_all,
328 fields(server_id = %server_info.id, tool_count = server_info.tools.len())
329 )]
330 pub fn generate_with_categories(
331 &self,
332 server_info: &ServerInfo,
333 categorizations: &HashMap<String, ToolCategorization>,
334 ) -> Result<GeneratedCode> {
335 if categorizations.is_empty() {
336 tracing::info!(
337 "Generating progressive loading code for server: {}",
338 server_info.name
339 );
340 } else {
341 tracing::info!(
342 "Generating progressive loading code with categorizations for server: {}",
343 server_info.name
344 );
345 }
346
347 enforce_tool_count_bound(server_info)?;
348
349 let mut code = GeneratedCode::new();
350 let mut total_bytes = 0usize;
351 let server_id = server_info.id.as_str();
352 let typescript_names = resolve_typescript_names(&server_info.tools);
353 let mut tool_metadata = Vec::with_capacity(server_info.tools.len());
354
355 for (idx, tool) in server_info.tools.iter().enumerate() {
357 let tool_name = tool.name.as_str();
358 let categorization = categorizations.get(tool_name);
359 let typescript_name = typescript_names.get(idx).cloned().unwrap_or_default();
360 let extracted_properties =
361 Self::extract_property_data(&tool.input_schema).map_err(|source| {
362 Self::wrap_tool_generation_error(tool, "extract property schema", source)
363 })?;
364 let properties_for_context = extracted_properties
365 .iter()
366 .map(|(info, _)| info.clone())
367 .collect();
368 let tool_context = Self::create_tool_context(
369 server_id,
370 tool,
371 categorization,
372 typescript_name.clone(),
373 properties_for_context,
374 );
375 let tool_code = self
376 .engine
377 .render("progressive/tool", &tool_context)
378 .map_err(|source| {
379 Self::wrap_tool_generation_error(tool, "render tool template", source)
380 })?;
381
382 add_tracked(
383 &mut code,
384 &mut total_bytes,
385 GeneratedFile {
386 path: format!("{}.ts", tool_context.typescript_name),
387 content: tool_code,
388 },
389 )
390 .map_err(|source| {
391 Self::wrap_tool_generation_error(tool, "track generated tool file", source)
392 })?;
393
394 tracing::debug!(
395 "Generated tool file: {}.ts (category: {:?})",
396 tool_context.typescript_name,
397 categorization.map(|c| &c.category)
398 );
399
400 tool_metadata.push(Self::create_tool_metadata(
401 tool,
402 categorization,
403 typescript_name,
404 extracted_properties,
405 ));
406 }
407
408 let index_context =
410 Self::create_index_context(server_info, Some(categorizations), &typescript_names);
411 let index_code = self.engine.render("progressive/index", &index_context)?;
412
413 add_tracked(
414 &mut code,
415 &mut total_bytes,
416 GeneratedFile {
417 path: "index.ts".to_string(),
418 content: index_code,
419 },
420 )?;
421
422 tracing::debug!(
423 "Generated index.ts with {} categorizations",
424 categorizations.len()
425 );
426
427 let bridge_context = BridgeContext::default();
429 let bridge_code = self
430 .engine
431 .render("progressive/runtime-bridge", &bridge_context)?;
432
433 add_tracked(
434 &mut code,
435 &mut total_bytes,
436 GeneratedFile {
437 path: "_runtime/mcp-bridge.ts".to_string(),
438 content: bridge_code,
439 },
440 )?;
441
442 tracing::debug!("Generated _runtime/mcp-bridge.ts");
443
444 add_tracked(
447 &mut code,
448 &mut total_bytes,
449 GeneratedFile {
450 path: "package.json".to_string(),
451 content: PACKAGE_JSON.to_string(),
452 },
453 )?;
454
455 tracing::debug!("Generated package.json");
456
457 add_tracked(
460 &mut code,
461 &mut total_bytes,
462 GeneratedFile {
463 path: "tsconfig.json".to_string(),
464 content: TSCONFIG_JSON.to_string(),
465 },
466 )?;
467
468 tracing::debug!("Generated tsconfig.json");
469
470 add_tracked(
472 &mut code,
473 &mut total_bytes,
474 Self::create_metadata_file(server_info, tool_metadata)?,
475 )?;
476
477 tracing::debug!("Generated {}", METADATA_FILE_NAME);
478
479 if categorizations.is_empty() {
480 tracing::info!(
481 "Successfully generated {} files for {} (progressive loading)",
482 code.file_count(),
483 server_info.name
484 );
485 } else {
486 tracing::info!(
487 "Successfully generated {} files for {} with categorizations (progressive loading)",
488 code.file_count(),
489 server_info.name
490 );
491 }
492
493 Ok(code)
494 }
495
496 fn create_tool_context(
514 server_id: &str,
515 tool: &mcp_execution_introspector::ToolInfo,
516 categorization: Option<&ToolCategorization>,
517 typescript_name: String,
518 properties: Vec<PropertyInfo>,
519 ) -> ToolContext {
520 let description = sanitize_jsdoc(&tool.description, 256);
521 let short_description = categorization.map_or_else(
524 || description.clone(),
525 |c| sanitize_jsdoc(&c.short_description, 256),
526 );
527
528 ToolContext {
529 server_id: sanitize_jsdoc(server_id, 256),
530 name: sanitize_jsdoc(tool.name.as_str(), 256),
531 name_literal: sanitize_ts_string_literal(tool.name.as_str()),
532 server_id_literal: sanitize_ts_string_literal(server_id),
533 typescript_name,
534 description,
535 input_schema: sanitize_schema_jsdoc_descriptions(tool.input_schema.clone()),
536 properties,
537 category: categorization.map(|c| sanitize_jsdoc(&c.category, 128)),
538 keywords: categorization.map(|c| render_keywords_for_jsdoc(&c.keywords)),
539 short_description,
540 }
541 }
542
543 fn create_index_context(
552 server_info: &ServerInfo,
553 categorizations: Option<&HashMap<String, ToolCategorization>>,
554 typescript_names: &[String],
555 ) -> IndexContext {
556 let tools: Vec<ToolSummary> = server_info
557 .tools
558 .iter()
559 .enumerate()
560 .map(|(idx, tool)| {
561 let tool_name = tool.name.as_str();
562 let cat = categorizations.and_then(|c| c.get(tool_name));
563 ToolSummary {
564 typescript_name: typescript_names.get(idx).cloned().unwrap_or_default(),
565 description: sanitize_jsdoc(&tool.description, 256),
566 category: cat.map(|c| sanitize_jsdoc(&c.category, 128)),
567 keywords: cat.map(|c| render_keywords_for_jsdoc(&c.keywords)),
568 short_description: cat.map(|c| sanitize_jsdoc(&c.short_description, 256)),
569 }
570 })
571 .collect();
572
573 let category_groups = categorizations.filter(|c| !c.is_empty()).map(|_| {
579 let mut groups: HashMap<String, Vec<ToolSummary>> = HashMap::new();
580
581 for tool in &tools {
582 let cat_name = tool
583 .category
584 .clone()
585 .unwrap_or_else(|| "uncategorized".to_string());
586 groups.entry(cat_name).or_default().push(tool.clone());
587 }
588
589 let mut result: Vec<CategoryInfo> = groups
590 .into_iter()
591 .map(|(name, tools)| CategoryInfo { name, tools })
592 .collect();
593
594 result.sort_by(|a, b| {
596 if a.name == "uncategorized" {
597 std::cmp::Ordering::Greater
598 } else if b.name == "uncategorized" {
599 std::cmp::Ordering::Less
600 } else {
601 a.name.cmp(&b.name)
602 }
603 });
604
605 result
606 });
607
608 IndexContext {
609 server_name: sanitize_jsdoc(&server_info.name, 256),
610 server_version: sanitize_jsdoc(&server_info.version, 64),
611 tool_count: server_info.tools.len(),
612 tools,
613 categories: category_groups,
614 }
615 }
616
617 fn wrap_tool_generation_error(tool: &ToolInfo, stage: &str, source: Error) -> Error {
630 Error::ScriptGenerationError {
631 tool: tool.name.as_str().to_string(),
632 message: format!("failed to {stage}"),
633 source: Some(Box::new(source)),
634 }
635 }
636
637 #[cfg(test)]
655 fn extract_property_infos(schema: &serde_json::Value) -> Result<Vec<PropertyInfo>> {
656 Ok(Self::extract_property_data(schema)?
657 .into_iter()
658 .map(|(info, _raw_description)| info)
659 .collect())
660 }
661
662 fn extract_property_data(
682 schema: &serde_json::Value,
683 ) -> Result<Vec<(PropertyInfo, Option<String>)>> {
684 let raw_properties = extract_properties(schema);
685
686 let mut properties = Vec::new();
687 let mut used_names = HashSet::new();
688 for prop in raw_properties {
689 let raw_name = prop["name"]
690 .as_str()
691 .ok_or_else(|| Error::ValidationError {
692 field: "name".to_string(),
693 reason: "Property name is not a string".to_string(),
694 })?
695 .to_string();
696
697 let typescript_type = prop["type"]
698 .as_str()
699 .ok_or_else(|| Error::ValidationError {
700 field: "type".to_string(),
701 reason: "Property type is not a string".to_string(),
702 })?
703 .to_string();
704
705 let required = prop["required"].as_bool().unwrap_or(false);
706
707 let raw_description = schema.as_object().and_then(|obj| {
710 obj.get("properties")
711 .and_then(|props| props.as_object())
712 .and_then(|props| props.get(&raw_name))
713 .and_then(|prop_schema| prop_schema.as_object())
714 .and_then(|obj| obj.get("description"))
715 .and_then(|desc| desc.as_str())
716 .map(str::to_string)
717 });
718 let description = raw_description
719 .as_deref()
720 .map(|desc| sanitize_jsdoc(desc, 256));
721
722 let base_name = sanitize_ts_identifier(&raw_name);
723 properties.push((
724 PropertyInfo {
725 name: disambiguate_identifier(&base_name, &mut used_names),
726 typescript_type,
727 description,
728 required,
729 },
730 raw_description,
731 ));
732 }
733
734 Ok(properties)
735 }
736
737 fn create_tool_metadata(
757 tool: &ToolInfo,
758 categorization: Option<&ToolCategorization>,
759 typescript_name: String,
760 properties: Vec<(PropertyInfo, Option<String>)>,
761 ) -> ToolMetadata {
762 let description = (!tool.description.is_empty()).then(|| tool.description.clone());
763 let category = categorization.map(|c| c.category.clone());
764 let keywords = categorization.map_or_else(Vec::new, |c| c.keywords.clone());
765
766 ToolMetadata {
767 name: tool.name.clone(),
768 typescript_name,
769 category,
770 keywords,
771 description,
772 parameters: properties
773 .into_iter()
774 .map(|(p, raw_description)| ParameterMetadata {
775 name: p.name,
776 typescript_type: p.typescript_type,
777 required: p.required,
778 description: raw_description,
779 })
780 .collect(),
781 }
782 }
783
784 fn create_metadata_file(
792 server_info: &ServerInfo,
793 tools: Vec<ToolMetadata>,
794 ) -> Result<GeneratedFile> {
795 let meta = ServerMetadata {
796 schema_version: METADATA_SCHEMA_VERSION,
797 server_id: server_info.id.clone(),
798 server_name: server_info.name.clone(),
799 server_version: server_info.version.clone(),
800 tools,
801 };
802
803 let content =
804 serde_json::to_string_pretty(&meta).map_err(|e| Error::SerializationError {
805 message: format!("failed to serialize {METADATA_FILE_NAME}"),
806 source: Some(e),
807 })?;
808
809 Ok(GeneratedFile {
810 path: METADATA_FILE_NAME.to_string(),
811 content,
812 })
813 }
814}
815
816fn enforce_tool_count_bound(server_info: &ServerInfo) -> Result<()> {
832 let projected_file_count = server_info.tools.len() + FIXED_FILE_COUNT;
833 if projected_file_count > MAX_GENERATED_FILES {
834 return Err(Error::ResourceLimitExceeded {
835 resource: ResourceKind::ToolCount {
836 server_id: server_info.id.clone(),
837 },
838 actual: server_info.tools.len(),
839 limit: MAX_GENERATED_FILES - FIXED_FILE_COUNT,
840 });
841 }
842 Ok(())
843}
844
845fn add_tracked(
864 code: &mut GeneratedCode,
865 total_bytes: &mut usize,
866 file: GeneratedFile,
867) -> Result<()> {
868 *total_bytes += file.content.len();
869 if *total_bytes > MAX_GENERATED_BYTES {
870 return Err(Error::ResourceLimitExceeded {
871 resource: ResourceKind::GeneratedOutputSize,
872 actual: *total_bytes,
873 limit: MAX_GENERATED_BYTES,
874 });
875 }
876
877 code.add_file(file)?;
878
879 if code.file_count() > MAX_GENERATED_FILES {
880 return Err(Error::ResourceLimitExceeded {
881 resource: ResourceKind::GeneratedFileCount,
882 actual: code.file_count(),
883 limit: MAX_GENERATED_FILES,
884 });
885 }
886
887 Ok(())
888}
889
890fn sanitize_jsdoc(s: &str, max_len: usize) -> String {
905 let neutralized = mcp_execution_core::untrusted::sanitize_untrusted_text(s, usize::MAX);
906 let sanitized = neutralized.replace("*/", "*\\/");
907 if sanitized.chars().count() > max_len {
908 sanitized.chars().take(max_len).collect()
909 } else {
910 sanitized
911 }
912}
913
914fn render_keywords_for_jsdoc(keywords: &[String]) -> String {
918 sanitize_jsdoc(&keywords.join(", "), 256)
919}
920
921pub(crate) fn sanitize_ts_string_literal(s: &str) -> String {
929 s.replace('\\', "\\\\")
930 .replace('\'', "\\'")
931 .replace('\r', "\\r")
932 .replace('\n', "\\n")
933 .replace('\u{2028}', "\\u2028")
934 .replace('\u{2029}', "\\u2029")
935}
936
937const RESERVED_WORDS: &[&str] = &[
943 "arguments",
944 "await",
945 "break",
946 "case",
947 "catch",
948 "class",
949 "const",
950 "continue",
951 "debugger",
952 "default",
953 "delete",
954 "do",
955 "else",
956 "enum",
957 "eval",
958 "export",
959 "extends",
960 "false",
961 "finally",
962 "for",
963 "function",
964 "if",
965 "implements",
966 "import",
967 "in",
968 "instanceof",
969 "interface",
970 "let",
971 "new",
972 "null",
973 "package",
974 "private",
975 "protected",
976 "public",
977 "return",
978 "static",
979 "super",
980 "switch",
981 "this",
982 "throw",
983 "true",
984 "try",
985 "typeof",
986 "var",
987 "void",
988 "while",
989 "with",
990 "yield",
991];
992
993const RESERVED_OUTPUT_NAMES: &[&str] = &["index"];
1001
1002fn resolve_typescript_names(tools: &[ToolInfo]) -> Vec<String> {
1031 let mut used_lower: HashSet<String> = RESERVED_OUTPUT_NAMES
1032 .iter()
1033 .map(|&s| s.to_ascii_lowercase())
1034 .collect();
1035 let mut resolved = Vec::with_capacity(tools.len());
1036
1037 for tool in tools {
1038 let base = sanitize_ts_identifier(&to_camel_case(tool.name.as_str()));
1039 resolved.push(disambiguate_output_filename(&base, &mut used_lower));
1040 }
1041
1042 resolved
1043}
1044
1045fn disambiguate_output_filename(base: &str, used_lower: &mut HashSet<String>) -> String {
1059 let mut candidate = base.to_string();
1060 let mut suffix = 2;
1061 loop {
1062 let is_reserved_word = RESERVED_WORDS.contains(&candidate.as_str());
1063 if !is_reserved_word && used_lower.insert(candidate.to_ascii_lowercase()) {
1064 return candidate;
1065 }
1066 candidate = format!("{base}_{suffix}");
1067 suffix += 1;
1068 }
1069}
1070
1071fn sanitize_schema_jsdoc_descriptions(mut value: serde_json::Value) -> serde_json::Value {
1072 let mut cap_hit = false;
1073 sanitize_schema_jsdoc_value(&mut value, 0, &mut cap_hit);
1074 if cap_hit {
1075 tracing::warn!(
1080 max_depth = MAX_SCHEMA_RECURSION_DEPTH,
1081 "schema nesting exceeded MAX_SCHEMA_RECURSION_DEPTH; descriptions beyond that depth \
1082 were left unsanitized"
1083 );
1084 }
1085 value
1086}
1087
1088fn sanitize_schema_jsdoc_value(value: &mut serde_json::Value, depth: usize, cap_hit: &mut bool) {
1102 if depth >= MAX_SCHEMA_RECURSION_DEPTH {
1103 *cap_hit = true;
1104 return;
1105 }
1106
1107 match value {
1108 serde_json::Value::Object(map) => {
1109 for (key, child) in map.iter_mut() {
1110 if key == "description" {
1111 if let Some(description) = child.as_str() {
1112 *child = serde_json::Value::String(sanitize_jsdoc(description, 256));
1113 } else {
1114 *child = serde_json::Value::Null;
1115 }
1116 } else {
1117 sanitize_schema_jsdoc_value(child, depth + 1, cap_hit);
1118 }
1119 }
1120 }
1121 serde_json::Value::Array(values) => {
1122 for child in values {
1123 sanitize_schema_jsdoc_value(child, depth + 1, cap_hit);
1124 }
1125 }
1126 _ => {}
1127 }
1128}
1129
1130#[cfg(test)]
1131mod tests {
1132 use super::*;
1133 use mcp_execution_core::{ServerId, ToolName};
1134 use mcp_execution_introspector::{ServerCapabilities, ToolInfo};
1135 use serde_json::json;
1136
1137 fn create_test_server_info() -> ServerInfo {
1138 ServerInfo {
1139 id: ServerId::new("test-server").unwrap(),
1140 name: "Test Server".to_string(),
1141 version: "1.0.0".to_string(),
1142 tools: vec![
1143 ToolInfo {
1144 name: ToolName::new("create_issue").unwrap(),
1145 description: "Creates a new issue".to_string(),
1146 input_schema: json!({
1147 "type": "object",
1148 "properties": {
1149 "title": {
1150 "type": "string",
1151 "description": "Issue title"
1152 },
1153 "body": {
1154 "type": "string",
1155 "description": "Issue body"
1156 }
1157 },
1158 "required": ["title"]
1159 }),
1160 output_schema: None,
1161 },
1162 ToolInfo {
1163 name: ToolName::new("update_issue").unwrap(),
1164 description: "Updates an existing issue".to_string(),
1165 input_schema: json!({
1166 "type": "object",
1167 "properties": {
1168 "id": {
1169 "type": "number"
1170 }
1171 },
1172 "required": ["id"]
1173 }),
1174 output_schema: None,
1175 },
1176 ],
1177 capabilities: ServerCapabilities {
1178 supports_tools: true,
1179 supports_resources: false,
1180 supports_prompts: false,
1181 },
1182 }
1183 }
1184
1185 #[test]
1186 fn test_progressive_generator_new() {
1187 let generator = ProgressiveGenerator::new();
1188 assert!(generator.is_ok());
1189 }
1190
1191 #[test]
1192 fn test_generate_progressive_files() {
1193 let generator = ProgressiveGenerator::new().unwrap();
1194 let server_info = create_test_server_info();
1195
1196 let code = generator.generate(&server_info).unwrap();
1197
1198 assert_eq!(code.file_count(), 7);
1206
1207 let tool_files: Vec<_> = code.files.iter().map(|f| f.path.as_str()).collect();
1209
1210 assert!(tool_files.contains(&"createIssue.ts"));
1211 assert!(tool_files.contains(&"updateIssue.ts"));
1212 assert!(tool_files.contains(&"index.ts"));
1213 assert!(tool_files.contains(&"_runtime/mcp-bridge.ts"));
1214 assert!(tool_files.contains(&"package.json"));
1215 assert!(tool_files.contains(&"tsconfig.json"));
1216 assert!(tool_files.contains(&"_meta.json"));
1217 }
1218
1219 #[test]
1225 fn test_generate_index_ts_has_no_category_grouping() {
1226 let generator = ProgressiveGenerator::new().unwrap();
1227 let server_info = create_test_server_info();
1228
1229 let code = generator.generate(&server_info).unwrap();
1230 let index_file = code.files.iter().find(|f| f.path == "index.ts").unwrap();
1231
1232 assert!(
1233 !index_file.content.contains("uncategorized"),
1234 "generate()'s index.ts must not contain category grouping: {}",
1235 index_file.content
1236 );
1237 }
1238
1239 #[test]
1246 fn test_generate_tsconfig_json_allows_ts_extension_imports() {
1247 let generator = ProgressiveGenerator::new().unwrap();
1248 let server_info = create_test_server_info();
1249
1250 let code = generator.generate(&server_info).unwrap();
1251
1252 let tsconfig_file = code
1253 .files
1254 .iter()
1255 .find(|f| f.path == "tsconfig.json")
1256 .expect("tsconfig.json not found");
1257
1258 let parsed: serde_json::Value =
1259 serde_json::from_str(&tsconfig_file.content).expect("tsconfig.json is not valid JSON");
1260 let compiler_options = &parsed["compilerOptions"];
1261
1262 assert_eq!(compiler_options["allowImportingTsExtensions"], true);
1263 assert_eq!(compiler_options["noEmit"], true);
1264 assert_eq!(compiler_options["types"], serde_json::json!(["node"]));
1265 }
1266
1267 #[test]
1272 fn test_generate_package_json_declares_types_node_dev_dependency() {
1273 let generator = ProgressiveGenerator::new().unwrap();
1274 let server_info = create_test_server_info();
1275
1276 let code = generator.generate(&server_info).unwrap();
1277
1278 let package_json_file = code
1279 .files
1280 .iter()
1281 .find(|f| f.path == "package.json")
1282 .expect("package.json not found");
1283
1284 let parsed: serde_json::Value = serde_json::from_str(&package_json_file.content)
1285 .expect("package.json is not valid JSON");
1286
1287 assert!(
1288 parsed["devDependencies"]["@types/node"].is_string(),
1289 "package.json is missing a @types/node devDependency: {parsed}"
1290 );
1291 }
1292
1293 #[test]
1294 fn test_generate_meta_json_preserves_parameter_descriptions() {
1295 let generator = ProgressiveGenerator::new().unwrap();
1299 let server_info = create_test_server_info();
1300
1301 let code = generator.generate(&server_info).unwrap();
1302 let meta_file = code.files.iter().find(|f| f.path == "_meta.json").unwrap();
1303 let meta: ServerMetadata = serde_json::from_str(&meta_file.content).unwrap();
1304
1305 assert_eq!(meta.schema_version, METADATA_SCHEMA_VERSION);
1306 assert_eq!(meta.server_id.as_str(), "test-server");
1307 assert_eq!(meta.server_name, "Test Server");
1308 assert_eq!(meta.server_version, "1.0.0");
1309 assert_eq!(meta.tools.len(), 2);
1310
1311 let create_issue = meta
1312 .tools
1313 .iter()
1314 .find(|t| t.name.as_str() == "create_issue")
1315 .unwrap();
1316 assert_eq!(create_issue.typescript_name, "createIssue");
1317 let title = create_issue
1318 .parameters
1319 .iter()
1320 .find(|p| p.name == "title")
1321 .unwrap();
1322 assert_eq!(title.description, Some("Issue title".to_string()));
1323 assert!(title.required);
1324 }
1325
1326 #[test]
1327 fn test_generate_with_categories_meta_json_includes_categorization() {
1328 let generator = ProgressiveGenerator::new().unwrap();
1329 let server_info = create_test_server_info();
1330
1331 let mut categorizations = HashMap::new();
1332 categorizations.insert(
1333 "create_issue".to_string(),
1334 ToolCategorization {
1335 category: "issues".to_string(),
1336 keywords: vec!["create".to_string(), "issue".to_string(), "new".to_string()],
1337 short_description: "Create a new issue".to_string(),
1338 },
1339 );
1340
1341 let code = generator
1342 .generate_with_categories(&server_info, &categorizations)
1343 .unwrap();
1344 let meta_file = code.files.iter().find(|f| f.path == "_meta.json").unwrap();
1345 let meta: ServerMetadata = serde_json::from_str(&meta_file.content).unwrap();
1346
1347 let create_issue = meta
1348 .tools
1349 .iter()
1350 .find(|t| t.name.as_str() == "create_issue")
1351 .unwrap();
1352 assert_eq!(create_issue.category, Some("issues".to_string()));
1353 assert_eq!(
1354 create_issue.keywords,
1355 vec!["create".to_string(), "issue".to_string(), "new".to_string()]
1356 );
1357
1358 let update_issue = meta
1359 .tools
1360 .iter()
1361 .find(|t| t.name.as_str() == "update_issue")
1362 .unwrap();
1363 assert!(update_issue.category.is_none());
1364 assert!(update_issue.keywords.is_empty());
1365 }
1366
1367 #[test]
1368 fn test_generate_meta_json_parameter_description_is_raw_not_jsdoc_sanitized() {
1369 let raw_description = format!(
1374 "Matches C-style /* */ comment blocks.\nSecond line follows. {}",
1375 "x".repeat(300)
1376 );
1377 assert!(raw_description.contains("*/"));
1378 assert!(raw_description.contains('\n'));
1379 assert!(raw_description.chars().count() > 256);
1380
1381 let server_info = ServerInfo {
1382 id: ServerId::new("test-server").unwrap(),
1383 name: "Test Server".to_string(),
1384 version: "1.0.0".to_string(),
1385 tools: vec![ToolInfo {
1386 name: ToolName::new("send_message").unwrap(),
1387 description: "Sends a message".to_string(),
1388 input_schema: json!({
1389 "type": "object",
1390 "properties": {
1391 "notes": {
1392 "type": "string",
1393 "description": raw_description
1394 }
1395 },
1396 "required": []
1397 }),
1398 output_schema: None,
1399 }],
1400 capabilities: ServerCapabilities {
1401 supports_tools: true,
1402 supports_resources: false,
1403 supports_prompts: false,
1404 },
1405 };
1406
1407 let generator = ProgressiveGenerator::new().unwrap();
1408 let code = generator.generate(&server_info).unwrap();
1409
1410 let meta_file = code.files.iter().find(|f| f.path == "_meta.json").unwrap();
1412 let meta: ServerMetadata = serde_json::from_str(&meta_file.content).unwrap();
1413 let send_message = meta
1414 .tools
1415 .iter()
1416 .find(|t| t.name.as_str() == "send_message")
1417 .unwrap();
1418 let notes = send_message
1419 .parameters
1420 .iter()
1421 .find(|p| p.name == "notes")
1422 .unwrap();
1423 assert_eq!(notes.description, Some(raw_description.clone()));
1424
1425 let ts_file = code
1428 .files
1429 .iter()
1430 .find(|f| f.path == "sendMessage.ts")
1431 .unwrap();
1432 assert!(
1433 !ts_file.content.contains(raw_description.as_str()),
1434 "the .ts file must not contain the raw, un-sanitized description verbatim"
1435 );
1436 assert!(
1437 ts_file.content.contains("*\\/"),
1438 "the .ts file must escape '*/' to avoid closing the JSDoc comment early"
1439 );
1440 assert!(
1441 !ts_file
1442 .content
1443 .contains("Matches C-style /* */ comment blocks.\nSecond"),
1444 "the .ts file must flatten newlines within the description to spaces"
1445 );
1446 }
1447
1448 #[test]
1449 fn test_create_tool_context() {
1450 let tool = ToolInfo {
1451 name: ToolName::new("send_message").unwrap(),
1452 description: "Sends a message".to_string(),
1453 input_schema: json!({
1454 "type": "object",
1455 "properties": {
1456 "text": {"type": "string"}
1457 },
1458 "required": ["text"]
1459 }),
1460 output_schema: None,
1461 };
1462
1463 let categorization = ToolCategorization {
1464 category: "messaging".to_string(),
1465 keywords: vec![
1466 "send".to_string(),
1467 "message".to_string(),
1468 "chat".to_string(),
1469 ],
1470 short_description: "Send a message".to_string(),
1471 };
1472 let properties = ProgressiveGenerator::extract_property_infos(&tool.input_schema).unwrap();
1473 let context = ProgressiveGenerator::create_tool_context(
1474 "test-server",
1475 &tool,
1476 Some(&categorization),
1477 "sendMessage".to_string(),
1478 properties,
1479 );
1480
1481 assert_eq!(context.server_id, "test-server");
1482 assert_eq!(context.name, "send_message");
1483 assert_eq!(context.name_literal, "send_message");
1484 assert_eq!(context.server_id_literal, "test-server");
1485 assert_eq!(context.typescript_name, "sendMessage");
1486 assert_eq!(context.description, "Sends a message");
1487 assert_eq!(context.properties.len(), 1);
1488 assert_eq!(context.properties[0].name, "text");
1489 assert_eq!(context.category, Some("messaging".to_string()));
1490 assert_eq!(context.keywords, Some("send, message, chat".to_string()));
1491 assert_eq!(context.short_description, "Send a message".to_string());
1492 }
1493
1494 #[test]
1495 fn test_wrap_tool_generation_error_preserves_tool_name_and_source() {
1496 let tool = ToolInfo {
1500 name: ToolName::new("send_message").unwrap(),
1501 description: String::new(),
1502 input_schema: json!({}),
1503 output_schema: None,
1504 };
1505 let source = Error::ValidationError {
1506 field: "type".to_string(),
1507 reason: "Property type is not a string".to_string(),
1508 };
1509
1510 let wrapped = ProgressiveGenerator::wrap_tool_generation_error(
1511 &tool,
1512 "extract property schema",
1513 source,
1514 );
1515
1516 match wrapped {
1517 Error::ScriptGenerationError {
1518 tool: tool_name,
1519 message,
1520 source,
1521 } => {
1522 assert_eq!(tool_name, "send_message");
1523 assert_eq!(message, "failed to extract property schema");
1526 let source = source.expect("source must be preserved for exit-code classification");
1527 assert!(source.to_string().contains("Property type is not a string"));
1528 }
1529 other => panic!("expected ScriptGenerationError, got {other:?}"),
1530 }
1531 }
1532
1533 #[test]
1534 fn test_wrap_tool_generation_error_covers_render_and_tracking_stages() {
1535 let tool = ToolInfo {
1539 name: ToolName::new("send_message").unwrap(),
1540 description: String::new(),
1541 input_schema: json!({}),
1542 output_schema: None,
1543 };
1544
1545 let render_failure = ProgressiveGenerator::wrap_tool_generation_error(
1546 &tool,
1547 "render tool template",
1548 Error::SerializationError {
1549 message: "Template rendering failed: boom".to_string(),
1550 source: None,
1551 },
1552 );
1553 assert!(render_failure.is_script_generation_error());
1554
1555 let tracking_failure = ProgressiveGenerator::wrap_tool_generation_error(
1556 &tool,
1557 "track generated tool file",
1558 Error::ResourceLimitExceeded {
1559 resource: ResourceKind::GeneratedOutputSize,
1560 actual: 10,
1561 limit: 5,
1562 },
1563 );
1564 match tracking_failure {
1565 Error::ScriptGenerationError {
1566 tool: tool_name,
1567 source,
1568 ..
1569 } => {
1570 assert_eq!(tool_name, "send_message");
1571 let source = source.expect("source must be preserved for exit-code classification");
1572 assert!(
1575 source
1576 .downcast_ref::<Error>()
1577 .unwrap()
1578 .is_resource_limit_exceeded()
1579 );
1580 }
1581 other => panic!("expected ScriptGenerationError, got {other:?}"),
1582 }
1583 }
1584
1585 #[test]
1586 fn test_create_tool_context_without_categorization_falls_back_to_description() {
1587 let generator = ProgressiveGenerator::new().unwrap();
1588 let tool = ToolInfo {
1589 name: ToolName::new("format_document").unwrap(),
1590 description: "Format document with language-specific rules".to_string(),
1591 input_schema: json!({
1592 "type": "object",
1593 "properties": {
1594 "text": {"type": "string"}
1595 },
1596 "required": ["text"]
1597 }),
1598 output_schema: None,
1599 };
1600
1601 let properties = ProgressiveGenerator::extract_property_infos(&tool.input_schema).unwrap();
1602 let context = ProgressiveGenerator::create_tool_context(
1603 "test-server",
1604 &tool,
1605 None,
1606 "formatDocument".to_string(),
1607 properties,
1608 );
1609
1610 assert_eq!(
1611 context.short_description,
1612 "Format document with language-specific rules".to_string()
1613 );
1614
1615 let rendered = generator
1617 .engine
1618 .render("progressive/tool", &context)
1619 .unwrap();
1620 assert!(rendered.contains("@description Format document with language-specific rules"));
1621 }
1622
1623 #[test]
1624 fn test_create_tool_context_input_schema_is_sanitized() {
1625 let tool = ToolInfo {
1626 name: ToolName::new("send_message").unwrap(),
1627 description: "Sends a message".to_string(),
1628 input_schema: json!({
1629 "type": "object",
1630 "description": "Schema */ injected\nnext",
1631 "properties": {
1632 "text": {"type": "string"}
1633 },
1634 "required": ["text"]
1635 }),
1636 output_schema: None,
1637 };
1638
1639 let properties = ProgressiveGenerator::extract_property_infos(&tool.input_schema).unwrap();
1640 let context = ProgressiveGenerator::create_tool_context(
1641 "test-server",
1642 &tool,
1643 None,
1644 "sendMessage".to_string(),
1645 properties,
1646 );
1647
1648 let expected = sanitize_schema_jsdoc_descriptions(tool.input_schema);
1649 assert_eq!(context.input_schema, expected);
1650 assert_eq!(
1651 context.input_schema["description"],
1652 json!("Schema *\\/ injected next")
1653 );
1654 }
1655
1656 #[test]
1657 fn test_create_index_context() {
1658 let server_info = create_test_server_info();
1659 let typescript_names = resolve_typescript_names(&server_info.tools);
1660
1661 let context =
1662 ProgressiveGenerator::create_index_context(&server_info, None, &typescript_names);
1663
1664 assert_eq!(context.server_name, "Test Server");
1665 assert_eq!(context.server_version, "1.0.0");
1666 assert_eq!(context.tool_count, 2);
1667 assert_eq!(context.tools.len(), 2);
1668 assert_eq!(context.tools[0].typescript_name, "createIssue");
1669 assert!(context.categories.is_none());
1670 }
1671
1672 #[test]
1677 fn test_tool_named_index_does_not_collide_with_index_ts() {
1678 let generator = ProgressiveGenerator::new().unwrap();
1679 let server_info = ServerInfo {
1680 id: ServerId::new("test-server").unwrap(),
1681 name: "Test Server".to_string(),
1682 version: "1.0.0".to_string(),
1683 tools: vec![ToolInfo {
1684 name: ToolName::new("index").unwrap(),
1685 description: "A tool literally named index".to_string(),
1686 input_schema: json!({
1687 "type": "object",
1688 "properties": {},
1689 "required": []
1690 }),
1691 output_schema: None,
1692 }],
1693 capabilities: ServerCapabilities {
1694 supports_tools: true,
1695 supports_resources: false,
1696 supports_prompts: false,
1697 },
1698 };
1699
1700 let code = generator.generate(&server_info).unwrap();
1701
1702 let typescript_names = resolve_typescript_names(&server_info.tools);
1703 assert_eq!(
1704 typescript_names[0], "index_2",
1705 "a tool named `index` must be disambiguated, not collide with the fixed index.ts"
1706 );
1707
1708 let tool_file = code
1709 .files
1710 .iter()
1711 .find(|f| f.path == "index_2.ts")
1712 .expect("the tool's own file must exist at its disambiguated path");
1713 assert!(
1714 tool_file.content.contains("A tool literally named index"),
1715 "the tool's own generated content must not have been lost: {}",
1716 tool_file.content
1717 );
1718
1719 let index_file = code
1720 .files
1721 .iter()
1722 .find(|f| f.path == "index.ts")
1723 .expect("the fixed index.ts re-export must still exist");
1724 assert!(
1725 index_file.content.contains("index_2"),
1726 "index.ts must re-export the tool's disambiguated identifier: {}",
1727 index_file.content
1728 );
1729 assert!(
1730 index_file
1731 .content
1732 .contains("export { callMCPTool } from './_runtime/mcp-bridge.ts';"),
1733 "index.ts must be the fixed re-export (with the runtime bridge re-export), \
1734 not the overwritten tool file: {}",
1735 index_file.content
1736 );
1737
1738 assert_eq!(
1740 code.files.iter().filter(|f| f.path == "index.ts").count(),
1741 1
1742 );
1743 assert_eq!(
1744 code.files.iter().filter(|f| f.path == "index_2.ts").count(),
1745 1
1746 );
1747 }
1748
1749 #[test]
1755 fn test_tool_named_index_with_different_case_is_disambiguated() {
1756 let generator = ProgressiveGenerator::new().unwrap();
1757 let server_info = ServerInfo {
1758 id: ServerId::new("test-server").unwrap(),
1759 name: "Test Server".to_string(),
1760 version: "1.0.0".to_string(),
1761 tools: vec![ToolInfo {
1762 name: ToolName::new("Index").unwrap(),
1763 description: "A tool literally named Index".to_string(),
1764 input_schema: json!({
1765 "type": "object",
1766 "properties": {},
1767 "required": []
1768 }),
1769 output_schema: None,
1770 }],
1771 capabilities: ServerCapabilities {
1772 supports_tools: true,
1773 supports_resources: false,
1774 supports_prompts: false,
1775 },
1776 };
1777
1778 let code = generator.generate(&server_info).unwrap();
1779
1780 let typescript_names = resolve_typescript_names(&server_info.tools);
1781 assert_eq!(
1782 typescript_names[0], "Index_2",
1783 "a tool named `Index` must be disambiguated case-insensitively against the \
1784 reserved `index` output name"
1785 );
1786
1787 assert!(
1788 code.files.iter().any(|f| f.path == "Index_2.ts"),
1789 "the tool's own file must exist at its disambiguated path: {:?}",
1790 code.files.iter().map(|f| &f.path).collect::<Vec<_>>()
1791 );
1792 assert_eq!(
1795 code.files.iter().filter(|f| f.path == "index.ts").count(),
1796 1
1797 );
1798 assert!(!code.files.iter().any(|f| f.path == "Index.ts"));
1799 }
1800
1801 #[test]
1807 fn test_tools_named_index_and_capital_index_do_not_collide_with_each_other() {
1808 let generator = ProgressiveGenerator::new().unwrap();
1809 let server_info = ServerInfo {
1810 id: ServerId::new("test-server").unwrap(),
1811 name: "Test Server".to_string(),
1812 version: "1.0.0".to_string(),
1813 tools: vec![
1814 ToolInfo {
1815 name: ToolName::new("Index").unwrap(),
1816 description: "Capitalized".to_string(),
1817 input_schema: json!({"type": "object", "properties": {}, "required": []}),
1818 output_schema: None,
1819 },
1820 ToolInfo {
1821 name: ToolName::new("index").unwrap(),
1822 description: "Lowercase".to_string(),
1823 input_schema: json!({"type": "object", "properties": {}, "required": []}),
1824 output_schema: None,
1825 },
1826 ],
1827 capabilities: ServerCapabilities {
1828 supports_tools: true,
1829 supports_resources: false,
1830 supports_prompts: false,
1831 },
1832 };
1833
1834 let code = generator.generate(&server_info).unwrap();
1835 let typescript_names = resolve_typescript_names(&server_info.tools);
1836
1837 assert_ne!(
1838 typescript_names[0].to_ascii_lowercase(),
1839 typescript_names[1].to_ascii_lowercase(),
1840 "the two tools' resolved names must not collide case-insensitively: {typescript_names:?}"
1841 );
1842
1843 let paths: Vec<_> = code.files.iter().map(|f| f.path.as_str()).collect();
1844 let mut lowercased_paths: Vec<String> =
1845 paths.iter().map(|p| p.to_ascii_lowercase()).collect();
1846 let before = lowercased_paths.len();
1847 lowercased_paths.sort();
1848 lowercased_paths.dedup();
1849 assert_eq!(
1850 lowercased_paths.len(),
1851 before,
1852 "no two generated file paths may be case-insensitive duplicates of each other: {paths:?}"
1853 );
1854 }
1855
1856 #[test]
1861 fn test_tools_differing_only_by_case_are_disambiguated_from_each_other() {
1862 let server_info = ServerInfo {
1863 id: ServerId::new("test-server").unwrap(),
1864 name: "Test Server".to_string(),
1865 version: "1.0.0".to_string(),
1866 tools: vec![
1867 ToolInfo {
1868 name: ToolName::new("get_user").unwrap(),
1869 description: "snake_case".to_string(),
1870 input_schema: json!({"type": "object", "properties": {}, "required": []}),
1871 output_schema: None,
1872 },
1873 ToolInfo {
1874 name: ToolName::new("GetUser").unwrap(),
1875 description: "PascalCase".to_string(),
1876 input_schema: json!({"type": "object", "properties": {}, "required": []}),
1877 output_schema: None,
1878 },
1879 ],
1880 capabilities: ServerCapabilities {
1881 supports_tools: true,
1882 supports_resources: false,
1883 supports_prompts: false,
1884 },
1885 };
1886
1887 let typescript_names = resolve_typescript_names(&server_info.tools);
1888
1889 assert_eq!(typescript_names[0], "getUser");
1890 assert_ne!(
1891 typescript_names[0].to_ascii_lowercase(),
1892 typescript_names[1].to_ascii_lowercase(),
1893 "getUser/GetUser must not collide case-insensitively: {typescript_names:?}"
1894 );
1895 }
1896
1897 #[test]
1898 fn test_extract_property_infos() {
1899 let schema = json!({
1900 "type": "object",
1901 "properties": {
1902 "name": {
1903 "type": "string",
1904 "description": "User name"
1905 },
1906 "age": {
1907 "type": "number"
1908 }
1909 },
1910 "required": ["name"]
1911 });
1912
1913 let props = ProgressiveGenerator::extract_property_infos(&schema).unwrap();
1914
1915 assert_eq!(props.len(), 2);
1916
1917 let name_prop = props.iter().find(|p| p.name == "name").unwrap();
1919 assert_eq!(name_prop.typescript_type, "string");
1920 assert_eq!(name_prop.description, Some("User name".to_string()));
1921 assert!(name_prop.required);
1922
1923 let age_prop = props.iter().find(|p| p.name == "age").unwrap();
1925 assert_eq!(age_prop.typescript_type, "number");
1926 assert!(!age_prop.required);
1927 }
1928
1929 #[test]
1930 fn test_extract_property_infos_sanitizes_malicious_property_name() {
1931 let schema = json!({
1932 "type": "object",
1933 "properties": {
1934 "x: string }; export const pwned = 1; interface J {": {
1935 "type": "string",
1936 "description": "Evil property"
1937 }
1938 },
1939 "required": []
1940 });
1941
1942 let props = ProgressiveGenerator::extract_property_infos(&schema).unwrap();
1943
1944 assert_eq!(props.len(), 1);
1945 assert!(!props[0].name.contains(['{', '}', ';', ':', ' ']));
1946 assert_eq!(props[0].description, Some("Evil property".to_string()));
1949 }
1950
1951 #[test]
1952 fn test_extract_property_infos_disambiguates_colliding_sibling_names() {
1953 let schema = json!({
1957 "type": "object",
1958 "properties": {
1959 "a-b": {"type": "string"},
1960 "a.b": {"type": "number"}
1961 },
1962 "required": []
1963 });
1964
1965 let props = ProgressiveGenerator::extract_property_infos(&schema).unwrap();
1966 let mut names: Vec<&str> = props.iter().map(|p| p.name.as_str()).collect();
1967 names.sort_unstable();
1968
1969 assert_eq!(names, vec!["a_b", "a_b_2"]);
1970 }
1971
1972 #[test]
1973 fn test_extract_property_infos_disambiguates_collision_introduced_by_collapsing() {
1974 let schema = json!({
1979 "type": "object",
1980 "properties": {
1981 "a-b": {"type": "string"},
1982 "a--b": {"type": "number"}
1983 },
1984 "required": []
1985 });
1986
1987 let props = ProgressiveGenerator::extract_property_infos(&schema).unwrap();
1988 let mut names: Vec<&str> = props.iter().map(|p| p.name.as_str()).collect();
1989 names.sort_unstable();
1990
1991 assert_eq!(names, vec!["a_b", "a_b_2"]);
1992 }
1993
1994 #[test]
1995 fn test_extract_property_infos_disambiguates_three_way_collision() {
1996 let schema = json!({
1997 "type": "object",
1998 "properties": {
1999 "a-b": {"type": "string"},
2000 "a.b": {"type": "number"},
2001 "a b": {"type": "boolean"}
2002 },
2003 "required": []
2004 });
2005
2006 let props = ProgressiveGenerator::extract_property_infos(&schema).unwrap();
2007 let mut names: Vec<&str> = props.iter().map(|p| p.name.as_str()).collect();
2008 names.sort_unstable();
2009
2010 assert_eq!(names, vec!["a_b", "a_b_2", "a_b_3"]);
2011 }
2012
2013 #[test]
2014 fn test_generate_disambiguates_colliding_top_level_params() {
2015 let generator = ProgressiveGenerator::new().unwrap();
2016 let mut server_info = create_test_server_info();
2017 server_info.tools[0].input_schema = json!({
2018 "type": "object",
2019 "properties": {
2020 "a-b": {"type": "string"},
2021 "a.b": {"type": "number"}
2022 },
2023 "required": []
2024 });
2025
2026 let code = generator.generate(&server_info).unwrap();
2027 let tool = code
2028 .files
2029 .iter()
2030 .find(|f| f.path == "createIssue.ts")
2031 .unwrap();
2032
2033 assert_eq!(
2034 tool.content.matches("a_b:").count() + tool.content.matches("a_b?:").count(),
2035 1,
2036 "field 'a_b' must appear exactly once in the Params interface: {}",
2037 tool.content
2038 );
2039 assert_eq!(
2040 tool.content.matches("a_b_2:").count() + tool.content.matches("a_b_2?:").count(),
2041 1,
2042 "disambiguated field 'a_b_2' must appear exactly once in the Params interface: {}",
2043 tool.content
2044 );
2045 }
2046
2047 #[test]
2048 fn test_generate_sanitizes_property_name_injection() {
2049 let generator = ProgressiveGenerator::new().unwrap();
2050 let mut server_info = create_test_server_info();
2051 server_info.tools[0].input_schema = json!({
2052 "type": "object",
2053 "properties": {
2054 "x: string }; export const pwned = evil(); interface J {": {"type": "string"}
2055 },
2056 "required": []
2057 });
2058
2059 let code = generator.generate(&server_info).unwrap();
2060 let tool = code
2061 .files
2062 .iter()
2063 .find(|f| f.path == "createIssue.ts")
2064 .unwrap();
2065
2066 assert!(
2067 !tool.content.contains("export const pwned"),
2068 "raw property name must not inject a top-level statement: {}",
2069 tool.content
2070 );
2071 }
2072
2073 #[test]
2074 fn test_sanitize_jsdoc_strips_comment_terminator() {
2075 assert_eq!(sanitize_jsdoc("Foo */ bar", 256), "Foo *\\/ bar");
2076 }
2077
2078 #[test]
2079 fn test_sanitize_jsdoc_replaces_newlines() {
2080 assert_eq!(
2081 sanitize_jsdoc("line1\nline2\r\nline3", 256),
2082 "line1 line2 line3"
2083 );
2084 }
2085
2086 #[test]
2087 fn test_sanitize_jsdoc_replaces_unicode_line_terminators() {
2088 assert_eq!(
2093 sanitize_jsdoc("line1\u{2028}line2\u{2029}line3", 256),
2094 "line1 line2 line3"
2095 );
2096 }
2097
2098 #[test]
2099 fn test_generate_with_categories_sanitizes_unicode_line_terminator_in_category() {
2100 let generator = ProgressiveGenerator::new().unwrap();
2101 let server_info = create_test_server_info();
2102
2103 let mut categorizations = HashMap::new();
2104 categorizations.insert(
2105 "create_issue".to_string(),
2106 ToolCategorization {
2107 category: "issues\u{2028}export const pwned = 1;".to_string(),
2108 keywords: vec![],
2109 short_description: "Create a new issue".to_string(),
2110 },
2111 );
2112
2113 let code = generator
2114 .generate_with_categories(&server_info, &categorizations)
2115 .unwrap();
2116 let index = code.files.iter().find(|f| f.path == "index.ts").unwrap();
2117
2118 assert!(
2123 index
2124 .content
2125 .contains("// --- issues export const pwned = 1; ---"),
2126 "sanitized category text should remain inert inside the comment: {}",
2127 index.content
2128 );
2129 assert!(
2130 !index.content.contains("\nexport const pwned"),
2131 "U+2028 must not terminate the `// --- {{category}} ---` line comment and \
2132 inject a live top-level statement: {}",
2133 index.content
2134 );
2135 }
2136
2137 #[test]
2138 fn test_sanitize_jsdoc_truncates() {
2139 let long = "a".repeat(300);
2140 assert_eq!(sanitize_jsdoc(&long, 256).chars().count(), 256);
2141 }
2142
2143 #[test]
2144 fn test_sanitize_jsdoc_passthrough() {
2145 assert_eq!(sanitize_jsdoc("Normal string", 256), "Normal string");
2146 }
2147
2148 #[test]
2149 fn test_sanitize_jsdoc_strips_ansi_escape_sequence() {
2150 let payload = "Innocuous \x1b[31mred text\x1b[0m looking description";
2155 let sanitized = sanitize_jsdoc(payload, 256);
2156 assert!(!sanitized.contains('\x1b'));
2157 assert_eq!(sanitized, "Innocuous [31mred text [0m looking description");
2158 }
2159
2160 #[test]
2161 fn test_sanitize_jsdoc_replaces_other_c0_control_characters_with_space() {
2162 assert_eq!(sanitize_jsdoc("a\u{0}b\u{7}c\u{7f}d", 256), "a b c d");
2163 }
2164
2165 #[test]
2172 fn test_sanitize_jsdoc_control_char_between_star_slash_cannot_reopen_comment() {
2173 for ctrl in ['\u{0}', '\u{7f}', '\u{1b}'] {
2174 let payload = format!("safe *{ctrl}/ export const pwned = 1; //");
2175 let sanitized = sanitize_jsdoc(&payload, 256);
2176 assert!(
2177 !sanitized.contains("*/"),
2178 "control char {ctrl:?} must not let a bare `*/` reappear: {sanitized:?}"
2179 );
2180 assert!(
2181 sanitized.contains("* /"),
2182 "the control char should be neutralized to a space, not deleted: {sanitized:?}"
2183 );
2184 }
2185 }
2186
2187 #[test]
2188 fn test_sanitize_ts_string_literal_escapes_quote_and_backslash() {
2189 assert_eq!(
2190 sanitize_ts_string_literal(r"it's a \test"),
2191 r"it\'s a \\test"
2192 );
2193 }
2194
2195 #[test]
2196 fn test_sanitize_ts_string_literal_escape_order_prevents_double_escaping() {
2197 assert_eq!(sanitize_ts_string_literal("\\'"), r"\\\'");
2200 }
2201
2202 #[test]
2203 fn test_sanitize_ts_string_literal_escapes_newlines() {
2204 assert_eq!(
2205 sanitize_ts_string_literal("line1\nline2\rline3"),
2206 "line1\\nline2\\rline3"
2207 );
2208 }
2209
2210 #[test]
2211 fn test_sanitize_ts_string_literal_escapes_unicode_line_terminators() {
2212 assert_eq!(
2215 sanitize_ts_string_literal("line1\u{2028}line2\u{2029}line3"),
2216 "line1\\u2028line2\\u2029line3"
2217 );
2218 }
2219
2220 #[test]
2225 fn test_sanitize_ts_identifier_passthrough_valid() {
2226 assert_eq!(sanitize_ts_identifier("sendMessage_1"), "sendMessage_1");
2227 }
2228
2229 #[test]
2230 fn test_generate_sanitizes_call_site_string_literal_injection() {
2231 let generator = ProgressiveGenerator::new().unwrap();
2232 let mut server_info = create_test_server_info();
2233 server_info.tools[0].name = ToolName::new("create_issue'); alert('pwned").unwrap();
2234
2235 let code = generator.generate(&server_info).unwrap();
2236 let tool = code
2237 .files
2238 .iter()
2239 .find(|f| {
2240 std::path::Path::new(&f.path)
2241 .extension()
2242 .is_some_and(|ext| ext.eq_ignore_ascii_case("ts"))
2243 && f.path != "index.ts"
2244 })
2245 .unwrap();
2246
2247 let call_site_line = tool
2258 .content
2259 .lines()
2260 .find(|line| line.contains("return (await callMCPTool("))
2261 .expect("generated tool file must contain a callMCPTool(...) invocation");
2262 assert!(
2263 !call_site_line.contains("'); alert('pwned"),
2264 "raw quote must not break out of the callMCPTool string literal: {call_site_line}"
2265 );
2266 }
2267
2268 #[test]
2269 fn test_resolve_typescript_names_disambiguates_collisions() {
2270 let tools = vec![
2271 ToolInfo {
2272 name: ToolName::new("foo-bar").unwrap(),
2273 description: String::new(),
2274 input_schema: json!({}),
2275 output_schema: None,
2276 },
2277 ToolInfo {
2278 name: ToolName::new("foo.bar").unwrap(),
2279 description: String::new(),
2280 input_schema: json!({}),
2281 output_schema: None,
2282 },
2283 ToolInfo {
2284 name: ToolName::new("foo bar").unwrap(),
2285 description: String::new(),
2286 input_schema: json!({}),
2287 output_schema: None,
2288 },
2289 ];
2290
2291 let resolved = resolve_typescript_names(&tools);
2292 let mut names: Vec<&String> = resolved.iter().collect();
2293 names.sort();
2294
2295 assert_eq!(resolved.len(), 3);
2297 let unique: HashSet<&String> = names.iter().copied().collect();
2298 assert_eq!(
2299 unique.len(),
2300 3,
2301 "collisions must be disambiguated: {names:?}"
2302 );
2303 assert_eq!(resolved[0], "foo_bar");
2304 }
2305
2306 #[test]
2307 fn test_resolve_typescript_names_disambiguates_identical_raw_names() {
2308 let tools = vec![
2312 ToolInfo {
2313 name: ToolName::new("dup").unwrap(),
2314 description: "First".to_string(),
2315 input_schema: json!({}),
2316 output_schema: None,
2317 },
2318 ToolInfo {
2319 name: ToolName::new("dup").unwrap(),
2320 description: "Second".to_string(),
2321 input_schema: json!({}),
2322 output_schema: None,
2323 },
2324 ];
2325
2326 let resolved = resolve_typescript_names(&tools);
2327
2328 assert_eq!(resolved, vec!["dup".to_string(), "dup_2".to_string()]);
2329 }
2330
2331 #[test]
2332 fn test_resolve_typescript_names_disambiguates_three_way_identical_raw_names() {
2333 let tools: Vec<ToolInfo> = (0..3)
2334 .map(|_| ToolInfo {
2335 name: ToolName::new("dup").unwrap(),
2336 description: String::new(),
2337 input_schema: json!({}),
2338 output_schema: None,
2339 })
2340 .collect();
2341
2342 let resolved = resolve_typescript_names(&tools);
2343
2344 assert_eq!(
2345 resolved,
2346 vec!["dup".to_string(), "dup_2".to_string(), "dup_3".to_string()]
2347 );
2348 }
2349
2350 #[test]
2351 fn test_resolve_typescript_names_disambiguates_reserved_words() {
2352 let reserved_tool_names = [
2353 "delete",
2354 "typeof",
2355 "class",
2356 "new",
2357 "import",
2358 "export",
2359 "in",
2360 "instanceof",
2361 "void",
2362 "enum",
2363 "eval",
2364 "arguments",
2365 ];
2366
2367 for name in reserved_tool_names {
2368 let tools = vec![ToolInfo {
2369 name: ToolName::new(name).unwrap(),
2370 description: String::new(),
2371 input_schema: json!({}),
2372 output_schema: None,
2373 }];
2374
2375 let resolved = resolve_typescript_names(&tools);
2376 let typescript_name = &resolved[0];
2377
2378 assert_ne!(
2379 typescript_name, name,
2380 "reserved word {name} must be disambiguated"
2381 );
2382 assert!(
2383 !RESERVED_WORDS.contains(&typescript_name.as_str()),
2384 "resolved name {typescript_name} for tool {name} must not be a reserved word"
2385 );
2386 }
2387 }
2388
2389 #[test]
2390 fn test_resolve_typescript_names_reserved_word_avoids_existing_collision() {
2391 let tools = vec![
2392 ToolInfo {
2393 name: ToolName::new("class").unwrap(),
2394 description: String::new(),
2395 input_schema: json!({}),
2396 output_schema: None,
2397 },
2398 ToolInfo {
2406 name: ToolName::new("class-2").unwrap(),
2407 description: String::new(),
2408 input_schema: json!({}),
2409 output_schema: None,
2410 },
2411 ];
2412
2413 let resolved = resolve_typescript_names(&tools);
2414
2415 assert_ne!(
2416 resolved[0], resolved[1],
2417 "a reserved-word tool's fallback name must not collide with an unrelated tool that already claims it"
2418 );
2419 assert!(!RESERVED_WORDS.contains(&resolved[0].as_str()));
2420 }
2421
2422 #[test]
2423 fn test_resolve_typescript_names_reserved_word_case_variant_is_not_suffixed() {
2424 let tools = vec![ToolInfo {
2428 name: ToolName::new("Delete").unwrap(),
2429 description: String::new(),
2430 input_schema: json!({}),
2431 output_schema: None,
2432 }];
2433
2434 let resolved = resolve_typescript_names(&tools);
2435
2436 assert_eq!(resolved[0], "Delete");
2437 }
2438
2439 #[test]
2440 fn test_resolve_typescript_names_exact_reserved_word_is_still_suffixed() {
2441 let tools = vec![ToolInfo {
2442 name: ToolName::new("delete").unwrap(),
2443 description: String::new(),
2444 input_schema: json!({}),
2445 output_schema: None,
2446 }];
2447
2448 let resolved = resolve_typescript_names(&tools);
2449
2450 assert_ne!(resolved[0], "delete");
2451 assert!(!RESERVED_WORDS.contains(&resolved[0].as_str()));
2452 }
2453
2454 #[test]
2455 fn test_resolve_typescript_names_delete_and_delete_case_variant_in_same_batch() {
2456 let make_tools = |first: &str, second: &str| {
2462 vec![
2463 ToolInfo {
2464 name: ToolName::new(first).unwrap(),
2465 description: String::new(),
2466 input_schema: json!({}),
2467 output_schema: None,
2468 },
2469 ToolInfo {
2470 name: ToolName::new(second).unwrap(),
2471 description: String::new(),
2472 input_schema: json!({}),
2473 output_schema: None,
2474 },
2475 ]
2476 };
2477
2478 let delete_first = resolve_typescript_names(&make_tools("delete", "Delete"));
2479 assert_ne!(delete_first[0], "delete");
2480 assert_eq!(delete_first[1], "Delete");
2481 assert_ne!(
2482 delete_first[0].to_ascii_lowercase(),
2483 delete_first[1].to_ascii_lowercase()
2484 );
2485
2486 let delete_second = resolve_typescript_names(&make_tools("Delete", "delete"));
2487 assert_eq!(delete_second[0], "Delete");
2488 assert_ne!(delete_second[1], "delete");
2489 assert_ne!(
2490 delete_second[0].to_ascii_lowercase(),
2491 delete_second[1].to_ascii_lowercase()
2492 );
2493 }
2494
2495 #[test]
2496 fn test_resolve_typescript_names_collapses_non_ascii_run() {
2497 let tools = vec![ToolInfo {
2500 name: ToolName::new("café_menu_日本語").unwrap(),
2501 description: String::new(),
2502 input_schema: json!({}),
2503 output_schema: None,
2504 }];
2505
2506 let resolved = resolve_typescript_names(&tools);
2507
2508 assert_eq!(resolved[0], "caf_Menu_");
2509 }
2510
2511 #[test]
2512 fn test_resolve_typescript_names_disambiguates_collision_introduced_by_collapsing() {
2513 let tools = vec![
2518 ToolInfo {
2519 name: ToolName::new("a-b").unwrap(),
2520 description: String::new(),
2521 input_schema: json!({}),
2522 output_schema: None,
2523 },
2524 ToolInfo {
2525 name: ToolName::new("a--b").unwrap(),
2526 description: String::new(),
2527 input_schema: json!({}),
2528 output_schema: None,
2529 },
2530 ];
2531
2532 let resolved = resolve_typescript_names(&tools);
2533
2534 assert_eq!(resolved, vec!["a_b", "a_b_2"]);
2535 }
2536
2537 #[test]
2538 fn test_generate_sanitizes_reserved_word_tool_name() {
2539 let generator = ProgressiveGenerator::new().unwrap();
2540 let mut server_info = create_test_server_info();
2541 server_info.tools = vec![ToolInfo {
2542 name: ToolName::new("delete").unwrap(),
2543 description: "Delete something".to_string(),
2544 input_schema: json!({}),
2545 output_schema: None,
2546 }];
2547
2548 let code = generator.generate(&server_info).unwrap();
2549 let tool_file = code.files.iter().find(|f| f.path == "delete_2.ts").unwrap();
2550
2551 assert!(!tool_file.content.contains("export async function delete("));
2552 assert!(
2553 tool_file
2554 .content
2555 .contains("export async function delete_2(")
2556 );
2557 }
2558
2559 #[test]
2560 fn test_generate_disambiguates_colliding_tool_names() {
2561 let generator = ProgressiveGenerator::new().unwrap();
2562 let mut server_info = create_test_server_info();
2563 server_info.tools = vec![
2564 ToolInfo {
2565 name: ToolName::new("foo-bar").unwrap(),
2566 description: "First".to_string(),
2567 input_schema: json!({}),
2568 output_schema: None,
2569 },
2570 ToolInfo {
2571 name: ToolName::new("foo.bar").unwrap(),
2572 description: "Second".to_string(),
2573 input_schema: json!({}),
2574 output_schema: None,
2575 },
2576 ];
2577
2578 let code = generator.generate(&server_info).unwrap();
2579
2580 let tool_files: Vec<&str> = code
2582 .files
2583 .iter()
2584 .filter(|f| f.path == "foo_bar.ts" || f.path == "foo_bar_2.ts")
2585 .map(|f| f.path.as_str())
2586 .collect();
2587 assert_eq!(
2588 tool_files.len(),
2589 2,
2590 "colliding names must not overwrite each other's file: {tool_files:?}"
2591 );
2592
2593 let index = code.files.iter().find(|f| f.path == "index.ts").unwrap();
2594 assert_eq!(
2595 index.content.matches("export { foo_bar,").count(),
2596 1,
2597 "index.ts must export the first tool's identifier exactly once"
2598 );
2599 assert_eq!(
2600 index.content.matches("export { foo_bar_2,").count(),
2601 1,
2602 "index.ts must export the disambiguated second identifier exactly once"
2603 );
2604 }
2605
2606 #[test]
2607 fn test_generate_disambiguates_identical_raw_tool_names() {
2608 let generator = ProgressiveGenerator::new().unwrap();
2612 let mut server_info = create_test_server_info();
2613 server_info.tools = vec![
2614 ToolInfo {
2615 name: ToolName::new("dup").unwrap(),
2616 description: "First".to_string(),
2617 input_schema: json!({}),
2618 output_schema: None,
2619 },
2620 ToolInfo {
2621 name: ToolName::new("dup").unwrap(),
2622 description: "Second".to_string(),
2623 input_schema: json!({}),
2624 output_schema: None,
2625 },
2626 ];
2627
2628 let code = generator.generate(&server_info).unwrap();
2629
2630 let dup_files: Vec<&str> = code
2631 .files
2632 .iter()
2633 .filter(|f| f.path == "dup.ts" || f.path == "dup_2.ts")
2634 .map(|f| f.path.as_str())
2635 .collect();
2636 assert_eq!(
2637 dup_files.len(),
2638 2,
2639 "identical raw tool names must not overwrite each other's file: {dup_files:?}"
2640 );
2641
2642 let index = code.files.iter().find(|f| f.path == "index.ts").unwrap();
2643 assert_eq!(
2644 index.content.matches("export { dup,").count(),
2645 1,
2646 "index.ts must export the first tool's identifier exactly once"
2647 );
2648 assert_eq!(
2649 index.content.matches("export { dup_2,").count(),
2650 1,
2651 "index.ts must export the disambiguated second identifier exactly once"
2652 );
2653 }
2654
2655 #[test]
2656 fn test_sanitize_schema_jsdoc_drops_non_string_descriptions() {
2657 let sanitized = sanitize_schema_jsdoc_descriptions(json!({
2658 "type": "object",
2659 "description": {"text": "Schema */ injected\nnext"},
2660 "properties": {
2661 "title": {
2662 "type": "string",
2663 "description": ["Title */ injected\nnext"]
2664 }
2665 }
2666 }));
2667
2668 assert!(sanitized["description"].is_null());
2669 assert!(sanitized["properties"]["title"]["description"].is_null());
2670 }
2671
2672 #[test]
2673 fn test_sanitize_schema_jsdoc_recurses_into_array_items() {
2674 let sanitized = sanitize_schema_jsdoc_descriptions(json!({
2675 "type": "object",
2676 "properties": {
2677 "tags": {
2678 "type": "array",
2679 "items": [
2680 {
2681 "type": "string",
2682 "description": "Tag */ injected\nnext"
2683 }
2684 ]
2685 }
2686 }
2687 }));
2688
2689 let description = sanitized["properties"]["tags"]["items"][0]["description"]
2690 .as_str()
2691 .unwrap();
2692
2693 assert_eq!(description, "Tag *\\/ injected next");
2694 }
2695
2696 fn nested_array_schema_with_descriptions(depth: usize, description: &str) -> serde_json::Value {
2708 let mut schema = json!({"type": "string", "description": description});
2709 for _ in 0..depth {
2710 let mut map = serde_json::Map::new();
2711 map.insert(
2712 "type".to_string(),
2713 serde_json::Value::String("array".to_string()),
2714 );
2715 map.insert("items".to_string(), schema);
2716 map.insert(
2717 "description".to_string(),
2718 serde_json::Value::String(description.to_string()),
2719 );
2720 schema = serde_json::Value::Object(map);
2721 }
2722 schema
2723 }
2724
2725 fn nth_level(value: &serde_json::Value, depth: usize) -> serde_json::Value {
2727 let mut v = value.clone();
2728 for _ in 0..depth {
2729 v = v["items"].clone();
2730 }
2731 v
2732 }
2733
2734 #[test]
2735 fn test_sanitize_schema_jsdoc_bounds_deeply_nested_schema() {
2736 const MALICIOUS: &str = "desc */ injected\nnext";
2744 let depth = MAX_SCHEMA_RECURSION_DEPTH + 10;
2745 let schema = nested_array_schema_with_descriptions(depth, MALICIOUS);
2746
2747 let sanitized = sanitize_schema_jsdoc_descriptions(schema);
2748
2749 let just_below_cap = nth_level(&sanitized, MAX_SCHEMA_RECURSION_DEPTH - 1);
2750 let below_description = just_below_cap["description"]
2751 .as_str()
2752 .expect("description below the cap must still be a string");
2753 assert!(
2754 !below_description.contains("*/"),
2755 "description one level below the cap must be sanitized: {below_description}"
2756 );
2757
2758 let at_cap = nth_level(&sanitized, MAX_SCHEMA_RECURSION_DEPTH);
2759 let at_cap_description = at_cap["description"]
2760 .as_str()
2761 .expect("description at the cap must still be a string");
2762 assert_eq!(
2763 at_cap_description, MALICIOUS,
2764 "description at the cap must be left untouched — the function must stop \
2765 recursing before reaching it"
2766 );
2767 }
2768
2769 #[test]
2770 fn test_sanitize_schema_jsdoc_survives_pathologically_deep_input() {
2771 std::thread::Builder::new()
2781 .stack_size(64 * 1024 * 1024)
2782 .spawn(|| {
2783 let schema = nested_array_schema_with_descriptions(5_000, "leaf");
2784 let sanitized = sanitize_schema_jsdoc_descriptions(schema);
2785 assert!(sanitized.is_object());
2786 })
2787 .expect("spawn test thread")
2788 .join()
2789 .expect("test thread panicked");
2790 }
2791
2792 #[test]
2793 fn test_sanitize_jsdoc_truncation_boundary_injection() {
2794 let max_len = 256;
2795 let payload = format!("{}*/{}", "a".repeat(max_len - 1), "trailer");
2799
2800 let sanitized = sanitize_jsdoc(&payload, max_len);
2801
2802 assert!(
2803 !sanitized.contains("*/"),
2804 "truncation must not re-open the JSDoc comment: {sanitized}"
2805 );
2806 assert_eq!(sanitized.chars().count(), max_len);
2807 }
2808
2809 #[test]
2825 fn test_generate_runtime_bridge_declares_forbidden_env_var_list() {
2826 let generator = ProgressiveGenerator::new().unwrap();
2827 let server_info = create_test_server_info();
2828
2829 let code = generator.generate(&server_info).unwrap();
2830 let bridge = code
2831 .files
2832 .iter()
2833 .find(|f| f.path == "_runtime/mcp-bridge.ts")
2834 .unwrap();
2835
2836 for forbidden_env in mcp_execution_core::forbidden_env_names() {
2837 assert!(
2838 bridge.content.contains(&format!("'{forbidden_env}'")),
2839 "runtime bridge must list forbidden env var {forbidden_env}: {}",
2840 bridge.content
2841 );
2842 }
2843
2844 for forbidden_char in mcp_execution_core::forbidden_chars() {
2845 let escaped = sanitize_ts_string_literal(&forbidden_char.to_string());
2846 assert!(
2847 bridge.content.contains(&format!("'{escaped}'")),
2848 "runtime bridge must list forbidden char {forbidden_char:?}: {}",
2849 bridge.content
2850 );
2851 }
2852
2853 assert!(
2854 bridge
2855 .content
2856 .contains(mcp_execution_core::forbidden_env_prefix()),
2857 "runtime bridge must reference the forbidden env prefix: {}",
2858 bridge.content
2859 );
2860 }
2861
2862 #[test]
2863 fn test_generate_preserves_benign_punctuation() {
2864 let generator = ProgressiveGenerator::new().unwrap();
2868 let mut server_info = create_test_server_info();
2869 server_info.tools[0].description =
2870 "Compares values: a < b && b > c, or use \"quotes\" & don't forget 'em".to_string();
2871
2872 let code = generator.generate(&server_info).unwrap();
2873 let tool = code
2874 .files
2875 .iter()
2876 .find(|f| f.path == "createIssue.ts")
2877 .unwrap();
2878
2879 assert!(
2880 tool.content
2881 .contains("a < b && b > c, or use \"quotes\" & don't forget 'em"),
2882 "benign punctuation must survive verbatim, not be HTML-escaped: {}",
2883 tool.content
2884 );
2885 for entity in ["<", ">", "&", """, "'", "'"] {
2886 assert!(
2887 !tool.content.contains(entity),
2888 "output must not contain HTML entity {entity}: {}",
2889 tool.content
2890 );
2891 }
2892 }
2893
2894 #[test]
2895 fn test_generate_sanitizes_jsdoc_injection() {
2896 let generator = ProgressiveGenerator::new().unwrap();
2897 let mut server_info = create_test_server_info();
2898 server_info.name = "Evil */ injection".to_string();
2899 server_info.version = "1.0\n<script>".to_string();
2900
2901 let code = generator.generate(&server_info).unwrap();
2902 let index = code.files.iter().find(|f| f.path == "index.ts").unwrap();
2903
2904 assert!(
2906 !index.content.contains("Evil */ injection"),
2907 "Server name should be sanitized in JSDoc"
2908 );
2909 assert!(
2910 !index.content.contains("1.0\n<script>"),
2911 "Server version should have newlines stripped"
2912 );
2913 }
2914
2915 #[test]
2916 fn test_generate_sanitizes_schema_and_category_jsdoc_injection() {
2917 let generator = ProgressiveGenerator::new().unwrap();
2918 let mut server_info = create_test_server_info();
2919 server_info.tools[0].input_schema = json!({
2920 "type": "object",
2921 "description": "Schema */ injected\nnext",
2922 "properties": {
2923 "title": {
2924 "type": "string",
2925 "description": "Title */ injected\nnext"
2926 }
2927 },
2928 "required": ["title"]
2929 });
2930
2931 let mut categorizations = HashMap::new();
2932 categorizations.insert(
2933 "create_issue".to_string(),
2934 ToolCategorization {
2935 category: "issues */ injected\nnext".to_string(),
2936 keywords: vec!["create,*/ injected\nnext".to_string()],
2937 short_description: "Create */ injected\nnext".to_string(),
2938 },
2939 );
2940
2941 let code = generator
2942 .generate_with_categories(&server_info, &categorizations)
2943 .unwrap();
2944 let tool = code
2945 .files
2946 .iter()
2947 .find(|f| f.path == "createIssue.ts")
2948 .unwrap();
2949
2950 for raw in [
2951 "Schema */ injected",
2952 "Title */ injected",
2953 "issues */ injected",
2954 "create,*/ injected",
2955 "Create */ injected",
2956 ] {
2957 assert!(
2958 !tool.content.contains(raw),
2959 "generated JSDoc should not contain raw injection text: {raw}"
2960 );
2961 }
2962
2963 assert!(tool.content.contains("Schema *\\/ injected next"));
2964 assert!(tool.content.contains("Title *\\/ injected next"));
2965 assert!(tool.content.contains("issues *\\/ injected next"));
2966 assert!(tool.content.contains("create,*\\/ injected next"));
2967 assert!(tool.content.contains("Create *\\/ injected next"));
2968 }
2969
2970 fn server_info_with_tool_count(count: usize) -> ServerInfo {
2973 ServerInfo {
2974 id: ServerId::new("bulk-server").unwrap(),
2975 name: "Bulk Server".to_string(),
2976 version: "1.0.0".to_string(),
2977 tools: (0..count)
2978 .map(|i| ToolInfo {
2979 name: ToolName::new(format!("tool{i}")).unwrap(),
2980 description: String::new(),
2981 input_schema: json!({}),
2982 output_schema: None,
2983 })
2984 .collect(),
2985 capabilities: ServerCapabilities {
2986 supports_tools: true,
2987 supports_resources: false,
2988 supports_prompts: false,
2989 },
2990 }
2991 }
2992
2993 #[test]
2994 fn test_generate_rejects_tool_count_that_would_exceed_max_generated_files() {
2995 let server_info = server_info_with_tool_count(MAX_GENERATED_FILES - FIXED_FILE_COUNT + 1);
2996 let generator = ProgressiveGenerator::new().unwrap();
2997
2998 let result = generator.generate(&server_info);
2999
3000 assert!(result.is_err());
3001 assert!(result.unwrap_err().is_resource_limit_exceeded());
3002 }
3003
3004 #[test]
3005 fn test_generate_accepts_tool_count_at_exact_max_generated_files() {
3006 let server_info = server_info_with_tool_count(MAX_GENERATED_FILES - FIXED_FILE_COUNT);
3007 let generator = ProgressiveGenerator::new().unwrap();
3008
3009 let code = generator.generate(&server_info).unwrap();
3010
3011 assert_eq!(code.file_count(), MAX_GENERATED_FILES);
3012 }
3013
3014 #[test]
3015 fn test_generate_with_categories_rejects_tool_count_that_would_exceed_max_generated_files() {
3016 let server_info = server_info_with_tool_count(MAX_GENERATED_FILES - FIXED_FILE_COUNT + 1);
3017 let generator = ProgressiveGenerator::new().unwrap();
3018
3019 let result = generator.generate_with_categories(&server_info, &HashMap::new());
3020
3021 assert!(result.is_err());
3022 assert!(result.unwrap_err().is_resource_limit_exceeded());
3023 }
3024
3025 #[test]
3026 fn test_add_tracked_rejects_oversized_total_bytes() {
3027 let mut code = GeneratedCode::new();
3028 let mut total_bytes = 0usize;
3029
3030 let result = add_tracked(
3031 &mut code,
3032 &mut total_bytes,
3033 GeneratedFile {
3034 path: "big.ts".to_string(),
3035 content: "a".repeat(MAX_GENERATED_BYTES + 1),
3036 },
3037 );
3038
3039 assert!(result.is_err());
3040 assert!(result.unwrap_err().is_resource_limit_exceeded());
3041 }
3042
3043 #[test]
3044 fn test_add_tracked_accepts_total_bytes_at_exact_max() {
3045 let mut code = GeneratedCode::new();
3046 let mut total_bytes = 0usize;
3047
3048 let result = add_tracked(
3049 &mut code,
3050 &mut total_bytes,
3051 GeneratedFile {
3052 path: "big.ts".to_string(),
3053 content: "a".repeat(MAX_GENERATED_BYTES),
3054 },
3055 );
3056
3057 assert!(result.is_ok());
3058 }
3059
3060 #[test]
3068 fn test_add_tracked_rejects_immediately_once_running_total_exceeds_max() {
3069 let mut code = GeneratedCode::new();
3070 let mut total_bytes = MAX_GENERATED_BYTES - 1;
3071
3072 let result = add_tracked(
3073 &mut code,
3074 &mut total_bytes,
3075 GeneratedFile {
3076 path: "second.ts".to_string(),
3077 content: "ab".to_string(),
3078 },
3079 );
3080
3081 assert!(result.is_err());
3082 assert!(result.unwrap_err().is_resource_limit_exceeded());
3083 assert_eq!(code.file_count(), 0);
3086 }
3087}