diff --git a/shellmcp/templates/server.py.j2 b/shellmcp/templates/server.py.j2 index 4a6858c..55b1d70 100644 --- a/shellmcp/templates/server.py.j2 +++ b/shellmcp/templates/server.py.j2 @@ -1,5 +1,7 @@ """Generated FastMCP server from YAML configuration.""" +import argparse +import fnmatch import os import subprocess import tempfile @@ -62,6 +64,59 @@ def render_template(template_str: str, **kwargs) -> str: except Exception as e: raise ValueError(f"Template rendering error: {e}") +def should_include_tool(tool_name: str, include_patterns: List[str]) -> bool: + """Check if a tool should be included based on the include patterns.""" + if not include_patterns: + return True # Include all tools if no patterns specified + + # Check if tool name matches any of the patterns + for pattern in include_patterns: + if fnmatch.fnmatch(tool_name, pattern): + return True + return False + +def parse_include_patterns(include_arg: Optional[str]) -> List[str]: + """Parse include patterns from command line argument.""" + if not include_arg: + return [] + + # Split by comma and strip whitespace + patterns = [pattern.strip() for pattern in include_arg.split(',')] + return [pattern for pattern in patterns if pattern] # Remove empty patterns + +def should_include_tool(tool_name: str, include_patterns: List[str]) -> bool: + """Check if a tool should be included based on the include patterns.""" + if not include_patterns: + return True # Include all tools if no patterns specified + + # Check if tool name matches any of the patterns + for pattern in include_patterns: + if fnmatch.fnmatch(tool_name, pattern): + return True + return False + +def parse_arguments(): + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="{{ config.server.desc }}", + formatter_class=argparse.RawDescriptionHelpFormatter + ) + + parser.add_argument( + '--include', + type=str, + help='Comma-separated list of tool name patterns to include. ' + 'Uses shell-style wildcards (*, ?, [seq], [!seq]). ' + 'Example: --include=get*,search* would include all tools starting with "get" or "search". ' + 'If not specified, all tools are included.' + ) + + return parser.parse_args() + +# Parse command line arguments +args = parse_arguments() +include_patterns = parse_include_patterns(args.include) + # Initialize FastMCP server mcp = FastMCP(name="{{ config.server.name }}") @@ -91,73 +146,75 @@ os.environ["{{ key }}"] = "{{ value }}" {% endfor %} {% set param_str = (params_without_defaults + params_with_defaults)|join(", ") %} -@mcp.tool() -def {{ func_name }}({{ param_str }}) -> Dict[str, Any]: - """ - {{ tool.desc }} +# Register tool only if it matches include patterns +if should_include_tool("{{ tool_name }}", include_patterns): + @mcp.tool() + def {{ func_name }}({{ param_str }}) -> Dict[str, Any]: + """ + {{ tool.desc }} {% if tool.help_cmd and help_outputs.get(tool_name) %} - - Help: -{{ help_outputs[tool_name].stdout|indent(4, first=True) }} + + Help: +{{ help_outputs[tool_name].stdout|indent(8, first=True) }} {% elif tool.help_cmd %} - - Help: Run `{{ tool.help_cmd }}` for more information. + + Help: Run `{{ tool.help_cmd }}` for more information. {% endif %} {% if resolved_args %} - - Parameters: + + Parameters: {% for arg in resolved_args %} - - {{ arg.name }} ({{ arg.type }}): {{ arg.help }} + - {{ arg.name }} ({{ arg.type }}): {{ arg.help }} {% if arg.default is not none %} - Default: {{ arg.default }} + Default: {{ arg.default }} {% endif %} {% if arg.choices %} - Allowed values: {{ arg.choices|join(', ') }} + Allowed values: {{ arg.choices|join(', ') }} {% endif %} {% if arg.pattern %} - Pattern: {{ arg.pattern }} + Pattern: {{ arg.pattern }} {% endif %} {% endfor %} {% endif %} - - Returns: - Dict[str, Any]: Command execution result with 'success', 'stdout', 'stderr', and 'returncode' fields. - """ - try: + + Returns: + Dict[str, Any]: Command execution result with 'success', 'stdout', 'stderr', and 'returncode' fields. + """ + try: {% for arg in resolved_args %} {% if arg.pattern %} - # Validate {{ arg.name }} pattern - import re - if not re.match(r"{{ arg.pattern }}", str({{ arg.name }})): - raise ValueError(f"Invalid {{ arg.name }}: must match pattern {{ arg.pattern }}") + # Validate {{ arg.name }} pattern + import re + if not re.match(r"{{ arg.pattern }}", str({{ arg.name }})): + raise ValueError(f"Invalid {{ arg.name }}: must match pattern {{ arg.pattern }}") {% endif %} {% if arg.choices %} - # Validate {{ arg.name }} choices - if {{ arg.name }} not in {{ arg.choices }}: - raise ValueError(f"Invalid {{ arg.name }}: must be one of {{ arg.choices }}") + # Validate {{ arg.name }} choices + if {{ arg.name }} not in {{ arg.choices }}: + raise ValueError(f"Invalid {{ arg.name }}: must be one of {{ arg.choices }}") {% endif %} {% endfor %} - - # Render command template - cmd = render_template("""{{ tool.cmd|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) - - # Execute command - env_vars = {} + + # Render command template + cmd = render_template("""{{ tool.cmd|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) + + # Execute command + env_vars = {} {% if tool.env %} {% for key, value in tool.env.items() %} - env_vars["{{ key }}"] = "{{ value }}" + env_vars["{{ key }}"] = "{{ value }}" {% endfor %} {% endif %} - result = execute_command(cmd, env_vars) - - return result - except Exception as e: - return { - "success": False, - "stdout": "", - "stderr": f"Error in {{ tool_name }}: {str(e)}", - "returncode": -1 - } + result = execute_command(cmd, env_vars) + + return result + except Exception as e: + return { + "success": False, + "stdout": "", + "stderr": f"Error in {{ tool_name }}: {str(e)}", + "returncode": -1 + } {% endfor %} @@ -165,6 +222,8 @@ def {{ func_name }}({{ param_str }}) -> Dict[str, Any]: # Resource handlers {% for resource_name, resource in config.resources.items() %} {% set func_name = resource_name.lower().replace('-', '_') %} +# Register resource only if it matches include patterns +if should_include_tool("{{ resource_name }}", include_patterns): {% set resolved_args = config.get_resolved_resource_arguments(resource_name) %} {% set params_with_defaults = [] %} {% set params_without_defaults = [] %} @@ -179,76 +238,76 @@ def {{ func_name }}({{ param_str }}) -> Dict[str, Any]: {% endfor %} {% set param_str = (params_without_defaults + params_with_defaults)|join(", ") %} -@mcp.resource("{{ resource.uri|replace('{{ ', '{')|replace(' }}', '}')|replace('{{', '{')|replace('}}', '}') }}") -def {{ func_name }}({{ param_str }}) -> str: - """ - {{ resource.description or resource.name }} + @mcp.resource("{{ resource.uri|replace('{{ ', '{')|replace(' }}', '}')|replace('{{', '{')|replace('}}', '}') }}") + def {{ func_name }}({{ param_str }}) -> str: + """ + {{ resource.description or resource.name }} {% if resolved_args %} - - Parameters: + + Parameters: {% for arg in resolved_args %} - - {{ arg.name }} ({{ arg.type }}): {{ arg.help }} + - {{ arg.name }} ({{ arg.type }}): {{ arg.help }} {% if arg.default is not none %} - Default: {{ arg.default }} + Default: {{ arg.default }} {% endif %} {% if arg.choices %} - Allowed values: {{ arg.choices|join(', ') }} + Allowed values: {{ arg.choices|join(', ') }} {% endif %} {% if arg.pattern %} - Pattern: {{ arg.pattern }} + Pattern: {{ arg.pattern }} {% endif %} {% endfor %} {% endif %} - - Returns: - str: The resource content. - """ - try: + + Returns: + str: The resource content. + """ + try: {% for arg in resolved_args %} {% if arg.pattern %} - # Validate {{ arg.name }} pattern - import re - if not re.match(r"{{ arg.pattern }}", str({{ arg.name }})): - raise ValueError(f"Invalid {{ arg.name }}: must match pattern {{ arg.pattern }}") + # Validate {{ arg.name }} pattern + import re + if not re.match(r"{{ arg.pattern }}", str({{ arg.name }})): + raise ValueError(f"Invalid {{ arg.name }}: must match pattern {{ arg.pattern }}") {% endif %} {% if arg.choices %} - # Validate {{ arg.name }} choices - if {{ arg.name }} not in {{ arg.choices }}: - raise ValueError(f"Invalid {{ arg.name }}: must be one of {{ arg.choices }}") + # Validate {{ arg.name }} choices + if {{ arg.name }} not in {{ arg.choices }}: + raise ValueError(f"Invalid {{ arg.name }}: must be one of {{ arg.choices }}") {% endif %} {% endfor %} - + {% if resource.text %} - # Use direct text content - content = render_template("""{{ resource.text|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) + # Use direct text content + content = render_template("""{{ resource.text|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) {% elif resource.file %} - # Read from file - file_path = render_template("""{{ resource.file|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) - try: - with open(file_path, 'r', encoding='utf-8') as f: - content = f.read() - except FileNotFoundError: - raise ValueError(f"Resource file not found: {file_path}") - except Exception as e: - raise ValueError(f"Error reading resource file {file_path}: {str(e)}") + # Read from file + file_path = render_template("""{{ resource.file|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) + try: + with open(file_path, 'r', encoding='utf-8') as f: + content = f.read() + except FileNotFoundError: + raise ValueError(f"Resource file not found: {file_path}") + except Exception as e: + raise ValueError(f"Error reading resource file {file_path}: {str(e)}") {% else %} - # Execute command - cmd = render_template("""{{ resource.cmd|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) - env_vars = {} + # Execute command + cmd = render_template("""{{ resource.cmd|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) + env_vars = {} {% if resource.env %} {% for key, value in resource.env.items() %} - env_vars["{{ key }}"] = "{{ value }}" + env_vars["{{ key }}"] = "{{ value }}" {% endfor %} {% endif %} - result = execute_command(cmd, env_vars) - if not result["success"]: - raise ValueError(f"Command failed: {result['stderr']}") - content = result["stdout"] + result = execute_command(cmd, env_vars) + if not result["success"]: + raise ValueError(f"Command failed: {result['stderr']}") + content = result["stdout"] {% endif %} - - return content - except Exception as e: - raise ValueError(f"Error in {{ resource_name }}: {str(e)}") + + return content + except Exception as e: + raise ValueError(f"Error in {{ resource_name }}: {str(e)}") {% endfor %} {% endif %} @@ -257,6 +316,8 @@ def {{ func_name }}({{ param_str }}) -> str: # Prompt handlers {% for prompt_name, prompt in config.prompts.items() %} {% set func_name = prompt_name.lower().replace('-', '_') %} +# Register prompt only if it matches include patterns +if should_include_tool("{{ prompt_name }}", include_patterns): {% set resolved_args = config.get_resolved_prompt_arguments(prompt_name) %} {% set params_with_defaults = [] %} {% set params_without_defaults = [] %} @@ -271,76 +332,76 @@ def {{ func_name }}({{ param_str }}) -> str: {% endfor %} {% set param_str = (params_without_defaults + params_with_defaults)|join(", ") %} -@mcp.prompt() -def {{ func_name }}({{ param_str }}) -> str: - """ - {{ prompt.description or prompt.name }} + @mcp.prompt() + def {{ func_name }}({{ param_str }}) -> str: + """ + {{ prompt.description or prompt.name }} {% if resolved_args %} - - Parameters: + + Parameters: {% for arg in resolved_args %} - - {{ arg.name }} ({{ arg.type }}): {{ arg.help }} + - {{ arg.name }} ({{ arg.type }}): {{ arg.help }} {% if arg.default is not none %} - Default: {{ arg.default }} + Default: {{ arg.default }} {% endif %} {% if arg.choices %} - Allowed values: {{ arg.choices|join(', ') }} + Allowed values: {{ arg.choices|join(', ') }} {% endif %} {% if arg.pattern %} - Pattern: {{ arg.pattern }} + Pattern: {{ arg.pattern }} {% endif %} {% endfor %} {% endif %} - - Returns: - str: The generated prompt content. - """ - try: + + Returns: + str: The generated prompt content. + """ + try: {% for arg in resolved_args %} {% if arg.pattern %} - # Validate {{ arg.name }} pattern - import re - if not re.match(r"{{ arg.pattern }}", str({{ arg.name }})): - raise ValueError(f"Invalid {{ arg.name }}: must match pattern {{ arg.pattern }}") + # Validate {{ arg.name }} pattern + import re + if not re.match(r"{{ arg.pattern }}", str({{ arg.name }})): + raise ValueError(f"Invalid {{ arg.name }}: must match pattern {{ arg.pattern }}") {% endif %} {% if arg.choices %} - # Validate {{ arg.name }} choices - if {{ arg.name }} not in {{ arg.choices }}: - raise ValueError(f"Invalid {{ arg.name }}: must be one of {{ arg.choices }}") + # Validate {{ arg.name }} choices + if {{ arg.name }} not in {{ arg.choices }}: + raise ValueError(f"Invalid {{ arg.name }}: must be one of {{ arg.choices }}") {% endif %} {% endfor %} - + {% if prompt.template %} - # Use direct template content - content = render_template("""{{ prompt.template|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) + # Use direct template content + content = render_template("""{{ prompt.template|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) {% elif prompt.file %} - # Read from file - file_path = render_template("""{{ prompt.file|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) - try: - with open(file_path, 'r', encoding='utf-8') as f: - content = f.read() - except FileNotFoundError: - raise ValueError(f"Prompt file not found: {file_path}") - except Exception as e: - raise ValueError(f"Error reading prompt file {file_path}: {str(e)}") + # Read from file + file_path = render_template("""{{ prompt.file|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) + try: + with open(file_path, 'r', encoding='utf-8') as f: + content = f.read() + except FileNotFoundError: + raise ValueError(f"Prompt file not found: {file_path}") + except Exception as e: + raise ValueError(f"Error reading prompt file {file_path}: {str(e)}") {% else %} - # Execute command - cmd = render_template("""{{ prompt.cmd|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) - env_vars = {} + # Execute command + cmd = render_template("""{{ prompt.cmd|escape_double_quotes }}""", {% for arg in resolved_args %}{{ arg.name }}={{ arg.name }}{% if not loop.last %}, {% endif %}{% endfor %}) + env_vars = {} {% if prompt.env %} {% for key, value in prompt.env.items() %} - env_vars["{{ key }}"] = "{{ value }}" + env_vars["{{ key }}"] = "{{ value }}" {% endfor %} {% endif %} - result = execute_command(cmd, env_vars) - if not result["success"]: - raise ValueError(f"Command failed: {result['stderr']}") - content = result["stdout"] + result = execute_command(cmd, env_vars) + if not result["success"]: + raise ValueError(f"Command failed: {result['stderr']}") + content = result["stdout"] {% endif %} - - return content - except Exception as e: - raise ValueError(f"Error in {{ prompt_name }}: {str(e)}") + + return content + except Exception as e: + raise ValueError(f"Error in {{ prompt_name }}: {str(e)}") {% endfor %} {% endif %} @@ -348,4 +409,32 @@ def {{ func_name }}({{ param_str }}) -> str: if __name__ == "__main__": print(f"Starting {SERVER_NAME} v{SERVER_VERSION}") print(f"Description: {SERVER_DESC}") + + # Show which tools are included + if include_patterns: + print(f"Include patterns: {', '.join(include_patterns)}") + included_tools = [] + {% for tool_name, tool in config.tools.items() %} + if should_include_tool("{{ tool_name }}", include_patterns): + included_tools.append("{{ tool_name }}") + {% endfor %} + {% if config.resources %} + {% for resource_name, resource in config.resources.items() %} + if should_include_tool("{{ resource_name }}", include_patterns): + included_tools.append("{{ resource_name }}") + {% endfor %} + {% endif %} + {% if config.prompts %} + {% for prompt_name, prompt in config.prompts.items() %} + if should_include_tool("{{ prompt_name }}", include_patterns): + included_tools.append("{{ prompt_name }}") + {% endfor %} + {% endif %} + if included_tools: + print(f"Included tools: {', '.join(included_tools)}") + else: + print("No tools match the include patterns") + else: + print("All tools included (no filtering)") + mcp.run() \ No newline at end of file