Skip to content

Commit 6c8482c

Browse files
committed
feat(managers): add tenant-scoped storage to ToolManager, ResourceManager, and PromptManager
Change internal storage from simple name-keyed dicts to composite (tenant_id, name) keyed dicts, enabling the same resource name to exist independently under different tenants. All public methods gain a keyword-only `tenant_id: str | None = None` parameter. Existing callers are unaffected — the default None scope preserves backward compatibility with single-tenant usage.
1 parent 0c36ac5 commit 6c8482c

5 files changed

Lines changed: 398 additions & 55 deletions

File tree

src/mcp/server/mcpserver/prompts/manager.py

Lines changed: 24 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,44 +15,55 @@
1515

1616

1717
class PromptManager:
18-
"""Manages MCPServer prompts."""
18+
"""Manages MCPServer prompts with optional tenant-scoped storage.
19+
20+
Prompts are stored in a dict keyed by ``(tenant_id, prompt_name)``.
21+
This allows the same prompt name to exist independently under different
22+
tenants. When ``tenant_id`` is ``None`` (the default), prompts live in
23+
a global scope, preserving backward compatibility with single-tenant usage.
24+
"""
1925

2026
def __init__(self, warn_on_duplicate_prompts: bool = True):
21-
self._prompts: dict[str, Prompt] = {}
27+
self._prompts: dict[tuple[str | None, str], Prompt] = {}
2228
self.warn_on_duplicate_prompts = warn_on_duplicate_prompts
2329

24-
def get_prompt(self, name: str) -> Prompt | None:
25-
"""Get prompt by name."""
26-
return self._prompts.get(name)
30+
def get_prompt(self, name: str, *, tenant_id: str | None = None) -> Prompt | None:
31+
"""Get prompt by name, optionally scoped to a tenant."""
32+
return self._prompts.get((tenant_id, name))
2733

28-
def list_prompts(self) -> list[Prompt]:
29-
"""List all registered prompts."""
30-
return list(self._prompts.values())
34+
def list_prompts(self, *, tenant_id: str | None = None) -> list[Prompt]:
35+
"""List all registered prompts for a given tenant scope."""
36+
return [prompt for (tid, _), prompt in self._prompts.items() if tid == tenant_id]
3137

3238
def add_prompt(
3339
self,
3440
prompt: Prompt,
41+
*,
42+
tenant_id: str | None = None,
3543
) -> Prompt:
36-
"""Add a prompt to the manager."""
44+
"""Add a prompt to the manager, optionally scoped to a tenant."""
3745

3846
# Check for duplicates
39-
existing = self._prompts.get(prompt.name)
47+
key = (tenant_id, prompt.name)
48+
existing = self._prompts.get(key)
4049
if existing:
4150
if self.warn_on_duplicate_prompts:
4251
logger.warning(f"Prompt already exists: {prompt.name}")
4352
return existing
4453

45-
self._prompts[prompt.name] = prompt
54+
self._prompts[key] = prompt
4655
return prompt
4756

4857
async def render_prompt(
4958
self,
5059
name: str,
5160
arguments: dict[str, Any] | None,
5261
context: Context[LifespanContextT, RequestT],
62+
*,
63+
tenant_id: str | None = None,
5364
) -> list[Message]:
54-
"""Render a prompt by name with arguments."""
55-
prompt = self.get_prompt(name)
65+
"""Render a prompt by name with arguments, optionally scoped to a tenant."""
66+
prompt = self.get_prompt(name, tenant_id=tenant_id)
5667
if not prompt:
5768
raise ValueError(f"Unknown prompt: {name}")
5869

src/mcp/server/mcpserver/resources/resource_manager.py

Lines changed: 44 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -20,18 +20,27 @@
2020

2121

2222
class ResourceManager:
23-
"""Manages MCPServer resources."""
23+
"""Manages MCPServer resources with optional tenant-scoped storage.
24+
25+
Resources and templates are stored in dicts keyed by
26+
``(tenant_id, uri_string)`` and ``(tenant_id, uri_template)``
27+
respectively. This allows the same URI to exist independently under
28+
different tenants. When ``tenant_id`` is ``None`` (the default),
29+
entries live in a global scope, preserving backward compatibility
30+
with single-tenant usage.
31+
"""
2432

2533
def __init__(self, warn_on_duplicate_resources: bool = True):
26-
self._resources: dict[str, Resource] = {}
27-
self._templates: dict[str, ResourceTemplate] = {}
34+
self._resources: dict[tuple[str | None, str], Resource] = {}
35+
self._templates: dict[tuple[str | None, str], ResourceTemplate] = {}
2836
self.warn_on_duplicate_resources = warn_on_duplicate_resources
2937

