Skip to content

Commit 2cc3e0e

Browse files
committed
Enables context injection into resources
1 parent 321498a commit 2cc3e0e

File tree

6 files changed

+347
-10
lines changed

6 files changed

+347
-10
lines changed

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

+14-3
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,19 @@
11
"""Resource manager functionality."""
22

33
from collections.abc import Callable
4-
from typing import Any
4+
from typing import TYPE_CHECKING, Any
55

66
from pydantic import AnyUrl
77

88
from mcp.server.fastmcp.resources.base import Resource
99
from mcp.server.fastmcp.resources.templates import ResourceTemplate
1010
from mcp.server.fastmcp.utilities.logging import get_logger
1111

12+
if TYPE_CHECKING:
13+
from mcp.server.fastmcp.server import Context
14+
from mcp.server.session import ServerSessionT
15+
from mcp.shared.context import LifespanContextT
16+
1217
logger = get_logger(__name__)
1318

1419

@@ -65,7 +70,11 @@ def add_template(
6570
self._templates[template.uri_template] = template
6671
return template
6772

68-
async def get_resource(self, uri: AnyUrl | str) -> Resource | None:
73+
async def get_resource(
74+
self,
75+
uri: AnyUrl | str,
76+
context: "Context[ServerSessionT, LifespanContextT] | None" = None,
77+
) -> Resource | None:
6978
"""Get resource by URI, checking concrete resources first, then templates."""
7079
uri_str = str(uri)
7180
logger.debug("Getting resource", extra={"uri": uri_str})
@@ -78,7 +87,9 @@ async def get_resource(self, uri: AnyUrl | str) -> Resource | None:
7887
for template in self._templates.values():
7988
if params := template.matches(uri_str):
8089
try:
81-
return await template.create_resource(uri_str, params)
90+
return await template.create_resource(
91+
uri_str, params, context=context
92+
)
8293
except Exception as e:
8394
raise ValueError(f"Error creating resource from template: {e}")
8495

src/mcp/server/fastmcp/resources/templates.py

+40-5
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,17 @@
55
import inspect
66
import re
77
from collections.abc import Callable
8-
from typing import Any
8+
from typing import TYPE_CHECKING, Any
99

10-
from pydantic import BaseModel, Field, TypeAdapter, validate_call
10+
from pydantic import BaseModel, Field, validate_call
1111

1212
from mcp.server.fastmcp.resources.types import FunctionResource, Resource
13+
from mcp.server.fastmcp.utilities.func_metadata import func_metadata
14+
15+
if TYPE_CHECKING:
16+
from mcp.server.fastmcp.server import Context
17+
from mcp.server.session import ServerSessionT
18+
from mcp.shared.context import LifespanContextT
1319

1420

1521
class ResourceTemplate(BaseModel):
@@ -27,6 +33,9 @@ class ResourceTemplate(BaseModel):
2733
parameters: dict[str, Any] = Field(
2834
description="JSON schema for function parameters"
2935
)
36+
context_kwarg: str | None = Field(
37+
None, description="Name of the kwarg that should receive context"
38+
)
3039

3140
@classmethod
3241
def from_function(
@@ -42,8 +51,24 @@ def from_function(
4251
if func_name == "<lambda>":
4352
raise ValueError("You must provide a name for lambda functions")
4453

45-
# Get schema from TypeAdapter - will fail if function isn't properly typed
46-
parameters = TypeAdapter(fn).json_schema()
54+
# Find context parameter if it exists
55+
context_kwarg = None
56+
sig = inspect.signature(fn)
57+
for param_name, param in sig.parameters.items():
58+
if (
59+
param.annotation.__name__ == "Context"
60+
if hasattr(param.annotation, "__name__")
61+
else False
62+
):
63+
context_kwarg = param_name
64+
break
65+
66+
# Get schema from func_metadata, excluding context parameter
67+
func_arg_metadata = func_metadata(
68+
fn,
69+
skip_names=[context_kwarg] if context_kwarg is not None else [],
70+
)
71+
parameters = func_arg_metadata.arg_model.model_json_schema()
4772

4873
# ensure the arguments are properly cast
4974
fn = validate_call(fn)
@@ -55,6 +80,7 @@ def from_function(
5580
mime_type=mime_type or "text/plain",
5681
fn=fn,
5782
parameters=parameters,
83+
context_kwarg=context_kwarg,
5884
)
5985

6086
def matches(self, uri: str) -> dict[str, Any] | None:
@@ -66,9 +92,18 @@ def matches(self, uri: str) -> dict[str, Any] | None:
6692
return match.groupdict()
6793
return None
6894

69-
async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource:
95+
async def create_resource(
96+
self,
97+
uri: str,
98+
params: dict[str, Any],
99+
context: Context[ServerSessionT, LifespanContextT] | None = None,
100+
) -> Resource:
70101
"""Create a resource from the template with the given parameters."""
71102
try:
103+
# Add context to params if needed
104+
if self.context_kwarg is not None and context is not None:
105+
params = {**params, self.context_kwarg: context}
106+
72107
# Call function and check if result is a coroutine
73108
result = self.fn(**params)
74109
if inspect.iscoroutine(result):

src/mcp/server/fastmcp/server.py

+24-2
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,8 @@ async def list_resource_templates(self) -> list[MCPResourceTemplate]:
230230
async def read_resource(self, uri: AnyUrl | str) -> Iterable[ReadResourceContents]:
231231
"""Read a resource by URI."""
232232

233-
resource = await self._resource_manager.get_resource(uri)
233+
context = self.get_context()
234+
resource = await self._resource_manager.get_resource(uri, context=context)
234235
if not resource:
235236
raise ResourceError(f"Unknown resource: {uri}")
236237

@@ -327,6 +328,10 @@ def resource(
327328
If the URI contains parameters (e.g. "resource://{param}") or the function
328329
has parameters, it will be registered as a template resource.
329330
331+
Resources can optionally request a Context object by adding a parameter with the
332+
Context type annotation. The context provides access to MCP capabilities like
333+
logging, progress reporting, and resource access.
334+
330335
Args:
331336
uri: URI for the resource (e.g. "resource://my-resource" or "resource://{param}")
332337
name: Optional name for the resource
@@ -351,6 +356,12 @@ def get_weather(city: str) -> str:
351356
async def get_weather(city: str) -> str:
352357
data = await fetch_weather(city)
353358
return f"Weather for {city}: {data}"
359+
360+
@server.resource("resource://{city}/weather")
361+
async def get_weather(city: str, ctx: Context) -> str:
362+
await ctx.info(f"Getting weather for {city}")
363+
data = await fetch_weather(city)
364+
return f"Weather for {city}: {data}"
354365
"""
355366
# Check if user passed function directly instead of calling decorator
356367
if callable(uri):
@@ -367,7 +378,18 @@ def decorator(fn: AnyFunction) -> AnyFunction:
367378
if has_uri_params or has_func_params:
368379
# Validate that URI params match function params
369380
uri_params = set(re.findall(r"{(\w+)}", uri))
370-
func_params = set(inspect.signature(fn).parameters.keys())
381+
382+
# Get all function params except 'ctx' or any parameter of type Context
383+
sig = inspect.signature(fn)
384+
func_params: set[str] = set()
385+
for param_name, param in sig.parameters.items():
386+
# Skip context parameters
387+
if param_name == "ctx" or (
388+
hasattr(param.annotation, "__name__")
389+
and param.annotation.__name__ == "Context"
390+
):
391+
continue
392+
func_params.add(param_name)
371393

372394
if uri_params != func_params:
373395
raise ValueError(

tests/server/fastmcp/resources/test_resource_manager.py

+101
Original file line numberDiff line numberDiff line change
@@ -139,3 +139,104 @@ def test_list_resources(self, temp_file: Path):
139139
resources = manager.list_resources()
140140
assert len(resources) == 2
141141
assert resources == [resource1, resource2]
142+
143+
@pytest.mark.anyio
144+
async def test_context_injection_in_template(self):
145+
"""Test that context is injected when getting a resource from a template."""
146+
from mcp.server.fastmcp import Context, FastMCP
147+
148+
manager = ResourceManager()
149+
150+
def resource_with_context(name: str, ctx: Context) -> str:
151+
assert isinstance(ctx, Context)
152+
return f"Hello, {name}!"
153+
154+
template = ResourceTemplate.from_function(
155+
fn=resource_with_context,
156+
uri_template="greet://{name}",
157+
name="greeter",
158+
)
159+
manager._templates[template.uri_template] = template
160+
161+
mcp = FastMCP()
162+
ctx = mcp.get_context()
163+
164+
resource = await manager.get_resource(AnyUrl("greet://world"), context=ctx)
165+
assert isinstance(resource, FunctionResource)
166+
content = await resource.read()
167+
assert content == "Hello, world!"
168+
169+
@pytest.mark.anyio
170+
async def test_context_injection_in_async_template(self):
171+
"""Test that context is properly injected in async template functions."""
172+
from mcp.server.fastmcp import Context, FastMCP
173+
174+
manager = ResourceManager()
175+
176+
async def async_resource(name: str, ctx: Context) -> str:
177+
assert isinstance(ctx, Context)
178+
return f"Async Hello, {name}!"
179+
180+
template = ResourceTemplate.from_function(
181+
fn=async_resource,
182+
uri_template="async-greet://{name}",
183+
name="async-greeter",
184+
)
185+
manager._templates[template.uri_template] = template
186+
187+
mcp = FastMCP()
188+
ctx = mcp.get_context()
189+
190+
resource = await manager.get_resource(
191+
AnyUrl("async-greet://world"), context=ctx
192+
)
193+
assert isinstance(resource, FunctionResource)
194+
content = await resource.read()
195+
assert content == "Async Hello, world!"
196+
197+
@pytest.mark.anyio
198+
async def test_optional_context_in_template(self):
199+
"""Test that context is optional when getting a resource from a template."""
200+
from mcp.server.fastmcp import Context
201+
202+
manager = ResourceManager()
203+
204+
def resource_with_optional_context(
205+
name: str, ctx: Context | None = None
206+
) -> str:
207+
return f"Hello, {name}!"
208+
209+
template = ResourceTemplate.from_function(
210+
fn=resource_with_optional_context,
211+
uri_template="greet://{name}",
212+
name="greeter",
213+
)
214+
manager._templates[template.uri_template] = template
215+
216+
resource = await manager.get_resource(AnyUrl("greet://world"))
217+
assert isinstance(resource, FunctionResource)
218+
content = await resource.read()
219+
assert content == "Hello, world!"
220+
221+
@pytest.mark.anyio
222+
async def test_context_error_handling_in_template(self):
223+
"""Test error handling when context injection fails in a template."""
224+
from mcp.server.fastmcp import Context, FastMCP
225+
226+
manager = ResourceManager()
227+
228+
def failing_resource(name: str, ctx: Context) -> str:
229+
raise ValueError("Test error")
230+
231+
template = ResourceTemplate.from_function(
232+
fn=failing_resource,
233+
uri_template="greet://{name}",
234+
name="greeter",
235+
)
236+
manager._templates[template.uri_template] = template
237+
238+
mcp = FastMCP()
239+
ctx = mcp.get_context()
240+
241+
with pytest.raises(ValueError, match="Error creating resource from template"):
242+
await manager.get_resource(AnyUrl("greet://world"), context=ctx)

0 commit comments

Comments
 (0)