Skip to content

Commit c96778d

Browse files
committed
fix: tool names use only allowed characters
1 parent e8dbc12 commit c96778d

5 files changed

Lines changed: 216 additions & 14 deletions

File tree

click_mcp/scanner.py

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,33 @@
99

1010
from .decorator import get_mcp_metadata
1111

12-
1312
# Dictionary to store positional arguments for each tool
1413
_tool_positional_args: Dict[str, List[str]] = {}
1514

15+
# Dictionary to store mapping between sanitized tool names and original paths
16+
_original_paths: Dict[str, str] = {}
17+
18+
19+
def sanitize_tool_name(name: str) -> str:
20+
"""
21+
Sanitize a tool name to conform to the regex pattern [a-zA-Z][a-zA-Z0-9_]*
22+
"""
23+
import re
24+
25+
sanitized = name.replace(".", "_")
26+
27+
if sanitized and not re.match(r"^[a-zA-Z]", sanitized):
28+
sanitized = "tool_" + sanitized
29+
30+
sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", sanitized)
31+
32+
return sanitized
33+
34+
35+
def get_original_path(sanitized_name: str) -> str:
36+
"""Get the original command path from a sanitized tool name."""
37+
return _original_paths.get(sanitized_name, sanitized_name)
38+
1639

1740
def scan_click_command(command: click.Group, parent_path: str = "") -> List[types.Tool]:
1841
"""
@@ -47,6 +70,7 @@ def scan_click_command(command: click.Group, parent_path: str = "") -> List[type
4770
cmd = command.get_command(ctx, name)
4871
if isinstance(cmd, click.Group):
4972
group_name = metadata.get("name", name)
73+
# Use underscore for path separator in sanitized paths
5074
group_path = (
5175
f"{parent_path}{group_name}." if parent_path else f"{group_name}."
5276
)
@@ -55,8 +79,13 @@ def scan_click_command(command: click.Group, parent_path: str = "") -> List[type
5579
# Handle command
5680
cmd = command.get_command(ctx, name)
5781
if cmd is not None:
82+
# Sanitize the command path to conform to the regex pattern
83+
sanitized_cmd_path = sanitize_tool_name(cmd_path)
84+
# Store the mapping between sanitized name and original path
85+
_original_paths[sanitized_cmd_path] = cmd_path
86+
5887
tool, positional_args = _convert_command_to_tool(
59-
cmd, cmd_info, cmd_path
88+
cmd, cmd_info, sanitized_cmd_path
6089
)
6190
tools.append(tool)
6291
# Store positional arguments in the global dictionary

click_mcp/server.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,11 @@ def _execute_command(
6969
self, tool_name: str, parameters: Dict[str, Any]
7070
) -> Dict[str, Any]:
7171
"""Execute a Click command and return its output."""
72-
command = self._find_command(self.cli_group, tool_name.split("."))
72+
# Get the original command path
73+
from .scanner import get_original_path
74+
75+
original_path = get_original_path(tool_name)
76+
command = self._find_command(self.cli_group, original_path.split("."))
7377
args = self._build_command_args(command, tool_name, parameters)
7478

7579
# Capture and return command output

tests/advanced_cli.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,13 @@ def get(key):
5252
click.echo(f"Value for {key}: example_value")
5353

5454

55+
@config.command()
56+
@click.argument("key")
57+
def get_value(key):
58+
"""Get a configuration value with underscore in command name."""
59+
click.echo(f"Value for {key} (from get_value): example_value")
60+
61+
5562
@cli.command()
5663
@click.option(
5764
"--format",

tests/test_mcp_server.py

Lines changed: 34 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -75,10 +75,10 @@ async def test_basic_server_tools(basic_mcp_session):
7575
assert "required" in greet_tool.inputSchema
7676
assert "name" in greet_tool.inputSchema["required"]
7777

78-
# Find the users.list command
78+
# Find the users_list command
7979
users_list_tool = None
8080
for tool in tools:
81-
if "users.list" in tool.name.lower():
81+
if "users_list" in tool.name.lower():
8282
users_list_tool = tool
8383
break
8484

@@ -136,15 +136,15 @@ async def test_invoke_greet_command(basic_mcp_session):
136136

137137
@pytest.mark.anyio
138138
async def test_invoke_users_list_command(basic_mcp_session):
139-
"""Test invoking the users.list command."""
139+
"""Test invoking the users_list command."""
140140
# Get tools
141141
result = await basic_mcp_session.list_tools()
142142
tools = result.tools
143143

144-
# Find the users.list command
144+
# Find the users_list command
145145
users_list_tool = None
146146
for tool in tools:
147-
if "users.list" in tool.name.lower():
147+
if "users_list" in tool.name.lower():
148148
users_list_tool = tool
149149
break
150150

@@ -244,10 +244,10 @@ async def test_advanced_server_tools(advanced_mcp_session):
244244
tools = result.tools
245245
assert len(tools) > 0
246246

247-
# Find the config.set command
247+
# Find the config_set command
248248
config_set_tool = None
249249
for tool in tools:
250-
if "config.set" in tool.name.lower():
250+
if "config_set" in tool.name.lower():
251251
config_set_tool = tool
252252
break
253253

@@ -324,10 +324,10 @@ async def test_invoke_advanced_commands(advanced_mcp_session):
324324
result = await advanced_mcp_session.list_tools()
325325
tools = result.tools
326326

327-
# Find the config.set command
327+
# Find the config_set command
328328
config_set_tool = None
329329
for tool in tools:
330-
if "config.set" in tool.name.lower():
330+
if "config_set" in tool.name.lower():
331331
config_set_tool = tool
332332
break
333333

@@ -342,10 +342,10 @@ async def test_invoke_advanced_commands(advanced_mcp_session):
342342
assert len(result.content) > 0
343343
assert result.content[0].text == "Setting test=value"
344344

345-
# Find the config.get command with positional argument
345+
# Find the config_get command with positional argument
346346
config_get_tool = None
347347
for tool in tools:
348-
if "config.get" in tool.name.lower():
348+
if "config_get" in tool.name.lower():
349349
config_get_tool = tool
350350
break
351351

@@ -365,6 +365,29 @@ async def test_invoke_advanced_commands(advanced_mcp_session):
365365
assert len(result.content) > 0
366366
assert "Value for test-key: example_value" in result.content[0].text
367367

368+
# Find the config_get_value command with underscore in name
369+
config_get_value_tool = None
370+
for tool in tools:
371+
if "get_value" in tool.name.lower():
372+
config_get_value_tool = tool
373+
break
374+
375+
assert config_get_value_tool is not None
376+
# Verify inputSchema structure for positional arguments
377+
assert config_get_value_tool.inputSchema["type"] == "object"
378+
assert "properties" in config_get_value_tool.inputSchema
379+
assert "key" in config_get_value_tool.inputSchema["properties"]
380+
assert "required" in config_get_value_tool.inputSchema
381+
assert "key" in config_get_value_tool.inputSchema["required"]
382+
383+
# Invoke the command with positional argument
384+
result = await advanced_mcp_session.call_tool(
385+
config_get_value_tool.name, {"key": "underscore-key"}
386+
)
387+
assert result is not None
388+
assert len(result.content) > 0
389+
assert "Value for underscore-key (from get_value)" in result.content[0].text
390+
368391
# Find the copy command with multiple positional arguments
369392
copy_tool = None
370393
for tool in tools:

tests/test_path_sanitization.py

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
"""
2+
Tests for path sanitization and restoration functionality.
3+
"""
4+
5+
import pytest
6+
7+
from click_mcp.scanner import _original_paths, get_original_path, sanitize_tool_name
8+
9+
10+
@pytest.fixture
11+
def clear_original_paths():
12+
"""Clear the original paths mapping before each test."""
13+
_original_paths.clear()
14+
yield
15+
_original_paths.clear() # Clean up after test
16+
17+
18+
def test_simple_path_sanitization(clear_original_paths):
19+
"""Test sanitization of simple paths."""
20+
# Store mapping manually since sanitize_tool_name doesn't do it anymore
21+
original = "command.subcommand"
22+
sanitized = sanitize_tool_name(original)
23+
_original_paths[sanitized] = original
24+
25+
assert sanitized == "command_subcommand"
26+
assert get_original_path(sanitized) == original
27+
28+
29+
def test_path_with_existing_underscores(clear_original_paths):
30+
"""Test sanitization of paths that already contain underscores."""
31+
# Store mapping manually
32+
original = "command.sub_command.action"
33+
sanitized = sanitize_tool_name(original)
34+
_original_paths[sanitized] = original
35+
36+
assert sanitized == "command_sub_command_action"
37+
assert get_original_path(sanitized) == original
38+
39+
40+
def test_path_with_special_characters(clear_original_paths):
41+
"""Test sanitization of paths with special characters."""
42+
# Store mapping manually
43+
original = "command.sub-command.action!"
44+
sanitized = sanitize_tool_name(original)
45+
_original_paths[sanitized] = original
46+
47+
assert sanitized == "command_sub_command_action_"
48+
assert get_original_path(sanitized) == original
49+
50+
51+
def test_path_starting_with_number(clear_original_paths):
52+
"""Test sanitization of paths that start with a number."""
53+
# Store mapping manually
54+
original = "1command.subcommand"
55+
sanitized = sanitize_tool_name(original)
56+
_original_paths[sanitized] = original
57+
58+
assert sanitized == "tool_1command_subcommand"
59+
assert get_original_path(sanitized) == original
60+
61+
62+
def test_fallback_to_sanitized_name(clear_original_paths):
63+
"""
64+
Test that get_original_path falls back to the sanitized name if no mapping exists.
65+
"""
66+
sanitized = "unknown_command"
67+
assert get_original_path(sanitized) == sanitized
68+
69+
70+
def test_multiple_sanitizations(clear_original_paths):
71+
"""Test that multiple sanitizations work correctly."""
72+
# First path
73+
original1 = "command.subcommand"
74+
sanitized1 = sanitize_tool_name(original1)
75+
_original_paths[sanitized1] = original1
76+
77+
# Second path
78+
original2 = "another.sub_command"
79+
sanitized2 = sanitize_tool_name(original2)
80+
_original_paths[sanitized2] = original2
81+
82+
assert sanitized1 == "command_subcommand"
83+
assert get_original_path(sanitized1) == original1
84+
85+
assert sanitized2 == "another_sub_command"
86+
assert get_original_path(sanitized2) == original2
87+
88+
89+
def test_scan_click_command_mapping():
90+
"""Test that scan_click_command correctly maps sanitized names to original paths."""
91+
import click
92+
93+
from click_mcp.scanner import scan_click_command
94+
95+
# Create a simple Click command group with nested commands
96+
@click.group()
97+
def cli():
98+
pass
99+
100+
@cli.command()
101+
def simple():
102+
pass
103+
104+
@cli.group()
105+
def nested():
106+
pass
107+
108+
@nested.command()
109+
def command():
110+
pass
111+
112+
@nested.command(name="with-dash")
113+
def with_dash():
114+
pass
115+
116+
@nested.command(name="under_score")
117+
def under_score():
118+
pass
119+
120+
# Clear the mapping before scanning
121+
_original_paths.clear()
122+
123+
# Scan the command group
124+
scan_click_command(cli)
125+
126+
# Check that the mapping contains the expected entries
127+
assert "simple" in _original_paths.values()
128+
assert "nested.command" in _original_paths.values()
129+
assert "nested.with-dash" in _original_paths.values()
130+
assert "nested.under_score" in _original_paths.values()
131+
132+
# Check that the sanitized names map to the correct original paths
133+
for sanitized, original in _original_paths.items():
134+
if original == "nested.with-dash":
135+
assert sanitized == "nested_with_dash"
136+
elif original == "nested.under_score":
137+
assert sanitized == "nested_under_score"
138+
elif original == "nested.command":
139+
assert sanitized == "nested_command"

0 commit comments

Comments
 (0)