30-
def add_resource(self, resource: Resource) -> Resource:
31-
"""Add a resource to the manager.
38+
def add_resource(self, resource: Resource, *, tenant_id: str | None = None) -> Resource:
39+
"""Add a resource to the manager, optionally scoped to a tenant.
3240
3341
Args:
3442
resource: A Resource instance to add
43+
tenant_id: Optional tenant scope for the resource
3544
3645
Returns:
3746
The added resource. If a resource with the same URI already exists,
@@ -45,12 +54,13 @@ def add_resource(self, resource: Resource) -> Resource:
4554
"resource_name": resource.name,
4655
},
4756
)
48-
existing = self._resources.get(str(resource.uri))
57+
key = (tenant_id, str(resource.uri))
58+
existing = self._resources.get(key)
4959
if existing:
5060
if self.warn_on_duplicate_resources:
5161
logger.warning(f"Resource already exists: {resource.uri}")
5262
return existing
53-
self._resources[str(resource.uri)] = resource
63+
self._resources[key] = resource
5464
return resource
5565

5666
def add_template(
@@ -64,8 +74,10 @@ def add_template(
6474
icons: list[Icon] | None = None,
6575
annotations: Annotations | None = None,
6676
meta: dict[str, Any] | None = None,
77+
*,
78+
tenant_id: str | None = None,
6779
) -> ResourceTemplate:
68-
"""Add a template from a function."""
80+
"""Add a template from a function, optionally scoped to a tenant."""
6981
template = ResourceTemplate.from_function(
7082
fn,
7183
uri_template=uri_template,
@@ -77,20 +89,28 @@ def add_template(
7789
annotations=annotations,
7890
meta=meta,
7991
)
80-
self._templates[template.uri_template] = template
92+
self._templates[(tenant_id, template.uri_template)] = template
8193
return template
8294

83-
async def get_resource(self, uri: AnyUrl | str, context: Context[LifespanContextT, RequestT]) -> Resource:
95+
async def get_resource(
96+
self,
97+
uri: AnyUrl | str,
98+
context: Context[LifespanContextT, RequestT],
99+
*,
100+
tenant_id: str | None = None,
101+
) -> Resource:
84102
"""Get resource by URI, checking concrete resources first, then templates."""
85103
uri_str = str(uri)
86104
logger.debug("Getting resource", extra={"uri": uri_str})
87105

88106
# First check concrete resources
89-
if resource := self._resources.get(uri_str):
107+
if resource := self._resources.get((tenant_id, uri_str)):
90108
return resource
91109

92-
# Then check templates
93-
for template in self._templates.values():
110+
# Then check templates for this tenant scope
111+
for (tid, _), template in self._templates.items():
112+
if tid != tenant_id:
113+
continue
94114
if params := template.matches(uri_str):
95115
try:
96116
return await template.create_resource(uri_str, params, context=context)
@@ -99,12 +119,14 @@ async def get_resource(self, uri: AnyUrl | str, context: Context[LifespanContext
99119

100120
raise ValueError(f"Unknown resource: {uri}")
101121

102-
def list_resources(self) -> list[Resource]:
103-
"""List all registered resources."""
104-
logger.debug("Listing resources", extra={"count": len(self._resources)})
105-
return list(self._resources.values())
106-
107-
def list_templates(self) -> list[ResourceTemplate]:
108-
"""List all registered templates."""
109-
logger.debug("Listing templates", extra={"count": len(self._templates)})
110-
return list(self._templates.values())
122+
def list_resources(self, *, tenant_id: str | None = None) -> list[Resource]:
123+
"""List all registered resources for a given tenant scope."""
124+
resources = [r for (tid, _), r in self._resources.items() if tid == tenant_id]
125+
logger.debug("Listing resources", extra={"count": len(resources)})
126+
return resources
127+
128+
def list_templates(self, *, tenant_id: str | None = None) -> list[ResourceTemplate]:
129+
"""List all registered templates for a given tenant scope."""
130+
templates = [t for (tid, _), t in self._templates.items() if tid == tenant_id]
131+
logger.debug("Listing templates", extra={"count": len(templates)})
132+
return templates

src/mcp/server/mcpserver/tools/tool_manager.py

Lines changed: 32 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -16,30 +16,37 @@
1616

1717

1818
class ToolManager:
19-
"""Manages MCPServer tools."""
19+
"""Manages MCPServer tools with optional tenant-scoped storage.
20+
21+
Tools are stored in a dict keyed by ``(tenant_id, tool_name)``.
22+
This allows the same tool name to exist independently under different
23+
tenants. When ``tenant_id`` is ``None`` (the default), tools live in
24+
a global scope, preserving backward compatibility with single-tenant usage.
25+
"""
2026

2127
def __init__(
2228
self,
2329
warn_on_duplicate_tools: bool = True,
2430
*,
2531
tools: list[Tool] | None = None,
2632
):
27-
self._tools: dict[str, Tool] = {}
33+
self._tools: dict[tuple[str | None, str], Tool] = {}
2834
if tools is not None:
2935
for tool in tools:
30-
if warn_on_duplicate_tools and tool.name in self._tools:
36+
key = (None, tool.name)
37+
if warn_on_duplicate_tools and key in self._tools:
3138
logger.warning(f"Tool already exists: {tool.name}")
32-
self._tools[tool.name] = tool
39+
self._tools[key] = tool
3340

3441
self.warn_on_duplicate_tools = warn_on_duplicate_tools
3542

36-
def get_tool(self, name: str) -> Tool | None:
37-
"""Get tool by name."""
38-
return self._tools.get(name)
43+
def get_tool(self, name: str, *, tenant_id: str | None = None) -> Tool | None:
44+
"""Get tool by name, optionally scoped to a tenant."""
45+
return self._tools.get((tenant_id, name))
3946

40-
def list_tools(self) -> list[Tool]:
41-
"""List all registered tools."""
42-
return list(self._tools.values())
47+
def list_tools(self, *, tenant_id: str | None = None) -> list[Tool]:
48+
"""List all registered tools for a given tenant scope."""
49+
return [tool for (tid, _), tool in self._tools.items() if tid == tenant_id]
4350

4451
def add_tool(
4552
self,
@@ -51,8 +58,10 @@ def add_tool(
5158
icons: list[Icon] | None = None,
5259
meta: dict[str, Any] | None = None,
5360
structured_output: bool | None = None,
61+
*,
62+
tenant_id: str | None = None,
5463
) -> Tool:
55-
"""Add a tool to the server."""
64+
"""Add a tool to the server, optionally scoped to a tenant."""
5665
tool = Tool.from_function(
5766
fn,
5867
name=name,
@@ -63,29 +72,33 @@ def add_tool(
6372
meta=meta,
6473
structured_output=structured_output,
6574
)
66-
existing = self._tools.get(tool.name)
75+
key = (tenant_id, tool.name)
76+
existing = self._tools.get(key)
6777
if existing:
6878
if self.warn_on_duplicate_tools:
6979
logger.warning(f"Tool already exists: {tool.name}")
7080
return existing
71-
self._tools[tool.name] = tool
81+
self._tools[key] = tool
7282
return tool
7383

74-
def remove_tool(self, name: str) -> None:
75-
"""Remove a tool by name."""
76-
if name not in self._tools:
84+
def remove_tool(self, name: str, *, tenant_id: str | None = None) -> None:
85+
"""Remove a tool by name, optionally scoped to a tenant."""
86+
key = (tenant_id, name)
87+
if key not in self._tools:
7788
raise ToolError(f"Unknown tool: {name}")
78-
del self._tools[name]
89+
del self._tools[key]
7990

8091
async def call_tool(
8192
self,
8293
name: str,
8394
arguments: dict[str, Any],
8495
context: Context[LifespanContextT, RequestT],
8596
convert_result: bool = False,
97+
*,
98+
tenant_id: str | None = None,
8699
) -> Any:
87-
"""Call a tool by name with arguments."""
88-
tool = self.get_tool(name)
100+
"""Call a tool by name with arguments, optionally scoped to a tenant."""
101+
tool = self.get_tool(name, tenant_id=tenant_id)
89102
if not tool:
90103
raise ToolError(f"Unknown tool: {name}")
91104

tests/server/mcpserver/resources/test_resource_manager.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ def greet(name: str) -> str:
103103
uri_template="greet://{name}",
104104
name="greeter",
105105
)
106-
manager._templates[template.uri_template] = template
106+
manager._templates[(None, template.uri_template)] = template
107107

108108
resource = await manager.get_resource(AnyUrl("greet://world"), Context())
109109
assert isinstance(resource, FunctionResource)

0 commit comments

Comments
 (0)