Skip to main content

mcp_execution_server/
service.rs

1//! MCP server implementation for progressive loading generation.
2//!
3//! The `GeneratorService` provides three main tools:
4//! 1. `introspect_server` - Connect to and introspect an MCP server
5//! 2. `save_categorized_tools` - Generate TypeScript files with categorization
6//! 3. `list_generated_servers` - List all servers with generated files
7
8use crate::clock::{Clock, SystemClock};
9use crate::state::StateManager;
10use crate::types::{
11    CategorizedTool, GeneratedServerInfo, IntrospectServerParams, IntrospectServerResult,
12    IntrospectedToolSummary, ListGeneratedServersParams, ListGeneratedServersResult,
13    PendingGeneration, SaveCategorizedToolsParams, SaveCategorizedToolsResult,
14};
15use mcp_execution_codegen::progressive::ProgressiveGenerator;
16use mcp_execution_core::{ServerConfig, ServerId};
17use mcp_execution_files::FilesBuilder;
18use mcp_execution_introspector::Introspector;
19use mcp_execution_skill::{
20    GenerateSkillParams, SaveSkillParams, SaveSkillResult, ScanError, build_skill_context,
21    extract_skill_metadata, scan_tools_directory, validate_server_id,
22};
23use rmcp::handler::server::ServerHandler;
24use rmcp::handler::server::tool::ToolRouter;
25use rmcp::handler::server::wrapper::Parameters;
26use rmcp::model::{
27    CallToolResult, ContentBlock, Implementation, ProtocolVersion, ServerCapabilities, ServerInfo,
28};
29use rmcp::{ErrorData as McpError, tool, tool_handler, tool_router};
30use std::collections::{HashMap, HashSet};
31use std::path::{Path, PathBuf};
32use std::sync::Arc;
33use tokio::sync::Mutex;
34
35/// Maximum SKILL.md content size in bytes (100KB).
36const MAX_SKILL_CONTENT_SIZE: usize = 100 * 1024;
37
38/// MCP server for progressive loading generation.
39///
40/// This service helps generate progressive loading TypeScript files for other
41/// MCP servers. Claude provides the categorization intelligence through natural
42/// language understanding - no separate LLM API needed.
43///
44/// # Workflow
45///
46/// 1. Call `introspect_server` to discover tools from a target MCP server
47/// 2. Claude analyzes the tools and assigns categories, keywords, descriptions
48/// 3. Call `save_categorized_tools` to generate TypeScript files
49/// 4. Use `list_generated_servers` to see all generated servers
50///
51/// # Examples
52///
53/// ```no_run
54/// use mcp_execution_server::service::GeneratorService;
55/// use rmcp::transport::stdio;
56///
57/// # async fn example() {
58/// let service = GeneratorService::new();
59/// // Service implements rmcp ServerHandler trait
60/// # }
61/// ```
62#[derive(Debug, Clone)]
63pub struct GeneratorService {
64    /// State manager for pending generations
65    state: Arc<StateManager>,
66
67    /// Per-server-id introspector locks.
68    ///
69    /// Keying the lock by [`ServerId`] means a slow or hung downstream MCP
70    /// server only blocks `introspect_server` calls for that same server id,
71    /// not for unrelated ids across all sessions. The outer map mutex is only
72    /// held long enough to fetch or insert the per-id handle - never across
73    /// the `discover_server` await point.
74    introspectors: Arc<Mutex<HashMap<ServerId, Arc<Mutex<Introspector>>>>>,
75
76    /// Per-output-directory export locks, keyed by the (uncanonicalized)
77    /// output path as supplied to `introspect_server` / stored on the
78    /// pending generation. Same rationale as `introspectors`: keying by the
79    /// contended resource means an export for one `output_dir` never blocks
80    /// an export for a different one.
81    exports: Arc<Mutex<HashMap<PathBuf, Arc<Mutex<()>>>>>,
82
83    /// Clock used to construct pending generations (shared with `state`)
84    clock: Arc<dyn Clock>,
85
86    /// Tool router for MCP protocol
87    // Only read via macro-expanded code generated by the `#[tool_router]` attribute
88    // macro, so the compiler's static dead-code analysis cannot see the usage.
89    #[allow(dead_code)]
90    tool_router: ToolRouter<Self>,
91}
92
93impl GeneratorService {
94    /// Creates a new generator service using the real system clock.
95    #[must_use]
96    pub fn new() -> Self {
97        Self::with_clock(Arc::new(SystemClock))
98    }
99
100    /// Creates a new generator service backed by a custom clock.
101    ///
102    /// Used in tests to inject a fake clock so session expiry can be
103    /// exercised deterministically.
104    fn with_clock(clock: Arc<dyn Clock>) -> Self {
105        Self {
106            state: Arc::new(StateManager::with_clock(Arc::clone(&clock))),
107            introspectors: Arc::new(Mutex::new(HashMap::new())),
108            exports: Arc::new(Mutex::new(HashMap::new())),
109            clock,
110            tool_router: Self::tool_router(),
111        }
112    }
113
114    /// Returns the per-server-id introspector handle, creating one if absent.
115    ///
116    /// The outer map lock is released before the returned handle is awaited
117    /// on, so discovery of unrelated server ids never contends on it.
118    async fn introspector_for(&self, server_id: &ServerId) -> Arc<Mutex<Introspector>> {
119        let mut introspectors = self.introspectors.lock().await;
120        introspectors
121            .entry(server_id.clone())
122            .or_insert_with(|| Arc::new(Mutex::new(Introspector::new())))
123            .clone()
124    }
125
126    /// Evicts the per-server-id introspector handle after use, but only if
127    /// the map still holds the exact handle the caller obtained.
128    ///
129    /// `server_id` values are caller-supplied, so without eviction the map
130    /// grows without bound as new ids are introspected. Called after
131    /// `discover_server` completes, regardless of outcome.
132    ///
133    /// A caller must pass the same `Arc<Mutex<Introspector>>` it received
134    /// from [`Self::introspector_for`]. Removing by `server_id` alone is a
135    /// TOCTOU bug: if another in-flight call for the same id already evicted
136    /// and a third call inserted a fresh handle, an unconditional `remove`
137    /// would prune that live handle out from under the third call. Comparing
138    /// with [`Arc::ptr_eq`] ensures a caller can only ever evict the entry it
139    /// created.
140    async fn evict_introspector(&self, server_id: &ServerId, handle: &Arc<Mutex<Introspector>>) {
141        let mut introspectors = self.introspectors.lock().await;
142        if let std::collections::hash_map::Entry::Occupied(entry) =
143            introspectors.entry(server_id.clone())
144            && Arc::ptr_eq(entry.get(), handle)
145        {
146            entry.remove();
147        }
148    }
149
150    /// Returns the per-output-directory export lock, creating one if absent.
151    ///
152    /// Mirrors [`Self::introspector_for`]: the outer map lock is released
153    /// before the returned handle is awaited on, so exports to unrelated
154    /// output directories never contend on it. Holding this lock across an
155    /// [`mcp_execution_files::FileSystem::export_to_filesystem`] call
156    /// serializes any two concurrent `save_categorized_tools` calls for the
157    /// same `output_dir` that overlap while holding the same handle,
158    /// narrowing the in-process trigger for the data-loss race described in
159    /// issue #169. This is not an unconditional guarantee across three or
160    /// more overlapping calls: a call that fetches a fresh handle only
161    /// after an earlier holder has already evicted its own can still run
162    /// concurrently with a still-in-flight call holding the stale handle
163    /// (same eviction-boundary gap as [`Self::evict_introspector`]). The
164    /// age-gated sweep in `mcp-execution-files` is what ultimately prevents
165    /// data loss if that happens.
166    async fn export_lock_for(&self, output_dir: &Path) -> Arc<Mutex<()>> {
167        let mut exports = self.exports.lock().await;
168        exports
169            .entry(output_dir.to_path_buf())
170            .or_insert_with(|| Arc::new(Mutex::new(())))
171            .clone()
172    }
173
174    /// Evicts the per-output-directory export lock after use, but only if
175    /// the map still holds the exact handle the caller obtained.
176    ///
177    /// Same identity-checked eviction as [`Self::evict_introspector`] and
178    /// for the same reason: `output_dir` values are caller-supplied, so
179    /// without eviction the map grows without bound, and an unconditional
180    /// `remove` keyed only by path would be a TOCTOU bug against a
181    /// concurrently inserted fresh handle.
182    async fn evict_export_lock(&self, output_dir: &Path, handle: &Arc<Mutex<()>>) {
183        let mut exports = self.exports.lock().await;
184        if let std::collections::hash_map::Entry::Occupied(entry) =
185            exports.entry(output_dir.to_path_buf())
186            && Arc::ptr_eq(entry.get(), handle)
187        {
188            entry.remove();
189        }
190    }
191}
192
193impl Default for GeneratorService {
194    fn default() -> Self {
195        Self::new()
196    }
197}
198
199#[tool_router]
200impl GeneratorService {
201    /// Introspect an MCP server and prepare for categorization.
202    ///
203    /// Connects to the target MCP server, discovers its tools, and returns
204    /// metadata for Claude to categorize. Returns a session ID for use with
205    /// `save_categorized_tools`.
206    #[tool(
207        description = "Connect to an MCP server, discover its tools, and return metadata for categorization. Returns a session ID for use with save_categorized_tools."
208    )]
209    async fn introspect_server(
210        &self,
211        Parameters(params): Parameters<IntrospectServerParams>,
212    ) -> Result<CallToolResult, McpError> {
213        // Validate server_id format
214        validate_server_id(&params.server_id).map_err(|e| McpError::invalid_params(e, None))?;
215
216        // Extract server_id before consuming params
217        let server_id_str = params.server_id;
218        let server_id = ServerId::new(&server_id_str);
219
220        // Determine output directory (needs server_id_str)
221        let output_dir = params.output_dir.unwrap_or_else(|| {
222            dirs::home_dir()
223                .unwrap_or_else(|| PathBuf::from("."))
224                .join(".claude")
225                .join("servers")
226                .join(&server_id_str)
227        });
228
229        // Build server config (consume args and env to avoid clones)
230        let mut config_builder = ServerConfig::builder().command(params.command);
231
232        for arg in params.args {
233            config_builder = config_builder.arg(arg);
234        }
235
236        for (key, value) in params.env {
237            config_builder = config_builder.env(key, value);
238        }
239
240        if let Some(secs) = params.connect_timeout_secs {
241            config_builder = config_builder.connect_timeout(std::time::Duration::from_secs(secs));
242        }
243
244        if let Some(secs) = params.discover_timeout_secs {
245            config_builder = config_builder.discover_timeout(std::time::Duration::from_secs(secs));
246        }
247
248        let config = config_builder.build();
249
250        // Connect and introspect, holding only the lock for this server_id
251        let introspector_handle = self.introspector_for(&server_id).await;
252        let discover_result = {
253            let mut introspector = introspector_handle.lock().await;
254            introspector
255                .discover_server(server_id.clone(), &config)
256                .await
257        };
258
259        // Evict the per-server-id handle regardless of outcome, so caller-supplied
260        // server_id values can't grow the introspectors map without bound. Only
261        // removes the entry if it is still this exact handle (see
262        // `evict_introspector` docs for why identity matters here).
263        self.evict_introspector(&server_id, &introspector_handle)
264            .await;
265
266        let server_info = discover_result.map_err(|e| {
267            if e.is_validation_error() {
268                McpError::invalid_params(e.to_string(), None)
269            } else {
270                McpError::internal_error(format!("Failed to introspect server: {e}"), None)
271            }
272        })?;
273
274        // Extract tool metadata for Claude
275        let tools: Vec<IntrospectedToolSummary> = server_info
276            .tools
277            .iter()
278            .map(|tool| {
279                let parameters = extract_parameter_names(&tool.input_schema);
280
281                IntrospectedToolSummary {
282                    name: tool.name.as_str().to_string(),
283                    description: tool.description.clone(),
284                    parameters,
285                }
286            })
287            .collect();
288
289        // Store pending generation
290        let pending = PendingGeneration::new(
291            server_id,
292            server_info.clone(),
293            config,
294            output_dir.clone(),
295            self.clock.as_ref(),
296        );
297
298        let session_id = self.state.store(pending.clone()).await;
299
300        // Build result
301        let result = IntrospectServerResult {
302            server_id: server_id_str,
303            server_name: server_info.name,
304            tools_found: tools.len(),
305            tools,
306            session_id,
307            expires_at: pending.expires_at,
308        };
309
310        Ok(CallToolResult::success(vec![ContentBlock::text(
311            serde_json::to_string_pretty(&result).map_err(|e| {
312                McpError::internal_error(format!("Failed to serialize result: {e}"), None)
313            })?,
314        )]))
315    }
316
317    /// Save categorized tools as TypeScript files.
318    ///
319    /// Generates progressive loading TypeScript files using Claude's
320    /// categorization. Requires `session_id` from a previous `introspect_server`
321    /// call.
322    #[tool(
323        description = "Generate progressive loading TypeScript files using Claude's categorization. Requires session_id from a previous introspect_server call."
324    )]
325    async fn save_categorized_tools(
326        &self,
327        Parameters(params): Parameters<SaveCategorizedToolsParams>,
328    ) -> Result<CallToolResult, McpError> {
329        // Retrieve pending generation
330        let pending = self.state.take(params.session_id).await.ok_or_else(|| {
331            McpError::invalid_params(
332                "Session not found or expired. Please run introspect_server again.",
333                None,
334            )
335        })?;
336
337        // Validate categorized tools match introspected tools
338        let introspected_names: HashSet<_> = pending
339            .server_info
340            .tools
341            .iter()
342            .map(|t| t.name.as_str())
343            .collect();
344
345        for cat_tool in &params.categorized_tools {
346            if !introspected_names.contains(cat_tool.name.as_str()) {
347                return Err(McpError::invalid_params(
348                    format!("Tool '{}' not found in introspected tools", cat_tool.name),
349                    None,
350                ));
351            }
352        }
353
354        // Build categorization map and category stats in single pass (avoid double iteration)
355        let tool_count = params.categorized_tools.len();
356        let mut categorization: HashMap<String, &CategorizedTool> =
357            HashMap::with_capacity(tool_count);
358        let mut categories: HashMap<String, usize> = HashMap::with_capacity(tool_count);
359
360        for tool in &params.categorized_tools {
361            categorization.insert(tool.name.clone(), tool);
362            *categories.entry(tool.category.clone()).or_default() += 1;
363        }
364
365        // Generate code with categorization
366        let generator = ProgressiveGenerator::new().map_err(|e| {
367            McpError::internal_error(format!("Failed to create generator: {e}"), None)
368        })?;
369
370        let code = generate_with_categorization(&generator, &pending.server_info, &categorization)
371            .map_err(|e| McpError::internal_error(format!("Failed to generate code: {e}"), None))?;
372
373        // Build virtual filesystem
374        let vfs = FilesBuilder::from_generated_code(code, "/")
375            .build()
376            .map_err(|e| McpError::internal_error(format!("Failed to build VFS: {e}"), None))?;
377
378        // Capture file count before moving vfs
379        let files_generated = vfs.file_count();
380
381        // Ensure the parent of the output directory exists (async). Only the
382        // parent is needed: `export_to_filesystem` publishes `output_dir`
383        // itself atomically (single rename on first generate, stage-then-swap
384        // on regeneration), so pre-creating it here would force the slower
385        // regeneration path even on a brand-new server.
386        if let Some(parent) = pending.output_dir.parent() {
387            tokio::fs::create_dir_all(parent).await.map_err(|e| {
388                McpError::internal_error(format!("Failed to create output directory: {e}"), None)
389            })?;
390        }
391
392        // Export to filesystem (blocking operation wrapped in spawn_blocking).
393        // Held across the export so a second concurrent call for the same
394        // output_dir blocks until the first finishes, rather than racing on
395        // the underlying staging/swap (see `export_lock_for`).
396        let export_lock = self.export_lock_for(&pending.output_dir).await;
397        let export_guard = export_lock.lock().await;
398
399        let output_dir = pending.output_dir.clone();
400        let export_result =
401            tokio::task::spawn_blocking(move || vfs.export_to_filesystem(&output_dir)).await;
402
403        drop(export_guard);
404        self.evict_export_lock(&pending.output_dir, &export_lock)
405            .await;
406
407        export_result
408            .map_err(|e| McpError::internal_error(format!("Task join error: {e}"), None))?
409            .map_err(|e| McpError::internal_error(format!("Failed to export files: {e}"), None))?;
410
411        let result = SaveCategorizedToolsResult {
412            success: true,
413            files_generated,
414            output_dir: pending.output_dir.display().to_string(),
415            categories,
416            errors: vec![],
417        };
418
419        Ok(CallToolResult::success(vec![ContentBlock::text(
420            serde_json::to_string_pretty(&result).map_err(|e| {
421                McpError::internal_error(format!("Failed to serialize result: {e}"), None)
422            })?,
423        )]))
424    }
425
426    /// List all servers with generated progressive loading files.
427    ///
428    /// Scans the output directory (default: `~/.claude/servers`) for servers
429    /// that have generated TypeScript files.
430    #[tool(
431        description = "List all MCP servers that have generated progressive loading files in ~/.claude/servers/"
432    )]
433    async fn list_generated_servers(
434        &self,
435        Parameters(params): Parameters<ListGeneratedServersParams>,
436    ) -> Result<CallToolResult, McpError> {
437        let base_dir = params.base_dir.map_or_else(
438            || {
439                dirs::home_dir()
440                    .unwrap_or_else(|| PathBuf::from("."))
441                    .join(".claude")
442                    .join("servers")
443            },
444            PathBuf::from,
445        );
446
447        // Scan directories (blocking operation wrapped in spawn_blocking)
448        let servers = tokio::task::spawn_blocking(move || {
449            let mut servers = Vec::new();
450
451            if base_dir.exists()
452                && base_dir.is_dir()
453                && let Ok(entries) = std::fs::read_dir(&base_dir)
454            {
455                for entry in entries.flatten() {
456                    if entry.path().is_dir() {
457                        let id = entry.file_name().to_string_lossy().to_string();
458
459                        // Count .ts files (excluding _runtime and starting with _)
460                        let tool_count = std::fs::read_dir(entry.path()).map_or(0, |e| {
461                            e.flatten()
462                                .filter(|f| {
463                                    let name = f.file_name();
464                                    let name = name.to_string_lossy();
465                                    name.ends_with(".ts") && !name.starts_with('_')
466                                })
467                                .count()
468                        });
469
470                        // Get modification time
471                        let generated_at = entry
472                            .metadata()
473                            .and_then(|m| m.modified())
474                            .ok()
475                            .map(chrono::DateTime::<chrono::Utc>::from);
476
477                        servers.push(GeneratedServerInfo {
478                            id,
479                            tool_count,
480                            generated_at,
481                            output_dir: entry.path().display().to_string(),
482                        });
483                    }
484                }
485            }
486
487            servers.sort_by(|a, b| a.id.cmp(&b.id));
488            servers
489        })
490        .await
491        .map_err(|e| McpError::internal_error(format!("Task join error: {e}"), None))?;
492
493        let result = ListGeneratedServersResult {
494            total_servers: servers.len(),
495            servers,
496        };
497
498        Ok(CallToolResult::success(vec![ContentBlock::text(
499            serde_json::to_string_pretty(&result).map_err(|e| {
500                McpError::internal_error(format!("Failed to serialize result: {e}"), None)
501            })?,
502        )]))
503    }
504
505    /// Generate context for creating a Claude Code skill.
506    ///
507    /// Analyzes generated TypeScript files and returns structured context
508    /// that Claude uses to generate an optimal SKILL.md file.
509    ///
510    /// # Workflow
511    ///
512    /// 1. Call `generate_skill` with `server_id`
513    /// 2. Claude receives context and `generation_prompt`
514    /// 3. Claude generates SKILL.md content
515    /// 4. Call `save_skill` with the generated content
516    #[tool(
517        description = "Analyze generated TypeScript files and return context for Claude to create a SKILL.md file. Returns tool metadata, categories, and a generation prompt."
518    )]
519    async fn generate_skill(
520        &self,
521        Parameters(params): Parameters<GenerateSkillParams>,
522    ) -> Result<CallToolResult, McpError> {
523        // Validate server_id format and length
524        validate_server_id(&params.server_id).map_err(|e| McpError::invalid_params(e, None))?;
525
526        // Determine servers directory
527        let servers_dir = params.servers_dir.unwrap_or_else(|| {
528            dirs::home_dir()
529                .unwrap_or_else(|| PathBuf::from("."))
530                .join(".claude")
531                .join("servers")
532        });
533
534        let server_dir = servers_dir.join(&params.server_id);
535
536        // Check if server directory exists
537        if !server_dir.exists() {
538            return Err(McpError::invalid_params(
539                format!(
540                    "Server directory not found: {}. Run generate first.",
541                    server_dir.display()
542                ),
543                None,
544            ));
545        }
546
547        // Scan and parse tool files. A missing or version-mismatched sidecar reflects the
548        // same "not generated / stale directory" caller situation as the `!server_dir.exists()`
549        // check above, so it is reported the same way (`invalid_params`), not as a server fault.
550        let scan_result = scan_tools_directory(&server_dir)
551            .await
552            .map_err(|e| match e {
553                ScanError::MissingMetadata { .. }
554                | ScanError::UnsupportedSchema { .. }
555                | ScanError::StaleMetadata { .. } => {
556                    McpError::invalid_params(format!("Failed to scan tools directory: {e}"), None)
557                }
558                ScanError::Io(_)
559                | ScanError::DirectoryNotFound { .. }
560                | ScanError::MetadataParse { .. }
561                | ScanError::TooManyFiles { .. }
562                | ScanError::FileTooLarge { .. } => {
563                    McpError::internal_error(format!("Failed to scan tools directory: {e}"), None)
564                }
565            })?;
566
567        if scan_result.tools.is_empty() {
568            return Err(McpError::invalid_params(
569                format!(
570                    "No tool files found in {}. Run generate first.",
571                    server_dir.display()
572                ),
573                None,
574            ));
575        }
576
577        // Build context
578        let mut result = build_skill_context(
579            &params.server_id,
580            &scan_result.tools,
581            params.use_case_hints.as_deref(),
582        );
583
584        // Surface non-fatal drift warnings (e.g. `.ts` files excluded for lacking
585        // a sidecar entry) in the structured response, not just server-side
586        // tracing output (issue #161).
587        result.warnings = scan_result.warnings;
588
589        // Override skill name if provided
590        if let Some(name) = params.skill_name {
591            result.skill_name = name;
592        }
593
594        Ok(CallToolResult::success(vec![ContentBlock::text(
595            serde_json::to_string_pretty(&result).map_err(|e| {
596                McpError::internal_error(format!("Failed to serialize result: {e}"), None)
597            })?,
598        )]))
599    }
600
601    /// Save a generated skill to the filesystem.
602    ///
603    /// Writes SKILL.md content to `~/.claude/skills/{server_id}/SKILL.md`.
604    /// Validates that the content contains required YAML frontmatter.
605    #[tool(
606        description = "Save generated SKILL.md content to ~/.claude/skills/{server_id}/. Use after Claude generates skill content from generate_skill context."
607    )]
608    async fn save_skill(
609        &self,
610        Parameters(params): Parameters<SaveSkillParams>,
611    ) -> Result<CallToolResult, McpError> {
612        // Validate server_id format and length
613        validate_server_id(&params.server_id).map_err(|e| McpError::invalid_params(e, None))?;
614
615        // Validate content size (DoS protection)
616        if params.content.len() > MAX_SKILL_CONTENT_SIZE {
617            return Err(McpError::invalid_params(
618                format!(
619                    "content too large: {} bytes exceeds {} limit",
620                    params.content.len(),
621                    MAX_SKILL_CONTENT_SIZE
622                ),
623                None,
624            ));
625        }
626
627        // Validate content has YAML frontmatter
628        if !params.content.starts_with("---") {
629            return Err(McpError::invalid_params(
630                "Content must start with YAML frontmatter (---)",
631                None,
632            ));
633        }
634
635        // Extract metadata from frontmatter
636        let metadata = extract_skill_metadata(&params.content)
637            .map_err(|e| McpError::invalid_params(format!("Invalid SKILL.md format: {e}"), None))?;
638
639        // Determine output path
640        let output_path = params.output_path.unwrap_or_else(|| {
641            dirs::home_dir()
642                .unwrap_or_else(|| PathBuf::from("."))
643                .join(".claude")
644                .join("skills")
645                .join(&params.server_id)
646                .join("SKILL.md")
647        });
648
649        // Check if file exists
650        let overwritten = output_path.exists();
651        if overwritten && !params.overwrite {
652            return Err(McpError::invalid_params(
653                format!(
654                    "Skill file already exists: {}. Use overwrite=true to replace.",
655                    output_path.display()
656                ),
657                None,
658            ));
659        }
660
661        // Create parent directory
662        if let Some(parent) = output_path.parent() {
663            tokio::fs::create_dir_all(parent).await.map_err(|e| {
664                McpError::internal_error(format!("Failed to create directory: {e}"), None)
665            })?;
666        }
667
668        // Write file
669        tokio::fs::write(&output_path, &params.content)
670            .await
671            .map_err(|e| McpError::internal_error(format!("Failed to write file: {e}"), None))?;
672
673        let result = SaveSkillResult {
674            success: true,
675            output_path: output_path.display().to_string(),
676            overwritten,
677            metadata,
678        };
679
680        Ok(CallToolResult::success(vec![ContentBlock::text(
681            serde_json::to_string_pretty(&result).map_err(|e| {
682                McpError::internal_error(format!("Failed to serialize result: {e}"), None)
683            })?,
684        )]))
685    }
686}
687
688#[tool_handler]
689impl ServerHandler for GeneratorService {
690    fn get_info(&self) -> ServerInfo {
691        let mut info = ServerInfo::default();
692        info.protocol_version = ProtocolVersion::V_2025_06_18;
693        info.capabilities = ServerCapabilities::builder().enable_tools().build();
694        info.server_info = Implementation::new(env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
695        info.instructions = Some(
696            "Generate progressive loading TypeScript files for MCP servers. \
697             Use introspect_server to discover tools, then save_categorized_tools \
698             with your categorization."
699                .to_string(),
700        );
701        info
702    }
703}
704
705// ============================================================================
706// Helper functions
707// ============================================================================
708
709/// Extracts parameter names from a JSON Schema.
710fn extract_parameter_names(schema: &serde_json::Value) -> Vec<String> {
711    schema
712        .get("properties")
713        .and_then(|p| p.as_object())
714        .map(|props| props.keys().cloned().collect())
715        .unwrap_or_default()
716}
717
718/// Generates code with categorization metadata.
719///
720/// Converts the categorization map to the format expected by the generator
721/// and calls `generate_with_categories`.
722fn generate_with_categorization(
723    generator: &ProgressiveGenerator,
724    server_info: &mcp_execution_introspector::ServerInfo,
725    categorization: &HashMap<String, &CategorizedTool>,
726) -> mcp_execution_core::Result<mcp_execution_codegen::GeneratedCode> {
727    use mcp_execution_codegen::progressive::ToolCategorization;
728
729    // Convert CategorizedTool map to ToolCategorization map
730    let categorizations: HashMap<String, ToolCategorization> = categorization
731        .iter()
732        .map(|(tool_name, cat_tool)| {
733            (
734                tool_name.clone(),
735                ToolCategorization {
736                    category: cat_tool.category.clone(),
737                    keywords: cat_tool.keywords.clone(),
738                    short_description: cat_tool.short_description.clone(),
739                },
740            )
741        })
742        .collect();
743
744    generator.generate_with_categories(server_info, &categorizations)
745}
746
747#[cfg(test)]
748mod tests {
749    use super::*;
750    use chrono::Utc;
751    use mcp_execution_core::ToolName;
752    use mcp_execution_introspector::{ServerCapabilities, ToolInfo};
753    use rmcp::model::ErrorCode;
754    use uuid::Uuid;
755
756    // ========================================================================
757    // Helper Functions Tests
758    // ========================================================================
759
760    #[test]
761    fn test_extract_parameter_names() {
762        let schema = serde_json::json!({
763            "type": "object",
764            "properties": {
765                "name": { "type": "string" },
766                "age": { "type": "number" }
767            }
768        });
769
770        let params = extract_parameter_names(&schema);
771        assert_eq!(params.len(), 2);
772        assert!(params.contains(&"name".to_string()));
773        assert!(params.contains(&"age".to_string()));
774    }
775
776    #[test]
777    fn test_extract_parameter_names_empty() {
778        let schema = serde_json::json!({
779            "type": "object"
780        });
781
782        let params = extract_parameter_names(&schema);
783        assert_eq!(params.len(), 0);
784    }
785
786    #[test]
787    fn test_extract_parameter_names_no_properties() {
788        let schema = serde_json::json!({
789            "type": "string"
790        });
791
792        let params = extract_parameter_names(&schema);
793        assert_eq!(params.len(), 0);
794    }
795
796    #[test]
797    fn test_extract_parameter_names_nested_object() {
798        let schema = serde_json::json!({
799            "type": "object",
800            "properties": {
801                "user": {
802                    "type": "object",
803                    "properties": {
804                        "name": { "type": "string" }
805                    }
806                },
807                "age": { "type": "number" }
808            }
809        });
810
811        let params = extract_parameter_names(&schema);
812        assert_eq!(params.len(), 2);
813        assert!(params.contains(&"user".to_string()));
814        assert!(params.contains(&"age".to_string()));
815    }
816
817    #[test]
818    fn test_generate_with_categorization() {
819        let generator = ProgressiveGenerator::new().unwrap();
820
821        let server_info = mcp_execution_introspector::ServerInfo {
822            id: ServerId::new("test"),
823            name: "Test Server".to_string(),
824            version: "1.0.0".to_string(),
825            capabilities: ServerCapabilities {
826                supports_tools: true,
827                supports_resources: false,
828                supports_prompts: false,
829            },
830            tools: vec![ToolInfo {
831                name: ToolName::new("test_tool"),
832                description: "Test tool description".to_string(),
833                input_schema: serde_json::json!({
834                    "type": "object",
835                    "properties": {
836                        "param1": { "type": "string" }
837                    }
838                }),
839                output_schema: None,
840            }],
841        };
842
843        let categorized_tool = CategorizedTool {
844            name: "test_tool".to_string(),
845            category: "testing".to_string(),
846            keywords: "test,tool".to_string(),
847            short_description: "Test tool for testing".to_string(),
848        };
849
850        let mut categorization = HashMap::new();
851        categorization.insert("test_tool".to_string(), &categorized_tool);
852
853        let result = generate_with_categorization(&generator, &server_info, &categorization);
854        assert!(result.is_ok());
855
856        let code = result.unwrap();
857        assert!(code.file_count() > 0, "Should generate at least one file");
858    }
859
860    #[test]
861    fn test_generate_with_categorization_multiple_tools() {
862        let generator = ProgressiveGenerator::new().unwrap();
863
864        let server_info = mcp_execution_introspector::ServerInfo {
865            id: ServerId::new("test"),
866            name: "Test Server".to_string(),
867            version: "1.0.0".to_string(),
868            capabilities: ServerCapabilities {
869                supports_tools: true,
870                supports_resources: false,
871                supports_prompts: false,
872            },
873            tools: vec![
874                ToolInfo {
875                    name: ToolName::new("tool1"),
876                    description: "First tool".to_string(),
877                    input_schema: serde_json::json!({"type": "object"}),
878                    output_schema: None,
879                },
880                ToolInfo {
881                    name: ToolName::new("tool2"),
882                    description: "Second tool".to_string(),
883                    input_schema: serde_json::json!({"type": "object"}),
884                    output_schema: None,
885                },
886            ],
887        };
888
889        let tool1 = CategorizedTool {
890            name: "tool1".to_string(),
891            category: "category1".to_string(),
892            keywords: "test".to_string(),
893            short_description: "Tool 1".to_string(),
894        };
895
896        let tool2 = CategorizedTool {
897            name: "tool2".to_string(),
898            category: "category2".to_string(),
899            keywords: "test".to_string(),
900            short_description: "Tool 2".to_string(),
901        };
902
903        let mut categorization = HashMap::new();
904        categorization.insert("tool1".to_string(), &tool1);
905        categorization.insert("tool2".to_string(), &tool2);
906
907        let result = generate_with_categorization(&generator, &server_info, &categorization);
908        assert!(result.is_ok());
909    }
910
911    #[test]
912    fn test_generate_with_categorization_empty_tools() {
913        let generator = ProgressiveGenerator::new().unwrap();
914
915        let server_id = ServerId::new("test");
916        let server_info = mcp_execution_introspector::ServerInfo {
917            id: server_id,
918            name: "Empty Server".to_string(),
919            version: "1.0.0".to_string(),
920            capabilities: ServerCapabilities {
921                supports_tools: true,
922                supports_resources: false,
923                supports_prompts: false,
924            },
925            tools: vec![],
926        };
927
928        let categorization = HashMap::new();
929
930        let result = generate_with_categorization(&generator, &server_info, &categorization);
931        assert!(result.is_ok());
932    }
933
934    // ========================================================================
935    // Service Tests
936    // ========================================================================
937
938    #[test]
939    fn test_generator_service_new() {
940        let service = GeneratorService::new();
941        assert!(service.introspectors.try_lock().is_ok());
942        assert!(service.exports.try_lock().is_ok());
943    }
944
945    #[test]
946    fn test_generator_service_default() {
947        let service = GeneratorService::default();
948        assert!(service.introspectors.try_lock().is_ok());
949        assert!(service.exports.try_lock().is_ok());
950    }
951
952    #[test]
953    fn test_get_info() {
954        let service = GeneratorService::new();
955        let info = service.get_info();
956
957        assert_eq!(info.protocol_version, ProtocolVersion::V_2025_06_18);
958        assert!(info.capabilities.tools.is_some());
959        assert!(info.instructions.is_some());
960        assert_eq!(info.server_info.name, env!("CARGO_PKG_NAME"));
961        assert_eq!(info.server_info.version, env!("CARGO_PKG_VERSION"));
962    }
963
964    // ========================================================================
965    // Input Validation Tests
966    // ========================================================================
967
968    #[tokio::test]
969    async fn test_introspect_server_invalid_server_id_uppercase() {
970        let service = GeneratorService::new();
971
972        let params = IntrospectServerParams {
973            server_id: "GitHub".to_string(), // Invalid: contains uppercase
974            command: "echo".to_string(),
975            args: vec![],
976            env: HashMap::new(),
977            output_dir: None,
978            connect_timeout_secs: None,
979            discover_timeout_secs: None,
980        };
981
982        let result = service.introspect_server(Parameters(params)).await;
983
984        assert!(result.is_err());
985        let err = result.unwrap_err();
986        assert_eq!(err.code, ErrorCode::INVALID_PARAMS); // Invalid params error code
987    }
988
989    #[tokio::test]
990    async fn test_introspect_server_invalid_server_id_underscore() {
991        let service = GeneratorService::new();
992
993        let params = IntrospectServerParams {
994            server_id: "git_hub".to_string(), // Invalid: contains underscore
995            command: "echo".to_string(),
996            args: vec![],
997            env: HashMap::new(),
998            output_dir: None,
999            connect_timeout_secs: None,
1000            discover_timeout_secs: None,
1001        };
1002
1003        let result = service.introspect_server(Parameters(params)).await;
1004
1005        assert!(result.is_err());
1006        let err = result.unwrap_err();
1007        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1008    }
1009
1010    #[tokio::test]
1011    async fn test_introspect_server_invalid_server_id_special_chars() {
1012        let service = GeneratorService::new();
1013
1014        let params = IntrospectServerParams {
1015            server_id: "git@hub".to_string(), // Invalid: contains @
1016            command: "echo".to_string(),
1017            args: vec![],
1018            env: HashMap::new(),
1019            output_dir: None,
1020            connect_timeout_secs: None,
1021            discover_timeout_secs: None,
1022        };
1023
1024        let result = service.introspect_server(Parameters(params)).await;
1025
1026        assert!(result.is_err());
1027    }
1028
1029    #[tokio::test]
1030    async fn test_introspect_server_valid_server_id_with_hyphens() {
1031        let service = GeneratorService::new();
1032
1033        let params = IntrospectServerParams {
1034            server_id: "git-hub-server".to_string(), // Valid
1035            command: "echo".to_string(),
1036            args: vec!["test".to_string()],
1037            env: HashMap::new(),
1038            output_dir: None,
1039            connect_timeout_secs: None,
1040            discover_timeout_secs: None,
1041        };
1042
1043        // This will fail because echo is not an MCP server, but validation should pass
1044        let result = service.introspect_server(Parameters(params)).await;
1045
1046        // Should fail with internal error (connection), not invalid params
1047        if let Err(err) = result {
1048            assert_ne!(
1049                err.code,
1050                ErrorCode::INVALID_PARAMS,
1051                "Should not be invalid params error"
1052            );
1053        }
1054    }
1055
1056    #[tokio::test]
1057    async fn test_introspect_server_valid_server_id_digits() {
1058        let service = GeneratorService::new();
1059
1060        let params = IntrospectServerParams {
1061            server_id: "server123".to_string(), // Valid: lowercase + digits
1062            command: "echo".to_string(),
1063            args: vec![],
1064            env: HashMap::new(),
1065            output_dir: None,
1066            connect_timeout_secs: None,
1067            discover_timeout_secs: None,
1068        };
1069
1070        let result = service.introspect_server(Parameters(params)).await;
1071
1072        // Should fail with internal error (connection), not invalid params
1073        if let Err(err) = result {
1074            assert_ne!(err.code, ErrorCode::INVALID_PARAMS);
1075        }
1076    }
1077
1078    /// A zero timeout is a client input error, not a server-side connection
1079    /// failure — it must surface as `INVALID_PARAMS`, matching the sibling
1080    /// `validate_server_id` behavior, not `internal_error`.
1081    #[tokio::test]
1082    async fn test_introspect_server_zero_connect_timeout_is_invalid_params() {
1083        let service = GeneratorService::new();
1084
1085        let params = IntrospectServerParams {
1086            server_id: "zero-timeout-test".to_string(),
1087            command: "echo".to_string(),
1088            args: vec![],
1089            env: HashMap::new(),
1090            output_dir: None,
1091            connect_timeout_secs: Some(0),
1092            discover_timeout_secs: None,
1093        };
1094
1095        let result = service.introspect_server(Parameters(params)).await;
1096
1097        let err = result.expect_err("zero connect_timeout must be rejected");
1098        assert_eq!(
1099            err.code,
1100            ErrorCode::INVALID_PARAMS,
1101            "zero timeout is a client input error, not an internal error"
1102        );
1103    }
1104
1105    // ========================================================================
1106    // Per-server-id locking Tests (issue #120)
1107    //
1108    // These test the exact `Arc<Mutex<Introspector>>` handles and keyed-lock
1109    // pattern that `introspect_server` relies on via `introspector_for`,
1110    // rather than driving a real (or fake) subprocess through
1111    // `discover_server`. This keeps the tests deterministic and
1112    // platform-independent while still exercising the production locking
1113    // primitive: `introspect_server` does nothing more than fetch a handle
1114    // via `introspector_for` and `.lock().await` it around the
1115    // `discover_server` call, so proving the handles behave correctly here
1116    // proves the concurrency property end to end.
1117    // ========================================================================
1118
1119    /// `introspect_server` must evict its per-server-id entry from the
1120    /// `introspectors` map once `discover_server` completes, regardless of
1121    /// outcome - otherwise caller-supplied `server_id`s would grow the map
1122    /// without bound.
1123    #[tokio::test]
1124    async fn test_introspect_server_evicts_map_entry_after_completion() {
1125        let service = GeneratorService::new();
1126
1127        let params = IntrospectServerParams {
1128            server_id: "evict-after-completion".to_string(),
1129            command: "echo".to_string(), // not an MCP server, discover_server fails fast
1130            args: vec![],
1131            env: HashMap::new(),
1132            output_dir: None,
1133            connect_timeout_secs: None,
1134            discover_timeout_secs: None,
1135        };
1136
1137        let result = service.introspect_server(Parameters(params)).await;
1138        assert!(
1139            result.is_err(),
1140            "echo is not an MCP server, expected a connection failure"
1141        );
1142
1143        assert!(
1144            service.introspectors.lock().await.is_empty(),
1145            "introspectors map should be empty after introspect_server completes, \
1146             regardless of success or failure"
1147        );
1148    }
1149
1150    /// Same `server_id` must resolve to the same introspector lock, so a
1151    /// second `introspect_server` call for that id cannot start
1152    /// `discover_server` until the first releases it.
1153    #[tokio::test]
1154    async fn test_introspector_for_same_id_shares_one_lock() {
1155        let service = GeneratorService::new();
1156        let server_id = ServerId::new("same-id-lock-test");
1157
1158        let handle_a = service.introspector_for(&server_id).await;
1159        let handle_b = service.introspector_for(&server_id).await;
1160
1161        assert!(
1162            Arc::ptr_eq(&handle_a, &handle_b),
1163            "the same server_id must reuse one introspector lock"
1164        );
1165    }
1166
1167    /// Different `server_id`s must resolve to independent introspector
1168    /// locks, so calls for unrelated ids never contend on the same mutex.
1169    #[tokio::test]
1170    async fn test_introspector_for_different_ids_get_independent_locks() {
1171        let service = GeneratorService::new();
1172
1173        let handle_a = service
1174            .introspector_for(&ServerId::new("diff-id-lock-a"))
1175            .await;
1176        let handle_b = service
1177            .introspector_for(&ServerId::new("diff-id-lock-b"))
1178            .await;
1179
1180        assert!(
1181            !Arc::ptr_eq(&handle_a, &handle_b),
1182            "different server_ids must get independent introspector locks"
1183        );
1184    }
1185
1186    /// Two holders of the *same* per-id lock (as returned by
1187    /// `introspector_for` for one `server_id`) must serialize: the second
1188    /// critical section cannot start until the first releases the lock, so
1189    /// total wall time is roughly additive (~2x the hold time).
1190    #[tokio::test]
1191    async fn test_same_id_lock_serializes_concurrent_holders() {
1192        let service = GeneratorService::new();
1193        let server_id = ServerId::new("same-id-timing-test");
1194        let hold_time = std::time::Duration::from_millis(150);
1195        let serialized_threshold = std::time::Duration::from_millis(250);
1196
1197        let handle_a = service.introspector_for(&server_id).await;
1198        let handle_b = service.introspector_for(&server_id).await;
1199
1200        let started = std::time::Instant::now();
1201        tokio::join!(
1202            async {
1203                let _guard = handle_a.lock().await;
1204                tokio::time::sleep(hold_time).await;
1205            },
1206            async {
1207                let _guard = handle_b.lock().await;
1208                tokio::time::sleep(hold_time).await;
1209            },
1210        );
1211        let elapsed = started.elapsed();
1212
1213        assert!(
1214            elapsed >= serialized_threshold,
1215            "holders of the same per-id lock should serialize \
1216             (expected >= {serialized_threshold:?}, i.e. two back-to-back {hold_time:?} \
1217             critical sections); took {elapsed:?}"
1218        );
1219    }
1220
1221    /// Two holders of *different* per-id locks (as returned by
1222    /// `introspector_for` for different `server_id`s) must not serialize:
1223    /// both critical sections run concurrently, so total wall time stays
1224    /// close to a single hold, not double it.
1225    #[tokio::test]
1226    async fn test_different_id_locks_do_not_serialize() {
1227        let service = GeneratorService::new();
1228        let hold_time = std::time::Duration::from_millis(150);
1229        let serialized_threshold = std::time::Duration::from_millis(250);
1230
1231        let handle_a = service
1232            .introspector_for(&ServerId::new("diff-id-timing-a"))
1233            .await;
1234        let handle_b = service
1235            .introspector_for(&ServerId::new("diff-id-timing-b"))
1236            .await;
1237
1238        let started = std::time::Instant::now();
1239        tokio::join!(
1240            async {
1241                let _guard = handle_a.lock().await;
1242                tokio::time::sleep(hold_time).await;
1243            },
1244            async {
1245                let _guard = handle_b.lock().await;
1246                tokio::time::sleep(hold_time).await;
1247            },
1248        );
1249        let elapsed = started.elapsed();
1250
1251        assert!(
1252            elapsed < serialized_threshold,
1253            "holders of different per-id locks should not serialize \
1254             (expected < {serialized_threshold:?}, i.e. close to a single {hold_time:?} hold); \
1255             took {elapsed:?}"
1256        );
1257    }
1258
1259    /// Regression test for the TOCTOU eviction bug (issue #130): eviction
1260    /// must be identity-checked, not just keyed by `server_id`.
1261    ///
1262    /// Simulates three overlapping callers for the same `server_id`:
1263    /// - A and B both call `introspector_for` while an entry already exists,
1264    ///   so (per `test_introspector_for_same_id_shares_one_lock`) they share
1265    ///   the exact same `Arc<Mutex<Introspector>>`.
1266    /// - A finishes first and evicts, removing the shared entry, while B is
1267    ///   still "in flight" (still holding its clone of that same `Arc`).
1268    /// - C then arrives, finds the map empty, and gets a brand-new `Arc` -
1269    ///   distinct from A/B's.
1270    /// - B finally finishes and attempts to evict using its (now stale)
1271    ///   handle. Because eviction is identity-checked via `Arc::ptr_eq`, this
1272    ///   must be a no-op: C's live entry must survive. Only C's own eviction
1273    ///   should remove it.
1274    #[tokio::test]
1275    async fn test_stale_eviction_does_not_remove_unrelated_entry() {
1276        let service = GeneratorService::new();
1277        let server_id = ServerId::new("toctou-abc-test");
1278
1279        // A and B both fetch the handle for the same id before either
1280        // evicts, so they end up sharing one Arc (mirrors the "shares one
1281        // lock" behavior already covered by
1282        // `test_introspector_for_same_id_shares_one_lock`).
1283        let handle_a = service.introspector_for(&server_id).await;
1284        let handle_b = service.introspector_for(&server_id).await;
1285        assert!(
1286            Arc::ptr_eq(&handle_a, &handle_b),
1287            "A and B must share one introspector handle for the same server_id"
1288        );
1289
1290        // A finishes first and evicts. B is still "in flight", holding its
1291        // clone of the now-removed shared Arc.
1292        service.evict_introspector(&server_id, &handle_a).await;
1293        assert!(
1294            service.introspectors.lock().await.is_empty(),
1295            "map should be empty right after A's eviction"
1296        );
1297
1298        // C arrives after A's eviction, finds the map empty, and gets a
1299        // fresh, distinct handle.
1300        let handle_c = service.introspector_for(&server_id).await;
1301        assert!(
1302            !Arc::ptr_eq(&handle_b, &handle_c),
1303            "C must get a handle distinct from A/B's stale one"
1304        );
1305
1306        // B finally finishes and tries to evict using its stale (A/B
1307        // shared) handle. This must be a no-op: C's live entry, keyed by
1308        // the same server_id, must survive because it is a different Arc.
1309        service.evict_introspector(&server_id, &handle_b).await;
1310        let introspectors = service.introspectors.lock().await;
1311        let current = introspectors
1312            .get(&server_id)
1313            .expect("C's entry must survive B's stale eviction attempt");
1314        assert!(
1315            Arc::ptr_eq(current, &handle_c),
1316            "the surviving entry must be C's handle, unaffected by B's stale eviction"
1317        );
1318        drop(introspectors);
1319
1320        // Only C's own eviction removes its entry.
1321        service.evict_introspector(&server_id, &handle_c).await;
1322        assert!(
1323            service.introspectors.lock().await.is_empty(),
1324            "map should be empty after C's own eviction"
1325        );
1326    }
1327
1328    // ========================================================================
1329    // Per-output-directory export locking Tests (issue #169)
1330    //
1331    // Mirrors the `introspector_for` tests above: these exercise the exact
1332    // `Arc<Mutex<()>>` handles and keyed-lock pattern that
1333    // `save_categorized_tools` relies on via `export_lock_for`, without
1334    // driving a real export through the filesystem.
1335    // ========================================================================
1336
1337    /// Same `output_dir` must resolve to the same export lock, so a second
1338    /// concurrent export for that directory cannot proceed until the first
1339    /// releases it.
1340    #[tokio::test]
1341    async fn test_export_lock_for_same_output_dir_shares_one_lock() {
1342        let service = GeneratorService::new();
1343        let output_dir = PathBuf::from("/tmp/same-output-dir-lock-test");
1344
1345        let handle_a = service.export_lock_for(&output_dir).await;
1346        let handle_b = service.export_lock_for(&output_dir).await;
1347
1348        assert!(
1349            Arc::ptr_eq(&handle_a, &handle_b),
1350            "the same output_dir must reuse one export lock"
1351        );
1352    }
1353
1354    /// Different `output_dir`s must resolve to independent export locks, so
1355    /// exports for unrelated directories never contend on the same mutex.
1356    #[tokio::test]
1357    async fn test_export_lock_for_different_output_dirs_get_independent_locks() {
1358        let service = GeneratorService::new();
1359
1360        let handle_a = service
1361            .export_lock_for(&PathBuf::from("/tmp/diff-output-dir-lock-a"))
1362            .await;
1363        let handle_b = service
1364            .export_lock_for(&PathBuf::from("/tmp/diff-output-dir-lock-b"))
1365            .await;
1366
1367        assert!(
1368            !Arc::ptr_eq(&handle_a, &handle_b),
1369            "different output_dirs must get independent export locks"
1370        );
1371    }
1372
1373    /// `evict_export_lock` must be identity-checked, not just keyed by
1374    /// `output_dir`, mirroring `test_stale_eviction_does_not_remove_unrelated_entry`.
1375    #[tokio::test]
1376    async fn test_export_lock_stale_eviction_does_not_remove_unrelated_entry() {
1377        let service = GeneratorService::new();
1378        let output_dir = PathBuf::from("/tmp/toctou-export-lock-test");
1379
1380        let handle_a = service.export_lock_for(&output_dir).await;
1381        let handle_b = service.export_lock_for(&output_dir).await;
1382        assert!(Arc::ptr_eq(&handle_a, &handle_b));
1383
1384        service.evict_export_lock(&output_dir, &handle_a).await;
1385        assert!(service.exports.lock().await.is_empty());
1386
1387        let handle_c = service.export_lock_for(&output_dir).await;
1388        assert!(!Arc::ptr_eq(&handle_b, &handle_c));
1389
1390        // B's stale eviction attempt must be a no-op: C's live entry survives.
1391        service.evict_export_lock(&output_dir, &handle_b).await;
1392        let exports = service.exports.lock().await;
1393        let current = exports
1394            .get(&output_dir)
1395            .expect("C's entry must survive B's stale eviction attempt");
1396        assert!(Arc::ptr_eq(current, &handle_c));
1397        drop(exports);
1398    }
1399
1400    // ========================================================================
1401    // save_categorized_tools Error Tests
1402    // ========================================================================
1403
1404    #[tokio::test]
1405    async fn test_save_categorized_tools_invalid_session() {
1406        let service = GeneratorService::new();
1407
1408        let params = SaveCategorizedToolsParams {
1409            session_id: Uuid::new_v4(), // Random UUID not in state
1410            categorized_tools: vec![],
1411        };
1412
1413        let result = service.save_categorized_tools(Parameters(params)).await;
1414
1415        assert!(result.is_err());
1416        let err = result.unwrap_err();
1417        assert_eq!(err.code, ErrorCode::INVALID_PARAMS); // Invalid params
1418        assert!(err.message.contains("Session not found"));
1419    }
1420
1421    #[tokio::test]
1422    async fn test_save_categorized_tools_tool_mismatch() {
1423        let service = GeneratorService::new();
1424
1425        // Create a pending generation with tool1
1426        let server_id = ServerId::new("test");
1427        let server_info = mcp_execution_introspector::ServerInfo {
1428            id: server_id.clone(),
1429            name: "Test".to_string(),
1430            version: "1.0.0".to_string(),
1431            capabilities: ServerCapabilities {
1432                supports_tools: true,
1433                supports_resources: false,
1434                supports_prompts: false,
1435            },
1436            tools: vec![ToolInfo {
1437                name: ToolName::new("tool1"),
1438                description: "Tool 1".to_string(),
1439                input_schema: serde_json::json!({"type": "object"}),
1440                output_schema: None,
1441            }],
1442        };
1443
1444        let pending = PendingGeneration::new(
1445            server_id,
1446            server_info,
1447            ServerConfig::builder().command("echo".to_string()).build(),
1448            PathBuf::from("/tmp/test"),
1449            &SystemClock,
1450        );
1451
1452        let session_id = service.state.store(pending).await;
1453
1454        // Try to save with tool2 (doesn't exist)
1455        let params = SaveCategorizedToolsParams {
1456            session_id,
1457            categorized_tools: vec![CategorizedTool {
1458                name: "tool2".to_string(), // Mismatch!
1459                category: "test".to_string(),
1460                keywords: "test".to_string(),
1461                short_description: "Test".to_string(),
1462            }],
1463        };
1464
1465        let result = service.save_categorized_tools(Parameters(params)).await;
1466
1467        assert!(result.is_err());
1468        let err = result.unwrap_err();
1469        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1470        assert!(err.message.contains("not found in introspected tools"));
1471    }
1472
1473    #[tokio::test]
1474    async fn test_save_categorized_tools_expired_session() {
1475        use crate::clock::TestClock;
1476        use chrono::Duration;
1477
1478        let service = GeneratorService::new();
1479
1480        // Create an expired pending generation
1481        let server_id = ServerId::new("test");
1482        let server_info = mcp_execution_introspector::ServerInfo {
1483            id: server_id.clone(),
1484            name: "Test".to_string(),
1485            version: "1.0.0".to_string(),
1486            capabilities: ServerCapabilities {
1487                supports_tools: true,
1488                supports_resources: false,
1489                supports_prompts: false,
1490            },
1491            tools: vec![],
1492        };
1493
1494        // Inject a clock fixed an hour in the past so `expires_at` is already
1495        // behind us, instead of rewinding `expires_at` after construction.
1496        let past_clock = TestClock::new(Utc::now() - Duration::hours(1));
1497        let pending = PendingGeneration::new(
1498            server_id,
1499            server_info,
1500            ServerConfig::builder().command("echo".to_string()).build(),
1501            PathBuf::from("/tmp/test"),
1502            &past_clock,
1503        );
1504
1505        let session_id = service.state.store(pending).await;
1506
1507        let params = SaveCategorizedToolsParams {
1508            session_id,
1509            categorized_tools: vec![],
1510        };
1511
1512        let result = service.save_categorized_tools(Parameters(params)).await;
1513
1514        assert!(result.is_err());
1515        let err = result.unwrap_err();
1516        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1517    }
1518
1519    /// Proves `GeneratorService::with_clock` actually drives session expiry end
1520    /// to end through `save_categorized_tools`: a session stored while the
1521    /// shared clock is fresh must become unreachable once that same clock (not
1522    /// the real wall clock) is advanced past the TTL. This exercises the
1523    /// `Arc<dyn Clock>` shared between `GeneratorService` and its
1524    /// `StateManager` (`with_clock` clones the same `Arc` into both).
1525    #[tokio::test]
1526    async fn test_shared_clock_drives_save_categorized_tools_expiry() {
1527        use crate::clock::TestClock;
1528        use chrono::Duration;
1529
1530        let start = Utc::now();
1531        let clock = Arc::new(TestClock::new(start));
1532        let service = GeneratorService::with_clock(Arc::clone(&clock) as Arc<dyn Clock>);
1533
1534        let server_id = ServerId::new("test");
1535        let server_info = mcp_execution_introspector::ServerInfo {
1536            id: server_id.clone(),
1537            name: "Test".to_string(),
1538            version: "1.0.0".to_string(),
1539            capabilities: ServerCapabilities {
1540                supports_tools: true,
1541                supports_resources: false,
1542                supports_prompts: false,
1543            },
1544            tools: vec![],
1545        };
1546
1547        let pending = PendingGeneration::new(
1548            server_id,
1549            server_info,
1550            ServerConfig::builder().command("echo".to_string()).build(),
1551            PathBuf::from("/tmp/test"),
1552            clock.as_ref(),
1553        );
1554
1555        let session_id = service.state.store(pending).await;
1556
1557        // Advance the service's own shared clock, not the real wall clock, past the TTL.
1558        clock.advance(
1559            Duration::minutes(PendingGeneration::DEFAULT_TIMEOUT_MINUTES) + Duration::seconds(1),
1560        );
1561
1562        let params = SaveCategorizedToolsParams {
1563            session_id,
1564            categorized_tools: vec![],
1565        };
1566
1567        let result = service.save_categorized_tools(Parameters(params)).await;
1568
1569        assert!(result.is_err());
1570        let err = result.unwrap_err();
1571        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1572    }
1573
1574    // ========================================================================
1575    // list_generated_servers Tests
1576    // ========================================================================
1577
1578    #[tokio::test]
1579    async fn test_list_generated_servers_nonexistent_dir() {
1580        let service = GeneratorService::new();
1581
1582        let params = ListGeneratedServersParams {
1583            base_dir: Some("/nonexistent/path/that/does/not/exist".to_string()),
1584        };
1585
1586        let result = service.list_generated_servers(Parameters(params)).await;
1587
1588        assert!(result.is_ok());
1589        let content = result.unwrap();
1590        let text_content = content.content[0].as_text().unwrap();
1591        let parsed: ListGeneratedServersResult = serde_json::from_str(&text_content.text).unwrap();
1592
1593        assert_eq!(parsed.total_servers, 0);
1594        assert_eq!(parsed.servers.len(), 0);
1595    }
1596
1597    #[tokio::test]
1598    async fn test_list_generated_servers_default_dir() {
1599        let service = GeneratorService::new();
1600
1601        let params = ListGeneratedServersParams { base_dir: None };
1602
1603        let result = service.list_generated_servers(Parameters(params)).await;
1604
1605        // Should succeed even if directory doesn't exist
1606        assert!(result.is_ok());
1607    }
1608
1609    // ========================================================================
1610    // generate_skill Error Tests
1611    // ========================================================================
1612
1613    #[tokio::test]
1614    async fn test_generate_skill_invalid_server_id_uppercase() {
1615        let service = GeneratorService::new();
1616
1617        let params = GenerateSkillParams {
1618            server_id: "GitHub".to_string(), // Invalid: uppercase
1619            skill_name: None,
1620            use_case_hints: None,
1621            servers_dir: None,
1622        };
1623
1624        let result = service.generate_skill(Parameters(params)).await;
1625
1626        assert!(result.is_err());
1627        let err = result.unwrap_err();
1628        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1629        assert!(err.message.contains("lowercase"));
1630    }
1631
1632    #[tokio::test]
1633    async fn test_generate_skill_invalid_server_id_special_chars() {
1634        let service = GeneratorService::new();
1635
1636        let params = GenerateSkillParams {
1637            server_id: "git@hub".to_string(), // Invalid: special chars
1638            skill_name: None,
1639            use_case_hints: None,
1640            servers_dir: None,
1641        };
1642
1643        let result = service.generate_skill(Parameters(params)).await;
1644
1645        assert!(result.is_err());
1646        let err = result.unwrap_err();
1647        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1648    }
1649
1650    #[tokio::test]
1651    async fn test_generate_skill_server_directory_not_found() {
1652        let service = GeneratorService::new();
1653
1654        let params = GenerateSkillParams {
1655            server_id: "nonexistent-server".to_string(),
1656            skill_name: None,
1657            use_case_hints: None,
1658            servers_dir: Some(PathBuf::from("/nonexistent/path")),
1659        };
1660
1661        let result = service.generate_skill(Parameters(params)).await;
1662
1663        assert!(result.is_err());
1664        let err = result.unwrap_err();
1665        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1666        assert!(err.message.contains("not found"));
1667    }
1668
1669    #[tokio::test]
1670    async fn test_generate_skill_missing_metadata_sidecar() {
1671        use tempfile::TempDir;
1672
1673        let service = GeneratorService::new();
1674        let temp_dir = TempDir::new().unwrap();
1675        let base_dir = temp_dir.path().to_path_buf();
1676
1677        // Create server directory but no `_meta.json` sidecar (e.g. a directory
1678        // generated by a pre-#141 version, or never generated at all).
1679        let target_dir = base_dir.join("test-server");
1680        tokio::fs::create_dir_all(&target_dir).await.unwrap();
1681
1682        let params = GenerateSkillParams {
1683            server_id: "test-server".to_string(),
1684            skill_name: None,
1685            use_case_hints: None,
1686            servers_dir: Some(base_dir),
1687        };
1688
1689        let result = service.generate_skill(Parameters(params)).await;
1690
1691        assert!(result.is_err());
1692        let err = result.unwrap_err();
1693        assert_eq!(
1694            err.code,
1695            ErrorCode::INVALID_PARAMS,
1696            "a missing sidecar is the same 'not generated' caller situation as a missing \
1697             server directory, and must be reported the same way"
1698        );
1699        assert!(err.message.contains("Failed to scan tools directory"));
1700    }
1701
1702    #[tokio::test]
1703    async fn test_generate_skill_stale_metadata_missing_ts_file() {
1704        use mcp_execution_core::metadata::{
1705            METADATA_FILE_NAME, METADATA_SCHEMA_VERSION, ParameterMetadata, ServerMetadata,
1706            ToolMetadata as SidecarToolMetadata,
1707        };
1708        use tempfile::TempDir;
1709
1710        let service = GeneratorService::new();
1711        let temp_dir = TempDir::new().unwrap();
1712        let base_dir = temp_dir.path().to_path_buf();
1713
1714        // Sidecar references a tool whose `.ts` file was never written (or was
1715        // deleted) — the drift `StaleMetadata` (issues #154/#155) exists to
1716        // catch, routed through the `generate_skill` MCP tool this time.
1717        let target_dir = base_dir.join("test-server");
1718        tokio::fs::create_dir_all(&target_dir).await.unwrap();
1719        let meta = ServerMetadata {
1720            schema_version: METADATA_SCHEMA_VERSION,
1721            server_id: "test-server".to_string(),
1722            server_name: "Test Server".to_string(),
1723            server_version: "1.0.0".to_string(),
1724            tools: vec![SidecarToolMetadata {
1725                name: "create_issue".to_string(),
1726                typescript_name: "createIssue".to_string(),
1727                category: None,
1728                keywords: vec![],
1729                description: None,
1730                parameters: vec![ParameterMetadata {
1731                    name: "title".to_string(),
1732                    typescript_type: "string".to_string(),
1733                    required: true,
1734                    description: None,
1735                }],
1736            }],
1737        };
1738        let content = serde_json::to_string_pretty(&meta).unwrap();
1739        tokio::fs::write(target_dir.join(METADATA_FILE_NAME), content)
1740            .await
1741            .unwrap();
1742        // Deliberately do not write `createIssue.ts`.
1743
1744        let params = GenerateSkillParams {
1745            server_id: "test-server".to_string(),
1746            skill_name: None,
1747            use_case_hints: None,
1748            servers_dir: Some(base_dir),
1749        };
1750
1751        let result = service.generate_skill(Parameters(params)).await;
1752
1753        assert!(result.is_err());
1754        let err = result.unwrap_err();
1755        assert_eq!(
1756            err.code,
1757            ErrorCode::INVALID_PARAMS,
1758            "stale metadata is the same 'not generated / drifted directory' caller situation \
1759             as a missing sidecar, and must be reported the same way"
1760        );
1761        assert!(err.message.contains("Failed to scan tools directory"));
1762        assert!(err.message.contains("create_issue"));
1763    }
1764
1765    #[tokio::test]
1766    async fn test_generate_skill_reports_orphan_ts_file_as_warning() {
1767        // Issue #161: a `.ts` file on disk with no matching `_meta.json` entry
1768        // is non-fatal, but must be surfaced in the structured JSON-RPC
1769        // response's `warnings` field, not just in server-side tracing output.
1770        use mcp_execution_core::metadata::{
1771            METADATA_FILE_NAME, METADATA_SCHEMA_VERSION, ParameterMetadata, ServerMetadata,
1772            ToolMetadata as SidecarToolMetadata,
1773        };
1774        use mcp_execution_skill::GenerateSkillResult;
1775        use tempfile::TempDir;
1776
1777        let service = GeneratorService::new();
1778        let temp_dir = TempDir::new().unwrap();
1779        let base_dir = temp_dir.path().to_path_buf();
1780
1781        let target_dir = base_dir.join("test-server");
1782        tokio::fs::create_dir_all(&target_dir).await.unwrap();
1783        let meta = ServerMetadata {
1784            schema_version: METADATA_SCHEMA_VERSION,
1785            server_id: "test-server".to_string(),
1786            server_name: "Test Server".to_string(),
1787            server_version: "1.0.0".to_string(),
1788            tools: vec![SidecarToolMetadata {
1789                name: "create_issue".to_string(),
1790                typescript_name: "createIssue".to_string(),
1791                category: None,
1792                keywords: vec![],
1793                description: None,
1794                parameters: vec![ParameterMetadata {
1795                    name: "title".to_string(),
1796                    typescript_type: "string".to_string(),
1797                    required: true,
1798                    description: None,
1799                }],
1800            }],
1801        };
1802        let content = serde_json::to_string_pretty(&meta).unwrap();
1803        tokio::fs::write(target_dir.join(METADATA_FILE_NAME), content)
1804            .await
1805            .unwrap();
1806        tokio::fs::write(target_dir.join("createIssue.ts"), "export {}")
1807            .await
1808            .unwrap();
1809        // Left over on disk with no sidecar entry — must not be fatal.
1810        tokio::fs::write(target_dir.join("orphanTool.ts"), "export {}")
1811            .await
1812            .unwrap();
1813
1814        let params = GenerateSkillParams {
1815            server_id: "test-server".to_string(),
1816            skill_name: None,
1817            use_case_hints: None,
1818            servers_dir: Some(base_dir),
1819        };
1820
1821        let result = service.generate_skill(Parameters(params)).await;
1822
1823        assert!(
1824            result.is_ok(),
1825            "an orphaned .ts file must not fail the call"
1826        );
1827        let content = result.unwrap();
1828        let text_content = content.content[0].as_text().unwrap();
1829        let parsed: GenerateSkillResult = serde_json::from_str(&text_content.text).unwrap();
1830
1831        assert_eq!(
1832            parsed.warnings.len(),
1833            1,
1834            "the orphaned .ts file must be surfaced as a warning"
1835        );
1836        assert!(
1837            parsed.warnings[0].contains("orphanTool.ts"),
1838            "warning must name the excluded file: {:?}",
1839            parsed.warnings[0]
1840        );
1841    }
1842
1843    // ========================================================================
1844    // save_skill Error Tests
1845    // ========================================================================
1846
1847    #[tokio::test]
1848    async fn test_save_skill_invalid_server_id() {
1849        let service = GeneratorService::new();
1850
1851        let params = SaveSkillParams {
1852            server_id: "Invalid_Server".to_string(), // Invalid: uppercase and underscore
1853            content: "---\nname: test\ndescription: test\n---\n# Test".to_string(),
1854            output_path: None,
1855            overwrite: false,
1856        };
1857
1858        let result = service.save_skill(Parameters(params)).await;
1859
1860        assert!(result.is_err());
1861        let err = result.unwrap_err();
1862        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1863        assert!(err.message.contains("lowercase"));
1864    }
1865
1866    #[tokio::test]
1867    async fn test_save_skill_missing_yaml_frontmatter() {
1868        let service = GeneratorService::new();
1869
1870        let params = SaveSkillParams {
1871            server_id: "test".to_string(),
1872            content: "# Test Skill\n\nNo YAML frontmatter here.".to_string(),
1873            output_path: None,
1874            overwrite: false,
1875        };
1876
1877        let result = service.save_skill(Parameters(params)).await;
1878
1879        assert!(result.is_err());
1880        let err = result.unwrap_err();
1881        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1882        assert!(err.message.contains("YAML frontmatter"));
1883    }
1884
1885    #[tokio::test]
1886    async fn test_save_skill_invalid_frontmatter_no_name() {
1887        let service = GeneratorService::new();
1888
1889        let params = SaveSkillParams {
1890            server_id: "test".to_string(),
1891            content: "---\ndescription: test\n---\n# Test".to_string(),
1892            output_path: None,
1893            overwrite: false,
1894        };
1895
1896        let result = service.save_skill(Parameters(params)).await;
1897
1898        assert!(result.is_err());
1899        let err = result.unwrap_err();
1900        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1901        assert!(err.message.contains("Invalid SKILL.md format"));
1902    }
1903
1904    #[tokio::test]
1905    async fn test_save_skill_invalid_frontmatter_no_description() {
1906        let service = GeneratorService::new();
1907
1908        let params = SaveSkillParams {
1909            server_id: "test".to_string(),
1910            content: "---\nname: test-skill\n---\n# Test".to_string(),
1911            output_path: None,
1912            overwrite: false,
1913        };
1914
1915        let result = service.save_skill(Parameters(params)).await;
1916
1917        assert!(result.is_err());
1918        let err = result.unwrap_err();
1919        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1920        assert!(err.message.contains("Invalid SKILL.md format"));
1921    }
1922
1923    #[tokio::test]
1924    async fn test_save_skill_file_exists_no_overwrite() {
1925        use tempfile::TempDir;
1926
1927        let service = GeneratorService::new();
1928        let temp_dir = TempDir::new().unwrap();
1929        let output_path = temp_dir.path().join("SKILL.md");
1930
1931        // Create existing file
1932        tokio::fs::write(&output_path, "existing content")
1933            .await
1934            .unwrap();
1935
1936        let params = SaveSkillParams {
1937            server_id: "test".to_string(),
1938            content: "---\nname: test\ndescription: test\n---\n# Test".to_string(),
1939            output_path: Some(output_path),
1940            overwrite: false,
1941        };
1942
1943        let result = service.save_skill(Parameters(params)).await;
1944
1945        assert!(result.is_err());
1946        let err = result.unwrap_err();
1947        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
1948        assert!(err.message.contains("already exists"));
1949        assert!(err.message.contains("overwrite=true"));
1950    }
1951
1952    #[tokio::test]
1953    async fn test_save_skill_file_exists_with_overwrite() {
1954        use tempfile::TempDir;
1955
1956        let service = GeneratorService::new();
1957        let temp_dir = TempDir::new().unwrap();
1958        let output_path = temp_dir.path().join("SKILL.md");
1959
1960        // Create existing file
1961        tokio::fs::write(&output_path, "existing content")
1962            .await
1963            .unwrap();
1964
1965        let params = SaveSkillParams {
1966            server_id: "test".to_string(),
1967            content: "---\nname: test\ndescription: test skill\n---\n# Test".to_string(),
1968            output_path: Some(output_path.clone()),
1969            overwrite: true,
1970        };
1971
1972        let result = service.save_skill(Parameters(params)).await;
1973
1974        assert!(result.is_ok());
1975        let content = result.unwrap();
1976        let text = content.content[0].as_text().unwrap();
1977        let parsed: SaveSkillResult = serde_json::from_str(&text.text).unwrap();
1978
1979        assert!(parsed.success);
1980        assert!(parsed.overwritten);
1981        assert_eq!(parsed.metadata.name, "test");
1982        assert_eq!(parsed.metadata.description, "test skill");
1983    }
1984
1985    #[tokio::test]
1986    async fn test_save_skill_valid_content() {
1987        use tempfile::TempDir;
1988
1989        let service = GeneratorService::new();
1990        let temp_dir = TempDir::new().unwrap();
1991        let output_path = temp_dir.path().join("SKILL.md");
1992
1993        let params = SaveSkillParams {
1994            server_id: "test".to_string(),
1995            content: "---\nname: test-skill\ndescription: A test skill\n---\n\n# Test Skill\n\n## Section 1\n\nContent here.".to_string(),
1996            output_path: Some(output_path.clone()),
1997            overwrite: false,
1998        };
1999
2000        let result = service.save_skill(Parameters(params)).await;
2001
2002        assert!(result.is_ok());
2003        let content = result.unwrap();
2004        let text = content.content[0].as_text().unwrap();
2005        let parsed: SaveSkillResult = serde_json::from_str(&text.text).unwrap();
2006
2007        assert!(parsed.success);
2008        assert!(!parsed.overwritten);
2009        assert_eq!(parsed.metadata.name, "test-skill");
2010        assert_eq!(parsed.metadata.description, "A test skill");
2011        assert!(parsed.metadata.section_count >= 1);
2012        assert!(parsed.metadata.word_count > 0);
2013
2014        // Verify file was written
2015        assert!(output_path.exists());
2016    }
2017}