use crate::clock::{Clock, SystemClock};
use crate::state::StateManager;
use crate::types::{
CategorizedTool, GeneratedServerInfo, IntrospectServerParams, IntrospectServerResult,
IntrospectedToolSummary, ListGeneratedServersParams, ListGeneratedServersResult,
PendingGeneration, SaveCategorizedToolsParams, SaveCategorizedToolsResult,
};
use mcp_execution_codegen::progressive::ProgressiveGenerator;
use mcp_execution_core::{ServerConfig, ServerId};
use mcp_execution_files::FilesBuilder;
use mcp_execution_introspector::Introspector;
use mcp_execution_skill::{
GenerateSkillParams, SaveSkillParams, SaveSkillResult, ScanError, build_skill_context,
extract_skill_metadata, scan_tools_directory, validate_server_id,
};
use rmcp::handler::server::ServerHandler;
use rmcp::handler::server::tool::ToolRouter;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::model::{
CallToolResult, ContentBlock, Implementation, ProtocolVersion, ServerCapabilities, ServerInfo,
};
use rmcp::{ErrorData as McpError, tool, tool_handler, tool_router};
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::Mutex;
const MAX_SKILL_CONTENT_SIZE: usize = 100 * 1024;
#[derive(Debug, Clone)]
pub struct GeneratorService {
state: Arc<StateManager>,
introspectors: Arc<Mutex<HashMap<ServerId, Arc<Mutex<Introspector>>>>>,
exports: Arc<Mutex<HashMap<PathBuf, Arc<Mutex<()>>>>>,
clock: Arc<dyn Clock>,
#[allow(dead_code)]
tool_router: ToolRouter<Self>,
}
impl GeneratorService {
#[must_use]
pub fn new() -> Self {
Self::with_clock(Arc::new(SystemClock))
}
fn with_clock(clock: Arc<dyn Clock>) -> Self {
Self {
state: Arc::new(StateManager::with_clock(Arc::clone(&clock))),
introspectors: Arc::new(Mutex::new(HashMap::new())),
exports: Arc::new(Mutex::new(HashMap::new())),
clock,
tool_router: Self::tool_router(),
}
}
async fn introspector_for(&self, server_id: &ServerId) -> Arc<Mutex<Introspector>> {
let mut introspectors = self.introspectors.lock().await;
introspectors
.entry(server_id.clone())
.or_insert_with(|| Arc::new(Mutex::new(Introspector::new())))
.clone()
}
async fn evict_introspector(&self, server_id: &ServerId, handle: &Arc<Mutex<Introspector>>) {
let mut introspectors = self.introspectors.lock().await;
if let std::collections::hash_map::Entry::Occupied(entry) =
introspectors.entry(server_id.clone())
&& Arc::ptr_eq(entry.get(), handle)
{
entry.remove();
}
}
async fn export_lock_for(&self, output_dir: &Path) -> Arc<Mutex<()>> {
let mut exports = self.exports.lock().await;
exports
.entry(output_dir.to_path_buf())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
}
async fn evict_export_lock(&self, output_dir: &Path, handle: &Arc<Mutex<()>>) {
let mut exports = self.exports.lock().await;
if let std::collections::hash_map::Entry::Occupied(entry) =
exports.entry(output_dir.to_path_buf())
&& Arc::ptr_eq(entry.get(), handle)
{
entry.remove();
}
}
}
impl Default for GeneratorService {
fn default() -> Self {
Self::new()
}
}
#[tool_router]
impl GeneratorService {
#[tool(
description = "Connect to an MCP server, discover its tools, and return metadata for categorization. Returns a session ID for use with save_categorized_tools."
)]
async fn introspect_server(
&self,
Parameters(params): Parameters<IntrospectServerParams>,
) -> Result<CallToolResult, McpError> {
validate_server_id(¶ms.server_id).map_err(|e| McpError::invalid_params(e, None))?;
let server_id_str = params.server_id;
let server_id = ServerId::new(&server_id_str);
let output_dir = params.output_dir.unwrap_or_else(|| {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".claude")
.join("servers")
.join(&server_id_str)
});
let mut config_builder = ServerConfig::builder().command(params.command);
for arg in params.args {
config_builder = config_builder.arg(arg);
}
for (key, value) in params.env {
config_builder = config_builder.env(key, value);
}
if let Some(secs) = params.connect_timeout_secs {
config_builder = config_builder.connect_timeout(std::time::Duration::from_secs(secs));
}
if let Some(secs) = params.discover_timeout_secs {
config_builder = config_builder.discover_timeout(std::time::Duration::from_secs(secs));
}
let config = config_builder.build();
let introspector_handle = self.introspector_for(&server_id).await;
let discover_result = {
let mut introspector = introspector_handle.lock().await;
introspector
.discover_server(server_id.clone(), &config)
.await
};
self.evict_introspector(&server_id, &introspector_handle)
.await;
let server_info = discover_result.map_err(|e| {
if e.is_validation_error() {
McpError::invalid_params(e.to_string(), None)
} else {
McpError::internal_error(format!("Failed to introspect server: {e}"), None)
}
})?;
let tools: Vec<IntrospectedToolSummary> = server_info
.tools
.iter()
.map(|tool| {
let parameters = extract_parameter_names(&tool.input_schema);
IntrospectedToolSummary {
name: tool.name.as_str().to_string(),
description: tool.description.clone(),
parameters,
}
})
.collect();
let pending = PendingGeneration::new(
server_id,
server_info.clone(),
config,
output_dir.clone(),
self.clock.as_ref(),
);
let session_id = self.state.store(pending.clone()).await;
let result = IntrospectServerResult {
server_id: server_id_str,
server_name: server_info.name,
tools_found: tools.len(),
tools,
session_id,
expires_at: pending.expires_at,
};
Ok(CallToolResult::success(vec![ContentBlock::text(
serde_json::to_string_pretty(&result).map_err(|e| {
McpError::internal_error(format!("Failed to serialize result: {e}"), None)
})?,
)]))
}
#[tool(
description = "Generate progressive loading TypeScript files using Claude's categorization. Requires session_id from a previous introspect_server call."
)]
async fn save_categorized_tools(
&self,
Parameters(params): Parameters<SaveCategorizedToolsParams>,
) -> Result<CallToolResult, McpError> {
let pending = self.state.take(params.session_id).await.ok_or_else(|| {
McpError::invalid_params(
"Session not found or expired. Please run introspect_server again.",
None,
)
})?;
let introspected_names: HashSet<_> = pending
.server_info
.tools
.iter()
.map(|t| t.name.as_str())
.collect();
for cat_tool in ¶ms.categorized_tools {
if !introspected_names.contains(cat_tool.name.as_str()) {
return Err(McpError::invalid_params(
format!("Tool '{}' not found in introspected tools", cat_tool.name),
None,
));
}
}
let tool_count = params.categorized_tools.len();
let mut categorization: HashMap<String, &CategorizedTool> =
HashMap::with_capacity(tool_count);
let mut categories: HashMap<String, usize> = HashMap::with_capacity(tool_count);
for tool in ¶ms.categorized_tools {
categorization.insert(tool.name.clone(), tool);
*categories.entry(tool.category.clone()).or_default() += 1;
}
let generator = ProgressiveGenerator::new().map_err(|e| {
McpError::internal_error(format!("Failed to create generator: {e}"), None)
})?;
let code = generate_with_categorization(&generator, &pending.server_info, &categorization)
.map_err(|e| McpError::internal_error(format!("Failed to generate code: {e}"), None))?;
let vfs = FilesBuilder::from_generated_code(code, "/")
.build()
.map_err(|e| McpError::internal_error(format!("Failed to build VFS: {e}"), None))?;
let files_generated = vfs.file_count();
if let Some(parent) = pending.output_dir.parent() {
tokio::fs::create_dir_all(parent).await.map_err(|e| {
McpError::internal_error(format!("Failed to create output directory: {e}"), None)
})?;
}
let export_lock = self.export_lock_for(&pending.output_dir).await;
let export_guard = export_lock.lock().await;
let output_dir = pending.output_dir.clone();
let export_result =
tokio::task::spawn_blocking(move || vfs.export_to_filesystem(&output_dir)).await;
drop(export_guard);
self.evict_export_lock(&pending.output_dir, &export_lock)
.await;
export_result
.map_err(|e| McpError::internal_error(format!("Task join error: {e}"), None))?
.map_err(|e| McpError::internal_error(format!("Failed to export files: {e}"), None))?;
let result = SaveCategorizedToolsResult {
success: true,
files_generated,
output_dir: pending.output_dir.display().to_string(),
categories,
errors: vec![],
};
Ok(CallToolResult::success(vec![ContentBlock::text(
serde_json::to_string_pretty(&result).map_err(|e| {
McpError::internal_error(format!("Failed to serialize result: {e}"), None)
})?,
)]))
}
#[tool(
description = "List all MCP servers that have generated progressive loading files in ~/.claude/servers/"
)]
async fn list_generated_servers(
&self,
Parameters(params): Parameters<ListGeneratedServersParams>,
) -> Result<CallToolResult, McpError> {
let base_dir = params.base_dir.map_or_else(
|| {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".claude")
.join("servers")
},
PathBuf::from,
);
let servers = tokio::task::spawn_blocking(move || {
let mut servers = Vec::new();
if base_dir.exists()
&& base_dir.is_dir()
&& let Ok(entries) = std::fs::read_dir(&base_dir)
{
for entry in entries.flatten() {
if entry.path().is_dir() {
let id = entry.file_name().to_string_lossy().to_string();
let tool_count = std::fs::read_dir(entry.path()).map_or(0, |e| {
e.flatten()
.filter(|f| {
let name = f.file_name();
let name = name.to_string_lossy();
name.ends_with(".ts") && !name.starts_with('_')
})
.count()
});
let generated_at = entry
.metadata()
.and_then(|m| m.modified())
.ok()
.map(chrono::DateTime::<chrono::Utc>::from);
servers.push(GeneratedServerInfo {
id,
tool_count,
generated_at,
output_dir: entry.path().display().to_string(),
});
}
}
}
servers.sort_by(|a, b| a.id.cmp(&b.id));
servers
})
.await
.map_err(|e| McpError::internal_error(format!("Task join error: {e}"), None))?;
let result = ListGeneratedServersResult {
total_servers: servers.len(),
servers,
};
Ok(CallToolResult::success(vec![ContentBlock::text(
serde_json::to_string_pretty(&result).map_err(|e| {
McpError::internal_error(format!("Failed to serialize result: {e}"), None)
})?,
)]))
}
#[tool(
description = "Analyze generated TypeScript files and return context for Claude to create a SKILL.md file. Returns tool metadata, categories, and a generation prompt."
)]
async fn generate_skill(
&self,
Parameters(params): Parameters<GenerateSkillParams>,
) -> Result<CallToolResult, McpError> {
validate_server_id(¶ms.server_id).map_err(|e| McpError::invalid_params(e, None))?;
let servers_dir = params.servers_dir.unwrap_or_else(|| {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".claude")
.join("servers")
});
let server_dir = servers_dir.join(¶ms.server_id);
if !server_dir.exists() {
return Err(McpError::invalid_params(
format!(
"Server directory not found: {}. Run generate first.",
server_dir.display()
),
None,
));
}
let scan_result = scan_tools_directory(&server_dir)
.await
.map_err(|e| match e {
ScanError::MissingMetadata { .. }
| ScanError::UnsupportedSchema { .. }
| ScanError::StaleMetadata { .. } => {
McpError::invalid_params(format!("Failed to scan tools directory: {e}"), None)
}
ScanError::Io(_)
| ScanError::DirectoryNotFound { .. }
| ScanError::MetadataParse { .. }
| ScanError::TooManyFiles { .. }
| ScanError::FileTooLarge { .. } => {
McpError::internal_error(format!("Failed to scan tools directory: {e}"), None)
}
})?;
if scan_result.tools.is_empty() {
return Err(McpError::invalid_params(
format!(
"No tool files found in {}. Run generate first.",
server_dir.display()
),
None,
));
}
let mut result = build_skill_context(
¶ms.server_id,
&scan_result.tools,
params.use_case_hints.as_deref(),
);
result.warnings = scan_result.warnings;
if let Some(name) = params.skill_name {
result.skill_name = name;
}
Ok(CallToolResult::success(vec![ContentBlock::text(
serde_json::to_string_pretty(&result).map_err(|e| {
McpError::internal_error(format!("Failed to serialize result: {e}"), None)
})?,
)]))
}
#[tool(
description = "Save generated SKILL.md content to ~/.claude/skills/{server_id}/. Use after Claude generates skill content from generate_skill context."
)]
async fn save_skill(
&self,
Parameters(params): Parameters<SaveSkillParams>,
) -> Result<CallToolResult, McpError> {
validate_server_id(¶ms.server_id).map_err(|e| McpError::invalid_params(e, None))?;
if params.content.len() > MAX_SKILL_CONTENT_SIZE {
return Err(McpError::invalid_params(
format!(
"content too large: {} bytes exceeds {} limit",
params.content.len(),
MAX_SKILL_CONTENT_SIZE
),
None,
));
}
if !params.content.starts_with("---") {
return Err(McpError::invalid_params(
"Content must start with YAML frontmatter (---)",
None,
));
}
let metadata = extract_skill_metadata(¶ms.content)
.map_err(|e| McpError::invalid_params(format!("Invalid SKILL.md format: {e}"), None))?;
let output_path = params.output_path.unwrap_or_else(|| {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".claude")
.join("skills")
.join(¶ms.server_id)
.join("SKILL.md")
});
let overwritten = output_path.exists();
if overwritten && !params.overwrite {
return Err(McpError::invalid_params(
format!(
"Skill file already exists: {}. Use overwrite=true to replace.",
output_path.display()
),
None,
));
}
if let Some(parent) = output_path.parent() {
tokio::fs::create_dir_all(parent).await.map_err(|e| {
McpError::internal_error(format!("Failed to create directory: {e}"), None)
})?;
}
tokio::fs::write(&output_path, ¶ms.content)
.await
.map_err(|e| McpError::internal_error(format!("Failed to write file: {e}"), None))?;
let result = SaveSkillResult {
success: true,
output_path: output_path.display().to_string(),
overwritten,
metadata,
};
Ok(CallToolResult::success(vec![ContentBlock::text(
serde_json::to_string_pretty(&result).map_err(|e| {
McpError::internal_error(format!("Failed to serialize result: {e}"), None)
})?,
)]))
}
}
#[tool_handler]
impl ServerHandler for GeneratorService {
fn get_info(&self) -> ServerInfo {
let mut info = ServerInfo::default();
info.protocol_version = ProtocolVersion::V_2025_06_18;
info.capabilities = ServerCapabilities::builder().enable_tools().build();
info.server_info = Implementation::new(env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
info.instructions = Some(
"Generate progressive loading TypeScript files for MCP servers. \
Use introspect_server to discover tools, then save_categorized_tools \
with your categorization."
.to_string(),
);
info
}
}
fn extract_parameter_names(schema: &serde_json::Value) -> Vec<String> {
schema
.get("properties")
.and_then(|p| p.as_object())
.map(|props| props.keys().cloned().collect())
.unwrap_or_default()
}
fn generate_with_categorization(
generator: &ProgressiveGenerator,
server_info: &mcp_execution_introspector::ServerInfo,
categorization: &HashMap<String, &CategorizedTool>,
) -> mcp_execution_core::Result<mcp_execution_codegen::GeneratedCode> {
use mcp_execution_codegen::progressive::ToolCategorization;
let categorizations: HashMap<String, ToolCategorization> = categorization
.iter()
.map(|(tool_name, cat_tool)| {
(
tool_name.clone(),
ToolCategorization {
category: cat_tool.category.clone(),
keywords: cat_tool.keywords.clone(),
short_description: cat_tool.short_description.clone(),
},
)
})
.collect();
generator.generate_with_categories(server_info, &categorizations)
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Utc;
use mcp_execution_core::ToolName;
use mcp_execution_introspector::{ServerCapabilities, ToolInfo};
use rmcp::model::ErrorCode;
use uuid::Uuid;
#[test]
fn test_extract_parameter_names() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"name": { "type": "string" },
"age": { "type": "number" }
}
});
let params = extract_parameter_names(&schema);
assert_eq!(params.len(), 2);
assert!(params.contains(&"name".to_string()));
assert!(params.contains(&"age".to_string()));
}
#[test]
fn test_extract_parameter_names_empty() {
let schema = serde_json::json!({
"type": "object"
});
let params = extract_parameter_names(&schema);
assert_eq!(params.len(), 0);
}
#[test]
fn test_extract_parameter_names_no_properties() {
let schema = serde_json::json!({
"type": "string"
});
let params = extract_parameter_names(&schema);
assert_eq!(params.len(), 0);
}
#[test]
fn test_extract_parameter_names_nested_object() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"user": {
"type": "object",
"properties": {
"name": { "type": "string" }
}
},
"age": { "type": "number" }
}
});
let params = extract_parameter_names(&schema);
assert_eq!(params.len(), 2);
assert!(params.contains(&"user".to_string()));
assert!(params.contains(&"age".to_string()));
}
#[test]
fn test_generate_with_categorization() {
let generator = ProgressiveGenerator::new().unwrap();
let server_info = mcp_execution_introspector::ServerInfo {
id: ServerId::new("test"),
name: "Test Server".to_string(),
version: "1.0.0".to_string(),
capabilities: ServerCapabilities {
supports_tools: true,
supports_resources: false,
supports_prompts: false,
},
tools: vec![ToolInfo {
name: ToolName::new("test_tool"),
description: "Test tool description".to_string(),
input_schema: serde_json::json!({
"type": "object",
"properties": {
"param1": { "type": "string" }
}
}),
output_schema: None,
}],
};
let categorized_tool = CategorizedTool {
name: "test_tool".to_string(),
category: "testing".to_string(),
keywords: "test,tool".to_string(),
short_description: "Test tool for testing".to_string(),
};
let mut categorization = HashMap::new();
categorization.insert("test_tool".to_string(), &categorized_tool);
let result = generate_with_categorization(&generator, &server_info, &categorization);
assert!(result.is_ok());
let code = result.unwrap();
assert!(code.file_count() > 0, "Should generate at least one file");
}
#[test]
fn test_generate_with_categorization_multiple_tools() {
let generator = ProgressiveGenerator::new().unwrap();
let server_info = mcp_execution_introspector::ServerInfo {
id: ServerId::new("test"),
name: "Test Server".to_string(),
version: "1.0.0".to_string(),
capabilities: ServerCapabilities {
supports_tools: true,
supports_resources: false,
supports_prompts: false,
},
tools: vec![
ToolInfo {
name: ToolName::new("tool1"),
description: "First tool".to_string(),
input_schema: serde_json::json!({"type": "object"}),
output_schema: None,
},
ToolInfo {
name: ToolName::new("tool2"),
description: "Second tool".to_string(),
input_schema: serde_json::json!({"type": "object"}),
output_schema: None,
},
],
};
let tool1 = CategorizedTool {
name: "tool1".to_string(),
category: "category1".to_string(),
keywords: "test".to_string(),
short_description: "Tool 1".to_string(),
};
let tool2 = CategorizedTool {
name: "tool2".to_string(),
category: "category2".to_string(),
keywords: "test".to_string(),
short_description: "Tool 2".to_string(),
};
let mut categorization = HashMap::new();
categorization.insert("tool1".to_string(), &tool1);
categorization.insert("tool2".to_string(), &tool2);
let result = generate_with_categorization(&generator, &server_info, &categorization);
assert!(result.is_ok());
}
#[test]
fn test_generate_with_categorization_empty_tools() {
let generator = ProgressiveGenerator::new().unwrap();
let server_id = ServerId::new("test");
let server_info = mcp_execution_introspector::ServerInfo {
id: server_id,
name: "Empty Server".to_string(),
version: "1.0.0".to_string(),
capabilities: ServerCapabilities {
supports_tools: true,
supports_resources: false,
supports_prompts: false,
},
tools: vec![],
};
let categorization = HashMap::new();
let result = generate_with_categorization(&generator, &server_info, &categorization);
assert!(result.is_ok());
}
#[test]
fn test_generator_service_new() {
let service = GeneratorService::new();
assert!(service.introspectors.try_lock().is_ok());
assert!(service.exports.try_lock().is_ok());
}
#[test]
fn test_generator_service_default() {
let service = GeneratorService::default();
assert!(service.introspectors.try_lock().is_ok());
assert!(service.exports.try_lock().is_ok());
}
#[test]
fn test_get_info() {
let service = GeneratorService::new();
let info = service.get_info();
assert_eq!(info.protocol_version, ProtocolVersion::V_2025_06_18);
assert!(info.capabilities.tools.is_some());
assert!(info.instructions.is_some());
assert_eq!(info.server_info.name, env!("CARGO_PKG_NAME"));
assert_eq!(info.server_info.version, env!("CARGO_PKG_VERSION"));
}
#[tokio::test]
async fn test_introspect_server_invalid_server_id_uppercase() {
let service = GeneratorService::new();
let params = IntrospectServerParams {
server_id: "GitHub".to_string(), command: "echo".to_string(),
args: vec![],
env: HashMap::new(),
output_dir: None,
connect_timeout_secs: None,
discover_timeout_secs: None,
};
let result = service.introspect_server(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS); }
#[tokio::test]
async fn test_introspect_server_invalid_server_id_underscore() {
let service = GeneratorService::new();
let params = IntrospectServerParams {
server_id: "git_hub".to_string(), command: "echo".to_string(),
args: vec![],
env: HashMap::new(),
output_dir: None,
connect_timeout_secs: None,
discover_timeout_secs: None,
};
let result = service.introspect_server(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
}
#[tokio::test]
async fn test_introspect_server_invalid_server_id_special_chars() {
let service = GeneratorService::new();
let params = IntrospectServerParams {
server_id: "git@hub".to_string(), command: "echo".to_string(),
args: vec![],
env: HashMap::new(),
output_dir: None,
connect_timeout_secs: None,
discover_timeout_secs: None,
};
let result = service.introspect_server(Parameters(params)).await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_introspect_server_valid_server_id_with_hyphens() {
let service = GeneratorService::new();
let params = IntrospectServerParams {
server_id: "git-hub-server".to_string(), command: "echo".to_string(),
args: vec!["test".to_string()],
env: HashMap::new(),
output_dir: None,
connect_timeout_secs: None,
discover_timeout_secs: None,
};
let result = service.introspect_server(Parameters(params)).await;
if let Err(err) = result {
assert_ne!(
err.code,
ErrorCode::INVALID_PARAMS,
"Should not be invalid params error"
);
}
}
#[tokio::test]
async fn test_introspect_server_valid_server_id_digits() {
let service = GeneratorService::new();
let params = IntrospectServerParams {
server_id: "server123".to_string(), command: "echo".to_string(),
args: vec![],
env: HashMap::new(),
output_dir: None,
connect_timeout_secs: None,
discover_timeout_secs: None,
};
let result = service.introspect_server(Parameters(params)).await;
if let Err(err) = result {
assert_ne!(err.code, ErrorCode::INVALID_PARAMS);
}
}
#[tokio::test]
async fn test_introspect_server_zero_connect_timeout_is_invalid_params() {
let service = GeneratorService::new();
let params = IntrospectServerParams {
server_id: "zero-timeout-test".to_string(),
command: "echo".to_string(),
args: vec![],
env: HashMap::new(),
output_dir: None,
connect_timeout_secs: Some(0),
discover_timeout_secs: None,
};
let result = service.introspect_server(Parameters(params)).await;
let err = result.expect_err("zero connect_timeout must be rejected");
assert_eq!(
err.code,
ErrorCode::INVALID_PARAMS,
"zero timeout is a client input error, not an internal error"
);
}
#[tokio::test]
async fn test_introspect_server_evicts_map_entry_after_completion() {
let service = GeneratorService::new();
let params = IntrospectServerParams {
server_id: "evict-after-completion".to_string(),
command: "echo".to_string(), args: vec![],
env: HashMap::new(),
output_dir: None,
connect_timeout_secs: None,
discover_timeout_secs: None,
};
let result = service.introspect_server(Parameters(params)).await;
assert!(
result.is_err(),
"echo is not an MCP server, expected a connection failure"
);
assert!(
service.introspectors.lock().await.is_empty(),
"introspectors map should be empty after introspect_server completes, \
regardless of success or failure"
);
}
#[tokio::test]
async fn test_introspector_for_same_id_shares_one_lock() {
let service = GeneratorService::new();
let server_id = ServerId::new("same-id-lock-test");
let handle_a = service.introspector_for(&server_id).await;
let handle_b = service.introspector_for(&server_id).await;
assert!(
Arc::ptr_eq(&handle_a, &handle_b),
"the same server_id must reuse one introspector lock"
);
}
#[tokio::test]
async fn test_introspector_for_different_ids_get_independent_locks() {
let service = GeneratorService::new();
let handle_a = service
.introspector_for(&ServerId::new("diff-id-lock-a"))
.await;
let handle_b = service
.introspector_for(&ServerId::new("diff-id-lock-b"))
.await;
assert!(
!Arc::ptr_eq(&handle_a, &handle_b),
"different server_ids must get independent introspector locks"
);
}
#[tokio::test]
async fn test_same_id_lock_serializes_concurrent_holders() {
let service = GeneratorService::new();
let server_id = ServerId::new("same-id-timing-test");
let hold_time = std::time::Duration::from_millis(150);
let serialized_threshold = std::time::Duration::from_millis(250);
let handle_a = service.introspector_for(&server_id).await;
let handle_b = service.introspector_for(&server_id).await;
let started = std::time::Instant::now();
tokio::join!(
async {
let _guard = handle_a.lock().await;
tokio::time::sleep(hold_time).await;
},
async {
let _guard = handle_b.lock().await;
tokio::time::sleep(hold_time).await;
},
);
let elapsed = started.elapsed();
assert!(
elapsed >= serialized_threshold,
"holders of the same per-id lock should serialize \
(expected >= {serialized_threshold:?}, i.e. two back-to-back {hold_time:?} \
critical sections); took {elapsed:?}"
);
}
#[tokio::test]
async fn test_different_id_locks_do_not_serialize() {
let service = GeneratorService::new();
let hold_time = std::time::Duration::from_millis(150);
let serialized_threshold = std::time::Duration::from_millis(250);
let handle_a = service
.introspector_for(&ServerId::new("diff-id-timing-a"))
.await;
let handle_b = service
.introspector_for(&ServerId::new("diff-id-timing-b"))
.await;
let started = std::time::Instant::now();
tokio::join!(
async {
let _guard = handle_a.lock().await;
tokio::time::sleep(hold_time).await;
},
async {
let _guard = handle_b.lock().await;
tokio::time::sleep(hold_time).await;
},
);
let elapsed = started.elapsed();
assert!(
elapsed < serialized_threshold,
"holders of different per-id locks should not serialize \
(expected < {serialized_threshold:?}, i.e. close to a single {hold_time:?} hold); \
took {elapsed:?}"
);
}
#[tokio::test]
async fn test_stale_eviction_does_not_remove_unrelated_entry() {
let service = GeneratorService::new();
let server_id = ServerId::new("toctou-abc-test");
let handle_a = service.introspector_for(&server_id).await;
let handle_b = service.introspector_for(&server_id).await;
assert!(
Arc::ptr_eq(&handle_a, &handle_b),
"A and B must share one introspector handle for the same server_id"
);
service.evict_introspector(&server_id, &handle_a).await;
assert!(
service.introspectors.lock().await.is_empty(),
"map should be empty right after A's eviction"
);
let handle_c = service.introspector_for(&server_id).await;
assert!(
!Arc::ptr_eq(&handle_b, &handle_c),
"C must get a handle distinct from A/B's stale one"
);
service.evict_introspector(&server_id, &handle_b).await;
let introspectors = service.introspectors.lock().await;
let current = introspectors
.get(&server_id)
.expect("C's entry must survive B's stale eviction attempt");
assert!(
Arc::ptr_eq(current, &handle_c),
"the surviving entry must be C's handle, unaffected by B's stale eviction"
);
drop(introspectors);
service.evict_introspector(&server_id, &handle_c).await;
assert!(
service.introspectors.lock().await.is_empty(),
"map should be empty after C's own eviction"
);
}
#[tokio::test]
async fn test_export_lock_for_same_output_dir_shares_one_lock() {
let service = GeneratorService::new();
let output_dir = PathBuf::from("/tmp/same-output-dir-lock-test");
let handle_a = service.export_lock_for(&output_dir).await;
let handle_b = service.export_lock_for(&output_dir).await;
assert!(
Arc::ptr_eq(&handle_a, &handle_b),
"the same output_dir must reuse one export lock"
);
}
#[tokio::test]
async fn test_export_lock_for_different_output_dirs_get_independent_locks() {
let service = GeneratorService::new();
let handle_a = service
.export_lock_for(&PathBuf::from("/tmp/diff-output-dir-lock-a"))
.await;
let handle_b = service
.export_lock_for(&PathBuf::from("/tmp/diff-output-dir-lock-b"))
.await;
assert!(
!Arc::ptr_eq(&handle_a, &handle_b),
"different output_dirs must get independent export locks"
);
}
#[tokio::test]
async fn test_export_lock_stale_eviction_does_not_remove_unrelated_entry() {
let service = GeneratorService::new();
let output_dir = PathBuf::from("/tmp/toctou-export-lock-test");
let handle_a = service.export_lock_for(&output_dir).await;
let handle_b = service.export_lock_for(&output_dir).await;
assert!(Arc::ptr_eq(&handle_a, &handle_b));
service.evict_export_lock(&output_dir, &handle_a).await;
assert!(service.exports.lock().await.is_empty());
let handle_c = service.export_lock_for(&output_dir).await;
assert!(!Arc::ptr_eq(&handle_b, &handle_c));
service.evict_export_lock(&output_dir, &handle_b).await;
let exports = service.exports.lock().await;
let current = exports
.get(&output_dir)
.expect("C's entry must survive B's stale eviction attempt");
assert!(Arc::ptr_eq(current, &handle_c));
drop(exports);
}
#[tokio::test]
async fn test_save_categorized_tools_invalid_session() {
let service = GeneratorService::new();
let params = SaveCategorizedToolsParams {
session_id: Uuid::new_v4(), categorized_tools: vec![],
};
let result = service.save_categorized_tools(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS); assert!(err.message.contains("Session not found"));
}
#[tokio::test]
async fn test_save_categorized_tools_tool_mismatch() {
let service = GeneratorService::new();
let server_id = ServerId::new("test");
let server_info = mcp_execution_introspector::ServerInfo {
id: server_id.clone(),
name: "Test".to_string(),
version: "1.0.0".to_string(),
capabilities: ServerCapabilities {
supports_tools: true,
supports_resources: false,
supports_prompts: false,
},
tools: vec![ToolInfo {
name: ToolName::new("tool1"),
description: "Tool 1".to_string(),
input_schema: serde_json::json!({"type": "object"}),
output_schema: None,
}],
};
let pending = PendingGeneration::new(
server_id,
server_info,
ServerConfig::builder().command("echo".to_string()).build(),
PathBuf::from("/tmp/test"),
&SystemClock,
);
let session_id = service.state.store(pending).await;
let params = SaveCategorizedToolsParams {
session_id,
categorized_tools: vec![CategorizedTool {
name: "tool2".to_string(), category: "test".to_string(),
keywords: "test".to_string(),
short_description: "Test".to_string(),
}],
};
let result = service.save_categorized_tools(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
assert!(err.message.contains("not found in introspected tools"));
}
#[tokio::test]
async fn test_save_categorized_tools_expired_session() {
use crate::clock::TestClock;
use chrono::Duration;
let service = GeneratorService::new();
let server_id = ServerId::new("test");
let server_info = mcp_execution_introspector::ServerInfo {
id: server_id.clone(),
name: "Test".to_string(),
version: "1.0.0".to_string(),
capabilities: ServerCapabilities {
supports_tools: true,
supports_resources: false,
supports_prompts: false,
},
tools: vec![],
};
let past_clock = TestClock::new(Utc::now() - Duration::hours(1));
let pending = PendingGeneration::new(
server_id,
server_info,
ServerConfig::builder().command("echo".to_string()).build(),
PathBuf::from("/tmp/test"),
&past_clock,
);
let session_id = service.state.store(pending).await;
let params = SaveCategorizedToolsParams {
session_id,
categorized_tools: vec![],
};
let result = service.save_categorized_tools(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
}
#[tokio::test]
async fn test_shared_clock_drives_save_categorized_tools_expiry() {
use crate::clock::TestClock;
use chrono::Duration;
let start = Utc::now();
let clock = Arc::new(TestClock::new(start));
let service = GeneratorService::with_clock(Arc::clone(&clock) as Arc<dyn Clock>);
let server_id = ServerId::new("test");
let server_info = mcp_execution_introspector::ServerInfo {
id: server_id.clone(),
name: "Test".to_string(),
version: "1.0.0".to_string(),
capabilities: ServerCapabilities {
supports_tools: true,
supports_resources: false,
supports_prompts: false,
},
tools: vec![],
};
let pending = PendingGeneration::new(
server_id,
server_info,
ServerConfig::builder().command("echo".to_string()).build(),
PathBuf::from("/tmp/test"),
clock.as_ref(),
);
let session_id = service.state.store(pending).await;
clock.advance(
Duration::minutes(PendingGeneration::DEFAULT_TIMEOUT_MINUTES) + Duration::seconds(1),
);
let params = SaveCategorizedToolsParams {
session_id,
categorized_tools: vec![],
};
let result = service.save_categorized_tools(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
}
#[tokio::test]
async fn test_list_generated_servers_nonexistent_dir() {
let service = GeneratorService::new();
let params = ListGeneratedServersParams {
base_dir: Some("/nonexistent/path/that/does/not/exist".to_string()),
};
let result = service.list_generated_servers(Parameters(params)).await;
assert!(result.is_ok());
let content = result.unwrap();
let text_content = content.content[0].as_text().unwrap();
let parsed: ListGeneratedServersResult = serde_json::from_str(&text_content.text).unwrap();
assert_eq!(parsed.total_servers, 0);
assert_eq!(parsed.servers.len(), 0);
}
#[tokio::test]
async fn test_list_generated_servers_default_dir() {
let service = GeneratorService::new();
let params = ListGeneratedServersParams { base_dir: None };
let result = service.list_generated_servers(Parameters(params)).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_generate_skill_invalid_server_id_uppercase() {
let service = GeneratorService::new();
let params = GenerateSkillParams {
server_id: "GitHub".to_string(), skill_name: None,
use_case_hints: None,
servers_dir: None,
};
let result = service.generate_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
assert!(err.message.contains("lowercase"));
}
#[tokio::test]
async fn test_generate_skill_invalid_server_id_special_chars() {
let service = GeneratorService::new();
let params = GenerateSkillParams {
server_id: "git@hub".to_string(), skill_name: None,
use_case_hints: None,
servers_dir: None,
};
let result = service.generate_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
}
#[tokio::test]
async fn test_generate_skill_server_directory_not_found() {
let service = GeneratorService::new();
let params = GenerateSkillParams {
server_id: "nonexistent-server".to_string(),
skill_name: None,
use_case_hints: None,
servers_dir: Some(PathBuf::from("/nonexistent/path")),
};
let result = service.generate_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
assert!(err.message.contains("not found"));
}
#[tokio::test]
async fn test_generate_skill_missing_metadata_sidecar() {
use tempfile::TempDir;
let service = GeneratorService::new();
let temp_dir = TempDir::new().unwrap();
let base_dir = temp_dir.path().to_path_buf();
let target_dir = base_dir.join("test-server");
tokio::fs::create_dir_all(&target_dir).await.unwrap();
let params = GenerateSkillParams {
server_id: "test-server".to_string(),
skill_name: None,
use_case_hints: None,
servers_dir: Some(base_dir),
};
let result = service.generate_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(
err.code,
ErrorCode::INVALID_PARAMS,
"a missing sidecar is the same 'not generated' caller situation as a missing \
server directory, and must be reported the same way"
);
assert!(err.message.contains("Failed to scan tools directory"));
}
#[tokio::test]
async fn test_generate_skill_stale_metadata_missing_ts_file() {
use mcp_execution_core::metadata::{
METADATA_FILE_NAME, METADATA_SCHEMA_VERSION, ParameterMetadata, ServerMetadata,
ToolMetadata as SidecarToolMetadata,
};
use tempfile::TempDir;
let service = GeneratorService::new();
let temp_dir = TempDir::new().unwrap();
let base_dir = temp_dir.path().to_path_buf();
let target_dir = base_dir.join("test-server");
tokio::fs::create_dir_all(&target_dir).await.unwrap();
let meta = ServerMetadata {
schema_version: METADATA_SCHEMA_VERSION,
server_id: "test-server".to_string(),
server_name: "Test Server".to_string(),
server_version: "1.0.0".to_string(),
tools: vec![SidecarToolMetadata {
name: "create_issue".to_string(),
typescript_name: "createIssue".to_string(),
category: None,
keywords: vec![],
description: None,
parameters: vec![ParameterMetadata {
name: "title".to_string(),
typescript_type: "string".to_string(),
required: true,
description: None,
}],
}],
};
let content = serde_json::to_string_pretty(&meta).unwrap();
tokio::fs::write(target_dir.join(METADATA_FILE_NAME), content)
.await
.unwrap();
let params = GenerateSkillParams {
server_id: "test-server".to_string(),
skill_name: None,
use_case_hints: None,
servers_dir: Some(base_dir),
};
let result = service.generate_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(
err.code,
ErrorCode::INVALID_PARAMS,
"stale metadata is the same 'not generated / drifted directory' caller situation \
as a missing sidecar, and must be reported the same way"
);
assert!(err.message.contains("Failed to scan tools directory"));
assert!(err.message.contains("create_issue"));
}
#[tokio::test]
async fn test_generate_skill_reports_orphan_ts_file_as_warning() {
use mcp_execution_core::metadata::{
METADATA_FILE_NAME, METADATA_SCHEMA_VERSION, ParameterMetadata, ServerMetadata,
ToolMetadata as SidecarToolMetadata,
};
use mcp_execution_skill::GenerateSkillResult;
use tempfile::TempDir;
let service = GeneratorService::new();
let temp_dir = TempDir::new().unwrap();
let base_dir = temp_dir.path().to_path_buf();
let target_dir = base_dir.join("test-server");
tokio::fs::create_dir_all(&target_dir).await.unwrap();
let meta = ServerMetadata {
schema_version: METADATA_SCHEMA_VERSION,
server_id: "test-server".to_string(),
server_name: "Test Server".to_string(),
server_version: "1.0.0".to_string(),
tools: vec![SidecarToolMetadata {
name: "create_issue".to_string(),
typescript_name: "createIssue".to_string(),
category: None,
keywords: vec![],
description: None,
parameters: vec![ParameterMetadata {
name: "title".to_string(),
typescript_type: "string".to_string(),
required: true,
description: None,
}],
}],
};
let content = serde_json::to_string_pretty(&meta).unwrap();
tokio::fs::write(target_dir.join(METADATA_FILE_NAME), content)
.await
.unwrap();
tokio::fs::write(target_dir.join("createIssue.ts"), "export {}")
.await
.unwrap();
tokio::fs::write(target_dir.join("orphanTool.ts"), "export {}")
.await
.unwrap();
let params = GenerateSkillParams {
server_id: "test-server".to_string(),
skill_name: None,
use_case_hints: None,
servers_dir: Some(base_dir),
};
let result = service.generate_skill(Parameters(params)).await;
assert!(
result.is_ok(),
"an orphaned .ts file must not fail the call"
);
let content = result.unwrap();
let text_content = content.content[0].as_text().unwrap();
let parsed: GenerateSkillResult = serde_json::from_str(&text_content.text).unwrap();
assert_eq!(
parsed.warnings.len(),
1,
"the orphaned .ts file must be surfaced as a warning"
);
assert!(
parsed.warnings[0].contains("orphanTool.ts"),
"warning must name the excluded file: {:?}",
parsed.warnings[0]
);
}
#[tokio::test]
async fn test_save_skill_invalid_server_id() {
let service = GeneratorService::new();
let params = SaveSkillParams {
server_id: "Invalid_Server".to_string(), content: "---\nname: test\ndescription: test\n---\n# Test".to_string(),
output_path: None,
overwrite: false,
};
let result = service.save_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
assert!(err.message.contains("lowercase"));
}
#[tokio::test]
async fn test_save_skill_missing_yaml_frontmatter() {
let service = GeneratorService::new();
let params = SaveSkillParams {
server_id: "test".to_string(),
content: "# Test Skill\n\nNo YAML frontmatter here.".to_string(),
output_path: None,
overwrite: false,
};
let result = service.save_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
assert!(err.message.contains("YAML frontmatter"));
}
#[tokio::test]
async fn test_save_skill_invalid_frontmatter_no_name() {
let service = GeneratorService::new();
let params = SaveSkillParams {
server_id: "test".to_string(),
content: "---\ndescription: test\n---\n# Test".to_string(),
output_path: None,
overwrite: false,
};
let result = service.save_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
assert!(err.message.contains("Invalid SKILL.md format"));
}
#[tokio::test]
async fn test_save_skill_invalid_frontmatter_no_description() {
let service = GeneratorService::new();
let params = SaveSkillParams {
server_id: "test".to_string(),
content: "---\nname: test-skill\n---\n# Test".to_string(),
output_path: None,
overwrite: false,
};
let result = service.save_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
assert!(err.message.contains("Invalid SKILL.md format"));
}
#[tokio::test]
async fn test_save_skill_file_exists_no_overwrite() {
use tempfile::TempDir;
let service = GeneratorService::new();
let temp_dir = TempDir::new().unwrap();
let output_path = temp_dir.path().join("SKILL.md");
tokio::fs::write(&output_path, "existing content")
.await
.unwrap();
let params = SaveSkillParams {
server_id: "test".to_string(),
content: "---\nname: test\ndescription: test\n---\n# Test".to_string(),
output_path: Some(output_path),
overwrite: false,
};
let result = service.save_skill(Parameters(params)).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
assert!(err.message.contains("already exists"));
assert!(err.message.contains("overwrite=true"));
}
#[tokio::test]
async fn test_save_skill_file_exists_with_overwrite() {
use tempfile::TempDir;
let service = GeneratorService::new();
let temp_dir = TempDir::new().unwrap();
let output_path = temp_dir.path().join("SKILL.md");
tokio::fs::write(&output_path, "existing content")
.await
.unwrap();
let params = SaveSkillParams {
server_id: "test".to_string(),
content: "---\nname: test\ndescription: test skill\n---\n# Test".to_string(),
output_path: Some(output_path.clone()),
overwrite: true,
};
let result = service.save_skill(Parameters(params)).await;
assert!(result.is_ok());
let content = result.unwrap();
let text = content.content[0].as_text().unwrap();
let parsed: SaveSkillResult = serde_json::from_str(&text.text).unwrap();
assert!(parsed.success);
assert!(parsed.overwritten);
assert_eq!(parsed.metadata.name, "test");
assert_eq!(parsed.metadata.description, "test skill");
}
#[tokio::test]
async fn test_save_skill_valid_content() {
use tempfile::TempDir;
let service = GeneratorService::new();
let temp_dir = TempDir::new().unwrap();
let output_path = temp_dir.path().join("SKILL.md");
let params = SaveSkillParams {
server_id: "test".to_string(),
content: "---\nname: test-skill\ndescription: A test skill\n---\n\n# Test Skill\n\n## Section 1\n\nContent here.".to_string(),
output_path: Some(output_path.clone()),
overwrite: false,
};
let result = service.save_skill(Parameters(params)).await;
assert!(result.is_ok());
let content = result.unwrap();
let text = content.content[0].as_text().unwrap();
let parsed: SaveSkillResult = serde_json::from_str(&text.text).unwrap();
assert!(parsed.success);
assert!(!parsed.overwritten);
assert_eq!(parsed.metadata.name, "test-skill");
assert_eq!(parsed.metadata.description, "A test skill");
assert!(parsed.metadata.section_count >= 1);
assert!(parsed.metadata.word_count > 0);
assert!(output_path.exists());
}
}