Merge pull request #10 from andylim-duo/fix/issue-8-manager-followups

refactor(managers): O(1) tenant lookups and missing APIs

Author: andylim-duo

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

@@ -17,23 +17,29 @@
 class PromptManager:
     """Manages MCPServer prompts with optional tenant-scoped storage.
 
-    Prompts are stored in a dict keyed by ``(tenant_id, prompt_name)``.
+    Prompts are stored in a nested dict: ``{tenant_id: {prompt_name: Prompt}}``.
     This allows the same prompt name to exist independently under different
-    tenants. When ``tenant_id`` is ``None`` (the default), prompts live in
-    a global scope, preserving backward compatibility with single-tenant usage.
+    tenants with O(1) lookups per tenant. When ``tenant_id`` is ``None``
+    (the default), prompts live in a global scope, preserving backward
+    compatibility with single-tenant usage.
+
+    Note: This class is not thread-safe. It is designed to run within a
+    single-threaded async event loop, where all synchronous mutations
+    execute atomically. Do not share instances across OS threads without
+    external synchronization.
     """
 
     def __init__(self, warn_on_duplicate_prompts: bool = True):
-        self._prompts: dict[tuple[str | None, str], Prompt] = {}
+        self._prompts: dict[str | None, dict[str, Prompt]] = {}
         self.warn_on_duplicate_prompts = warn_on_duplicate_prompts
 
     def get_prompt(self, name: str, *, tenant_id: str | None = None) -> Prompt | None:
         """Get prompt by name, optionally scoped to a tenant."""
-        return self._prompts.get((tenant_id, name))
+        return self._prompts.get(tenant_id, {}).get(name)
 
     def list_prompts(self, *, tenant_id: str | None = None) -> list[Prompt]:
         """List all registered prompts for a given tenant scope."""
-        return [prompt for (tid, _), prompt in self._prompts.items() if tid == tenant_id]
+        return list(self._prompts.get(tenant_id, {}).values())
 
     def add_prompt(
         self,
@@ -42,18 +48,25 @@ def add_prompt(
         tenant_id: str | None = None,
     ) -> Prompt:
         """Add a prompt to the manager, optionally scoped to a tenant."""
-
-        # Check for duplicates
-        key = (tenant_id, prompt.name)
-        existing = self._prompts.get(key)
+        scope = self._prompts.setdefault(tenant_id, {})
+        existing = scope.get(prompt.name)
         if existing:
             if self.warn_on_duplicate_prompts:
                 logger.warning(f"Prompt already exists: {prompt.name}")
             return existing
 
-        self._prompts[key] = prompt
+        scope[prompt.name] = prompt
         return prompt
 
+    def remove_prompt(self, name: str, *, tenant_id: str | None = None) -> None:
+        """Remove a prompt by name, optionally scoped to a tenant."""
+        scope = self._prompts.get(tenant_id, {})
+        if name not in scope:
+            raise ValueError(f"Unknown prompt: {name}")
+        del scope[name]
+        if not scope and tenant_id in self._prompts:
+            del self._prompts[tenant_id]
+
     async def render_prompt(
         self,
         name: str,

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

@@ -22,17 +22,23 @@
 class ResourceManager:
     """Manages MCPServer resources with optional tenant-scoped storage.
 
-    Resources and templates are stored in dicts keyed by
-    ``(tenant_id, uri_string)`` and ``(tenant_id, uri_template)``
-    respectively. This allows the same URI to exist independently under
-    different tenants. When ``tenant_id`` is ``None`` (the default),
+    Resources and templates are stored in nested dicts:
+    ``{tenant_id: {uri_string: Resource}}`` and
+    ``{tenant_id: {uri_template: ResourceTemplate}}`` respectively.
+    This allows the same URI to exist independently under different tenants
+    with O(1) lookups per tenant. When ``tenant_id`` is ``None`` (the default),
     entries live in a global scope, preserving backward compatibility
     with single-tenant usage.
+
+    Note: This class is not thread-safe. It is designed to run within a
+    single-threaded async event loop, where all synchronous mutations
+    execute atomically. Do not share instances across OS threads without
+    external synchronization.
     """
 
     def __init__(self, warn_on_duplicate_resources: bool = True):
-        self._resources: dict[tuple[str | None, str], Resource] = {}
-        self._templates: dict[tuple[str | None, str], ResourceTemplate] = {}
+        self._resources: dict[str | None, dict[str, Resource]] = {}
+        self._templates: dict[str | None, dict[str, ResourceTemplate]] = {}
         self.warn_on_duplicate_resources = warn_on_duplicate_resources
 
     def add_resource(self, resource: Resource, *, tenant_id: str | None = None) -> Resource:
@@ -54,13 +60,14 @@ def add_resource(self, resource: Resource, *, tenant_id: str | None = None) -> R
                 "resource_name": resource.name,
             },
         )
