Skip to content

Commit b097306

Browse files
rfrneo4ja-s-g93
andauthored
memory - add namespacing (#183)
* added namespacing to server tools, updated docs and change logs, added unit testing * fixed readme * update utils, tests --------- Co-authored-by: runfourestrun <[email protected]> Co-authored-by: alex <[email protected]>
1 parent fa0ea3a commit b097306

File tree

7 files changed

+312
-13
lines changed

7 files changed

+312
-13
lines changed

servers/mcp-neo4j-memory/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
### Changed
66

77
### Added
8+
* Add namespacing support for multi-tenant deployments with `--namespace` CLI argument and `NEO4J_NAMESPACE` environment variable
89

910
## v0.4.0
1011

servers/mcp-neo4j-memory/README.md

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,15 @@ Alternatively, you can set environment variables:
147147
}
148148
```
149149

150+
#### Namespacing
151+
For multi-tenant deployments, add `--namespace` to prefix tool names:
152+
```json
153+
"args": [ "[email protected]", "--namespace", "myapp", "--db-url", "..." ]
154+
```
155+
Tools become: `myapp-read_graph`, `myapp-create_entities`, etc.
156+
157+
Can also use `NEO4J_NAMESPACE` environment variable.
158+
150159
### 🌐 HTTP Transport Mode
151160

152161
The server supports HTTP transport for web-based deployments and microservices:
@@ -166,6 +175,7 @@ export NEO4J_TRANSPORT=http
166175
export NEO4J_MCP_SERVER_HOST=127.0.0.1
167176
export NEO4J_MCP_SERVER_PORT=8080
168177
export NEO4J_MCP_SERVER_PATH=/api/mcp/
178+
export NEO4J_NAMESPACE=myapp
169179
mcp-neo4j-memory
170180
```
171181

@@ -298,6 +308,7 @@ docker run --rm -p 8000:8000 \
298308
| `NEO4J_MCP_SERVER_PATH` | `/mcp/` | Path for accessing MCP server |
299309
| `NEO4J_MCP_SERVER_ALLOW_ORIGINS` | _(empty - secure by default)_ | Comma-separated list of allowed CORS origins |
300310
| `NEO4J_MCP_SERVER_ALLOWED_HOSTS` | `localhost,127.0.0.1` | Comma-separated list of allowed hosts (DNS rebinding protection) |
311+
| `NEO4J_NAMESPACE` | _(empty - no prefix)_ | Namespace prefix for tool names (e.g., `myapp-read_graph`) |
301312

302313
### 🌐 SSE Transport for Legacy Web Access
303314

servers/mcp-neo4j-memory/src/mcp_neo4j_memory/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ def main():
1515
parser.add_argument('--username', default=None, help='Neo4j username')
1616
parser.add_argument('--password', default=None, help='Neo4j password')
1717
parser.add_argument("--database", default=None, help="Neo4j database name")
18+
parser.add_argument("--namespace", default=None, help="Tool namespace prefix")
1819
parser.add_argument("--transport", default=None, help="Transport type (stdio, sse, http)")
1920
parser.add_argument("--server-host", default=None, help="HTTP host (default: 127.0.0.1)")
2021
parser.add_argument("--server-port", type=int, default=None, help="HTTP port (default: 8000)")
@@ -23,6 +24,7 @@ def main():
2324
parser.add_argument("--allowed-hosts", default=None, help="Comma-separated list of allowed hosts for DNS rebinding protection")
2425

2526
args = parser.parse_args()
27+
2628
config = process_config(args)
2729
asyncio.run(server.main(**config))
2830

servers/mcp-neo4j-memory/src/mcp_neo4j_memory/server.py

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -15,18 +15,25 @@
1515
from mcp.types import ToolAnnotations
1616

1717
from .neo4j_memory import Neo4jMemory, Entity, Relation, ObservationAddition, ObservationDeletion, KnowledgeGraph
18+
from .utils import format_namespace
1819

1920
# Set up logging
2021
logger = logging.getLogger('mcp_neo4j_memory')
2122
logger.setLevel(logging.INFO)
2223

2324

24-
def create_mcp_server(memory: Neo4jMemory) -> FastMCP:
25+
26+
27+
28+
def create_mcp_server(memory: Neo4jMemory, namespace: str = "") -> FastMCP:
2529
"""Create an MCP server instance for memory management."""
2630

31+
namespace_prefix = format_namespace(namespace)
2732
mcp: FastMCP = FastMCP("mcp-neo4j-memory", dependencies=["neo4j", "pydantic"], stateless_http=True)
2833

