From 9fdcb8ac93a2f308afa8d0285e77981a123dc652 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 15 Apr 2025 10:04:41 -0400 Subject: [PATCH] Update matching logic for resources --- src/fastmcp/resources/resource_manager.py | 9 +- src/fastmcp/resources/template.py | 36 +++-- tests/resources/test_resource_manager.py | 161 ++++++++++++++++++++++ 3 files changed, 194 insertions(+), 12 deletions(-) diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 6b170ae77..1dc21a45e 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -8,7 +8,7 @@ from pydantic import AnyUrl from fastmcp.exceptions import ResourceError from fastmcp.resources import FunctionResource, Resource -from fastmcp.resources.template import ResourceTemplate +from fastmcp.resources.template import ResourceTemplate, match_uri_template from fastmcp.settings import DuplicateBehavior from fastmcp.utilities.logging import get_logger @@ -207,9 +207,10 @@ class ResourceManager: if resource := self._resources.get(uri_str): return resource - # Then check templates - for template in self._templates.values(): - if params := template.matches(uri_str): + # Then check templates - use the utility function to match against storage keys + for storage_key, template in self._templates.items(): + # Try to match against the storage key (which might be a custom key) + if params := match_uri_template(uri_str, storage_key): try: return await template.create_resource(uri_str, params) except Exception as e: diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index aadae3260..7d05aa057 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -7,12 +7,37 @@ import re from collections.abc import Callable from typing import Annotated, Any -from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call +from pydantic import ( + AnyUrl, + BaseModel, + BeforeValidator, + Field, + TypeAdapter, + validate_call, +) from fastmcp.resources.types import FunctionResource, Resource from fastmcp.utilities.types import _convert_set_defaults +def match_uri_template(uri: str, uri_template: str) -> dict[str, Any] | None: + """Match a URI against a template and extract parameters. + + Args: + uri: The URI to match against the template + uri_template: The URI template to match against + + Returns: + A dictionary of extracted parameters if there's a match, or None if no match + """ + # Convert template to regex pattern + pattern = uri_template.replace("{", "(?P<").replace("}", ">[^/]+)") + match = re.match(f"^{pattern}$", uri) + if match: + return match.groupdict() + return None + + class MyModel(BaseModel): key: str value: int @@ -94,12 +119,7 @@ class ResourceTemplate(BaseModel): def matches(self, uri: str) -> dict[str, Any] | None: """Check if URI matches template and extract parameters.""" - # Convert template to regex pattern - pattern = self.uri_template.replace("{", "(?P<").replace("}", ">[^/]+)") - match = re.match(f"^{pattern}$", uri) - if match: - return match.groupdict() - return None + return match_uri_template(uri, self.uri_template) async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: """Create a resource from the template with the given parameters.""" @@ -110,7 +130,7 @@ class ResourceTemplate(BaseModel): result = await result return FunctionResource( - uri=uri, # type: ignore + uri=AnyUrl(uri), # Explicitly convert to AnyUrl name=self.name, description=self.description, mime_type=self.mime_type, diff --git a/tests/resources/test_resource_manager.py b/tests/resources/test_resource_manager.py index ede30eec4..28e0b3903 100644 --- a/tests/resources/test_resource_manager.py +++ b/tests/resources/test_resource_manager.py @@ -688,3 +688,164 @@ class TestImports: # Verify both resource types were imported with prefixes assert "test+data://resource" in target_manager._resources assert "test+data://template/{id}" in target_manager._templates + + +class TestCustomResourceKeys: + """Test adding resources and templates with custom keys.""" + + def test_add_resource_with_custom_key(self, temp_file: Path): + """Test adding a resource with a custom key different from its URI.""" + manager = ResourceManager() + original_uri = f"file://{temp_file}" + custom_key = "custom://resource/key" + + resource = FileResource( + uri=FileUrl(original_uri), + name="test_resource", + path=temp_file, + ) + + manager.add_resource(resource, key=custom_key) + + # Resource should be accessible via custom key + assert custom_key in manager._resources + # But not via its original URI + assert original_uri not in manager._resources + # The resource's internal URI remains unchanged + assert str(manager._resources[custom_key].uri) == original_uri + + def test_add_template_with_custom_key(self): + """Test adding a template with a custom key different from its URI template.""" + manager = ResourceManager() + + def template_fn(id: str) -> str: + return f"Template {id}" + + original_uri_template = "test://{id}" + custom_key = "custom://{id}/template" + + template = ResourceTemplate.from_function( + fn=template_fn, + uri_template=original_uri_template, + name="test_template", + ) + + manager.add_template(template, key=custom_key) + + # Template should be accessible via custom key + assert custom_key in manager._templates + # But not via its original URI template + assert original_uri_template not in manager._templates + # The template's internal URI template remains unchanged + assert str(manager._templates[custom_key].uri_template) == original_uri_template + + @pytest.mark.anyio + async def test_get_resource_with_custom_key(self, temp_file: Path): + """Test that get_resource works with resources added with custom keys.""" + manager = ResourceManager() + original_uri = f"file://{temp_file}" + custom_key = "custom://resource/path" + + resource = FileResource( + uri=FileUrl(original_uri), + name="test_resource", + path=temp_file, + ) + + manager.add_resource(resource, key=custom_key) + + # Should be retrievable by the custom key + retrieved = await manager.get_resource(custom_key) + assert retrieved is not None + assert str(retrieved.uri) == original_uri + + # Should NOT be retrievable by the original URI + with pytest.raises(ResourceError, match="Unknown resource"): + await manager.get_resource(original_uri) + + @pytest.mark.anyio + async def test_get_resource_from_template_with_custom_key(self): + """Test that templates with custom keys can create resources.""" + manager = ResourceManager() + + def greet(name: str) -> str: + return f"Hello, {name}!" + + original_template = "greet://{name}" + custom_key = "custom://greet/{name}" + + template = ResourceTemplate.from_function( + fn=greet, + uri_template=original_template, + name="custom_greeter", + ) + + manager.add_template(template, key=custom_key) + + # Using a URI that matches the custom key pattern + resource = await manager.get_resource("custom://greet/world") + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Hello, world!" + + # Shouldn't work with the original template pattern + with pytest.raises(ResourceError, match="Unknown resource"): + await manager.get_resource("greet://world") + + def test_import_resources_with_custom_keys(self): + """Test that import_resources properly uses custom keys.""" + source_manager = ResourceManager() + target_manager = ResourceManager() + + # Add a resource to source manager + async def resource_fn(): + return "Resource data" + + original_uri = "data://original" + resource = FunctionResource( + uri=AnyUrl(original_uri), + name="original_resource", + fn=resource_fn, + ) + source_manager.add_resource(resource) + + # Import with prefix which creates a new key + prefix = "imported+" + target_manager.import_resources(source_manager, prefix) + + # Resource should be in target manager with prefixed URI as key + prefixed_uri = f"{prefix}{original_uri}" + assert prefixed_uri in target_manager._resources + + # The resource's internal URI should remain unchanged + stored_resource = target_manager._resources[prefixed_uri] + assert str(stored_resource.uri) == original_uri + + def test_import_templates_with_custom_keys(self): + """Test that import_templates properly uses custom keys.""" + source_manager = ResourceManager() + target_manager = ResourceManager() + + # Add a template to source manager + async def template_fn(id: str): + return f"Template {id}" + + original_template = "template://{id}" + template = ResourceTemplate.from_function( + fn=template_fn, + uri_template=original_template, + name="original_template", + ) + source_manager.add_template(template) + + # Import with prefix which creates a new key + prefix = "imported+" + target_manager.import_templates(source_manager, prefix) + + # Template should be in target manager with prefixed URI template as key + prefixed_template = f"{prefix}{original_template}" + assert prefixed_template in target_manager._templates + + # The template's internal URI template should remain unchanged + stored_template = target_manager._templates[prefixed_template] + assert str(stored_template.uri_template) == original_template