-        key = (tenant_id, str(resource.uri))
-        existing = self._resources.get(key)
+        scope = self._resources.setdefault(tenant_id, {})
+        uri_str = str(resource.uri)
+        existing = scope.get(uri_str)
         if existing:
             if self.warn_on_duplicate_resources:
                 logger.warning(f"Resource already exists: {resource.uri}")
             return existing
-        self._resources[key] = resource
+        scope[uri_str] = resource
         return resource
 
     def add_template(
@@ -77,7 +84,12 @@ def add_template(
         *,
         tenant_id: str | None = None,
     ) -> ResourceTemplate:
-        """Add a template from a function, optionally scoped to a tenant."""
+        """Add a template from a function, optionally scoped to a tenant.
+
+        Returns:
+            The added template. If a template with the same URI template already
+            exists, returns the existing template.
+        """
         template = ResourceTemplate.from_function(
             fn,
             uri_template=uri_template,
@@ -89,9 +101,25 @@ def add_template(
             annotations=annotations,
             meta=meta,
         )
-        self._templates[(tenant_id, template.uri_template)] = template
+        scope = self._templates.setdefault(tenant_id, {})
+        existing = scope.get(template.uri_template)
+        if existing:
+            if self.warn_on_duplicate_resources:
+                logger.warning(f"Resource template already exists: {template.uri_template}")
+            return existing
+        scope[template.uri_template] = template
         return template
 
+    def remove_resource(self, uri: AnyUrl | str, *, tenant_id: str | None = None) -> None:
+        """Remove a resource by URI, optionally scoped to a tenant."""
+        uri_str = str(uri)
+        scope = self._resources.get(tenant_id, {})
+        if uri_str not in scope:
+            raise ValueError(f"Unknown resource: {uri}")
+        del scope[uri_str]
+        if not scope and tenant_id in self._resources:
+            del self._resources[tenant_id]
+
     async def get_resource(
         self,
         uri: AnyUrl | str,
@@ -104,13 +132,12 @@ async def get_resource(
         logger.debug("Getting resource", extra={"uri": uri_str})
 
         # First check concrete resources
-        if resource := self._resources.get((tenant_id, uri_str)):
+        resource = self._resources.get(tenant_id, {}).get(uri_str)
+        if resource:
             return resource
 
         # Then check templates for this tenant scope
-        for (tid, _), template in self._templates.items():
-            if tid != tenant_id:
-                continue
+        for template in self._templates.get(tenant_id, {}).values():
             if params := template.matches(uri_str):
                 try:
                     return await template.create_resource(uri_str, params, context=context)
@@ -121,12 +148,12 @@ async def get_resource(
 
     def list_resources(self, *, tenant_id: str | None = None) -> list[Resource]:
         """List all registered resources for a given tenant scope."""
-        resources = [r for (tid, _), r in self._resources.items() if tid == tenant_id]
+        resources = list(self._resources.get(tenant_id, {}).values())
         logger.debug("Listing resources", extra={"count": len(resources)})
         return resources
 
     def list_templates(self, *, tenant_id: str | None = None) -> list[ResourceTemplate]:
         """List all registered templates for a given tenant scope."""
-        templates = [t for (tid, _), t in self._templates.items() if tid == tenant_id]
+        templates = list(self._templates.get(tenant_id, {}).values())
         logger.debug("Listing templates", extra={"count": len(templates)})
         return templates

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

@@ -18,10 +18,16 @@
 class ToolManager:
     """Manages MCPServer tools with optional tenant-scoped storage.
 
-    Tools are stored in a dict keyed by ``(tenant_id, tool_name)``.
+    Tools are stored in a nested dict: ``{tenant_id: {tool_name: Tool}}``.
     This allows the same tool name to exist independently under different
-    tenants. When ``tenant_id`` is ``None`` (the default), tools live in
-    a global scope, preserving backward compatibility with single-tenant usage.
+    tenants with O(1) lookups per tenant. When ``tenant_id`` is ``None``
+    (the default), tools live in a global scope, preserving backward
+    compatibility with single-tenant usage.
+
+    Note: This class is not thread-safe. It is designed to run within a
+    single-threaded async event loop, where all synchronous mutations
+    execute atomically. Do not share instances across OS threads without
+    external synchronization.
     """
 
     def __init__(
@@ -30,23 +36,23 @@ def __init__(
         *,
         tools: list[Tool] | None = None,
     ):
-        self._tools: dict[tuple[str | None, str], Tool] = {}
+        self._tools: dict[str | None, dict[str, Tool]] = {}
         if tools is not None:
+            scope = self._tools.setdefault(None, {})
             for tool in tools:
-                key = (None, tool.name)
-                if warn_on_duplicate_tools and key in self._tools:
+                if warn_on_duplicate_tools and tool.name in scope:
                     logger.warning(f"Tool already exists: {tool.name}")
-                self._tools[key] = tool
+                scope[tool.name] = tool
 
         self.warn_on_duplicate_tools = warn_on_duplicate_tools
 
     def get_tool(self, name: str, *, tenant_id: str | None = None) -> Tool | None:
         """Get tool by name, optionally scoped to a tenant."""
-        return self._tools.get((tenant_id, name))
+        return self._tools.get(tenant_id, {}).get(name)
 
     def list_tools(self, *, tenant_id: str | None = None) -> list[Tool]:
         """List all registered tools for a given tenant scope."""
-        return [tool for (tid, _), tool in self._tools.items() if tid == tenant_id]
+        return list(self._tools.get(tenant_id, {}).values())
 
     def add_tool(
         self,
@@ -72,21 +78,23 @@ def add_tool(
             meta=meta,
             structured_output=structured_output,
         )
-        key = (tenant_id, tool.name)
-        existing = self._tools.get(key)
+        scope = self._tools.setdefault(tenant_id, {})
+        existing = scope.get(tool.name)
         if existing:
             if self.warn_on_duplicate_tools:
                 logger.warning(f"Tool already exists: {tool.name}")
             return existing
-        self._tools[key] = tool
+        scope[tool.name] = tool
         return tool
 
     def remove_tool(self, name: str, *, tenant_id: str | None = None) -> None:
         """Remove a tool by name, optionally scoped to a tenant."""
-        key = (tenant_id, name)
-        if key not in self._tools:
+        scope = self._tools.get(tenant_id, {})
+        if name not in scope:
             raise ToolError(f"Unknown tool: {name}")
-        del self._tools[key]
+        del scope[name]
+        if not scope and tenant_id in self._tools:
+            del self._tools[tenant_id]
 
     async def call_tool(
         self,

tests/server/mcpserver/conftest.py

@@ -0,0 +1,22 @@
+from collections.abc import Callable
+from typing import Any
+
+import pytest
+
+from mcp.server.mcpserver.context import Context
+
+MakeContext = Callable[..., Context[Any, Any]]
+
+
+@pytest.fixture
+def make_context() -> MakeContext:
+    """Factory fixture for creating Context instances in tests.
+
+    Centralizes Context construction so that tests don't break if the
+    Context.__init__ signature changes in later iterations.
+    """
+
+    def _make(**kwargs: Any) -> Context[Any, Any]:
+        return Context(**kwargs)
+
+    return _make

tests/server/mcpserver/resources/test_resource_manager.py

@@ -4,8 +4,8 @@
 import pytest
 from pydantic import AnyUrl
 
-from mcp.server.mcpserver import Context
-from mcp.server.mcpserver.resources import FileResource, FunctionResource, ResourceManager, ResourceTemplate
+from mcp.server.mcpserver.resources import FileResource, FunctionResource, ResourceManager
+from tests.server.mcpserver.conftest import MakeContext
 
 
 @pytest.fixture
@@ -78,7 +78,7 @@ def test_disable_warn_on_duplicate_resources(self, temp_file: Path, caplog: pyte
         assert "Resource already exists" not in caplog.text
 
     @pytest.mark.anyio
-    async def test_get_resource(self, temp_file: Path):
+    async def test_get_resource(self, temp_file: Path, make_context: MakeContext):
         """Test getting a resource by URI."""
         manager = ResourceManager()
         resource = FileResource(
@@ -87,35 +87,30 @@ async def test_get_resource(self, temp_file: Path):
             path=temp_file,
         )
         manager.add_resource(resource)
-        retrieved = await manager.get_resource(resource.uri, Context())
+        retrieved = await manager.get_resource(resource.uri, make_context())
         assert retrieved == resource
 
     @pytest.mark.anyio
-    async def test_get_resource_from_template(self):
+    async def test_get_resource_from_template(self, make_context: MakeContext):
         """Test getting a resource through a template."""
         manager = ResourceManager()
 
         def greet(name: str) -> str:
             return f"Hello, {name}!"
 
-        template = ResourceTemplate.from_function(
-            fn=greet,
-            uri_template="greet://{name}",
-            name="greeter",
-        )
-        manager._templates[(None, template.uri_template)] = template
+        manager.add_template(fn=greet, uri_template="greet://{name}", name="greeter")
 
-        resource = await manager.get_resource(AnyUrl("greet://world"), Context())
+        resource = await manager.get_resource(AnyUrl("greet://world"), make_context())
         assert isinstance(resource, FunctionResource)
         content = await resource.read()
         assert content == "Hello, world!"
 
     @pytest.mark.anyio
-    async def test_get_unknown_resource(self):
+    async def test_get_unknown_resource(self, make_context: MakeContext):
         """Test getting a non-existent resource."""
         manager = ResourceManager()
         with pytest.raises(ValueError, match="Unknown resource"):
-            await manager.get_resource(AnyUrl("unknown://test"), Context())
+            await manager.get_resource(AnyUrl("unknown://test"), make_context())
 
     def test_list_resources(self, temp_file: Path):
         """Test listing all resources."""
@@ -175,3 +170,37 @@ def get_item(id: str) -> str:  # pragma: no cover
         )
 
         assert template.meta is None
+
+    def test_add_duplicate_template(self):
+        """Test adding the same template twice returns the existing one."""
+        manager = ResourceManager()
+
+        def get_item(id: str) -> str:  # pragma: no cover
+            return f"Item {id}"
+
+        first = manager.add_template(fn=get_item, uri_template="resource://items/{id}")
+        second = manager.add_template(fn=get_item, uri_template="resource://items/{id}")
+        assert first is second
+        assert len(manager.list_templates()) == 1
+
+    def test_warn_on_duplicate_template(self, caplog: pytest.LogCaptureFixture):
+        """Test warning on duplicate template."""
+        manager = ResourceManager()
+
+        def get_item(id: str) -> str:  # pragma: no cover
+            return f"Item {id}"
+
+        manager.add_template(fn=get_item, uri_template="resource://items/{id}")
+        manager.add_template(fn=get_item, uri_template="resource://items/{id}")
+        assert "Resource template already exists" in caplog.text
+
+    def test_disable_warn_on_duplicate_template(self, caplog: pytest.LogCaptureFixture):
+        """Test disabling warning on duplicate template."""
+        manager = ResourceManager(warn_on_duplicate_resources=False)
+
+        def get_item(id: str) -> str:  # pragma: no cover
+            return f"Item {id}"
+
+        manager.add_template(fn=get_item, uri_template="resource://items/{id}")
+        manager.add_template(fn=get_item, uri_template="resource://items/{id}")
+        assert "Resource template already exists" not in caplog.text