29-
@mcp.tool(annotations=ToolAnnotations(title="Read Graph",
34+
@mcp.tool(
35+
name=namespace_prefix + "read_graph",
36+
annotations=ToolAnnotations(title="Read Graph",
3037
readOnlyHint=True,
3138
destructiveHint=False,
3239
idempotentHint=True,
@@ -45,7 +52,9 @@ async def read_graph() -> KnowledgeGraph:
4552
logger.error(f"Error reading full knowledge graph: {e}")
4653
raise ToolError(f"Error reading full knowledge graph: {e}")
4754

48-
@mcp.tool(annotations=ToolAnnotations(title="Create Entities",
55+
@mcp.tool(
56+
name=namespace_prefix + "create_entities",
57+
annotations=ToolAnnotations(title="Create Entities",
4958
readOnlyHint=False,
5059
destructiveHint=False,
5160
idempotentHint=True,
@@ -65,7 +74,9 @@ async def create_entities(entities: list[Entity] = Field(..., description="List
6574
logger.error(f"Error creating entities: {e}")
6675
raise ToolError(f"Error creating entities: {e}")
6776

68-
@mcp.tool(annotations=ToolAnnotations(title="Create Relations",
77+
@mcp.tool(
78+
name=namespace_prefix + "create_relations",
79+
annotations=ToolAnnotations(title="Create Relations",
6980
readOnlyHint=False,
7081
destructiveHint=False,
7182
idempotentHint=True,
@@ -85,7 +96,9 @@ async def create_relations(relations: list[Relation] = Field(..., description="L
8596
logger.error(f"Error creating relations: {e}")
8697
raise ToolError(f"Error creating relations: {e}")
8798

88-
@mcp.tool(annotations=ToolAnnotations(title="Add Observations",
99+
@mcp.tool(
100+
name=namespace_prefix + "add_observations",
101+
annotations=ToolAnnotations(title="Add Observations",
89102
readOnlyHint=False,
90103
destructiveHint=False,
91104
idempotentHint=True,
@@ -105,7 +118,9 @@ async def add_observations(observations: list[ObservationAddition] = Field(...,
105118
logger.error(f"Error adding observations: {e}")
106119
raise ToolError(f"Error adding observations: {e}")
107120

108-
@mcp.tool(annotations=ToolAnnotations(title="Delete Entities",
121+
@mcp.tool(
122+
name=namespace_prefix + "delete_entities",
123+
annotations=ToolAnnotations(title="Delete Entities",
109124
readOnlyHint=False,
110125
destructiveHint=True,
111126
idempotentHint=True,
@@ -124,7 +139,9 @@ async def delete_entities(entityNames: list[str] = Field(..., description="List
124139
logger.error(f"Error deleting entities: {e}")
125140
raise ToolError(f"Error deleting entities: {e}")
126141

127-
@mcp.tool(annotations=ToolAnnotations(title="Delete Observations",
142+
@mcp.tool(
143+
name=namespace_prefix + "delete_observations",
144+
annotations=ToolAnnotations(title="Delete Observations",
128145
readOnlyHint=False,
129146
destructiveHint=True,
130147
idempotentHint=True,
@@ -144,7 +161,9 @@ async def delete_observations(deletions: list[ObservationDeletion] = Field(...,
144161
logger.error(f"Error deleting observations: {e}")
145162
raise ToolError(f"Error deleting observations: {e}")
146163

147-
@mcp.tool(annotations=ToolAnnotations(title="Delete Relations",
164+
@mcp.tool(
165+
name=namespace_prefix + "delete_relations",
166+
annotations=ToolAnnotations(title="Delete Relations",
148167
readOnlyHint=False,
149168
destructiveHint=True,
150169
idempotentHint=True,
@@ -164,7 +183,9 @@ async def delete_relations(relations: list[Relation] = Field(..., description="L
164183
logger.error(f"Error deleting relations: {e}")
165184
raise ToolError(f"Error deleting relations: {e}")
166185

167-
@mcp.tool(annotations=ToolAnnotations(title="Search Memories",
186+
@mcp.tool(
187+
name=namespace_prefix + "search_memories",
188+
annotations=ToolAnnotations(title="Search Memories",
168189
readOnlyHint=True,
169190
destructiveHint=False,
170191
idempotentHint=True,
@@ -183,7 +204,9 @@ async def search_memories(query: str = Field(..., description="Search query for
183204
logger.error(f"Error searching memories: {e}")
184205
raise ToolError(f"Error searching memories: {e}")
185206

186-
@mcp.tool(annotations=ToolAnnotations(title="Find Memories by Name",
207+
@mcp.tool(
208+
name=namespace_prefix + "find_memories_by_name",
209+
annotations=ToolAnnotations(title="Find Memories by Name",
187210
readOnlyHint=True,
188211
destructiveHint=False,
189212
idempotentHint=True,
@@ -211,6 +234,7 @@ async def main(
211234
neo4j_password: str,
212235
neo4j_database: str,
213236
transport: Literal["stdio", "sse", "http"] = "stdio",
237+
namespace: str = "",
214238
host: str = "127.0.0.1",
215239
port: int = 8000,
216240
path: str = "/mcp/",
@@ -255,7 +279,7 @@ async def main(
255279
]
256280

257281
# Create MCP server
258-
mcp = create_mcp_server(memory)
282+
mcp = create_mcp_server(memory, namespace)
259283
logger.info("MCP server created")
260284

261285
# Run the server with the specified transport

servers/mcp-neo4j-memory/src/mcp_neo4j_memory/utils.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,16 @@
66
logger = logging.getLogger("mcp_neo4j_memory")
77
logger.setLevel(logging.INFO)
88

9+
def format_namespace(namespace: str) -> str:
10+
"""Format namespace by ensuring it ends with a hyphen if not empty."""
11+
if namespace:
12+
if namespace.endswith("-"):
13+
return namespace
14+
else:
15+
return namespace + "-"
16+
else:
17+
return ""
18+
919
def process_config(args: argparse.Namespace) -> dict[str, Union[str, int, None]]:
1020
"""
1121
Process the command line arguments and environment variables to create a config dictionary.
@@ -165,5 +175,17 @@ def process_config(args: argparse.Namespace) -> dict[str, Union[str, int, None]]
165175
"Info: No allowed hosts provided. Defaulting to secure mode - only localhost and 127.0.0.1 allowed."
166176
)
167177
config["allowed_hosts"] = ["localhost", "127.0.0.1"]
178+
179+
# namespace configuration
180+
if args.namespace is not None:
181+
logger.info(f"Info: Namespace provided for tools: {args.namespace}")
182+
config["namespace"] = args.namespace
183+
else:
184+
if os.getenv("NEO4J_NAMESPACE") is not None:
185+
logger.info(f"Info: Namespace provided for tools: {os.getenv('NEO4J_NAMESPACE')}")
186+
config["namespace"] = os.getenv("NEO4J_NAMESPACE")
187+
else:
188+
logger.info("Info: No namespace provided for tools. No namespace will be used.")
189+
config["namespace"] = ""
168190

169191
return config
Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
import pytest
2+
from unittest.mock import Mock, AsyncMock
3+
4+
from mcp_neo4j_memory.server import format_namespace, create_mcp_server
5+
from mcp_neo4j_memory.neo4j_memory import Neo4jMemory, KnowledgeGraph
6+
7+
8+
class TestFormatNamespace:
9+
"""Test the format_namespace function behavior."""
10+
11+
def testformat_namespace_empty_string(self):
12+
"""Test format_namespace with empty string returns empty string."""
13+
assert format_namespace("") == ""
14+
15+
def testformat_namespace_no_hyphen(self):
16+
"""Test format_namespace adds hyphen when not present."""
17+
assert format_namespace("myapp") == "myapp-"
18+
19+
def testformat_namespace_with_hyphen(self):
20+
"""Test format_namespace returns string as-is when hyphen already present."""
21+
assert format_namespace("myapp-") == "myapp-"
22+
23+
def testformat_namespace_complex_name(self):
24+
"""Test format_namespace with complex namespace names."""
25+
assert format_namespace("company.product") == "company.product-"
26+
assert format_namespace("app_v2") == "app_v2-"
27+
28+
29+
class TestNamespacing:
30+
"""Test namespacing functionality."""
31+
32+
@pytest.fixture
33+
def mock_memory(self):
34+
"""Create a mock Neo4jMemory for testing."""
35+
memory = Mock(spec=Neo4jMemory)
36+
# Mock all the async methods that the tools will call
37+
knowledge_graph = KnowledgeGraph(entities=[], relations=[])
38+
memory.read_graph = AsyncMock(return_value=knowledge_graph)
39+
memory.create_entities = AsyncMock(return_value=[])
40+
memory.create_relations = AsyncMock(return_value=[])
41+
memory.add_observations = AsyncMock(return_value=[])
42+
memory.delete_entities = AsyncMock(return_value=None)
43+
memory.delete_observations = AsyncMock(return_value=None)
44+
memory.delete_relations = AsyncMock(return_value=None)
45+
memory.search_memories = AsyncMock(return_value=knowledge_graph)
46+
memory.find_memories_by_name = AsyncMock(return_value=knowledge_graph)
47+
return memory
48+
49+
@pytest.mark.asyncio
50+
async def test_namespace_tool_prefixes(self, mock_memory):
51+
"""Test that tools are correctly prefixed with namespace."""
52+
# Test with namespace
53+
namespaced_server = create_mcp_server(mock_memory, namespace="test-ns")
54+
tools = await namespaced_server.get_tools()
55+
56+
expected_tools = [
57+
"test-ns-read_graph",
58+
"test-ns-create_entities",
59+
"test-ns-create_relations",
60+
"test-ns-add_observations",
61+
"test-ns-delete_entities",
62+
"test-ns-delete_observations",
63+
"test-ns-delete_relations",
64+
"test-ns-search_memories",
65+
"test-ns-find_memories_by_name"
66+
]
67+
68+
for expected_tool in expected_tools:
69+
assert expected_tool in tools.keys(), f"Tool {expected_tool} not found in tools"
70+
71+
# Test without namespace (default tools)
72+
default_server = create_mcp_server(mock_memory)
73+
default_tools = await default_server.get_tools()
74+
75+
expected_default_tools = [
76+
"read_graph",
77+
"create_entities",
78+
"create_relations",
79+
"add_observations",
80+
"delete_entities",
81+
"delete_observations",
82+
"delete_relations",
83+
"search_memories",
84+
"find_memories_by_name"
85+
]
86+
87+
for expected_tool in expected_default_tools:
88+
assert expected_tool in default_tools.keys(), f"Default tool {expected_tool} not found"
89+
90+
@pytest.mark.asyncio
91+
async def test_namespace_tool_functionality(self, mock_memory):
92+
"""Test that namespaced tools function correctly."""
93+
namespaced_server = create_mcp_server(mock_memory, namespace="test")
94+
tools = await namespaced_server.get_tools()
95+
96+
# Test that a namespaced tool exists and works
97+
read_tool = tools.get("test-read_graph")
98+
assert read_tool is not None
99+
100+
# Call the tool function and verify it works
101+
result = await read_tool.fn()
102+
mock_memory.read_graph.assert_called_once()
103+
104+
@pytest.mark.asyncio
105+
async def test_multiple_namespace_isolation(self, mock_memory):
106+
"""Test that different namespaces create isolated tool sets."""
107+
server_a = create_mcp_server(mock_memory, namespace="app-a")
108+
server_b = create_mcp_server(mock_memory, namespace="app-b")
109+
110+
tools_a = await server_a.get_tools()
111+
tools_b = await server_b.get_tools()
112+
113+
# Verify app-a tools exist in server_a but not server_b
114+
assert "app-a-read_graph" in tools_a.keys()
115+
assert "app-a-read_graph" not in tools_b.keys()
116+
117+
# Verify app-b tools exist in server_b but not server_a
118+
assert "app-b-read_graph" in tools_b.keys()
119+
assert "app-b-read_graph" not in tools_a.keys()
120+
121+
# Verify both servers have the same number of tools
122+
assert len(tools_a) == len(tools_b)
123+
124+
@pytest.mark.asyncio
125+
async def test_namespace_hyphen_handling(self, mock_memory):
126+
"""Test namespace hyphen handling edge cases."""
127+
# Namespace already ending with hyphen
128+
server_with_hyphen = create_mcp_server(mock_memory, namespace="myapp-")
129+
tools_with_hyphen = await server_with_hyphen.get_tools()
130+
assert "myapp-read_graph" in tools_with_hyphen.keys()
131+
132+
# Namespace without hyphen
133+
server_without_hyphen = create_mcp_server(mock_memory, namespace="myapp")
134+
tools_without_hyphen = await server_without_hyphen.get_tools()
135+
assert "myapp-read_graph" in tools_without_hyphen.keys()
136+
137+
# Both should result in identical tool names
138+
assert set(tools_with_hyphen.keys()) == set(tools_without_hyphen.keys())
139+
140+
@pytest.mark.asyncio
141+
async def test_complex_namespace_names(self, mock_memory):
142+
"""Test complex namespace naming scenarios."""
143+
complex_namespaces = [
144+
"company.product",
145+
"app_v2",
146+
"client-123",
147+
"test.env.staging"
148+
]
149+
150+
for namespace in complex_namespaces:
151+
server = create_mcp_server(mock_memory, namespace=namespace)
152+
tools = await server.get_tools()
153+
154+
# Verify tools are properly prefixed
155+
expected_tool = f"{namespace}-read_graph"
156+
assert expected_tool in tools.keys(), f"Tool {expected_tool} not found for namespace '{namespace}'"
157+
158+
@pytest.mark.asyncio
159+
async def test_namespace_tool_count_consistency(self, mock_memory):
160+
"""Test that namespaced and default servers have the same number of tools."""
161+
default_server = create_mcp_server(mock_memory)
162+
namespaced_server = create_mcp_server(mock_memory, namespace="test")
163+
164+
default_tools = await default_server.get_tools()
165+
namespaced_tools = await namespaced_server.get_tools()
166+
167+
# Should have the same number of tools
168+
assert len(default_tools) == len(namespaced_tools)
169+
170+
# Verify we have the expected number of tools (9 tools based on the server implementation)
171+
assert len(default_tools) == 9
172+
assert len(namespaced_tools) == 9

0 commit comments

Comments
 (0)