From f8c68afd55b540a3c6061cf9a46f8254be491283 Mon Sep 17 00:00:00 2001 From: jonvet Date: Sat, 8 Mar 2025 18:22:53 +0000 Subject: [PATCH 1/5] fix: enables context injection into resources --- .../fastmcp/resources/resource_manager.py | 13 +- src/mcp/server/fastmcp/resources/templates.py | 40 +++++- src/mcp/server/fastmcp/server.py | 26 +++- .../resources/test_resource_manager.py | 101 ++++++++++++++ .../resources/test_resource_template.py | 128 ++++++++++++++++++ tests/server/fastmcp/test_server.py | 40 ++++++ 6 files changed, 338 insertions(+), 10 deletions(-) diff --git a/src/mcp/server/fastmcp/resources/resource_manager.py b/src/mcp/server/fastmcp/resources/resource_manager.py index d27e6ac12..72f55d049 100644 --- a/src/mcp/server/fastmcp/resources/resource_manager.py +++ b/src/mcp/server/fastmcp/resources/resource_manager.py @@ -1,7 +1,7 @@ """Resource manager functionality.""" from collections.abc import Callable -from typing import Any +from typing import TYPE_CHECKING, Any from pydantic import AnyUrl @@ -9,6 +9,9 @@ from mcp.server.fastmcp.resources.templates import ResourceTemplate from mcp.server.fastmcp.utilities.logging import get_logger +if TYPE_CHECKING: + from mcp.server.fastmcp.server import Context + logger = get_logger(__name__) @@ -65,7 +68,9 @@ def add_template( self._templates[template.uri_template] = template return template - async def get_resource(self, uri: AnyUrl | str) -> Resource | None: + async def get_resource( + self, uri: AnyUrl | str, context: "Context | None" = None + ) -> Resource | None: """Get resource by URI, checking concrete resources first, then templates.""" uri_str = str(uri) logger.debug("Getting resource", extra={"uri": uri_str}) @@ -78,7 +83,9 @@ async def get_resource(self, uri: AnyUrl | str) -> Resource | None: for template in self._templates.values(): if params := template.matches(uri_str): try: - return await template.create_resource(uri_str, params) + return await template.create_resource( + uri_str, params, context=context + ) except Exception as e: raise ValueError(f"Error creating resource from template: {e}") diff --git a/src/mcp/server/fastmcp/resources/templates.py b/src/mcp/server/fastmcp/resources/templates.py index a30b18253..e283c2906 100644 --- a/src/mcp/server/fastmcp/resources/templates.py +++ b/src/mcp/server/fastmcp/resources/templates.py @@ -5,11 +5,15 @@ import inspect import re from collections.abc import Callable -from typing import Any +from typing import TYPE_CHECKING, Any -from pydantic import BaseModel, Field, TypeAdapter, validate_call +from pydantic import BaseModel, Field, validate_call from mcp.server.fastmcp.resources.types import FunctionResource, Resource +from mcp.server.fastmcp.utilities.func_metadata import func_metadata + +if TYPE_CHECKING: + from mcp.server.fastmcp.server import Context class ResourceTemplate(BaseModel): @@ -27,6 +31,9 @@ class ResourceTemplate(BaseModel): parameters: dict[str, Any] = Field( description="JSON schema for function parameters" ) + context_kwarg: str | None = Field( + None, description="Name of the kwarg that should receive context" + ) @classmethod def from_function( @@ -42,8 +49,24 @@ def from_function( if func_name == "": raise ValueError("You must provide a name for lambda functions") - # Get schema from TypeAdapter - will fail if function isn't properly typed - parameters = TypeAdapter(fn).json_schema() + # Find context parameter if it exists + context_kwarg = None + sig = inspect.signature(fn) + for param_name, param in sig.parameters.items(): + if ( + param.annotation.__name__ == "Context" + if hasattr(param.annotation, "__name__") + else False + ): + context_kwarg = param_name + break + + # Get schema from func_metadata, excluding context parameter + func_arg_metadata = func_metadata( + fn, + skip_names=[context_kwarg] if context_kwarg is not None else [], + ) + parameters = func_arg_metadata.arg_model.model_json_schema() # ensure the arguments are properly cast fn = validate_call(fn) @@ -55,6 +78,7 @@ def from_function( mime_type=mime_type or "text/plain", fn=fn, parameters=parameters, + context_kwarg=context_kwarg, ) def matches(self, uri: str) -> dict[str, Any] | None: @@ -66,9 +90,15 @@ def matches(self, uri: str) -> dict[str, Any] | None: return match.groupdict() return None - async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: + async def create_resource( + self, uri: str, params: dict[str, Any], context: Context | None = None + ) -> Resource: """Create a resource from the template with the given parameters.""" try: + # Add context to params if needed + if self.context_kwarg is not None and context is not None: + params = {**params, self.context_kwarg: context} + # Call function and check if result is a coroutine result = self.fn(**params) if inspect.iscoroutine(result): diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index bf0ce880a..8c770ff1e 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -230,7 +230,8 @@ async def list_resource_templates(self) -> list[MCPResourceTemplate]: async def read_resource(self, uri: AnyUrl | str) -> Iterable[ReadResourceContents]: """Read a resource by URI.""" - resource = await self._resource_manager.get_resource(uri) + context = self.get_context() + resource = await self._resource_manager.get_resource(uri, context=context) if not resource: raise ResourceError(f"Unknown resource: {uri}") @@ -327,6 +328,10 @@ def resource( If the URI contains parameters (e.g. "resource://{param}") or the function has parameters, it will be registered as a template resource. + Resources can optionally request a Context object by adding a parameter with the + Context type annotation. The context provides access to MCP capabilities like + logging, progress reporting, and resource access. + Args: uri: URI for the resource (e.g. "resource://my-resource" or "resource://{param}") name: Optional name for the resource @@ -351,6 +356,12 @@ def get_weather(city: str) -> str: async def get_weather(city: str) -> str: data = await fetch_weather(city) return f"Weather for {city}: {data}" + + @server.resource("resource://{city}/weather") + async def get_weather(city: str, ctx: Context) -> str: + await ctx.info(f"Getting weather for {city}") + data = await fetch_weather(city) + return f"Weather for {city}: {data}" """ # Check if user passed function directly instead of calling decorator if callable(uri): @@ -367,7 +378,18 @@ def decorator(fn: AnyFunction) -> AnyFunction: if has_uri_params or has_func_params: # Validate that URI params match function params uri_params = set(re.findall(r"{(\w+)}", uri)) - func_params = set(inspect.signature(fn).parameters.keys()) + + # Get all function params except 'ctx' or any parameter of type Context + sig = inspect.signature(fn) + func_params = set() + for param_name, param in sig.parameters.items(): + # Skip context parameters + if param_name == "ctx" or ( + hasattr(param.annotation, "__name__") + and param.annotation.__name__ == "Context" + ): + continue + func_params.add(param_name) if uri_params != func_params: raise ValueError( diff --git a/tests/server/fastmcp/resources/test_resource_manager.py b/tests/server/fastmcp/resources/test_resource_manager.py index 4423e5315..9e0c050c7 100644 --- a/tests/server/fastmcp/resources/test_resource_manager.py +++ b/tests/server/fastmcp/resources/test_resource_manager.py @@ -139,3 +139,104 @@ def test_list_resources(self, temp_file: Path): resources = manager.list_resources() assert len(resources) == 2 assert resources == [resource1, resource2] + + @pytest.mark.anyio + async def test_context_injection_in_template(self): + """Test that context is injected when getting a resource from a template.""" + from mcp.server.fastmcp import Context, FastMCP + + manager = ResourceManager() + + def resource_with_context(name: str, ctx: Context) -> str: + assert isinstance(ctx, Context) + return f"Hello, {name}!" + + template = ResourceTemplate.from_function( + fn=resource_with_context, + uri_template="greet://{name}", + name="greeter", + ) + manager._templates[template.uri_template] = template + + mcp = FastMCP() + ctx = mcp.get_context() + + resource = await manager.get_resource(AnyUrl("greet://world"), context=ctx) + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Hello, world!" + + @pytest.mark.anyio + async def test_context_injection_in_async_template(self): + """Test that context is properly injected in async template functions.""" + from mcp.server.fastmcp import Context, FastMCP + + manager = ResourceManager() + + async def async_resource(name: str, ctx: Context) -> str: + assert isinstance(ctx, Context) + return f"Async Hello, {name}!" + + template = ResourceTemplate.from_function( + fn=async_resource, + uri_template="async-greet://{name}", + name="async-greeter", + ) + manager._templates[template.uri_template] = template + + mcp = FastMCP() + ctx = mcp.get_context() + + resource = await manager.get_resource( + AnyUrl("async-greet://world"), context=ctx + ) + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Async Hello, world!" + + @pytest.mark.anyio + async def test_optional_context_in_template(self): + """Test that context is optional when getting a resource from a template.""" + from mcp.server.fastmcp import Context + + manager = ResourceManager() + + def resource_with_optional_context( + name: str, ctx: Context | None = None + ) -> str: + return f"Hello, {name}!" + + template = ResourceTemplate.from_function( + fn=resource_with_optional_context, + uri_template="greet://{name}", + name="greeter", + ) + manager._templates[template.uri_template] = template + + resource = await manager.get_resource(AnyUrl("greet://world")) + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Hello, world!" + + @pytest.mark.anyio + async def test_context_error_handling_in_template(self): + """Test error handling when context injection fails in a template.""" + from mcp.server.fastmcp import Context, FastMCP + + manager = ResourceManager() + + def failing_resource(name: str, ctx: Context) -> str: + raise ValueError("Test error") + + template = ResourceTemplate.from_function( + fn=failing_resource, + uri_template="greet://{name}", + name="greeter", + ) + manager._templates[template.uri_template] = template + + mcp = FastMCP() + ctx = mcp.get_context() + + with pytest.raises(ValueError, match="Error creating resource from template"): + await manager.get_resource(AnyUrl("greet://world"), context=ctx) diff --git a/tests/server/fastmcp/resources/test_resource_template.py b/tests/server/fastmcp/resources/test_resource_template.py index 09bc600d0..1fad8e632 100644 --- a/tests/server/fastmcp/resources/test_resource_template.py +++ b/tests/server/fastmcp/resources/test_resource_template.py @@ -186,3 +186,131 @@ def get_data(value: str) -> CustomData: assert isinstance(resource, FunctionResource) content = await resource.read() assert content == "hello" + + def test_context_parameter_detection(self): + """Test that context params are detected in ResourceTemplate.from_function().""" + from mcp.server.fastmcp import Context + + def resource_with_context(key: str, ctx: Context) -> str: + return f"Key: {key}" + + template = ResourceTemplate.from_function( + fn=resource_with_context, + uri_template="test://{key}", + name="test", + ) + assert template.context_kwarg == "ctx" + + def resource_without_context(key: str) -> str: + return f"Key: {key}" + + template = ResourceTemplate.from_function( + fn=resource_without_context, + uri_template="test://{key}", + name="test", + ) + assert template.context_kwarg is None + + @pytest.mark.anyio + async def test_context_injection(self): + """Test that context is properly injected during resource creation.""" + from mcp.server.fastmcp import Context, FastMCP + + def resource_with_context(key: str, ctx: Context) -> str: + assert isinstance(ctx, Context) + return f"Key: {key}" + + template = ResourceTemplate.from_function( + fn=resource_with_context, + uri_template="test://{key}", + name="test", + ) + + mcp = FastMCP() + ctx = mcp.get_context() + resource = await template.create_resource( + "test://value", {"key": "value"}, context=ctx + ) + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Key: value" + + @pytest.mark.anyio + async def test_context_injection_async(self): + """Test that context is properly injected in async resource functions.""" + from mcp.server.fastmcp import Context, FastMCP + + async def async_resource(key: str, ctx: Context) -> str: + assert isinstance(ctx, Context) + return f"Async Key: {key}" + + template = ResourceTemplate.from_function( + fn=async_resource, + uri_template="test://{key}", + name="test", + ) + + mcp = FastMCP() + ctx = mcp.get_context() + resource = await template.create_resource( + "test://value", {"key": "value"}, context=ctx + ) + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Async Key: value" + + @pytest.mark.anyio + async def test_context_optional(self): + """Test that context is optional when creating resources.""" + from mcp.server.fastmcp import Context + + def resource_with_optional_context(key: str, ctx: Context | None = None) -> str: + return f"Key: {key}" + + template = ResourceTemplate.from_function( + fn=resource_with_optional_context, + uri_template="test://{key}", + name="test", + ) + + resource = await template.create_resource("test://value", {"key": "value"}) + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Key: value" + + @pytest.mark.anyio + async def test_context_error_handling(self): + """Test error handling when context injection fails.""" + from mcp.server.fastmcp import Context, FastMCP + + def resource_with_context(key: str, ctx: Context) -> str: + raise ValueError("Test error") + + template = ResourceTemplate.from_function( + fn=resource_with_context, + uri_template="test://{key}", + name="test", + ) + + mcp = FastMCP() + ctx = mcp.get_context() + with pytest.raises(ValueError, match="Error creating resource from template"): + await template.create_resource( + "test://value", {"key": "value"}, context=ctx + ) + + def test_context_arg_excluded_from_schema(self): + """Test that context parameters are excluded from the JSON schema.""" + from mcp.server.fastmcp import Context + + def resource_with_context(a: str, ctx: Context) -> str: + return a + + template = ResourceTemplate.from_function( + fn=resource_with_context, + uri_template="test://{key}", + name="test", + ) + + assert "ctx" not in json.dumps(template.parameters) + assert "Context" not in json.dumps(template.parameters) diff --git a/tests/server/fastmcp/test_server.py b/tests/server/fastmcp/test_server.py index e76e59c52..a59dde328 100644 --- a/tests/server/fastmcp/test_server.py +++ b/tests/server/fastmcp/test_server.py @@ -465,6 +465,46 @@ def get_data(name: str) -> str: result = await resource.read() assert result == "Data for test" + @pytest.mark.anyio + async def test_context_injection_in_resource(self): + """Test that context is properly injected in resource functions.""" + mcp = FastMCP() + context_was_injected = False + + @mcp.resource("resource://{name}/data") + def get_data_with_context(name: str, ctx: Context) -> str: + nonlocal context_was_injected + assert isinstance(ctx, Context) + context_was_injected = True + return f"Data for {name} with context" + + assert len(mcp._resource_manager._templates) == 1 + + result = list(await mcp.read_resource("resource://test/data")) + assert len(result) == 1 + assert result[0].content == "Data for test with context" + assert context_was_injected + + @pytest.mark.anyio + async def test_context_injection_in_async_resource(self): + """Test that context is properly injected in async resource functions.""" + mcp = FastMCP() + context_was_injected = False + + @mcp.resource("resource://{name}/async-data") + async def get_async_data_with_context(name: str, ctx: Context) -> str: + nonlocal context_was_injected + assert isinstance(ctx, Context) + context_was_injected = True + return f"Async data for {name} with context" + + assert len(mcp._resource_manager._templates) == 1 + + result = list(await mcp.read_resource("resource://test/async-data")) + assert len(result) == 1 + assert result[0].content == "Async data for test with context" + assert context_was_injected + class TestContextInjection: """Test context injection in tools.""" From c78d1ac76b9017377495dfa65e2eaaa52e8c8921 Mon Sep 17 00:00:00 2001 From: fubuloubu <3859395+fubuloubu@users.noreply.github.com> Date: Wed, 2 Apr 2025 13:09:29 -0400 Subject: [PATCH 2/5] fix: add type params to generic Context --- .../fastmcp/resources/resource_manager.py | 6 ++++- src/mcp/server/fastmcp/resources/templates.py | 26 +++++++++++-------- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/src/mcp/server/fastmcp/resources/resource_manager.py b/src/mcp/server/fastmcp/resources/resource_manager.py index 72f55d049..03b2f5843 100644 --- a/src/mcp/server/fastmcp/resources/resource_manager.py +++ b/src/mcp/server/fastmcp/resources/resource_manager.py @@ -11,6 +11,8 @@ if TYPE_CHECKING: from mcp.server.fastmcp.server import Context + from mcp.server.session import ServerSessionT + from mcp.shared.context import LifespanContextT logger = get_logger(__name__) @@ -69,7 +71,9 @@ def add_template( return template async def get_resource( - self, uri: AnyUrl | str, context: "Context | None" = None + self, + uri: AnyUrl | str, + context: "Context[ServerSessionT, LifespanContextT] | None" = None, ) -> Resource | None: """Get resource by URI, checking concrete resources first, then templates.""" uri_str = str(uri) diff --git a/src/mcp/server/fastmcp/resources/templates.py b/src/mcp/server/fastmcp/resources/templates.py index e283c2906..d9f1448c9 100644 --- a/src/mcp/server/fastmcp/resources/templates.py +++ b/src/mcp/server/fastmcp/resources/templates.py @@ -14,6 +14,8 @@ if TYPE_CHECKING: from mcp.server.fastmcp.server import Context + from mcp.server.session import ServerSessionT + from mcp.shared.context import LifespanContextT class ResourceTemplate(BaseModel): @@ -43,6 +45,7 @@ def from_function( name: str | None = None, description: str | None = None, mime_type: str | None = None, + context_kwarg: str | None = None, ) -> ResourceTemplate: """Create a template from a function.""" func_name = name or fn.__name__ @@ -50,16 +53,14 @@ def from_function( raise ValueError("You must provide a name for lambda functions") # Find context parameter if it exists - context_kwarg = None - sig = inspect.signature(fn) - for param_name, param in sig.parameters.items(): - if ( - param.annotation.__name__ == "Context" - if hasattr(param.annotation, "__name__") - else False - ): - context_kwarg = param_name - break + if context_kwarg is None: + from mcp.server.fastmcp.server import Context + + sig = inspect.signature(fn) + for param_name, param in sig.parameters.items(): + if param.annotation is Context: + context_kwarg = param_name + break # Get schema from func_metadata, excluding context parameter func_arg_metadata = func_metadata( @@ -91,7 +92,10 @@ def matches(self, uri: str) -> dict[str, Any] | None: return None async def create_resource( - self, uri: str, params: dict[str, Any], context: Context | None = None + self, + uri: str, + params: dict[str, Any], + context: Context[ServerSessionT, LifespanContextT] | None = None, ) -> Resource: """Create a resource from the template with the given parameters.""" try: From 5920bd2593c7945bcb7be7d2f18ada413ee7540e Mon Sep 17 00:00:00 2001 From: fubuloubu <3859395+fubuloubu@users.noreply.github.com> Date: Wed, 2 Apr 2025 13:22:10 -0400 Subject: [PATCH 3/5] fix: add type hint to set variable --- src/mcp/server/fastmcp/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index 8c770ff1e..d12d448be 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -381,7 +381,7 @@ def decorator(fn: AnyFunction) -> AnyFunction: # Get all function params except 'ctx' or any parameter of type Context sig = inspect.signature(fn) - func_params = set() + func_params: set[str] = set() for param_name, param in sig.parameters.items(): # Skip context parameters if param_name == "ctx" or ( From 079030554f57f812e49f5ba4490738c839ba17a7 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Wed, 7 May 2025 10:53:17 +0200 Subject: [PATCH 4/5] fix tests --- src/mcp/server/fastmcp/resources/resource_manager.py | 4 +++- tests/server/fastmcp/resources/test_resource_template.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/mcp/server/fastmcp/resources/resource_manager.py b/src/mcp/server/fastmcp/resources/resource_manager.py index 03b2f5843..67119f3c8 100644 --- a/src/mcp/server/fastmcp/resources/resource_manager.py +++ b/src/mcp/server/fastmcp/resources/resource_manager.py @@ -1,5 +1,7 @@ """Resource manager functionality.""" +from __future__ import annotations as _annotations + from collections.abc import Callable from typing import TYPE_CHECKING, Any @@ -73,7 +75,7 @@ def add_template( async def get_resource( self, uri: AnyUrl | str, - context: "Context[ServerSessionT, LifespanContextT] | None" = None, + context: Context[ServerSessionT, LifespanContextT] | None = None, ) -> Resource | None: """Get resource by URI, checking concrete resources first, then templates.""" uri_str = str(uri) diff --git a/tests/server/fastmcp/resources/test_resource_template.py b/tests/server/fastmcp/resources/test_resource_template.py index 1fad8e632..ff5e28599 100644 --- a/tests/server/fastmcp/resources/test_resource_template.py +++ b/tests/server/fastmcp/resources/test_resource_template.py @@ -185,7 +185,7 @@ def get_data(value: str) -> CustomData: assert isinstance(resource, FunctionResource) content = await resource.read() - assert content == "hello" + assert content == '"hello"' def test_context_parameter_detection(self): """Test that context params are detected in ResourceTemplate.from_function().""" From cb17c8be4a53664c8cd8b0c7f2d531549b16f746 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Wed, 7 May 2025 11:16:54 +0200 Subject: [PATCH 5/5] Allow generic parameters to Context on resources --- src/mcp/server/fastmcp/resources/templates.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/mcp/server/fastmcp/resources/templates.py b/src/mcp/server/fastmcp/resources/templates.py index d9f1448c9..f403b6077 100644 --- a/src/mcp/server/fastmcp/resources/templates.py +++ b/src/mcp/server/fastmcp/resources/templates.py @@ -5,7 +5,7 @@ import inspect import re from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, get_origin from pydantic import BaseModel, Field, validate_call @@ -58,7 +58,9 @@ def from_function( sig = inspect.signature(fn) for param_name, param in sig.parameters.items(): - if param.annotation is Context: + if get_origin(param.annotation) is not None: + continue + if issubclass(param.annotation, Context): context_kwarg = param_name break