Implement tool name restrictions from the MCP spec (#97)

Ickerday · web-flow · commit 733527fc32d6 · 2026-03-23T16:10:25.000+01:00
diff --git a/splunklib/ai/registry.py b/splunklib/ai/registry.py
@@ -11,9 +11,11 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations
 # under the License.
+
 import asyncio
 import inspect
 import logging
+import string
 from collections.abc import Callable, Sequence
 from dataclasses import asdict, dataclass
 from logging import Logger
@@ -416,6 +418,12 @@ def wrapper(func: Callable[_P, _R]) -> Callable[_P, _R]:
             if name is None:
                 name = func.__name__
 
+            if not is_tool_name_valid(name):
+                raise ToolRegistryRuntimeError(
+                    f"Tool name {name} doesn't conform to MCP spec, see: "
+                    + "https://modelcontextprotocol.io/specification/latest/server/tools#tool-names"
+                )
+
             if self._executing:
                 raise ToolRegistryRuntimeError(
                     "ToolRegistry is already running, cannot define new tools"
@@ -497,3 +505,16 @@ def _drop_type_annotations_of(
     new_func.__annotations__ = new_annotations
 
     return new_func
+
+
+MCP_ALLOWED_CHARS = string.ascii_letters + string.digits + "_-."
+
+
+def is_tool_name_valid(name: str) -> bool:
+    """Checks compliance with the MCP spec restrictions, see:
+    https://modelcontextprotocol.io/specification/latest/server/tools#tool-names
+    """
+    if not (1 <= len(name) <= 128):
+        return False
+
+    return set(name).issubset(MCP_ALLOWED_CHARS)
diff --git a/tests/unit/ai/test_registry_unit.py b/tests/unit/ai/test_registry_unit.py
@@ -15,6 +15,7 @@
 # pyright: reportPrivateUsage=false, reportUnusedFunction=false, reportUnusedParameter=false
 
 import os
+import string
 import sys
 import unittest
 from collections.abc import AsyncGenerator
@@ -27,7 +28,12 @@
 from mcp.client.stdio import stdio_client
 from mcp.types import TextContent
 
-from splunklib.ai.registry import ToolContext, ToolRegistry, ToolRegistryRuntimeError
+from splunklib.ai.registry import (
+    ToolContext,
+    ToolRegistry,
+    ToolRegistryRuntimeError,
+    is_tool_name_valid,
+)
 
 
 class TestJSONSchemaInference(unittest.TestCase):
@@ -407,6 +413,38 @@ def tool(foo: int) -> int:
             register_name(r)
 
 
+@pytest.mark.parametrize(
+    argnames="name",
+    argvalues=[
+        ".",
+        "." * 128,
+        "func.tool-name_v2",
+        string.ascii_letters + string.digits,
+    ],
+)
+def test_valid_name_passes(name: str) -> None:
+    assert is_tool_name_valid(name)
+
+
+@pytest.mark.parametrize(
+    argnames="name",
+    argvalues=[
+        "",
+        "—",
+        "." * 129,
+        "tool^name+=|/",
+        string.punctuation,
+    ],
+)
+def test_tool_decorator_raises_on_invalid_name(name: str) -> None:
+    reg = ToolRegistry()
+
+    with pytest.raises(ToolRegistryRuntimeError, match=r"Tool name .*"):
+
+        @reg.tool(name)
+        def mock_tool() -> None: ...
+
+
 class TestRegistryTestCase(unittest.IsolatedAsyncioTestCase):
     @asynccontextmanager
     async def connect(self, name: str) -> AsyncGenerator[ClientSession, Any]:
